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 6, 21:45 UTC
AITER MLA backend: Avoid CPU sync in _build_decode (#35765)
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
pschlan-amd committed
Mar 11, 2026 at 08:25 UTC
eac2dc2b410dc11af4b424802e86ef9d36bac28a
1 file changed
+46
-15
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+46
-15
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
17
MLACommonMetadataBuilder,
18
QueryLenSupport,
19
)
20
+from vllm.triton_utils import tl, triton
21
from vllm.v1.attention.backend import AttentionCGSupport, AttentionLayer, MultipleOf
22
from vllm.v1.kv_cache_interface import AttentionSpec
23
@@ -108,13 +109,16 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
109
max_num_reqs, dtype=torch.int32, device=device
110
)
111
112
+ # Persistent buffer for paged_kv_indices to avoid blocking boolean mask
113
+ # indexing (block_table_tensor[mask]) which has data-dependent output size.
114
+ self.paged_kv_indices = torch.zeros(
115
+ max_num_pages, dtype=torch.int32, device=device
116
+ )
117
+
118
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
119
self.paged_kv_indptr = torch.zeros(
120
max_num_reqs + 1, dtype=torch.int32, device=device
121
)
115
- self.paged_kv_indices = torch.zeros(
116
- max_num_pages, dtype=torch.int32, device=device
117
- )
122
123
self.qo_indptr = torch.zeros(
124
max_num_reqs + 1, dtype=torch.int32, device=device
@@ -134,11 +138,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
138
device = self.device
139
num_reqs = seq_lens_device.size(0)
140
137
- mask = torch.arange(
138
- block_table_tensor.size(1), dtype=block_table_tensor.dtype, device=device
139
- ).unsqueeze(0) < seq_lens_device.unsqueeze(1)
140
- paged_kv_indices = block_table_tensor[mask]
141
-
141
# kernel block size is always 1, so each page has exactly 1 token.
142
# last_page_len is always 1 - just slice the pre-initialized buffer.
143
paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
@@ -153,14 +152,17 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
152
max_qo_len = qo_len.max().item()
153
154
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
156
- num_actual_pages = paged_kv_indices.size(0)
157
-
158
- self.paged_kv_indices[:num_actual_pages].copy_(
159
- paged_kv_indices, non_blocking=True
160
- )
161
- self.paged_kv_indices[num_actual_pages:].fill_(-1)
162
- paged_kv_indices = self.paged_kv_indices[:num_actual_pages]
155
+ self.paged_kv_indices.fill_(-1)
156
+ _copy_page_indices_kernel[(num_reqs,)](
157
+ self.paged_kv_indices,
158
+ block_table_tensor,
159
+ block_table_tensor.stride(0),
160
+ paged_kv_indptr,
161
+ BLOCK_SIZE=1024,
162
+ )
163
+ paged_kv_indices = self.paged_kv_indices
164
165
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
166
self.paged_kv_indptr[: 1 + num_reqs].copy_(
167
paged_kv_indptr, non_blocking=True
168
)
@@ -196,6 +198,35 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
198
return attn_metadata
199
200
201
+@triton.jit
202
+def _copy_page_indices_kernel(
203
+ page_indices,
204
+ block_table,
205
+ block_table_stride,
206
+ cu_num_blocks,
207
+ BLOCK_SIZE: tl.constexpr,
208
+):
209
+ """Copy block table rows into a flat page_indices buffer using indptr.
210
+ Avoids blocking boolean mask indexing (tensor[mask]) which has
211
+ data-dependent output size and forces sync.
212
+ This is the same kernel as introduced in backends/flashinfer.py.
213
+ """
214
+ req_idx = tl.program_id(0)
215
+ row_ptr = block_table + req_idx * block_table_stride
216
+ start_idx = tl.load(cu_num_blocks + req_idx)
217
+ end_idx = tl.load(cu_num_blocks + req_idx + 1)
218
+ num_blocks = end_idx - start_idx
219
+
220
+ offset = tl.arange(0, BLOCK_SIZE)
221
+ for i in tl.range(0, num_blocks, BLOCK_SIZE):
222
+ block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks)
223
+ tl.store(
224
+ page_indices + start_idx + i + offset,
225
+ block_ids,
226
+ mask=i + offset < num_blocks,
227
+ )
228
+
229
+
230
class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
231
def __init__(
232
self,