diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 8811ba5..84d8fa6 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -101,6 +101,8 @@ class envs: DEEPINFER_ACCELERATOR = EnvVar.str(default="nvidia") # Neuron only: "neuron" runs on the NeuronCore, "cpu" runs the same pure-torch kernels on the host. DEEPINFER_NEURON_DEVICE = EnvVar.str(default="neuron") + # Neuron only, debug: "1" records float8 cat / advanced-indexing call sites (summary at exit), "raise" raises. + DEEPINFER_NEURON_FP8_GUARD = EnvVar.str(default="") # -- Runtime -- DEEPINFER_COMPUTE_SUM_SMALLER_PROBS = EnvVar.bool() diff --git a/deepinfer/model_executor/models/v4_modules/utils.py b/deepinfer/model_executor/models/v4_modules/utils.py index d8c9e95..d165a74 100644 --- a/deepinfer/model_executor/models/v4_modules/utils.py +++ b/deepinfer/model_executor/models/v4_modules/utils.py @@ -49,7 +49,12 @@ def convert_rope( def merge_linear_weights(weights: list[nn.Parameter], dim: int = 0) -> nn.Parameter: - gathered_weights = nn.Parameter(torch.cat(weights, dim=dim)) + if weights[0].dtype.is_floating_point and weights[0].element_size() == 1: + # FP8: concatenate the raw bytes (bit-identical; torch.cat of float8 is wrong on the Neuron device). + merged = torch.cat([w.data.view(torch.uint8) for w in weights], dim=dim).view(weights[0].dtype) + else: + merged = torch.cat(weights, dim=dim) + gathered_weights = nn.Parameter(merged) if gathered_weights.element_size() == 1: for weight in weights[:-1]: # NOTE The shape of the last weight does not have to be a multiple of 128. assert weight.shape[dim] % 128 == 0, f"{weight.shape=}" diff --git a/deepinfer/neuron_ops/NEURON_PORT.md b/deepinfer/neuron_ops/NEURON_PORT.md index 5ecc57d..53863b5 100644 --- a/deepinfer/neuron_ops/NEURON_PORT.md +++ b/deepinfer/neuron_ops/NEURON_PORT.md @@ -43,6 +43,19 @@ norms/rope, KV cache store/load, basic sampling). is packed two per byte — unpack with bit ops; ue8m0 scales are exponent bytes (value = 2**(b-127)). 6. Correctness over speed. Match the reference numerics (fp32 accumulation where the kernel accumulates in fp32). 7. Each file has a short module docstring saying which upstream API it mirrors. +8. **float8 on the neuron device**: torch-neuronx 2.14 returns wrong data for `torch.cat`/`stack` of float8 + tensors (e4m3fn and e8m0fnu) and for advanced indexing (`t[idx_tensor, ...]`) on non-contiguous float8 views + (tests/neuron/device/test_device_primitives.py); the same ops on `uint8` views of the same bytes are correct. + Move 1-byte floats as bytes, for every concat / gather / scatter / index_put, contiguous or not: + `torch.cat([t.view(torch.uint8) for t in parts], dim).view(dtype)`, `t.view(torch.uint8)[idx].view(dtype)`, + `dst.view(torch.uint8)[idx] = src.view(torch.uint8)` (helpers: `deep_ep/utils/indexing.py`). This is + bit-identical on every platform, so loader/model code can use it unconditionally (`read_concat`, + `merge_linear_weights`). Basic slicing, dtype casts and elementwise math followed by an upcast are fine. + To find offending sites without a NeuronCore, run with `DEEPINFER_NEURON_FP8_GUARD=1` (also in CPU mode): + `deepinfer/neuron_ops/fp8_guard.py` records the call site (file:line, dtype/shape/contiguity) of every float8 + cat/stack/advanced-index/gather/scatter/index_put/where and prints a deduplicated summary at exit + (`=raise` raises at the first one; `DEEPINFER_NEURON_FP8_GUARD_LOG=` appends the summary to a file). + The tiny e2e run and the CPU unit stages must report 0 sites. ## Testing CPU mode (same code, tensors on the host): diff --git a/deepinfer/neuron_ops/fp8_guard.py b/deepinfer/neuron_ops/fp8_guard.py new file mode 100644 index 0000000..a788b02 --- /dev/null +++ b/deepinfer/neuron_ops/fp8_guard.py @@ -0,0 +1,214 @@ +"""Debug guard: report every float8 concatenation / advanced-indexing call site. + +torch-neuronx 2.14 returns wrong data on the ``neuron`` device for ``torch.cat`` of float8 tensors and for +advanced indexing on non-contiguous float8 views, while the same ops on ``uint8`` views of the same bytes are +correct (tests/neuron/device/test_device_primitives.py). Kernel and loader code must therefore move float8 bytes +through ``uint8`` views. This guard finds the sites that do not, **without a NeuronCore**: it also works with +``DEEPINFER_NEURON_DEVICE=cpu``, where the ops are correct and only the call sites are recorded. + + DEEPINFER_NEURON_FP8_GUARD=1 record, print a deduplicated summary at exit + DEEPINFER_NEURON_FP8_GUARD=raise raise RuntimeError at the first offending call + DEEPINFER_NEURON_FP8_GUARD_LOG= also append the summary to (one block per process) + +Installed by ``PlatformNeuron.__init__`` only when the variable is set. It monkey-patches the Python entry points +(``torch.cat``/``stack``/``concat``/..., ``Tensor.__getitem__``/``__setitem__`` with an advanced index, +``gather``/``index_select``/``take*``/``embedding``, ``index_put(_)``/``scatter*``/``index_copy_``/..., ``where``) +and records, for calls that involve a 1-byte float tensor, the first stack frame outside torch and outside this +module together with dtype / shape / contiguity. Basic indexing (ints, slices, None, Ellipsis) and dtype casts are +views or copies that behave correctly and are not reported. Debug only: the overhead is not tuned. +""" + +from __future__ import annotations + +import atexit +import numbers +import os +import sys +import threading +from collections import OrderedDict +from typing import Any, Callable + +import torch + +_MODE: str | None = None # None (not installed), "record" or "raise" +_LOCK = threading.Lock() +_TLS = threading.local() +_SITES: "OrderedDict[tuple[str, str], dict[str, Any]]" = OrderedDict() +_MAX_SIGS = 4 # distinct dtype/shape/contiguity signatures kept per site +_SKIP_DIRS = (os.path.dirname(torch.__file__) + os.sep,) +_THIS_FILE = os.path.abspath(__file__) +_SKIP_FILES = (_THIS_FILE, os.path.join(os.path.dirname(_THIS_FILE), "transfer_to_neuron.py")) # wrappers, not sites + + +def _is_fp8(x: Any) -> bool: + return isinstance(x, torch.Tensor) and x.dtype.is_floating_point and x.dtype.itemsize == 1 + + +def _call_site() -> str: + f = sys._getframe(2) + while f is not None: + fn = os.path.abspath(f.f_code.co_filename) + if fn not in _SKIP_FILES and not fn.startswith(_SKIP_DIRS): + return f"{fn}:{f.f_lineno} ({f.f_code.co_name})" + f = f.f_back + return "" + + +def _desc(t: torch.Tensor) -> str: + dt = str(t.dtype).replace("torch.", "") + return f"{dt}{list(t.shape)}{'' if t.is_contiguous() else ' non-contig'}" + + +def _report(op: str, tensors: list[torch.Tensor]) -> None: + site = _call_site() + sig = ", ".join(_desc(t) for t in tensors[:4]) + (", ..." if len(tensors) > 4 else "") + with _LOCK: + entry = _SITES.setdefault((op, site), {"count": 0, "sigs": []}) + entry["count"] += 1 + if sig not in entry["sigs"] and len(entry["sigs"]) < _MAX_SIGS: + entry["sigs"].append(sig) + if _MODE == "raise": + raise RuntimeError(f"DEEPINFER_NEURON_FP8_GUARD: float8 {op} at {site}: {sig} " + "(route the bytes through .view(torch.uint8))") + + +def _tensors_in(obj: Any) -> list[torch.Tensor]: + if isinstance(obj, torch.Tensor): + return [obj] + if isinstance(obj, (list, tuple)): + return [t for t in obj if isinstance(t, torch.Tensor)] + return [] + + +def _is_advanced_index(idx: Any) -> bool: + """Tensor / list / array index anywhere in the index tuple (a top-level tuple is basic indexing).""" + items = idx if isinstance(idx, tuple) else (idx,) + for it in items: + if isinstance(it, (torch.Tensor, list)): + return True + if hasattr(it, "__array__") and not isinstance(it, numbers.Integral) and getattr(it, "ndim", 1) > 0: + return True # numpy array (numpy integer scalars are basic indexing) + return False + + +def _wrap(fn: Callable, op: str, pick: Callable[..., list[torch.Tensor]]) -> Callable: + if getattr(fn, "_fp8_guard_wrapped", False): + return fn + + def wrapper(*args, **kwargs): + if not getattr(_TLS, "busy", False): + _TLS.busy = True + try: + involved = [t for t in pick(*args, **kwargs) if _is_fp8(t)] + except Exception: # never let the guard change behaviour of a bad call; the real call raises below + involved = [] + finally: + _TLS.busy = False + if involved: + _report(op, involved) + return fn(*args, **kwargs) + + wrapper._fp8_guard_wrapped = True # type: ignore[attr-defined] + wrapper.__wrapped__ = fn # type: ignore[attr-defined] + wrapper.__name__ = getattr(fn, "__name__", op) + wrapper.__doc__ = getattr(fn, "__doc__", None) + return wrapper + + +# --- operand pickers: return the tensors whose float8-ness makes the call suspicious ------------------------------- +def _p_seq(tensors=None, *a, **k): # cat / stack / concat(tensors, dim) + return _tensors_in(tensors if tensors is not None else k.get("tensors")) + + +def _p_self(self, *a, **k): # gather / index_select / take ...: the source tensor + return [self] if isinstance(self, torch.Tensor) else [] + + +def _p_getitem(self, idx): + return [self] if _is_advanced_index(idx) else [] + + +def _p_setitem(self, idx, value): + return [self, *_tensors_in(value)] if _is_advanced_index(idx) else [] + + +def _p_all(*a, **k): # every tensor operand (index_put / scatter / where ...) + out: list[torch.Tensor] = [] + for x in list(a) + list(k.values()): + out.extend(_tensors_in(x)) + return out + + +def _p_embedding(input, weight, *a, **k): + return [weight] + + +_TORCH_FNS = { + "cat": _p_seq, "concat": _p_seq, "concatenate": _p_seq, "stack": _p_seq, + "hstack": _p_seq, "vstack": _p_seq, "dstack": _p_seq, "column_stack": _p_seq, + "gather": _p_self, "index_select": _p_self, "take": _p_self, "take_along_dim": _p_self, + "index_put": _p_all, "index_put_": _p_all, "scatter": _p_all, "scatter_add": _p_all, + "scatter_reduce": _p_all, "index_copy": _p_all, "index_add": _p_all, "index_fill": _p_all, + "masked_scatter": _p_all, "masked_select": _p_all, "where": _p_all, +} +_TENSOR_METHODS = { + "gather": _p_self, "index_select": _p_self, "take": _p_self, "take_along_dim": _p_self, + "index_put": _p_all, "index_put_": _p_all, "scatter": _p_all, "scatter_": _p_all, + "scatter_add": _p_all, "scatter_add_": _p_all, "scatter_reduce": _p_all, "scatter_reduce_": _p_all, + "index_copy": _p_all, "index_copy_": _p_all, "index_add": _p_all, "index_add_": _p_all, + "index_fill": _p_all, "index_fill_": _p_all, "masked_scatter": _p_all, "masked_scatter_": _p_all, + "masked_select": _p_all, "where": _p_all, +} + + +def _summary() -> str: + lines = [f"===== DEEPINFER_NEURON_FP8_GUARD (pid {os.getpid()}): {len(_SITES)} float8 cat/advanced-index site(s)"] + for (op, site), e in _SITES.items(): + lines.append(f" {op:<20} x{e['count']:<7} {site}") + for sig in e["sigs"]: + lines.append(f" {sig}") + return "\n".join(lines) + + +def _at_exit() -> None: + text = _summary() + print(text, file=sys.stderr, flush=True) + path = os.getenv("DEEPINFER_NEURON_FP8_GUARD_LOG") + if path: + with open(path, "a") as f: + f.write(text + "\n") + + +def sites() -> dict[tuple[str, str], dict[str, Any]]: + """Recorded (op, call site) -> {count, sigs}; for tests.""" + with _LOCK: + return {k: {"count": v["count"], "sigs": list(v["sigs"])} for k, v in _SITES.items()} + + +def reset() -> None: + with _LOCK: + _SITES.clear() + + +def install(mode: str | None = None) -> None: + """Install the guard. ``mode``: ``"raise"`` or anything truthy for record mode (default: from the env var).""" + global _MODE + if mode is None: + mode = os.getenv("DEEPINFER_NEURON_FP8_GUARD", "") + if not mode or mode == "0" or _MODE is not None: + return + _MODE = "raise" if mode == "raise" else "record" + for name, pick in _TORCH_FNS.items(): + if hasattr(torch, name): + setattr(torch, name, _wrap(getattr(torch, name), f"torch.{name}", pick)) + for name, pick in _TENSOR_METHODS.items(): + if hasattr(torch.Tensor, name): + setattr(torch.Tensor, name, _wrap(getattr(torch.Tensor, name), f"Tensor.{name}", pick)) + torch.Tensor.__getitem__ = _wrap(torch.Tensor.__getitem__, "getitem[adv]", _p_getitem) + torch.Tensor.__setitem__ = _wrap(torch.Tensor.__setitem__, "setitem[adv]", _p_setitem) + import torch.nn.functional as F + + F.embedding = _wrap(F.embedding, "F.embedding", _p_embedding) + atexit.register(_at_exit) + print(f"[deepinfer] DEEPINFER_NEURON_FP8_GUARD={_MODE}: recording float8 cat/advanced-index sites", + file=sys.stderr, flush=True) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py b/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py index bab9c1d..d4f8aca 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py @@ -10,11 +10,10 @@ from __future__ import annotations import torch -_BYTE_PATH_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) - - def _needs_byte_path(dtype: torch.dtype) -> bool: - return dtype in _BYTE_PATH_DTYPES + # 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: diff --git a/deepinfer/neuron_ops/vendor/deep_gemm/attention.py b/deepinfer/neuron_ops/vendor/deep_gemm/attention.py index 7797fee..ece906b 100644 --- a/deepinfer/neuron_ops/vendor/deep_gemm/attention.py +++ b/deepinfer/neuron_ops/vendor/deep_gemm/attention.py @@ -71,6 +71,14 @@ def _chunk_columns(num_q: int, num_heads: int, gather_keys: bool, multiple: int return max(multiple, chunk // multiple * multiple) +def _take_rows(t: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """``t[index]``; 1-byte float rows are gathered as ``uint8`` bytes (bit-identical, and advanced indexing on + float8 is wrong on the Neuron device, torch-neuronx 2.14).""" + if t.dtype.is_floating_point and t.dtype.itemsize == 1: + return t.view(torch.uint8)[index].view(t.dtype) + return t[index] + + def _decode_rows(data: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: """Dequantize MQA rows ``[..., head_dim_storage]`` with UE8M0 scales.""" values = decode_data(data) @@ -358,7 +366,7 @@ def fp8_fp4_sparse_mqa_logits( valid = (columns < (num_blocks * sparse_block_kv)[:, None]) & (tokens >= ks[:, None]) & (tokens < ke[:, None]) safe = tokens.clamp(0, max(0, num_kv_tokens - 1)) flat = safe.reshape(-1) - keys = _decode_rows(kv_data[flat], kv_sf[flat]).reshape(safe.shape[0], safe.shape[1], HEAD_DIM) + keys = _decode_rows(_take_rows(kv_data, flat), kv_sf[flat]).reshape(safe.shape[0], safe.shape[1], HEAD_DIM) values = _scores_per_query_keys(q_f32, keys, weights_f32) logits[:, c0:c1] = torch.where(valid, values, torch.full_like(values, _NEG_INF)).to(torch.bfloat16) return logits diff --git a/deepinfer/neuron_ops/vendor/deep_select/interface.py b/deepinfer/neuron_ops/vendor/deep_select/interface.py index 51399e6..48b4fbe 100644 --- a/deepinfer/neuron_ops/vendor/deep_select/interface.py +++ b/deepinfer/neuron_ops/vendor/deep_select/interface.py @@ -114,7 +114,10 @@ def topk( if sorted_index: # Keep the valid prefix in ascending index order; the padding stays at the end. - big = torch.full_like(idx, torch.iinfo(torch.int64).max) + # Any value above every real index works as the sort sentinel. Keep it in int32 range: the neuron + # compiler narrows int64 (--implicit-integer-downcast), which turned int64.max into -1 and sorted the + # padding in front of the selected indices for rows with end < topk. + big = torch.full_like(idx, vocab_size) idx = torch.sort(torch.where(selected_mask, idx, big), dim=1).values idx = torch.where(selected_mask, idx, torch.zeros_like(idx)) diff --git a/deepinfer/platform.py b/deepinfer/platform.py index 15c923c..745e736 100644 --- a/deepinfer/platform.py +++ b/deepinfer/platform.py @@ -271,6 +271,12 @@ class PlatformNeuron(Platform): import deepinfer.neuron_ops # noqa: F401 + if envs.DEEPINFER_NEURON_FP8_GUARD not in ("", "0"): + # Debug only: report float8 cat / advanced-indexing call sites (wrong on the device), see fp8_guard.py. + from deepinfer.neuron_ops import fp8_guard + + fp8_guard.install(envs.DEEPINFER_NEURON_FP8_GUARD) + def get_distributed_options(self) -> object: return None diff --git a/examples/neuron/device_diff.py b/examples/neuron/device_diff.py index a6a0ac1..83d01f8 100644 --- a/examples/neuron/device_diff.py +++ b/examples/neuron/device_diff.py @@ -31,7 +31,11 @@ _ACTIVE = [True] def _to_host(x: Any) -> Any: if isinstance(x, torch.Tensor): # Keep strides: a few kernels depend on the caller's (e.g. column-major) layout. - return torch.empty_strided(x.shape, x.stride(), dtype=x.dtype, device="cpu").copy_(x.cpu()) if x.device.type != "cpu" else x.clone() + h = torch.empty_strided(x.shape, x.stride(), dtype=x.dtype, device="cpu").copy_(x.cpu()) if x.device.type != "cpu" else x.clone() + # Metadata tensors carry Python state (e.g. deep_gemm's paged sparse MQA FIRST/REUSE state and KV snapshot). + for k, v in getattr(x, "__dict__", {}).items(): + setattr(h, k, _to_host(v)) + return h if isinstance(x, (list, tuple)): return type(x)(_to_host(v) for v in x) if isinstance(x, dict): diff --git a/examples/neuron/probe_layer0.py b/examples/neuron/probe_layer0.py index acf1229..6cb03a8 100644 --- a/examples/neuron/probe_layer0.py +++ b/examples/neuron/probe_layer0.py @@ -22,8 +22,11 @@ import sys from pathlib import Path HERE = Path(__file__).resolve().parent -LAYERS = (0, 1, 2) # layer 2 only to catch layer 1's FFN hc_post (fused into layer 2's pre in DeepInfer) -COMPARE_LAYERS = (0, 1) +# PROBE_LAYERS=N probes layers 0..N-1 (default 2); one extra layer is captured to catch the previous layer's FFN hc_post +# (fused into the next layer's pre in DeepInfer). +_N_PROBE = int(os.environ.get("PROBE_LAYERS", "2")) +LAYERS = tuple(range(_N_PROBE + 1)) +COMPARE_LAYERS = tuple(range(_N_PROBE)) class Cap: @@ -264,7 +267,16 @@ def run_deepinfer(args): cap.put("kv_rope_quant", o.reshape(o.shape[0], -1)) wrap(attn_mod, "fused_norm_rope_hadamard_cast_store_kvcache", fused_out) - wrap(attn_mod, "prefill_attn_varlen_sparse", lambda o, a, k: cap.put("attn_core", o[0])) + def attn_in_out(o, a, k): + cap.put("attn_core", o[0]) + cap.put("attn_in.q", a[0]) + cap.put("attn_in.kv", a[1].reshape(a[1].shape[0], -1)) + cap.put("attn_in.idx", a[2].float()) + idx_len = k.get("index_length", a[5] if len(a) > 5 else None) + if idx_len is not None: + cap.put("attn_in.idx_len", idx_len.float()) + + wrap(attn_mod, "prefill_attn_varlen_sparse", attn_in_out) import tile_kernels.quant.swiglu_forward_kernel as swk orig_swiglu = mlp_mod.swiglu_forward_and_per_token_cast diff --git a/tests/neuron/device/test_device_primitives.py b/tests/neuron/device/test_device_primitives.py index 7295b87..f5a9d5b 100644 --- a/tests/neuron/device/test_device_primitives.py +++ b/tests/neuron/device/test_device_primitives.py @@ -67,6 +67,11 @@ def main() -> int: lambda: same((s.to(DEV).to(torch.int32) << 23).view(torch.float32), (s.to(torch.int32) << 23).view(torch.float32))) u = torch.tensor([1, 2**40 + 7, 2**63 + 5], dtype=torch.uint64) check("uint64 round trip", "known", lambda: same(u.to(DEV), u)) + i64 = torch.tensor([[2, 0, 1, 7]]) + m64 = torch.tensor([[True, True, True, False]]) + big = torch.iinfo(torch.int64).max + check("int64.max sentinel through torch.where (implicit int32 downcast)", "known", lambda: same( + torch.where(m64.to(DEV), i64.to(DEV), torch.full_like(i64.to(DEV), big)), torch.where(m64, i64, torch.full_like(i64, big)))) check("uint64 add", "known", lambda: same(u.to(DEV) + 3, u + 3)) # views / in-place semantics @@ -89,10 +94,19 @@ def main() -> int: def paged(t: torch.Tensor) -> torch.Tensor: # KV-cache style strided record view return t.as_strided((2, 129, 512), (902880, 528, 1)) + # Same layout without the e4m3 NaN encodings (0x7F/0xFF), so torch.equal compares real values. + nbuf = torch.where((buf & 0x7F) == 0x7F, torch.zeros_like(buf), buf) + nbufd = nbuf.to(DEV) + check("fp8 strided view -> bf16 cast (KV record layout)", "ok", + lambda: same(paged(nbufd).view(torch.float8_e4m3fn).to(torch.bfloat16), paged(nbuf).view(torch.float8_e4m3fn).to(torch.bfloat16))) + check("fp8 strided view -> fp32 cast", "ok", + lambda: same(paged(nbufd)[:, :7].view(torch.float8_e4m3fn).float(), paged(nbuf)[:, :7].view(torch.float8_e4m3fn).float())) check("uint8 strided view advanced indexing", "ok", lambda: same(paged(bufd)[blk.to(DEV), row.to(DEV)], paged(buf)[blk, row])) - check("fp8 strided view advanced indexing", "known", + check("fp8 strided view advanced indexing (incl. NaN bytes)", "known", lambda: same(paged(bufd).view(torch.float8_e4m3fn)[blk.to(DEV), row.to(DEV)], paged(buf).view(torch.float8_e4m3fn)[blk, row])) + check("fp8 strided view advanced indexing (no NaN bytes)", "known", + lambda: same(paged(nbufd).view(torch.float8_e4m3fn)[blk.to(DEV), row.to(DEV)], paged(nbuf).view(torch.float8_e4m3fn)[blk, row])) idx = torch.tensor([1, 1, 4]) check("index_add_ (duplicates)", "ok", lambda: same(torch.zeros(11, 512, dtype=torch.bfloat16, device=DEV).index_add_(0, idx.to(DEV), xd[:3]), torch.zeros(11, 512, dtype=torch.bfloat16).index_add_(0, idx, x[:3])))