From efaca9547dbdfd2d5269f274aaf3798c1f730435 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Fri, 9 Oct 2026 00:30:18 -0700 Subject: [PATCH] [Perf][DSv4.1] Exchange packed FP8 Engram rows across DP and feed wkv MXFP8 (#59923) Signed-off-by: Juntian Liu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_engram_dp_shard.py | 73 ++++++ tests/kernels/test_engram.py | 13 +- vllm/models/deepseek_v41/common/engram.py | 223 +++++++++++++++++- .../deepseek_v41/common/ops/query_quant.py | 24 +- vllm/models/deepseek_v41/nvidia/model.py | 30 ++- .../deepseek_v41/nvidia/ops/engram_peer.py | 173 ++++++++++++++ 6 files changed, 505 insertions(+), 31 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/engram_peer.py diff --git a/tests/distributed/test_engram_dp_shard.py b/tests/distributed/test_engram_dp_shard.py index 637d179a5d..8b7d3e8cec 100644 --- a/tests/distributed/test_engram_dp_shard.py +++ b/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 diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 15bdc9162c..3663108d9d 100644 --- a/tests/kernels/test_engram.py +++ b/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) diff --git a/vllm/models/deepseek_v41/common/engram.py b/vllm/models/deepseek_v41/common/engram.py index d4810ee4ea..0e160c8b06 100644 --- a/vllm/models/deepseek_v41/common/engram.py +++ b/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: diff --git a/vllm/models/deepseek_v41/common/ops/query_quant.py b/vllm/models/deepseek_v41/common/ops/query_quant.py index 6094d1980c..97714e29f1 100644 --- a/vllm/models/deepseek_v41/common/ops/query_quant.py +++ b/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)) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 18f8f27843..8ec0155f73 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/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: diff --git a/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py b/vllm/models/deepseek_v41/nvidia/ops/engram_peer.py new file mode 100644 index 0000000000..4f9423bb1e --- /dev/null +++ b/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