mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
Merge 0a9059c516f9b458f2b271d207a49ac80388763e into 553741557a57df44a46fb3ab72b6aafeaf945337
This commit is contained in:
wukehu
GitHub
commit
07f9209891
@ -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()
|
||||
|
||||
@ -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
404
deepinfer/model_executor/parallel_utils/engram_local.py
Normal 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
|
||||
@ -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
74
tests/worker/test_engram_local.py
Normal 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]
|
||||
Loading…
x
Reference in New Issue
Block a user