Eight weight codes per thread in the W8 row-split staging: the changed kernel 1.12x, production GEMM shapes 1.09x

[MichaelDementii]

w8_rowsplit_gemm_mma's dequant_w gave each lane one 16-bit code pair per quantisation group and
broadcast the group scales across the warp with shfl_sync, because a lane's pair could belong to a
group whose scale another lane had loaded. A row of a weight tile therefore cost 32 four-byte shared
stores on the BK == 64 schedules and 64 on the BK == 128 one. The loop is now indexed by thread
over eight-code chunks: a thread owns a whole chunk inside one quantisation group, reads its own
scale halfword directly, decodes eight codes out of one uint2 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 projection half; the MoE half was #106. 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, one commit on ce09aee5. Every number here was measured at 1fc1cb76, its
parent; the single commit between them cancels abandoned streaming requests in src/serve and is
on no path measured here. dev reviewed: it is two commits ahead (11e76d8d, 0b5b7c9a), adding
a perplexity corpus and a target_logprobs op, and touches neither this header nor its consumers.

This header is instantiated by six public Op families - linear, linear_add, linear_pair,
linear_swiglu, attn_input_proj and gdn_input_proj - so the change reaches all six. Note the
kernel name: the GROUPS == 4 schedule MmaR64x16C48K128A1 is selected at t in [41,48] for
N=34816 and [34,48] for N=248320 (src/ops/linear/w8/w8_dispatch.cpp:30-37), so this is not a
prefill-only change; those extents are swept below.

What changes

One function, src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh. The staging loop is indexed by thread
rather than by row pair, all three shfl_sync call sites go - one on the GROUPS == 2 branch and
two on GROUPS == 4 - and the store goes through W8Bf16x8Bits, the 16-byte union this file
already keeps for that shape (w8_rowsplit_gemm_mma.cuh:22).

Eight codes is the widest chunk that stays contiguous in As for every row: w8g32_swz64 permutes
whole eight-element runs, and a wider run is contiguous only when row & 7 is even. What the body
needs beyond that is now stated where the loop lives, because the GROUPS == 4 branch that used to
assert it is gone with the branch:

constexpr int kChunksPerRow = BK / 8;
static_assert(BK % 32 == 0 && (8 % GROUPS) == 0,
              "an eight-code chunk must lie inside one W8G32 group and the scale cache "
              "must hold whole tiles");

The scale load changes shape rather than growing. Every thread now reads its own halfword instead of
two or four lanes per warp reading a word and broadcasting it, and a warp's 32 threads land on four
distinct 32-bit words in four distinct banks - SCALE_CACHE_BYTES is 16, so consecutive rows sit
four banks apart - with same-address broadcast inside each. LDS over the changed bodies falls
48.4% in the census below.

Operator benchmarks

Existing benchmarks, product shapes, --repeat 50, cold L2 through bench::measure_cold_launch.

./build/bench/ninfer_w8_linear_add_bench --production-only --repeat 50
./build/bench/ninfer_linear_bench --qtype W8 --n 12288 --k 2048 --t 4096 --repeat 50
./build/bench/ninfer_linear_bench --suite all --repeat 50
./build/bench/ninfer_linear_bench --qtype W8 --n 34816 --k 5120 --sweep 30:56 --repeat 50

w8_linear_add --production-only covers four different kernels and only one of them is this one, so
the rows are grouped by the kernel each shape reaches:

kernel	rows	T	median, pass 1 / 2	best	worst
w8_rowsplit_gemm_mma	18	129-1024	1.123 / 1.123	1.180	1.018
w8_small_t_mma (splitk8)	21	2-48	1.000 / 1.000	1.066	0.929
medium_splitk	10	63-128	0.999 / 1.000	1.043	0.959
decode	1	1	0.857 / 1.143	-	-

The three kernels this change does not touch give this arm's noise floor, and it is wide at the
narrow end: the single T=1 decode row reads 0.857 in one pass and 1.143 in the other, and
w8_small_t_mma spans 0.929 to 1.066 while its median stays exactly 1.000. The changed kernel's
18 rows, by contrast, land on the same median to three decimals in both passes and its worst row,
1.018, repeats. Read the worst case against that floor rather than against 1.000.

The GROUPS == 4 band is the one the kernel name points at and the one no benchmark reaches by
default, so it is swept token by token, t from 30 to 56, at both registered N:

N	band	rows	median	best	worst	outside the band
34816	t in [41,48]	8	1.079	1.100	1.074	median 1.000
248320	t in [34,48]	15	1.163	1.186	1.090	median 1.085

At N=34816 the band is bounded on both sides by rows that do not move, because t <= 40 goes to
w8_small_t and this kernel does not appear again until t > 48. At N=248320 the rows above the
band also move: t > 48 selects MmaR32C64 and then MmaR64C128, which are the same header.

