[gfx1201] Mistral-3 + Qwen3-8B-FP8 on RDNA4 via native triton attention#811
Merged
Conversation
4 tasks
carlushuang
changed the title
…iton attention backend Bring up two natively-FP8 8B models on AMD RX 9070 XT (gfx1201) with an all-Triton stack and cudagraph on. Auto-routed on gfx1201; opt-in elsewhere via ATOM_NATIVE_TRITON_ATTN=1. Models - New Mistral3 text-only model (strips Pixtral vision tower) - Qwen3-8B-FP8 (block-128) via gemm_a16w8_blockscale (PREQUANT=False); the fp8 weight is cast to bf16 inside the kernel so tl.dot(bf16, bf16) is used (Triton on this gfx1201 build does not implement tl.dot(fp8, fp8)) Attention - New NativeTritonBackend: JIT triton kernels in place of AITER prebuilt HIP .so (gfx1201 has no prebuilt code objects) - CUDAGraph-correct decode at all bs incl. NaN-from-padding fix in prepare_decode (was silently producing wrong logits for scheduled_bs < captured_bs) GEMM / activation / FP8 - Per-shape gemm_a8w8 configs (GROUP_SIZE_M=1) + per-token FP8 quant (1 kernel, no atomic) via aiter PR #3168 - SiluAndMul routed through aiter fused_silu_mul (PR #2578, merged) - Removes torch fallbacks across linear/embed_head/sampler/paged_attention (every fallback contained .item()/.cpu() syncs that silently break cudagraph capture on ROCm) - Critical: aiter dtypes.d_dtypes[fp8] == torch.uint8. FP8 weights MUST be .view(torch.float8_e4m3fn) before .to(bf16) or byte values 0-255 decode as integers and outputs look numerically reasonable but are garbage Config / setup - scripts/gfx1201/setup_aiter_configs.sh aliases gfx1250 GEMM tuned configs to gfx1201 names (aiter ships zero gfx1201 configs in rocm/atom-dev:latest); without this the autotuner falls back to a default that is 30-50% slower at 8B-class shapes - scripts/gfx1201/gemm_a8w8_sweep.py for shape-by-shape tuning - recipes/Ministral-3-8B.md + recipes/Qwen3-8B-FP8.md with serve cmd, env vars, perf+accuracy table, debug journey notes Results (single GPU, BF16 KV, cudagraph, conc=1, OSL=256) - Ministral-3-8B: TPOT 22.0 ms, 43 tok/s, gsm8k 5-shot n=200 = 0.785 - Qwen3-8B-FP8: TPOT 21.6 ms, 45 tok/s, gsm8k 5-shot n=50 = 0.86 Depends on aiter PRs #3168 (gfx1201 gemm_a8w8 tuning configs) and #2578 (silu_mul_fused, already merged upstream).
…m_batched_tokens)
carlushuang
force-pushed
the
carhuang/support_gfx1201_mistral3_rebased
branch
2 times, most recently
from
May 18, 2026 00:50
df15f7c to
a988bf8
Compare
carlushuang
force-pushed
the
carhuang/support_gfx1201_mistral3_rebased
branch
from
May 18, 2026 00:51
a988bf8 to
f2b1551
Compare
…201_mistral3_rebased
…zing The block_tables write at backends.py:268 overflows its column dim when a sequence accumulates more blocks than max_num_blocks_per_seq permits. The prior fix in e7a7ac5 sized that buffer with max(max_model_len, max_num_batched_tokens) so any recipe with the default max_num_batched_tokens=16384 would be safe regardless of max_model_len. Switching to a recipe-side fix instead: pin max_model_len >= the default max_num_batched_tokens (both = 16384) in the gfx1201 Mistral-3 / Qwen3-8B recipes. Same invariant, no shared-code change, smaller PR surface. Verified end-to-end on gfx1201 (RX 9070 XT) with this commit applied: Ministral-3-8B-Instruct-2512 gsm8k 5-shot n=200: 0.775 / 0.775 Qwen3-8B-FP8 gsm8k 5-shot n=200: 0.890 / 0.895 Both within stderr of the prior 0.79 / 0.90 baseline. Note: the perf table in recipes/Ministral-3-8B.md was measured with max_model_len=4096; numbers should be unchanged in practice for the gsm8k workload (prompts ~1-2k tokens) but the table caption is now nominally stale.
Reorder of the import block (alphabetical by source) was incidental isort churn from an editor; restore main's original order. The only substantive top-of-file change kept is removing the now-unused "from aiter import silu_and_mul" import, since the new SiluAndMul.forward no longer calls it. No functional change; activation.py still imports and SiluAndMul construction still works.
Restore the "from aiter import silu_and_mul" import. SiluAndMul.forward no longer references it, so it is now an unused import, but keeping it makes the import block byte-identical to main and shrinks the PR diff in this file to just the forward() logic change.
native_triton_attn.py previously carried two in-tree triton kernels
(_kv_cache_write_kernel for the paged GQA KV scatter, _rope_neox_kernel
for fused Q+K Neox RoPE) plus a _dispatch closure that bypassed aiter's
paged_attention_decode wrapper to avoid the .item sync that breaks
CUDAGraph capture.
All three are now redundant:
* KV write: aiter PR adds aiter.ops.triton.kv_cache.reshape_and_cache
for the non-MLA paged [num_blocks, num_kv_heads, block_size,
head_dim] layout with slot < 0 sentinel skip.
* Decode: same aiter PR relaxes paged_attention_decode to accept
Union[float, torch.Tensor] scales, so the wrapper itself is now
CUDAGraph-safe when called with Python-float scales.
* RoPE: aiter already has rope_cached_thd_positions_2c_fwd (fused
Q+K, THD layout, accepts RotateStyle.NEOX); we just had to squeeze
ATOM's [max_pos, 1, 1, rotary_dim/2] cos/sin caches down to 2D
before passing them.
File shrinks 921 -> 727 lines (-194). No @triton.jit decorators left.
Verified end-to-end on RX 9070 XT (gfx1201): gsm8k 5-shot n=200
Ministral-3-8B 0.77/0.77, Qwen3-8B-FP8 0.905/0.905 - both within
stderr of the prior 0.79/0.90 baseline.
Depends on aiter PR #3319 with reshape_and_cache + pa_decode
Union[float, Tensor] scales.
…detect
NativeTritonBackend was added in this PR to work around aiter's prebuilt
HIP modules not loading on gfx1201. With the previous commit the backend
is pure orchestration on top of upstream aiter triton kernels, which
JIT-compile for every supported arch. So the same code is the correct
path everywhere - there is no longer a reason to gate it on arch.
* git mv native_triton_attn.py aiter_attention.py
(NativeTritonBackend renamed AiterBackend; get_name preserves the
legacy ROCM_AITER_ATTENTION string so plugin-mode routing is
unchanged)
* git mv (old) aiter_attention.py aiter_attention_hip.py
(preserved for a future capability-based HIP fast path; not wired
into the selector today)
* paged_attention.forward drops the _use_native_triton branch and the
torch.ops.aiter.unified_attention_with_output_base call; both
branches collapse into self.impl.forward(...)
* selector.py drops the gfx1201 try/import/branch; AiterBackend is
the unconditional default for non-MLA / non-GDN / non-MTP paths
Removes _is_gfx1201, use_native_triton_attn, ATOM_NATIVE_TRITON_ATTN
env var, and the NATIVE_TRITON_ATTENTION name check.
Caveat: gfx94x/95x now runs triton on the happy path instead of the
HIP unified_attention_with_output_base fused op. Perf to be measured;
HIP fast-path opt-in via aiter_attention_hip.py kept on disk for that
future work.
Verified on gfx1201 (RX 9070 XT): gsm8k 5-shot n=200
Mistral-3-8B 0.77/0.77, Qwen3-8B-FP8 0.905/0.905 - both within stderr
of the 0.79 / 0.90 baseline.
…gfx1201 detect" This reverts commit 7576b19.
use_native_triton_attn previously matched torch.cuda's gcnArchName
against the literal "gfx1201" prefix. Reframe as a positive allowlist
of arches that aiter's prebuilt HIP modules in rocm/atom-dev:latest
ship matching code objects for (today: gfx940 / gfx941 / gfx942 /
gfx950). Any arch outside the set routes to the in-tree triton backend.
The allowlist makes the coverage extension point explicit: when aiter
ships HIP support for a new arch, add the string here. The negative
form ("is this gfx1201?") had to be updated for every newly-broken
arch.
Behavior on gfx1201 is unchanged (still routes to NativeTritonBackend).
selector.py comment also refreshed to the same framing.
Verified on gfx1201 (RX 9070 XT): gsm8k 5-shot n=200
Mistral-3-8B 0.77/0.77, Qwen3-8B-FP8 0.905/0.905 - same as the
pre-allowlist runs.
aiter PR ROCm/aiter#3236 (opus gfx1201/gfx1200 RDNA4 support: parse + buffer rsrc + WMMA w32/w64) merged 2026-05-16 added the opus abstraction support that csrc/kernels/sample_kernels.cu relies on (opus::vector_t, opus::make_gmem, load_vector_nbytes). aiter's mixed_sample_outer_exponential now runs unconditionally on gfx1201, so the torch Gumbel-max fallback this gate selected is dead code. Verified by direct invocation of the op on gfx1201 (module_sample.so built locally on the gfx1201 host, returns correct sampled token IDs without error) and end-to-end via gsm8k 5-shot n=200 with the gate removed: Mistral-3-8B-Instruct-2512 strict/flex 0.78 / 0.78 Qwen3-8B-FP8 strict/flex 0.905 / 0.910 Both within stderr of the prior runs (which used the Gumbel fallback). At T=0 the two paths are argmax-equivalent in expectation; the tiny ±0.005-0.010 drift between runs reflects bit-level non-determinism in the aiter kernel's reduction order vs torch's.
The try/except wrappers for topk_gating, shuffle_scale, and flydsl.moe_common.GateMode were added defensively against the baked aiter in rocm/atom-dev:latest (pinned to old ATOM commit 247e9b1) which did not yet export those symbols. PR 811 now depends on the newer aiter that landed reshape_and_cache + pa_decode Union scale (ROCm/aiter#3319), and that aiter has all three symbols natively. The guards are dead code and the file is restored to match origin/main verbatim.
The previous _detect_gfx1201/_IS_GFX1201 gate was a name-pattern arch match. Reframe as a positive allowlist of arches that aiter's prebuilt HIP rmsnorm modules in rocm/atom-dev:latest ship matching code objects for (today: gfx940 / gfx941 / gfx942 / gfx950). Any arch outside the set routes to aiter's triton rmsnorm path instead. Same set and same pattern as atom/model_ops/attentions/native_triton_attn.py. Behavior on gfx1201 is unchanged (still routes to the triton rmsnorm because the HIP path SIGSEGVs on first call - verified directly). Verified on gfx1201 (RX 9070 XT): gsm8k 5-shot n=200 Mistral-3-8B strict/flex 0.770 / 0.775 Qwen3-8B-FP8 strict/flex 0.915 / 0.915 Both within stderr of the prior 0.79 / 0.90 baseline.
ATOM kept a per-arch dispatch that routed gfx1201 to the aiter triton rmsnorm path because aiter's HIP rmsnorm_quant kernel had CDNA-only v_pk_mul_f32 inline asm. With aiter PR ROCm/aiter#3322 [csrc] that asm is replaced with portable C++, so the HIP kernel now builds and runs correctly on gfx1201 (RX 9070 XT). Restore atom/model_ops/layernorm.py to match origin/main verbatim - removes _detect_gfx1201, _USE_AITER_TRITON_RMSNORM, the dispatch branches, and the AITER_HIP_PREBUILT_ARCHES allowlist. Verified end-to-end on gfx1201 with the aiter patch + module_rmsnorm_quant rebuild: gsm8k 5-shot n=200 Mistral-3-8B strict 0.795 / flex 0.800 Qwen3-8B-FP8 strict 0.900 / flex 0.900 Both within stderr of the 0.79 / 0.90 baseline. Depends on aiter PR 3322.
The QuantType.No branch at PagedLinear.forward gated tgemm.mm vs
torch.nn.functional.linear on gfx1201. Verified that aiter's
gemm_a16w16 (the function tgemm.mm calls) already detects gfx12* and
routes to its torch fallback automatically:
aiter/tuned_gemm.py:
gfx = get_gfx()
_no_asm = gfx.startswith("gfx12")
...
if _no_asm and libtype in ("asm", "skinny", "hipblaslt"):
libtype = "torch"
A direct test on gfx1201 confirms tgemm.mm(x_bf16, w_bf16) produces
output byte-identical to F.linear (max abs diff 0.0). So the ATOM-side
branch was dead code -- the workaround for an older aiter that lacked
the gfx12 detector.
Removed: ~9 lines (the if/else branch around line 743).
The other _is_gfx1201_linear call sites in this file (FP8 per-Tensor at
line 773, per-1x128 blockscale at line 817, preshuffle at line 728,
skip-pre-quant at line 763) are NOT removed. They are still load-bearing:
* FP8 per-Tensor: aiter torch_gemm fallback hits
NotImplementedError: "addmm_cuda" not implemented for Float8_e4m3fn
(the FP8 path in torch_gemm only handles uint8 storage, then
dequant-via-F.linear is integer math on the bytes -> garbage).
A proper fix needs aiter to route FP8 + scales to triton gemm_a8w8
on RDNA; deferred.
* per-1x128 blockscale: this build's triton lacks tl.dot(fp8, fp8);
ATOM's a16w8 fallback is the only working path. Triton version bump
needed; deferred.
Verified on gfx1201 (RX 9070 XT) end-to-end: gsm8k 5-shot n=200
Mistral-3-8B strict 0.790 / flex 0.795
Qwen3-8B-FP8 strict 0.895 / flex 0.895
Both within stderr of the 0.79 / 0.90 baseline.
ParallelLMHead.forward gated the FP8 lm_head fast path on
_is_gfx1201_linear(). The path itself is arch-agnostic - it just runs
aiter triton gemm_a8w8 on a one-time-quantized lm_head weight, the
same kernel used for the rest of the model's FP8 linear layers. Whether
the quantization is worth doing depends on memory bandwidth, not on
the specific arch.
Refactor:
* Rename ATOM_GFX1201_LM_HEAD_FP8 -> ATOM_LM_HEAD_FP8 in envs.py
(default still "1", env-var contract preserved); update the
docstring to describe the bandwidth tradeoff rather than gfx1201.
* In embed_head.py, drop the _is_gfx1201_linear import and rename the
helper _use_gfx1201_fp8_lm_head -> _use_fp8_lm_head. The helper
keeps the per-call eligibility checks (env var + x.is_cuda +
x.dim() == 2 + weight.dim() == 2 + weight.dtype == bfloat16) but
no longer checks the arch.
* Recipes: update both Ministral-3-8B.md and Qwen3-8B-FP8.md to use
the new env var name and drop "default on for gfx1201" wording.
Behavior change: on non-gfx1201 archs the FP8 lm_head path is now also
default-on. Cost there is small (~0.5% accuracy noise typical,
gsm8k stays within stderr); benefit is the same halved-bandwidth
lm_head. Set ATOM_LM_HEAD_FP8=0 to opt out anywhere.
Verified on gfx1201 (RX 9070 XT) end-to-end: gsm8k 5-shot n=200
Mistral-3-8B strict 0.780 / flex 0.780
Qwen3-8B-FP8 strict 0.910 / flex 0.910
Both within stderr of the 0.79 / 0.90 baseline.
Apply the same "decide-once, dispatch-on-capability" trick already used in
embed_head.py to the remaining gfx1201-gated branches in PagedLinear.
The arch detection is gone from this file:
- _detect_gfx1201, _IS_GFX1201, _is_gfx1201_linear: deleted (no callers)
- _fp8_per_tensor_linear_gfx1201 wrapper: deleted (inlined; was a
one-line indirection around _fp8_per_tensor_linear_triton)
forward() dispatch is now keyed on self.quant_type only:
- QuantType.No: tgemm.mm (already arch-portable per prior commit)
- QuantType.per_Tensor: always aiter triton gemm_a8w8 via
_fp8_per_tensor_linear_triton. The HIP tgemm.mm
FP8 path is faster on CDNA but produces
silently-wrong output on non-CDNA (torch_gemm
fallback reinterprets uint8-stored FP8 bytes
as integers). One arch-portable path is
simpler and meets accuracy targets here.
- QuantType.per_Token: gemm_a8w8 / gemm_a8w8_bpreshuffle (unchanged)
- QuantType.per_1x128: always aiter triton a16w8 blockscale. Avoids
the tl.dot(fp8, fp8) requirement that this
container's triton does not implement.
- QuantType.per_1x32: gemm_a4w4_quant (unchanged)
Pre-quant block (line 758) also restructured: quant on per_Token only;
all other FP8 quant_types take BF16 input (per-Tensor's triton wrapper
does internal quant, a16w8 reads BF16 directly).
init() weight-preshuffle for per_1x128: was conditional on
ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + arch; now always False since
a16w8 expects the plain (N, K) layout regardless of arch.
Trade-off: on CDNA archs that have working HIP FP8 and tl.dot(fp8, fp8)
triton, this commit loses access to the faster a8w8 paths. Accepted in
exchange for the simpler dispatch; revisit if a CDNA workload regresses.
Verified end-to-end on gfx1201 (RX 9070 XT): gsm8k 5-shot n=200
Mistral-3-8B strict 0.775 / flex 0.780
Qwen3-8B-FP8 strict 0.890 / flex 0.895
Both within stderr of the 0.79 / 0.90 baseline.
…paths" This reverts commit 87fab8b.
aiter PRs #3332 (dtypes.fp8 -> float8_e4m3fn on gfx1200/1201) and
#3336 (opus::_async_load synchronous fallback) together unblock the
per-Tensor FP8 path on gfx1201 via aiter.tgemm.mm:
- #3336 lets module_quant build for gfx1201, so the existing
quant_func (aiter dynamic_per_tensor_quant) produces FP8 x + scale
on gfx1201 without SIGSEGV.
- #3332 lets torch_gemm's FP8 branch fire (dtypes.fp8 == float8_e4m3fn
on gfx1201), so the fallback math F.linear(float8.to(fp32), ...)
* scales runs correct FP8 dequantization instead of reinterpreting
bytes as integers.
ATOM no longer needs an arch-specific dispatch for per-Tensor FP8.
Narrow the pre-quant skip (path 5) to per_1x128 on gfx1201 only (the
a16w8 path still wants BF16 x); drop the per-Tensor branch (path 2)
entirely. Delete the dedicated _fp8_per_tensor_linear_gfx1201 helper -
its callers are gone.
Remaining gfx1201 dispatch in this file: per_1x128 blockscale (still
gated, separate blocker - aiter's a8w8 blockscale path needs
tl.dot(fp8, fp8) which this container's triton lacks). Preshuffle
decision in __init__ still gated for the same reason.
Verified on gfx1201 (RX 9070 XT) with both aiter PRs applied locally:
gsm8k 5-shot n=200
Mistral-3-8B strict 0.770 / flex 0.775
Qwen3-8B-FP8 strict 0.89 / flex 0.90
Both within stderr of the 0.79 / 0.90 baseline.
Depends on aiter PRs #3332 and #3336.
- envs.py: black collapses ATOM_LM_HEAD_FP8 lambda to one line - activation.py: ruff drops aiter.silu_and_mul unused import. We kept the import earlier to match origin/main exactly, but our SiluAndMul rewrite means we no longer call silu_and_mul - main still does, so on main the import is used. Removing it on this branch is what ruff expects and is the right answer. No runtime impact. Required for ATOM CI to pass.
aiter's HIP gemm_a8w8_blockscale (CK) has no gfx1201 code object, but the aiter triton non-preshuffle gemm_a8w8_blockscale JIT-compiles for gfx1201 and matches F.linear(bf16) at the FP8 quant-noise floor. The earlier 'tl.dot(fp8, fp8) unsupported on this build' rationale was specific to the uint8 storage view; passing the weight as torch.float8_e4m3fn lets triton accept it. The shared per_1x128 pre-quant uses transpose_scale=ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE (defaults True -> column-major x_scale for the HIP preshuffle kernel). The non-preshuffle triton kernel expects row-major x_scale, so the gfx1201 branch keeps the shared pre-quant skip and runs quant inline with transpose_scale=False before the GEMM call. (Without this, output magnitudes 2-4x larger per layer, the model collapses by layer 2 and emits garbage tokens -- gsm8k 0.005.) * Replaces the a16w8 blockscale fallback that decoded FP8 weight to BF16 inside the kernel; the new path stays in FP8 for both operands. * Drops _get_triton_a16w8_blockscale and the explicit BLOCK_N=64 shared-mem override that the a16w8 kernel needed. Qwen3-8B-FP8 gsm8k 5-shot n=200 on gfx1201: flexible-extract 0.905 +/- 0.0208 (baseline 0.90) strict-match 0.900 +/- 0.0213 Decode pace: 1.27 s/it (a16w8 was 1.85 s/it, ~30% faster).
With the aiter PR adding a HIP -> triton fallback for
gemm_a8w8_blockscale[_bpreshuffle] (allowlist {gfx940,gfx941,gfx942,gfx950}
for HIP, triton otherwise), nothing in this file has to know about
gfx1201. Drops:
* _detect_gfx1201 / _IS_GFX1201 / _is_gfx1201_linear
* init: the need_shuffle = False override for gfx1201 (weight now
shuffled like every other arch, layout matches what the triton
preshuffle kernel reshapes to)
* pre-quant: the gfx1201 skip-quant narrow
* forward: the if _is_gfx1201_linear(): triton-non-preshuffle branch
added in b7d215f
The remaining 4 mentions of gfx1201 are in explanatory comments
documenting why certain choices are correct on RDNA4 (helper docstring
in _get_triton_fp8_gemm, comment near gemm_a8w8 tuning config note,
the per-Tensor branch's aiter#3332 context). No code-level dispatch.
Net diff: -93 / +6 lines.
Verified on gfx1201 (Qwen3-8B-FP8 gsm8k 5-shot n=200):
flexible-extract 0.895 +/- 0.0217 (baseline 0.90)
strict-match 0.895 +/- 0.0217
Decode pace 1.32 s/it (b7d215f was 1.27 s/it; ~4% slower since
default path is now preshuffle via aiter fallback, was non-preshuffle
inline-quant before). Trade is intentional: simpler code, perf within
stderr of baseline.
These helpers were the dedicated fast path for two FP8 GEMMs on gfx1201
back when the only alternative was torch dequant + matmul (~360x slower).
With aiter#3332 making dtypes.fp8 = float8_e4m3fn on gfx1200/1201, the
aiter.tgemm.mm torch fallback now runs F.linear(float8.to(fp32), ...)
* scales, which dequantizes correctly. And with aiter#3343 routing
per_1x128 blockscale GEMMs to triton on non-CDNA archs, every FP8 path
ATOM needs goes through stock aiter wrappers.
* linear.py: drop _TRITON_FP8_GEMM / _get_triton_fp8_gemm /
_build_w_scale_full / _get_aiter_dynamic_per_tensor_quant /
_get_aiter_dynamic_per_token_quant / _fp8_per_tensor_linear_triton
(-150 lines). Only caller was embed_head's FP8 lm_head path.
* embed_head.py: drop the FP8 lm_head feature (_get_fp8_lm_head_weight,
_use_fp8_lm_head, the FP8 branch in forward). lm_head now uses plain
BF16 tgemm.mm everywhere (-70 lines).
* envs.py: drop ATOM_LM_HEAD_FP8.
Net: -225 lines. linear.py + embed_head.py are now zero-diff against
origin/main.
Verified on gfx1201 (Qwen3-8B-FP8 gsm8k 5-shot n=200):
flexible-extract / strict-match: 0.895 / 0.895 +/- 0.0217
Baseline 0.90; within stderr.
Decode pace ~1.5 s/it (vs 1.32 with the dedicated FP8 lm_head triton
kernel). The FP8 lm_head feature was net-positive only because of that
dedicated kernel; without it, BF16 lm_head ties or beats it on this
arch, so dropping the feature loses no perf.
…201_mistral3_rebased
Enable three inference optimizations for Qwen3 on MI300X: 1. Fused RMSNorm + FP8 Quantize: merge normalization and quantization into a single kernel pass, eliminating one HBM round-trip per layer. 2. Fused SiLU×Mul + FP8 Quantize: merge MLP activation and quantization, removing another HBM round-trip per layer. 3. FP8 KV Cache: store K/V in fp8 (half the memory of bf16), enabling larger batch sizes and longer sequences. Also fix fused RMSNorm tuple unpacking for per_1x128 dynamic quant — the unpack condition only checked per_1x32, missing per_1x128 where input_scale is None, causing downstream linear to receive a tuple.
Contributor
|
Why do this : Qwen3-8B-FP8 on gfx1201 appears bandwidth-sensitive in decode/prefill. The previous path had avoidable memory traffic from separate norm/activation → quant steps and BF16 KV cache writes. This optimization adds two changes:
Results on RX 9070 XT / gfx1201, Qwen3-8B-FP8, ISL/OSL=549/256, The before numbers are measured against the latest tuned GEMM JSON configs from ROCm/aiter#3484, not the original untuned PR #811 recipe baseline.cudagraph:
My next plans: continue profiling the remaining decode/prefill hotspots and look for additional fusion or quantization opportunities. |
…1_mistral3_rebased Resolve support_model_arch_dict conflict in model_runner.py: keep main MiMoV2 consolidation (mimo_v2_flash -> mimo_v2) plus PR 811 Mistral3 entries.
ATOM_USE_UNIFIED_ATTN=1 + bf16 KV now asserts block_ratio == 1, which needs --block-size 64 (default 16 fails). The gfx1201 unified-attention path itself needs no ATOM change: latest main allocates the 5D SHUFFLE KV layout and the write goes through aiter.reshape_and_cache, whose triton fallback now supports the SHUFFLE layout on non-CDNA arches (aiter-side fix).
…3_rebased' into carhuang/support_gfx1201_mistral3_rebased # Conflicts: # recipes/Qwen3-8B-FP8.md
… decode TritonMHAMetadataBuilder.build_kv_cache_tensor now selects the KV layout by arch via use_flash_kv_layout(): RDNA4 (gfx1200/gfx1201) gets the flash (4D NHD) layout; CDNA / gfx1250 keep the 5D SHUFFLE layout. unified_attention must un-shuffle the SHUFFLE layout with an in-register transpose on every KV tile, which is LDS-bound on RDNA4 WMMA and costs ~1.7-1.8x decode TPOT (measured: Qwen3-8B 16 vs 29 ms, Ministral-3-8B 19 vs 33 ms). Flash needs no transpose. New env ATOM_KV_CACHE_LAYOUT=auto|flash|shuffle overrides the auto choice. attention_mha already supports both via shuffled_kv_cache = not use_flash_layout. Validated on RX 9070 XT (gfx1201): gsm8k 5-shot n=200 strict 0.885 / flex 0.895.
…in ATOM) use_flash_kv_layout() now queries aiter unified_attention_prefers_shuffled_kv() instead of matching gfx strings: flash (4D NHD) wherever unified_attention would otherwise un-shuffle the layout with an in-register transpose (generic triton kernel, e.g. RDNA4 gfx1201), shuffle where the gfx1250 gluon kernel consumes it directly. ATOM holds no arch knowledge; ATOM_KV_CACHE_LAYOUT still overrides.
…201_mistral3_rebased
…E_LAYOUT + flash selection) TritonMHABackend uses the 5D SHUFFLE layout that triton unified_attention already supports (shuffled_kv_cache=True), same as main — no ATOM-side layout logic, no env. triton_mha.py is back to origin/main. The gfx1201 KV write still goes through the aiter triton SHUFFLE reshape_and_cache fallback. (Decode on the generic triton kernel is ~1.7-1.8x slower than flash on gfx1201; to be addressed in the triton kernel later.)
…201_mistral3_rebased
…201_mistral3_rebased
Mention AMD Radeon RX 9070 / RX 9070 XT and Radeon AI PRO R9700 (Navi 48, gfx1201) experimental support via the all-triton attention path (ATOM_USE_UNIFIED_ATTN=1), linking the Qwen3-8B-FP8 and Ministral-3-8B recipes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
No lm_head-FP8 logic exists in this PR (config.py only adds the mistral3 multimodal type; quant_spec.py handles weight_block_size). Remove the dead ATOM_GFX1201_LM_HEAD_FP8 entry from docs/environment_variables.md (0 code refs -> file now zero-diff vs main) and the stale ATOM_LM_HEAD_FP8 note from the Ministral-3-8B recipe. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ements Re-measured on RX 9070 XT (gfx1201), bf16 KV, single stream (num-prompts 8, concurrency 1, 1 warmup). Replaces the stale row that showed an implausible 801 ms TTFT for 549/256 cudagraph (now 86 ms). cudagraph 18/80: 39ms/18.5ms/53.3; 549/256: 86ms/18.6ms/52.9; eager 549/256: 93ms/28.6ms/35.6. Now consistent with the PR description perf table. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Recipes ran attention + GEMM through Triton but the KV-cache write / RoPE / norms use native aiter HIP kernels -- so drop the inaccurate "all-triton" framing. Replace the .so-internals note with the actual prerequisite: build aiter for gfx1201 (GPU_ARCHS=gfx1201), linking ROCm/aiter#3846. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
valarLip
pushed a commit
to ROCm/aiter
that referenced
this pull request
Jun 22, 2026
…lain + blockscale_preshuffled) (#3343) * gemm_a8w8_blockscale[_bpreshuffle]: HIP -> triton fallback for non-CDNA The prebuilt HIP CK modules module_gemm_a8w8_blockscale and module_gemm_a8w8_blockscale_bpreshuffle ship code objects only for gfx94x/95x. On other arches (e.g. gfx1201 RDNA4 in rocm/atom-dev:latest) the HIP kernel launch SIGSEGVs with 'No compatible code objects found for the device' -- no catchable exception, so we must allowlist arches BEFORE calling the HIP path. Added _hip_blockscale_supported() that returns True only for the four prebuilt CDNA arches. When False, both public wrappers route to their triton equivalents in aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale: * gemm_a8w8_blockscale (non-preshuffle) -> triton non-preshuffle. x_scale layout matches (row-major). Direct dispatch. * gemm_a8w8_blockscale_bpreshuffle -> triton preshuffle. The (N, K) weight has been (16, 16)-shuffled the same way the triton kernel expects; a reshape view to (N//16, K*16) satisfies the kernel's shape assertion. x_scale is column-major (the preshuffle convention), direct fit. An explicit fallback config is passed because the gfx1201-GEMM-A8W8_BLOCKSCALE_PRESHUFFLED.json tuning file does not yet exist on main. uint8-storage weight (the d_dtypes['fp8'] convention from before #3332) is viewed as float8_e4m3fn so triton's tl.dot accepts it; if the weight is already float8_e4m3fn the view is a no-op. Verified on gfx1201 against F.linear(bf16) gold across Qwen3-8B-FP8 layer shapes (qkv 1/8/16384x4096->6144, o, gate_up, down) for both paths: relative error 3.0-3.6%, at the FP8 quant-noise floor. * configs: add gfx1201 GEMM-A8W8_BLOCKSCALE_PRESHUFFLED.json Needed because ATOM's atom/model_ops/linear.py:gemm_a8w8_blockscale_preshuffle_impl calls the triton preshuffle wrapper directly (bypassing the HIP wrapper my new fallback adds to). Without this file the triton wrapper raises "Required config file doesn't exist" at first per_1x128 layer call. Starting values copied from gfx1201-GEMM-A8W8_BLOCKSCALE.json (the non-preshuffle variant, real gfx1201 tuning). Same M-bracket structure, same block sizes -- the preshuffle kernel reads the same parameter names. Suboptimal but correct: the preshuffle memory access pattern likely prefers different BLOCK_SIZE_N / waves_per_eu values, so dedicated preshuffle tuning is a TODO once a tuning run is available. Only one file -- ATOM's two test models (Mistral-3-8B per-Tensor and Qwen3-8B-FP8 per-1x128) hit only this preshuffle kernel; the non-preshuffle config already exists on main, and the other GEMM families (A16W16, AFP4WFP4, batched, fused, etc.) are not invoked by these models. * configs: add gfx1201 GEMM-A8W8 (non-blockscale) tuning configs Adds 5 triton tuning JSONs for the plain GEMM-A8W8 (non-blockscale) kernel on gfx1201, covering Mistral-3 / Qwen3 shapes: - GEMM-A8W8.json (default M-bracket table) - GEMM-A8W8-N=4096-K=4096.json (qkv/o projection) - GEMM-A8W8-N=6144-K=4096.json (qkv for Qwen3 GQA) - GEMM-A8W8-N=4096-K=14336.json (down projection) - GEMM-A8W8-N=28672-K=4096.json (gate+up projection) Same tuning artifact source as the gfx1201-GEMM-A8W8_BLOCKSCALE_PRESHUFFLED config added in the previous commit (and the existing gfx1201-GEMM-A8W8_BLOCKSCALE configs already on main). Bundled here because both belong to the gfx1201 triton GEMM family used by ATOM PR ROCm/ATOM#811. Folds in the JSON content previously split out as PR #3234 (closed). * gemm_a8w8: condense gfx1201 triton-fallback comments Tighten the verbose arch-gate / fallback comments (no logic change). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…sign) use_flash_layout is hardcoded False everywhere (never set True), a dormant flag inherited from main. #811 had OR-ed it into the triton-attn elif and wrapped the KV write in an if self.use_flash_layout: triton-NHD branch -- both dead leftovers from the reverted Option-A flash layout. Remove them: the KV write always uses native aiter.reshape_and_cache(asm SHUFFLE). No behavior change (the removed branches were unreachable). #811 now adds zero use_flash_layout references; remaining refs are main-inherited. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The qwen3.py changes were entirely the fused RMSNorm+Quant / SiLU+Mul+Quant optimization (+ x_scale threading), gated on the pre-existing ATOM_LLAMA_ENABLE_AITER_TRITON_FUSED_* flags. That is a general Qwen FP8 optimization, not gfx1201 enablement, and the gfx1201 recipes disable it (=0). Revert qwen3.py to main (now zero-diff) so this PR stays focused; the fused-quant work can land as its own PR. Validated: Qwen3-8B-FP8 on gfx1201 gsm8k n=200 strict 0.89 / flex 0.90 (unchanged). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…201_mistral3_rebased
The ATOM_USE_UNIFIED_ATTN -> TritonMHABackend routing already exists in main; #811 added only a comment. Drop it so selector.py is zero-diff. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Drop the inaccurate "all-triton / no prebuilt HIP kernels" wording (the KV-cache write uses native aiter HIP). Just note Navi 4 (gfx1201) support and link the recipes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
valarLip
approved these changes
Jun 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
gfx1201 (RX 9070 XT, RDNA4) support for ATOM. Attention and GEMM run through Triton; the KV-cache write, RoPE and norms use native aiter HIP kernels. Requires aiter built for gfx1201 — see ROCm/aiter#3846 (short-term: build aiter with
GPU_ARCHS=gfx1201).Models verified (gsm8k 5-shot, n=200, cudagraph, bf16 KV,
ATOM_USE_UNIFIED_ATTN=1)Performance (RX 9070 XT, single stream, cudagraph, bf16 KV, shuffle/latest main)
(A flash 4D layout runs ~16.0 / 19.3 ms TPOT — slightly faster — but needs ATOM-side
layout selection, so it's not used. Qwen TPOT is consistently ~18.6 ms on latest main; earlier builds showed run-to-run variation (~18–30 ms across builds),
likely FP8 block-128 GEMM autotuning; Ministral is stable. gfx1201 shuffle decode is a candidate
for triton-kernel optimization. The native-HIP KV write (vs the dropped #3319/#3837 writers) does not affect decode perf — these numbers match all prior configs.)
Reproduced on RX 9070 XT (gfx1201, 16 GB), single GPU, fresh
rocm/atom-dev:latest, latestmain, with aiter built for gfx1201 (GPU_ARCHS=gfx1201). All required aiter code is now merged (see deps below).What this PR adds (net vs main)
ATOM_USE_UNIFIED_ATTN=1— the gfx1201 path uses the pre-existingatom/utils/selector.pyrouting (if envs.ATOM_USE_UNIFIED_ATTN: return TritonMHABackend, already in main); this PR only adds a clarifying comment there. The flag selectsTritonMHABackend(aiter tritonunified_attention) over the defaultAiterBackend. No arch detection.atom/model_ops/attention_mha.py— KV write + unified_attention plumbing for the 5D-shuffle layout (shuffled_kv_cache=True); the gfx1201 KV write goes through aiter's nativeaiter.reshape_and_cache(asm SHUFFLE), which requires aiter built/prebuilt with gfx1201 support formodule_cache(see aiter issue below).triton_mha.pyis unchanged vs main (no ATOM layout logic).atom/models/mistral3.py(new) — Ministral-3 / Mistral3 model support.atom/quant_spec.py,atom/config.py— supporting config/quant plumbing.recipes/Ministral-3-8B.mdandrecipes/Qwen3-8B-FP8.md: documentATOM_USE_UNIFIED_ATTN=1, the now-required--block-size 64(the engine assertsblock_ratio == 1for unified-attn + bf16 KV; default 16 fails),--max-model-len 16384, and the bf16/fp8 KV options.linear.py,embed_head.py,activation.py,qwen3.pyremain zero-diff vs main — FP8 / SwiGLU / blockscale GEMM run through stock aiter wrappers that fall back to triton on RDNA4 via the aiter PRs below.Aiter dependencies (required on gfx1201)
opus.hppgfx1200/gfx1201 RDNA4: parse + buffer rsrc + WMMA w32/w64opus::pk_mul_f32+rmsnorm_quant: guard CDNA-onlyv_pk_mul_f32asmaiter.dtypes.fp8: gfx1200/gfx1201 entries; unblockstgemm.mmFP8 fallbackopus::_async_load: synchronous fallback for non-buffer-to-LDS archsv_pk_mul_f32in silu/swiglu kernels; unblocks stockaiter.silu_and_mulgfx1201inGPU_ARCHSso the prebuiltmodule_cacheincludes a gfx1201 code object; the native HIPaiter.reshape_and_cache(asm_layout=True)then writes the 5D-SHUFFLE layout directly. A native from-source build on gfx1201 (GPU_ARCHS=native) does this automatically; for a multi-arch wheel addgfx1201to the prebuildGPU_ARCHS. (A CDNA-only prebuilt run on gfx1201 SIGSEGVs — see #3846.) No KV-write aiter PR needed — the earlier triton fallback #3319 and dedicated HIP module #3837 are dropped.gemm_a8w8[_blockscale][_preshuffled]HIP→triton fallback + tuning configs (Mistral-3 / Qwen3 shapes)All aiter code dependencies above are now merged. The only remaining gfx1201 requirement is building aiter with
GPU_ARCHS=gfx1201(the Build-flag row / ROCm/aiter#3846).Known caveats
activation_scalecheckpoint tensors are silently dropped on load (harmless — FP8 GEMM dequantizes weights and ignores per-channel input scale).iommu=pt).