4 Commits

Author SHA1 Message Date
whn09
8a5209ee65 Neuron basic mode — part 3: end-to-end check vs the official reference
- examples/neuron: tiny random V4.1 checkpoint in the real HF layout,
  pure-torch replacements for the official reference kernels, golden
  generation, a teacher-forced DeepInfer runner (LLM offline API) and a
  comparison tool. CPU mode: 25/27 greedy tokens match, worst logits
  cos 0.9985 / rel 0.056 against a reference self-noise floor of 0.032-0.060.
- model_loader: check parameters against the platform device, not is_cuda.
- deep_gemm paged MQA logits: do not dereference block-table slots past the
  context.
- bind_cpu: skip in Neuron CPU mode (1 core per NeuronCore share would
  serialize host compute).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-20 20:20:02 +08:00
whn09
f4144fc3fc Neuron basic mode — part 2: deep_ep, sampling kernels, fixes
- vendor/deep_ep: EPBuffer dispatch/combine with torch.distributed uneven
  all_to_all (host-synced split sizes), EngramBuffer sharded lookup,
  BucketBuffer all_gather (CP-only reduce paths stubbed); verified against
  the DeepEP-Ascend references at world size 1/2/4.
- tilelang_ops/neuron sampling: random_sampling (asc two-level fixed-point
  sampler, Philox), fused_sample (logprobs / do_sample), sampler metadata,
  sum_smaller_probs, structured-output bitmask; DSpark/MTP/varlen kernels
  raise NotImplementedError. All 42 asc *_tl factories now have a Neuron
  counterpart.
- vendor/flash_attn: SDPA varlen stand-in for the ViT (mirrors npu_ops).
- transfer_to_neuron: keep `with torch.device(...)` working on patched
  factories; fix manual_seed recursion in host mode.
- engram: gather the packed ue8m0 sf as bytes on Neuron.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-20 18:57:27 +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