@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f43e1d26e3

[ROCm] Enable AITER and FP8 inference on GFX120x (#43615)

Signed-off-by: skysnow2001 <skysnow9285@gmail.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

skysnow2001 committed Aug 3, 2026 at 17:45 UTC f43e1d26e3e6b40398be27b218cf2f2786432028

12 files changed +172 -44

vllm/_aiter_ops.py

+84 -14

@@ -72,6 +72,21 @@ def is_aiter_found_and_supported() -> bool:

72 return False

75 +def is_aiter_found_and_supported_on_rdna4() -> bool:

76 + """RDNA4 (gfx12) analog of `is_aiter_found_and_supported()`.

77 +

78 + gfx12 has no aiter CK build, so this deliberately stays off the gfx9

79 + `@if_aiter_supported` umbrella; it reports only that aiter's Triton

80 + kernels are usable here. Like its gfx9 counterpart it checks platform +

81 + arch + library availability and does not check environment variables.

82 + """

83 + if current_platform.is_rocm() and IS_AITER_FOUND:

84 + from vllm.platforms.rocm import on_rdna4

85 +

86 + return on_rdna4()

87 + return False

88 +

89 +

90 @functools.cache

91 def _load_gemm_tuned_configs(

92 q_dtype_w: torch.dtype, csv_path: str

@@ -1702,6 +1717,24 @@ class rocm_aiter_ops:

1717 def is_enabled(cls) -> bool:

1718 return cls._AITER_ENABLED

1719

1720 + @classmethod

1721 + def is_rdna_aiter_enabled(cls) -> bool:

1722 + """AITER on RDNA4 (gfx12): library present, arch is rdna4, and the user

1723 + enabled aiter. Only aiter's Triton kernels exist on gfx12 (no CK build),

1724 + so this deliberately stays off the gfx9/CK `@if_aiter_supported` umbrella

1725 + and gates only the Triton paths rdna4 uses. The gfx12 analog of

1726 + `is_enabled()`."""

1727 + if not current_platform.is_rocm() or not IS_AITER_FOUND:

1728 + return False

1729 + from vllm.platforms.rocm import on_rdna4

1730 +

1731 + return on_rdna4() and cls._AITER_ENABLED

1732 +

1733 + @classmethod

1734 + def is_rdna_linear_enabled(cls) -> bool:

1735 + """RDNA4 (gfx12) analog of is_linear_enabled() (aiter Triton blockscale)."""

1736 + return cls.is_rdna_aiter_enabled() and cls._LINEAR_ENABLED

1737 +

1738 @classmethod

1739 @if_aiter_supported

1740 def is_linear_enabled(cls) -> bool:

@@ -1864,17 +1897,8 @@ class rocm_aiter_ops:

1897 aiter_ar_comm if isinstance(aiter_ar_comm, AiterCustomAllreduce) else None

1898 )

1899

1867 - @classmethod

1868 - @if_aiter_supported

1869 - def are_gdn_triton_kernels_available(cls) -> bool:

1870 - """Check if AITER Triton kernels for GDN attention are importable.

1871 -

1872 - These are optional Triton kernels (conv1d fast-path, gated delta net)

1873 - used by GatedDeltaNetAttention's decode fast-path. They may be absent

1874 - in older aiter builds.

1875 - """

1876 - if not cls._AITER_ENABLED:

1877 - return False

1900 + @staticmethod

1901 + def _gdn_triton_kernels_importable() -> bool:

1902 try:

1903 import aiter.ops.triton.causal_conv1d_update_single_token # noqa: F401

1904 import aiter.ops.triton.gated_delta_net # noqa: F401

@@ -1886,6 +1910,22 @@ class rocm_aiter_ops:

1910 except (ImportError, ModuleNotFoundError):

1911 return False

1912

1913 + @classmethod

1914 + @if_aiter_supported

1915 + def are_gdn_triton_kernels_available(cls) -> bool:

1916 + """Check if AITER Triton kernels for GDN attention are importable.

1917 +

1918 + These are optional Triton kernels (conv1d fast-path, gated delta net)

1919 + used by GatedDeltaNetAttention's decode fast-path. They may be absent

1920 + in older aiter builds.

1921 + """

1922 + return cls._AITER_ENABLED and cls._gdn_triton_kernels_importable()

1923 +

1924 + @classmethod

1925 + def is_rdna_gdn_triton_kernels_available(cls) -> bool:

1926 + """RDNA4 (gfx12) analog of are_gdn_triton_kernels_available()."""

1927 + return cls.is_rdna_aiter_enabled() and cls._gdn_triton_kernels_importable()

1928 +

1929 @classmethod

1930 @if_aiter_supported

1931 @functools.cache

@@ -1902,8 +1942,12 @@ class rocm_aiter_ops:

1942 return "gate_mode" in inspect.signature(fused_moe).parameters

1943

1944 @staticmethod

1905 - @if_aiter_supported

1945 def register_ops_once() -> None:

1946 + if not (

1947 + is_aiter_found_and_supported() or is_aiter_found_and_supported_on_rdna4()

1948 + ):

1949 + return

1950 +

1951 global _OPS_REGISTERED

1952 if not _OPS_REGISTERED:

1953 # register all the custom ops here

@@ -2774,7 +2818,11 @@ class rocm_aiter_ops:

2818

2819 @staticmethod

2820 def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:

2777 - return (n, k) in [

2821 + if not current_platform.is_rocm():

2822 + return False

2823 + from vllm.platforms.rocm import on_gfx950, on_rdna4

2824 +

2825 + gfx950_tuned = {

2826 (1024, 8192),

2827 (2112, 7168),

2828 (3072, 1536),

@@ -2786,7 +2834,29 @@ class rocm_aiter_ops:

2834 (7168, 256),

2835 (8192, 1024),

2836 (8192, 32768),

2789 - ]

2837 + }

2838 + rdna4_tuned = gfx950_tuned | {

2839 + (2048, 2048),

2840 + (2624, 6144),

2841 + (3072, 6144),

2842 + (3584, 512),

2843 + (4096, 512),

2844 + (6144, 1536),

2845 + (6144, 2048),

2846 + (7168, 2304),

2847 + (7168, 16384),

2848 + (7168, 18432),

2849 + (8192, 8192),

2850 + (16384, 1536),

2851 + (24576, 1536),

2852 + (32768, 512),

2853 + (36864, 7168),

2854 + }

2855 + if on_rdna4():

2856 + return (n, k) in rdna4_tuned

2857 + if on_gfx950():

2858 + return (n, k) in gfx950_tuned

2859 + return False

2860

2861 @staticmethod

2862 def is_triton_gemm_afp4wfp4_presh_ws_tuned(n: int, k: int) -> bool:

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

+21 -8

@@ -581,6 +581,9 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):

