Skip to content

Commit 9af9bf4

Browse files
Xiaodong Yeclaude
Xiaodong Ye
and
claude
committed
MUSA-0203: PR vllm-project#34880 backport — llm_base_proposer wiring (Step 3)
Wires the CUDAGraphWrapper(FULL) draft-model capture path. Three new patch files target the v0.20.0-dev EagleProposer implementation (which lives in llm_base_proposer.py, not eagle.py as in PR vllm-project#34880's base). vllm__v1__spec_decode__llm_base_proposer.patch.py (11 hunks): 1. Import CUDAGraphWrapper + BatchDescriptor. 2. CudagraphDispatcher(self.vllm_config) -> (..., for_draft_model=True). 3. initialize_cudagraph_keys: pass target's cudagraph_mode through directly (was restricted to PIECEWISE). With for_draft_model=True the dispatcher now also registers FULL-mode keys at qdl=1 when target has FULL captures. 4. load_model: wrap self.model with CUDAGraphWrapper(runtime_mode=FULL) when target uses FULL captures. The wrapper is a no-op when the runtime mode at call time doesn't match FULL — safe with PIECEWISE target (our current SOTA config) until target switches to FULL_AND_PIECEWISE. 5. _determine_batch_execution_and_padding returns 4-tuple now (adds batch_desc slot); threads uniform_decode + uniform_decode_query_len through. 6. propose()'s first set_forward_context call passes batch_descriptor=batch_desc. 7. propose()'s decode-loop set_forward_context call passes batch_descriptor=batch_desc with uniform_decode_query_len=1. 8. dummy_run() 4-tuple unpack. vllm__v1__spec_decode__dflash.patch.py (1 hunk): DFlashProposer.dummy_run 4-tuple unpack. vllm__v1__spec_decode__utils.patch.py (2 hunks): eagle_prepare_inputs_padded_kernel: early-exit guard for query_len==0 rows + reuse q_end_idx. (Independent kernel safety fix from same PR.) Skipped from PR vllm-project#34880 (lower-priority for MUSA's PIECEWISE workload, can be added later if FULL_AND_PIECEWISE perf tuning needs them): - target_model_batch_desc parameter to propose() — requires gpu_model_runner caller change which adds ~12 more hunks. - dummy_run two-pass capture variant (matters only at boot capture when target uses FULL_AND_PIECEWISE). - prepare_next_token_ids_padded / prepare_inputs_padded padding edge cases (matters for padded-drafter batches > 1). With these patches MUSA's PIECEWISE-mode SOTA (current production config) is unchanged at runtime — the new code paths only activate when target uses FULL_AND_PIECEWISE or FULL_DECODE_ONLY. The infrastructure is in place for follow-up FULL-mode validation (MUSA-0205 candidate). Total Step 3 diff: +178 lines across 3 patch files (5 patches total for Steps 2+3, 23 hunks). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 30d3ef4 commit 9af9bf4

3 files changed

Lines changed: 285 additions & 0 deletions

File tree

‎vllm_musa/patches/vllm__v1__spec_decode__dflash.patch.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3+
"""MUSA-0203: backport of vllm-project/vllm#34880 — dflash.py tuple shape.
4+
5+
After ``vllm__v1__spec_decode__llm_base_proposer`` patches
6+
``_determine_batch_execution_and_padding`` to return 4-tuple,
7+
``DFlashProposer.dummy_run`` must unpack the new shape too.
8+
"""
9+
10+
_OLD = """ num_query_tokens = min(num_tokens, self.max_query_tokens)
11+
cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
12+
self._determine_batch_execution_and_padding(
13+
num_query_tokens, use_cudagraphs=use_cudagraphs
14+
)"""
15+
16+
_NEW = """ num_query_tokens = min(num_tokens, self.max_query_tokens)
17+
cudagraph_runtime_mode, _, num_input_tokens, num_tokens_across_dp = (
18+
self._determine_batch_execution_and_padding(
19+
num_query_tokens, use_cudagraphs=use_cudagraphs
20+
)"""
21+
22+
PATCHES = [(_OLD, _NEW)]

‎vllm_musa/patches/vllm__v1__spec_decode__llm_base_proposer.patch.py‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3+
"""MUSA-0203: backport of vllm-project/vllm#34880 — llm_base_proposer.py.
4+
5+
In upstream PR #34880 these hunks target ``vllm/v1/spec_decode/eagle.py``.
6+
On v0.20.0-dev the EagleProposer implementation lives in
7+
``vllm/v1/spec_decode/llm_base_proposer.py`` (eagle.py is a 22-line
8+
stub). The hunks below re-target the same semantic changes to the
9+
v0.20.0-dev file structure.
10+
11+
Changes:
12+
1. Import ``CUDAGraphWrapper`` + ``BatchDescriptor``.
13+
2. ``CudagraphDispatcher(self.vllm_config)`` -> ``(..., for_draft_model=True)``.
14+
3. ``initialize_cudagraph_keys``: pass through target's
15+
``cudagraph_mode`` directly (was restricted to PIECEWISE).
16+
4. ``load_model``: wrap ``self.model`` with
17+
``CUDAGraphWrapper(runtime_mode=FULL)`` when target has FULL
18+
captures (no-op when target is PIECEWISE-only).
19+
5. ``_determine_batch_execution_and_padding``: return 4-tuple
20+
(adds ``batch_desc``); thread ``uniform_decode`` +
21+
``uniform_decode_query_len`` through.
22+
6. ``propose``: pass ``batch_descriptor=batch_desc`` to both
23+
``set_forward_context`` calls (first forward + per-step loop).
24+
25+
Skipped from PR #34880 in this version (lower-priority for MUSA
26+
PIECEWISE workload — can be added later):
27+
- ``target_model_batch_desc`` propose() parameter (gpu_model_runner
28+
caller change required).
29+
- ``dummy_run`` two-pass capture variant.
30+
- ``prepare_next_token_ids_padded`` / ``prepare_inputs_padded``
31+
padding fixes.
32+
33+
When upstream PR #34880 merges (or v0.20.0-dev rebases onto it), this
34+
patch's anchors will stop matching and the file should be deleted.
35+
"""
36+
37+
# ---- Hunk 1: import CUDAGraphWrapper + BatchDescriptor ----
38+
_OLD_IMPORTS = """from vllm.distributed.parallel_state import get_pp_group
39+
from vllm.forward_context import set_forward_context"""
40+
41+
_NEW_IMPORTS = """from vllm.compilation.cuda_graph import CUDAGraphWrapper
42+
from vllm.distributed.parallel_state import get_pp_group
43+
from vllm.forward_context import BatchDescriptor, set_forward_context"""
44+
45+
# ---- Hunk 2: dispatcher init — for_draft_model=True ----
46+
_OLD_DISPATCHER_INIT = """ self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)"""
47+
48+
_NEW_DISPATCHER_INIT = """ # MUSA-0203 / PR #34880: dedicated dispatcher for the draft model so
49+
# FULL-mode capture keys are registered with uniform_decode_query_len=1
50+
# (each Eagle decode step is 1 token per request).
51+
self.cudagraph_dispatcher = CudagraphDispatcher(
52+
self.vllm_config, for_draft_model=True
53+
)"""
54+
55+
# ---- Hunk 3: initialize_cudagraph_keys — pass target mode through ----
56+
_OLD_INIT_KEYS = """ \"\"\"Initialize cudagraph dispatcher keys for eagle.
57+
58+
Eagle only supports PIECEWISE cudagraphs (via mixed_mode).
59+
This should be called after adjust_cudagraph_sizes_for_spec_decode.
60+
\"\"\"
61+
if (
62+
not self.speculative_config.enforce_eager
63+
and cudagraph_mode.mixed_mode()
64+
in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
65+
):
66+
eagle_cudagraph_mode = CUDAGraphMode.PIECEWISE
67+
else:
68+
eagle_cudagraph_mode = CUDAGraphMode.NONE
69+
70+
self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)"""
71+
72+
_NEW_INIT_KEYS = """ \"\"\"Initialize cudagraph dispatcher keys for eagle.
73+
74+
MUSA-0203 / PR #34880: pass the target's cudagraph_mode through
75+
directly. With ``for_draft_model=True`` the dispatcher now
76+
registers both the legacy mixed-mode (PIECEWISE) keys AND the
77+
new FULL-mode keys at uniform_decode_query_len=1, so the draft
78+
model can dispatch FULL captures when the target uses FULL or
79+
FULL_AND_PIECEWISE.
80+
\"\"\"
81+
if self.speculative_config.enforce_eager:
82+
cudagraph_mode = CUDAGraphMode.NONE
83+
self.cudagraph_dispatcher.initialize_cudagraph_keys(cudagraph_mode)"""
84+
85+
# ---- Hunk 4: load_model — wrap with CUDAGraphWrapper ----
86+
_OLD_LOAD_MODEL = """self.model = self._get_model()
87+
88+
# Find draft layers"""
89+
90+
_NEW_LOAD_MODEL = """self.model = self._get_model()
91+
# MUSA-0203 / PR #34880: wrap the draft model with FULL-mode
92+
# CUDAGraphWrapper when target uses FULL captures. CUDAGraphWrapper
93+
# is a no-op when the runtime mode at call time doesn't match FULL,
94+
# so this is safe when target is PIECEWISE-only.
95+
cudagraph_mode = self.compilation_config.cudagraph_mode
96+
if (
97+
cudagraph_mode.has_full_cudagraphs()
98+
and not self.vllm_config.parallel_config.use_ubatching
99+
and not self.speculative_config.disable_padded_drafter_batch
100+
):
101+
self.model = CUDAGraphWrapper(
102+
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
103+
)
104+
105+
# Find draft layers"""
106+
107+
# ---- Hunk 5: _determine_batch_execution_and_padding signature + dispatch + return ----
108+
_OLD_DETERMINE = """ def _determine_batch_execution_and_padding(
109+
self,
110+
num_tokens: int,
111+
use_cudagraphs: bool = True,
112+
) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
113+
cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
114+
num_tokens,
115+
valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),"""
116+
117+
_NEW_DETERMINE = """ def _determine_batch_execution_and_padding(
118+
self,
119+
num_tokens: int,
120+
uniform_decode: bool = False,
121+
use_cudagraphs: bool = True,
122+
uniform_decode_query_len: int | None = None,
123+
) -> tuple[CUDAGraphMode, BatchDescriptor, int, torch.Tensor | None]:
124+
cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
125+
num_tokens,
126+
uniform_decode,
127+
valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),
128+
uniform_decode_query_len=uniform_decode_query_len,"""
129+
130+
# Need to also patch the return statement of _determine_batch_execution_and_padding
131+
# to 4-tuple. Find the unique anchor.
132+
_OLD_DETERMINE_RETURN = """ return cudagraph_mode, num_tokens_padded, num_tokens_across_dp"""
133+
134+
_NEW_DETERMINE_RETURN = """ return cudagraph_mode, batch_desc, num_tokens_padded, num_tokens_across_dp"""
135+
136+
# ---- Hunk 6: propose() first-forward _determine_batch unpack ----
137+
# In v0.20.0-dev the propose() first forward is:
138+
# cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
139+
# self._determine_batch_execution_and_padding(num_tokens)
140+
# )
141+
# We need 4-tuple unpack now.
142+
_OLD_PROPOSE_FIRST_DET = """ cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
143+
self._determine_batch_execution_and_padding(num_tokens)
144+
)"""
145+
146+
_NEW_PROPOSE_FIRST_DET = """ cudagraph_runtime_mode, batch_desc, num_input_tokens, num_tokens_across_dp = (
147+
self._determine_batch_execution_and_padding(num_tokens)
148+
)"""
149+
150+
# ---- Hunk 7: propose() first-forward set_forward_context — thread batch_desc ----
151+
_OLD_FFC_FIRST = """ cudagraph_runtime_mode=cudagraph_runtime_mode,
152+
slot_mapping=self._get_slot_mapping(
153+
slot_mapping_size, common_attn_metadata.slot_mapping
154+
),"""
155+
156+
_NEW_FFC_FIRST = """ cudagraph_runtime_mode=cudagraph_runtime_mode,
157+
batch_descriptor=batch_desc,
158+
slot_mapping=self._get_slot_mapping(
159+
slot_mapping_size, common_attn_metadata.slot_mapping
160+
),"""
161+
162+
# ---- Hunk 8: propose() loop _determine_batch unpack ----
163+
# The second occurrence is inside the loop, with `batch_size` not `num_tokens`.
164+
_OLD_PROPOSE_LOOP_DET = """ cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = (
165+
self._determine_batch_execution_and_padding(batch_size)
166+
)"""
167+
168+
_NEW_PROPOSE_LOOP_DET = """ # MUSA-0203 / PR #34880: subsequent decode passes dispatch with
169+
# (batch_size, uniform=True, qdl=1) so the draft's FULL-mode capture
170+
# keys are exercised.
171+
cudagraph_runtime_mode, batch_desc, input_batch_size, batch_size_across_dp = (
172+
self._determine_batch_execution_and_padding(
173+
batch_size, uniform_decode=True, uniform_decode_query_len=1
174+
)
175+
)"""
176+
177+
# ---- Hunk 9: propose() loop set_forward_context — thread batch_desc ----
178+
_OLD_FFC_LOOP = """ cudagraph_runtime_mode=cudagraph_runtime_mode,
179+
slot_mapping=self._get_slot_mapping(input_batch_size),"""
180+
181+
_NEW_FFC_LOOP = """ cudagraph_runtime_mode=cudagraph_runtime_mode,
182+
batch_descriptor=batch_desc,
183+
slot_mapping=self._get_slot_mapping(input_batch_size),"""
184+
185+
# ---- Hunk 10: dummy_run() unpack — match 4-tuple shape ----
186+
_OLD_DUMMY_UNPACK = """ cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
187+
self._determine_batch_execution_and_padding(
188+
num_tokens, use_cudagraphs=use_cudagraphs
189+
)
190+
)"""
191+
192+
_NEW_DUMMY_UNPACK = """ cudagraph_runtime_mode, _, num_input_tokens, num_tokens_across_dp = (
193+
self._determine_batch_execution_and_padding(
194+
num_tokens, use_cudagraphs=use_cudagraphs
195+
)
196+
)"""
197+
198+
PATCHES = [
199+
(_OLD_IMPORTS, _NEW_IMPORTS),
200+
(_OLD_DISPATCHER_INIT, _NEW_DISPATCHER_INIT),
201+
(_OLD_INIT_KEYS, _NEW_INIT_KEYS),
202+
(_OLD_LOAD_MODEL, _NEW_LOAD_MODEL),
203+
(_OLD_DETERMINE, _NEW_DETERMINE),
204+
(_OLD_DETERMINE_RETURN, _NEW_DETERMINE_RETURN),
205+
(_OLD_PROPOSE_FIRST_DET, _NEW_PROPOSE_FIRST_DET),
206+
(_OLD_FFC_FIRST, _NEW_FFC_FIRST),
207+
(_OLD_PROPOSE_LOOP_DET, _NEW_PROPOSE_LOOP_DET),
208+
(_OLD_FFC_LOOP, _NEW_FFC_LOOP),
209+
(_OLD_DUMMY_UNPACK, _NEW_DUMMY_UNPACK),
210+
]

‎vllm_musa/patches/vllm__v1__spec_decode__utils.patch.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3+
"""MUSA-0203: backport of vllm-project/vllm#34880 — utils.py kernel safety.
4+
5+
Adds early-exit guard for ``query_len == 0`` rows in
6+
``eagle_prepare_inputs_padded_kernel``. The padded-drafter batch can
7+
contain empty padding rows whose ``query_start_loc[i] == query_start_loc[i+1]``.
8+
Without this guard the kernel writes a stale ``q_last_tok_idx`` for
9+
those rows and can corrupt downstream draft token sampling.
10+
11+
NOTE: this is the upstream ``vllm/v1/spec_decode/utils.py`` —
12+
distinct from ``vllm_musa/v1/spec_decode/utils.py`` which lives at
13+
the same path under vllm-musa and installs the MUSA-Triton-adapted
14+
``eagle_prepare_next_token_padded_kernel``.
15+
"""
16+
17+
# Hunk 1: insert query_len == 0 guard at top of kernel body
18+
_OLD = """ if req_idx >= num_reqs:
19+
return
20+
21+
# Calculate num_draft_tokens from cu_num_draft_tokens, which is an inclusive
22+
# cumulative sum (first entry is the first value, not zero).
23+
cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx)"""
24+
25+
_NEW = """ if req_idx >= num_reqs:
26+
return
27+
28+
q_start_idx = tl.load(query_start_loc_gpu_ptr + req_idx)
29+
q_end_idx = tl.load(query_start_loc_gpu_ptr + req_idx + 1)
30+
query_len = q_end_idx - q_start_idx
31+
32+
if query_len == 0:
33+
tl.store(token_indices_to_sample_ptr + req_idx, q_end_idx - 1)
34+
tl.store(num_rejected_tokens_gpu_ptr + req_idx, 0)
35+
return
36+
37+
# Calculate num_draft_tokens from cu_num_draft_tokens, which is an inclusive
38+
# cumulative sum (first entry is the first value, not zero).
39+
cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx)"""
40+
41+
# Hunk 2: reuse q_end_idx instead of redundant tl.load
42+
_OLD_LAST_TOK = """ # query_start_loc[req_idx + 1] is the start position of the next request,
43+
# which is one past the last token of this request.
44+
q_last_tok_idx = tl.load(query_start_loc_gpu_ptr + req_idx + 1) - 1"""
45+
46+
_NEW_LAST_TOK = """ # query_start_loc[req_idx + 1] is the start position of the next request,
47+
# which is one past the last token of this request.
48+
q_last_tok_idx = q_end_idx - 1"""
49+
50+
PATCHES = [
51+
(_OLD, _NEW),
52+
(_OLD_LAST_TOK, _NEW_LAST_TOK),
53+
]

0 commit comments

Comments
 (0)