Eight weight codes per thread in the MoE prefill staging: routed MoE 1.10x, prefill up to +4.6%

[MichaelDementii]

sparse_moe_prefill_q4_gate_up_kernel and sparse_moe_prefill_qx_down_kernel staged their weight
tiles one bf16 pair per lane: each lane decoded one packed byte and issued a 4-byte store_vec, so
a row of the tile cost 32 shared-memory stores. Each thread now decodes the eight codes of one chunk
and stores them with a single 16-byte store_vec. Per weight the arithmetic is unchanged.

When you closed #96 you named "separate MoE and projection kernel optimizations" as work that
belongs in its own submission. This is the MoE half. The projection half - the same widening in
w8_rowsplit_gemm_mma, which six unrelated Op families instantiate - is a different contract and a
different claim, so it is not in this branch; I will send it separately. The two branches are
independent by file and were each built and tested on their own; I did not build their union.

Targeted at master. dev is two commits ahead (3c937178, 921698b7), both in sampling and the
frontend; neither touches anything here.

What changes

• Q4MmaDecodeAtom::decode_eight builds four bf16 pairs from one packed word with byte_perm.
• Q5MmaDecodeAtom::decode_eight and Q6MmaDecodeAtom::decode_eight sit beside decode_pair,
  evaluating the same expression for the four bytes a thread now owns. decode_pair stays: it is
  still the entry point for src/ops/common/rowsplit_grouped_mma.cuh and the Q5/Q6 row-split GEMMs.
• Q5RowSplitStorage and Q6RowSplitStorage publish kHighBytesPerChunk, derived from the group
  constants they already own, so the kernel no longer infers the high-bit stride by division.
• Q5DownMma::decode and Q6DownMma::decode lose their last caller and go with it.

The staging loop is indexed by thread rather than by warp, and its row width comes from
kCodeBytesPerGroup with a static_assert that a row of codes decodes to exactly the tile's k
width.

Operator benchmark

ninfer_sparse_moe_bench, product codecs, cold L2, median of 50, --execution eager because
prefill is not graph-captured - the only capture site is src/core/decode_graph.cpp and the prefill
chunk loop is outside it.

./build/bench/ninfer_sparse_moe_bench --codec all --tokens 1024 --cache cold --execution eager --repeat 50
./build/bench/ninfer_sparse_moe_bench --codec all --tokens 4096 --cache cold --execution eager --repeat 50
./build/bench/ninfer_sparse_moe_bench --codec all --tokens 8192 --cache cold --execution eager --repeat 50

full_sparse_moe_device_body, speedup against master:

codec	1024	4096	8192
q4-q5	1.050	1.100	1.100
q4-q6	1.052	1.099	1.100
w8-w8	0.999	1.002	0.998

The w8-w8 routed codec is a separate kernel pair that does not take these decoders; it is a
harness-stability control, not a control for the change.

All three token counts are above kSparseMoePrefillWideMin (768), so only the wide <8,64>
instantiation is timed. The narrow <4,32> form appears in the SASS census below and in the Op
tests, but has no measurement here.

The five W8 benchmarks and the pair benchmark are unchanged by this branch, which is the expected
result and is included as evidence that the change is confined to the two MoE kernels:

benchmark	this branch
w8_linear_add, the 18 row-split MMA shapes	median 1.000, worst 1.000
w8_linear_add, all 50 shapes	two move: the T=1 decode row at 1.164 and a medium_splitk row at 1.030
w8_linear_swiglu	5 of 39 rows move, median 1.038
linear --suite all	5 of 68 rows move, 30 of 33 W8 rows flat
gdn_input_proj, w8	1.000 at T=128, 256, 512; 0.993 at 1024
attn_input_proj, w8-qkv	1.000 at all four extents
attn_input_proj, w8-qgkv	1.000, except 1.012 at T=512
linear_pair, T 64-1024	1.000, except 0.961 at T=768

Those excursions bound the benchmarks' own noise: on a branch that touches none of those kernels,
single rows still move by 16% at T=1 and by 3.9% at T=768. Read the MoE speedups above against that
floor rather than against 1.000.

Resources

cuobjdump --dump-resource-usage, this branch against master:

kernel family	instantiations	registers max	registers mean	shared	spills
sparse_moe_prefill*	19	116 -> 116	70.5 -> 70.6	unchanged	none either side
w8_rowsplit_gemm_mma	158	128 -> 128	101.3 -> 101.3	unchanged	none either side

The first row pools every sparse_moe_prefill kernel, including the scan, gather, reduce and router
kernels this change does not touch. Only two instantiations move at all:
sparse_moe_prefill_q4_gate_up_kernel<8,64> 71 -> 75 registers, against the 85 available under
__launch_bounds__(256, 3), so it keeps three blocks per SM; and
sparse_moe_prefill_qx_down_kernel<Q5DownMma,8,64> 78 -> 77. The W8 row is there because that
kernel is the sibling PR's subject: this branch leaves it untouched, register for register.

