mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
0016315f80
Both drain the device stream in torch-neuronx (a cast round-trips through the host, a scalar operand is a host->device copy): int32 ids index directly and out-of-range rows are cleared with masked_fill. Tiny cp=2 nextn=5 on device: 9/9 golden tokens. 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 (torch.compiler.is_compiling() or 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
|
|
# No scalar tensor operands: on the neuron device a host scalar is a host->device copy, which drains the
|
|
# stream.
|
|
rows = _gather_rows(x[:, :D], index.clamp(0, n - 1))
|
|
y.copy_(rows.masked_fill(((index < 0) | (index >= n)).view(m, 1), 0))
|
|
|
|
return _gather
|