dstair/qwen3.5-moe-fp8-trainium2

Qwen3.5-35B-A3B FP8 on AWS Trainium2 — loader + MoE kernel

Loader and NKI kernel code for running the official Qwen3.5-35B-A3B FP8 checkpoint on AWS Trainium2 via PyTorch Native.

This repository does not contain model weights. It contains the code that makes the existing Qwen FP8 release usable on Trn2. Download the weights directly from Qwen:

• Weights (FP8): https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8
• Base model (BF16): https://huggingface.co/Qwen/Qwen3.5-35B-A3B

Both are Apache-2.0.

Why this exists

Qwen's official FP8 checkpoint stores expert weights as OCP E4M3FN bytes with 128×128 BF16 block scales (weight_scale_inv). Two things make it non-trivial to run on Trainium2:

1. Trn2's Tensor Engine does not accept E4M3FN as an nc_matmul stationary dtype, and neuronx-cc rejects torch.float8_e4m3fn on Trn2 (NCC_EVRF051). Trn2's matmul wants legacy float8_e4m3, which torch does not expose as a dtype.
2. The Native DLC image ships neither safetensors nor transformers, so the FP8 bytes and their block scales must be read and decoded by hand.

This code closes both gaps.

Contents

File	What it is
st_reader.py	Dependency-free .safetensors reader (mmap + manual header parse). Keeps F8_E4M3 tensors as raw uint8 bytes because the DLC has no FP8 torch dtype and Trn2 needs legacy E4M3, not E4M3FN.
moe_w8.py	CPU-side block-W8 weight prep for the block_ob_coalesced path: decode E4M3FN (decode_e4m3fn), encode/decode Trn2 legacy E4M3 (encode_legacy_e4m3 / decode_legacy_e4m3, LEGACY_E4M3_MAX = 240.0), block-scaled dequant (dequantize_official_fp8, BLOCK_SIZE = 128), per-output-block requantization + coalesced scale packing (requantize_official_fp8_output_block, pack_coalesced_block_scales), and routing (build_local_affinities). Pure torch.
kernels/moe_fused_w8_35b.py	Fused all-expert MoE NKI kernel (gate/up GEMM → SwiGLU → down GEMM → affinity scale → expert sum) that never materializes [experts, batch, hidden]. Converts weight bytes tile-by-tile: E4M3FN codes < 0x78 via the legacy-E4M3 converter, the finite E4M3FN-only codes 0x78..0x7e patched exactly.
kernels/moe_fused_w8_35b_ops.py	torch.ops.moe_w8.* registration for the two coalesced variants; fused_fp8_block_coalesced_ob (per-output-block PSUM post-scale) is the fastest.
examples/prepare_fp8_moe_weights.py	Runnable example: official FP8 checkpoint → the six kernel operands (one layer, one TP rank) + routing affinities; --run-kernel calls the op end-to-end on Trn2.

Example

examples/prepare_fp8_moe_weights.py runs the full path — official FP8 checkpoint → kernel operands (w8_gate_up/w8_down + coalesced scale tables) → routing affinities, with an optional on-device kernel call:

# weight prep only (CPU; needs the FP8 checkpoint + this repo on PYTHONPATH)
python3 examples/prepare_fp8_moe_weights.py \
    --checkpoint /path/to/Qwen3.5-35B-A3B-FP8 --layer 0 --world-size 8 --rank 0

# end-to-end on a Trn2 device (inside the PyTorch-Native DLC)
python3 examples/prepare_fp8_moe_weights.py \
    --checkpoint /path/to/Qwen3.5-35B-A3B-FP8 \
    --layer 0 --world-size 8 --rank 0 --run-kernel --batch 128

The kernel computes only this rank's local experts (expert-parallel); in a real serving loop you all-reduce the [B, H] output across the TP group, then add the shared expert. The router and shared expert stay in the surrounding graph.

The FP8-on-Trn2 recipe (what took iterations to get right)

• Read FP8 as bytes, not as a float8 dtype. st_reader.py maps F8_E4M3 → uint8.
• Decode E4M3FN from the bit layout (decode_e4m3fn) rather than relying on a torch FP8 dtype — E4M3FN uses exponent 15 for finite values through 448; only exponent=15, mantissa=7 is NaN.
• Re-encode to Trn2 legacy E4M3 for the Tensor Engine (encode_legacy_e4m3), which is finite through 240 and rejects the reserved-exponent Inf/NaN codes (validate_legacy_e4m3_bytes). Inside the kernel, weights are stored as nl.int8 and .view(nl.float8_e4m3) so the HLO operand stays int8 (dodges NCC_EVRF051).
• Apply the 128×128 BF16 block scales on dequant (dequantize_official_fp8).
• The router and shared expert stay in the surrounding PyTorch graph; only the local routed experts are fused.

⚠️ FP8 here is a memory / capacity lever (it roughly halves resident expert bytes, unlocking higher batch at long context), not a single-stream decode-latency win. See the source repo's notes.

Requirements

AWS Trainium2, a pre-GA PyTorch Native DLC, and (for the kernel) nkilib. CPU-side conversion utilities in moe_w8.py / st_reader.py run anywhere with PyTorch.

Source

Developed as part of a full Qwen3.5-35B-A3B reference implementation: https://github.com/dstair/neuron-pytorch-native (contrib/qwen3.6-35b-a3b).

License

Apache-2.0. See LICENSE and NOTICE. Original code only; no model weights are redistributed. Model weights are © the Qwen team, released under Apache-2.0 at the links above.