Prompt path: an MMQ group's experts gathered in one launch, after one wait, released by one event (-9% on a 32K prompt, same output)

[sergqwer]

Every prompt chunk on a native pack runs the streamed walk: for each expert the compute stream waits for its copy, gathers it into the MMQ group buffer and records an event that gives its ring slot back. The gather takes ~3 us. Under WDDM, each wait and record also leaves ~10 us of GPU idle, even when the copy landed long before.

An nsys trace of a 32K prompt shows:

• 226 ms of idle between 22,659 back-to-back gathers;
• the gather kernels themselves take 71 ms;
• the host is 3-11 ms ahead of the GPU the whole time, so the idle is not the host being late.

What changes:

• One launch per group: gather_native_group / copy16_group_kernel copy a group's 16 experts in one launch, the same bytes to the same slots.
• One wait per group: on the group's last streamed copy. The copy stream is in order, so it covers the others.
• One event per group: it releases all of the group's ring slots. The issuer waits on used[used_of[slot]], where used_of maps each slot to the event that released it. A later record of that event only waits longer.
• Skipped ring entries: an entry the routing skipped inside an open group first gathers what the group holds so far. So no more than a group's entries are ever held back from the issuer, and a 96-slot ring cannot stall.
• Unchanged: the products, their order, and the MMQ tail memsets. The staged path (chunks below 1024) and the FP16 path are as before.
• Opt-out: STRATA_PREFILL_GROUP_GATHER=0 gathers, waits and releases one expert at a time.

Same output

First-token logits are bit-identical with the switch on and off at 2K, 8K, 16K and 32K, and with a 96-slot ring (STRATA_PREFILL_RING=96) at 8K:

• They were dumped with diagnostics: STRATA_DUMP_FIRST_LOGITS writes the first generated token's logits #276's STRATA_DUMP_FIRST_LOGITS in a test build.
• The runs used a fixed --expert-cache 14900. With auto the slot count moves by ~100 from run to run with the desktop's free VRAM, and that alone changes a first token now and then, on main too.

Speed

Setup: IQ2_XS (ISTA), RTX 5090, 9950X3D, 128 GB, Windows 11, --max-context 262144 --kv int8 --prefill auto --vram-reserve-mib 1500, cache auto. The runs alternated between the switch off and on.

prompt	0.1.31	switch off	this PR
32K	6,184 ms	6,106 / 6,098	5,563 / 5,598 (-8.6%)
8K	1,741 ms	1,822 / 1,750	1,596 / 1,596 (-10%)
2K (fixed cache)		872	762 (-12.6%)
8K, 96-slot ring (fixed cache)		1,718	1,561

On an NVFP4 pack (our fork, larger experts): 32K -5.2%, 16K -7.9%, 8K -11.1%.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[sergqwer]

Rebased on 0.1.32 (c499bd1), no conflicts.

Checked on 0.1.32: IQ2_XS (ISTA), RTX 5090, a fixed --expert-cache, --vram-reserve-mib 1500. The first generated token's logits (STRATA_DUMP_FIRST_LOGITS, now in main) are byte-identical to main's after a 2K and a 32K prompt. Every target builds with STRATA_BUILD_TESTS.

The gain on 0.1.32, the same runs (prompt reading, one run each):

prompt	main	this PR
2K	858.8 ms	722.9 ms (-16%)
32K (4 chunks of 8192)	5,757 ms	5,150 ms (-10.5%)

0.1.32's 32 stager threads and 128-deep ring (a GGUF read in place) need nothing here: the group follows kRing at
run time, as it did with STRATA_STAGER_RING=96.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[bsorensen110]

Tested this on the R9700. It clearly wins here too, and the per-phase timings show exactly why.

Setup

