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, 13:30 UTC
[FEAT][ROCm]: Support AITER MLA (#15893)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: qli88 <qiang.li2@amd.com>
vllmellm committed
Apr 23, 2025 at 00:31 UTC
30bc3e0f665eabae365355641aa86a29ae445967
9 files changed
+667
-29
tests/kernels/test_attention_selector.py
+128
-21
@@ -19,45 +19,152 @@ def clear_cache():
19
_cached_get_attn_backend.cache_clear()
20
21
22
-@pytest.mark.parametrize(
23
- "name", ["TORCH_SDPA", "ROCM_FLASH", "XFORMERS", "FLASHINFER"])
22
+# Define MLA and non-MLA backends separately
23
+DEVICE_MLA_BACKENDS = {
24
+ "cuda": ["TRITON_MLA", "FLASHMLA"],
25
+ "hip": ["TRITON_MLA", "ROCM_AITER_MLA"],
26
+ "cpu": [],
27
+}
28
+
29
+DEVICE_REGULAR_ATTN_BACKENDS = {
30
+ "cuda": ["XFORMERS", "FLASHINFER"],
31
+ "hip": ["ROCM_FLASH"],
32
+ "cpu": ["TORCH_SDPA"],
33
+}
34
+
35
+DEVICE_MLA_BLOCK_SIZES = {
36
+ "cuda": [16, 64], # CUDA supports both standard and extended block sizes
37
+ "hip": [16, 1], # HIP requires special handling for block_size=1
38
+ "cpu": [16] # CPU uses fixed block size from test cases
39
+}
40
+
41
+
42
+def generate_params():
43
+ params = []
44
+ for use_mla in [True, False]:
45
+ for device in ["cuda", "hip", "cpu"]:
46
+ backends = DEVICE_MLA_BACKENDS[
47
+ device] if use_mla else DEVICE_REGULAR_ATTN_BACKENDS[device]
48
+ for name in backends:
49
+ block_sizes = DEVICE_MLA_BLOCK_SIZES[device] if use_mla else [
50
+ 16
51
+ ]
52
+ for block_size in block_sizes:
53
+ params.append(
54
+ pytest.param(
55
+ device,
56
+ name,
57
+ use_mla,
58
+ block_size,
59
+ id=
60
+ f"{device}_{name}_mla_{str(use_mla)[0]}_blks{block_size}"
61
+ ))
62
+ return params
63
+
64
+
65
+@pytest.mark.parametrize("device, name, use_mla, block_size",
66
+ generate_params())
67
@pytest.mark.parametrize("use_v1", [True, False])
25
-@pytest.mark.parametrize("device", ["cpu", "hip", "cuda"])
68
def test_env(
69
+ device: str,
70
name: str,
71
+ use_mla: bool,
72
+ block_size: int,
73
use_v1: bool,
29
- device: str,
74
monkeypatch: pytest.MonkeyPatch,
75
):
32
- """Test that the attention selector can be set via environment variable.
33
- Note that we do not test FlashAttn because it is the default backend.
34
- """
35
-
76
+ """Test attention backend selection with valid device-backend pairs."""
77
with monkeypatch.context() as m:
78
m.setenv("VLLM_USE_V1", "1" if use_v1 else "0")
79
m.setenv(STR_BACKEND_ENV_VAR, name)
80
+ m.setenv("VLLM_MLA_DISABLE", "1" if use_mla else "0")
81
82
if device == "cpu":
83
with patch("vllm.attention.selector.current_platform",
84
CpuPlatform()):
85
backend = get_attn_backend(16, torch.float16, torch.float16,
44
- 16, False)
86
+ block_size, False)
87
assert backend.get_name() == "TORCH_SDPA"
88
+
89
elif device == "hip":
90
with patch("vllm.attention.selector.current_platform",
91
RocmPlatform()):
49
- backend = get_attn_backend(16, torch.float16, torch.float16,
50
- 16, False)
51
- EXPECTED = "TRITON_ATTN_VLLM_V1" if use_v1 else "ROCM_FLASH"
52
- assert backend.get_name() == EXPECTED
53
- else:
54
- if name in ["XFORMERS", "FLASHINFER"]:
55
- with patch("vllm.attention.selector.current_platform",
56
- CudaPlatform()):
57
- backend = get_attn_backend(16, torch.float16,
58
- torch.float16, 16, False)
59
- EXPECTED = "FLASH_ATTN_VLLM_V1" if use_v1 else name
60
- assert backend.get_name() == EXPECTED
92
+ if use_mla:
93
+ # Validate HIP MLA backend-block_size combinations
94
+ valid_combination = (
95
+ (name == "TRITON_MLA" and block_size != 1)
96
+ or (name == "ROCM_AITER_MLA" and block_size == 1))
97
+
98
+ if valid_combination:
99
+ backend = get_attn_backend(16,
100
+ torch.float16,
101
+ torch.float16,
102
+ block_size,
103
+ False,
104
+ use_mla=use_mla)
105
+ assert backend.get_name() == name
106
+ else:
107
+ with pytest.raises(ValueError) as exc_info:
108
+ get_attn_backend(16,
109
+ torch.float16,
110
+ torch.float16,
111
+ block_size,
112
+ False,
113
+ use_mla=use_mla)
114
+ assert f"The selected backend, {name}" in str(
115
+ exc_info.value)
116
+ else:
117
+ backend = get_attn_backend(16,
118
+ torch.float16,
119
+ torch.float16,
120
+ block_size,
121
+ False,
122
+ use_mla=use_mla)
123
+ expected = "TRITON_ATTN_VLLM_V1" if use_v1 else "ROCM_FLASH"
124
+ assert backend.get_name() == expected
125
+
126
+ elif device == "cuda":
127
+ with patch("vllm.attention.selector.current_platform",
128
+ CudaPlatform()):
129
+ if use_mla:
130
+ if name == "FLASHMLA" and block_size == 64:
131
+ from vllm.attention.backends.flashmla import (
132
+ is_flashmla_supported)
133
+
134
+ # only on cuda platforms with specific capability.
135
+ is_supported, _ = is_flashmla_supported()
136
+
137
+ if not is_supported:
138
+ # if platform is not supported then skip this case.
139
+ pytest.skip()
140
+ else:
141
+ backend = get_attn_backend(16,
142
+ torch.float16,
143
+ torch.float16,
144
+ block_size,
145
+ False,
146
+ use_mla=use_mla)
147
+ expected = f"{name}_VLLM_V1" if use_v1 else name
148
+ assert backend.get_name() == expected
149
+ else:
150
+ backend = get_attn_backend(16,
151
+ torch.float16,
152
+ torch.float16,
153
+ block_size,
154
+ False,
155
+ use_mla=use_mla)
156
+ expected = ("TRITON_MLA_VLLM_V1"
157
+ if use_v1 else "TRITON_MLA")
158
+ assert backend.get_name() == expected
159
+ else:
160
+ backend = get_attn_backend(16,
161
+ torch.float16,
162
+ torch.float16,
163
+ block_size,
164
+ False,
165
+ use_mla=use_mla)
166
+ expected = "FLASH_ATTN_VLLM_V1" if use_v1 else name
167
+ assert backend.get_name() == expected
168
169
170
def test_flash_attn(monkeypatch: pytest.MonkeyPatch):
tests/kernels/test_rocm_attention_selector.py
+28
-1
@@ -28,7 +28,34 @@ def test_selector(monkeypatch: pytest.MonkeyPatch):
28
assert (backend.get_name() == "ROCM_FLASH"
29
or backend.get_name() == "TRITON_ATTN_VLLM_V1")
30
31
- # mla test for deepseek related
31
+ # MLA test for deepseek related
32
+
33
+ # change the attention backend to triton MLA
34
+ m.setenv(STR_BACKEND_ENV_VAR, "TRITON_MLA")
35
+ backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
36
+ False, True)
37
+ assert backend.get_name() == "TRITON_MLA"
38
+
39
+ # If attention backend is None
40
+ # If use_mla is true
41
+ # The selected backend is triton MLA
42
+ m.setenv(STR_BACKEND_ENV_VAR, None)
43
backend = get_attn_backend(576, torch.bfloat16, "auto", 16, False,
44
False, True)
45
assert backend.get_name() == "TRITON_MLA"
46
+
47
+ # change the attention backend to AITER MLA
48
+ m.setenv(STR_BACKEND_ENV_VAR, "ROCM_AITER_MLA")
49
+ backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False,
50
+ False, True)
51
+ assert backend.get_name() == "ROCM_AITER_MLA"
52
+
53
+ # If attention backend is None
54
+ # If use_mla is true
55
+ # If VLLM_ROCM_USE_AITER is enabled
56
+ # The selected backend is ROCM_AITER_MLA
57
+ m.setenv(STR_BACKEND_ENV_VAR, None)
58
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
59
+ backend = get_attn_backend(576, torch.bfloat16, "auto", 1, False,
60
+ False, True)
61
+ assert backend.get_name() == "ROCM_AITER_MLA"
vllm/attention/backends/mla/common.py
+18
-3
@@ -711,12 +711,24 @@ class MLACommonMetadata(AttentionMetadata):
711
self.seq_lens[i] += 1
712
self.max_decode_seq_len = max(self.seq_lens)
713
714
+ self._ops_advance_step(num_seqs=num_seqs,
715
+ num_queries=num_queries,
716
+ block_size=block_size,
717
+ input_tokens=model_input.input_tokens,
718
+ sampled_token_ids=sampled_token_ids,
719
+ input_positions=model_input.input_positions)
720
+
721
+ def _ops_advance_step(self, num_seqs: int, num_queries: int,
722
+ block_size: int, input_tokens: torch.Tensor,
723
+ sampled_token_ids: torch.Tensor,
724
+ input_positions: torch.Tensor) -> None:
725
+ # here we use advance_step_flashinfo to update the paged_kv_* tensors
726
ops.advance_step_flashattn(num_seqs=num_seqs,
727
num_queries=num_queries,
728
block_size=block_size,
717
- input_tokens=model_input.input_tokens,
729
+ input_tokens=input_tokens,
730
sampled_token_ids=sampled_token_ids,
719
- input_positions=model_input.input_positions,
731
+ input_positions=input_positions,
732
seq_lens=self.seq_lens_tensor,
733
slot_mapping=self.slot_mapping,
734
block_tables=self.block_tables)
@@ -727,6 +739,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
739
NOTE: Please read the comment at the top of the file before trying to
740
understand this class
741
"""
742
+ BLOCK_TABLE_EXTENDER: list[list[int]] = []
743
744
def __init__(self, input_builder: "ModelInputForGPUBuilder"):
745
self.input_builder = input_builder
@@ -877,8 +890,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):
890
num_seqs = len(seq_lens)
891
if use_captured_graph:
892
self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
880
- self.block_tables.extend([] * cuda_graph_pad_size)
893
+ self.block_tables.extend(self.__class__.BLOCK_TABLE_EXTENDER *
894
+ cuda_graph_pad_size)
895
num_decode_tokens = batch_size - self.num_prefill_tokens
896
+
897
block_tables = self._get_graph_runner_block_tables(
898
num_seqs, self.block_tables)
899
else:
vllm/attention/backends/rocm_aiter_mla.py
new
+412
@@ -0,0 +1,412 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from contextlib import contextmanager
4
+from dataclasses import dataclass
5
+from typing import TYPE_CHECKING, Any, Optional, Type, Union
6
+
7
+import torch
8
+
9
+import vllm._custom_ops as ops
10
+import vllm.envs as envs
11
+from vllm.attention.backends.mla.common import (MLACommonBackend,
12
+ MLACommonImpl,
13
+ MLACommonMetadata,
14
+ MLACommonMetadataBuilder,
15
+ MLACommonState)
16
+from vllm.attention.backends.utils import (compute_slot_mapping,
17
+ compute_slot_mapping_start_idx,
18
+ is_block_tables_empty)
19
+from vllm.attention.ops.rocm_aiter_mla import (aiter_mla_decode_fwd,
20
+ get_aiter_mla_metadata)
21
+
22
+if TYPE_CHECKING:
23
+ from vllm.worker.model_runner import ModelInputForGPUBuilder
24
+
25
+
26
+def is_aiter_mla_enabled() -> bool:
27
+ return envs.VLLM_ROCM_USE_AITER \
28
+ and envs.VLLM_ROCM_USE_AITER_MLA
29
+
30
+
31
+class AiterMLABackend(MLACommonBackend):
32
+
33
+ @staticmethod
34
+ def get_name() -> str:
35
+ return "ROCM_AITER_MLA"
36
+
37
+ @staticmethod
38
+ def get_impl_cls() -> Type["AiterMLAImpl"]:
39
+ return AiterMLAImpl
40
+
41
+ @staticmethod
42
+ def get_metadata_cls() -> Type["AiterMLAMetadata"]:
43
+ return AiterMLAMetadata
44
+
45
+ @staticmethod
46
+ def get_builder_cls() -> Type["AiterMLAMetadataBuilder"]:
47
+ return AiterMLAMetadataBuilder
48
+
49
+ @staticmethod
50
+ def get_state_cls() -> Type["AiterMLAState"]:
51
+ return AiterMLAState
52
+
53
+
54
+@dataclass
55
+class AiterMLAMetadata(MLACommonMetadata):
56
+ # The following 4 tensors are for current version of AITER MLA
57
+ block_table_bound: Optional[torch.Tensor] = None
58
+ # The indptr of the paged kv cache, shape: [batch_size + 1]
59
+ paged_kv_indptr: Optional[torch.Tensor] = None
60
+ # The page indices of the paged kv cache
61
+ paged_kv_indices: Optional[torch.Tensor] = None
62
+ # The number of entries in the last page of each request in
63
+ # the paged kv cache, shape: [batch_size]
64
+ paged_kv_last_page_lens: Optional[torch.Tensor] = None
65
+
66
+ @property
67
+ def prefill_metadata(self):
68
+ prefill_metadata = super().prefill_metadata
69
+ self._cached_prefill_metadata = prefill_metadata
70
+
71
+ if prefill_metadata is not None:
72
+ prefill_metadata.paged_kv_indptr = self.paged_kv_indptr
73
+ prefill_metadata.paged_kv_indices = self.paged_kv_indices
74
+ prefill_metadata\
75
+ .paged_kv_last_page_lens = self.paged_kv_last_page_lens
76
+ prefill_metadata.block_table_bound = self.block_table_bound
77
+
78
+ # update the cache
79
+ self._cached_prefill_metadata = self.__class__(
80
+ **prefill_metadata.__dict__)
81
+
82
+ return self._cached_prefill_metadata
83
+
84
+ @property
85
+ def decode_metadata(self):
86
+ decode_metadata = super().decode_metadata
87
+
88
+ self._cached_decode_metadata = decode_metadata
89
+
90
+ if decode_metadata is not None:
91
+ decode_metadata.paged_kv_indptr = self.paged_kv_indptr
92
+ decode_metadata.paged_kv_indices = self.paged_kv_indices
93
+ decode_metadata\
94
+ .paged_kv_last_page_lens = self.paged_kv_last_page_lens
95
+ decode_metadata.block_table_bound = self.block_table_bound
96
+
97
+ # update the cache
98
+ self._cached_decode_metadata = self.__class__(
99
+ **decode_metadata.__dict__)
100
+
101
+ return self._cached_decode_metadata
102
+
103
+ def _ops_advance_step(self, num_seqs: int, num_queries: int,
104
+ block_size: int, input_tokens: torch.Tensor,
105
+ sampled_token_ids: torch.Tensor,
106
+ input_positions: torch.Tensor) -> None:
107
+
108
+ ops.advance_step_flashinfer(
109
+ num_seqs=num_seqs,
110
+ num_queries=num_queries,
111
+ block_size=block_size,
112
+ input_tokens=input_tokens,
113
+ sampled_token_ids=sampled_token_ids,
114
+ input_positions=input_positions,
115
+ seq_lens=self.seq_lens_tensor,
116
+ slot_mapping=self.slot_mapping,
117
+ block_tables=self.block_tables,
118
+ paged_kv_indices=self.paged_kv_indices,
119
+ paged_kv_indptr=self.paged_kv_indptr,
120
+ paged_kv_last_page_lens=self.paged_kv_last_page_lens,
121
+ block_table_bound=self.block_table_bound)
122
+
123
+
124
+class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
125
+ BLOCK_TABLE_EXTENDER: list[list[int]] = [[]]
126
+
127
+ def __init__(self, input_builder: "ModelInputForGPUBuilder"):
128
+ super().__init__(input_builder)
129
+ assert self.runner.model_config.max_model_len == 32768,\
130
+ "AITER MLA requires max model len to be set to 32768"
131
+ assert self.block_size == 1, "AITER MLA requires only block size 1."
132
+
133
+ def prepare(self):
134
+ super().prepare()
135
+ self.paged_kv_indices: list[int] = []
136
+ self.paged_kv_indptr: list[int] = [0]
137
+ self.paged_kv_last_page_lens: list[int] = []
138
+ self.total_blocks = 0
139
+
140
+ def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,
141
+ prefix_cache_hit: bool):
142
+ """Add a sequence group to the metadata. Specifically update/append
143
+ 1. context length.
144
+ 2. block table.
145
+ 3. slot mapping.
146
+ """
147
+ is_prompt = inter_data.is_prompt
148
+ block_tables = inter_data.block_tables
149
+
150
+ for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
151
+ curr_sliding_window_block, input_positions) in zip(
152
+ inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
153
+ inter_data.orig_seq_lens, inter_data.seq_lens,
154
+ inter_data.query_lens, inter_data.context_lens,
155
+ inter_data.curr_sliding_window_blocks,
156
+ inter_data.input_positions):
157
+ self.input_positions.extend(input_positions)
158
+ self.context_lens.append(context_len)
159
+ if is_prompt:
160
+ self.num_prefills += 1
161
+ self.num_prefill_tokens += token_len
162
+ self.prefill_seq_lens.append(seq_len)
163
+ else:
164
+ self.num_decode_tokens += query_len
165
+ self.curr_seq_lens.append(curr_seq_len)
166
+
167
+ # Compute block table.
168
+ # TODO(sang): Combine chunked prefill and prefix caching by
169
+ # only allowing multiple of block_size chunk size.
170
+ # NOTE: This only works for oooooooxxx style attention.
171
+ block_table = []
172
+ if prefix_cache_hit:
173
+ # NOTE(woosuk): For flash-attn, the block table should
174
+ # include the entries for the incoming prefill tokens.
175
+ block_table = block_tables[seq_id]
176
+ elif ((chunked_prefill_enabled or not is_prompt)
177
+ and block_tables is not None):
178
+ if curr_sliding_window_block == 0:
179
+ block_table = block_tables[seq_id]
180
+ else:
181
+ block_table = block_tables[seq_id][
182
+ -curr_sliding_window_block:]
183
+ self.block_tables.append(block_table)
184
+
185
+ # Compute slot mapping.
186
+ is_profile_run = is_block_tables_empty(block_tables)
187
+ start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
188
+ context_len,
189
+ self.sliding_window)
190
+ compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
191
+ seq_len, context_len, start_idx,
192
+ self.block_size, inter_data.block_tables)
193
+ if is_profile_run:
194
+ return
195
+
196
+ # Update paged_kv_* tensors only for non-profile run
197
+ block_table = block_tables[seq_id]
198
+ self._update_paged_kv_tensors(block_table, seq_len)
199
+
200
+ def _update_paged_kv_tensors(self, block_table: list[int], seq_len: int):
201
+ # Get the number of valid blocks based on sequence length.
202
+ # If seq_len = 16, block_size = 16,
203
+ # block_table_bound is 1 with 1 valid block.
204
+ # If seq_len = 15, block_size = 16,
205
+ # block_table_bound is 0 + 1 with 1 valid block.
206
+ self.total_blocks += len(block_table)
207
+ block_table_bound = seq_len // self.block_size + 1 \
208
+ if seq_len % self.block_size != 0 \
209
+ else seq_len // self.block_size
210
+ self.paged_kv_indices.extend(block_table[:block_table_bound])
211
+ self.paged_kv_indptr.append(self.paged_kv_indptr[-1] +
212
+ block_table_bound)
213
+
214
+ last_page_len = seq_len % self.block_size
215
+ if last_page_len == 0:
216
+ last_page_len = self.block_size
217
+ self.paged_kv_last_page_lens.append(last_page_len)
218
+
219
+ def build(self, seq_lens: list[int], query_lens: list[int],
220
+ cuda_graph_pad_size: int, batch_size: int) -> AiterMLAMetadata:
221
+ metadata = super().build(seq_lens, query_lens, cuda_graph_pad_size,
222
+ batch_size)
223
+ device = self.runner.device
224
+ use_captured_graph = cuda_graph_pad_size != -1
225
+
226
+ if use_captured_graph:
227
+ last_paged_kv_indptr = self.paged_kv_indptr[-1]
228
+ self.paged_kv_indptr.extend([last_paged_kv_indptr] *
229
+ cuda_graph_pad_size)
230
+ self.paged_kv_last_page_lens.extend([0] * cuda_graph_pad_size)
231
+
232
+ # For current version of AITER MLA
233
+ if len(self.paged_kv_indptr) > 0:
234
+ # extend to the maximum number of blocks as returned by the
235
+ # scheduler
236
+ self.paged_kv_indices.extend(
237
+ [0] * (self.total_blocks - len(self.paged_kv_indices)))
238
+ paged_kv_indices_tensor = torch.tensor(self.paged_kv_indices,
239
+ device=device,
240
+ dtype=torch.int)
241
+ paged_kv_indptr_tensor = torch.tensor(self.paged_kv_indptr,
242
+ device=device,
243
+ dtype=torch.int)
244
+ paged_kv_last_page_lens_tensor = torch.tensor(
245
+ self.paged_kv_last_page_lens, device=device, dtype=torch.int)
246
+ block_table_bound_tensor = torch.zeros(len(self.paged_kv_indptr) -
247
+ 1,
248
+ device=device,
249
+ dtype=torch.int)
250
+ else:
251
+ paged_kv_indices_tensor = None
252
+ paged_kv_indptr_tensor = None
253
+ paged_kv_last_page_lens_tensor = None
254
+ block_table_bound_tensor = None
255
+
256
+ metadata.paged_kv_indptr = paged_kv_indptr_tensor
257
+ metadata.paged_kv_indices = paged_kv_indices_tensor
258
+ metadata.paged_kv_last_page_lens = paged_kv_last_page_lens_tensor
259
+ metadata.block_table_bound = block_table_bound_tensor
260
+
261
+ return metadata
262
+
263
+
264
+class AiterMLAState(MLACommonState[AiterMLAMetadata]):
265
+
266
+ @contextmanager
267
+ def graph_capture(self, max_batch_size: int):
268
+ kv_indices, kv_indptr, last_page_lens = get_aiter_mla_metadata(
269
+ max_batch_size=max_batch_size,
270
+ block_size=self.runner.block_size,
271
+ max_block_per_batch=self.runner.get_max_block_per_batch(),
272
+ device=self.runner.device)
273
+ self._paged_kv_indices_tensor = kv_indices
274
+ self._paged_kv_indptr_tensor = kv_indptr
275
+ self._paged_kv_last_page_lens_tensor = last_page_lens
276
+
277
+ with super().graph_capture(max_batch_size):
278
+ yield
279
+
280
+ del self._paged_kv_indices_tensor
281
+ del self._paged_kv_indptr_tensor
282
+ del self._paged_kv_last_page_lens_tensor
283
+
284
+ def graph_capture_get_metadata_for_batch(
285
+ self,
286
+ batch_size: int,
287
+ is_encoder_decoder_model: bool = False) -> AiterMLAMetadata:
288
+
289
+ metadata = super().graph_capture_get_metadata_for_batch(
290
+ batch_size, is_encoder_decoder_model)
291
+
292
+ paged_kv_indptr = self._paged_kv_indptr_tensor[:batch_size + 1]
293
+ paged_kv_indices = self._paged_kv_indices_tensor
294
+ paged_kv_last_page_lens = self._paged_kv_last_page_lens_tensor[:
295
+ batch_size]
296
+
297
+ metadata.paged_kv_indptr = paged_kv_indptr
298
+ metadata.paged_kv_indices = paged_kv_indices
299
+ metadata.paged_kv_last_page_lens = paged_kv_last_page_lens
300
+
301
+ return metadata
302
+
303
+ def get_graph_input_buffers(self,
304
+ attn_metadata: AiterMLAMetadata,
305
+ is_encoder_decoder_model: bool = False):
306
+ input_buffers = super().get_graph_input_buffers(
307
+ attn_metadata, is_encoder_decoder_model)
308
+ input_buffers[
309
+ 'paged_kv_indptr'] = attn_metadata.decode_metadata.paged_kv_indptr
310
+ input_buffers[
311
+ "paged_kv_indices"] = attn_metadata.\
312
+ decode_metadata.paged_kv_indices
313
+ input_buffers[
314
+ "paged_kv_last_page_lens"] = attn_metadata.\
315
+ decode_metadata.paged_kv_last_page_lens
316
+
317
+ return input_buffers
318
+
319
+ def prepare_graph_input_buffers(self,
320
+ input_buffers,
321
+ attn_metadata: AiterMLAMetadata,
322
+ is_encoder_decoder_model: bool = False):
323
+ super().prepare_graph_input_buffers(input_buffers, attn_metadata,
324
+ is_encoder_decoder_model)
325
+
326
+ num_total_blocks = attn_metadata.decode_metadata.paged_kv_indices.shape[
327
+ 0]
328
+ input_buffers["paged_kv_indptr"].copy_(
329
+ attn_metadata.decode_metadata.paged_kv_indptr, non_blocking=True)
330
+ input_buffers["paged_kv_indices"][:num_total_blocks].copy_(
331
+ attn_metadata.decode_metadata.paged_kv_indices, non_blocking=True)
332
+ input_buffers["paged_kv_last_page_lens"].copy_(
333
+ attn_metadata.decode_metadata.paged_kv_last_page_lens,
334
+ non_blocking=True)
335
+
336
+
337
+class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
338
+
339
+ def __init__(
340
+ self,
341
+ num_heads: int,
342
+ head_size: int,
343
+ scale: float,
344
+ num_kv_heads: int,
345
+ alibi_slopes: Optional[list[float]],
346
+ sliding_window: Optional[int],
347
+ kv_cache_dtype: str,
348
+ blocksparse_params: Optional[dict[str, Any]],
349
+ logits_soft_cap: Optional[float],
350
+ attn_type: str,
351
+ # MLA Specific Arguments
352
+ **mla_args) -> None:
353
+ super().__init__(num_heads, head_size, scale, num_kv_heads,
354
+ alibi_slopes, sliding_window, kv_cache_dtype,
355
+ blocksparse_params, logits_soft_cap, attn_type,
356
+ **mla_args)
357
+
358
+ unsupported_features = [
359
+ alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
360
+ ]
361
+ if any(unsupported_features):
362
+ raise NotImplementedError(
363
+ "Aiter MLA does not support one of the following: "
364
+ "alibi_slopes, sliding_window, blocksparse_params, "
365
+ "logits_soft_cap")
366
+
367
+ from aiter import flash_attn_varlen_func
368
+ self.flash_attn_varlen_func = flash_attn_varlen_func
369
+
370
+ def _flash_attn_varlen_diff_headdims(
371
+ self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
372
+ softmax_scale: float, return_softmax_lse: bool,
373
+ **kwargs) -> Union[tuple[torch.Tensor, ...], torch.Tensor]:
374
+ output = self.flash_attn_varlen_func(
375
+ q=q,
376
+ k=k,
377
+ v=v,
378
+ softmax_scale=softmax_scale,
379
+ return_lse=return_softmax_lse,
380
+ **kwargs,
381
+ )
382
+
383
+ return output
384
+
385
+ def _forward_decode(
386
+ self,
387
+ q_nope: torch.Tensor,
388
+ q_pe: torch.Tensor,
389
+ kv_c_and_k_pe_cache: torch.Tensor,
390
+ attn_metadata: AiterMLAMetadata,
391
+ ) -> torch.Tensor:
392
+ assert kv_c_and_k_pe_cache.numel() > 0
393
+
394
+ decode_meta = attn_metadata.decode_metadata
395
+ assert decode_meta is not None
396
+ B = q_nope.shape[0]
397
+
398
+ q = torch.cat([q_nope, q_pe], dim=-1)
399
+ o = torch.zeros(B,
400
+ self.num_heads,
401
+ self.kv_lora_rank,
402
+ dtype=q.dtype,
403
+ device=q.device)
404
+
405
+ kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
406
+
407
+ aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,
408
+ attn_metadata.paged_kv_indptr,
409
+ attn_metadata.paged_kv_indices,
410
+ attn_metadata.paged_kv_last_page_lens)
411
+
412
+ return self._v_up_proj_and_o_proj(o)
vllm/attention/ops/rocm_aiter_mla.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from typing import Optional
4
+
5
+import torch
6
+
7
+
8
+def get_aiter_mla_metadata(max_batch_size: int, block_size: int,
9
+ max_block_per_batch: int,
10
+ device: torch.device) -> tuple[torch.Tensor, ...]:
11
+ paged_kv_indices = torch.zeros(max_batch_size * max_block_per_batch,
12
+ dtype=torch.int32,
13
+ device=device)
14
+ paged_kv_indptr = torch.zeros(max_batch_size + 1,
15
+ dtype=torch.int32,
16
+ device=device)
17
+ paged_kv_last_page_lens = torch.full((max_batch_size, ),
18
+ block_size,
19
+ dtype=torch.int32)
20
+ return paged_kv_indices, paged_kv_indptr, paged_kv_last_page_lens
21
+
22
+
23
+def aiter_mla_decode_fwd(
24
+ q: torch.Tensor,
25
+ kv_buffer: torch.Tensor,
26
+ o: torch.Tensor,
27
+ sm_scale: float,
28
+ kv_indptr: Optional[torch.Tensor] = None,
29
+ kv_indices: Optional[torch.Tensor] = None,
30
+ kv_last_page_lens: Optional[torch.Tensor] = None,
31
+ logit_cap: float = 0.0,
32
+):
33
+ from aiter.mla import mla_decode_fwd
34
+
35
+ mla_decode_fwd(q,
36
+ kv_buffer.view(-1, 1, 1, q.shape[-1]),
37
+ o,
38
+ kv_indptr,
39
+ kv_indices,
40
+ kv_last_page_lens,
41
+ sm_scale=sm_scale,
42
+ logit_cap=logit_cap)
vllm/config.py
+1
-1
@@ -1248,7 +1248,7 @@ class ModelConfig:
1248
or getattr(self.hf_config, "is_matryoshka", False))
1249
1250
1251
-BlockSize = Literal[8, 16, 32, 64, 128]
1251
+BlockSize = Literal[1, 8, 16, 32, 64, 128]
1252
CacheDType = Literal["auto", "fp8", "fp8_e4m3", "fp8_e5m2"]
1253
PrefixCachingHashAlgo = Literal["builtin", "sha256"]
1254
vllm/envs.py
+6
@@ -79,6 +79,7 @@ if TYPE_CHECKING:
79
VLLM_ROCM_USE_AITER_LINEAR: bool = True
80
VLLM_ROCM_USE_AITER_MOE: bool = True
81
VLLM_ROCM_USE_AITER_RMSNORM: bool = True
82
+ VLLM_ROCM_USE_AITER_MLA: bool = True
83
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
84
VLLM_ROCM_FP8_PADDING: bool = True
85
VLLM_ROCM_MOE_PADDING: bool = True
@@ -558,6 +559,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
559
lambda: (os.getenv("VLLM_ROCM_USE_AITER_RMSNORM", "True").lower() in
560
("true", "1")),
561
562
+ # Whether to use aiter mla ops.
563
+ # By default is enabled.
564
+ "VLLM_ROCM_USE_AITER_MLA":
565
+ lambda: (os.getenv("VLLM_ROCM_USE_AITER_MLA", "True").lower() in
566
+ ("true", "1")),
567
# use rocm skinny gemms
568
"VLLM_ROCM_USE_SKINNY_GEMM":
569
lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in
vllm/platforms/interface.py
+1
@@ -39,6 +39,7 @@ class _Backend(enum.Enum):
39
TRITON_ATTN_VLLM_V1 = enum.auto()
40
XFORMERS = enum.auto()
41
ROCM_FLASH = enum.auto()
42
+ ROCM_AITER_MLA = enum.auto()
43
TORCH_SDPA = enum.auto()
44
FLASHINFER = enum.auto()
45
TRITON_MLA = enum.auto() # Supported by V1
vllm/platforms/rocm.py
+31
-3
@@ -141,8 +141,36 @@ class RocmPlatform(Platform):
141
kv_cache_dtype, block_size, use_v1,
142
use_mla) -> str:
143
if use_mla:
144
- logger.info("Using Triton MLA backend.")
145
- return "vllm.attention.backends.triton_mla.TritonMLABackend"
144
+ from vllm.attention.backends.rocm_aiter_mla import (
145
+ is_aiter_mla_enabled)
146
+
147
+ if selected_backend is None:
148
+ selected_backend = (_Backend.ROCM_AITER_MLA if
149
+ is_aiter_mla_enabled() or block_size == 1
150
+ else _Backend.TRITON_MLA)
151
+
152
+ if selected_backend == _Backend.TRITON_MLA:
153
+ if block_size != 1:
154
+ logger.info("Using Triton MLA backend.")
155
+ return "vllm.attention.backends.triton_mla.TritonMLABackend" # noqa: E501
156
+ else:
157
+ raise ValueError(
158
+ f" The selected backend, {selected_backend.name},"
159
+ f"does not support block size {block_size}.")
160
+ elif selected_backend == _Backend.ROCM_AITER_MLA:
161
+ if block_size == 1:
162
+ logger.info("Using AITER MLA backend.")
163
+ return "vllm.attention.backends.rocm_aiter_mla.AiterMLABackend" # noqa: E501
164
+ else:
165
+ raise ValueError(
166
+ f" The selected backend, {selected_backend.name},"
167
+ f"does not support block size {block_size}."
168
+ "(currently only supports block size 1)")
169
+ else:
170
+ raise ValueError(
171
+ f" The selected backend, {selected_backend.name},"
172
+ f"is not MLA type while requested for MLA backend.")
173
+
174
selected_backend = (_Backend.ROCM_FLASH if selected_backend
175
== _Backend.FLASH_ATTN else selected_backend)
176
if envs.VLLM_USE_V1:
@@ -317,4 +345,4 @@ class RocmPlatform(Platform):
345
@classmethod
346
def get_cu_count(cls, device_id: int = 0) -> int:
347
return torch.cuda.get_device_properties(
320
- device_id).multi_processor_count
\ No newline at end of file
348
+ device_id).multi_processor_count