Skip to content

[GLM-5.3-Flash] Restore opt-in in-graph verify metadata - #38852

Closed
Fridge003 wants to merge 1 commit into
mainfrom
codex/glm53-ingraph-verify-metadata
Closed

Fridge003 wants to merge 1 commit into
mainfrom
codex/glm53-ingraph-verify-metadata

Conversation

@Fridge003

@Fridge003 Fridge003 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Restore SGLANG_EXPERIMENTAL_DSA_INGRAPH_VERIFY_METADATA=1, removed by #38071. With KPool fusion enabled, target verification can refresh its metadata inside the model CUDA graph instead of launching that sequence eagerly before each replay.

Builds on PR #38845, now merged into main. KPool metadata fusion is enabled by default for supported CUDA geometry; this PR adds the separate, default-off in-graph verify flag.

The follow-up DeepGEMM option is PR #38858. This PR is rebased onto main at a66451c058dcce0799f5c6d1620b527371db4828 and targets main; its diff contains only the in-graph optimization.

Modifications

  • Keep eligibility, capture, replay and recorded state in the dedicated dsa_metadata_ingraph.py mixin; the backend contains only initialization and lifecycle hooks.
  • Capture the fused verify metadata, DeepGEMM schedules, top-k plan and KPool write-plan refreshes. Preserve destination addresses and the captured static input identities.
  • Reject replay with replaced input buffers; clear recorded state on recapture. Unsupported configurations retain the ordinary path.
  • Restore an optional preallocated output for plan_topk_v2, with shape/dtype/device/layout validation.
  • Add GPU coverage for updated sequence lengths and request mappings, non-KPool operation, input identity checks and recapture.

Accuracy Tests

Rebased head: 3b8b48279969253e2d0fad43109328b94e766f6d. The rebase preserves main's consolidated DSAMetadataManagementMixin and default-on KPool fusion, while retaining the default-off in-graph flag. Changed-file pre-commit checks and git diff --check pass. The B200 correctness and benchmark results below are the preserved pre-rebase measurements; GPU tests were not rerun for this rebase.

Draft: the unchanged full-model accuracy gate did not pass. On baizhou-dev-b200 (4× B200), commit 65ec643b4707d13200edd7474445cf60fa2f1271 with both restoration flags enabled:

SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1 \
SGLANG_EXPERIMENTAL_DSA_INGRAPH_VERIFY_METADATA=1 PYTHONPATH=python \
  python test/registered/e2e/models/test_glm53_flash_b200.py -v
# Ran 4 tests in 687.185 s: FAILED (errors=1)

The test file, 500 GSM8K examples, 20 shots and 0.930 floor are unchanged. The test wrapper reports the failing accuracy assertion as an error.

Recipe This PR GSM8K Result
DFlash2 0.936 Pass
HighThroughput 0.944 Pass
LowLatency 0.926 Fail: below 0.930

The LowLatency speed/acceptance test passes: 353.51 tok/s, 4.605 accept length (gates: >250 tok/s and >4.0). The in-graph verify path is active in the real model logs.

Focused in-graph GPU regressions pass 4/4:

PYTHONPATH=python python test/registered/kernel/attention/test_dsa_ingraph_metadata.py

Unmodified main also failed the LowLatency accuracy floor: 0.928 initially and 0.920 in one controlled repeat. The parent KPool PR passed its full run (LowLatency 0.936). Baseline variability does not establish that this change is harmless. All original results and thresholds are retained. One controlled repeat of the unchanged LowLatency class passed 2/2 in 290.203 s: GSM8K 0.936, speed 350.79 tok/s, accept length 4.422. Run the same command above with TestGLM53FlashB200LowLatency appended to select that class. The initial 0.926 failure remains part of the evidence. This PR stays draft because the intermittent full-model accuracy failure remains unexplained; a passing repeat does not resolve its cause.

An additional real-GPU metadata diagnostic passed 16 configurations × 17 rounds = 272 graph replays: bs=1/4/16/48, draft lengths 2/4/6/8, changing sequence lengths across pool/page/top-k boundaries through 100001 tokens, and remapped requests. It interleaves live graphs and compares ordinary versus optimized metadata and stable destination addresses. This isolates metadata construction; it is not a substitute for the full-model accuracy gate. The script, exact heads/flags and logs are in the evidence link below.

Speed Tests and Profiling

Same four B200 GPUs on baizhou-dev-b200, TP4/EP4, CUDA 13.0, driver 580.173.02, PyTorch 2.13.0+cu130, Triton 3.7.1, DeepGEMM 0.1.7 and FlashInfer 0.6.18. The fixed main baseline is 9a2f17f41d185614a17f740c006c4aca84c016bf.

Model revisions: zai-org/GLM-5.3-Flash@eb9eb208eb0d988989d07a6a12d0fdeb5f52574a and incoai/GLM-5.3-Flash-DFlash2@bf582e4eacc1810f76656d1811693ff6c6737d2a. Every weight shard was SHA256-verified.

Serving benchmarks use the test file's LowLatency EAGLE recipe, with radix caching disabled consistently for measurement. bs means --max-concurrency; each run sends 4 * bs requests under that concurrency limit. Inputs are exactly 8192 token IDs, outputs exactly 1024 tokens, seed 42, temperature 0, one warmup request and a cache flush. EOS is ignored by the current CLI default. Each cell has three runs; throughput is the median of the three run-level rates, and latency is the median of the three run means; throughput ranges are retained. All measured requests completed with the expected token counts.

bs Main tok/s Parent PR tok/s This PR tok/s (min–max) Speedup vs parent Speedup vs main
1 511.93 537.30 523.69 (516.19–527.25) 0.975× 1.023×
4 1319.90 1366.22 1354.75 (1351.42–1369.60) 0.992× 1.026×
16 2100.45 2178.85 2159.02 (2152.35–2164.28) 0.991× 1.028×

This workload shows no incremental throughput gain over the parent PR: the measured median throughput is lower by 2.5%, 0.8% and 0.9% at bs=1/4/16. The cumulative results remain above the fixed main baseline.

Latency medians for this PR:

bs Mean TTFT (ms) Mean TPOT (ms) Mean E2E (ms)
1 262.041 1.633 1944.816
4 481.251 2.403 2958.062
16 2024.215 5.320 7466.935

Reproduce from this PR's checkout (other restored flags unset):

export PYTHONPATH=python
export SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1
export SGLANG_EXPERIMENTAL_DSA_INGRAPH_VERIFY_METADATA=1

python -m sglang.launch_server --model-path zai-org/GLM-5.3-Flash --tp-size 4 --ep-size 4 --dsa-prefill-backend trtllm --dsa-decode-backend trtllm --kv-cache-dtype fp8_e4m3 --moe-runner-backend deep_gemm --reasoning-parser glm45 --tool-call-parser glm47 --speculative-algorithm EAGLE --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6 --speculative-adaptive --disable-radix-cache --host 127.0.0.1 --port 31080

In another shell after the server is ready:

export PYTHONPATH=python
for rep in 1 2 3; do
  for bs in 1 4 16; do
    python -m sglang.bench_serving --backend sglang \
      --base-url http://127.0.0.1:31080 --model zai-org/GLM-5.3-Flash \
      --dataset-name random --num-prompts "$((4 * bs))" --max-concurrency "$bs" \
      --random-input-len 8192 --random-output-len 1024 --random-range-ratio 1 \
      --seed 42 --tokenize-prompt --warmup-requests 1 --flush-cache \
      --output-file "bs${bs}-rep${rep}.jsonl"
  done
done

Raw run-level measurements, accuracy records and a standalone recomputation script are available together.

Checklist

  • Changed-file pre-commit checks pass.
  • GPU regression tests are registered in base-b-kernel-unit.
  • Feature behavior and activation are documented above.
  • Full-model correctness and serving benchmark results are recorded.

CI States

Latest PR Test (Base): ❌ Run #34724482723
Latest PR Test (Extra): ❌ Run #34724482748
Latest PR Test (AMD ROCm 10): ❌ Run #34724482716

@Fridge003
Fridge003 force-pushed the codex/glm53-ingraph-verify-metadata branch from 65ec643 to 879fdc8 Compare September 10, 2026 19:59
@Fridge003
Fridge003 force-pushed the codex/glm53-kpool-metadata-fusion branch from b327ef9 to e586d07 Compare September 10, 2026 19:59
@Fridge003 Fridge003 changed the title [GLM-5.3] Restore opt-in in-graph verify metadata [GLM-5.3-Flash] Restore opt-in in-graph verify metadata Sep 11, 2026
Base automatically changed from codex/glm53-kpool-metadata-fusion to main September 12, 2026 23:01
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: zRzRzRzRzRzRzR <Yuxuan.Zhang2@liverpool.ac.uk>
Co-authored-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
Co-authored-by: zanes-ops <zanes@nvidia.com>
@Fridge003
Fridge003 force-pushed the codex/glm53-ingraph-verify-metadata branch from 879fdc8 to 3b8b482 Compare September 12, 2026 23:05
@Fridge003 Fridge003 closed this Sep 13, 2026
@Fridge003

Copy link
Copy Markdown
Collaborator Author

This optimization is not valid

@Jiminator
Jiminator deleted the codex/glm53-ingraph-verify-metadata branch September 14, 2026 04:41
@alexnails
alexnails restored the codex/glm53-ingraph-verify-metadata branch September 14, 2026 05:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant