[DSA] Enable tilelang fp8_e4m3 KV cache on CUDA (raw layout) — verified on sm_121

[beastllama]

Stacked on #36507; mergeable once that lands. Opening now for visibility per #36830.

Summary

The tilelang DSA backend already ships a complete raw-fp8 sparse attention kernel (sparse_mla_fwd_decode_partial_fp8 — generic TileLang, no HIP intrinsics). It is unreachable on CUDA because three plumbing sites are HIP-gated, and CUDA is always handed the scaled 528 B/token pool layout the kernel cannot parse. This PR routes the raw 512 B/token layout on CUDA when both DSA backends are tilelang, relaxes the arg gate accordingly (SM89+, hard error on mixed backends), and keys the raw fused-quant write on layout rather than platform (strictly more correct on HIP as well).

Changes (5 files + test)

1. srt/arg_groups/overrides.py — _check_tilelang_dsa_fp8_kv: allow CUDA when prefill+decode backends are both tilelang and SM>=89; explicit ValueError on mixed backends; warning log on enablement (serves as a point-of-effect check)
2. srt/mem_cache/kv_cache_configurator.py — calculate_mla_kv_cache_dim: CUDA + both-tilelang -> raw dim -> raw fp8 pool
3. srt/mem_cache/memory_pool.py — _write_mla_kv_buffer: raw-write branch keyed on not dsa_kv_cache_store_fp8 instead of _is_hip; None-guard for cache_k_rope (NoPE models: GLM-5.3-Flash has qk_rope_head_dim=0)
4. kernels/ops/attention/dsa/tilelang_kernel.py — dispatch to the fp8 partial kernel on CUDA under the same condition, with fp8-sized tiles for the ~100KB dynamic-smem class. The broader GB10 tile retune is deliberately NOT in this PR (separate arch-conditional proposal per our GLM-5.3-Flash support #36507 comment) — this diff is fp8 routing only.
5. srt/layers/attention/dsa_backend.py — supports_mha_one_shot=False for raw-fp8 (the one-shot fp8 dequant helpers assume the scaled layout)
6. test/registered/attention/test_dsa_tilelang_fp8_kv_cuda.py — one-hot exactness vs gathered reference, spread case within the analytic fp8-prob budget, and a scrambled-index negative control that must fail (guards against a vacuous comparator). Executed on sm_121: 2 passed.

Hardware coverage — stated plainly

Validated end-to-end ONLY on sm_121 (2x DGX Spark GB10). The SM89+ gate follows the kernel's own requirements, but SM89/SM90 are untested by us. Review/CI verification on other archs invited.

Validation (2x DGX Spark GB10 sm_121, GLM-5.3-Flash NVFP4, TP=2, DFlash2 drafter bf16)

• Kernel: one-hot exact rel_err 0.000000; scrambled-index negative control fails at 91%; spread case 2.35% = analytic fp8-prob GEMM noise
• Boot-gate negative control (tilelang + flashmla_sparse + fp8) dies with the new error, as required
• Decode n=5 medians: 29.2 code vs 29.8 same-image bf16 (parity); temp-0 outputs 4/5 exactly identical; 32k-depth recall probe pass (re [Bug] [ROCm] DeepSeek-V4-Flash FP8 on MI300X (gfx942): output degrades with context length — coherent <100 tokens, garbage by ~300, NIAH 0/3 #36390); TTFT@16k warm 6.6 s vs 7.9 s bf16 (~17% faster prefill)
• Concurrency note: an earlier "raw layout resolves single-stream" caveat was wrong — root cause was the mamba state cache capping ALL DFLASH configs at 2 streams, filed separately as [DFLASH] Mamba state cache silently caps effective concurrency on hybrid-KDA models — 5 state slots/request, info-level log only (+75% throughput hidden behind one flag) #36889 (+75% aggregate at c12 once lifted). fp8-KV and 8-way concurrency coexist.

Credits

Kernel-path analysis and port from a 2x DGX Spark homelab; validation protocol and full experiment ladder: https://github.com/beastllama/GLM-5.3-Flash-DFlash2-SGLang-2x-DGX-Spark

Addresses #36830. Concurrency companion: #36889. Relevant to #36390 (ROCm fp8 accuracy) and the sm_121 bring-up in #36845/#36806.

---

CI States

Latest PR Test (Base): ❌ Run #33206743796
Latest PR Test (Extra): ❌ Run #33206743584
Latest PR Test (AMD ROCm 7.2): ❌ Run #33206743786

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 06a1db7b6f

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve DCP mapping in raw-FP8 cache writes

When CUDA TileLang FP8 is combined with --dcp-size > 1, loc contains widened logical token IDs, but set_mla_kv_buffer_triton_fp8_quant writes directly to dst_buffer[loc] without filtering by DCP ownership or dividing by the DCP world size. The normal MLA writer passes DCP_RANK/DCP_WORLD_SIZE and performs both operations, so this newly enabled branch can write the wrong row or access beyond the rank-local pool before attention runs. Add the same DCP mapping to the fused quant writer or reject this configuration.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Support or reject DCP before routing CUDA FP8

With --dcp-size > 1, both DSA forward paths request LSE during decode and target verification via _should_return_dsa_dcp_lse, but CUDA FP8 now enters this partial-and-combine branch and immediately asserts that return_lse is false. Consequently, an otherwise accepted CUDA TileLang FP8 configuration aborts on its first DCP decode or verify step. The FP8 path must return LSE for cross-rank reduction, or argument validation must reject TileLang FP8 with DCP.

Useful? React with 👍 / 👎.

[beastllama]

Friendly ping for a maintainer with triage rights — CI has never executed on this PR.

All 9 red checks are the same fast-failing gate job, not a test failure:

PR Labels: [jit-kernel, memory-pool]
Require run-ci: true
Require run-ci-extra: true
Missing required label 'run-ci'.

Every gate exits in 2–7 s and 71 downstream jobs are SKIPPED, so nothing has actually been
compiled or run yet. I only have pull permission on the repo, so I can't add the labels myself.

Could someone add run-ci and run-ci-extra?

Review feedback so far is addressed in 436a89b:

• black formatting on the three touched files
• registered the fp8 test in the CUDA CI registry
• answered the Codex-bot P1 by rejecting tilelang-fp8 + DCP in argument validation (the raw
  writer has no DCP rank filter, so the combination is now refused rather than silently wrong)

Worth noting for whoever reviews: the registered test exercises the fp8 kernel on the H100 CI
runner, so merging this also buys SM90 coverage for a path that until now was only validated on
sm_121 (GB10) hardware I have locally. Happy to iterate on anything the run turns up.

[Hansen523]

Validated on SM90 (8x H20-3e) — works, ~1.8x KV capacity at identical footprint

Following up from #36830, where you offered to share early. Your gate is SM89+, but your validation is sm_121, so we ran this on Hopper to close that gap. It works.

Setup

8x NVIDIA H20-3e (sm_90, 140 GB each), single node. ZhipuAI/GLM-5.3-Flash FP8 weights.

Because this PR branches from xinyuan/glm-5.3-flash-support at 08-28 14:56, it predates #36884 and #36885. Testing it as-is would have shown their symptoms (silently wrong output at tp==dp==ep) and I'd have misattributed them to this PR — so I merged your branch onto 7fa1924c0448 (that branch's head, both fixes included) first. Clean merge, no conflicts, all six of your files intact.

--tp-size 8 --data-parallel-size 8 --enable-dp-attention --expert-parallel-size 8 \
--quantization fp8 --disable-shared-experts-fusion --moe-runner-backend deep_gemm \
--dsa-prefill-backend tilelang --dsa-decode-backend tilelang \
--context-length 32768 --max-running-requests 256 --mem-fraction-static 0.80 \
--mm-enable-dp-encoder --kv-cache-dtype {bfloat16 | fp8_e4m3}

Everything identical between the two arms except --kv-cache-dtype.

The gate admits sm_90 and logs the expected line:

Enabling the tilelang DSA fp8_e4m3 KV cache on CUDA: KV is stored as raw fp8
(no per-tile scales), matching the ROCm tilelang path.

Result: same bytes, 1.795x the tokens

	bf16	fp8_e4m3
KV Cache is allocated dtype	torch.bfloat16	torch.float8_e4m3fn
#tokens	2,742,144	4,922,240
KV size	32.48 GB	32.47 GB
max_mamba_cache_size	211	211

1.795x more tokens in the same 32.5 GB. Short of the 2x the layout implies, which is expected: only the 11 full-attention layers carry the halved KV, and the DSA indexer cells appear not to halve with it. The mamba pool is untouched at 211 either way, so the gain is purely the KV side.

Output parity at temperature=0

Final answers (text after </think>, --reasoning-parser glm45):

prompt	bf16	fp8_e4m3
"中国的首都是哪里?只回答城市名。"	北京 (57 tok)	北京 (59 tok)
"2+3等于几?只回答数字。"	5 (76 tok)	5 (76 tok)
"9.11和9.8哪个大?只回答数字。"	9.8 (195 tok)	9.8 (141 tok)

3/3 final answers identical. Completion length differs on 2 of 3 — the reasoning trajectory shifts slightly, which is what one would expect from fp8 KV, and it converges on the same answer each time. Worth noting the 9.11-vs-9.8 case, which is a known trip-up for many models: fp8 got there in 141 tokens against bf16's 195.

I have not attempted a broad accuracy suite; this is parity spot-checking, not a GSM8K run.

On your open item — max_running_requests resolving lower than requested

I think this may not be the raw layout. On an earlier run I passed --max-running-requests 8 and the engine resolved 1; that is pool_configurator.py:714:

num_reqs = get_schedule().max_running_requests // kvc.ps.attn_dp_size

8 // 8 = 1. The logged value is per DP rank, not global. Raising the flag to 256 in the runs above made it a non-issue, with both arms behaving identically.

There is a second, independent cap on this model worth ruling out before attributing anything to the layout — the mamba state cache:

max_running_requests is capped to N by the mamba state cache
(max_mamba_cache_size=M, R state slots per request)

We measured the per-slot state scaling exactly with attn_tp_size (conv 1.166 → 2.329 → 4.695 MB, ssm 33.24 → 66.54 → 133.4 MB at dp2/dp4/dp8), and --mamba-ssm-dtype bfloat16 roughly doubles that pool with byte-identical output at temperature=0. If the mamba pool is what is binding, that flag moves it and the KV layout does not — an easy way to tell the two apart. (This is also what #36889 is about, I think.)

Summary

sm_90 behaves as your gate assumes. Happy to run anything else that would help this land — a longer prompt set, deeper context, or the boot-gate negative control on Hopper. Full logs available on request.

[marknx]

Tested this PR on 2x DGX Spark (GB10, sm_121) over 10GbE TCP, TP=2, serving GLM-5.3-Flash
(Glm5NextForConditionalGeneration, NVFP4-mixed self-quant) with DFLASH speculative decoding
(block 6), --dsa-prefill-backend tilelang --dsa-decode-backend tilelang.

1. Bug + fix: crashes for models with qk_rope_head_dim = 0 (GLM-5.3-Flash)

As-is, boot fails during CUDA graph capture / first kernel build with:

ValueError: Unsupported k_dim 0

Root cause: sparse_mla_fwd_decode_partial_fp8 emits the tail path unconditionally
(q_tail/k_tail allocs, tail load loop, tail T.gemm). GLM-5.3-Flash has qk_rope_head_dim = 0
→ d_tail = 0 → tilelang MMA emitter rejects a zero-width GEMM. The bf16 kernels avoid this via
the v1 (no-tail) / v2 split.

Fix that works for us — compile-time guards, mirroring the bf16 specialization
(4 spots inside sparse_mla_fwd_decode_partial_fp8):

if d_tail > 0:
    q_tail_buf = T.alloc_shared([h_per_block, d_tail], fp8_dtype)
    k_tail_shared = T.alloc_shared([BI, d_tail], fp8_dtype)
...
if d_tail > 0:
    T.copy(q_fp8[b_i, s_i, H0:H1, d_v:], q_tail_buf)
...
if d_tail > 0:
    for bi_i, j in T.Parallel(BI, d_tail):
        page = page_idx_shared[bi_i]
        k_tail_shared[bi_i, j] = kv_fp8[b_i, page, g_i, rope_offset_fp8 + j]
...
if d_tail > 0:
    T.gemm(q_tail_buf, k_tail_shared, acc_s, transpose_B=True,
           policy=T.GemmWarpPolicy.FullCol)

With these guards the server boots and serves correctly on sm_121.

2. Results on sm_121 (2x GB10)

• KV pool: 85,376 → 161,600 tokens (+89 %) at identical settings
  (mem-fraction 0.94, ctx 57,344, max-running-requests 4, DFLASH block 6, mamba-ssm bf16)
• Decode speed: unchanged within noise (34–47 tok/s across 16k–48k prompts)
• Needle retrieval: 6/6 across 16k/48k x begin/mid/end — no retrieval regressions
• GSM8K (200 tasks, temp 0): 98.5 % vs 98.0 % bf16 baseline — no regression

3. Quality caveat: code generation drops measurably with raw-layout FP8 KV

• HumanEval (164 tasks): 85.4 % with fp8 KV vs 90.2 % bf16 baseline (same box, same day,
  same server config otherwise). Confirmed by a second independent run: 85.4 % again (140/164 both
  times — deterministic, not noise).
• GSM8K and needle retrieval are unaffected; the drop is specific to long precise code
  generation, which is consistent with unscaled (raw) e4m3 KV rounding noise accumulating.
• For comparison, the vLLM-side GB10 community setup using the scaled fp8_ds_mla packed format
  reports near-identical KLD to bf16. Might be worth a note in the PR description that raw layout
  trades some generation quality for simplicity, and/or a follow-up for a scaled layout.

4. Long-context status on sm_121

• ctx 131,072 with fp8 KV + DFLASH: fully functional — needle retrieval 6/6 at 100k/125k depth,
  decode still 29–36 tok/s on 125k prompts. Pool reached our --max-total-tokens cap (196,608).
• ctx 262,144 boots (pool 254,784), but the first ~200k-token prefill kills the worker
  (HTTP 500, then connection refused) — this reproduces [GB10] Long prefill (>40k tokens) exhausts unified memory and silently kills the worker rank — no traceback, no OOM record; cross-stack control passes at 54k #36941 exactly. Note: with
  vm.min_free_kbytes=5GB + vfs_cache_pressure=200 host hardening, the box itself now survives
  (worker dies cleanly instead of hard-freezing the node), but the process-level exhaustion
  remains. Happy to provide logs there.

Happy to re-run any specific configs on this hardware — we have 2x GB10 available for testing.