581

582 gated_norm_shapes.add((num_v_heads // layer.tp_size, head_v_dim))

583

584 + # RDNA4 uses native quant ops and supports only Triton replacements.

585 + match_aiter_quant_op = not rocm_aiter_ops.is_rdna_aiter_enabled()

586 +

587 # Make sure fused add patterns are before simple rms norm,

588 # as the latter is a subset of the former in torch ops.

589 # The DoubleQuant patterns handle 1 rms_norm -> 2 group_fp8_quant

@@ -592,22 +595,22 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):

595 for epsilon in [1e-5, 1e-6]:

596 # Fuse aiter rms_norm + 2x aiter group fp8 quant

597 DoubleAiterRMSFp8GroupQuantPattern(

595 - epsilon, FP8_DTYPE, GroupShape(1, 128)

598 + epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op

599 ).register(self.patterns)

600

601 # View-tolerant sibling for DSv3.2 q_c norm fan-out

602 DoubleAiterRMSFp8GroupQuantViewPattern(

600 - epsilon, FP8_DTYPE, GroupShape(1, 128)

603 + epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op

604 ).register(self.patterns)

605

606 # Fuse aiter rms_norm + aiter dynamic group fp8 quant

607 AiterRMSFp8GroupQuantPattern(

605 - epsilon, FP8_DTYPE, GroupShape(1, 128)

608 + epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op

609 ).register(self.patterns)

610

611 # Fuse aiter fused_add_rms_norm + aiter dynamic group fp8 quant

612 AiterFusedAddRMSFp8GroupQuantPattern(

610 - epsilon, FP8_DTYPE, GroupShape(1, 128)

613 + epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op

614 ).register(self.patterns)

615

616 # When quant_fp8 custom ops are disabled, both AITER and native

@@ -620,6 +623,8 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):

623 match_aiter_quant_options = (

624 [True, False] if is_quant_fp8_enabled else [False]

625 )

626 + if not match_aiter_quant_op:

627 + match_aiter_quant_options = []

