Skip to content

[Fix] glm5.2: Align PD HiSparse prefill KV layout - #44256

Merged
ispobock merged 2 commits into
sgl-project:mainfrom
alphabetc1:fix/glm52-hisparse-pd-kv-layout
Oct 14, 2026
Merged

ispobock merged 2 commits into
sgl-project:mainfrom
alphabetc1:fix/glm52-hisparse-pd-kv-layout

Conversation

@alphabetc1

@alphabetc1 alphabetc1 commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

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 selects flashmla_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

  • Explicitly select --dsa-prefill-backend flashmla_kv for this test's prefill worker so its KV layout matches decode.
  • Document why the backend must be explicit.
  • Remove both the CI registry disabled marker and the class-level unittest.skip added 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 at 4560da58eebe5660dc2e8833178351085d6f9dac.

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.

Configuration Result
Original backend defaults, real 24-shot GSM8K prompt First 4,198-token request reproduces NaN detected! verify: target model logits
Original layout mismatch with CUDA Graph disabled Also fails with an early attention NaN
Both workers using the FlashMLA KV layout, CUDA Graph enabled, diagnostic instrumentation removed The same request succeeds; another 16 GSM8K requests at concurrency 8 all succeed, each generating 64 tokens; harness exits 0

The aligned experiment explicitly set both DSA backend flags to flashmla_kv; on the prefill worker, calculate_mla_kv_cache_dim determines 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

  • Follow the repository's existing formatting and code style; git diff --check passes.
  • Provide real-weight E2E reproduction and comparison results above.
  • Validate the re-enabled full registered suite on 8 x B300 in PR CI.

CI States

Latest PR Test (Base): ❌ Run #38559260628
Latest PR Test (Extra): ❌ Run #38559260366
Latest PR Test (AMD ROCm 10): ❌ Run #38559260466

@alphabetc1

Copy link
Copy Markdown
Collaborator Author

/rerun-test -c

@github-actions

github-actions Bot commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

🚀 8-gpu-b300 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_dsa_glm52_hisparse.py

@ispobock
ispobock merged commit 7b53a13 into sgl-project:main Oct 14, 2026
111 of 123 checks passed
@alphabetc1
alphabetc1 deleted the fix/glm52-hisparse-pd-kv-layout branch October 14, 2026 12:41
@alphabetc1

Copy link
Copy Markdown
Collaborator Author

/rerun-test -c

@github-actions

github-actions Bot commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

🚀 8-gpu-b300 (1 test): ❌ View workflow run

cd test/ && python3 registered/e2e/models/test_dsa_glm52_hisparse.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants