[Kimi-K3] Allow DSPARK verify on cutedsl_mla (fold_sq)

[yhyang201]

_dspark_verify_on_decode_backend caps cutedsl_mla at q_len <= 4, from when the cute-dsl MLA decode kernel rejected q_len >= 5. flashinfer's monolithic MLA decode now folds seq_len_q into the head dim (fold_sq, added in flashinfer #3309, gated-fixed in #3664, present in the pinned 0.6.15.post1), so q_len > 4 is supported.

K3 DSPARK verify uses q_len = block_size(7) + 1 = 8, which exceeded the cap, so verify silently fell back to trtllm_mla — the fold-less path that re-reads the KV per query row and is slow at long context.

Dropping the cap (cute-dsl serves any verify width via fold_sq) routes verify to the fold path. Measured on Kimi-K3, TP8, 8xB300, bf16 kvcache, bs=1 isl=900000 osl=2048 (decode backend the only variable):

	trtllm_mla	cutedsl_mla
decode ITL	10.18 ms	5.07 ms
decode output throughput	98.2 tok/s	197.4 tok/s
per spec-step	~44 ms	~22 ms
acc_length	4.34	4.26

~2x decode throughput at 900k context, acc_length unchanged.

GPU timeline (nsys, B300, bf16, q=8, 900k) — isolated MLA decode kernel

Isolated flashinfer MLA decode call profiled with nsys (--trace=cuda), same shapes, decode backend the only variable.

trtllm_mla — one fold-less fmhaSm100f…HQk576… kernel, 1.263 ms:

cutedsl_mla — …monolithic mla_decode_fp16 Blackwell… fold_sq kernel + a small split-KV reduction, ~320 µs total (~3.9x faster kernel):

The pure MLA attention kernel is ~4x faster; MLA is roughly half of a decode step, so the e2e decode speedup is ~2x.

     if backend == "cutedsl_mla":
-        # The cute-dsl kernel rejects q_len >= 5 with no fallback.
-        return q_len <= 4
+        # cute-dsl monolithic MLA decode folds the verify tokens into the head
+        # dim (fold_sq), so it serves any DSPARK verify width. Needs flashinfer
+        # >= 0.6.15 (older builds reject q_len >= 5).
+        return True

Which decode backend? (isolated kernel microbenchmark)

Isolated flashinfer MLA decode on B300, flashinfer 0.6.15.post1, µs/call (median). H = query heads per rank (TP1=96, TP4=24, TP8=12, TP16=6; MQA, 1 KV head). q = spec verify tokens.

Typical K3 @ 1M-context configs (per-rank, µs/call, bold = pick)

config (per-rank)	q	bf16 trtllm	bf16 cutedsl	fp8 trtllm	fp8 cutedsl
TP1+DCP8 (H=96, ctx~125k)	1	ERR	129.3	ERR	122.2
TP1+DCP8 (H=96, ctx~125k)	8	ERR	170.3	ERR	125.6
TP8 (H=12, ctx=1M)	1	254.4	330.7	130.4	164.3
TP8 (H=12, ctx=1M)	8	1413.0	359.4	666.6	168.2

For K3 DSPARK verify (q>1) at long context, cute-dsl is either the only option (TP1+DCP8) or ~4x faster (TP8 @ 1M) — hence this PR relaxes the guard to allow it.

Decision rule

1. q = 1 → trtllm-gen. Always faster, even at 1M ctx (bf16 254 vs 331; fp8 130 vs 164). fold_sq needs q>1 to help.

2. trtllm-gen unavailable → cute-dsl (only option). trtllm-gen fails with computeCtaAndClusterConfig: numHeadsQ/numHeadsKv not supported at TP4 (H=24, all q/ctx) and TP1 (H=96) at long ctx or q≥7. TP8/TP16 always work; cute-dsl works everywhere. Dtype-independent.

3. q > 1 and both available (TP8/TP16): pick by context length — cute-dsl wins above:

   q	bf16	fp8
   2	~400k	~800k
   4	~180k	~350k
   8	~100k	~200k

   Below the threshold, trtllm-gen is ~30% faster wherever it runs.

Why: trtllm-gen re-reads KV per query row → time ≈ O(ctx × q). cute-dsl folds the q verify tokens into the head/MMA tile (F = largest divisor of q with H·F ≤ 128) → flat ~120µs up to ~256k, rising to only ~360µs bf16 / ~170µs fp8 at 1M.

Raw q×ctx sweep, TP8 / H=12 (trtllm | cutedsl µs/call; TP16/H=6 is within noise)

bf16:

q	131k	500k	1000k
1	94.1 / 121.8	142.7 / 176.6	254.4 / 330.7
2	96.6 / 126.5	215.8 / 189.8	407.5 / 357.0
4	114.7 / 124.8	376.5 / 190.7	732.7 / 358.1
8	200.9 / 123.9	717.1 / 192.6	1413.0 / 359.4

fp8:

q	131k	500k	1000k
1	92.7 / 120.6	96.5 / 121.2	130.4 / 164.3
2	94.1 / 124.3	110.9 / 121.2	198.2 / 165.6
4	95.7 / 121.6	183.7 / 122.6	347.2 / 167.0
8	103.4 / 123.6	343.3 / 123.2	666.6 / 168.2

TODO / follow-ups (draft)

• Context/TP-aware backend choice (beyond this eligibility guard): the fastest backend is (TP, q, context, kv_dtype)-dependent — where both run and q>1, cutedsl only wins above a context threshold (bf16 q=8 ~100k … q=2 ~400k; fp8 ~2x higher), below which trtllm-gen is ~30% faster. sglang can't switch decode backend after launch (backends are fixed at init and baked into the decode CUDA graphs), so a per-step context-aware choice would need capturing graphs for both backends + a per-batch seq_len dispatcher.
• Evaluate making cutedsl_mla the default decode backend for Kimi-K3 (currently trtllm_mla) -- needs correctness + no-regression, especially short context / low q where trtllm-gen is faster.

---

CI States

Latest PR Test (Base): ❌ Run #30986145615
Latest PR Test (Extra): ❌ Run #30986145542

[yhyang201]

A quick summary on picking the decode/verify attention backend for K3 MLA (cute-dsl vs trtllm), based on kernel microbenchmarks on B300:

DCP: cute-dsl only — trtllm's decode kernel doesn't implement DCP and errors out; already forced, no choice.

Pure TP8 (no DCP):

• Plain decode (q=1) → trtllm, faster all the way to 1M in both bf16/fp8; cute-dsl has no edge.
• DSPARK verify (q=8) → depends on context length:
  • bf16: >~100k use cute-dsl, otherwise trtllm
  • fp8: >~200k use cute-dsl, otherwise trtllm

Why: trtllm slows down ~linearly with (context × q), while cute-dsl stays roughly flat thanks to fold_sq — so the longer the context, the more cute-dsl wins (at 1M the MLA kernel is ~4x faster, e2e decode ~2x). fp8 halves the KV bytes so trtllm holds up longer, pushing the crossover out ~2x.