CUDA: MMQ tile configuration and FP4 accumulation for GB10 (sm_121)

[danielhanchen]

Prompt processing on GB10 (DGX Spark, sm_121a) is matmul bound, and MMQ is the kernel that runs every
quantized mul_mat. An audit of the sm_121a SASS put the MMQ mainloop at about 31 % of the int8
tensor-core peak and 20 to 27 % of the FP4 peak at prompt-sized batches, against 65 to 71 % for
CUTLASS on the same GPU. Two of the causes are cheap to remove, and both are confined to the
Blackwell configuration.

1. The dense tile is one block per SM. At J = 128 the 128-row tile needs 61.5 KB of shared memory
and the kernel asks for 256 threads, so exactly one block is resident and nothing overlaps the
register-staged global loads or the two __syncthreads per K step. The MMA layout ties I to 16 rows
per warp, so a 64-row tile at 128 threads is the only way to get two resident blocks (40.5 KB per
block at J = 128, two fit in the 100 KB per SM of GB10). Same kernels, same numerics, more resident
warps.

mul_mat_id does not want this. The MoE launch picks J from ncols_max = the total token count
while each expert only holds ncols_max*n_expert_used/n_experts columns, so halving the tile height
only doubles the number of row tiles that pay for a mostly empty J tile: measured at -12 to -21 % on
MUL_MAT_ID at 2048 tokens. The table therefore keeps the stock 128-row tile in its
fallback == true entries and mul_mat_id asks for that configuration on Blackwell. Requesting the
bounds-checked variant cannot change results, it only enables the src0->ne[1] range check in the
tile loads and the write-back.

2. The FP4 accumulator is copied. mma.sync...m16n8k64 block-scaled already takes C as an
accumulator input, but ggml_cuda_mmq_vec_dot_fp4_fp4_mma zeroed a temporary tile per instruction
and added it into sum[] afterwards. Accumulating into sum[] directly removes 805 FADD and all of
the register spilling (26 STL / 30 LDL) from the NVFP4 J = 128 kernel; 224 -> 256 OMMA visible in the
loop.

The second commit fixes a latent bug found on the way: mul_mat_q_case re-derived fallback from
args.nrows_x % 128 and ignored the flag ggml_cuda_mul_mat_q had already computed and already used
to size the src1 padding through ggml_cuda_mmq_get_J_max. The two agreed by construction while
the only input was ne01, but any other reason to want the bounds-checked configuration was silently
dropped. The decision is now carried in mmq_args, so the configuration the padding was sized for is
the configuration that runs.

Measurements

DGX Spark GB10, sm_121a, nvidia-smi -lgc 300,2100, every cell at 2086 to 2093 MHz, single node with
the GPU lock held, no other compute process resident. test-backend-ops perf -b CUDA0, each cell
bracketed by a base cell before and after; the two base cells agree to within 2 % on every dense
case and to 0.1 % end to end.

Dense MUL_MAT, TFLOP/s (N x M x K, M = tokens)

shape	type	base	patched
34816 x 512 x 5120	q4_K	62.03	68.28	+10.1 %
34816 x 2048 x 5120	q4_K	62.49	69.19	+10.7 %
34816 x 512 x 5120	q5_K	60.61	65.68	+8.4 %
34816 x 2048 x 5120	q5_K	61.11	66.52	+8.9 %
34816 x 512 x 5120	q6_K	48.36	52.52	+8.6 %
34816 x 2048 x 5120	q6_K	46.89	52.18	+11.3 %
34816 x 512 x 5120	mxfp4	105.03	124.04	+18.1 %
34816 x 2048 x 5120	mxfp4	103.53	123.13	+18.9 %
34816 x 512 x 5120	nvfp4	123.90	136.67	+10.3 %
34816 x 2048 x 5120	nvfp4	126.20	142.12	+12.6 %
8192 x 512 x 5120	q4_K	52.00	61.57	+18.4 %
8192 x 512 x 5120	nvfp4	68.53	117.76	+71.8 %
5120 x 512 x 17408	q4_K	44.95	57.04	+26.9 %
5120 x 512 x 17408	nvfp4	53.16	96.83	+82.1 %
5120 x 2048 x 2048	nvfp4	77.39	89.01	+15.0 %
34816 x 32 x 5120	q4_K	19.07	18.98	-0.5 %
34816 x 512 x 5120	q8_0	66.89	66.44	-0.7 %

