Run FP8 models on AMD GPUs (ROCm/HIP) — including Strix Halo (Ryzen AI Max, RDNA 3.5 / gfx1151). RDNA 3.5 has no FP8 tensor instructions: torch._scaled_mm (PyTorch's FP8 GEMM) is implemented only on CDNA3 (MI300, gfx94x) and RDNA4, so an FP8-trained or FP8-quantized checkpoint has no native e4m3 matmul to run on the iGPU. This package fills that gap with bit-exact FP8 emulation in pure PyTorch tensor ops (log2/round/exp2/clamp), so the same model that needs an NVIDIA Hopper GPU can run — numerically faithfully — on an AMD APU.

It is emulation, not hardware FP8 (torch._scaled_mm is confirmed MI300+/Hopper only — it raises on gfx1151). On CDNA3 / RDNA4 (native GPU FP8) the quantizer is the seam to swap for a real torch._scaled_mm.

Fp8TELinear takes compute_dtype controlling the emulated GEMM:

• torch.bfloat16 (default) — scale/round/matmul in bf16. Uses the iGPU's matrix cores (~15× faster GEMM than fp32) and is the faithful path for FP8-trained checkpoints whose scales were calibrated on a bf16 forward. On Evo 2 1B this is measurably more accurate end-to-end (78.7% vs 68.7% for the older mixed-precision emulation, ~0.9pp off the H100 reference) — fp32 hurts such checkpoints (38.6%).
• torch.float32 — bit-exact vs a native FP8 GEMM (H100/MI300 accumulate in fp32), ~15× slower on this iGPU. Use for hardware-FP8 fidelity rather than the bf16 checkpoint path.

The scaled, e4m3-rounded weight is precomputed once at load (not re-quantized every forward).

This is the AMD/ROCm sibling of FP8-MPS (Apple Silicon); it is the FP8 layer that Evo2StrixHalo builds on, the same way FP8-MPS backs the evo2Mac port.

FP8 checkpoints ship in two incompatible on-disk conventions; this package handles both:

pip install -e .          # from a checkout
python tests/test_fp8.py  # all should pass

This needs a ROCm PyTorch build. For Strix Halo (gfx1151) install the AMD nightly wheels, e.g. pip install --index-url https://rocm.nightlies.amd.com/v2/gfx1151/ torch.

import torch
from fp8_rocm import quantize_e4m3, Fp8TELinear, Fp8PTQLinear, best_device

dev = best_device()  # "cuda" on a ROCm GPU (AMD reports through the CUDA API)

# Bit-exact e4m3 rounding that runs on the AMD GPU:
q = quantize_e4m3(torch.randn(4, device=dev))

# TE format: bf16 weight + per-tensor act/weight scales from the checkpoint:
lin = Fp8TELinear(weight_bf16, bias, act_scale=50.0, weight_scale=1500.0).to(dev)

# PTQ format: pre-quantized e4m3 weight + scale_inv (per-tensor or per-block):
lin = Fp8PTQLinear(weight_fp8, weight_scale_inv, bias, block=128).to(dev)

quantize_e4m3 is verified bit-exact against torch.float8_e4m3fn across 100k random values (it saturates above 448 rather than producing NaN, matching what FP8 GEMM paths expect after their pre-scale clamp). quantize_e5m2 is provided too.

Swap a whole model's FP8 layers in place, then run it on the AMD GPU:

from fp8_rocm import apply_ptq_emulation, apply_te_emulation

# PTQ models (Nemotron, DeepSeek-V3, Qwen3-FP8, …): auto-detects e4m3 weights
# with a sibling weight_scale_inv and replaces them.
n = apply_ptq_emulation(model)            # returns # layers swapped

# TE models (bf16 weight + per-tensor scales recovered from _extra_state):
n = apply_te_emulation(model, scales)     # scales = {path: {"act":…, "weight":…}}

Two scripts check the emulation against ground truth:

• scripts/validate_te_evo2.py [--model evo2_1b_base] — self-contained per-layer check: replays each FP8 layer's TE GEMM with Fp8TELinear on the AMD GPU and diffs it against PyTorch's native float8_e4m3fn (the H100 math), using the real checkpoint scales. No full model needed.
• scripts/validate_evo2_h100.py --model evo2_1b_base (and evo2_7b) — full next-token loss/accuracy on the bundled prompts, compared to the published H100 FP8 numbers from Evo 2's test_evo2.py. Shows the bf16-fallback floor vs the e4m3-emulated result for the FP8-trained 1B, and that the bf16-robust 7B lands on the reference without emulation.

quantize_e4m3 is bit-exact vs torch.float8_e4m3fn; both linear formats and both adapters are unit-tested. On Strix Halo this is the FP8 path under Evo2StrixHalo (whose bundled fp8_emulation.py this package generalizes), and the validation mirrors the Apple-Silicon FP8-MPS: Evo 2 1B per-layer ≈1.5–2.2e-3 vs native e4m3, and the e4m3 emulation recovers full-model next-token accuracy to 78.68%, within 0.87 pp of the H100 reference (79.556%), up from a 32.6% bf16-fallback floor.

Scope: this works for models whose FP8 footprint is the linear layers (the 1B's input projections). It does not make Evo 2's 20B/40B usable: those are FP8 across the whole operator stack, and emulating their 117 linears — confirmed against a CPU native-float8_e4m3fn reference — leaves them at chance. Their FP8 dependence is in the attention/fused stack, not the linears (see docs/PAPER.md §5.6). Such models need true FP8 hardware.