[Spec] Extract shared draft worker construction and generalize draft sampler capture

[hnyls2002]

Behavior-preserving DFLASH refactor, one logical change per commit:

• rename DFlashVerifyInput.num_tokens_per_batch to num_tokens_per_req, matching the field name the attention backends and EagleVerifyInput already use
• require a CUDA device for DFLASH at argument resolution instead of failing mid-startup
• extract draft TpModelWorker construction into draft_worker_common (build_draft_tp_worker, make_draft_input_v2, make_draft_block_spec_info, build_block_pos_offsets); this also dedups two identical draft-input constructions inside the DFLASH worker
• replace the runner's hardcoded dflash_draft_sampler special case with a generic model_runner.capture_tail_hooks list; the DFLASH sampler registers itself as a hook with the same fail-loudly semantics
• extract build_dflash_verify_target_probs from the sampling accept helper

Verification: full test_dflash.py run has an identical pass/fail signature to the main baseline on the same machine (48 passed; the same 3 pre-existing environment-specific failures on both).

---

CI States

Latest PR Test (Base): ❌ Run #29162877753
Latest PR Test (Extra): 🚫 Run #29162877678

[hnyls2002]

/rerun-test test_dflash.py test_basic_sanity_dflash.py

[github-actions]

Results for /rerun-test test_dflash.py test_basic_sanity_dflash.py:

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

cd test/ && python3 registered/spec/dflash/test_dflash.py
cd test/ && python3 registered/core/test_basic_sanity_dflash.py

[gemini-code-assist]

Code Review

This pull request refactors the speculative decoding draft worker setup by extracting common utilities into a new draft_worker_common.py module and generalizing the tail capture process using a list of hooks (capture_tail_hooks) in the model runner. It also adds a CUDA device check for DFLASH speculative decoding and extracts target probability building logic in dflash_utils.py. The review feedback highlights several improvement opportunities: wrapping the draft worker initialization in a try...finally block to guarantee the restoration of global server arguments, using the idiomatic is_hip() helper instead of torch.version.hip, and resetting or assigning capture_tail_hooks instead of appending to prevent duplicate hook accumulation if CUDA graphs are initialized multiple times.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[hnyls2002]

/tag-and-rerun-ci extra
Both red CI states are only the pr-gate failing for the missing run-ci label; no test stage has run yet.
Opting into the extra suite for the DFLASH pressure coverage: test_gemma4_dflash_31b_extra.py, test_kimi_k26_nvfp4_dflash.py.
Key targeted tests to watch once CI runs: test_basic_sanity_dflash.py, and test_pcg_with_speculative_decoding_dflash.py (direct coverage for the capture_tail_hooks change in the decode graph capture path).

[hnyls2002]

/rerun-test test_pcg_with_speculative_decoding_dflash.py test_kimi_k26_nvfp4_dflash.py test_gemma4_dflash_31b_extra.py

[github-actions]

Results for /rerun-test test_pcg_with_speculative_decoding_dflash.py test_kimi_k26_nvfp4_dflash.py test_gemma4_dflash_31b_extra.py:

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

cd test/ && python3 registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py

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

cd test/ && python3 registered/quant/test_kimi_k26_nvfp4_dflash.py

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

cd test/ && python3 registered/spec/test_gemma4_dflash_31b_extra.py

[hnyls2002]

Latest PR Test (Base): ✅ Run #29149096757
Latest PR Test (Extra): ❌ Run #29149096704

[hnyls2002]

/rerun-test test_kimi_k26_nvfp4_dflash.py

[github-actions]

Results for /rerun-test test_kimi_k26_nvfp4_dflash.py:

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

cd test/ && python3 registered/quant/test_kimi_k26_nvfp4_dflash.py