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 6, 00:45 UTC
[ROCm] Utilize persistent MLA kernel from AITER (#36574)
Signed-off-by: Sathish Sanjeevi <sathish.krishnan.p.s@gmail.com>
Sathish Sanjeevi committed
Mar 25, 2026 at 12:00 UTC
978fc18bf000dbe36d361e32995b7178dbdde20c
2 files changed
+154
-2
vllm/_aiter_ops.py
+47
@@ -386,6 +386,12 @@ def _rocm_aiter_mla_decode_fwd_impl(
386
logit_cap: float = 0.0,
387
q_scale: torch.Tensor | None = None,
388
kv_scale: torch.Tensor | None = None,
389
+ work_meta_data: torch.Tensor | None = None,
390
+ work_indptr: torch.Tensor | None = None,
391
+ work_info_set: torch.Tensor | None = None,
392
+ reduce_indptr: torch.Tensor | None = None,
393
+ reduce_final_map: torch.Tensor | None = None,
394
+ reduce_partial_map: torch.Tensor | None = None,
395
) -> None:
396
from aiter.mla import mla_decode_fwd
397
@@ -399,6 +405,29 @@ def _rocm_aiter_mla_decode_fwd_impl(
405
kwargs["q_scale"] = q_scale
406
kwargs["kv_scale"] = kv_scale
407
408
+ if work_meta_data is not None:
409
+ assert work_indptr is not None, (
410
+ "work_indptr must be provided with work_meta_data"
411
+ )
412
+ assert work_info_set is not None, (
413
+ "work_info_set must be provided with work_meta_data"
414
+ )
415
+ assert reduce_indptr is not None, (
416
+ "reduce_indptr must be provided with work_meta_data"
417
+ )
418
+ assert reduce_final_map is not None, (
419
+ "reduce_final_map must be provided with work_meta_data"
420
+ )
421
+ assert reduce_partial_map is not None, (
422
+ "reduce_partial_map must be provided with work_meta_data"
423
+ )
424
+ kwargs["work_meta_data"] = work_meta_data
425
+ kwargs["work_indptr"] = work_indptr
426
+ kwargs["work_info_set"] = work_info_set
427
+ kwargs["reduce_indptr"] = reduce_indptr
428
+ kwargs["reduce_final_map"] = reduce_final_map
429
+ kwargs["reduce_partial_map"] = reduce_partial_map
430
+
431
mla_decode_fwd(
432
q,
433
kv_buffer.view(-1, 1, 1, q.shape[-1]),
@@ -425,6 +454,12 @@ def _rocm_aiter_mla_decode_fwd_fake(
454
logit_cap: float = 0.0,
455
q_scale: torch.Tensor | None = None,
456
kv_scale: torch.Tensor | None = None,
457
+ work_meta_data: torch.Tensor | None = None,
458
+ work_indptr: torch.Tensor | None = None,
459
+ work_info_set: torch.Tensor | None = None,
460
+ reduce_indptr: torch.Tensor | None = None,
461
+ reduce_final_map: torch.Tensor | None = None,
462
+ reduce_partial_map: torch.Tensor | None = None,
463
) -> None:
464
pass
465
@@ -1659,6 +1694,12 @@ class rocm_aiter_ops:
1694
logit_cap: float = 0.0,
1695
q_scale: torch.Tensor | None = None,
1696
kv_scale: torch.Tensor | None = None,
1697
+ work_meta_data: torch.Tensor | None = None,
1698
+ work_indptr: torch.Tensor | None = None,
1699
+ work_info_set: torch.Tensor | None = None,
1700
+ reduce_indptr: torch.Tensor | None = None,
1701
+ reduce_final_map: torch.Tensor | None = None,
1702
+ reduce_partial_map: torch.Tensor | None = None,
1703
):
1704
torch.ops.vllm.rocm_aiter_mla_decode_fwd(
1705
q,
@@ -1673,6 +1714,12 @@ class rocm_aiter_ops:
1714
logit_cap=logit_cap,
1715
q_scale=q_scale,
1716
kv_scale=kv_scale,
1717
+ work_meta_data=work_meta_data,
1718
+ work_indptr=work_indptr,
1719
+ work_info_set=work_info_set,
1720
+ reduce_indptr=reduce_indptr,
1721
+ reduce_final_map=reduce_final_map,
1722
+ reduce_partial_map=reduce_partial_map,
1723
)
1724
1725
@staticmethod
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+107
-2
@@ -18,7 +18,12 @@ from vllm.model_executor.layers.attention.mla_attention import (
18
QueryLenSupport,
19
)
20
from vllm.triton_utils import tl, triton
21
-from vllm.v1.attention.backend import AttentionCGSupport, AttentionLayer, MultipleOf
21
+from vllm.v1.attention.backend import (
22
+ AttentionCGSupport,
23
+ AttentionLayer,
24
+ CommonAttentionMetadata,
25
+ MultipleOf,
26
+)
27
from vllm.v1.kv_cache_interface import AttentionSpec
28
29
@@ -71,8 +76,14 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
76
max_qo_len: int | None = None
77
78
79
+@dataclass
80
class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
75
- pass
81
+ work_meta_data: torch.Tensor | None = None
82
+ work_indptr: torch.Tensor | None = None
83
+ work_info_set: torch.Tensor | None = None
84
+ reduce_indptr: torch.Tensor | None = None
85
+ reduce_final_map: torch.Tensor | None = None
86
+ reduce_partial_map: torch.Tensor | None = None
87
88
89
class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
@@ -116,6 +127,55 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
127
max_num_pages, dtype=torch.int32, device=device
128
)
129
130
+ from aiter import dtypes, get_mla_metadata_info_v1
131
+
132
+ self._num_attention_heads = vllm_config.model_config.get_num_attention_heads(
133
+ vllm_config.parallel_config
134
+ )
135
+ q_dtype = self.decode_attn_out_dtype
136
+ kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
137
+ if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
138
+ kv_cache_dtype_str = "fp8"
139
+ else:
140
+ kv_cache_dtype_str = "bf16"
141
+ kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16)
142
+ (
143
+ (work_meta_data_size, work_meta_data_type),
144
+ (work_indptr_size, work_indptr_type),
145
+ (work_info_set_size, work_info_set_type),
146
+ (reduce_indptr_size, reduce_indptr_type),
147
+ (reduce_final_map_size, reduce_final_map_type),
148
+ (reduce_partial_map_size, reduce_partial_map_type),
149
+ ) = get_mla_metadata_info_v1(
150
+ max_num_reqs,
151
+ 1,
152
+ self._num_attention_heads,
153
+ q_dtype,
154
+ kv_dtype,
155
+ is_sparse=False,
156
+ fast_mode=True,
157
+ )
158
+ self._mla_work_meta_data = torch.empty(
159
+ work_meta_data_size, dtype=work_meta_data_type, device=device
160
+ )
161
+ self._mla_work_indptr = torch.empty(
162
+ work_indptr_size, dtype=work_indptr_type, device=device
163
+ )
164
+ self._mla_work_info_set = torch.empty(
165
+ work_info_set_size, dtype=work_info_set_type, device=device
166
+ )
167
+ self._mla_reduce_indptr = torch.empty(
168
+ reduce_indptr_size, dtype=reduce_indptr_type, device=device
169
+ )
170
+ self._mla_reduce_final_map = torch.empty(
171
+ reduce_final_map_size, dtype=reduce_final_map_type, device=device
172
+ )
173
+ self._mla_reduce_partial_map = torch.empty(
174
+ reduce_partial_map_size,
175
+ dtype=reduce_partial_map_type,
176
+ device=device,
177
+ )
178
+
179
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
180
self.paged_kv_indptr = torch.zeros(
181
max_num_reqs + 1, dtype=torch.int32, device=device
@@ -184,6 +244,28 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
244
0, num_reqs + 1, step=1, dtype=torch.int32, device=device
245
)
246
247
+ from aiter import get_mla_metadata_v1
248
+
249
+ get_mla_metadata_v1(
250
+ qo_indptr,
251
+ paged_kv_indptr,
252
+ paged_kv_last_page_len,
253
+ self._num_attention_heads,
254
+ 1,
255
+ True,
256
+ self._mla_work_meta_data,
257
+ self._mla_work_info_set,
258
+ self._mla_work_indptr,
259
+ self._mla_reduce_indptr,
260
+ self._mla_reduce_final_map,
261
+ self._mla_reduce_partial_map,
262
+ page_size=1,
263
+ kv_granularity=16,
264
+ max_seqlen_qo=max_qo_len,
265
+ uni_seqlen_qo=max_qo_len,
266
+ fast_mode=True,
267
+ )
268
+
269
attn_metadata = AiterMLADecodeMetadata(
270
block_table=block_table_tensor,
271
seq_lens=seq_lens_device,
@@ -198,6 +280,23 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
280
281
return attn_metadata
282
283
+ def build(
284
+ self,
285
+ common_prefix_len: int,
286
+ common_attn_metadata: CommonAttentionMetadata,
287
+ fast_build: bool = False,
288
+ ) -> AiterMLAMetadata:
289
+ attn_metadata = super().build(
290
+ common_prefix_len, common_attn_metadata, fast_build
291
+ )
292
+ attn_metadata.work_meta_data = self._mla_work_meta_data
293
+ attn_metadata.work_indptr = self._mla_work_indptr
294
+ attn_metadata.work_info_set = self._mla_work_info_set
295
+ attn_metadata.reduce_indptr = self._mla_reduce_indptr
296
+ attn_metadata.reduce_final_map = self._mla_reduce_final_map
297
+ attn_metadata.reduce_partial_map = self._mla_reduce_partial_map
298
+ return attn_metadata
299
+
300
301
@triton.jit
302
def _copy_page_indices_kernel(
@@ -338,6 +437,12 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
437
attn_metadata.decode.paged_kv_last_page_len,
438
q_scale=layer._q_scale,
439
kv_scale=layer._k_scale,
440
+ work_meta_data=attn_metadata.work_meta_data,
441
+ work_indptr=attn_metadata.work_indptr,
442
+ work_info_set=attn_metadata.work_info_set,
443
+ reduce_indptr=attn_metadata.reduce_indptr,
444
+ reduce_final_map=attn_metadata.reduce_final_map,
445
+ reduce_partial_map=attn_metadata.reduce_partial_map,
446
)
447
448
if self._needs_head_repeat: