| |
| |
| @@ -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, |
| |
| |
| @@ -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 |
| |
| |
| @@ -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] |
| |
| |
| @@ -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, |
|
|
