Skip to content

Commit afd11cd

Browse files
Xiaodong Yeclaude
andcommitted
MUSA-0203: gate CUDAGraphWrapper(FULL) wrap behind VLLM_MUSA_DRAFT_FULL_WRAP
Step 4's load_model CUDAGraphWrapper(FULL) wrap is functionally correct but causes accept-rate regression (47.95% -> 19.5% on M2.5+Eagle3 yeahdongcn70). Suspected cause: captured draft FA kernel reads stale seq_lens/slot_mapping from dummy_run's synthetic batch. Default the wrap to OFF (matches Step 1+2+3 boot behavior with FULL captures of TARGET only, draft remains eager). Opt back in with `VLLM_MUSA_DRAFT_FULL_WRAP=1` once the stale-state issue is debugged. The dummy_run/dispatcher/gpu_input_batch/utils patches remain active — they're additive and harmless without the wrap. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 3d43c68 commit afd11cd

1 file changed

Lines changed: 7 additions & 1 deletion

File tree

vllm_musa/patches/vllm__v1__spec_decode__llm_base_proposer.patch.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,9 +92,15 @@
9292
# CUDAGraphWrapper when target uses FULL captures. CUDAGraphWrapper
9393
# is a no-op when the runtime mode at call time doesn't match FULL,
9494
# so this is safe when target is PIECEWISE-only.
95+
# Gated by VLLM_MUSA_DRAFT_FULL_WRAP=1 (default OFF). The wrap
96+
# makes the captured replay produce stale draft tokens (accept
97+
# rate drops 47% -> 19%) until we plumb target_model_batch_desc
98+
# through the full PR #34880 caller chain.
99+
import os as _os
95100
cudagraph_mode = self.compilation_config.cudagraph_mode
96101
if (
97-
cudagraph_mode.has_full_cudagraphs()
102+
_os.environ.get("VLLM_MUSA_DRAFT_FULL_WRAP", "0") == "1"
103+
and cudagraph_mode.has_full_cudagraphs()
98104
and not self.vllm_config.parallel_config.use_ubatching
99105
and not self.speculative_config.disable_padded_drafter_batch
100106
):

0 commit comments

Comments
 (0)