[Perf][DSv4.1] Exchange packed FP8 Engram rows across DP and feed wkv MXFP8 (#59923)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/distributed/test_engram_dp_shard.py

@@ -39,9 +39,16 @@ from vllm.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    convert_swizzled_to_linear,
+)
 from vllm.models.deepseek_v41.common import engram as engram_ops
 from vllm.models.deepseek_v41.common.engram import (
     Engram,
+    EngramBatch,
     EngramLayout,
     ParallelEngramEmbedding,
     engram_head_shard_rank,
@@ -207,10 +214,75 @@ def _worker(rank: int, tp_size: int, port: int) -> None:
                 n_heads,
                 sequence_parallel,
             )
+        if tp_size == 1:
+            _check_fp8_batch(vllm_config)
     finally:
         cleanup_dist_env_and_memory()
 
 
+def _check_fp8_batch(vllm_config):
+    """MXFP8 staging matches the tables across peer/all-to-all slots, idle
+    replicas and graph replays."""
+    from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
+
+    dp_rank = vllm_config.parallel_config.data_parallel_rank
+    layers = []
+    for index in range(2):
+        weight, scale = _full_table(sum(HEAD_SIZES))
+        weight = (weight.float() * (1 - 2 * index)).to(torch.float8_e4m3fn)
+        table = ParallelEngramEmbedding(
+            sum(HEAD_SIZES), DIM, HEAD_SIZES, cpu_offload=True
+        )
+        table.weight.weight_loader(table.weight, weight)
+        table.weight_scale_inv.weight_loader(table.weight_scale_inv, scale)
+        staged = torch.empty(0, device="cuda")
+        reference = (weight.cuda(), scale.cuda())
+        layers.append(SimpleNamespace(embed_tokens=table, staged_rows=staged))
+        layers[-1].layer_hash_index, layers[-1].reference = index, reference
+    batch = EngramBatch(layers, 257)
+
+    def check(ids):
+        tokens = ids.shape[0]
+        for layer in layers:
+            x = batch.wkv_input(layer, tokens)
+            sf = convert_swizzled_to_linear(x.scale, tokens, x.data.shape[1], BLOCK)
+            actual = dequant_mxfp8_to_bf16(x.data, sf).unflatten(-1, (-1, DIM))
+            weight, scale = layer.reference
+            expected = _reference(weight, scale, ids[:, layer.layer_hash_index])
+            torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+    for counts in ((7, 0, 5, 1), (127,) * 4, (128,) * 4, (257, 0, 129, 1)):
+        tokens = counts[dp_rank]
+        seed = 700 + 10 * dp_rank
+        ids = torch.stack([_make_ids(HEAD_SIZES, tokens, seed + i) for i in (0, 1)], 1)
+        with set_forward_context(
+            None,
+            vllm_config,
+            num_tokens=tokens,
+            num_tokens_across_dp=torch.tensor(counts, dtype=torch.int32),
+        ):
+            batch.prepare_embeddings(ids)
+            check(ids)
+            for capture in ("full", "breakable") if len(set(counts)) == 1 else ():
+                warmup = torch.cuda.Stream()
+                warmup.wait_stream(torch.cuda.current_stream())
+                if capture == "full":
+                    graph = torch.cuda.CUDAGraph()
+                    with torch.cuda.graph(graph, stream=warmup):
+                        batch.prepare_embeddings(ids)
+                else:
+                    graph = BreakableCUDAGraphCapture()
+                    with torch.cuda.stream(warmup), graph:
+                        batch.prepare_embeddings(ids)
+                torch.cuda.current_stream().wait_stream(warmup)
+                for step in (2, 3):
+                    replay_ids = _make_ids(HEAD_SIZES, tokens, seed + step)
+                    ids.copy_(torch.stack([replay_ids] * 2, 1))
+                    graph.replay()
+                    check(ids)
+                del graph
+
+
 def _check_shared_storage_lifetime_and_failures(failure):
     """One peer's failure must roll back every mapping/registration, without hangs."""
     import gc
