DS4 H200: rocmfp2_mix MoE matvec 3.31x + ROCmFP4 dense matvec 1.53x (+42.6% decode)

[DeanoC]

Accumulating branch for DS4 CUDA/H200 kernel optimizations found by the
geo-evo kernel-optimization loop. Further improved kernels
land here as additional commits, so this can be tested and merged once there are a few.

Base is h200-serving, the branch the H200 production server actually runs. It existed only
locally, so it was pushed alongside this PR — and the two CUDA prerequisites it was missing
(__ballot portability, joyai-llm pre-tokenizer) have now been landed on the base, where they
belong: they were already live in the production binary, and without the ballot fix the base does not
compile for CUDA at all. This PR is therefore pure kernel work — eight commits, two kernels, no
build fixes mixed in.

Headline

DS4 decode goes 43.57 → 62.15 tok/s (+42.6%) on one H200 (sm_90, CUDA 13.0), across two kernels,
with the greedy token stream byte-identical throughout:

kernel	before	after	share of decode when targeted
mix_matvec_rocmfp2_moe_kernel (MoE gate/up expert matvec)	—	3.31x	35.0%
mul_mat_vec_q<(ggml_type)101> (dense ROCmFP4_FAST matvec)	—	1.53x	32.1%

Total GPU kernel time per decode is down 28.7% (5399.6 → 3847.4 ms).

Commits

On mix_matvec_rocmfp2_moe_kernel (the MoE gate/up expert matvec, 35.0% of decode when targeted):

commit	what	cumulative kernel
87f165b	stage MoE activations once per block	1.976x
a5587fe	issue all row loads before any FMA	2.089x
85dc3f5	stage activations in LDS, coalesced + swizzled	2.920x
0ee7bda	nibble-keyed, bank-replicated codebook table	3.314x
683f75d	static_assert the hand-computed nibble shift	3.314x

Then, on the dense ROCmFP4 matvec — which the first campaign promoted to the hottest kernel:

commit	what	cumulative kernel
e6036de	table-free ROCmFP4 codebook expansion via prmt	1.399x (+8.9% decode)
84a8c6b	aligned dword qs loads via __funnelshift_r	1.431x (no e2e claim — see below)
25b9a75	split the small-K block into two row groups	1.526x (+1.7% decode)

A more aggressive four-row-group form of the last one — removing the cross-warp reduction entirely —
was tried and rejected by the correctness gate on sha mismatch. Two groups is the largest split
that preserves the fold exactly; the commit message records this so nobody re-attempts it blind.

The first two are prerequisites that had never been committed anywhere — they lived only in the
working tree of the H200 box. Without the ballot fix dflash_server does not build for CUDA (raw
__ballot() is HIP-only; ptxas rejects a non-.sync vote on sm_70+), so no remote branch built on
sm_90 before this. They are separate commits and can be split onto their own PR.

Why this kernel, and what was actually wrong

mix_matvec_rocmfp2_moe_kernel (qtype-106 Q2_1_ROCMFP2_MIX, MoE gate/up expert matvec, batch-1
decode) was the largest single consumer of DS4 decode time on H200: 35.0% of GPU kernel time
(1891.4 ms of 5399.6 ms over 255 greedy decode steps, 16182 launches) per
nsys --cuda-graph-trace=node.

It was not bandwidth-bound. ncu at real decode geometry (in=7168, out=2048, top-k 4, ntok=1)
put DRAM at 3.9% while L1/TEX sat at 89.7% — bound on L1 tag-lookup wavefronts. The
lane→block map gives lane L blocks L, L+32, …, so consecutive lanes read activations 128 B apart and
every float4 activation request touched 32 distinct sectors while using 16 of each 32 B
(27.2 sectors/request, near the worst case). The kernel then paid that request four times per
block, because mix_block_accum() was called once per fold and each call re-staged the same 32
activations from the same address — a source comment asserted nvcc would CSE them, and it does not
(288 static LDG.E.128 for a loop needing 72).

This contradicts the gfx1151 finding that this kernel family is DRAM-bandwidth-bound. Both
conclusions are correct on their own hardware: HBM3e moves that roofline by more than an order of
magnitude, so what starves on ~256 GB/s unified LPDDR5X is latency- and L1-bound here.

Correctness

The first optimization is bit-exact by construction, not merely by test:
mix_block_accum(b, xc, col0, …) is defined as
mix_load_x32(xc + col0, xv); mix_block_accum_x(b, xv, …), so hoisting the shared stage cannot
reassociate anything.

Beyond that, every step was gated on the sha256 of the greedy completion (256 tokens,
temperature=0) matching a frozen reference — which is also byte-identical to the pre-change
production binary's output (6c1f0278…c912). The standalone harness additionally checks output
hashes per code path; all four combinations were bit-identical at each step:

