coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-28 23:07:09 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/deepinfer/model_executor/layers/tilelang_ops/neuron
History
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
..
__init__.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
_unsupported.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
candidate_blocks.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
compress.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
compute_seqlens_k_zos.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
dflash_copy.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
dflash.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
flash_mla_utils.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
fused_sample.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
fused_transform.py Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24) 2026-10-22 11:59:37 +08:00
gather_engram_tokens_zos.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
gather_input_ids_zos.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
gather.py Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24) 2026-10-22 11:59:37 +08:00
get_kv_index.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
image_embed_ops.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
kv_index_cat.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
kv_quant.py Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24) 2026-10-22 11:59:37 +08:00
load_paged_kv_cache.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
norm.py Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24) 2026-10-22 11:59:37 +08:00
prepare_cp_mirror.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
prepare_v4_decode_metadata.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
random_sampling.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
rstd_1ulp.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
sampler.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
schedule_mtp_varlen.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
state_cache.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
store_kv_cache.py Neuron: static decode graphs, static-shape MoE, trn2 e4m3 fp8 (issues #9, #12, #24) 2026-10-22 11:59:37 +08:00
store_mtp_info.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
structure_output.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
sum_smaller_probs.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
swa_index.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00