3 Commits

Author SHA1 Message Date
whn09
3629d113fb Neuron: prompt chunks through the decode region graphs (issue #26)
- PrepareArgs.prompt_as_decode: decode-layout metadata for a prompt chunk;
  input_ids are the prompt's own tokens instead of last token + MTP drafts.
- DEEPINFER_NEURON_PREFILL_BUCKET adds one graph size; the model runner pads
  a prompt chunk to it with the decode placeholders (every EP rank takes the
  same bucket, so the static collectives match). Sampling still sees a prompt
  step.
- nki_attention.paged_mqa_logits: one gather + one kernel call for a chunk
  (one sequence) instead of a T-way loop.
- bf16 linear for the LM head eagerly on a NeuronCore (b.float() was a
  2.6 GiB copy a call); DEEPINFER_NEURON_REGION_TIMING; --num_prompts.

Host mode, tiny V4.1, chunk 128: same tokens as eager prefill (short 9/9,
long 715 tokens in 6 chunks 8/9 against golden).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 03:19:56 +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