| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Qwen4Exp n-gram embeddings with device and pinned-host storage.""" |
| 4 | |
| 5 | from collections.abc import Iterable |
| 6 | |
| 7 | import torch |
| 8 | from torch import nn |
| 9 | from transformers import Qwen4ExpTextConfig |
| 10 | |
| 11 | from vllm.config import get_current_vllm_config |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.layers.quantization.base_config import ( |
| 14 | QuantizationConfig, |
| 15 | ) |
| 16 | from vllm.model_executor.models.utils import AutoWeightsLoader |
| 17 | |
| 18 | from ..common.ngram_embedding import ( |
| 19 | Qwen4ExpPLEDeviceEmbedding, |
| 20 | Qwen4ExpPLEEmbedding, |
| 21 | Qwen4ExpPLEEmbeddingMethod, |
| 22 | Qwen4ExpPLEFp8EmbeddingMethod, |
| 23 | Qwen4ExpPLEPinnedHostEmbedding, |
| 24 | Qwen4ExpPLEUnquantizedEmbeddingMethod, |
| 25 | ) |
| 26 | from .ops.ple import ple_ngram_ids |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | __all__ = [ |
| 31 | "Qwen4ExpPLEDeviceEmbedding", |
| 32 | "Qwen4ExpPLEEmbedding", |
| 33 | "Qwen4ExpPLEEmbeddingMethod", |
| 34 | "Qwen4ExpPLEFp8EmbeddingMethod", |
| 35 | "Qwen4ExpPLEPinnedHostEmbedding", |
| 36 | "Qwen4ExpPLEUnquantizedEmbeddingMethod", |
| 37 | "Qwen4ExpNGramEmbedding", |
| 38 | ] |
| 39 | |
| 40 | |
| 41 | class Qwen4ExpNGramEmbedding(nn.Module): |
| 42 | _MASK64 = (1 << 64) - 1 |
| 43 | _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 |
| 44 | _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 |
| 45 | _SPLITMIX_M2 = 0x94D049BB133111EB |
| 46 | _PLE_LAYER_PRIME = 10007 |
| 47 | |
| 48 | @classmethod |
| 49 | def _splitmix64(cls, value: int) -> int: |
| 50 | """Mix an integer into a deterministic unsigned 64-bit value.""" |
| 51 | value = (value + cls._SPLITMIX_GAMMA) & cls._MASK64 |
| 52 | value = ((value ^ (value >> 30)) * cls._SPLITMIX_M1) & cls._MASK64 |
| 53 | value = ((value ^ (value >> 27)) * cls._SPLITMIX_M2) & cls._MASK64 |
| 54 | return (value ^ (value >> 31)) & cls._MASK64 |
| 55 | |
| 56 | @staticmethod |
| 57 | def _is_prime_64(value: int) -> bool: |
| 58 | """Return whether a 64-bit integer is prime.""" |
| 59 | if value < 2: |
| 60 | return False |
| 61 | for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): |
| 62 | if value % prime == 0: |
| 63 | return value == prime |
| 64 | exponent = value - 1 |
| 65 | shifts = 0 |
| 66 | while exponent % 2 == 0: |
| 67 | exponent //= 2 |
| 68 | shifts += 1 |
| 69 | for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022): |
| 70 | if base % value == 0: |
| 71 | continue |
| 72 | witness = pow(base, exponent, value) |
| 73 | if witness in (1, value - 1): |
| 74 | continue |
| 75 | for _ in range(shifts - 1): |
| 76 | witness = pow(witness, 2, value) |
| 77 | if witness == value - 1: |
| 78 | break |
| 79 | else: |
| 80 | return False |
| 81 | return True |
| 82 | |
| 83 | @classmethod |
| 84 | def _nth_prime_after(cls, start: int, count: int) -> int: |
| 85 | """Return the ``count``-th prime strictly greater than ``start``.""" |
| 86 | prime = int(start) |
| 87 | for _ in range(count): |
| 88 | candidate = prime + 1 |
| 89 | if candidate <= 2: |
| 90 | prime = 2 |
| 91 | continue |
| 92 | if candidate % 2 == 0: |
| 93 | candidate += 1 |
| 94 | while not cls._is_prime_64(candidate): |
| 95 | candidate += 2 |
| 96 | prime = candidate |
| 97 | return prime |
| 98 | |
| 99 | @classmethod |
| 100 | def _make_layer_multipliers( |
| 101 | cls, |
| 102 | *, |
| 103 | ngram_size: int, |
| 104 | unigram_vocab_size: int, |
| 105 | seed: int, |
| 106 | ple_dense_layer_id: int, |
| 107 | ) -> list[int]: |
| 108 | """Build deterministic hash multipliers for one PLE layer.""" |
| 109 | max_multiplier = ((1 << 63) - 1) // unigram_vocab_size |
| 110 | half_bound = max(1, max_multiplier // 2) |
| 111 | base_seed = seed + cls._PLE_LAYER_PRIME * ple_dense_layer_id |
| 112 | multipliers = [] |
| 113 | for index in range(ngram_size): |
| 114 | value = base_seed + cls._SPLITMIX_GAMMA * (index + 1) |
| 115 | multipliers.append(2 * (cls._splitmix64(value) % half_bound) + 1) |
| 116 | return multipliers |
| 117 | |
| 118 | @classmethod |
| 119 | def _make_vocab_layout( |
| 120 | cls, |
| 121 | *, |
| 122 | ngram_vocab_size_base: int, |
| 123 | ngram_heads: int, |
| 124 | ple_dense_layer_id: int, |
| 125 | ) -> tuple[list[int], list[int], int]: |
| 126 | """Build per-head vocabulary sizes, offsets, and total row count.""" |
| 127 | sizes: list[int] = [] |
| 128 | offsets: list[int] = [] |
| 129 | offset = 0 |
| 130 | for local_head in range(ngram_heads): |
| 131 | global_head = ple_dense_layer_id * ngram_heads + local_head |
| 132 | size = cls._nth_prime_after(ngram_vocab_size_base - 1, global_head + 1) |
| 133 | sizes.append(size) |
| 134 | offsets.append(offset) |
| 135 | offset += size |
| 136 | return sizes, offsets, offset |
| 137 | |
| 138 | def __init__( |
| 139 | self, |
| 140 | config: Qwen4ExpTextConfig, |
| 141 | embedding_dim: int, |
| 142 | ple_dense_layer_id: int, |
| 143 | max_total_tokens: int, |
| 144 | *, |
| 145 | data_parallel_rank: int, |
| 146 | prefix: str, |
| 147 | quant_config: QuantizationConfig | None = None, |
| 148 | params_dtype: torch.dtype | None = None, |
| 149 | ) -> None: |
| 150 | super().__init__() |
| 151 | self.embedding_dim = embedding_dim |
| 152 | self.ngram_size = int(config.ngram_size) |
| 153 | self.heads_per_ngram = int(config.heads_per_ngram) |
| 154 | self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram |
| 155 | if self.ngram_size < 2: |
| 156 | raise ValueError(f"ngram_size must be >= 2, got {self.ngram_size}") |
| 157 | if self.heads_per_ngram <= 0: |
| 158 | raise ValueError(f"heads_per_ngram must be > 0, got {self.heads_per_ngram}") |
| 159 | if embedding_dim % self.ngram_heads: |
| 160 | raise ValueError( |
| 161 | "ple_embed_dim must be divisible by total ngram heads: " |
| 162 | f"{embedding_dim} % {self.ngram_heads} != 0" |
| 163 | ) |
| 164 | self.head_dim = embedding_dim // self.ngram_heads |
| 165 | self.eos_token_id = int(config.eos_token_id) |
| 166 | self.unigram_vocab_size = int(config.vocab_size) |
| 167 | self.split_ngram_parts = int(getattr(config, "split_ngram_parts", 512)) |
| 168 | if self.split_ngram_parts <= 0: |
| 169 | raise ValueError("split_ngram_parts must be positive") |
| 170 | |
| 171 | multipliers = self._make_layer_multipliers( |
| 172 | ngram_size=self.ngram_size, |
| 173 | unigram_vocab_size=self.unigram_vocab_size, |
| 174 | seed=int(getattr(config, "seed", 1234)), |
| 175 | ple_dense_layer_id=ple_dense_layer_id, |
| 176 | ) |
| 177 | self.register_buffer( |
| 178 | "layer_multipliers", |
| 179 | torch.tensor(multipliers, dtype=torch.long), |
| 180 | persistent=True, |
| 181 | ) |
| 182 | |
| 183 | sizes, offsets, total_vocab_size = self._make_vocab_layout( |
| 184 | ngram_vocab_size_base=int(config.ngram_vocab_size_base), |
| 185 | ngram_heads=self.ngram_heads, |
| 186 | ple_dense_layer_id=ple_dense_layer_id, |
| 187 | ) |
| 188 | self.register_buffer( |
| 189 | "ngram_heads_vocab_sizes", |
| 190 | torch.tensor(sizes, dtype=torch.long), |
| 191 | persistent=True, |
| 192 | ) |
| 193 | self.register_buffer( |
| 194 | "ngram_heads_offsets", |
| 195 | torch.tensor(offsets, dtype=torch.long), |
| 196 | persistent=True, |
| 197 | ) |
| 198 | divisor = int(config.make_ngram_vocab_size_divisible_by) |
| 199 | padded_vocab_size = ((total_vocab_size + divisor - 1) // divisor) * divisor |
| 200 | embedding_prefix = f"{prefix}.ngram_embedding" |
| 201 | embedding_quant_method = Qwen4ExpPLEEmbeddingMethod.from_quant_config( |
| 202 | quant_config, |
| 203 | embedding_prefix, |
| 204 | getattr(config, "ple_embedding_dtype", None), |
| 205 | ) |
| 206 | if params_dtype is None: |
| 207 | params_dtype = torch.get_default_dtype() |
| 208 | engram_config = get_current_vllm_config().engram_config |
| 209 | embedding_cls = ( |
| 210 | Qwen4ExpPLEPinnedHostEmbedding |
| 211 | if engram_config is not None and engram_config.cpu_offload |
| 212 | else Qwen4ExpPLEDeviceEmbedding |
| 213 | ) |
| 214 | self.ngram_embedding = embedding_cls( |
| 215 | padded_vocab_size, |
| 216 | self.head_dim, |
| 217 | params_dtype=params_dtype, |
| 218 | padding_size=divisor, |
| 219 | prefix=embedding_prefix, |
| 220 | embedding_method=embedding_quant_method, |
| 221 | num_ngram_heads=self.ngram_heads, |
| 222 | max_total_tokens=max_total_tokens, |
| 223 | data_parallel_rank=data_parallel_rank, |
| 224 | ) |
| 225 | if self.ngram_embedding.supports_prefetch: |
| 226 | # The side-stream lookup outlives eager-break args, whose |
| 227 | # graph-pool storage later segments may reuse. |
| 228 | self._prefetch_ids = torch.empty( |
| 229 | max_total_tokens, self.ngram_heads, dtype=torch.long |
| 230 | ) |
| 231 | weight = self.ngram_embedding.weight |
| 232 | logger.info( |
| 233 | "Initialized PLE embedding %s: quantization_method=%s, " |
| 234 | "weight_dtype=%s, weight_device=%s, pinned=%s", |
| 235 | embedding_prefix, |
| 236 | type(embedding_quant_method).__name__, |
| 237 | weight.dtype, |
| 238 | weight.device, |
| 239 | weight.is_pinned(), |
| 240 | ) |
| 241 | |
| 242 | @staticmethod |
| 243 | def _shift_precompute( |
| 244 | tokens: torch.Tensor, eos_token_id: int |
| 245 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 246 | if tokens.dim() != 2: |
| 247 | raise ValueError("tokens must be a 2D tensor") |
| 248 | batch_size, seq_len = tokens.shape |
| 249 | positions = torch.arange(seq_len, device=tokens.device, dtype=torch.int64) |
| 250 | eos_positions = torch.where(tokens == eos_token_id, positions, -1) |
| 251 | previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values |
| 252 | previous_eos = torch.cat( |
| 253 | [ |
| 254 | eos_positions.new_full((batch_size, 1), -1), |
| 255 | previous_eos_inclusive[:, :-1], |
| 256 | ], |
| 257 | dim=1, |
| 258 | ) |
| 259 | return positions, positions.unsqueeze(0) - previous_eos - 1 |
| 260 | |
| 261 | @staticmethod |
| 262 | def _shift_apply( |
| 263 | tokens: torch.Tensor, |
| 264 | positions: torch.Tensor, |
| 265 | position_in_segment: torch.Tensor, |
| 266 | shift: int, |
| 267 | eos_token_id: int, |
| 268 | ) -> torch.Tensor: |
| 269 | if shift == 0: |
| 270 | return tokens |
| 271 | source = positions - shift |
| 272 | gather_indices = source.clamp_min(0).unsqueeze(0).expand(tokens.shape[0], -1) |
| 273 | shifted = tokens.gather(1, gather_indices) |
| 274 | valid = (source.unsqueeze(0) >= 0) & (position_in_segment >= shift) |
| 275 | return torch.where(valid, shifted, tokens.new_full((), eos_token_id)) |
| 276 | |
| 277 | def compute_ngram_ids( |
| 278 | self, |
| 279 | input_ids: torch.Tensor, |
| 280 | query_start_loc: torch.Tensor, |
| 281 | ngram_context: torch.Tensor, |
| 282 | output: torch.Tensor | None = None, |
| 283 | ) -> torch.Tensor: |
| 284 | """Compute n-gram embedding indices for the current request layout.""" |
| 285 | input_ids = input_ids.reshape(-1) |
| 286 | num_reqs = query_start_loc.numel() - 1 |
| 287 | num_tokens = input_ids.shape[0] |
| 288 | |
| 289 | if input_ids.is_cuda: |
| 290 | return ple_ngram_ids( |
| 291 | input_ids=input_ids, |
| 292 | query_start_loc=query_start_loc, |
| 293 | ngram_context=ngram_context, |
| 294 | layer_multipliers=self.layer_multipliers, |
| 295 | ngram_heads_vocab_sizes=self.ngram_heads_vocab_sizes, |
| 296 | ngram_heads_offsets=self.ngram_heads_offsets, |
| 297 | eos_token_id=self.eos_token_id, |
| 298 | heads_per_ngram=self.heads_per_ngram, |
| 299 | output=output, |
| 300 | ) |
| 301 | input_ids = input_ids.long() |
| 302 | query_start_loc = query_start_loc.long() |
| 303 | positions = torch.arange(num_tokens, device=input_ids.device, dtype=torch.int64) |
| 304 | packed = torch.full( |
| 305 | (num_reqs, num_tokens), |
| 306 | self.eos_token_id, |
| 307 | device=input_ids.device, |
| 308 | dtype=torch.int64, |
| 309 | ) |
| 310 | request_indices = torch.searchsorted(query_start_loc, positions, right=True) - 1 |
| 311 | request_indices.clamp_(max=num_reqs - 1) |
| 312 | columns = (positions - query_start_loc[request_indices]).clamp( |
| 313 | 0, packed.shape[1] - 1 |
| 314 | ) |
| 315 | packed[request_indices, columns] = input_ids |
| 316 | ngram_context = ngram_context[:num_reqs].to( |
| 317 | device=input_ids.device, dtype=torch.long |
| 318 | ) |
| 319 | |
| 320 | context = torch.cat([ngram_context, packed], dim=-1) |
| 321 | positions_2d, position_in_segment = self._shift_precompute( |
| 322 | context, self.eos_token_id |
| 323 | ) |
| 324 | shifted = [context] |
| 325 | for shift in range(1, self.ngram_size): |
| 326 | shifted.append( |
| 327 | self._shift_apply( |
| 328 | context, |
| 329 | positions_2d, |
| 330 | position_in_segment, |
| 331 | shift, |
| 332 | self.eos_token_id, |
| 333 | ) |
| 334 | ) |
| 335 | adjusted_columns = columns + self.ngram_size - 1 |
| 336 | id_blocks = [] |
| 337 | for ngram in range(2, self.ngram_size + 1): |
| 338 | start = (ngram - 2) * self.heads_per_ngram |
| 339 | end = start + self.heads_per_ngram |
| 340 | mixed = shifted[0] * self.layer_multipliers[0] |
| 341 | for index in range(1, ngram): |
| 342 | mixed = torch.bitwise_xor( |
| 343 | mixed, shifted[index] * self.layer_multipliers[index] |
| 344 | ) |
| 345 | sizes = self.ngram_heads_vocab_sizes[start:end] |
| 346 | offsets = self.ngram_heads_offsets[start:end] |
| 347 | ids = torch.remainder(mixed.unsqueeze(-1), sizes) + offsets |
| 348 | id_blocks.append(ids[request_indices, adjusted_columns]) |
| 349 | return torch.cat(id_blocks, dim=-1) |
| 350 | |
| 351 | def forward( |
| 352 | self, |
| 353 | hidden_states: torch.Tensor, |
| 354 | input_ids: torch.Tensor, |
| 355 | query_start_loc: torch.Tensor, |
| 356 | ngram_context: torch.Tensor, |
| 357 | ) -> torch.Tensor: |
| 358 | embedding = self.ngram_embedding |
| 359 | if embedding.supports_prefetch: |
| 360 | return embedding(hidden_states) |
| 361 | ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context) |
| 362 | return self.ngram_embedding(ngram_ids).flatten(-2) |
| 363 | |
| 364 | def start_prefetch( |
| 365 | self, |
| 366 | hidden_states: torch.Tensor, |
| 367 | input_ids: torch.Tensor, |
| 368 | query_start_loc: torch.Tensor, |
| 369 | ngram_context: torch.Tensor, |
| 370 | ) -> None: |
| 371 | """Start the pinned lookup while the preceding decoder layer runs.""" |
| 372 | embedding = self.ngram_embedding |
| 373 | if not embedding.supports_prefetch: |
| 374 | return |
| 375 | ngram_ids = self.compute_ngram_ids( |
| 376 | input_ids, |
| 377 | query_start_loc, |
| 378 | ngram_context, |
| 379 | output=self._prefetch_ids[: input_ids.numel()], |
| 380 | ) |
| 381 | embedding.start_prefetch(hidden_states, ngram_ids) |
| 382 | |
| 383 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 384 | """Load hash buffers and checkpoint-split embedding rows.""" |
| 385 | persistent_buffers = { |
| 386 | "layer_multipliers": self.layer_multipliers, |
| 387 | "ngram_heads_offsets": self.ngram_heads_offsets, |
| 388 | "ngram_heads_vocab_sizes": self.ngram_heads_vocab_sizes, |
| 389 | } |
| 390 | loaded: set[str] = set() |
| 391 | regular_weights: list[tuple[str, torch.Tensor]] = [] |
| 392 | shard_prefix = "ngram_embedding.shard_" |
| 393 | |
| 394 | for name, loaded_weight in weights: |
| 395 | leaf_name = name.rsplit(".", 1)[-1] |
| 396 | if leaf_name.startswith("hashstats_") or leaf_name == "token_lookup": |
| 397 | continue |
| 398 | if name in persistent_buffers: |
| 399 | buffer = persistent_buffers[name] |
| 400 | if buffer.shape != loaded_weight.shape: |
| 401 | raise ValueError( |
| 402 | f"Shape mismatch for {name}: expected " |
| 403 | f"{tuple(buffer.shape)}, got {tuple(loaded_weight.shape)}" |
| 404 | ) |
| 405 | buffer.copy_(loaded_weight.to(device=buffer.device, dtype=buffer.dtype)) |
| 406 | loaded.add(name) |
| 407 | continue |
| 408 | if name.startswith(shard_prefix) and name.endswith(".weight"): |
| 409 | shard_text = name[len(shard_prefix) : -len(".weight")] |
| 410 | if not shard_text.isdigit(): |
| 411 | regular_weights.append((name, loaded_weight)) |
| 412 | continue |
| 413 | shard_index = int(shard_text) |
| 414 | if shard_index >= self.split_ngram_parts: |
| 415 | raise ValueError( |
| 416 | f"PLE embedding shard index {shard_index} exceeds " |
| 417 | f"split_ngram_parts={self.split_ngram_parts}" |
| 418 | ) |
| 419 | embedding = self.ngram_embedding |
| 420 | shard_size = ( |
| 421 | embedding.org_vocab_size + self.split_ngram_parts - 1 |
| 422 | ) // self.split_ngram_parts |
| 423 | checkpoint_start = shard_index * shard_size |
| 424 | expected_rows = max( |
| 425 | 0, |
| 426 | min(shard_size, embedding.org_vocab_size - checkpoint_start), |
| 427 | ) |
| 428 | expected_shape = (expected_rows, embedding.embedding_dim) |
| 429 | if tuple(loaded_weight.shape) != expected_shape: |
| 430 | raise ValueError( |
| 431 | f"Shape mismatch for PLE embedding shard {shard_index}: " |
| 432 | f"expected {expected_shape}, got " |
| 433 | f"{tuple(loaded_weight.shape)}" |
| 434 | ) |
| 435 | embedding.weight.weight_loader( |
| 436 | embedding.weight, |
| 437 | loaded_weight, |
| 438 | checkpoint_start=checkpoint_start, |
| 439 | ) |
| 440 | loaded.add("ngram_embedding.weight") |
| 441 | continue |
| 442 | regular_weights.append((name, loaded_weight)) |
| 443 | |
| 444 | if regular_weights: |
| 445 | loaded.update(AutoWeightsLoader(self).load_weights(regular_weights)) |
| 446 | return loaded |
| 447 | |
| 448 | |
| 449 | __all__ = [ |
| 450 | "Qwen4ExpNGramEmbedding", |
| 451 | ] |