4 Commits

Author SHA1 Message Date
whn09
4bc8ed3eaf Neuron: separate state regions, chunk-wide indexer context (issue #26)
- The state buffer's column regions get their own allocations on Neuron:
  a layer's uint8 SWA cache and fp32 compressor state are both region graph
  inputs, and aot_autograd compares offsets of differently typed views of one
  storage in their own element units (false overlap -> 'input mutations on
  views with different dtypes').
- paged_mqa_logits chunk path: gather the keys the chunk's last token sees
  (lens.max()), not the first token's.
- Debug: DEEPINFER_NEURON_ALIAS_DEBUG prints aliased graph inputs; a cProfile
  of decode step 8 under DEEPINFER_NEURON_REGION_TIMING.

Tiny V4.1 on a NeuronCore, bucket 128: short 8/9 (= eager), long 715 tokens
in 6 chunks 9/9 (eager 8/9) in 26 s instead of 978 s of eager prefill.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 03:46:56 +08:00
whn09
c24574bfd6 Neuron: full DeepSeek-V4.1-Flash load + decode region graphs at EP=64 (issues #21, #25)
- Layer-major KV / state layout so per-layer views are contiguous graph inputs
  (the executor copied non-contiguous inputs every call).
- One compiled entry per (class, method), rank as a registered buffer, no
  layer_idx guard: every layer reuses the same graphs and loaded NEFFs.
- fp8 weight rescale on the host (numpy LUT, ~2 min instead of 35+ min of
  eager per-shape compiles); engram sf packing on the host.
- int16 ue8m0 packing in fp32 (int32 interleaved slices narrowed to int16
  crash neuronx-cc, NCC_IBCG902).
- Token embedding and the text model's vision tower stay in host memory on
  Neuron (2.5 GiB of HBM a rank).
- offline example: --max_num_seqs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 01:39:22 +08:00
whn09
02e099fff8 Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24)
- NeuronGraphRunner: the CudaGraphRunner contract over torch.compile(dynamic=False), one graph per
  batch size (DEEPINFER_ENFORCE_EAGER=0, backend DEEPINFER_NEURON_COMPILE_BACKEND). Layer cache views are
  re-sliced from their root allocations inside the graph, so AOT never regenerates aliased mutated
  inputs with as_strided.
- Decode MoE with fixed shapes: packed all-gather, dense local experts, fp32 reduce-scatter
  (neuron_ops/static_moe.py, deep_gemm.m_grouped_fp8_fp4_gemm_nt_shared_a).
- Graph-break and compile fixes: hide torch.cuda during compile, torch.device equality for FX codegen,
  hash-based force_random routing, Engram hash/fetch on the host (exact int64), env vars resolved before
  tracing, Movable annotation, no .item()/nonzero/as_strided writes on the decode path, KV stores route
  padding slots to a scratch block kept past the pool, native-dtype gather under compile.
- fp8: trn2 reads float8_e4m3fn as IEEE e4m3 (max 240). Activations/KV quantize to 224
  (DEEPINFER_NEURON_FP8_MAX), checkpoint fp8 weights are rescaled at load (bytes / 2, ue8m0 + 1), and
  neuronx-cc gets --experimental-unsafe-fp8e4m3fn-as-fp8e4m3 via --internal-hlo2tensorizer-options.
  The uint8-carrier backend (neuron_ops/neuron_compile.py, "deepinfer_neuron") stays as a fallback.
- Tests: tests/neuron/compile; run_all.sh runs kernel unit tests at the e4m3fn range (448) and the
  end-to-end stages at 224. Docs: vllm-neuron PR #71 analysis and compile notes.

Host: tiny model 9/9 vs golden (all prompts) eager, compiled path (aot_eager) matches eager; CPU suite green.
Device: the whole decode graph still crashes neuronx-cc (NCC_IINL902, being bisected).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-22 11:59:37 +08:00
Jiashi Li
d585e6542c initial commit 2026-10-20 12:09:42 +08:00