Skip to content

[DCP] Restore GLM-5.3 Flash decode CP and fix virtual KV addressing - #39117

Open
kpham-sgl wants to merge 12 commits into
sgl-project:mainfrom
kpham-sgl:codex/restore-glm53-flash-dcp
Open

kpham-sgl wants to merge 12 commits into
sgl-project:mainfrom
kpham-sgl:codex/restore-glm53-flash-dcp

Conversation

@kpham-sgl

@kpham-sgl kpham-sgl commented Sep 11, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Restore GLM-5.3-Flash decode context parallelism after its temporary removal in #37487, with the virtual-addressing defects in #36886 fixed. Previously, NoPE attention KV writes and sparse reads used DCP virtual locations against per-rank buffers; crossing the physical allocation watermark could corrupt adjacent buffers and degrade accuracy or crash. The replicated index-K cache also needs to cover the full virtual address space.

Fixes #36886.

Modifications

  • Shard NoPE KV writes by owner rank and localize sparse decode/target-verify selections before reading attention KV. Produce base-2 LSE for the DCP merge, including zero output and -inf LSE for ranks owning no selected tokens.
  • Gather cached prefixes into the per-request ragged layout for sparse extend; restore GLM's DCP metadata hook and handle the absent RoPE component.
  • Size replicated index-K and draft pools for virtual locations, preserve the physical DSA page size, and account for replicated storage in memory budgeting. Validate indexer locations against index-K capacity.
  • Keep EAGLE draft KV replicated and bypass the target's Q/LSE collectives and prefix planner for draft execution.
  • Add a GPU kernel regression comparing merged DCP partials against PyTorch, and an 8×H200 GLM integration test covering ordinary decode, EAGLE, allocation beyond per-rank capacity, and prefix reuse.

The restored DSA DCP path requires CUDA NoPE MLA, BF16 KV, TileLang prefill/decode, and fused top-k. Unsupported combinations, including HiSparse and prefill CP, are rejected by the existing DCP argument validator after model-specific defaults resolve. TileLang’s existing validation enforces the BF16 KV requirement.

Accuracy Tests

4×B300, same model/source and equal global KV capacity across all four configurations. Model: zai-org/GLM-5.3-Flash@03eb5366286afd40d2221b1d9c63a6dd1ba4832e, original FP8 weights. Benchmarked at snapshot d4d53e97af74f905848b176a19b181f29fc90a8d over main 94ce940ff8. Subsequent cleanup, the merge from main, and argument-validation changes are covered by the checks below; the four-configuration benchmark was not repeated on the latest commit.

Configuration AIME2026 correct / 480 AIME2026 pass@1 ± SEM AIME truncations GSM8K correct / 1314 GSM8K
TP4/EP4 baseline 442 92.08% ± 1.21 25 1258 95.74%
TP4/EP4 + EAGLE 437 91.04% ± 0.85 25 1259 95.81%
DCP4 447 93.13% ± 0.83 23 1258 95.74%
DCP4 + EAGLE 441 91.88% ± 1.01 26 1256 95.59%

All eight evaluations completed with zero request errors, zero evaluator retries, and no cache retractions. All GSM8K runs had zero truncations. AIME truncations remained wrong in the full denominator.

DCP4+EAGLE is −0.21 points versus baseline, +0.83 versus EAGLE, and −1.25 versus DCP4. Each difference is smaller than its combined repeat SEM; these runs do not establish statistical equivalence or token identity. The earlier single-pass AIME2025 check had only 30 responses, so one changed answer moved the score by 3.33 points.

  • AIME2026: all 30 questions × 16 sampled repeats, temperature 1.0, top-p 0.95, maximum reasoning effort, 65536-token output cap, 64 threads. sgl-eval==0.1.0, vendored NeMo-Skills 645cf567ff08c0ae9cc3fc8e1edbb975b3067816, default AIME26 prompt/grader. SEM is across repeat accuracies on the fixed question bank.
  • GSM8K: OpenAI test split, first five rows used as fixed demonstrations and excluded from scoring; 1314 questions, five-shot chat, temperature 0, top-p 1, output cap 4096, 64 threads, no retries. Use get_few_shot_examples, get_one_example, and get_answer_value from sglang.test.simple_eval_mixed_prefix_gsm8k. Explicit reasoning_effort=low because GLM's template ignores thinking=false.
  • Regression evidence: DCP4 reached 1,684,736 live KV tokens and DCP4+EAGLE reached 1,587,456, both above the 1,048,576-token per-rank pool, without illegal access. A separate smaller-pool run passed repeated long-prompt/prefix-reuse checks in both modes.
  • Tests: All five selective rerun groups passed at a530fe448c (CPU, H100, 8×H200, 4×B200, and 8×B200), including all four cases in the new 8×H200 GLM Flash DCP integration test. Rerun results. At 6c8b78f581, the existing server-argument and resolution suites passed (244 tests, 68 subtests), 17 configuration checks verified supported settings and rejection paths, and all pre-commit checks passed.
Reproduction settings

Common server flags:

