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, 23:30 UTC
[Bugfix] Massage MLA's usage of flash attn for RoCM (#13310)
Tyler Michael Smith committed
Feb 15, 2025 at 00:33 UTC
97a3d6d99522179d6adac188216bb942622a6d04
1 file changed
+11
-2
vllm/attention/backends/mla/utils.py
+11
-2
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import functools
4
from abc import abstractmethod
5
from dataclasses import dataclass
6
from typing import Any, Dict, Generic, List, Optional, Tuple
@@ -183,6 +184,15 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
184
self.o_proj = o_proj
185
self.vllm_flash_attn_version = get_flash_attn_version()
186
187
+ # Handle the differences between the flash_attn_varlen from flash_attn
188
+ # and the one from vllm_flash_attn. The former is used on RoCM and the
189
+ # latter has an additional parameter to control FA2 vs FA3
190
+ self.flash_attn_varlen_func = flash_attn_varlen_func
191
+ if self.vllm_flash_attn_version is not None:
192
+ self.flash_attn_varlen_func = \
193
+ functools.partial(flash_attn_varlen_func,
194
+ fa_version=self.vllm_flash_attn_version)
195
+
196
def _v_up_proj_and_o_proj(self, x):
197
if envs.VLLM_MLA_PERFORM_MATRIX_ABSORPTION:
198
if is_fp8(self.W_UV_O):
@@ -487,7 +497,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
497
v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
498
value=0)
499
490
- attn_output = flash_attn_varlen_func(
500
+ attn_output = self.flash_attn_varlen_func(
501
q=q,
502
k=k,
503
v=v_padded,
@@ -497,7 +507,6 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
507
max_seqlen_k=max_prefill_seq_len,
508
softmax_scale=self.scale,
509
causal=True,
500
- fa_version=self.vllm_flash_attn_version,
510
)
511
attn_output = attn_output\
512
.view(-1, self.num_heads, q.shape[-1])[..., :v.shape[-1]]\