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, 11:00 UTC
[BugFix][AMD] Compatible patch for AITER lib after 04/20 (#17912)
Signed-off-by: Qiang Li <qiang.li2@amd.com>
qli88 committed
May 14, 2025 at 01:05 UTC
4f8b37322561ca07ba3e3e450ecff0a231d0d8f8
3 files changed
+54
-17
vllm/attention/backends/rocm_aiter_mla.py
+37
-12
@@ -53,7 +53,7 @@ class AiterMLABackend(MLACommonBackend):
53
54
@dataclass
55
class AiterMLAMetadata(MLACommonMetadata):
56
- # The following 4 tensors are for current version of AITER MLA
56
+ # The following 5 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
@@ -63,6 +63,10 @@ class AiterMLAMetadata(MLACommonMetadata):
63
# the paged kv cache, shape: [batch_size]
64
paged_kv_last_page_lens: Optional[torch.Tensor] = None
65
66
+ # This is just to make new AITER MLA API work
67
+ # -- MTP support is not added yet.
68
+ qo_indptr: Optional[torch.Tensor] = None
69
+
70
@property
71
def prefill_metadata(self):
72
prefill_metadata = super().prefill_metadata
@@ -74,6 +78,7 @@ class AiterMLAMetadata(MLACommonMetadata):
78
prefill_metadata\
79
.paged_kv_last_page_lens = self.paged_kv_last_page_lens
80
prefill_metadata.block_table_bound = self.block_table_bound
81
+ prefill_metadata.qo_indptr = self.qo_indptr
82
83
# update the cache
84
self._cached_prefill_metadata = self.__class__(
@@ -93,6 +98,7 @@ class AiterMLAMetadata(MLACommonMetadata):
98
decode_metadata\
99
.paged_kv_last_page_lens = self.paged_kv_last_page_lens
100
decode_metadata.block_table_bound = self.block_table_bound
101
+ decode_metadata.qo_indptr = self.qo_indptr
102
103
# update the cache
104
self._cached_decode_metadata = self.__class__(
@@ -136,6 +142,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
142
self.paged_kv_indptr: list[int] = [0]
143
self.paged_kv_last_page_lens: list[int] = []
144
self.total_blocks = 0
145
+ self.qo_indptr: list[int] = [0]
146
147
def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,
148
prefix_cache_hit: bool):
@@ -208,6 +215,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
215
self.paged_kv_indices.extend(block_table[:block_table_bound])
216
self.paged_kv_indptr.append(self.paged_kv_indptr[-1] +
217
block_table_bound)
218
+ self.qo_indptr.append(self.qo_indptr[-1] + 1)
219
220
last_page_len = seq_len % self.block_size
221
if last_page_len == 0:
@@ -226,6 +234,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
234
self.paged_kv_indptr.extend([last_paged_kv_indptr] *
235
cuda_graph_pad_size)
236
self.paged_kv_last_page_lens.extend([0] * cuda_graph_pad_size)
237
+ last_qo_indptr = self.qo_indptr[-1]
238
+ self.qo_indptr.extend([last_qo_indptr] * cuda_graph_pad_size)
239
240
# For current version of AITER MLA
241
if len(self.paged_kv_indptr) > 0:
@@ -245,16 +255,22 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
255
1,
256
device=device,
257
dtype=torch.int)
258
+
259
+ qo_indptr = torch.tensor(self.qo_indptr,
260
+ device=device,
261
+ dtype=torch.int)
262
else:
263
paged_kv_indices_tensor = None
264
paged_kv_indptr_tensor = None
265
paged_kv_last_page_lens_tensor = None
266
block_table_bound_tensor = None
267
+ qo_indptr = None
268
269
metadata.paged_kv_indptr = paged_kv_indptr_tensor
270
metadata.paged_kv_indices = paged_kv_indices_tensor
271
metadata.paged_kv_last_page_lens = paged_kv_last_page_lens_tensor
272
metadata.block_table_bound = block_table_bound_tensor
273
+ metadata.qo_indptr = qo_indptr
274
275
return metadata
276
@@ -263,14 +279,17 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
279
280
@contextmanager
281
def graph_capture(self, max_batch_size: int):
266
- kv_indices, kv_indptr, last_page_lens = get_aiter_mla_metadata(
267
- max_batch_size=max_batch_size,
268
- block_size=self.runner.block_size,
269
- max_block_per_batch=self.runner.get_max_block_per_batch(),
270
- device=self.runner.device)
282
+ kv_indices, kv_indptr, last_page_lens, qo_indptr = \
283
+ get_aiter_mla_metadata(
284
+ max_batch_size=max_batch_size,
285
+ block_size=self.runner.block_size,
286
+ max_block_per_batch=\
287
+ self.runner.get_max_block_per_batch(),
288
+ device=self.runner.device)
289
self._paged_kv_indices_tensor = kv_indices
290
self._paged_kv_indptr_tensor = kv_indptr
291
self._paged_kv_last_page_lens_tensor = last_page_lens
292
+ self._qo_indptr_tensor = qo_indptr
293
294
with super().graph_capture(max_batch_size):
295
yield
@@ -278,6 +297,7 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
297
del self._paged_kv_indices_tensor
298
del self._paged_kv_indptr_tensor
299
del self._paged_kv_last_page_lens_tensor
300
+ del self._qo_indptr_tensor
301
302
def graph_capture_get_metadata_for_batch(
303
self,
@@ -291,10 +311,12 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
311
paged_kv_indices = self._paged_kv_indices_tensor
312
paged_kv_last_page_lens = self._paged_kv_last_page_lens_tensor[:
313
batch_size]
314
+ qo_indptr = self._qo_indptr_tensor[:batch_size + 1]
315
316
metadata.paged_kv_indptr = paged_kv_indptr
317
metadata.paged_kv_indices = paged_kv_indices
318
metadata.paged_kv_last_page_lens = paged_kv_last_page_lens
319
+ metadata.qo_indptr = qo_indptr
320
321
return metadata
322
@@ -311,6 +333,7 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
333
input_buffers[
334
"paged_kv_last_page_lens"] = attn_metadata.\
335
decode_metadata.paged_kv_last_page_lens
336
+ input_buffers['qo_indptr'] = attn_metadata.qo_indptr
337
338
return input_buffers
339
@@ -330,6 +353,8 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
353
input_buffers["paged_kv_last_page_lens"].copy_(
354
attn_metadata.decode_metadata.paged_kv_last_page_lens,
355
non_blocking=True)
356
+ input_buffers["qo_indptr"].copy_(
357
+ attn_metadata.decode_metadata.qo_indptr, non_blocking=True)
358
359
360
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
@@ -370,11 +395,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
395
softmax_scale: float, return_softmax_lse: bool,
396
**kwargs) -> Union[tuple[torch.Tensor, ...], torch.Tensor]:
397
output = self.flash_attn_varlen_func(
373
- q=q,
374
- k=k,
375
- v=v,
376
- softmax_scale=softmax_scale,
377
- return_lse=return_softmax_lse,
398
+ q,
399
+ k,
400
+ v,
401
**kwargs,
402
)
403
@@ -394,7 +417,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
417
B = q_nope.shape[0]
418
419
q = torch.cat([q_nope, q_pe], dim=-1)
397
- o = torch.zeros(B,
420
+ o = torch.empty(B,
421
self.num_heads,
422
self.kv_lora_rank,
423
dtype=q.dtype,
@@ -403,6 +426,8 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
426
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
427
428
aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,
429
+ attn_metadata.qo_indptr,
430
+ attn_metadata.max_query_len,
431
attn_metadata.paged_kv_indptr,
432
attn_metadata.paged_kv_indices,
433
attn_metadata.paged_kv_last_page_lens)
vllm/attention/ops/rocm_aiter_mla.py
+12
-1
@@ -20,7 +20,8 @@ def get_aiter_mla_metadata(max_batch_size: int, block_size: int,
20
paged_kv_last_page_lens = torch.full((max_batch_size, ),
21
block_size,
22
dtype=torch.int32)
23
- return paged_kv_indices, paged_kv_indptr, paged_kv_last_page_lens
23
+ qo_indptr = torch.zeros(max_batch_size + 1, dtype=torch.int, device=device)
24
+ return paged_kv_indices, paged_kv_indptr, paged_kv_last_page_lens, qo_indptr
25
26
27
def aiter_mla_decode_fwd(
@@ -28,6 +29,8 @@ def aiter_mla_decode_fwd(
29
kv_buffer: torch.Tensor,
30
o: torch.Tensor,
31
sm_scale: float,
32
+ qo_indptr: torch.Tensor,
33
+ max_seqlen_qo: int,
34
kv_indptr: Optional[torch.Tensor] = None,
35
kv_indices: Optional[torch.Tensor] = None,
36
kv_last_page_lens: Optional[torch.Tensor] = None,
@@ -38,6 +41,8 @@ def aiter_mla_decode_fwd(
41
kv_buffer.view(
42
-1, 1, 1, q.shape[-1]),
43
o,
44
+ qo_indptr,
45
+ max_seqlen_qo,
46
kv_indptr,
47
kv_indices,
48
kv_last_page_lens,
@@ -49,6 +54,8 @@ def mla_decode_fwd_impl(
54
q: torch.Tensor,
55
kv_buffer: torch.Tensor,
56
o: torch.Tensor,
57
+ qo_indptr: torch.Tensor,
58
+ max_seqlen_qo: int,
59
kv_indptr: Optional[torch.Tensor] = None,
60
kv_indices: Optional[torch.Tensor] = None,
61
kv_last_page_lens: Optional[torch.Tensor] = None,
@@ -60,9 +67,11 @@ def mla_decode_fwd_impl(
67
mla_decode_fwd(q,
68
kv_buffer.view(-1, 1, 1, q.shape[-1]),
69
o,
70
+ qo_indptr,
71
kv_indptr,
72
kv_indices,
73
kv_last_page_lens,
74
+ max_seqlen_qo,
75
sm_scale=sm_scale,
76
logit_cap=logit_cap)
77
@@ -71,6 +80,8 @@ def mla_decode_fwd_fake(
80
q: torch.Tensor,
81
kv_buffer: torch.Tensor,
82
o: torch.Tensor,
83
+ qo_indptr: torch.Tensor,
84
+ max_seqlen_qo: int,
85
kv_indptr: Optional[torch.Tensor] = None,
86
kv_indices: Optional[torch.Tensor] = None,
87
kv_last_page_lens: Optional[torch.Tensor] = None,
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+5
-4
@@ -123,10 +123,11 @@ def rocm_aiter_fmoe_fp8_blockscale_g1u1_impl(
123
124
fmoe_fp8_blockscale_g1u1(out_asm, a1, w1, w2, sorted_token_ids,
125
sorted_weight_buf, sorted_expert_ids,
126
- num_valid_ids, topk, w1_scale.view(local_E, -1),
127
- w2_scale.view(local_E, -1),
128
- a1_scale.t().contiguous(), *block_shape,
129
- smooth_scale)
126
+ num_valid_ids, topk,
127
+ a1_scale.t().contiguous(),
128
+ w1_scale.view(local_E, -1),
129
+ w2_scale.view(local_E,
130
+ -1), *block_shape, smooth_scale)
131
132
return out_asm
133