[GLM-5.3 Flash] Restore and enable KPool metadata fusion

[Fridge003]

Motivation

Restore the KPool metadata fusion removed by #38071 and enable it by default for supported CUDA KPool geometry. GLM-5.3-Flash currently builds pool-aware metadata through the ordinary path; this restores fused decode, target-verify and draft-extend construction, plus reuse of derived metadata between compatible MTP draft backends.

This is the first of three dependent restoration PRs. SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION now defaults to true; set it to 0 to restore ordinary metadata. Platform and geometry checks still gate activation.

Restoration stack: #38845 → #38852 → #38858. Each PR has a separate feature diff; dependent PRs must be rebased and retargeted to main after their parent merges.

Modifications

• Keep KPool kernels in kernels/ops/attention/dsa_kpool_metadata/{decode,verify,draft_extend,scan}.py; retain ordinary non-KPool kernels in their existing module.
• Consolidate feature selection, supported-geometry checks and buffer-copy/sibling reuse in dsa_metadata_manager.py, under DSAMetadataManagementMixin, as requested in review.
• Enable KPool-aware fused construction in both decode and verify MTP precompute. Reuse sibling metadata only when the source and destination layouts are compatible; preserve captured buffer addresses.
• Add real GPU regressions for pool/page/top-k boundaries, live tails, request remapping, draft acceptance, derived schedule refresh and CUDA graph replay.

Mechanical move — reproducible

The existing precomputed buffer-copy body is extracted and moved in two separate mechanical commits. Preparation and feature restoration are separate semantic commits within this PR. Both mechanical commits reproduce byte-for-byte; their scripts, proof engine and chain report are available together.

gh gist clone 2c48fb0e78509398c81e5d81e4d9dbd7 /tmp/glm53-copy-proof
# Run from a checkout containing this PR's commits and pre-commit dependencies.
python /tmp/glm53-copy-proof/886d7f36f.py
python /tmp/glm53-copy-proof/749c191a6.py

Review follow-up — metadata manager consolidation

Commit 988bebf725f3be5d2fb7e25bbac058e91dca22b8 merges the two metadata modules and classes without changing their five method bodies. Changed-file pre-commit checks pass. Both focused GPU test files pass on baizhou-dev-b200 (4/4 tests total). The consolidation script and byte-for-byte reproduction proof are available separately from the earlier extraction proof:

gh gist clone 8c11ec96acdb75613bb13d9c8e9054ef /tmp/glm53-manager-proof
python3 /tmp/glm53-manager-proof/988bebf725.py

The requested AIME26 evaluation is complete (details below). A separate semantic commit enables the flag by default; the mechanical consolidation proof above is unchanged.

AIME26 review validation (2026-09-12)

Evaluated commit 988bebf725f3be5d2fb7e25bbac058e91dca22b8 on 4× B200 with TP4/EP4 and only SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1. Each run completed 30 problems × 16 repeats (480 samples), using sgl-eval 0.1.1, thinking, temperature 1.0, top-p 0.95, max-tokens 200000 and the evaluator's default 64 threads. Main/draft model revisions are the pinned revisions recorded below.

Recipe	Mean pass@1 over 16 repeats	Correct	Token-cap endings	Request errors	Elapsed
EAGLE adaptive	95.42%	458/480	5/480	0/480	2529.93 s
DFLASH	94.58%	454/480	8/480	0/480	2295.05 s

Both runs have pass@16 and majority@16 of 100%. No answer was extracted for 7 EAGLE and 11 DFLASH samples; these remain included in the reported scores. The token-cap endings are retained, not discarded or rerun. All 16 output files per recipe were audited for 30 unique problem IDs and matching problem sets; independently recomputed accuracy matches metrics.json. Full generations, server/evaluator logs, manifests, metrics and SHA256 audit records are retained locally. These are fusion-on accuracy measurements, not a matched fusion-off comparison.

After this evaluation, a separate default-on commit changes only the environment default and its comment. Explicit =0 and platform/geometry fallbacks are preserved. Changed-file pre-commit and unset/false/true environment checks pass. The six requested CI reruns will validate the final PR head.

Reproduce each server recipe in a separate run:

# EAGLE
SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1 sglang serve --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 --host 0.0.0.0 --port 30000 --speculative-algorithm EAGLE --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6 --speculative-adaptive

