mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +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>
67 lines
5.1 KiB
Markdown
67 lines
5.1 KiB
Markdown
# DeepInfer on AWS Trainium (native torch-neuronx, no XLA)
|
|
|
|
## Stack
|
|
- torch 2.14 + torch-neuronx 2.14.3 (PrivateUse1 device `neuron`, eager + `torch.compile(backend="neuron")`),
|
|
Neuron SDK 2.33 (runtime 2.35). No torch-xla, no libtorch-neuronx-lite.
|
|
- Target: trn2.48xlarge (64 logical NeuronCores, LNC=2, 24 GB HBM per core), Neuron DLAMI.
|
|
- Python env: a venv layered on the DLAMI's `/opt/aws_neuronx_venv_torch_neuronx_2_14_0` (add its `bin/`
|
|
to `PATH` so `neuronx-cc` is found), plus `requirements.txt` minus tilelang / tile-kernels / torch.
|
|
- Build: `DEEPINFER_ACCELERATOR=neuron ./scripts/build.sh build` (host-only C++ extension, `-DWITH_NEURON`).
|
|
|
|
## Platform design
|
|
- `PlatformNeuron` (deepinfer/platform.py): `is_npu` is **true** (Neuron follows the Ascend kernel
|
|
layouts / code paths), `is_neuron` is true, `tilelang_ops_backend = "neuron"`.
|
|
- `deepinfer/neuron_ops/transfer_to_neuron.py` maps `"cuda"` devices and `torch.cuda.*` onto the target
|
|
device. `DEEPINFER_NEURON_DEVICE=neuron|cpu` selects the target; `cpu` runs everything on the host.
|
|
- `deepinfer/neuron_ops/vendor/` is put at the front of `sys.path` (deepinfer/__init__.py) so
|
|
`import deep_gemm`, `flash_mla`, `deep_ep`, `deep_select`, `tile_kernels`, `tilelang` resolve to the
|
|
pure-torch packages there.
|
|
- `deepinfer/model_executor/layers/tilelang_ops/neuron/` is the kernel backend: every `xxx_tl` factory
|
|
that `asc/` provides, same name and signature, returning a callable with the same call signature
|
|
and the same in-place output semantics.
|
|
|
|
## Scope: basic adaptation only
|
|
Optional optimizations stay OFF on Neuron (defaults set in deepinfer/__init__.py): CUDA graphs (eager),
|
|
DSpark/MTP drafting (`DEEPINFER_USE_NEXTN_DECODE=0`), zero-overhead scheduling, MTP varlen, EPLB, CP>1,
|
|
YOCO bounded replay, batch-invariant mode, GPU-only mega kernels, VL/vision. Kernels reachable only from
|
|
those features: export the name with the right signature but raise
|
|
`NotImplementedError("<name>: not supported on Neuron (basic mode)")`. Fully implement what the model
|
|
architecture needs (MLA attention, DSA indexer/top-k, compressor, mHC, Engram, MoE with FP8/FP4 dequant,
|
|
norms/rope, KV cache store/load, basic sampling).
|
|
|
|
## Rules for kernel code (vendor packages and tilelang_ops/neuron)
|
|
1. Pure torch, device-agnostic: allocate on the input tensors' device; never hardcode `cuda`/`npu`/`neuron`.
|
|
Must run with tensors on CPU (correctness now) and on `neuron` (later).
|
|
2. Semantics follow the **Ascend** variant (DeepGEMM-Ascend, DeepEP-Ascend, `tilelang_ops/asc/`, the
|
|
`PLATFORM.is_npu` branches); where Ascend has no variant, follow the CUDA one.
|
|
3. Output semantics must match exactly: if a kernel writes into caller-provided tensors, write in place
|
|
(`out.copy_(...)`, indexed assignment); keep the caller-visible shape/dtype/stride contract.
|
|
4. Prefer static shapes and no host syncs (`.item()`, `.tolist()`, `nonzero`, boolean-mask indexing) in
|
|
per-step model code — the decode step will later be compiled to one static graph per batch bucket.
|
|
Where a data-dependent shape is unavoidable, isolate it and leave a `# NEURON: host sync` comment.
|
|
5. Neuron dtype notes: prefer int32 over int64 for indices on device; fp8 is `torch.float8_e4m3fn`; FP4 (e2m1)
|
|
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=<file>` 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):
|
|
```
|
|
export DEEPINFER_ACCELERATOR=neuron DEEPINFER_NEURON_DEVICE=cpu
|
|
```
|
|
Compare against the upstream references (`*_ref` functions next to each `tilelang_ops` interface,
|
|
`tile_kernels.torch`, and the reference implementations inside `vendor/<Repo>/tests`).
|