| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GPU-resident Qwen4Exp position-learning enhancement layers.""" |
| 4 | |
| 5 | import math |
| 6 | from collections.abc import Iterable, Sequence |
| 7 | |
| 8 | import torch |
| 9 | import torch.nn.functional as F |
| 10 | from torch import nn |
| 11 | from transformers import Qwen4ExpTextConfig |
| 12 | |
| 13 | from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config |
| 14 | from vllm.forward_context import get_forward_context |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 17 | from vllm.model_executor.layers.mamba.abstract import MambaBase |
| 18 | from vllm.model_executor.layers.mamba.mamba_utils import ( |
| 19 | MambaStateDtypeCalculator, |
| 20 | MambaStateShapeCalculator, |
| 21 | is_conv_state_dim_first, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.quantization.base_config import ( |
| 24 | QuantizationConfig, |
| 25 | ) |
| 26 | from vllm.model_executor.models.utils import AutoWeightsLoader |
| 27 | from vllm.utils.torch_utils import direct_register_custom_op |
| 28 | from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum |
| 29 | from vllm.v1.attention.backends.short_conv_attn import ( |
| 30 | PleShortConvAttentionBackend, |
| 31 | PleShortConvAttentionMetadata, |
| 32 | ) |
| 33 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 34 | |
| 35 | from ..common.ngram_embedding import ( |
| 36 | Qwen4ExpPLEDeviceEmbedding, |
| 37 | Qwen4ExpPLEEmbeddingMethod, |
| 38 | Qwen4ExpPLEPinnedHostEmbedding, |
| 39 | ) |
| 40 | |
| 41 | logger = init_logger(__name__) |
| 42 | |
| 43 | |
| 44 | class Qwen4ExpPLEGroupedNorm(nn.Module): |
| 45 | def __init__( |
| 46 | self, |
| 47 | hidden_size: int, |
| 48 | eps: float, |
| 49 | group_size: int | None, |
| 50 | dtype: torch.dtype | None, |
| 51 | ) -> None: |
| 52 | super().__init__() |
| 53 | if group_size is not None and hidden_size % group_size: |
| 54 | raise ValueError( |
| 55 | f"hidden_size ({hidden_size}) must be divisible by " |
| 56 | f"group_size ({group_size})" |
| 57 | ) |
| 58 | self.eps = eps |
| 59 | self.group_size = group_size |
| 60 | self.weight = nn.Parameter(torch.zeros(hidden_size, dtype=dtype)) |
| 61 | |
| 62 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 63 | input_dtype = hidden_states.dtype |
| 64 | hidden_states = hidden_states.float() |
| 65 | if self.group_size is None: |
| 66 | variance = hidden_states.square().mean(dim=-1, keepdim=True) |
| 67 | normalized = hidden_states * torch.rsqrt(variance + self.eps) |
| 68 | else: |
| 69 | grouped = hidden_states.unflatten( |
| 70 | -1, (hidden_states.shape[-1] // self.group_size, self.group_size) |
| 71 | ) |
| 72 | variance = grouped.square().mean(dim=-1, keepdim=True) |
| 73 | normalized = (grouped * torch.rsqrt(variance + self.eps)).flatten(-2) |
| 74 | return (normalized * (1.0 + self.weight.float())).to(input_dtype) |
| 75 | |
| 76 | |
| 77 | class Qwen4ExpNGramEmbedding(nn.Module): |
| 78 | _MASK64 = (1 << 64) - 1 |
| 79 | _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 |
| 80 | _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 |
| 81 | _SPLITMIX_M2 = 0x94D049BB133111EB |
| 82 | _PLE_LAYER_PRIME = 10007 |
| 83 | |
| 84 | @classmethod |
| 85 | def _splitmix64(cls, value: int) -> int: |
| 86 | """Mix an integer into a deterministic unsigned 64-bit value.""" |
| 87 | value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64 |
| 88 | value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64 |
| 89 | value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64 |
| 90 | return (value ^ (value >> 31)) & cls._MASK64 |
| 91 | |
| 92 | @staticmethod |
| 93 | def _is_prime_64(value: int) -> bool: |
| 94 | """Return whether a 64-bit integer is prime.""" |
| 95 | if value < 2: |
| 96 | return False |
| 97 | for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): |
| 98 | if value % prime == 0: |
| 99 | return value == prime |
| 100 | exponent = value - 1 |
| 101 | shifts = 0 |
| 102 | while exponent % 2 == 0: |
| 103 | exponent //= 2 |
| 104 | shifts += 1 |
| 105 | for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): |
| 106 | if base % value == 0: |
| 107 | continue |
| 108 | witness = pow(base, exponent, value) |
| 109 | if witness in (1, value - 1): |
| 110 | continue |
| 111 | for _ in range(shifts - 1): |
| 112 | witness = pow(witness, 2, value) |
| 113 | if witness == value - 1: |
| 114 | break |
| 115 | else: |
| 116 | return False |
| 117 | return True |
| 118 | |
| 119 | @classmethod |
| 120 | def _nth_prime_after(cls, start: int, count: int) -> int: |
| 121 | """Return the ``count``-th prime strictly greater than ``start``.""" |
| 122 | prime = int(start) |
| 123 | for _ in range(count): |
| 124 | candidate = prime + 1 |
| 125 | if candidate <= 2: |
| 126 | prime = 2 |
| 127 | continue |
| 128 | if candidate % 2 == 0: |
| 129 | candidate += 1 |
| 130 | while not cls._is_prime_64(candidate): |
| 131 | candidate += 2 |
| 132 | prime = candidate |
| 133 | return prime |
| 134 | |
| 135 | @classmethod |
| 136 | def _make_layer_multipliers( |
| 137 | cls, |
| 138 | *, |
| 139 | ngram_size: int, |
| 140 | unigram_vocab_size: int, |
| 141 | seed: int, |
| 142 | ple_dense_layer_id: int, |
| 143 | ) -> list[int]: |
| 144 | """Build deterministic hash multipliers for one PLE layer.""" |
| 145 | max_multiplier = ((1 << 63) - 1) // unigram_vocab_size |
| 146 | half_bound = max(1, max_multiplier // 2) |
| 147 | base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id |
| 148 | multipliers = [] |
| 149 | for index in range(ngram_size): |
| 150 | value = base_seed + cls._SPLITMIX_GAMMA * (index + 1) |
| 151 | multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1) |
| 152 | return multipliers |
| 153 | |
| 154 | @classmethod |
| 155 | def _make_vocab_layout( |
| 156 | cls, |
| 157 | *, |
| 158 | ngram_vocab_size_base: int, |
| 159 | ngram_heads: int, |
| 160 | ple_dense_layer_id: int, |
| 161 | ) -> tuple[list[int], list[int], int]: |
| 162 | """Build per-head vocabulary sizes, offsets, and total row count.""" |
| 163 | sizes: list[int] = [] |
| 164 | offsets: list[int] = [] |
| 165 | offset = 0 |
| 166 | for local_head in range(ngram_heads): |
| 167 | global_head = ple_dense_layer_id * ngram_heads + local_head |
| 168 | size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) |
| 169 | sizes.append(size) |
| 170 | offsets.append(offset) |
| 171 | offset += size |
| 172 | return sizes, offsets, offset |
| 173 | |
| 174 | def __init__( |
| 175 | self, |
| 176 | config: Qwen4ExpTextConfig, |
| 177 | embedding_dim: int, |
| 178 | ple_dense_layer_id: int, |
| 179 | max_total_tokens: int, |
| 180 | max_num_reqs: int, |
| 181 | prefix: str, |
| 182 | layer_name: str, |
| 183 | *, |
| 184 | data_parallel_rank: int = 0, |
| 185 | quant_config: QuantizationConfig | None = None, |
| 186 | params_dtype: torch.dtype | None = None, |
| 187 | ) -> None: |
| 188 | super().__init__() |
| 189 | self.embedding_dim = embedding_dim |
| 190 | self.layer_name = layer_name |
| 191 | self.ngram_size = int(config.ngram_size) |
| 192 | self.heads_per_ngram = int(config.heads_per_ngram) |
| 193 | self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram |
| 194 | if self.ngram_size < 2: |
| 195 | raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") |
| 196 | if self.heads_per_ngram <= 0: |
| 197 | raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") |
| 198 | if embedding_dim % self.ngram_heads: |
| 199 | raise ValueError( |
| 200 | "ple_embed_dim must be divisible by total ngram heads: " |
| 201 | f"{embedding_dim} % {self.ngram_heads} != 0" |
| 202 | ) |
| 203 | self.head_dim = embedding_dim // self.ngram_heads |
| 204 | self.eos_token_id = int(config.eos_token_id) |
| 205 | self.unigram_vocab_size = int(config.vocab_size) |
| 206 | self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) |
| 207 | if self.split_ngram_parts <= 0: |
| 208 | raise ValueError("split_ngram_parts must be positive") |
| 209 | |
| 210 | multipliers = self._make_layer_multipliers( |
| 211 | ngram_size=self.ngram_size, |
| 212 | unigram_vocab_size=self.unigram_vocab_size, |
| 213 | seed=int(getattr(config, "seed", 1234)), |
| 214 | ple_dense_layer_id=ple_dense_layer_id, |
| 215 | ) |
| 216 | self.register_buffer( |
| 217 | "layer_multipliers", |
| 218 | torch.tensor(multipliers, dtype=torch.long), |
| 219 | persistent=True, |
| 220 | ) |
| 221 | |
| 222 | sizes, offsets, total_vocab_size = self._make_vocab_layout( |
| 223 | ngram_vocab_size_base=int(config.ngram_vocab_size_base), |
| 224 | ngram_heads=self.ngram_heads, |
| 225 | ple_dense_layer_id=ple_dense_layer_id, |
| 226 | ) |
| 227 | self.register_buffer( |
| 228 | "ngram_heads_vocab_sizes", |
| 229 | torch.tensor(sizes, dtype=torch.long), |
| 230 | persistent=True, |
| 231 | ) |
| 232 | self.register_buffer( |
| 233 | "ngram_heads_offsets", |
| 234 | torch.tensor(offsets, dtype=torch.long), |
| 235 | persistent=True, |
| 236 | ) |
| 237 | divisor = int(config.make_ngram_vocab_size_divisible_by) |
| 238 | padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor |
| 239 | if params_dtype is None: |
| 240 | params_dtype = torch.get_default_dtype() |
| 241 | embedding_prefix = f"{prefix}.ngram_embedding" |
| 242 | embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( |
| 243 | quant_config, |
| 244 | embedding_prefix, |
| 245 | getattr(config, "ple_embedding_dtype", None), |
| 246 | ) |
| 247 | engram_config = get_current_vllm_config().engram_config |
| 248 | embedding_cls = ( |
| 249 | Qwen4ExpPLEPinnedHostEmbedding |
| 250 | if engram_config is not None and engram_config.cpu_offload |
| 251 | else Qwen4ExpPLEDeviceEmbedding |
| 252 | ) |
| 253 | self.ngram_embedding = embedding_cls( |
| 254 | padded_vocab_size, |
| 255 | self.head_dim, |
| 256 | params_dtype=params_dtype, |
| 257 | padding_size=divisor, |
| 258 | prefix=embedding_prefix, |
| 259 | embedding_method=embedding_quant_method, |
| 260 | num_ngram_heads=self.ngram_heads, |
| 261 | max_total_tokens=max_total_tokens, |
| 262 | data_parallel_rank=data_parallel_rank, |
| 263 | ) |
| 264 | logger.info( |
| 265 | "Initialized AMD PLE embedding %s: quantization_method=%s, " |
| 266 | "weight_dtype=%s, weight_device=%s, pinned=%s", |
| 267 | embedding_prefix, |
| 268 | type(embedding_quant_method).__name__, |
| 269 | self.ngram_embedding.weight.dtype, |
| 270 | self.ngram_embedding.weight.device, |
| 271 | self.ngram_embedding.weight.is_pinned(), |
| 272 | ) |
| 273 | self.register_buffer( |
| 274 | "positions_buffer", |
| 275 | torch.arange(max_total_tokens, dtype=torch.int64), |
| 276 | persistent=False, |
| 277 | ) |
| 278 | self.register_buffer( |
| 279 | "padded_buffer", |
| 280 | torch.full( |
| 281 | (max_num_reqs, max_total_tokens), |
| 282 | self.eos_token_id, |
| 283 | dtype=torch.int64, |
| 284 | ), |
| 285 | persistent=False, |
| 286 | ) |
| 287 | |
| 288 | @staticmethod |
| 289 | def _shift_precompute( |
| 290 | tokens: torch.Tensor, eos_token_id: int |
| 291 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 292 | if tokens.dim() != 2: |
| 293 | raise ValueError("tokens must be a 2D tensor") |
| 294 | batch_size, seq_len = tokens.shape |
| 295 | positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) |
| 296 | eos_positions = torch.where(tokens == eos_token_id, positions, -1) |
| 297 | previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values |
| 298 | previous_eos = torch.cat( |
| 299 | [ |
| 300 | eos_positions.new_full((batch_size, 1), -1), |
| 301 | previous_eos_inclusive[:, :-1], |
| 302 | ], |
| 303 | dim=1, |
| 304 | ) |
| 305 | return positions, positions.unsqueeze(0) - previous_eos - 1 |
| 306 | |
| 307 | @staticmethod |
| 308 | def _shift_apply( |
| 309 | tokens: torch.Tensor, |
| 310 | positions: torch.Tensor, |
| 311 | position_in_segment: torch.Tensor, |
| 312 | shift: int, |
| 313 | eos_token_id: int, |
| 314 | ) -> torch.Tensor: |
| 315 | if shift == 0: |
| 316 | return tokens |
| 317 | source = positions - shift |
| 318 | gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) |
| 319 | shifted = tokens.gather(1, gather_indices) |
| 320 | valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) |
| 321 | return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) |
| 322 | |
| 323 | def forward( |
| 324 | self, |
| 325 | input_ids: torch.Tensor, |
| 326 | query_start_loc: torch.Tensor, |
| 327 | ngram_context: torch.Tensor, |
| 328 | ) -> torch.Tensor: |
| 329 | input_ids = input_ids.reshape(-1).long() |
| 330 | query_start_loc = query_start_loc.long() |
| 331 | num_reqs = query_start_loc.numel() - 1 |
| 332 | num_tokens = input_ids.shape[0] |
| 333 | if num_tokens > self.positions_buffer.numel(): |
| 334 | raise ValueError( |
| 335 | f"PLE received {num_tokens} tokens, but its workspace supports " |
| 336 | f"at most {self.positions_buffer.numel()}" |
| 337 | ) |
| 338 | if num_reqs > self.padded_buffer.shape[0]: |
| 339 | raise ValueError( |
| 340 | f"PLE received {num_reqs} requests, but its workspace supports " |
| 341 | f"at most {self.padded_buffer.shape[0]}" |
| 342 | ) |
| 343 | |
| 344 | positions = self.positions_buffer[:num_tokens] |
| 345 | packed = self.padded_buffer[:num_reqs] |
| 346 | packed.fill_(self.eos_token_id) |
| 347 | request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 |
| 348 | request_indices.clamp_(max=num_reqs - 1) |
| 349 | columns = (positions - query_start_loc[request_indices]).clamp( |
| 350 | 0, packed.shape[1] - 1 |
| 351 | ) |
| 352 | packed[request_indices, columns] = input_ids |
| 353 | ngram_context = ngram_context[:num_reqs].to( |
| 354 | device=input_ids.device, dtype=torch.long |
| 355 | ) |
| 356 | |
| 357 | context = torch.cat([ngram_context, packed], dim=-1) |
| 358 | positions_2d, position_in_segment = self._shift_precompute( |
| 359 | context, self.eos_token_id |
| 360 | ) |
| 361 | shifted = [context] |
| 362 | for shift in range(1, self.ngram_size): |
| 363 | shifted.append( |
| 364 | self._shift_apply( |
| 365 | context, |
| 366 | positions_2d, |
| 367 | position_in_segment, |
| 368 | shift, |
| 369 | self.eos_token_id, |
| 370 | ) |
| 371 | ) |
| 372 | adjusted_columns = columns + self.ngram_size - 1 |
| 373 | id_blocks = [] |
| 374 | for ngram in range(2, self.ngram_size + 1): |
| 375 | start = (ngram - 2) * self.heads_per_ngram |
| 376 | end = start + self.heads_per_ngram |
| 377 | mixed = shifted[0] * self.layer_multipliers[0] |
| 378 | for index in range(1, ngram): |
| 379 | mixed = torch.bitwise_xor( |
| 380 | mixed, shifted[index] * self.layer_multipliers[index] |
| 381 | ) |
| 382 | sizes = self.ngram_heads_vocab_sizes[start:end] |
| 383 | offsets = self.ngram_heads_offsets[start:end] |
| 384 | ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets |
| 385 | id_blocks.append(ids[request_indices, adjusted_columns]) |
| 386 | ngram_ids = torch.cat(id_blocks, dim=-1) |
| 387 | embedding = self.ngram_embedding |
| 388 | if embedding.supports_prefetch: |
| 389 | output = ngram_ids.new_empty( |
| 390 | (ngram_ids.shape[0], self.embedding_dim), |
| 391 | dtype=embedding.weight.dtype, |
| 392 | ) |
| 393 | torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding_pinned( |
| 394 | ngram_ids, |
| 395 | output, |
| 396 | self.layer_name, |
| 397 | ) |
| 398 | return output |
| 399 | output = ngram_ids.new_empty( |
| 400 | (ngram_ids.shape[0], self.embedding_dim), |
| 401 | dtype=embedding.params_dtype, |
| 402 | ) |
| 403 | torch.ops.vllm.qwen4_exp_amd_ple_ngram_embedding( |
| 404 | ngram_ids, |
| 405 | output, |
| 406 | self.layer_name, |
| 407 | ) |
| 408 | return output |
| 409 | |
| 410 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 411 | """Load hash buffers and checkpoint-split embedding rows.""" |
| 412 | persistent_buffers = { |
| 413 | "layer_multipliers": self.layer_multipliers, |
| 414 | "ngram_heads_offsets": self.ngram_heads_offsets, |
| 415 | "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, |
| 416 | } |
| 417 | loaded: set[str] = set() |
| 418 | regular_weights: list[tuple[str, torch.Tensor]] = [] |
| 419 | shard_prefix = "ngram_embedding.shard_" |
| 420 | |
| 421 | for name, loaded_weight in weights: |
| 422 | leaf_name = name.rsplit(".", 1)[-1] |
| 423 | if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": |
| 424 | continue |
| 425 | if name in persistent_buffers: |
| 426 | buffer = persistent_buffers[name] |
| 427 | if buffer.shape != loaded_weight.shape: |
| 428 | raise ValueError( |
| 429 | f"Shape mismatch for {name}: expected " |
| 430 | f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" |
| 431 | ) |
| 432 | buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) |
| 433 | loaded.add(name) |
| 434 | continue |
| 435 | if name.startswith(shard_prefix) and name.endswith(".weight"): |
| 436 | shard_text = name[len(shard_prefix) : -len(".weight")] |
| 437 | if not shard_text.isdigit(): |
| 438 | regular_weights.append((name, loaded_weight)) |
| 439 | continue |
| 440 | shard_index = int(shard_text) |
| 441 | if shard_index >= self.split_ngram_parts: |
| 442 | raise ValueError( |
| 443 | f"PLE embedding shard index {shard_index} exceeds " |
| 444 | f"split_ngram_parts={self.split_ngram_parts}" |
| 445 | ) |
| 446 | embedding = self.ngram_embedding |
| 447 | shard_size = ( |
| 448 | embedding.org_vocab_size + self.split_ngram_parts - 1 |
| 449 | ) // self.split_ngram_parts |
| 450 | checkpoint_start = shard_index * shard_size |
| 451 | expected_rows = max( |
| 452 | 0, |
| 453 | min(shard_size, embedding.org_vocab_size - checkpoint_start), |
| 454 | ) |
| 455 | expected_shape = (expected_rows, embedding.embedding_dim) |
| 456 | if tuple(loaded_weight.shape) != expected_shape: |
| 457 | raise ValueError( |
| 458 | f"Shape mismatch for PLE embedding shard {shard_index}: " |
| 459 | f"expected {expected_shape}, got " |
| 460 | f"{tuple(loaded_weight.shape)}" |
| 461 | ) |
| 462 | embedding.weight.weight_loader( |
| 463 | embedding.weight, |
| 464 | loaded_weight, |
| 465 | checkpoint_start=checkpoint_start, |
| 466 | ) |
| 467 | loaded.add("ngram_embedding.weight") |
| 468 | continue |
| 469 | regular_weights.append((name, loaded_weight)) |
| 470 | |
| 471 | if regular_weights: |
| 472 | loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) |
| 473 | return loaded |
| 474 | |
| 475 | |
| 476 | class Qwen4ExpPLELayer(nn.Module, MambaBase): |
| 477 | def __init__( |
| 478 | self, |
| 479 | config: Qwen4ExpTextConfig, |
| 480 | vllm_config: VllmConfig, |
| 481 | layer_idx: int = 0, |
| 482 | ple_dense_layer_id: int | None = None, |
| 483 | prefix: str = "", |
| 484 | ) -> None: |
| 485 | super().__init__() |
| 486 | model_config = vllm_config.model_config |
| 487 | cache_config = vllm_config.cache_config |
| 488 | quant_config = vllm_config.quant_config |
| 489 | self.model_config: ModelConfig = model_config |
| 490 | self.cache_config: CacheConfig = cache_config |
| 491 | self.layer_idx = layer_idx |
| 492 | self.ple_dense_layer_id = ( |
| 493 | int(ple_dense_layer_id) |
| 494 | if ple_dense_layer_id is not None |
| 495 | else int(layer_idx) |
| 496 | ) |
| 497 | self.prefix = prefix |
| 498 | self.hidden_size = int(config.hidden_size) |
| 499 | self.hc_count = config.hc_count |
| 500 | self.hc_hidden_size = self.hidden_size * self.hc_count |
| 501 | self.conv_kernel_size = int(config.ple_conv_kernel_size) |
| 502 | self.short_conv_dilation = int(config.ngram_size) |
| 503 | self.conv_state_len = (self.conv_kernel_size - 1) * self.short_conv_dilation |
| 504 | self.num_spec_tokens = vllm_config.num_speculative_tokens |
| 505 | self.activation = "silu" |
| 506 | self.ple_embedding: nn.Module = Qwen4ExpNGramEmbedding( |
| 507 | config, |
| 508 | int(config.ple_embed_dim), |
| 509 | self.ple_dense_layer_id, |
| 510 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 511 | vllm_config.scheduler_config.max_num_seqs, |
| 512 | f"{prefix}.ple_embedding", |
| 513 | prefix, |
| 514 | data_parallel_rank=vllm_config.parallel_config.data_parallel_rank, |
| 515 | quant_config=quant_config, |
| 516 | params_dtype=model_config.dtype, |
| 517 | ) |
| 518 | self.key_proj = ReplicatedLinear( |
| 519 | int(config.ple_embed_dim), |
| 520 | self.hc_hidden_size, |
| 521 | bias=False, |
| 522 | quant_config=quant_config, |
| 523 | prefix=f"{prefix}.key_proj", |
| 524 | ) |
| 525 | self.value_proj = ReplicatedLinear( |
| 526 | int(config.ple_embed_dim), |
| 527 | self.hidden_size, |
| 528 | bias=False, |
| 529 | quant_config=quant_config, |
| 530 | prefix=f"{prefix}.value_proj", |
| 531 | ) |
| 532 | norm_args = ( |
| 533 | self.hc_hidden_size, |
| 534 | config.rms_norm_eps, |
| 535 | self.hidden_size, |
| 536 | model_config.dtype, |
| 537 | ) |
| 538 | self.norm_key = Qwen4ExpPLEGroupedNorm(*norm_args) |
| 539 | self.norm_query = Qwen4ExpPLEGroupedNorm(*norm_args) |
| 540 | self.norm_conv = Qwen4ExpPLEGroupedNorm(*norm_args) |
| 541 | self.conv1d = nn.Conv1d( |
| 542 | self.hc_hidden_size, |
| 543 | self.hc_hidden_size, |
| 544 | self.conv_kernel_size, |
| 545 | groups=self.hc_hidden_size, |
| 546 | padding=self.conv_state_len, |
| 547 | dilation=self.short_conv_dilation, |
| 548 | bias=False, |
| 549 | dtype=model_config.dtype, |
| 550 | ) |
| 551 | nn.init.zeros_(self.conv1d.weight) |
| 552 | self.conv1d.weight._no_reinit = True |
| 553 | self.kv_cache = (torch.tensor([]),) |
| 554 | compilation_config = get_current_vllm_config().compilation_config |
| 555 | if prefix in compilation_config.static_forward_context: |
| 556 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 557 | compilation_config.static_forward_context[prefix] = self |
| 558 | |
| 559 | @property |
| 560 | def mamba_type(self) -> MambaAttentionBackendEnum: |
| 561 | return MambaAttentionBackendEnum.SHORT_CONV |
| 562 | |
| 563 | @property |
| 564 | def is_kv_cache_tp_replicated(self) -> bool: |
| 565 | return True |
| 566 | |
| 567 | def get_attn_backend(self) -> type[PleShortConvAttentionBackend]: |
| 568 | return PleShortConvAttentionBackend |
| 569 | |
| 570 | def get_state_dtype(self) -> tuple[torch.dtype, ...]: |
| 571 | return MambaStateDtypeCalculator.short_conv_state_dtype( |
| 572 | self.model_config.dtype, self.cache_config.mamba_cache_dtype |
| 573 | ) |
| 574 | |
| 575 | def get_state_shape(self) -> Sequence[tuple[int, ...]]: |
| 576 | return MambaStateShapeCalculator.short_conv_state_shape( |
| 577 | tp_world_size=1, |
| 578 | intermediate_size=self.hc_hidden_size, |
| 579 | conv_kernel=self.conv_state_len + 1, |
| 580 | num_spec=self.num_spec_tokens, |
| 581 | ) |
| 582 | |
| 583 | def _apply_norm( |
| 584 | self, norm: Qwen4ExpPLEGroupedNorm, hidden_states: torch.Tensor |
| 585 | ) -> torch.Tensor: |
| 586 | shape = hidden_states.shape |
| 587 | return norm(hidden_states.flatten(-2)).reshape(shape) |
| 588 | |
| 589 | def _short_conv_fallback(self, inputs: torch.Tensor) -> torch.Tensor: |
| 590 | # Profiling / CUDA graph capture only; conv state is not updated. |
| 591 | inputs_t = inputs.transpose(0, 1).unsqueeze(0) |
| 592 | output = self.conv1d(inputs_t)[..., : inputs_t.size(-1)] |
| 593 | return F.silu(output).squeeze(0).transpose(0, 1) |
| 594 | |
| 595 | def _short_conv_dilated_decode_batched( |
| 596 | self, |
| 597 | x_d: torch.Tensor, |
| 598 | conv_state: torch.Tensor, |
| 599 | conv_weights: torch.Tensor, |
| 600 | state_indices_tensor_d: torch.Tensor, |
| 601 | has_initial_states_d: torch.Tensor | None, |
| 602 | ) -> torch.Tensor: |
| 603 | state_indices = state_indices_tensor_d.to( |
| 604 | device=conv_state.device, dtype=torch.int64 |
| 605 | ) |
| 606 | # FULL cudagraph padded decode rows use NULL_BLOCK_ID. Remap them to |
| 607 | # slot 0 for a safe gather, then zero output and skip write-back. |
| 608 | valid_state = state_indices != NULL_BLOCK_ID |
| 609 | state_indices = torch.where( |
| 610 | valid_state, state_indices, torch.zeros_like(state_indices) |
| 611 | ) |
| 612 | if has_initial_states_d is None: |
| 613 | has_initial_state = valid_state |
| 614 | else: |
| 615 | if has_initial_states_d.numel() < state_indices_tensor_d.numel(): |
| 616 | raise ValueError( |
| 617 | "has_initial_states_d size mismatch: " |
| 618 | f"got {has_initial_states_d.numel()}, " |
| 619 | f"need >= {state_indices_tensor_d.numel()}." |
| 620 | ) |
| 621 | has_initial_state = has_initial_states_d[ |
| 622 | : state_indices_tensor_d.numel() |
| 623 | ].to(device=conv_state.device, dtype=torch.bool) |
| 624 | has_initial_state = has_initial_state & valid_state |
| 625 | |
| 626 | cached_state = conv_state.index_select(0, state_indices) |
| 627 | state = cached_state[..., : self.conv_state_len].to(x_d.dtype) |
| 628 | if self.conv_state_len > 0: |
| 629 | initial_state = torch.where( |
| 630 | has_initial_state.view(-1, 1, 1), |
| 631 | state, |
| 632 | torch.zeros_like(state), |
| 633 | ) |
| 634 | history = torch.cat((initial_state, x_d.unsqueeze(-1)), dim=-1) |
| 635 | else: |
| 636 | history = x_d.unsqueeze(-1) |
| 637 | |
| 638 | conv_output = F.conv1d( |
| 639 | history, |
| 640 | conv_weights.unsqueeze(1).contiguous(), |
| 641 | groups=history.size(1), |
| 642 | dilation=self.short_conv_dilation, |
| 643 | ).squeeze(-1) |
| 644 | output = F.silu(conv_output) |
| 645 | output = output * valid_state.view(-1, 1).to(output.dtype) |
| 646 | |
| 647 | if self.conv_state_len > 0: |
| 648 | next_state = history[..., -self.conv_state_len :] |
| 649 | # Padded rows are remapped to the reserved null slot. Preserve its |
| 650 | # existing value while writing the new states for valid rows. |
| 651 | existing_base_state = cached_state[..., : self.conv_state_len] |
| 652 | safe_next_state = torch.where( |
| 653 | valid_state.view(-1, 1, 1), |
| 654 | next_state.to(conv_state.dtype), |
| 655 | existing_base_state, |
| 656 | ) |
| 657 | cached_state[..., : self.conv_state_len] = safe_next_state |
| 658 | conv_state.index_copy_(0, state_indices, cached_state) |
| 659 | |
| 660 | return output |
| 661 | |
| 662 | def _short_conv_dilated_prefill_batched( |
| 663 | self, |
| 664 | x_p: torch.Tensor, |
| 665 | metadata: PleShortConvAttentionMetadata, |
| 666 | conv_state: torch.Tensor, |
| 667 | conv_weights: torch.Tensor, |
| 668 | state_indices_tensor_p: torch.Tensor, |
| 669 | num_prefills: int, |
| 670 | num_decode_tokens: int, |
| 671 | num_prefill_tokens: int, |
| 672 | ) -> torch.Tensor: |
| 673 | query_start_loc_p = metadata.query_start_loc_p |
| 674 | if query_start_loc_p is None: |
| 675 | raise ValueError("query_start_loc is required for prefill short-conv") |
| 676 | # The metadata builder guarantees that the prefill query offsets start |
| 677 | # at 0 and end at num_prefill_tokens. Avoid reading those values here, |
| 678 | # since doing so would force a device-to-host synchronization. |
| 679 | has_initial_states_p = metadata.has_initial_states_p |
| 680 | if has_initial_states_p is None: |
| 681 | raise ValueError("has_initial_states_p is required for prefill short-conv") |
| 682 | |
| 683 | output = torch.empty_like(x_p) |
| 684 | q_starts = query_start_loc_p.to(torch.int64) |
| 685 | if state_indices_tensor_p.numel() < num_prefills: |
| 686 | raise ValueError( |
| 687 | "state_indices_tensor_p size mismatch: " |
| 688 | f"got {state_indices_tensor_p.numel()}, " |
| 689 | f"need >= {num_prefills}." |
| 690 | ) |
| 691 | if has_initial_states_p.numel() < num_prefills: |
| 692 | raise ValueError( |
| 693 | "has_initial_states_p size mismatch: " |
| 694 | f"got {has_initial_states_p.numel()}, " |
| 695 | f"need >= {num_prefills}." |
| 696 | ) |
| 697 | if num_prefills == 0 or x_p.numel() == 0: |
| 698 | return output |
| 699 | lengths = q_starts[1:] - q_starts[:-1] |
| 700 | # Use the CPU-computed packing width from the metadata builder instead |
| 701 | # of synchronizing on lengths.max(). |
| 702 | max_len = metadata.max_prefill_query_len |
| 703 | if max_len <= 0: |
| 704 | return output |
| 705 | |
| 706 | hidden_size = x_p.shape[1] |
| 707 | positions = torch.arange( |
| 708 | num_prefill_tokens, device=x_p.device, dtype=torch.int64 |
| 709 | ) |
| 710 | req_indices = torch.searchsorted(q_starts[1:], positions, right=True) |
| 711 | col_indices = positions - q_starts[req_indices] |
| 712 | |
| 713 | packed_tokens = x_p.new_zeros((num_prefills, max_len, hidden_size)) |
| 714 | packed_tokens[req_indices, col_indices] = x_p |
| 715 | packed_tokens = packed_tokens.transpose(1, 2).contiguous() |
| 716 | |
| 717 | state_indices = state_indices_tensor_p[:num_prefills].to( |
| 718 | device=conv_state.device, dtype=torch.int64 |
| 719 | ) |
| 720 | valid_state = state_indices != NULL_BLOCK_ID |
| 721 | state_indices = torch.where( |
| 722 | valid_state, state_indices, torch.zeros_like(state_indices) |
| 723 | ) |
| 724 | has_initial = has_initial_states_p[:num_prefills].to( |
| 725 | device=conv_state.device, dtype=torch.bool |
| 726 | ) |
| 727 | if self.conv_state_len > 0: |
| 728 | if conv_state.shape[0] == 0: |
| 729 | state = conv_state.new_zeros( |
| 730 | (num_prefills, hidden_size, self.conv_state_len), |
| 731 | dtype=x_p.dtype, |
| 732 | ) |
| 733 | else: |
| 734 | state = conv_state.index_select(0, state_indices)[ |
| 735 | ..., : self.conv_state_len |
| 736 | ].to(x_p.dtype) |
| 737 | use_initial_mask = (valid_state & has_initial).view(num_prefills, 1, 1) |
| 738 | initial_state = torch.where( |
| 739 | use_initial_mask, |
| 740 | state, |
| 741 | torch.zeros_like(state), |
| 742 | ) |
| 743 | history = torch.cat((initial_state, packed_tokens), dim=-1) |
| 744 | else: |
| 745 | history = packed_tokens |
| 746 | |
| 747 | conv_output = F.conv1d( |
| 748 | history, |
| 749 | conv_weights.unsqueeze(1).contiguous(), |
| 750 | groups=history.size(1), |
| 751 | dilation=self.short_conv_dilation, |
| 752 | ) |
| 753 | conv_output = F.silu(conv_output).transpose(1, 2).contiguous() |
| 754 | |
| 755 | token_positions = torch.arange(max_len, device=x_p.device, dtype=torch.int64) |
| 756 | valid_tokens = token_positions.view(1, max_len) < lengths.view(num_prefills, 1) |
| 757 | valid_output_mask = valid_tokens & valid_state.to(device=x_p.device).view( |
| 758 | num_prefills, 1 |
| 759 | ) |
| 760 | conv_output.masked_fill_(~valid_output_mask.unsqueeze(-1), 0) |
| 761 | output.copy_(conv_output[req_indices, col_indices]) |
| 762 | |
| 763 | if self.conv_state_len > 0 and conv_state.shape[0] > 0: |
| 764 | state_starts = lengths.to(device=history.device, dtype=torch.int64).view( |
| 765 | num_prefills, 1, 1 |
| 766 | ) |
| 767 | state_offsets = torch.arange( |
| 768 | self.conv_state_len, device=history.device, dtype=torch.int64 |
| 769 | ).view(1, 1, self.conv_state_len) |
| 770 | next_state = history.gather( |
| 771 | dim=2, |
| 772 | index=(state_starts + state_offsets).expand(-1, history.size(1), -1), |
| 773 | ) |
| 774 | # Write back without a host synchronization. Valid, non-empty rows |
| 775 | # receive their new state; padding and zero-length rows keep the |
| 776 | # current cache value. |
| 777 | existing_state = conv_state.index_select(0, state_indices) |
| 778 | existing_base_state = existing_state[..., : self.conv_state_len] |
| 779 | update_mask = valid_state & (lengths.to(device=conv_state.device) > 0) |
| 780 | safe_next_state = torch.where( |
| 781 | update_mask.view(num_prefills, 1, 1), |
| 782 | next_state.to(conv_state.dtype), |
| 783 | existing_base_state, |
| 784 | ) |
| 785 | existing_state[..., : self.conv_state_len] = safe_next_state |
| 786 | conv_state.index_copy_(0, state_indices, existing_state) |
| 787 | return output |
| 788 | |
| 789 | def _short_conv_dilated_spec_batched( |
| 790 | self, |
| 791 | x_spec: torch.Tensor, |
| 792 | conv_state: torch.Tensor, |
| 793 | conv_weights: torch.Tensor, |
| 794 | spec_state_indices_tensor: torch.Tensor, |
| 795 | spec_query_start_loc: torch.Tensor, |
| 796 | num_accepted_tokens: torch.Tensor, |
| 797 | spec_query_len: int, |
| 798 | ) -> torch.Tensor: |
| 799 | """Dilated short-conv for speculative-decode (MTP) requests. |
| 800 | |
| 801 | Each spec request feeds multiple (draft + 1) query tokens. The conv |
| 802 | outputs are computed causally after rolling back the previous draft |
| 803 | state by ``num_accepted_tokens - 1``. The current candidate inputs stay |
| 804 | in the extended cache for the next forward, matching |
| 805 | ``causal_conv1d_update``. |
| 806 | |
| 807 | ``spec_query_len`` (== num_speculative_tokens + 1) is the maximum query |
| 808 | length and is a Python int, so no host synchronization is needed; this |
| 809 | keeps the path safe for full CUDA-graph capture/replay where the buffers |
| 810 | are padded at the request level. |
| 811 | """ |
| 812 | num_reqs = spec_state_indices_tensor.numel() |
| 813 | hidden_size = x_spec.size(-1) |
| 814 | # Use a fixed packing width instead of synchronizing on lengths.max(). |
| 815 | max_len = spec_query_len |
| 816 | # Full CUDA graphs can pad these buffers. Only the first num_reqs |
| 817 | # accepted-token counts belong to actual speculative requests. |
| 818 | num_accepted_tokens = num_accepted_tokens[:num_reqs] |
| 819 | q_starts = spec_query_start_loc[: num_reqs + 1].to(torch.int64) |
| 820 | # Keep the number of real speculative tokens on the device. |
| 821 | total_real_tokens = q_starts[num_reqs] |
| 822 | |
| 823 | state_indices = spec_state_indices_tensor.to( |
| 824 | device=conv_state.device, dtype=torch.int64 |
| 825 | ) |
| 826 | valid_state = state_indices != NULL_BLOCK_ID |
| 827 | state_indices = torch.where( |
| 828 | valid_state, state_indices, torch.zeros_like(state_indices) |
| 829 | ) |
| 830 | positions = torch.arange( |
| 831 | x_spec.size(0), device=x_spec.device, dtype=torch.int64 |
| 832 | ) |
| 833 | # Route graph-padded token rows to the discarded dummy request so that |
| 834 | # they cannot overwrite real packed data. |
| 835 | req_indices = torch.searchsorted(q_starts[1:], positions, right=True) |
| 836 | valid_tokens = (positions < total_real_tokens) & (req_indices < num_reqs) |
| 837 | clamped_req_indices = req_indices.clamp_max(max(num_reqs - 1, 0)) |
| 838 | col_indices = (positions - q_starts[clamped_req_indices]).clamp_(0, max_len - 1) |
| 839 | pack_req_indices = torch.where( |
| 840 | valid_tokens, |
| 841 | clamped_req_indices, |
| 842 | torch.full_like(req_indices, num_reqs), |
| 843 | ) |
| 844 | pack_col_indices = torch.where( |
| 845 | valid_tokens, col_indices, torch.zeros_like(col_indices) |
| 846 | ) |
| 847 | |
| 848 | # The last request row is the dummy sink for graph padding. |
| 849 | packed = x_spec.new_zeros((num_reqs + 1, max_len, hidden_size)) |
| 850 | packed[pack_req_indices, pack_col_indices] = x_spec |
| 851 | packed = packed.transpose(1, 2).contiguous() |
| 852 | |
| 853 | if self.conv_state_len > 0: |
| 854 | cached_state = conv_state.index_select(0, state_indices) |
| 855 | rollback_offsets = num_accepted_tokens.to( |
| 856 | device=conv_state.device, dtype=torch.int64 |
| 857 | ).sub(1) |
| 858 | rollback_offsets = torch.where( |
| 859 | valid_state, |
| 860 | rollback_offsets.clamp_(0, max_len - 1), |
| 861 | torch.zeros_like(rollback_offsets), |
| 862 | ) |
| 863 | state_offsets = torch.arange( |
| 864 | self.conv_state_len, device=conv_state.device, dtype=torch.int64 |
| 865 | ).view(1, 1, self.conv_state_len) |
| 866 | rollback_indices = rollback_offsets.view(-1, 1, 1) + state_offsets |
| 867 | state = cached_state.gather( |
| 868 | 2, rollback_indices.expand(-1, hidden_size, -1) |
| 869 | ).to(x_spec.dtype) |
| 870 | state = torch.where( |
| 871 | valid_state.view(num_reqs, 1, 1), |
| 872 | state, |
| 873 | torch.zeros_like(state), |
| 874 | ) |
| 875 | # Append a zeroed dummy-row state to match the [num_reqs + 1] pack. |
| 876 | dummy_state = state.new_zeros((1, hidden_size, self.conv_state_len)) |
| 877 | state_full = torch.cat((state, dummy_state), dim=0) |
| 878 | history = torch.cat((state_full, packed), dim=-1) |
| 879 | else: |
| 880 | history = packed |
| 881 | |
| 882 | conv_output = F.conv1d( |
| 883 | history, |
| 884 | conv_weights.unsqueeze(1).contiguous(), |
| 885 | groups=history.size(1), |
| 886 | dilation=self.short_conv_dilation, |
| 887 | ) |
| 888 | conv_output = F.silu(conv_output).transpose(1, 2).contiguous() |
| 889 | |
| 890 | output = conv_output[pack_req_indices, pack_col_indices] |
| 891 | output = output * valid_tokens.view(-1, 1).to(output.dtype) |
| 892 | |
| 893 | # Keep all current candidate inputs in the extended state. On the next |
| 894 | # target forward, ``num_accepted_tokens - 1`` selects the rollback |
| 895 | # window before processing the newly scheduled tokens. |
| 896 | if self.conv_state_len > 0: |
| 897 | state_capacity = self.conv_state_len + max_len - 1 |
| 898 | if conv_state.size(-1) < state_capacity: |
| 899 | raise RuntimeError( |
| 900 | "PLE short-conv cache cannot retain speculative tokens: " |
| 901 | f"got {conv_state.size(-1)}, need {state_capacity}." |
| 902 | ) |
| 903 | candidate_state = history[:num_reqs, :, 1 : state_capacity + 1] |
| 904 | query_lengths = q_starts[1:] - q_starts[:-1] |
| 905 | state_positions = torch.arange( |
| 906 | state_capacity, device=history.device, dtype=torch.int64 |
| 907 | ).view(1, 1, state_capacity) |
| 908 | update_lengths = (self.conv_state_len + query_lengths - 1).view( |
| 909 | num_reqs, 1, 1 |
| 910 | ) |
| 911 | update_mask = valid_state.view(num_reqs, 1, 1) & ( |
| 912 | state_positions < update_lengths |
| 913 | ) |
| 914 | existing_state = cached_state[..., :state_capacity] |
| 915 | next_state = torch.where( |
| 916 | update_mask, |
| 917 | candidate_state.to(conv_state.dtype), |
| 918 | existing_state, |
| 919 | ) |
| 920 | cached_state[..., :state_capacity] = next_state |
| 921 | conv_state.index_copy_(0, state_indices, cached_state) |
| 922 | |
| 923 | return output |
| 924 | |
| 925 | def _short_conv_dilated_dispatch( |
| 926 | self, |
| 927 | inputs: torch.Tensor, |
| 928 | metadata: PleShortConvAttentionMetadata, |
| 929 | conv_state: torch.Tensor, |
| 930 | conv_weights: torch.Tensor, |
| 931 | ) -> torch.Tensor: |
| 932 | num_prefills = metadata.num_prefills |
| 933 | num_decodes = metadata.num_decodes |
| 934 | num_decode_tokens = metadata.num_decode_tokens |
| 935 | num_prefill_tokens = metadata.num_prefill_tokens |
| 936 | has_prefill = num_prefills > 0 |
| 937 | has_decode = num_decodes > 0 |
| 938 | has_spec = metadata.spec_sequence_masks is not None |
| 939 | x = inputs[: metadata.num_actual_tokens] |
| 940 | |
| 941 | # Split spec / non-spec tokens. |
| 942 | if has_spec: |
| 943 | if has_prefill or has_decode: |
| 944 | assert metadata.spec_token_indx is not None |
| 945 | assert metadata.non_spec_token_indx is not None |
| 946 | x_spec = x.index_select(0, metadata.spec_token_indx.long()) |
| 947 | x_non_spec = x.index_select(0, metadata.non_spec_token_indx.long()) |
| 948 | else: |
| 949 | x_spec = x |
| 950 | x_non_spec = None |
| 951 | else: |
| 952 | x_spec = None |
| 953 | x_non_spec = x |
| 954 | |
| 955 | spec_output = None |
| 956 | # 1. Run the multi-query speculative-decode part. |
| 957 | if has_spec: |
| 958 | assert metadata.spec_state_indices_tensor is not None |
| 959 | assert metadata.spec_query_start_loc is not None |
| 960 | assert metadata.num_accepted_tokens is not None |
| 961 | spec_output = self._short_conv_dilated_spec_batched( |
| 962 | x_spec=x_spec, |
| 963 | conv_state=conv_state, |
| 964 | conv_weights=conv_weights, |
| 965 | spec_state_indices_tensor=metadata.spec_state_indices_tensor[ |
| 966 | : metadata.num_spec_decodes |
| 967 | ], |
| 968 | spec_query_start_loc=metadata.spec_query_start_loc, |
| 969 | num_accepted_tokens=metadata.num_accepted_tokens, |
| 970 | spec_query_len=metadata.spec_query_len, |
| 971 | ) |
| 972 | |
| 973 | # 2. Run regular decode and prefill requests. |
| 974 | conv_out_non_spec = None |
| 975 | state_indices_tensor = metadata.state_indices_tensor |
| 976 | if x_non_spec is not None: |
| 977 | assert state_indices_tensor is not None |
| 978 | if has_prefill: |
| 979 | state_indices_tensor_d, state_indices_tensor_p = torch.split( |
| 980 | state_indices_tensor, |
| 981 | [num_decodes, num_prefills], |
| 982 | dim=0, |
| 983 | ) |
| 984 | x_d, x_p = torch.split( |
| 985 | x_non_spec, |
| 986 | [num_decode_tokens, num_prefill_tokens], |
| 987 | dim=0, |
| 988 | ) |
| 989 | non_spec_parts: list[torch.Tensor] = [] |
| 990 | if has_decode: |
| 991 | non_spec_parts.append( |
| 992 | self._short_conv_dilated_decode_batched( |
| 993 | x_d=x_d, |
| 994 | conv_state=conv_state, |
| 995 | conv_weights=conv_weights, |
| 996 | state_indices_tensor_d=state_indices_tensor_d, |
| 997 | has_initial_states_d=metadata.has_initial_states_d, |
| 998 | ) |
| 999 | ) |
| 1000 | non_spec_parts.append( |
| 1001 | self._short_conv_dilated_prefill_batched( |
| 1002 | x_p=x_p, |
| 1003 | metadata=metadata, |
| 1004 | conv_state=conv_state, |
| 1005 | conv_weights=conv_weights, |
| 1006 | state_indices_tensor_p=state_indices_tensor_p, |
| 1007 | num_prefills=num_prefills, |
| 1008 | num_decode_tokens=num_decode_tokens, |
| 1009 | num_prefill_tokens=num_prefill_tokens, |
| 1010 | ) |
| 1011 | ) |
| 1012 | conv_out_non_spec = torch.vstack(non_spec_parts) |
| 1013 | else: |
| 1014 | conv_out_non_spec = self._short_conv_dilated_decode_batched( |
| 1015 | x_d=x_non_spec, |
| 1016 | conv_state=conv_state, |
| 1017 | conv_weights=conv_weights, |
| 1018 | state_indices_tensor_d=state_indices_tensor[: x_non_spec.size(0)], |
| 1019 | has_initial_states_d=metadata.has_initial_states_d, |
| 1020 | ) |
| 1021 | |
| 1022 | # 3. Merge both parts back into the original token order. |
| 1023 | if has_spec and conv_out_non_spec is not None: |
| 1024 | assert metadata.spec_token_indx is not None |
| 1025 | assert metadata.non_spec_token_indx is not None |
| 1026 | assert spec_output is not None |
| 1027 | output = x.new_empty((metadata.num_actual_tokens, x.size(-1))) |
| 1028 | output.index_copy_(0, metadata.spec_token_indx, spec_output) |
| 1029 | output.index_copy_(0, metadata.non_spec_token_indx, conv_out_non_spec) |
| 1030 | return output |
| 1031 | elif has_spec: |
| 1032 | assert spec_output is not None |
| 1033 | return spec_output |
| 1034 | if conv_out_non_spec is None: |
| 1035 | return x |
| 1036 | return conv_out_non_spec |
| 1037 | |
| 1038 | def _short_conv(self, inputs: torch.Tensor) -> torch.Tensor: |
| 1039 | forward_context = get_forward_context() |
| 1040 | attn_metadata = forward_context.attn_metadata |
| 1041 | if attn_metadata is None: |
| 1042 | return self._short_conv_fallback(inputs) |
| 1043 | |
| 1044 | if not isinstance(attn_metadata, dict): |
| 1045 | raise RuntimeError( |
| 1046 | "PLE short-conv expects per-layer attention metadata dict " |
| 1047 | f"during inference, got {type(attn_metadata).__name__}." |
| 1048 | ) |
| 1049 | |
| 1050 | layer_attn_metadata = attn_metadata.get(self.prefix) |
| 1051 | if layer_attn_metadata is None: |
| 1052 | # MRV2 omits Mamba-family metadata during profile warmup. |
| 1053 | return self._short_conv_fallback(inputs) |
| 1054 | if not isinstance(layer_attn_metadata, PleShortConvAttentionMetadata): |
| 1055 | raise TypeError( |
| 1056 | "Expected PleShortConvAttentionMetadata for layer " |
| 1057 | f"'{self.prefix}', got " |
| 1058 | f"{type(layer_attn_metadata).__name__}." |
| 1059 | ) |
| 1060 | |
| 1061 | conv_state = self.kv_cache[0] |
| 1062 | if not is_conv_state_dim_first(): |
| 1063 | conv_state = conv_state.transpose(-1, -2) |
| 1064 | conv_weights = self.conv1d.weight.squeeze(1) |
| 1065 | |
| 1066 | state_capacity = self.conv_state_len + self.num_spec_tokens |
| 1067 | if state_capacity > 0: |
| 1068 | if conv_state.size(-1) < state_capacity: |
| 1069 | raise RuntimeError( |
| 1070 | "PLE short-conv cache is smaller than expected for " |
| 1071 | f"dilated convolution: got {conv_state.size(-1)}, " |
| 1072 | f"expect at least {state_capacity}." |
| 1073 | ) |
| 1074 | conv_state = conv_state[..., -state_capacity:] |
| 1075 | return self._short_conv_dilated_dispatch( |
| 1076 | inputs, |
| 1077 | layer_attn_metadata, |
| 1078 | conv_state, |
| 1079 | conv_weights.to(dtype=inputs.dtype), |
| 1080 | ) |
| 1081 | |
| 1082 | def forward( |
| 1083 | self, |
| 1084 | hidden_states: torch.Tensor, |
| 1085 | input_ids: torch.Tensor, |
| 1086 | query_start_loc: torch.Tensor, |
| 1087 | ngram_context: torch.Tensor, |
| 1088 | ) -> torch.Tensor: |
| 1089 | input_ids = input_ids.reshape(-1) |
| 1090 | if input_ids.shape[0] != hidden_states.shape[0]: |
| 1091 | raise ValueError( |
| 1092 | "PLE expects input_ids and hidden_states to have the same " |
| 1093 | f"token length, got {input_ids.shape[0]} and " |
| 1094 | f"{hidden_states.shape[0]}" |
| 1095 | ) |
| 1096 | embeddings = self.ple_embedding(input_ids, query_start_loc, ngram_context) |
| 1097 | embeddings = self.ple_embedding.ngram_embedding.dequantize( |
| 1098 | embeddings, |
| 1099 | self.ple_embedding.ngram_embedding.params_dtype, |
| 1100 | ) |
| 1101 | key, _ = self.key_proj(embeddings) |
| 1102 | value, _ = self.value_proj(embeddings) |
| 1103 | token_count = hidden_states.shape[0] |
| 1104 | key = key.reshape(token_count, self.hc_count, self.hidden_size) |
| 1105 | query = hidden_states.reshape(token_count, self.hc_count, self.hidden_size) |
| 1106 | key = self._apply_norm(self.norm_key, key) |
| 1107 | query = self._apply_norm(self.norm_query, query) |
| 1108 | gate = (key * query).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) |
| 1109 | gate = torch.sigmoid(gate.sign() * gate.abs().clamp_min(1e-6).sqrt()) |
| 1110 | gated_value = gate * value.unsqueeze(-2) |
| 1111 | normalized = self._apply_norm(self.norm_conv, gated_value).flatten(-2) |
| 1112 | conv_output = torch.zeros_like(normalized) |
| 1113 | torch.ops.vllm.qwen4_exp_ple_short_conv( |
| 1114 | normalized, |
| 1115 | conv_output, |
| 1116 | self.prefix, |
| 1117 | ) |
| 1118 | return gated_value.flatten(-2) + conv_output |
| 1119 | |
| 1120 | |
| 1121 | def qwen4_exp_amd_ple_ngram_embedding( |
| 1122 | ngram_ids: torch.Tensor, |
| 1123 | output: torch.Tensor, |
| 1124 | layer_name: str, |
| 1125 | ) -> None: |
| 1126 | """Run the large PLE embedding lookup outside Inductor's FX graph. |
| 1127 | |
| 1128 | Keeping the embedding weight in ``static_forward_context`` prevents AOT |
| 1129 | compile-time autotuning from materializing a synthetic copy of the weight. |
| 1130 | """ |
| 1131 | layer = get_forward_context().no_compile_layers[layer_name] |
| 1132 | if not isinstance(layer, Qwen4ExpPLELayer): |
| 1133 | raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner") |
| 1134 | result = layer.ple_embedding.ngram_embedding(ngram_ids).flatten(-2) |
| 1135 | output.copy_(result) |
| 1136 | |
| 1137 | |
| 1138 | def qwen4_exp_amd_ple_ngram_embedding_pinned( |
| 1139 | ngram_ids: torch.Tensor, |
| 1140 | output: torch.Tensor, |
| 1141 | layer_name: str, |
| 1142 | ) -> None: |
| 1143 | """Run the pinned PLE UVA lookup outside Inductor's FX graph. |
| 1144 | |
| 1145 | Same rationale as the device-path escape: keeping the large embedding weight |
| 1146 | out of the graph prevents AOT compile-time autotuning from materializing a |
| 1147 | synthetic copy of the weight. |
| 1148 | """ |
| 1149 | layer = get_forward_context().no_compile_layers[layer_name] |
| 1150 | if not isinstance(layer, Qwen4ExpPLELayer): |
| 1151 | raise TypeError(f"{layer_name} is not a Qwen4Exp PLE owner") |
| 1152 | result = layer.ple_embedding.ngram_embedding.sync_lookup(ngram_ids).flatten(-2) |
| 1153 | output.copy_(result) |
| 1154 | |
| 1155 | |
| 1156 | def qwen4_exp_ple_short_conv( |
| 1157 | inputs: torch.Tensor, |
| 1158 | output: torch.Tensor, |
| 1159 | layer_name: str, |
| 1160 | ) -> None: |
| 1161 | layer = get_forward_context().no_compile_layers[layer_name] |
| 1162 | result = layer._short_conv(inputs) |
| 1163 | output[: result.shape[0]].copy_(result) |
| 1164 | |
| 1165 | |
| 1166 | direct_register_custom_op( |
| 1167 | op_name="qwen4_exp_amd_ple_ngram_embedding", |
| 1168 | op_func=qwen4_exp_amd_ple_ngram_embedding, |
| 1169 | mutates_args=["output"], |
| 1170 | ) |
| 1171 | |
| 1172 | |
| 1173 | direct_register_custom_op( |
| 1174 | op_name="qwen4_exp_amd_ple_ngram_embedding_pinned", |
| 1175 | op_func=qwen4_exp_amd_ple_ngram_embedding_pinned, |
| 1176 | mutates_args=["output"], |
| 1177 | ) |
| 1178 | |
| 1179 | |
| 1180 | direct_register_custom_op( |
| 1181 | op_name="qwen4_exp_ple_short_conv", |
| 1182 | op_func=qwen4_exp_ple_short_conv, |
| 1183 | mutates_args=["output"], |
| 1184 | ) |
| 1185 | |
| 1186 | |
| 1187 | __all__ = [ |
| 1188 | "Qwen4ExpNGramEmbedding", |
| 1189 | "Qwen4ExpPLEGroupedNorm", |
| 1190 | "Qwen4ExpPLELayer", |
| 1191 | ] |