@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4c2d277643

[Perf][DSv4.1] Faster Engram host lookups: sorted rows, inline big lookups (#59327)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Juntian Liu committed Sep 30, 2026 at 20:54 UTC 4c2d277643e217344056e1d2c42115d5f005912f
6 files changed +547 -556
tests/distributed/test_engram_dp_shard.py
+24 -5
@@ -39,10 +39,10 @@ from vllm.distributed.parallel_state import (
39 initialize_model_parallel,
40 )
41 from vllm.forward_context import set_forward_context
42 -from vllm.models.deepseek_v41.common.engram import EngramLayout
43 -from vllm.models.deepseek_v41.nvidia import engram as engram_ops
44 -from vllm.models.deepseek_v41.nvidia.engram import (
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,
@@ -686,11 +686,30 @@ def test_engram_dp_shared_memory_runtime_requirements(
686 )
687
688
689 -def test_engram_tables_too_large_for_shm_are_not_shared(monkeypatch):
690 - """A /dev/shm smaller than the tables must fall back instead of failing startup."""
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(
tests/kernels/test_engram.py
+25 -25
@@ -9,13 +9,10 @@ import torch
9
10 from vllm.models.deepseek_v41.common import engram as engram_ops
11 from vllm.models.deepseek_v41.common.engram import (
12 - Engram as CommonEngram,
13 -)
14 -from vllm.models.deepseek_v41.common.engram import (
12 + Engram,
13 NgramHashState,
14 + ParallelEngramEmbedding,
15 )
17 -from vllm.models.deepseek_v41.nvidia import engram as nvidia_engram_ops
18 -from vllm.models.deepseek_v41.nvidia.engram import Engram, ParallelEngramEmbedding
16 from vllm.platforms import current_platform
17
18
@@ -111,7 +108,7 @@ def test_fused_engram_post_wkv_matches_reference(
108 )
109 monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: tp_rank)
110
114 - module = CommonEngram.__new__(CommonEngram)
111 + module = Engram.__new__(Engram)
112 torch.nn.Module.__init__(module)
113 module.dim = dim
114 module.hc_mult = hc_mult
@@ -121,6 +118,7 @@ def test_fused_engram_post_wkv_matches_reference(
118 module.embed_tokens = torch.nn.Identity()
119 # `forward` reads rows staged by `prepare_embeddings`, so inject kv there.
120 module.embed_tokens.tp_size = module.embed_tokens.n_hash_cols = 1
121 + module.embed_tokens.dp_size = 1
122 module.staged_rows = kv.unsqueeze(1)
123 if use_sequence_parallel:
124 padded = torch.nn.functional.pad(kv, (0, 0, 0, (-num_kv_tokens) % tp_size))
@@ -592,9 +590,8 @@ def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch
590 monkeypatch.setattr(
591 engram_ops, "get_tensor_model_parallel_world_size", lambda: tp_size
592 )
595 - monkeypatch.setattr(nvidia_engram_ops, "get_engram_dp_size", lambda: dp_size)
593 + monkeypatch.setattr(engram_ops, "get_engram_dp_size", lambda: dp_size)
594 monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
597 - monkeypatch.setattr(nvidia_engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
595 with pytest.raises(AssertionError, match="ranks without hash heads"):
596 ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads)
597
@@ -602,10 +599,19 @@ def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch
599 @pytest.mark.skipif(
600 not current_platform.is_cuda_alike(), reason="CUDA or ROCm required"
601 )
605 -@pytest.mark.parametrize("cpu_offload", [False, True])
602 +@pytest.mark.parametrize(
603 + "cpu_offload,sort_rows", [(False, False), (True, False), (True, True)]
604 +)
605 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
607 -def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
608 - """Keep complete buckets and reconstruct head order, including TP padding."""
606 +def test_engram_head_shards_reconstruct_checkpoint(
607 + cpu_offload, sort_rows, tp_size, monkeypatch
608 +):
609 + """Keep complete buckets and reconstruct head order, including TP padding,
610 + whether or not the host lookup sorts its rows by table offset."""
611 + if sort_rows:
612 + monkeypatch.setattr(
613 + engram_ops, "_engram_lookup_thresholds", lambda _: (0, None)
614 + )
615 head_sizes = (17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73)
616 num_rows, dim = sum(head_sizes), 64
617 torch.manual_seed(0)
@@ -796,10 +802,8 @@ def test_engram_shared_prefetch_stream_waits_per_layer():
802
803
804 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
799 -@pytest.mark.parametrize(
800 - "backend,cpu_offload", [("nvidia", None), ("nvidia", False), ("common", False)]
801 -)
802 -def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
805 +@pytest.mark.parametrize("cpu_offload", [None, False])
806 +def test_engram_constructor_honors_offload(monkeypatch, cpu_offload):
807 """Default offload uses pinned storage and the given stream; False uses HBM."""
808 from vllm.config import EngramConfig
809
@@ -817,23 +821,18 @@ def test_engram_constructor_honors_offload(monkeypatch, backend, cpu_offload):
821 else EngramConfig(cpu_offload=cpu_offload),
822 scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
823 )
820 - offloaded = backend == "nvidia" and vllm_config.engram_config.cpu_offload
821 - if backend == "common":
822 - vllm_config.engram_config = None
824 + offloaded = vllm_config.engram_config.cpu_offload
825 monkeypatch.setattr(engram_ops, "get_current_vllm_config", lambda: vllm_config)
824 - monkeypatch.setattr(
825 - nvidia_engram_ops, "get_current_vllm_config", lambda: vllm_config
826 - )
826 monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_world_size", lambda: 1)
827 monkeypatch.setattr(engram_ops, "get_tensor_model_parallel_rank", lambda: 0)
828 monkeypatch.setattr(
829 engram_ops, "ColumnParallelLinear", lambda *a, **k: torch.nn.Identity()
830 )
831 stream = torch.cuda.Stream()
833 - kwargs = {"prefetch_stream": stream} if backend == "nvidia" else {}
832 with torch.device("cuda"):
835 - cls = Engram if backend == "nvidia" else CommonEngram
836 - module = cls(config, None, layout, 0, False, "engram", **kwargs)
833 + module = Engram(
834 + config, None, layout, 0, False, "engram", prefetch_stream=stream
835 + )
836 layer = module.embed_tokens
837 if offloaded:
838 assert layer.weight.device.type == "cpu" and layer.weight.is_pinned()
@@ -897,7 +896,8 @@ def test_engram_prefetch_detects_missing_dependency(monkeypatch, missing_depende
896 "The allocator reuse negative control requires the native allocator"
897 )
898
900 - def start(self, hash_ids, rows, stream):
899 + def start(self, hash_ids, rows):
900 + stream = self._prefetch_stream
901 if missing_dependency != "producer_wait":
902 stream.wait_stream(torch.cuda.current_stream())
903 if missing_dependency != "record_stream":
vllm/models/deepseek_v41/amd/model.py
+1 -5
@@ -65,13 +65,9 @@ from vllm.utils.math_utils import cdiv
65 from vllm.v1.attention.backends.registry import AttentionBackendEnum
66 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
67
68 -from ..common.engram import EngramLayout, NgramHashState
68 +from ..common.engram import Engram, EngramLayout, NgramHashState
69 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
70
71 -# Engram host offload and its prefetch stream are neither ROCm- nor
72 -# NVIDIA-specific, so they are imported rather than duplicated.
73 -from ..nvidia.engram import Engram
74 -
71 if typing.TYPE_CHECKING:
72 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
73
vllm/models/deepseek_v41/common/engram.py
+490 -22
@@ -34,23 +34,41 @@ chunk-by-chunk while an n-gram at position ``p`` needs the token ids at
34 cache for the rest.
35 """
36
37 +import ctypes
38 +import mmap
39 +import os
40 +import tempfile
41 import weakref
42 +from contextlib import ExitStack
43
44 import numpy as np
45 import torch
46 from torch import nn
47
48 +from vllm.compilation.breakable_cudagraph import eager_break_during_capture
49 from vllm.config import VllmConfig, get_current_vllm_config
50 from vllm.distributed import (
51 + get_dp_group,
52 + get_engram_dp_group,
53 + get_engram_dp_size,
54 get_tensor_model_parallel_rank,
55 get_tensor_model_parallel_world_size,
56 tensor_model_parallel_all_gather,
57 )
58 +from vllm.distributed.device_communicators.shm_broadcast import (
59 + SHM_PATH,
60 + check_shm_free_space,
61 +)
62 +from vllm.distributed.parallel_state import GroupCoordinator
63 +from vllm.forward_context import get_forward_context, is_forward_context_available
64 from vllm.logger import init_logger
65 from vllm.model_executor.layers.linear import ColumnParallelLinear
66 from vllm.model_executor.layers.quantization import QuantizationConfig
67 from vllm.model_executor.utils import set_weight_attrs
68 +from vllm.platforms import current_platform
69 from vllm.triton_utils import tl, triton
70 +from vllm.utils.platform_utils import is_uva_available
71 +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
72
73 logger = init_logger(__name__)
74
@@ -58,6 +76,18 @@ logger = init_logger(__name__)
76 DEAD_ID = -1
77
78
79 +def _engram_lookup_thresholds(device: torch.device) -> tuple[int | None, int | None]:
80 + """Host lookup thresholds on `device`: the rows from which to sort, and the
81 + tokens (max over DP ranks) from which to run inline. None disables that path."""
82 + if current_platform.is_cuda() and current_platform.is_device_capability_family(
83 + 100, device.index or 0
84 + ):
85 + # Measured on GB200.
86 + return 36864, 4096
87 + # TODO: Add hand-tuned thresholds for other GPU families.
88 + return None, None
89 +
90 +
91 def _is_prime(n: int) -> bool:
92 """Deterministic Miller-Rabin for n < 2**32 (avoids a sympy import)."""
93 if n < 2:
@@ -583,12 +613,14 @@ def _engram_head_shard_weight_loader(
613 "ids_stride_t",
614 "ids_stride_h",
615 "GRID",
616 + "SORTED",
617 ]
618 )
619 def _engram_lookup_kernel(
620 weight,
621 scales,
622 ids,
623 + sorted_dst,
624 out,
625 vocab_start,
626 vocab_end,
@@ -602,27 +634,37 @@ def _engram_lookup_kernel(
634 QUANT_BLOCK: tl.constexpr,
635 BLOCK_R: tl.constexpr,
636 GRID,
637 + SORTED,
638 ):
639 """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
640
641 Only this rank's heads are read; padded heads write zeros for all-gather.
609 - `weight`/`scales` may address pinned host memory through UVA.
642 + `weight`/`scales` may address pinned host memory through UVA. If SORTED,
643 + `ids` holds rows relative to vocab_start in table order and `sorted_dst`
644 + their output rows.
645 """
646 cols = tl.arange(0, DIM)
647 scale_cols = cols // QUANT_BLOCK
648 for base in tl.range(tl.program_id(0) * BLOCK_R, num_rows, GRID * BLOCK_R):
649 rows = base + tl.arange(0, BLOCK_R)
650 valid = rows < num_rows
616 - head = HEAD_START + rows % LOCAL_HEADS
617 - token = (rows // LOCAL_HEADS).to(tl.int64)
618 - index = tl.load(
619 - ids + token * ids_stride_t + head * ids_stride_h,
620 - mask=valid & (head < TOTAL_HEADS),
621 - other=-1,
622 - ).to(tl.int64)
623 - owned = valid & (head < TOTAL_HEADS)
624 - owned &= (index >= vocab_start) & (index < vocab_end)
625 - local = tl.where(owned, index - vocab_start, 0)
651 + if SORTED:
652 + local = tl.load(ids + rows, mask=valid, other=-1).to(tl.int64)
653 + dst = tl.load(sorted_dst + rows, mask=valid, other=0).to(tl.int32)
654 + owned = (local >= 0) & (local < vocab_end - vocab_start)
655 + else:
656 + head = HEAD_START + rows % LOCAL_HEADS
657 + token = (rows // LOCAL_HEADS).to(tl.int64)
658 + index = tl.load(
659 + ids + token * ids_stride_t + head * ids_stride_h,
660 + mask=valid & (head < TOTAL_HEADS),
661 + other=-1,
662 + ).to(tl.int64)
663 + dst = rows
664 + owned = valid & (head < TOTAL_HEADS)
665 + owned &= (index >= vocab_start) & (index < vocab_end)
666 + local = index - vocab_start
667 + local = tl.where(owned, local, 0)
668 values = tl.load(
669 weight + local[:, None] * DIM + cols[None, :],
670 mask=owned[:, None],
@@ -636,16 +678,240 @@ def _engram_lookup_kernel(
678 # ue8m0 is a power of two, so its byte *is* the fp32 exponent field.
679 scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
680 tl.store(
639 - out + rows[:, None] * DIM + cols[None, :],
681 + out + dst[:, None] * DIM + cols[None, :],
682 (values.to(tl.float32) * scale).to(tl.bfloat16),
683 mask=valid[:, None],
684 )
685
686
687 +def engram_head_shard_rank() -> int:
688 + """This rank's slot among the hash-head shards of one engram table.
689 +
690 + TP-major, so the shards a DP gather brings in are contiguous heads and
691 + the following TP gather completes the head order.
692 + """
693 + dp_group = get_engram_dp_group()
694 + dp_size = dp_group.world_size if dp_group is not None else 1
695 + dp_rank = dp_group.rank_in_group if dp_group is not None else 0
696 + return get_tensor_model_parallel_rank() * dp_size + dp_rank
697 +
698 +
699 +def engram_gathered_num_tokens() -> int:
700 + """Per-replica token slot for the node-local Engram DP group."""
701 + dp_metadata = get_forward_context().dp_metadata
702 + if dp_metadata is None:
703 + raise RuntimeError("a DP-shared engram table needs DP token metadata")
704 + group = get_engram_dp_group()
705 + assert group is not None
706 + # Engram groups are contiguous slices of the full DP group.
707 + start = get_dp_group().rank_in_group - group.rank_in_group
708 + return int(
709 + dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max()
710 + )
711 +
712 +
713 +def gather_engram_hashes(
714 + hash_ids: torch.Tensor, *, dp_shared_memory: bool = False
715 +) -> torch.Tensor:
716 + """Collect the n-gram ids of every DP replica sharing one table.
717 +
718 + Replicas are padded to a common token slot, so the gathered shape is
719 + static under CUDA graph capture (where DP already pads alike).
720 + """
721 + dp_group = get_engram_dp_group()
722 + if dp_group is None or dp_shared_memory:
723 + return hash_ids
724 + slot = engram_gathered_num_tokens()
725 + if hash_ids.shape[0] > slot:
726 + raise ValueError("Engram token count exceeds the DP token slot")
727 + if hash_ids.shape[0] < slot:
728 + pad = hash_ids.new_full(
729 + (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID
730 + )
731 + hash_ids = torch.cat((hash_ids, pad))
732 + return dp_group.all_gather(hash_ids, dim=0)
733 +
734 +
735 +class DPSharedEngramStorage:
736 + """Registered host weights shared by a node-local DP group with one writer."""
737 +
738 + def __init__(
739 + self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator
740 + ) -> None:
741 + self.group = group
742 + weight_bytes = num_rows * dim
743 + storage = self._allocate(weight_bytes + weight_bytes // block_size)
744 + self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim)
745 + self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size)
746 + self._views: tuple[torch.Tensor, torch.Tensor] | None = None
747 +
748 + def _allocate(self, num_bytes: int) -> torch.Tensor:
749 + """Map and register one physical allocation across a node-local DP group."""
750 + group = self.group
751 + with ExitStack() as stack:
752 + path = error = None
753 + if group.rank_in_group == 0:
754 + try:
755 + check_shm_free_space(num_bytes)
756 + backing_file = stack.enter_context(
757 + tempfile.NamedTemporaryFile(prefix="vllm_engram_", dir=SHM_PATH)
758 + )
759 + backing_file.truncate(num_bytes)
760 + path = backing_file.name
761 + except Exception as exc:
762 + error = f"{type(exc).__name__}: {exc}"
763 + path, error = group.broadcast_object((path, error))
764 + if error is not None:
765 + raise RuntimeError(
766 + "Engram shared-memory creation failed on EDP rank 0: " + error
767 + )
768 +
769 + mapping = owner = tensor = None
770 + try:
771 + with open(path, "r+b") as file:
772 + mapping = mmap.mmap(file.fileno(), num_bytes, flags=mmap.MAP_SHARED)
773 + owner = np.frombuffer(mapping, dtype=np.uint8)
774 + pointer = owner.ctypes.data
775 + tensor = torch.from_numpy(owner)
776 + result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
777 + if result.value != 0:
778 + raise RuntimeError(f"cudaHostRegister failed: {result}")
779 + finalizer = weakref.finalize(owner, self._unregister, mapping, pointer)
780 + finalizer.atexit = False # type: ignore[misc]
781 + # The UVA helper otherwise allocates a private pinned copy.
782 + if not tensor.is_pinned():
783 + raise RuntimeError(
784 + "cudaHostRegister did not pin the shared Engram mapping"
785 + )
786 + except Exception as exc:
787 + error = f"{type(exc).__name__}: {exc}"
788 +
789 + # Also fences peer mappings before the leader unlinks the file.
790 + errors: list[str | None] = [None] * group.world_size
791 + torch.distributed.all_gather_object(errors, error, group=group.cpu_group)
792 + failures = "; ".join(
793 + f"EDP rank {rank}: {error}"
794 + for rank, error in enumerate(errors)
795 + if error is not None
796 + )
797 + if failures:
798 + # Dropping the owner runs the finalizer; the mapping then unmaps.
799 + del tensor, owner, mapping
800 + raise RuntimeError(
801 + "Engram shared-memory initialization failed: " + failures
802 + )
803 + assert tensor is not None
804 + return tensor
805 +
806 + @staticmethod
807 + def _unregister(mapping: mmap.mmap, pointer: int) -> None:
808 + # Torch storage retains the numpy owner, including through cached UVA views.
809 + # Keep its mmap alive until CUDA has released the registration.
810 + result = torch.cuda.cudart().cudaHostUnregister(pointer)
811 + if result.value != 0:
812 + logger.warning("Engram cudaHostUnregister failed: %s", result)
813 +
814 + def load_weight(
815 + self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
816 + ) -> None:
817 + if self.group.rank_in_group == 0:
818 + _engram_head_shard_weight_loader(param, loaded_weight)
819 + # Read order may differ across ranks. Equal load counts ensure all shared
820 + # weights are ready after the last weight-loader call returns.
821 + torch.distributed.barrier(group=self.group.cpu_group)
822 +
823 + def get_views(
824 + self, weight: torch.Tensor, scales: torch.Tensor
825 + ) -> tuple[torch.Tensor, torch.Tensor]:
826 + if (weight.data_ptr(), scales.data_ptr()) != (
827 + self.weight.data_ptr(),
828 + self.weight_scale_inv.data_ptr(),
829 + ):
830 + raise RuntimeError("Shared Engram parameter storage must not be replaced")
831 + if self._views is None:
832 + self._views = (
833 + get_accelerator_view_from_cpu_tensor(self.weight),
834 + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv),
835 + )
836 + return self._views
837 +
838 +
839 +def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
840 + """Whether co-located DP replicas exist and /dev/shm can hold the full tables."""
841 + if get_engram_dp_size() == 1:
842 + logger.warning_once(
843 + "Engram DP replicas are not co-located on one node; "
844 + "storing the offloaded tables per rank instead of sharing them."
845 + )
846 + return False
847 + num_bytes = sum(layout.num_embeddings) * (
848 + layout.head_dim + layout.head_dim // block_size
849 + )
850 + group = get_engram_dp_group()
851 + assert group is not None
852 + # Only the leader allocates, so every replica follows its check.
853 + error = None
854 + if group.rank_in_group == 0:
855 + if not os.path.isdir(SHM_PATH):
856 + error = f"{SHM_PATH} is not mounted"
857 + else:
858 + try:
859 + check_shm_free_space(num_bytes, allocation_name="Engram tables")
860 + except Exception as exc:
861 + error = f"{type(exc).__name__}: {exc}"
862 + error = group.broadcast_object(error)
863 + if error is not None:
864 + logger.warning_once(
865 + "Sharding the offloaded Engram tables across DP replicas: %s", error
866 + )
867 + return error is None
868 +
869 +
870 +def _allocate_huge_page_storage(num_bytes: int) -> torch.Tensor | None:
871 + """Register prefaulted huge pages, or return None for pinned-memory fallback."""
872 + try:
873 + mapping = mmap.mmap(-1, num_bytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
874 + mapping.madvise(mmap.MADV_HUGEPAGE)
875 + except OSError as exc:
876 + logger.warning(
877 + "Engram huge-page allocation failed (%s); using pinned memory.", exc
878 + )
879 + return None
880 +
881 + owner = np.frombuffer(mapping, dtype=np.uint8)
882 + # Fault the pages in before CUDA pins them, so they can be huge.
883 + owner[:: mmap.PAGESIZE] = 0
884 + tensor = torch.from_numpy(owner)
885 + pointer = tensor.data_ptr()
886 + result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
887 + if result.value != 0:
888 + logger.warning(
889 + "Engram cudaHostRegister failed (%s); using pinned memory.", result
890 + )
891 + return None
892 + # Tensor views retain owner; its finalizer retains mapping until unregister.
893 + finalizer = weakref.finalize(
894 + owner, DPSharedEngramStorage._unregister, mapping, pointer
895 + )
896 + finalizer.atexit = False # type: ignore[misc]
897 + if not tensor.is_pinned():
898 + logger.warning(
899 + "CUDA did not recognize Engram registration; using pinned memory."
900 + )
901 + return None
902 + return tensor
903 +
904 +
905 class ParallelEngramEmbedding(nn.Module):
646 - """TP-sharded hash heads with FP8 rows and per-block E8M0 scales."""
906 + """Hash heads with FP8 rows and per-block E8M0 scales.
907 +
908 + Heads are TP-sharded, and additionally DP-sharded across a node-local
909 + group unless the table lives in (DP-shared) pinned host memory.
910 + """
911
912 _weight_loader = staticmethod(_engram_head_shard_weight_loader)
913 + _shared_memory: DPSharedEngramStorage | None = None
914 + _packed: torch.Tensor | None = None
915
916 def __init__(
917 self,
@@ -653,10 +919,32 @@ class ParallelEngramEmbedding(nn.Module):
919 dim: int,
920 head_sizes: tuple[int, ...],
921 block_size: int = 32,
922 + cpu_offload: bool = False,
923 + dp_shared_memory: bool = False,
924 + use_thp: bool = False,
925 ) -> None:
926 super().__init__()
927 assert head_sizes and all(size > 0 for size in head_sizes)
928 assert sum(head_sizes) <= num_embeddings
929 + self.cpu_offload = cpu_offload
930 + self.dp_shared_memory = dp_shared_memory
931 + self.use_thp = use_thp
932 + self.dp_size = get_engram_dp_size()
933 + if dp_shared_memory:
934 + if not cpu_offload:
935 + raise ValueError("dp_shared_memory requires cpu_offload=True")
936 + if self.dp_size <= 1:
937 + raise ValueError(
938 + "dp_shared_memory requires a node-local Engram DP "
939 + f"group with size > 1; effective Engram DP size is {self.dp_size}. "
940 + "Check that the node layout and rank placement allow complete "
941 + "DP replicas to be co-located."
942 + )
943 + self.dp_size = 1
944 + if cpu_offload and not is_uva_available():
945 + raise RuntimeError("Engram CPU offload requires UVA support")
946 + self._views: tuple[torch.Tensor, torch.Tensor] | None = None
947 + self._view_src: tuple[int, int] | None = None
948 self.num_embeddings = num_embeddings
949 self.dim = dim
950 self.block_size = block_size
@@ -688,20 +976,99 @@ class ParallelEngramEmbedding(nn.Module):
976 "engram_vocab_start": self.vocab_start_idx,
977 },
978 )
979 + if cpu_offload:
980 + # Constant dummy values avoid randomizing huge CPU lookup tables.
981 + set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
982 + # The ue8m0 encoding of scale 1.0 is exponent byte 127.
983 + set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
984 + logger.info(
985 + "Engram table offloaded to pinned host memory: %d rows x %d, "
986 + "%.2f GiB %s",
987 + self.part_num_embeddings,
988 + dim,
989 + self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
990 + "shared across DP replicas" if dp_shared_memory else "per rank",
991 + )
992
993 def _get_shard_info(self) -> tuple[int, int]:
693 - return self.tp_size, get_tensor_model_parallel_rank()
994 + if self.dp_size == 1:
995 + return self.tp_size, get_tensor_model_parallel_rank()
996 + return self.tp_size * self.dp_size, engram_head_shard_rank()
997
998 def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
999 + if self.dp_shared_memory:
1000 + group = get_engram_dp_group()
1001 + assert group is not None
1002 + storage = DPSharedEngramStorage(
1003 + self.part_num_embeddings, self.dim, self.block_size, group
1004 + )
1005 + self._shared_memory = storage
1006 + self._weight_loader = storage.load_weight
1007 + return storage.weight, storage.weight_scale_inv
1008 + scale_dim = self.dim // self.block_size
1009 + if not self.cpu_offload:
1010 + return (
1011 + torch.empty(
1012 + self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn
1013 + ),
1014 + torch.empty(self.part_num_embeddings, scale_dim, dtype=torch.uint8),
1015 + )
1016 + if self.use_thp:
1017 + weight_bytes = self.part_num_embeddings * self.dim
1018 + packed = _allocate_huge_page_storage(
1019 + weight_bytes + weight_bytes // self.block_size
1020 + )
1021 + if packed is not None:
1022 + self._packed = packed
1023 + return (
1024 + packed[:weight_bytes].view(torch.float8_e4m3fn).view(-1, self.dim),
1025 + packed[weight_bytes:].view(-1, scale_dim),
1026 + )
1027 + # Model initialization may be inside a CUDA device context.
1028 return (
697 - torch.empty(self.part_num_embeddings, self.dim, dtype=torch.float8_e4m3fn),
1029 torch.empty(
699 - self.part_num_embeddings, self.dim // self.block_size, dtype=torch.uint8
1030 + self.part_num_embeddings,
1031 + self.dim,
1032 + dtype=torch.float8_e4m3fn,
1033 + device="cpu",
1034 + pin_memory=True,
1035 + ),
1036 + torch.empty(
1037 + self.part_num_embeddings,
1038 + scale_dim,
1039 + dtype=torch.uint8,
1040 + device="cpu",
1041 + pin_memory=True,
1042 ),
1043 )
1044
1045 + def collapse_huge_pages(self) -> None:
1046 + """Best-effort MADV_COLLAPSE (Linux >= 6.1) of pages that faulted small."""
1047 + if self._packed is None:
1048 + return
1049 + addr, num_bytes = self._packed.data_ptr(), self._packed.nbytes
1050 + libc = ctypes.CDLL(None, use_errno=True)
1051 + libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int)
1052 + if libc.madvise(addr, num_bytes, 25) != 0:
1053 + logger.warning(
1054 + "Engram MADV_COLLAPSE failed; keeping existing pages: %s",
1055 + os.strerror(ctypes.get_errno()),
1056 + )
1057 +
1058 def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
704 - return self.weight.data, self.weight_scale_inv.data
1059 + if self._shared_memory is not None:
1060 + return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
1061 + if not self.cpu_offload:
1062 + return self.weight.data, self.weight_scale_inv.data
1063 + src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
1064 + if self._view_src != src:
1065 + self._views = (
1066 + get_accelerator_view_from_cpu_tensor(self.weight.data),
1067 + get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
1068 + )
1069 + self._view_src = src
1070 + assert self._views is not None
1071 + return self._views
1072
1073 def lookup(
1074 self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
@@ -714,6 +1081,24 @@ class ParallelEngramEmbedding(nn.Module):
1081 if not rows:
1082 return
1083 weight, scales = self._storage()
1084 + ids_stride_t, ids_stride_h = indices.stride()
1085 + sort_min_rows, _ = _engram_lookup_thresholds(indices.device)
1086 + sort = (
1087 + self.cpu_offload
1088 + and self._packed is None
1089 + and sort_min_rows is not None
1090 + and rows >= sort_min_rows
1091 + )
1092 + dst = indices
1093 + if sort:
1094 + # Small-page host tables miss the TLB per row; in table order, rows on
1095 + # one page share it.
1096 + head_end = self.head_start + self.part_n_hash_cols
1097 + local = indices[:, self.head_start : head_end] - self.vocab_start_idx
1098 + if (pad := self.part_n_hash_cols - local.shape[1]) > 0:
1099 + local = nn.functional.pad(local, (0, pad), value=-1)
1100 + indices, dst = local.flatten().sort()
1101 + dst = dst.int()
1102 # The table dwarfs TLB reach, so a persistent grid near the SM count
1103 # beats one program per row; halve it to leave SMs for the main stream.
1104 tiles = triton.cdiv(rows, 16)
@@ -722,12 +1107,13 @@ class ParallelEngramEmbedding(nn.Module):
1107 weight,
1108 scales,
1109 indices,
1110 + dst,
1111 out,
1112 self.vocab_start_idx,
1113 self.vocab_end_idx,
1114 rows,
729 - indices.stride(0),
730 - indices.stride(1),
1115 + ids_stride_t,
1116 + ids_stride_h,
1117 HEAD_START=self.head_start,
1118 LOCAL_HEADS=self.part_n_hash_cols,
1119 TOTAL_HEADS=self.n_hash_cols,
@@ -735,17 +1121,23 @@ class ParallelEngramEmbedding(nn.Module):
1121 QUANT_BLOCK=self.block_size,
1122 BLOCK_R=16,
1123 GRID=grid,
1124 + SORTED=sort,
1125 )
1126
1127 def forward(self, indices: torch.Tensor) -> torch.Tensor:
1128 """indices: [num_tokens, n_hash_cols] -> [num_tokens, n_hash_cols, dim]
1129 bf16, gathered from all shards for this replica's tokens."""
1130 + num_tokens = indices.shape[0]
1131 + if self.dp_size > 1:
1132 + indices = gather_engram_hashes(indices)
1133 out = torch.empty(
1134 (indices.shape[0], self.part_n_hash_cols, self.dim),
1135 dtype=torch.bfloat16,
1136 device=indices.device,
1137 )
1138 self.lookup(indices, out)
1139 + if self.dp_size > 1:
1140 + out = _gather_engram_rows(out, num_tokens)
1141 if self.tp_size > 1:
1142 out = tensor_model_parallel_all_gather(out, dim=1)
1143 return out[:, : self.n_hash_cols]
@@ -891,6 +1283,25 @@ def _engram_select_rows(
1283 )
1284
1285
1286 +def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor:
1287 + """Exchange DP tokens for heads, retaining only this replica's tokens."""
1288 + dp_group = get_engram_dp_group()
1289 + assert dp_group is not None
1290 + slot, remainder = divmod(staged.shape[0], dp_group.world_size)
1291 + assert remainder == 0 and 0 <= num_tokens <= slot
1292 + gathered = dp_group.all_gather(staged, dim=0)
1293 + local_heads, dim = staged.shape[1:]
1294 + rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim))
1295 + _engram_select_rows(
1296 + gathered,
1297 + rows,
1298 + staged.shape[0],
1299 + dp_group.rank_in_group * slot,
1300 + local_heads * dim,
1301 + )
1302 + return rows
1303 +
1304 +
1305 class Engram(nn.Module):
1306 """Writes an n-gram lookup into the residual stream, gated by how well it
1307 matches that stream.
@@ -901,6 +1312,9 @@ class Engram(nn.Module):
1312 training kernel).
1313 """
1314
1315 + _prefetch_stream: torch.cuda.Stream | None = None
1316 + _prefetch_done: torch.cuda.Event | None = None
1317 +
1318 def __init__(
1319 self,
1320 config,
@@ -909,8 +1323,12 @@ class Engram(nn.Module):
1323 layer_hash_index: int,
1324 use_sequence_parallel: bool,
1325 prefix: str,
1326 + prefetch_stream: torch.cuda.Stream | None = None,
1327 ) -> None:
1328 super().__init__()
1329 + # Layers sharing one stream serialize their offloaded lookups, so they
1330 + # take turns instead of jointly starving decoder compute of SMs.
1331 + self._prefetch_stream = prefetch_stream
1332 self.layer_hash_index = layer_hash_index
1333 self.dim = config.hidden_size
1334 self.hc_mult = config.hc_mult
@@ -950,28 +1368,78 @@ class Engram(nn.Module):
1368 def _create_embedding(
1369 self, layout: EngramLayout, layer_hash_index: int
1370 ) -> ParallelEngramEmbedding:
1371 + engram_config = get_current_vllm_config().engram_config
1372 + assert engram_config is not None
1373 return ParallelEngramEmbedding(
1374 layout.num_embeddings[layer_hash_index],
1375 layout.head_dim,
1376 tuple(size for order in layout.primes[layer_hash_index] for size in order),
1377 + cpu_offload=engram_config.cpu_offload,
1378 + dp_shared_memory=bool(engram_config.dp_shared_memory),
1379 + use_thp=engram_config.use_thp,
1380 )
1381
1382 def _init_staging(self, max_tokens: int, head_dim: int) -> None:
1383 # Persistent storage keeps lookup addresses stable across graph replays.
1384 self.staged_rows = torch.empty(
962 - max_tokens,
1385 + max_tokens * self.embed_tokens.dp_size,
1386 self.embed_tokens.part_n_hash_cols,
1387 head_dim,
1388 dtype=torch.bfloat16,
1389 )
1390 + if not self.embed_tokens.cpu_offload:
1391 + self._prefetch_stream = None
1392 + return
1393 + if self._prefetch_stream is None:
1394 + self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
1395 + self._prefetch_done = torch.cuda.Event()
1396
1397 def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
969 - """Look up rows before the decoder layers consume them."""
1398 + """Look up, or prefetch from host memory, rows before the decoder layers
1399 + consume them; DP-sharded tables take the group's gathered hash IDs."""
1400 rows = self.staged_rows[: hash_ids.shape[0]]
1401 assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
972 - self.embed_tokens.lookup(hash_ids, rows)
1402 + if self._prefetch_stream is None:
1403 + self.embed_tokens.lookup(hash_ids, rows)
1404 + else:
1405 + self._start_prefetch(hash_ids, rows)
1406 +
1407 + @eager_break_during_capture
1408 + def _start_prefetch(self, hash_ids: torch.Tensor, rows: torch.Tensor) -> None:
1409 + # Eager boundaries let the lookup span piecewise graph segments, and
1410 + # decide per replay (on the replay stream) whether it runs inline.
1411 + num_tokens = hash_ids.shape[0]
1412 + # All DP ranks must agree, or EP collectives wait on the slowest one.
1413 + if is_forward_context_available() and (dp := get_forward_context().dp_metadata):
1414 + num_tokens = int(dp.num_tokens_across_dp_cpu.max())
1415 + current = stream = torch.cuda.current_stream()
1416 + # Big lookups stall persistent main-stream kernels anyway: run them inline.
1417 + _, inline_min_tokens = _engram_lookup_thresholds(hash_ids.device)
1418 + if inline_min_tokens is not None and num_tokens >= inline_min_tokens:
1419 + self.embed_tokens.lookup(hash_ids, rows)
1420 + else:
1421 + stream = self._prefetch_stream
1422 + assert stream is not None
1423 + stream.wait_stream(current)
1424 + # Keep temporary hash storage alive until lookup finishes reading it.
1425 + hash_ids.record_stream(stream)
1426 + with torch.cuda.stream(stream):
1427 + self.embed_tokens.lookup(hash_ids, rows, background=True)
1428 + assert self._prefetch_done is not None
1429 + self._prefetch_done.record(stream)
1430 +
1431 + @eager_break_during_capture
1432 + def _finish_prefetch(self, event: torch.cuda.Event) -> None:
1433 + # Wait for this layer's lookup only, not for later layers on the stream.
1434 + torch.cuda.current_stream().wait_event(event)
1435
1436 def _ready_rows(self, num_tokens: int) -> torch.Tensor:
1437 + if self._prefetch_done is not None:
1438 + self._finish_prefetch(self._prefetch_done)
1439 + if self.embed_tokens.dp_size > 1:
1440 + slot = engram_gathered_num_tokens()
1441 + staged = self.staged_rows[: slot * self.embed_tokens.dp_size]
1442 + return _gather_engram_rows(staged, num_tokens)
1443 return self.staged_rows[:num_tokens]
1444
1445 def embed(self, hash_ids: torch.Tensor) -> torch.Tensor:
vllm/models/deepseek_v41/nvidia/engram.py deleted
-497
@@ -1,497 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""NVIDIA Engram DP sharding, shared host storage, and asynchronous prefetch."""
4 -
5 -import ctypes
6 -import mmap
7 -import os
8 -import shutil
9 -import tempfile
10 -import weakref
11 -from contextlib import ExitStack
12 -
13 -import numpy as np
14 -import torch
15 -
16 -from vllm.compilation.breakable_cudagraph import eager_break_during_capture
17 -from vllm.config import get_current_vllm_config
18 -from vllm.distributed import (
19 - get_dp_group,
20 - get_engram_dp_group,
21 - get_engram_dp_size,
22 - get_tensor_model_parallel_rank,
23 - tensor_model_parallel_all_gather,
24 -)
25 -from vllm.distributed.device_communicators.shm_broadcast import (
26 - SHM_PATH,
27 - check_shm_free_space,
28 -)
29 -from vllm.distributed.parallel_state import GroupCoordinator
30 -from vllm.forward_context import get_forward_context
31 -from vllm.logger import init_logger
32 -from vllm.model_executor.utils import set_weight_attrs
33 -from vllm.models.deepseek_v41.common.engram import (
34 - DEAD_ID,
35 - EngramLayout,
36 - _engram_head_shard_weight_loader,
37 - _engram_select_rows,
38 -)
39 -from vllm.models.deepseek_v41.common.engram import (
40 - Engram as BaseEngram,
41 -)
42 -from vllm.models.deepseek_v41.common.engram import (
43 - ParallelEngramEmbedding as BaseParallelEngramEmbedding,
44 -)
45 -from vllm.utils.platform_utils import is_uva_available
46 -from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
47 -
48 -logger = init_logger(__name__)
49 -
50 -
51 -def _allocate_huge_page_storage(num_bytes: int) -> torch.Tensor | None:
52 - """Register prefaulted huge pages, or return None for pinned-memory fallback."""
53 - try:
54 - mapping = mmap.mmap(-1, num_bytes, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS)
55 - mapping.madvise(mmap.MADV_HUGEPAGE)
56 - except OSError as exc:
57 - logger.warning(
58 - "Engram huge-page allocation failed (%s); using pinned memory.", exc
59 - )
60 - return None
61 -
62 - owner = np.frombuffer(mapping, dtype=np.uint8)
63 - # Fault the pages in before CUDA pins them, so they can be huge.
64 - owner[:: mmap.PAGESIZE] = 0
65 - tensor = torch.from_numpy(owner)
66 - pointer = tensor.data_ptr()
67 - result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
68 - if result.value != 0:
69 - logger.warning(
70 - "Engram cudaHostRegister failed (%s); using pinned memory.", result
71 - )
72 - return None
73 - # Tensor views retain owner; its finalizer retains mapping until unregister.
74 - finalizer = weakref.finalize(
75 - owner, DPSharedEngramStorage._unregister, mapping, pointer
76 - )
77 - finalizer.atexit = False # type: ignore[misc]
78 - if not tensor.is_pinned():
79 - logger.warning(
80 - "CUDA did not recognize Engram registration; using pinned memory."
81 - )
82 - return None
83 - return tensor
84 -
85 -
86 -def engram_head_shard_rank() -> int:
87 - """This rank's slot among the hash-head shards of one engram table.
88 -
89 - TP-major, so the shards a DP gather brings in are contiguous heads and
90 - the following TP gather completes the head order.
91 - """
92 - dp_group = get_engram_dp_group()
93 - dp_size = dp_group.world_size if dp_group is not None else 1
94 - dp_rank = dp_group.rank_in_group if dp_group is not None else 0
95 - return get_tensor_model_parallel_rank() * dp_size + dp_rank
96 -
97 -
98 -def engram_gathered_num_tokens() -> int:
99 - """Per-replica token slot for the node-local Engram DP group."""
100 - dp_metadata = get_forward_context().dp_metadata
101 - if dp_metadata is None:
102 - raise RuntimeError("a DP-shared engram table needs DP token metadata")
103 - group = get_engram_dp_group()
104 - assert group is not None
105 - # Engram groups are contiguous slices of the full DP group.
106 - start = get_dp_group().rank_in_group - group.rank_in_group
107 - return int(
108 - dp_metadata.num_tokens_across_dp_cpu[start : start + group.world_size].max()
109 - )
110 -
111 -
112 -def gather_engram_hashes(
113 - hash_ids: torch.Tensor, *, dp_shared_memory: bool = False
114 -) -> torch.Tensor:
115 - """Collect the n-gram ids of every DP replica sharing one table.
116 -
117 - Replicas are padded to a common token slot, so the gathered shape is
118 - static under CUDA graph capture (where DP already pads alike).
119 - """
120 - dp_group = get_engram_dp_group()
121 - if dp_group is None or dp_shared_memory:
122 - return hash_ids
123 - slot = engram_gathered_num_tokens()
124 - if hash_ids.shape[0] > slot:
125 - raise ValueError("Engram token count exceeds the DP token slot")
126 - if hash_ids.shape[0] < slot:
127 - pad = hash_ids.new_full(
128 - (slot - hash_ids.shape[0], *hash_ids.shape[1:]), DEAD_ID
129 - )
130 - hash_ids = torch.cat((hash_ids, pad))
131 - return dp_group.all_gather(hash_ids, dim=0)
132 -
133 -
134 -def can_share_engram_tables(layout: EngramLayout, block_size: int = 32) -> bool:
135 - """Whether co-located DP replicas exist and /dev/shm can hold the full tables."""
136 - if not os.path.isdir(SHM_PATH):
137 - logger.warning_once(
138 - "%s is not mounted; storing the offloaded Engram tables per rank "
139 - "instead of sharing them across DP replicas.",
140 - SHM_PATH,
141 - )
142 - return False
143 - if get_engram_dp_size() == 1:
144 - logger.warning_once(
145 - "Engram DP replicas are not co-located on one node; "
146 - "storing the offloaded tables per rank instead of sharing them."
147 - )
148 - return False
149 - num_bytes = sum(layout.num_embeddings) * (
150 - layout.head_dim + layout.head_dim // block_size
151 - )
152 - if shutil.disk_usage(SHM_PATH).total < num_bytes:
153 - logger.warning_once(
154 - "Sharing Engram tables across DP replicas needs %.1f GiB of %s "
155 - "(--shm-size or --ipc=host); sharding them across replicas instead.",
156 - num_bytes / 1024**3,
157 - SHM_PATH,
158 - )
159 - return False
160 - return True
161 -
162 -
163 -class DPSharedEngramStorage:
164 - """Registered host weights shared by a node-local DP group with one writer."""
165 -
166 - def __init__(
167 - self, num_rows: int, dim: int, block_size: int, group: GroupCoordinator
168 - ) -> None:
169 - self.group = group
170 - weight_bytes = num_rows * dim
171 - storage = self._allocate(weight_bytes + weight_bytes // block_size)
172 - self.weight = storage[:weight_bytes].view(torch.float8_e4m3fn).view(-1, dim)
173 - self.weight_scale_inv = storage[weight_bytes:].view(-1, dim // block_size)
174 - self._views: tuple[torch.Tensor, torch.Tensor] | None = None
175 -
176 - def _allocate(self, num_bytes: int) -> torch.Tensor:
177 - """Map and register one physical allocation across a node-local DP group."""
178 - group = self.group
179 - with ExitStack() as stack:
180 - path = error = None
181 - if group.rank_in_group == 0:
182 - try:
183 - check_shm_free_space(num_bytes)
184 - backing_file = stack.enter_context(
185 - tempfile.NamedTemporaryFile(prefix="vllm_engram_", dir=SHM_PATH)
186 - )
187 - backing_file.truncate(num_bytes)
188 - path = backing_file.name
189 - except Exception as exc:
190 - error = f"{type(exc).__name__}: {exc}"
191 - path, error = group.broadcast_object((path, error))
192 - if error is not None:
193 - raise RuntimeError(
194 - "Engram shared-memory creation failed on EDP rank 0: " + error
195 - )
196 -
197 - mapping = owner = tensor = None
198 - try:
199 - with open(path, "r+b") as file:
200 - mapping = mmap.mmap(file.fileno(), num_bytes, flags=mmap.MAP_SHARED)
201 - owner = np.frombuffer(mapping, dtype=np.uint8)
202 - pointer = owner.ctypes.data
203 - tensor = torch.from_numpy(owner)
204 - result = torch.cuda.cudart().cudaHostRegister(pointer, num_bytes, 0)
205 - if result.value != 0:
206 - raise RuntimeError(f"cudaHostRegister failed: {result}")
207 - finalizer = weakref.finalize(owner, self._unregister, mapping, pointer)
208 - finalizer.atexit = False # type: ignore[misc]
209 - # The UVA helper otherwise allocates a private pinned copy.
210 - if not tensor.is_pinned():
211 - raise RuntimeError(
212 - "cudaHostRegister did not pin the shared Engram mapping"
213 - )
214 - except Exception as exc:
215 - error = f"{type(exc).__name__}: {exc}"
216 -
217 - # Also fences peer mappings before the leader unlinks the file.
218 - errors: list[str | None] = [None] * group.world_size
219 - torch.distributed.all_gather_object(errors, error, group=group.cpu_group)
220 - failures = "; ".join(
221 - f"EDP rank {rank}: {error}"
222 - for rank, error in enumerate(errors)
223 - if error is not None
224 - )
225 - if failures:
226 - # Dropping the owner runs the finalizer; the mapping then unmaps.
227 - del tensor, owner, mapping
228 - raise RuntimeError(
229 - "Engram shared-memory initialization failed: " + failures
230 - )
231 - assert tensor is not None
232 - return tensor
233 -
234 - @staticmethod
235 - def _unregister(mapping: mmap.mmap, pointer: int) -> None:
236 - # Torch storage retains the numpy owner, including through cached UVA views.
237 - # Keep its mmap alive until CUDA has released the registration.
238 - result = torch.cuda.cudart().cudaHostUnregister(pointer)
239 - if result.value != 0:
240 - logger.warning("Engram cudaHostUnregister failed: %s", result)
241 -
242 - def load_weight(
243 - self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
244 - ) -> None:
245 - if self.group.rank_in_group == 0:
246 - _engram_head_shard_weight_loader(param, loaded_weight)
247 - # Read order may differ across ranks. Equal load counts ensure all shared
248 - # weights are ready after the last weight-loader call returns.
249 - torch.distributed.barrier(group=self.group.cpu_group)
250 -
251 - def get_views(
252 - self, weight: torch.Tensor, scales: torch.Tensor
253 - ) -> tuple[torch.Tensor, torch.Tensor]:
254 - if (weight.data_ptr(), scales.data_ptr()) != (
255 - self.weight.data_ptr(),
256 - self.weight_scale_inv.data_ptr(),
257 - ):
258 - raise RuntimeError("Shared Engram parameter storage must not be replaced")
259 - if self._views is None:
260 - self._views = (
261 - get_accelerator_view_from_cpu_tensor(self.weight),
262 - get_accelerator_view_from_cpu_tensor(self.weight_scale_inv),
263 - )
264 - return self._views
265 -
266 -
267 -class ParallelEngramEmbedding(BaseParallelEngramEmbedding):
268 - """Extend TP lookup with DP head sharding or shared, CPU-offloaded TP slices."""
269 -
270 - _shared_memory: DPSharedEngramStorage | None = None
271 -
272 - def __init__(
273 - self,
274 - num_embeddings: int,
275 - dim: int,
276 - head_sizes: tuple[int, ...],
277 - block_size: int = 32,
278 - cpu_offload: bool = False,
279 - dp_shared_memory: bool = False,
280 - use_thp: bool = False,
281 - ) -> None:
282 - self.cpu_offload = cpu_offload
283 - self.dp_shared_memory = dp_shared_memory
284 - self.use_thp = use_thp
285 - self._packed: torch.Tensor | None = None
286 - self.dp_size = get_engram_dp_size()
287 - if dp_shared_memory:
288 - if not cpu_offload:
289 - raise ValueError("dp_shared_memory requires cpu_offload=True")
290 - if self.dp_size <= 1:
291 - raise ValueError(
292 - "dp_shared_memory requires a node-local Engram DP "
293 - f"group with size > 1; effective Engram DP size is {self.dp_size}. "
294 - "Check that the node layout and rank placement allow complete "
295 - "DP replicas to be co-located."
296 - )
297 - self.dp_size = 1
298 - if cpu_offload and not is_uva_available():
299 - raise RuntimeError("Engram CPU offload requires UVA support")
300 - self._views: tuple[torch.Tensor, torch.Tensor] | None = None
301 - self._view_src: tuple[int, int] | None = None
302 - super().__init__(num_embeddings, dim, head_sizes, block_size)
303 - if cpu_offload:
304 - # Constant dummy values avoid randomizing huge CPU lookup tables.
305 - set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
306 - # The ue8m0 encoding of scale 1.0 is exponent byte 127.
307 - set_weight_attrs(self.weight_scale_inv, {"dummy_weight_value": 127})
308 - logger.info(
309 - "Engram table offloaded to pinned host memory: %d rows x %d, "
310 - "%.2f GiB %s",
311 - self.part_num_embeddings,
312 - dim,
313 - self.part_num_embeddings * (dim + dim // block_size) / 1024**3,
314 - "shared across DP replicas" if dp_shared_memory else "per rank",
315 - )
316 -
317 - def _get_shard_info(self) -> tuple[int, int]:
318 - if self.dp_size == 1:
319 - return super()._get_shard_info()
320 - return self.tp_size * self.dp_size, engram_head_shard_rank()
321 -
322 - def _allocate_weights(self) -> tuple[torch.Tensor, torch.Tensor]:
323 - if self.dp_shared_memory:
324 - group = get_engram_dp_group()
325 - assert group is not None
326 - storage = DPSharedEngramStorage(
327 - self.part_num_embeddings, self.dim, self.block_size, group
328 - )
329 - self._shared_memory = storage
330 - self._weight_loader = storage.load_weight
331 - return storage.weight, storage.weight_scale_inv
332 - if not self.cpu_offload:
333 - return super()._allocate_weights()
334 - if self.use_thp:
335 - weight_bytes = self.part_num_embeddings * self.dim
336 - packed = _allocate_huge_page_storage(
337 - weight_bytes + weight_bytes // self.block_size
338 - )
339 - if packed is not None:
340 - self._packed = packed
341 - return (
342 - packed[:weight_bytes].view(torch.float8_e4m3fn).view(-1, self.dim),
343 - packed[weight_bytes:].view(-1, self.dim // self.block_size),
344 - )
345 - # Model initialization may be inside a CUDA device context.
346 - return (
347 - torch.empty(
348 - self.part_num_embeddings,
349 - self.dim,
350 - dtype=torch.float8_e4m3fn,
351 - device="cpu",
352 - pin_memory=True,
353 - ),
354 - torch.empty(
355 - self.part_num_embeddings,
356 - self.dim // self.block_size,
357 - dtype=torch.uint8,
358 - device="cpu",
359 - pin_memory=True,
360 - ),
361 - )
362 -
363 - def collapse_huge_pages(self) -> None:
364 - """Best-effort MADV_COLLAPSE (Linux >= 6.1) of pages that faulted small."""
365 - if self._packed is None:
366 - return
367 - addr, num_bytes = self._packed.data_ptr(), self._packed.nbytes
368 - libc = ctypes.CDLL(None, use_errno=True)
369 - libc.madvise.argtypes = (ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int)
370 - if libc.madvise(addr, num_bytes, 25) != 0:
371 - logger.warning(
372 - "Engram MADV_COLLAPSE failed; keeping existing pages: %s",
373 - os.strerror(ctypes.get_errno()),
374 - )
375 -
376 - def _storage(self) -> tuple[torch.Tensor, torch.Tensor]:
377 - if self._shared_memory is not None:
378 - return self._shared_memory.get_views(self.weight, self.weight_scale_inv)
379 - if not self.cpu_offload:
380 - return super()._storage()
381 - src = (self.weight.data_ptr(), self.weight_scale_inv.data_ptr())
382 - if self._view_src != src:
383 - self._views = (
384 - get_accelerator_view_from_cpu_tensor(self.weight.data),
385 - get_accelerator_view_from_cpu_tensor(self.weight_scale_inv.data),
386 - )
387 - self._view_src = src
388 - assert self._views is not None
389 - return self._views
390 -
391 - def forward(self, indices: torch.Tensor) -> torch.Tensor:
392 - if self.dp_size == 1:
393 - return super().forward(indices)
394 - num_tokens = indices.shape[0]
395 - indices = gather_engram_hashes(indices)
396 - out = torch.empty(
397 - (indices.shape[0], self.part_n_hash_cols, self.dim),
398 - dtype=torch.bfloat16,
399 - device=indices.device,
400 - )
401 - self.lookup(indices, out)
402 - out = _gather_engram_rows(out, num_tokens)
403 - if self.tp_size > 1:
404 - out = tensor_model_parallel_all_gather(out, dim=1)
405 - return out[:, : self.n_hash_cols]
406 -
407 -
408 -def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor:
409 - """Exchange DP tokens for heads, retaining only this replica's tokens."""
410 - dp_group = get_engram_dp_group()
411 - assert dp_group is not None
412 - slot, remainder = divmod(staged.shape[0], dp_group.world_size)
413 - assert remainder == 0 and 0 <= num_tokens <= slot
414 - gathered = dp_group.all_gather(staged, dim=0)
415 - local_heads, dim = staged.shape[1:]
416 - rows = staged.new_empty((num_tokens, dp_group.world_size * local_heads, dim))
417 - _engram_select_rows(
418 - gathered,
419 - rows,
420 - staged.shape[0],
421 - dp_group.rank_in_group * slot,
422 - local_heads * dim,
423 - )
424 - return rows
425 -
426 -
427 -class Engram(BaseEngram):
428 - """NVIDIA Engram with asynchronous offload and node-local DP lookup."""
429 -
430 - _prefetch_stream: torch.cuda.Stream | None = None
431 - _prefetch_done: torch.cuda.Event | None = None
432 -
433 - def __init__(
434 - self, *args, prefetch_stream: torch.cuda.Stream | None = None, **kwargs
435 - ) -> None:
436 - # Layers sharing one stream serialize their offloaded lookups, so they
437 - # take turns instead of jointly starving decoder compute of SMs.
438 - self._prefetch_stream = prefetch_stream
439 - super().__init__(*args, **kwargs)
440 -
441 - def _create_embedding(
442 - self, layout: EngramLayout, layer_hash_index: int
443 - ) -> ParallelEngramEmbedding:
444 - engram_config = get_current_vllm_config().engram_config
445 - assert engram_config is not None
446 - return ParallelEngramEmbedding(
447 - layout.num_embeddings[layer_hash_index],
448 - layout.head_dim,
449 - tuple(size for order in layout.primes[layer_hash_index] for size in order),
450 - cpu_offload=engram_config.cpu_offload,
451 - dp_shared_memory=bool(engram_config.dp_shared_memory),
452 - use_thp=engram_config.use_thp,
453 - )
454 -
455 - def _init_staging(self, max_tokens: int, head_dim: int) -> None:
456 - super()._init_staging(max_tokens * self.embed_tokens.dp_size, head_dim)
457 - if not self.embed_tokens.cpu_offload:
458 - self._prefetch_stream = None
459 - return
460 - if self._prefetch_stream is None:
461 - self._prefetch_stream = torch.cuda.Stream(device=self.staged_rows.device)
462 - self._prefetch_done = torch.cuda.Event()
463 -
464 - def prepare_embeddings(self, hash_ids: torch.Tensor) -> None:
465 - """Prefetch local shared rows or the DP group's gathered hash IDs."""
466 - if self._prefetch_stream is None:
467 - return super().prepare_embeddings(hash_ids)
468 - rows = self.staged_rows[: hash_ids.shape[0]]
469 - assert rows.shape[0] == hash_ids.shape[0], "engram staging buffer too small"
470 - self._start_prefetch(hash_ids, rows, self._prefetch_stream)
471 -
472 - @eager_break_during_capture
473 - def _start_prefetch(
474 - self, hash_ids: torch.Tensor, rows: torch.Tensor, stream: torch.cuda.Stream
475 - ) -> None:
476 - # Eager boundaries let the lookup span piecewise graph segments.
477 - stream.wait_stream(torch.cuda.current_stream())
478 - # Keep temporary hash storage alive until lookup finishes reading it.
479 - hash_ids.record_stream(stream)
480 - with torch.cuda.stream(stream):
481 - self.embed_tokens.lookup(hash_ids, rows, background=True)
482 - assert self._prefetch_done is not None
483 - self._prefetch_done.record(stream)
484 -
485 - @eager_break_during_capture
486 - def _finish_prefetch(self, event: torch.cuda.Event) -> None:
487 - # Wait for this layer's lookup only, not for later layers on the stream.
488 - torch.cuda.current_stream().wait_event(event)
489 -
490 - def _ready_rows(self, num_tokens: int) -> torch.Tensor:
491 - if self._prefetch_done is not None:
492 - self._finish_prefetch(self._prefetch_done)
493 - if self.embed_tokens.dp_size > 1:
494 - slot = engram_gathered_num_tokens()
495 - staged = self.staged_rows[: slot * self.embed_tokens.dp_size]
496 - return _gather_engram_rows(staged, num_tokens)
497 - return super()._ready_rows(num_tokens)
vllm/models/deepseek_v41/nvidia/model.py
+7 -2
@@ -98,9 +98,14 @@ from vllm.utils.math_utils import cdiv
98 from vllm.v1.attention.backends.registry import AttentionBackendEnum
99 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
100
101 -from ..common.engram import EngramLayout, NgramHashState
101 +from ..common.engram import (
102 + Engram,
103 + EngramLayout,
104 + NgramHashState,
105 + can_share_engram_tables,
106 + gather_engram_hashes,
107 +)
108 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
103 -from .engram import Engram, can_share_engram_tables, gather_engram_hashes
109 from .ops.mhc import (
110 MHC_OVERLAP_MAX_TOKENS,
111 init_mhc_all_reduce,