The four production W8 GEMM shapes, every one of them the changed kernel:

N, K	T=1024	2048	4096	8192
12288, 2048	1.080	1.086	1.087	1.086
9216, 2048	1.093	1.090	1.088	1.088
2048, 4096	1.091	1.083	1.096	1.089
2048, 16384	1.089	1.118	1.097	1.092

Sixteen cells, median 1.089, spread 1.080 to 1.118 - narrower than the w8_linear_add rows above,
which sweep T through shapes where the staging loop is a smaller share of the body.

linear --suite all carries Q4, Q5, Q6 and W8 rows, so the formats this change cannot touch sit in
the same table as the ones it does. Two independent passes, so a row that moves can be told from a
row that is noisy:

rows	pass 1 median	range	pass 2 median	range
W8, 33 rows	1.057	0.998-1.179	1.065	0.859-1.165
Q4/Q5/Q6, 35 rows	1.000	0.996-1.006	1.000	0.997-1.002

18 of the 33 W8 rows move by more than 1% in both passes. No W8 row is slower in both passes -
the 0.859 in pass 2 is a single row that reads 1.03 in pass 1, which is why I am reporting two
passes rather than one. The 35 untouched rows span 0.996 to 1.006 across both passes, and that is
the floor the W8 numbers should be read against.

The other four Op families

Every benchmark below was run twice, so a row that moves can be told from a row that is noisy.

gdn_input_proj and attn_input_proj each measure one operator in several weight formats, so the
change and the formats it cannot touch sit in the same table (pass 1 / pass 2):

operator	format	T=64	128	256	512	1024
gdn_input_proj	w8	1.000/1.000	1.065/1.065	1.109/1.108	1.097/1.097	1.066/1.066
attn_input_proj	w8-qkv	1.000/1.000	1.105/1.105	1.160/1.161	1.102/1.075	1.085/1.085
attn_input_proj	w8-qgkv	1.000/1.000	1.044/1.044	1.067/1.067	1.056/1.075	1.093/1.093
both	fp8, nvfp4, q4q5, bf16	0.987-1.032	0.987-1.000	0.991-1.031	0.995-1.000	0.997-1.000

The untouched formats span 0.987 to 1.032 over both passes and 40 cells; every W8 row above T=64
sits outside that band in both passes.

linear_pair puts the route boundary on display, because kK2048Routes selects the first
ConcatMma schedule that instantiates this header at T=193:

T	64	128	192	193	256	384	512	768	1024
pass 1	0.984	1.000	1.000	1.200	1.133	1.200	1.118	1.042	1.073
pass 2	1.000	1.000	0.994	1.200	1.133	1.199	1.118	1.083	1.068

w8_linear_swiglu --production-only, 39 rows: 17 faster in both passes, median 1.107, range
1.038 to 1.143; none slower in both passes; 19 flat in both. Three rows disagree between passes -
all three are linear_swiglu.w8.splitk.mma.pair.exact_t at T=3, 11 and 17, reading 0.912/1.100,
0.917/1.091 and 1.000/1.074, which is what this benchmark's scatter looks like at those token
counts rather than a route effect.

Finally, the MoE prefill benchmark is the confinement control - a different Op, which the sibling
PR #106 changed and this one does not touch:

codec	q4-q5	q4-q6	w8-w8
T=4096	0.998	0.999	0.998
T=8192	0.999	1.000	1.000

Resources

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

kernel	instantiations	registers max	registers mean	shared	spills
w8_rowsplit_gemm_mma	120	127 -> 117	102.4 -> 96.2	unchanged	none either side

Registers fall in 103 of the 120 and rise in 9; the 22 distinct shared-memory sizes are the same
list before and after; local, spill and stack are zero on both sides. No allocation, workspace
boundary, graph node or public contract changes; the diff is one file.

Instruction stream, cuobjdump -sass on the two ninfer binaries, compared body by body. Of 2937
device functions 2899 are comparable and 120 of those changed; every one of the 120 is an
instantiation of this kernel. The 38 that cannot be compared are the internal-linkage instantiations
from src/ops/linear_pair/w8/w8_pair_gemm_concat.cu, whose symbols carry a per-compilation module
id that differs between builds. Over the 120 changed bodies the instruction count falls
179,948 -> 161,238 (-10.4%), per body -33.3% to -1.6%, median -4.9%.

What moves, summed over those bodies:

opcode	master -> branch
SHFL	1150 -> 0	the three broadcast sites
WARPSYNC	1360 -> 0
STS	2884 -> 1724	-40.2%, the four-byte stores becoming one 16-byte store
LDS	2934 -> 1514	-48.4%
I2F	3560 -> 1240	replaced by I2FP 0 -> 3720, the packed form
HMMA	15120 -> 15120	identical, and it moves in 0 of the 120 bodies