python3 -m sglang.launch_server \
  --model-path /path/to/pinned/GLM-5.3-Flash \
  --served-model-name glm53-flash --tp-size 4 --ep-size 4 \
  --kv-cache-dtype bfloat16 \
  --dsa-prefill-backend tilelang --dsa-decode-backend tilelang \
  --moe-runner-backend deep_gemm --mem-fraction-static 0.75 \
  --max-total-tokens 4194304 --max-running-requests 64 \
  --max-mamba-cache-size 512 --chunked-prefill-size 4096 \
  --context-length 131072 --cuda-graph-max-bs-decode 64 \
  --reasoning-parser glm45 --tool-call-parser glm47 --port 30000

For DCP, add --dcp-size 4 --dcp-comm-backend ag_rs and set --max-total-tokens 1048576, preserving global capacity at 4,194,304. For EAGLE, add --speculative-algorithm EAGLE --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6. Radix and decode graphs are enabled. Flush the cache between benchmarks.

Environment: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, SGLANG_ENABLE_ASYNC_ASSERT=true, SGLANG_JIT_DEEPGEMM_FAST_WARMUP=true, NCCL_NVLS_ENABLE=0.

sgl-eval run aime26 --base-url http://127.0.0.1:30000/v1 \
  --model glm53-flash --num-examples 30 --num-threads 64 --n-repeats 16 \
  --thinking --reasoning-effort max --temperature 1.0 --top-p 0.95 \
  --max-tokens 65536 --out-dir /path/to/results

Speed Tests and Profiling

Observed AIME2026 evaluation timing, excluding server startup:

Configuration Completion tokens Client seconds Output tokens/s
baseline 5,930,849 2260.84 2623
EAGLE 6,545,197 1870.63 3499
DCP4 5,962,534 2651.32 2249
DCP4 + EAGLE 6,292,787 1964.58 3203

These are correctness-run measurements with different sampled output lengths, not a dedicated throughput comparison. No separate profiling run was performed.

Checklist

  • Format code with repository pre-commit hooks.
  • Add registered kernel and model integration tests.
  • Document supported settings and reproduction commands above.
  • Provide accuracy results and observed evaluation timing, with limitations.
  • Follow repository code style.

CI States

Latest PR Test (Base): ❌ Run #34915797264
Latest PR Test (Extra): ❌ Run #34915796902
Latest PR Test (AMD ROCm 10): ❌ Run #34915797059

@kpham-sgl

Copy link
Copy Markdown
Collaborator Author

/rerun-test test_dcp_layout_unit.py test_dsv31_dcp8_gsm8k.py test_kimi_linear_dcp4.py test_kimi_linear_dcp_dspark4.py test_qwen3p5_triton_dcp.py test_reduce_scatter_along_dim.py test_trtllm_mla_family_dcp_metadata.py test_dsa_dcp_kernels.py test_dcp_lse_combine.py test_pd_dcp_gather.py test_glm53_flash_h200.py test_glm53_flash_b200.py test_glm53_flash_dcp.py test_pool_configurator.py test_kimi_linear_unified_memory_dcp_blackwell.py test_unified_radix_cache_kl_dcp.py test_kimi_linear_pd_dcp4.py

@github-actions

github-actions Bot commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-test test_dcp_layout_unit.py test_dsv31_dcp8_gsm8k.py test_kimi_linear_dcp4.py test_kimi_linear_dcp_dspark4.py test_qwen3p5_triton_dcp.py test_reduce_scatter_along_dim.py test_trtllm_mla_family_dcp_metadata.py test_dsa_dcp_kernels.py test_dcp_lse_combine.py test_pd_dcp_gather.py test_glm53_flash_h200.py test_glm53_flash_b200.py test_glm53_flash_dcp.py test_pool_configurator.py test_kimi_linear_unified_memory_dcp_blackwell.py test_unified_radix_cache_kl_dcp.py test_kimi_linear_pd_dcp4.py:

🚀 ubuntu-latest (2 tests): ✅ View workflow run

cd test/ && python3 registered/dcp/test_dcp_layout_unit.py
cd test/ && python3 registered/unit/model_executor/test_pool_configurator.py

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

cd test/ && python3 registered/dcp/test_dsv31_dcp8_gsm8k.py
cd test/ && python3 registered/dcp/test_reduce_scatter_along_dim.py
cd test/ && python3 registered/e2e/models/test_glm53_flash_h200.py
cd test/ && python3 registered/e2e/dcp/test_glm53_flash_dcp.py

🚀 4-gpu-b200 (7 tests): ✅ View workflow run

cd test/ && python3 registered/dcp/test_kimi_linear_dcp4.py
cd test/ && python3 registered/dcp/test_kimi_linear_dcp_dspark4.py
cd test/ && python3 registered/dcp/test_qwen3p5_triton_dcp.py
cd test/ && python3 registered/dcp/test_trtllm_mla_family_dcp_metadata.py
cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py
cd test/ && python3 registered/e2e/models/test_kimi_linear_unified_memory_dcp_blackwell.py
cd test/ && python3 registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dcp.py

🚀 1-gpu-h100 (3 tests): ✅ View workflow run

cd test/ && python3 registered/kernel/dcp/test_dsa_dcp_kernels.py
cd test/ && python3 registered/kernels/test_dcp_lse_combine.py
cd test/ && python3 registered/kernels/ops/kvcache/test_pd_dcp_gather.py

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

cd test/ && python3 registered/disaggregation/test_kimi_linear_pd_dcp4.py

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

Projects

None yet

1 participant