[mmangkad]

Motivation

GLM-5.3-Flash MTP 5/1/6 decode spends about 5% of each step with the GPU outside CUDA graphs. The largest gap is between the draft-extend and draft graphs, where DeepseekSparseAttnMultiStepBackend fills each of the 4 draft-step backends separately, though every step ends up holding identical metadata.

Modifications

• Draft steps 1.. now alias step 0's captured metadata. Replay refreshes step 0 once from the precomputed input, and steps 1.. no longer allocate their own wide page_table and seqlen buffers.
• DeepseekSparseAttnBackend gets two protected helpers: _init_cuda_graph_state_shared sizes only the multi-CTA KV counter, which _forward_trtllm reads per backend, and _share_cuda_graph_metadata aliases step 0's metadata.
• Sharing does not depend on KPool fusion, so it replaces both the KPool sibling copies and the multi-destination copy block.
• test_dsa_metadata_replay.py captures and replays the steps through the wrapper for index_kpool 4 (KPool fusion on and off) and 1, and checks that steps share step 0's buffers, only step 0 refreshes, and the result matches the ordinary path.

Sharing is safe because every step consumes the same precomputed input and the draft forward only reads the metadata. All writes to DSAMetadata / KPoolWritePlan are in the init and replay paths; the one forward-time writer (dsa_npu_indexer.py) is prefill-only.

Only GLM-5.3-Flash sets index_kpool (4); DeepSeek-V3.2 and GLM-5 through 5.3 use 1. With more than 3 draft steps, those models used the multi-destination copy block, which never refreshed paged_mqa_schedule_metadata / topk_v2_plan for the first three steps; they now get the full refresh.

Some cleanups this change enables can be done as a follow-up.

Accuracy Tests

sgl-eval 0.1.2 GSM8K, 200 examples, temperature 1.0, top_p 0.95, thinking on.

Model	GSM8K before	GSM8K after	Accept length (conc. 1) before	Accept length after
GLM-5.3-Flash	99.0%	98.0%	4.49	4.49
GLM-5.3	98.5%	99.5%	4.19	4.19

test_dsa_metadata_replay.py: 3 passed (6 subtests) on GB300.

Speed Tests and Profiling

GLM-5.3-Flash, 4× GB300, --tp-size 4, trtllm DSA, fp8_e4m3 KV, flashinfer_trtllm MoE, EAGLE 5/1/6, greedy, 2,048 output tokens. Accept length pinned with SGLANG_SIMULATE_ACC_LEN=5 SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token (5.00 on every run).

	Before	After
fwd_occupancy, concurrency 1 / 8 / 32	93.19 / 94.62 / 93.13%	94.46 / 95.63 / 94.64%
Per-request decode, concurrency 1 / 8 / 32	630 / 372 / 235 tok/s	639 / 376 / 239 tok/s
Draft decode graph memory per GPU	1.02 GB	0.36 GB

nsys, batch size 1, medians over ~165–185 steps:

	Before	After
Gap: draft-extend → draft	286 µs	175 µs
Gap: draft → verify	684 µs	363 µs
GPU idle per step	884 µs	619 µs
fused_metadata_copy_kernel per step	4	0
Eager memcpys per step	51	24

The draft → verify gap also shrinks because the verify graph launch waits on host work: with less of it during draft replay, the launch lands sooner.

The draft-extend → draft gap of the median step (by that gap), same 400 µs scale, starting 40 µs before the draft-extend graph ends:

Before: the copy ladder, including 4 fused_metadata_copy_kernel, runs until the draft graph starts at ~326 µs.

After: only step 0's refresh remains; the draft graph starts at ~215 µs.

---

CI States

Latest PR Test (Base): ❌ Run #38664087872
Latest PR Test (Extra): ❌ Run #38664087888
Latest PR Test (AMD ROCm 10): ❌ Run #38664087886