| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Engram tables shared by co-located DP replicas. |
| 4 | |
| 5 | Four GPUs on one node split the hash heads four ways, whether as DEP |
| 6 | (TP1 x DP4, the single-node MoE layout) or TP2 x DP2. A lookup only returns |
| 7 | the right rows if the ids of every replica are gathered, each rank looks up |
| 8 | its own heads, and each replica trades those tokens back for the heads its |
| 9 | peers own. |
| 10 | |
| 11 | With shared host storage, DP replicas instead map the same TP slice and |
| 12 | prefetch only their own tokens without DP gathers. |
| 13 | """ |
| 14 | |
| 15 | from itertools import product |
| 16 | from pathlib import Path |
| 17 | from tempfile import TemporaryDirectory |
| 18 | from threading import Event |
| 19 | from types import SimpleNamespace |
| 20 | |
| 21 | import pytest |
| 22 | import torch |
| 23 | import torch.multiprocessing as mp |
| 24 | |
| 25 | from vllm.config import ( |
| 26 | EngramConfig, |
| 27 | VllmConfig, |
| 28 | get_current_vllm_config_or_none, |
| 29 | set_current_vllm_config, |
| 30 | ) |
| 31 | from vllm.config.load import LoadConfig |
| 32 | from vllm.config.parallel import ParallelConfig |
| 33 | from vllm.config.scheduler import SchedulerConfig |
| 34 | from vllm.distributed import cleanup_dist_env_and_memory, parallel_state |
| 35 | from vllm.distributed.parallel_state import ( |
| 36 | get_engram_dp_group, |
| 37 | get_engram_dp_size, |
| 38 | init_distributed_environment, |
| 39 | initialize_model_parallel, |
| 40 | ) |
| 41 | from vllm.forward_context import set_forward_context |
| 42 | from vllm.models.deepseek_v41.common import engram as engram_ops |
| 43 | from vllm.models.deepseek_v41.common.engram import ( |
| 44 | Engram, |
| 45 | EngramLayout, |
| 46 | ParallelEngramEmbedding, |
| 47 | engram_head_shard_rank, |
| 48 | gather_engram_hashes, |
| 49 | ) |
| 50 | from vllm.platforms import current_platform |
| 51 | from vllm.utils.network_utils import get_open_port |
| 52 | from vllm.utils.system_utils import update_environment_variables |
| 53 | |
| 54 | WORLD_SIZE = 4 |
| 55 | DIM, BLOCK = 64, 32 |
| 56 | # Prime-sized buckets, as the model's layout builds them. |
| 57 | HEAD_SIZES = (97, 101, 103, 107, 109, 113, 127, 131) |
| 58 | |
| 59 | |
| 60 | def _token_counts(dp_size: int) -> tuple[tuple[int, ...], ...]: |
| 61 | """Even loads, then uneven ones with an idle replica to pad around.""" |
| 62 | return ((8,) * dp_size, (7, 0, 5, 1)[:dp_size], (7,) * dp_size) |
| 63 | |
| 64 | |
| 65 | def _full_table(rows: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 66 | """The unsharded table, identical on every rank.""" |
| 67 | torch.manual_seed(0) |
| 68 | weight = (torch.randn(rows, DIM) * 4).to(torch.float8_e4m3fn) |
| 69 | scale_inv = torch.randint(120, 134, (rows, DIM // BLOCK), dtype=torch.uint8) |
| 70 | return weight, scale_inv |
| 71 | |
| 72 | |
| 73 | def _make_ids(head_sizes: tuple[int, ...], num_tokens: int, seed: int) -> torch.Tensor: |
| 74 | """Ids of each head, drawn from that head's own bucket range.""" |
| 75 | torch.manual_seed(seed) |
| 76 | offset = 0 |
| 77 | columns = [] |
| 78 | for size in head_sizes: |
| 79 | columns.append(torch.randint(offset, offset + size, (num_tokens, 1))) |
| 80 | offset += size |
| 81 | return torch.cat(columns, dim=1).to(torch.int32).cuda() |
| 82 | |
| 83 | |
| 84 | def _reference(weight, scale_inv, ids) -> torch.Tensor: |
| 85 | """Dequantized lookup against the whole table.""" |
| 86 | valid = ids >= 0 |
| 87 | safe_ids = ids.clamp_min(0).long() |
| 88 | values = torch.nn.functional.embedding(safe_ids, weight).float() |
| 89 | scales = torch.nn.functional.embedding(safe_ids, scale_inv) |
| 90 | scales = (scales.to(torch.int32) << 23).view(torch.float32) |
| 91 | values = values.unflatten(-1, (-1, BLOCK)) * scales.unsqueeze(-1) |
| 92 | return values.flatten(-2).to(torch.bfloat16).masked_fill(~valid.unsqueeze(-1), 0) |
| 93 | |
| 94 | |
| 95 | @pytest.mark.parametrize( |
| 96 | "architecture,options,edp_ranks,etp_ranks", |
| 97 | [ |
| 98 | ("DeepseekV41ForCausalLM", {}, [0, 2], [0, 1]), |
| 99 | ("Qwen4ExpForCausalLM", {}, None, [0, 1]), |
| 100 | ( |
| 101 | "Qwen4ExpForCausalLM", |
| 102 | {"embedding_across_dp": True}, |
| 103 | None, |
| 104 | [0, 1, 2, 3], |
| 105 | ), |
| 106 | ("Qwen4ExpForConditionalGeneration", {}, None, [0, 1]), |
| 107 | ], |
| 108 | ) |
| 109 | def test_engram_group_creation_honors_model_and_options( |
| 110 | monkeypatch, architecture, options, edp_ranks, etp_ranks |
| 111 | ): |
| 112 | config = SimpleNamespace( |
| 113 | model_config=SimpleNamespace(architecture=architecture, is_moe=False), |
| 114 | parallel_config=ParallelConfig( |
| 115 | tensor_parallel_size=2, |
| 116 | data_parallel_size=2, |
| 117 | distributed_executor_backend="mp", |
| 118 | ), |
| 119 | engram_config=EngramConfig(**options), |
| 120 | ) |
| 121 | for name in ( |
| 122 | "_TP", |
| 123 | "_ETP", |
| 124 | "_ENGRAM_DP", |
| 125 | "_DCP", |
| 126 | "_PCP", |
| 127 | "_PP", |
| 128 | "_DP", |
| 129 | "_EP", |
| 130 | "_EPLB", |
| 131 | ): |
| 132 | monkeypatch.setattr(parallel_state, name, None) |
| 133 | monkeypatch.setattr(parallel_state, "_WORLD", SimpleNamespace(local_rank=0)) |
| 134 | monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) |
| 135 | monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4) |
| 136 | monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) |
| 137 | monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1) |
| 138 | |
| 139 | def make_group(group_ranks, *args, **kwargs): |
| 140 | ranks = next(ranks for ranks in group_ranks if 0 in ranks) |
| 141 | return SimpleNamespace(ranks=ranks, rank_in_group=0, world_size=len(ranks)) |
| 142 | |
| 143 | monkeypatch.setattr(parallel_state, "init_model_parallel_group", make_group) |
| 144 | with set_current_vllm_config(config): |
| 145 | initialize_model_parallel(tensor_model_parallel_size=2, backend="nccl") |
| 146 | group = get_engram_dp_group() |
| 147 | assert (group.ranks if group is not None else None) == edp_ranks |
| 148 | assert parallel_state.get_etp_group().ranks == etp_ranks |
| 149 | |
| 150 | |
| 151 | def _worker(rank: int, tp_size: int, port: int) -> None: |
| 152 | torch.accelerator.set_device_index(rank) |
| 153 | update_environment_variables( |
| 154 | { |
| 155 | "RANK": str(rank), |
| 156 | "LOCAL_RANK": str(rank), |
| 157 | "WORLD_SIZE": str(WORLD_SIZE), |
| 158 | "MASTER_ADDR": "localhost", |
| 159 | "MASTER_PORT": str(port), |
| 160 | } |
| 161 | ) |
| 162 | dp_size = WORLD_SIZE // tp_size |
| 163 | dp_rank, tp_rank = divmod(rank, tp_size) |
| 164 | vllm_config = VllmConfig( |
| 165 | parallel_config=ParallelConfig( |
| 166 | tensor_parallel_size=tp_size, |
| 167 | data_parallel_size=dp_size, |
| 168 | data_parallel_rank=dp_rank, |
| 169 | ), |
| 170 | scheduler_config=SchedulerConfig( |
| 171 | max_model_len=8, |
| 172 | is_encoder_decoder=False, |
| 173 | max_num_batched_tokens=8, |
| 174 | max_num_seqs=8, |
| 175 | ), |
| 176 | ) |
| 177 | # Supply the resolved Engram settings without constructing a full model. |
| 178 | vllm_config.engram_config = EngramConfig() |
| 179 | vllm_config.model_config = SimpleNamespace( |
| 180 | architecture="DeepseekV41ForCausalLM", is_moe=True |
| 181 | ) |
| 182 | try: |
| 183 | init_distributed_environment() |
| 184 | with set_current_vllm_config(vllm_config): |
| 185 | initialize_model_parallel(tensor_model_parallel_size=tp_size) |
| 186 | assert get_engram_dp_size() == dp_size |
| 187 | assert get_engram_dp_group().ranks == [ |
| 188 | tp_rank + replica * tp_size for replica in range(dp_size) |
| 189 | ] |
| 190 | assert engram_head_shard_rank() == tp_rank * dp_size + dp_rank |
| 191 | |
| 192 | for failure in ("create", "open", "mmap", "register", "pinned", None): |
| 193 | _check_shared_storage_lifetime_and_failures(failure) |
| 194 | |
| 195 | # Reuse the processes and communication groups across storage modes. |
| 196 | for (cpu_offload, dp_shared_memory), n_heads, sequence_parallel in product( |
| 197 | [(False, False), (True, False), (True, True)], |
| 198 | [8, 7], |
| 199 | [False, True] if tp_size > 1 else [False], |
| 200 | ): |
| 201 | _check_table( |
| 202 | vllm_config, |
| 203 | cpu_offload, |
| 204 | dp_shared_memory, |
| 205 | n_heads, |
| 206 | sequence_parallel, |
| 207 | ) |
| 208 | finally: |
| 209 | cleanup_dist_env_and_memory() |
| 210 | |
| 211 | |
| 212 | def _check_shared_storage_lifetime_and_failures(failure): |
| 213 | """One peer's failure must roll back every mapping/registration, without hangs.""" |
| 214 | import gc |
| 215 | import weakref |
| 216 | |
| 217 | group = get_engram_dp_group() |
| 218 | runtime = torch.cuda.cudart() |
| 219 | make_file = engram_ops.tempfile.NamedTemporaryFile |
| 220 | map_file = engram_ops.mmap.mmap |
| 221 | paths, mappings, registrations, unregistrations = [], [], [], [] |
| 222 | failed_rank = 0 if failure == "create" else 1 |
| 223 | |
| 224 | def fail(): |
| 225 | raise OSError(f"injected {failure} failure") |
| 226 | |
| 227 | def temporary_file(*args, **kwargs): |
| 228 | if failure == "create": |
| 229 | fail() |
| 230 | file = make_file(*args, **kwargs) |
| 231 | paths.append(Path(file.name)) |
| 232 | return file |
| 233 | |
| 234 | def open_file(*args, **kwargs): |
| 235 | if failure == "open" and group.rank_in_group == failed_rank: |
| 236 | fail() |
| 237 | return open(*args, **kwargs) |
| 238 | |
| 239 | def mapping(*args, **kwargs): |
| 240 | if failure == "mmap" and group.rank_in_group == failed_rank: |
| 241 | fail() |
| 242 | result = map_file(*args, **kwargs) |
| 243 | mappings.append(weakref.ref(result)) |
| 244 | return result |
| 245 | |
| 246 | def register(pointer, size, flags): |
| 247 | if failure == "register" and group.rank_in_group == failed_rank: |
| 248 | return SimpleNamespace(value=1) |
| 249 | result = runtime.cudaHostRegister(pointer, size, flags) |
| 250 | assert result.value == 0 |
| 251 | registrations.append(pointer) |
| 252 | return result |
| 253 | |
| 254 | def unregister(pointer): |
| 255 | unregistrations.append(pointer) |
| 256 | result = runtime.cudaHostUnregister(pointer) |
| 257 | assert result.value == 0 |
| 258 | return result |
| 259 | |
| 260 | with pytest.MonkeyPatch.context() as patch: |
| 261 | patch.setattr( |
| 262 | engram_ops, |
| 263 | "tempfile", |
| 264 | SimpleNamespace(NamedTemporaryFile=temporary_file), |
| 265 | ) |
| 266 | patch.setattr(engram_ops, "open", open_file, raising=False) |
| 267 | patch.setattr( |
| 268 | engram_ops, |
| 269 | "mmap", |
| 270 | SimpleNamespace(mmap=mapping, MAP_SHARED=engram_ops.mmap.MAP_SHARED), |
| 271 | ) |
| 272 | patch.setattr( |
| 273 | torch.cuda, |
| 274 | "cudart", |
| 275 | lambda: SimpleNamespace( |
| 276 | cudaHostRegister=register, |
| 277 | cudaHostUnregister=unregister, |
| 278 | ), |
| 279 | ) |
| 280 | if failure == "pinned" and group.rank_in_group == failed_rank: |
| 281 | patch.setattr(torch.Tensor, "is_pinned", lambda self: False) |
| 282 | if failure is not None: |
| 283 | stage = "cudaHostRegister" if failure in ("register", "pinned") else failure |
| 284 | with pytest.raises(RuntimeError, match=f"EDP rank {failed_rank}.*{stage}"): |
| 285 | engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group) |
| 286 | assert unregistrations == registrations |
| 287 | else: |
| 288 | storage = engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group) |
| 289 | assert all(not path.exists() for path in paths) |
| 290 | storage_ref = weakref.ref(storage) |
| 291 | parameter = torch.nn.Parameter(storage.weight, requires_grad=False) |
| 292 | views = storage.get_views(parameter, storage.weight_scale_inv) |
| 293 | del storage |
| 294 | gc.collect() |
| 295 | assert storage_ref() is None |
| 296 | assert not unregistrations |
| 297 | del parameter |
| 298 | gc.collect() |
| 299 | assert not unregistrations # UVA views still retain the CPU allocation. |
| 300 | del views |
| 301 | gc.collect() |
| 302 | assert unregistrations == registrations |
| 303 | for ref in mappings: |
| 304 | mapped = ref() |
| 305 | assert mapped is None or mapped.closed |
| 306 | assert all(not path.exists() for path in paths) |
| 307 | |
| 308 | |
| 309 | def _check_table( |
| 310 | vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel |
| 311 | ): |
| 312 | parallel = vllm_config.parallel_config |
| 313 | dp_size, tp_size = parallel.data_parallel_size, parallel.tensor_parallel_size |
| 314 | dp_rank = parallel.data_parallel_rank |
| 315 | tp_rank = engram_ops.get_tensor_model_parallel_rank() |
| 316 | multithread = ( |
| 317 | dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1) |
| 318 | ) |
| 319 | load_config = ( |
| 320 | LoadConfig( |
| 321 | load_format="safetensors", |
| 322 | use_tqdm_on_load=False, |
| 323 | model_loader_extra_config={ |
| 324 | "enable_multithread_load": True, |
| 325 | "num_threads": 2, |
| 326 | }, |
| 327 | ) |
| 328 | if multithread |
| 329 | else vllm_config.load_config |
| 330 | ) |
| 331 | with pytest.MonkeyPatch.context() as m: |
| 332 | head_sizes = HEAD_SIZES[:n_heads] |
| 333 | rows = sum(head_sizes) |
| 334 | config = SimpleNamespace( |
| 335 | hidden_size=DIM, |
| 336 | hc_mult=2, |
| 337 | rms_norm_eps=1e-6, |
| 338 | engram_layer_ids=[0], |
| 339 | engram_num_embeddings=[rows], |
| 340 | engram_max_ngram_size=2, |
| 341 | engram_n_heads=n_heads, |
| 342 | engram_head_dim=DIM, |
| 343 | engram_compressed_vocab_size=32, |
| 344 | engram_pad_token_id=0, |
| 345 | engram_vocab_size=HEAD_SIZES[0], |
| 346 | ) |
| 347 | with set_current_vllm_config(vllm_config): |
| 348 | weight, scale_inv = _full_table(rows) |
| 349 | layout = EngramLayout(config) |
| 350 | # Stub configuration inputs, not Engram's initialization or execution. |
| 351 | component_config = SimpleNamespace( |
| 352 | engram_config=EngramConfig( |
| 353 | cpu_offload=cpu_offload, |
| 354 | dp_shared_memory=dp_shared_memory, |
| 355 | ), |
| 356 | scheduler_config=vllm_config.scheduler_config, |
| 357 | load_config=load_config, |
| 358 | ) |
| 359 | with m.context() as init_patch, torch.device("cuda"): |
| 360 | init_patch.setattr( |
| 361 | engram_ops, "get_current_vllm_config", lambda: component_config |
| 362 | ) |
| 363 | engram = Engram( |
| 364 | config, |
| 365 | quant_config=None, |
| 366 | layout=layout, |
| 367 | layer_hash_index=0, |
| 368 | use_sequence_parallel=sequence_parallel, |
| 369 | prefix="model.layers.0.engram", |
| 370 | ) |
| 371 | layer = engram.embed_tokens |
| 372 | if multithread: |
| 373 | _load_shared_multithread(layer, weight, scale_inv, load_config) |
| 374 | else: |
| 375 | # Only the leader has valid checkpoint payload in shared mode. |
| 376 | loader_weight = weight if not dp_shared_memory or dp_rank == 0 else None |
| 377 | loader_scale = ( |
| 378 | scale_inv if not dp_shared_memory or dp_rank == 0 else None |
| 379 | ) |
| 380 | layer.weight.weight_loader(layer.weight, loader_weight) |
| 381 | layer.weight_scale_inv.weight_loader( |
| 382 | layer.weight_scale_inv, loader_scale |
| 383 | ) |
| 384 | if dp_shared_memory: |
| 385 | assert layer.dp_size == 1 |
| 386 | assert layer.head_start == tp_rank * layer.part_n_hash_cols |
| 387 | assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned() |
| 388 | |
| 389 | def unexpected_collective(*args, **kwargs): |
| 390 | raise AssertionError("shared lookup must not use DP collectives") |
| 391 | |
| 392 | m.setattr(get_engram_dp_group(), "all_gather", unexpected_collective) |
| 393 | |
| 394 | token_counts = _token_counts(dp_size) |
| 395 | |
| 396 | # Production forward runs after the construction-time config has exited. |
| 397 | assert get_current_vllm_config_or_none() is None |
| 398 | for batch_idx, counts in enumerate(token_counts): |
| 399 | num_tokens = counts[dp_rank] |
| 400 | # Ids differ per replica but must agree across its TP ranks. |
| 401 | ids = _make_ids(head_sizes, num_tokens, seed=100 + dp_rank) |
| 402 | if batch_idx == 2 and dp_rank == 1: |
| 403 | ids.fill_(engram_ops.DEAD_ID) |
| 404 | with set_forward_context( |
| 405 | None, |
| 406 | vllm_config, |
| 407 | num_tokens=num_tokens, |
| 408 | num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32), |
| 409 | ): |
| 410 | gathered = gather_engram_hashes( |
| 411 | ids, dp_shared_memory=layer.dp_shared_memory |
| 412 | ) |
| 413 | expected_tokens = ( |
| 414 | num_tokens if dp_shared_memory else max(counts) * dp_size |
| 415 | ) |
| 416 | assert gathered.shape[0] == expected_tokens |
| 417 | engram.prepare_embeddings(gathered) |
| 418 | looked_up = engram.embed(ids) |
| 419 | standalone = layer(ids) |
| 420 | expected = _reference(weight.cuda(), scale_inv.cuda(), ids) |
| 421 | torch.testing.assert_close(standalone, expected, atol=0, rtol=0) |
| 422 | if sequence_parallel: |
| 423 | # SP hands back only this rank's token shard, zero-padded. |
| 424 | chunk = -(-num_tokens // tp_size) |
| 425 | expected = torch.nn.functional.pad( |
| 426 | expected, (0, 0, 0, 0, 0, chunk * tp_size - num_tokens) |
| 427 | ) |
| 428 | expected = expected[tp_rank * chunk : (tp_rank + 1) * chunk] |
| 429 | torch.testing.assert_close(looked_up, expected, atol=0, rtol=0) |
| 430 | |
| 431 | if ( |
| 432 | cpu_offload |
| 433 | and n_heads == 7 |
| 434 | and not sequence_parallel |
| 435 | and (not dp_shared_memory or tp_size == 1) |
| 436 | ): |
| 437 | _check_dummy_hash_model_forward( |
| 438 | vllm_config, |
| 439 | engram, |
| 440 | head_sizes, |
| 441 | weight, |
| 442 | scale_inv, |
| 443 | dp_shared_memory, |
| 444 | ) |
| 445 | |
| 446 | replay = n_heads == 7 and ( |
| 447 | (dp_shared_memory and sequence_parallel == (tp_size > 1)) |
| 448 | or ( |
| 449 | cpu_offload |
| 450 | and not dp_shared_memory |
| 451 | and sequence_parallel == (tp_size > 1) |
| 452 | ) |
| 453 | ) |
| 454 | if replay: |
| 455 | captures = ("full", "breakable") |
| 456 | for capture in captures: |
| 457 | _check_prefetch_replay( |
| 458 | engram, |
| 459 | head_sizes, |
| 460 | weight, |
| 461 | scale_inv, |
| 462 | dp_shared_memory, |
| 463 | capture, |
| 464 | ) |
| 465 | if dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1): |
| 466 | with m.context() as storage_patch: |
| 467 | storage_patch.setattr(layer.weight, "data", layer.weight.data.clone()) |
| 468 | with pytest.raises(RuntimeError, match="storage must not be replaced"): |
| 469 | layer(ids) |
| 470 | |
| 471 | |
| 472 | def _check_dummy_hash_model_forward( |
| 473 | vllm_config, engram, head_sizes, weight, scales, dp_shared_memory |
| 474 | ): |
| 475 | """A metadata-less replica joins DP sharding but skips a shared lookup.""" |
| 476 | from vllm.models.deepseek_v41.common.engram import NgramHashState |
| 477 | from vllm.models.deepseek_v41.nvidia import model as model_ops |
| 478 | |
| 479 | dp_rank = vllm_config.parallel_config.data_parallel_rank |
| 480 | if dp_shared_memory and dp_rank != 1: |
| 481 | return |
| 482 | tokens = 4 |
| 483 | ids = _make_ids(head_sizes, tokens, seed=300 + dp_rank) |
| 484 | state = NgramHashState.__new__(NgramHashState) |
| 485 | torch.nn.Module.__init__(state) |
| 486 | state.multipliers = torch.empty(1, 2, dtype=torch.int64, device="cuda") |
| 487 | state.primes = torch.empty(1, 1, len(head_sizes), dtype=torch.int64, device="cuda") |
| 488 | state.ensure_cache = lambda: True |
| 489 | # Hash arithmetic is covered separately; keep the real model's branch and |
| 490 | # NgramHashState.dummy_hashes, plus real preparation/lookup collectives. |
| 491 | state.forward = lambda *args: ids.unsqueeze(1) |
| 492 | |
| 493 | class Decoder(SimpleNamespace): |
| 494 | # Engram hashes and the mask stay the last positional arguments; |
| 495 | # **kwargs absorbs the decoder's keyword-only flags. |
| 496 | def __call__(self, hidden, positions, input_ids, *args, **kwargs): |
| 497 | hashes, keep = args[-2:] |
| 498 | if dp_shared_memory and dp_rank == 1: |
| 499 | assert hashes is None and keep is None |
| 500 | output = hidden.new_zeros( |
| 501 | tokens, len(head_sizes), DIM, dtype=torch.bfloat16 |
| 502 | ) |
| 503 | else: |
| 504 | assert hashes is not None and keep is not None |
| 505 | assert bool(keep.all()) == (dp_rank != 1) |
| 506 | if dp_rank == 1: |
| 507 | assert torch.all(hashes == engram_ops.DEAD_ID) |
| 508 | output = engram.embed(hashes[:, 0]) |
| 509 | # Trailing None is previous_aux; this stub captures no aux states. |
| 510 | return output, hidden, hidden, hidden, hidden, None |
| 511 | |
| 512 | model = model_ops.DeepseekV4Model.__new__(model_ops.DeepseekV4Model) |
| 513 | torch.nn.Module.__init__(model) |
| 514 | model.use_mega_moe = False |
| 515 | model.use_sequence_parallel = False |
| 516 | model.fuse_mhc_all_reduce = False |
| 517 | model.engram_hash = state |
| 518 | model.engram_swa_prefix = "swa" |
| 519 | model.engram_dp_shared_memory = dp_shared_memory |
| 520 | model.layers = [Decoder(engram=engram)] |
| 521 | model.start_layer = 0 |
| 522 | model.end_layer = 1 |
| 523 | model.decoder_replay_start = model.end_layer |
| 524 | model.decoder_replay_layers = None |
| 525 | model.aux_hidden_state_layers = () |
| 526 | metadata = ( |
| 527 | None |
| 528 | if dp_rank == 1 |
| 529 | else { |
| 530 | "swa": SimpleNamespace( |
| 531 | query_start_loc=None, slot_mapping=None, block_table=None |
| 532 | ) |
| 533 | } |
| 534 | ) |
| 535 | with pytest.MonkeyPatch.context() as patch: |
| 536 | patch.setattr( |
| 537 | model_ops, |
| 538 | "get_pp_group", |
| 539 | lambda: SimpleNamespace( |
| 540 | is_first_rank=True, |
| 541 | is_last_rank=False, |
| 542 | ), |
| 543 | ) |
| 544 | patch.setattr(model_ops, "mhc_post_tilelang", lambda hidden, *args: hidden) |
| 545 | with set_forward_context( |
| 546 | metadata, |
| 547 | vllm_config, |
| 548 | num_tokens=tokens, |
| 549 | num_tokens_across_dp=torch.full( |
| 550 | (vllm_config.parallel_config.data_parallel_size,), |
| 551 | tokens, |
| 552 | dtype=torch.int32, |
| 553 | ), |
| 554 | ): |
| 555 | output = model( |
| 556 | torch.arange(tokens, device="cuda"), |
| 557 | torch.arange(tokens, device="cuda"), |
| 558 | intermediate_tensors=None, |
| 559 | inputs_embeds=torch.zeros(tokens, DIM, device="cuda"), |
| 560 | lookback_token_ids=torch.full((1, 1), -1, device="cuda"), |
| 561 | )["hidden_states"] |
| 562 | expected_ids = torch.full_like(ids, engram_ops.DEAD_ID) if dp_rank == 1 else ids |
| 563 | expected = _reference(weight.cuda(), scales.cuda(), expected_ids) |
| 564 | torch.testing.assert_close(output, expected, atol=0, rtol=0) |
| 565 | |
| 566 | |
| 567 | def _load_shared_multithread(layer, weight, scale_inv, load_config): |
| 568 | """Opposite shard completion orders must leave every shared parameter ready.""" |
| 569 | from safetensors.torch import save_file |
| 570 | |
| 571 | from vllm.model_executor.model_loader import weight_utils |
| 572 | from vllm.model_executor.model_loader.default_loader import DefaultModelLoader |
| 573 | |
| 574 | tensors = { |
| 575 | "weight": weight, |
| 576 | "weight_scale_inv": scale_inv.view(torch.float8_e8m0fnu), |
| 577 | } |
| 578 | order = list(tensors) |
| 579 | if get_engram_dp_group().rank_in_group % 2: |
| 580 | order.reverse() |
| 581 | first_loaded = Event() |
| 582 | original_load = weight_utils.load_file |
| 583 | |
| 584 | def load_file(path, **kwargs): |
| 585 | if Path(path).stem != order[0]: |
| 586 | assert first_loaded.wait(timeout=30), "First shared weight load stalled" |
| 587 | return original_load(path, **kwargs) |
| 588 | |
| 589 | with TemporaryDirectory() as directory, pytest.MonkeyPatch.context() as m: |
| 590 | for name, tensor in tensors.items(): |
| 591 | save_file({name: tensor}, str(Path(directory) / f"{name}.safetensors")) |
| 592 | m.setattr(weight_utils, "load_file", load_file) |
| 593 | loader = DefaultModelLoader(load_config) |
| 594 | source = DefaultModelLoader.Source(directory, revision=None) |
| 595 | loaded = [] |
| 596 | for name, tensor in loader._get_weights_iterator(source): |
| 597 | param = getattr(layer, name) |
| 598 | param.weight_loader(param, tensor) |
| 599 | loaded.append(name) |
| 600 | first_loaded.set() |
| 601 | assert loaded == order |
| 602 | |
| 603 | # Check before any further collective could mask incomplete shared writes. |
| 604 | for name, tensor in tensors.items(): |
| 605 | param = getattr(layer, name) |
| 606 | expected = tensor.narrow(0, param.engram_vocab_start, param.shape[0]) |
| 607 | torch.testing.assert_close( |
| 608 | param.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0 |
| 609 | ) |
| 610 | |
| 611 | |
| 612 | def _check_prefetch_replay( |
| 613 | engram, head_sizes, weight, scale_inv, dp_shared_memory, capture |
| 614 | ): |
| 615 | """DP-sharded and shared prefetches must replay with new IDs.""" |
| 616 | from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture |
| 617 | |
| 618 | ids = _make_ids(head_sizes, 8, seed=501) |
| 619 | tp_rank = engram_ops.get_tensor_model_parallel_rank() |
| 620 | tokens = 8 // engram.embed_tokens.tp_size if engram.use_sequence_parallel else 8 |
| 621 | output = torch.empty( |
| 622 | tokens, len(head_sizes), DIM, device="cuda", dtype=torch.bfloat16 |
| 623 | ) |
| 624 | |
| 625 | def step(): |
| 626 | gathered = gather_engram_hashes(ids.clone(), dp_shared_memory=dp_shared_memory) |
| 627 | engram.prepare_embeddings(gathered) |
| 628 | output.copy_(engram.embed(ids)) |
| 629 | |
| 630 | with pytest.MonkeyPatch.context() as patch: |
| 631 | patch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 8) |
| 632 | warmup = torch.cuda.Stream() |
| 633 | warmup.wait_stream(torch.cuda.current_stream()) |
| 634 | with torch.cuda.stream(warmup): |
| 635 | step() |
| 636 | if capture == "breakable": |
| 637 | graph = BreakableCUDAGraphCapture() |
| 638 | with torch.cuda.stream(warmup), graph: |
| 639 | step() |
| 640 | assert graph.num_eager_breaks == 2 |
| 641 | else: |
| 642 | graph = torch.cuda.CUDAGraph() |
| 643 | with torch.cuda.graph(graph, stream=warmup): |
| 644 | step() |
| 645 | torch.cuda.current_stream().wait_stream(warmup) |
| 646 | for iteration in range(2): |
| 647 | ids.copy_(_make_ids(head_sizes, 8, seed=600 + iteration)) |
| 648 | graph.replay() |
| 649 | expected = _reference(weight.cuda(), scale_inv.cuda(), ids) |
| 650 | if engram.use_sequence_parallel: |
| 651 | expected = expected[tp_rank * tokens : (tp_rank + 1) * tokens] |
| 652 | torch.testing.assert_close(output, expected, atol=0, rtol=0) |
| 653 | torch.accelerator.synchronize() |
| 654 | |
| 655 | |
| 656 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") |
| 657 | @pytest.mark.parametrize("tp_size", [1, 2]) |
| 658 | def test_engram_table_shared_across_dp_replicas(tp_size, monkeypatch): |
| 659 | """Check all storage modes with one process launch per TP/DP topology.""" |
| 660 | if torch.accelerator.device_count() < WORLD_SIZE: |
| 661 | pytest.skip(f"Need {WORLD_SIZE} GPUs to run the test.") |
| 662 | monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1") |
| 663 | mp.spawn(_worker, args=(tp_size, get_open_port()), nprocs=WORLD_SIZE) |
| 664 | cleanup_dist_env_and_memory() |
| 665 | |
| 666 | |
| 667 | @pytest.mark.parametrize( |
| 668 | "cpu_offload,dp_size,error", |
| 669 | [ |
| 670 | (False, 4, "requires cpu_offload=True"), |
| 671 | (True, 1, "effective Engram DP size is 1"), |
| 672 | ], |
| 673 | ) |
| 674 | def test_engram_dp_shared_memory_runtime_requirements( |
| 675 | monkeypatch, cpu_offload, dp_size, error |
| 676 | ): |
| 677 | """Report unmet storage/topology requirements before allocating CUDA memory.""" |
| 678 | monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size) |
| 679 | with pytest.raises(ValueError, match=error): |
| 680 | ParallelEngramEmbedding( |
| 681 | sum(HEAD_SIZES), |
| 682 | DIM, |
| 683 | HEAD_SIZES, |
| 684 | cpu_offload=cpu_offload, |
| 685 | dp_shared_memory=True, |
| 686 | ) |
| 687 | |
| 688 | |
| 689 | @pytest.mark.parametrize("query_fails", [False, True]) |
| 690 | def test_engram_tables_too_large_for_shm_are_not_shared(monkeypatch, query_fails): |
| 691 | """A /dev/shm that is too small or cannot be queried must fall back, and the |
| 692 | leader must still broadcast that decision so its peers do not hang.""" |
| 693 | broadcasts = [] |
| 694 | |
| 695 | def broadcast(obj): |
| 696 | broadcasts.append(obj) |
| 697 | return obj |
| 698 | |
| 699 | def query(*args, **kwargs): |
| 700 | raise OSError("statvfs failed") |
| 701 | |
| 702 | monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: 2) |
| 703 | monkeypatch.setattr( |
| 704 | engram_ops, |
| 705 | "get_engram_dp_group", |
| 706 | lambda: SimpleNamespace(rank_in_group=0, broadcast_object=broadcast), |
| 707 | ) |
| 708 | if query_fails: |
| 709 | monkeypatch.setattr(engram_ops, "check_shm_free_space", query) |
| 710 | layout = SimpleNamespace(num_embeddings=(1 << 50,), head_dim=DIM) |
| 711 | assert not engram_ops.can_share_engram_tables(layout) |
| 712 | assert len(broadcasts) == 1 and broadcasts[0] is not None |
| 713 | |
| 714 | |
| 715 | @pytest.mark.parametrize( |
| 716 | "node_ids,dp_size,replica_size,expected", |
| 717 | [ |
| 718 | ([0] * 4, 1, 4, 1), |
| 719 | ([0] * 4, 4, 1, 4), |
| 720 | ([0] * 4, 2, 2, 2), |
| 721 | ([0] * 8 + [1] * 8, 8, 2, 4), |
| 722 | ([0] * 6 + [1] * 6, 4, 1, 2), |
| 723 | ([0] * 3 + [1] * 4, 7, 1, 1), |
| 724 | ([0] * 8 + [1] * 8 + [2] * 8, 8, 3, 1), |
| 725 | ([0, 1] * 4, 4, 2, 1), |
| 726 | ([0] * 3 + [1] * 5, 4, 2, 1), |
| 727 | ], |
| 728 | ) |
| 729 | def test_engram_dp_shard_size_respects_node_boundaries( |
| 730 | node_ids, dp_size, replica_size, expected, monkeypatch |
| 731 | ): |
| 732 | """Sharing must fall back when replicas or node rank ranges do not align.""" |
| 733 | monkeypatch.setattr(parallel_state, "get_node_count", lambda: len(set(node_ids))) |
| 734 | monkeypatch.setattr( |
| 735 | parallel_state, "get_world_group", lambda: SimpleNamespace(cpu_group=None) |
| 736 | ) |
| 737 | monkeypatch.setattr( |
| 738 | parallel_state, |
| 739 | "in_the_same_node_as", |
| 740 | lambda group, src: [node == node_ids[src] for node in node_ids], |
| 741 | ) |
| 742 | assert ( |
| 743 | parallel_state._engram_dp_shard_size(len(node_ids), dp_size, replica_size) |
| 744 | == expected |
| 745 | ) |
| 746 | |
| 747 | |
| 748 | @pytest.mark.parametrize("num_tokens", [0, 2, 4, 5]) |
| 749 | def test_engram_hash_padding_has_no_valid_rows(num_tokens, monkeypatch): |
| 750 | """Idle/padded slots must not query row zero; oversized batches are rejected.""" |
| 751 | monkeypatch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 4) |
| 752 | monkeypatch.setattr( |
| 753 | engram_ops, |
| 754 | "get_engram_dp_group", |
| 755 | lambda: SimpleNamespace(all_gather=lambda ids, dim: torch.cat([ids, ids], dim)), |
| 756 | ) |
| 757 | ids = torch.full((num_tokens, 2, 3), 17, dtype=torch.int32) |
| 758 | if num_tokens > 4: |
| 759 | with pytest.raises(ValueError, match="exceeds the DP token slot"): |
| 760 | gather_engram_hashes(ids) |
| 761 | return |
| 762 | gathered = gather_engram_hashes(ids).reshape(2, 4, 2, 3) |
| 763 | for replica in gathered: |
| 764 | torch.testing.assert_close(replica[:num_tokens], ids) |
| 765 | assert torch.all(replica[num_tokens:] == engram_ops.DEAD_ID) |
| 766 | |
| 767 | |
| 768 | @pytest.mark.parametrize("tp_size", [1, 2]) |
| 769 | @pytest.mark.parametrize("uniform", [False, True]) |
| 770 | def test_engram_hash_padding_ignores_other_nodes(tp_size, uniform, monkeypatch): |
| 771 | """Two four-GPU nodes pad only to their own maximum, including TP x DP.""" |
| 772 | group_size = 4 // tp_size |
| 773 | counts = ( |
| 774 | [4] * (8 // tp_size) |
| 775 | if uniform |
| 776 | else [4096] + [2] * (group_size - 1) + [1, 3] + [0] * (group_size - 2) |
| 777 | ) |
| 778 | monkeypatch.setattr( |
| 779 | engram_ops, |
| 780 | "get_forward_context", |
| 781 | lambda: SimpleNamespace( |
| 782 | dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(counts)) |
| 783 | ), |
| 784 | ) |
| 785 | for dp_rank in range(8 // tp_size): |
| 786 | slot = 4 if uniform else (4096 if dp_rank < group_size else 3) |
| 787 | ids = torch.full((counts[dp_rank], 2, 3), 17, dtype=torch.int32) |
| 788 | monkeypatch.setattr( |
| 789 | engram_ops, |
| 790 | "get_dp_group", |
| 791 | lambda rank=dp_rank: SimpleNamespace(rank_in_group=rank), |
| 792 | ) |
| 793 | monkeypatch.setattr( |
| 794 | engram_ops, |
| 795 | "get_engram_dp_group", |
| 796 | lambda rank=dp_rank: SimpleNamespace( |
| 797 | rank_in_group=rank % group_size, |
| 798 | world_size=group_size, |
| 799 | all_gather=lambda tensor, dim: torch.cat([tensor] * group_size, dim), |
| 800 | ), |
| 801 | ) |
| 802 | gathered = gather_engram_hashes(ids) |
| 803 | expected = ids.new_full((slot, 2, 3), engram_ops.DEAD_ID) |
| 804 | expected[: ids.shape[0]] = ids |
| 805 | torch.testing.assert_close( |
| 806 | gathered, expected.repeat(group_size, 1, 1), msg=f"DP rank {dp_rank}" |
| 807 | ) |