628

629 for match_aiter_quant in match_aiter_quant_options:

630 # Fuse aiter rms_norm + (aiter / vllm built-in)

@@ -637,7 +642,10 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):

642 # Fuse decomposed RMSNormGated + group fp8 quant.

643 # The replacement op (fused_rms_gated_fp8_group_quant) requires

644 # an aiter version that includes the GDN triton kernel renames.

640 - if gated_norm_shapes and rocm_aiter_ops.are_gdn_triton_kernels_available():

645 + if gated_norm_shapes and (

646 + rocm_aiter_ops.are_gdn_triton_kernels_available()

647 + or rocm_aiter_ops.is_rdna_gdn_triton_kernels_available()

648 + ):

649 for num_heads, head_dim in gated_norm_shapes:

650 if head_dim != 128:

651 continue

@@ -647,6 +655,7 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):

655 GroupShape(1, 128),

656 num_heads=num_heads,

657 head_dim=head_dim,

658 + match_aiter_quant=match_aiter_quant_op,

659 ).register(self.patterns)

660

661 self.dump_patterns(config, self.patterns)

@@ -679,10 +688,10 @@ class AiterSiluMulFp8GroupQuantPattern(VllmPatternReplacement):

688

689 FUSED_SILU_MUL_QUANT_OP = rocm_aiter_ops.get_act_mul_fused_fp8_group_quant_op()

690

682 - def __init__(self) -> None:

691 + def __init__(self, match_aiter_quant_op: bool = True) -> None:

692 self.silu_and_mul_matcher = MatcherSiluAndMul()

693 self.quant_matcher = MatcherQuantFP8(

685 - quant_key=kFp8Dynamic128Sym, match_rocm_aiter=True

694 + quant_key=kFp8Dynamic128Sym, match_rocm_aiter=match_aiter_quant_op

695 )

696

697 def get_inputs(self) -> list[torch.Tensor]:

@@ -725,7 +734,11 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmFusionPatternMatcherPass):

734 def __init__(self, config: VllmConfig) -> None:

735 super().__init__(config, "rocm_aiter_silu_mul_fp8_group_quant_fusion_pass")

736

728 - self.register(AiterSiluMulFp8GroupQuantPattern())

737 + match_aiter_quant_op = not rocm_aiter_ops.is_rdna_aiter_enabled()

738 +

739 + self.register(

740 + AiterSiluMulFp8GroupQuantPattern(match_aiter_quant_op=match_aiter_quant_op)

741 + )

742

743 self.dump_patterns(config, self.pm_pass)

744

vllm/compilation/passes/pass_manager.py

+12 -4

@@ -18,7 +18,7 @@ from .ir.clone_elimination import UnsafeCloneEliminationPass

18 from .ir.lowering_pass import VllmIRLoweringPass

19 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

21 -if rocm_aiter_ops.is_enabled():

21 +if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():

22 from .fusion.allreduce_rms_fusion import (

23 RocmAiterAllReduceFusionPass,

24 )

@@ -166,7 +166,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

166 if enable_transformers_norm_canonicalization:

167 self.passes += [AddRMSNormFusionPass(config)]

168

169 - if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():

169 + if self.pass_config.fuse_act_padding and (

170 + rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled()

171 + ):

172 # Run the more specific RMSNorm+router-pad fusion before

173 # AR+RMS, since both consume fused_add_rms_norm.

174 self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]

@@ -183,7 +185,10 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

185 self.passes += [RMSNormReshapeFusionPass(config)]

186

187 if self.pass_config.fuse_norm_quant:

186 - if rocm_aiter_ops.is_enabled():

188 + if (

189 + rocm_aiter_ops.is_enabled()

190 + or rocm_aiter_ops.is_rdna_aiter_enabled()

191 + ):

192 self.passes += [

193 RocmAiterRMSNormQuantFusionPass(config),

194 ]

@@ -191,7 +196,10 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

196

197 if self.pass_config.fuse_act_quant:

198 self.passes += [ActivationQuantFusionPass(config)]

194 - if rocm_aiter_ops.is_enabled():

199 + if (

200 + rocm_aiter_ops.is_enabled()

201 + or rocm_aiter_ops.is_rdna_aiter_enabled()

202 + ):

203 self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]

204

205 if self.pass_config.fuse_qk_norm_rope_kvcache:

vllm/model_executor/kernels/linear/scaled_mm/aiter.py

