| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import replace |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | from typing_extensions import override |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.forward_context import set_forward_context |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 14 | from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer |
| 15 | from vllm.v1.spec_decode.utils import ( |
| 16 | _copy_and_expand_dflash_inputs, |
| 17 | next_power_of_2, |
| 18 | ) |
| 19 | |
| 20 | logger = init_logger(__name__) |
| 21 | |
| 22 | |
| 23 | class DFlashProposer(SpecDecodeBaseProposer): |
| 24 | def __init__( |
| 25 | self, |
| 26 | vllm_config: VllmConfig, |
| 27 | device: torch.device, |
| 28 | runner=None, |
| 29 | ): |
| 30 | assert vllm_config.speculative_config is not None |
| 31 | assert vllm_config.speculative_config.method == "dflash" |
| 32 | super().__init__( |
| 33 | vllm_config=vllm_config, |
| 34 | device=device, |
| 35 | pass_hidden_states_to_model=True, |
| 36 | runner=runner, |
| 37 | ) |
| 38 | |
| 39 | # Only next_token_ids and mask tokens are query tokens, all other context is K/V |
| 40 | self.max_query_tokens = self.max_batch_size * (1 + self.num_speculative_tokens) |
| 41 | self.max_padded_query_tokens = max( |
| 42 | self.max_query_tokens, |
| 43 | vllm_config.compilation_config.max_cudagraph_capture_size or 0, |
| 44 | ) |
| 45 | # Positions covers both context states + query states |
| 46 | self.max_positions = self.max_num_tokens + self.max_padded_query_tokens |
| 47 | |
| 48 | # Separate context buffers to keep query buffer addresses stable for CUDA graphs |
| 49 | self._context_slot_mapping_buffer = torch.zeros( |
| 50 | self.max_num_tokens, |
| 51 | dtype=torch.int64, |
| 52 | device=device, |
| 53 | ) |
| 54 | self._slot_mapping_buffer = torch.zeros( |
| 55 | self.max_padded_query_tokens, |
| 56 | dtype=torch.int64, |
| 57 | device=device, |
| 58 | ) |
| 59 | self._context_positions_buffer = torch.zeros( |
| 60 | self.max_num_tokens, |
| 61 | dtype=torch.int64, |
| 62 | device=device, |
| 63 | ) |
| 64 | self.positions = torch.zeros( |
| 65 | self.max_padded_query_tokens, |
| 66 | dtype=torch.int64, |
| 67 | device=device, |
| 68 | ) |
| 69 | |
| 70 | self.arange = torch.arange( |
| 71 | self.max_positions + 1, device=device, dtype=torch.int32 |
| 72 | ) |
| 73 | |
| 74 | # For DFlash we use the input embeddings to embed the mask token |
| 75 | self.parallel_drafting_hidden_state_tensor = None |
| 76 | |
| 77 | from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal |
| 78 | |
| 79 | self.dflash_causal = not dflash_has_any_non_causal( |
| 80 | self.draft_model_config.hf_config |
| 81 | ) |
| 82 | _copy_and_expand_dflash_inputs.register_warmup( |
| 83 | block_table_stride=(self.max_model_len + self.block_size - 1) |
| 84 | // self.block_size, |
| 85 | parallel_drafting_token_id=self.parallel_drafting_token_id, |
| 86 | block_size=self.block_size, |
| 87 | num_speculative_tokens=self.num_speculative_tokens, |
| 88 | ) |
| 89 | |
| 90 | @override |
| 91 | def _create_draft_vllm_config(self) -> VllmConfig: |
| 92 | base = super()._create_draft_vllm_config() |
| 93 | # The draft model is text-only — clear the target's multimodal |
| 94 | # flag so flash_attn is not rejected for mm_prefix support. |
| 95 | arch = base.model_config.model_arch_config |
| 96 | if arch.is_mm_prefix_lm: |
| 97 | base.model_config.model_arch_config = replace(arch, is_mm_prefix_lm=False) |
| 98 | return replace( |
| 99 | base, |
| 100 | attention_config=replace( |
| 101 | base.attention_config, |
| 102 | use_non_causal=not self.dflash_causal, |
| 103 | ), |
| 104 | ) |
| 105 | |
| 106 | @override |
| 107 | def _warn_if_multimodal(self): |
| 108 | # Override to allow multimodal inputs since DFlash supports Qwen3.5 models |
| 109 | pass |
| 110 | |
| 111 | @override |
| 112 | def set_inputs_first_pass( |
| 113 | self, |
| 114 | target_token_ids: torch.Tensor, |
| 115 | next_token_ids: torch.Tensor, |
| 116 | target_positions: torch.Tensor, |
| 117 | target_hidden_states: torch.Tensor, |
| 118 | token_indices_to_sample: torch.Tensor | None, |
| 119 | cad: CommonAttentionMetadata, |
| 120 | num_rejected_tokens_gpu: torch.Tensor | None, |
| 121 | ) -> tuple[int, torch.Tensor, CommonAttentionMetadata]: |
| 122 | # DFlash cross-attention: context K/V from target hidden states, |
| 123 | # Q from query embeddings (bonus + mask tokens). |
| 124 | batch_size = cad.batch_size() |
| 125 | num_context = target_token_ids.shape[0] |
| 126 | num_query_per_req = 1 + self.num_speculative_tokens |
| 127 | num_query_total = batch_size * num_query_per_req |
| 128 | |
| 129 | # Store for build_model_inputs_first_pass to use |
| 130 | self._dflash_num_context = num_context |
| 131 | |
| 132 | # We don't need to copy into a buffer here since the context preprocessing |
| 133 | # does not run in a CUDA graph |
| 134 | self._dflash_hidden_states = target_hidden_states |
| 135 | |
| 136 | token_indices_to_sample = torch.empty( |
| 137 | batch_size * self.num_speculative_tokens, |
| 138 | dtype=torch.int32, |
| 139 | device=self.device, |
| 140 | ) |
| 141 | |
| 142 | # Launch fused triton kernel for input_ids, positions, slot_mapping, |
| 143 | # and token_indices_to_sample |
| 144 | max_ctx_per_req = cad.max_query_len |
| 145 | max_tokens_per_req = max_ctx_per_req + num_query_per_req |
| 146 | BLOCK_SIZE = min(256, next_power_of_2(max_tokens_per_req)) |
| 147 | |
| 148 | has_num_rejected = num_rejected_tokens_gpu is not None |
| 149 | _copy_and_expand_dflash_inputs( |
| 150 | # Inputs |
| 151 | next_token_ids_ptr=next_token_ids, |
| 152 | target_positions_ptr=target_positions, |
| 153 | # Outputs |
| 154 | out_input_ids_ptr=self.input_ids, |
| 155 | out_context_positions_ptr=self._context_positions_buffer, |
| 156 | out_query_positions_ptr=self.positions, |
| 157 | out_context_slot_mapping_ptr=self._context_slot_mapping_buffer, |
| 158 | out_query_slot_mapping_ptr=self._slot_mapping_buffer, |
| 159 | out_token_indices_ptr=token_indices_to_sample, |
| 160 | # Block table |
| 161 | block_table_ptr=cad.block_table_tensor, |
| 162 | block_table_stride=cad.block_table_tensor.stride(0), |
| 163 | # Metadata |
| 164 | query_start_loc_ptr=cad.query_start_loc, |
| 165 | num_rejected_tokens_ptr=( |
| 166 | num_rejected_tokens_gpu if has_num_rejected else 0 |
| 167 | ), |
| 168 | # Scalars |
| 169 | parallel_drafting_token_id=self.parallel_drafting_token_id, |
| 170 | block_size=self.block_size, |
| 171 | num_query_per_req=num_query_per_req, |
| 172 | num_speculative_tokens=self.num_speculative_tokens, |
| 173 | total_input_tokens=num_context, |
| 174 | max_tokens_per_req=max_tokens_per_req, |
| 175 | triton_block_size=BLOCK_SIZE, |
| 176 | has_num_rejected=has_num_rejected, |
| 177 | ) |
| 178 | |
| 179 | query_slot_mapping = self._slot_mapping_buffer[:num_query_total] |
| 180 | new_query_start_loc = self.arange[: batch_size + 1] * num_query_per_req |
| 181 | |
| 182 | # In padded mode, cad.seq_lens includes rejected tokens. Subtract |
| 183 | # them so attention only sees the valid prefix of context states. |
| 184 | effective_seq_lens = cad.seq_lens |
| 185 | if has_num_rejected: |
| 186 | effective_seq_lens = effective_seq_lens - num_rejected_tokens_gpu |
| 187 | |
| 188 | # Skip num_rejected_tokens (GPU-only); overestimating is fine here. |
| 189 | new_seq_lens_cpu_upper_bound = ( |
| 190 | cad.seq_lens_cpu_upper_bound + num_query_per_req |
| 191 | if cad.seq_lens_cpu_upper_bound is not None |
| 192 | else None |
| 193 | ) |
| 194 | new_cad = CommonAttentionMetadata( |
| 195 | query_start_loc=new_query_start_loc, |
| 196 | seq_lens=effective_seq_lens + num_query_per_req, |
| 197 | query_start_loc_cpu=( |
| 198 | torch.from_numpy(self.token_arange_np[: batch_size + 1]).clone() |
| 199 | * num_query_per_req |
| 200 | ), |
| 201 | seq_lens_cpu_upper_bound=new_seq_lens_cpu_upper_bound, |
| 202 | num_reqs=cad.num_reqs, |
| 203 | num_actual_tokens=num_query_total, |
| 204 | max_query_len=num_query_per_req, |
| 205 | max_seq_len=cad.max_seq_len + num_query_per_req, |
| 206 | block_table_tensor=cad.block_table_tensor, |
| 207 | slot_mapping=query_slot_mapping, |
| 208 | causal=self.dflash_causal, |
| 209 | ) |
| 210 | |
| 211 | return num_query_total, token_indices_to_sample, new_cad |
| 212 | |
| 213 | @override |
| 214 | @torch.inference_mode() |
| 215 | def dummy_run( |
| 216 | self, |
| 217 | num_tokens: int, |
| 218 | use_cudagraphs: bool = True, |
| 219 | is_graph_capturing: bool = False, |
| 220 | slot_mappings: dict[str, torch.Tensor] | None = None, |
| 221 | ) -> None: |
| 222 | """Key differences to default dummy_run: |
| 223 | - Only one forward pass due to parallel drafting |
| 224 | - DFlash uses context states as unpadded metadata, so hidden_states will |
| 225 | use the unpadded num_tokens instead of num_input_tokens |
| 226 | - max_query_tokens is quite small, DFlash only sees spec tokens as queries |
| 227 | - Multimodal inputs are not currently supported |
| 228 | """ |
| 229 | num_query_tokens = min(num_tokens, self.max_query_tokens) |
| 230 | cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = ( |
| 231 | self._determine_batch_execution_and_padding( |
| 232 | num_query_tokens, use_cudagraphs=use_cudagraphs |
| 233 | ) |
| 234 | ) |
| 235 | |
| 236 | # Slot mapping sized to num_input_tokens (query only), matching |
| 237 | # the K/V tensor size from the model forward. Context KVs are |
| 238 | # pre-inserted separately and don't flow through the model. |
| 239 | if ( |
| 240 | self._draft_attn_layer_names |
| 241 | and slot_mappings is not None |
| 242 | and next(iter(self._draft_attn_layer_names)) in slot_mappings |
| 243 | ): |
| 244 | slot_mapping_dict = self._get_slot_mapping(num_input_tokens) |
| 245 | else: |
| 246 | slot_mapping_dict = slot_mappings or {} |
| 247 | |
| 248 | # Context and query positions use separate buffers; no copy needed. |
| 249 | context_positions = self._context_positions_buffer[:num_tokens] |
| 250 | # Context states will be passed directly to the precomputation without |
| 251 | # going through the buffer, since no CUDA graph is used for the precomputation. |
| 252 | # For the dummy run, we use the dummy buffer. |
| 253 | context_states = self.hidden_states[:num_tokens] |
| 254 | |
| 255 | # Run the KV projection (GEMM + norms + RoPE) for memory profiling, |
| 256 | self.model.precompute_and_store_context_kv(context_states, context_positions) |
| 257 | with set_forward_context( |
| 258 | None, |
| 259 | self.vllm_config, |
| 260 | num_tokens=num_input_tokens, |
| 261 | num_tokens_across_dp=num_tokens_across_dp, |
| 262 | cudagraph_runtime_mode=cudagraph_runtime_mode, |
| 263 | slot_mapping=slot_mapping_dict, |
| 264 | ): |
| 265 | self.model( |
| 266 | input_ids=self.input_ids[:num_input_tokens], |
| 267 | positions=self._get_positions(num_input_tokens), |
| 268 | inputs_embeds=None, |
| 269 | ) |
| 270 | |
| 271 | @override |
| 272 | def build_model_inputs_first_pass( |
| 273 | self, |
| 274 | num_tokens: int, |
| 275 | num_input_tokens: int, |
| 276 | mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None, |
| 277 | ) -> tuple[dict[str, Any], int]: |
| 278 | # Context and query positions/slots were written to separate |
| 279 | # buffers by the kernel — no copy needed. |
| 280 | num_context = self._dflash_num_context |
| 281 | |
| 282 | # Pre-insert context KVs directly into cache |
| 283 | self.model.precompute_and_store_context_kv( |
| 284 | self._dflash_hidden_states, # Shape is already [num_context, hidden_size] |
| 285 | self._context_positions_buffer[:num_context], |
| 286 | self._context_slot_mapping_buffer[:num_context], |
| 287 | ) |
| 288 | return ( |
| 289 | dict( |
| 290 | input_ids=self.input_ids[:num_input_tokens], |
| 291 | positions=self._get_positions(num_input_tokens), |
| 292 | inputs_embeds=None, |
| 293 | ), |
| 294 | num_input_tokens, |
| 295 | ) |
| 296 | |
| 297 | @override |
| 298 | def build_per_group_and_layer_attn_metadata( |
| 299 | self, cad: CommonAttentionMetadata, draft_index: int = 0 |
| 300 | ) -> tuple[list[object], dict[str, object]]: |
| 301 | per_group, per_layer = super().build_per_group_and_layer_attn_metadata( |
| 302 | cad, draft_index |
| 303 | ) |
| 304 | if not self.dflash_causal: |
| 305 | # Require all layers to support non-causal attention when required by DFlash |
| 306 | for layer_name, attn_metadata in per_layer.items(): |
| 307 | assert getattr(attn_metadata, "causal", None) is False, ( |
| 308 | f"Attention metadata for layer {layer_name} does not have" |
| 309 | " non-causal support, which is required for DFlash." |
| 310 | " Consider using a different attention backend, e.g FlashAttention." |
| 311 | ) |
| 312 | return per_group, per_layer |
| 313 | |
| 314 | @override |
| 315 | def _get_eagle3_use_aux_hidden_state_from_config(self): |
| 316 | return self.dflash_config.get("use_aux_hidden_state", True) |
| 317 | |
| 318 | @property |
| 319 | def dflash_config(self): |
| 320 | return getattr(self.draft_model_config.hf_config, "dflash_config", None) or {} |