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 15, 03:30 UTC
[BugFix][V1][ROCm] Triton MLA uses V0 backend on V1 engine (#19067)
Signed-off-by: Tianyuan Wu <Tianyuan.Wu@amd.com>
TY-AMD committed
Jul 1, 2025 at 16:12 UTC
96453cfa831340788ef72c42bc2a1a2b4496a27f
5 files changed
+78
-10
tests/kernels/attention/test_attention_selector.py
+2
-4
@@ -106,10 +106,8 @@ def test_env(
106
block_size,
107
False,
108
use_mla=use_mla)
109
- if use_v1 and name != "TRITON_MLA":
110
- assert backend.get_name() == f"{name}_VLLM_V1"
111
- else:
112
- assert backend.get_name() == name
109
+ expected = f"{name}_VLLM_V1" if use_v1 else name
110
+ assert backend.get_name() == expected
111
else:
112
with pytest.raises(ValueError) as exc_info:
113
get_attn_backend(16,
tests/kernels/attention/test_rocm_attention_selector.py
+4
-2
@@ -35,7 +35,8 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
35
m.setenv(STR_BACKEND_ENV_VAR, "TRITON_MLA")
36
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
37
False, True)
38
- assert backend.get_name() == "TRITON_MLA"
38
+ assert (backend.get_name() == "TRITON_MLA"
39
+ or backend.get_name() == "TRITON_MLA_VLLM_V1")
40
41
# If attention backend is None
42
# If use_mla is true
@@ -43,7 +44,8 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
44
m.setenv(STR_BACKEND_ENV_VAR, None)
45
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
46
False, True)
46
- assert backend.get_name() == "TRITON_MLA"
47
+ assert (backend.get_name() == "TRITON_MLA"
48
+ or backend.get_name() == "TRITON_MLA_VLLM_V1")
49
50
# change the attention backend to AITER MLA
51
m.setenv(STR_BACKEND_ENV_VAR, "ROCM_AITER_MLA")
vllm/platforms/rocm.py
+8
-2
@@ -186,8 +186,14 @@ class RocmPlatform(Platform):
186
187
if selected_backend == _Backend.TRITON_MLA:
188
if block_size != 1:
189
- logger.info("Using Triton MLA backend.")
190
- return "vllm.attention.backends.triton_mla.TritonMLABackend" # noqa: E501
189
+ if use_v1:
190
+ logger.info_once(
191
+ "Using Triton MLA backend on V1 engine.")
192
+ return ("vllm.v1.attention.backends.mla."
193
+ "triton_mla.TritonMLABackend")
194
+ else:
195
+ logger.info("Using Triton MLA backend.")
196
+ return "vllm.attention.backends.triton_mla.TritonMLABackend" # noqa: E501
197
else:
198
raise ValueError(
199
f" The selected backend, {selected_backend.name},"
vllm/v1/attention/backends/mla/common.py
+7
-2
@@ -640,7 +640,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
640
self.qk_head_dim = qk_head_dim
641
self.v_head_dim = v_head_dim
642
self.kv_b_proj = kv_b_proj
643
- self.vllm_flash_attn_version = get_flash_attn_version()
643
644
# Handle the differences between the flash_attn_varlen from flash_attn
645
# and the one from vllm_flash_attn. The former is used on RoCM and the
@@ -672,11 +671,17 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
671
maybe_padded_v = torch.nn.functional.pad(
672
v, [0, q.shape[-1] - v.shape[-1]], value=0)
673
674
+ if is_vllm_fa:
675
+ kwargs["return_softmax_lse"] = return_softmax_lse
676
+ else:
677
+ # ROCm leverages the upstream flash_attn, which takes a parameter
678
+ # called "return_attn_probs" instead of return_softmax_lse
679
+ kwargs["return_attn_probs"] = return_softmax_lse
680
+
681
attn_out = self.flash_attn_varlen_func(
682
q=q,
683
k=k,
684
v=maybe_padded_v,
679
- return_softmax_lse=return_softmax_lse,
685
softmax_scale=softmax_scale,
686
**kwargs,
687
)
vllm/v1/attention/backends/mla/triton_mla.py
+57
@@ -5,10 +5,14 @@ from typing import Any, Optional
5
6
import torch
7
8
+from vllm import envs
9
from vllm.attention.backends.abstract import (AttentionType,
10
is_quantized_kv_cache)
11
from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
12
+from vllm.attention.ops.triton_flash_attention import triton_attention
13
from vllm.logger import init_logger
14
+from vllm.platforms import current_platform
15
+from vllm.triton_utils import HAS_TRITON
16
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
17
MLACommonImpl,
18
MLACommonMetadata)
@@ -68,6 +72,59 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
72
raise NotImplementedError(
73
"TritonMLA V1 with FP8 KV cache not yet supported")
74
75
+ self.use_triton_flash_attn = envs.VLLM_USE_TRITON_FLASH_ATTN
76
+ self.triton_fa_func = triton_attention if HAS_TRITON else None
77
+
78
+ def _flash_attn_varlen_diff_headdims_rocm(self,
79
+ q,
80
+ k,
81
+ v,
82
+ softmax_scale=None,
83
+ **kwargs):
84
+ assert self.triton_fa_func is not None
85
+
86
+ # Triton Attention requires a padded V
87
+ padded_v = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
88
+ value=0)
89
+ # The output of triton_attention is a tuple of
90
+ # [output_tensor, encoded_softmax] where encoded_softmax is always None
91
+ output_tensor, _ = self.triton_fa_func(
92
+ q,
93
+ k,
94
+ padded_v,
95
+ None, # output
96
+ kwargs["cu_seqlens_q"],
97
+ kwargs["cu_seqlens_k"],
98
+ kwargs["max_seqlen_q"],
99
+ kwargs["max_seqlen_k"],
100
+ kwargs["causal"],
101
+ softmax_scale,
102
+ None, # bias
103
+ )
104
+
105
+ return output_tensor
106
+
107
+ def _flash_attn_varlen_diff_headdims(self,
108
+ q,
109
+ k,
110
+ v,
111
+ return_softmax_lse=False,
112
+ softmax_scale=None,
113
+ **kwargs):
114
+ if current_platform.is_rocm() \
115
+ and self.use_triton_flash_attn \
116
+ and not return_softmax_lse:
117
+ return self._flash_attn_varlen_diff_headdims_rocm(
118
+ q, k, v, softmax_scale=softmax_scale, **kwargs)
119
+ else:
120
+ return super()._flash_attn_varlen_diff_headdims(
121
+ q,
122
+ k,
123
+ v,
124
+ return_softmax_lse=return_softmax_lse,
125
+ softmax_scale=softmax_scale,
126
+ **kwargs)
127
+
128
def _forward_decode(
129
self,
130
q_nope: torch.Tensor,