@@ -516,6 +588,7 @@ def _check_dummy_hash_model_forward(
     model.use_native_mega_moe = False
     model.use_sequence_parallel = False
     model.fuse_mhc_all_reduce = False
+    model.engram_batch = None
     model.engram_hash = state
     model.engram_swa_prefix = "swa"
     model.engram_dp_shared_memory = dp_shared_memory

tests/kernels/test_engram.py

@@ -7,6 +7,9 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    dequant_mxfp8_to_bf16,
+)
 from vllm.models.deepseek_v41.common import engram as engram_ops
 from vllm.models.deepseek_v41.common.engram import (
     Engram,
@@ -714,8 +717,9 @@ def test_engram_lookup_reuses_jit_across_token_shapes():
 )
 @pytest.mark.parametrize("cpu_offload", [False, True])
 @pytest.mark.parametrize("background", [False, True])
+@pytest.mark.parametrize("packed", [False, True])
 @pytest.mark.parametrize("num_tokens", [1, 7, 256])
-def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
+def test_engram_lookup_matches_torch(cpu_offload, background, packed, num_tokens):
     """The fused gather must be bit-exact with the torch dequant path it
     replaces, from HBM and from pinned host memory alike, and must contribute
     zeros for rows another TP rank owns."""
@@ -729,8 +733,13 @@ def test_engram_lookup_matches_torch(cpu_offload, background, num_tokens):
         layer.vocab_start_idx,
         layer.vocab_end_idx,
     )
-    out = torch.empty(num_tokens, cols, layer.dim, dtype=torch.bfloat16, device="cuda")
+    width = layer.dim + layer.dim // 32 if packed else layer.dim
+    dtype = torch.uint8 if packed else torch.bfloat16
+    out = torch.empty(num_tokens, cols, width, dtype=dtype, device="cuda")
     layer.lookup(ids, out, background=background)
+    if packed:
+        values = out[..., : layer.dim].view(torch.float8_e4m3fn)
+        out = dequant_mxfp8_to_bf16(values, out[..., layer.dim :])
     assert torch.equal(out, expected)
 
 

vllm/models/deepseek_v41/common/engram.py

