[BugFix] Avoid allocating KV cache for skipped indexers (vllm-project#11065)

## What this PR does

GLM-5.2 reuses top-k indices on most backbone layers, so those layers no
longer initialize a local Indexer. The model runner still described and
allocated an Indexer KV cache for every sparse MLA layer, wasting device
memory and underestimating the number of KV blocks that fit in the
available memory.

During KV cache grouping, the Ascend MLA merge path also accepted the
704- and 576-dimension layouts as identical and expanded every layer
back to the larger layout. This change keeps those layouts heterogeneous
through block planning. This PR builds a smaller per-layer KV cache spec
for sparse MLA layers that do not own an Indexer.

## How was this patch tested?


- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@1f486d9

Signed-off-by: ZYang6263 <zy626375@gmail.com>
Signed-off-by: ZYang6263 <50876451+ZYang6263@users.noreply.github.com>

Author: ZYang6263

tests/ut/worker/a2/test_model_runner_v1.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import torch
+from vllm.model_executor.layers.attention import MLAAttention
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
 
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
@@ -85,6 +86,58 @@ def test_reshape_kv_cache_uses_layer_spec_for_draft_gqa(self):
         self.assertEqual(k_cache.shape, (2, 16, 8, 64))
         self.assertEqual(v_cache.shape, (2, 16, 8, 64))
 
+    @patch("vllm_ascend.worker.model_runner_v1.has_ec_transfer", return_value=False)
+    @patch("vllm_ascend.worker.model_runner_v1.get_layers_from_vllm_config")
+    def test_sparse_layer_without_indexer_allocates_only_mla_kv_cache(
+        self,
+        mock_get_layers,
+        _mock_has_ec_transfer,
+    ):
+        runner = self._build_runner()
+        runner.use_sparse = True
+        runner.block_size = 16
+        runner.sparse_head_dim = (512, 64, 128)
+        runner.kv_cache_dtype = torch.bfloat16
+        runner.shared_kv_cache_layers = {}
+        runner.ascend_config = MagicMock()
+        runner.ascend_config.is_sparse_c8_layer.return_value = False
+        runner.model_config.hf_text_config = SimpleNamespace(
+            kv_lora_rank=512,
+            qk_rope_head_dim=64,
+        )
+        runner.vllm_config.cache_config.cache_dtype = "auto"
+
+        attn_module = MLAAttention.__new__(MLAAttention)
+        torch.nn.Module.__init__(attn_module)
+        attn_module.impl = SimpleNamespace(has_indexer=False)
+        layer_name = "model.layers.1.self_attn.attn"
+        mock_get_layers.return_value = {layer_name: attn_module}
+
+        spec = runner.get_kv_cache_spec()[layer_name]
+        self.assertEqual(spec.sparse_head_dim, (512, 64, 0))
+
+        kv_cache_config = KVCacheConfig(
+            num_blocks=2,
+            kv_cache_tensors=[
+                KVCacheTensor(
+                    size=spec.page_size_bytes * 2,
+                    shared_by=[layer_name],
+                )
+            ],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    layer_names=[layer_name],
+                    kv_cache_spec=spec,
+                )
+            ],
+        )
+
+        raw_caches = runner._allocate_kv_cache_tensors(kv_cache_config)
+        raw_k_cache, raw_v_cache = raw_caches[layer_name]
+
+        self.assertEqual(raw_k_cache.numel(), 2 * 16 * 512 * 2)
+        self.assertEqual(raw_v_cache.numel(), 2 * 16 * 64 * 2)
+
 
 class TestNPUModelRunnerOutputTokenIds(unittest.TestCase):
     def _build_runner(self):

vllm_ascend/core/kv_cache_interface.py

@@ -162,6 +162,21 @@ def merge(cls, specs: list[Self]) -> Self:
         assert all(isinstance(spec, MLAAttentionSpec) for spec in specs), (
             "All attention layers in the same KV cache group must be MLAAttentionSpec."
         )
