coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-30 15:57:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add opt-in local Engram fetch over NVLink

Browse Source 
Signed-off-by: wukehu <wukehu3@gmail.com>
This commit is contained in:
wukehu 2026-10-28 20:38:41 +08:00
parent 553741557a
commit 0a9059c516
5 changed files with 526 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
deepinfer/environ.py
View File

@ -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()

15
deepinfer/model_executor/models/v4_modules/engram.py
View File

@ -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()

404
deepinfer/model_executor/parallel_utils/engram_local.py Normal file
View File

@ -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

36
deepinfer/model_executor/parallel_utils/userbuffer.py
View File

@ -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,

74
tests/worker/test_engram_local.py Normal file
View File

@ -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]