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 12, 22:15 UTC
[AMD] Use Decoupled Kernel Block Size to Support AITER MLA block_size=1 (#27715)
Signed-off-by: chiangzhang <chiangzhang@tencent.com>
Qiang Zhang committed
Nov 20, 2025 at 10:11 UTC
3fb0d90999887949629d1e9bac4d98336a35c475
2 files changed
+13
-46
vllm/attention/backends/abstract.py
+6
-8
@@ -119,14 +119,12 @@ class AttentionBackend(ABC):
119
return True
120
121
for supported_size in cls.supported_kernel_block_sizes:
122
- is_multiple_of = (
123
- isinstance(supported_size, MultipleOf)
124
- and block_size % supported_size.base == 0
125
- )
126
- is_int_equal = (
127
- isinstance(supported_size, int) and block_size == supported_size
128
- )
129
- if is_multiple_of or is_int_equal:
122
+ if isinstance(supported_size, MultipleOf):
123
+ supported_size = supported_size.base
124
+ # With hybrid_blocks feature, the framework-level block size
125
+ # only needs to be a multiple of the kernel's requirement,
126
+ # even if the kernel requires a fixed block_size.
127
+ if block_size % supported_size == 0:
128
return True
129
return False
130
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+7
-38
@@ -7,9 +7,8 @@ from typing import ClassVar
7
import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
-from vllm.attention.backends.abstract import AttentionLayer
10
+from vllm.attention.backends.abstract import AttentionLayer, MultipleOf
11
from vllm.config import VllmConfig
12
-from vllm.utils.math_utils import cdiv
12
from vllm.v1.attention.backends.mla.common import (
13
MLACommonBackend,
14
MLACommonDecodeMetadata,
@@ -22,6 +21,8 @@ from vllm.v1.kv_cache_interface import AttentionSpec
21
22
23
class AiterMLABackend(MLACommonBackend):
24
+ supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [1]
25
+
26
@staticmethod
27
def get_name() -> str:
28
return "ROCM_AITER_MLA"
@@ -71,9 +72,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
72
)
73
74
self.compilation_config = vllm_config.compilation_config
74
- max_num_pages_per_req = cdiv(
75
- vllm_config.model_config.max_model_len, self.kv_cache_spec.block_size
76
- )
75
+ # kernel block size is always 1.
76
+ max_num_pages_per_req = vllm_config.model_config.max_model_len
77
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
78
max_num_pages = max_num_reqs * max_num_pages_per_req
79
@@ -82,11 +82,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
82
# so we can only use the persistent buffer if a cudagraph is actually
83
# being used.
84
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
85
- self.block_table_remapping = torch.zeros(
86
- [max_num_reqs, max_num_pages_per_req * self.kv_cache_spec.block_size],
87
- dtype=torch.int32,
88
- device=device,
89
- )
85
self.paged_kv_indptr = torch.zeros(
86
max_num_reqs + 1, dtype=torch.int32, device=device
87
)
@@ -111,36 +106,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
106
num_decode_tokens: int,
107
dcp_tot_seq_lens_device: torch.Tensor | None,
108
) -> AiterMLADecodeMetadata:
114
- page_size = self.kv_cache_spec.block_size
109
+ # kernel block size is always 1, although the kv block size is not 1.
110
device = self.device
111
num_reqs = seq_lens_device.size(0)
117
- bs, _ = block_table_tensor.shape
118
- block_table_tensor = (
119
- block_table_tensor.unsqueeze(-1).expand(-1, -1, page_size) * page_size
120
- )
121
- block_table_tensor = (
122
- block_table_tensor
123
- + torch.arange(
124
- 0,
125
- page_size,
126
- device=block_table_tensor.device,
127
- dtype=block_table_tensor.dtype,
128
- )[None, None, :]
129
- )
130
- block_table_tensor = block_table_tensor.view(bs, -1)
112
132
- # after remapping, we assume the block size already equals to 1
133
-
134
- max_blk_size_per_req = block_table_tensor.shape[-1]
113
mask = torch.arange(
114
block_table_tensor.size(1), dtype=block_table_tensor.dtype, device=device
115
).unsqueeze(0) < seq_lens_device.unsqueeze(1)
116
paged_kv_indices = block_table_tensor[mask]
117
140
- paged_kv_last_page_len = seq_lens_device % page_size
141
- paged_kv_last_page_len = torch.where(
142
- paged_kv_last_page_len == 0, page_size, paged_kv_last_page_len
143
- )
118
+ paged_kv_last_page_len = torch.where(seq_lens_device == 0, 1, seq_lens_device)
119
120
paged_kv_indptr = torch.cat(
121
[
@@ -151,12 +126,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
126
127
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
128
num_actual_pages = paged_kv_indices.size(0)
154
- self.block_table_remapping[:num_reqs, :max_blk_size_per_req].copy_(
155
- block_table_tensor, non_blocking=True
156
- )
157
- block_table_tensor = self.block_table_remapping[
158
- :num_reqs, :max_blk_size_per_req
159
- ]
129
130
self.paged_kv_indices[:num_actual_pages].copy_(
131
paged_kv_indices, non_blocking=True