mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
cc90926a1c
Device-memory fixes found bringing up the released checkpoint (16.45 GiB of weights per 24 GB core): - Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the all-gather of the replicated SF table (~3 GB per layer, allocated on the device -> "Failed to allocate resource"). - gather_tl: index_select / advanced indexing of an int8 table allocates several times the table on the device (+4.9 GiB for the 1.3 GB embedding); gather through F.embedding on an int32 / fp16 view instead (bit-exact, no extra memory). - deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements and one expert at a time instead of whole projections / all local experts in fp32 (OOM on a fragmented pool); bit-identical results. - transfer_to_neuron: import torch.utils._device explicitly (missing when torch_neuronx is not auto-loaded). examples/neuron/offline_inference_neuron.py: offline inference in basic mode (text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for the candidate-block filter), one process per NeuronCore. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
43 lines
1.7 KiB
Python
43 lines
1.7 KiB
Python
"""Mirrors ``tilelang_ops/asc/gather.py`` (``gather_tl``): bytewise row gather with zero-fill."""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import torch
|
|
|
|
|
|
def _gather_rows(x: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
|
|
"""``x[index]`` for a byte table ``x [n, D]``.
|
|
|
|
On the neuron device, index_select / advanced indexing of an int8 table allocates several times the table
|
|
(+4.9 GiB for the 1.3 GB V4.1 embedding, which ran the full model out of HBM), while F.embedding through a
|
|
4- or 2-byte view copies only the selected rows and is bit-exact (verified on arbitrary bytes, NaN patterns
|
|
included).
|
|
"""
|
|
if x.is_contiguous():
|
|
for elem, dtype in ((4, torch.int32), (2, torch.float16)):
|
|
if x.shape[1] % elem == 0 and x.storage_offset() % elem == 0:
|
|
return torch.nn.functional.embedding(index, x.view(dtype)).view(torch.int8)
|
|
return torch.index_select(x, 0, index)
|
|
|
|
|
|
def gather_tl(D: int, dtype: torch.dtype, index_dtype: torch.dtype, x_stride: int) -> Callable[..., None]:
|
|
"""Kernel args: ``x [n, D]`` (row stride ``x_stride``), ``index [m]``, ``y [m, D]`` contiguous.
|
|
|
|
``y[i] = x[index[i]]`` if ``0 <= index[i] < n`` else 0; written in place.
|
|
"""
|
|
|
|
def _gather(x: torch.Tensor, index: torch.Tensor, y: torch.Tensor) -> None:
|
|
n = x.shape[0]
|
|
m = index.shape[0]
|
|
if m == 0:
|
|
return
|
|
if n == 0:
|
|
y.zero_()
|
|
return
|
|
valid = (index >= 0) & (index < n)
|
|
safe = torch.where(valid, index, torch.zeros_like(index))
|
|
rows = _gather_rows(x[:, :D], safe)
|
|
y.copy_(torch.where(valid.view(m, 1), rows, torch.zeros((), dtype=rows.dtype, device=rows.device)))
|
|
|
|
return _gather
|