From f4e862ac1793afab35f216c719f57a3ca0fc6b82 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Mon, 16 Mar 2026 19:52:28 +0800
Subject: [PATCH 01/17] Fix kvcache allocate for different target/draft models.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/worker/model_runner_v1.py | 105 ++++++++++++++++----------
 1 file changed, 65 insertions(+), 40 deletions(-)

diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index 017ca2d32db..76a2fcdb7a8 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -2612,6 +2612,34 @@ def initialize_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[st
         bind_kv_cache(kv_caches, self.compilation_config.static_forward_context, self.kv_caches, num_attn_module)
         return kv_caches
 
+    def _get_layer_kv_cache_specs(self, kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
+        layer_kv_cache_spec: dict[str, KVCacheSpec] = {}
+        for group_kv_cache_spec in kv_cache_config.kv_cache_groups:
+            group_spec = group_kv_cache_spec.kv_cache_spec
+            for layer_name in group_kv_cache_spec.layer_names:
+                if isinstance(group_spec, UniformTypeKVCacheSpecs):
+                    layer_kv_cache_spec[layer_name] = group_spec.kv_cache_specs[layer_name]
+                else:
+                    layer_kv_cache_spec[layer_name] = group_spec
+        return layer_kv_cache_spec
+
+    def _get_attention_kv_cache_dims(self, layer_name: str, kv_cache_spec: AttentionSpec) -> tuple[int, int]:
+        if isinstance(kv_cache_spec, MLAAttentionSpec):
+            attn_layers = get_layers_from_vllm_config(
+                self.vllm_config,
+                AttentionLayerBase,
+                [layer_name],
+            )
+            attn_layer = attn_layers[layer_name]
+            if not isinstance(attn_layer, MLAAttention):
+                raise TypeError(
+                    f"Expected MLAAttention layer for {layer_name}, got {type(attn_layer).__name__}."
+                )
+            return attn_layer.kv_lora_rank, attn_layer.qk_rope_head_dim
+
+        head_size_v = kv_cache_spec.head_size_v if hasattr(kv_cache_spec, "head_size_v") else kv_cache_spec.head_size
+        return kv_cache_spec.head_size, head_size_v
+
     def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str, torch.Tensor]:
         """
         Initializes the KV cache buffer with the correct size. The buffer needs
@@ -2632,10 +2660,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
         kv_cache_raw_tensors: dict[str, torch.Tensor | torch.Tensor | None] = {}
         # prefill disaggregation need the addr of cache tensor be aligned with 2M
         alignment = 2 * 1024 * 1024
-        layer_kv_cache_spec: dict[str, KVCacheSpec] = {}
-        for group_kv_cache_spec in kv_cache_config.kv_cache_groups:
-            for layer_name in group_kv_cache_spec.layer_names:
-                layer_kv_cache_spec[layer_name] = group_kv_cache_spec.kv_cache_spec
+        layer_kv_cache_spec = self._get_layer_kv_cache_specs(kv_cache_config)
         # If some tensors are shared by linear layers and attention layers,
         # the same tensor format must be maintained even if some layers
         # have only linear or attention layers, for example, the mtp layer.
@@ -2670,19 +2695,12 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                     # as it only support the 0-dim of kv_cache is `num_blocks`.
                     # For deepseek mla, we need to spilt cache tensor accrodding to the nope head dim
                     # and rope head dim.
-                    if self.model_config.use_mla:
-                        head_size = (
-                            self.model_config.hf_text_config.qk_rope_head_dim
-                            + self.model_config.hf_text_config.kv_lora_rank
-                        )
+                    current_kv_cache_spec = layer_kv_cache_spec[layer_name]
+                    assert isinstance(current_kv_cache_spec, AttentionSpec)
 
                     dsa_k_cache_factor = None
                     dsa_k_cache_size = None
-                    if not self.model_config.use_mla:
-                        # for non-mla model, use FullAttentionSpec
-                        k_tensor_split_factor = 2
-                        v_tensor_split_factor = 2
-                    elif self.use_sparse:
+                    if self.use_sparse:
                         # for deepseek v3.2, we split the kv cache according to the corresponding ratio
                         sparse_sum_head_size = sum(self._get_sparse_kv_cache_ratio())
                         k_tensor_split_factor, v_tensor_split_factor, dsa_k_cache_factor = [  # type: ignore
@@ -2690,9 +2708,10 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                         ]
                         dsa_k_cache_size = int(kv_cache_tensor.size // dsa_k_cache_factor)
                     else:
-                        # for other deepseek models, use MLAAttentionSpec
-                        k_tensor_split_factor = head_size / self.model_config.hf_text_config.kv_lora_rank
-                        v_tensor_split_factor = head_size / self.model_config.hf_text_config.qk_rope_head_dim
+                        k_dim, v_dim = self._get_attention_kv_cache_dims(layer_name, current_kv_cache_spec)
+                        total_dim = k_dim + v_dim
+                        k_tensor_split_factor = total_dim / k_dim
+                        v_tensor_split_factor = total_dim / v_dim
 
                     k_tensor_size = int(kv_cache_tensor.size // k_tensor_split_factor)
                     v_tensor_size = int(kv_cache_tensor.size // v_tensor_split_factor)
@@ -2752,10 +2771,7 @@ def _reshape_kv_cache_tensors(
             corresponding memory buffer for KV cache.
         """
         kv_caches: dict[str, torch.Tensor] = {}
-        layer_kv_cache_spec = {}
-        for group in kv_cache_config.kv_cache_groups:
-            for layer_name in group.layer_names:
-                layer_kv_cache_spec[layer_name] = group.kv_cache_spec
+        layer_kv_cache_spec = self._get_layer_kv_cache_specs(kv_cache_config)
         for group in self._kv_cache_spec_attn_group_iterator():
             kv_cache_spec = group.kv_cache_spec
             attn_backend = group.backend
@@ -2763,9 +2779,11 @@ def _reshape_kv_cache_tensors(
                 if layer_name in self.runner_only_attn_layers:
                     continue
 
+                current_kv_cache_spec = layer_kv_cache_spec[layer_name]
+
                 # TODO: remove this after the OOM issue is located and fixed, otherwise, some model may
                 # encounter OOM issue
-                if isinstance(kv_cache_spec, AttentionSpec):
+                if isinstance(current_kv_cache_spec, AttentionSpec):
                     raw_dsa_k_tensor = None
                     if self.use_sparse:
                         raw_k_tensor, raw_v_tensor, raw_dsa_k_tensor = kv_cache_raw_tensors[  # type: ignore
@@ -2785,8 +2803,8 @@ def _reshape_kv_cache_tensors(
                         sum_page_size_bytes = raw_k_tensor.numel() + raw_v_tensor.numel()
                     assert raw_k_tensor is not None
                     assert raw_v_tensor is not None
-                    assert sum_page_size_bytes % kv_cache_spec.page_size_bytes == 0
-                    num_blocks = sum_page_size_bytes // kv_cache_spec.page_size_bytes
+                    assert sum_page_size_bytes % current_kv_cache_spec.page_size_bytes == 0
+                    num_blocks = sum_page_size_bytes // current_kv_cache_spec.page_size_bytes
 
                     # `num_blocks` is the number of blocks the model runner can use.
                     # `kv_cache_config.num_blocks` is the number of blocks that
@@ -2800,16 +2818,16 @@ def _reshape_kv_cache_tensors(
                     if hasattr(attn_backend, "get_supported_kernel_block_sizes") and self.use_hybrid_blocks:
                         block_size = attn_backend.get_supported_kernel_block_sizes()[0]
 
-                        block_size_chunk = kv_cache_spec.block_size // block_size
+                        block_size_chunk = current_kv_cache_spec.block_size // block_size
                         kv_cache_shape = attn_backend.get_kv_cache_shape(
                             num_blocks * block_size_chunk,
                             block_size,
-                            kv_cache_spec.num_kv_heads,
-                            kv_cache_spec.head_size,
+                            current_kv_cache_spec.num_kv_heads,
+                            current_kv_cache_spec.head_size,
                         )
                         if self.hybrid_with_attn_and_mamba:
                             attn_tensor_page_size = int(np.prod(kv_cache_shape[1:])) * get_dtype_size(
-                                kv_cache_spec.dtype
+                                current_kv_cache_spec.dtype
                             )
                             conv_block_padding_size = raw_k_tensor.numel() - attn_tensor_page_size * 2
                             raw_kv_tensor = raw_k_tensor[conv_block_padding_size:]
@@ -2817,26 +2835,33 @@ def _reshape_kv_cache_tensors(
                             raw_v_tensor = raw_kv_tensor[attn_tensor_page_size:]
                     else:
                         kv_cache_shape = self.attn_backend.get_kv_cache_shape(
-                            num_blocks, kv_cache_spec.block_size, kv_cache_spec.num_kv_heads, kv_cache_spec.head_size
+                            num_blocks,
+                            current_kv_cache_spec.block_size,
+                            current_kv_cache_spec.num_kv_heads,
+                            current_kv_cache_spec.head_size,
                         )
-                    dtype = kv_cache_spec.dtype
-                    if not self.model_config.use_mla:
+                    dtype = current_kv_cache_spec.dtype
+                    if not isinstance(current_kv_cache_spec, MLAAttentionSpec):
                         k_shape = kv_cache_shape[1:]
-                        v_shape = k_shape
+                        if hasattr(current_kv_cache_spec, "head_size_v"):
+                            v_shape = (*kv_cache_shape[1:-1], current_kv_cache_spec.head_size_v)
+                        else:
+                            v_shape = k_shape
                     else:
                         # k_cache: nope_cache    v_cache: rope_cache
                         mla_num_blocks, mla_block_size, num_kv_heads, _ = kv_cache_shape
+                        k_dim, v_dim = self._get_attention_kv_cache_dims(layer_name, current_kv_cache_spec)
                         k_shape = [
                             mla_num_blocks,
                             mla_block_size,
                             num_kv_heads,
-                            self.model_config.hf_text_config.kv_lora_rank,
+                            k_dim,
                         ]
                         v_shape = [
                             mla_num_blocks,
                             mla_block_size,
                             num_kv_heads,
-                            self.model_config.hf_text_config.qk_rope_head_dim,
+                            v_dim,
                         ]
                     k_cache = raw_k_tensor.view(dtype).view(k_shape)
                     v_cache = raw_v_tensor.view(dtype).view(v_shape)
@@ -2845,19 +2870,19 @@ def _reshape_kv_cache_tensors(
                         index_head_dim = self._get_sparse_kv_cache_ratio()[-1]
                         dsa_k_cache_shape = (
                             num_blocks,
-                            kv_cache_spec.block_size,
-                            kv_cache_spec.num_kv_heads,
+                            current_kv_cache_spec.block_size,
+                            current_kv_cache_spec.num_kv_heads,
                             index_head_dim,
                         )
                         dsa_k_cache = raw_dsa_k_tensor.view(dtype).view(dsa_k_cache_shape)
                         kv_caches[layer_name] = (k_cache, v_cache, dsa_k_cache)
                     else:
                         kv_caches[layer_name] = (k_cache, v_cache)
-                elif isinstance(kv_cache_spec, MambaSpec):
+                elif isinstance(current_kv_cache_spec, MambaSpec):
                     raw_tensor = kv_cache_raw_tensors[layer_name]
                     assert raw_tensor is not None
-                    assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
-                    num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
+                    assert raw_tensor.numel() % current_kv_cache_spec.page_size_bytes == 0
+                    num_blocks = raw_tensor.numel() // current_kv_cache_spec.page_size_bytes
                     assert num_blocks >= kv_cache_config.num_blocks
 
                     # `num_blocks` is the number of blocks the model runner can use.
@@ -2877,7 +2902,7 @@ def _reshape_kv_cache_tensors(
                     # tensor1: [(kv_padding), conv           , ...]
                     # tensor2: [k           , ssm            , ...]
                     # tensor3: [v           , (mamba_padding), ...]
-                    for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
+                    for shape, dtype in zip(current_kv_cache_spec.shapes, current_kv_cache_spec.dtypes):
                         # normally, there is conv state and ssm state in this loop. And there is only
                         # a conv state in some special models.
                         target_shape = (num_blocks, *shape)

From d2f09d576fa9416260f8dfdca23ec286329a21d4 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Mon, 16 Mar 2026 20:26:05 +0800
Subject: [PATCH 02/17] Fix0.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/worker/model_runner_v1.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index 3e7089294f3..4466f56cecb 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -2806,7 +2806,6 @@ def _reshape_kv_cache_tensors(
         kv_caches: dict[str, torch.Tensor] = {}
         layer_kv_cache_spec = self._get_layer_kv_cache_specs(kv_cache_config)
         for group in self._kv_cache_spec_attn_group_iterator():
-            kv_cache_spec = group.kv_cache_spec
             attn_backend = group.backend
             for layer_name in group.layer_names:
                 if layer_name in self.runner_only_attn_layers:
@@ -2905,8 +2904,8 @@ def _reshape_kv_cache_tensors(
                             num_kv_heads,
                             v_dim,
                         ]
-                    k_cache = raw_k_tensor.view(kv_cache_spec.dtype).view(k_shape)
-                    v_cache = raw_v_tensor.view(kv_cache_spec.dtype).view(v_shape)
+                    k_cache = raw_k_tensor.view(current_kv_cache_spec.dtype).view(k_shape)
+                    v_cache = raw_v_tensor.view(current_kv_cache_spec.dtype).view(v_shape)
 
                     if self.use_sparse:
                         dsa_k_cache_shape = (
@@ -2921,8 +2920,8 @@ def _reshape_kv_cache_tensors(
                             # dsa_k_scale
                             dsa_k_scale_cache_shape = (
                                 num_blocks,
-                                kv_cache_spec.block_size,
-                                kv_cache_spec.num_kv_heads,
+                                current_kv_cache_spec.block_size,
+                                current_kv_cache_spec.num_kv_heads,
                                 1,
                             )
                             assert raw_dsa_k_scale_tensor is not None
@@ -2934,7 +2933,7 @@ def _reshape_kv_cache_tensors(
                             kv_caches[layer_name] = (k_cache, v_cache, dsa_k_cache, dsa_k_scale_cache)
                         else:
                             # dsa_k
-                            dsa_k_cache = raw_dsa_k_tensor.view(kv_cache_spec.dtype).view(dsa_k_cache_shape)
+                            dsa_k_cache = raw_dsa_k_tensor.view(current_kv_cache_spec.dtype).view(dsa_k_cache_shape)
                             kv_caches[layer_name] = (k_cache, v_cache, dsa_k_cache)
                     else:
                         kv_caches[layer_name] = (k_cache, v_cache)

From 2f780723e08199ad3e1f2087480a19cab7afeb4f Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Mon, 16 Mar 2026 20:36:52 +0800
Subject: [PATCH 03/17] Add kv dim validation.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/worker/model_runner_v1.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index 4466f56cecb..909ed2c30b2 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -2721,6 +2721,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str
                         dsa_k_scale_tensor_split_factor = sparse_kv_cache_ratio[3]
                     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
                         total_dim = k_dim + v_dim
                         k_tensor_split_factor = total_dim / k_dim
                         v_tensor_split_factor = total_dim / v_dim

From 711eef1935e3871be3ca7f3c1ee3fba909b65ae5 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Tue, 17 Mar 2026 09:20:05 +0800
Subject: [PATCH 04/17] Fix1.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/spec_decode/eagle_proposer.py | 4 ++--
 vllm_ascend/worker/model_runner_v1.py     | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/vllm_ascend/spec_decode/eagle_proposer.py b/vllm_ascend/spec_decode/eagle_proposer.py
index aa86823fe45..306337e3040 100644
--- a/vllm_ascend/spec_decode/eagle_proposer.py
+++ b/vllm_ascend/spec_decode/eagle_proposer.py
@@ -388,7 +388,7 @@ def dummy_run(
                     : num_reqs * self.decode_threshold
                 ]
 
-            builder = self.runner.attn_groups[0][0].get_metadata_builder()
+            builder = self.draft_attn_groups[0].get_metadata_builder()
             # update the tensor's address for each step.
             for draft_step in range(self.num_speculative_tokens):
                 common_attn_metadata = self.shallow_copy_metadata(common_attn_metadata)
@@ -598,7 +598,7 @@ def _propose(
         common_attn_metadata.slot_mapping = self.slot_mapping_group[0]
         common_attn_metadata.num_input_tokens = num_input_tokens
         # FIXME(woosuk): The below two ops cause synchronization. Optimize.
-        builder = self.runner.attn_groups[0][0].get_metadata_builder()
+        builder = self.draft_attn_groups[0].get_metadata_builder()
         attn_metadata = builder.build(0, common_attn_metadata, self.runner.get_model())
 
         if self.uses_mrope:
diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index 909ed2c30b2..2a520d21c94 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -2877,7 +2877,7 @@ def _reshape_kv_cache_tensors(
                             raw_k_tensor = raw_kv_tensor[:attn_tensor_page_size]
                             raw_v_tensor = raw_kv_tensor[attn_tensor_page_size:]
                     else:
-                        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
+                        kv_cache_shape = attn_backend.get_kv_cache_shape(
                             num_blocks,
                             current_kv_cache_spec.block_size,
                             current_kv_cache_spec.num_kv_heads,

From 66b6379db5368c72d6e107c7c0fc8dff69ed6e8b Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Tue, 17 Mar 2026 09:54:54 +0800
Subject: [PATCH 05/17] Fix for lint.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/worker/model_runner_v1.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
index d026ad5c14a..c0dd83c1853 100644
--- a/vllm_ascend/worker/model_runner_v1.py
+++ b/vllm_ascend/worker/model_runner_v1.py
@@ -2901,18 +2901,18 @@ def _reshape_kv_cache_tensors(
                         # k_cache: nope_cache    v_cache: rope_cache
                         mla_num_blocks, mla_block_size, num_kv_heads, _ = kv_cache_shape
                         k_dim, v_dim = self._get_attention_kv_cache_dims(layer_name, current_kv_cache_spec)
-                        k_shape = [
+                        k_shape = (
                             mla_num_blocks,
                             mla_block_size,
                             num_kv_heads,
                             k_dim,
-                        ]
-                        v_shape = [
+                        )
+                        v_shape = (
                             mla_num_blocks,
                             mla_block_size,
                             num_kv_heads,
                             v_dim,
-                        ]
+                        )
                     k_cache_dtype = v_cache_dtype = current_kv_cache_spec.dtype
                     if self.is_kv_consumer and self.vllm_config.quant_config is not None:
                         k_cache_dtype, v_cache_dtype = self.vllm_config.quant_config.get_kv_quant_dtype(

From bf4aaa5111b9214fed4d0f4f6e0e0d8a23eccff1 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Tue, 17 Mar 2026 11:27:01 +0800
Subject: [PATCH 06/17] Fix for old version.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/spec_decode/eagle_proposer.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/vllm_ascend/spec_decode/eagle_proposer.py b/vllm_ascend/spec_decode/eagle_proposer.py
index 51d80585cb3..69c645dd1c0 100644
--- a/vllm_ascend/spec_decode/eagle_proposer.py
+++ b/vllm_ascend/spec_decode/eagle_proposer.py
@@ -390,7 +390,10 @@ def dummy_run(
                     : num_reqs * self.decode_threshold
                 ]
 
-            builder = self.draft_attn_groups[0].get_metadata_builder()
+            if hasattr(self, "draft_attn_groups") and len(self.draft_attn_groups) > 0:
+                builder = self.draft_attn_groups[0].get_metadata_builder()
+            else:
+                builder = self.runner.attn_groups[0][0].get_metadata_builder()
             # update the tensor's address for each step.
             for draft_step in range(self.num_speculative_tokens):
                 common_attn_metadata = self.shallow_copy_metadata(common_attn_metadata)
@@ -602,7 +605,10 @@ def _propose(
         common_attn_metadata.slot_mapping = self.slot_mapping_group[0]
         common_attn_metadata.num_input_tokens = num_input_tokens
         # FIXME(woosuk): The below two ops cause synchronization. Optimize.
-        builder = self.draft_attn_groups[0].get_metadata_builder()
+        if hasattr(self, "draft_attn_groups") and len(self.draft_attn_groups) > 0:
+            builder = self.draft_attn_groups[0].get_metadata_builder()
+        else:
+            builder = self.runner.attn_groups[0][0].get_metadata_builder()
         attn_metadata = builder.build(0, common_attn_metadata, self.runner.get_model())
 
         if self.uses_mrope:

From 849a16558de56e07f8f3e209553fdac7b3fc2e9a Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Tue, 17 Mar 2026 16:46:02 +0800
Subject: [PATCH 07/17] Fix review comments.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/spec_decode/eagle_proposer.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/vllm_ascend/spec_decode/eagle_proposer.py b/vllm_ascend/spec_decode/eagle_proposer.py
index 69c645dd1c0..d5deafd9ea3 100644
--- a/vllm_ascend/spec_decode/eagle_proposer.py
+++ b/vllm_ascend/spec_decode/eagle_proposer.py
@@ -390,7 +390,8 @@ def dummy_run(
                     : num_reqs * self.decode_threshold
                 ]
 
-            if hasattr(self, "draft_attn_groups") and len(self.draft_attn_groups) > 0:
+            if vllm_version_is("0.17.0"):
+                assert len(self.draft_attn_groups) > 0
                 builder = self.draft_attn_groups[0].get_metadata_builder()
             else:
                 builder = self.runner.attn_groups[0][0].get_metadata_builder()
@@ -605,7 +606,8 @@ def _propose(
         common_attn_metadata.slot_mapping = self.slot_mapping_group[0]
         common_attn_metadata.num_input_tokens = num_input_tokens
         # FIXME(woosuk): The below two ops cause synchronization. Optimize.
-        if hasattr(self, "draft_attn_groups") and len(self.draft_attn_groups) > 0:
+        if vllm_version_is("0.17.0"): 
+            assert len(self.draft_attn_groups) > 0
             builder = self.draft_attn_groups[0].get_metadata_builder()
         else:
             builder = self.runner.attn_groups[0][0].get_metadata_builder()

From 325decb8e447a00e301b747c3c8edd4050bea7bd Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Tue, 17 Mar 2026 17:12:51 +0800
Subject: [PATCH 08/17] Fix for lint.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/spec_decode/eagle_proposer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm_ascend/spec_decode/eagle_proposer.py b/vllm_ascend/spec_decode/eagle_proposer.py
index d5deafd9ea3..8a3625acdd6 100644
--- a/vllm_ascend/spec_decode/eagle_proposer.py
+++ b/vllm_ascend/spec_decode/eagle_proposer.py
@@ -606,7 +606,7 @@ def _propose(
         common_attn_metadata.slot_mapping = self.slot_mapping_group[0]
         common_attn_metadata.num_input_tokens = num_input_tokens
         # FIXME(woosuk): The below two ops cause synchronization. Optimize.
-        if vllm_version_is("0.17.0"): 
+        if vllm_version_is("0.17.0"):
             assert len(self.draft_attn_groups) > 0
             builder = self.draft_attn_groups[0].get_metadata_builder()
         else:

From 36716d65f80a8039d5a18f89b861fac18b67ae53 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Tue, 17 Mar 2026 20:44:19 +0800
Subject: [PATCH 09/17] Add e2e test for mla target and gqa draft.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .../singlecard/spec_decode/test_mtp_eagle_correctness.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
index ef89d328947..671d7809796 100644
--- a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
+++ b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
@@ -38,9 +38,14 @@
     "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
     "RedHatAI/Qwen3-8B-speculator.eagle3"
 ]
-MODELS_MAIN = ["LLM-Research/Meta-Llama-3.1-8B-Instruct", "Qwen/Qwen3-8B"]
+MODELS_MAIN = [
+    "LLM-Research/Meta-Llama-3.1-8B-Instruct",
+    "Qwen/Qwen3-8B",
+    "wemaster/deepseek_mtp_main_random_bf16"]
 VALID_COMBINATIONS = {("eagle", "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
                        "LLM-Research/Meta-Llama-3.1-8B-Instruct"),
+                      ("eagle", "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
+                       "wemaster/deepseek_mtp_main_random_bf16"),
                       ("eagle3", "RedHatAI/Qwen3-8B-speculator.eagle3",
                        "Qwen/Qwen3-8B")}
 
@@ -127,7 +132,7 @@ def test_deepseek_mtp_correctness(model_name: str, num_speculative_tokens: int,
 @pytest.mark.parametrize("disable_padded_drafter_batch", [True, False])
 @pytest.mark.parametrize("async_scheduling", [True, False])
 @pytest.mark.parametrize("draft_tensor_parallel_size", [None, 1])
-def test_llama_qwen3_eagle_correctness(
+def test_llama_qwen3_deepseek_eagle_correctness(
         model_name: str, model_name_main: str, num_speculative_tokens: int,
         method: str, disable_padded_drafter_batch: bool,
         async_scheduling: bool, draft_tensor_parallel_size: Union[None, int]):

From 5d05577457fbe003f5356ff0f857ad3347e92e67 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Wed, 18 Mar 2026 01:28:11 +0800
Subject: [PATCH 10/17] Fix for CI.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .../singlecard/spec_decode/test_mtp_eagle_correctness.py  | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
index 671d7809796..1010dd9b891 100644
--- a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
+++ b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
@@ -48,6 +48,7 @@
                        "wemaster/deepseek_mtp_main_random_bf16"),
                       ("eagle3", "RedHatAI/Qwen3-8B-speculator.eagle3",
                        "Qwen/Qwen3-8B")}
+DEEPSEEK_MAIN = "wemaster/deepseek_mtp_main_random_bf16"
 
 
 @pytest.mark.parametrize("model_name", MODELS)
@@ -150,6 +151,13 @@ def test_llama_qwen3_deepseek_eagle_correctness(
             f"Invalid combination: method={method}, model_name={model_name}, model_name_main={model_name_main}, or case not support yet"
         )
 
+    # Keep the new deepseek-main + eagle coverage, but limit matrix fan-out
+    # to reduce long-running CI instability on singlecard-full part 0.
+    if model_name_main == DEEPSEEK_MAIN and (
+        num_speculative_tokens != 1 or async_scheduling or draft_tensor_parallel_size is not None
+    ):
+        pytest.skip("Limit deepseek-eagle CI matrix to a stable core subset.")
+
     sampling_params = SamplingParams(
         max_tokens=300,
         temperature=0.0,

From 15771f91e650933cf80a76af943fa703c0186220 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Wed, 18 Mar 2026 09:26:48 +0800
Subject: [PATCH 11/17] Fix for CI.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .../spec_decode/test_mtp_eagle_correctness.py  | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
index 1010dd9b891..771fa76af50 100644
--- a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
+++ b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
@@ -30,6 +30,7 @@
 from vllm.config import CompilationConfig
 
 from tests.e2e.conftest import VllmRunner, cleanup_dist_env_and_memory
+from vllm_ascend.utils import vllm_version_is
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -48,7 +49,11 @@
                        "wemaster/deepseek_mtp_main_random_bf16"),
                       ("eagle3", "RedHatAI/Qwen3-8B-speculator.eagle3",
                        "Qwen/Qwen3-8B")}
-DEEPSEEK_MAIN = "wemaster/deepseek_mtp_main_random_bf16"
+VERSION_GATED_COMBINATION = (
+    "eagle",
+    "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
+    "wemaster/deepseek_mtp_main_random_bf16",
+)
 
 
 @pytest.mark.parametrize("model_name", MODELS)
@@ -151,12 +156,13 @@ def test_llama_qwen3_deepseek_eagle_correctness(
             f"Invalid combination: method={method}, model_name={model_name}, model_name_main={model_name_main}, or case not support yet"
         )
 
-    # Keep the new deepseek-main + eagle coverage, but limit matrix fan-out
-    # to reduce long-running CI instability on singlecard-full part 0.
-    if model_name_main == DEEPSEEK_MAIN and (
-        num_speculative_tokens != 1 or async_scheduling or draft_tensor_parallel_size is not None
+    if (
+        (method, model_name, model_name_main) == VERSION_GATED_COMBINATION
+        and not vllm_version_is("0.17.0")
     ):
-        pytest.skip("Limit deepseek-eagle CI matrix to a stable core subset.")
+        pytest.skip(
+            "Skip deepseek-main + eagle coverage for non-v0.17.0 vLLM versions."
+        )
 
     sampling_params = SamplingParams(
         max_tokens=300,

From 255e5d556c8694512cb3e01002b2ca6393ae483c Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Wed, 18 Mar 2026 10:51:25 +0800
Subject: [PATCH 12/17] Fix lint.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 vllm_ascend/spec_decode/eagle_proposer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm_ascend/spec_decode/eagle_proposer.py b/vllm_ascend/spec_decode/eagle_proposer.py
index 665afaa0624..869cd28735e 100644
--- a/vllm_ascend/spec_decode/eagle_proposer.py
+++ b/vllm_ascend/spec_decode/eagle_proposer.py
@@ -46,7 +46,7 @@
 from vllm_ascend.compilation.acl_graph import ACLGraphWrapper, update_full_graph_params
 from vllm_ascend.ops.triton.spec_decode.utils import prepare_inputs_padded_kernel
 from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
-from vllm_ascend.utils import enable_sp, lmhead_tp_enable, shared_expert_dp_enabled
+from vllm_ascend.utils import enable_sp, lmhead_tp_enable, shared_expert_dp_enabled, vllm_version_is
 
 # Currently we will fix block size to a small one since `num_reqs` can't be too large
 _PREPARE_INPUTS_BLOCK_SIZE = 4

From 40a6b2cb7bd6ada5fe0ca3f101088fd8977a2ab4 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Wed, 18 Mar 2026 14:29:41 +0800
Subject: [PATCH 13/17] Fix CI.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .github/workflows/pr_test_full.yaml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/.github/workflows/pr_test_full.yaml b/.github/workflows/pr_test_full.yaml
index b7f35825cef..9fad0dd2e5d 100644
--- a/.github/workflows/pr_test_full.yaml
+++ b/.github/workflows/pr_test_full.yaml
@@ -48,7 +48,6 @@ jobs:
       - name: Setup git proxy
         run: |
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
-          git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf https://github.com/
       # NOTE: Do not update the version of checkout, there have some issue on self_hosted runner with the higher version
       - uses: actions/checkout@v6
       - uses: dorny/paths-filter@v3

From 62bf55882356befc5970f7d84267eb2c399a7d7c Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Thu, 19 Mar 2026 11:37:26 +0800
Subject: [PATCH 14/17] Fix CI.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .../spec_decode/test_mtp_eagle_correctness.py       | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
index 771fa76af50..c8723d4f938 100644
--- a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
+++ b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
@@ -54,6 +54,11 @@
     "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
     "wemaster/deepseek_mtp_main_random_bf16",
 )
+EAGER_VALID_COMBINATION = (
+    "eagle",
+    "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
+    "wemaster/deepseek_mtp_main_random_bf16",
+)
 
 
 @pytest.mark.parametrize("model_name", MODELS)
@@ -164,6 +169,12 @@ def test_llama_qwen3_deepseek_eagle_correctness(
             "Skip deepseek-main + eagle coverage for non-v0.17.0 vLLM versions."
         )
 
+    # TODO(slippersss): After supporting graph mode, remove the enforce_eager restriction.
+    if (method, model_name, model_name_main) == EAGER_VALID_COMBINATION:
+        enforce_eager = True
+    else:
+        enforce_eager = False
+
     sampling_params = SamplingParams(
         max_tokens=300,
         temperature=0.0,
@@ -188,6 +199,7 @@ def test_llama_qwen3_deepseek_eagle_correctness(
                         draft_tensor_parallel_size,
                         "max_model_len": 128,
                     },
+                    enforce_eager=enforce_eager,
                     compilation_config=CompilationConfig(
                         cudagraph_mode="FULL",
                         cudagraph_capture_sizes=[5, 12])) as llm:
@@ -203,6 +215,7 @@ def test_llama_qwen3_deepseek_eagle_correctness(
                     max_model_len=4096,
                     seed=1024,
                     async_scheduling=async_scheduling,
+                    enforce_eager=enforce_eager,
                     compilation_config=CompilationConfig(
                         cudagraph_mode="FULL_DECODE_ONLY",
                         cudagraph_capture_sizes=[12])) as llm:

From 9f1e4059c2e5bc624c291a646fa890b09dd6e3d3 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Thu, 19 Mar 2026 15:13:22 +0800
Subject: [PATCH 15/17] Remove inappropriate e2e test cases, replace with
 module unit tests.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .../spec_decode/test_mtp_eagle_correctness.py | 36 +-------
 tests/ut/worker/test_model_runner_v1.py       | 87 +++++++++++++++++++
 2 files changed, 89 insertions(+), 34 deletions(-)
 create mode 100644 tests/ut/worker/test_model_runner_v1.py

diff --git a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
index c8723d4f938..ef89d328947 100644
--- a/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
+++ b/tests/e2e/singlecard/spec_decode/test_mtp_eagle_correctness.py
@@ -30,7 +30,6 @@
 from vllm.config import CompilationConfig
 
 from tests.e2e.conftest import VllmRunner, cleanup_dist_env_and_memory
-from vllm_ascend.utils import vllm_version_is
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -39,26 +38,11 @@
     "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
     "RedHatAI/Qwen3-8B-speculator.eagle3"
 ]
-MODELS_MAIN = [
-    "LLM-Research/Meta-Llama-3.1-8B-Instruct",
-    "Qwen/Qwen3-8B",
-    "wemaster/deepseek_mtp_main_random_bf16"]
+MODELS_MAIN = ["LLM-Research/Meta-Llama-3.1-8B-Instruct", "Qwen/Qwen3-8B"]
 VALID_COMBINATIONS = {("eagle", "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
                        "LLM-Research/Meta-Llama-3.1-8B-Instruct"),
-                      ("eagle", "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
-                       "wemaster/deepseek_mtp_main_random_bf16"),
                       ("eagle3", "RedHatAI/Qwen3-8B-speculator.eagle3",
                        "Qwen/Qwen3-8B")}
-VERSION_GATED_COMBINATION = (
-    "eagle",
-    "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
-    "wemaster/deepseek_mtp_main_random_bf16",
-)
-EAGER_VALID_COMBINATION = (
-    "eagle",
-    "vllm-ascend/EAGLE-LLaMA3.1-Instruct-8B",
-    "wemaster/deepseek_mtp_main_random_bf16",
-)
 
 
 @pytest.mark.parametrize("model_name", MODELS)
@@ -143,7 +127,7 @@ def test_deepseek_mtp_correctness(model_name: str, num_speculative_tokens: int,
 @pytest.mark.parametrize("disable_padded_drafter_batch", [True, False])
 @pytest.mark.parametrize("async_scheduling", [True, False])
 @pytest.mark.parametrize("draft_tensor_parallel_size", [None, 1])
-def test_llama_qwen3_deepseek_eagle_correctness(
+def test_llama_qwen3_eagle_correctness(
         model_name: str, model_name_main: str, num_speculative_tokens: int,
         method: str, disable_padded_drafter_batch: bool,
         async_scheduling: bool, draft_tensor_parallel_size: Union[None, int]):
@@ -161,20 +145,6 @@ def test_llama_qwen3_deepseek_eagle_correctness(
             f"Invalid combination: method={method}, model_name={model_name}, model_name_main={model_name_main}, or case not support yet"
         )
 
-    if (
-        (method, model_name, model_name_main) == VERSION_GATED_COMBINATION
-        and not vllm_version_is("0.17.0")
-    ):
-        pytest.skip(
-            "Skip deepseek-main + eagle coverage for non-v0.17.0 vLLM versions."
-        )
-
-    # TODO(slippersss): After supporting graph mode, remove the enforce_eager restriction.
-    if (method, model_name, model_name_main) == EAGER_VALID_COMBINATION:
-        enforce_eager = True
-    else:
-        enforce_eager = False
-
     sampling_params = SamplingParams(
         max_tokens=300,
         temperature=0.0,
@@ -199,7 +169,6 @@ def test_llama_qwen3_deepseek_eagle_correctness(
                         draft_tensor_parallel_size,
                         "max_model_len": 128,
                     },
-                    enforce_eager=enforce_eager,
                     compilation_config=CompilationConfig(
                         cudagraph_mode="FULL",
                         cudagraph_capture_sizes=[5, 12])) as llm:
@@ -215,7 +184,6 @@ def test_llama_qwen3_deepseek_eagle_correctness(
                     max_model_len=4096,
                     seed=1024,
                     async_scheduling=async_scheduling,
-                    enforce_eager=enforce_eager,
                     compilation_config=CompilationConfig(
                         cudagraph_mode="FULL_DECODE_ONLY",
                         cudagraph_capture_sizes=[12])) as llm:
diff --git a/tests/ut/worker/test_model_runner_v1.py b/tests/ut/worker/test_model_runner_v1.py
new file mode 100644
index 00000000000..621103c7a53
--- /dev/null
+++ b/tests/ut/worker/test_model_runner_v1.py
@@ -0,0 +1,87 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
+
+from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
+
+
+class TestNPUModelRunnerKVCache(unittest.TestCase):
+
+    def _build_runner(self):
+        runner = NPUModelRunner.__new__(NPUModelRunner)
+        runner.device = torch.device("cpu")
+        runner.use_sparse = False
+        runner.use_hybrid_blocks = False
+        runner.hybrid_with_attn_and_mamba = False
+        runner.runner_only_attn_layers = set()
+        runner.vllm_config = MagicMock()
+        runner.vllm_config.kv_transfer_config = None
+        runner.model_config = MagicMock()
+        runner.model_config.use_mla = True
+        backend = MagicMock()
+        backend.get_kv_cache_shape.side_effect = lambda num_blocks, block_size, num_kv_heads, head_size: (
+            2,
+            num_blocks,
+            block_size,
+            num_kv_heads,
+            head_size,
+        )
+        runner.attn_backend = backend
+        return runner
+
+    def test_allocate_kv_cache_uses_layer_spec_for_draft_gqa(self):
+        runner = self._build_runner()
+        kv_cache_spec = FullAttentionSpec(
+            block_size=16,
+            num_kv_heads=8,
+            head_size=64,
+            head_size_v=64,
+            dtype=torch.float16,
+        )
+        kv_cache_config = KVCacheConfig(
+            num_blocks=2,
+            kv_cache_tensors=[KVCacheTensor(size=kv_cache_spec.page_size_bytes * 2, shared_by=["draft_attn"])],
+            kv_cache_groups=[KVCacheGroupSpec(layer_names=["draft_attn"], kv_cache_spec=kv_cache_spec)],
+        )
+
+        kv_cache_raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config)
+        k_cache_raw, v_cache_raw = kv_cache_raw_tensors["draft_attn"]
+
+        self.assertEqual(k_cache_raw.numel(), kv_cache_spec.page_size_bytes)
+        self.assertEqual(v_cache_raw.numel(), kv_cache_spec.page_size_bytes)
+
+    def test_reshape_kv_cache_uses_layer_spec_for_draft_gqa(self):
+        runner = self._build_runner()
+        kv_cache_spec = FullAttentionSpec(
+            block_size=16,
+            num_kv_heads=8,
+            head_size=64,
+            head_size_v=64,
+            dtype=torch.float16,
+        )
+        kv_cache_config = KVCacheConfig(
+            num_blocks=2,
+            kv_cache_tensors=[KVCacheTensor(size=kv_cache_spec.page_size_bytes * 2, shared_by=["draft_attn"])],
+            kv_cache_groups=[KVCacheGroupSpec(layer_names=["draft_attn"], kv_cache_spec=kv_cache_spec)],
+        )
+        kv_cache_raw_tensors = runner._allocate_kv_cache_tensors(kv_cache_config)
+        runner._kv_cache_spec_attn_group_iterator = lambda: [
+            SimpleNamespace(
+                kv_cache_spec=kv_cache_spec,
+                backend=runner.attn_backend,
+                layer_names=["draft_attn"],
+            )
+        ]
+
+        kv_caches = runner._reshape_kv_cache_tensors(kv_cache_config, kv_cache_raw_tensors)
+        k_cache, v_cache = kv_caches["draft_attn"]
+
+        self.assertEqual(k_cache.shape, (2, 16, 8, 64))
+        self.assertEqual(v_cache.shape, (2, 16, 8, 64))
+
+
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file

From 5dda25582873292e041e2cc461edd33faa82c2ff Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Thu, 19 Mar 2026 15:36:00 +0800
Subject: [PATCH 16/17] Fix for ut.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 tests/ut/worker/test_model_runner_v1.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/ut/worker/test_model_runner_v1.py b/tests/ut/worker/test_model_runner_v1.py
index 621103c7a53..0cb09fafdd0 100644
--- a/tests/ut/worker/test_model_runner_v1.py
+++ b/tests/ut/worker/test_model_runner_v1.py
@@ -14,9 +14,11 @@ def _build_runner(self):
         runner = NPUModelRunner.__new__(NPUModelRunner)
         runner.device = torch.device("cpu")
         runner.use_sparse = False
+        runner.use_sparse_c8_indexer = False
         runner.use_hybrid_blocks = False
         runner.hybrid_with_attn_and_mamba = False
         runner.runner_only_attn_layers = set()
+        runner.is_kv_consumer = False
         runner.vllm_config = MagicMock()
         runner.vllm_config.kv_transfer_config = None
         runner.model_config = MagicMock()

From cef2642f73b996474a313e2a0adcc61999725b27 Mon Sep 17 00:00:00 2001
From: SidaoY <1024863041@qq.com>
Date: Fri, 20 Mar 2026 09:36:46 +0800
Subject: [PATCH 17/17] Fix.

Signed-off-by: SidaoY <1024863041@qq.com>
---
 .github/workflows/pr_test_full.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/pr_test_full.yaml b/.github/workflows/pr_test_full.yaml
index 5de636be51c..ca15a121d8a 100644
--- a/.github/workflows/pr_test_full.yaml
+++ b/.github/workflows/pr_test_full.yaml
@@ -48,6 +48,7 @@ jobs:
       - name: Setup git proxy
         run: |
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
+          git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf https://github.com/
       # NOTE: Do not update the version of checkout, there have some issue on self_hosted runner with the higher version
       - uses: actions/checkout@v6
       - uses: dorny/paths-filter@v3