+        layout_set = {
+            (
+                spec.block_size,
+                spec.num_kv_heads,
+                spec.head_size,
+                spec.scale_dim,
+                spec.scale_dtype,
+                spec.sparse_head_dim,
+                spec.dtype,
+            )
+            for spec in specs
+        }
+        assert len(layout_set) == 1, (
+            "All attention layers in the same KV cache group must use the same KV cache layout."
+        )
         cache_dtype_str_set = set(spec.cache_dtype_str for spec in specs)
         assert len(cache_dtype_str_set) == 1, (
             "All attention layers in the same KV cache group must use the same quantization method."

vllm_ascend/utils.py

@@ -41,6 +41,7 @@
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import AttentionSpec
 else:
     VllmConfig = None
 
@@ -1693,6 +1694,11 @@ def kv_cache_spec_uses_sparse_c8(kv_cache_spec) -> bool:
     return isinstance(kv_cache_spec, AscendMLAAttentionSpec) and bool(getattr(kv_cache_spec, "cache_sparse_c8", False))
 
 
+def sparse_kv_cache_has_indexer(kv_cache_spec: AttentionSpec) -> bool:
+    sparse_head_dim = getattr(kv_cache_spec, "sparse_head_dim", None)
+    return sparse_head_dim is not None and len(sparse_head_dim) == 3 and sparse_head_dim[2] > 0
+
+
 def is_hidden_state_cache_spec(spec) -> bool:
     """Whether ``spec`` marks an ``extract_hidden_states`` cache-only layer."""
     from vllm.v1.kv_cache_interface import HiddenStateCacheSpec

vllm_ascend/worker/model_runner_v1.py

@@ -158,6 +158,7 @@
     set_potential_max_tokens,
     set_weight_prefetch_method,
     should_skip_allreduce_across_dp_group,
+    sparse_kv_cache_has_indexer,
     vllm_version_is,
 )
 from vllm_ascend.worker.npu_input_batch import NPUInputBatch
@@ -4260,20 +4261,29 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                         # for deepseek v3.2, we split the kv cache according to the corresponding ratio
                         kv_cache_spec = layer_kv_cache_spec[layer_name]
                         current_sparse_c8 = kv_cache_spec_uses_sparse_c8(kv_cache_spec)
-                        sparse_kv_cache_ratio = kv_cache_spec.sparse_kv_cache_ratio
-                        
-                        # A5 sparse C8: (ckv_ratio, qli_ratio, qli_scale_ratio, None)
-                        # A3 sparse C8: (k_ratio, v_ratio, qli_ratio, qli_scale_ratio)
-                        if current_sparse_c8 and get_ascend_device_type() == AscendDeviceType.A5:
-                            k_tensor_split_factor = sparse_kv_cache_ratio[0]  # ckv
-                            v_tensor_split_factor = None  # merged
-                            dsa_k_tensor_split_factor = sparse_kv_cache_ratio[1]  # qli_tensor
-                            dsa_k_scale_tensor_split_factor = sparse_kv_cache_ratio[2]  # qli_scale
+                        has_indexer_cache = sparse_kv_cache_has_indexer(kv_cache_spec)
+
+                        if has_indexer_cache:
+                            sparse_kv_cache_ratio = kv_cache_spec.sparse_kv_cache_ratio
+
+                            # A5 sparse C8: (ckv_ratio, qli_ratio, qli_scale_ratio, None)
+                            # A3 sparse C8: (k_ratio, v_ratio, qli_ratio, qli_scale_ratio)
+                            if current_sparse_c8 and get_ascend_device_type() == AscendDeviceType.A5:
+                                k_tensor_split_factor = sparse_kv_cache_ratio[0]  # ckv
+                                v_tensor_split_factor = None  # merged
+                                dsa_k_tensor_split_factor = sparse_kv_cache_ratio[1]  # qli_tensor
+                                dsa_k_scale_tensor_split_factor = sparse_kv_cache_ratio[2]  # qli_scale
+                            else:
+                                k_tensor_split_factor = sparse_kv_cache_ratio[0]
+                                v_tensor_split_factor = sparse_kv_cache_ratio[1]
+                                dsa_k_tensor_split_factor = sparse_kv_cache_ratio[2]
+                                dsa_k_scale_tensor_split_factor = (
+                                    sparse_kv_cache_ratio[3] if current_sparse_c8 else None
+                                )
                         else:
-                            k_tensor_split_factor = sparse_kv_cache_ratio[0]
-                            v_tensor_split_factor = sparse_kv_cache_ratio[1]
-                            dsa_k_tensor_split_factor = sparse_kv_cache_ratio[2]
-                            dsa_k_scale_tensor_split_factor = sparse_kv_cache_ratio[3] if current_sparse_c8 else None
+                            assert not current_sparse_c8
+                            k_dim, v_dim, _ = kv_cache_spec.sparse_head_dim
+                            k_tensor_split_factor, v_tensor_split_factor = calc_split_factor([k_dim, v_dim])
                     else:
                         k_dim, v_dim = self._get_attention_kv_cache_dims(layer_name, current_kv_cache_spec)
                         assert k_dim > 0 and v_dim > 0
