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 13, 18:15 UTC
[ROCm][MLA] Support block-size > 1 for AITER MLA backend (#27224)
Signed-off-by: ganyi <ygan@amd.com> Co-authored-by: wuhuikx <hattie.wu@amd.com>
Pleaplusone committed
Nov 5, 2025 at 23:43 UTC
6cae1e53326ab89551352b926854af7cd96b688a
3 files changed
+34
-24
tests/kernels/attention/test_attention_selector.py
-7
@@ -104,13 +104,6 @@ def test_env(
104
16, torch.float16, None, block_size, use_mla=use_mla
105
)
106
assert f"The selected backend, {name}" in str(exc_info.value)
107
- elif name == "ROCM_AITER_MLA" and block_size != 1:
108
- # ROCM_AITER_MLA only supports block_size == 1
109
- with pytest.raises(ValueError) as exc_info:
110
- get_attn_backend(
111
- 16, torch.float16, None, block_size, use_mla=use_mla
112
- )
113
- assert f"The selected backend, {name}" in str(exc_info.value)
107
else:
108
# Valid backend-block_size combination
109
backend = get_attn_backend(
vllm/platforms/rocm.py
+3
-10
@@ -252,16 +252,9 @@ class RocmPlatform(Platform):
252
f"does not support block size {block_size}."
253
)
254
if selected_backend == _Backend.ROCM_AITER_MLA:
255
- if block_size == 1:
256
- logger.info("Using AITER MLA backend.")
257
- return (
258
- "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501
259
- )
260
- raise ValueError(
261
- f" The selected backend, {selected_backend.name},"
262
- f"does not support block size {block_size}."
263
- "(currently only supports block size 1)"
264
- )
255
+ logger.info("Using AITER MLA backend.")
256
+ return "vllm.v1.attention.backends.mla.rocm_aiter_mla.AiterMLABackend" # noqa: E501
257
+
258
raise ValueError(
259
f" The selected backend, {selected_backend.name},"
260
f"is not MLA type while requested for MLA backend."
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+31
-7
@@ -78,9 +78,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
78
super().__init__(
79
kv_cache_spec, layer_names, vllm_config, device, AiterMLAMetadata
80
)
81
- assert self.kv_cache_spec.block_size == 1, (
82
- "AITER MLAonly supports block size 1."
83
- )
81
82
self.compilation_config = vllm_config.compilation_config
83
max_num_pages_per_req = cdiv(
@@ -94,6 +91,11 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
91
# so we can only use the persistent buffer if a cudagraph is actually
92
# being used.
93
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
94
+ self.block_table_remapping = torch.zeros(
95
+ [max_num_reqs, max_num_pages_per_req * self.kv_cache_spec.block_size],
96
+ dtype=torch.int32,
97
+ device=device,
98
+ )
99
self.paged_kv_indptr = torch.zeros(
100
max_num_reqs + 1, dtype=torch.int32, device=device
101
)
@@ -119,13 +121,29 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
121
dcp_tot_seq_lens_device: torch.Tensor | None,
122
) -> AiterMLADecodeMetadata:
123
page_size = self.kv_cache_spec.block_size
122
- block_table_bounds = (seq_lens_device + page_size - 1) // page_size
124
device = self.device
125
num_reqs = seq_lens_device.size(0)
126
+ bs, _ = block_table_tensor.shape
127
+ block_table_tensor = (
128
+ block_table_tensor.unsqueeze(-1).expand(-1, -1, page_size) * page_size
129
+ )
130
+ block_table_tensor = (
131
+ block_table_tensor
132
+ + torch.arange(
133
+ 0,
134
+ page_size,
135
+ device=block_table_tensor.device,
136
+ dtype=block_table_tensor.dtype,
137
+ )[None, None, :]
138
+ )
139
+ block_table_tensor = block_table_tensor.view(bs, -1)
140
141
+ # after remapping, we assume the block size already equals to 1
142
+
143
+ max_blk_size_per_req = block_table_tensor.shape[-1]
144
mask = torch.arange(
145
block_table_tensor.size(1), dtype=block_table_tensor.dtype, device=device
128
- ).unsqueeze(0) < block_table_bounds.unsqueeze(1)
146
+ ).unsqueeze(0) < seq_lens_device.unsqueeze(1)
147
paged_kv_indices = block_table_tensor[mask]
148
149
paged_kv_last_page_len = seq_lens_device % page_size
@@ -135,13 +153,19 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
153
154
paged_kv_indptr = torch.cat(
155
[
138
- torch.zeros(1, dtype=block_table_bounds.dtype, device=device),
139
- block_table_bounds.cumsum(dim=0, dtype=torch.int32),
156
+ torch.zeros(1, dtype=seq_lens_device.dtype, device=device),
157
+ seq_lens_device.cumsum(dim=0, dtype=torch.int32),
158
]
159
)
160
161
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
162
num_actual_pages = paged_kv_indices.size(0)
163
+ self.block_table_remapping[:num_reqs, :max_blk_size_per_req].copy_(
164
+ block_table_tensor, non_blocking=True
165
+ )
166
+ block_table_tensor = self.block_table_remapping[
167
+ :num_reqs, :max_blk_size_per_req
168
+ ]
169
170
self.paged_kv_indices[:num_actual_pages].copy_(
171
paged_kv_indices, non_blocking=True