Sampler._verification runs as one compiled region for greedy batches without structured output: the native
sampling structs enter as Python mirrors (neuron_ops/sampling_mirror.py), the flags are read once per batch
eagerly (prime_host_flags), and greedy mirrors carry no seeds / positions. Untested on device yet.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- fused_sample: _host_flags/_greedy_verify flags persist while the batch's requests are unchanged
(each device read drained the queue and waited for the whole forward). Tiny golden 9/9; full model
step time unchanged (660 ms), driver dflash launch 71 -> 10.5 ms host.
- bind_cpu on Neuron: the device's NUMA node (DEEPINFER_NEURON_CPU_BIND=node|core|legacy); the generic
split gave each rank one CPU and ranks 0-15 the wrong node. Neutral on TPOT (228.2 ms).
- profiled followers exit through sys.exit so the runtime writes their trace.
- op_census: trace_syncs / census_sites (DEEPINFER_NEURON_TRACE_SYNCS=1|2|3).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Both drain the device stream in torch-neuronx (a cast round-trips through the
host, a scalar operand is a host->device copy): int32 ids index directly and
out-of-range rows are cleared with masked_fill. Tiny cp=2 nextn=5 on device:
9/9 golden tokens.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- fused_sample._host_flags: the top-p/top-k and greedy flags of the step's
sampling parameters are read from the device once a step (a DSpark step
samples the same tensors six times); reset_host_flags at each sample step.
- NEURON_RT_INSPECT_* stays on the driver (64 ranks would overwrite one
output directory).
- VocabParallelEmbedding.ids_in_range: one F.embedding instead of the
zero-filling gather (~5 eager NEFF launches).
- bench33 (512 in, TP=64, nextn=5): TPOT 272.6 -> 268.2 ms.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- fused_sample reject mode for greedy rows (accept iff the draft is the
argmax), get_accept_result / finalize_verification (static final rows),
store_mtp_info (static pool writes, slot 0 kept), confidence_head_fused,
dflash_input_store, get_fused_dflash_get_input: the frontend refs with the
kernels' argument lists.
- CP followers run the draft forward (its MoE has collectives); the driver's
accepted offsets reach them through one all-reduce on Neuron (no eager
broadcast), the request first rows through the graph metadata.
- run_deepinfer_tiny.py --nextn: free-running greedy with drafts. Host tiny,
nextn=5, eager: short 9/9 golden tokens; long matches up to its known near
tie at step 5.
- bench17: the two-level MoE all-reduce does not move TPOT (433 ms).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The bisection is 32 rounds over the vocabulary: ~0.5 s a decode step at 129280
even compiled, and the default parameters (top_p 1, top_k -1) filter nothing.
Such batches take the plain temperature log-softmax: 529.6 -> 1.3 ms on one
core, same tokens, logprobs within 1e-9.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The op census of a decode step (DEEPINFER_NEURON_REGION_TIMING) found 85 eager
aten ops in the forward outside the region graphs but 1905 in sampling (the
top-p/top-k bisection and the int64 philox draw), ~0.6 s of NEFF launches.
- The log-softmax / top-p / top-k part runs as one compiled graph: 1.4 ms
instead of 46.7 ms on one core for the [1, 4096] force_topk row, logprobs
within 1e-9.
- Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so
the two-level draw returns their argmax; an all-greedy batch takes the argmax
from the graph and skips the draw. The philox draw stays eager for sampled
rows: neuronx-cc graphs reject 64-bit constants outside the int32 range
(NCC_ESFH001), an int32 philox is #3.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
torch.remainder ran ~140 ms a call eagerly on the device and was 554 of the
640 ms a decode step spent sampling. a - (a // b) * b gives the same result for
the non-negative operands (checked on 100k random values).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- 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>
Fixes found by reducing neuronx-cc NCC_IINL902 crashes to minimal reproducers (crash rates over repeated
compiles, the compiler is nondeterministic on these graphs):
- uint8 byte assembly of narrowed values (KV records, e2m1/e4m3 codes, ue8m0 exponents) crashed or
silently corrupted bytes: producers keep int32 codes under compile and stores narrow once
(fused_transform, kv_quant, store_kv_cache; neuron_compile.encode_codes).
- int16 packed scales written through a uint8 view and read back in the same graph: the Neuron kernel
now receives the int16 storage and packs byte pairs arithmetically (_alloc_quant, return_quant).
- fp8 legalization pass treated plain float->bf16 conversions as fp8 carriers (wrong values + crash).
- Projection stage split further (nullcontext instead of a disabled CP session).
Device (tiny model, regions mode): every decode graph compiles; argmax vs golden short 9/9,
medium 8/9 (step 6 near-tie, margin 0.065), long 9/9, cos 0.997-0.9995.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- 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>
- 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>
Device-memory fixes found bringing up the released checkpoint (16.45 GiB of
weights per 24 GB core):
- Engram: the Neuron EngramBuffer keeps only this rank's SF rows, so skip the
all-gather of the replicated SF table (~3 GB per layer, allocated on the
device -> "Failed to allocate resource").
- gather_tl: index_select / advanced indexing of an int8 table allocates
several times the table on the device (+4.9 GiB for the 1.3 GB embedding);
gather through F.embedding on an int32 / fp16 view instead (bit-exact, no
extra memory).
- deep_gemm FP8/FP4 GEMMs: dequantize the weight in N chunks of 32M elements
and one expert at a time instead of whole projections / all local experts
in fp32 (OOM on a fragmented pool); bit-identical results.
- transfer_to_neuron: import torch.utils._device explicitly (missing when
torch_neuronx is not auto-loaded).
examples/neuron/offline_inference_neuron.py: offline inference in basic mode
(text prompts, no MTP / EPLB, fixed KV pool, max_model_len >= 2048 * 8 for
the candidate-block filter), one process per NeuronCore.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- 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>
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>
