[JackChuang]

Summary

Enables: correct output for GLM-5.3-Flash with an FP8 KV cache and the TileLang DSA backends on current main. Since #42689, about half of its answers under concurrent requests turn into gibberish (a run of ! up to the token limit), and GSM8K drops to 50.0%; with this PR it scores 98.0%.

Since #36003 the MLA KV writers skip the slot reserved for padding rows (slot 0). The zero-RoPE writer set_mla_kv_buffer_kernel_norope, which every GLM-5.3-Flash KV write goes through (qk_rope_head_dim = 0), still writes it. This PR makes it skip the slot the same way.

• Accuracy: GLM-5.3-Flash, GSM8K 200 questions: main 50.0% → this PR 98.0%, and 78.0% → 98.0% with MTP 5/1/6; answers cut off at the token limit: 51.0% → 0.5%.
• Cost: no measurable slowdown. End to end, TTFT and TPOT with 1 request and prefill throughput at saturation are within 1.5% of main; the writer kernel is at most 0.07 µs slower than on main from 1 to 8,192 rows. No new flags or environment variables.
• Speed: the attention kernels are left as they are, so decode speed is unchanged; masking the FP8 decode kernel's loads instead (as #45988 does) makes TTFT up to 1.7× as long and decode up to 8% slower in the setup below.

Motivation

Padding rows (CUDA graph capture and padded replays) write to out_cache_loc = 0 with undefined inputs. Since #42689, per-token FP8 quantization keeps BF16 NaNs instead of clamping them, so the zero-RoPE writer can now store NaN rows in slot 0. The FP8 TileLang sparse decode maps unused top-k entries to row 0, and masking the scores does not stop a NaN (-inf + NaN, 0 * NaN), so prompts shorter than index_topk decode NaN logits.

#45988 hardens that reader by masking every KV load in the FP8 decode kernel, which makes that kernel slower: at GLM-5.3-Flash TP4 shapes (16 heads per GPU, top-k 2048) it takes 1.3× as long for one decode token and 2.7–3.3× as long for 16 to 256 (see Speed Tests). This PR stops the write instead, so slot 0 stays clean for every kernel that reads it and the decode kernel is unchanged. The two fixes are independent: with both applied, GSM8K scores 98.5% in the setup below.

The NaN rows come from padded outputs that are left uninitialized, such as those of the KDA/GDN recurrent kernel, which #46225 zeroes. This PR keeps slot 0 clean whatever the padding rows hold, as the RoPE writers already do, so it does not rely on every kernel zeroing its padded outputs; the two can land in either order.

The same skip is part of the DCP rework in #39117 / #40433, which also adds DCP ownership to this kernel. This PR lands only the non-DCP part, so main stops producing NaN now; #40433 can rebase by adding its DCP arguments on top. cc @kpham-sgl @Fridge003 @BBuf

Modifications

• python/sglang/kernels/ops/kvcache/mla_buffer.py: set_mla_kv_buffer_kernel_norope takes reserved_skip_index and masks its store exactly like set_mla_kv_buffer_kernel; _set_mla_kv_buffer_impl passes it. The torch-native fallback set_mla_kv_buffer_naive now skips the slot for zero RoPE too, so the two paths stay bit-identical.
• test/registered/kernels/ops/kvcache/test_set_mla_kv_buffer.py: the two existing Triton reserved-slot tests also run with rope_dim=0. The skip case fails on main and passes with this PR.

