Skip to content

Commit 04a7022

Browse files
maleksan85fsx950223gshtras
authored
6 10 rc1 (#563)
* 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> --------- Signed-off-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: fsx950223 <fsx950223@outlook.com> Co-authored-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
1 parent 307d8bc commit 04a7022

4 files changed

Lines changed: 675 additions & 28 deletions

File tree

‎vllm/envs.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@
8787
VLLM_ROCM_USE_AITER_MOE: bool = True
8888
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
8989
VLLM_ROCM_USE_AITER_MLA: bool = True
90+
VLLM_ROCM_USE_AITER_MHA: bool = True
9091
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
9192
VLLM_ROCM_FP8_PADDING: bool = True
9293
VLLM_ROCM_MOE_PADDING: bool = True
@@ -658,6 +659,13 @@ def get_vllm_port() -> Optional[int]:
658659
"VLLM_ROCM_USE_AITER_MLA":
659660
lambda: (os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in
660661
("true", "1")),
662+
663+
# Whether to use aiter mha ops.
664+
# By default is enabled.
665+
"VLLM_ROCM_USE_AITER_MHA":
666+
lambda: (os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in
667+
("true", "1")),
668+
661669
# use rocm skinny gemms
662670
"VLLM_ROCM_USE_SKINNY_GEMM":
663671
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
@@ -198,9 +198,15 @@ def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
198198
selected_backend = (_Backend.ROCM_FLASH if selected_backend
199199
== _Backend.FLASH_ATTN else selected_backend)
200200
if envs.VLLM_USE_V1:
201-
logger.info("Using Triton Attention backend on V1 engine.")
202-
return ("vllm.v1.attention.backends."
203-
"triton_attn.TritonAttentionBackend")
201+
if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA \
202+
and on_mi250_mi300():
203+
logger.info("Using Flash Attention backend on V1 engine.")
204+
return ("vllm.v1.attention.backends."
205+
"rocm_aiter_fa.AiterFlashAttentionBackend")
206+
else:
207+
logger.info("Using Triton Attention backend on V1 engine.")
208+
return ("vllm.v1.attention.backends."
209+
"triton_attn.TritonAttentionBackend")
204210
if selected_backend == _Backend.ROCM_FLASH:
205211
if not cls.has_device_capability(90):
206212
# not Instinct series GPUs.

0 commit comments

Comments
 (0)