# DFLASH
SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1 sglang serve --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 --host 0.0.0.0 --port 30000 --speculative-algorithm DFLASH --speculative-draft-model-path incoai/GLM-5.3-Flash-DFlash2 --speculative-draft-attention-backend fa4

After the server becomes healthy:

sgl-eval run aime26 --model zai-org/GLM-5.3-Flash --api-key EMPTY \
  --n-repeats 16 --max-tokens 200000 --temperature 1.0 --top-p 0.95 \
  --thinking --out-dir /sgl-workspace/logs/aime26-RECIPE \
  --base-url http://localhost:30000/v1

Use a distinct output directory for each recipe.

Accuracy Tests

Earlier GSM8K validation on baizhou-dev-b200 (4× B200), commit b327ef982f321c652f88f127de39b2ad968c4e12 with only the KPool fusion flag enabled:

SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=1 PYTHONPATH=python \
  python test/registered/e2e/models/test_glm53_flash_b200.py -v
# Ran 4 tests: OK (727.023 s)

The test file is unchanged from the fixed main baseline. Each recipe uses 500 GSM8K examples and 20 shots; the original 0.930 floor is unchanged.

The table also includes AIME26 mean pass@1 over 16 repeats, measured separately at commit 988bebf725f3be5d2fb7e25bbac058e91dca22b8 with TP4/EP4 and fusion enabled. Token-cap endings remain included in the score; full AIME26 settings and audit details are recorded above.

Benchmark	Recipe	Accuracy	Samples	Result
GSM8K	DFlash2	93.80%	500	Pass
GSM8K	HighThroughput	95.00%	500	Pass
GSM8K	LowLatency	93.60%	500	Pass
AIME26	EAGLE adaptive	95.42% (458/480)	30 × 16	Completed; 5 token-cap endings, 0 request errors
AIME26	DFLASH	94.58% (454/480)	30 × 16	Completed; 8 token-cap endings, 0 request errors

The LowLatency speed/acceptance test also passes: 320.06 tok/s, 4.422 accept length (original gates: >250 tok/s and >4.0).

Focused GPU regressions pass 4/4:

PYTHONPATH=python python test/registered/kernel/attention/test_dsa_kpool_metadata_fusion.py
PYTHONPATH=python python test/registered/kernel/attention/test_dsa_metadata_replay.py

Baseline evidence is retained: unmodified main scored DFlash2 0.932, HighThroughput 0.930, LowLatency 0.928, failing only the LowLatency GSM8K floor. One controlled repeat of the unchanged main LowLatency class scored 0.920 and again failed that floor; its speed/acceptance gates passed on both runs. The same 500 questions produced 28 correctness changes between the two main LowLatency runs. Both baseline failures are retained, with the original threshold unchanged. This PR's full run passes all four checks.

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	This PR tok/s (min–max)	Speedup vs main
1	511.93	537.30 (536.17–538.12)	1.050×
4	1319.90	1366.22 (1317.93–1374.38)	1.035×
16	2100.45	2178.85 (2111.45–2190.20)	1.037×

Latency medians for this PR:

bs	Mean TTFT (ms)	Mean TPOT (ms)	Mean E2E (ms)
1	262.281	1.600	1894.105
4	427.608	2.439	2939.365
16	2018.392	5.377	7466.242

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

export PYTHONPATH=python
export SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION=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 #34714472356
Latest PR Test (Extra): ❌ Run #34714472142
Latest PR Test (AMD ROCm 10): ❌ Run #34714472186

[Fridge003]

/rerun-test test/registered/kernel/attention/test_dsa_kpool_metadata_fusion.py test/registered/kernel/attention/test_dsa_metadata_replay.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_glm53_flash_h200.py test/registered/e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py test/registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py

[github-actions]

Results for /rerun-test test/registered/kernel/attention/test_dsa_kpool_metadata_fusion.py test/registered/kernel/attention/test_dsa_metadata_replay.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_glm53_flash_h200.py test/registered/e2e/models/test_dsa_glm52_pd_mtp_cp_layersplit.py test/registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py:

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

cd test/ && python3 registered/kernel/attention/test_dsa_kpool_metadata_fusion.py
cd test/ && python3 registered/kernel/attention/test_dsa_metadata_replay.py

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

cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py
cd test/ && python3 registered/e2e/models/test_dsa_glm52_nvfp4_tp_mtp.py

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

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

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

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