@@ -62,9 +62,16 @@ from vllm.distributed.device_communicators.shm_broadcast import (
 from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
 from vllm.model_executor.utils import set_weight_attrs
+from vllm.models.deepseek_v41.common.ops.query_quant import (
+    can_fuse_query_quant,
+    mxfp8_scale_bytes,
+    mxfp8_scale_offset,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.platform_utils import is_uva_available
@@ -72,6 +79,10 @@ from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
 
 logger = init_logger(__name__)
 
+# Per-rank token slot from which the Engram DP exchange switches from peer
+# reads to all-to-all.
+ENGRAM_A2A_MIN_SLOT = 128
+
 # Cache value for tokens that take no part in an n-gram (image spans).
 DEAD_ID = -1
 
@@ -644,8 +655,10 @@ def _engram_lookup_kernel(
     BLOCK_R: tl.constexpr,
     GRID,
     SORTED: tl.constexpr if _SORTED_IS_CONSTEXPR else None,  # type: ignore[valid-type]
+    PACKED: tl.constexpr,
+    OUT_STRIDE: tl.constexpr,
 ):
-    """Gather fp8 rows, apply their ue8m0 block scales, write bf16.
+    """Gather FP8 rows as packed values/scales or dequantized BF16.
 
     Only this rank's heads are read; padded heads write zeros for all-gather.
     `weight`/`scales` may address pinned host memory through UVA. If SORTED,
@@ -684,13 +697,28 @@ def _engram_lookup_kernel(
             mask=owned[:, None],
             other=0,
         )
-        # ue8m0 is a power of two, so its byte *is* the fp32 exponent field.
-        scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
-        tl.store(
-            out + dst[:, None] * DIM + cols[None, :],
-            (values.to(tl.float32) * scale).to(tl.bfloat16),
-            mask=valid[:, None],
-        )
+        if PACKED:
+            width: tl.constexpr = DIM + DIM // QUANT_BLOCK
+            dst = dst.to(tl.int64)
+            dst = dst // LOCAL_HEADS * OUT_STRIDE + dst % LOCAL_HEADS * width
+            tl.store(
+                out + dst[:, None] + cols[None, :],
+                values.to(tl.uint8, bitcast=True),
+                mask=valid[:, None],
+            )
+            tl.store(
+                out + dst[:, None] + DIM + scale_cols[None, :],
+                scale,
+                mask=valid[:, None] & (cols[None, :] % QUANT_BLOCK == 0),
+            )
+        else:
+            # ue8m0 is a power of two, so its byte *is* the fp32 exponent field.
+            scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
+            tl.store(
+                out + dst[:, None] * DIM + cols[None, :],
+                (values.to(tl.float32) * scale).to(tl.bfloat16),
+                mask=valid[:, None],
+            )
 
 
 def engram_head_shard_rank() -> int:
@@ -1082,7 +1110,7 @@ class ParallelEngramEmbedding(nn.Module):
     def lookup(
         self, indices: torch.Tensor, out: torch.Tensor, background: bool = False
     ) -> None:
-        """Look up local heads of [T, heads] into [T, local_heads, dim] bf16.
+        """Look up local heads into BF16, or packed FP8 bytes followed by scales.
 
         `background` limits the grid to leave SMs for concurrent work.
         """
@@ -1131,6 +1159,8 @@ class ParallelEngramEmbedding(nn.Module):
             BLOCK_R=16,
             GRID=grid,
             SORTED=sort,
+            PACKED=out.dtype == torch.uint8,
+            OUT_STRIDE=out.stride(0),
         )
 
     def forward(self, indices: torch.Tensor) -> torch.Tensor:
@@ -1311,6 +1341,175 @@ def _gather_engram_rows(staged: torch.Tensor, num_tokens: int) -> torch.Tensor:
     return rows
 
 
+@triton.jit(do_not_specialize=["num_tokens"])
+def _engram_unpack_fp8_kernel(
+    packed,
+    values,
+    scales,
+    num_tokens,
+    values_stride,
+    scales_stride,
+    HEADS: tl.constexpr,
+    LOCAL_HEADS: tl.constexpr,
+    LAYERS: tl.constexpr,
+    RANK: tl.constexpr,
+    DIM: tl.constexpr,
+    BLOCK: tl.constexpr,
+    PEER: tl.constexpr,
+):
+    """Split `[rank][token][layer][local head]` packed rows into MXFP8 values and
+    F8_128x4 scales, one layer per grid column; `packed` holds per-rank
+    pointers if PEER. Scales past `num_tokens` zero-pad the 128-row tile.
+    """
+    row = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    layer = tl.program_id(1).to(tl.int64)
+    token, head = row // HEADS, row % HEADS
+    if PEER:
+        src = tl.load(packed + head // LOCAL_HEADS)
+        src = tl.multiple_of(src, DIM // 32).to(tl.pointer_type(tl.uint8))
+        source = RANK * num_tokens + token
+    else:
+        src = packed
+        source = head // LOCAL_HEADS * num_tokens + token
+    source = (source * LAYERS + layer) * LOCAL_HEADS + head % LOCAL_HEADS
+    # Packed rows are only scale-width aligned.
+    src += tl.multiple_of(source * (DIM + DIM // 32), DIM // 32)
+    valid = token < num_tokens
+    cols = tl.arange(0, DIM)
+    data = tl.load(src[:, None] + cols[None, :], mask=valid[:, None], other=0)
+    values += layer * values_stride + row * DIM
+    tl.store(values[:, None] + cols[None, :], data, mask=valid[:, None])
+    scale_cols = tl.arange(0, DIM // 32)
+    scale = tl.load(
+        src[:, None] + DIM + scale_cols[None, :], mask=valid[:, None], other=0
+    )
+    group = head[:, None] * (DIM // 32) + scale_cols[None, :]
+    offset = mxfp8_scale_offset(token[:, None], group, HEADS * DIM // 32)
+    tl.store(scales + layer * scales_stride + offset, scale)
+
+
+class EngramBatch(nn.Module):
+    """Stage every Engram layer's WKV input as MXFP8 with one exchange.
+
+    A module so the model owns the peer exchange's pointer buffers.
+    """
+
+    def __init__(self, layers: list["Engram"], max_tokens: int) -> None:
+        super().__init__()
+        first = layers[0].layer_hash_index
+        hash_indices = [layer.layer_hash_index for layer in layers]
+        assert hash_indices == list(range(first, first + len(layers)))
+        self.layers = layers
+        table = layers[0].embed_tokens
+        device = layers[0].staged_rows.device
+        self.packed_rows = torch.empty(
+            max_tokens * table.dp_size,
+            len(layers),
+            table.part_n_hash_cols,
+            table.dim + table.dim // 32,
+            dtype=torch.uint8,
+            device=device,
+        )
+        width = table.n_hash_cols * table.dim
+        self.fp8_values = torch.empty(
+            len(layers), max_tokens, width, dtype=torch.uint8, device=device
+        )
+        self.fp8_scales = torch.empty(
+            len(layers),
+            mxfp8_scale_bytes(max_tokens, width),
+            dtype=torch.uint8,
+            device=device,
+        )
+        max_peer_slot = min(max_tokens, ENGRAM_A2A_MIN_SLOT - 1)
+        self.peer_exchange = None
+        if table.cpu_offload:
+            from vllm.models.deepseek_v41.nvidia.ops.engram_peer import (
+                EngramPeerExchange,
+            )
+
+            self.peer_exchange = EngramPeerExchange.create(layers, max_peer_slot)
+        for layer in layers:
+            # Owned by the model; keep it out of each layer's submodules.
+            object.__setattr__(layer, "_batch", self)
+            # The layers read this batch's staging instead of their BF16 rows.
+            layer.staged_rows = layer.staged_rows.new_empty(
+                (0, *layer.staged_rows.shape[1:])
+            )
+
+    @classmethod
+    def create(cls, layers: list["Engram"], max_tokens: int) -> "EngramBatch | None":
+        """Batch DP-sharded tables whose WKV takes FlashInfer MXFP8 input."""
+        if not layers:
+            return None
+        table = layers[0].embed_tokens
+        if (
+            table.dp_size == 1
+            or table.tp_size > 1
+            or table.dim % 128
+            or get_current_vllm_config().lora_config is not None
+            or not can_fuse_query_quant([layer.wkv for layer in layers])
+        ):
+            return None
+        return cls(layers, max_tokens)
+
+    def prepare_embeddings(self, hashes: torch.Tensor) -> None:
+        """Stage every layer's WKV input before the decoder layers.
+
+        DP-sharded tables decide on the group's slot, so every rank agrees.
+        """
+        slot = engram_gathered_num_tokens()
+        peer = self.peer_exchange is not None and slot < ENGRAM_A2A_MIN_SLOT
+        self._stage(hashes if peer else gather_engram_hashes(hashes), slot, peer)
+
+    @eager_break_during_capture
+    def _stage(self, hashes: torch.Tensor, slot: int, peer: bool) -> None:
+        if slot == 0:
+            return
+        if peer:
+            assert self.peer_exchange is not None
+            rows = self.peer_exchange.exchange(hashes, slot)
+            rank = self.peer_exchange.dp_rank
+        else:
+            rows, rank = self.packed_rows[: hashes.shape[0]], 0
+            for index, layer in enumerate(self.layers):
+                layer.embed_tokens.lookup(
+                    hashes[:, layer.layer_hash_index], rows[:, index]
+                )
+            group = get_engram_dp_group()
+            assert group is not None
+            sent, rows = rows, torch.empty_like(rows)
+            torch.distributed.all_to_all_single(rows, sent, group=group.device_group)
+        table = self.layers[0].embed_tokens
+        heads = table.n_hash_cols
+        grid = (triton.cdiv(slot, 128) * 128 * heads // 16, len(self.layers))
+        _engram_unpack_fp8_kernel[grid](
+            rows,
+            self.fp8_values,
+            self.fp8_scales,
+            slot,
+            self.fp8_values.stride(0),
+            self.fp8_scales.stride(0),
+            HEADS=heads,
+            LOCAL_HEADS=table.part_n_hash_cols,
+            LAYERS=len(self.layers),
+            RANK=rank,
+            DIM=table.dim,
+            BLOCK=16,
+            PEER=peer,
+        )
+
+    def wkv_input(self, layer: "Engram", num_tokens: int) -> QuantizedActivation:
+        index = layer.layer_hash_index - self.layers[0].layer_hash_index
+        values = self.fp8_values[index, :num_tokens].view(torch.float8_e4m3fn)
+        return QuantizedActivation(
+            values,
+            self.fp8_scales[index, : mxfp8_scale_bytes(num_tokens, values.shape[1])],
+            torch.bfloat16,
+            values.shape,
+            kMxfp8Dynamic,
+        )
+
+
 class Engram(nn.Module):
     """Writes an n-gram lookup into the residual stream, gated by how well it
     matches that stream.
@@ -1323,6 +1522,7 @@ class Engram(nn.Module):
 
     _prefetch_stream: torch.cuda.Stream | None = None
     _prefetch_done: torch.cuda.Event | None = None
+    _batch: EngramBatch | None = None
 
     def __init__(
         self,
@@ -1482,7 +1682,10 @@ class Engram(nn.Module):
         """hidden_states: [T, hc_mult, dim]; hash_ids: [T, n_hash_cols] (all
         tokens, pre sequence-parallel shard); token_mask: [T], False shuts
         the gate so those positions pass through untouched."""
-        kv = self.wkv(self.embed(hash_ids).flatten(-2))
+        if self._batch is not None:
+            kv = self.wkv(self._batch.wkv_input(self, hash_ids.shape[0]))
+        else:
+            kv = self.wkv(self.embed(hash_ids).flatten(-2))
         num_kv_tokens = hash_ids.shape[0]
         assert token_mask is None or token_mask.shape == (num_kv_tokens,)
         if self.use_sequence_parallel:

vllm/models/deepseek_v41/common/ops/query_quant.py

@@ -21,6 +21,18 @@ def _norm_row(x, weight, row, stride, eps, SIZE: tl.constexpr, BLOCK: tl.constex
     return values * rrms * w
 
 
+@triton.jit
+def mxfp8_scale_offset(row, group, GROUPS: tl.constexpr):
+    """F8_128x4: [row/128, group/4, row%32, row%128/32, group%4]."""
+    offset = (row // 128 * triton.cdiv(GROUPS, 4) + group // 4) * 512
+    return offset + row % 32 * 16 + row % 128 // 32 * 4 + group % 4
+
+
+def mxfp8_scale_bytes(tokens: int, width: int) -> int:
+    """Bytes of F8_128x4 scales for `tokens` rows of `width` MXFP8 values."""
+    return triton.cdiv(tokens, 128) * 128 * triton.cdiv(width // 32, 4) * 4
+
+
 @triton.jit(do_not_specialize=["num_tokens"])
 def _q_kv_norm_quant_kernel(
     q,
@@ -68,14 +80,7 @@ def _q_kv_norm_quant_kernel(
             sf = tl.full((BLOCK // 32,), 0, tl.uint32)
         padded_groups: tl.constexpr = triton.cdiv(Q_SIZE // 32, 4) * 4
         sf = tl.where(groups < Q_SIZE // 32, sf, 0)
-        # F8_128x4: [row/128, group/4, row%32, row%128/32, group%4].
-        offsets = (
-            row // 128 * (128 * padded_groups)
-            + groups // 4 * 512
-            + row % 32 * 16
-            + row % 128 // 32 * 4
-            + groups % 4
-        )
+        offsets = mxfp8_scale_offset(row, groups, Q_SIZE // 32)
         tl.store(scales + offsets, sf, groups < padded_groups)
     elif row < num_tokens:
         y = _norm_row(kv, kvw, row, kv_stride, eps, KV_SIZE, BLOCK)
@@ -99,9 +104,8 @@ def fused_q_kv_rmsnorm_quant(
     qo = torch.empty(qr.shape, dtype=torch.float8_e4m3fn, device=qr.device)
     kvo = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
     padded_tokens = triton.cdiv(tokens, 128) * 128
-    padded_groups = triton.cdiv(q_size // 32, 4) * 4
     scales = torch.empty(
-        padded_tokens * padded_groups, dtype=torch.uint8, device=qr.device
+        mxfp8_scale_bytes(tokens, q_size), dtype=torch.uint8, device=qr.device
     )
     if tokens:
         block = triton.next_power_of_2(max(q_size, kv_size))

vllm/models/deepseek_v41/nvidia/model.py

@@ -100,6 +100,7 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
 from ..common.engram import (
     Engram,
+    EngramBatch,
     EngramLayout,
     NgramHashState,
     can_share_engram_tables,
@@ -777,6 +778,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             ),
             prefix=f"{prefix}.layers",
         )
+        self.engram_batch = EngramBatch.create(
+            [
+                layer.engram
+                for layer in islice(self.layers, self.start_layer, self.end_layer)
+                if getattr(layer, "engram", None) is not None
+            ],
+            vllm_config.scheduler_config.max_num_batched_tokens,
+        )
         if self.fuse_mhc_all_reduce:
             # A MoE's top-k finalize folds into the next layer's first mHC
             # boundary, which the last local layer lacks and an engram layer
@@ -971,15 +980,18 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
             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(
-                            gathered_hashes[:, engram.layer_hash_index]
-                        )
+                if self.engram_batch is not None:
+                    self.engram_batch.prepare_embeddings(engram_hashes)
+                else:
+                    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(
+                                gathered_hashes[:, engram.layer_hash_index]
+                            )
 
         full_num_tokens = positions.shape[0]
         if self.use_sequence_parallel:

vllm/models/deepseek_v41/nvidia/ops/engram_peer.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Symmetric-memory Engram exchange for slots below ENGRAM_A2A_MIN_SLOT."""
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+from torch import nn
+
+from vllm.distributed import get_engram_dp_group
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+if TYPE_CHECKING:
+    from vllm.models.deepseek_v41.common.engram import Engram
+
+
+@triton.jit(do_not_specialize=["num_tokens"])
+def _peer_lookup_kernel(
+    metadata,
+    id_ptrs,
+    output,
+    num_tokens,
+    HEADS: tl.constexpr,
+    LAYERS: tl.constexpr,
+    DP: tl.constexpr,
+    RANK: tl.constexpr,
+    DIM: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    """Look up this rank's heads for every peer's tokens into packed rows."""
+    row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    local_heads: tl.constexpr = HEADS // DP
+    valid = row < num_tokens * LAYERS * HEADS
+    token = row // (LAYERS * local_heads)
+    layer, head = row // local_heads % LAYERS, row % local_heads
+    ids = tl.load(id_ptrs + token // num_tokens, valid, other=0).to(
+        tl.pointer_type(tl.int32)
+    )
+    index = tl.load(
+        ids + (token % num_tokens * LAYERS + layer) * HEADS + RANK * local_heads + head,
+        valid,
+        other=-1,
+    ).to(tl.int64)
+    # Per-layer metadata: weight pointer, scale pointer, vocab start, vocab end.
+    table = metadata + layer * 4
+    start, end = tl.load(table + 2), tl.load(table + 3)
+    owned = valid & (index >= start) & (index < end)
+    local = tl.where(owned, index - start, 0)
+    weight = tl.multiple_of(tl.load(table), DIM).to(tl.pointer_type(tl.uint8))
+    scales = tl.load(table + 1)
+    scales = tl.multiple_of(scales, DIM // 32).to(tl.pointer_type(tl.uint8))
+    col = tl.arange(0, DIM)
+    value = tl.load(
+        weight[:, None] + local[:, None] * DIM + col[None, :], owned[:, None], other=0
+    )
+    dst = row.to(tl.int64) * (DIM + DIM // 32)
+    tl.store(output + dst[:, None] + col[None, :], value, valid[:, None])
+    scale_cols = tl.arange(0, DIM // 32)
+    scale = tl.load(
+        scales[:, None] + local[:, None] * (DIM // 32) + scale_cols[None, :],
+        owned[:, None],
+        other=0,
+    )
+    tl.store(output + dst[:, None] + DIM + scale_cols[None, :], scale, valid[:, None])
+
+
+class EngramPeerExchange(nn.Module):
+    """Two phase barriers fence reuse of each other's buffers across forwards.
+
+    Pointer tables are buffers so level-2 sleep restores them on wake-up.
+    """
+
+    def __init__(self, layers: list["Engram"], max_tokens: int):
+        import torch.distributed._symmetric_memory as symm_mem
+
+        super().__init__()
+        group = get_engram_dp_group()
+        assert group is not None
+        self.dp_size, self.dp_rank = group.world_size, group.rank_in_group
+        device = layers[0].staged_rows.device
+        self.num_layers = len(layers)
+        self.hash_start = layers[0].layer_hash_index
+        table = layers[0].embed_tokens
+        self.heads, self.dim = table.n_hash_cols, table.dim
+        self.ids = symm_mem.empty(
+            (max_tokens, self.num_layers, self.heads),
+            dtype=torch.int32,
+            device=device,
+        )
+        self.rows = symm_mem.empty(
+            (
+                max_tokens * self.dp_size,
+                self.num_layers,
+                self.heads // self.dp_size,
+                self.dim + self.dim // 32,
+            ),
+            dtype=torch.uint8,
+            device=device,
+        )
+        self.id_handle = symm_mem.rendezvous(self.ids, group.device_group)
+        self.row_handle = symm_mem.rendezvous(self.rows, group.device_group)
+        metadata = []
+        for layer in layers:
+            table = layer.embed_tokens
+            weight, scale = table._storage()
+            metadata.append(
+                (
+                    weight.data_ptr(),
+                    scale.data_ptr(),
+                    table.vocab_start_idx,
+                    table.vocab_end_idx,
+                )
+            )
+        for name, data in (
+            ("metadata", metadata),
+            ("id_ptrs", self.id_handle.buffer_ptrs),
+            ("row_ptrs", self.row_handle.buffer_ptrs),
+        ):
+            self.register_buffer(
+                name,
+                torch.tensor(data, dtype=torch.int64, device=device),
+                persistent=False,
+            )
+
+    @classmethod
+    def create(
+        cls, layers: list["Engram"], max_tokens: int
+    ) -> "EngramPeerExchange | None":
+        # The Engram DP group is node-local and every rank agrees on these.
+        group = get_engram_dp_group()
+        assert group is not None
+        if not (
+            current_platform.is_device_capability_family(100)
+            and group.world_size in (2, 4, 8)
+            and layers[0].embed_tokens.n_hash_cols % group.world_size == 0
+        ):
+            return None
+        physical_id = current_platform.visible_device_id_to_physical_device_id(
+            torch.accelerator.current_device_index()
+        )
+        physical_ids = [None] * group.world_size
+        dist.all_gather_object(physical_ids, physical_id, group=group.cpu_group)
+        if not current_platform.is_fully_connected(physical_ids):
+            return None
+        return cls(layers, max_tokens)
+
+    def exchange(self, hashes: torch.Tensor, slot: int) -> torch.Tensor:
+        """Look up this rank's heads for every peer; return per-rank row pointers."""
+        tokens = hashes.shape[0]
+        assert 0 < slot <= self.ids.shape[0] and tokens <= slot
+        hash_end = self.hash_start + self.num_layers
+        self.ids[:tokens].copy_(hashes[:, self.hash_start : hash_end])
+        if tokens < slot:
+            self.ids[tokens:slot].fill_(-1)
+        # The preceding rows barrier completed every peer's hash reads. This
+        # barrier completes their previous output reads before rows are reused.
+        self.id_handle.barrier()
+        _peer_lookup_kernel[(triton.cdiv(slot * self.num_layers * self.heads, 16),)](
+            self.metadata,
+            self.id_ptrs,
+            self.rows,
+            slot,
+            HEADS=self.heads,
+            LAYERS=self.num_layers,
+            DP=self.dp_size,
+            RANK=self.dp_rank,
+            DIM=self.dim,
+            BLOCK=16,
+        )
+        self.row_handle.barrier()
+        return self.row_ptrs