From 0a9059c516f9b458f2b271d207a49ac80388763e Mon Sep 17 00:00:00 2001 From: wukehu Date: Wed, 28 Oct 2026 20:38:41 +0800 Subject: [PATCH] feat: add opt-in local Engram fetch over NVLink Signed-off-by: wukehu --- deepinfer/environ.py | 3 + .../models/v4_modules/engram.py | 15 +- .../parallel_utils/engram_local.py | 404 ++++++++++++++++++ .../parallel_utils/userbuffer.py | 36 +- tests/worker/test_engram_local.py | 74 ++++ 5 files changed, 526 insertions(+), 6 deletions(-) create mode 100644 deepinfer/model_executor/parallel_utils/engram_local.py create mode 100644 tests/worker/test_engram_local.py diff --git a/deepinfer/environ.py b/deepinfer/environ.py index e3a55bb..17c4c03 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -152,6 +152,9 @@ class envs: # -- Model Loading -- DEEPINFER_ENGRAM_PAD_TOKEN = EnvVar.int(default=2) DEEPINFER_ENGRAM_FORCE_GPU_STORAGE = EnvVar(lambda v: _bool(v) or envs.DEEPINFER_ACCELERATOR == "ascend") + # Serve Engram from GPU shards mapped across the NVLink node (torch symmetric memory) instead of + # DeepEP's EngramBuffer, which needs NCCL GIN (GDAKI) RDMA. Single node only; see engram_local.py. + DEEPINFER_ENGRAM_LOCAL = EnvVar.bool() # -- Sampling -- DEEPINFER_SAMPLER_FORCE_TOPK = EnvVar.int() diff --git a/deepinfer/model_executor/models/v4_modules/engram.py b/deepinfer/model_executor/models/v4_modules/engram.py index 123f2ac..a625b6d 100644 --- a/deepinfer/model_executor/models/v4_modules/engram.py +++ b/deepinfer/model_executor/models/v4_modules/engram.py @@ -20,6 +20,7 @@ from deepinfer.engine.engine_type import EngineType from deepinfer.model_executor.layers.linear import ReplicatedLinear from deepinfer.model_executor.models.v4_modules.engram_metadata import CompressedVocabulary from deepinfer.model_executor.models.v4_modules.engram_parameters import EngramEmbedding +from deepinfer.model_executor.parallel_utils.engram_local import LocalEngramBuffer from deepinfer.model_executor.parallel_utils.parallel_state import PG from deepinfer.model_executor.parallel_utils.userbuffer import userbuffer @@ -165,10 +166,10 @@ class EngramRuntime: def __init__(self, layers: Sequence[Engram]) -> None: self._layers = layers assert self._layers - self._buffer: EngramBuffer | None = userbuffer.engram + self._buffer: EngramBuffer | LocalEngramBuffer | None = userbuffer.engram @staticmethod - def allocate(embeddings: Sequence[NgramEmbedding]) -> EngramBuffer: + def allocate(embeddings: Sequence[NgramEmbedding]) -> EngramBuffer | LocalEngramBuffer: reference_embedding = embeddings[0] # The sf of the FP8 table packs 4/2 ue8m0 values into one int32 (CUDA) / int16 (ASC) (pack_ue8m0_to_int in gather_sf), if reference_embedding.fp8_embed: @@ -205,8 +206,14 @@ class EngramRuntime: embeddings = [layer.ngram_embedding for layer in self._layers] if self._buffer is None: self._buffer = self.allocate(embeddings) - weights = [cuda_view_for_write(embedding.weight) for embedding in embeddings] - sfs = [embedding.gather_sf() for embedding in embeddings] if embeddings[0].sf is not None else None + if isinstance(self._buffer, LocalEngramBuffer): + # The local buffer copies CPU or GPU tables itself and keeps only this rank's scale rows, + # so the scales are packed locally instead of all-gathered. + weights = [embedding.weight for embedding in embeddings] + sfs = [_pack_ue8m0_to_int(embedding.sf) for embedding in embeddings] if embeddings[0].sf is not None else None + else: + weights = [cuda_view_for_write(embedding.weight) for embedding in embeddings] + sfs = [embedding.gather_sf() for embedding in embeddings] if embeddings[0].sf is not None else None self._buffer.write(weights, sfs=sfs) for embedding in embeddings: embedding.release_storage() diff --git a/deepinfer/model_executor/parallel_utils/engram_local.py b/deepinfer/model_executor/parallel_utils/engram_local.py new file mode 100644 index 0000000..5dd3ed3 --- /dev/null +++ b/deepinfer/model_executor/parallel_utils/engram_local.py @@ -0,0 +1,404 @@ +"""Engram fetch without RDMA, for one NVLink node (enabled by DEEPINFER_ENGRAM_LOCAL=1). + +Author: Intellifusion Cluster Inference Team + +DeepEP's `EngramBuffer` reads remote table rows with NCCL GIN (GDAKI) RDMA gets, so it needs RDMA verbs +devices even when every rank sits in the same NVLink domain. `LocalEngramBuffer` is a drop-in replacement +with the same `write()` / `fetch()` contract: + +* Each rank keeps its own shard of every layer's table, plus its own rows of the packed scale factors, + in one flat byte buffer allocated with torch symmetric memory (the mechanism DeepGEMM's mega-MoE + `SymmBuffer` and DeepInfer's intranode EPLB shuffle already use). Every rank maps the shards of all + peers, so a fetch is a set of plain gathers that read peer memory over NVLink (one-sided, like the + RDMA gets): no collective is issued per step, ranks do not have to call `fetch` in lockstep, and the + token counts may differ between ranks. +* A global row id `g` of layer `l` belongs to rank `g // n_l` at local row `g % n_l`, where `n_l` is the + per-rank row count of that layer (the same rule as DeepEP's `engram_fetch.cuh`). +* `fetch` returns one hook per layer. A hook returns exactly what DeepEP's hook returns: + - BF16 tables: `data` of shape `[num_tokens, num_entries_per_token * hidden]`, contiguous. + - FP8 tables: `(data, sf)`. `data` is FP8 E4M3 with the shape above. `sf` holds the packed scale + factors of the requested rows, `[num_tokens, num_entries_per_token * num_sf_packs]`, in the dtype the + scales were written with (int32 packed UE8M0 in DeepInfer). With `use_tma_aligned_col_major_sf=True` + it is MN-major and TMA aligned: strides `(1, align(num_tokens, 4))`, a view into a backing tensor of + shape `[num_layers, num_entries_per_token * num_sf_packs * align(num_tokens, 4)]` (layer `l` starts + at storage offset `l * layer_stride`), which is the layout DeepGEMM's `fp8_fp4_gemm_nt` accepts for + SFA without a transform (`check_sf_layout(..., tma_stride_check=True)`). +* Entries with a negative index (image tokens, see `engram_hash`) or an index outside the table are + returned as zero bytes (data and scale factors). DeepEP leaves those rows unwritten (stale bytes). + +The fetch path uses only device-side torch ops with shapes that depend on `num_tokens` alone, so it does +not synchronize the host and can be captured in a CUDA graph. + +The module imports nothing beyond torch so that it can be tested on CPU with the gloo backend: tests +inject a `ShardAllocator` that returns shared CPU tensors instead of symmetric CUDA memory. +""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +# Byte alignment of every region inside a shard buffer. +_REGION_ALIGNMENT_BYTES = 256 +# DeepEP `kNumAlignedSFPacks = 16 / sizeof(sf_pack_t)` (compiled.cuh) and DeepGEMM +# `get_tma_aligned_size(mn, 4)` (utils/math.hpp): the MN stride of 4-byte scale factors is a multiple of 4. +_NUM_ALIGNED_SF_PACKS = 4 +# Same limits as DeepEP's `EngramBuffer.set_config` (engram.hpp) and `EngramWorkspace::kNumMaxLayers`. +_NUM_MAX_LAYERS = 128 +_NUM_RDMA_ALIGNMENT_BYTES = 64 +# Host to device copies of the table are split into chunks of about this size. +_WRITE_CHUNK_BYTES = 1 << 30 + +# (group, num_bytes, device) -> (local flat uint8 buffer, flat uint8 views of every rank's buffer +# indexed by group rank, an object to keep alive while the views are in use) +ShardAllocator = Callable[[dist.ProcessGroup, int, torch.device], tuple[torch.Tensor, list[torch.Tensor], object]] + +_shard_allocator_override: ShardAllocator | None = None + +logger = logging.getLogger(__name__) + + +def set_shard_allocator(allocator: ShardAllocator | None) -> None: + """Override how shard buffers are allocated (tests use shared CPU tensors); None restores the default.""" + global _shard_allocator_override + _shard_allocator_override = allocator + + +def _align(x: int, alignment: int) -> int: + return (x + alignment - 1) // alignment * alignment + + +def _env_flag(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None: + return default + return value not in ("0", "", "false", "False") + + +def symmetric_memory_allocator( + group: dist.ProcessGroup, + num_bytes: int, + device: torch.device, +) -> tuple[torch.Tensor, list[torch.Tensor], object]: + """Allocate `num_bytes` of torch symmetric memory on every rank and map all peers' buffers. + + This is the allocation used by DeepGEMM's `SymmBuffer` (deep_gemm/mega/__init__.py) and by + DeepInfer's intranode EPLB shuffle (eplb/__init__.py `_initialize_shuffle`). With one process per GPU + and only that GPU visible, DeepInfer sets TORCH_SYMM_MEM_ALLOW_OVERLAPPING_DEVICES=1 (deepinfer/__init__.py). + A single-rank group also goes through symmetric memory (so that a one-GPU smoke test exercises the + same code) unless DEEPINFER_ENGRAM_LOCAL_SYMM_MEM=0, which then uses a plain allocation. + """ + rank, num_ranks = group.rank(), group.size() + if num_ranks == 1 and not _env_flag("DEEPINFER_ENGRAM_LOCAL_SYMM_MEM", True): + local = torch.empty(num_bytes, dtype=torch.uint8, device=device) + return local, [local], None + + import torch.distributed._symmetric_memory as symm_mem + + local = symm_mem.empty(num_bytes, dtype=torch.uint8, device=device) + handle = symm_mem.rendezvous(local, group) + peers = [local if peer == rank else handle.get_buffer(peer, (num_bytes,), torch.uint8) for peer in range(num_ranks)] + return local, peers, handle + + +def _gather_view_dtype(num_bytes_per_row: int) -> torch.dtype: + """Widest integer dtype that tiles a row, so gathers move 8-byte words instead of single bytes.""" + for dtype in (torch.int64, torch.int32, torch.int16): + if num_bytes_per_row % dtype.itemsize == 0: + return dtype + return torch.uint8 + + +@dataclass(frozen=True) +class LocalEngramLayout: + """Byte layout of one rank's shard buffer; identical on every rank of the engram group. + + `[table of layer 0] ... [table of layer L-1] [scale rows of layer 0] ... [scale rows of layer L-1]`, + every region aligned to `_REGION_ALIGNMENT_BYTES`. Table rows are `hidden * elem_size` bytes, scale + rows `num_sf_packs * sf_elem_size` bytes, both for this rank's `num_entries_per_layer[l]` rows only. + """ + + num_entries_per_layer: tuple[int, ...] + row_bytes: int + sf_row_bytes: int + table_offsets: tuple[int, ...] + sf_offsets: tuple[int, ...] + num_bytes: int + + @staticmethod + def build( + num_entries_per_layer: Sequence[int], + hidden: int, + elem_size: int, + num_sf_packs: int, + sf_elem_size: int = 4, + ) -> LocalEngramLayout: + row_bytes = hidden * elem_size + sf_row_bytes = num_sf_packs * sf_elem_size + offset = 0 + table_offsets, sf_offsets = [], [] + for num_entries in num_entries_per_layer: + table_offsets.append(offset) + offset = _align(offset + num_entries * row_bytes, _REGION_ALIGNMENT_BYTES) + for num_entries in num_entries_per_layer: + sf_offsets.append(offset) + offset = _align(offset + num_entries * sf_row_bytes, _REGION_ALIGNMENT_BYTES) + return LocalEngramLayout( + num_entries_per_layer=tuple(int(n) for n in num_entries_per_layer), + row_bytes=row_bytes, + sf_row_bytes=sf_row_bytes, + table_offsets=tuple(table_offsets), + sf_offsets=tuple(sf_offsets), + num_bytes=max(offset, _REGION_ALIGNMENT_BYTES), + ) + + +class LocalEngramBuffer: + """Drop-in replacement for `deep_ep.EngramBuffer` on one NVLink node; see the module docstring.""" + + def __init__( + self, + group: dist.ProcessGroup, + num_entries_per_layer: Sequence[int], + hidden: int, + num_max_tokens: int, + num_entries_per_token: int, + dtype: torch.dtype = torch.bfloat16, + num_sf_packs: int = 0, + *, + device: torch.device | str | None = None, + allocator: ShardAllocator | None = None, + ) -> None: + # Same preconditions as DeepEP's EngramBuffer.set_config (csrc/buffers/engram.hpp). + if dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise ValueError(f"unsupported Engram table dtype {dtype}") + if (dtype == torch.float8_e4m3fn) != (num_sf_packs > 0): + raise ValueError(f"FP8 tables need num_sf_packs > 0 and BF16 tables need 0, got {dtype=} {num_sf_packs=}") + if num_max_tokens <= 0 or num_entries_per_token <= 0 or hidden <= 0: + raise ValueError(f"invalid sizes: {num_max_tokens=} {num_entries_per_token=} {hidden=}") + if not 0 < len(num_entries_per_layer) <= _NUM_MAX_LAYERS or any(n <= 0 for n in num_entries_per_layer): + raise ValueError(f"invalid {num_entries_per_layer=}") + if hidden * dtype.itemsize % _NUM_RDMA_ALIGNMENT_BYTES != 0: + raise ValueError(f"row size {hidden * dtype.itemsize} bytes is not a multiple of {_NUM_RDMA_ALIGNMENT_BYTES}") + + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) if torch.cuda.is_available() else torch.device("cpu") + self.device = torch.device(device) + self.group = group + self.rank_idx = group.rank() + self.num_ranks = group.size() + self.num_entries_per_layer = [int(n) for n in num_entries_per_layer] + self.num_layers = len(self.num_entries_per_layer) + self.hidden = hidden + self.dtype = dtype + self.elem_size = dtype.itemsize + self.num_max_tokens = num_max_tokens + self.num_entries_per_token = num_entries_per_token + self.num_sf_packs = num_sf_packs + self.use_fp8 = num_sf_packs > 0 + # DeepInfer writes int32 packed UE8M0 (PLATFORM.fp8_sf_type); the dtype of the written scales wins. + self.sf_dtype = torch.int32 + self.layout = LocalEngramLayout.build(self.num_entries_per_layer, hidden, self.elem_size, num_sf_packs) + self.num_bytes = self.layout.num_bytes + self._check_same_config_on_all_ranks() + + allocator = allocator or _shard_allocator_override or symmetric_memory_allocator + local, peers, self._keepalive = allocator(group, self.num_bytes, self.device) + assert len(peers) == self.num_ranks and peers[self.rank_idx].data_ptr() == local.data_ptr() + assert all(p.dtype == torch.uint8 and p.numel() >= self.num_bytes and p.is_contiguous() for p in peers) + self._local = local + self._peers = peers + + layout = self.layout + self._table_word = _gather_view_dtype(layout.row_bytes) + self._sf_word = _gather_view_dtype(layout.sf_row_bytes) if self.use_fp8 else torch.uint8 + # Per layer, the local table and scale regions as uint8 [rows, row_bytes] (written by `write`). + entries = list(enumerate(self.num_entries_per_layer)) + self._local_tables = [self._region(local, layout.table_offsets[i], n, layout.row_bytes) for i, n in entries] + self._local_sfs = [self._region(local, layout.sf_offsets[i], n, layout.sf_row_bytes) for i, n in entries] if self.use_fp8 else [] + # Per layer, per peer rank: word views used by the gathers in `fetch`. + self._peer_tables = [ + [self._region(peer, layout.table_offsets[i], n, layout.row_bytes).view(self._table_word) for peer in peers] + for i, n in enumerate(self.num_entries_per_layer) + ] + self._peer_sfs = ( + [ + [self._region(peer, layout.sf_offsets[i], n, layout.sf_row_bytes).view(self._sf_word) for peer in peers] + for i, n in enumerate(self.num_entries_per_layer) + ] + if self.use_fp8 + else [] + ) + self._logged_first_fetch = False + self._barrier() + + # ----- helpers ----- + + @staticmethod + def _region(buffer: torch.Tensor, offset: int, num_rows: int, row_bytes: int) -> torch.Tensor: + return buffer[offset : offset + num_rows * row_bytes].view(num_rows, row_bytes) + + def _barrier(self) -> None: + """Make all prior writes of every rank visible to every rank (host-synchronous, init time only).""" + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + token = torch.zeros(1, dtype=torch.int32, device=self.device) + dist.all_reduce(token, group=self.group) + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + + def _check_same_config_on_all_ranks(self) -> None: + mine = torch.tensor( + [self.num_bytes, self.hidden, self.elem_size, self.num_sf_packs, self.num_entries_per_token, *self.num_entries_per_layer], + dtype=torch.int64, + device=self.device, + ) + gathered = torch.empty(self.num_ranks * mine.numel(), dtype=torch.int64, device=self.device) + dist.all_gather_into_tensor(gathered, mine, group=self.group) + gathered = gathered.view(self.num_ranks, mine.numel()) + if not bool((gathered == mine.unsqueeze(0)).all()): + raise RuntimeError(f"Engram local buffer config differs across ranks: {gathered.tolist()}") + + # ----- DeepEP EngramBuffer API ----- + + def destroy(self) -> None: + """Drop the shard buffer and the peer mappings.""" + self._local_tables = self._local_sfs = self._peer_tables = self._peer_sfs = [] + self._peers = [] + self._local = None + self._keepalive = None + + def write(self, storages: list[torch.Tensor], sfs: list[torch.Tensor] | None = None) -> None: + """Copy this rank's table shards (CPU or GPU, `[n_l, hidden]`) and their scale factors into the buffer. + + `sfs[l]` may be this rank's rows only (`[n_l, num_sf_packs]`) or the scale table gathered over the + engram group (`[num_ranks * n_l, num_sf_packs]`, what DeepEP takes); in the second case this rank's + slice is used, as DeepEP does (`nvl_rank_idx * num_shard_bytes`). + """ + self._barrier() + if len(storages) != self.num_layers: + raise ValueError(f"expected {self.num_layers} storages, got {len(storages)}") + for layer_idx, storage in enumerate(storages): + num_entries = self.num_entries_per_layer[layer_idx] + if tuple(storage.shape) != (num_entries, self.hidden): + raise ValueError(f"layer {layer_idx}: storage shape {tuple(storage.shape)} != {(num_entries, self.hidden)}") + if storage.dtype != self.dtype or not storage.is_contiguous(): + raise ValueError(f"layer {layer_idx}: storage must be a contiguous {self.dtype} tensor, got {storage.dtype}") + source = storage.view(torch.uint8).view(num_entries, self.layout.row_bytes) + destination = self._local_tables[layer_idx] + rows_per_chunk = max(1, _WRITE_CHUNK_BYTES // self.layout.row_bytes) + for start in range(0, num_entries, rows_per_chunk): + stop = min(start + rows_per_chunk, num_entries) + destination[start:stop].copy_(source[start:stop]) + + if self.use_fp8: + if sfs is None or len(sfs) != self.num_layers: + raise ValueError("FP8 Engram tables need one scale-factor tensor per layer") + self.sf_dtype = sfs[0].dtype + for layer_idx, sf in enumerate(sfs): + num_entries = self.num_entries_per_layer[layer_idx] + if sf.dim() != 2 or sf.size(1) != self.num_sf_packs or sf.element_size() != 4 or not sf.is_contiguous(): + raise ValueError( + f"layer {layer_idx}: scale factors must be contiguous [rows, {self.num_sf_packs}] with 4-byte elements" + ) + if sf.dtype != self.sf_dtype: + raise ValueError(f"layer {layer_idx}: scale-factor dtype {sf.dtype} != {self.sf_dtype}") + if sf.size(0) == self.num_ranks * num_entries and self.num_ranks > 1: + sf = sf[self.rank_idx * num_entries : (self.rank_idx + 1) * num_entries] + elif sf.size(0) != num_entries: + raise ValueError( + f"layer {layer_idx}: {sf.size(0)} scale rows, expected {num_entries} or {self.num_ranks * num_entries}" + ) + self._local_sfs[layer_idx].copy_(sf.view(torch.uint8).view(num_entries, self.layout.sf_row_bytes)) + elif sfs is not None: + raise ValueError("BF16 Engram tables take no scale factors") + self._barrier() + + def fetch(self, indices: torch.Tensor, num_qps: int = 0, use_tma_aligned_col_major_sf: bool = True) -> list[Callable]: + """Gather the requested rows of every layer and return one hook per layer. + + Arguments: + indices: `[num_layers, num_tokens, num_entries_per_token]` int32 global row ids. + num_qps: accepted for API compatibility with DeepEP and ignored (no QPs here). + use_tma_aligned_col_major_sf: lay out the FP8 scale factors MN-major and TMA aligned. + """ + del num_qps + if indices.dim() != 3 or indices.dtype != torch.int32 or not indices.is_contiguous(): + raise ValueError( + f"indices must be a contiguous int32 [layers, tokens, entries] tensor, got {indices.dtype} {tuple(indices.shape)}" + ) + if indices.device.type != self.device.type: + raise ValueError(f"indices on {indices.device}, buffer on {self.device}") + num_layers, num_tokens, num_entries_per_token = indices.shape + if num_layers != self.num_layers or num_entries_per_token != self.num_entries_per_token: + raise ValueError( + f"indices shape {tuple(indices.shape)} does not match {self.num_layers} layers x {self.num_entries_per_token} entries" + ) + if num_tokens > self.num_max_tokens: + raise ValueError(f"{num_tokens=} exceeds num_max_tokens={self.num_max_tokens}") + + num_entries = num_tokens * num_entries_per_token + sf_width = num_entries_per_token * self.num_sf_packs + if self.use_fp8: + if use_tma_aligned_col_major_sf: + sf_token_stride, sf_hidden_stride = 1, _align(num_tokens, _NUM_ALIGNED_SF_PACKS) + sf_layer_stride = sf_width * sf_hidden_stride + else: + sf_token_stride, sf_hidden_stride = sf_width, 1 + sf_layer_stride = sf_width * num_tokens + sf_backing = torch.zeros((num_layers, sf_layer_stride), dtype=self.sf_dtype, device=self.device) + + results = [] + for layer_idx in range(num_layers): + owner, local_row = self._route(indices[layer_idx].reshape(num_entries), self.num_entries_per_layer[layer_idx]) + rows = self._gather(self._peer_tables[layer_idx], owner, local_row) + data = rows.view(self.dtype).view(num_tokens, num_entries_per_token * self.hidden) + if self.use_fp8: + sf_rows = self._gather(self._peer_sfs[layer_idx], owner, local_row) + sf = sf_backing.as_strided((num_tokens, sf_width), (sf_token_stride, sf_hidden_stride), layer_idx * sf_layer_stride) + sf.copy_(sf_rows.view(self.sf_dtype).view(num_tokens, sf_width)) + results.append((data, sf)) + else: + results.append(data) + if not self._logged_first_fetch: + self._logged_first_fetch = True + first = results[0] + data, sf = first if self.use_fp8 else (first, None) + logger.warning( + f"Engram local fetch in use (rank {self.rank_idx}/{self.num_ranks}): {num_layers} layers, {num_tokens=}, " + f"data {data.dtype} {tuple(data.shape)}" + + (f", sf {sf.dtype} {tuple(sf.shape)} strides {sf.stride()}" if sf is not None else "") + ) + return [_make_hook(result) for result in results] + + # ----- fetch internals ----- + + def _route(self, global_idx: torch.Tensor, num_entries_per_rank: int) -> tuple[torch.Tensor, torch.Tensor]: + """Map global row ids to (owner rank or -1, local row); invalid ids get owner -1 and row 0.""" + global_idx = global_idx.to(torch.int64) + valid = (global_idx >= 0) & (global_idx < self.num_ranks * num_entries_per_rank) + owner = torch.where(valid, torch.div(global_idx, num_entries_per_rank, rounding_mode="floor"), -1) + local_row = torch.where(valid, global_idx - owner * num_entries_per_rank, 0) + return owner, local_row + + @staticmethod + def _gather(views: Sequence[torch.Tensor], owner: torch.Tensor, local_row: torch.Tensor) -> torch.Tensor: + """Rows `views[owner[i]][local_row[i]]`, zero where owner is -1. Reads peer memory directly.""" + out = torch.zeros((owner.numel(), views[0].size(1)), dtype=views[0].dtype, device=owner.device) + for peer, view in enumerate(views): + mine = owner == peer + rows = view.index_select(0, torch.where(mine, local_row, 0)) + out = torch.where(mine.unsqueeze(1), rows, out) + return out + + +def _make_hook(result: object) -> Callable[[], object]: + def hook() -> object: + return result + + return hook diff --git a/deepinfer/model_executor/parallel_utils/userbuffer.py b/deepinfer/model_executor/parallel_utils/userbuffer.py index 3dcffd8..1d9f087 100644 --- a/deepinfer/model_executor/parallel_utils/userbuffer.py +++ b/deepinfer/model_executor/parallel_utils/userbuffer.py @@ -14,6 +14,7 @@ from deep_ep import EngramBuffer, EPBuffer from deepinfer import PLATFORM, envs from deepinfer.logger import init_logger from deepinfer.model_executor.parallel_utils.cp import CPManager +from deepinfer.model_executor.parallel_utils.engram_local import LocalEngramBuffer logger = init_logger(__name__) @@ -57,7 +58,7 @@ class _UserBufferManager: def __init__(self) -> None: self.ep: EPBuffer | None = None self.cp: CPManager | None = None - self.engram: EngramBuffer | None = None + self.engram: EngramBuffer | LocalEngramBuffer | None = None self._sym_buffers: dict[_SymBufferSpec, deep_gemm.SymmBuffer] = {} def allocate_ep( @@ -133,8 +134,12 @@ class _UserBufferManager: num_entries_per_token: int, dtype: torch.dtype, num_sf_packs: int = 0, - ) -> EngramBuffer: + ) -> EngramBuffer | LocalEngramBuffer: assert self.engram is None, "allocate_engram can only be called once" + if envs.DEEPINFER_ENGRAM_LOCAL: + return self._allocate_local_engram( + engram_group, num_entries_per_layer, hidden_size, max_num_tokens, num_entries_per_token, dtype, num_sf_packs + ) num_gpu_bytes, num_rdma_storage_size = EngramBuffer.get_storage_size_hint( engram_group, num_entries_per_layer, @@ -180,6 +185,33 @@ class _UserBufferManager: ) return self.engram + def _allocate_local_engram( + self, + engram_group: dist.ProcessGroup, + num_entries_per_layer: list[int], + hidden_size: int, + max_num_tokens: int, + num_entries_per_token: int, + dtype: torch.dtype, + num_sf_packs: int, + ) -> LocalEngramBuffer: + """No-RDMA Engram: shards in GPU symmetric memory, fetched with NVLink peer reads (one node only).""" + self.engram = LocalEngramBuffer( + engram_group, + num_entries_per_layer, + hidden_size, + max_num_tokens, + num_entries_per_token, + dtype, + num_sf_packs, + ) + logger.info( + f"Created local (no RDMA) Engram buffer: {_gib(self.engram.num_bytes)} GPU shard per rank, " + f"({len(num_entries_per_layer)} layers, Engram {engram_group.rank()}/{engram_group.size()}), " + f"{num_entries_per_layer=}, {hidden_size=}, {max_num_tokens=}, {num_entries_per_token=}, {dtype=}, {num_sf_packs=}" + ) + return self.engram + def allocate_sym_buffer( self, ep_group: dist.ProcessGroup, diff --git a/tests/worker/test_engram_local.py b/tests/worker/test_engram_local.py new file mode 100644 index 0000000..ed49127 --- /dev/null +++ b/tests/worker/test_engram_local.py @@ -0,0 +1,74 @@ +"""Focused LocalEngram correctness tests. + +Author: Intellifusion Cluster Inference Team +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +import deepinfer.model_executor.parallel_utils.engram_local as local_module +from deepinfer.model_executor.parallel_utils.engram_local import LocalEngramBuffer, LocalEngramLayout + + +class _Group: + @staticmethod + def rank() -> int: + return 0 + + @staticmethod + def size() -> int: + return 1 + + +def _allocator(_group: object, num_bytes: int, device: torch.device): + local = torch.empty(num_bytes, dtype=torch.uint8, device=device) + return local, [local], SimpleNamespace() + + +@pytest.fixture +def cpu_dist(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(local_module.dist, "all_reduce", lambda tensor, group=None: tensor) + monkeypatch.setattr(local_module.dist, "all_gather_into_tensor", lambda output, input, group=None: output.copy_(input)) + + +def test_layout_regions_are_aligned() -> None: + layout = LocalEngramLayout.build((3, 5), hidden=32, elem_size=2, num_sf_packs=4) + assert layout.row_bytes == 64 + assert layout.sf_row_bytes == 16 + assert all(offset % 256 == 0 for offset in (*layout.table_offsets, *layout.sf_offsets, layout.num_bytes)) + assert layout.table_offsets[0] < layout.table_offsets[1] < layout.sf_offsets[0] < layout.sf_offsets[1] + + +def test_route_and_gather_zero_invalid_rows() -> None: + buffer = object.__new__(LocalEngramBuffer) + buffer.num_ranks = 2 + owner, local_row = buffer._route(torch.tensor([-1, 0, 3, 6], dtype=torch.int32), num_entries_per_rank=3) + assert owner.tolist() == [-1, 0, 1, -1] + assert local_row.tolist() == [0, 0, 0, 0] + views = [torch.tensor([[10, 11], [20, 21], [30, 31]]), torch.tensor([[40, 41], [50, 51], [60, 61]])] + assert LocalEngramBuffer._gather(views, owner, local_row).tolist() == [[0, 0], [10, 11], [40, 41], [0, 0]] + + +def test_bf16_write_fetch_and_invalid_index(cpu_dist: None) -> None: + # BF16 rows must satisfy the same 64-byte alignment as the RDMA layout. + buffer = LocalEngramBuffer(_Group(), [4], 32, 2, 2, dtype=torch.bfloat16, device="cpu", allocator=_allocator) + buffer.write([torch.arange(128, dtype=torch.bfloat16).view(4, 32).contiguous()]) + result = buffer.fetch(torch.tensor([[[0, 1], [3, -1]]], dtype=torch.int32))[0]() + assert result.shape == (2, 64) + assert result[0, :4].tolist() == [0.0, 1.0, 2.0, 3.0] + assert result[1, :4].tolist() == [96.0, 97.0, 98.0, 99.0] + assert torch.count_nonzero(result[1, 32:]).item() == 0 + + +def test_fp8_fetch_preserves_scale_layout(cpu_dist: None) -> None: + buffer = LocalEngramBuffer(_Group(), [4], 64, 2, 1, dtype=torch.float8_e4m3fn, num_sf_packs=1, device="cpu", allocator=_allocator) + buffer.write([torch.ones((4, 64), dtype=torch.float8_e4m3fn)], [torch.arange(4, dtype=torch.int32).view(4, 1).contiguous()]) + data, sf = buffer.fetch(torch.tensor([[[0], [3]]], dtype=torch.int32))[0]() + assert data.shape == (2, 64) + assert sf.shape == (2, 1) + assert sf.stride() == (1, 4) + assert sf[:, 0].tolist() == [0, 3]