cuda: grouped-expert MUL_MAT_ID for small MoE verify batches

[davide221]

What

Opt-in grouped-expert path for MUL_MAT_ID on small token batches (2-16), targeting speculative-verify workloads on quantized MoE models. Enable with DFLASH_MMID_GROUPED=1.

Why

Consecutive draft tokens route to heavily overlapping expert sets, but mul_mat_vec_q_moe reads each (token, slot) pair's expert weights independently. Measured on a 256-expert top-8 MoE (Laguna XS 2.1):

verify rows	expert reads today	distinct experts (union)	duplication
8	64	39.4	1.62x
15	120	58.9	2.04x

How

• A tiny single-block prep kernel rank-sorts the batch's (token, slot) pairs by routed expert ([TAG_MMID_GROUPED] in mmvq.cu). Capture-safe: ids never leave the device, all launch shapes are static.
• The main kernel is mul_mat_vec_q_moe with one change: warp w processes sorted pair blockIdx.y*8 + w instead of (slot = blockIdx.y, token = w). Same-expert pairs are adjacent after sorting, so the warps of a block share an expert and duplicate weight reads hit L1/L2 instead of DRAM.
• Launch shape, per-warp structure, vec_dot sequence and warp reduction are identical to mul_mat_vec_q_moe → bit-exact outputs, CUDA-graph capture unaffected.
• DFLASH_MMID_GROUPED_TYPES bitmask: 1 = Q4_K (default), 2 = Q6_K (also lifts Q6_K's MUL_MAT_ID MMVQ ceiling to 16, keeping CUDA graphs on for those batches), 4 = Q4_0/Q8_0/Q5_K.
• Model-agnostic: any MoE with n_expert_used <= 16 (asserted).

Results

RTX 3090 (sm_86), Laguna XS 2.1 Q4_K_M + DFlash drafter, greedy HumanEval prompts. Outputs byte-identical to baseline in all cases (verified per-prompt, identical accept trajectories):

config	baseline tok/s	grouped tok/s	gain
verify width 8	212.0	219.9	+3.7%
verify width 15	~157	173.3	~+10%

Gains grow with verify width (more overlap to dedup), which lowers the marginal cost of wide/hedged verify configurations.

Notes

• Off by default; zero impact when the env is unset (single branch check in ggml_cuda_mul_mat_vec_q).
• Prep adds one ~µs kernel launch per MUL_MAT_ID call; included in all numbers above.
• Kernel design history: a register-blocked variant (tokens accumulated in registers per warp) lost 28% to spills; a warp-per-group variant lost 7% to occupancy collapse from early-exiting warps. The shipped design keeps the tuned kernel's exact occupancy profile and takes the dedup purely from cache locality.

---

Summary by cubic

Adds an opt-in grouped-expert path for MUL_MAT_ID on small MoE verify batches (2–16 tokens) to reduce duplicate expert weight reads and speed up speculative verify. Outputs remain bit-exact and CUDA graph capture stays supported, with gains up to ~10% at width 15.

• New Features
  • Enable with DFLASH_MMID_GROUPED=1; control quant types via DFLASH_MMID_GROUPED_TYPES (1=Q4_K default, 2=Q6_K, 4=Q4_0/Q8_0/Q5_K).
  • Tiny device prep kernel rank-sorts (token, slot) pairs by expert; static launch shapes; capture-safe.
  • Grouped mul_mat_vec_q_moe processes sorted pairs so warps share experts; same launch/warp structure → bit-exact outputs.
  • For enabled types, raises MUL_MAT_ID batch ceiling to 16 and keeps CUDA graphs on; asserted for n_expert_used <= 16; off by default.
  • Perf (RTX 3090, Laguna XS 2.1 Q4_K_M): width 8 +3.7% tok/s; width 15 ~+10%.

^{Written for commit 15b5b1b. Summary will update on new commits.}

[cubic-dev-ai]

2 issues found across 1 file

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="ggml/src/ggml-cuda/mmvq.cu">

<violation number="1" location="ggml/src/ggml-cuda/mmvq.cu:253">
P3: The `MMID_META_GS` macro is defined but never read or written by any kernel. It creates a 257-int (~1 KB) gap in the meta allocation that serves no purpose. This inflates the per-call scratch allocation and makes the layout harder to follow. Consider removing the dead constant and collapsing the layout so `MMID_META_PT` directly follows `MMID_META_GE + MMID_GROUPED_MAX_PAIRS`, which would reduce `MMID_META_INTS` from 1026 to 769.</violation>

<violation number="2" location="ggml/src/ggml-cuda/mmvq.cu:1577">
P2: Enabling `DFLASH_MMID_GROUPED` can switch small MUL_MAT_ID calls on non-Turing+/non-NVIDIA devices into the new grouped CUDA path even though `get_mmvq_mmid_max_batch()` only enables this feature for NVIDIA Turing+. Consider applying the same compute-capability guard here so the opt-in path matches the routing/graph eligibility logic.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P2: Enabling DFLASH_MMID_GROUPED can switch small MUL_MAT_ID calls on non-Turing+/non-NVIDIA devices into the new grouped CUDA path even though get_mmvq_mmid_max_batch() only enables this feature for NVIDIA Turing+. Consider applying the same compute-capability guard here so the opt-in path matches the routing/graph eligibility logic.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At ggml/src/ggml-cuda/mmvq.cu, line 1577:

<comment>Enabling `DFLASH_MMID_GROUPED` can switch small MUL_MAT_ID calls on non-Turing+/non-NVIDIA devices into the new grouped CUDA path even though `get_mmvq_mmid_max_batch()` only enables this feature for NVIDIA Turing+. Consider applying the same compute-capability guard here so the opt-in path matches the routing/graph eligibility logic.</comment>

<file context>
@@ -1294,6 +1572,25 @@ void ggml_cuda_mul_mat_vec_q(
 
+    // [TAG_MMID_GROUPED] grouped-expert path for small MUL_MAT_ID batches.
+    if (ids && ncols_dst >= 2 && ncols_dst <= MMVQ_MAX_MOE_BATCH_SIZE &&
+        mmid_grouped_env() && mmid_grouped_type_ok(src0->type)) {
+        const int np = (int) (nchannels_dst*ncols_dst);
+        GGML_ASSERT(np <= MMID_GROUPED_MAX_PAIRS && "DFLASH_MMID_GROUPED supports n_expert_used <= 16");
</file context>

Suggested change

	mmid_grouped_env() && mmid_grouped_type_ok(src0->type)) {
	mmid_grouped_env() && mmid_grouped_type_ok(src0->type) &&
	GGML_CUDA_CC_IS_NVIDIA(ggml_cuda_info().devices[ggml_cuda_get_device()].cc) &&
	ggml_cuda_info().devices[ggml_cuda_get_device()].cc >= GGML_CUDA_CC_TURING) {

[cubic-dev-ai]

P3: The MMID_META_GS macro is defined but never read or written by any kernel. It creates a 257-int (~1 KB) gap in the meta allocation that serves no purpose. This inflates the per-call scratch allocation and makes the layout harder to follow. Consider removing the dead constant and collapsing the layout so MMID_META_PT directly follows MMID_META_GE + MMID_GROUPED_MAX_PAIRS, which would reduce MMID_META_INTS from 1026 to 769.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At ggml/src/ggml-cuda/mmvq.cu, line 253:

<comment>The `MMID_META_GS` macro is defined but never read or written by any kernel. It creates a 257-int (~1 KB) gap in the meta allocation that serves no purpose. This inflates the per-call scratch allocation and makes the layout harder to follow. Consider removing the dead constant and collapsing the layout so `MMID_META_PT` directly follows `MMID_META_GE + MMID_GROUPED_MAX_PAIRS`, which would reduce `MMID_META_INTS` from 1026 to 769.</comment>

<file context>
@@ -233,8 +233,67 @@ static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna4(ggml_type
+#define MMID_GROUPED_MAX_TPG   8
+#define MMID_META_NG 0
+#define MMID_META_GE 1
+#define MMID_META_GS (MMID_META_GE + MMID_GROUPED_MAX_PAIRS)
+#define MMID_META_PT (MMID_META_GS + MMID_GROUPED_MAX_PAIRS + 1)
+#define MMID_META_PS (MMID_META_PT + MMID_GROUPED_MAX_PAIRS)
</file context>

[davide221]

Superseded: ggml is vendored into lucebox-hub now, PR moved there.