| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import bisect |
| 4 | import inspect |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.models.deepseek_v41.common import engram as engram_ops |
| 11 | from vllm.models.deepseek_v41.common.engram import ( |
| 12 | Engram, |
| 13 | NgramHashState, |
| 14 | ParallelEngramEmbedding, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | |
| 18 | |
| 19 | def _reference_engram_post_wkv( |
| 20 | hidden_states, |
| 21 | kv, |
| 22 | q_weight, |
| 23 | k_weight, |
| 24 | token_mask, |
| 25 | eps, |
| 26 | clamp_value, |
| 27 | kv_start, |
| 28 | ): |
| 29 | num_tokens, hc_mult, dim = hidden_states.shape |
| 30 | source = torch.arange(num_tokens, device=kv.device) + kv_start |
| 31 | valid = source < kv.shape[0] |
| 32 | local_kv = kv.new_zeros(num_tokens, kv.shape[1]) |
| 33 | local_kv[valid] = kv[source[valid]] |
| 34 | |
| 35 | key = local_kv[:, : hc_mult * dim].float().view(num_tokens, hc_mult, dim) |
| 36 | value = local_kv[:, hc_mult * dim :].float() |
| 37 | h = hidden_states.float() |
| 38 | weight = q_weight.float() * k_weight.float() |
| 39 | rstd = torch.rsqrt(h.square().mean(-1) + eps) * torch.rsqrt( |
| 40 | key.square().mean(-1) + eps |
| 41 | ) |
| 42 | dot = (h * weight * key).sum(-1) * rstd * dim**-0.5 |
| 43 | gate = torch.sigmoid(torch.copysign(dot.abs().clamp_min(clamp_value).sqrt(), dot)) |
| 44 | if token_mask is not None: |
| 45 | local_mask = torch.zeros(num_tokens, dtype=torch.bool, device=kv.device) |
| 46 | local_mask[valid] = token_mask[source[valid]] |
| 47 | gate = gate.masked_fill(~local_mask.unsqueeze(-1), 0) |
| 48 | return (h + gate.unsqueeze(-1) * value.unsqueeze(-2)).to(hidden_states.dtype) |
| 49 | |
| 50 | |
| 51 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 52 | @pytest.mark.parametrize( |
| 53 | "num_tokens,num_kv_tokens,hc_mult,dim,tp_size,tp_rank,use_mask", |
| 54 | [ |
| 55 | (1, 1, 4, 5120, 1, 0, False), |
| 56 | (3, 5, 4, 96, 2, 0, True), |
| 57 | (3, 5, 2, 257, 2, 1, False), |
| 58 | (3, 5, 4, 96, 2, 1, True), |
| 59 | (1, 1, 4, 96, 4, 3, True), |
| 60 | (0, 0, 4, 96, 4, 0, False), |
| 61 | ], |
| 62 | ) |
| 63 | def test_fused_engram_post_wkv_matches_reference( |
| 64 | num_tokens, |
| 65 | num_kv_tokens, |
| 66 | hc_mult, |
| 67 | dim, |
| 68 | tp_size, |
| 69 | tp_rank, |
| 70 | use_mask, |
| 71 | monkeypatch, |
| 72 | ): |
| 73 | """Preserve gate math, WKV layout, masking, and SP padding semantics.""" |
| 74 | torch.manual_seed(0) |
| 75 | device = "cuda" |
| 76 | hidden_states = torch.randn( |
| 77 | num_tokens, hc_mult, dim, dtype=torch.bfloat16, device=device |
| 78 | ) |
| 79 | kv = torch.randn( |
| 80 | num_kv_tokens, |
| 81 | (hc_mult + 1) * dim, |
| 82 | dtype=torch.bfloat16, |
| 83 | device=device, |
| 84 | ) |
| 85 | q_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device) |
| 86 | k_weight = torch.randn(hc_mult, dim, dtype=torch.bfloat16, device=device) |
| 87 | token_mask = None |
| 88 | if use_mask: |
| 89 | token_mask = torch.arange(num_kv_tokens, device=device) % 2 == 0 |
| 90 | eps = 1e-20 |
| 91 | clamp_value = 1e-6 |
| 92 | use_sequence_parallel = tp_size > 1 |
| 93 | shard_size = (num_kv_tokens + tp_size - 1) // tp_size |
| 94 | kv_start = tp_rank * shard_size if use_sequence_parallel else 0 |
| 95 | |
| 96 | expected = _reference_engram_post_wkv( |
| 97 | hidden_states, |
| 98 | kv, |
| 99 | q_weight, |
| 100 | k_weight, |
| 101 | token_mask, |
| 102 | eps, |
| 103 | clamp_value, |
| 104 | kv_start, |
| 105 | ) |
| 106 | monkeypatch.setattr( |
| 107 | engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size |
| 108 | ) |
| 109 | monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) |
| 110 | |
| 111 | module = Engram.__new__(Engram) |
| 112 | torch.nn.Module.__init__(module) |
| 113 | module.dim = dim |
| 114 | module.hc_mult = hc_mult |
| 115 | module.eps = eps |
| 116 | module.clamp_value = clamp_value |
| 117 | module.use_sequence_parallel = use_sequence_parallel |
| 118 | module.embed_tokens = torch.nn.Identity() |
| 119 | # `forward` reads rows staged by `prepare_embeddings`, so inject kv there. |
| 120 | module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1 |
| 121 | module.embed_tokens.dp_size = 1 |
| 122 | module.staged_rows = kv.unsqueeze(1) |
| 123 | if use_sequence_parallel: |
| 124 | padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size)) |
| 125 | module.staged_rows = padded[ |
| 126 | tp_rank * shard_size : (tp_rank + 1) * shard_size |
| 127 | ].unsqueeze(1) |
| 128 | module.wkv = torch.nn.Identity() |
| 129 | |
| 130 | def check_local_projection(_, args): |
| 131 | assert args[0].shape[0] == num_tokens |
| 132 | |
| 133 | module.wkv.register_forward_pre_hook(check_local_projection) |
| 134 | module.q_weight = torch.nn.Parameter(q_weight, requires_grad=False) |
| 135 | module.k_weight = torch.nn.Parameter(k_weight, requires_grad=False) |
| 136 | |
| 137 | actual = module(hidden_states, kv.unsqueeze(1), token_mask) |
| 138 | |
| 139 | assert actual.shape == hidden_states.shape |
| 140 | assert actual.dtype == hidden_states.dtype |
| 141 | torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) |
| 142 | |
| 143 | |
| 144 | def _hash_state(use_slot_cache: bool) -> NgramHashState: |
| 145 | state = NgramHashState.__new__(NgramHashState) |
| 146 | torch.nn.Module.__init__(state) |
| 147 | state.block_size = 64 |
| 148 | state.use_slot_cache = use_slot_cache |
| 149 | state._cache = None |
| 150 | state._kv_cache_ref = None |
| 151 | state.swa_cache_module = torch.nn.Module() |
| 152 | state.swa_cache_module.kv_cache = torch.empty(0) |
| 153 | return state |
| 154 | |
| 155 | |
| 156 | @pytest.mark.parametrize("num_tokens", [0, 7]) |
| 157 | def test_engram_dummy_hashes_leave_history_untouched(num_tokens): |
| 158 | """Dummy forwards have no valid table IDs and do not alter cached history.""" |
| 159 | state = _hash_state(use_slot_cache=True) |
| 160 | state.multipliers = torch.empty(2, 4, dtype=torch.int64) |
| 161 | state.primes = torch.empty(2, 3, 8, dtype=torch.int64) |
| 162 | state._cache = torch.arange(16, dtype=torch.int32) |
| 163 | history = state._cache.clone() |
| 164 | |
| 165 | hashes, keep = state.dummy_hashes(torch.arange(num_tokens)) |
| 166 | |
| 167 | assert hashes.shape == (num_tokens, 2, 24) |
| 168 | assert hashes.dtype == torch.int32 |
| 169 | assert torch.all(hashes == engram_ops.DEAD_ID) |
| 170 | assert keep.shape == (num_tokens,) and keep.dtype == torch.bool |
| 171 | assert not keep.any() |
| 172 | torch.testing.assert_close(state._cache, history) |
| 173 | |
| 174 | |
| 175 | @pytest.mark.parametrize("num_blocks", [4, 100]) |
| 176 | def test_engram_cache_follows_kv_cache_binding(num_blocks): |
| 177 | """Discard profiling history when the real KV cache replaces its storage.""" |
| 178 | state = _hash_state(use_slot_cache=True) |
| 179 | assert not state.ensure_cache() |
| 180 | |
| 181 | state.swa_cache_module.kv_cache = torch.empty(4, 64, 1) |
| 182 | assert state.ensure_cache() |
| 183 | profiling_cache = state._cache |
| 184 | profiling_cache.fill_(7) |
| 185 | assert state.ensure_cache() |
| 186 | assert state._cache is profiling_cache |
| 187 | assert torch.all(state._cache == 7) |
| 188 | |
| 189 | state.swa_cache_module.kv_cache = torch.empty(num_blocks, 64, 1) |
| 190 | assert state.ensure_cache() |
| 191 | assert state._cache is not profiling_cache |
| 192 | assert state._cache.shape == (num_blocks * 64,) |
| 193 | assert torch.count_nonzero(state._cache) == 0 |
| 194 | |
| 195 | state.swa_cache_module.kv_cache = torch.empty(0) |
| 196 | assert not state.ensure_cache() |
| 197 | assert state._cache is None |
| 198 | |
| 199 | |
| 200 | def test_engram_without_slot_cache_only_gates_on_kv_binding(): |
| 201 | """The V2 runner supplies every lookback, so no slot cache is allocated.""" |
| 202 | state = _hash_state(use_slot_cache=False) |
| 203 | assert not state.ensure_cache() |
| 204 | state.swa_cache_module.kv_cache = torch.empty(4, 64, 1) |
| 205 | assert state.ensure_cache() |
| 206 | assert state._cache is None |
| 207 | |
| 208 | |
| 209 | def _reference_hashes( |
| 210 | ids, |
| 211 | positions, |
| 212 | slots, |
| 213 | starts, |
| 214 | tables, |
| 215 | dead, |
| 216 | cache, |
| 217 | token_map, |
| 218 | multipliers, |
| 219 | primes, |
| 220 | offsets, |
| 221 | ): |
| 222 | """Scalar n-gram oracle with persistent, physically addressed token history.""" |
| 223 | for token, slot, is_dead in zip(ids, slots, dead): |
| 224 | if slot >= 0: |
| 225 | cache[slot] = -1 if is_dead else token_map[token] |
| 226 | result = [] |
| 227 | for i, position in enumerate(positions): |
| 228 | req = min(bisect.bisect_right(starts[1:], i), len(tables) - 1) |
| 229 | history, blocked = [], False |
| 230 | for shift in range(4): |
| 231 | lookback = position - shift |
| 232 | p = max(0, min(lookback, len(tables[req]) * 16 - 1)) |
| 233 | source = cache[tables[req][p // 16] * 16 + p % 16] |
| 234 | blocked |= lookback < 0 or source == -1 |
| 235 | history.append(0 if blocked else source) |
| 236 | layers = [] |
| 237 | for layer, row in enumerate(multipliers): |
| 238 | rolling, hashes = history[0] * row[0], [] |
| 239 | for j, head_primes in enumerate(primes[layer], start=1): |
| 240 | rolling ^= history[j] * row[j] |
| 241 | for head, prime in enumerate(head_primes): |
| 242 | col = (j - 1) * len(head_primes) + head |
| 243 | hashes.append(rolling % prime + offsets[layer][col]) |
| 244 | layers.append(hashes) |
| 245 | result.append(layers) |
| 246 | return torch.tensor(result, dtype=torch.int32) |
| 247 | |
| 248 | |
| 249 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 250 | @pytest.mark.parametrize("query_len", [1, 17, 257]) |
| 251 | @pytest.mark.parametrize("capture", [False, True]) |
| 252 | @pytest.mark.parametrize("strided_slots", [False, True]) |
| 253 | @pytest.mark.parametrize("large_hashes", [False, True]) |
| 254 | def test_engram_hash_cache_replay(query_len, capture, strided_slots, large_hashes): |
| 255 | """Preserve history, dead tokens and padding as inputs change across replay.""" |
| 256 | num_tokens = 2 * query_len + 3 |
| 257 | blocks = (3 * query_len + 15) // 16 |
| 258 | tables = [list(range(blocks, 2 * blocks)), list(range(blocks))] |
| 259 | starts = [0, query_len, 2 * query_len] |
| 260 | |
| 261 | def tensor(data, dtype=torch.int64): |
| 262 | return torch.tensor(data, dtype=dtype, device="cuda") |
| 263 | |
| 264 | input_ids = tensor([0] * num_tokens) |
| 265 | positions = torch.zeros_like(input_ids) |
| 266 | slots = tensor([-1] * (num_tokens * (2 if strided_slots else 1))) |
| 267 | if strided_slots: |
| 268 | slots = slots[::2] |
| 269 | dead_mask = tensor([False] * num_tokens, torch.bool) |
| 270 | cache = tensor([-1] * (2 * blocks * 16), torch.int32) |
| 271 | token_map = [i * 3091 if large_hashes else i % 7 for i in range(32)] |
| 272 | multipliers = ( |
| 273 | engram_ops.compute_hash_multipliers((1, 14), 4, 99092).tolist() |
| 274 | if large_hashes |
| 275 | else [[3, 5, 7, 9], [11, 13, 15, 17]] |
| 276 | ) |
| 277 | primes = ( |
| 278 | [ |
| 279 | [ |
| 280 | [16000057 + 2 * (layer * 24 + j * 8 + h) for h in range(8)] |
| 281 | for j in range(3) |
| 282 | ] |
| 283 | for layer in range(2) |
| 284 | ] |
| 285 | if large_hashes |
| 286 | else [[[19], [23], [29]]] * 2 |
| 287 | ) |
| 288 | offsets = [] |
| 289 | for layer in primes: |
| 290 | row, offset = [], 0 |
| 291 | for heads in layer: |
| 292 | for prime in heads: |
| 293 | row.append(offset) |
| 294 | offset += prime |
| 295 | offsets.append(row) |
| 296 | state = NgramHashState.__new__(NgramHashState) |
| 297 | torch.nn.Module.__init__(state) |
| 298 | state.token_map = tensor(token_map, torch.int32) |
| 299 | state.multipliers = tensor(multipliers) |
| 300 | state.primes = tensor(primes) |
| 301 | state.offsets = tensor(offsets) |
| 302 | state._cache = cache |
| 303 | state.use_slot_cache = True |
| 304 | state.pad_id = 0 |
| 305 | state.block_size = 16 |
| 306 | lookback = tensor([[-1] * 3] * 2, torch.int32) |
| 307 | args = ( |
| 308 | input_ids, |
| 309 | positions, |
| 310 | tensor(starts, torch.int32), |
| 311 | dead_mask, |
| 312 | lookback, |
| 313 | torch.zeros_like(lookback, dtype=torch.bool), |
| 314 | slots, |
| 315 | tensor(tables, torch.int32), |
| 316 | ) |
| 317 | op = state |
| 318 | graph = None |
| 319 | if capture: |
| 320 | stream = torch.cuda.Stream() |
| 321 | stream.wait_stream(torch.cuda.current_stream()) |
| 322 | with torch.cuda.stream(stream): |
| 323 | for _ in range(3): |
| 324 | op(*args) |
| 325 | torch.cuda.current_stream().wait_stream(stream) |
| 326 | graph = torch.cuda.CUDAGraph() |
| 327 | with torch.cuda.graph(graph, stream=stream): |
| 328 | output = op(*args) |
| 329 | torch.cuda.current_stream().wait_stream(stream) |
| 330 | cache.fill_(-1) |
| 331 | expected_cache = [-1] * cache.numel() |
| 332 | |
| 333 | for step in range(3): |
| 334 | ids = [(i + 3 * step) % 32 for i in range(num_tokens)] |
| 335 | pos = list(range(step * query_len, (step + 1) * query_len)) * 2 |
| 336 | physical_slots = [ |
| 337 | tables[i // query_len][p // 16] * 16 + p % 16 for i, p in enumerate(pos) |
| 338 | ] |
| 339 | pos += [0] * 3 |
| 340 | physical_slots += [-1] * 3 |
| 341 | # A fully padded replay must leave every cache slot untouched. |
| 342 | if step == 2: |
| 343 | physical_slots = [-1] * num_tokens |
| 344 | dead = [(i + step) % 5 == 0 for i in range(num_tokens)] |
| 345 | input_ids.copy_(tensor(ids)) |
| 346 | positions.copy_(tensor(pos)) |
| 347 | slots.copy_(tensor(physical_slots)) |
| 348 | dead_mask.copy_(tensor(dead, torch.bool)) |
| 349 | if graph is None: |
| 350 | output = op(*args) |
| 351 | else: |
| 352 | graph.replay() |
| 353 | expected = _reference_hashes( |
| 354 | ids, |
| 355 | pos, |
| 356 | physical_slots, |
| 357 | starts, |
| 358 | tables, |
| 359 | dead, |
| 360 | expected_cache, |
| 361 | token_map, |
| 362 | multipliers, |
| 363 | primes, |
| 364 | offsets, |
| 365 | ) |
| 366 | if step != 2: |
| 367 | real = 2 * query_len |
| 368 | torch.testing.assert_close( |
| 369 | output[:real].cpu(), expected[:real], rtol=0, atol=0 |
| 370 | ) |
| 371 | assert cache.cpu().tolist() == expected_cache |
| 372 | |
| 373 | empty_output = state( |
| 374 | input_ids[:0], positions[:0], args[2], dead_mask[:0], *args[4:] |
| 375 | ) |
| 376 | assert empty_output.shape == (0, 2, len(offsets[0])) |
| 377 | assert cache.cpu().tolist() == expected_cache |
| 378 | |
| 379 | |
| 380 | def _hash_ids(requests, tables_meta, block_table, cache=None, capture=False): |
| 381 | """Run the op on a batch of (token_ids, start, window) requests; `window` |
| 382 | is the runner's [depth] lookback for that request (None = all unknown). |
| 383 | Token IDs 14 and 22 stand in for image tokens that break n-grams. |
| 384 | Returns one [len(token_ids), ...] hash tensor per request.""" |
| 385 | token_map, multipliers, primes, offsets, block_size = tables_meta |
| 386 | depth = multipliers.shape[1] - 1 |
| 387 | |
| 388 | def tensor(data, dtype=torch.int64): |
| 389 | return torch.tensor(data, dtype=dtype, device="cuda") |
| 390 | |
| 391 | input_ids, positions, slots, windows, starts = [], [], [], [], [0] |
| 392 | for req, (token_ids, start, window) in enumerate(requests): |
| 393 | pos = list(range(start, start + len(token_ids))) |
| 394 | input_ids += token_ids |
| 395 | positions += pos |
| 396 | slots += [ |
| 397 | int(block_table[req, p // block_size]) * block_size + p % block_size |
| 398 | for p in pos |
| 399 | ] |
| 400 | windows.append(window or [-1] * depth) |
| 401 | starts.append(starts[-1] + len(token_ids)) |
| 402 | window = tensor(windows, torch.int32) |
| 403 | state = _hash_state(use_slot_cache=cache is not None) |
| 404 | state.token_map = token_map |
| 405 | state.multipliers = multipliers |
| 406 | state.primes = primes |
| 407 | state.offsets = offsets |
| 408 | state.pad_id = 0 |
| 409 | state.block_size = block_size |
| 410 | state._cache = cache |
| 411 | input_tensor = tensor(input_ids, torch.int32) |
| 412 | args = ( |
| 413 | input_tensor, |
| 414 | tensor(positions), |
| 415 | tensor(starts, torch.int32), |
| 416 | (input_tensor == 14) | (input_tensor == 22), |
| 417 | window, |
| 418 | (window == 14) | (window == 22), |
| 419 | tensor(slots) if cache is not None else None, |
| 420 | block_table if cache is not None else None, |
| 421 | ) |
| 422 | out = state(*args) |
| 423 | if capture: |
| 424 | graph = torch.cuda.CUDAGraph() |
| 425 | with torch.cuda.graph(graph): |
| 426 | out = state(*args) |
| 427 | graph.replay() |
| 428 | return [out[a:b] for a, b in zip(starts, starts[1:])] |
| 429 | |
| 430 | |
| 431 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 432 | @pytest.mark.parametrize("runner", ["v2_window_only", "v1_prompt_window_plus_cache"]) |
| 433 | @pytest.mark.parametrize("capture", [False, True]) |
| 434 | def test_lookback_window_reproduces_single_instance(runner, capture): |
| 435 | """Decoding on a fresh instance whose prompt KV came from elsewhere (P/D, |
| 436 | offload) must hash like the instance that processed the whole sequence, |
| 437 | with two requests batched and chunks of 1-2 tokens so in-batch, window |
| 438 | and (V1) cache lookbacks mix within one step. |
| 439 | |
| 440 | V2 supplies every lookback from its device token history and needs no |
| 441 | slot cache. V1 supplies prompt positions only and reads generated |
| 442 | positions from the slot cache it fills itself. Without any window the |
| 443 | first decode token hashes stale slots (negative control).""" |
| 444 | block_size = 4 |
| 445 | block_table = torch.tensor( |
| 446 | [[3, 1, 5, 0], [2, 4, 6, 7]], dtype=torch.int32, device="cuda" |
| 447 | ) |
| 448 | tables_meta = ( |
| 449 | torch.arange(50, dtype=torch.int32, device="cuda"), |
| 450 | torch.tensor([[3, 5, 7, 9]], device="cuda"), |
| 451 | torch.tensor([[[97, 89], [83, 79], [73, 71]]], device="cuda"), |
| 452 | torch.tensor([[0, 97, 186, 269, 348, 421]], device="cuda"), |
| 453 | block_size, |
| 454 | ) |
| 455 | depth = 3 |
| 456 | prompts = [[11, 12, 13, 14, 15, 16], [31, 32, 33, 34, 35]] |
| 457 | decodes = [[21, 22, 23, 24], [41, 42, 43, 44]] |
| 458 | chunk_sizes = [1, 2, 1] |
| 459 | v2 = runner.startswith("v2") |
| 460 | |
| 461 | def fresh_cache(): |
| 462 | return torch.zeros(8 * block_size, dtype=torch.int32, device="cuda") |
| 463 | |
| 464 | reference = _hash_ids( |
| 465 | [(p + d, 0, None) for p, d in zip(prompts, decodes)], |
| 466 | tables_meta, |
| 467 | block_table, |
| 468 | cache=fresh_cache(), |
| 469 | ) |
| 470 | |
| 471 | histories = [list(p) for p in prompts] |
| 472 | cache = None if v2 else fresh_cache() |
| 473 | consumed = 0 |
| 474 | for chunk in chunk_sizes: |
| 475 | batch = [] |
| 476 | for history, prompt, decode in zip(histories, prompts, decodes): |
| 477 | start = len(history) |
| 478 | tokens = decode[consumed : consumed + chunk] |
| 479 | window = [ |
| 480 | history[start - 1 - j] if v2 or start - 1 - j < len(prompt) else -1 |
| 481 | for j in range(depth) |
| 482 | ] |
| 483 | batch.append((tokens, start, window)) |
| 484 | got = _hash_ids(batch, tables_meta, block_table, cache=cache, capture=capture) |
| 485 | for req, (tokens, start, _) in enumerate(batch): |
| 486 | assert torch.equal(got[req], reference[req][start : start + chunk]), ( |
| 487 | f"{runner} request {req} chunk at {start}" |
| 488 | ) |
| 489 | histories[req] += tokens |
| 490 | consumed += chunk |
| 491 | |
| 492 | unseeded = _hash_ids( |
| 493 | [(decodes[0][:1], len(prompts[0]), None)], |
| 494 | tables_meta, |
| 495 | block_table, |
| 496 | cache=fresh_cache(), |
| 497 | ) |
| 498 | assert not torch.equal(unseeded[0][0], reference[0][len(prompts[0])]) |
| 499 | |
| 500 | |
| 501 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 502 | def test_v2_model_state_gathers_lookback_window(): |
| 503 | """The window is gathered on device from the runner's token history.""" |
| 504 | from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState |
| 505 | |
| 506 | depth, max_num_reqs, max_model_len = 3, 4, 16 |
| 507 | state = DeepseekV41ModelState.__new__(DeepseekV41ModelState) |
| 508 | state.rope_state = None |
| 509 | state.supports_mm_inputs = False |
| 510 | state.prompt_embeds_state = None |
| 511 | state.lookback_token_ids = torch.full( |
| 512 | (max_num_reqs, depth), -1, dtype=torch.int32, device="cuda" |
| 513 | ) |
| 514 | |
| 515 | all_token_ids = torch.zeros(max_num_reqs, max_model_len, dtype=torch.int32) |
| 516 | all_token_ids[2, :8] = torch.tensor([11, 12, 13, 14, 15, 16, 21, 22]) |
| 517 | all_token_ids[0, :5] = torch.arange(1, 6) |
| 518 | req_states = SimpleNamespace( |
| 519 | num_computed_tokens=SimpleNamespace( |
| 520 | gpu=torch.tensor([0, 9, 8, 6], dtype=torch.int32, device="cuda") |
| 521 | ), |
| 522 | all_token_ids=SimpleNamespace(gpu=all_token_ids.cuda()), |
| 523 | ) |
| 524 | # Batch rows -> request state rows: a request two decode steps in, a fresh |
| 525 | # request, and one at the end of a 5-token prompt. |
| 526 | input_batch = SimpleNamespace( |
| 527 | idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda") |
| 528 | ) |
| 529 | window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"] |
| 530 | assert window.cpu().tolist() == [ |
| 531 | [22, 21, 16], |
| 532 | [-1, -1, -1], |
| 533 | [0, 0, 0], |
| 534 | [-1, -1, -1], |
| 535 | ] |
| 536 | # Graph capture runs on the dummy inputs and replays read the same buffer. |
| 537 | dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"] |
| 538 | assert dummy is window and torch.all(dummy == -1) |
| 539 | |
| 540 | |
| 541 | def _reference_lookup(weight, scale_inv, ids, start, end, block=32): |
| 542 | """The torch expression the fused kernel replaces.""" |
| 543 | mask = (ids < start) | (ids >= end) |
| 544 | local = (ids - start).masked_fill(mask, 0).long() |
| 545 | values = torch.nn.functional.embedding(local, weight) |
| 546 | scales = torch.nn.functional.embedding(local, scale_inv) |
| 547 | scales = (scales.to(torch.int32) << 23).view(torch.float32) |
| 548 | values = values.float().unflatten(-1, (-1, block)) |
| 549 | values = (values * scales.unsqueeze(-1)).flatten(-2).to(torch.bfloat16) |
| 550 | return values.masked_fill(mask.unsqueeze(-1), 0) |
| 551 | |
| 552 | |
| 553 | def _make_embedding(cpu_offload, rows=4096, dim=256, block=32): |
| 554 | layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding) |
| 555 | torch.nn.Module.__init__(layer) |
| 556 | layer.dim, layer.block_size, layer.tp_size = dim, block, 1 |
| 557 | layer.dp_size = 1 |
| 558 | layer.n_hash_cols = layer.part_n_hash_cols = 24 |
| 559 | layer.head_start = 0 |
| 560 | layer.cpu_offload = cpu_offload |
| 561 | layer.dp_shared_memory = False |
| 562 | layer.part_num_embeddings = rows |
| 563 | # A window strictly inside the table, so unowned rows are exercised too. |
| 564 | layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2 |
| 565 | layer._views = layer._view_src = None |
| 566 | layer._num_sms = torch.cuda.get_device_properties( |
| 567 | torch.accelerator.current_device_index() |
| 568 | ).multi_processor_count |
| 569 | kwargs = ( |
| 570 | {"device": "cpu", "pin_memory": True} if cpu_offload else {"device": "cuda"} |
| 571 | ) |
| 572 | owned = layer.vocab_end_idx - layer.vocab_start_idx |
| 573 | layer.weight = torch.empty(owned, dim, dtype=torch.float8_e4m3fn, **kwargs) |
| 574 | layer.weight_scale_inv = torch.empty( |
| 575 | owned, dim // block, dtype=torch.uint8, **kwargs |
| 576 | ) |
| 577 | torch.manual_seed(0) |
| 578 | layer.weight.copy_((torch.randn(owned, dim) * 4).to(torch.float8_e4m3fn)) |
| 579 | layer.weight_scale_inv.copy_( |
| 580 | torch.randint(120, 134, (owned, dim // block), dtype=torch.uint8) |
| 581 | ) |
| 582 | return layer |
| 583 | |
| 584 | |
| 585 | @pytest.mark.parametrize( |
| 586 | "tp_size,dp_size,n_heads", [(1, 4, 5), (2, 2, 5), (4, 1, 6), (8, 1, 6)] |
| 587 | ) |
| 588 | def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch): |
| 589 | """Reject empty owners before allocating weights or accessing CUDA.""" |
| 590 | monkeypatch.setattr( |
| 591 | engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size |
| 592 | ) |
| 593 | monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size) |
| 594 | monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) |
| 595 | with pytest.raises(AssertionError, match="ranks without hash heads"): |
| 596 | ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads) |
| 597 | |
| 598 | |
| 599 | @pytest.mark.skipif( |
| 600 | not current_platform.is_cuda_alike(), reason="CUDA or ROCm required" |
| 601 | ) |
| 602 | @pytest.mark.parametrize( |
| 603 | "cpu_offload,sort_rows", [(False, False), (True, False), (True, True)] |
| 604 | ) |
| 605 | @pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) |
| 606 | def test_engram_head_shards_reconstruct_checkpoint( |
| 607 | cpu_offload, sort_rows, tp_size, monkeypatch |
| 608 | ): |
| 609 | """Keep complete buckets and reconstruct head order, including TP padding, |
| 610 | whether or not the host lookup sorts its rows by table offset.""" |
| 611 | if sort_rows: |
| 612 | monkeypatch.setattr( |
| 613 | engram_ops, "_engram_lookup_thresholds", lambda _: (0, None) |
| 614 | ) |
| 615 | head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73) |
| 616 | num_rows, dim = sum(head_sizes), 64 |
| 617 | torch.manual_seed(0) |
| 618 | weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn) |
| 619 | scales = torch.randint(120, 134, (num_rows + 7, dim // 32), dtype=torch.uint8) |
| 620 | hashes = torch.empty(7, 2, len(head_sizes), dtype=torch.int32, device="cuda") |
| 621 | start = 0 |
| 622 | for head, size in enumerate(head_sizes): |
| 623 | hashes[:, :, head].random_(start, start + size) |
| 624 | hashes[0, :, head] = start |
| 625 | hashes[-1, :, head] = start + size - 1 |
| 626 | start += size |
| 627 | ids = hashes[:, 1] |
| 628 | expected = _reference_lookup(weight.cuda(), scales.cuda(), ids, 0, num_rows) |
| 629 | monkeypatch.setattr( |
| 630 | engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size |
| 631 | ) |
| 632 | shards = [] |
| 633 | layers = [] |
| 634 | for rank in range(tp_size): |
| 635 | monkeypatch.setattr( |
| 636 | engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank |
| 637 | ) |
| 638 | with torch.device("cuda"): |
| 639 | layer = ParallelEngramEmbedding( |
| 640 | num_rows + 7, dim, head_sizes, cpu_offload=cpu_offload |
| 641 | ) |
| 642 | layer.weight.weight_loader(layer.weight, weight) |
| 643 | layer.weight_scale_inv.weight_loader( |
| 644 | layer.weight_scale_inv, scales.view(torch.float8_e8m0fnu) |
| 645 | ) |
| 646 | out = torch.empty( |
| 647 | len(ids), |
| 648 | layer.part_n_hash_cols, |
| 649 | dim, |
| 650 | device="cuda", |
| 651 | dtype=torch.bfloat16, |
| 652 | ) |
| 653 | layer.lookup(ids, out) |
| 654 | shards.append(out) |
| 655 | layers.append(layer) |
| 656 | gathered = torch.cat(shards, dim=1) |
| 657 | torch.testing.assert_close(gathered[:, : len(head_sizes)], expected, rtol=0, atol=0) |
| 658 | assert torch.count_nonzero(gathered[:, len(head_sizes) :]) == 0 |
| 659 | assert sum(layer.part_num_embeddings for layer in layers) == num_rows |
| 660 | |
| 661 | def gather(local, dim): |
| 662 | torch.testing.assert_close(local, shards[0], rtol=0, atol=0) |
| 663 | return torch.cat(shards, dim=dim) |
| 664 | |
| 665 | monkeypatch.setattr(engram_ops, "tensor_model_parallel_all_gather", gather) |
| 666 | torch.testing.assert_close(layers[0](ids), expected, rtol=0, atol=0) |
| 667 | module = Engram.__new__(Engram) |
| 668 | torch.nn.Module.__init__(module) |
| 669 | module.embed_tokens = layers[0] |
| 670 | module.use_sequence_parallel = False |
| 671 | module.staged_rows = torch.empty_like(shards[0]) |
| 672 | module.prepare_embeddings(ids) |
| 673 | torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0) |
| 674 | # Slicing before head reordering must preserve padded heads and tokens. |
| 675 | module.use_sequence_parallel = True |
| 676 | chunk = (len(ids) + tp_size - 1) // tp_size |
| 677 | padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size)) |
| 678 | for rank in range(tp_size): |
| 679 | monkeypatch.setattr( |
| 680 | engram_ops, "get_tensor_model_parallel_rank", lambda rank=rank: rank |
| 681 | ) |
| 682 | torch.testing.assert_close( |
| 683 | module.embed(ids), padded[rank * chunk : (rank + 1) * chunk], rtol=0, atol=0 |
| 684 | ) |
| 685 | |
| 686 | |
| 687 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 688 | def test_engram_lookup_reuses_jit_across_token_shapes(): |
| 689 | """Runtime token counts and launch grids must share one JIT variant.""" |
| 690 | layer = _make_embedding(cpu_offload=False) |
| 691 | kernel = engram_ops._engram_lookup_kernel |
| 692 | kernel_cache = kernel.device_caches[torch.accelerator.current_device_index()][0] |
| 693 | kernel_cache.clear() |
| 694 | |
| 695 | cache_sizes = [] |
| 696 | for background in (False, True): |
| 697 | for num_tokens in (1, 7, 256): |
| 698 | ids = torch.zeros(num_tokens, 24, dtype=torch.int32, device="cuda") |
| 699 | out = torch.empty( |
| 700 | num_tokens, |
| 701 | 24, |
| 702 | layer.dim, |
| 703 | dtype=torch.bfloat16, |
| 704 | device="cuda", |
| 705 | ) |
| 706 | layer.lookup(ids, out, background=background) |
| 707 | cache_sizes.append(len(kernel_cache)) |
| 708 | |
| 709 | assert cache_sizes == [1] * len(cache_sizes) |
| 710 | |
| 711 | |
| 712 | @pytest.mark.skipif( |
| 713 | not current_platform.is_cuda_alike(), reason="CUDA or ROCm required" |
| 714 | ) |
| 715 | @pytest.mark.parametrize("cpu_offload", [False, True]) |
| 716 | @pytest.mark.parametrize("background", [False, True]) |
| 717 | @pytest.mark.parametrize("num_tokens", [1, 7, 256]) |
| 718 | def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): |
| 719 | """The fused gather must be bit-exact with the torch dequant path it |
| 720 | replaces, from HBM and from pinned host memory alike, and must contribute |
| 721 | zeros for rows another TP rank owns.""" |
| 722 | layer = _make_embedding(cpu_offload) |
| 723 | cols, rows = 24, layer.part_num_embeddings |
| 724 | ids = torch.randint(0, rows, (num_tokens, cols), dtype=torch.int32, device="cuda") |
| 725 | expected = _reference_lookup( |
| 726 | layer.weight.cuda(), |
| 727 | layer.weight_scale_inv.cuda(), |
| 728 | ids, |
| 729 | layer.vocab_start_idx, |
| 730 | layer.vocab_end_idx, |
| 731 | ) |
| 732 | out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda") |
| 733 | layer.lookup(ids, out, background=background) |
| 734 | assert torch.equal(out, expected) |
| 735 | |
| 736 | |
| 737 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 738 | @pytest.mark.parametrize("register_fails", [False, True]) |
| 739 | def test_engram_use_thp_lookup_and_fallback(monkeypatch, register_fails): |
| 740 | """Huge-page tables and their pinned fallback both give exact lookup rows.""" |
| 741 | monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1) |
| 742 | monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) |
| 743 | if register_fails: |
| 744 | monkeypatch.setattr( |
| 745 | torch.cuda.cudart(), "cudaHostRegister", lambda *_: SimpleNamespace(value=1) |
| 746 | ) |
| 747 | rows, dim = 32769, 64 |
| 748 | with torch.device("cuda"): |
| 749 | layer = ParallelEngramEmbedding( |
| 750 | rows, dim, (rows,), cpu_offload=True, use_thp=True |
| 751 | ) |
| 752 | assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned() |
| 753 | same_storage = layer.weight.untyped_storage().data_ptr() == ( |
| 754 | layer.weight_scale_inv.untyped_storage().data_ptr() |
| 755 | ) |
| 756 | assert same_storage != register_fails |
| 757 | layer.weight.data.fill_(2) |
| 758 | layer.weight_scale_inv.data.fill_(127) |
| 759 | ids = torch.tensor( |
| 760 | [[0], [rows - 1], [-1], [rows]], device="cuda", dtype=torch.int32 |
| 761 | ) |
| 762 | out = torch.empty(4, 1, dim, device="cuda", dtype=torch.bfloat16) |
| 763 | layer.lookup(ids, out, background=True) |
| 764 | expected = torch.zeros_like(out) |
| 765 | expected[:2] = 2 |
| 766 | torch.testing.assert_close(out, expected, rtol=0, atol=0) |
| 767 | |
| 768 | |
| 769 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 770 | def test_engram_shared_prefetch_stream_waits_per_layer(): |
| 771 | """Consuming a layer must not wait for later layers on the shared stream.""" |
| 772 | stream = torch.cuda.Stream() |
| 773 | modules = [] |
| 774 | for _ in range(2): |
| 775 | module = Engram.__new__(Engram) |
| 776 | torch.nn.Module.__init__(module) |
| 777 | module._prefetch_stream = stream |
| 778 | module.embed_tokens = _make_embedding(cpu_offload=True) |
| 779 | module.use_sequence_parallel = False |
| 780 | with torch.device("cuda"): |
| 781 | module._init_staging(8, module.embed_tokens.dim) |
| 782 | assert module._prefetch_stream is stream |
| 783 | modules.append(module) |
| 784 | hashes = torch.randint(1024, 3072, (8, 2, 24), device="cuda", dtype=torch.int32) |
| 785 | modules[0].prepare_embeddings(hashes[:, 0]) |
| 786 | with torch.cuda.stream(stream): |
| 787 | torch.cuda._sleep(100_000_000) # ~50 ms ahead of the second lookup |
| 788 | modules[1].prepare_embeddings(hashes[:, 1]) |
| 789 | rows = modules[0].embed(hashes[:, 0]) |
| 790 | torch.cuda.current_stream().synchronize() |
| 791 | assert not modules[1]._prefetch_done.query(), "waited for the second lookup" |
| 792 | layer = modules[0].embed_tokens |
| 793 | expected = _reference_lookup( |
| 794 | layer.weight.cuda(), |
| 795 | layer.weight_scale_inv.cuda(), |
| 796 | hashes[:, 0], |
| 797 | layer.vocab_start_idx, |
| 798 | layer.vocab_end_idx, |
| 799 | ) |
| 800 | torch.testing.assert_close(rows, expected, rtol=0, atol=0) |
| 801 | torch.accelerator.synchronize() |
| 802 | |
| 803 | |
| 804 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 805 | @pytest.mark.parametrize("cpu_offload", [None, False]) |
| 806 | def test_engram_constructor_honors_offload(monkeypatch, cpu_offload): |
| 807 | """Default offload uses pinned storage and the given stream; False uses HBM.""" |
| 808 | from vllm.config import EngramConfig |
| 809 | |
| 810 | config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6) |
| 811 | layout = SimpleNamespace( |
| 812 | num_embeddings=(72,), |
| 813 | head_dim=32, |
| 814 | max_ngram_size=2, |
| 815 | n_heads=24, |
| 816 | primes=(((3,) * 24,),), |
| 817 | ) |
| 818 | vllm_config = SimpleNamespace( |
| 819 | engram_config=EngramConfig() |
| 820 | if cpu_offload is None |
| 821 | else EngramConfig(cpu_offload=cpu_offload), |
| 822 | scheduler_config=SimpleNamespace(max_num_batched_tokens=8), |
| 823 | ) |
| 824 | offloaded = vllm_config.engram_config.cpu_offload |
| 825 | monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config) |
| 826 | monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1) |
| 827 | monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) |
| 828 | monkeypatch.setattr( |
| 829 | engram_ops, "ColumnParallelLinear", lambda *a, **k: torch.nn.Identity() |
| 830 | ) |
| 831 | stream = torch.cuda.Stream() |
| 832 | with torch.device("cuda"): |
| 833 | module = Engram( |
| 834 | config, None, layout, 0, False, "engram", prefetch_stream=stream |
| 835 | ) |
| 836 | layer = module.embed_tokens |
| 837 | if offloaded: |
| 838 | assert layer.weight.device.type == "cpu" and layer.weight.is_pinned() |
| 839 | else: |
| 840 | assert layer.weight.device.type == "cuda" |
| 841 | layer.weight.data.copy_(torch.randn(layer.weight.shape).to(torch.float8_e4m3fn)) |
| 842 | layer.weight_scale_inv.data.fill_(127) |
| 843 | streams = [] |
| 844 | lookup = layer.lookup |
| 845 | |
| 846 | def track_lookup(ids, out, background=False): |
| 847 | streams.append(torch.cuda.current_stream()) |
| 848 | lookup(ids, out, background=background) |
| 849 | |
| 850 | layer.lookup = track_lookup |
| 851 | main = torch.cuda.current_stream() |
| 852 | ids = torch.randint( |
| 853 | layer.vocab_start_idx, |
| 854 | layer.vocab_end_idx, |
| 855 | (8, 24), |
| 856 | device="cuda", |
| 857 | dtype=torch.int32, |
| 858 | ) |
| 859 | module.prepare_embeddings(ids) |
| 860 | expected = _reference_lookup( |
| 861 | layer.weight.cuda(), |
| 862 | layer.weight_scale_inv.cuda(), |
| 863 | ids, |
| 864 | layer.vocab_start_idx, |
| 865 | layer.vocab_end_idx, |
| 866 | ) |
| 867 | torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0) |
| 868 | if offloaded: |
| 869 | assert streams == [stream] |
| 870 | else: |
| 871 | assert streams == [main] |
| 872 | |
| 873 | |
| 874 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 875 | @pytest.mark.parametrize( |
| 876 | "cpu_offload,delay", |
| 877 | [(False, None), (True, None), (True, "producer"), (True, "lookup")], |
| 878 | ) |
| 879 | @pytest.mark.parametrize("capture", ["eager", "full", "breakable"]) |
| 880 | def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, delay): |
| 881 | """Temporary lookup IDs survive allocator reuse and graph replay.""" |
| 882 | _run_engram_prepared_rows(cpu_offload, capture, delay=delay) |
| 883 | |
| 884 | |
| 885 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 886 | @pytest.mark.parametrize( |
| 887 | "missing_dependency", ["record_stream", "producer_wait", "consumer_wait"] |
| 888 | ) |
| 889 | def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_dependency): |
| 890 | """Negative controls must expose incorrect rows when a dependency is removed.""" |
| 891 | if ( |
| 892 | missing_dependency == "record_stream" |
| 893 | and torch.cuda.memory.get_allocator_backend() != "native" |
| 894 | ): |
| 895 | pytest.skip( |
| 896 | "The allocator reuse negative control requires the native allocator" |
| 897 | ) |
| 898 | |
| 899 | def start(self, hash_ids, rows): |
| 900 | stream = self._prefetch_stream |
| 901 | if missing_dependency != "producer_wait": |
| 902 | stream.wait_stream(torch.cuda.current_stream()) |
| 903 | if missing_dependency != "record_stream": |
| 904 | hash_ids.record_stream(stream) |
| 905 | with torch.cuda.stream(stream): |
| 906 | self.embed_tokens.lookup(hash_ids, rows, background=True) |
| 907 | self._prefetch_done.record(stream) |
| 908 | |
| 909 | def finish(self, event): |
| 910 | pass |
| 911 | |
| 912 | monkeypatch.setattr(Engram, "_start_prefetch", start) |
| 913 | if missing_dependency == "consumer_wait": |
| 914 | monkeypatch.setattr(Engram, "_finish_prefetch", finish) |
| 915 | delay = "producer" if missing_dependency == "producer_wait" else "lookup" |
| 916 | try: |
| 917 | with pytest.raises( |
| 918 | AssertionError, match=r"Tensor-likes are not (?:equal|close)!" |
| 919 | ): |
| 920 | _run_engram_prepared_rows(True, "eager", delay=delay) |
| 921 | finally: |
| 922 | torch.accelerator.synchronize() |
| 923 | |
| 924 | |
| 925 | def _run_engram_prepared_rows( |
| 926 | cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, delay=None |
| 927 | ): |
| 928 | from vllm.compilation.breakable_cudagraph import ( |
| 929 | BreakableCUDAGraphCapture, |
| 930 | eager_break_during_capture, |
| 931 | ) |
| 932 | |
| 933 | layer = _make_embedding(cpu_offload) |
| 934 | if delay == "lookup": |
| 935 | lookup = layer.lookup |
| 936 | |
| 937 | def delayed_lookup(indices, out, background=False): |
| 938 | torch.cuda._sleep(2_000_000) |
| 939 | lookup(indices, out, background=background) |
| 940 | |
| 941 | layer.lookup = delayed_lookup |
| 942 | cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64) |
| 943 | layer.n_hash_cols = cols |
| 944 | layer.tp_size = tp_size |
| 945 | layer.part_n_hash_cols = (cols + tp_size - 1) // tp_size |
| 946 | layer.head_start = rank * layer.part_n_hash_cols |
| 947 | engram = Engram.__new__(Engram) |
| 948 | torch.nn.Module.__init__(engram) |
| 949 | engram.embed_tokens = layer |
| 950 | engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None |
| 951 | engram._prefetch_done = torch.cuda.Event() if cpu_offload else None |
| 952 | # Exercise the production eager boundaries even when the test process |
| 953 | # imported Engram before breakable graphs were enabled. |
| 954 | if capture == "breakable": |
| 955 | import vllm.envs as envs |
| 956 | |
| 957 | with pytest.MonkeyPatch.context() as patch: |
| 958 | patch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", True) |
| 959 | for name in ("_start_prefetch", "_finish_prefetch"): |
| 960 | fn = eager_break_during_capture(inspect.unwrap(getattr(Engram, name))) |
| 961 | setattr(engram, name, fn.__get__(engram, Engram)) |
| 962 | engram.use_sequence_parallel = use_sequence_parallel |
| 963 | engram.staged_rows = torch.empty( |
| 964 | num_tokens, |
| 965 | layer.part_n_hash_cols, |
| 966 | layer.dim, |
| 967 | dtype=torch.bfloat16, |
| 968 | device="cuda", |
| 969 | ) |
| 970 | # Match the non-contiguous per-layer slice of the model hash tensor. |
| 971 | hashes = torch.randint( |
| 972 | 0, |
| 973 | layer.part_num_embeddings, |
| 974 | (num_tokens, 2, cols), |
| 975 | dtype=torch.int32, |
| 976 | device="cuda", |
| 977 | ) |
| 978 | src = hashes[:, 1] |
| 979 | local_tokens = ( |
| 980 | (num_tokens + tp_size - 1) // tp_size if use_sequence_parallel else num_tokens |
| 981 | ) |
| 982 | out = torch.empty( |
| 983 | local_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda" |
| 984 | ) |
| 985 | embed = engram.embed |
| 986 | if capture == "compiled": |
| 987 | embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True) |
| 988 | |
| 989 | def step(cap=None): |
| 990 | if delay == "producer": |
| 991 | torch.cuda._sleep(2_000_000) |
| 992 | # Drop the last reference to a non-contiguous input after launching lookup. |
| 993 | engram.prepare_embeddings(hashes.clone()[:, 1]) |
| 994 | # Exercise same-size allocator reuse before consuming the prefetched rows. |
| 995 | torch.empty_like(hashes).fill_(layer.part_num_embeddings - 1) |
| 996 | if cap is not None: |
| 997 | cap.add_eager(lambda: None) |
| 998 | out.copy_(embed(src)) |
| 999 | |
| 1000 | warmup = torch.cuda.Stream() |
| 1001 | warmup.wait_stream(torch.cuda.current_stream()) |
| 1002 | with torch.cuda.stream(warmup): |
| 1003 | step() |
| 1004 | torch.cuda.current_stream().wait_stream(warmup) |
| 1005 | |
| 1006 | graph = None |
| 1007 | if capture == "full": |
| 1008 | graph = torch.cuda.CUDAGraph() |
| 1009 | with torch.cuda.graph(graph): |
| 1010 | step() |
| 1011 | elif capture == "breakable": |
| 1012 | graph = BreakableCUDAGraphCapture() |
| 1013 | with torch.cuda.stream(warmup), graph: |
| 1014 | step(graph) |
| 1015 | torch.cuda.current_stream().wait_stream(warmup) |
| 1016 | assert graph.num_eager_breaks == (3 if cpu_offload else 1) |
| 1017 | |
| 1018 | for _ in range(3): |
| 1019 | # Eager prefill can overwrite staging rows between decode replays. |
| 1020 | hashes.random_(0, layer.part_num_embeddings) |
| 1021 | step() |
| 1022 | hashes.random_(0, layer.part_num_embeddings) |
| 1023 | if graph is None: |
| 1024 | step() |
| 1025 | else: |
| 1026 | graph.replay() |
| 1027 | expected = _reference_lookup( |
| 1028 | layer.weight.cuda(), |
| 1029 | layer.weight_scale_inv.cuda(), |
| 1030 | src, |
| 1031 | layer.vocab_start_idx, |
| 1032 | layer.vocab_end_idx, |
| 1033 | ) |
| 1034 | if use_sequence_parallel: |
| 1035 | expected = torch.nn.functional.pad( |
| 1036 | expected, (0, 0, 0, 0, 0, (-num_tokens) % tp_size) |
| 1037 | )[rank * local_tokens : (rank + 1) * local_tokens] |
| 1038 | torch.testing.assert_close(out, expected, rtol=0, atol=0) |
| 1039 | |
| 1040 | |
| 1041 | def _engram_tp_worker(rank, tp_size, port): |
| 1042 | from tests.utils import init_test_distributed_environment |
| 1043 | from vllm.distributed import cleanup_dist_env_and_memory |
| 1044 | |
| 1045 | torch.accelerator.set_device_index(rank) |
| 1046 | init_test_distributed_environment(tp_size, 1, rank, str(port), local_rank=rank) |
| 1047 | try: |
| 1048 | for cpu_offload in (False, True): |
| 1049 | for sp in (False, True): |
| 1050 | for capture in ("eager", "compiled", "full", "breakable"): |
| 1051 | _run_engram_prepared_rows(cpu_offload, capture, tp_size, rank, sp) |
| 1052 | finally: |
| 1053 | cleanup_dist_env_and_memory() |
| 1054 | |
| 1055 | |
| 1056 | @pytest.mark.distributed(num_gpus=2) |
| 1057 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 1058 | def test_engram_head_collectives_survive_graph_breaks(): |
| 1059 | """All-gather preserves head order and local SP tokens across graph replay.""" |
| 1060 | from vllm.utils.network_utils import get_open_port |
| 1061 | |
| 1062 | if torch.accelerator.device_count() < 2: |
| 1063 | pytest.skip("Requires two GPUs") |
| 1064 | torch.multiprocessing.spawn(_engram_tp_worker, args=(2, get_open_port()), nprocs=2) |