@@ -4296,9 +4306,9 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                     dsa_k_tensor_size = None
                     dsa_k_scale_tensor_size = None
                     #### for deepseek sparse attention
-                    if self.use_sparse:
+                    if self.use_sparse and has_indexer_cache:
                         dsa_k_tensor_size = int(kv_cache_tensor.size // dsa_k_tensor_split_factor)
-                    if self.use_sparse and current_sparse_c8:
+                    if self.use_sparse and has_indexer_cache and current_sparse_c8:
                         dsa_k_scale_tensor_size = int(kv_cache_tensor.size // dsa_k_scale_tensor_split_factor)
 
                     # Allocate raw int8 tensors. Even bf16/fp16 KV cache entries
@@ -4317,7 +4327,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                             alignment,
                         )
 
-                    if self.use_sparse:
+                    if self.use_sparse and has_indexer_cache:
                         assert dsa_k_tensor_size is not None
 
                         if current_sparse_c8:
@@ -4342,7 +4352,9 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                         # shared the attn kvcache for all shared layers
                         if "attn" in layer_name_inner and "linear_attn" not in layer_name_inner:
                             if self.use_sparse:
-                                if current_sparse_c8:
+                                if not has_indexer_cache:
+                                    kv_cache_raw_tensors[layer_name_inner] = (k_tensor, v_tensor)
+                                elif current_sparse_c8:
                                     if get_ascend_device_type() == AscendDeviceType.A5:
                                         kv_cache_raw_tensors[layer_name_inner] = (
                                             k_tensor, dsa_k_tensor, dsa_k_scale_tensor
@@ -4487,8 +4499,9 @@ def _reshape_kv_cache_tensors(
                     # _allocate_kv_cache_tensors; route them to the dedicated
                     # elif branch below before the sparse branch tries to
                     # unpack them as a (k, v, dsa_k[, scale]) tuple.
-                    if self.use_sparse and "cache_only_layers" not in layer_name:
-                        current_sparse_c8 = kv_cache_spec_uses_sparse_c8(current_kv_cache_spec)
+                    current_sparse_c8 = kv_cache_spec_uses_sparse_c8(current_kv_cache_spec)
+                    has_indexer_cache = sparse_kv_cache_has_indexer(current_kv_cache_spec)
+                    if self.use_sparse and has_indexer_cache and "cache_only_layers" not in layer_name:
                         if current_sparse_c8:
                             if get_ascend_device_type() == AscendDeviceType.A5:
                                 raw_k_tensor, raw_dsa_k_tensor, raw_dsa_k_scale_tensor = kv_cache_raw_tensors[  # type: ignore
@@ -4672,7 +4685,7 @@ def _reshape_kv_cache_tensors(
                     else:
                         v_cache = raw_v_tensor.view(v_cache_dtype).view(v_shape)
 
-                    if self.use_sparse:
+                    if self.use_sparse and has_indexer_cache:
                         dsa_k_cache_shape = (
                             num_blocks,
                             current_kv_cache_spec.block_size,
@@ -4974,14 +4987,25 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
 
             elif isinstance(attn_module, MLAAttention):
                 if self.use_sparse:
+                    has_indexer = attn_module.impl.has_indexer
+                    if has_indexer:
+                        sparse_head_dim = self.sparse_head_dim
+                    else:
+                        # Layers that reuse another layer's top-k indices only
+                        # need the MLA latent and RoPE caches.
+                        sparse_head_dim = (
+                            self.model_config.hf_text_config.kv_lora_rank,
+                            self.model_config.hf_text_config.qk_rope_head_dim,
+                            0,
+                        )
                     kv_cache_spec[layer_name] = AscendMLAAttentionSpec(
                         block_size=self.block_size,
                         num_kv_heads=1,
-                        head_size=sum(self.sparse_head_dim),
-                        sparse_head_dim=self.sparse_head_dim,
+                        head_size=sum(sparse_head_dim),
+                        sparse_head_dim=sparse_head_dim,
                         dtype=self.kv_cache_dtype,
                         cache_dtype_str=self.vllm_config.cache_config.cache_dtype,
-                        cache_sparse_c8=self.ascend_config.is_sparse_c8_layer(layer_name),
+                        cache_sparse_c8=has_indexer and self.ascend_config.is_sparse_c8_layer(layer_name),
                     )
                 elif spec := attn_module.get_kv_cache_spec(self.vllm_config):
                     if getattr(attn_module.impl, "fa_quant_layer", False):