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 10, 19:00 UTC
[Bugfix][Hardware][AMD] Fix last_page_len calculation in AITER MLA decode (#31282)
Signed-off-by: c0de128 <kevin.mckay@outlook.com>
Kevin McKay committed
Jan 1, 2026 at 23:14 UTC
825c2dc133d490c468a35c573500107f44b32ac8
1 file changed
+12
-9
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+12
-9
@@ -88,6 +88,13 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
88
# TODO: we can disambiguate between decode and mixed-prefill decode here
89
# so we can only use the persistent buffer if a cudagraph is actually
90
# being used.
91
+
92
+ # paged_kv_last_page_len is always 1s (kernel block size is always 1),
93
+ # so we create it once and reuse slices in both eager and cudagraph modes.
94
+ self.paged_kv_last_page_len = torch.ones(
95
+ max_num_reqs, dtype=torch.int32, device=device
96
+ )
97
+
98
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
99
self.paged_kv_indptr = torch.zeros(
100
max_num_reqs + 1, dtype=torch.int32, device=device
@@ -95,9 +102,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
102
self.paged_kv_indices = torch.zeros(
103
max_num_pages, dtype=torch.int32, device=device
104
)
98
- self.paged_kv_last_page_len = torch.zeros(
99
- max_num_reqs, dtype=torch.int32, device=device
100
- )
105
106
self.qo_indptr = torch.zeros(
107
max_num_reqs + 1, dtype=torch.int32, device=device
@@ -122,7 +126,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
126
).unsqueeze(0) < seq_lens_device.unsqueeze(1)
127
paged_kv_indices = block_table_tensor[mask]
128
125
- paged_kv_last_page_len = torch.where(seq_lens_device == 0, 1, seq_lens_device)
129
+ # kernel block size is always 1, so each page has exactly 1 token.
130
+ # last_page_len is always 1 - just slice the pre-initialized buffer.
131
+ paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
132
133
paged_kv_indptr = torch.cat(
134
[
@@ -148,11 +154,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
154
self.paged_kv_indptr[1 + num_reqs :].fill_(paged_kv_indptr[-1])
155
paged_kv_indptr = self.paged_kv_indptr[: 1 + num_reqs]
156
151
- self.paged_kv_last_page_len[:num_reqs].copy_(
152
- paged_kv_last_page_len, non_blocking=True
153
- )
154
- self.paged_kv_last_page_len[num_reqs:].fill_(1)
155
- paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
157
+ # paged_kv_last_page_len already uses the pre-initialized buffer slice
158
+ # (set above), so no copy needed - buffer is always 1s.
159
160
self.qo_indptr[: 1 + num_reqs].copy_(
161
query_start_loc_device, non_blocking=True