File size: 12,948 Bytes
ee3e28a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 | --- /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/v1/spec_decode/llm_base_proposer.py.orig 2026-08-16 17:03:01.543167000 +0800
+++ /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/v1/spec_decode/llm_base_proposer.py 2026-08-16 17:03:18.492170389 +0800
@@ -186,6 +186,16 @@
self.mrope_positions = torch.zeros(
(3, self.max_positions + 1), dtype=torch.int64, device=device
)
+ # Draft-model path: slot mapping must be computed from SEQUENCE
+ # indices, which differ from (compressed) M-RoPE position values
+ # whenever the prompt contains an image. self.positions holds the
+ # sequence indices; seq_step tracks them across drafting steps.
+ self.positions = torch.zeros(
+ self.max_positions, dtype=torch.int64, device=device
+ )
+ self.seq_step = torch.zeros(
+ self.max_batch_size, dtype=torch.int64, device=device
+ )
elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
self.xdrope_positions = torch.zeros(
(self.uses_xdrope_dim, self.max_positions + 1),
@@ -334,9 +344,13 @@
)
def _raise_if_mrope(self):
- if self.draft_model_config.uses_mrope:
+ # M-RoPE is supported on the draft-model path: slot mapping is computed
+ # from true sequence indices (self.positions doubles as a seq-index
+ # buffer, see set_inputs_first_pass) while M-RoPE positions are
+ # scattered separately into self.mrope_positions.
+ if self.draft_model_config.uses_mrope and self.parallel_drafting:
raise NotImplementedError(
- "Speculative Decoding with draft models or parallel drafting "
+ "Speculative Decoding with parallel drafting "
"does not support M-RoPE yet"
)
@@ -628,6 +642,11 @@
if self.uses_mrope:
positions = self.mrope_positions[:, token_indices_to_sample]
+ if self.needs_extra_input_slots:
+ # Track true sequence indices alongside M-RoPE positions for
+ # per-step slot mapping (self.positions holds sequence indices
+ # on this path; see set_inputs_first_pass).
+ self.seq_step[:batch_size] = self.positions[token_indices_to_sample]
else:
positions = self.positions[token_indices_to_sample]
hidden_states = hidden_states[token_indices_to_sample]
@@ -776,12 +795,22 @@
) -> torch.Tensor:
"""Update positions, slot mappings, and sequence metadata for the
next draft step. Returns the updated positions tensor."""
- positions_1d = positions[0] if self.uses_mrope else positions
- if self.uses_mrope:
+ mrope_draft_model = self.uses_mrope and self.needs_extra_input_slots
+ if mrope_draft_model:
+ # Slot mapping must advance by SEQUENCE index, not by (compressed)
+ # M-RoPE position value. seq_step is advanced in place by the
+ # fused kernel; M-RoPE positions advance by +1 on all three rows
+ # separately below (drafted tokens are always text-region).
+ positions_1d = self.seq_step[:batch_size]
+ out_pos = self.seq_step[:batch_size]
+ elif self.uses_mrope:
+ positions_1d = positions[0]
out_pos = self.mrope_positions[0, :batch_size]
elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
+ positions_1d = positions
out_pos = self.xdrope_positions[0, :batch_size]
else:
+ positions_1d = positions
out_pos = self.positions[:batch_size]
eagle_step_update_slot_mapping_and_metadata(
positions_1d=positions_1d,
@@ -794,7 +823,14 @@
input_batch_size=input_batch_size,
)
common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
- if self.uses_mrope:
+ if mrope_draft_model:
+ torch.clamp(
+ positions + 1,
+ max=self.max_model_len - 1,
+ out=self.mrope_positions[:, :batch_size],
+ )
+ positions = self.mrope_positions[:, :batch_size]
+ elif self.uses_mrope:
self.mrope_positions[1:, :batch_size] = self.mrope_positions[0, :batch_size]
positions = self.mrope_positions[:, :batch_size]
elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
@@ -898,14 +934,34 @@
if num_rejected_tokens_gpu is not None:
query_end_loc = query_end_loc - num_rejected_tokens_gpu
+ kernel_positions_in = target_positions
+ if self.uses_mrope:
+ # The kernel synthesizes positions as start_pos + j, reading
+ # only target_positions_ptr[query_start_loc[r]] per request.
+ # Slot mapping (step 2 below) consumes self.positions as
+ # SEQUENCE indices, which diverge from compressed M-RoPE
+ # position values on image prompts. Feed the kernel per-request
+ # context lengths so self.positions ends up holding sequence
+ # indices; true M-RoPE positions are scattered into
+ # self.mrope_positions right after the kernel.
+ naive_qlens = cad.query_start_loc[1:] - cad.query_start_loc[:-1]
+ seq_scratch = torch.zeros(
+ total_num_input_tokens, dtype=torch.int64, device=self.device
+ )
+ seq_scratch[cad.query_start_loc[:-1].to(torch.int64)] = (
+ cad.seq_lens - naive_qlens
+ ).to(torch.int64)
+ kernel_positions_in = seq_scratch
+
copy_and_expand_eagle_inputs_kernel[grid](
# (Padded) Inputs from the target model
target_token_ids_ptr=target_token_ids,
- target_positions_ptr=target_positions,
+ target_positions_ptr=kernel_positions_in,
next_token_ids_ptr=next_token_ids, # sampled tokens, one per request
# Outputs to the drafting buffers
out_input_ids_ptr=self.input_ids,
- out_positions_ptr=self.positions, # Doesn't support mrope for now
+ # For M-RoPE this receives sequence indices (see above)
+ out_positions_ptr=self.positions,
out_is_rejected_token_mask_ptr=self.is_rejected_token_mask,
out_is_masked_token_mask_ptr=self.is_masked_token_mask,
out_new_token_indices_ptr=token_indices_to_sample,
@@ -922,6 +978,31 @@
shift_input_ids=self.pass_hidden_states_to_model,
BLOCK_SIZE_TOKENS=BLOCK_SIZE_TOKENS,
)
+
+ self._mm_expand_map = None
+ if self.uses_mrope:
+ # Scatter true M-RoPE positions into the expanded layout.
+ # shift_input_ids is False on this path, so input token i of
+ # request r lands at output slot i + r * extra_slots_per_request.
+ # ORDERING IS LOAD-BEARING: under this formula the rejected-
+ # region inputs collide with the bonus slot, so the bonus-slot
+ # write below must come last to overwrite them. Trailing
+ # rejected output slots keep stale values, which is harmless
+ # (their KV writes are masked and outputs ignored).
+ req_ids = torch.repeat_interleave(
+ self.arange[:batch_size].to(torch.int64), naive_qlens
+ )
+ out_map = (
+ self.arange[:total_num_input_tokens]
+ + req_ids * self.extra_slots_per_request
+ )
+ self.mrope_positions[:, out_map] = target_positions[
+ :, :total_num_input_tokens
+ ]
+ self.mrope_positions[:, token_indices_to_sample.to(torch.int64)] = (
+ target_positions[:, query_end_loc.to(torch.int64)] + 1
+ )
+ self._mm_expand_map = out_map
if self.pass_hidden_states_to_model:
assert self.parallel_drafting_hidden_state_tensor is not None
self.hidden_states[out_hidden_state_mapping] = target_hidden_states
@@ -968,6 +1049,21 @@
if self.supports_mm_inputs:
mm_embeds, is_mm_embed = mm_embed_inputs or (None, None)
+ if (
+ is_mm_embed is not None
+ and getattr(self, "_mm_expand_map", None) is not None
+ ):
+ # Draft-model path: re-align the multimodal mask from the
+ # runner's token layout to the expanded (bonus-slot-inserted)
+ # drafting layout. Bonus/rejected slots are text -> False.
+ expanded = torch.zeros(
+ num_tokens, dtype=is_mm_embed.dtype, device=is_mm_embed.device
+ )
+ expanded[self._mm_expand_map] = is_mm_embed[
+ : self._mm_expand_map.shape[0]
+ ]
+ is_mm_embed = expanded
+
self.inputs_embeds[:num_tokens] = self.model.embed_input_ids(
self.input_ids[:num_tokens],
multimodal_embeddings=mm_embeds,
--- /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/v1/worker/gpu_model_runner.py.orig 2026-08-16 17:03:01.546349000 +0800
+++ /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/v1/worker/gpu_model_runner.py 2026-08-16 17:03:23.167490000 +0800
@@ -5260,7 +5260,7 @@
if self.supports_mm_inputs and self.drafter.supports_mm_inputs:
mm_embed_inputs = self._gather_mm_embeddings(
scheduler_output,
- shift_computed_tokens=1,
+ shift_computed_tokens=(1 if self.drafter.pass_hidden_states_to_model else 0),
)
else:
mm_embed_inputs = None
--- /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/model_executor/layers/quantization/compressed_tensors/utils.py.orig 2026-08-16 17:08:26.513544000 +0800
+++ /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/model_executor/layers/quantization/compressed_tensors/utils.py 2026-08-16 17:08:26.531649000 +0800
@@ -55,6 +55,12 @@
if layer_name is None:
return False
+ # Draft models in speculative decoding are built under the "draft_model."
+ # prefix (see DraftModelProposer._get_model), which is unknown to the
+ # checkpoint ignore list -- strip it before matching.
+ if layer_name.startswith("draft_model."):
+ layer_name = layer_name[len("draft_model."):]
+
# layer_name = model.layers.0.self_attn.qkv_proj
# proj_name = qkv_proj
proj_name = layer_name.split(".")[-1]
--- /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/v1/sample/rejection_sampler.py.orig 2026-08-16 17:51:54.716036000 +0800
+++ /testessfs10/users/zeyu.zhang/wangyu_ssd/vllm_env/lib/python3.11/site-packages/vllm/v1/sample/rejection_sampler.py 2026-08-16 17:51:54.737130000 +0800
@@ -8,6 +8,10 @@
from typing import TYPE_CHECKING
import torch
+
+import os
+
+_RELAX_LOGTAU = float(os.environ.get("VLLM_SPEC_RELAX_LOGTAU", "0") or "0")
import torch.nn as nn
from vllm.config.model import PROCESSED_LOGPROBS_MODES
@@ -453,6 +457,22 @@
if not sampling_metadata.all_random:
# Rejection sampling for greedy sampling requests.
target_argmax = target_logits.argmax(dim=-1)
+ if _RELAX_LOGTAU > 0.0 and num_tokens > 0 and sampling_metadata.all_greedy:
+ # Relaxed acceptance gate (logit-ratio, CSD-style): accept draft
+ # token d when z(d) >= z(argmax) - logtau (logtau = -ln tau).
+ # Implemented by overwriting target_argmax at gate-passing
+ # positions, so the greedy kernel realizes the relaxed policy;
+ # correction/bonus tokens are unaffected (they only come from
+ # gate-failing positions or the bonus slot).
+ _zmax = target_logits.gather(1, target_argmax.unsqueeze(1)).squeeze(1)
+ _zd = target_logits.gather(
+ 1, draft_token_ids.to(torch.int64).unsqueeze(1)
+ ).squeeze(1)
+ target_argmax = torch.where(
+ _zd >= _zmax - _RELAX_LOGTAU,
+ draft_token_ids.to(target_argmax.dtype),
+ target_argmax,
+ )
rejection_greedy_sample_kernel[(batch_size,)](
output_token_ids,
cu_num_draft_tokens,
|