8 Commits

Author SHA1 Message Date
whn09
f67fab6bdd Neuron MoE: two-level all-reduce (device, then across devices) (#27)
Microbenchmark at 64 ranks, [1, 5120] fp32 in a compiled graph: world
all-reduce 3.37 ms, 8 consecutive ranks 0.41 ms, 4-core device group then a
16-device group (one core each) 2.01 ms; a strided 8-device group is rejected
by the runtime. The replicated-token MoE all-reduce (one a layer) now takes
the two-level form (DEEPINFER_NEURON_HIER_ALLREDUCE=0 restores one world
all-reduce). Host tiny cp=8 matches golden.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-24 11:05:25 +08:00
whn09
d8f7a2e898 Neuron: PR #147 OpenAI server (+ /v1/completions, Neuron worker setup), offline bench, MoE fixes (#26, #27, #29)
- serve/: the minimal OpenAI-compatible server of deepseek-ai/DeepInfer#147,
  with a Neuron worker setup (basic mode env, per-rank NeuronCore, fixed KV
  pool via --num-kv-blocks, no MTP lookahead) and a /v1/completions endpoint
  (prompt text or token ids, no chat template): what vllm bench serve
  --backend openai drives, as vllm-neuron#71's run_bench.py does.
- examples/neuron/bench_neuron.py: the same TTFT/TPOT matrix offline.
- static MoE: one packed fp32 all-gather instead of three (each 64-rank
  collective is ~2 ms of fixed latency a layer); a capacity-grouped expert
  path for prefill chunks (DEEPINFER_NEURON_MOE_CAPACITY, dense rerun on
  overflow; not yet validated on device).
- Sub-region timings at split() boundaries under DEEPINFER_NEURON_REGION_TIMING.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 16:23:28 +08:00
whn09
32abff6662 Neuron: skip the vacuous candidate-block filter; mask padding rows in static MoE (#26)
- When runtime_max_model_len <= candidate_topk_blocks * candidate_block_size,
  the candidate blocks hold every position, so the filter keeps them all and
  the dense indexer logits give the same top-k. On Neuron the candidate path
  gathers all candidate keys per query token: [128, 16384] rows a layer for a
  prompt chunk, the NEFF that still ran HBM out in rg16.
- static_moe_forward takes the decode metadata's valid_token_mask: padding rows
  select no expert.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 05:25:28 +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
9c8c85d83d Neuron: per-layer region graphs + vllm-neuron#71 NKI kernels (issues #10, #11, #13, #25)
- Decode runs eagerly and compiles per-layer regions (attention half, MoE) with stage boundaries
  (neuron_ops/regions.py, DEEPINFER_NEURON_GRAPH_MODE=regions); layer KV views are sliced from the
  roots inside the graph (_LazyKVCache).
- vllm-neuron#71 NKI kernels (Apache-2.0, vendored under neuron_ops/nki/ with attribution), called via
  torch_neuronx.nki_hop.wrap_nki: blockquant (fp8/fp4 GEMMs incl. grouped einsum, MoE and shared
  expert on bf16 activations), sparse_attn / index_score / compress_pool (nki_attention.py),
  hc_mixes / hc_expand / hc_collapse / hc_attn_out / rmsnorm / swiglu (nki_hc.py, nki_norm.py).
- Compile-safe formulations: row-major scales, no dtype bitcasts of graph inputs (uint8 weight
  tensors), whole-row KV stores, byte assembly in int32 narrowed once at the store, arithmetic int16
  scale packing, plain fma/rsqrt and bf16 F.linear under compile; default backend deepinfer_neuron
  (torch fp8 carried as uint8, avoiding NCC_EOCP001 with the kernels' OCP fp8).
- Host: tiny model 9/9 vs golden on all prompts; device: kernel tests pass (13 attention, 88 hc/norm);
  region decode graph compiles except the attention projection stage (being reduced).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-22 17:36:54 +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
whn09
ed49792a0c Add AWS Trainium (Neuron) platform, basic mode — part 1
Run DeepInfer on the native torch-neuronx device (PrivateUse1 `neuron`,
torch-neuronx 2.14, no torch-xla), following the Ascend port's structure.

- deep_cluster: Neuron discovery (one DeepInfer device = one logical
  NeuronCore, NEURON_RT_VISIBLE_CORES, NUMA from neuron-ls).
- PlatformNeuron: takes the Ascend (is_npu) kernel layouts, adds is_neuron;
  DEEPINFER_NEURON_DEVICE=cpu runs the same code on the host for
  correctness work.
- neuron_ops/transfer_to_neuron: map "cuda" devices, torch.device and
  torch.cuda.* onto neuron:0 (counterpart of torch_npu transfer_to_npu).
- C++ extension: -DWITH_NEURON host-only build; torch-allocated buffers and
  tensor-view slot copies instead of raw device-pointer copies.
- Pure-torch drop-in vendor packages (Ascend API semantics): deep_gemm,
  flash_mla, deep_select, tile_kernels, tilelang (constants only).
- tilelang_ops/neuron backend: KV cache / metadata / norm / compress kernels;
  CP, ZOS and VL entry points raise NotImplementedError.
- Basic mode defaults on Neuron: eager, no DSpark/MTP drafting, no
  zero-overhead scheduling; MoE avoids the GPU-only mega kernel; mHC
  prenorm GEMM uses the Ascend single-split reduction.

Not yet included: deep_ep, the sampling kernels, end-to-end harness.

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