• R9700 (gfx1201), ROCm 10.2.0a nightly / hipBLASLt 1.5.0, Linux (so the per-launch idle is smaller than WDDM — the gain is proportionally smaller than your −8.6%).
• Swift 1.5 IQ3_XXS (native pack), --prefill 32768 --kv int8 --max-context 262144, MMQ ON (STRATA_PREFILL_MMQ=ON), STRATA_HIP_WMMA=1 + STRATA_SELECT_WMMA=1.
• Cherry-picked cleanly onto current main (base c499bd1 = this PR's base); build 53/53 ctest.
• A/B is the same binary flipping only STRATA_PREFILL_GROUP_GATHER (unset vs =0), so the arms differ in nothing but the gather mode.

Kernel launches (rocprofv3 --kernel-trace --stats, a 30,870-token prompt):

• copy16_kernel 24,059 → copy16_group_kernel 1,959 (12.3× fewer), total kernels 51,650 → 29,550 (−43%).

GPU prefill timeline (STRATA_PREFILL_TIMING=1 phase table, an 8,823-token chunk):

• GPU timeline 3,444 → 3,239 ms (−6.0%).
• Per phase: wait copy 108 → 27 ms (−81 ms — the per-expert cross-stream sync this removes), dequant 457 → 333 ms (−124 ms), gather 44 → 43 ms (unchanged, as expected), gemm gate/up 611 → 619 ms (noise). The win is real and it lands where you said it would.

Wall clock (needle bench, fresh prompts, needle found in every arm):

• 30,870 tokens: 12.8 → 12.3 s (+4%)
• 186,669 tokens: 67.8 → 65.2 s (+3.8%)
• 4,210 tokens fresh (bench_prefill): 1,833 → 1,952 tok/s (+6.5%)

Bit-identical output we didn't re-verify (only needle recall + parity through ctest), but your own logit check plus these numbers make this an easy yes for us — we're promoting it locally. Thanks for the clean change.

[BlueKingMuch]

Data point from Windows (WDDM): RTX 4080 SUPER 32 GB (sm_89, 80 SMs), Ryzen 7 5800X3D, PCIe 3.0, Windows 11, CUDA 13.3; IQ3_S at --max-context 262144, int8 K/V, --prefill 32768, 11055 experts in VRAM (7605 slots lent to the prompt path). Build: 0.1.32 (+#363, #413, #346) with this PR merged together with #385, #374, #378 and #407, no conflicts; this PR against STRATA_PREFILL_GROUP_GATHER=0 in the same build, fresh prompts from source code, 2 runs each, STRATA_PREFILL_TIMING=1:

	STRATA_PREFILL_GROUP_GATHER=0	this PR
32K prompt, two runs	10.76 / 10.26 s (3047 t/s)	10.32 / 9.84 s (3178 t/s), +4.3 %
64K prompt, two runs	22.19 / 22.09 s (2891 t/s)	21.38 / 21.29 s (3000 t/s), +3.8 %
wait copy in a 32K chunk	259 ms	19-20 ms
dequant in a 32K chunk	664-672 ms	422-423 ms

The 8 greedy tokens after each prompt are the same with and without. A short piece is not affected: the 64K prompt is read here as 63035 + 923 tokens, and the 923-token piece waits 980-990 ms for copies either way, half of its 1.9 s. For agent turns, which add a few thousand tokens to a cached prefix, this copy wait is a large part of the time to first token on this PCIe 3.0 machine.

[sergqwer]

@bsorensen110 @BlueKingMuch thank you both for testing this on your hardware (R9700 on Linux, RTX 4080 SUPER on Windows) and for the detailed numbers.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[sergqwer]

Rebased on 0.1.36 (936fc67).

• Conflict: prefill.cpp (the host grouping is now the non-fused branch); the PR's lines are the same there.
• One addition for the fused path: a slot gathered in a group is released by its group's last event (used_of). 0.1.36's fused experts record one event per ring slot, so they now also reset used_of for it. Otherwise a fused layer after a grouped MMQ layer could leave a slot pointing at an older event: the issuer would wait on that one and refill the slot while the fused kernel still reads it. This only arises with STRATA_PF_FUSED=1 on a native pack whose layers are partly covered; the Q2_0 pack never groups (it is not native).
• Checks (IQ2_XS, fixed --expert-cache 14900): the first token's logits are byte-identical to v0.1.36 after 2K and 32K. In that fixed-cache run the 32K prompt read in 5.20 s, against 5.80 s for v0.1.36 (-10%). Three later interleaved runs on another 32K prompt scattered too much to add to that: 6.1-10.3 s for v0.1.36 itself, 9.1-10.0 s here.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZy1yKaDDiA8a7svdwaHio

[Niko1221]

Released in 0.1.38 as is: a group's experts gathered in one launch, one wait, one event. Same bits here too (the prompt path's residuals at 4K and 20K are byte-identical with STRATA_PREFILL_GROUP_GATHER=0). Merged under your name (GitHub doesn't show it as merged because it went in by hand). Thank you!