| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | from torch import nn |
| 7 | |
| 8 | from vllm.compilation.backends import set_model_tag |
| 9 | from vllm.compilation.decorators import support_torch_compile |
| 10 | from vllm.config import CacheConfig, VllmConfig |
| 11 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 12 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 13 | from vllm.model_executor.layers.quantization.base_config import QuantizationConfig |
| 14 | from vllm.triton_utils import tl, triton |
| 15 | from vllm.utils.torch_utils import direct_register_custom_op |
| 16 | |
| 17 | from .qwen3_dflash import ( |
| 18 | DFlashQwen3DecoderLayer, |
| 19 | DFlashQwen3ForCausalLM, |
| 20 | DFlashQwen3Model, |
| 21 | ) |
| 22 | from .utils import maybe_prefix |
| 23 | |
| 24 | |
| 25 | @triton.jit |
| 26 | def _dflash2_grouped_conv_kernel( |
| 27 | x_ptr, |
| 28 | delta_ptr, |
| 29 | base_ptr, |
| 30 | output_ptr, |
| 31 | x_stride_row, |
| 32 | delta_stride_row, |
| 33 | delta_stride_tap, |
| 34 | base_stride_tap, |
| 35 | output_stride_row, |
| 36 | NUM_CHANNELS: tl.constexpr, |
| 37 | BLOCK_SIZE: tl.constexpr, |
| 38 | GROUP_SIZE: tl.constexpr, |
| 39 | TAPS: tl.constexpr, |
| 40 | ELEMENT_BLOCK: tl.constexpr, |
| 41 | ) -> None: |
| 42 | row = tl.program_id(0) // triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK) |
| 43 | col_block = tl.program_id(0) % triton.cdiv(NUM_CHANNELS, ELEMENT_BLOCK) |
| 44 | channels = col_block * ELEMENT_BLOCK + tl.arange(0, ELEMENT_BLOCK) |
| 45 | mask = channels < NUM_CHANNELS |
| 46 | groups = channels // GROUP_SIZE |
| 47 | position = row % BLOCK_SIZE |
| 48 | |
| 49 | delta_row = delta_ptr + row * delta_stride_row |
| 50 | x_row = x_ptr + row * x_stride_row |
| 51 | accumulator = ( |
| 52 | tl.load(base_ptr + channels, mask=mask, other=0.0).to(tl.float32) |
| 53 | + tl.load(delta_row + groups, mask=mask, other=0.0).to(tl.float32) |
| 54 | ) * tl.load(x_row + channels, mask=mask, other=0.0).to(tl.float32) |
| 55 | |
| 56 | for tap in tl.static_range(1, TAPS): |
| 57 | tap_mask = mask & (position >= tap) |
| 58 | coefficient = tl.load( |
| 59 | base_ptr + tap * base_stride_tap + channels, |
| 60 | mask=tap_mask, |
| 61 | other=0.0, |
| 62 | ).to(tl.float32) + tl.load( |
| 63 | delta_row + tap * delta_stride_tap + groups, |
| 64 | mask=tap_mask, |
| 65 | other=0.0, |
| 66 | ).to(tl.float32) |
| 67 | x = tl.load( |
| 68 | x_ptr + (row - tap) * x_stride_row + channels, |
| 69 | mask=tap_mask, |
| 70 | other=0.0, |
| 71 | ).to(tl.float32) |
| 72 | accumulator += coefficient * x |
| 73 | |
| 74 | tl.store( |
| 75 | output_ptr + row * output_stride_row + channels, |
| 76 | accumulator, |
| 77 | mask=mask, |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | def dflash2_grouped_conv_impl( |
| 82 | x: torch.Tensor, |
| 83 | delta: torch.Tensor, |
| 84 | base: torch.Tensor, |
| 85 | block_size: int, |
| 86 | group_size: int, |
| 87 | ) -> torch.Tensor: |
| 88 | """DFlash2 grouped convolution adapted from LightSeek TokenSpeed PR #1399.""" |
| 89 | num_rows, num_channels = x.shape |
| 90 | output = torch.empty_like(x) |
| 91 | if num_rows == 0: |
| 92 | return output |
| 93 | |
| 94 | element_block = 1024 if num_rows >= 128 and num_channels % 1024 == 0 else 512 |
| 95 | grid = (num_rows * triton.cdiv(num_channels, element_block),) |
| 96 | _dflash2_grouped_conv_kernel[grid]( |
| 97 | x, |
| 98 | delta, |
| 99 | base, |
| 100 | output, |
| 101 | x.stride(0), |
| 102 | delta.stride(0), |
| 103 | delta.stride(1), |
| 104 | base.stride(0), |
| 105 | output.stride(0), |
| 106 | NUM_CHANNELS=num_channels, |
| 107 | BLOCK_SIZE=block_size, |
| 108 | GROUP_SIZE=group_size, |
| 109 | TAPS=base.shape[0], |
| 110 | ELEMENT_BLOCK=element_block, |
| 111 | num_warps=4, |
| 112 | ) |
| 113 | return output |
| 114 | |
| 115 | |
| 116 | def dflash2_grouped_conv_fake( |
| 117 | x: torch.Tensor, |
| 118 | delta: torch.Tensor, |
| 119 | base: torch.Tensor, |
| 120 | block_size: int, |
| 121 | group_size: int, |
| 122 | ) -> torch.Tensor: |
| 123 | return torch.empty_like(x) |
| 124 | |
| 125 | |
| 126 | direct_register_custom_op( |
| 127 | op_name="dflash2_grouped_conv", |
| 128 | op_func=dflash2_grouped_conv_impl, |
| 129 | fake_impl=dflash2_grouped_conv_fake, |
| 130 | tags=(torch.Tag.needs_fixed_stride_order,), |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | def dflash2_grouped_conv( |
| 135 | x: torch.Tensor, |
| 136 | delta: torch.Tensor, |
| 137 | base: torch.Tensor, |
| 138 | block_size: int, |
| 139 | group_size: int, |
| 140 | ) -> torch.Tensor: |
| 141 | return torch.ops.vllm.dflash2_grouped_conv(x, delta, base, block_size, group_size) |
| 142 | |
| 143 | |
| 144 | def _grouped_conv( |
| 145 | hidden_states: torch.Tensor, |
| 146 | delta: torch.Tensor, |
| 147 | base: torch.Tensor, |
| 148 | block_size: int, |
| 149 | num_groups: int, |
| 150 | group_size: int, |
| 151 | taps: int, |
| 152 | ) -> torch.Tensor: |
| 153 | if hidden_states.is_cuda: |
| 154 | return dflash2_grouped_conv(hidden_states, delta, base, block_size, group_size) |
| 155 | |
| 156 | blocks = hidden_states.unflatten(-1, (num_groups, group_size)) |
| 157 | coefficients = base.view(1, taps, num_groups, group_size) + delta.unsqueeze(-1) |
| 158 | output = coefficients[:, 0] * blocks |
| 159 | position = torch.arange(hidden_states.shape[0], device=hidden_states.device) |
| 160 | if block_size & (block_size - 1) == 0: |
| 161 | position = position & (block_size - 1) |
| 162 | else: |
| 163 | position = position % block_size |
| 164 | for tap in range(1, taps): |
| 165 | shifted = F.pad(blocks[:-tap], (0, 0, 0, 0, tap, 0)) |
| 166 | output += coefficients[:, tap] * shifted * (position >= tap).view(-1, 1, 1) |
| 167 | return output.flatten(-2) |
| 168 | |
| 169 | |
| 170 | class DFlashGroupedConv(nn.Module): |
| 171 | def __init__( |
| 172 | self, |
| 173 | hidden_size: int, |
| 174 | taps: int, |
| 175 | group_size: int, |
| 176 | block_size: int, |
| 177 | params_dtype: torch.dtype, |
| 178 | prefix: str, |
| 179 | quant_config: QuantizationConfig | None = None, |
| 180 | ) -> None: |
| 181 | super().__init__() |
| 182 | if hidden_size % group_size: |
| 183 | raise ValueError( |
| 184 | f"conv_group_size={group_size} must divide hidden_size={hidden_size}." |
| 185 | ) |
| 186 | self.block_size = block_size |
| 187 | self.taps = taps |
| 188 | self.group_size = group_size |
| 189 | self.num_groups = hidden_size // group_size |
| 190 | self.base_kernel = nn.Parameter( |
| 191 | torch.empty(2, taps, hidden_size, dtype=params_dtype), |
| 192 | requires_grad=False, |
| 193 | ) |
| 194 | self.kernel_projection = ReplicatedLinear( |
| 195 | hidden_size, |
| 196 | 2 * taps * self.num_groups, |
| 197 | bias=False, |
| 198 | params_dtype=params_dtype, |
| 199 | quant_config=quant_config, |
| 200 | prefix=maybe_prefix(prefix, "kernel_projection"), |
| 201 | return_bias=False, |
| 202 | ) |
| 203 | |
| 204 | def _convolve( |
| 205 | self, hidden_states: torch.Tensor, delta: torch.Tensor, side: int |
| 206 | ) -> torch.Tensor: |
| 207 | return _grouped_conv( |
| 208 | hidden_states, |
| 209 | delta, |
| 210 | self.base_kernel[side], |
| 211 | self.block_size, |
| 212 | self.num_groups, |
| 213 | self.group_size, |
| 214 | self.taps, |
| 215 | ) |
| 216 | |
| 217 | def prepare(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 218 | coefficients = self.kernel_projection(hidden_states).reshape( |
| 219 | hidden_states.shape[0], 2, self.taps, self.num_groups |
| 220 | ) |
| 221 | return self._convolve(hidden_states, coefficients[:, 0], 0), coefficients[:, 1] |
| 222 | |
| 223 | def finish( |
| 224 | self, hidden_states: torch.Tensor, coefficients: torch.Tensor |
| 225 | ) -> torch.Tensor: |
| 226 | return self._convolve(hidden_states, coefficients, 1) |
| 227 | |
| 228 | |
| 229 | class DFlash2Qwen3DecoderLayer(DFlashQwen3DecoderLayer): |
| 230 | def __init__( |
| 231 | self, |
| 232 | vllm_config: VllmConfig, |
| 233 | *, |
| 234 | config, |
| 235 | layer_idx: int, |
| 236 | cache_config: CacheConfig | None = None, |
| 237 | quant_config: QuantizationConfig | None = None, |
| 238 | prefix: str = "", |
| 239 | ) -> None: |
| 240 | super().__init__( |
| 241 | vllm_config, |
| 242 | config=config, |
| 243 | layer_idx=layer_idx, |
| 244 | cache_config=cache_config, |
| 245 | quant_config=quant_config, |
| 246 | prefix=prefix, |
| 247 | ) |
| 248 | draft_config = config.dflash_config |
| 249 | speculative_config = vllm_config.speculative_config |
| 250 | assert speculative_config is not None |
| 251 | conv_args = dict( |
| 252 | hidden_size=config.hidden_size, |
| 253 | taps=int(draft_config["conv_kernel_size"]), |
| 254 | group_size=int(draft_config["conv_group_size"]), |
| 255 | # Query tokens per request: the bonus token plus the mask tokens. |
| 256 | block_size=1 + speculative_config.num_speculative_tokens, |
| 257 | params_dtype=vllm_config.model_config.dtype, |
| 258 | quant_config=quant_config, |
| 259 | ) |
| 260 | self.attention_conv = DFlashGroupedConv( |
| 261 | **conv_args, prefix=maybe_prefix(prefix, "attention_conv") |
| 262 | ) |
| 263 | self.mlp_conv = DFlashGroupedConv( |
| 264 | **conv_args, prefix=maybe_prefix(prefix, "mlp_conv") |
| 265 | ) |
| 266 | |
| 267 | def forward( |
| 268 | self, |
| 269 | positions: torch.Tensor, |
| 270 | hidden_states: torch.Tensor, |
| 271 | residual: torch.Tensor | None, |
| 272 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 273 | if residual is None: |
| 274 | residual = hidden_states |
| 275 | hidden_states = self.input_layernorm(hidden_states) |
| 276 | else: |
| 277 | hidden_states, residual = self.input_layernorm(hidden_states, residual) |
| 278 | |
| 279 | hidden_states, coefficients = self.attention_conv.prepare(hidden_states) |
| 280 | hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states) |
| 281 | hidden_states = self.attention_conv.finish(hidden_states, coefficients) |
| 282 | |
| 283 | hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) |
| 284 | hidden_states, coefficients = self.mlp_conv.prepare(hidden_states) |
| 285 | hidden_states = self.mlp(hidden_states) |
| 286 | hidden_states = self.mlp_conv.finish(hidden_states, coefficients) |
| 287 | return hidden_states, residual |
| 288 | |
| 289 | |
| 290 | def _score_edges( |
| 291 | predecessor_table: torch.Tensor, |
| 292 | successor_table: torch.Tensor, |
| 293 | candidate_ids: torch.Tensor, |
| 294 | unary_logits: torch.Tensor, |
| 295 | hidden: torch.Tensor, |
| 296 | anchor_token_ids: torch.Tensor, |
| 297 | top_k: int, |
| 298 | ) -> torch.Tensor: |
| 299 | successors = successor_table[candidate_ids] |
| 300 | predecessor_ids = torch.cat( |
| 301 | ( |
| 302 | anchor_token_ids[:, None, None].expand(-1, 1, top_k), |
| 303 | candidate_ids[:, :-1], |
| 304 | ), |
| 305 | dim=1, |
| 306 | ) |
| 307 | predecessors = predecessor_table[predecessor_ids] |
| 308 | return unary_logits[:, :, None] + torch.einsum( |
| 309 | "blpr,blcr->blpc", predecessors * hidden[:, :, None], successors |
| 310 | ) |
| 311 | |
| 312 | |
| 313 | @support_torch_compile |
| 314 | class CandidateSelector(nn.Module): |
| 315 | def __init__( |
| 316 | self, |
| 317 | hidden_size: int, |
| 318 | vocab_size: int, |
| 319 | rank: int, |
| 320 | top_k: int, |
| 321 | params_dtype: torch.dtype, |
| 322 | prefix: str, |
| 323 | ) -> None: |
| 324 | super().__init__() |
| 325 | self.top_k = top_k |
| 326 | self.predecessor_codebook = nn.Parameter( |
| 327 | torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False |
| 328 | ) |
| 329 | self.successor_codebook = nn.Parameter( |
| 330 | torch.empty(vocab_size, rank, dtype=params_dtype), requires_grad=False |
| 331 | ) |
| 332 | self.hidden_projection = ReplicatedLinear( |
| 333 | hidden_size, |
| 334 | rank, |
| 335 | bias=False, |
| 336 | params_dtype=params_dtype, |
| 337 | quant_config=None, |
| 338 | prefix=maybe_prefix(prefix, "hidden_projection"), |
| 339 | return_bias=False, |
| 340 | ) |
| 341 | |
| 342 | def forward( |
| 343 | self, |
| 344 | candidate_ids: torch.Tensor, |
| 345 | unary_logits: torch.Tensor, |
| 346 | hidden_states: torch.Tensor, |
| 347 | anchor_token_ids: torch.Tensor, |
| 348 | ) -> torch.Tensor: |
| 349 | hidden = self.hidden_projection(hidden_states) |
| 350 | return _score_edges( |
| 351 | self.predecessor_codebook, |
| 352 | self.successor_codebook, |
| 353 | candidate_ids, |
| 354 | unary_logits, |
| 355 | hidden, |
| 356 | anchor_token_ids, |
| 357 | self.top_k, |
| 358 | ) |
| 359 | |
| 360 | |
| 361 | class DFlash2Qwen3Model(DFlashQwen3Model): |
| 362 | decoder_layer_cls = DFlash2Qwen3DecoderLayer |
| 363 | |
| 364 | def __init__( |
| 365 | self, |
| 366 | *, |
| 367 | vllm_config: VllmConfig, |
| 368 | start_layer_id: int = 0, |
| 369 | prefix: str = "", |
| 370 | ) -> None: |
| 371 | super().__init__( |
| 372 | vllm_config=vllm_config, |
| 373 | start_layer_id=start_layer_id, |
| 374 | prefix=prefix, |
| 375 | ) |
| 376 | draft_config = self.config.dflash_config |
| 377 | self.input_embedding_scale = float( |
| 378 | draft_config.get("input_embedding_scale", 1.0) |
| 379 | ) |
| 380 | # Without its own tag the selector shares the draft head's compile cache. |
| 381 | with set_model_tag("dflash2_candidate_selector"): |
| 382 | self.candidate_selector = CandidateSelector( |
| 383 | hidden_size=self.config.hidden_size, |
| 384 | vocab_size=self.config.vocab_size, |
| 385 | rank=int(draft_config["selector_rank"]), |
| 386 | top_k=int(draft_config["selector_top_k"]), |
| 387 | params_dtype=vllm_config.model_config.dtype, |
| 388 | prefix=maybe_prefix(prefix, "candidate_selector"), |
| 389 | ) |
| 390 | |
| 391 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 392 | return super().embed_input_ids(input_ids) * self.input_embedding_scale |
| 393 | |
| 394 | |
| 395 | class DFlash2Qwen3ForCausalLM(DFlashQwen3ForCausalLM): |
| 396 | model_cls = DFlash2Qwen3Model |
| 397 | |
| 398 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 399 | super().__init__(vllm_config=vllm_config, prefix=prefix) |
| 400 | draft_config = self.config.dflash_config |
| 401 | softcap = float(draft_config.get("final_logit_softcapping") or 0.0) |
| 402 | self.candidate_logits_processor = LogitsProcessor( |
| 403 | vllm_config.model_config.get_vocab_size(), |
| 404 | scale=float(draft_config.get("output_multiplier", 1.0)), |
| 405 | soft_cap=softcap if softcap > 0 else None, |
| 406 | ) |
| 407 | |
| 408 | def compute_candidates( |
| 409 | self, hidden_states: torch.Tensor |
| 410 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 411 | return self.candidate_logits_processor.get_top_k_tokens( |
| 412 | self.lm_head, hidden_states, self.model.candidate_selector.top_k |
| 413 | ) |
| 414 | |
| 415 | |
| 416 | EntryClass = DFlash2Qwen3ForCausalLM |