No allocation, workspace boundary, graph node or public contract changes; the diff is four files,
all under src/.

Instruction stream, cuobjdump -sass on the two ninfer binaries, compared body by body. Of 2927
device functions, 6 changed and 2850 are identical; 71 could not be compared, because they are
internal-linkage symbols carrying a per-compilation module id that differs between builds. All six
changed bodies are instantiations of the two kernels above:

instantiation	instructions	STS	LDS	HMMA
qx_down_kernel<Q5DownMma,8,64>	1888 -> 1568 (-16.9%)	32 -> 8	96 -> 24	128 -> 128
q4_gate_up_kernel<4,32>	720 -> 672 (-6.7%)	4 -> 1	12 -> 9	16 -> 16

Across all six the reduction is 6.4% to 16.9%, median 16.3%. HMMA is identical in every one -
the MMA stream is exactly what it was, and what shrinks is the work around it.

Numerics

Per weight each decoder computes the value the path it replaces computed:

• Q4: a nibble n decodes to (n ^ 8) - 8, an integer in [-8, 7] that bf16 represents exactly, so
  forming 128 + (n ^ 8) and subtracting 136 lands on the same bits the
  nibble -> int -> float -> bf16 path produced.
• Q5/Q6: the expression decode_pair evaluates. With lane = 4 * chunk + i the old high-bit index
  lane >> 2 is chunk (Q5) and lane >> 1 is 2 * chunk + (i >> 1) (Q6), which is what the new
  chunk pointer and kHighBytesPerChunk reproduce.
• The 16-byte store lands on the same eight shared slots in the same order as the four 4-byte stores
  it replaces: gemm_swz64(row, 8c) is a multiple of 8 elements, so the byte offset is 16-byte
  aligned, and c ^ (row & 7) is a bijection on the row's chunks, so the runs tile each row exactly
  once.

The shipped Op tests compare against tolerances rather than bits, so they qualify the routes but do
not by themselves establish bit-identity. What does: the HMMA census above, and greedy output
byte-identical to master at every measured point. tests/ops/test_sparse_moe.cpp reaches both
instantiations of both kernels - its token list straddles 768 - for both down codecs.

cd build && ctest -j1

92 pass, 1 skipped, 1 fails - on this branch and on master built in the same directory in the same
campaign. The skip is 27b_load_plan, which requires both real 27B artifacts and only one is on this
box. The failure is ninfer_qwen3_6_27b_prefix_real_test, with the same message on both arms:
Host checkpoint restore changed greedy output: restored=64,1248, baseline=64,56127 ... transfers=1/1/3/3. It is reported separately as #105.

End to end

Confirmation only. Arms alternate inside each round, every point its own process, greedy, four
rounds. The spread column is the largest round-to-round spread among the arms in that row.

./build/apps/ninfer <artifact> --messages <prompt> --max-new 32 --greedy --no-thinking \
  --max-context 131072 --prefill-chunk 8192 --kv-dtype bf16

model	chunk	tokens	n	prefill	decode	spread, %
Qwen3.6-35B-A3B	8192	8 515	4	+4.55%	+0.01%	0.90
Qwen3.6-35B-A3B	8192	33 031	4	+3.72%	+0.05%	0.74
Qwen3.6-35B-A3B	1024	33 031	4	+3.47%	+0.05%	0.56
Qwen3.6-27B dense	8192	33 031	4	-0.08%	+0.02%	0.47

All 16 generations are byte-identical to master. The dense Qwen3.6-27B "has a dense SwiGLU MLP in
every layer and does not contain MoE experts" (docs/maintainer/qwen3.6-27b-model.md:41), so that
row is a null control for this branch rather than a result.

Checks not run

• No ncu counters: RmProfilingAdminOnly is set on this host. Kernel-level claims come from the
  benchmark and from cuobjdump, not from hardware counters.
• The narrow <4,32> instantiation is qualified by the Op tests but not benchmarked.
• docs/performance.md publishes prefill throughput for qwen3_6_35b_a3b measured at a 1,024-token
  prefill chunk, which is a path this change is on, so those figures would move. I have not re-run
  that harness - it uses INT8 group-64 KV, prefix reuse, CUDA Graphs and five seeds through the serve
  corpus - and I am not proposing edits to the document.
• Speculation is off in every run here, so the MTP leaves are not exercised.
• Qwen3.8-27B NVFP4 is not measured; it takes neither kernel.

RTX 5090, sm_120a, driver 580.105.08, CUDA 13.1.115, Release, -DCMAKE_CUDA_ARCHITECTURES=120a.
Base is 6e8b2e2a.