[chromecast56]

Motivation

With DFlash/DSpark, every captured decode CUDA graph keeps its own packed aux hidden-state output alive, so reserved graph memory grows with the sum of captured sizes. One buffer sized for the largest graph is enough: the worker writes target hidden states into draft KV before the next target forward on that stream.

Modifications

• AuxHiddenStatePacker (Improve CUDA graph and speculative execution output handling #37329) takes an optional out= buffer, checks its shape and dtype on the first append, and writes captures into it.
• DecodeCudaGraphRunner keeps one [max_num_token, K * hidden] buffer per stream (so PD-multiplexed streams don't alias) and gives each captured size a row slice through ForwardBatch.aux_hidden_states_buffer. The buffer is used only for DFlash-family target runners with pp_size == 1, and only for models implementing SupportsSharedAuxHiddenStates.get_packed_aux_hidden_size().
• Kimi K3 now captures through AuxHiddenStatePacker (PP stages keep the list, since inherited captures arrive already concatenated) and is the first model to opt in. Other packer users (DeepSeek V2, MiMo V2) and EAGLE3 are unchanged.

Tests

• test_spec_aux_hidden_state.py: packing into an external buffer gives the same result as the list-path concat, lands in the caller's storage, and raises on a buffer that doesn't fit.
• test_decode_cuda_graph_runner.py: graph sizes 64/8/32 share one buffer; a second stream gets its own; non-DFlash, draft, PP and non-opt-in models get no shared buffer.

GPU check, Kimi K3 colocated TP8 on 8xB300 with DFLASH (block size 8), GSM8K 5-shot on all 1,319 questions. The comparison is main + #40794 (K3 DFLASH) with and without this PR. Target-verify CUDA-graph memory per GPU goes from 2.215 GB to 1.756 GB (−0.46 GB, −21%), and startup free memory rises by the same amount. Accuracy and acceptance are unchanged: 95.0% / average accept length 5.07 without this PR, 95.6% / 5.06 with it.

---

CI States

Latest PR Test (Base): ⏳ Run #36628674548
Latest PR Test (Extra): ❌ Run #36628673923
Latest PR Test (AMD ROCm 10): ⏳ Run #36628674234

[chromecast56]

/tag-and-rerun-ci

[hnyls2002]

/rerun-test test_kimi_k3_b300.py test_kimi_k3_b300_low_latency.py test_basic_sanity_dflash.py test_basic_sanity_dspark.py

[github-actions]

Results for /rerun-test test_kimi_k3_b300.py test_kimi_k3_b300_low_latency.py test_basic_sanity_dflash.py test_basic_sanity_dspark.py:

🚀 8-gpu-b300 (2 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_kimi_k3_b300.py
cd test/ && python3 registered/e2e/models/test_kimi_k3_b300_low_latency.py

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

cd test/ && python3 registered/core/test_basic_sanity_dflash.py

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

cd test/ && python3 registered/core/test_basic_sanity_dspark.py

[hnyls2002]

Pushed fixes for:

• TBO capture crash: filter_batch rejects the new aux_hidden_states_buffer field on the parent batch. Children now get None (066dd09).
• Per-stream buffers: unnecessary. pdmux never replays two decode graphs at once, so there is now one buffer per runner (08a5a40).
• Separate buffer plumbing: the buffer now lives in DecodeInputBuffers and is sliced in capture_prepare, the same way as the PP proxy buffer. Its width comes from misc_utils.resolve_aux_hidden_states_width (08a5a40).
• Gating split across three places: the runner only checks DFlash target / non-draft; the PP check lives only in K3's packs_aux_hidden_states (08a5a40).