instantiation	mode	pristine	final
<true>	1	151.85 us	73.41 → further reduced by the LDS and codebook steps
<true>	0	129.16 us	62.51 → "
<false>	1	81.66 us	47.58 → "
<false>	0	75.08 us	39.94 → "

881ebd9 adds a compile-time assertion on the hand-computed nibble shift, and an earlier commit in
the campaign added coverage for the misaligned fill arm the LDS staging introduced.

Measurement discipline

Each step was measured twice (independent runs, model reloaded each time) against a frozen baseline,
with GPU clocks pinned and the GPU exclusive. Observed repeat spreads: microbench 0.003–0.08%,
decode 0.02–0.65%. Throughput is measured in a process with no profiler attached, because nsys
injection costs ~11% of decode tok/s even with the capture stopped.

Per-step, against the bound implied by the kernel's share of remaining decode time
(total / (total - share·(1 − 1/speedup))):

step	kernel	step bound	observed	residual
stage once per block	1.976x	+20.9%	+19.0%	−1.6%
loads before FMAs	2.089x	+1.2%	~0%	−1.2%
LDS staging	2.920x	+6.2%	+7.1%	+0.9%
nibble codebook	3.314x	+1.9%	+1.7%	−0.2%

Cumulatively 3.314x predicts 57.68 tok/s against 56.44 observed (−2.15%). Sitting slightly under
the share-derived bound is the expected signature; overshooting it would indicate measurement error.
Note the "loads before FMAs" step is a genuine +5.7% kernel gain whose e2e effect was not
measurable — its step bound (+1.2%) was below the decode noise floor. It is included because the
kernel-level signal is solid and it is a prerequisite for the LDS step, not because it moved
throughput on its own.

How to test

The GPU must be exclusive: the model needs ~100.7 GB of 143.8 GB, so ds4-flash.service cannot be
running while measuring.

cmake -S server -B build-90 -G Ninja -DCMAKE_BUILD_TYPE=Release \
      -DDFLASH27B_GPU_BACKEND=cuda -DCMAKE_CUDA_ARCHITECTURES=90 \
      -DDFLASH27B_FA_ALL_QUANTS=ON -DDFLASH27B_SERVER=ON
cmake --build build-90 -j --target dflash_server

A standalone driver that exercises this kernel at decode geometry with no model load (~2 s per
measurement instead of a 100 GB server start) is in geo-evo at
tools/stack_kernels/ds4_h200_rocmfp2_mix_bench.cu; the full attempt log, ncu tables and rejected
ideas are in handoffs/progress-ds4_h200_mmvq_rocmfp2_mix.md.

Caveats for a reviewer

• One GPU, one arch. Measured only on H200/sm_90. The source is arch-neutral, but the motivation
  is H200-specific and the LDS staging changes shared-memory pressure, which is worth a look on
  gfx1151/gfx1201 before assuming it carries.
• One geometry. Correctness coverage is both instantiations × both quant modes at
  in=7168/out=2048/top-k 4/batch-1 decode, plus the 256-token greedy stream. Other shapes and the
  prefill path are not exercised by this evidence.
• Ratios, not absolutes. +29.7% is with clocks pinned; pinning costs ~5% in absolute terms on
  this workload (45.82 tok/s unpinned vs 43.57 pinned at baseline). Trust the ratio.
• One commit deliberately makes no throughput claim. 84a8c6b is a genuine +2.3% kernel gain
  (23x the rig's ~0.1% microbench floor) but removes only 22.1 ms of 3926.6 ms, bounding decode at
  +0.57% — below the rig's ~0.7% decode noise floor. Measured decode was unchanged within noise.
  It is included because it is bit-exact, composes, and cuts real memory traffic. Judge it on the
  kernel signal, not on tok/s. The same is true of a5587fe in the first campaign.
• 84a8c6b reads slightly outside the nominal block (up to 2 bytes past the last block, and the
  aligned dword below addr). Its commit message carries the in-bounds argument, which rests on
  ggml_backend_cuda_buffer_type_get_alloc_size() padding ne0 to a multiple of
  MATRIX_ROW_PADDING. That is the one place a reviewer should push back hardest, because it is an
  allocator invariant rather than a local property.
• What is now hot. After both campaigns: mul_mat_vec_q<qtype 101> ~24.8%,
  mix_matvec_rocmfp3_moe_kernel ~15%, mix_matvec_rocmfp2_moe_kernel ~15.6%. quantize_q8_1 is
  still unfused at ~4.1% with one launch per matvec (118854 of each), which looks like the cheapest
  remaining structural win. Shares move a lot as campaigns land — re-profile before choosing a
  target.

🤖 Generated with Claude Code