new MHA v4 entrypoint ; hosts a spectrum of quantized & sparse attention kernels

[jcaraban]

Motivation

⚠️ These MHA kernels are mainly tested with Diffusion Inference workloads.

Introduce an extensible, inference-only MHA v4 entrypoint for mixed low-precision ASM attention kernels, without overwhelming the existing aiter.ops.mha module. In the future we can consider merging into the common entrypoint, but these kernels are best kept independent as they evolve.

Q: why v4? A: IMO this is the closest thing to FAv4 on cdna3/4, also doing 1) conditional-softmax-rescaling, 2) exponential-approx/emu, 3) optimized pipelines for valu/mfma/loads utilization. Though I'm open for naming suggestions.

Replaces #3495 and #3186

Technical Details

• Introduces optimized low-precision dense attention for gfx950, with initial gfx942 support.
• Adds INT8, FP8, MXFP6 E2M3, and MXFP4 E2M1 operand combinations with BF16 output.
• Includes the binaries, quantization and packing routines, host dispatch, JIT integration, and benchmark coverage.
• Establishes aiter.ops.mha_v4 as an extensible interface for future architectures, formats, output types, sparse attention, GQA, causal attention, head dimensions, etc
• Preserves efficient torch.compile execution and distributed overlap by separating Q/K/V preprocessing and safely encapsulating specialized packed layouts.

Test Plan

• Run focused MHA v4 pytest coverage.
• Validate eager versus full-graph compiled parity.
• Exercise allocator churn and downstream tensor consumption.
• Run the unified 8K benchmark across all kernels.

Test Result

• 19 focused tests passed.
• All six raw paths matched eager and compiled execution.
• gfx942/gfx950 manifest generation passed.
• The 8K all-kernel benchmark completed successfully.

Current mi355x long-sequence dense ASM kernel throughput, excluding Q/K/V preprocessing:

Q/K	V	O	TFLOPs
INT8	FP8	BF16	2315
FP8	FP8	BF16	3118
mxFP6	FP8	BF16	3430
mxFP4	FP8	BF16	3540
mxFP6	FP4	BF16	3790
mxFP4	FP4	BF16	4000

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

	softmax_scale: Optional[float] = None,
	softmax_scale: float | None = None,

[jcaraban]

No! This exact change led to e2e regressions when it was applied to aiter.ops.mha recently. It messes with Inductor and no test is capturing the regression yet.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4627 --add-label <label>

[github-actions]

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

Suggested change

	from enum import IntEnum
	from typing import Optional
	import torch
	import triton
	from torch import Tensor
	from aiter import dtypes
	from aiter.ops.triton._triton_kernels.quant.sage_attention_quant import (
	sage_quant_v_amax_finalize_kernel,
	sage_quant_v_amax_partial_kernel,
	sage_quant_v_kernel,
	)
	from aiter.ops.triton.quant.mxfp6_fmha_pack import (
	fp6_k_raw_buffer_sizes,
	fp6_k_lds_order_views_from_raw,
	reorder_fp6_k_lds_order_triton,
	)
	from aiter.ops.triton.quant.sage_attention_quant_wrappers import (
	fp4_v_padded_sequence,
	fp4_v_raw_buffer_size,
	pack_v_mxfp4_colmajor_raw,
	)
	from ..jit.core import compile_ops
	from ..jit.utils.chip_info import get_gfx
	MHA_V4_LOG2E = 1.4426950408889634
	from enum import IntEnum
	from typing import Optional
	import torch
	import triton
	from torch import Tensor
	from aiter import dtypes
	from aiter.ops.triton._triton_kernels.quant.sage_attention_quant import (
	sage_quant_v_amax_finalize_kernel,
	sage_quant_v_amax_partial_kernel,
	sage_quant_v_kernel,
	)
	from aiter.ops.triton.quant.mxfp6_fmha_pack import (
	fp6_k_lds_order_views_from_raw,
	fp6_k_raw_buffer_sizes,
	reorder_fp6_k_lds_order_triton,
	)
	from aiter.ops.triton.quant.sage_attention_quant_wrappers import (
	fp4_v_padded_sequence,
	fp4_v_raw_buffer_size,
	pack_v_mxfp4_colmajor_raw,
	)
	from ..jit.core import compile_ops
	from ..jit.utils.chip_info import get_gfx
	MHA_V4_LOG2E = 1.4426950408889634

