A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 10:00 UTC
[ROCm] [AITER] [DOC] Add usage description about check functions in `_aiter_ops` (#30586)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Dec 16, 2025 at 21:50 UTC
d0fb5729298eb0b3683445dac566764b093e5854
2 files changed
+88
-18
vllm/_aiter_ops.py
+88
-15
@@ -642,48 +642,130 @@ _OPS_REGISTERED = False
642
643
644
class rocm_aiter_ops:
645
+ """ROCm AITER operations wrapper for AMD GPU acceleration in vLLM.
646
+
647
+ This class centralizes the import and registration of AITER ops,
648
+ and provides a unified interface for checking if AITER is enabled.
649
+ Operations are only available on supported gfx9
650
+ architectures when aiter is installed.
651
+
652
+ The class uses environment variables to control which features are enabled,
653
+ allowing fine-grained control over which AITER optimizations are used.
654
+
655
+ Environment Variables:
656
+ VLLM_ROCM_USE_AITER: Main toggle for all AITER operations.
657
+ VLLM_ROCM_USE_AITER_LINEAR: Controls GEMM and quantization ops.
658
+ VLLM_ROCM_USE_AITER_RMSNORM: Controls RMSNorm operations.
659
+ VLLM_ROCM_USE_AITER_MOE: Controls MoE (Mixture of Experts) ops.
660
+ VLLM_ROCM_USE_AITER_MLA: Controls MLA (Multi-head Latent Attention) ops.
661
+ VLLM_ROCM_USE_AITER_MHA: Controls MHA ops including flash_attn_varlen.
662
+ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: Controls Triton unified attention.
663
+ VLLM_ROCM_USE_AITER_FP8BMM: Controls FP8 batched matrix multiply.
664
+ VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: Controls FP4 assembly GEMM.
665
+ VLLM_ROCM_USE_AITER_TRITON_ROPE: Controls Triton rotary embeddings.
666
+ VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: Controls shared expert fusion.
667
+ VLLM_ROCM_USE_AITER_TRITON_GEMM: Controls Triton unquantized GEMM.
668
+
669
+ Note:
670
+ The environment variables are assigned when the module is imported,
671
+ so you can't change the environment variables after the module is imported.
672
+ This is done out of performance consideration. Accessing environment variables
673
+ is expensive as described in issue https://github.com/vllm-project/vllm/issues/17067
674
+ so we don't want to do it repeatedly, especially in the hot path (the forward pass).
675
+ You can call the refresh_env_variables() function to reload the env variables
676
+ after monkey patching the env variables in the unit test.
677
+
678
+ Check Functions:
679
+ All check functions (is_*_enabled) are decorated with @if_aiter_supported,
680
+ which verifies: (1) platform is ROCm, (2) device arch is gfx9, and
681
+ (3) aiter library is installed. The check function then also verifies
682
+ the corresponding environment variable is enabled.
683
+ i.e. ___
684
+ is_enabled() == current_platform.is_rocm() and | checked by
685
+ current_platform.is_on_gfx9() and | @if_aiter_supported
686
+ IS_AITER_FOUND and _______________|
687
+ cls._AITER_ENABLED -----> Check by the logic in `is_enabled()`
688
+
689
+ Example:
690
+ from vllm._aiter_ops import rocm_aiter_ops
691
+
692
+ # Check if aiter is enabled before using operations
693
+ if rocm_aiter_ops.is_enabled():
694
+ result = rocm_aiter_ops.rms_norm(x, weight, epsilon)
695
+
696
+ Operations:
697
+ - RMS normalization: rms_norm, rms_norm2d_with_add
698
+ - GEMM operations: gemm_a8w8, gemm_a8w8_blockscale
699
+ - Fused MoE: fused_moe, asm_moe_tkw1
700
+ - Routing: topk_softmax, biased_grouped_topk, grouped_topk
701
+ - MLA decode: mla_decode_fwd
702
+ - Quantization: per_tensor_quant, per_token_quant, group_fp8_quant
703
+ - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
704
+ """
705
+
706
+ # Check if the env variable is set
707
_AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
708
_LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
709
_RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
710
_FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
711
_MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
650
- _PG_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_PAGED_ATTN
712
_MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
713
_TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
714
+ # TODO: Consolidate under _LINEAR_ENABLED
715
_FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
716
+ # TODO: Consolidate under _LINEAR_ENABLED
717
_FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
718
+ # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE
719
_TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
720
_MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
721
+ # TODO: Consolidate under _LINEAR_ENABLED
722
_TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
723
724
+ @classmethod
725
+ def refresh_env_variables(cls):
726
+ """
727
+ Since the environment variables are assigned when the module is imported,
728
+ This is a helper function to reload all the env variables from
729
+ the environment variables.
730
+ for example, after monkey patching the env variables in the unit test,
731
+ you can call this function to reload the env variables.
732
+ """
733
+ cls._AITER_ENABLED = envs.VLLM_ROCM_USE_AITER
734
+ cls._LINEAR_ENABLED = envs.VLLM_ROCM_USE_AITER_LINEAR
735
+ cls._RMSNORM_ENABLED = envs.VLLM_ROCM_USE_AITER_RMSNORM
736
+ cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE
737
+ cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA
738
+ cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA
739
+ cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
740
+ cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM
741
+ cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM
742
+ cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE
743
+ cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS
744
+ cls._TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
745
+
746
@classmethod
747
@if_aiter_supported
748
def is_enabled(cls) -> bool:
662
- """Verifies device specs and availability of aiter main env variable."""
749
return cls._AITER_ENABLED
750
751
@classmethod
752
@if_aiter_supported
753
def is_linear_enabled(cls) -> bool:
668
- """ "Verifies device specs and availability of env variable."""
754
return cls._AITER_ENABLED and cls._LINEAR_ENABLED
755
756
@classmethod
757
@if_aiter_supported
758
def is_linear_fp8_enaled(cls) -> bool:
674
- """ "Verifies device specs and availability of env variable."""
759
return cls.is_linear_enabled()
760
761
@classmethod
762
@if_aiter_supported
763
def is_rmsnorm_enabled(cls) -> bool:
680
- """ "Verifies device specs and availability of env variable."""
764
return cls._AITER_ENABLED and cls._RMSNORM_ENABLED
765
766
@classmethod
767
@if_aiter_supported
768
def is_fused_moe_enabled(cls) -> bool:
686
- """ "Verifies device specs and availability of env variable."""
769
return cls._AITER_ENABLED and cls._FMOE_ENABLED
770
771
@classmethod
@@ -694,25 +776,16 @@ class rocm_aiter_ops:
776
@classmethod
777
@if_aiter_supported
778
def is_mla_enabled(cls) -> bool:
697
- """ "Verifies device specs and availability of env variable."""
779
return cls._AITER_ENABLED and cls._MLA_ENABLED
780
781
@classmethod
782
@if_aiter_supported
783
def is_mha_enabled(cls) -> bool:
703
- """ "Verifies device specs and availability of env variable."""
784
return cls._AITER_ENABLED and cls._MHA_ENABLED
785
706
- @classmethod
707
- @if_aiter_supported
708
- def is_pa_attn_enabled(cls) -> bool:
709
- """ "Verifies device specs and availability of env variable."""
710
- return cls._AITER_ENABLED and cls._PG_ATTN_ENABLED
711
-
786
@classmethod
787
@if_aiter_supported
788
def is_triton_unified_attn_enabled(cls) -> bool:
715
- """ "Verifies device specs and availability of env variable."""
789
return cls._AITER_ENABLED and cls._TRITON_UNIFIED_ATTN_ENABLED
790
791
@classmethod
vllm/platforms/rocm.py
-3
@@ -124,8 +124,6 @@ def use_rocm_custom_paged_attention(
124
alibi_slopes: torch.Tensor | None = None,
125
sinks: torch.Tensor | None = None,
126
) -> bool:
127
- from vllm._aiter_ops import rocm_aiter_ops
128
-
127
GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
128
ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
129
ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
@@ -141,7 +139,6 @@ def use_rocm_custom_paged_attention(
139
and (gqa_ratio >= 1 and gqa_ratio <= 16)
140
and max_seq_len <= 128 * 1024
141
and (envs.VLLM_ROCM_CUSTOM_PAGED_ATTN)
144
- and not (rocm_aiter_ops.is_pa_attn_enabled())
142
and sinks is None
143
)
144