feat(mix): batched (MMQ) path for qtypes 105/106, routed on batch width - #607
feat(mix): batched (MMQ) path for qtypes 105/106, routed on batch width#607geometric[bot] wants to merge 5 commits into
Conversation
…the dense matvec The qtype 105/106 matvec read activations one scalar at a time and handled one output row per warp. Vectorizing the loads and blocking two rows per warp measured 2.06x for dense decode of these qtypes on an H200, with BIT-IDENTICAL output — the per-block accumulation order is unchanged. On AMD (gfx1151, gfx1201) the same change is neutral, so it is not gated by backend. Methodology note: that 2.06x was taken with a decode bench belonging to a dense model line that is NOT part of this PR, so it is not reproducible from this tree alone. The change is bit-exact, so correctness is covered by the existing test_rocmfp_mix_slice_matvec gate; happy to have maintainers re-run throughput on the RTX 3090 or Ryzen 395 AI Max per CONTRIBUTING, or to add a standalone bench if you would rather have one in-tree. Dense consumers are what this helps: DeepSeek-V4's own 105/106 tensors are MoE experts, which take the mul_mat_id path rather than this one.
…and route it on batch width
The MMQ kernels for both mix qtypes already existed and already plumbed the
out-of-band codebooks, but were reachable only through an env var
(DFLASH_DS4_MIX_MMQ_PREFILL) and gated to RDNA, so ne11 > 1 fell back to
dequantize-to-bf16 + dense GEMM — 48% of a 16-token batch's GPU time inside
dequantize_rocmfp{2,3}_mix_kernel (nsys), i.e. the multiply discarded the
sub-4 bpw artifact and ran 16-bit for its duration.
The toggle becomes a real API (ggml_cuda_mix_mmq_enabled / set / clear /
env_pinned in ggml-cuda.h) with precedence explicit-call > env > compiled
default, renamed DFLASH_MIX_MMQ since it was never prefill-specific; the old
spelling is still honoured. NVIDIA is no longer excluded: the DP4A/MMA tile
these types declare is portable, and a correctness gate now covers it rather
than an assumption.
MMQ is NOT uniformly better, so the choice is per multiply rather than per
process. It wins on narrow batches and loses on wide ones, because the dequant
path hands a wide N to a well-tiled dense GEMM that the MMQ kernel does not
out-tile. Measured out of tree on a dense 3.3 bpw artifact, prefill tok/s,
off -> on:
ne11 8 16 64 256 1024 2048
gfx1151 1.80x 1.77x 1.65x 1.05x 0.89x 0.86x
gfx1201 5.11x 4.02x 3.26x 1.83x 1.13x 0.98x
so ggml_cuda_should_use_mmq width-gates the mix qtypes at mix_mmq_max_ne11
(1024 on RDNA 4, 256 elsewhere; DFLASH_MIX_MMQ_MAX_NE11 overrides). One request
then gets the narrow-batch win on speculative verify AND the dense-GEMM win on
prefill with nothing to configure. Re-measured with the gate in place, the
wide-batch regressions are gone (gfx1151 2048: 0.86x -> 1.00x) and every
narrow-batch win is retained. Decode is unchanged in all arms, which is the
control: ne11 == 1 takes the MMV kernel and never reaches this gate.
Default stays OFF with no in-tree opt-in caller, deliberately: DS4's 105/106
tensors are MoE experts on mul_mat_id, so this path cannot fire for them. That
is measured, not assumed — six serving configs byte-identical either way, a
3974-token prefill at 184.1 s +/-0.1%, and a controlled decode A/B (one server
at a time, hard teardown, arms ordered 0,1,1,0) at 21.72 / 21.72 / 21.74 tok/s
on gfx1151. A dense consumer of these qtypes opts in for itself and shows its
own measurement.
Methodology note: the width sweep is AMD (gfx1151/gfx1201), n=1 per cell on an
idle host, and the crossover is bracketed rather than resolved (256-1024 on RDNA
3.5, 1024-2048 on RDNA 4). NVIDIA has no width sweep, so it takes the
conservative RDNA 3.5 bound — that keeps every measured NVIDIA win and declines
only widths nothing has measured there. Power limit was not pinned; happy to
re-run under your methodology.
The batched path had no correctness test on either vendor: it was reachable only behind an env var and gated to RDNA, so it had never executed on CUDA at all. Enabling it (previous commit) without a gate would be a numerics change backed by nothing. Compares MMQ against the already-validated matvec kernel for both qtypes across ne11 1/4/16/64 and expert counts 1 (dense mul_mat) and >1 (MoE), asserting err(MMQ) <= tolerance * err(dequant) plus an absolute 1%-of-|ref| bar. On an H200 MMQ lands 0.09-0.14% of |ref| and is 5-7x CLOSER to reference than the dequant path — dequant rounds through bf16 where MMQ keeps integer dot products, so the batched path is faster and more faithful at once. Green both on the default path and under DFLASH_MIX_MMQ=0, which covers the toggle's precedence.
Guarded by if(EXISTS) like its neighbours, and backend-general: the kernels declare cudaStream_t and ggml's vendors/hip.h maps the cuda* spellings onto hip*, so one source builds for both. Kept separate from the test source per CONTRIBUTING (build config in its own commit).
…oE in the gate
Two defects, both mine, found by extending the gate to the dispatch DeepSeek-V4
actually uses.
The width cap was wrong for MoE. Declining MMQ on a DENSE multiply hands the
work to a well-tiled dense GEMM, which is why the cap exists and why it measured
well. Declining it on a mul_mat_id does something else entirely: it falls to the
host-synchronised sort path — a cudaStreamSynchronize, a CPU-side id sort, an
expand, and no CUDA-graph capture. Same threshold, opposite consequence.
Measured on H200, 8 experts / top_k 2, q2_1_rocmfp2_mix at 512 tokens: 0.412 ms
with the cap applied to MoE vs 0.202 ms without — the cap cost 2.04x on exactly
the widths a prefill uses. The cap is now dense-only; mul_mat_id is keyed off
n_experts > 1 (the dense call sites pass 0, or 1 for the fused gate/up pair).
With MMQ reaching MoE at width, both qtypes beat the fallback at every width
where the toggle changes the kernel: 105 at 64 tokens 0.116 vs 0.249 ms, at 512
0.212 vs 0.370; 106 at 64 0.106 vs 0.247, at 512 0.202 vs 0.359.
The test claimed coverage it did not have. Its comment said 'expert counts 1
(dense mul_mat) and >1 (MoE)' while expert_set was {1}: a 3-D src0 through
ggml_mul_mat does not reach MMQ, so the multi-expert case had been removed as
vacuous and the claim was never corrected. It now exercises ggml_mul_mat_id
directly across 1/4/64/512 tokens, cross-checking MMQ against dequantize+GEMM
through the same call, reporting whether the toggle changed the result at all
(so a silently-inert case cannot pass), and timing both arms. That report is
also what located the defect above.
Fixture note: top-k ids are now sampled WITHOUT replacement. The runtime's sort
records one entry per (expert, token) and breaks on the first match, so a
duplicated expert makes it emit fewer rows than ne12*n_expert_used and trip its
own assert. Real top-k routing cannot select an expert twice; sampling with
replacement produced an abort that looked like a kernel bug and was bad data.
|
Pushed The width cap was wrong for MoEThe cap is right for dense multiplies: declining MMQ there hands the work to a well-tiled dense GEMM, which is exactly what the measurements in the original commit show. Declining it on a Measured on H200, 8 experts / top_k 2,
So the cap was costing 2.04x on precisely the widths a prefill uses. It is now dense-only, keyed off With MMQ reaching MoE at width, it beats the fallback at every width where the toggle changes the kernel at all:
At 1 and 4 tokens the toggle changes nothing, because the bespoke fused mix The test claimed coverage it did not haveThe commit message said the gate covers "expert counts 1 (dense The gate now exercises One note on the fixture, since the failure was instructiveTop-k ids are sampled without replacement. The runtime's sort records one entry per (expert, token) and Still open, and not addressed hereThe |
1 participant
Summary
The MMQ kernels for the mix qtypes (105/106) already existed and already
plumbed the out-of-band codebooks, but were reachable only through an env var
(
DFLASH_DS4_MIX_MMQ_PREFILL) and gated to RDNA. So everyne11 > 1multiplyfell back to dequantize-to-bf16 + dense GEMM — 48% of a 16-token batch's GPU
time sat inside
dequantize_rocmfp{2,3}_mix_kernel(nsys), i.e. the multiplydiscarded the sub-4 bpw artifact and ran 16-bit for its duration.
This makes that path reachable, gives it a correctness gate it never had, and
routes it per multiply on batch width — because MMQ is not uniformly better.
It wins on narrow batches and loses on wide ones, since the dequant path hands
a wide N to a well-tiled dense GEMM that the MMQ kernel does not out-tile. A
single process-wide switch therefore cannot be right: one served request wants
the narrow-batch win on its speculative-verify steps and the dense-GEMM win
on its prefill.
Also included: a dense-matvec improvement for the same qtypes (vectorized
activation loads, two output rows per warp), which is bit-exact.
CORRECTED (see eab505c and the comment below). An earlier revision of this
description said this change is dormant for DeepSeek-V4 because its 105/106
tensors are MoE experts reached through
ggml_mul_mat_id, "which does not takethis path". That reason is wrong.
ggml_cuda_mul_mat_iddoes consultggml_cuda_should_use_mmqand does dispatch to MMQ: measured on H200 with 8experts / top_k 2, MMQ engages from ~64 tokens and runs 1.7-2.3x faster than the
alternative. The gate in this PR now covers that dispatch directly.
What remains true is the observation that motivated the default: an A/B on a
real DS4 artifact (gfx1151) came back byte-identical with the flag on and off,
with prefill at 184.1 s +/-0.1%. Byte-identical output is strong evidence MMQ did
not run anywhere in that graph, since MMQ quantizes activations to int8 and
provably changes results. But I do not currently have a verified explanation
for why it did not run, and that A/B predates both the width gate and eab505c,
so it cannot be cited as this branch's behaviour. Treat DS4 impact here as
UNMEASURED pending a re-run, not as established zero.
The default therefore stays OFF with no in-tree opt-in caller — as a
conservative choice while that is unresolved, not because the path is
unreachable. The value this PR is offered on is the batched path being correct,
tested (now including MoE), and width-routed.
Impact
Implementation
perf(mix): vectorized activation loads and two output rows per warp in the dense 105/106 matvec. Bit-identical output — the per-block accumulation order is unchanged.feat(mix): the toggle becomes a real API (ggml_cuda_mix_mmq_enabled/set/clear/env_pinnedinggml-cuda.h), precedence explicit-call > env > compiled default. RenamedDFLASH_MIX_MMQsince it was never prefill-specific; the old spelling is still honoured.feat(mix): NVIDIA is no longer excluded — the DP4A/MMA tile these types declare is portable, and it now has a correctness gate rather than an assumption.feat(mix):ggml_cuda_should_use_mmqwidth-gates the mix qtypes atmix_mmq_max_ne11(1024 on RDNA 4, 256 elsewhere). Decode never reaches this gate:ne11 == 1takes the MMV kernel.test(mix): MMQ vs the already-validated matvec kernel acrossne111/4/16/64 and expert counts 1 (densemul_mat) and >1 (MoE).chore(build): wires the gate into the test targets, guarded byif(EXISTS)like its neighbours.Runtime Flags / Configuration
Default-on paths:
Disable path:
DFLASH_MIX_MMQ=0: forces the batched path off; wins over a backend's programmatic opt-in, so an operator can always override.DFLASH_DS4_MIX_MMQ_PREFILL: the previous spelling, still read as a fallback so existing scripts keep working.Debug/profiling-only flags:
DFLASH_MIX_MMQ_MAX_NE11=<n>: overrides the per-architecture width threshold, for re-sweeping the crossover.DFLASH_MIX_MMQ=1enables the batched path process-wide. Precedence isexplicit API call > environment > compiled default, so an operator running
DFLASH_MIX_MMQ=0is never silently overridden by a model backend.Validation
test_rocmfp_mix_mmq(new) on H200, default pathtest_rocmfp_mix_mmqunderDFLASH_MIX_MMQ=0test_rocmfp_mix_slice_matvec(covers the bit-exact matvec change)ggml-cuda,test_rocmfp_mix_mmqanddflash_serverfrom this branch onmainggml_mul_mat_id(MoE) with mix qtypes, H200, 8 experts / top_k 2Notes
Methodology, stated honestly.
2.06xmatvec figure was taken with a decode bench belonging to a densemodel line that is not part of this PR, so it is not reproducible from
this tree alone. The change is bit-exact, so correctness is covered by the
existing
test_rocmfp_mix_slice_matvec. Happy to have maintainers re-runthroughput on the RTX 3090 or Ryzen 395 AI Max per CONTRIBUTING, or to add a
standalone in-tree bench if you would prefer one.
host. The crossover is bracketed, not resolved: 256-1024 on RDNA 3.5 and
1024-2048 on RDNA 4. Power limit was not pinned; happy to re-run under your
methodology.
ne114..16 verify batches. It therefore takes the conservative RDNA 3.5 bound,
which keeps every measured NVIDIA win and declines only widths nothing has
measured there. Raise with
DFLASH_MIX_MMQ_MAX_NE11once swept.unrelated compile on the same box moved an identical request by ~5x with the
step count unchanged. All figures above were taken on an idle host with
arm order reversed and a warmup discarded.
Retraction carried in the history. An earlier revision of this work
reported a 35% DS4 regression under the flag. That was a harness artifact — a
failed arm left the previous server's ~91 GiB unified allocation live, so the
next arm measured contention. The controlled re-run is the 21.72 / 21.72 /
21.74 tok/s above. The original "inert on DS4" evidence was byte-identical
outputs, which cannot distinguish "same kernels ran" from "the path declined to
engage"; that conclusion held, but it now rests on throughput instead.
One open decision for maintainers: whether you want this landed with the
default off while the DS4 question is re-measured, or held until that re-run is
done. I would rather flag the uncertainty than let the earlier "unused by DS4"
framing stand, since it rested on a mechanism I have since disproved.
Mentioning @Luce-Org/maintainers per CONTRIBUTING now that this is ready for
review.