+18 -3

@@ -397,10 +397,14 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

397

398 @classmethod

399 def is_supported(cls, compute_capability=None):

400 + if (

401 + rocm_aiter_ops.is_linear_enabled()

402 + or rocm_aiter_ops.is_rdna_linear_enabled()

403 + ):

404 + return True, None

405 return (

401 - rocm_aiter_ops.is_linear_enabled(),

402 - "Only supported on ROCm platform \

403 - with aiter package installed.",

406 + False,

407 + "Only supported on ROCm platform with aiter package installed.",

408 )

409

410 @classmethod

@@ -416,6 +420,17 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

420 "Supports only dynamic per token group activation "

421 "quantization with group_shape=(1,128).",

422 )

423 +

424 + # RDNA4 (gfx12) only has the aiter Triton blockscale backend, which

425 + # needs a per-(N,K) tune. Reject untuned shapes so the dispatcher falls

426 + # through to the generic backend.

427 + if rocm_aiter_ops.is_rdna_linear_enabled():

428 + n, k = config.weight_shape

429 + if not rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k):

430 + return (

431 + False,

432 + f"(N={n}, K={k}) is not in the aiter Triton blockscale tuned list.",

433 + )

434 return True, None

435

436 def apply_block_scaled_mm(

vllm/model_executor/kernels/linear/scaled_mm/pytorch.py

+3 -3

@@ -108,10 +108,10 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):

108 if not current_platform.is_rocm():

109 return False, "requires ROCm."

110

111 - from vllm.platforms.rocm import get_cdna_version

111 + from vllm.platforms.rocm import get_cdna_version, on_rdna4

112

113 - if get_cdna_version() <= 2:

114 - return False, "requires CDNA3+"

113 + if get_cdna_version() <= 2 and not on_rdna4():

114 + return False, "requires CDNA3+ or RDNA4"

115

116 if compute_capability is not None and compute_capability < 94:

117 return False, "requires compute capability 94 and above."

vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py

+4 -4

@@ -814,13 +814,13 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):

814 ) -> bool:

815 p = current_platform

816 if p.is_rocm():

817 - from vllm.platforms.rocm import get_cdna_version

817 + from vllm.platforms.rocm import get_cdna_version, on_rdna4

818

819 - _rocm_support_fp8 = get_cdna_version() > 2

819 + is_rocm_fp8_supported = get_cdna_version() > 2 or on_rdna4()

820 else:

821 - _rocm_support_fp8 = False

821 + is_rocm_fp8_supported = False

822