What changes: only writes aimed at the reserved slot are dropped, as the RoPE writers already do; real token rows are written as before. It reaches every MLA model with qk_rope_head_dim = 0 (GLM-5.3-Flash, BF16 or FP8 KV); the KV allocators never hand out slot 0 and the pool zero-initializes it, so the slot now stays zero. DCP handling in this kernel is unchanged (#36886, #40433).

Accuracy Tests

GLM-5.3-Flash on 4x H200 with the server arguments of the H200 CI tests (test_glm53_flash_h200_*.py), at TP4 and with an FP8 KV cache. Prefill CUDA graphs are off; decode CUDA graphs stay on. The eval uses the sampling settings of those CI tests. The MTP 5/1/6 rows add the speculative flags of test_glm53_flash_h200_low_latency.py (--speculative-algorithm EAGLE --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6).

python -m sglang.launch_server --model-path zai-org/GLM-5.3-Flash --tp-size 4 \
  --dsa-prefill-backend tilelang --dsa-decode-backend tilelang --kv-cache-dtype fp8_e4m3 \
  --moe-runner-backend deep_gemm --reasoning-parser auto --tool-call-parser auto \
  --cuda-graph-backend-prefill disabled
sgl-eval run gsm8k --base-url http://localhost:30000/v1 --num-examples 200 --num-threads 256 \
  --thinking --max-tokens 8192 --temperature 1.0 --top-p 0.95

	main (f9257d6c7c)	this PR
GSM8K accuracy, 200 questions	50.0%	98.0%
Answers cut off at the 8192-token limit	51.0% (each ends in a run of !)	0.5%
GSM8K accuracy with MTP 5/1/6	78.0%	98.0%
Answers cut off with MTP 5/1/6	21.5%	0.5%

Speed Tests and Profiling

Kernel time of the zero-RoPE writer (set_mla_kv_buffer_triton with an empty RoPE part) on H200, timed with CUDA graph replay:

Rows written per call	main	this PR
1,024	2.00 µs	1.99 µs
8,192	10.74 µs	10.74 µs
8,192, a quarter of them padding rows aimed at slot 0	24.20 µs	10.72 µs

The rows are BF16 with 512 values; FP8 rows of 512 to 656 bytes behave the same. Without padding rows, this PR is at most 0.07 µs slower than main at any size from 1 to 8,192 rows. With padding rows it is faster, because main writes all of them to the same slot.

End to end, with a BF16 KV cache (on main the FP8 output is not valid) and otherwise the setup of the accuracy tests. Rows with 1 request or at saturation use bench_serving --dataset-name random --random-range-ratio 1 (fixed prompt and output lengths); the rows at 16 and 64 concurrent requests use prompts of about 8k tokens. Medians:

	main	this PR
TTFT, 1 request, 1k / 4k / 16k-token prompt	150 / 187 / 609 ms	148 / 184 / 611 ms
TPOT, 1 request, 1k-token prompt, 512 output tokens	7.38 ms	7.38 ms
TPOT, 16 / 64 concurrent requests, ~8k-token prompts, 512 output tokens	22.60 / 47.46 ms	22.60 / 47.44 ms
Prefill throughput at saturation, 8k-token prompts, 32 to 128 concurrent requests	28.7k tok/s	28.7k tok/s

The reader-side alternative (#45988) slows the FP8 decode kernel instead. tilelang_sparse_fwd with an FP8 KV cache at GLM-5.3-Flash TP4 shapes (16 heads, d_v 512, top-k 2048, every index valid), H200:

Decode tokens	main	with #45988
1	19.9 µs	26.5 µs
16	24.7 µs	66.1 µs
64	84.2 µs	254.5 µs
256	305.5 µs	1016.2 µs

End to end with an FP8 KV cache and the same workloads, medians. Prefill uses the same kernel, so TTFT grows too:

	this PR	with #45988
TTFT, 1 request, 1k / 4k / 16k-token prompt	156 / 193 / 656 ms	157 / 307 / 1,114 ms
TTFT, 8 / 32 concurrent requests, ~8k-token prompts	1.74 / 4.44 s	2.98 / 7.68 s
Prefill throughput at saturation, 8k-token prompts, 32 to 128 concurrent requests	26.6k tok/s	15.3k tok/s
TPOT, 1 request, 1k-token prompt, 512 output tokens	6.60 ms	6.74 ms
TPOT, 16 concurrent requests, ~8k-token prompts, 512 output tokens	22.1 ms	23.9 ms
TPOT / output throughput, 64 concurrent requests, 1k-token prompts, 512 output tokens	36.5 ms / 1,608 tok/s	39.4 ms / 1,425 tok/s

CI

test_set_mla_kv_buffer.py passes, including the new zero-RoPE cases, and pre-commit run passes on the changed files. On H200, in the per-PR NVIDIA suites (base-a and base-b on 1 and 2 GPUs, base-c on 4 and 8 GPUs) and the five CPU suites, every test that fails with this PR also fails the same way on the parent commit (535aeae2a7). The H200 GLM-5.3-Flash tests (BF16 KV cache, TileLang DSA), whose KV writes all go through the changed kernel, pass with and without this PR: GSM8K 97.3% with this PR vs 97.2% on main with EAGLE (test_glm53_flash_h200_low_latency.py), and 96.9% vs 97.4% without it (test_glm53_flash_h200_high_throughput.py). Not run here: the PD disaggregation, Mooncake and NIXL tests and hicache/test_pp_with_hicache.py; Blackwell GPUs; ROCm.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #40650688282
Latest PR Test (Extra): ❌ Run #40650688166
Latest PR Test (AMD ROCm 10): ❌ Run #40650688212

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-test test/registered/e2e/models/test_glm53_flash_h200_low_latency.py test/registered/e2e/models/test_glm53_flash_h200_high_throughput.py

Targeted validation for 159443c4443c.

[github-actions]

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_h200_low_latency.py
cd test/ && python3 registered/e2e/models/test_glm53_flash_h200_high_throughput.py

[BBuf]

/rerun-failed-ci

[BBuf]

LGTM. Zero-RoPE writes now preserve the reserved slot; wait for related CUDA/AMD kernel CI and H200 model checks.

[BBuf]

/rerun-test test/registered/e2e/models/test_glm53_flash_h200_low_latency.py test/registered/e2e/models/test_glm53_flash_h200_high_throughput.py

Targeted validation for 159443c4443c.

[BBuf]

/tag-and-rerun-ci

[github-actions]

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_h200_low_latency.py
cd test/ && python3 registered/e2e/models/test_glm53_flash_h200_high_throughput.py

[BBuf]

/tag-and-rerun-ci

[tanth47]

Cross-link: this covers the remaining writer from #36207 (MLA KV writers overwriting the reserved padding slot 0 → NaN-poisoned outputs). #36003 fixed the other MLA writers; that issue was auto-closed as inactive before the zero-RoPE path was addressed.