Batch 32 and q8_0 are unchanged: q8_0, q4_0, q5_0, iq4_xs and iq4_nl are deliberately left on the
stock tile, since the narrow tile was worth under 2 % for them at large N and cost 23 % at N = 512.
The types that gain are the ones whose mainloop is rescale bound (the K-quants) or FP4.

MoE MUL_MAT_ID, TFLOP/s (256 experts, 8 used)

Two base cells are shown because the first MoE cell of a window is cold; the patched cell sits inside
the base bracket on every case, i.e. no regression.

shape	type	base	patched	base (bracket)
1024 x 512 x 2048	q4_K	5.64	6.28	6.47
1024 x 2048 x 2048	q4_K	15.41	16.55	17.01
1024 x 2048 x 2048	q8_0	12.20	12.40	12.46
1024 x 2048 x 2048	mxfp4	18.75	20.91	20.29
1024 x 2048 x 2048	nvfp4	19.86	20.55	20.50
2048 x 2048 x 512	q4_K	10.33	11.83	12.12
2048 x 2048 x 512	nvfp4	11.05	14.30	13.27

Without the mul_mat_id guard the same build gives q4_K 1024 x 2048 x 2048 at 13.64 (-18.6 % against
a 16.75 base) and q8_0 at 9.80 (-20.9 %), which is what the guard is there to avoid.

End to end

llama-batched-bench -m Qwen3.8-27B-UD-Q4_K_XL.gguf -c 4608 -b 2048 -ub 2048 -ngl 99 -fa on -npp 512,2048 -ntg 16 -npl 1, base / patched / base in one lock window at one clock state.

	base	patched	base (bracket)
pp512 t/s	768.03	835.63	768.61	+8.8 %
pp2048 t/s	800.35	830.09	799.96	+3.7 %

Things that were tried and did not work

• cp.async for the y tile. Both K halves fetched with cp.async.cg into a double buffer, one
  cp.async.wait_all plus one __syncthreads replacing the two register copy loops and one of the
  two barriers per K step. 21 to 24 % slower on every dense shape (q4_K 62.4 -> 48.6 TFLOP/s at
  M = 512). The extra y tile in shared memory costs more occupancy than the asynchronous copy buys.
• 192-row / 384-thread tiles. ptxas caps the K-quant kernels at 168 registers and spills (136 B
  stack on q4_K J = 128); the build faults at runtime on MUL_MAT_ID.
• Hoisting the Q4_K block rescale. Both operand scales change every 32 k, so the product
  dA_ik * dB_jk is a different scalar for every (i, j, k-block) and no partial sum can stay in
  int32 across k-blocks. Accumulating in int32 across a 256-wide superblock, which is what
  CUTLASS-style block-scaled kernels do, requires quantizing the activations with one scale per 256
  values instead of per 32. That changes results and needs a new quantize_mmq layout; it is not
  part of this PR.

Proof that nothing else is touched

• The diff is four files, all under ggml/src/ggml-cuda/: the Blackwell tile table, the FP4 vec_dot,
  and the two lines of mmq.cu / mmq.cuh that carry the tile choice. No RPC, backend-registry,
  CPU, Metal, Vulkan, HIP or SYCL file is touched.
• Every table entry that changed is inside ggml_cuda_mmq_get_config_blackwell, which is only
  reached under blackwell_mma_available(cc) on the host and BLACKWELL_MMA_AVAILABLE in device
  code. The mul_mat_id term in fallback is gated on blackwell_mma_available(cc) as well. Every
  other GPU selects exactly the configuration it selected before.
• The FP4 vec_dot change is inside ggml_cuda_mmq_vec_dot_fp4_fp4_mma, which only exists for the
  block-scaled FP4 MMA path.
• cmake -DGGML_CUDA=OFF -DGGML_RPC=ON -DGGML_NATIVE=ON builds clean with the patch applied
  (test-backend-ops and llama-bench both link).
• test-backend-ops test -b CUDA0 on the patched build: 13572/13572 tests passed, Backend CUDA0:
  OK. Restricted runs: -o MUL_MAT 1193/1193, -o MUL_MAT_ID 869/869.

Still open

The CUTLASS W4A4 path of ggml-org#26704 reaches roughly 250 TFLOP/s on this class of GPU,
still about 1.75x what this gets to, so the tile table is not the end of the story. ggml-org#26977
(routed-MoE tile width from the mean columns per expert) and ggml-org#26159 (compact expert/tile job
list for the NVFP4 MoE grid) cherry-pick cleanly onto this and are a real MoE win here (+37 % q4_K
experts at 128 tokens, +39 % NVFP4 at 2048); they are upstream PRs, so they are not duplicated in
this one.