Repository navigation
[Fix] glm5.2: Align PD HiSparse prefill KV layout - #44256
Merged
ispobock merged 2 commits intoOct 14, 2026
Merged
Conversation
Collaborator
Author
|
/rerun-test -c |
Contributor
|
🚀 |
5 tasks
huangtingwei9988
approved these changes
Oct 14, 2026
ispobock
approved these changes
Oct 14, 2026
Collaborator
Author
|
/rerun-test -c |
Contributor
|
🚀 |
3 of 4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
The GLM-5.2 PD + HiSparse + MTP test introduced in #37771 fails on its first GSM8K request with
NaN detected! verify: target model logits:https://github.com/sgl-project/sglang/actions/runs/38354570978/job/115154648003
The prefill worker leaves its DSA backend unspecified. On Blackwell with FP8 KV, this defaults to
trtllm, whose raw MLA KV layout is 576 bytes/token (512 FP8 NoPE + 64 FP8 RoPE). The decode worker explicitly selectsflashmla_kv, which consumes 656 bytes/token (512 FP8 NoPE + 16 bytes of FP32 scales + 128 bytes of BF16 RoPE).The original CI logs confirm this backend mismatch. The ordinary Mooncake MLA transfer copies bytes without converting the formats and computes destination page offsets using the sender's item length. Decode therefore interprets incompatible KV bytes as FlashMLA's packed layout. Captured inputs in the reproduction contained negative quantization scales, and layer-level assertions caught NaNs in early attention computation.
Modifications
--dsa-prefill-backend flashmla_kvfor this test's prefill worker so its KV layout matches decode.disabledmarker and the class-levelunittest.skipadded in [CI] glm5.2: Temporarily disable PD HiSparse+MTP B300 test #44192, restoring the registered 8-GPU B300 suite and selective reruns with the corrected KV layout.Accuracy Tests
Real-weight end-to-end reproduction on 4 x GB300, using
nvidia/GLM-5.2-NVFP4, prefill TP2, decode TP2 + attention DP2, Mooncake TCP, FlashInfer 0.7.1, and Mooncake 0.3.13.post1. The experiment used main at4560da58eebe5660dc2e8833178351085d6f9dac.Retained FP8 KV, EAGLE MTP (3 steps, top-k 1, 4 draft tokens), rejection sampling, HiSparse (top_k=2048, device_buffer_size=4096, host_to_device_ratio=5), and CUDA Graph.
NaN detected! verify: target model logitsThe aligned experiment explicitly set both DSA backend flags to
flashmla_kv; on the prefill worker,calculate_mla_kv_cache_dimdetermines the pool layout from the active prefill backend. This patch sets that active backend; the test already explicitly configures the decode worker's active backend.These are reproduction and NaN-regression checks, not a full GSM8K accuracy or speculative acceptance-length evaluation. The complete 500-question suite, long-context cases, and original 8 x B300 / RDMA topology have not been rerun. Eight-GPU capacity was unavailable during the investigation.
Speed Tests and Profiling
No performance benchmark was run. This change updates test configuration.
Checklist
git diff --checkpasses.CI States
Latest PR Test (Base): ❌ Run #38559260628
Latest PR Test (Extra): ❌ Run #38559260366
Latest PR Test (AMD ROCm 10): ❌ Run #38559260466