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, 04:45 UTC
[Bugfix] Enable PP with AITER+V1 (#19822)
Signed-off-by: Qiang Li <qiang.li2@amd.com>
qli88 committed
Jun 19, 2025 at 23:43 UTC
e3a3e4db463d5fc45def4d39d256ccf42fb70044
2 files changed
+3
-11
vllm/model_executor/layers/layernorm.py
-1
@@ -45,7 +45,6 @@ def fused_add_rms_norm(
45
46
def rocm_aiter_rms_norm(x: torch.Tensor, weight: torch.Tensor,
47
variance_epsilon: float) -> torch.Tensor:
48
-
48
import aiter as rocm_aiter
49
if x.dim() > 2:
50
x_original_shape = x.shape
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+3
-10
@@ -201,16 +201,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
201
202
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
203
204
- if self.num_heads == 16:
205
- # AITER MLA decode kernel only supports
206
- # max_seqlen_q=1 when using 16 heads.
207
- max_seqlen_qo = 1
208
- else:
209
- # AITER MLA decode Kernel handles arbitrary
210
- # max_seqlen_q values when using 128 heads.
211
- assert attn_metadata.prefill is not None
212
- max_seqlen_qo = attn_metadata.prefill.max_query_len
213
-
204
+ # max_seqlen_qo must be 1 except for MTP
205
+ # TODO: Find the best value for MTP
206
+ max_seqlen_qo = 1
207
aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,
208
attn_metadata.decode.qo_indptr, max_seqlen_qo,
209
attn_metadata.decode.paged_kv_indptr,