Skip to content

Commit f0c8e63

Browse files
fsx950223fsx950223
committed
Fa upstream3 (#561)
* integrate aiter Signed-off-by: fsx950223 <fsx950223@outlook.com> * add env variable Signed-off-by: fsx950223 <fsx950223@outlook.com> * rename function Signed-off-by: fsx950223 <fsx950223@outlook.com> * optimize kernels with small query lens Signed-off-by: fsx950223 <fsx950223@outlook.com> * change condition Signed-off-by: fsx950223 <fsx950223@outlook.com> * add rocm aiter backend Signed-off-by: fsx950223 <fsx950223@outlook.com> * new fa impl Signed-off-by: fsx950223 <fsx950223@outlook.com> * update api Signed-off-by: fsx950223 <fsx950223@outlook.com> * optimize performance Signed-off-by: fsx950223 <fsx950223@outlook.com> * remove try catch Signed-off-by: fsx950223 <fsx950223@outlook.com> * clean code Signed-off-by: fsx950223 <fsx950223@outlook.com> --------- Signed-off-by: fsx950223 <fsx950223@outlook.com>
1 parent 1068556 commit f0c8e63

4 files changed

Lines changed: 638 additions & 28 deletions

File tree

‎vllm/envs.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
VLLM_ROCM_USE_AITER_MOE: bool = True
8383
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
8484
VLLM_ROCM_USE_AITER_MLA: bool = True
85+
VLLM_ROCM_USE_AITER_MHA: bool = True
8586
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
8687
VLLM_ROCM_FP8_PADDING: bool = True
8788
VLLM_ROCM_MOE_PADDING: bool = True
@@ -623,6 +624,13 @@ def get_vllm_port() -> Optional[int]:
623624
"VLLM_ROCM_USE_AITER_MLA":
624625
lambda: (os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in
625626
("true", "1")),
627+
628+
# Whether to use aiter mha ops.
629+
# By default is enabled.
630+
"VLLM_ROCM_USE_AITER_MHA":
631+
lambda: (os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in
632+
("true", "1")),
633+
626634
# use rocm skinny gemms
627635
"VLLM_ROCM_USE_SKINNY_GEMM":
628636
lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in

‎vllm/model_executor/layers/layernorm.py‎

Lines changed: 51 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import vllm.envs as envs
99
from vllm.model_executor.custom_op import CustomOp
1010
from vllm.platforms import current_platform
11+
from vllm.utils import direct_register_custom_op
1112

1213

1314
def is_rocm_aiter_rmsnorm_enabled() -> bool:
@@ -42,46 +43,71 @@ def fused_add_rms_norm(
4243
return x, residual
4344

4445

45-
def rocm_aiter_rms_norm(x: torch.Tensor, weight: torch.Tensor,
46-
variance_epsilon: float) -> torch.Tensor:
46+
if is_rocm_aiter_rmsnorm_enabled():
4747

48-
import aiter as rocm_aiter
49-
if x.dim() > 2:
50-
x_original_shape = x.shape
51-
x = x.reshape(-1, x_original_shape[-1])
52-
x = rocm_aiter.rms_norm(x, weight, variance_epsilon)
53-
return x.reshape(x_original_shape)
48+
def rocm_aiter_rms_norm_impl(x: torch.Tensor, weight: torch.Tensor,
49+
variance_epsilon: float) -> torch.Tensor:
5450

55-
return rocm_aiter.rms_norm(x, weight, variance_epsilon)
51+
import aiter as rocm_aiter
52+
if x.dim() > 2:
53+
x_original_shape = x.shape
54+
x = x.reshape(-1, x_original_shape[-1])
55+
x = rocm_aiter.rms_norm(x, weight, variance_epsilon)
56+
return x.reshape(x_original_shape)
5657

58+
return rocm_aiter.rms_norm(x, weight, variance_epsilon)
5759

58-
def rocm_aiter_fused_add_rms_norm(
59-
x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
60-
variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
60+
def rocm_aiter_rms_norm_fake(input: torch.Tensor, weight: torch.Tensor,
61+
variance_epsilon: float) -> torch.Tensor:
62+
return input.clone()
6163

62-
import aiter as rocm_aiter
64+
direct_register_custom_op(
65+
op_name="rocm_aiter_rms_norm",
66+
op_func=rocm_aiter_rms_norm_impl,
67+
mutates_args=[],
68+
fake_impl=rocm_aiter_rms_norm_fake,
69+
dispatch_key=current_platform.dispatch_key,
70+
)
6371

64-
residual_out = torch.empty_like(residual)
65-
output = torch.empty_like(x)
66-
rocm_aiter.rmsnorm2d_fwd_with_add(
67-
output, # output
68-
x, # input
69-
residual, # residual input
70-
residual_out, # residual output
71-
weight,
72-
variance_epsilon,
72+
def rocm_aiter_fused_add_rms_norm_impl(
73+
x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
74+
variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
75+
76+
import aiter as rocm_aiter
77+
residual_out = torch.empty_like(residual)
78+
output = torch.empty_like(x)
79+
rocm_aiter.rmsnorm2d_fwd_with_add(
80+
output, # output
81+
x, # input
82+
residual, # residual input
83+
residual_out, # residual output
84+
weight,
85+
variance_epsilon,
86+
)
87+
return output, residual_out
88+
89+
def rocm_aiter_fused_add_rms_norm_fake(
90+
x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
91+
variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
92+
return x.clone(), residual.clone()
93+
94+
direct_register_custom_op(
95+
op_name="rocm_aiter_fused_add_rms_norm",
96+
op_func=rocm_aiter_fused_add_rms_norm_impl,
97+
mutates_args=[],
98+
fake_impl=rocm_aiter_fused_add_rms_norm_fake,
99+
dispatch_key=current_platform.dispatch_key,
73100
)
74-
return output, residual_out
75101

76102

77103
def dispatch_cuda_rmsnorm_func(add_residual: bool):
78104
if add_residual:
79105
if is_rocm_aiter_rmsnorm_enabled():
80-
return rocm_aiter_fused_add_rms_norm
106+
return torch.ops.vllm.rocm_aiter_fused_add_rms_norm
81107
return fused_add_rms_norm
82108

83109
if is_rocm_aiter_rmsnorm_enabled():
84-
return rocm_aiter_rms_norm
110+
return torch.ops.vllm.rocm_aiter_rms_norm
85111
return rms_norm
86112

87113

‎vllm/platforms/rocm.py‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -197,9 +197,15 @@ def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
197197
selected_backend = (_Backend.ROCM_FLASH if selected_backend
198198
== _Backend.FLASH_ATTN else selected_backend)
199199
if envs.VLLM_USE_V1:
200-
logger.info("Using Triton Attention backend on V1 engine.")
201-
return ("vllm.v1.attention.backends."
202-
"triton_attn.TritonAttentionBackend")
200+
if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA \
201+
and on_mi250_mi300():
202+
logger.info("Using Flash Attention backend on V1 engine.")
203+
return ("vllm.v1.attention.backends."
204+
"rocm_aiter_fa.AiterFlashAttentionBackend")
205+
else:
206+
logger.info("Using Triton Attention backend on V1 engine.")
207+
return ("vllm.v1.attention.backends."
208+
"triton_attn.TritonAttentionBackend")
203209
if selected_backend == _Backend.ROCM_FLASH:
204210
if not cls.has_device_capability(90):
205211
# not Instinct series GPUs.

0 commit comments

Comments
 (0)