@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 095d2f87e8

Fork 0

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 4, 16:15 UTC

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

Signed-off-by: Qiang Li <qiang.li2@amd.com>

qli88 committed Apr 24, 2026 at 14:54 UTC 095d2f87e8519de27f1fc39d9d22b299efdf0010

3 files changed +68 -27

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+55 -24

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from dataclasses import dataclass

5 -from typing import ClassVar

5 +from typing import ClassVar, Final

7 import torch

@@ -389,6 +389,53 @@ def _expand_page_indices_kernel(

389 )

390

391

392 +class AiterMLAHelper:

393 + """

394 + AITER MLA implementation requires num_heads >= 16. If num_heads < 16 and

395 + 16 % num_heads == 0, we can pad q to 16 heads; otherwise AITER has to fail.

396 + """

397 +

398 + _AITER_MIN_MLA_HEADS: Final = 16

399 +

400 + @staticmethod

401 + def check_num_heads_validity(num_heads: int):

402 + assert AiterMLAHelper.is_valid_num_heads(num_heads), (

403 + f"Aiter MLA requires that num_heads be multiples or divisors of 16, "

404 + f"but provided {num_heads} number of heads.\n"

405 + f"Try adjusting tensor_parallel_size value."

406 + )

407 +

408 + @staticmethod

409 + def is_valid_num_heads(num_heads: int) -> bool:

410 + return (

411 + num_heads % AiterMLAHelper._AITER_MIN_MLA_HEADS == 0

412 + if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS

413 + else AiterMLAHelper._AITER_MIN_MLA_HEADS % num_heads == 0

414 + )

415 +

416 + @staticmethod

417 + def get_actual_mla_num_heads(num_heads: int) -> int:

418 + return max(num_heads, AiterMLAHelper._AITER_MIN_MLA_HEADS)

419 +

420 + @staticmethod

421 + def get_mla_padded_q(num_heads: int, q: torch.Tensor) -> torch.Tensor:

422 + return (

423 + q

424 + if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS

425 + else q.repeat_interleave(

426 + AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, dim=1

427 + )

428 + )

429 +

430 + @staticmethod

431 + def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:

432 + return (

433 + o

434 + if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS

435 + else o[:, :: AiterMLAHelper._AITER_MIN_MLA_HEADS // num_heads, :]

436 + )

437 +

438 +

439 class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

440 def __init__(

441 self,

@@ -418,17 +465,8 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

465 kv_sharing_target_layer_name,

466 **mla_args,

467 )

421 - _valid_heads = num_heads in (4, 8) or (

422 - num_heads % 16 == 0 and 16 <= num_heads <= 128

423 - )

424 - assert _valid_heads, (

425 - f"Aiter MLA supports num_heads of 4, 8, or multiples of 16 "

426 - f"in [16, 128].\n"

427 - f"Provided {num_heads} number of heads.\n"

428 - "Try adjusting tensor_parallel_size value."

429 - )

430 - self._needs_head_repeat = num_heads < 16

431 - self._head_repeat_factor = 16 // num_heads if num_heads < 16 else 1

468 + AiterMLAHelper.check_num_heads_validity(num_heads)

469 +

470 unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]

471 if any(unsupported_features):

472 raise NotImplementedError(

@@ -471,15 +509,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

509 assert isinstance(q, torch.Tensor)

510 B = q.shape[0]

511

474 - if self._needs_head_repeat:

475 - q = q.repeat_interleave(self._head_repeat_factor, dim=1)

476 - kernel_num_heads = 16

477 - else:

478 - kernel_num_heads = self.num_heads

479 -

512 + mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)

513 + mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)

514 o = torch.empty(

515 B,

482 - kernel_num_heads,

516 + mla_num_heads,

517 self.kv_lora_rank,

518 dtype=attn_metadata.decode.attn_out_dtype,

519 device=q.device,

@@ -506,7 +540,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

540 )

541

542 rocm_aiter_ops.mla_decode_fwd(

509 - q,

543 + mla_padded_q,

544 kv_buffer,

545 o,

546 self.scale,

@@ -518,7 +552,4 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

552 **mla_kwargs,

553 )

554

521 - if self._needs_head_repeat:

522 - o = o[:, :: self._head_repeat_factor, :]

523 -

524 - return o, None

555 + return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, o), None

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

+10 -3

@@ -28,6 +28,9 @@ from vllm.v1.attention.backend import (

28 from vllm.v1.attention.backends.mla.flashmla_sparse import (

29 triton_convert_req_index_to_global_index,

30 )

31 +from vllm.v1.attention.backends.mla.rocm_aiter_mla import (

32 + AiterMLAHelper,

33 +)

34 from vllm.v1.kv_cache_interface import AttentionSpec

36 if TYPE_CHECKING:

@@ -299,6 +302,8 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])

302 indexer: "Indexer | None" = None,

303 **mla_args,

304 ) -> None:

305 + AiterMLAHelper.check_num_heads_validity(num_heads)

306 +

307 self.num_heads = num_heads

308 self.head_size = head_size

309 self.scale = float(scale)

@@ -317,8 +322,9 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])

322 attn_metadata: ROCMAiterMLASparseMetadata,

323 ) -> torch.Tensor:

324 num_tokens = q.shape[0]

325 + mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)

326 output = torch.empty(

321 - [num_tokens, self.num_heads, self.kv_lora_rank],

327 + [num_tokens, mla_num_heads, self.kv_lora_rank],

328 dtype=q.dtype,

329 device=q.device,

330 )

@@ -344,7 +350,7 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])

350 attn_metadata.paged_kv_last_page_len,

351 )

352

347 - return output[:, : self.num_heads, :]

353 + return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, output)

354

355 def forward_mqa(

356 self,

@@ -374,8 +380,9 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])

380 NUM_TOPK_TOKENS=attn_metadata.topk_tokens,

381 )

382

383 + mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)

384 attn_out = self._forward_bf16_kv(

378 - q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata

385 + mla_padded_q, kv_c_and_k_pe_cache, topk_indices_global, attn_metadata

386 )

387

388 return attn_out, None

vllm/v1/attention/ops/rocm_aiter_mla_sparse.py

@@ -339,6 +339,8 @@ def rocm_fp8_paged_mqa_logits(

339 device="cuda",

340 dtype=torch.float32,

341 )

342 + # TODO: 1. Replace _stage1 and out_qk.sum with another fused variant;

343 + # 2. Remove ChunkQ when AITER PR #2891 merged

344 deepgemm_fp8_paged_mqa_logits_stage1(

345 q_fp8,

346 kv_cache_fp8,

@@ -347,6 +349,7 @@ def rocm_fp8_paged_mqa_logits(

349 context_lens,

350 block_tables,

351 max_model_len,

352 + ChunkQ=heads,

353 )

354 return out_qk.sum(dim=0)

355 else: