fix(qtip): gate the trellis grouped GEMM on m-tile occupancy, not token count

[heydryft]

The regression

The trellis grouped GEMM does replace the gather — qtip_gather_gemv_warp_kernel disappears entirely and qtip_lut_grouped_gemm_kernel takes 62.4% of GPU kernel time at B=32. But aggregate throughput collapsed:

leg	B	ragged	agg tok/s	QTIP kernel selected	GPU busy	launches/tok
R8	8	on	34.19	gather_gemv_warp 38.5%	67.7%	1659
R32	32	on	6.00	lut_grouped_gemm 62.4%	90.8%	1675

Widening the batch 4× made serving 5.7× worse. Memory controller 3.5%, SM issue 22.5%, 48.2% unallocated warps — saturated doing nothing. The kernel ran 7,888 µs/call, 133× above its own bandwidth bound (nsys, exclusive H200, 97a65d643).

The cause is a threshold, not a kernel defect

The gate was n_tokens > DECODE_REGIME_MAX_TOKENS with DECODE_REGIME_MAX_TOKENS = 8, so it flipped at n = 9.

The grouped GEMM sorts pairs by expert and walks sum_e ceil(count_e / GROUPED_TILE_M) m-tiles, staging each woken expert's packed bytes once per m-tile. Its whole advantage over the per-pair GEMV is therefore the average tile fill:

pairs per woken expert = n*k / (E * (1 - (1 - k/E)^n))
fill                   = that / GROUPED_TILE_M

At V4's top-6-of-256:

n	pairs	woken experts	pairs/expert	tile fill
32	192	136.1	1.41	8.8%
128	768	243.7	3.15	19.7%
512	3072	256.0	12.00	75.0%
683	4098	256.0	16.01	100.0%
1024	6144	256.0	24.00	150.0%

The in-tree A/B brackets exactly that curve: 1.00× at N=128 and N=512, 2.41× at N=1024.

The kernel is fine. It was switched on roughly two orders of magnitude too early.

The fix

Gate on the property the kernel actually needs — grouped_gemm_tiles_amortize: pairs per woken expert ≥ GROUPED_TILE_M, i.e. an average m-tile is full.

Chosen over raising the constant to ~683 because a hardcoded number rots the moment top_k, num_experts or GROUPED_TILE_M moves — and the last constant that was not re-derived is this regression. The predicate needs no device data: it is a closed form in three host-side scalars, so it is still safe inside a CUDA-graph capture.

Uniform routing is the conservative assumption: a skewed real router wakes fewer distinct experts and so fills tiles more than this predicts, so the gate errs toward switching late, never early.

DECODE_REGIME_MAX_TOKENS survives as the RUN-161 capture-safety floor, and ARC_QTIP_ONDEVICE_MOE_MAX_TOKENS still pins the GEMV arm outright so a harness can A/B the two kernels.

Applied to both rungs

The bitshift (qtip2b) rung carried the identical defect at bitshift.rs:1841, behind a comment arguing the cap of 8 was correct there because its over-boundary path is the amortizing kernel. True — but only once the tiles are full; at 9 tokens they are not. Same predicate, same reasons. The bitshift rung is not measured in this PR — the checkpoint under test is qtip2.

Also: --v4-ragged-decode documented a false equivalence

mistralrs-cli/src/args/mod.rs claimed the flag was "equivalent to ARC_V4_XS_PER_SEQ=1, which remains supported". It was not. v4_ragged_decode defaults to false, so every build path called with_v4_ragged_decode(false) → request_xs_per_sequence(false), which latched the OnceLock before the first read could ever consult the environment. Instrumentation on the box showed batch_can_be_ragged=false xs_per_seq_enabled=false in a run with the variable set — the D32 leg above.

Not asking for a feature is not the same as asking for it to be off, so only true latches now and the documented fallback works. A caller that genuinely wants to force it off against the environment can still call mistralrs_core::request_xs_per_sequence(false) directly.

Tests

cargo test -p mistralrs-quant --lib qtip::gather_policy — 9 passed.

• tile_fill_matches_the_v4_curve_in_the_module_docs pins the fill table the derivation is written from, so the docs cannot go stale silently.
• grouped_gate_refuses_decode_widths_and_accepts_full_tiles asserts the regression directly: no selection at n = 9…512, selection at 683+ — and that the same token count flips the other way under different routing (32 tokens × top-6 of 8 experts), which is the property a constant cannot have.
• grouped_min_tokens_is_the_exact_flip_point checks the reported boundary against the predicate at n and n−1 across 30 routing shapes.
• grouped_boundary_is_the_fused_gemv_boundary pins that the grouped and fused-GEMV boundaries leave no token range where dispatch falls through to the dequantize-materialize path. They are set by two independent constants (DEQUANT_TRAFFIC_RATIO, GROUPED_TILE_M) that both happen to be 16; if either moves, this fails loudly instead of opening that window silently.

Runs now also log which MoE kernel was selected and at what tile fill (log_grouped_gemm_engaged_once), so a results artifact can be read back to the dispatch that produced it.

---

Re-measurement on this branch — and a retraction of the premise

Same box, same instrument, same harness rule, 89ec93140, binary 85aca0165f58…. Legs F8 / F32 against the prior R8 / R32, both with --v4-ragged-decode.

The fix does what it says. The new log_grouped_gemm_engaged_once line settles it directly:

QTIP MoE: trellis grouped GEMM engaged at 1572 tokens x top-6 of 256 experts —
m-tiles 230% full (36.84 pairs per woken expert against GROUPED_TILE_M=16);
the amortization boundary for this routing shape is 683 tokens.

and the profile shows the dispatch split that was not there before — qtip_gather_gemv_warp_kernel reappears at B=32 for the 32-token decode steps (534 calls) while the grouped GEMM keeps the 1,572-token prefill chunks (618 calls):

