Summary

This PR introduces the Neuron NKI stack into ZML and wires it into the LLM attention path.

The main goal is to make Neuron attention run through NKI kernels.
It adds the compiler/runtime plumbing needed to embed NKI kernels as Neuron custom-native calls, then uses that path for SDPA-style prefill/decode attention. Paged Attention will come later.

On inf2.8xlarge it runs :

• 25 tok/s for Llama 3.1 8B BF16
• 63.5 tok/s for LFM2.5-1.2B-Instruct BF16
• 19 tok/s for Qwen3.5-9B BF16

Neuron NKI compiler integration

Adds a new zml.ops.neuronNki(...) API that lowers a Zig tensor operation to a StableHLO custom call with:

• call_target_name = "AwsNeuronCustomNativeKernel"
• generated operand/result layouts
• input/output tensor signatures derived from ZML shapes and dtypes
• a compiled NKI backend_config attached to the custom call

The compiler bridge lives in:

• platforms/neuron/nki_kernel.zig
• platforms/neuron/nki_kernel_compiler.py
• platforms/neuron/nki_compiler_launcher.zig
• platforms/neuron/python_launcher.zig

At graph construction time, ZML resolves the NKI source file from Bazel runfiles, writes a JSON compile request, invokes the sandboxed nki-cc launcher, and reads back the base64 backend config expected by Neuron’s custom-native kernel path.

The compiler target is inferred from the actual Neuron instance through nrt_get_instance_info, mapping supported families to targets.

Neuron runtime

The Neuron platform packaging is updated.

The previous setup carried a local libneuronxla.zig shim and initialized Python directly from the PJRT proxy so ZML could expose the compiler pieces expected by Neuron. With libneuronxla v0.3, the AWS package now owns more of that compiler/runtime surface, free of python, so the sandbox is rebuilt around the packaged libneuronxla payload instead of ZML’s local replacement.

• adds nki>=0.4.0,<0.5
• includes aws-neuronx-tools, libgcc_s1, and libstdcpp6 required as needed by the so library
• exposes sandboxed neuronx-cc and nki-cc binaries
• moves Python initialization out of the PJRT shim and into reusable launcher code
• configures Neuron compiler/runtime env from the PJRT wrapper, including NEURON_CC_FLAGS, sandbox PATH etc.
• use Shardy instead of GSPMD

NKI attention backend

Adds a new attention backend:

• zml.attention.attention.Backend.nki
• selected automatically for .neuron
• exposed in examples/llm --backend=nki
• wired into Llama metadata/session handling

⚠️ We have an issue for Paged Attention, kernels will be added in another PR.

StableHLO/lowering compatibility fixes for Neuron

I added shims in a few Neuron-sensitive lowering paths:

• keep scalar integer broadcasts as scalar broadcasts before StableHLO emission
• preserve gather fill semantics by replacing inactive sentinel lanes before backend indirect-memory lowering
• preserve scatter drop semantics by replacing inactive sentinel indices and masking inactive additive updates to zero
• keep downstream resharding after the scalar gather fast path, preventing Neuron SPMD from hoisting an all-gather before a narrow dynamic slice

These shims are isolated in zml.ops.LoweringCompatibility so model code can continue expressing normal StableHLO semantics. The gather/scatter sentinel handling is related to upstream Neuron issue aws-neuron/aws-neuron-sdk#1335

Also for topk a specific branch has been added until the Neuron team fix aws-neuron/aws-neuron-sdk#1339 . And I asked for a in graph level data movement API: aws-neuron/aws-neuron-sdk#1340

Neuron profiling workflow

Adds a repo-level Neuron profiling workflow with documentation:

• --config=neuron-profile
• //tools/neuron:profile
• //tools/neuron:server
• //tools/neuron:ingest
• //tools/neuron:summary-json
• //tools/neuron:summary-txt
• //tools/neuron:summary-perfetto

Known follow-ups

• MoE NKI kernels
• Test kernels performance / adapt for trn2 and trn3

Related doc

• https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/deep-dives/nki-compiler.html
• https://github.com/aws-neuron/nkipy/blob/main/nkipy/src/nkipy/core/backend/hlo.py