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 2, 11:00 UTC
[Bugfix] Corrupted MLA + linear attention (#43961)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
May 29, 2026 at 20:00 UTC
d2889722ff3a28ab7a49acd1b5926fcff9adf029
1 file changed
+10
-2
vllm/v1/core/single_type_kv_cache_manager.py
+10
-2
@@ -237,7 +237,11 @@ class SingleTypeKVCacheManager(ABC):
237
cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks)
238
)
239
req_blocks.extend(allocated_blocks)
240
- if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec):
240
+ if type(self.kv_cache_spec) in (
241
+ FullAttentionSpec,
242
+ TQFullAttentionSpec,
243
+ MLAAttentionSpec,
244
+ ):
245
self.new_block_ids.extend(b.block_id for b in allocated_blocks)
246
247
def allocate_new_blocks(
@@ -265,7 +269,11 @@ class SingleTypeKVCacheManager(ABC):
269
else:
270
new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
271
req_blocks.extend(new_blocks)
268
- if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec):
272
+ if type(self.kv_cache_spec) in (
273
+ FullAttentionSpec,
274
+ TQFullAttentionSpec,
275
+ MLAAttentionSpec,
276
+ ):
277
self.new_block_ids.extend(b.block_id for b in new_blocks)
278
return new_blocks
279