From d2889722ff3a28ab7a49acd1b5926fcff9adf029 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Fri, 29 May 2026 20:00:51 +0800 Subject: [PATCH] [Bugfix] Corrupted MLA + linear attention (#43961) Signed-off-by: Thien Tran --- vllm/v1/core/single_type_kv_cache_manager.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index e29919022e..cd000dc849 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -237,7 +237,11 @@ class SingleTypeKVCacheManager(ABC): cdiv(num_total_computed_tokens, self.block_size) - len(req_blocks) ) req_blocks.extend(allocated_blocks) - if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec): + if type(self.kv_cache_spec) in ( + FullAttentionSpec, + TQFullAttentionSpec, + MLAAttentionSpec, + ): self.new_block_ids.extend(b.block_id for b in allocated_blocks) def allocate_new_blocks( @@ -265,7 +269,11 @@ class SingleTypeKVCacheManager(ABC): else: new_blocks = self.block_pool.get_new_blocks(num_new_blocks) req_blocks.extend(new_blocks) - if type(self.kv_cache_spec) in (FullAttentionSpec, TQFullAttentionSpec): + if type(self.kv_cache_spec) in ( + FullAttentionSpec, + TQFullAttentionSpec, + MLAAttentionSpec, + ): self.new_block_ids.extend(b.block_id for b in new_blocks) return new_blocks