[lyesrock]

Problem and result

The chunked prompt decodes an expert's packed weights again for every 64 of its rows. A popular
expert of a long prompt has hundreds of rows, so its weights are decoded and multiplied several
times over. DS4_FAT_ROWS=N (default 0: the path is unchanged, the packed tiles run exactly as
before) takes the experts of N rows or more onto one other route: their weights decoded once to the
FP16 values the tiles multiply, the rows gathered, one cuBLAS FP16 GEMM, clamped SwiGLU in FP16.

Measured on the RTX 5090 (sysai2, driver 615.71.10) with the IST-DAS 2.25 bpw pack, this 0.2.7.1
base, ctx 16K, one user, DS4_FAT_ROWS=64, ABBA protocol (four fresh boots alternating off/on, a
fresh random 13K-token prompt for every measurement, conversation slots and the page cache cleared
between boots, nice=0, the RAM copy verified at every boot, 11 563 - 11 617 experts):

• first (cold) 13K prompt: patched-off 1 144 / 1 145, on 1 170 / 1 206; and against the
  released 0.1.52 binary itself (a second ABBA, official/on/official/on): official 1 206 /
  1 162, on 1 221 / 1 199. The on side wins every boot: +2 - 4%
• further 13K prompts: patched-off 1 561 - 1 594, on 1 642 - 1 689; official 1 579 - 1 623,
  on 1 662 - 1 710 - +4 - 6% on both comparisons
• decode: unchanged (short 256: 30 - 32 both; warm ramp 49 - 51 both)

The on side wins all four boots in both regimes. An earlier draft of this PR quoted +39% at 13K
(1 400 to 1 932 - 1 959): that compared a warmed on run against a single cold off run. Corrected
above; sorry for the noise.

The same products, the sums' order differs: a greedy answer can diverge after the first rounding.
12 greedy prompts at temperature 0, boots side by side: off vs the stock 0.1.52 10 of 12 outputs
byte-identical, on vs stock 9 to 10 of 12 - but the stock engine against itself, two boots, gives
9 of 12 on these free-form prompts, so the differences stay inside the stock run-to-run noise (all
differing answers were equivalent wording, same facts, same arithmetic).

Approach

One expert's weights as FP16 once (iq::f16_rows), the rows gathered (gather_f16),
cublasGemmEx, clamped SwiGLU in FP16 (swiglu_f16); SetArgs skip_ge leaves those experts out
of the tiles. skip_ge sits last in the CUDA/HIP structs, before the Metal-only fields, so the
CUDA/HIP kernel parameter layout keeps upstream's rule (the Metal fields go last). Default 0
changes nothing.

Validation

Passed: the CUDA build (sm_120) on this base; the ABBA prefill numbers above; the 12-prompt greedy
comparison; decode unchanged on the same boots.

Failed: nothing.

Not run: the HIP and the Metal build of this branch (no ROCm and no Mac on this box) - the touched
files include kernel sources, so a maintainer build there is worth having; a KL/ppl run on this
base (the quality evidence here is the 12-prompt greedy comparison above).

Limits

The gain is modest (+4 - 6%) at 13K tokens on sm_120: upstream's tiles are themselves much faster
when warm than in the first draft's comparison, so what is left to win is the repeated decode of a
popular expert's weights, not a stalled GEMM. Longer prompts (32K+) and cards with slower
dequant likely keep the larger share of the gain, but that was not measured here.

[Niko1221]

Thank you for this, and for the careful write-up. I rebased the patch onto the current main (0.2.7.1) and kept your commit as it is, with you as the author. Before it goes in I made four small changes in a separate commit:

• the three new launches go through strata_ds4_launch (the Mac build needs that);
• the path stays off on Metal, as on HIP;
• it is only taken for expert types that f16_rows decodes (a pack with a Q5_K/Q6_K or trellis down type would otherwise stop a request with "gate / up not decoded");
• N is clamped to 17 or more, since the light kernels already own the experts of 16 rows or fewer.

What we measured on two RTX A6000 with the IST-DAS 2.25 pack, with the arms interleaved in one engine process:

