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
Revert "[BugFix][AMD] Compatible patch for latest AITER(05/07/2025)" (#17910)
Michael Goin committed
May 9, 2025 at 09:58 UTC
85b72cb7b12cb7a76a83c3a9f825196d798d63e7
4 files changed
+23
-54
vllm/attention/backends/mla/common.py
+6
-6
@@ -1213,9 +1213,9 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1213
1214
attn_output, attn_softmax_lse = \
1215
self._flash_attn_varlen_diff_headdims(
1216
- q,
1217
- k,
1218
- v,
1216
+ q=q,
1217
+ k=k,
1218
+ v=v,
1219
cu_seqlens_q=prefill_metadata.query_start_loc,
1220
cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],
1221
max_seqlen_q=prefill_metadata.max_query_len,
@@ -1267,9 +1267,9 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1267
k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1268
1269
output = self._flash_attn_varlen_diff_headdims(
1270
- q,
1271
- k,
1272
- v,
1270
+ q=q,
1271
+ k=k,
1272
+ v=v,
1273
cu_seqlens_q=prefill_metadata.query_start_loc,
1274
cu_seqlens_k=prefill_metadata.query_start_loc,
1275
max_seqlen_q=prefill_metadata.max_prefill_seq_len,
vllm/attention/backends/rocm_aiter_mla.py
+12
-37
@@ -53,7 +53,7 @@ class AiterMLABackend(MLACommonBackend):
53
54
@dataclass
55
class AiterMLAMetadata(MLACommonMetadata):
56
- # The following 5 tensors are for current version of AITER MLA
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
@@ -63,10 +63,6 @@ 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
-
66
@property
67
def prefill_metadata(self):
68
prefill_metadata = super().prefill_metadata
@@ -78,7 +74,6 @@ class AiterMLAMetadata(MLACommonMetadata):
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
81
- prefill_metadata.qo_indptr = self.qo_indptr
77
78
# update the cache
79
self._cached_prefill_metadata = self.__class__(
@@ -98,7 +93,6 @@ class AiterMLAMetadata(MLACommonMetadata):
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
101
- decode_metadata.qo_indptr = self.qo_indptr
96
97
# update the cache
98
self._cached_decode_metadata = self.__class__(
@@ -142,7 +136,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
136
self.paged_kv_indptr: list[int] = [0]
137
self.paged_kv_last_page_lens: list[int] = []
138
self.total_blocks = 0
145
- self.qo_indptr: list[int] = [0]
139
140
def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,
141
prefix_cache_hit: bool):
@@ -215,7 +208,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
208
self.paged_kv_indices.extend(block_table[:block_table_bound])
209
self.paged_kv_indptr.append(self.paged_kv_indptr[-1] +
210
block_table_bound)
218
- self.qo_indptr.append(self.qo_indptr[-1] + 1)
211
212
last_page_len = seq_len % self.block_size
213
if last_page_len == 0:
@@ -234,8 +226,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
226
self.paged_kv_indptr.extend([last_paged_kv_indptr] *
227
cuda_graph_pad_size)
228
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)
229
230
# For current version of AITER MLA
231
if len(self.paged_kv_indptr) > 0:
@@ -255,22 +245,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
245
1,
246
device=device,
247
dtype=torch.int)
258
-
259
- qo_indptr = torch.tensor(self.qo_indptr,
260
- device=device,
261
- dtype=torch.int)
248
else:
249
paged_kv_indices_tensor = None
250
paged_kv_indptr_tensor = None
251
paged_kv_last_page_lens_tensor = None
252
block_table_bound_tensor = None
267
- qo_indptr = None
253
254
metadata.paged_kv_indptr = paged_kv_indptr_tensor
255
metadata.paged_kv_indices = paged_kv_indices_tensor
256
metadata.paged_kv_last_page_lens = paged_kv_last_page_lens_tensor
257
metadata.block_table_bound = block_table_bound_tensor
273
- metadata.qo_indptr = qo_indptr
258
259
return metadata
260
@@ -279,17 +263,14 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
263
264
@contextmanager
265
def graph_capture(self, max_batch_size: int):
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)
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)
271
self._paged_kv_indices_tensor = kv_indices
272
self._paged_kv_indptr_tensor = kv_indptr
273
self._paged_kv_last_page_lens_tensor = last_page_lens
292
- self._qo_indptr_tensor = qo_indptr
274
275
with super().graph_capture(max_batch_size):
276
yield
@@ -297,7 +278,6 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
278
del self._paged_kv_indices_tensor
279
del self._paged_kv_indptr_tensor
280
del self._paged_kv_last_page_lens_tensor
300
- del self._qo_indptr_tensor
281
282
def graph_capture_get_metadata_for_batch(
283
self,
@@ -311,12 +291,10 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
291
paged_kv_indices = self._paged_kv_indices_tensor
292
paged_kv_last_page_lens = self._paged_kv_last_page_lens_tensor[:
293
batch_size]
314
- qo_indptr = self._qo_indptr_tensor[:batch_size + 1]
294
295
metadata.paged_kv_indptr = paged_kv_indptr
296
metadata.paged_kv_indices = paged_kv_indices
297
metadata.paged_kv_last_page_lens = paged_kv_last_page_lens
319
- metadata.qo_indptr = qo_indptr
298
299
return metadata
300
@@ -333,7 +311,6 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
311
input_buffers[
312
"paged_kv_last_page_lens"] = attn_metadata.\
313
decode_metadata.paged_kv_last_page_lens
336
- input_buffers['qo_indptr'] = attn_metadata.qo_indptr
314
315
return input_buffers
316
@@ -353,8 +330,6 @@ class AiterMLAState(MLACommonState[AiterMLAMetadata]):
330
input_buffers["paged_kv_last_page_lens"].copy_(
331
attn_metadata.decode_metadata.paged_kv_last_page_lens,
332
non_blocking=True)
356
- input_buffers["qo_indptr"].copy_(
357
- attn_metadata.decode_metadata.qo_indptr, non_blocking=True)
333
334
335
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
@@ -395,9 +370,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
370
softmax_scale: float, return_softmax_lse: bool,
371
**kwargs) -> Union[tuple[torch.Tensor, ...], torch.Tensor]:
372
output = self.flash_attn_varlen_func(
398
- q,
399
- k,
400
- v,
373
+ q=q,
374
+ k=k,
375
+ v=v,
376
+ softmax_scale=softmax_scale,
377
+ return_lse=return_softmax_lse,
378
**kwargs,
379
)
380
@@ -417,7 +394,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
394
B = q_nope.shape[0]
395
396
q = torch.cat([q_nope, q_pe], dim=-1)
420
- o = torch.empty(B,
397
+ o = torch.zeros(B,
398
self.num_heads,
399
self.kv_lora_rank,
400
dtype=q.dtype,
@@ -426,8 +403,6 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
403
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
404
405
aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,
429
- attn_metadata.qo_indptr,
430
- attn_metadata.max_query_len,
406
attn_metadata.paged_kv_indptr,
407
attn_metadata.paged_kv_indices,
408
attn_metadata.paged_kv_last_page_lens)
vllm/attention/ops/rocm_aiter_mla.py
+1
-6
@@ -20,8 +20,7 @@ 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
- 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
23
+ return paged_kv_indices, paged_kv_indptr, paged_kv_last_page_lens
24
25
26
def aiter_mla_decode_fwd(
@@ -29,8 +28,6 @@ def aiter_mla_decode_fwd(
28
kv_buffer: torch.Tensor,
29
o: torch.Tensor,
30
sm_scale: float,
32
- qo_indptr: torch.Tensor,
33
- max_seqlen_qo: int,
31
kv_indptr: Optional[torch.Tensor] = None,
32
kv_indices: Optional[torch.Tensor] = None,
33
kv_last_page_lens: Optional[torch.Tensor] = None,
@@ -63,11 +60,9 @@ def mla_decode_fwd_impl(
60
mla_decode_fwd(q,
61
kv_buffer.view(-1, 1, 1, q.shape[-1]),
62
o,
66
- qo_indptr,
63
kv_indptr,
64
kv_indices,
65
kv_last_page_lens,
70
- max_seqlen_qo,
66
sm_scale=sm_scale,
67
logit_cap=logit_cap)
68
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+4
-5
@@ -123,11 +123,10 @@ 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,
127
- a1_scale.t().contiguous(),
128
- w1_scale.view(local_E, -1),
129
- w2_scale.view(local_E,
130
- -1), *block_shape, smooth_scale)
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)
130
131
return out_asm
132