Not everything falls: FMUL +35.2%, F2FP +17.5% and SHF +37.5% are the index arithmetic and the
packed conversion the new loop does per thread instead of per lane pair. The two extremes:

instantiation	instructions	STS	LDS	I2F	HMMA
deepest reduction	1922 -> 1282 (-33.3%)	72 -> 37	61 -> 11	80 -> 10	160 -> 160
largest body	2913 -> 2466 (-15.3%)	46 -> 21	61 -> 26	60 -> 10	160 -> 160

The census does not reach the 38 concat instantiations, so for those the evidence is the greedy
output gate below rather than an instruction census.

Numerics

Per weight the arithmetic is the same int8 sign extension, the same multiply by the same group
scale, and the same __floats2bfloat162_rn. Only the route by which the scale reaches the thread
changes. For GROUPS == 2 the old code loaded a uint32 at the scale tile offset and extracted
>> ((gg & 1) * 16); for GROUPS == 4 it shuffled two uint32 from +0 and +4. Both resolve to
the halfword at scale_tile_offset + gg * 2, which is what the thread now reads directly, and the
highest byte index touched is 15, the last of the 16 in SCALE_CACHE_BYTES, as before.

The 16-byte store lands on the same eight shared slots in the same order as the four 4-byte stores
it replaces: w8g32_swz64(row, 8c) is a multiple of eight elements, so the byte offset is 16-byte
aligned, and c ^ (row & 7) is a bijection on a row's chunks. Removing the shuffles is safe because
they sat in a warp-uniform loop - the old trip count depended only on warp, so all 32 lanes were
converged.

The shipped Op tests compare against tolerances rather than bits. What establishes bit-identity is
the HMMA census above and greedy output byte-identical to master. The GROUPS == 4 schedule is
reached by tests/ops/linear/test_w8_a16.cpp through src/ops/linear/w8/w8_dispatch.cpp:30-37.

cd build && ctest -j1

100% tests passed, 0 tests failed out of 94, 1 skipped, on this branch and on master built in
the same directory in the same run. The skip is 27b_load_plan, which needs both real 27B artifacts
and only one is on this box. ninfer_qwen3_6_27b_prefix_real_test, which I reported as #105 and
which failed on a clean tree at the previous base, passes here - fd48e2fa fixed it.

End to end

Confirmation only; the operator benchmarks above are the claim. Arms alternate inside each round,
every point its own process, greedy, four rounds.

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

model	chunk	prompt tokens	per-round branch/master	prefill	decode
Qwen3.6-35B-A3B	8192	8,515	1.0294 1.0211 1.0247 1.0233	+2.46%	+0.01%
Qwen3.6-35B-A3B	8192	33,031	1.0198 1.0199 1.0216 1.0184	+1.99%	-0.04%
Qwen3.6-35B-A3B	8192	65,882	1.0151 1.0163 1.0153 1.0152	+1.55%	-0.01%
Qwen3.6-35B-A3B	1024	33,031	1.0177 1.0170 1.0177 1.0163	+1.72%	+0.02%
Qwen3.6-27B dense	8192	33,031	0.9974 1.0019 0.9973 1.0030	-0.01%	+0.06%

I quote the paired per-round ratios rather than a mean of means because the absolute rate drifts
between rounds on both arms; the ratios do not. The dense 27B row is the null control and behaves
like one: its four ratios straddle 1.000 and average to -0.01%. Its W8 row-split tensors are the MTP
and vision entries, which a text-only greedy run with speculation off does not exercise.

All 20 generations are byte-identical to master.

Checks not run

• No ncu counters: RmProfilingAdminOnly is set on this host. Kernel-level claims come from the
  benchmarks and from cuobjdump, not from hardware counters.
• docs/performance.md publishes prefill throughput for qwen3_6_35b_a3b at a 1,024-token prefill
  chunk, 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.
• The GROUPS == 2 schedules are exercised throughout, but I did not separate their rows from the
  GROUPS == 4 ones outside the two swept bands.
• Qwen3.8-27B NVFP4 is not measured; it takes no W8 row-split route.

RTX 5090, sm_120a, driver 580.105.08, CUDA 13.1.115, Release, -DCMAKE_CUDA_ARCHITECTURES=120a.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-08-29T11:31:58.224388Z	e591725	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[Neroued]

Great work. One suggestion for future performance reports: besides the speedup over the previous implementation, a roofline analysis, e.g. achieved memory bandwidth and tensor core peak relative to the hardware limits would be very useful. It helps clarify how much optimization headroom remains.