| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Engram: n-gram hash lookups gated into the hyper-connection stream. |
| 4 | |
| 5 | Port of the reference ``inference/engram.py`` + ``Engram`` / |
| 6 | ``ParallelEngramEmbedding`` from ``inference/model.py`` (DeepSeek V4.1 |
| 7 | checkpoint layout). Engram modules live on the backbone layers listed in |
| 8 | ``engram_layer_ids`` only. |
| 9 | |
| 10 | Two pieces of cross-forward state are needed because vLLM streams tokens |
| 11 | chunk-by-chunk while an n-gram at position ``p`` needs the token ids at |
| 12 | ``p-1..p-3``: |
| 13 | |
| 14 | - ``token_map``: token id -> compressed vocab id, built once from the |
| 15 | model's tokenizer at init (deterministic; asserted against |
| 16 | ``engram_compressed_vocab_size``). |
| 17 | - ``hash_cache``: one int32 slot per KV slot of the first local layer's |
| 18 | sliding-window cache, holding the compressed id (or DEAD) of the token |
| 19 | last written to that slot. Slots are stable per (request, position) — |
| 20 | the block table pins a position to a physical slot, prefix-cache hits |
| 21 | reuse both the physical blocks and the identical token ids, and |
| 22 | spec-decode rollbacks rewrite the same slots — so lookbacks read back |
| 23 | exactly what the owning request wrote. Lookback depth (3) is far inside |
| 24 | the sliding window (128), so window eviction never frees a block a |
| 25 | live lookback still needs. |
| 26 | |
| 27 | Slots are not part of the KV cache, so KV loaded from another instance |
| 28 | (P/D, offload connectors) leaves them unwritten. The runner therefore |
| 29 | passes ``lookback_token_ids``, the ids just before each request's chunk |
| 30 | start, which take precedence over the slots. The V2 runner reads them |
| 31 | from its device-resident token history and needs no slot cache; the V1 |
| 32 | runner's CPU token table holds placeholders for generated tokens under |
| 33 | async scheduling, so it passes prompt positions only and keeps the slot |
| 34 | cache for the rest. |
| 35 | """ |
| 36 | |
| 37 | import ctypes |
| 38 | import mmap |
| 39 | import os |
| 40 | import tempfile |
| 41 | import weakref |
| 42 | from contextlib import ExitStack |
| 43 | |
| 44 | import numpy as np |
| 45 | import torch |
| 46 | from torch import nn |
| 47 | |
| 48 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 49 | from vllm.config import VllmConfig, get_current_vllm_config |
| 50 | from vllm.distributed import ( |
| 51 | get_dp_group, |
| 52 | get_engram_dp_group, |
| 53 | get_engram_dp_size, |
| 54 | get_tensor_model_parallel_rank, |
| 55 | get_tensor_model_parallel_world_size, |
| 56 | tensor_model_parallel_all_gather, |
| 57 | ) |
| 58 | from vllm.distributed.device_communicators.shm_broadcast import ( |
| 59 | SHM_PATH, |
| 60 | check_shm_free_space, |
| 61 | ) |
| 62 | from vllm.distributed.parallel_state import GroupCoordinator |
| 63 | from vllm.forward_context import get_forward_context, is_forward_context_available |
| 64 | from vllm.logger import init_logger |
| 65 | from vllm.model_executor.layers.linear import ColumnParallelLinear |
| 66 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 67 | from vllm.model_executor.utils import set_weight_attrs |
| 68 | from vllm.platforms import current_platform |
| 69 | from vllm.triton_utils import tl, triton |
| 70 | from vllm.utils.platform_utils import is_uva_available |
| 71 | from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor |
| 72 | |
| 73 | logger = init_logger(__name__) |
| 74 | |
| 75 | # Cache value for tokens that take no part in an n-gram (image spans). |
| 76 | DEAD_ID = -1 |
| 77 | |
| 78 | |
| 79 | def _engram_lookup_thresholds(device: torch.device) -> tuple[int | None, int | None]: |
| 80 | """Host lookup thresholds on `device`: the rows from which to sort, and the |
| 81 | tokens (max over DP ranks) from which to run inline. None disables that path.""" |
| 82 | if current_platform.is_cuda() and current_platform.is_device_capability_family( |
| 83 | 100, device.index or 0 |
| 84 | ): |
| 85 | # Measured on GB200. |
| 86 | return 36864, 4096 |
| 87 | # TODO: Add hand-tuned thresholds for other GPU families. |
| 88 | return None, None |
| 89 | |
| 90 | |
| 91 | def _is_prime(n: int) -> bool: |
| 92 | """Deterministic Miller-Rabin for n < 2**32 (avoids a sympy import).""" |
| 93 | if n < 2: |
| 94 | return False |
| 95 | for p in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37): |
| 96 | if n % p == 0: |
| 97 | return n == p |
| 98 | d = n - 1 |
| 99 | r = 0 |
| 100 | while d % 2 == 0: |
| 101 | d //= 2 |
| 102 | r += 1 |
| 103 | for a in (2, 7, 61): |
| 104 | x = pow(a, d, n) |
| 105 | if x in (1, n - 1): |
| 106 | continue |
| 107 | for _ in range(r - 1): |
| 108 | x = x * x % n |
| 109 | if x == n - 1: |
| 110 | break |
| 111 | else: |
| 112 | return False |
| 113 | return True |
| 114 | |
| 115 | |
| 116 | def find_next_prime(start: int, seen_primes: set[int]) -> int: |
| 117 | """The smallest prime above `start` that has not been handed out yet.""" |
| 118 | candidate = start + 1 |
| 119 | while not _is_prime(candidate) or candidate in seen_primes: |
| 120 | candidate += 1 |
| 121 | return candidate |
| 122 | |
| 123 | |
| 124 | def build_compressed_token_map(tokenizer) -> tuple[list[int], int]: |
| 125 | """Map every token id onto a smaller id space where tokens that normalize |
| 126 | alike collapse together. |
| 127 | |
| 128 | N-grams are hashed over these compressed ids, so " The", "the" and "THE" |
| 129 | all hash the same way. The compressed size matters beyond bounds checking: |
| 130 | every hash multiplier is derived from it. |
| 131 | """ |
| 132 | from tokenizers import Regex, normalizers |
| 133 | |
| 134 | # A private-use char, so a token that is exactly one space survives |
| 135 | # Strip() instead of collapsing to the empty string and merging with |
| 136 | # unrelated tokens. |
| 137 | sentinel = "\ue000" |
| 138 | normalizer = normalizers.Sequence( |
| 139 | [ |
| 140 | normalizers.NFKC(), |
| 141 | normalizers.NFD(), |
| 142 | normalizers.StripAccents(), |
| 143 | normalizers.Lowercase(), |
| 144 | normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), |
| 145 | normalizers.Replace(Regex(r"^ $"), sentinel), |
| 146 | normalizers.Strip(), |
| 147 | normalizers.Replace(sentinel, " "), |
| 148 | ] |
| 149 | ) |
| 150 | |
| 151 | # The raw Rust tokenizer, matching what training decodes with |
| 152 | # (no clean_up_tokenization_spaces). |
| 153 | backend = tokenizer.backend_tokenizer |
| 154 | key_to_new: dict[str, int] = {} |
| 155 | lookup = [0] * len(tokenizer) |
| 156 | for token_id in range(len(tokenizer)): |
| 157 | text = backend.decode([token_id], skip_special_tokens=False) |
| 158 | if "\ufffd" in text: |
| 159 | # A partial UTF-8 byte token: nothing to normalize, so key it |
| 160 | # by its raw form. |
| 161 | key = backend.id_to_token(token_id) |
| 162 | else: |
| 163 | normalized = normalizer.normalize_str(text) |
| 164 | key = normalized if normalized else text |
| 165 | |
| 166 | new_id = key_to_new.get(key) |
| 167 | if new_id is None: |
| 168 | new_id = len(key_to_new) |
| 169 | key_to_new[key] = new_id |
| 170 | lookup[token_id] = new_id |
| 171 | |
| 172 | return lookup, len(key_to_new) |
| 173 | |
| 174 | |
| 175 | def compute_hash_multipliers( |
| 176 | layer_ids: tuple[int, ...], max_ngram_size: int, compressed_vocab_size: int |
| 177 | ) -> torch.Tensor: |
| 178 | """One multiplier per (layer, lookback), from a per-layer RNG so layers |
| 179 | hash differently. Kept odd and bounded so `token_id * multiplier` cannot |
| 180 | overflow int64. |
| 181 | """ |
| 182 | max_long = np.iinfo(np.int64).max |
| 183 | multiplier_bound = max(1, (max_long // compressed_vocab_size) // 2) |
| 184 | rows = [] |
| 185 | for layer_id in layer_ids: |
| 186 | generator = np.random.default_rng(10007 * layer_id) |
| 187 | values = generator.integers( |
| 188 | low=0, |
| 189 | high=multiplier_bound, |
| 190 | size=(max_ngram_size,), |
| 191 | dtype=np.int64, |
| 192 | ) |
| 193 | rows.append(torch.tensor(values * 2 + 1)) |
| 194 | return torch.stack(rows) |
| 195 | |
| 196 | |
| 197 | class EngramLayout: |
| 198 | """Bucket layout of the n-gram hash tables. |
| 199 | |
| 200 | A position is hashed as `max_ngram_size - 1` n-grams (2-gram .. max), each |
| 201 | split over `n_heads` heads. Every (n-gram size, head) pair owns its own |
| 202 | prime-sized bucket range in the layer's table; the primes are drawn in |
| 203 | order and never reused, which keeps the ranges disjoint. |
| 204 | """ |
| 205 | |
| 206 | def __init__(self, config) -> None: |
| 207 | self.layer_ids: tuple[int, ...] = tuple(config.engram_layer_ids) |
| 208 | self.num_embeddings: tuple[int, ...] = tuple(config.engram_num_embeddings) |
| 209 | self.max_ngram_size: int = config.engram_max_ngram_size |
| 210 | self.n_heads: int = config.engram_n_heads |
| 211 | self.head_dim: int = config.engram_head_dim |
| 212 | self.compressed_vocab_size: int = config.engram_compressed_vocab_size |
| 213 | self.pad_token_id: int = config.engram_pad_token_id |
| 214 | assert len(self.layer_ids) == len(self.num_embeddings) |
| 215 | |
| 216 | primes = [] |
| 217 | seen: set[int] = set() |
| 218 | for _ in self.layer_ids: |
| 219 | per_ngram = [] |
| 220 | for _ in range(self.max_ngram_size - 1): |
| 221 | sizes, current = [], config.engram_vocab_size - 1 |
| 222 | for _ in range(self.n_heads): |
| 223 | current = find_next_prime(current, seen) |
| 224 | seen.add(current) |
| 225 | sizes.append(current) |
| 226 | per_ngram.append(tuple(sizes)) |
| 227 | primes.append(tuple(per_ngram)) |
| 228 | self.primes: tuple[tuple[tuple[int, ...], ...], ...] = tuple(primes) |
| 229 | self.n_hash_cols = (self.max_ngram_size - 1) * self.n_heads |
| 230 | flat = [[p for per_ngram in layer for p in per_ngram] for layer in primes] |
| 231 | offsets = [np.cumsum([0, *sizes[:-1]]) for sizes in flat] |
| 232 | self.offsets = torch.tensor(np.array(offsets)) # [n_layers, n_hash_cols] |
| 233 | |
| 234 | @classmethod |
| 235 | def from_config(cls, config) -> "EngramLayout | None": |
| 236 | if not getattr(config, "engram_layer_ids", None): |
| 237 | return None |
| 238 | return cls(config) |
| 239 | |
| 240 | |
| 241 | @triton.jit(do_not_specialize=["num_tokens"]) |
| 242 | def _write_hash_cache_kernel( |
| 243 | input_ids, |
| 244 | token_map, |
| 245 | dead_mask, |
| 246 | slot_mapping, |
| 247 | cache, |
| 248 | num_tokens, |
| 249 | input_stride, |
| 250 | mask_stride, |
| 251 | slot_stride, |
| 252 | BLOCK_SIZE: tl.constexpr, |
| 253 | dead_id, |
| 254 | ): |
| 255 | token_idx = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 256 | slot = tl.load( |
| 257 | slot_mapping + token_idx * slot_stride, token_idx < num_tokens, other=-1 |
| 258 | ).to(tl.int64) |
| 259 | valid = (token_idx < num_tokens) & (slot >= 0) |
| 260 | token = tl.load(input_ids + token_idx * input_stride, valid, other=0) |
| 261 | value = tl.load(token_map + token, valid, other=0) |
| 262 | dead = tl.load(dead_mask + token_idx * mask_stride, valid, other=False) |
| 263 | value = tl.where(dead, dead_id, value) |
| 264 | tl.store(cache + slot, value, valid) |
| 265 | |
| 266 | |
| 267 | # Keep request shapes out of the cache key so one warmed variant covers runtime batches. |
| 268 | @triton.jit( |
| 269 | do_not_specialize=[ |
| 270 | "num_tokens", |
| 271 | "num_slots", |
| 272 | "num_query_rows", |
| 273 | "num_table_rows", |
| 274 | "max_blocks", |
| 275 | ] |
| 276 | ) |
| 277 | def _hash_ids_kernel( |
| 278 | input_ids, |
| 279 | token_map, |
| 280 | dead_mask, |
| 281 | positions, |
| 282 | block_table, |
| 283 | query_start_loc, |
| 284 | multipliers, |
| 285 | primes, |
| 286 | offsets, |
| 287 | cache, |
| 288 | lookback_token_ids, |
| 289 | lookback_dead_mask, |
| 290 | output, |
| 291 | num_tokens, |
| 292 | num_slots, |
| 293 | pad_id, |
| 294 | input_stride, |
| 295 | mask_stride, |
| 296 | position_stride, |
| 297 | table_stride, |
| 298 | table_col_stride, |
| 299 | query_stride, |
| 300 | num_query_rows, |
| 301 | num_table_rows, |
| 302 | max_blocks, |
| 303 | cache_block_size, |
| 304 | MAX_NGRAM: tl.constexpr, |
| 305 | num_heads, |
| 306 | BLOCK_T: tl.constexpr, |
| 307 | BLOCK_H: tl.constexpr, |
| 308 | dead_id, |
| 309 | lookback_depth, |
| 310 | lookback_row_stride, |
| 311 | lookback_col_stride, |
| 312 | lookback_mask_row_stride, |
| 313 | lookback_mask_col_stride, |
| 314 | ): |
| 315 | token = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) |
| 316 | layer = tl.program_id(1) |
| 317 | num_layers = tl.num_programs(1) |
| 318 | valid = token < num_tokens |
| 319 | # Upper bound in query_start_loc[1:], including repeated padding boundaries. |
| 320 | lo = tl.full((BLOCK_T,), 0, tl.int32) |
| 321 | hi = tl.full((BLOCK_T,), num_query_rows, tl.int32) |
| 322 | while tl.sum((lo < hi).to(tl.int32), 0) > 0: |
| 323 | mid = (lo + hi) // 2 |
| 324 | end = tl.load( |
| 325 | query_start_loc + (mid + 1) * query_stride, |
| 326 | lo < hi, |
| 327 | other=0, |
| 328 | ) |
| 329 | right = token >= end |
| 330 | active = lo < hi |
| 331 | lo = tl.where(active & right, mid + 1, lo) |
| 332 | hi = tl.where(active & ~right, mid, hi) |
| 333 | req = tl.minimum(lo, num_query_rows - 1).to(tl.int64) |
| 334 | chunk_idx = tl.load(query_start_loc + req * query_stride) |
| 335 | chunk_idx = tl.minimum(chunk_idx, num_tokens - 1).to(tl.int64) |
| 336 | chunk_start = tl.load(positions + chunk_idx * position_stride) |
| 337 | position = tl.load(positions + token * position_stride, valid, other=0).to(tl.int64) |
| 338 | head = tl.arange(0, BLOCK_H) |
| 339 | blocked = tl.full((BLOCK_T,), False, tl.int1) |
| 340 | rolling = tl.full((BLOCK_T,), 0, tl.int64) |
| 341 | for shift in tl.static_range(MAX_NGRAM): |
| 342 | lookback = position - shift |
| 343 | in_batch = lookback >= chunk_start |
| 344 | batch_idx = tl.maximum(token - shift, 0) |
| 345 | batch_token = tl.load( |
| 346 | input_ids + batch_idx * input_stride, valid & in_batch, other=0 |
| 347 | ) |
| 348 | batch_source = tl.load(token_map + batch_token, valid & in_batch, other=0) |
| 349 | batch_dead = tl.load( |
| 350 | dead_mask + batch_idx * mask_stride, valid & in_batch, other=False |
| 351 | ) |
| 352 | batch_source = tl.where(batch_dead, dead_id, batch_source) |
| 353 | |
| 354 | col = chunk_start - 1 - lookback |
| 355 | in_window = valid & ~in_batch & (col >= 0) & (col < lookback_depth) |
| 356 | col = tl.minimum(tl.maximum(col, 0), lookback_depth - 1) |
| 357 | window_token = tl.load( |
| 358 | lookback_token_ids + req * lookback_row_stride + col * lookback_col_stride, |
| 359 | in_window, |
| 360 | other=-1, |
| 361 | ) |
| 362 | known = in_window & (window_token >= 0) |
| 363 | window_source = tl.load(token_map + window_token, known, other=0) |
| 364 | window_dead = tl.load( |
| 365 | lookback_dead_mask |
| 366 | + req * lookback_mask_row_stride |
| 367 | + col * lookback_mask_col_stride, |
| 368 | known, |
| 369 | other=False, |
| 370 | ) |
| 371 | window_source = tl.where(window_dead, dead_id, window_source) |
| 372 | |
| 373 | if cache is not None: |
| 374 | clamped = tl.minimum( |
| 375 | tl.maximum(lookback, 0), max_blocks * cache_block_size - 1 |
| 376 | ) |
| 377 | block_row = tl.minimum(req, num_table_rows - 1) |
| 378 | needs_cache = valid & ~in_batch & ~known |
| 379 | block = tl.load( |
| 380 | block_table |
| 381 | + block_row * table_stride |
| 382 | + (clamped // cache_block_size) * table_col_stride, |
| 383 | needs_cache, |
| 384 | other=0, |
| 385 | ).to(tl.int64) |
| 386 | slot = tl.minimum( |
| 387 | tl.maximum(block * cache_block_size + clamped % cache_block_size, 0), |
| 388 | num_slots - 1, |
| 389 | ) |
| 390 | fallback = tl.load(cache + slot, needs_cache, other=0) |
| 391 | else: |
| 392 | fallback = tl.full((BLOCK_T,), pad_id, tl.int32) |
| 393 | source = tl.where( |
| 394 | in_batch, batch_source, tl.where(known, window_source, fallback) |
| 395 | ).to(tl.int64) |
| 396 | blocked |= (lookback < 0) | (source == dead_id) |
| 397 | value = tl.where(blocked, pad_id, source) |
| 398 | multiplier = tl.load(multipliers + layer * MAX_NGRAM + shift) |
| 399 | rolling ^= value * multiplier |
| 400 | if shift > 0: |
| 401 | col = (shift - 1) * num_heads + head |
| 402 | param_offset = layer * (MAX_NGRAM - 1) * num_heads + col |
| 403 | prime = tl.load(primes + param_offset, head < num_heads, other=1) |
| 404 | offset = tl.load(offsets + param_offset, head < num_heads, other=0) |
| 405 | hashed = rolling[:, None] % prime[None, :] + offset[None, :] |
| 406 | out_offset = (token.to(tl.int64) * num_layers + layer)[:, None] * ( |
| 407 | (MAX_NGRAM - 1) * num_heads |
| 408 | ) + col[None, :] |
| 409 | tl.store(output + out_offset, hashed, valid[:, None] & (head < num_heads)) |
| 410 | |
| 411 | |
| 412 | class NgramHashState(nn.Module): |
| 413 | """Maps each position to the hash ids of the n-grams ending there. |
| 414 | |
| 415 | Stateless on the V2 runner, which supplies every lookback token id. On |
| 416 | the V1 runner it also keeps `hash_cache`, the slot-keyed rolling store |
| 417 | of compressed ids (see module docstring), for generated tokens. |
| 418 | """ |
| 419 | |
| 420 | def __init__( |
| 421 | self, |
| 422 | vllm_config: VllmConfig, |
| 423 | layout: EngramLayout, |
| 424 | swa_cache_module: nn.Module, |
| 425 | ) -> None: |
| 426 | super().__init__() |
| 427 | self.layout = layout |
| 428 | self.swa_cache_module = swa_cache_module |
| 429 | self.block_size: int = swa_cache_module.block_size |
| 430 | self.lookback_depth: int = layout.max_ngram_size - 1 |
| 431 | self.use_slot_cache: bool = not vllm_config.use_v2_model_runner |
| 432 | self._cache: torch.Tensor | None = None |
| 433 | self._kv_cache_ref: weakref.ReferenceType[torch.Tensor] | None = None |
| 434 | |
| 435 | model_config = vllm_config.model_config |
| 436 | from transformers import AutoTokenizer |
| 437 | |
| 438 | tokenizer = AutoTokenizer.from_pretrained( |
| 439 | model_config.tokenizer, |
| 440 | trust_remote_code=model_config.trust_remote_code, |
| 441 | revision=model_config.revision, |
| 442 | ) |
| 443 | token_map, vocab_size = build_compressed_token_map(tokenizer) |
| 444 | if vocab_size != layout.compressed_vocab_size: |
| 445 | raise ValueError( |
| 446 | f"Compressed vocab size mismatch: built {vocab_size} from the " |
| 447 | f"tokenizer, config expects {layout.compressed_vocab_size}; " |
| 448 | "every hash multiplier derives from it, so the engram tables " |
| 449 | "would be silently rehashed." |
| 450 | ) |
| 451 | self.pad_id = token_map[layout.pad_token_id] |
| 452 | multipliers = compute_hash_multipliers( |
| 453 | layout.layer_ids, layout.max_ngram_size, vocab_size |
| 454 | ) |
| 455 | self.register_buffer( |
| 456 | "token_map", torch.tensor(token_map, dtype=torch.int32), persistent=False |
| 457 | ) |
| 458 | self.register_buffer("primes", torch.tensor(layout.primes), persistent=False) |
| 459 | self.register_buffer("offsets", layout.offsets, persistent=False) |
| 460 | self.register_buffer("multipliers", multipliers, persistent=False) |
| 461 | logger.info( |
| 462 | "Built engram token map (%d -> %d ids) for layers %s", |
| 463 | len(token_map), |
| 464 | vocab_size, |
| 465 | layout.layer_ids, |
| 466 | ) |
| 467 | |
| 468 | def ensure_cache(self) -> bool: |
| 469 | """Lazily size the slot-keyed cache from the bound SWA KV cache. |
| 470 | |
| 471 | Returns False while the KV cache is unbound (profile run); the caller |
| 472 | skips engram hashing then. Without the slot cache only that check |
| 473 | remains. |
| 474 | """ |
| 475 | kv_cache = self.swa_cache_module.kv_cache |
| 476 | if kv_cache.numel() == 0: |
| 477 | self._cache = None |
| 478 | self._kv_cache_ref = None |
| 479 | return False |
| 480 | if not self.use_slot_cache: |
| 481 | return True |
| 482 | if self._kv_cache_ref is not None and self._kv_cache_ref() is kv_cache: |
| 483 | return True |
| 484 | # Graph memory profiling binds a temporary, smaller KV cache first. |
| 485 | # Rebinding must discard its hash history without retaining KV storage. |
| 486 | self._cache = torch.zeros( |
| 487 | kv_cache.shape[0] * self.block_size, |
| 488 | dtype=torch.int32, |
| 489 | device=kv_cache.device, |
| 490 | ) |
| 491 | self._kv_cache_ref = weakref.ref(kv_cache) |
| 492 | return True |
| 493 | |
| 494 | def dummy_hashes( |
| 495 | self, input_ids: torch.Tensor |
| 496 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 497 | """Participate in DP lookups without valid rows or hash-cache updates.""" |
| 498 | num_tokens = input_ids.shape[0] |
| 499 | num_layers, max_ngram = self.multipliers.shape |
| 500 | num_heads = self.primes.shape[-1] |
| 501 | hashes = input_ids.new_full( |
| 502 | (num_tokens, num_layers, (max_ngram - 1) * num_heads), |
| 503 | DEAD_ID, |
| 504 | dtype=torch.int32, |
| 505 | ) |
| 506 | keep = torch.zeros(num_tokens, dtype=torch.bool, device=input_ids.device) |
| 507 | return hashes, keep |
| 508 | |
| 509 | def forward( |
| 510 | self, |
| 511 | input_ids: torch.Tensor, |
| 512 | positions: torch.Tensor, |
| 513 | query_start_loc: torch.Tensor, |
| 514 | dead_mask: torch.Tensor, |
| 515 | lookback_token_ids: torch.Tensor, |
| 516 | lookback_dead_mask: torch.Tensor, |
| 517 | slot_mapping: torch.Tensor | None, |
| 518 | block_table: torch.Tensor | None, |
| 519 | ) -> torch.Tensor: |
| 520 | """Compute [tokens, layers, hash columns] int32 n-gram hashes. |
| 521 | |
| 522 | History comes from the current chunk, then the runner's lookback |
| 523 | window, then the optional V1 slot cache. V2 needs only one launch. |
| 524 | """ |
| 525 | cache = self._cache if self.use_slot_cache else None |
| 526 | num_tokens = input_ids.shape[0] |
| 527 | num_layers, max_ngram = self.multipliers.shape |
| 528 | num_heads = self.primes.shape[-1] |
| 529 | output = input_ids.new_empty( |
| 530 | (num_tokens, num_layers, (max_ngram - 1) * num_heads), dtype=torch.int32 |
| 531 | ) |
| 532 | if num_tokens == 0: |
| 533 | return output |
| 534 | if self.use_slot_cache: |
| 535 | assert cache is not None and slot_mapping is not None |
| 536 | assert block_table is not None |
| 537 | # Finish writes before other thread blocks read fallback history. |
| 538 | _write_hash_cache_kernel[(triton.cdiv(num_tokens, 256),)]( |
| 539 | input_ids, |
| 540 | self.token_map, |
| 541 | dead_mask, |
| 542 | slot_mapping, |
| 543 | cache, |
| 544 | num_tokens, |
| 545 | input_ids.stride(0), |
| 546 | dead_mask.stride(0), |
| 547 | slot_mapping.stride(0), |
| 548 | 256, |
| 549 | DEAD_ID, |
| 550 | ) |
| 551 | _hash_ids_kernel[(triton.cdiv(num_tokens, 32), num_layers)]( |
| 552 | input_ids, |
| 553 | self.token_map, |
| 554 | dead_mask, |
| 555 | positions, |
| 556 | block_table, |
| 557 | query_start_loc, |
| 558 | self.multipliers, |
| 559 | self.primes, |
| 560 | self.offsets, |
| 561 | cache, |
| 562 | lookback_token_ids, |
| 563 | lookback_dead_mask, |
| 564 | output, |
| 565 | num_tokens, |
| 566 | cache.shape[0] if cache is not None else 0, |
| 567 | self.pad_id, |
| 568 | input_stride=input_ids.stride(0), |
| 569 | mask_stride=dead_mask.stride(0), |
| 570 | position_stride=positions.stride(0), |
| 571 | table_stride=block_table.stride(0) if block_table is not None else 0, |
| 572 | table_col_stride=block_table.stride(1) if block_table is not None else 0, |
| 573 | query_stride=query_start_loc.stride(0), |
| 574 | num_query_rows=query_start_loc.numel() - 1, |
| 575 | num_table_rows=block_table.shape[0] if block_table is not None else 0, |
| 576 | max_blocks=block_table.shape[1] if block_table is not None else 0, |
| 577 | cache_block_size=self.block_size, |
| 578 | MAX_NGRAM=max_ngram, |
| 579 | num_heads=num_heads, |
| 580 | BLOCK_T=32, |
| 581 | BLOCK_H=triton.next_power_of_2(num_heads), |
| 582 | dead_id=DEAD_ID, |
| 583 | lookback_depth=lookback_token_ids.shape[1], |
| 584 | lookback_row_stride=lookback_token_ids.stride(0), |
| 585 | lookback_col_stride=lookback_token_ids.stride(1), |
| 586 | lookback_mask_row_stride=lookback_dead_mask.stride(0), |
| 587 | lookback_mask_col_stride=lookback_dead_mask.stride(1), |
| 588 | num_warps=4, |
| 589 | ) |
| 590 | return output |
| 591 | |
| 592 | |
| 593 | def _engram_head_shard_weight_loader( |
| 594 | param: torch.nn.Parameter, loaded_weight: torch.Tensor |
| 595 | ) -> None: |
| 596 | """Load this rank's complete head buckets. ue8m0 scales arrive as |
| 597 | float8_e8m0fnu; keep the raw bytes (the param stores uint8).""" |
| 598 | part_rows = param.shape[0] |
| 599 | if loaded_weight.dtype == torch.float8_e8m0fnu: |
| 600 | loaded_weight = loaded_weight.view(torch.uint8) |
| 601 | shard = loaded_weight.narrow(0, param.engram_vocab_start, part_rows) |
| 602 | assert shard.shape == param.shape, ( |
| 603 | f"engram shard {tuple(shard.shape)} does not fit param {tuple(param.shape)}" |
| 604 | ) |
| 605 | param.data.copy_(shard) |
| 606 | |
| 607 | |
| 608 | # Branching on SORTED at runtime intermittently crashes Triton 3.8's |
| 609 | # RemoveLayoutConversions. |
| 610 | # TODO: Remove the 3.8 check once Triton ships |
| 611 | # https://github.com/triton-lang/triton/pull/10706. |
| 612 | _SORTED_IS_CONSTEXPR = tl.constexpr( |
| 613 | triton.__version__.startswith("3.8") and current_platform.is_rocm() |
| 614 | ) |
| 615 | |
| 616 | |
| 617 | @triton.jit( |
| 618 | do_not_specialize=[ |
| 619 | "vocab_start", |
| 620 | "vocab_end", |
| 621 | "num_rows", |
| 622 | "ids_stride_t", |
| 623 | "ids_stride_h", |
| 624 | "GRID", |
| 625 | *(() if _SORTED_IS_CONSTEXPR else ("SORTED",)), |
| 626 | ] |
| 627 | ) |
| 628 | def _engram_lookup_kernel( |
| 629 | weight, |
| 630 | scales, |
| 631 | ids, |
| 632 | sorted_dst, |
| 633 | out, |
| 634 | vocab_start, |
| 635 | vocab_end, |
| 636 | num_rows, |
| 637 | ids_stride_t, |
| 638 | ids_stride_h, |
| 639 | HEAD_START: tl.constexpr, |
| 640 | LOCAL_HEADS: tl.constexpr, |
| 641 | TOTAL_HEADS: tl.constexpr, |
| 642 | DIM: tl.constexpr, |
| 643 | QUANT_BLOCK: tl.constexpr, |
| 644 | BLOCK_R: tl.constexpr, |
| 645 | GRID, |
| 646 | SORTED: tl.constexpr if _SORTED_IS_CONSTEXPR else None, # type: ignore[valid-type] |
| 647 | ): |
| 648 | """Gather fp8 rows, apply their ue8m0 block scales, write bf16. |
| 649 | |
| 650 | Only this rank's heads are read; padded heads write zeros for all-gather. |
| 651 | `weight`/`scales` may address pinned host memory through UVA. If SORTED, |
| 652 | `ids` holds rows relative to vocab_start in table order and `sorted_dst` |
| 653 | their output rows. |
| 654 | """ |
| 655 | cols = tl.arange(0, DIM) |
| 656 | scale_cols = cols // QUANT_BLOCK |
| 657 | for base in tl.range(tl.program_id(0) * BLOCK_R, num_rows, GRID * BLOCK_R): |
| 658 | rows = base + tl.arange(0, BLOCK_R) |
| 659 | valid = rows < num_rows |
| 660 | if SORTED: |
| 661 | local = tl.load(ids + rows, mask=valid, other=-1).to(tl.int64) |
| 662 | dst = tl.load(sorted_dst + rows, mask=valid, other=0).to(tl.int32) |
| 663 | owned = (local >= 0) & (local < vocab_end - vocab_start) |
| 664 | else: |
| 665 | head = HEAD_START + rows % LOCAL_HEADS |
| 666 | token = (rows // LOCAL_HEADS).to(tl.int64) |
| 667 | index = tl.load( |
| 668 | ids + token * ids_stride_t + head * ids_stride_h, |
| 669 | mask=valid & (head < TOTAL_HEADS), |
| 670 | other=-1, |
| 671 | ).to(tl.int64) |
| 672 | dst = rows |
| 673 | owned = valid & (head < TOTAL_HEADS) |
| 674 | owned &= (index >= vocab_start) & (index < vocab_end) |
| 675 | local = index - vocab_start |
| 676 | local = tl.where(owned, local, 0) |
| 677 | values = tl.load( |
| 678 | weight + local[:, None] * DIM + cols[None, :], |
| 679 | mask=owned[:, None], |
| 680 | other=0.0, |
| 681 | ) |
| 682 | scale = tl.load( |
| 683 | scales + local[:, None] * (DIM // QUANT_BLOCK) + scale_cols[None, :], |
| 684 | mask=owned[:, None], |
| 685 | other=0, |
| 686 | ) |
| 687 | # ue8m0 is a power of two, so its byte *is* the fp32 exponent field. |
| 688 | scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True) |
| 689 | tl.store( |
| 690 | out + dst[:, None] * DIM + cols[None, :], |
| 691 | (values.to(tl.float32) * scale).to(tl.bfloat16), |
| 692 | mask=valid[:, None], |
| 693 | ) |
| 694 | |
| 695 | |
| 696 | def engram_head_shard_rank() -> int: |
| 697 | """This rank's slot among the hash-head shards of one engram table. |
| 698 | |
| 699 | TP-major, so the shards a DP gather brings in are contiguous heads and |
| 700 | the following TP gather completes the head order. |
| 701 | """ |
| 702 | dp_group = get_engram_dp_group() |
| 703 | dp_size = dp_group.world_size if dp_group is not None else 1 |
| 704 | dp_rank = dp_group.rank_in_group if dp_group is not None else 0 |
| 705 | return get_tensor_model_parallel_rank() * dp_size + dp_rank |
| 706 | |
| 707 | |
| 708 | def engram_gathered_num_tokens() -> int: |
| 709 | """Per-replica token slot for the node-local Engram DP group.""" |
| 710 | dp_metadata = get_forward_context().dp_metadata |
| 711 | if dp_metadata is None: |
| 712 | raise RuntimeError("a DP-shared engram table needs DP token metadata") |
| 713 | group = get_engram_dp_group() |
| 714 | assert group is not None |
| 715 | # Engram groups are contiguous slices of the full DP group. |
| 716 | start = get_dp_group().rank_in_group - group.rank_in_group |
| 717 | return int( |
| 718 | dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max() |
| 719 | ) |
| 720 | |
| 721 | |
| 722 | def gather_engram_hashes( |
| 723 | hash_ids: torch.Tensor, *, dp_shared_memory: bool = False |
| 724 | ) -> torch.Tensor: |
| 725 | """Collect the n-gram ids of every DP replica sharing one table. |
| 726 | |
| 727 | Replicas are padded to a common token slot, so the gathered shape is |
| 728 | static under CUDA graph capture (where DP already pads alike). |
| 729 | """ |
| 730 | dp_group = get_engram_dp_group() |
| 731 | if dp_group is None or dp_shared_memory: |
| 732 | return hash_ids |
| 733 | slot = engram_gathered_num_tokens() |
| 734 | if hash_ids.shape[0] > slot: |
| 735 | raise ValueError("Engram token count exceeds the DP token slot") |
| 736 | if hash_ids.shape[0] < slot: |
| 737 | pad = hash_ids.new_full( |
| 738 | (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID |
| 739 | ) |
| 740 | hash_ids = torch.cat((hash_ids, pad)) |
| 741 | return dp_group.all_gather(hash_ids, dim=0) |
| 742 | |
| 743 | |
| 744 | class DPSharedEngramStorage: |
| 745 | """Registered host weights shared by a node-local DP group with one writer.""" |
| 746 | |
| 747 | def __init__( |
| 748 | self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator |
| 749 | ) -> None: |
| 750 | self.group = group |
| 751 | weight_bytes = num_rows * dim |
| 752 | storage = self._allocate(weight_bytes + weight_bytes // block_size) |
| 753 | self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim) |
| 754 | self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size) |
| 755 | self._views: tuple[torch.Tensor, torch.Tensor] | None = None |
| 756 | |
| 757 | def _allocate(self, num_bytes: int) -> torch.Tensor: |
| 758 | """Map and register one physical allocation across a node-local DP group.""" |
| 759 | group = self.group |
| 760 | with ExitStack() as stack: |
| 761 | path = error = None |
| 762 | if group.rank_in_group == 0: |
| 763 | try: |
| 764 | check_shm_free_space(num_bytes) |
| 765 | backing_file = stack.enter_context( |
| 766 | tempfile.NamedTemporaryFile(prefix="vllm_engram_", dir=SHM_PATH) |
| 767 | ) |
| 768 | backing_file.truncate(num_bytes) |
| 769 | path = backing_file.name |
| 770 | except Exception as exc: |
| 771 | error = f"{type(exc).__name__}: {exc}" |
| 772 | path, error = group.broadcast_object((path, error)) |
| 773 | if error is not None: |
| 774 | raise RuntimeError( |
| 775 | "Engram shared-memory creation failed on EDP rank 0: " + error |
| 776 | ) |
| 777 | |
| 778 | mapping = owner = tensor = None |
| 779 | try: |
| 780 | with open(path, "r+b") as file: |
| 781 | mapping = mmap.mmap(file.fileno(), num_bytes, flags=mmap.MAP_SHARED) |
| 782 | owner = np.frombuffer(mapping, dtype=np.uint8) |
| 783 | pointer = owner.ctypes.data |
| 784 | tensor = torch.from_numpy(owner) |
| 785 | result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0) |
| 786 | if result.value != 0: |
| 787 | raise RuntimeError(f"cudaHostRegister failed: {result}") |
| 788 | finalizer = weakref.finalize(owner, self._unregister, mapping, pointer) |
| 789 | finalizer.atexit = False # type: ignore[misc] |
| 790 | # The UVA helper otherwise allocates a private pinned copy. |
| 791 | if not tensor.is_pinned(): |
| 792 | raise RuntimeError( |
| 793 | "cudaHostRegister did not pin the shared Engram mapping" |
| 794 | ) |
| 795 | except Exception as exc: |
| 796 | error = f"{type(exc).__name__}: {exc}" |
| 797 | |
| 798 | # Also fences peer mappings before the leader unlinks the file. |
| 799 | errors: list[str | None] = [None] * group.world_size |
| 800 | torch.distributed.all_gather_object(errors, error, group=group.cpu_group) |
| 801 | failures = "; ".join( |
| 802 | f"EDP rank {rank}: {error}" |
| 803 | for rank, error in enumerate(errors) |
| 804 | if error is not None |
| 805 | ) |
| 806 | if failures: |
| 807 | # Dropping the owner runs the finalizer; the mapping then unmaps. |
| 808 | del tensor, owner, mapping |
| 809 | raise RuntimeError( |
| 810 | "Engram shared-memory initialization failed: " + failures |
| 811 | ) |
| 812 | assert tensor is not None |
| 813 | return tensor |
| 814 | |
| 815 | @staticmethod |
| 816 | def _unregister(mapping: mmap.mmap, pointer: int) -> None: |
| 817 | # Torch storage retains the numpy owner, including through cached UVA views. |
| 818 | # Keep its mmap alive until CUDA has released the registration. |
| 819 | result = torch.cuda.cudart().cudaHostUnregister(pointer) |
| 820 | if result.value != 0: |
| 821 | logger.warning("Engram cudaHostUnregister failed: %s", result) |
| 822 | |
| 823 | def load_weight( |
| 824 | self, param: torch.nn.Parameter, loaded_weight: torch.Tensor |
| 825 | ) -> None: |
| 826 | if self.group.rank_in_group == 0: |
| 827 | _engram_head_shard_weight_loader(param, loaded_weight) |
| 828 | # Read order may differ across ranks. Equal load counts ensure all shared |
| 829 | # weights are ready after the last weight-loader call returns. |
| 830 | torch.distributed.barrier(group=self.group.cpu_group) |
| 831 | |
| 832 | def get_views( |
| 833 | self, weight: torch.Tensor, scales: torch.Tensor |
| 834 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 835 | if (weight.data_ptr(), scales.data_ptr()) != ( |
| 836 | self.weight.data_ptr(), |
| 837 | self.weight_scale_inv.data_ptr(), |
| 838 | ): |
| 839 | raise RuntimeError("Shared Engram parameter storage must not be replaced") |
| 840 | if self._views is None: |
| 841 | self._views = ( |
| 842 | get_accelerator_view_from_cpu_tensor(self.weight), |
| 843 | get_accelerator_view_from_cpu_tensor(self.weight_scale_inv), |
| 844 | ) |
| 845 | return self._views |
| 846 | |
| 847 | |
| 848 | def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool: |
| 849 | """Whether co-located DP replicas exist and /dev/shm can hold the full tables.""" |
| 850 | if get_engram_dp_size() == 1: |
| 851 | logger.warning_once( |
| 852 | "Engram DP replicas are not co-located on one node; " |
| 853 | "storing the offloaded tables per rank instead of sharing them." |
| 854 | ) |
| 855 | return False |
| 856 | num_bytes = sum(layout.num_embeddings) * ( |
| 857 | layout.head_dim + layout.head_dim // block_size |
| 858 | ) |
| 859 | group = get_engram_dp_group() |
| 860 | assert group is not None |
| 861 | # Only the leader allocates, so every replica follows its check. |
| 862 | error = None |
| 863 | if group.rank_in_group == 0: |
| 864 | if not os.path.isdir(SHM_PATH): |
| 865 | error = f"{SHM_PATH} is not mounted" |
| 866 | else: |
| 867 | try: |
| 868 | check_shm_free_space(num_bytes, allocation_name="Engram tables") |
| 869 | except Exception as exc: |
| 870 | error = f"{type(exc).__name__}: {exc}" |
| 871 | error = group.broadcast_object(error) |
| 872 | if error is not None: |
| 873 | logger.warning_once( |
| 874 | "Sharding the offloaded Engram tables across DP replicas: %s", error |
| 875 | ) |
| 876 | return error is None |
| 877 | |
| 878 | |
| 879 | def _allocate_huge_page_storage(num_bytes: int) -> torch.Tensor | None: |
| 880 | """Register prefaulted huge pages, or return None for pinned-memory fallback.""" |
| 881 | try: |
| 882 | mapping = mmap.mmap(-1, num_bytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS) |
| 883 | mapping.madvise(mmap.MADV_HUGEPAGE) |
| 884 | except OSError as exc: |
| 885 | logger.warning( |
| 886 | "Engram huge-page allocation failed (%s); using pinned memory.", exc |
| 887 | ) |
| 888 | return None |
| 889 | |
| 890 | owner = np.frombuffer(mapping, dtype=np.uint8) |
| 891 | # Fault the pages in before CUDA pins them, so they can be huge. |
| 892 | owner[:: mmap.PAGESIZE] = 0 |
| 893 | tensor = torch.from_numpy(owner) |
| 894 | pointer = tensor.data_ptr() |
| 895 | result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0) |
| 896 | if result.value != 0: |
| 897 | logger.warning( |
| 898 | "Engram cudaHostRegister failed (%s); using pinned memory.", result |
| 899 | ) |
| 900 | return None |
| 901 | # Tensor views retain owner; its finalizer retains mapping until unregister. |
| 902 | finalizer = weakref.finalize( |
| 903 | owner, DPSharedEngramStorage._unregister, mapping, pointer |
| 904 | ) |
| 905 | finalizer.atexit = False # type: ignore[misc] |
| 906 | if not tensor.is_pinned(): |
| 907 | logger.warning( |
| 908 | "CUDA did not recognize Engram registration; using pinned memory." |
| 909 | ) |
| 910 | return None |
| 911 | return tensor |
| 912 | |
| 913 | |
| 914 | class ParallelEngramEmbedding(nn.Module): |
| 915 | """Hash heads with FP8 rows and per-block E8M0 scales. |
| 916 | |
| 917 | Heads are TP-sharded, and additionally DP-sharded across a node-local |
| 918 | group unless the table lives in (DP-shared) pinned host memory. |
| 919 | """ |
| 920 | |
| 921 | _weight_loader = staticmethod(_engram_head_shard_weight_loader) |
| 922 | _shared_memory: DPSharedEngramStorage | None = None |
| 923 | _packed: torch.Tensor | None = None |
| 924 | |
| 925 | def __init__( |
| 926 | self, |
| 927 | num_embeddings: int, |
| 928 | dim: int, |
| 929 | head_sizes: tuple[int, ...], |
| 930 | block_size: int = 32, |
| 931 | cpu_offload: bool = False, |
| 932 | dp_shared_memory: bool = False, |
| 933 | use_thp: bool = False, |
| 934 | ) -> None: |
| 935 | super().__init__() |
| 936 | assert head_sizes and all(size > 0 for size in head_sizes) |
| 937 | assert sum(head_sizes) <= num_embeddings |
| 938 | self.cpu_offload = cpu_offload |
| 939 | self.dp_shared_memory = dp_shared_memory |
| 940 | self.use_thp = use_thp |
| 941 | self.dp_size = get_engram_dp_size() |
| 942 | if dp_shared_memory: |
| 943 | if not cpu_offload: |
| 944 | raise ValueError("dp_shared_memory requires cpu_offload=True") |
| 945 | if self.dp_size <= 1: |
| 946 | raise ValueError( |
| 947 | "dp_shared_memory requires a node-local Engram DP " |
| 948 | f"group with size > 1; effective Engram DP size is {self.dp_size}. " |
| 949 | "Check that the node layout and rank placement allow complete " |
| 950 | "DP replicas to be co-located." |
| 951 | ) |
| 952 | self.dp_size = 1 |
| 953 | if cpu_offload and not is_uva_available(): |
| 954 | raise RuntimeError("Engram CPU offload requires UVA support") |
| 955 | self._views: tuple[torch.Tensor, torch.Tensor] | None = None |
| 956 | self._view_src: tuple[int, int] | None = None |
| 957 | self.num_embeddings = num_embeddings |
| 958 | self.dim = dim |
| 959 | self.block_size = block_size |
| 960 | self.n_hash_cols = len(head_sizes) |
| 961 | self.tp_size = get_tensor_model_parallel_world_size() |
| 962 | num_shards, head_rank = self._get_shard_info() |
| 963 | self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, num_shards) |
| 964 | # TODO: Support row-wise sharding when there are too few hash heads. |
| 965 | assert (num_shards - 1) * self.part_n_hash_cols < self.n_hash_cols, ( |
| 966 | f"Engram sharding leaves ranks without hash heads: " |
| 967 | f"{self.n_hash_cols} heads over {num_shards} shards" |
| 968 | ) |
| 969 | self.head_start = head_rank * self.part_n_hash_cols |
| 970 | head_end = self.head_start + self.part_n_hash_cols |
| 971 | self.vocab_start_idx = sum(head_sizes[: self.head_start]) |
| 972 | self.vocab_end_idx = sum(head_sizes[:head_end]) |
| 973 | self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx |
| 974 | self._num_sms = torch.cuda.get_device_properties( |
| 975 | torch.accelerator.current_device_index() |
| 976 | ).multi_processor_count |
| 977 | weight, scales = self._allocate_weights() |
| 978 | self.weight = nn.Parameter(weight, requires_grad=False) |
| 979 | self.weight_scale_inv = nn.Parameter(scales, requires_grad=False) |
| 980 | for param in (self.weight, self.weight_scale_inv): |
| 981 | set_weight_attrs( |
| 982 | param, |
| 983 | { |
| 984 | "weight_loader": self._weight_loader, |
| 985 | "engram_vocab_start": self.vocab_start_idx, |
| 986 | }, |
| 987 | ) |
| 988 | if cpu_offload: |
| 989 | # Constant dummy values avoid randomizing huge CPU lookup tables. |
| 990 | set_weight_attrs(self.weight, {"dummy_weight_value": 1.0}) |
| 991 | # The ue8m0 encoding of scale 1.0 is exponent byte 127. |
| 992 | set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127}) |
| 993 | logger.info( |
| 994 | "Engram table offloaded to pinned host memory: %d rows x %d, " |
| 995 | "%.2f GiB %s", |
| 996 | self.part_num_embeddings, |
| 997 | dim, |
| 998 | self.part_num_embeddings * (dim + dim // block_size) / 1024**3, |
| 999 | "shared across DP replicas" if dp_shared_memory else "per rank", |
| 1000 | ) |
| 1001 | |
| 1002 | def _get_shard_info(self) -> tuple[int, int]: |
| 1003 | if self.dp_size == 1: |
| 1004 | return self.tp_size, get_tensor_model_parallel_rank() |
| 1005 | return self.tp_size * self.dp_size, engram_head_shard_rank() |
| 1006 | |
| 1007 | def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]: |
| 1008 | if self.dp_shared_memory: |
| 1009 | group = get_engram_dp_group() |
| 1010 | assert group is not None |
| 1011 | storage = DPSharedEngramStorage( |
| 1012 | self.part_num_embeddings, self.dim, self.block_size, group |
| 1013 | ) |
| 1014 | self._shared_memory = storage |
| 1015 | self._weight_loader = storage.load_weight |
| 1016 | return storage.weight, storage.weight_scale_inv |
| 1017 | scale_dim = self.dim // self.block_size |
| 1018 | if not self.cpu_offload: |
| 1019 | return ( |
| 1020 | torch.empty( |
| 1021 | self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn |
| 1022 | ), |
| 1023 | torch.empty(self.part_num_embeddings, scale_dim, dtype=torch.uint8), |
| 1024 | ) |
| 1025 | if self.use_thp: |
| 1026 | weight_bytes = self.part_num_embeddings * self.dim |
| 1027 | packed = _allocate_huge_page_storage( |
| 1028 | weight_bytes + weight_bytes // self.block_size |
| 1029 | ) |
| 1030 | if packed is not None: |
| 1031 | self._packed = packed |
| 1032 | return ( |
| 1033 | packed[:weight_bytes].view(torch.float8_e4m3fn).view(-1, self.dim), |
| 1034 | packed[weight_bytes:].view(-1, scale_dim), |
| 1035 | ) |
| 1036 | # Model initialization may be inside a CUDA device context. |
| 1037 | return ( |
| 1038 | torch.empty( |
| 1039 | self.part_num_embeddings, |
| 1040 | self.dim, |
| 1041 | dtype=torch.float8_e4m3fn, |
| 1042 | device="cpu", |
| 1043 | pin_memory=True, |
| 1044 | ), |
| 1045 | torch.empty( |
| 1046 | self.part_num_embeddings, |
| 1047 | scale_dim, |
| 1048 | dtype=torch.uint8, |
| 1049 | device="cpu", |
| 1050 | pin_memory=True, |
| 1051 | ), |
| 1052 | ) |
| 1053 | |
| 1054 | def collapse_huge_pages(self) -> None: |
| 1055 | """Best-effort MADV_COLLAPSE (Linux >= 6.1) of pages that faulted small.""" |
| 1056 | if self._packed is None: |
| 1057 | return |
| 1058 | addr, num_bytes = self._packed.data_ptr(), self._packed.nbytes |
| 1059 | libc = ctypes.CDLL(None, use_errno=True) |
| 1060 | libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int) |
| 1061 | if libc.madvise(addr, num_bytes, 25) != 0: |
| 1062 | logger.warning( |
| 1063 | "Engram MADV_COLLAPSE failed; keeping existing pages: %s", |
| 1064 | os.strerror(ctypes.get_errno()), |
| 1065 | ) |
| 1066 | |
| 1067 | def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: |
| 1068 | if self._shared_memory is not None: |
| 1069 | return self._shared_memory.get_views(self.weight, self.weight_scale_inv) |
| 1070 | if not self.cpu_offload: |
| 1071 | return self.weight.data, self.weight_scale_inv.data |
| 1072 | src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr()) |
| 1073 | if self._view_src != src: |
| 1074 | self._views = ( |
| 1075 | get_accelerator_view_from_cpu_tensor(self.weight.data), |
| 1076 | get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data), |
| 1077 | ) |
| 1078 | self._view_src = src |
| 1079 | assert self._views is not None |
| 1080 | return self._views |
| 1081 | |
| 1082 | def lookup( |
| 1083 | self, indices: torch.Tensor, out: torch.Tensor, background: bool = False |
| 1084 | ) -> None: |
| 1085 | """Look up local heads of [T, heads] into [T, local_heads, dim] bf16. |
| 1086 | |
| 1087 | `background` limits the grid to leave SMs for concurrent work. |
| 1088 | """ |
| 1089 | rows = indices.shape[0] * self.part_n_hash_cols |
| 1090 | if not rows: |
| 1091 | return |
| 1092 | weight, scales = self._storage() |
| 1093 | ids_stride_t, ids_stride_h = indices.stride() |
| 1094 | sort_min_rows, _ = _engram_lookup_thresholds(indices.device) |
| 1095 | sort = ( |
| 1096 | self.cpu_offload |
| 1097 | and self._packed is None |
| 1098 | and sort_min_rows is not None |
| 1099 | and rows >= sort_min_rows |
| 1100 | ) |
| 1101 | dst = indices |
| 1102 | if sort: |
| 1103 | # Small-page host tables miss the TLB per row; in table order, rows on |
| 1104 | # one page share it. |
| 1105 | head_end = self.head_start + self.part_n_hash_cols |
| 1106 | local = indices[:, self.head_start : head_end] - self.vocab_start_idx |
| 1107 | if (pad := self.part_n_hash_cols - local.shape[1]) > 0: |
| 1108 | local = nn.functional.pad(local, (0, pad), value=-1) |
| 1109 | indices, dst = local.flatten().sort() |
| 1110 | dst = dst.int() |
| 1111 | # The table dwarfs TLB reach, so a persistent grid near the SM count |
| 1112 | # beats one program per row; halve it to leave SMs for the main stream. |
| 1113 | tiles = triton.cdiv(rows, 16) |
| 1114 | grid = min(tiles, self._num_sms // 2 if background else self._num_sms) |
| 1115 | _engram_lookup_kernel[(grid,)]( |
| 1116 | weight, |
| 1117 | scales, |
| 1118 | indices, |
| 1119 | dst, |
| 1120 | out, |
| 1121 | self.vocab_start_idx, |
| 1122 | self.vocab_end_idx, |
| 1123 | rows, |
| 1124 | ids_stride_t, |
| 1125 | ids_stride_h, |
| 1126 | HEAD_START=self.head_start, |
| 1127 | LOCAL_HEADS=self.part_n_hash_cols, |
| 1128 | TOTAL_HEADS=self.n_hash_cols, |
| 1129 | DIM=self.dim, |
| 1130 | QUANT_BLOCK=self.block_size, |
| 1131 | BLOCK_R=16, |
| 1132 | GRID=grid, |
| 1133 | SORTED=sort, |
| 1134 | ) |
| 1135 | |
| 1136 | def forward(self, indices: torch.Tensor) -> torch.Tensor: |
| 1137 | """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim] |
| 1138 | bf16, gathered from all shards for this replica's tokens.""" |
| 1139 | num_tokens = indices.shape[0] |
| 1140 | if self.dp_size > 1: |
| 1141 | indices = gather_engram_hashes(indices) |
| 1142 | out = torch.empty( |
| 1143 | (indices.shape[0], self.part_n_hash_cols, self.dim), |
| 1144 | dtype=torch.bfloat16, |
| 1145 | device=indices.device, |
| 1146 | ) |
| 1147 | self.lookup(indices, out) |
| 1148 | if self.dp_size > 1: |
| 1149 | out = _gather_engram_rows(out, num_tokens) |
| 1150 | if self.tp_size > 1: |
| 1151 | out = tensor_model_parallel_all_gather(out, dim=1) |
| 1152 | return out[:, : self.n_hash_cols] |
| 1153 | |
| 1154 | |
| 1155 | @triton.jit(do_not_specialize=["num_kv_tokens"]) |
| 1156 | def _fused_engram_post_wkv_kernel( |
| 1157 | hidden_states, |
| 1158 | kv, |
| 1159 | q_weight, |
| 1160 | k_weight, |
| 1161 | token_mask, |
| 1162 | output, |
| 1163 | num_kv_tokens, |
| 1164 | hidden_stride_t, |
| 1165 | hidden_stride_h, |
| 1166 | hidden_stride_d, |
| 1167 | kv_stride_t, |
| 1168 | kv_stride_d, |
| 1169 | q_stride_h, |
| 1170 | q_stride_d, |
| 1171 | k_stride_h, |
| 1172 | k_stride_d, |
| 1173 | mask_stride, |
| 1174 | output_stride_t, |
| 1175 | output_stride_h, |
| 1176 | output_stride_d, |
| 1177 | eps, |
| 1178 | clamp_value, |
| 1179 | DIM: tl.constexpr, |
| 1180 | HC_MULT: tl.constexpr, |
| 1181 | BLOCK_SIZE: tl.constexpr, |
| 1182 | HAS_MASK: tl.constexpr, |
| 1183 | ): |
| 1184 | program_idx = tl.program_id(0) |
| 1185 | token_idx = program_idx // HC_MULT |
| 1186 | hc_idx = program_idx % HC_MULT |
| 1187 | token_idx = token_idx.to(tl.int64) |
| 1188 | source_idx = token_idx |
| 1189 | source_valid = source_idx < num_kv_tokens |
| 1190 | |
| 1191 | dim_offsets = tl.arange(0, BLOCK_SIZE) |
| 1192 | dim_valid = dim_offsets < DIM |
| 1193 | hidden = tl.load( |
| 1194 | hidden_states |
| 1195 | + token_idx * hidden_stride_t |
| 1196 | + hc_idx * hidden_stride_h |
| 1197 | + dim_offsets * hidden_stride_d, |
| 1198 | mask=dim_valid, |
| 1199 | other=0.0, |
| 1200 | ).to(tl.float32) |
| 1201 | key = tl.load( |
| 1202 | kv + source_idx * kv_stride_t + (hc_idx * DIM + dim_offsets) * kv_stride_d, |
| 1203 | mask=source_valid & dim_valid, |
| 1204 | other=0.0, |
| 1205 | ).to(tl.float32) |
| 1206 | q = tl.load( |
| 1207 | q_weight + hc_idx * q_stride_h + dim_offsets * q_stride_d, |
| 1208 | mask=dim_valid, |
| 1209 | other=0.0, |
| 1210 | ).to(tl.float32) |
| 1211 | k = tl.load( |
| 1212 | k_weight + hc_idx * k_stride_h + dim_offsets * k_stride_d, |
| 1213 | mask=dim_valid, |
| 1214 | other=0.0, |
| 1215 | ).to(tl.float32) |
| 1216 | |
| 1217 | hidden_rms = tl.rsqrt(tl.sum(hidden * hidden, axis=0) / DIM + eps) |
| 1218 | key_rms = tl.rsqrt(tl.sum(key * key, axis=0) / DIM + eps) |
| 1219 | dot = tl.sum(hidden * q * k * key, axis=0) |
| 1220 | dot *= hidden_rms * key_rms * tl.rsqrt(DIM * 1.0) |
| 1221 | gate_input = tl.sqrt(tl.maximum(tl.abs(dot), clamp_value)) |
| 1222 | gate_input = tl.where(dot < 0.0, -gate_input, gate_input) |
| 1223 | gate = tl.sigmoid(gate_input) |
| 1224 | if HAS_MASK: |
| 1225 | active = tl.load( |
| 1226 | token_mask + source_idx * mask_stride, |
| 1227 | mask=source_valid, |
| 1228 | other=0, |
| 1229 | ) |
| 1230 | gate = tl.where(active, gate, 0.0) |
| 1231 | |
| 1232 | value = tl.load( |
| 1233 | kv + source_idx * kv_stride_t + (HC_MULT * DIM + dim_offsets) * kv_stride_d, |
| 1234 | mask=source_valid & dim_valid, |
| 1235 | other=0.0, |
| 1236 | ).to(tl.float32) |
| 1237 | tl.store( |
| 1238 | output |
| 1239 | + token_idx * output_stride_t |
| 1240 | + hc_idx * output_stride_h |
| 1241 | + dim_offsets * output_stride_d, |
| 1242 | hidden + gate * value, |
| 1243 | mask=dim_valid, |
| 1244 | ) |
| 1245 | |
| 1246 | |
| 1247 | @triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"]) |
| 1248 | def _engram_select_rows_kernel( |
| 1249 | gathered, |
| 1250 | output, |
| 1251 | num_tokens, |
| 1252 | token_start, |
| 1253 | num_elements, |
| 1254 | LOCAL_WIDTH: tl.constexpr, |
| 1255 | WIDTH: tl.constexpr, |
| 1256 | BLOCK_SIZE: tl.constexpr, |
| 1257 | ): |
| 1258 | offsets = tl.program_id(0).to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 1259 | tokens = token_start + offsets // WIDTH |
| 1260 | cols = offsets % WIDTH |
| 1261 | source = (cols // LOCAL_WIDTH * num_tokens + tokens) * LOCAL_WIDTH |
| 1262 | source += cols % LOCAL_WIDTH |
| 1263 | values = tl.load( |
| 1264 | gathered + source, (offsets < num_elements) & (tokens < num_tokens), other=0 |
| 1265 | ) |
| 1266 | tl.store(output + offsets, values, offsets < num_elements) |
| 1267 | |
| 1268 | |
| 1269 | def _engram_select_rows( |
| 1270 | gathered: torch.Tensor, |
| 1271 | output: torch.Tensor, |
| 1272 | source_tokens: int, |
| 1273 | token_start: int, |
| 1274 | local_width: int, |
| 1275 | ) -> None: |
| 1276 | """Copy one token window out of a rank-major gathered buffer. |
| 1277 | |
| 1278 | Both gathers land rank-major ([rank][token][local width]); this walks the |
| 1279 | window the rank keeps and lays its ranks out side by side as width. |
| 1280 | """ |
| 1281 | if output.numel() == 0: |
| 1282 | return |
| 1283 | _engram_select_rows_kernel[(triton.cdiv(output.numel(), 1024),)]( |
| 1284 | gathered, |
| 1285 | output, |
| 1286 | source_tokens, |
| 1287 | token_start, |
| 1288 | output.numel(), |
| 1289 | local_width, |
| 1290 | output.shape[1] * output.shape[2], |
| 1291 | BLOCK_SIZE=1024, |
| 1292 | ) |
| 1293 | |
| 1294 | |
| 1295 | def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor: |
| 1296 | """Exchange DP tokens for heads, retaining only this replica's tokens.""" |
| 1297 | dp_group = get_engram_dp_group() |
| 1298 | assert dp_group is not None |
| 1299 | slot, remainder = divmod(staged.shape[0], dp_group.world_size) |
| 1300 | assert remainder == 0 and 0 <= num_tokens <= slot |
| 1301 | gathered = dp_group.all_gather(staged, dim=0) |
| 1302 | local_heads, dim = staged.shape[1:] |
| 1303 | rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim)) |
| 1304 | _engram_select_rows( |
| 1305 | gathered, |
| 1306 | rows, |
| 1307 | staged.shape[0], |
| 1308 | dp_group.rank_in_group * slot, |
| 1309 | local_heads * dim, |
| 1310 | ) |
| 1311 | return rows |
| 1312 | |
| 1313 | |
| 1314 | class Engram(nn.Module): |
| 1315 | """Writes an n-gram lookup into the residual stream, gated by how well it |
| 1316 | matches that stream. |
| 1317 | |
| 1318 | The hash ids fetch `n_hash_cols` rows; `wkv` turns them into one key per |
| 1319 | hc copy plus a shared value. The gate is a normalized dot product of the |
| 1320 | stream against the key, signed-sqrt'ed before the sigmoid (matching the |
| 1321 | training kernel). |
| 1322 | """ |
| 1323 | |
| 1324 | _prefetch_stream: torch.cuda.Stream | None = None |
| 1325 | _prefetch_done: torch.cuda.Event | None = None |
| 1326 | |
| 1327 | def __init__( |
| 1328 | self, |
| 1329 | config, |
| 1330 | quant_config: QuantizationConfig | None, |
| 1331 | layout: EngramLayout, |
| 1332 | layer_hash_index: int, |
| 1333 | use_sequence_parallel: bool, |
| 1334 | prefix: str, |
| 1335 | prefetch_stream: torch.cuda.Stream | None = None, |
| 1336 | ) -> None: |
| 1337 | super().__init__() |
| 1338 | # Layers sharing one stream serialize their offloaded lookups, so they |
| 1339 | # take turns instead of jointly starving decoder compute of SMs. |
| 1340 | self._prefetch_stream = prefetch_stream |
| 1341 | self.layer_hash_index = layer_hash_index |
| 1342 | self.dim = config.hidden_size |
| 1343 | self.hc_mult = config.hc_mult |
| 1344 | self.eps = config.rms_norm_eps |
| 1345 | self.clamp_value = 1e-6 |
| 1346 | self.use_sequence_parallel = use_sequence_parallel |
| 1347 | |
| 1348 | # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight`` |
| 1349 | # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight`` |
| 1350 | # suffix rule. |
| 1351 | self.embed_tokens = self._create_embedding(layout, layer_hash_index) |
| 1352 | n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads |
| 1353 | # Without sequence parallelism every TP rank holds every token, so |
| 1354 | # shard the output columns instead of replicating the projection. |
| 1355 | self.wkv = ColumnParallelLinear( |
| 1356 | n_hash_cols * layout.head_dim, |
| 1357 | self.dim * (self.hc_mult + 1), |
| 1358 | bias=False, |
| 1359 | gather_output=True, |
| 1360 | quant_config=quant_config, |
| 1361 | return_bias=False, |
| 1362 | prefix=f"{prefix}.wkv", |
| 1363 | disable_tp=use_sequence_parallel, |
| 1364 | ) |
| 1365 | self.q_weight = nn.Parameter( |
| 1366 | torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16), |
| 1367 | requires_grad=False, |
| 1368 | ) |
| 1369 | self.k_weight = nn.Parameter( |
| 1370 | torch.empty(self.hc_mult, self.dim, dtype=torch.bfloat16), |
| 1371 | requires_grad=False, |
| 1372 | ) |
| 1373 | |
| 1374 | max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens |
| 1375 | self._init_staging(max_tokens, layout.head_dim) |
| 1376 | |
| 1377 | def _create_embedding( |
| 1378 | self, layout: EngramLayout, layer_hash_index: int |
| 1379 | ) -> ParallelEngramEmbedding: |
| 1380 | engram_config = get_current_vllm_config().engram_config |
| 1381 | assert engram_config is not None |
| 1382 | return ParallelEngramEmbedding( |
| 1383 | layout.num_embeddings[layer_hash_index], |
| 1384 | layout.head_dim, |
| 1385 | tuple(size for order in layout.primes[layer_hash_index] for size in order), |
| 1386 | cpu_offload=engram_config.cpu_offload, |
| 1387 | dp_shared_memory=bool(engram_config.dp_shared_memory), |
| 1388 | use_thp=engram_config.use_thp, |
| 1389 | ) |
| 1390 | |
| 1391 | def _init_staging(self, max_tokens: int, head_dim: int) -> None: |
| 1392 | # Persistent storage keeps lookup addresses stable across graph replays. |
| 1393 | self.staged_rows = torch.empty( |
| 1394 | max_tokens * self.embed_tokens.dp_size, |
| 1395 | self.embed_tokens.part_n_hash_cols, |
| 1396 | head_dim, |
| 1397 | dtype=torch.bfloat16, |
| 1398 | ) |
| 1399 | if not self.embed_tokens.cpu_offload: |
| 1400 | self._prefetch_stream = None |
| 1401 | return |
| 1402 | if self._prefetch_stream is None: |
| 1403 | self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device) |
| 1404 | self._prefetch_done = torch.cuda.Event() |
| 1405 | |
| 1406 | def prepare_embeddings(self, hash_ids: torch.Tensor) -> None: |
| 1407 | """Look up, or prefetch from host memory, rows before the decoder layers |
| 1408 | consume them; DP-sharded tables take the group's gathered hash IDs.""" |
| 1409 | rows = self.staged_rows[: hash_ids.shape[0]] |
| 1410 | assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small" |
| 1411 | if self._prefetch_stream is None: |
| 1412 | self.embed_tokens.lookup(hash_ids, rows) |
| 1413 | else: |
| 1414 | self._start_prefetch(hash_ids, rows) |
| 1415 | |
| 1416 | @eager_break_during_capture |
| 1417 | def _start_prefetch(self, hash_ids: torch.Tensor, rows: torch.Tensor) -> None: |
| 1418 | # Eager boundaries let the lookup span piecewise graph segments, and |
| 1419 | # decide per replay (on the replay stream) whether it runs inline. |
| 1420 | num_tokens = hash_ids.shape[0] |
| 1421 | # All DP ranks must agree, or EP collectives wait on the slowest one. |
| 1422 | if is_forward_context_available() and (dp := get_forward_context().dp_metadata): |
| 1423 | num_tokens = int(dp.num_tokens_across_dp_cpu.max()) |
| 1424 | current = stream = torch.cuda.current_stream() |
| 1425 | # Big lookups stall persistent main-stream kernels anyway: run them inline. |
| 1426 | _, inline_min_tokens = _engram_lookup_thresholds(hash_ids.device) |
| 1427 | if inline_min_tokens is not None and num_tokens >= inline_min_tokens: |
| 1428 | self.embed_tokens.lookup(hash_ids, rows) |
| 1429 | else: |
| 1430 | stream = self._prefetch_stream |
| 1431 | assert stream is not None |
| 1432 | stream.wait_stream(current) |
| 1433 | # Keep temporary hash storage alive until lookup finishes reading it. |
| 1434 | hash_ids.record_stream(stream) |
| 1435 | with torch.cuda.stream(stream): |
| 1436 | self.embed_tokens.lookup(hash_ids, rows, background=True) |
| 1437 | assert self._prefetch_done is not None |
| 1438 | self._prefetch_done.record(stream) |
| 1439 | |
| 1440 | @eager_break_during_capture |
| 1441 | def _finish_prefetch(self, event: torch.cuda.Event) -> None: |
| 1442 | # Wait for this layer's lookup only, not for later layers on the stream. |
| 1443 | torch.cuda.current_stream().wait_event(event) |
| 1444 | |
| 1445 | def _ready_rows(self, num_tokens: int) -> torch.Tensor: |
| 1446 | if self._prefetch_done is not None: |
| 1447 | self._finish_prefetch(self._prefetch_done) |
| 1448 | if self.embed_tokens.dp_size > 1: |
| 1449 | slot = engram_gathered_num_tokens() |
| 1450 | staged = self.staged_rows[: slot * self.embed_tokens.dp_size] |
| 1451 | return _gather_engram_rows(staged, num_tokens) |
| 1452 | return self.staged_rows[:num_tokens] |
| 1453 | |
| 1454 | def embed(self, hash_ids: torch.Tensor) -> torch.Tensor: |
| 1455 | """Gather heads, returning only local tokens when SP is enabled.""" |
| 1456 | rows = self._ready_rows(hash_ids.shape[0]) |
| 1457 | if self.embed_tokens.tp_size == 1: |
| 1458 | return rows[:, : self.embed_tokens.n_hash_cols] |
| 1459 | if self.use_sequence_parallel: |
| 1460 | tp_size = self.embed_tokens.tp_size |
| 1461 | num_tokens, local_heads, dim = rows.shape |
| 1462 | gathered = tensor_model_parallel_all_gather(rows, dim=0) |
| 1463 | chunk = (num_tokens + tp_size - 1) // tp_size |
| 1464 | out = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim)) |
| 1465 | _engram_select_rows( |
| 1466 | gathered, |
| 1467 | out, |
| 1468 | num_tokens, |
| 1469 | get_tensor_model_parallel_rank() * chunk, |
| 1470 | local_heads * dim, |
| 1471 | ) |
| 1472 | return out |
| 1473 | rows = tensor_model_parallel_all_gather(rows, dim=1) |
| 1474 | return rows[:, : self.embed_tokens.n_hash_cols] |
| 1475 | |
| 1476 | def forward( |
| 1477 | self, |
| 1478 | hidden_states: torch.Tensor, |
| 1479 | hash_ids: torch.Tensor, |
| 1480 | token_mask: torch.Tensor | None = None, |
| 1481 | ) -> torch.Tensor: |
| 1482 | """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all |
| 1483 | tokens, pre sequence-parallel shard); token_mask: [T], False shuts |
| 1484 | the gate so those positions pass through untouched.""" |
| 1485 | kv = self.wkv(self.embed(hash_ids).flatten(-2)) |
| 1486 | num_kv_tokens = hash_ids.shape[0] |
| 1487 | assert token_mask is None or token_mask.shape == (num_kv_tokens,) |
| 1488 | if self.use_sequence_parallel: |
| 1489 | tp_size = get_tensor_model_parallel_world_size() |
| 1490 | tp_rank = get_tensor_model_parallel_rank() |
| 1491 | shard_size = (num_kv_tokens + tp_size - 1) // tp_size |
| 1492 | assert hidden_states.shape[0] == shard_size |
| 1493 | start = min(tp_rank * shard_size, num_kv_tokens) |
| 1494 | num_kv_tokens = min(shard_size, num_kv_tokens - start) |
| 1495 | if token_mask is not None: |
| 1496 | token_mask = token_mask[start : start + num_kv_tokens] |
| 1497 | |
| 1498 | num_tokens, hc_mult, dim = hidden_states.shape |
| 1499 | assert hc_mult == self.hc_mult and dim == self.dim |
| 1500 | assert kv.ndim == 2 and kv.shape[1] == (hc_mult + 1) * dim |
| 1501 | output = torch.empty_like(hidden_states) |
| 1502 | if num_tokens == 0: |
| 1503 | return output |
| 1504 | |
| 1505 | block_size = triton.next_power_of_2(dim) |
| 1506 | num_warps = 8 if block_size >= 2048 else 4 |
| 1507 | mask = token_mask if token_mask is not None else hidden_states |
| 1508 | _fused_engram_post_wkv_kernel[(num_tokens * hc_mult,)]( |
| 1509 | hidden_states, |
| 1510 | kv, |
| 1511 | self.q_weight, |
| 1512 | self.k_weight, |
| 1513 | mask, |
| 1514 | output, |
| 1515 | num_kv_tokens, |
| 1516 | hidden_states.stride(0), |
| 1517 | hidden_states.stride(1), |
| 1518 | hidden_states.stride(2), |
| 1519 | kv.stride(0), |
| 1520 | kv.stride(1), |
| 1521 | self.q_weight.stride(0), |
| 1522 | self.q_weight.stride(1), |
| 1523 | self.k_weight.stride(0), |
| 1524 | self.k_weight.stride(1), |
| 1525 | token_mask.stride(0) if token_mask is not None else 0, |
| 1526 | output.stride(0), |
| 1527 | output.stride(1), |
| 1528 | output.stride(2), |
| 1529 | self.eps, |
| 1530 | self.clamp_value, |
| 1531 | DIM=dim, |
| 1532 | HC_MULT=hc_mult, |
| 1533 | BLOCK_SIZE=block_size, |
| 1534 | HAS_MASK=token_mask is not None, |
| 1535 | num_warps=num_warps, |
| 1536 | ) |
| 1537 | return output |