823 - device_supports_fp8 = _rocm_support_fp8 or (

823 + device_supports_fp8 = is_rocm_fp8_supported or (

824 p.is_cuda() and p.has_device_capability((8, 9))

825 )

826

vllm/model_executor/layers/fused_moe/oracle/fp8.py

+8 -2

@@ -379,8 +379,14 @@ def select_fp8_moe_backend(

379

380 # Handle explicit AITER FP8 configuration.

381 if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"):

382 - if not envs.VLLM_ROCM_USE_AITER or not envs.VLLM_ROCM_USE_AITER_MOE:

383 - AVAILABLE_BACKENDS.remove(Fp8MoeBackend.AITER)

382 + skip_aiter_moe = (

383 + not envs.VLLM_ROCM_USE_AITER

384 + or not envs.VLLM_ROCM_USE_AITER_MOE

385 + or rocm_aiter_ops.is_rdna_aiter_enabled()

386 + )

387 + if skip_aiter_moe:

388 + if Fp8MoeBackend.AITER in AVAILABLE_BACKENDS:

389 + AVAILABLE_BACKENDS.remove(Fp8MoeBackend.AITER)

390 else:

391 backend = Fp8MoeBackend.AITER

392 return _return_or_raise(

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

+6 -1

@@ -298,7 +298,12 @@ def select_unquantized_moe_backend(

298

299 # Handle explicit AITER FP8 configuration.

300 if envs.is_set("VLLM_ROCM_USE_AITER") or envs.is_set("VLLM_ROCM_USE_AITER_MOE"):

301 - if not envs.VLLM_ROCM_USE_AITER or not envs.VLLM_ROCM_USE_AITER_MOE:

301 + skip_aiter_moe = (

302 + not envs.VLLM_ROCM_USE_AITER

303 + or not envs.VLLM_ROCM_USE_AITER_MOE

304 + or rocm_aiter_ops.is_rdna_aiter_enabled()

305 + )

306 + if skip_aiter_moe:

307 if UnquantizedMoeBackend.AITER in AVAILABLE_BACKENDS:

308 AVAILABLE_BACKENDS.remove(UnquantizedMoeBackend.AITER)

309 else:

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

+4 -1

@@ -69,7 +69,10 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

69 # Availability is checked centrally via rocm_aiter_ops; the actual function

70 # references are imported here so that they can be called without per-call

71 # import overhead.

72 -GDN_AITER_TRITON_AVAILABLE = rocm_aiter_ops.are_gdn_triton_kernels_available()

72 +GDN_AITER_TRITON_AVAILABLE = (

73 + rocm_aiter_ops.are_gdn_triton_kernels_available()

74 + or rocm_aiter_ops.is_rdna_gdn_triton_kernels_available()

75 +)

77 if GDN_AITER_TRITON_AVAILABLE:

78 from aiter.ops.triton.causal_conv1d_update_single_token import (

vllm/model_executor/layers/sparse_attn_indexer.py

+1 -1

@@ -850,7 +850,7 @@ class SparseAttnIndexer(CustomOp):

850 assert isinstance(q_quant, torch.Tensor), (

851 "AMD sparse_attn_indexer expects a single FP8 q_quant tensor"

852 )

853 - if rocm_aiter_ops.is_enabled():

853 + if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():

854 return torch.ops.vllm.rocm_aiter_sparse_attn_indexer(

855 hidden_states,

856 _encode_layer_name(self.k_cache.prefix),

vllm/platforms/rocm.py

+9 -1

@@ -218,6 +218,7 @@ _ON_GFX1250 = "gfx1250" in _GCN_ARCH

218 _ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9", "gfx1250"])

219 # RDNA = gfx11/gfx12 minus the CDNA-classified gfx1250.

220 _ON_RDNA = _ON_GFX1X and not _ON_CDNA

221 +_ON_RDNA4 = any(arch in _GCN_ARCH for arch in ["gfx1200", "gfx1201"])

222

223

224 def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:

@@ -315,6 +316,10 @@ def on_gfx1250() -> bool:

316 return _ON_GFX1250

317

318

319 +def on_rdna4() -> bool:

320 + return _ON_RDNA4

321 +

322 +

323 def on_mi3xx() -> bool:

324 return _ON_MI3XX

325

@@ -471,6 +476,8 @@ def _get_backend_priorities(

476 backends.append(AttentionBackendEnum.ROCM_AITER_FA)

477 if is_aiter_found_and_supported():

478 backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)

479 + elif rocm_aiter_ops.is_rdna_aiter_enabled():

480 + backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)

481 backends.append(AttentionBackendEnum.TRITON_ATTN)

482 backends.append(AttentionBackendEnum.TURBOQUANT)

483

@@ -918,7 +925,7 @@ class RocmPlatform(Platform):

925

926 @classmethod

927 def supports_fp8(cls) -> bool:

921 - return on_cdna() or on_gfx12x()

928 + return on_cdna() or on_rdna4()

929

930 @classmethod

931 def is_fp8_fnuz(cls) -> bool:

@@ -1065,6 +1072,7 @@ class RocmPlatform(Platform):

1072 cc.cudagraph_mode != CUDAGraphMode.NONE

1073 and envs.VLLM_ROCM_USE_AITER

1074 and envs.VLLM_ROCM_USE_AITER_RMSNORM

1075 + and not on_rdna4()

1076 ):

1077 rms_norm = ["aiter"] + default

1078 else:

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

+2 -2

@@ -417,7 +417,7 @@ def rocm_fp8_paged_mqa_logits(

417 batch_size, next_n = q_fp8.shape[:2]

418 block_size = kv_cache_fp8.shape[1]

419

420 - if rocm_aiter_ops.is_enabled():

420 + if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():

421 aiter_paged_mqa_logits_module = paged_mqa_logits_module()

422

423 if aiter_paged_mqa_logits_module is not None:

@@ -582,7 +582,7 @@ def rocm_fp8_mqa_logits(

582 )

583

584 aiter_mqa_logits_module = None

585 - if rocm_aiter_ops.is_enabled():

585 + if rocm_aiter_ops.is_enabled() or rocm_aiter_ops.is_rdna_aiter_enabled():

586 aiter_mqa_logits_module = mqa_logits_module()

587

588 if aiter_mqa_logits_module is not None: