karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Bug] Fix GLM-5.1 running error on ROCm platform (#40763)

Browse Source 
Signed-off-by: Qiang Li <qiang.li2@amd.com>
This commit is contained in:
qli88
2026-04-24 19:54:40 +00:00
committed by GitHub
parent 21792520e7
commit 095d2f87e8
3 changed files with 68 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+55 -24
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import ClassVar
from typing import ClassVar, Final
import torch
@@ -389,6 +389,53 @@ def _expand_page_indices_kernel(
)
class AiterMLAHelper:
"""
AITER MLA implementation requires num_heads >= 16. If num_heads < 16 and
16 % num_heads == 0, we can pad q to 16 heads; otherwise AITER has to fail.
"""
_AITER_MIN_MLA_HEADS: Final = 16
@staticmethod
def check_num_heads_validity(num_heads: int):
assert AiterMLAHelper.is_valid_num_heads(num_heads), (
f"Aiter MLA requires that num_heads be multiples or divisors of 16, "
f"but provided {num_heads} number of heads.\n"
f"Try adjusting tensor_parallel_size value."
)
@staticmethod
def is_valid_num_heads(num_heads: int) -> bool:
return (
num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
else AiterMLAHelper._AITER_MIN_MLA_HEADS % num_heads == 0
)
@staticmethod
def get_actual_mla_num_heads(num_heads: int) -> int:
return max(num_heads, AiterMLAHelper._AITER_MIN_MLA_HEADS)
@staticmethod
def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:
return (
q
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
else q.repeat_interleave(
AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, dim=1
)
)
@staticmethod
def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
return (
o
if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :]
)
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
def __init__(
self,
@@ -418,17 +465,8 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
kv_sharing_target_layer_name,
**mla_args,
)
_valid_heads = num_heads in (4, 8) or (
num_heads % 16 == 0 and 16 <= num_heads <= 128
)
assert _valid_heads, (
f"Aiter MLA supports num_heads of 4, 8, or multiples of 16 "
f"in [16, 128].\n"
f"Provided {num_heads} number of heads.\n"
"Try adjusting tensor_parallel_size value."
)
self._needs_head_repeat = num_heads < 16
self._head_repeat_factor = 16 // num_heads if num_heads < 16 else 1
AiterMLAHelper.check_num_heads_validity(num_heads)
unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
if any(unsupported_features):
raise NotImplementedError(
@@ -471,15 +509,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
assert isinstance(q, torch.Tensor)
B = q.shape[0]
if self._needs_head_repeat:
q = q.repeat_interleave(self._head_repeat_factor, dim=1)
kernel_num_heads = 16
else:
kernel_num_heads = self.num_heads
mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
o = torch.empty(
B,
kernel_num_heads,
mla_num_heads,
self.kv_lora_rank,
dtype=attn_metadata.decode.attn_out_dtype,
device=q.device,
@@ -506,7 +540,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
)
rocm_aiter_ops.mla_decode_fwd(
q,
mla_padded_q,
kv_buffer,
o,
self.scale,
@@ -518,7 +552,4 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
**mla_kwargs,
)
if self._needs_head_repeat:
o = o[:, :: self._head_repeat_factor, :]
return o, None
return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, o), None
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+10 -3
View File
@@ -28,6 +28,9 @@ from vllm.v1.attention.backend import (
from vllm.v1.attention.backends.mla.flashmla_sparse import (
triton_convert_req_index_to_global_index,
)
from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
AiterMLAHelper,
)
from vllm.v1.kv_cache_interface import AttentionSpec
if TYPE_CHECKING:
@@ -299,6 +302,8 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
indexer: "Indexer | None" = None,
**mla_args,
) -> None:
AiterMLAHelper.check_num_heads_validity(num_heads)
self.num_heads = num_heads
self.head_size = head_size
self.scale = float(scale)
@@ -317,8 +322,9 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
attn_metadata: ROCMAiterMLASparseMetadata,
) -> torch.Tensor:
num_tokens = q.shape[0]
mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
output = torch.empty(
[num_tokens, self.num_heads, self.kv_lora_rank],
[num_tokens, mla_num_heads, self.kv_lora_rank],
dtype=q.dtype,
device=q.device,
)
@@ -344,7 +350,7 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
attn_metadata.paged_kv_last_page_len,
)
return output[:, : self.num_heads, :]
return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)
def forward_mqa(
self,
@@ -374,8 +380,9 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
)
mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
attn_out = self._forward_bf16_kv(
q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata
mla_padded_q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata
)
return attn_out, None
vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
+3
View File
@@ -339,6 +339,8 @@ def rocm_fp8_paged_mqa_logits(
device="cuda",
dtype=torch.float32,
)
# TODO: 1. Replace _stage1 and out_qk.sum with another fused variant;
# 2. Remove ChunkQ when AITER PR #2891 merged
deepgemm_fp8_paged_mqa_logits_stage1(
q_fp8,
kv_cache_fp8,
@@ -347,6 +349,7 @@ def rocm_fp8_paged_mqa_logits(
context_lens,
block_tables,
max_model_len,
ChunkQ=heads,
)
return out_qk.sum(dim=0)
else: