diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index dfe900d827..070664ac7a 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -206,6 +206,9 @@ steps: - tests/v1/distributed - tests/v1/engine/test_engine_core_client.py - tests/distributed/test_utils + - tests/distributed/test_engram_dp_shard.py + - vllm/config/engram.py + - vllm/models/deepseek_v4_1/ commands: # https://github.com/NVIDIA/nccl/issues/1838 - export NCCL_CUMEM_HOST_ENABLE=0 @@ -217,6 +220,7 @@ steps: - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp - pytest -v -s distributed/test_utils.py + - pytest -v -s distributed/test_engram_dp_shard.py - label: ":nvidia: (L4) Distributed Compile + Comm" key: distributed-compile-comm-4-gpus diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py new file mode 100644 index 0000000000..ebb0254289 --- /dev/null +++ b/tests/distributed/test_engram_dp_shard.py @@ -0,0 +1,776 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Engram tables shared by co-located DP replicas. + +Four GPUs on one node split the hash heads four ways, whether as DEP +(TP1 x DP4, the single-node MoE layout) or TP2 x DP2. A lookup only returns +the right rows if the ids of every replica are gathered, each rank looks up +its own heads, and each replica trades those tokens back for the heads its +peers own. + +With shared host storage, DP replicas instead map the same TP slice and +prefetch only their own tokens without DP gathers. +""" + +from itertools import product +from pathlib import Path +from tempfile import TemporaryDirectory +from threading import Event +from types import SimpleNamespace + +import pytest +import torch +import torch.multiprocessing as mp + +from vllm.config import ( + EngramConfig, + VllmConfig, + get_current_vllm_config_or_none, + set_current_vllm_config, +) +from vllm.config.load import LoadConfig +from vllm.config.parallel import ParallelConfig +from vllm.config.scheduler import SchedulerConfig +from vllm.distributed import cleanup_dist_env_and_memory, parallel_state +from vllm.distributed.parallel_state import ( + get_engram_dp_group, + get_engram_dp_size, + init_distributed_environment, + initialize_model_parallel, +) +from vllm.forward_context import set_forward_context +from vllm.models.deepseek_v4_1.common.engram import EngramLayout +from vllm.models.deepseek_v4_1.nvidia import engram as engram_ops +from vllm.models.deepseek_v4_1.nvidia.engram import ( + Engram, + ParallelEngramEmbedding, + engram_head_shard_rank, + gather_engram_hashes, +) +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port +from vllm.utils.system_utils import update_environment_variables + +WORLD_SIZE = 4 +DIM, BLOCK = 64, 32 +# Prime-sized buckets, as the model's layout builds them. +HEAD_SIZES = (97, 101, 103, 107, 109, 113, 127, 131) + + +def _token_counts(dp_size: int) -> tuple[tuple[int, ...], ...]: + """Even loads, then uneven ones with an idle replica to pad around.""" + return ((8,) * dp_size, (7, 0, 5, 1)[:dp_size], (7,) * dp_size) + + +def _full_table(rows: int) -> tuple[torch.Tensor, torch.Tensor]: + """The unsharded table, identical on every rank.""" + torch.manual_seed(0) + weight = (torch.randn(rows, DIM) * 4).to(torch.float8_e4m3fn) + scale_inv = torch.randint(120, 134, (rows, DIM // BLOCK), dtype=torch.uint8) + return weight, scale_inv + + +def _make_ids(head_sizes: tuple[int, ...], num_tokens: int, seed: int) -> torch.Tensor: + """Ids of each head, drawn from that head's own bucket range.""" + torch.manual_seed(seed) + offset = 0 + columns = [] + for size in head_sizes: + columns.append(torch.randint(offset, offset + size, (num_tokens, 1))) + offset += size + return torch.cat(columns, dim=1).to(torch.int32).cuda() + + +def _reference(weight, scale_inv, ids) -> torch.Tensor: + """Dequantized lookup against the whole table.""" + valid = ids >= 0 + safe_ids = ids.clamp_min(0).long() + values = torch.nn.functional.embedding(safe_ids, weight).float() + scales = torch.nn.functional.embedding(safe_ids, scale_inv) + scales = (scales.to(torch.int32) << 23).view(torch.float32) + values = values.unflatten(-1, (-1, BLOCK)) * scales.unsqueeze(-1) + return values.flatten(-2).to(torch.bfloat16).masked_fill(~valid.unsqueeze(-1), 0) + + +@pytest.mark.parametrize( + "architecture,options,edp_ranks,etp_ranks", + [ + ("DeepseekV41ForCausalLM", {}, [0, 2], [0, 1]), + ("Qwen4ExpForCausalLM", {}, None, [0, 1]), + ( + "Qwen4ExpForCausalLM", + {"embedding_across_dp": True}, + None, + [0, 1, 2, 3], + ), + ("Qwen4ExpForConditionalGeneration", {}, None, [0, 1]), + ], +) +def test_engram_group_creation_honors_model_and_options( + monkeypatch, architecture, options, edp_ranks, etp_ranks +): + config = SimpleNamespace( + model_config=SimpleNamespace(architecture=architecture, is_moe=False), + parallel_config=ParallelConfig( + tensor_parallel_size=2, + data_parallel_size=2, + distributed_executor_backend="mp", + ), + engram_config=EngramConfig(**options), + ) + for name in ( + "_TP", + "_ETP", + "_ENGRAM_DP", + "_DCP", + "_PCP", + "_PP", + "_DP", + "_EP", + "_EPLB", + ): + monkeypatch.setattr(parallel_state, name, None) + monkeypatch.setattr(parallel_state, "_WORLD", SimpleNamespace(local_rank=0)) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 4) + monkeypatch.setattr(torch.distributed, "get_rank", lambda: 0) + monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1) + + def make_group(group_ranks, *args, **kwargs): + ranks = next(ranks for ranks in group_ranks if 0 in ranks) + return SimpleNamespace(ranks=ranks, rank_in_group=0, world_size=len(ranks)) + + monkeypatch.setattr(parallel_state, "init_model_parallel_group", make_group) + with set_current_vllm_config(config): + initialize_model_parallel(tensor_model_parallel_size=2, backend="nccl") + group = get_engram_dp_group() + assert (group.ranks if group is not None else None) == edp_ranks + assert parallel_state.get_etp_group().ranks == etp_ranks + + +def _worker(rank: int, tp_size: int, port: int) -> None: + torch.accelerator.set_device_index(rank) + update_environment_variables( + { + "RANK": str(rank), + "LOCAL_RANK": str(rank), + "WORLD_SIZE": str(WORLD_SIZE), + "MASTER_ADDR": "localhost", + "MASTER_PORT": str(port), + } + ) + dp_size = WORLD_SIZE // tp_size + dp_rank, tp_rank = divmod(rank, tp_size) + vllm_config = VllmConfig( + parallel_config=ParallelConfig( + tensor_parallel_size=tp_size, + data_parallel_size=dp_size, + data_parallel_rank=dp_rank, + ), + scheduler_config=SchedulerConfig( + max_model_len=8, + is_encoder_decoder=False, + max_num_batched_tokens=8, + max_num_seqs=8, + ), + ) + # Supply the resolved Engram settings without constructing a full model. + vllm_config.engram_config = EngramConfig() + vllm_config.model_config = SimpleNamespace( + architecture="DeepseekV41ForCausalLM", is_moe=True + ) + try: + init_distributed_environment() + with set_current_vllm_config(vllm_config): + initialize_model_parallel(tensor_model_parallel_size=tp_size) + assert get_engram_dp_size() == dp_size + assert get_engram_dp_group().ranks == [ + tp_rank + replica * tp_size for replica in range(dp_size) + ] + assert engram_head_shard_rank() == tp_rank * dp_size + dp_rank + + for failure in ("create", "open", "mmap", "register", "pinned", None): + _check_shared_storage_lifetime_and_failures(failure) + + # Reuse the processes and communication groups across storage modes. + for (cpu_offload, dp_shared_memory), n_heads, sequence_parallel in product( + [(False, False), (True, False), (True, True)], + [8, 7], + [False, True] if tp_size > 1 else [False], + ): + _check_table( + vllm_config, + cpu_offload, + dp_shared_memory, + n_heads, + sequence_parallel, + ) + finally: + cleanup_dist_env_and_memory() + + +def _check_shared_storage_lifetime_and_failures(failure): + """One peer's failure must roll back every mapping/registration, without hangs.""" + import gc + import weakref + + group = get_engram_dp_group() + runtime = torch.cuda.cudart() + make_file = engram_ops.tempfile.NamedTemporaryFile + map_file = engram_ops.mmap.mmap + paths, mappings, registrations, unregistrations = [], [], [], [] + failed_rank = 0 if failure == "create" else 1 + + def fail(): + raise OSError(f"injected {failure} failure") + + def temporary_file(*args, **kwargs): + if failure == "create": + fail() + file = make_file(*args, **kwargs) + paths.append(Path(file.name)) + return file + + def open_file(*args, **kwargs): + if failure == "open" and group.rank_in_group == failed_rank: + fail() + return open(*args, **kwargs) + + def mapping(*args, **kwargs): + if failure == "mmap" and group.rank_in_group == failed_rank: + fail() + result = map_file(*args, **kwargs) + mappings.append(weakref.ref(result)) + return result + + def register(pointer, size, flags): + if failure == "register" and group.rank_in_group == failed_rank: + return SimpleNamespace(value=1) + result = runtime.cudaHostRegister(pointer, size, flags) + assert result.value == 0 + registrations.append(pointer) + return result + + def unregister(pointer): + unregistrations.append(pointer) + result = runtime.cudaHostUnregister(pointer) + assert result.value == 0 + return result + + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + engram_ops, + "tempfile", + SimpleNamespace(NamedTemporaryFile=temporary_file), + ) + patch.setattr(engram_ops, "open", open_file, raising=False) + patch.setattr( + engram_ops, + "mmap", + SimpleNamespace(mmap=mapping, MAP_SHARED=engram_ops.mmap.MAP_SHARED), + ) + patch.setattr( + torch.cuda, + "cudart", + lambda: SimpleNamespace( + cudaHostRegister=register, + cudaHostUnregister=unregister, + ), + ) + if failure == "pinned" and group.rank_in_group == failed_rank: + patch.setattr(torch.Tensor, "is_pinned", lambda self: False) + if failure is not None: + stage = "cudaHostRegister" if failure in ("register", "pinned") else failure + with pytest.raises(RuntimeError, match=f"EDP rank {failed_rank}.*{stage}"): + engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group) + assert unregistrations == registrations + else: + storage = engram_ops.DPSharedEngramStorage(128, DIM, BLOCK, group) + assert all(not path.exists() for path in paths) + storage_ref = weakref.ref(storage) + parameter = torch.nn.Parameter(storage.weight, requires_grad=False) + views = storage.get_views(parameter, storage.weight_scale_inv) + del storage + gc.collect() + assert storage_ref() is None + assert not unregistrations + del parameter + gc.collect() + assert not unregistrations # UVA views still retain the CPU allocation. + del views + gc.collect() + assert unregistrations == registrations + for ref in mappings: + mapped = ref() + assert mapped is None or mapped.closed + assert all(not path.exists() for path in paths) + + +def _check_table( + vllm_config, cpu_offload, dp_shared_memory, n_heads, sequence_parallel +): + parallel = vllm_config.parallel_config + dp_size, tp_size = parallel.data_parallel_size, parallel.tensor_parallel_size + dp_rank = parallel.data_parallel_rank + tp_rank = engram_ops.get_tensor_model_parallel_rank() + multithread = ( + dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1) + ) + load_config = ( + LoadConfig( + load_format="safetensors", + use_tqdm_on_load=False, + model_loader_extra_config={ + "enable_multithread_load": True, + "num_threads": 2, + }, + ) + if multithread + else vllm_config.load_config + ) + with pytest.MonkeyPatch.context() as m: + head_sizes = HEAD_SIZES[:n_heads] + rows = sum(head_sizes) + config = SimpleNamespace( + hidden_size=DIM, + hc_mult=2, + rms_norm_eps=1e-6, + engram_layer_ids=[0], + engram_num_embeddings=[rows], + engram_max_ngram_size=2, + engram_n_heads=n_heads, + engram_head_dim=DIM, + engram_compressed_vocab_size=32, + engram_pad_token_id=0, + engram_vocab_size=HEAD_SIZES[0], + ) + with set_current_vllm_config(vllm_config): + weight, scale_inv = _full_table(rows) + layout = EngramLayout(config) + # Stub configuration inputs, not Engram's initialization or execution. + component_config = SimpleNamespace( + engram_config=EngramConfig( + cpu_offload=cpu_offload, + dp_shared_memory=dp_shared_memory, + ), + scheduler_config=vllm_config.scheduler_config, + load_config=load_config, + ) + with m.context() as init_patch, torch.device("cuda"): + init_patch.setattr( + engram_ops, "get_current_vllm_config", lambda: component_config + ) + engram = Engram( + config, + quant_config=None, + layout=layout, + layer_hash_index=0, + use_sequence_parallel=sequence_parallel, + prefix="model.layers.0.engram", + ) + layer = engram.embed_tokens + if multithread: + _load_shared_multithread(layer, weight, scale_inv, load_config) + else: + # Only the leader has valid checkpoint payload in shared mode. + loader_weight = weight if not dp_shared_memory or dp_rank == 0 else None + loader_scale = ( + scale_inv if not dp_shared_memory or dp_rank == 0 else None + ) + layer.weight.weight_loader(layer.weight, loader_weight) + layer.weight_scale_inv.weight_loader( + layer.weight_scale_inv, loader_scale + ) + if dp_shared_memory: + assert layer.dp_size == 1 + assert layer.head_start == tp_rank * layer.part_n_hash_cols + assert layer.weight.is_pinned() and layer.weight_scale_inv.is_pinned() + + def unexpected_collective(*args, **kwargs): + raise AssertionError("shared lookup must not use DP collectives") + + m.setattr(get_engram_dp_group(), "all_gather", unexpected_collective) + + token_counts = _token_counts(dp_size) + + # Production forward runs after the construction-time config has exited. + assert get_current_vllm_config_or_none() is None + for batch_idx, counts in enumerate(token_counts): + num_tokens = counts[dp_rank] + # Ids differ per replica but must agree across its TP ranks. + ids = _make_ids(head_sizes, num_tokens, seed=100 + dp_rank) + if batch_idx == 2 and dp_rank == 1: + ids.fill_(engram_ops.DEAD_ID) + with set_forward_context( + None, + vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32), + ): + gathered = gather_engram_hashes( + ids, dp_shared_memory=layer.dp_shared_memory + ) + expected_tokens = ( + num_tokens if dp_shared_memory else max(counts) * dp_size + ) + assert gathered.shape[0] == expected_tokens + engram.prepare_embeddings(gathered) + looked_up = engram.embed(ids) + standalone = layer(ids) + expected = _reference(weight.cuda(), scale_inv.cuda(), ids) + torch.testing.assert_close(standalone, expected, atol=0, rtol=0) + if sequence_parallel: + # SP hands back only this rank's token shard, zero-padded. + chunk = -(-num_tokens // tp_size) + expected = torch.nn.functional.pad( + expected, (0, 0, 0, 0, 0, chunk * tp_size - num_tokens) + ) + expected = expected[tp_rank * chunk : (tp_rank + 1) * chunk] + torch.testing.assert_close(looked_up, expected, atol=0, rtol=0) + + if ( + cpu_offload + and n_heads == 7 + and not sequence_parallel + and (not dp_shared_memory or tp_size == 1) + ): + _check_dummy_hash_model_forward( + vllm_config, + engram, + head_sizes, + weight, + scale_inv, + dp_shared_memory, + ) + + replay = n_heads == 7 and ( + (dp_shared_memory and sequence_parallel == (tp_size > 1)) + or ( + cpu_offload + and not dp_shared_memory + and sequence_parallel == (tp_size > 1) + ) + ) + if replay: + captures = ("full", "breakable") + for capture in captures: + _check_prefetch_replay( + engram, + head_sizes, + weight, + scale_inv, + dp_shared_memory, + capture, + ) + if dp_shared_memory and n_heads == 7 and sequence_parallel == (tp_size > 1): + with m.context() as storage_patch: + storage_patch.setattr(layer.weight, "data", layer.weight.data.clone()) + with pytest.raises(RuntimeError, match="storage must not be replaced"): + layer(ids) + + +def _check_dummy_hash_model_forward( + vllm_config, engram, head_sizes, weight, scales, dp_shared_memory +): + """A metadata-less replica joins DP sharding but skips a shared lookup.""" + from vllm.models.deepseek_v4_1.common.engram import NgramHashState + from vllm.models.deepseek_v4_1.nvidia import model as model_ops + + dp_rank = vllm_config.parallel_config.data_parallel_rank + if dp_shared_memory and dp_rank != 1: + return + tokens = 4 + ids = _make_ids(head_sizes, tokens, seed=300 + dp_rank) + state = NgramHashState.__new__(NgramHashState) + torch.nn.Module.__init__(state) + state.multipliers = torch.empty(1, 2, dtype=torch.int64, device="cuda") + state.primes = torch.empty(1, 1, len(head_sizes), dtype=torch.int64, device="cuda") + state.ensure_cache = lambda: True + # Hash arithmetic is covered separately; keep the real model's branch and + # NgramHashState.dummy_hashes, plus real preparation/lookup collectives. + state.forward = lambda *args: ids.unsqueeze(1) + + class Decoder(SimpleNamespace): + def __call__(self, hidden, positions, input_ids, *args): + hashes, keep = args[-2:] + if dp_shared_memory and dp_rank == 1: + assert hashes is None and keep is None + output = hidden.new_zeros( + tokens, len(head_sizes), DIM, dtype=torch.bfloat16 + ) + else: + assert hashes is not None and keep is not None + assert bool(keep.all()) == (dp_rank != 1) + if dp_rank == 1: + assert torch.all(hashes == engram_ops.DEAD_ID) + output = engram.embed(hashes[:, 0]) + return output, None, None, None, None + + model = SimpleNamespace( + use_mega_moe=False, + use_sequence_parallel=False, + engram_hash=state, + engram_swa_prefix="swa", + engram_dp_shared_memory=dp_shared_memory, + layers=[Decoder(engram=engram)], + start_layer=0, + end_layer=1, + aux_hidden_state_layers=(), + ) + metadata = ( + None + if dp_rank == 1 + else { + "swa": SimpleNamespace( + query_start_loc=None, slot_mapping=None, block_table=None + ) + } + ) + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + model_ops, + "get_pp_group", + lambda: SimpleNamespace( + is_first_rank=True, + is_last_rank=False, + ), + ) + patch.setattr(model_ops, "mhc_post_tilelang", lambda hidden, *args: hidden) + with set_forward_context( + metadata, + vllm_config, + num_tokens=tokens, + num_tokens_across_dp=torch.full( + (vllm_config.parallel_config.data_parallel_size,), + tokens, + dtype=torch.int32, + ), + ): + output = model_ops.DeepseekV4Model.forward( + model, + torch.arange(tokens, device="cuda"), + torch.arange(tokens, device="cuda"), + intermediate_tensors=None, + inputs_embeds=torch.zeros(tokens, DIM, device="cuda"), + lookback_token_ids=torch.full((1, 1), -1, device="cuda"), + )["hidden_states"] + expected_ids = torch.full_like(ids, engram_ops.DEAD_ID) if dp_rank == 1 else ids + expected = _reference(weight.cuda(), scales.cuda(), expected_ids) + torch.testing.assert_close(output, expected, atol=0, rtol=0) + + +def _load_shared_multithread(layer, weight, scale_inv, load_config): + """Opposite shard completion orders must leave every shared parameter ready.""" + from safetensors.torch import save_file + + from vllm.model_executor.model_loader import weight_utils + from vllm.model_executor.model_loader.default_loader import DefaultModelLoader + + tensors = { + "weight": weight, + "weight_scale_inv": scale_inv.view(torch.float8_e8m0fnu), + } + order = list(tensors) + if get_engram_dp_group().rank_in_group % 2: + order.reverse() + first_loaded = Event() + original_load = weight_utils.load_file + + def load_file(path, **kwargs): + if Path(path).stem != order[0]: + assert first_loaded.wait(timeout=30), "First shared weight load stalled" + return original_load(path, **kwargs) + + with TemporaryDirectory() as directory, pytest.MonkeyPatch.context() as m: + for name, tensor in tensors.items(): + save_file({name: tensor}, str(Path(directory) / f"{name}.safetensors")) + m.setattr(weight_utils, "load_file", load_file) + loader = DefaultModelLoader(load_config) + source = DefaultModelLoader.Source(directory, revision=None) + loaded = [] + for name, tensor in loader._get_weights_iterator(source): + param = getattr(layer, name) + param.weight_loader(param, tensor) + loaded.append(name) + first_loaded.set() + assert loaded == order + + # Check before any further collective could mask incomplete shared writes. + for name, tensor in tensors.items(): + param = getattr(layer, name) + expected = tensor.narrow(0, param.engram_vocab_start, param.shape[0]) + torch.testing.assert_close( + param.view(torch.uint8), expected.view(torch.uint8), atol=0, rtol=0 + ) + + +def _check_prefetch_replay( + engram, head_sizes, weight, scale_inv, dp_shared_memory, capture +): + """DP-sharded and shared prefetches must replay with new IDs.""" + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + + ids = _make_ids(head_sizes, 8, seed=501) + tp_rank = engram_ops.get_tensor_model_parallel_rank() + tokens = 8 // engram.embed_tokens.tp_size if engram.use_sequence_parallel else 8 + output = torch.empty( + tokens, len(head_sizes), DIM, device="cuda", dtype=torch.bfloat16 + ) + + def step(): + gathered = gather_engram_hashes(ids.clone(), dp_shared_memory=dp_shared_memory) + engram.prepare_embeddings(gathered) + output.copy_(engram.embed(ids)) + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 8) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + step() + if capture == "breakable": + graph = BreakableCUDAGraphCapture() + with torch.cuda.stream(warmup), graph: + step() + assert graph.num_eager_breaks == 2 + else: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=warmup): + step() + torch.cuda.current_stream().wait_stream(warmup) + for iteration in range(2): + ids.copy_(_make_ids(head_sizes, 8, seed=600 + iteration)) + graph.replay() + expected = _reference(weight.cuda(), scale_inv.cuda(), ids) + if engram.use_sequence_parallel: + expected = expected[tp_rank * tokens : (tp_rank + 1) * tokens] + torch.testing.assert_close(output, expected, atol=0, rtol=0) + torch.accelerator.synchronize() + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_engram_table_shared_across_dp_replicas(tp_size, monkeypatch): + """Check all storage modes with one process launch per TP/DP topology.""" + if torch.accelerator.device_count() < WORLD_SIZE: + pytest.skip(f"Need {WORLD_SIZE} GPUs to run the test.") + monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1") + mp.spawn(_worker, args=(tp_size, get_open_port()), nprocs=WORLD_SIZE) + cleanup_dist_env_and_memory() + + +@pytest.mark.parametrize( + "cpu_offload,dp_size,error", + [ + (False, 4, "requires cpu_offload=True"), + (True, 1, "effective Engram DP size is 1"), + ], +) +def test_engram_dp_shared_memory_runtime_requirements( + monkeypatch, cpu_offload, dp_size, error +): + """Report unmet storage/topology requirements before allocating CUDA memory.""" + monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size) + with pytest.raises(ValueError, match=error): + ParallelEngramEmbedding( + sum(HEAD_SIZES), + DIM, + HEAD_SIZES, + cpu_offload=cpu_offload, + dp_shared_memory=True, + ) + + +@pytest.mark.parametrize( + "node_ids,dp_size,replica_size,expected", + [ + ([0] * 4, 1, 4, 1), + ([0] * 4, 4, 1, 4), + ([0] * 4, 2, 2, 2), + ([0] * 8 + [1] * 8, 8, 2, 4), + ([0] * 6 + [1] * 6, 4, 1, 2), + ([0] * 3 + [1] * 4, 7, 1, 1), + ([0] * 8 + [1] * 8 + [2] * 8, 8, 3, 1), + ([0, 1] * 4, 4, 2, 1), + ([0] * 3 + [1] * 5, 4, 2, 1), + ], +) +def test_engram_dp_shard_size_respects_node_boundaries( + node_ids, dp_size, replica_size, expected, monkeypatch +): + """Sharing must fall back when replicas or node rank ranges do not align.""" + monkeypatch.setattr(parallel_state, "get_node_count", lambda: len(set(node_ids))) + monkeypatch.setattr( + parallel_state, "get_world_group", lambda: SimpleNamespace(cpu_group=None) + ) + monkeypatch.setattr( + parallel_state, + "in_the_same_node_as", + lambda group, src: [node == node_ids[src] for node in node_ids], + ) + assert ( + parallel_state._engram_dp_shard_size(len(node_ids), dp_size, replica_size) + == expected + ) + + +@pytest.mark.parametrize("num_tokens", [0, 2, 4, 5]) +def test_engram_hash_padding_has_no_valid_rows(num_tokens, monkeypatch): + """Idle/padded slots must not query row zero; oversized batches are rejected.""" + monkeypatch.setattr(engram_ops, "engram_gathered_num_tokens", lambda: 4) + monkeypatch.setattr( + engram_ops, + "get_engram_dp_group", + lambda: SimpleNamespace(all_gather=lambda ids, dim: torch.cat([ids, ids], dim)), + ) + ids = torch.full((num_tokens, 2, 3), 17, dtype=torch.int32) + if num_tokens > 4: + with pytest.raises(ValueError, match="exceeds the DP token slot"): + gather_engram_hashes(ids) + return + gathered = gather_engram_hashes(ids).reshape(2, 4, 2, 3) + for replica in gathered: + torch.testing.assert_close(replica[:num_tokens], ids) + assert torch.all(replica[num_tokens:] == engram_ops.DEAD_ID) + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("uniform", [False, True]) +def test_engram_hash_padding_ignores_other_nodes(tp_size, uniform, monkeypatch): + """Two four-GPU nodes pad only to their own maximum, including TP x DP.""" + group_size = 4 // tp_size + counts = ( + [4] * (8 // tp_size) + if uniform + else [4096] + [2] * (group_size - 1) + [1, 3] + [0] * (group_size - 2) + ) + monkeypatch.setattr( + engram_ops, + "get_forward_context", + lambda: SimpleNamespace( + dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=torch.tensor(counts)) + ), + ) + for dp_rank in range(8 // tp_size): + slot = 4 if uniform else (4096 if dp_rank < group_size else 3) + ids = torch.full((counts[dp_rank], 2, 3), 17, dtype=torch.int32) + monkeypatch.setattr( + engram_ops, + "get_dp_group", + lambda rank=dp_rank: SimpleNamespace(rank_in_group=rank), + ) + monkeypatch.setattr( + engram_ops, + "get_engram_dp_group", + lambda rank=dp_rank: SimpleNamespace( + rank_in_group=rank % group_size, + world_size=group_size, + all_gather=lambda tensor, dim: torch.cat([tensor] * group_size, dim), + ), + ) + gathered = gather_engram_hashes(ids) + expected = ids.new_full((slot, 2, 3), engram_ops.DEAD_ID) + expected[: ids.shape[0]] = ids + torch.testing.assert_close( + gathered, expected.repeat(group_size, 1, 1), msg=f"DP rank {dp_rank}" + ) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 0262b42d11..d6ce5b36b5 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -89,7 +89,10 @@ def test_watermark_config_cli(): @pytest.mark.parametrize( "options", [ - ["--engram-config", '{"cpu_offload": false, "embedding_across_dp": true}'], + [ + "--engram-config", + '{"cpu_offload": false, "embedding_across_dp": true}', + ], [ "--engram-config.cpu_offload", "false", @@ -98,9 +101,8 @@ def test_watermark_config_cli(): ], ], ) -def test_engram_config_cli(options, monkeypatch): - """CLI settings take precedence over the legacy offload environment.""" - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1") +def test_engram_config_cli(options): + """JSON and dotted CLI options independently control Engram settings.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = EngineArgs.from_cli_args(parser.parse_args(options)) assert args.engram_config is not None @@ -109,14 +111,26 @@ def test_engram_config_cli(options, monkeypatch): @pytest.mark.parametrize( - "options, provided", - [([], False), (["--engram-config", "{}"], True)], + "options,provided,dp_shared_memory", + [ + ([], False, False), + (["--engram-config", "{}"], True, False), + ( + ["--engram-config", '{"dp_shared_memory": true}'], + True, + True, + ), + (["--engram-config.dp_shared_memory", "true"], True, True), + ], ) -def test_engram_config_cli_optional(options, provided): - """An explicit empty config must remain distinct from an omitted config.""" +def test_engram_config_cli_optional(options, provided, dp_shared_memory): + """Explicit configs honor defaults and the JSON/dotted DP shared-memory flag.""" parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) args = EngineArgs.from_cli_args(parser.parse_args(options)) assert (args.engram_config is not None) == provided + if provided: + assert args.engram_config.cpu_offload is True + assert args.engram_config.dp_shared_memory is dp_shared_memory @pytest.mark.parametrize( diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index f7aa6dd03e..c7e08ced82 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import bisect +import inspect from types import SimpleNamespace import pytest @@ -8,10 +9,13 @@ import torch from vllm.models.deepseek_v4_1.common import engram as engram_ops from vllm.models.deepseek_v4_1.common.engram import ( - Engram, - NgramHashState, - ParallelEngramEmbedding, + Engram as CommonEngram, ) +from vllm.models.deepseek_v4_1.common.engram import ( + NgramHashState, +) +from vllm.models.deepseek_v4_1.nvidia import engram as nvidia_engram_ops +from vllm.models.deepseek_v4_1.nvidia.engram import Engram, ParallelEngramEmbedding from vllm.platforms import current_platform @@ -107,7 +111,7 @@ def test_fused_engram_post_wkv_matches_reference( ) monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank) - module = Engram.__new__(Engram) + module = CommonEngram.__new__(CommonEngram) torch.nn.Module.__init__(module) module.dim = dim module.hc_mult = hc_mult @@ -116,7 +120,7 @@ def test_fused_engram_post_wkv_matches_reference( module.use_sequence_parallel = use_sequence_parallel module.embed_tokens = torch.nn.Identity() # `forward` reads rows staged by `prepare_embeddings`, so inject kv there. - module.embed_tokens.tp_size = 1 + module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1 module.staged_rows = kv.unsqueeze(1) if use_sequence_parallel: padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size)) @@ -151,6 +155,25 @@ def _hash_state(use_slot_cache: bool) -> NgramHashState: return state +@pytest.mark.parametrize("num_tokens", [0, 7]) +def test_engram_dummy_hashes_leave_history_untouched(num_tokens): + """Dummy forwards have no valid table IDs and do not alter cached history.""" + state = _hash_state(use_slot_cache=True) + state.multipliers = torch.empty(2, 4, dtype=torch.int64) + state.primes = torch.empty(2, 3, 8, dtype=torch.int64) + state._cache = torch.arange(16, dtype=torch.int32) + history = state._cache.clone() + + hashes, keep = state.dummy_hashes(torch.arange(num_tokens)) + + assert hashes.shape == (num_tokens, 2, 24) + assert hashes.dtype == torch.int32 + assert torch.all(hashes == engram_ops.DEAD_ID) + assert keep.shape == (num_tokens,) and keep.dtype == torch.bool + assert not keep.any() + torch.testing.assert_close(state._cache, history) + + @pytest.mark.parametrize("num_blocks", [4, 100]) def test_engram_cache_follows_kv_cache_binding(num_blocks): """Discard profiling history when the real KV cache replaces its storage.""" @@ -533,9 +556,11 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32): layer = ParallelEngramEmbedding.__new__(ParallelEngramEmbedding) torch.nn.Module.__init__(layer) layer.dim, layer.block_size, layer.tp_size = dim, block, 1 + layer.dp_size = 1 layer.n_hash_cols = layer.part_n_hash_cols = 24 layer.head_start = 0 layer.cpu_offload = cpu_offload + layer.dp_shared_memory = False layer.part_num_embeddings = rows # A window strictly inside the table, so unowned rows are exercised too. layer.vocab_start_idx, layer.vocab_end_idx = rows // 4, rows // 4 + rows // 2 @@ -559,12 +584,27 @@ def _make_embedding(cpu_offload, rows=4096, dim=256, block=32): return layer +@pytest.mark.parametrize( + "tp_size,dp_size,n_heads", [(1, 4, 5), (2, 2, 5), (4, 1, 6), (8, 1, 6)] +) +def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch): + """Reject empty owners before allocating weights or accessing CUDA.""" + monkeypatch.setattr( + engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr(nvidia_engram_ops, "get_engram_dp_size", lambda: dp_size) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(nvidia_engram_ops, "get_tensor_model_parallel_rank", lambda: 0) + with pytest.raises(AssertionError, match="ranks without hash heads"): + ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch): """Keep complete buckets and reconstruct head order, including TP padding.""" - head_sizes = (17, 19, 23, 29, 31, 37) + head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73) num_rows, dim = sum(head_sizes), 64 torch.manual_seed(0) weight = torch.randn(num_rows + 7, dim).to(torch.float8_e4m3fn) @@ -623,7 +663,7 @@ def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeyp module.staged_rows = torch.empty_like(shards[0]) module.prepare_embeddings(ids) torch.testing.assert_close(module.embed(ids), expected, rtol=0, atol=0) - # Slicing before head reordering must preserve padded heads and empty owners. + # Slicing before head reordering must preserve padded heads and tokens. module.use_sequence_parallel = True chunk = (len(ids) + tp_size - 1) // tp_size padded = torch.nn.functional.pad(expected, (0, 0, 0, 0, 0, (-len(ids)) % tp_size)) @@ -636,6 +676,31 @@ def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeyp ) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +def test_engram_lookup_reuses_jit_across_token_shapes(): + """Runtime token counts and launch grids must share one JIT variant.""" + layer = _make_embedding(cpu_offload=False) + kernel = engram_ops._engram_lookup_kernel + kernel_cache = kernel.device_caches[torch.accelerator.current_device_index()][0] + kernel_cache.clear() + + cache_sizes = [] + for background in (False, True): + for num_tokens in (1, 7, 256): + ids = torch.zeros(num_tokens, 24, dtype=torch.int32, device="cuda") + out = torch.empty( + num_tokens, + 24, + layer.dim, + dtype=torch.bfloat16, + device="cuda", + ) + layer.lookup(ids, out, background=background) + cache_sizes.append(len(kernel_cache)) + + assert cache_sizes == [1] * len(cache_sizes) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("background", [False, True]) @@ -660,19 +725,147 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") -@pytest.mark.parametrize("cpu_offload", [False, True]) +@pytest.mark.parametrize( + "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)] +) +def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload): + """Default offload uses pinned storage and a side stream; False uses HBM.""" + from vllm.config import EngramConfig + + config = SimpleNamespace(hidden_size=16, hc_mult=1, rms_norm_eps=1e-6) + layout = SimpleNamespace( + num_embeddings=(72,), + head_dim=32, + max_ngram_size=2, + n_heads=24, + primes=(((3,) * 24,),), + ) + vllm_config = SimpleNamespace( + engram_config=EngramConfig() + if cpu_offload is None + else EngramConfig(cpu_offload=cpu_offload), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + ) + offloaded = backend == "nvidia" and vllm_config.engram_config.cpu_offload + if backend == "common": + vllm_config.engram_config = None + monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config) + monkeypatch.setattr( + nvidia_engram_ops, "get_current_vllm_config", lambda: vllm_config + ) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + engram_ops, "ReplicatedLinear", lambda *a, **k: torch.nn.Identity() + ) + with torch.device("cuda"): + cls = Engram if backend == "nvidia" else CommonEngram + module = cls(config, None, layout, 0, False, "engram") + layer = module.embed_tokens + if offloaded: + assert layer.weight.device.type == "cpu" and layer.weight.is_pinned() + else: + assert layer.weight.device.type == "cuda" + layer.weight.data.copy_(torch.randn(layer.weight.shape).to(torch.float8_e4m3fn)) + layer.weight_scale_inv.data.fill_(127) + streams = [] + lookup = layer.lookup + + def track_lookup(ids, out, background=False): + streams.append(torch.cuda.current_stream()) + lookup(ids, out, background=background) + + layer.lookup = track_lookup + main = torch.cuda.current_stream() + ids = torch.randint( + layer.vocab_start_idx, + layer.vocab_end_idx, + (8, 24), + device="cuda", + dtype=torch.int32, + ) + module.prepare_embeddings(ids) + expected = _reference_lookup( + layer.weight.cuda(), + layer.weight_scale_inv.cuda(), + ids, + layer.vocab_start_idx, + layer.vocab_end_idx, + ) + torch.testing.assert_close(module.embed(ids), expected, atol=0, rtol=0) + if offloaded: + assert streams == [module._prefetch_stream] + assert streams[0] != main + else: + assert streams == [main] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize( + "cpu_offload,delay", + [(False, None), (True, None), (True, "producer"), (True, "lookup")], +) @pytest.mark.parametrize("capture", ["eager", "full", "breakable"]) -def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture): - """Consume early lookup results after a break, with fresh IDs each replay.""" - _run_engram_prepared_rows(cpu_offload, capture) +def test_engram_prepared_rows_survive_graph_breaks(cpu_offload, capture, delay): + """Temporary lookup IDs survive allocator reuse and graph replay.""" + _run_engram_prepared_rows(cpu_offload, capture, delay=delay) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.parametrize( + "missing_dependency", ["record_stream", "producer_wait", "consumer_wait"] +) +def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_dependency): + """Negative controls must expose incorrect rows when a dependency is removed.""" + if ( + missing_dependency == "record_stream" + and torch.cuda.memory.get_allocator_backend() != "native" + ): + pytest.skip( + "The allocator reuse negative control requires the native allocator" + ) + + def start(self, hash_ids, rows, stream): + if missing_dependency != "producer_wait": + stream.wait_stream(torch.cuda.current_stream()) + if missing_dependency != "record_stream": + hash_ids.record_stream(stream) + with torch.cuda.stream(stream): + self.embed_tokens.lookup(hash_ids, rows, background=True) + + def finish(self, stream): + pass + + monkeypatch.setattr(Engram, "_start_prefetch", start) + if missing_dependency == "consumer_wait": + monkeypatch.setattr(Engram, "_finish_prefetch", finish) + delay = "producer" if missing_dependency == "producer_wait" else "lookup" + try: + with pytest.raises( + AssertionError, match=r"Tensor-likes are not (?:equal|close)!" + ): + _run_engram_prepared_rows(True, "eager", delay=delay) + finally: + torch.accelerator.synchronize() def _run_engram_prepared_rows( - cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False + cpu_offload, capture, tp_size=1, rank=0, use_sequence_parallel=False, delay=None ): - from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture + from vllm.compilation.breakable_cudagraph import ( + BreakableCUDAGraphCapture, + eager_break_during_capture, + ) layer = _make_embedding(cpu_offload) + if delay == "lookup": + lookup = layer.lookup + + def delayed_lookup(indices, out, background=False): + torch.cuda._sleep(2_000_000) + lookup(indices, out, background=background) + + layer.lookup = delayed_lookup cols, num_tokens = (23, 65) if use_sequence_parallel else (24, 64) layer.n_hash_cols = cols layer.tp_size = tp_size @@ -681,6 +874,17 @@ def _run_engram_prepared_rows( engram = Engram.__new__(Engram) torch.nn.Module.__init__(engram) engram.embed_tokens = layer + engram._prefetch_stream = torch.cuda.Stream() if cpu_offload else None + # Exercise the production eager boundaries even when the test process + # imported Engram before breakable graphs were enabled. + if capture == "breakable": + import vllm.envs as envs + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(envs, "VLLM_USE_BREAKABLE_CUDAGRAPH", True) + for name in ("_start_prefetch", "_finish_prefetch"): + fn = eager_break_during_capture(inspect.unwrap(getattr(Engram, name))) + setattr(engram, name, fn.__get__(engram, Engram)) engram.use_sequence_parallel = use_sequence_parallel engram.staged_rows = torch.empty( num_tokens, @@ -709,7 +913,12 @@ def _run_engram_prepared_rows( embed = torch.compile(embed, backend="eager", fullgraph=True, dynamic=True) def step(cap=None): - engram.prepare_embeddings(src) + if delay == "producer": + torch.cuda._sleep(2_000_000) + # Drop the last reference to a non-contiguous input after launching lookup. + engram.prepare_embeddings(hashes.clone()[:, 1]) + # Exercise same-size allocator reuse before consuming the prefetched rows. + torch.empty_like(hashes).fill_(layer.part_num_embeddings - 1) if cap is not None: cap.add_eager(lambda: None) out.copy_(embed(src)) @@ -730,9 +939,12 @@ def _run_engram_prepared_rows( with torch.cuda.stream(warmup), graph: step(graph) torch.cuda.current_stream().wait_stream(warmup) - assert graph.num_graphs == 2 + assert graph.num_eager_breaks == (3 if cpu_offload else 1) for _ in range(3): + # Eager prefill can overwrite staging rows between decode replays. + hashes.random_(0, layer.part_num_embeddings) + step() hashes.random_(0, layer.part_num_embeddings) if graph is None: step() diff --git a/tests/test_config.py b/tests/test_config.py index 218ec61445..c2629a75ab 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1158,31 +1158,68 @@ def test_engram_tensor_parallel_size(dp_size: int, across_dp: bool, expected: in assert config.get_parallel_size(parallel) == expected -def test_engram_rejects_elastic_cross_dp(): +@pytest.mark.parametrize("option", ["embedding_across_dp", "dp_shared_memory"]) +def test_engram_rejects_elastic_cross_dp(option): parallel = ParallelConfig( tensor_parallel_size=4, data_parallel_size=2, distributed_executor_backend="mp", ) parallel.enable_elastic_ep = True - with pytest.raises(ValueError, match="embedding_across_dp.*elastic EP"): - EngramConfig(embedding_across_dp=True).verify_parallel_config(parallel) + with pytest.raises(ValueError, match=f"{option}.*elastic EP"): + EngramConfig(**{option: True}).verify_parallel_config(parallel) -@pytest.mark.parametrize("legacy", [None, "0", "1"]) -def test_engram_cpu_offload_environment_fallback(monkeypatch, legacy): - """Explicit settings must override the legacy environment fallback.""" - monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False) - if legacy is not None: - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy) - assert EngramConfig().cpu_offload == (legacy == "1") - assert EngramConfig(cpu_offload=False).cpu_offload is False - assert EngramConfig(cpu_offload=True).cpu_offload is True +def test_engram_dp_shared_memory_requires_cpu_offload(): + with pytest.raises(ValueError, match="requires cpu_offload"): + EngramConfig(cpu_offload=False, dp_shared_memory=True) + + +@pytest.mark.parametrize( + "dp_size,load_format,multithread,error", + [ + (1, "auto", False, "requires data_parallel_size > 1"), + (2, "dummy", False, "requires load_format"), + (2, "sharded_state", False, "requires load_format"), + (2, "auto", False, None), + (2, "safetensors", True, None), + (2, "pt", True, None), + ], +) +def test_engram_dp_shared_memory_config_validation( + monkeypatch, dp_size, load_format, multithread, error +): + """Reject invalid shared configs before distributed init; allow threaded loads.""" + monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + config = cast( + VllmConfig, + SimpleNamespace( + model_config=SimpleNamespace( + architecture="DeepseekV41ForCausalLM", + hf_text_config=SimpleNamespace(engram_layer_ids=[1]), + ), + speculative_config=None, + engram_config=EngramConfig(cpu_offload=True, dp_shared_memory=True), + parallel_config=ParallelConfig(data_parallel_size=dp_size), + load_config=LoadConfig( + load_format=load_format, + model_loader_extra_config={"enable_multithread_load": multithread}, + ), + ), + ) + if error: + with pytest.raises(ValueError, match=error): + VllmConfig._resolve_and_verify_engram_config(config) + else: + VllmConfig._resolve_and_verify_engram_config(config) @pytest.mark.parametrize( "architecture, ple_layers, cuda, supported", [ + ("DeepseekV41ForCausalLM", [1], True, True), + ("DeepseekV41ForCausalLM", [], True, False), + ("DeepseekV41ForCausalLM", [1], False, False), ("Qwen4ExpForCausalLM", [1], True, True), ("Qwen4ExpForConditionalGeneration", [1], True, True), ("Qwen4ExpForCausalLM", [], True, False), @@ -1201,7 +1238,9 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo ModelConfig, SimpleNamespace( architecture=architecture, - hf_text_config=SimpleNamespace(ple_layer_ids=ple_layers), + hf_text_config=SimpleNamespace( + ple_layer_ids=ple_layers, engram_layer_ids=ple_layers + ), ), ) if architecture is not None @@ -1214,55 +1253,73 @@ def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, suppo with pytest.raises(ValueError, match="requires a model with supported Engram"): config.verify_model_config(model) + resolved = cast( + VllmConfig, + SimpleNamespace( + model_config=model, + speculative_config=None, + parallel_config=ParallelConfig(), + load_config=LoadConfig(load_format="dummy"), + engram_config=None, + ), + ) + VllmConfig._resolve_and_verify_engram_config(resolved) + assert (resolved.engram_config is not None) == supported + if supported: + assert resolved.engram_config.cpu_offload is True -def test_engram_config_defaults_to_none(monkeypatch): + +@pytest.mark.parametrize( + ("value", "expected"), [(None, True), ("0", False), ("1", True)] +) +def test_engram_cpu_offload_environment_default(monkeypatch, value, expected): monkeypatch.delenv("VLLM_PLE_CPU_OFFLOAD", raising=False) + if value is not None: + monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", value) + assert EngramConfig().cpu_offload is expected + assert EngramConfig(cpu_offload=False).cpu_offload is False + assert EngramConfig(cpu_offload=True).cpu_offload is True + + +def test_engram_config_defaults_to_none(): config = VllmConfig() assert config.engram_config is None assert config.compute_hash() -@pytest.mark.parametrize("legacy", ["0", "1"]) -def test_engram_none_resolves_legacy_offload(monkeypatch, legacy): - """Legacy enablement materializes a config before model validation.""" - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy) - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("enable_dbo,ubatch_size", [(True, 0), (False, 2)]) +def test_deepseek_engram_rejects_microbatching(explicit, enable_dbo, ubatch_size): + """Reject overlapping Engram staging even without an explicit EngramConfig.""" config = cast( VllmConfig, SimpleNamespace( model_config=SimpleNamespace( - architecture="Qwen4ExpForCausalLM", - hf_text_config=SimpleNamespace(ple_layer_ids=[1]), + architecture="DeepseekV41ForCausalLM", + hf_text_config=SimpleNamespace(engram_layer_ids=[1]), ), speculative_config=None, - engram_config=None, - parallel_config=ParallelConfig(), + engram_config=EngramConfig(cpu_offload=False) if explicit else None, + parallel_config=ParallelConfig( + enable_dbo=enable_dbo, ubatch_size=ubatch_size + ), ), ) - VllmConfig._resolve_and_verify_engram_config(config) - if legacy == "1": - assert config.engram_config is not None - assert config.engram_config.cpu_offload is True - else: - assert config.engram_config is None + with pytest.raises( + ValueError, match="Engram does not support DBO or microbatching" + ): + VllmConfig._resolve_and_verify_engram_config(config) -@pytest.mark.parametrize("legacy", ["0", "1"]) -def test_engram_explicit_config_requires_supported_model(monkeypatch, legacy): +def test_engram_explicit_config_requires_supported_model(): """Explicit all-false settings still opt into model validation.""" - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", legacy) with pytest.raises(ValueError, match="requires a model with supported Engram"): VllmConfig(engram_config=EngramConfig(cpu_offload=False)) -def test_engram_legacy_offload_requires_supported_model(monkeypatch): - monkeypatch.setenv("VLLM_PLE_CPU_OFFLOAD", "1") - with pytest.raises(ValueError, match="requires a model with supported Engram"): - VllmConfig() - - @pytest.mark.parametrize("target_has_ple", [False, True]) -def test_engram_draft_config_validates_target(monkeypatch, target_has_ple): +@pytest.mark.parametrize("explicit", [False, True]) +def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit): """MTP may inherit cross-DP sharding without having its own PLE layers.""" monkeypatch.setattr(current_platform, "is_cuda", lambda: True) target = SimpleNamespace( @@ -1277,24 +1334,29 @@ def test_engram_draft_config_validates_target(monkeypatch, target_has_ple): speculative_config=SimpleNamespace( draft_model_config=draft, target_model_config=target ), - engram_config=EngramConfig(embedding_across_dp=True), + engram_config=EngramConfig(embedding_across_dp=True) if explicit else None, parallel_config=ParallelConfig(), + load_config=LoadConfig(load_format="dummy"), ), ) - if target_has_ple: - VllmConfig._resolve_and_verify_engram_config(config) - else: + if explicit and not target_has_ple: with pytest.raises(ValueError, match="requires a model with supported Engram"): VllmConfig._resolve_and_verify_engram_config(config) + else: + VllmConfig._resolve_and_verify_engram_config(config) + assert (config.engram_config is not None) == target_has_ple + if target_has_ple: + assert config.engram_config.cpu_offload is True -def test_engram_hash_tracks_storage_and_sharding(): +def test_engram_hash_tracks_execution_options(): configs = [ - EngramConfig(cpu_offload=offload, embedding_across_dp=across_dp) - for offload in (False, True) - for across_dp in (False, True) + EngramConfig(cpu_offload=True), + EngramConfig(cpu_offload=False), + EngramConfig(cpu_offload=True, embedding_across_dp=True), + EngramConfig(cpu_offload=True, dp_shared_memory=True), ] - assert len({config.compute_hash() for config in configs}) == 4 + assert len({config.compute_hash() for config in configs}) == len(configs) @pytest.mark.parametrize("port", [1, 29550, 65535]) diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 2aede1868d..5abb2bd34a 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -3,18 +3,17 @@ from typing import TYPE_CHECKING -from pydantic import Field +from pydantic import Field, model_validator +from typing_extensions import Self import vllm.envs as envs from vllm.config.utils import config, get_hash_factors, hash_factors -from vllm.logger import init_logger if TYPE_CHECKING: + from vllm.config.load import LoadConfig from vllm.config.model import ModelConfig from vllm.config.parallel import ParallelConfig -logger = init_logger(__name__) - # Architecture -> the hf_text_config field naming its n-gram layers. A model is # only configurable here if it actually has such layers to store. _NGRAM_LAYER_FIELDS = { @@ -25,28 +24,44 @@ _NGRAM_LAYER_FIELDS = { def _default_cpu_offload() -> bool: - """Honor the legacy environment variable only when the field is omitted.""" - if envs.is_set("VLLM_PLE_CPU_OFFLOAD"): - logger.warning_once( - "VLLM_PLE_CPU_OFFLOAD is a legacy setting and may be removed in a " - "future release. Use --engram-config.cpu_offload instead." - ) return envs.VLLM_PLE_CPU_OFFLOAD +def model_has_engram_layers(model_config: "ModelConfig | None") -> bool: + """Whether the model carries n-gram embedding layers.""" + if model_config is None: + return False + field = _NGRAM_LAYER_FIELDS.get(model_config.architecture) + if field is None: + return False + return bool(getattr(model_config.hf_text_config, field, None)) + + @config class EngramConfig: """Configuration for Engram embedding storage and sharding.""" cpu_offload: bool = Field(default_factory=_default_cpu_offload) """Store embedding weights in pinned CPU memory for UVA lookup. - Defaults to False, or VLLM_PLE_CPU_OFFLOAD when set for compatibility. - An explicit value takes precedence over the legacy environment variable.""" + Defaults to VLLM_PLE_CPU_OFFLOAD, which is enabled by default. An explicit + value takes precedence over the environment variable.""" embedding_across_dp: bool = False """Shard embeddings across TP and all DP ranks when enabled. Otherwise, each DP rank has a separate TP-sharded embedding replica.""" + dp_shared_memory: bool = False + """Share CPU-offloaded embedding weights between co-located + DP replicas. Each node stores one copy of every TP shard, reducing host + memory without per-step Engram DP collectives. Requires sufficient + /dev/shm capacity and a shared IPC namespace.""" + + @model_validator(mode="after") + def _validate_shared_memory(self) -> Self: + if self.dp_shared_memory and not self.cpu_offload: + raise ValueError("dp_shared_memory requires cpu_offload=True") + return self + def verify_model_config(self, model_config: "ModelConfig | None") -> None: """Reject Engram configuration for models without n-gram embeddings.""" from vllm.platforms import current_platform @@ -64,13 +79,16 @@ class EngramConfig: ): raise ValueError( "EngramConfig requires a model with supported Engram " - "embeddings. Currently only the CUDA Qwen4Exp and DeepSeek " - "V4.1 implementations with non-empty n-gram layer ids are " - "supported." + "embeddings, non-empty n-gram layer ids, and CUDA." ) def verify_parallel_config(self, parallel_config: "ParallelConfig") -> None: """Reject unsupported embedding parallel topologies.""" + if self.dp_shared_memory: + if parallel_config.data_parallel_size <= 1: + raise ValueError("dp_shared_memory requires data_parallel_size > 1.") + if parallel_config.enable_elastic_ep: + raise ValueError("dp_shared_memory is not supported with elastic EP.") if ( self.embedding_across_dp and parallel_config.data_parallel_size > 1 @@ -80,6 +98,18 @@ class EngramConfig: "Engram embedding_across_dp is not supported with elastic EP yet." ) + def verify_load_config(self, load_config: "LoadConfig") -> None: + """Shared tables require a loader that invokes parameter weight callbacks.""" + if self.dp_shared_memory and load_config.load_format not in ( + "auto", + "safetensors", + "pt", + ): + raise ValueError( + "dp_shared_memory requires load_format 'auto', " + f"'safetensors' or 'pt'; got {load_config.load_format!r}." + ) + def get_parallel_size(self, parallel_config: "ParallelConfig") -> int: """Derive the embedding group size from the parallel configuration.""" size = parallel_config.tensor_parallel_size diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6fab6f0c71..7f0a3f0d87 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -34,7 +34,7 @@ from .device import DeviceConfig from .diffusion import DiffusionConfig from .ec_manager_config import EncoderCacheManagerConfig from .ec_transfer import ECTransferConfig -from .engram import EngramConfig +from .engram import EngramConfig, model_has_engram_layers from .kernel import KernelConfig from .kv_events import KVEventsConfig from .kv_transfer import KVTransferConfig @@ -381,7 +381,7 @@ class VllmConfig: attention_config: AttentionConfig = Field(default_factory=AttentionConfig) """Attention configuration.""" engram_config: EngramConfig | None = None - """Optional Engram configuration, only valid for supported PLE models.""" + """N-gram embedding storage and sharding settings.""" mamba_config: MambaConfig = Field(default_factory=MambaConfig) """Mamba configuration.""" kernel_config: KernelConfig = Field(default_factory=KernelConfig) @@ -1218,22 +1218,37 @@ class VllmConfig: ) def _resolve_and_verify_engram_config(self) -> None: - """Resolve legacy offload settings and validate model and parallel configs.""" - if self.engram_config is None: - if not envs.VLLM_PLE_CPU_OFFLOAD: - return - self.engram_config = EngramConfig() + """Resolve defaults and validate n-gram embedding settings.""" + from vllm.platforms import current_platform + model_config = self.model_config speculative_config = self.speculative_config # Draft configs inherit the target's communication groups and settings. - # Qwen4Exp MTP itself disables PLE, so validate its target instead. + # Validate the target because the draft may disable n-gram embeddings. if ( speculative_config is not None and model_config is speculative_config.draft_model_config ): model_config = speculative_config.target_model_config + if ( + model_config is not None + and model_config.architecture == "DeepseekV41ForCausalLM" + and getattr(model_config.hf_text_config, "engram_layer_ids", None) + and self.parallel_config.use_ubatching + ): + raise ValueError( + "DeepSeek V4.1 Engram does not support DBO or microbatching. " + "Disable --enable-dbo and set --ubatch-size to 0." + ) + if self.engram_config is None: + if not current_platform.is_cuda() or not model_has_engram_layers( + model_config + ): + return + self.engram_config = EngramConfig() self.engram_config.verify_model_config(model_config) self.engram_config.verify_parallel_config(self.parallel_config) + self.engram_config.verify_load_config(self.load_config) logger.info_once("Resolved Engram configuration: %s", str(self.engram_config)) def __post_init__(self): diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 0466779c34..dc3022a16f 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -32,6 +32,7 @@ from collections.abc import Callable from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import dataclass from datetime import timedelta +from math import gcd from multiprocessing import shared_memory from typing import TYPE_CHECKING, Any, Protocol from unittest.mock import patch @@ -1558,6 +1559,23 @@ def get_etp_group() -> GroupCoordinator: return _ETP +_ENGRAM_DP: GroupCoordinator | None = None + + +def get_engram_dp_group() -> GroupCoordinator | None: + """Return the DP replicas that share one engram embedding table. + + None when every replica holds a full (TP-sharded) copy, which is the case + for models without engram layers and for replicas that span nodes. + """ + return _ENGRAM_DP + + +def get_engram_dp_size() -> int: + """Number of DP replicas one engram embedding table is sharded over.""" + return _ENGRAM_DP.world_size if _ENGRAM_DP is not None else 1 + + _DCP: GroupCoordinator | None = None @@ -1933,6 +1951,29 @@ def _elastic_join_warmup_ctx() -> AbstractContextManager[Any]: return nullcontext() +def _engram_dp_shard_size( + world_size: int, + data_parallel_size: int, + replica_size: int, +) -> int: + """Size of DP shards aligned with complete replicas and node boundaries.""" + node_count = get_node_count() + replicas_per_node, remainder = divmod(world_size, node_count * replica_size) + if remainder: + return 1 + # Shards must divide both the DP group and each node's replica count. + shard_size = gcd(data_parallel_size, replicas_per_node) + if shard_size == 1 or node_count == 1: + return shard_size + ranks_per_node = replicas_per_node * replica_size + for start in range(0, world_size, ranks_per_node): + same_node = in_the_same_node_as(get_world_group().cpu_group, start) + node_ranks = [rank for rank, local in enumerate(same_node) if local] + if node_ranks != list(range(start, start + ranks_per_node)): + return 1 + return shard_size + + def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, @@ -2055,6 +2096,40 @@ def initialize_model_parallel( group_name="etp", ) + # Build the node-local DP group used to shard or share an Engram table. + global _ENGRAM_DP + assert _ENGRAM_DP is None, "engram data parallel group is already initialized" + engram_dp_size = 1 + engram_config = config.engram_config + if ( + engram_config is not None + and config.model_config is not None + and config.model_config.architecture == "DeepseekV41ForCausalLM" + and not enable_elastic_ep + ): + engram_dp_size = _engram_dp_shard_size( + world_size, + data_parallel_size, + tensor_model_parallel_size + * pipeline_model_parallel_size + * prefill_context_model_parallel_size, + ) + if engram_dp_size > 1: + group_ranks = ( + all_ranks.permute(0, 2, 3, 4, 1).reshape(-1, engram_dp_size).unbind(0) + ) + _ENGRAM_DP = init_model_parallel_group( + [x.tolist() for x in group_ranks], + get_world_group().local_rank, + backend, + group_name="edp", + ) + logger.info_once( + "Engram node-local DP group size: %d (TP=%d)", + engram_dp_size, + tensor_model_parallel_size, + ) + # Build the DCP model-parallel groups. global _DCP assert _DCP is None, "decode context model parallel group is already initialized" @@ -2304,6 +2379,11 @@ def destroy_model_parallel(): _TP.destroy() _TP = None + global _ENGRAM_DP + if _ENGRAM_DP: + _ENGRAM_DP.destroy() + _ENGRAM_DP = None + global _DCP if _DCP: _DCP.destroy() diff --git a/vllm/envs.py b/vllm/envs.py index aa32fc38ba..0269a96409 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -154,7 +154,7 @@ if TYPE_CHECKING: VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False VLLM_ENABLE_V1_MULTIPROCESSING: bool = True VLLM_LOG_BATCHSIZE_INTERVAL: float = -1 - VLLM_PLE_CPU_OFFLOAD: bool = False + VLLM_PLE_CPU_OFFLOAD: bool = True VLLM_DISABLE_COMPILE_CACHE: bool = False VLLM_REPLICATE_EMBED: bool = False VLLM_USE_LAYERNAME: bool = True @@ -2059,11 +2059,8 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_LOG_MODEL_INSPECTION": lambda: bool( int(os.getenv("VLLM_LOG_MODEL_INSPECTION", "0")) ), - # Keep Qwen4Exp PLE embedding tables in pinned CPU memory and gather their - # rows through UVA on a dedicated CUDA stream. - # Legacy fallback for EngramConfig.cpu_offload, which takes precedence. - # This environment variable may be removed in a future release. - "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "0"))), + # Store n-gram embedding tables in pinned CPU memory for UVA lookup. + "VLLM_PLE_CPU_OFFLOAD": lambda: bool(int(os.getenv("VLLM_PLE_CPU_OFFLOAD", "1"))), # Debug logging for --enable-mfu-metrics "VLLM_DEBUG_MFU_METRICS": lambda: bool( int(os.getenv("VLLM_DEBUG_MFU_METRICS", "0")) diff --git a/vllm/models/deepseek_v4_1/common/engram.py b/vllm/models/deepseek_v4_1/common/engram.py index 4cfba842b0..c415bbaa6e 100644 --- a/vllm/models/deepseek_v4_1/common/engram.py +++ b/vllm/models/deepseek_v4_1/common/engram.py @@ -51,8 +51,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.utils import set_weight_attrs from vllm.triton_utils import tl, triton -from vllm.utils.platform_utils import is_uva_available -from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor logger = init_logger(__name__) @@ -236,6 +234,7 @@ def _write_hash_cache_kernel( tl.store(cache + slot, value, valid) +# Keep request shapes out of the cache key so one warmed variant covers runtime batches. @triton.jit( do_not_specialize=[ "num_tokens", @@ -462,6 +461,21 @@ class NgramHashState(nn.Module): self._kv_cache_ref = weakref.ref(kv_cache) return True + def dummy_hashes( + self, input_ids: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Participate in DP lookups without valid rows or hash-cache updates.""" + num_tokens = input_ids.shape[0] + num_layers, max_ngram = self.multipliers.shape + num_heads = self.primes.shape[-1] + hashes = input_ids.new_full( + (num_tokens, num_layers, (max_ngram - 1) * num_heads), + DEAD_ID, + dtype=torch.int32, + ) + keep = torch.zeros(num_tokens, dtype=torch.bool, device=input_ids.device) + return hashes, keep + def forward( self, input_ids: torch.Tensor, @@ -561,7 +575,16 @@ def _engram_head_shard_weight_loader( param.data.copy_(shard) -@triton.jit +@triton.jit( + do_not_specialize=[ + "vocab_start", + "vocab_end", + "num_rows", + "ids_stride_t", + "ids_stride_h", + "GRID", + ] +) def _engram_lookup_kernel( weight, scales, @@ -578,7 +601,7 @@ def _engram_lookup_kernel( DIM: tl.constexpr, QUANT_BLOCK: tl.constexpr, BLOCK_R: tl.constexpr, - GRID: tl.constexpr, + GRID, ): """Gather fp8 rows, apply their ue8m0 block scales, write bf16. @@ -620,12 +643,9 @@ def _engram_lookup_kernel( class ParallelEngramEmbedding(nn.Module): - """The n-gram hash table, sharded by complete hash heads over TP ranks. - Rows stay fp8 and are dequantized with ue8m0 per-32 scales on lookup. + """TP-sharded hash heads with FP8 rows and per-block E8M0 scales.""" - With `cpu_offload` the shard lives in pinned host memory and is read over - UVA instead of HBM; the TP sharding is unchanged either way. - """ + _weight_loader = staticmethod(_engram_head_shard_weight_loader) def __init__( self, @@ -633,89 +653,55 @@ class ParallelEngramEmbedding(nn.Module): dim: int, head_sizes: tuple[int, ...], block_size: int = 32, - cpu_offload: bool = False, - ): + ) -> None: super().__init__() - tp_size = get_tensor_model_parallel_world_size() - tp_rank = get_tensor_model_parallel_rank() assert head_sizes and all(size > 0 for size in head_sizes) assert sum(head_sizes) <= num_embeddings - if cpu_offload and not is_uva_available(): - raise RuntimeError("Engram CPU offload requires UVA support") self.num_embeddings = num_embeddings self.dim = dim self.block_size = block_size self.n_hash_cols = len(head_sizes) - self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, tp_size) - self.head_start = tp_rank * self.part_n_hash_cols + self.tp_size = get_tensor_model_parallel_world_size() + num_shards, head_rank = self._get_shard_info() + self.part_n_hash_cols = triton.cdiv(self.n_hash_cols, num_shards) + # TODO: Support row-wise sharding when there are too few hash heads. + assert (num_shards - 1) * self.part_n_hash_cols < self.n_hash_cols, ( + f"Engram sharding leaves ranks without hash heads: " + f"{self.n_hash_cols} heads over {num_shards} shards" + ) + self.head_start = head_rank * self.part_n_hash_cols head_end = self.head_start + self.part_n_hash_cols self.vocab_start_idx = sum(head_sizes[: self.head_start]) self.vocab_end_idx = sum(head_sizes[:head_end]) self.part_num_embeddings = self.vocab_end_idx - self.vocab_start_idx - self.tp_size = tp_size - self.cpu_offload = cpu_offload - self._views: tuple[torch.Tensor, torch.Tensor] | None = None - self._view_src: tuple[int, int] | None = None self._num_sms = torch.cuda.get_device_properties( torch.accelerator.current_device_index() ).multi_processor_count - - # Explicit device: model init runs under a `torch.device("cuda")` - # context, which would otherwise put the shard in HBM. - kwargs = {"device": "cpu", "pin_memory": True} if cpu_offload else {} - self.weight = nn.Parameter( - torch.empty( - self.part_num_embeddings, dim, dtype=torch.float8_e4m3fn, **kwargs - ), - requires_grad=False, - ) - self.weight_scale_inv = nn.Parameter( - torch.empty( - self.part_num_embeddings, - dim // block_size, - dtype=torch.uint8, - **kwargs, - ), - requires_grad=False, - ) + weight, scales = self._allocate_weights() + self.weight = nn.Parameter(weight, requires_grad=False) + self.weight_scale_inv = nn.Parameter(scales, requires_grad=False) for param in (self.weight, self.weight_scale_inv): set_weight_attrs( param, { - "weight_loader": _engram_head_shard_weight_loader, + "weight_loader": self._weight_loader, "engram_vocab_start": self.vocab_start_idx, }, ) - if cpu_offload: - # Constant dummy values avoid randomizing huge CPU lookup tables. - set_weight_attrs(self.weight, {"dummy_weight_value": 1.0}) - # The ue8m0 encoding of scale 1.0 is exponent byte 127. - set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127}) - logger.info( - "Engram table offloaded to pinned host memory: %d rows x %d, " - "%.2f GiB per rank", - self.part_num_embeddings, - dim, - self.part_num_embeddings * (dim + dim // block_size) / 1024**3, - ) + + def _get_shard_info(self) -> tuple[int, int]: + return self.tp_size, get_tensor_model_parallel_rank() + + def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.empty(self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn), + torch.empty( + self.part_num_embeddings, self.dim // self.block_size, dtype=torch.uint8 + ), + ) def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: - """Parameters when resident, else cached UVA views of the pinned shard. - - Rebuilt if anything swaps `.data`, so a stale device pointer cannot - survive silently. - """ - if not self.cpu_offload: - return self.weight.data, self.weight_scale_inv.data - src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr()) - if self._view_src != src: - self._views = ( - get_accelerator_view_from_cpu_tensor(self.weight.data), - get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data), - ) - self._view_src = src - assert self._views is not None - return self._views + return self.weight.data, self.weight_scale_inv.data def lookup( self, indices: torch.Tensor, out: torch.Tensor, background: bool = False @@ -753,7 +739,7 @@ class ParallelEngramEmbedding(nn.Module): def forward(self, indices: torch.Tensor) -> torch.Tensor: """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim] - bf16, gathered from all TP shards.""" + bf16, gathered from all shards for this replica's tokens.""" out = torch.empty( (indices.shape[0], self.part_n_hash_cols, self.dim), dtype=torch.bfloat16, @@ -762,8 +748,7 @@ class ParallelEngramEmbedding(nn.Module): self.lookup(indices, out) if self.tp_size > 1: out = tensor_model_parallel_all_gather(out, dim=1) - out = out[:, : self.n_hash_cols] - return out + return out[:, : self.n_hash_cols] @triton.jit(do_not_specialize=["num_kv_tokens"]) @@ -859,7 +844,7 @@ def _fused_engram_post_wkv_kernel( @triton.jit(do_not_specialize=["num_tokens", "token_start", "num_elements"]) -def _engram_sp_rows_kernel( +def _engram_select_rows_kernel( gathered, output, num_tokens, @@ -880,6 +865,32 @@ def _engram_sp_rows_kernel( tl.store(output + offsets, values, offsets < num_elements) +def _engram_select_rows( + gathered: torch.Tensor, + output: torch.Tensor, + source_tokens: int, + token_start: int, + local_width: int, +) -> None: + """Copy one token window out of a rank-major gathered buffer. + + Both gathers land rank-major ([rank][token][local width]); this walks the + window the rank keeps and lays its ranks out side by side as width. + """ + if output.numel() == 0: + return + _engram_select_rows_kernel[(triton.cdiv(output.numel(), 1024),)]( + gathered, + output, + source_tokens, + token_start, + output.numel(), + local_width, + output.shape[1] * output.shape[2], + BLOCK_SIZE=1024, + ) + + class Engram(nn.Module): """Writes an n-gram lookup into the residual stream, gated by how well it matches that stream. @@ -910,13 +921,7 @@ class Engram(nn.Module): # Named ``embed_tokens`` so the checkpoint's ``engram.embed.weight`` # survives the mapper's ``embed.weight`` -> ``embed_tokens.weight`` # suffix rule. - engram_config = get_current_vllm_config().engram_config - self.embed_tokens = ParallelEngramEmbedding( - layout.num_embeddings[layer_hash_index], - layout.head_dim, - tuple(size for order in layout.primes[layer_hash_index] for size in order), - cpu_offload=engram_config.cpu_offload if engram_config else True, - ) + self.embed_tokens = self._create_embedding(layout, layer_hash_index) n_hash_cols = (layout.max_ngram_size - 1) * layout.n_heads self.wkv = ReplicatedLinear( n_hash_cols * layout.head_dim, @@ -936,40 +941,54 @@ class Engram(nn.Module): ) max_tokens = get_current_vllm_config().scheduler_config.max_num_batched_tokens - # Keep lookup results alive across breakable graph segments. + self._init_staging(max_tokens, layout.head_dim) + + def _create_embedding( + self, layout: EngramLayout, layer_hash_index: int + ) -> ParallelEngramEmbedding: + return ParallelEngramEmbedding( + layout.num_embeddings[layer_hash_index], + layout.head_dim, + tuple(size for order in layout.primes[layer_hash_index] for size in order), + ) + + def _init_staging(self, max_tokens: int, head_dim: int) -> None: + # Persistent storage keeps lookup addresses stable across graph replays. self.staged_rows = torch.empty( max_tokens, self.embed_tokens.part_n_hash_cols, - layout.head_dim, + head_dim, dtype=torch.bfloat16, ) def prepare_embeddings(self, hash_ids: torch.Tensor) -> None: - """Gather this layer's rows on the main stream before decoder layers.""" - self.embed_tokens.lookup(hash_ids, self.staged_rows[: hash_ids.shape[0]]) + """Look up rows before the decoder layers consume them.""" + rows = self.staged_rows[: hash_ids.shape[0]] + assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small" + self.embed_tokens.lookup(hash_ids, rows) + + def _ready_rows(self, num_tokens: int) -> torch.Tensor: + return self.staged_rows[:num_tokens] def embed(self, hash_ids: torch.Tensor) -> torch.Tensor: """Gather heads, returning only local tokens when SP is enabled.""" - rows = self.staged_rows[: hash_ids.shape[0]] + rows = self._ready_rows(hash_ids.shape[0]) if self.embed_tokens.tp_size == 1: - return rows + return rows[:, : self.embed_tokens.n_hash_cols] if self.use_sequence_parallel: tp_size = self.embed_tokens.tp_size num_tokens, local_heads, dim = rows.shape gathered = tensor_model_parallel_all_gather(rows, dim=0) chunk = (num_tokens + tp_size - 1) // tp_size - rows = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim)) - _engram_sp_rows_kernel[(triton.cdiv(rows.numel(), 1024),)]( + out = rows.new_empty((chunk, self.embed_tokens.n_hash_cols, dim)) + _engram_select_rows( gathered, - rows, + out, num_tokens, get_tensor_model_parallel_rank() * chunk, - rows.numel(), local_heads * dim, - self.embed_tokens.n_hash_cols * dim, - BLOCK_SIZE=1024, ) - return rows + return out rows = tensor_model_parallel_all_gather(rows, dim=1) return rows[:, : self.embed_tokens.n_hash_cols] diff --git a/vllm/models/deepseek_v4_1/nvidia/engram.py b/vllm/models/deepseek_v4_1/nvidia/engram.py new file mode 100644 index 0000000000..b141b6c349 --- /dev/null +++ b/vllm/models/deepseek_v4_1/nvidia/engram.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch.""" + +import mmap +import tempfile +import weakref +from contextlib import ExitStack + +import numpy as np +import torch + +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_dp_group, + get_engram_dp_group, + get_engram_dp_size, + get_tensor_model_parallel_rank, + tensor_model_parallel_all_gather, +) +from vllm.distributed.parallel_state import GroupCoordinator +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.utils import set_weight_attrs +from vllm.models.deepseek_v4_1.common.engram import ( + DEAD_ID, + EngramLayout, + _engram_head_shard_weight_loader, + _engram_select_rows, +) +from vllm.models.deepseek_v4_1.common.engram import ( + Engram as BaseEngram, +) +from vllm.models.deepseek_v4_1.common.engram import ( + ParallelEngramEmbedding as BaseParallelEngramEmbedding, +) +from vllm.utils.platform_utils import is_uva_available +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + +logger = init_logger(__name__) + + +def engram_head_shard_rank() -> int: + """This rank's slot among the hash-head shards of one engram table. + + TP-major, so the shards a DP gather brings in are contiguous heads and + the following TP gather completes the head order. + """ + dp_group = get_engram_dp_group() + dp_size = dp_group.world_size if dp_group is not None else 1 + dp_rank = dp_group.rank_in_group if dp_group is not None else 0 + return get_tensor_model_parallel_rank() * dp_size + dp_rank + + +def engram_gathered_num_tokens() -> int: + """Per-replica token slot for the node-local Engram DP group.""" + dp_metadata = get_forward_context().dp_metadata + if dp_metadata is None: + raise RuntimeError("a DP-shared engram table needs DP token metadata") + group = get_engram_dp_group() + assert group is not None + # Engram groups are contiguous slices of the full DP group. + start = get_dp_group().rank_in_group - group.rank_in_group + return int( + dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max() + ) + + +def gather_engram_hashes( + hash_ids: torch.Tensor, *, dp_shared_memory: bool = False +) -> torch.Tensor: + """Collect the n-gram ids of every DP replica sharing one table. + + Replicas are padded to a common token slot, so the gathered shape is + static under CUDA graph capture (where DP already pads alike). + """ + dp_group = get_engram_dp_group() + if dp_group is None or dp_shared_memory: + return hash_ids + slot = engram_gathered_num_tokens() + if hash_ids.shape[0] > slot: + raise ValueError("Engram token count exceeds the DP token slot") + if hash_ids.shape[0] < slot: + pad = hash_ids.new_full( + (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID + ) + hash_ids = torch.cat((hash_ids, pad)) + return dp_group.all_gather(hash_ids, dim=0) + + +class DPSharedEngramStorage: + """Registered host weights shared by a node-local DP group with one writer.""" + + def __init__( + self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator + ) -> None: + self.group = group + weight_bytes = num_rows * dim + storage = self._allocate(weight_bytes + weight_bytes // block_size) + self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim) + self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size) + self._views: tuple[torch.Tensor, torch.Tensor] | None = None + + def _allocate(self, num_bytes: int) -> torch.Tensor: + """Map and register one physical allocation across a node-local DP group.""" + from vllm.distributed.device_communicators.shm_broadcast import ( + check_shm_free_space, + ) + + group = self.group + with ExitStack() as stack: + path, error = None, None + if group.rank_in_group == 0: + try: + check_shm_free_space(num_bytes) + backing_file = stack.enter_context( + tempfile.NamedTemporaryFile( + prefix="vllm_engram_", dir="/dev/shm" + ) + ) + backing_file.truncate(num_bytes) + path = backing_file.name + except Exception as exc: + error = f"{type(exc).__name__}: {exc}" + path, error = group.broadcast_object((path, error)) + if error is not None: + raise RuntimeError( + "Engram shared-memory creation failed on EDP rank 0: " + error + ) + + mapping = owner = tensor = finalizer = None + stage = "open" + try: + try: + with open(path, "r+b") as file: + stage = "mmap" + mapping = mmap.mmap( + file.fileno(), num_bytes, flags=mmap.MAP_SHARED + ) + stage = "cudaHostRegister" + owner = np.frombuffer(mapping, dtype=np.uint8) + pointer = owner.ctypes.data + tensor = torch.from_numpy(owner) + result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0) + if result.value != 0: + raise RuntimeError(f"cudaHostRegister failed: {result}") + finalizer = weakref.finalize( + owner, self._unregister, mapping, pointer + ) + finalizer.atexit = False # type: ignore[misc] + # The UVA helper otherwise allocates a private pinned copy. + if not tensor.is_pinned(): + raise RuntimeError( + "CUDA did not recognize the shared Engram registration" + ) + except Exception as exc: + error = f"{stage}: {type(exc).__name__}: {exc}" + + errors: list[str | None] = [None] * group.world_size + # Also fences peer mappings before the leader unlinks the file. + torch.distributed.all_gather_object( + errors, error, group=group.cpu_group + ) + failures = "; ".join( + f"EDP rank {rank}: {error}" + for rank, error in enumerate(errors) + if error is not None + ) + if failures: + raise RuntimeError( + "Engram shared-memory initialization failed: " + failures + ) + assert tensor is not None + return tensor + except Exception: + if finalizer is not None: + finalizer() + tensor = owner = None + if mapping is not None: + mapping.close() + raise + + @staticmethod + def _unregister(mapping: mmap.mmap, pointer: int) -> None: + # Torch storage retains the numpy owner, including through cached UVA views. + # Keep its mmap alive until CUDA has released the registration. + result = torch.cuda.cudart().cudaHostUnregister(pointer) + if result.value != 0: + logger.warning("Engram cudaHostUnregister failed: %s", result) + + def load_weight( + self, param: torch.nn.Parameter, loaded_weight: torch.Tensor + ) -> None: + if self.group.rank_in_group == 0: + _engram_head_shard_weight_loader(param, loaded_weight) + # Read order may differ across ranks. Equal load counts ensure all shared + # weights are ready after the last weight-loader call returns. + torch.distributed.barrier(group=self.group.cpu_group) + + def get_views( + self, weight: torch.Tensor, scales: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if (weight.data_ptr(), scales.data_ptr()) != ( + self.weight.data_ptr(), + self.weight_scale_inv.data_ptr(), + ): + raise RuntimeError("Shared Engram parameter storage must not be replaced") + if self._views is None: + self._views = ( + get_accelerator_view_from_cpu_tensor(self.weight), + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv), + ) + return self._views + + +class ParallelEngramEmbedding(BaseParallelEngramEmbedding): + """Extend TP lookup with DP head sharding or shared, CPU-offloaded TP slices.""" + + _shared_memory: DPSharedEngramStorage | None = None + + def __init__( + self, + num_embeddings: int, + dim: int, + head_sizes: tuple[int, ...], + block_size: int = 32, + cpu_offload: bool = False, + dp_shared_memory: bool = False, + ) -> None: + self.cpu_offload = cpu_offload + self.dp_shared_memory = dp_shared_memory + self.dp_size = get_engram_dp_size() + if dp_shared_memory: + if not cpu_offload: + raise ValueError("dp_shared_memory requires cpu_offload=True") + if self.dp_size <= 1: + raise ValueError( + "dp_shared_memory requires a node-local Engram DP " + f"group with size > 1; effective Engram DP size is {self.dp_size}. " + "Check that the node layout and rank placement allow complete " + "DP replicas to be co-located." + ) + self.dp_size = 1 + if cpu_offload and not is_uva_available(): + raise RuntimeError("Engram CPU offload requires UVA support") + self._views: tuple[torch.Tensor, torch.Tensor] | None = None + self._view_src: tuple[int, int] | None = None + super().__init__(num_embeddings, dim, head_sizes, block_size) + if cpu_offload: + # Constant dummy values avoid randomizing huge CPU lookup tables. + set_weight_attrs(self.weight, {"dummy_weight_value": 1.0}) + # The ue8m0 encoding of scale 1.0 is exponent byte 127. + set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127}) + logger.info( + "Engram table offloaded to pinned host memory: %d rows x %d, " + "%.2f GiB %s", + self.part_num_embeddings, + dim, + self.part_num_embeddings * (dim + dim // block_size) / 1024**3, + "shared across DP replicas" if dp_shared_memory else "per rank", + ) + + def _get_shard_info(self) -> tuple[int, int]: + if self.dp_size == 1: + return super()._get_shard_info() + return self.tp_size * self.dp_size, engram_head_shard_rank() + + def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.dp_shared_memory: + group = get_engram_dp_group() + assert group is not None + storage = DPSharedEngramStorage( + self.part_num_embeddings, self.dim, self.block_size, group + ) + self._shared_memory = storage + self._weight_loader = storage.load_weight + return storage.weight, storage.weight_scale_inv + if not self.cpu_offload: + return super()._allocate_weights() + # Model initialization may be inside a CUDA device context. + return ( + torch.empty( + self.part_num_embeddings, + self.dim, + dtype=torch.float8_e4m3fn, + device="cpu", + pin_memory=True, + ), + torch.empty( + self.part_num_embeddings, + self.dim // self.block_size, + dtype=torch.uint8, + device="cpu", + pin_memory=True, + ), + ) + + def _storage(self) -> tuple[torch.Tensor, torch.Tensor]: + if self._shared_memory is not None: + return self._shared_memory.get_views(self.weight, self.weight_scale_inv) + if not self.cpu_offload: + return super()._storage() + src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr()) + if self._view_src != src: + self._views = ( + get_accelerator_view_from_cpu_tensor(self.weight.data), + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data), + ) + self._view_src = src + assert self._views is not None + return self._views + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + if self.dp_size == 1: + return super().forward(indices) + num_tokens = indices.shape[0] + indices = gather_engram_hashes(indices) + out = torch.empty( + (indices.shape[0], self.part_n_hash_cols, self.dim), + dtype=torch.bfloat16, + device=indices.device, + ) + self.lookup(indices, out) + out = _gather_engram_rows(out, num_tokens) + if self.tp_size > 1: + out = tensor_model_parallel_all_gather(out, dim=1) + return out[:, : self.n_hash_cols] + + +def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor: + """Exchange DP tokens for heads, retaining only this replica's tokens.""" + dp_group = get_engram_dp_group() + assert dp_group is not None + slot, remainder = divmod(staged.shape[0], dp_group.world_size) + assert remainder == 0 and 0 <= num_tokens <= slot + gathered = dp_group.all_gather(staged, dim=0) + local_heads, dim = staged.shape[1:] + rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim)) + _engram_select_rows( + gathered, + rows, + staged.shape[0], + dp_group.rank_in_group * slot, + local_heads * dim, + ) + return rows + + +class Engram(BaseEngram): + """NVIDIA Engram with asynchronous offload and node-local DP lookup.""" + + _prefetch_stream: torch.cuda.Stream | None = None + + def _create_embedding( + self, layout: EngramLayout, layer_hash_index: int + ) -> ParallelEngramEmbedding: + engram_config = get_current_vllm_config().engram_config + assert engram_config is not None + return ParallelEngramEmbedding( + layout.num_embeddings[layer_hash_index], + layout.head_dim, + tuple(size for order in layout.primes[layer_hash_index] for size in order), + cpu_offload=engram_config.cpu_offload, + dp_shared_memory=engram_config.dp_shared_memory, + ) + + def _init_staging(self, max_tokens: int, head_dim: int) -> None: + super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim) + if self.embed_tokens.cpu_offload: + self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device) + + def prepare_embeddings(self, hash_ids: torch.Tensor) -> None: + """Prefetch local shared rows or the DP group's gathered hash IDs.""" + if self._prefetch_stream is None: + return super().prepare_embeddings(hash_ids) + rows = self.staged_rows[: hash_ids.shape[0]] + assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small" + self._start_prefetch(hash_ids, rows, self._prefetch_stream) + + @eager_break_during_capture + def _start_prefetch( + self, hash_ids: torch.Tensor, rows: torch.Tensor, stream: torch.cuda.Stream + ) -> None: + # Eager boundaries let the lookup span piecewise graph segments. + stream.wait_stream(torch.cuda.current_stream()) + # Keep temporary hash storage alive until lookup finishes reading it. + hash_ids.record_stream(stream) + with torch.cuda.stream(stream): + self.embed_tokens.lookup(hash_ids, rows, background=True) + + @eager_break_during_capture + def _finish_prefetch(self, stream: torch.cuda.Stream) -> None: + torch.cuda.current_stream().wait_stream(stream) + + def _ready_rows(self, num_tokens: int) -> torch.Tensor: + if self._prefetch_stream is not None: + self._finish_prefetch(self._prefetch_stream) + if self.embed_tokens.dp_size > 1: + slot = engram_gathered_num_tokens() + staged = self.staged_rows[: slot * self.embed_tokens.dp_size] + return _gather_engram_rows(staged, num_tokens) + return super()._ready_rows(num_tokens) diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index 6d803736ec..0868089da6 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v4_1/nvidia/model.py @@ -12,6 +12,7 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.config.kernel import MEGA_MOE_BACKENDS from vllm.distributed import ( + get_engram_dp_size, get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -74,8 +75,9 @@ from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id -from ..common.engram import Engram, EngramLayout, NgramHashState +from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask +from .engram import Engram, gather_engram_hashes if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -466,6 +468,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): # layer's sliding-window KV cache. Only PP ranks owning an engram # layer need it. self.engram_hash: NgramHashState | None = None + self.engram_dp_shared_memory = bool( + vllm_config.engram_config and vllm_config.engram_config.dp_shared_memory + ) self.engram_swa_prefix: str | None = None if self.engram_layout is not None: local_engram = any( @@ -594,12 +599,21 @@ class DeepseekV4Model(nn.Module, EagleModelMixin): swa_metadata.slot_mapping, swa_metadata.block_table, ) + elif not self.engram_dp_shared_memory and get_engram_dp_size() > 1: + # DP-sharded lookups are collective, so a replica skipping the + # hash still has to reach them. + engram_hashes, engram_mask = self.engram_hash.dummy_hashes(input_ids) + if engram_hashes is not None: # Gather all Engram rows before entering the decoder layers. + # One gather feeds every layer sharing the DP-split table. + gathered_hashes = gather_engram_hashes( + engram_hashes, dp_shared_memory=self.engram_dp_shared_memory + ) for layer in islice(self.layers, self.start_layer, self.end_layer): engram = getattr(layer, "engram", None) if engram is not None: engram.prepare_embeddings( - engram_hashes[:, engram.layer_hash_index] + gathered_hashes[:, engram.layer_hash_index] ) full_num_tokens = positions.shape[0]