• With DS4_FAT_ROWS unset, the scores are byte-identical to main.
• With 64: prefill +0.1% at 5K tokens, +2.1% at 13K, +2.4% at 32K, and +5.0% on 8K of code. Decode is unchanged, and VRAM grows by 74 MB per card.
• The 13K sweep is flat from 64 upward and does nothing at 32.
• After a 12K-token prefix, the KL against stock is 0.24. Changing --prefill-chunk on stock moves it by the same amount (0.26 and 0.28), so that path is simply sensitive to any rounding change.
• The quality suite (19 tasks) is unchanged, apart from one flaky summary check that stock also fails.
• The HIP build of the three files compiles.

So it will go in as an opt-in, default off, in the next update (0.2.8). This pull request closes when that reaches main.

Your +39% on the 5090 is much larger than what we see. Our cards keep most of the pack in RAM, so the prompt is bound by the expert flow there, which probably explains part of the gap. If you can send the same interleaved numbers for the 5090 (stock against 64 at 5K / 13K / 32K) and the free VRAM left at the end of a load, we can look at an automatic rule for the cards where it pays. One thing to know for that: the buffers are a plain cudaMalloc outside the engine's VRAM budget, so for a default they have to move into the load-time budget.

[lyesrock]

Thank you - for the rebase, the four adjustments and the measurements in one pass. All four make sense to me: the path through strata_ds4_launch is something I should have used from the start, and clamping N to 17+ is the right edge, since below it the light kernels already own the rows.

On the +39%: I corrected it in the body - that draft compared a warm on-arm with a cold off-arm. The honest numbers for this 5090, strict A/B with four alternating fresh boots and a fresh random prompt every measurement, are +2-4% cold and +4-6% on repeat prompts at 13K, decode unchanged - against the released 0.1.52 binary itself, not only against the patched off-arm. That lands on your A6000 band (+2.1-2.4% at 13K), which is reassuring.

The regime here is the one you suspect: --mmap-experts off, the CPU lane's 113 GB of experts copied into RAM, so the prompt is bound by the expert flow - and the gain still shows. I will run the interleaved one-process sweep you asked for (stock against 64 at 5K / 13K / 32K, fresh text per arm) and post the table with the free VRAM left at the end of a load in a follow-up comment this week.

Noted for the default rule: the buffers are a plain cudaMalloc outside the load-time budget and must move inside it before anything automatic. On this card the 74 MB a card class of cost is small beside the ~550 MB the load leaves over, so a budget-aware rule should be safe here.

[lyesrock]

As promised, the interleaved sweep on the 5090. One note first: my box cannot interleave the arms inside one process (my branch toggles at load), so this is four fresh boots, alternating ON/OFF/ON/OFF, --mmap-experts off (all 113 GB of the lane's experts in RAM), context 65536, --prefill-chunk 8192, fresh random text every measurement, slots and page cache cleared between boots, nice=0. The first prefill of a boot is cold and its spread is wide (705-922 at 5K) - I kept both pairs and the honest reading is the pair of medians, not the single boot that flattered it.

prompt	stock (off)	DS4_FAT_ROWS=64	delta
5K	686.8 / 707.5	922.5 / 705.0	+0 to +2 % (boot 1 was an outlier)
13K	1541.9 / 1582.1	1557.0 / 1652.9	+1.0 / +4.5 %
32K	2069.4 / 2101.9	2156.5 / 2230.8	+4.2 / +6.1 %

Decode unchanged (warm 47-51 both sides). VRAM: at the end of a load the free VRAM is ~950 MiB (context 65536; ~550 MiB at 16K). The FAT_ROWS buffers add ~150 MiB on top after the first long prompt (one RTX 5090, no per-card doubling) - your 74 MB per card and my ~150 MB agree on order.

So on this card the gain grows with the prompt, not flat: 5K ~0, 13K +2-4, 32K +5. A rule keyed on the free VRAM left after the load and on chunk length would fit: the card here pays exactly in the regime you describe, and the buffers fit in the ~550-950 MiB that the load leaves over.

[Niko1221]

Thanks lyesrock. Your change is in 0.2.8 as an opt-in (DS4_FAT_ROWS, off by default), with a follow-up that limits the path to the pack types it supports and clamps the row count. On 2x RTX A6000 it measured +2 to +5% on long prompts, which is under our bar for a default, and the default path is bit-identical to before. Closing this PR because it is now in main; a table from a 5090, 3090 or 4090 would let us make it automatic there.