leg	B	gather_gemv	lut_grouped_gemm
R32 (before)	32	absent	1292 calls, 62.4%
F32 (after)	32	534 calls, 2.7%	618 calls, 58.2%

But it does not restore batch scaling, and the premise it was filed against does not survive contact with the data.

The B=32 leg never decodes. Aggregate tok/s in 2 s bins over each whole leg:

leg	trajectory	end state
R8 / F8	climbs to ~36 tok/s at t≈14 s / t≈34 s, then flat	8 seqs, 46–47 tokens each — genuine steady state
R32 / F32	monotonically climbing to the last bin, 1 → 26 / 1 → 21 sequences	2–4 tokens per sequence — still admitting

At 18 s of nsys, B=32 executed roughly 4–6 decode steps total (534 GEMV calls ÷ ~120 MoE calls per forward). The rest was prefill. So the "34.19 tok/s at B=8 → 5.99 at B=32" comparison placed a steady-state decode window beside a prefill-admission window. It is not a batch-scaling measurement, and the 5.7× is not a decode regression.

Matched-depth check, last 10 s of each token stream:

B	before	after
8	36.80 tok/s (8 seqs, 46 tok each)	37.60 tok/s (8 seqs, 47 tok each)	identical, as predicted — the gate is unchanged at n ≤ 8
32	9.90 tok/s (26 seqs, 2–4 tok each)	6.40 tok/s (21 seqs, 2–4 tok each)	both mid-ramp, neither is decode

And tile occupancy does not explain the kernel's inefficiency either. The grouped GEMM's mean call cost rose to 15,600 µs once the cheap decode calls were routed away from it and only prefill remained — at 230% tile fill, still ~60× above its bandwidth bound, with DRAM read at 1.9% and 48.4% unallocated warps. The 7,888 µs/call figure in the original report was an average of expensive 1,572-token prefill calls and cheap 32-token decode calls; separating them makes the prefill call look worse, not better. Whatever is costing 60–130× on this kernel is not m-tile emptiness.

What this PR is, then

A correct gate, landed on its own merits: n_tokens > 8 is indefensible as a boundary for a kernel whose amortization begins at 683 tokens for this routing, it is not the property the kernel needs, and it silently rots when top_k / num_experts / GROUPED_TILE_M move. B=8 is unchanged, and B=32 decode steps now reach the GEMV instead of a 9%-full tile.

It is not the fix for the B=32 throughput number. That number is V4's known prefill disease (memory/mission/BUDGET_V4_PREFILL.md: 11.8 ms per prompt token, ~99.5% GPU-busy) reached through a harness whose waitsteady rule declares steady state as soon as the decode cohort is wide, which at B=32 happens while ~30 sequences are still queued for prefill. Two follow-ups worth filing separately: (1) the grouped GEMM is 60×+ off its bandwidth bound at full occupancy; (2) waitsteady.py needs a per-sequence token-depth floor, not just a cohort-width gate, or every B≥16 leg will keep reporting prefill as decode.

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     73        25129        18043         4319         2767
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     74         4600         4597            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  145        15139        12482          811         1846
 Shell                    37         8630         5732         2323          575
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1498         1294           54          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                203        44777            0        34735        10042
 |- BASH                  72         1654         1202          331          121
 |- C                      3           17           17            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  18          708          708            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  66         2051         1716           77          258
 |- TOML                   6          207          164            0           43
 |- YAML                   5           41           36            5            0
 (Total)                            50548         4687        35277        10584
─────────────────────────────────────────────────────────────────────────────────
 Rust                    671       324626       280130        15983        28513
 |- Markdown             489        27661          471        23865         3325
 (Total)                           352287       280601        39848        31838
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1317       484719       346246        87004        51469
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

Post-merge note, recorded because this PR is not inert and the merge summary should not imply it was.

What changes for a user who sets no flags. The gate at bitshift.rs:1850 is now n_tokens <= DECODE_REGIME_MAX_TOKENS || !gather_policy::grouped_gemm_tiles_amortize(...). Working the arithmetic:

• b=1 decode (n=1): unchanged. 1 <= 8 short-circuits true on both sides -> fused gather GEMV, as before.
• Prefill n≈2000: unchanged. expected_pairs_per_distinct_expert(2000,6,256) = 46.9 >= GROUPED_TILE_M(16) -> amortize -> grouped GEMM, as before.
• The middle band moved, and it is real: grouped_gemm_min_tokens(6,256) == 683, so for 9 <= n_tokens <= 682 the gate now prefers gather-GEMV where it previously chose the grouped GEMM. That covers batched decode at B>=9 and every prefill chunk under 683 tokens.

So the framing 'decode and prefill are unchanged' is true only at the two endpoints. Anyone batching above 8 gets a different MoE kernel with no flag set.

Why I am not reverting it: both paths compute the same thing — this is a dispatch heuristic, not a math change — so the exposure is performance, not correctness. And the new gate is the more principled one (actual m-tile occupancy beats a token-count proxy).

⚠️ But it is unmeasured on the band it moved, and there is reason for doubt. The gather-GEMV path is documented as a no-dedup GEMV with 3.15x redundant reads, so preferring it from 9 to 682 tokens could be slower for mid-size prefill chunks than the grouped GEMM it displaced. That is exactly the band chunked prefill lives in.

Action for the next session with a box: A/B the 9-682 band specifically — B=16/32/64 decode and 128/256/512-token prefill chunks — before treating this gate as settled. Related: #99 raises GROUPED_TILE_M to 64 on SM90+, which moves this threshold ~4x; that PR must make the gate read grouped_tile_m_for_cc rather than the constant, or the gate and the kernel will disagree about tile size.