[Hardware][AMD] integrate aiter chunked prefill into vllm

vllm/envs.py

@@ -82,6 +82,7 @@
     VLLM_ROCM_USE_AITER_MOE: bool = True
     VLLM_ROCM_USE_AITER_RMSNORM: bool = True
     VLLM_ROCM_USE_AITER_MLA: bool = True
+    VLLM_ROCM_USE_AITER_MHA: bool = True
     VLLM_ROCM_USE_SKINNY_GEMM: bool = True
     VLLM_ROCM_FP8_PADDING: bool = True
     VLLM_ROCM_MOE_PADDING: bool = True
@@ -624,6 +625,13 @@ def get_vllm_port() -> Optional[int]:
     "VLLM_ROCM_USE_AITER_MLA":
     lambda: (os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in
              ("true", "1")),
+
+    # Whether to use aiter mha ops.
+    # By default is enabled.
+    "VLLM_ROCM_USE_AITER_MHA":
+    lambda: (os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in
+             ("true", "1")),
+
     # use rocm skinny gemms
     "VLLM_ROCM_USE_SKINNY_GEMM":
     lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in

vllm/model_executor/layers/layernorm.py

@@ -8,6 +8,7 @@
 import vllm.envs as envs
 from vllm.model_executor.custom_op import CustomOp
 from vllm.platforms import current_platform
+from vllm.utils import direct_register_custom_op
 
 
 def is_rocm_aiter_rmsnorm_enabled() -> bool:
@@ -42,46 +43,71 @@ def fused_add_rms_norm(
     return x, residual
 
 
-def rocm_aiter_rms_norm(x: torch.Tensor, weight: torch.Tensor,
-                        variance_epsilon: float) -> torch.Tensor:
+if is_rocm_aiter_rmsnorm_enabled():
 
-    import aiter as rocm_aiter
-    if x.dim() > 2:
-        x_original_shape = x.shape
-        x = x.reshape(-1, x_original_shape[-1])
-        x = rocm_aiter.rms_norm(x, weight, variance_epsilon)
-        return x.reshape(x_original_shape)
+    def rocm_aiter_rms_norm_impl(x: torch.Tensor, weight: torch.Tensor,
+                                 variance_epsilon: float) -> torch.Tensor:
 
-    return rocm_aiter.rms_norm(x, weight, variance_epsilon)
+        import aiter as rocm_aiter
+        if x.dim() > 2:
+            x_original_shape = x.shape
+            x = x.reshape(-1, x_original_shape[-1])
+            x = rocm_aiter.rms_norm(x, weight, variance_epsilon)
+            return x.reshape(x_original_shape)
 
+        return rocm_aiter.rms_norm(x, weight, variance_epsilon)
 
-def rocm_aiter_fused_add_rms_norm(
-        x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
-        variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
+    def rocm_aiter_rms_norm_fake(input: torch.Tensor, weight: torch.Tensor,
+                                 variance_epsilon: float) -> torch.Tensor:
+        return input.clone()
 
-    import aiter as rocm_aiter
+    direct_register_custom_op(
+        op_name="rocm_aiter_rms_norm",
+        op_func=rocm_aiter_rms_norm_impl,
+        mutates_args=[],
+        fake_impl=rocm_aiter_rms_norm_fake,
+        dispatch_key=current_platform.dispatch_key,
+    )
 
-    residual_out = torch.empty_like(residual)
-    output = torch.empty_like(x)
-    rocm_aiter.rmsnorm2d_fwd_with_add(
-        output,  # output
-        x,  # input
-        residual,  # residual input
-        residual_out,  # residual output
-        weight,
-        variance_epsilon,
+    def rocm_aiter_fused_add_rms_norm_impl(
+            x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+            variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
+
+        import aiter as rocm_aiter
+        residual_out = torch.empty_like(residual)
+        output = torch.empty_like(x)
+        rocm_aiter.rmsnorm2d_fwd_with_add(
+            output,  # output
+            x,  # input
+            residual,  # residual input
+            residual_out,  # residual output
+            weight,
+            variance_epsilon,
+        )
+        return output, residual_out
+
+    def rocm_aiter_fused_add_rms_norm_fake(
+            x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+            variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
+        return x.clone(), residual.clone()
+
+    direct_register_custom_op(
+        op_name="rocm_aiter_fused_add_rms_norm",
+        op_func=rocm_aiter_fused_add_rms_norm_impl,
+        mutates_args=[],
+        fake_impl=rocm_aiter_fused_add_rms_norm_fake,
+        dispatch_key=current_platform.dispatch_key,
     )
-    return output, residual_out
 
 
 def dispatch_cuda_rmsnorm_func(add_residual: bool):
     if add_residual:
         if is_rocm_aiter_rmsnorm_enabled():
-            return rocm_aiter_fused_add_rms_norm
+            return torch.ops.vllm.rocm_aiter_fused_add_rms_norm
         return fused_add_rms_norm
 
     if is_rocm_aiter_rmsnorm_enabled():
-        return rocm_aiter_rms_norm
+        return torch.ops.vllm.rocm_aiter_rms_norm
     return rms_norm
 
 

vllm/platforms/rocm.py

@@ -197,9 +197,15 @@ def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
         selected_backend = (_Backend.ROCM_FLASH if selected_backend
                             == _Backend.FLASH_ATTN else selected_backend)
         if envs.VLLM_USE_V1:
-            logger.info("Using Triton Attention backend on V1 engine.")
-            return ("vllm.v1.attention.backends."
-                    "triton_attn.TritonAttentionBackend")
+            if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA \
+                and on_mi250_mi300():
+                logger.info("Using Flash Attention backend on V1 engine.")
+                return ("vllm.v1.attention.backends."
+                        "rocm_aiter_fa.AiterFlashAttentionBackend")
+            else:
+                logger.info("Using Triton Attention backend on V1 engine.")
+                return ("vllm.v1.attention.backends."
+                        "triton_attn.TritonAttentionBackend")
         if selected_backend == _Backend.ROCM_FLASH:
             if not cls.has_device_capability(90):
                 # not Instinct series GPUs.