mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-09 22:59:44 +08:00
8935f80230
Three torch-neuronx 2.14 device behaviours, each worked around in a bit-identical way: - torch.cat of float8 tensors is wrong on the device: merge_linear_weights (fused q_a/kv, q_b, ... projections) now concatenates the raw bytes, like the HF loader's read_concat. - int64 constants are narrowed to int32 by the compiler (--implicit-integer-downcast): deep_select.topk used int64.max as the sort sentinel, which became -1 and pushed the selected indices out of rows with end < topk -- the first compressed layer then attended to no compressed KV for short prompts. The sentinel is now vocab_size. - advanced indexing of non-contiguous float8 views is wrong: the remaining site (deep_gemm sparse MQA logits) goes through uint8 views; deep_ep's byte path also covers e8m0. DEEPINFER_NEURON_FP8_GUARD=1|raise (neuron_ops/fp8_guard.py) reports every cat/stack/advanced-index/gather/scatter/where on a float8 tensor with its call site; the tiny end-to-end run and all CPU unit tests report zero sites. Tiny checkpoint on one NeuronCore, teacher-forced vs the official reference: short 9/9 (min cos 0.9985), medium 8/9 (0.9991), long 8/9 (0.9978); logits rel <= 0.066, within the reference's own noise floor -- same as CPU mode. Also: probe_layer0 PROBE_LAYERS=N and attention-input capture, device_diff copies tensor attributes (replays deep_gemm sparse MQA metadata), device primitive checks for the narrowing and strided fp8 casts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
58 lines
2.5 KiB
Python
58 lines
2.5 KiB
Python
"""Row gather/scatter helpers that also work for narrow dtypes (FP8).
|
|
|
|
FP8 (`torch.float8_e4m3fn`) tensors do not support every indexing / fill op on CPU or on
|
|
`neuron`, and DeepEP only ever moves FP8 rows around bit-exactly. These helpers therefore
|
|
route such dtypes through a `uint8` view of the row bytes; everything else is indexed
|
|
directly. Mirrors the "move bytes, never reinterpret" behaviour of the Ascend kernels.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
def _needs_byte_path(dtype: torch.dtype) -> bool:
|
|
# Every 1-byte float (e4m3fn, e5m2, e8m0fnu, ...): advanced indexing / cat of float8 is wrong on the Neuron
|
|
# device (torch-neuronx 2.14), the same ops on uint8 views are correct.
|
|
return dtype.is_floating_point and dtype.itemsize == 1
|
|
|
|
|
|
def byte_rows(x: torch.Tensor) -> torch.Tensor:
|
|
"""`[rows, cols]` -> `[rows, cols * itemsize]` uint8 view (requires contiguity).
|
|
|
|
`view(torch.uint8)` already rescales the last dimension, which keeps this correct for
|
|
zero-row tensors (an extra `view(rows, -1)` would be ambiguous there).
|
|
"""
|
|
assert x.dim() == 2 and x.is_contiguous()
|
|
return x.view(torch.uint8)
|
|
|
|
|
|
def zeros_rows(num_rows: int, num_cols: int, dtype: torch.dtype, device) -> torch.Tensor:
|
|
"""Zero-filled `[num_rows, num_cols]` tensor, FP8-safe."""
|
|
if _needs_byte_path(dtype):
|
|
return torch.zeros((num_rows, num_cols * dtype.itemsize), dtype=torch.uint8,
|
|
device=device).view(dtype)
|
|
return torch.zeros((num_rows, num_cols), dtype=dtype, device=device)
|
|
|
|
|
|
def empty_rows(num_rows: int, num_cols: int, dtype: torch.dtype, device) -> torch.Tensor:
|
|
"""Uninitialized `[num_rows, num_cols]` tensor, FP8-safe."""
|
|
if _needs_byte_path(dtype):
|
|
return torch.empty((num_rows, num_cols * dtype.itemsize), dtype=torch.uint8,
|
|
device=device).view(dtype)
|
|
return torch.empty((num_rows, num_cols), dtype=dtype, device=device)
|
|
|
|
|
|
def take_rows(src: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
|
|
"""`src[index]` on the row axis, FP8-safe."""
|
|
if not _needs_byte_path(src.dtype):
|
|
return src[index]
|
|
return byte_rows(src.contiguous())[index].contiguous().view(src.dtype)
|
|
|
|
|
|
def put_rows(dst: torch.Tensor, index: torch.Tensor, src: torch.Tensor) -> None:
|
|
"""`dst[index] = src` on the row axis, FP8-safe (`dst` must be contiguous for FP8)."""
|
|
if not _needs_byte_path(dst.dtype):
|
|
dst[index] = src
|
|
return
|
|
byte_rows(dst)[index] = byte_rows(src.contiguous())
|