[github-actions]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

	softmax_scale: Optional[float] = None,
	softmax_scale: float | None = None,

[github-actions]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

	out: Optional[Tensor] = None,
	out: Tensor | None = None,

[github-actions]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

	softmax_scale: Optional[float] = None,
	softmax_scale: float | None = None,

[github-actions]

⚠️ [ruff] <UP045> _{reported by reviewdog 🐶}
Use X | None for type annotations

Suggested change

	out: Optional[Tensor] = None,
	out: Tensor | None = None,

[github-actions]

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

aiter/aiter/ops/triton/quant/sage_attention_quant_wrappers.py

Lines 1 to 22 in 50c3dfc

	import functools
	import torch
	import triton
	import aiter
	from aiter.ops.triton._triton_kernels.attention.fav3_sage_attention import (
	map_dims,
	)
	from aiter.ops.triton._triton_kernels.quant.sage_attention_quant import (
	_compute_delta_s_kernel,
	_q_smooth_int8_kernel,
	_rot_k_only_kernel,
	_rot_q_kernel,
	_rotate_quantize_k_kernel,
	_rotate_quantize_q_kernel,
	sage_quant_kernel,
	sage_quant_v_fp4_colmajor_kernel,
	sage_quant_v_mxfp4_colmajor_kernel,
	sage_quant_v_kernel,
	)
	from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp

[github-actions]

⚠️ [ruff] <I001> _{reported by reviewdog 🐶}
Import block is un-sorted or un-formatted

aiter/op_tests/op_benchmarks/triton/bench_sage.py

Lines 1 to 71 in 50c3dfc

	from __future__ import annotations
	import argparse
	import csv
	import glob
	import json
	import logging
	import os
	import re
	import sys
	import tempfile
	import time
	from dataclasses import dataclass
	from typing import Any, Literal
	import torch
	import triton
	import aiter
	from aiter.ops.mha import (
	flash_attn_func,
	)
	from aiter.ops.mha_v4 import (
	AttentionFormat,
	MHA_V4_LOG2E,
	mha_v4,
	mha_v4_packed,
	mxfp4_k_view,
	native_fp8_format,
	quantize_mxfp4_k,
	rotate_activation_mxfp4_quant,
	rotate_activation_mxfp6_quant,
	scale_modes_for_formats,
	)
	from aiter.ops.triton._triton_kernels.flash_attn_triton_amd import flash_attn_3
	from aiter.ops.triton._triton_kernels.quant.sage_attention_quant import (
	sage_quant_v_amax_finalize_kernel,
	sage_quant_v_amax_partial_kernel,
	sage_quant_v_kernel,
	)
	from aiter.ops.triton.attention.fav3_sage import (
	fav3_sage_func,
	fav3_sage_wrapper_func,
	get_sage_fwd_configs,
	)
	from aiter.ops.triton.attention.fav3_sage_attention_mxfp4_wrapper import (
	fav3_sage_mxfp4_func,
	fav3_sage_mxfp4_wrapper,
	get_sage_fwd_configs_mxfp4,
	)
	from aiter.ops.triton.attention.mha_v3 import _quantize_bshd
	from aiter.ops.triton.attention.utils import block_attn_mask_to_ragged_lut
	from aiter.ops.triton.quant.mxfp6_fmha_pack import (
	reorder_fp6_k_lds_order_triton,
	)
	from aiter.ops.triton.quant.sage_attention_quant_wrappers import (
	create_hadamard_matrix,
	sage_quant,
	sage_quant_f4f4,
	sage_quant_mxfp4,
	sage_quant_mxfp6,
	sage_quant_v_mxfp4,
	)
	from aiter.test_mha_common import attention_ref, attention_ref_block_sparse
	from op_tests.op_benchmarks.triton.utils.benchmark_utils import (
	get_caller_name_no_ext,
	)
	from op_tests.triton_tests.attention.test_fav3_sage import (
	check_attention_outputs,
	compare_accuracy,
	)