init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

tests/ut/_310p/attention/test_attention_mask_310.py

@@ -0,0 +1,45 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310
+
+
+class TestAttentionMaskBuilder310(TestBase):
+    def setUp(self):
+        self.max_seqlen = 4096
+        self.attention_mask_builder = AttentionMaskBuilder310(torch.device("cpu"), self.max_seqlen)
+
+    @patch("torch_npu.npu_format_cast")
+    def test_get_attention_mask_310(self, mock_format_cast):
+        mock_format_cast.side_effect = lambda x, y: x
+        self.attention_mask_builder.support_compressed_mask = False
+        model_config = MagicMock()
+        attn_mask = self.attention_mask_builder.get_attention_mask(causal=True, model_config=model_config)
+        self.assertEqual(attn_mask.shape, (1, self.max_seqlen // 16, self.max_seqlen, 16))
+        self.assertEqual(attn_mask[0][-1][0][-1], torch.tensor(float("-inf"), dtype=torch.float16))
+
+    @patch("torch_npu.npu_format_cast")
+    def test_get_splitfuse_attn_mask_310(self, mock_format_cast):
+        mock_format_cast.side_effect = lambda x, y: x
+        attn_metadata = MagicMock()
+        attn_metadata.query_start_loc = torch.tensor([0, 1, 5])
+        attn_metadata.seq_lens = torch.tensor([7, 4])
+        attn_mask = self.attention_mask_builder.get_splitfuse_mask(attn_metadata, torch.device("cpu"))
+        self.assertEqual(attn_mask.shape, (1, self.max_seqlen // 16, 16, 16))

tests/ut/_310p/attention/test_attention_v1_310.py

@@ -0,0 +1,292 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.attention.attention_v1 import (
+    AscendAttentionBackend310,
+    AscendAttentionBackendImpl310,
+    AscendAttentionMetadataBuilder310,
+    AscendAttentionState,
+)
+from vllm_ascend._310p.attention.metadata_builder import (
+    AscendAttentionMetadataBuilder310 as AscendMetadataBuilder310Direct,
+)
+
+
+class TestAscendAttentionBackend310(TestBase):
+    def setUp(self):
+        self.mock_config = MagicMock()
+        self.utils_patcher = patch("vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_config)
+        self.utils_patcher.start()
+
+    def test_get_impl_cls(self):
+        self.assertEqual(AscendAttentionBackend310.get_impl_cls(), AscendAttentionBackendImpl310)
+
+    def test_get_builder_cls(self):
+        self.assertEqual(AscendAttentionBackend310.get_builder_cls(), AscendAttentionMetadataBuilder310)
+
+    def test_get_kv_cache_shape_not(self):
+        result = AscendAttentionBackend310.get_kv_cache_shape(10, 20, 30, 40)
+        self.assertEqual(result, (2, 10, 75, 20, 16))
+
+
+class TestAscendAttentionBackendImpl310(TestBase):
+    def setUp(self):
+        self.attention_type = MagicMock()
+        self.attention_type.DECODER = "decoder"
+        self.attention_type.ENCODER = "encoder"
+        self.attn_metadata = MagicMock()
+        self.attn_metadata.return_value = "1"
+        self.mock_vllm_config = MagicMock()
+        self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
+        self.layer_no_quant.layer_name = "test_layer"
+        self.layer_no_quant._k_scale_float = 1.0
+        self.layer_no_quant._v_scale_float = 1.0
+        self.config_patcher = patch(
+            "vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.config_patcher.start()
+        self.impl = AscendAttentionBackendImpl310(
+            num_heads=8,
+            head_size=128,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("torch_npu._npu_flash_attention")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_forward_prefill_310(self, mock_get_forward_context, mock_npu_flash_attention, mock_npu_reshape_and_cache):
+        """Test forward pass in PrefillNoCache state."""
+        query = torch.randn(10, 8, 64)
+        key = torch.randn(10, 8, 64)
+        value = torch.randn(10, 8, 64)
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.PrefillNoCache
+        metadata.attn_mask = torch.randn(1, 1, 10, 10)
+        metadata.query_lens = torch.tensor([10])
+        metadata.seq_lens = torch.tensor([10])
+        metadata.actual_seq_lengths_q = [10]
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 10
+        metadata.num_decode_tokens = 0
+        metadata.num_decodes = 0
+        metadata.num_prefills = 10
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+
+        self.impl.support_compressed_mask = False
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+        mock_npu_flash_attention.return_value = torch.ones(10, 8, 64)
+        result = self.impl.forward_impl(query, key, value, None, metadata, output)
+
+        mock_npu_flash_attention.assert_called_once()
+        _, kwargs = mock_npu_flash_attention.call_args
+        self.assertIs(kwargs["query"], query)
+        self.assertIs(kwargs["key"], key)
+        self.assertIs(kwargs["value"], value)
+        self.assertIs(kwargs["mask"], metadata.attn_mask)
+        self.assertIs(kwargs["seq_len"], metadata.seq_lens)
+        self.assertEqual(kwargs["scale_value"], self.impl.scale)
+        self.assertEqual(kwargs["num_heads"], self.impl.num_heads)
+        self.assertEqual(kwargs["num_kv_heads"], self.impl.num_kv_heads)
+        self.assertIs(kwargs["out"], output)
+        self.assertIs(result, output)
+
+    @patch("torch_npu.npu_format_cast", return_value=torch.randn((1, 128, 16, 16), dtype=torch.float16))
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("torch_npu._npu_paged_attention_splitfuse")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_forward_chunked_prefill_310(
+        self,
+        mock_get_forward_context,
+        mock_npu_paged_attention_splitfuse,
+        mock_npu_reshape_and_cache,
+        mock_format_cast,
+    ):
+        """Test forward pass in ChunkedPrefill state"""
+        query = torch.randn(5, 8, 64)
+        key, value = None, None
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.ChunkedPrefill
+        metadata.attn_mask = torch.randn(1, 128, 16, 16)
+        metadata.query_lens = torch.tensor([5])
+        metadata.seq_lens = torch.tensor([1, 4])
+        metadata.query_start_loc = torch.tensor([0, 1, 5])
+        metadata.actual_seq_lengths_q = [5]
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 10
+        metadata.num_decode_tokens = 0
+        metadata.num_decodes = 0
+        metadata.num_prefills = 10
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+
+        self.impl.support_compressed_mask = False
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+        mock_npu_paged_attention_splitfuse.return_value = torch.ones(5, 8, 64)
+        output = self.impl.forward_impl(query, key, value, None, metadata, output)
+
+        mock_npu_paged_attention_splitfuse.assert_called_once()
+
+    @patch("torch_npu.npu_format_cast", return_value=torch.randn((1, 128, 16, 16), dtype=torch.float16))
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("torch_npu._npu_paged_attention_splitfuse")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_forward_prefill_cache_hit_310(
+        self,
+        mock_get_forward_context,
+        mock_npu_paged_attention_splitfuse,
+        mock_npu_reshape_and_cache,
+        mock_format_cast,
+    ):
+        """Test forward pass in PrefillCacheHit state"""
+        query = torch.randn(5, 8, 64)
+        key, value = None, None
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.PrefillCacheHit
+        metadata.attn_mask = torch.randn(1, 128, 16, 16)
+        metadata.query_lens = torch.tensor([5])
+        metadata.seq_lens = torch.tensor([1, 4])
+        metadata.query_start_loc = torch.tensor([0, 1, 5])
+        metadata.actual_seq_lengths_q = [5]
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 10
+        metadata.num_decode_tokens = 0
+        metadata.num_decodes = 0
+        metadata.num_prefills = 10
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+
+        self.impl.support_compressed_mask = False
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+        mock_npu_paged_attention_splitfuse.return_value = torch.ones(5, 8, 64)
+        output = self.impl.forward_impl(query, key, value, None, metadata, output)
+
+        mock_npu_paged_attention_splitfuse.assert_called_once()
+
+    @patch("vllm_ascend.attention.attention_v1.using_paged_attention")
+    @patch("torch_npu._npu_paged_attention", create=True)
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_forward_paged_attention_310(
+        self, mock_get_forward_context, mock_npu_reshape_and_cache, mock_paged_attention, mock_using_paged_attention
+    ):
+        """Test forward pass in DecodeOnly state"""
+        query = torch.randn(4, 8 * 64)
+        key, value = None, None
+        output = torch.empty_like(query)
+
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.seq_lens = torch.tensor([4])
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 4
+        metadata.slot_mapping = torch.zeros(4, dtype=torch.long)
+        metadata.num_decodes = 4
+        metadata.num_prefills = 0
+        mock_using_paged_attention.return_value = True
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        output = self.impl.forward_impl(query, key, value, None, metadata, output)
+
+        mock_paged_attention.assert_called_once()
+
+    @patch("vllm_ascend._310p.attention.attention_v1.AscendAttentionBackendImpl310.forward_chunked_prefill_310")
+    def test_forward_mtp_310(self, mock_chunked_prefill):
+        query = torch.randn(4, 8 * 64)
+        key, value = None, None
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.SpecDecoding
+        mock_chunked_prefill.return_value = output
+
+        result = self.impl.forward_impl(query, key, value, None, metadata, output)
+
+        mock_chunked_prefill.assert_called_once_with(query, metadata, output)
+        self.assertIs(result, output)
+
+
+class TestAscendAttentionMetadataBuilder310(TestBase):
+    def test_fill_query_lens_cpu_without_buffer(self):
+        builder = AscendMetadataBuilder310Direct.__new__(AscendMetadataBuilder310Direct)
+        builder._query_lens_cpu_buffer = None
+        query_start_loc_cpu = torch.tensor([0, 1, 5, 11, 20], dtype=torch.int32)
+        result = builder._fill_query_lens_cpu(num_reqs=3, query_start_loc_cpu=query_start_loc_cpu, is_drafting=False)
+        expected = torch.tensor([1, 4, 6], dtype=torch.int32)
+        torch.testing.assert_close(result, expected)
+
+    def test_fill_query_lens_cpu_with_buffer_not_drafting(self):
+        builder = AscendMetadataBuilder310Direct.__new__(AscendMetadataBuilder310Direct)
+        builder._query_lens_cpu_buffer = torch.zeros(10, dtype=torch.int32, device="cpu")
+        query_start_loc_cpu = torch.tensor([0, 1, 5, 11, 20], dtype=torch.int32)
+        result = builder._fill_query_lens_cpu(num_reqs=3, query_start_loc_cpu=query_start_loc_cpu, is_drafting=False)
+        expected = torch.tensor([1, 4, 6], dtype=torch.int32)
+        torch.testing.assert_close(result, expected)
+        assert result.data_ptr() == builder._query_lens_cpu_buffer[:3].data_ptr()
+
+    def test_fill_query_lens_cpu_with_buffer_is_drafting(self):
+        builder = AscendMetadataBuilder310Direct.__new__(AscendMetadataBuilder310Direct)
+        builder._query_lens_cpu_buffer = torch.zeros(10, dtype=torch.int32, device="cpu")
+        query_start_loc_cpu = torch.tensor([0, 1, 5, 11, 20], dtype=torch.int32)
+        result1 = builder._fill_query_lens_cpu(num_reqs=3, query_start_loc_cpu=query_start_loc_cpu, is_drafting=True)
+        result2 = builder._fill_query_lens_cpu(num_reqs=3, query_start_loc_cpu=query_start_loc_cpu, is_drafting=True)
+        expected = torch.tensor([1, 4, 6], dtype=torch.int32)
+        torch.testing.assert_close(result1, expected)
+        torch.testing.assert_close(result2, expected)
+        assert result1.data_ptr() != builder._query_lens_cpu_buffer[:3].data_ptr()
+        assert result2.data_ptr() != builder._query_lens_cpu_buffer[:3].data_ptr()
+
+    def test_build_for_drafting_calls_build_with_is_drafting_true(self):
+        builder = object.__new__(AscendMetadataBuilder310Direct)
+        builder._query_lens_cpu_buffer = torch.zeros(10, dtype=torch.int32, device="cpu")
+        builder.device = torch.device("cpu")
+        from vllm.v1.kv_cache_interface import AttentionSpec
+
+        from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310
+
+        builder.attn_mask_builder = AttentionMaskBuilder310(torch.device("cpu"), 4096)
+        builder.kv_cache_spec = AttentionSpec(
+            block_size=128,
+            num_kv_heads=2,
+            head_size=64,
+            dtype=torch.float16,
+        )
+        builder.layer_names = []
+        builder.vllm_config = MagicMock()
+        builder.vllm_config.model_config.max_model_len = 4096
+        builder.vllm_config.scheduler_config.max_num_seqs = 8
+
+        common_attn_metadata = MagicMock()
+        common_attn_metadata.num_reqs = 2
+        common_attn_metadata.query_start_loc = torch.tensor([0, 1, 3])
+        common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 1, 3])
+        common_attn_metadata.seq_lens = torch.tensor([1, 2])
+
+        with patch.object(AscendMetadataBuilder310Direct.__bases__[0], "build", return_value=MagicMock()) as mock_build:
+            result = builder.build_for_drafting(common_attn_metadata=common_attn_metadata, draft_index=0)
+            mock_build.assert_called_once_with(0, common_attn_metadata, True)
+            assert result is not None

tests/ut/_310p/fused_moe/test_experts_selector_310.py

@@ -0,0 +1,85 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm_ascend._310p.fused_moe.experts_selector import select_experts
+
+
+class TestExpertsSelector310:
+    @pytest.mark.parametrize("global_num_experts", [256, 128])
+    def test_select_experts(self, global_num_experts):
+        hidden_states = torch.randn(8, 16)
+        router_logits = torch.randn(8, 8)
+
+        with patch("torch_npu.npu_moe_gating_top_k_softmax") as mock_npu:
+            mock_npu.return_value = (
+                torch.randn(8, 2),
+                torch.randint(0, 8, (8, 2), dtype=torch.int32),
+                None,
+            )
+
+            topk_weights, topk_ids = select_experts(
+                hidden_states=hidden_states,
+                router_logits=router_logits,
+                top_k=2,
+                use_grouped_topk=False,
+                renormalize=True,
+                topk_group=None,
+                num_expert_group=None,
+                custom_routing_function=None,
+                scoring_func="softmax",
+                e_score_correction_bias=None,
+                global_num_experts=global_num_experts,
+            )
+
+            mock_npu.assert_called_once()
+
+        assert topk_weights.shape == (8, 2)
+        assert topk_ids.shape == (8, 2)
+
+    def test_select_experts_chunks_large_token_batch(self):
+        num_tokens = 2050
+        hidden_states = torch.randn(num_tokens, 16)
+        router_logits = torch.randn(num_tokens, 8)
+
+        def mock_gating(logits, k):
+            return (
+                torch.ones(logits.shape[0], k),
+                torch.zeros(logits.shape[0], k, dtype=torch.int32),
+                None,
+            )
+
+        with patch(
+            "torch_npu.npu_moe_gating_top_k_softmax",
+            side_effect=mock_gating,
+        ) as mock_npu:
+            topk_weights, topk_ids = select_experts(
+                hidden_states=hidden_states,
+                router_logits=router_logits,
+                top_k=2,
+                use_grouped_topk=False,
+                renormalize=True,
+                custom_routing_function=None,
+                scoring_func="softmax",
+            )
+
+        assert [call.args[0].shape[0] for call in mock_npu.call_args_list] == [1024, 1024, 2]
+        assert topk_weights.shape == (num_tokens, 2)
+        assert topk_ids.shape == (num_tokens, 2)
+        assert torch.all(topk_weights == 0.5)

tests/ut/_310p/fused_moe/test_moe_mlp_310.py

@@ -0,0 +1,179 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, call, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.fused_moe.moe_comm_method import AllGatherCommImpl310
+from vllm_ascend._310p.fused_moe.moe_mlp import unified_apply_mlp
+from vllm_ascend.ops.fused_moe.moe_runtime_args import (
+    MoEMlpComputeInput,
+    MoEQuantParams,
+    MoEWeights,
+)
+from vllm_ascend.quantization.quant_type import QuantType
+
+
+def build_mlp_compute_input_fixture(
+    *,
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    group_list: torch.Tensor,
+    with_quant: bool,
+    w1_scale: torch.Tensor | None = None,
+    w2_scale: torch.Tensor | None = None,
+    group_list_type: int = 1,
+) -> MoEMlpComputeInput:
+    return MoEMlpComputeInput(
+        hidden_states=hidden_states,
+        group_list=group_list,
+        group_list_type=group_list_type,
+        dynamic_scale=None,
+        topk_scales=None,
+        weights=MoEWeights(w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale),
+        quant=MoEQuantParams(quant_type=QuantType.W8A8 if with_quant else QuantType.NONE),
+        fusion=False,
+        activation="silu",
+        need_trans=False,
+        dynamic_eplb=False,
+    )
+
+
+class TestUnifiedApplyMLP310(TestBase):
+    @patch("vllm_ascend._310p.fused_moe.moe_comm_method.unified_apply_mlp")
+    def test_all_gather_apply_mlp_returns_common_tuple_contract(self, mock_unified_apply_mlp):
+        mlp_compute_input = MagicMock(spec=MoEMlpComputeInput)
+        mlp_output = torch.randn(10, 20, dtype=torch.float16)
+        mock_unified_apply_mlp.return_value = mlp_output
+
+        comm_impl = AllGatherCommImpl310.__new__(AllGatherCommImpl310)
+
+        output, before_gmm2_evt = comm_impl._apply_mlp(mlp_compute_input)
+
+        self.assertIs(output, mlp_output)
+        self.assertIsNone(before_gmm2_evt)
+        mock_unified_apply_mlp.assert_called_once_with(mlp_compute_input=mlp_compute_input)
+
+    @patch("torch_npu.npu_grouped_matmul", create=True)
+    @patch("torch_npu.npu_swiglu")
+    def test_unified_apply_mlp_without_quantization_310(self, mock_npu_swiglu, mock_npu_grouped_matmul):
+        mock_gmm1_out = torch.randn(10, 40, dtype=torch.float16)
+        mock_gmm2_out = torch.randn(10, 20, dtype=torch.float16)
+        mock_npu_grouped_matmul.side_effect = [[mock_gmm1_out], [mock_gmm2_out]]
+
+        mock_npu_swiglu_output = torch.randn(10, 40, dtype=torch.float16)
+        mock_npu_swiglu.return_value = mock_npu_swiglu_output
+
+        hidden_states = torch.randn(10, 20, dtype=torch.float16)
+        w1 = torch.randn(5, 20, 40, dtype=torch.float16)
+        w2 = torch.randn(5, 40, 20, dtype=torch.float16)
+        group_list = torch.tensor([2, 4, 6, 8, 10], dtype=torch.int64)
+
+        result = unified_apply_mlp(
+            mlp_compute_input=build_mlp_compute_input_fixture(
+                hidden_states=hidden_states,
+                w1=w1,
+                w2=w2,
+                group_list=group_list,
+                with_quant=False,
+            )
+        )
+
+        self.assertEqual(mock_npu_grouped_matmul.call_count, 2)
+        mock_npu_grouped_matmul.assert_has_calls(
+            [
+                call(
+                    x=[hidden_states], weight=[w1], split_item=2, group_list_type=1, group_type=0, group_list=group_list
+                ),
+                call(
+                    x=[mock_npu_swiglu_output],
+                    weight=[w2],
+                    split_item=2,
+                    group_list_type=1,
+                    group_type=0,
+                    group_list=group_list,
+                ),
+            ],
+            any_order=True,
+        )
+        mock_npu_swiglu.assert_called_once()
+        mock_npu_swiglu.assert_called_with(mock_gmm1_out)
+
+        self.assertEqual(result.shape, hidden_states.shape)
+        self.assertEqual(result.dtype, torch.float16)
+
+    @patch("torch.cumsum")
+    @patch("torch_npu.npu_quant_grouped_matmul_dequant", create=True)
+    @patch("torch_npu.npu_swiglu")
+    def test_unified_apply_mlp_with_quantization_310(
+        self, mock_npu_swiglu, mock_npu_quant_grouped_matmul_dequant, mock_cumsum
+    ):
+        mock_cumsum_out = torch.arange(0, 10, dtype=torch.int64)
+        mock_cumsum.return_value = mock_cumsum_out
+        mock_gmm1_out = torch.randn(10, 40, dtype=torch.float16)
+        mock_gmm2_out = torch.randn(10, 20, dtype=torch.float16)
+        mock_npu_quant_grouped_matmul_dequant.side_effect = [mock_gmm1_out, mock_gmm2_out]
+
+        mock_npu_swiglu_output = torch.randn(10, 40, dtype=torch.float16)
+        mock_npu_swiglu.return_value = mock_npu_swiglu_output
+
+        hidden_states = torch.randn(10, 20, dtype=torch.float16)
+        w1 = torch.randn(5, 20, 40, dtype=torch.float16)
+        w1_scale = torch.rand(5, 40, dtype=torch.float32)
+        w2 = torch.randn(5, 40, 20, dtype=torch.float16)
+        w2_scale = torch.rand(5, 40, dtype=torch.float32)
+        group_list = torch.tensor([2, 4, 6, 8, 10], dtype=torch.int64)
+
+        result = unified_apply_mlp(
+            mlp_compute_input=build_mlp_compute_input_fixture(
+                hidden_states=hidden_states,
+                w1=w1,
+                w2=w2,
+                group_list=group_list,
+                with_quant=True,
+                w1_scale=w1_scale,
+                w2_scale=w2_scale,
+            )
+        )
+
+        mock_cumsum.assert_called_once()
+        self.assertEqual(mock_npu_quant_grouped_matmul_dequant.call_count, 2)
+        mock_npu_quant_grouped_matmul_dequant.assert_has_calls(
+            [
+                call(
+                    x=hidden_states,
+                    quantized_weight=w1,
+                    weight_scale=w1_scale,
+                    group_list=mock_cumsum_out,
+                    quant_mode="pertoken",
+                ),
+                call(
+                    x=mock_npu_swiglu_output,
+                    quantized_weight=w2,
+                    weight_scale=w2_scale,
+                    group_list=mock_cumsum_out,
+                    quant_mode="pertoken",
+                ),
+            ],
+            any_order=True,
+        )
+        mock_npu_swiglu.assert_called_once()
+        mock_npu_swiglu.assert_called_with(mock_gmm1_out)
+
+        self.assertEqual(result.shape, hidden_states.shape)
+        self.assertEqual(result.dtype, torch.float16)

tests/ut/_310p/fused_moe/test_shared_fused_moe_310.py

@@ -0,0 +1,109 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import torch
+import torch.nn.functional as F
+
+from vllm_ascend._310p.fused_moe.fused_moe import (
+    AscendFusedMoE310,
+)
+
+
+class _DummyGate(torch.nn.Module):
+    def forward(self, hidden_states: torch.Tensor):
+        # Keep gate output deterministic: sigmoid(0)=0.5.
+        return torch.zeros(
+            hidden_states.shape[0],
+            1,
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        ), None
+
+
+class _DummySharedExperts(torch.nn.Module):
+    def __init__(self, with_gate: bool):
+        super().__init__()
+        self.expert_gate = _DummyGate() if with_gate else None
+
+    def forward(self, hidden_states: torch.Tensor):
+        out = hidden_states * 2.0 + 1.0
+        if self.expert_gate is not None:
+            gate_out, _ = self.expert_gate(hidden_states)
+            out = F.sigmoid(gate_out) * out
+        return out
+
+
+def _build_layer(shared_experts: torch.nn.Module | None) -> AscendFusedMoE310:
+    layer = AscendFusedMoE310.__new__(AscendFusedMoE310)
+    # The test bypasses full layer init with __new__, so we must initialize
+    # nn.Module internals before assigning child modules.
+    torch.nn.Module.__init__(layer)
+    layer._shared_experts = shared_experts
+    return layer
+
+
+def test_forward_shared_experts_without_gate_310():
+    layer = _build_layer(_DummySharedExperts(with_gate=False))
+    hidden_states = torch.randn(4, 8)
+    output = layer._forward_shared_experts(hidden_states)
+    expected = hidden_states * 2.0 + 1.0
+    torch.testing.assert_close(output, expected)
+
+
+def test_forward_shared_experts_with_gate_310():
+    layer = _build_layer(_DummySharedExperts(with_gate=True))
+    hidden_states = torch.randn(4, 8)
+    output = layer._forward_shared_experts(hidden_states)
+    expected = 0.5 * (hidden_states * 2.0 + 1.0)
+    torch.testing.assert_close(output, expected)
+
+
+def test_forward_impl_with_shared_experts_returns_tuple_310():
+    layer = _build_layer(_DummySharedExperts(with_gate=True))
+    hidden_states = torch.randn(3, 8)
+    router_logits = torch.randn(3, 8)
+    routed_out = torch.randn(3, 8)
+
+    with patch.object(AscendFusedMoE310, "forward_impl", return_value=routed_out):
+        shared_out, routed = layer.shared_forward_impl(hidden_states, router_logits)
+
+    expected_shared = 0.5 * (hidden_states * 2.0 + 1.0)
+    torch.testing.assert_close(shared_out, expected_shared)
+    torch.testing.assert_close(routed, routed_out)
+
+
+def test_forward_impl_without_shared_experts_integration_310():
+    layer = _build_layer(None)
+    hidden_states = torch.randn(3, 8)
+    assert layer._forward_shared_experts(hidden_states) is None
+
+
+def test_forward_impl_without_shared_experts_returns_routed_only_310():
+    layer = _build_layer(None)
+    hidden_states = torch.randn(3, 8)
+    router_logits = torch.randn(3, 8)
+    routed_out = torch.randn(3, 8)
+
+    with patch.object(AscendFusedMoE310, "forward_impl", return_value=routed_out):
+        output = layer.shared_forward_impl(hidden_states, router_logits)
+
+    torch.testing.assert_close(output, routed_out)
+
+
+def test_is_internal_router_is_false_310():
+    layer = _build_layer(_DummySharedExperts(with_gate=True))
+    assert layer.is_internal_router is False

tests/ut/_310p/ops/test_chunk_gated_delta_rule_310.py

@@ -0,0 +1,147 @@
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch_npu
+
+from vllm_ascend._310p.ops.fla.chunk_gated_delta_rule import chunk_gated_delta_rule_pytorch
+
+
+def _cpu_rms_norm(x, weight, eps):
+    """CPU fallback for torch_npu.npu_rms_norm used by l2norm_310p on CPU runners."""
+    orig_dtype = x.dtype
+    x32 = x.float()
+    var = x32.pow(2).mean(-1, keepdim=True)
+    x32 = x32 * torch.rsqrt(var + eps)
+    out = (x32 * weight.float()).to(orig_dtype)
+    return out, None
+
+
+@pytest.fixture(autouse=True)
+def _mock_npu_rms_norm():
+    # conftest stubs npu_rms_norm with a bare MagicMock(); override with a CPU impl.
+    with patch.object(torch_npu, "npu_rms_norm", side_effect=_cpu_rms_norm, create=True):
+        yield
+
+
+def test_chunk_gated_delta_rule_310_output_shape_and_dtype():
+    torch.manual_seed(0)
+
+    bsz = 2
+    total_tokens = 7
+    num_qk_heads = 2
+    num_v_heads = 4
+    kdim = 16
+    vdim = 12
+
+    q = torch.randn(bsz, total_tokens, num_qk_heads, kdim, dtype=torch.float16)
+    k = torch.randn(bsz, total_tokens, num_qk_heads, kdim, dtype=torch.float16)
+    v = torch.randn(bsz, total_tokens, num_v_heads, vdim, dtype=torch.float16)
+    g = -0.2 * torch.rand(bsz, total_tokens, num_v_heads, dtype=torch.float32)
+    beta = (0.15 + 0.35 * torch.rand(bsz, total_tokens, num_v_heads, dtype=torch.float32)).to(torch.float16)
+    initial_state = torch.randn(bsz, num_v_heads, vdim, kdim, dtype=torch.float16)
+
+    out, final_state = chunk_gated_delta_rule_pytorch(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        output_final_state=True,
+        cu_seqlens=None,
+        head_first=False,
+        use_qk_l2norm_in_kernel=True,
+    )
+
+    assert out.shape == v.shape
+    assert out.dtype == v.dtype
+    assert final_state is not None
+    assert final_state.shape == initial_state.shape
+    assert final_state.dtype == torch.float32
+
+
+def test_chunk_gated_delta_rule_310_varlen_path():
+    torch.manual_seed(0)
+
+    bsz = 1
+    total_tokens = 9
+    num_qk_heads = 2
+    num_v_heads = 4
+    kdim = 16
+    vdim = 12
+
+    q = torch.randn(bsz, total_tokens, num_qk_heads, kdim, dtype=torch.float16)
+    k = torch.randn(bsz, total_tokens, num_qk_heads, kdim, dtype=torch.float16)
+    v = torch.randn(bsz, total_tokens, num_v_heads, vdim, dtype=torch.float16)
+    g = -0.2 * torch.rand(bsz, total_tokens, num_v_heads, dtype=torch.float32)
+    beta = (0.15 + 0.35 * torch.rand(bsz, total_tokens, num_v_heads, dtype=torch.float32)).to(torch.float16)
+    cu_seqlens = torch.tensor([0, 4, 9], dtype=torch.long)
+    initial_state = torch.randn(2, num_v_heads, vdim, kdim, dtype=torch.float16)
+
+    out, final_state = chunk_gated_delta_rule_pytorch(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        initial_state=initial_state,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+        head_first=False,
+        use_qk_l2norm_in_kernel=False,
+    )
+
+    assert out.shape == v.shape
+    assert final_state is not None
+    assert final_state.shape == initial_state.shape
+
+
+def test_chunk_gated_delta_rule_310_varlen_tnd_path():
+    torch.manual_seed(0)
+
+    total_tokens = 9
+    num_qk_heads = 2
+    num_v_heads = 4
+    kdim = 16
+    vdim = 12
+
+    q_tnd = torch.randn(total_tokens, num_qk_heads, kdim, dtype=torch.float16)
+    k_tnd = torch.randn(total_tokens, num_qk_heads, kdim, dtype=torch.float16)
+    v_tnd = torch.randn(total_tokens, num_v_heads, vdim, dtype=torch.float16)
+    g_tnd = -0.2 * torch.rand(total_tokens, num_v_heads, dtype=torch.float32)
+    beta_tnd = (0.15 + 0.35 * torch.rand(total_tokens, num_v_heads, dtype=torch.float32)).to(torch.float16)
+    cu_seqlens = torch.tensor([0, 4, 9], dtype=torch.long)
+    initial_state = torch.randn(2, num_v_heads, vdim, kdim, dtype=torch.float16)
+
+    out_tnd, final_state_tnd = chunk_gated_delta_rule_pytorch(
+        q=q_tnd,
+        k=k_tnd,
+        v=v_tnd,
+        g=g_tnd,
+        beta=beta_tnd,
+        initial_state=initial_state,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+        head_first=False,
+        use_qk_l2norm_in_kernel=False,
+    )
+
+    out_bthd, final_state_bthd = chunk_gated_delta_rule_pytorch(
+        q=q_tnd.unsqueeze(0),
+        k=k_tnd.unsqueeze(0),
+        v=v_tnd.unsqueeze(0),
+        g=g_tnd.unsqueeze(0),
+        beta=beta_tnd.unsqueeze(0),
+        initial_state=initial_state,
+        output_final_state=True,
+        cu_seqlens=cu_seqlens,
+        head_first=False,
+        use_qk_l2norm_in_kernel=False,
+    )
+
+    assert out_tnd.shape == v_tnd.shape
+    torch.testing.assert_close(out_tnd, out_bthd[0], rtol=1e-3, atol=1e-3)
+    assert final_state_tnd is not None
+    assert final_state_bthd is not None
+    torch.testing.assert_close(final_state_tnd, final_state_bthd, rtol=1e-4, atol=1e-4)

tests/ut/_310p/ops/test_conv_310.py

@@ -0,0 +1,35 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from vllm.config import set_current_vllm_config
+from vllm.model_executor.layers.conv import Conv3dLayer
+
+from vllm_ascend._310p.ops.conv import AscendConv3dLayer310
+
+
+@pytest.fixture(autouse=True)
+def default_vllm_config():
+    mock_config = MagicMock()
+    mock_config.compilation_config.custom_ops = ["all"]
+    with set_current_vllm_config(mock_config):
+        yield mock_config
+
+
+def test_conv3d_310_forward_oot_uses_forward_native():
+    layer = AscendConv3dLayer310(
+        in_channels=2,
+        out_channels=4,
+        kernel_size=(2, 2, 2),
+        stride=(2, 2, 2),
+        bias=True,
+        params_dtype=torch.float32,
+    )
+    x = torch.randn(1, 2, 4, 4, 4, dtype=torch.float32)
+    expected = torch.randn(1, 4, 2, 2, 2, dtype=torch.float32)
+
+    with patch.object(Conv3dLayer, "forward_native", autospec=True, return_value=expected) as mock_forward_native:
+        out = layer.forward_oot(x)
+
+    mock_forward_native.assert_called_once_with(layer, x)
+    assert out is expected

tests/ut/_310p/ops/test_gdn_310.py

@@ -0,0 +1,141 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from types import SimpleNamespace
+
+import torch
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+
+from vllm_ascend._310p.ops.fla.gdn_310 import (
+    AscendGatedDeltaNetAttention310,
+    _mask_padded_recurrent_accepted_tokens,
+    _zero_padded_tokens,
+)
+from vllm_ascend._310p.ops.gdn_attn_builder_310 import (
+    AscendGDNAttentionBackend310,
+    AscendGDNAttentionMetadataBuilder310,
+)
+
+
+def test_ascend_gdn_attention_310_uses_310p_backend():
+    assert AscendGatedDeltaNetAttention310.get_attn_backend(object()) is AscendGDNAttentionBackend310
+    assert AscendGDNAttentionBackend310.get_builder_cls() is AscendGDNAttentionMetadataBuilder310
+
+
+def test_zero_padded_tokens_masks_only_padded_token_positions():
+    tensor = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
+
+    masked = _zero_padded_tokens(tensor, torch.tensor(2), token_dim=1)
+
+    torch.testing.assert_close(masked[:, :2], tensor[:, :2])
+    assert torch.count_nonzero(masked[:, 2:]) == 0
+
+
+def test_mask_padded_recurrent_accepted_tokens_zeros_dummy_requests():
+    accepted_tokens = torch.tensor([2, 3, 4], dtype=torch.int64)
+    actual_seq_lengths = torch.tensor([4, 0, 1], dtype=torch.int32)
+
+    masked = _mask_padded_recurrent_accepted_tokens(
+        accepted_tokens,
+        actual_seq_lengths,
+    )
+
+    assert masked.dtype == torch.int32
+    assert masked.tolist() == [2, 0, 4]
+
+
+def test_builder310_pads_spec_decode_metadata_with_dummy_requests():
+    builder = object.__new__(AscendGDNAttentionMetadataBuilder310)
+    builder.spec_state_indices_tensor = torch.full((4, 2), -1, dtype=torch.int32)
+    builder.spec_sequence_masks = torch.empty(4, dtype=torch.bool)
+    builder.non_spec_token_indx = torch.empty(0, dtype=torch.int32)
+    builder.spec_token_indx = torch.empty(8, dtype=torch.int32)
+    builder.spec_query_start_loc = torch.empty(5, dtype=torch.int32)
+    builder.num_accepted_tokens = torch.empty(4, dtype=torch.int32)
+    builder.spec_actual_seq_lengths = torch.empty(5, dtype=torch.int32)
+    builder.use_full_cuda_graph = True
+    attn_metadata = SimpleNamespace(
+        num_prefills=0,
+        num_decodes=0,
+        num_spec_decodes=2,
+        spec_state_indices_tensor=torch.tensor(
+            [[3, 30], [4, 40]],
+            dtype=torch.int32,
+        ),
+        spec_sequence_masks=torch.tensor([True, True]),
+        spec_query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32),
+        num_accepted_tokens=torch.tensor([2, 3], dtype=torch.int32),
+        non_spec_token_indx=torch.empty(0, dtype=torch.int32),
+        spec_token_indx=torch.arange(8, dtype=torch.int32),
+    )
+
+    builder._pad_spec_decode_metadata(attn_metadata, graph_batch_size=4)
+
+    assert attn_metadata.spec_state_indices_tensor.tolist() == [
+        [3, 30],
+        [4, 40],
+        [NULL_BLOCK_ID, NULL_BLOCK_ID],
+        [NULL_BLOCK_ID, NULL_BLOCK_ID],
+    ]
+    assert attn_metadata.spec_sequence_masks.tolist() == [True, True, False, False]
+    assert attn_metadata.spec_query_start_loc.tolist() == [0, 4, 8, 8, 8]
+    assert attn_metadata.num_accepted_tokens.tolist() == [2, 3, 0, 0]
+    spec_meta = attn_metadata.spec_decode_metadata.spec_causal_conv1d
+    assert spec_meta.query_start_loc.data_ptr() == attn_metadata.spec_query_start_loc.data_ptr()
+    assert spec_meta.cache_indices.data_ptr() == attn_metadata.spec_state_indices_tensor.data_ptr()
+    assert spec_meta.num_accepted_tokens.data_ptr() == attn_metadata.num_accepted_tokens.data_ptr()
+    assert attn_metadata.spec_decode_metadata.actual_seq_lengths.tolist() == [0, 4, 4, 0, 0]
+
+
+def test_builder310_refreshes_non_spec_decode_graph_metadata():
+    builder = object.__new__(AscendGDNAttentionMetadataBuilder310)
+    builder.non_spec_state_indices_tensor = torch.full((4,), 77, dtype=torch.int32)
+    builder.non_spec_query_start_loc = torch.full((5,), 77, dtype=torch.int32)
+    builder.non_spec_actual_seq_lengths = torch.full((5,), 77, dtype=torch.int32)
+    builder.use_full_cuda_graph = True
+    attn_metadata = SimpleNamespace(
+        num_prefills=0,
+        num_decodes=4,
+        num_decode_tokens=2,
+        num_spec_decodes=0,
+        non_spec_state_indices_tensor=torch.tensor(
+            [10, 11, 98, 99],
+            dtype=torch.int32,
+        ),
+        non_spec_query_start_loc=torch.tensor(
+            [0, 1, 2, 2, 2],
+            dtype=torch.int32,
+        ),
+    )
+
+    builder._pad_decode_metadata(attn_metadata, graph_batch_size=4)
+
+    assert attn_metadata.non_spec_state_indices_tensor.tolist() == [
+        10,
+        11,
+        NULL_BLOCK_ID,
+        NULL_BLOCK_ID,
+    ]
+    assert attn_metadata.non_spec_query_start_loc.tolist() == [0, 1, 2, 2, 2]
+    assert attn_metadata.non_spec_state_indices_tensor.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()
+    assert attn_metadata.non_spec_query_start_loc.data_ptr() == builder.non_spec_query_start_loc.data_ptr()
+    decode_meta = attn_metadata.non_spec_decode_metadata
+    conv_meta = decode_meta.causal_conv1d
+    assert conv_meta.query_start_loc.data_ptr() == attn_metadata.non_spec_query_start_loc.data_ptr()
+    assert conv_meta.cache_indices.data_ptr() == attn_metadata.non_spec_state_indices_tensor.data_ptr()
+    assert decode_meta.actual_seq_lengths.data_ptr() == builder.non_spec_actual_seq_lengths.data_ptr()
+    assert decode_meta.actual_seq_lengths.tolist() == [0, 1, 1, 0, 0]

tests/ut/_310p/ops/test_layernorm_310.py

@@ -0,0 +1,70 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from vllm.config import set_current_vllm_config
+from vllm.model_executor.layers.layernorm import RMSNormGated
+
+from vllm_ascend._310p.ops.layernorm import AscendRMSNormGated310
+
+
+@pytest.fixture(autouse=True)
+def default_vllm_config():
+    mock_config = MagicMock()
+    mock_config.compilation_config.custom_ops = ["all"]
+    with set_current_vllm_config(mock_config):
+        yield mock_config
+
+
+@patch("torch.nn.functional.silu", side_effect=lambda tensor: tensor + 1)
+@patch("torch_npu.npu_rms_norm")
+def test_rmsnorm_gated_310_forward_oot_uses_rmsnorm_activation_mul(mock_rms_norm, mock_silu):
+    layer = AscendRMSNormGated310(hidden_size=8, eps=1e-5, norm_before_gate=True)
+    x = torch.randn(2, 8, dtype=torch.float32)
+    z = torch.randn(2, 8, dtype=torch.float32)
+    normed = torch.randn(2, 8, dtype=torch.float32)
+    mock_rms_norm.return_value = (normed, None)
+
+    with patch.object(RMSNormGated, "forward_native", autospec=True) as mock_forward_native:
+        out = layer.forward_oot(x, z)
+
+    mock_forward_native.assert_not_called()
+    mock_rms_norm.assert_called_once()
+    rms_norm_args = mock_rms_norm.call_args.args
+    assert rms_norm_args[0] is x
+    assert rms_norm_args[1] is layer.weight
+    assert rms_norm_args[2] == layer.eps
+    mock_silu.assert_called_once_with(z)
+    assert torch.allclose(out, normed * (z + 1))
+
+
+@patch("torch_npu.npu_rms_norm")
+def test_rmsnorm_gated_310_forward_oot_uses_rmsnorm_without_gate(mock_rms_norm):
+    layer = AscendRMSNormGated310(hidden_size=8, eps=1e-5)
+    x = torch.randn(2, 8, dtype=torch.float32)
+    expected = torch.randn(2, 8, dtype=torch.float32)
+    mock_rms_norm.return_value = (expected, None)
+
+    with patch.object(RMSNormGated, "forward_native", autospec=True, return_value=expected) as mock_forward_native:
+        out = layer.forward_oot(x, None)
+
+    mock_forward_native.assert_not_called()
+    mock_rms_norm.assert_called_once()
+    rms_norm_args = mock_rms_norm.call_args.args
+    assert rms_norm_args[0] is x
+    assert rms_norm_args[1] is layer.weight
+    assert rms_norm_args[2] == layer.eps
+    assert out is expected
+
+
+def test_rmsnorm_gated_310_forward_oot_keeps_native_for_group_norm():
+    layer = AscendRMSNormGated310(hidden_size=8, eps=1e-5, group_size=4)
+    x = torch.randn(2, 8, dtype=torch.float32)
+    z = torch.randn(2, 8, dtype=torch.float32)
+    expected = torch.randn(2, 8, dtype=torch.float32)
+
+    with patch.object(RMSNormGated, "forward_native", autospec=True, return_value=expected) as mock_forward_native:
+        out = layer.forward_oot(x, z)
+
+    mock_forward_native.assert_called_once_with(layer, x, z)
+    assert out is expected

tests/ut/_310p/ops/test_mm_encoder_attention_310.py

@@ -0,0 +1,81 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest import mock
+
+import torch
+
+from vllm_ascend import utils
+from vllm_ascend._310p.ops.mm_encoder_attention import AscendMMEncoderAttention310
+
+
+def test_register_customop_overrides_mm_encoder_attention_for_310p():
+    original_registered = utils._ASCEND_CUSTOMOP_IS_REIGISTERED
+    try:
+        utils._ASCEND_CUSTOMOP_IS_REIGISTERED = False
+        with (
+            mock.patch("vllm.model_executor.custom_op.CustomOp.register_oot"),
+            mock.patch("vllm_ascend.utils.is_310p", return_value=True),
+        ):
+            utils.register_ascend_customop()
+
+        assert utils.REGISTERED_ASCEND_OPS["MMEncoderAttention"] is AscendMMEncoderAttention310
+    finally:
+        utils._ASCEND_CUSTOMOP_IS_REIGISTERED = original_registered
+
+
+def test_mm_encoder_attention_310_forward_oot_with_padding():
+    layer = AscendMMEncoderAttention310.__new__(AscendMMEncoderAttention310)
+    layer.num_heads = 4
+    layer.num_kv_heads = 2
+    layer.head_size = 80
+    layer.enable_pad = True
+    layer.scale_value = layer.head_size**-0.5
+    layer.support_approximate_calculation = False
+
+    bsz, q_len, kv_len = 2, 3, 3
+    query = torch.randn(bsz, q_len, layer.num_heads, layer.head_size)
+    key = torch.randn(bsz, kv_len, layer.num_kv_heads, layer.head_size)
+    value = torch.randn(bsz, kv_len, layer.num_kv_heads, layer.head_size)
+
+    capture = {}
+
+    def fake_flash_attention_unpad(*, query, key, value, seq_len, scale_value, num_heads, num_kv_heads, out):
+        capture["query_shape"] = query.shape
+        capture["key_shape"] = key.shape
+        capture["value_shape"] = value.shape
+        capture["seq_len"] = seq_len
+        capture["scale_value"] = scale_value
+        capture["num_heads"] = num_heads
+        capture["num_kv_heads"] = num_kv_heads
+        out.copy_(query + 1.0)
+
+    with mock.patch(
+        "vllm_ascend._310p.ops.mm_encoder_attention.torch_npu._npu_flash_attention_unpad",
+        side_effect=fake_flash_attention_unpad,
+        create=True,
+    ):
+        out = layer.forward_oot(query, key, value)
+
+    assert capture["query_shape"] == (bsz * q_len, layer.num_heads, 128)
+    assert capture["key_shape"] == (bsz * kv_len, layer.num_heads, 128)
+    assert capture["value_shape"] == (bsz * kv_len, layer.num_heads, 128)
+    assert capture["seq_len"].device.type == "cpu"
+    torch.testing.assert_close(capture["seq_len"], torch.tensor([q_len, q_len], dtype=torch.int32))
+    assert capture["num_heads"] == layer.num_heads
+    assert capture["num_kv_heads"] == layer.num_kv_heads
+
+    assert out.shape == query.shape
+    torch.testing.assert_close(out, query + 1.0)

tests/ut/_310p/ops/test_rotary_embedding_310.py

@@ -0,0 +1,85 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import torch
+
+from vllm_ascend._310p.ops import rotary_embedding as rotary_310
+from vllm_ascend._310p.ops.rotary_embedding import (
+    AscendMRotaryEmbedding310,
+    AscendRotaryEmbedding310,
+    set_mrope_apply_rotary_slices,
+)
+
+
+def _reset_mrope_globals():
+    rotary_310._mrope_cos_slice = None
+    rotary_310._mrope_sin_slice = None
+
+
+def _build_mrope_embedding() -> AscendMRotaryEmbedding310:
+    emb = AscendMRotaryEmbedding310.__new__(AscendMRotaryEmbedding310)
+    emb.mrope_section = [2, 2, 2]
+    emb.mrope_interleaved = False
+    emb.cos_sin_cache = torch.randn(64, 12, dtype=torch.float32)
+    return emb
+
+
+def test_set_mrope_apply_rotary_slices_populates_globals():
+    _reset_mrope_globals()
+    emb = _build_mrope_embedding()
+    positions = torch.randint(0, emb.cos_sin_cache.shape[0], (3, 4), dtype=torch.long)
+    set_mrope_apply_rotary_slices(
+        emb.cos_sin_cache,
+        positions,
+        mrope_section=emb.mrope_section,
+        mrope_interleaved=emb.mrope_interleaved,
+    )
+
+    assert rotary_310._mrope_cos_slice is not None
+    assert rotary_310._mrope_sin_slice is not None
+    assert rotary_310._mrope_cos_slice.shape[1] == positions.shape[-1]
+
+
+def test_set_mrope_apply_rotary_slices_reuses_buffer_address():
+    _reset_mrope_globals()
+    emb = _build_mrope_embedding()
+    positions = torch.randint(0, emb.cos_sin_cache.shape[0], (3, 4), dtype=torch.long)
+
+    set_mrope_apply_rotary_slices(
+        emb.cos_sin_cache,
+        positions,
+        mrope_section=emb.mrope_section,
+        mrope_interleaved=emb.mrope_interleaved,
+    )
+    first_ptr = rotary_310._mrope_cos_slice.data_ptr()
+
+    set_mrope_apply_rotary_slices(
+        emb.cos_sin_cache,
+        positions,
+        mrope_section=emb.mrope_section,
+        mrope_interleaved=emb.mrope_interleaved,
+    )
+    second_ptr = rotary_310._mrope_cos_slice.data_ptr()
+
+    assert first_ptr == second_ptr
+
+
+def test_ascend_rotary_embedding_310_drafting_flag():
+    assert hasattr(AscendRotaryEmbedding310, "_is_drafting_update_enabled")
+    assert AscendRotaryEmbedding310._is_drafting_update_enabled is False
+    AscendRotaryEmbedding310.set_rope_position_flag_310p(True)
+    assert AscendRotaryEmbedding310._is_drafting_update_enabled is True
+    AscendRotaryEmbedding310.set_rope_position_flag_310p(False)
+    assert AscendRotaryEmbedding310._is_drafting_update_enabled is False

tests/ut/_310p/quantization/test_modelslim_config_310.py

@@ -0,0 +1,138 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, patch
+
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig, FusedMoEParallelConfig
+from vllm.model_executor.layers.linear import LinearBase
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.fused_moe.fused_moe import AscendUnquantizedFusedMoEMethod310
+from vllm_ascend._310p.quantization.modelslim_config import AscendModelSlimConfig310
+from vllm_ascend.ops.linear import AscendUnquantizedLinearMethod
+from vllm_ascend.utils import vllm_version_is
+
+if vllm_version_is("0.23.0"):
+    from vllm.model_executor.layers.fused_moe import FusedMoE
+else:
+    from vllm.model_executor.layers.fused_moe import RoutedExperts
+
+
+class TestAscendModelSlimConfig310(TestBase):
+    def setUp(self):
+        self.sample_config = {
+            "weight": "INT8",
+            "layer1.weight": "INT8",
+            "layer2.weight": "FLOAT",
+            "fused_layer.weight": "FLOAT",
+            "fused_layer.shard1.weight": "FLOAT",
+            "fused_layer.shard2.weight": "FLOAT",
+            "shard1.weight": "FLOAT",
+            "shard2.weight": "FLOAT",
+        }
+        self.ascend_config = AscendModelSlimConfig310(self.sample_config)
+        self.ascend_config.packed_modules_mapping = None
+
+    def test_get_quant_method_for_linear_310(self):
+        mock_config = MagicMock()
+        mock_config.model_config.hf_config.model_type = None
+        linear_layer = MagicMock(spec=LinearBase)
+        # Test skipped layer
+        with (
+            patch("vllm_ascend._310p.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=True),
+        ):
+            method = self.ascend_config.get_quant_method(linear_layer, ".attn")
+            self.assertIsInstance(method, AscendUnquantizedLinearMethod)
+
+        # Test quantized layer
+        mock_scheme = MagicMock()
+        with (
+            patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=False),
+            patch("vllm_ascend._310p.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch("vllm_ascend._310p.quantization.modelslim_config.create_scheme_for_layer", return_value=mock_scheme),
+            patch(
+                "vllm_ascend._310p.quantization.modelslim_config.AscendLinearMethod", return_value=MagicMock()
+            ) as mock_ascend_linear,
+        ):
+            method = self.ascend_config.get_quant_method(linear_layer, ".attn")
+            self.assertIs(method, mock_ascend_linear.return_value)
+            mock_ascend_linear.assert_called_once_with(mock_scheme)
+
+    def test_get_quant_method_maps_lm_head_prefix_310(self):
+        config = AscendModelSlimConfig310({"language_model.lm_head.weight": "INT8"})
+        linear_layer = MagicMock(spec=LinearBase)
+        mock_config = MagicMock()
+        mock_config.model_config.hf_config.model_type = "qwen3_5_moe"
+        mock_scheme = MagicMock()
+
+        with (
+            patch("vllm_ascend._310p.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch(
+                "vllm_ascend._310p.quantization.modelslim_config.create_scheme_for_layer",
+                return_value=mock_scheme,
+            ) as mock_create_scheme,
+            patch("vllm_ascend._310p.quantization.modelslim_config.AscendLinearMethod", return_value=MagicMock()),
+        ):
+            config.get_quant_method(linear_layer, "lm_head")
+
+        mock_create_scheme.assert_called_once_with(
+            quant_description=config.quant_description,
+            prefix="language_model.lm_head",
+            layer_type="linear",
+            packed_modules_mapping=config.packed_modules_mapping,
+        )
+
+    def test_get_quant_method_for_fused_moe_310(self):
+        if vllm_version_is("0.23.0"):
+            fused_moe_cls = FusedMoE
+        else:
+            fused_moe_cls = RoutedExperts
+        fused_moe_layer = MagicMock(spec=fused_moe_cls)
+        fused_moe_layer.moe = MagicMock(spec=FusedMoEConfig)
+        fused_moe_layer.moe_config = MagicMock(spec=FusedMoEConfig)
+        fused_moe_layer.moe_config.moe_backend = "auto"
+        fused_moe_layer.moe_config.moe_parallel_config = MagicMock(spec=FusedMoEParallelConfig)
+        fused_moe_layer.moe_config.moe_parallel_config.use_ep = True
+        fused_moe_layer.moe_config.moe_parallel_config.dp_size = 1
+        mock_config = MagicMock()
+        mock_config.model_config.hf_config.model_type = None
+        mock_config.compilation_config.custom_ops = ["all"]
+        mock_scheme = MagicMock()
+        # Test skipped layer
+        with (
+            patch("vllm.config.vllm.get_current_vllm_config", return_value=mock_config),
+            patch("vllm_ascend._310p.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=True),
+        ):
+            method = self.ascend_config.get_quant_method(fused_moe_layer, ".moe")
+            self.assertIsInstance(method, AscendUnquantizedFusedMoEMethod310)
+
+        # Test quantized layer
+        mock_scheme = MagicMock()
+        with (
+            patch.object(self.ascend_config, "is_layer_skipped_ascend", return_value=False),
+            patch("vllm.config.vllm.get_current_vllm_config", return_value=mock_config),
+            patch("vllm_ascend._310p.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch("vllm_ascend.quantization.modelslim_config.get_current_vllm_config", return_value=mock_config),
+            patch("vllm_ascend._310p.quantization.modelslim_config.create_scheme_for_layer", return_value=mock_scheme),
+            patch(
+                "vllm_ascend._310p.quantization.modelslim_config.AscendFusedMoEMethod", return_value=MagicMock()
+            ) as fused_moe_method,
+        ):
+            method = self.ascend_config.get_quant_method(fused_moe_layer, ".moe")
+            self.assertIs(method, fused_moe_method.return_value)
+            fused_moe_method.assert_called_once_with(mock_scheme, fused_moe_layer.moe_config)

tests/ut/_310p/quantization/test_w8a8_dynamic_310.py

@@ -0,0 +1,145 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, Mock, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.quantization.methods.w8a8_dynamic import (
+    AscendW8A8DynamicFusedMoEMethod310,
+    AscendW8A8DynamicLinearMethod310,
+)
+
+
+class TestAscendW8A8FusedMoEMethod310(TestBase):
+    num_experts = 8
+    hidden_size = 128
+    intermediate_size = 128
+
+    @patch("vllm_ascend._310p.quantization.methods.w8a8_dynamic.get_ep_group")
+    def setUp(self, mock_get_ep_group):
+        with patch(
+            "vllm_ascend._310p.quantization.methods.w8a8_dynamic.get_current_vllm_config"
+        ) as mock_get_current_vllm_config:
+            mock_vllm_config = Mock()
+            mock_vllm_config.quant_config = Mock(quant_description={"group_size": 0})
+            mock_vllm_config.scheduler_config = Mock(
+                max_num_batched_tokens=2048, max_model_len=2048, enable_chunked_prefill=False
+            )
+            mock_get_current_vllm_config.return_value = mock_vllm_config
+            mock_ep_group = Mock()
+            mock_get_ep_group.return_value = mock_ep_group
+            mock_ascend_config = Mock()
+
+            mock_ascend_config.enable_chunked_prefill = False
+
+            self.quant_method = AscendW8A8DynamicFusedMoEMethod310()
+
+    def test_get_weight_310(self):
+        param_dict = self.quant_method.get_weight(
+            self.num_experts, self.intermediate_size, self.hidden_size, torch.float16
+        )
+        self.assertEqual(param_dict["w13_weight"].dtype, torch.int8)
+        self.assertEqual(
+            param_dict["w13_weight"].shape, (self.num_experts, 2 * self.intermediate_size, self.hidden_size)
+        )
+        self.assertEqual(param_dict["w2_weight"].dtype, torch.int8)
+        self.assertEqual(param_dict["w2_weight"].shape, (self.num_experts, self.hidden_size, self.intermediate_size))
+
+    def test_get_dynamic_quant_param_310(self):
+        param_dict = self.quant_method.get_dynamic_quant_param(
+            self.num_experts, self.intermediate_size, self.hidden_size, torch.float16
+        )
+        self.assertEqual(param_dict["w13_weight_scale"].dtype, torch.float32)
+        self.assertEqual(param_dict["w13_weight_scale"].shape, (self.num_experts, 2 * self.intermediate_size, 1))
+        self.assertEqual(param_dict["w2_weight_scale"].dtype, torch.float32)
+        self.assertEqual(param_dict["w2_weight_scale"].shape, (self.num_experts, self.hidden_size, 1))
+
+
+class TestAscendW8A8DynamicLinearMethod310(TestBase):
+    def setUp(self):
+        self.method = AscendW8A8DynamicLinearMethod310()
+
+    def test_get_weight_310(self):
+        weight = self.method.get_weight(10, 20)
+        self.assertEqual(weight["weight"].dtype, torch.int8)
+        self.assertEqual(weight["weight"].shape, (20, 10))
+
+    def test_get_perchannel_param_310(self):
+        params = self.method.get_perchannel_param(10, torch.float32)
+
+        self.assertEqual(params["weight_scale"].dtype, torch.float32)
+        self.assertEqual(params["weight_offset"].dtype, torch.float32)
+
+        self.assertEqual(params["weight_scale"].shape, (10, 1))
+        self.assertEqual(params["weight_offset"].shape, (10, 1))
+
+    @patch("torch_npu.npu_dynamic_quant", create=True)
+    @patch("torch_npu.npu_quant_matmul")
+    def test_apply_310(self, mock_npu_quant_matmul, mock_npu_dynamic_quantize):
+        layer = MagicMock()
+        layer.weight = torch.randn(128, 256, dtype=torch.float16)
+        layer.weight_scale = torch.randn(128, dtype=torch.float32)
+        layer.params_dtype = torch.float16
+
+        x = torch.randn(32, 128, dtype=torch.float16)
+        expect_x_output = torch.randint(-128, 127, x.shape, dtype=torch.int8)
+        expect_pertoken_scale_output = torch.randn(x.shape[0], dtype=torch.float32)
+        mock_npu_dynamic_quantize.return_value = expect_x_output, expect_pertoken_scale_output
+
+        expected_y_output = torch.randn(32, 256)
+        mock_npu_quant_matmul.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_npu_dynamic_quantize.assert_called_with(x)
+        mock_npu_quant_matmul.assert_called_once()
+        (args, kwargs) = mock_npu_quant_matmul.call_args
+
+        # positional args
+        self.assertTrue(torch.equal(args[0], expect_x_output))
+        self.assertTrue(torch.equal(args[1], layer.weight.data))
+        self.assertTrue(torch.equal(args[2], layer.weight_scale))
+
+        # kwargs
+        self.assertTrue(torch.equal(kwargs["pertoken_scale"], expect_pertoken_scale_output))
+        self.assertTrue(kwargs["bias"] is None)
+        self.assertEqual(kwargs["output_dtype"], layer.params_dtype)
+
+        self.assertTrue(torch.equal(output, expected_y_output))
+
+    @patch("vllm_ascend.utils.is_310p", return_value=True)
+    @patch("torch_npu.npu_format_cast")
+    def test_process_weights_after_loading_calls_nz_format_cast_310p(self, mock_npu_format_cast, _mock_is_310p):
+        mock_npu_format_cast.side_effect = lambda x, fmt: x
+
+        layer = MagicMock()
+
+        # Attributes used by process_weights_after_loading()
+        layer.weight = MagicMock()
+        layer.weight_scale = MagicMock()
+        layer.weight_offset = MagicMock()
+
+        layer.weight.data = torch.randint(-127, 128, (128, 256), dtype=torch.int8)
+
+        layer.weight_scale.data = torch.randn(128, 1, dtype=torch.bfloat16)
+        layer.weight_offset.data = torch.randn(128, 1, dtype=torch.bfloat16)
+        # w2_weight_offset is reshaped to (N, -1); any (N, 1) is fine
+        layer.w2_weight_offset.data = torch.randn(128, 1, dtype=torch.bfloat16)
+
+        self.method.process_weights_after_loading(layer)
+
+        mock_npu_format_cast.assert_called_once()

tests/ut/_310p/quantization/test_w8a8_static_310.py

@@ -0,0 +1,151 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.quantization.methods.w8a8_static import AscendW8A8LinearMethod310
+
+
+class TestAscendW8A8LinearMethod310(TestBase):
+    def setUp(self):
+        self.method = AscendW8A8LinearMethod310()
+
+    def test_get_weight_310(self):
+        weight = self.method.get_weight(10, 20)
+        self.assertEqual(weight["weight"].dtype, torch.int8)
+        self.assertEqual(weight["weight"].shape, (20, 10))
+
+    def test_get_pertensor_param_310(self):
+        params = self.method.get_pertensor_param(torch.float16)
+        self.assertEqual(params["input_scale"].dtype, torch.float16)
+        self.assertEqual(params["input_offset"].dtype, torch.int8)
+        self.assertEqual(params["input_scale"].shape, (1,))
+        self.assertEqual(params["input_offset"].shape, (1,))
+
+    def test_get_perchannel_param_310(self):
+        params = self.method.get_perchannel_param(10, torch.float16)
+
+        self.assertEqual(params["quant_bias"].dtype, torch.int32)
+        self.assertEqual(params["deq_scale"].dtype, torch.int64)
+        self.assertEqual(params["weight_scale"].dtype, torch.float16)
+        self.assertEqual(params["weight_offset"].dtype, torch.float16)
+
+        self.assertEqual(params["quant_bias"].shape, (10,))
+        self.assertEqual(params["deq_scale"].shape, (10,))
+        self.assertEqual(params["weight_scale"].shape, (10, 1))
+        self.assertEqual(params["weight_offset"].shape, (10, 1))
+
+    @patch("torch.ops.vllm.quantize")
+    @patch("torch_npu.npu_quant_matmul")
+    def test_apply_with_x_not_int8_310(self, mock_npu_quant_matmul, mock_quantize):
+        layer = MagicMock()
+        layer.aclnn_input_scale = torch.randn(256)
+        layer.aclnn_input_scale_reciprocal = 1.0 / layer.aclnn_input_scale
+        layer.aclnn_input_offset = torch.randint(-128, 127, (256,), dtype=torch.int8)
+        layer.weight = torch.randn(128, 256)
+        layer.deq_scale = torch.randn(128)
+        layer.quant_bias = torch.randint(-128, 127, (256,))
+        layer.params_dtype = torch.float16
+
+        x = torch.randn(32, 128)
+        expect_x_output = torch.randint(-128, 127, x.shape, dtype=torch.int8)
+        mock_quantize.return_value = expect_x_output
+
+        expected_y_output = torch.randn(32, 256)
+        mock_npu_quant_matmul.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_quantize.assert_called_with(
+            x,
+            layer.aclnn_input_scale,
+            layer.aclnn_input_scale_reciprocal,
+            layer.aclnn_input_offset,
+        )
+        mock_npu_quant_matmul.assert_called_once()
+        (args, kwargs) = mock_npu_quant_matmul.call_args
+
+        # positional args
+        self.assertTrue(torch.equal(args[0], expect_x_output))
+        self.assertTrue(torch.equal(args[1], layer.weight.data))
+        self.assertTrue(torch.equal(args[2], layer.deq_scale))
+
+        # kwargs
+        self.assertTrue(torch.equal(kwargs["bias"], layer.quant_bias))
+        self.assertEqual(kwargs["output_dtype"], layer.params_dtype)
+
+        self.assertTrue(torch.equal(output, expected_y_output))
+
+    @patch("torch.ops.vllm.quantize")
+    @patch("torch_npu.npu_quant_matmul")
+    def test_apply_with_x_is_int8_310(self, mock_npu_quant_matmul, mock_quantize):
+        layer = MagicMock()
+        layer.aclnn_input_scale = torch.randn(256)
+        layer.aclnn_input_offset = torch.randint(-128, 127, (256,), dtype=torch.int8)
+        layer.weight = torch.randn(128, 256)
+        layer.deq_scale = torch.randn(128)
+        layer.quant_bias = torch.randint(-128, 127, (256,))
+        layer.params_dtype = torch.float16
+
+        x = torch.randint(-128, 127, (32, 128), dtype=torch.int8)
+
+        expected_y_output = torch.randn(32, 256)
+        mock_npu_quant_matmul.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_quantize.assert_not_called()
+        mock_npu_quant_matmul.assert_called_once()
+        (args, kwargs) = mock_npu_quant_matmul.call_args
+
+        self.assertTrue(torch.equal(args[0], x))
+        self.assertTrue(torch.equal(args[1], layer.weight.data))
+        self.assertTrue(torch.equal(args[2], layer.deq_scale))
+
+        self.assertTrue(torch.equal(kwargs["bias"], layer.quant_bias))
+        self.assertEqual(kwargs["output_dtype"], layer.params_dtype)
+
+        self.assertTrue(torch.equal(output, expected_y_output))
+
+    @patch("vllm_ascend.utils.is_310p", return_value=True)
+    @patch("torch_npu.npu_format_cast")
+    def test_process_weights_after_loading_calls_nz_format_cast_310p(self, mock_npu_format_cast, _mock_is_310p):
+        mock_npu_format_cast.side_effect = lambda x, fmt: x
+
+        layer = MagicMock()
+
+        # Attributes used by process_weights_after_loading()
+        layer.weight = MagicMock()
+        layer.input_scale = MagicMock()
+        layer.input_offset = MagicMock()
+        layer.weight_scale = MagicMock()
+        layer.weight_offset = MagicMock()
+        layer.w2_weight_offset = MagicMock()
+
+        layer.weight.data = torch.randint(-127, 128, (128, 256), dtype=torch.int8)
+        layer.input_scale.data = torch.tensor([0.1], dtype=torch.float16)
+        layer.input_offset.data = torch.tensor([0], dtype=torch.int8)
+
+        layer.weight_scale.data = torch.randn(128, 1, dtype=torch.bfloat16)
+        layer.weight_offset.data = torch.randn(128, 1, dtype=torch.bfloat16)
+        # w2_weight_offset is reshaped to (N, -1); any (N, 1) is fine
+        layer.w2_weight_offset.data = torch.randn(128, 1, dtype=torch.bfloat16)
+
+        self.method.process_weights_after_loading(layer)
+
+        mock_npu_format_cast.assert_called_once()

tests/ut/_310p/quantization/test_w8a8s_310.py

@@ -0,0 +1,93 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.quantization.methods.w8a8s import AscendW8A8SLinearMethod310
+
+
+class TestAscendW8A8SLinearMethod310(TestBase):
+    def setUp(self):
+        self.method = AscendW8A8SLinearMethod310()
+
+    def test_get_weight_310(self):
+        weight = self.method.get_weight(10, 20)
+        self.assertEqual(weight["weight"].dtype, torch.int8)
+        self.assertEqual(weight["weight"].shape, (20, 10))
+
+    def test_get_pertensor_param_310(self):
+        params = self.method.get_pertensor_param(torch.float16)
+        self.assertEqual(params["input_scale"].dtype, torch.float16)
+        self.assertEqual(params["input_offset"].dtype, torch.int8)
+        self.assertEqual(params["input_scale"].shape, (1,))
+        self.assertEqual(params["input_offset"].shape, (1,))
+
+    def test_get_perchannel_param_310(self):
+        params = self.method.get_perchannel_param(10, torch.float16)
+
+        self.assertEqual(params["quant_bias"].dtype, torch.int32)
+        self.assertEqual(params["deq_scale"].dtype, torch.int64)
+        self.assertEqual(params["quant_bias"].shape, (10,))
+        self.assertEqual(params["deq_scale"].shape, (10,))
+
+    @patch("torch.ops.vllm.quantize")
+    @patch("torch_npu.npu_quant_matmul")
+    def test_apply_with_x_not_int8_310(self, mock_npu_quant_matmul, mock_quantize):
+        layer = MagicMock()
+        layer.aclnn_input_scale = torch.randn(256)
+        layer.aclnn_input_scale_reciprocal = 1.0 / layer.aclnn_input_scale
+        layer.aclnn_input_offset = torch.randint(-128, 127, (256,), dtype=torch.int8)
+        layer.weight = torch.randn(128, 256)
+        layer.deq_scale = torch.randn(128)
+        layer.quant_bias = torch.randint(-128, 127, (256,))
+        layer.params_dtype = torch.float16
+
+        x = torch.randn(32, 128)
+        expect_x_output = torch.randint(-128, 127, x.shape, dtype=torch.int8)
+        mock_quantize.return_value = expect_x_output
+
+        expected_y_output = torch.randn(32, 256)
+        mock_npu_quant_matmul.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_quantize.assert_called_with(
+            x, layer.aclnn_input_scale, layer.aclnn_input_scale_reciprocal, layer.aclnn_input_offset
+        )
+        self.assertTrue(torch.equal(output, expected_y_output))
+
+    @patch("torch.ops.vllm.quantize")
+    @patch("torch_npu.npu_quant_matmul")
+    def test_apply_with_x_is_int8_310(self, mock_npu_quant_matmul, mock_quantize):
+        layer = MagicMock()
+        layer.aclnn_input_scale = torch.randn(256)
+        layer.aclnn_input_offset = torch.randint(-128, 127, (256,), dtype=torch.int8)
+        layer.weight = torch.randn(128, 256)
+        layer.deq_scale = torch.randn(128)
+        layer.quant_bias = torch.randint(-128, 127, (256,))
+        layer.params_dtype = torch.float16
+
+        x = torch.randint(-128, 127, (32, 128), dtype=torch.int8)
+
+        expected_y_output = torch.randn(32, 256)
+        mock_npu_quant_matmul.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_quantize.assert_not_called()
+        self.assertTrue(torch.equal(output, expected_y_output))

tests/ut/_310p/quantization/test_w8a8sc_310.py

@@ -0,0 +1,108 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.quantization.methods.w8a8sc import AscendW8A8SCLinearMethod310
+
+
+class TestAscendW8A8SCLinearMethod310(TestBase):
+    def setUp(self):
+        self.method = AscendW8A8SCLinearMethod310()
+
+    def test_get_weight_310(self):
+        weight = self.method.get_weight(10, 20)
+        self.assertEqual(weight["weight"].dtype, torch.int8)
+        self.assertEqual(weight["weight"].shape, (10 * 20,))
+        self.assertEqual(weight["index"].dtype, torch.int8)
+        index_len = math.ceil(10 / 256) * math.ceil(20 / 128) * 8
+        self.assertEqual(weight["index"].shape, (index_len,))
+        self.assertEqual(weight["info"].dtype, torch.int64)
+        self.assertEqual(weight["info"].shape, (5,))
+
+    def test_get_pertensor_param_310(self):
+        params = self.method.get_pertensor_param(torch.float16)
+        self.assertEqual(params["input_scale"].dtype, torch.float16)
+        self.assertEqual(params["input_offset"].dtype, torch.int8)
+        self.assertEqual(params["input_scale"].shape, (1,))
+        self.assertEqual(params["input_offset"].shape, (1,))
+
+    def test_get_perchannel_param_310(self):
+        params = self.method.get_perchannel_param(10, torch.float16)
+
+        self.assertEqual(params["quant_bias"].dtype, torch.int32)
+        self.assertEqual(params["deq_scale"].dtype, torch.int64)
+        self.assertEqual(params["quant_bias"].shape, (10,))
+        self.assertEqual(params["deq_scale"].shape, (10,))
+
+    @pytest.mark.skip("Skip as npu_matmul_compress_dequant will be supported in PTA 26.0.0.")
+    @patch("torch.ops.vllm.quantize")
+    @patch("torch_npu.npu_matmul_compress_dequant")
+    def test_apply_with_x_not_int8_310(self, mock_matmul_compress_dequant, mock_quantize):
+        layer = MagicMock()
+        layer.aclnn_input_scale = torch.randn(256)
+        layer.aclnn_input_scale_reciprocal = 1.0 / layer.aclnn_input_scale
+        layer.aclnn_input_offset = torch.randint(-128, 127, (256,), dtype=torch.int8)
+        layer.weight = torch.randint(-128, 127, (256 * 128,), dtype=torch.int8)
+        layer.index = torch.randint(-128, 127, (8,), dtype=torch.int8)
+        layer.deq_scale = torch.randn(128)
+        layer.quant_bias = torch.randint(-128, 127, (256,))
+        layer.params_dtype = torch.float16
+
+        x = torch.randn(32, 128)
+        expect_x_output = torch.randint(-128, 127, x.shape, dtype=torch.int8)
+        mock_quantize.return_value = expect_x_output
+
+        expected_y_output = torch.randn(32, 256)
+        mock_matmul_compress_dequant.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_quantize.assert_called_with(
+            x, layer.aclnn_input_scale, layer.aclnn_input_scale_reciprocal, layer.aclnn_input_offset
+        )
+        mock_matmul_compress_dequant.assert_called_with(
+            expect_x_output, layer.weight, layer.index, layer.quant_bias, layer.deq_scale
+        )
+        self.assertTrue(torch.equal(output, expected_y_output))
+
+    @pytest.mark.skip("Skip as npu_matmul_compress_dequant will be supported in PTA 26.0.0.")
+    @patch("torch.ops.vllm.quantize")
+    @patch("torch_npu.npu_matmul_compress_dequant")
+    def test_apply_with_x_is_int8_310(self, mock_matmul_compress_dequant, mock_quantize):
+        layer = MagicMock()
+        layer.aclnn_input_scale = torch.randn(256)
+        layer.aclnn_input_offset = torch.randint(-128, 127, (256,), dtype=torch.int8)
+        layer.weight = torch.randint(-128, 127, (256 * 128,), dtype=torch.int8)
+        layer.index = torch.randint(-128, 127, (8,), dtype=torch.int8)
+        layer.deq_scale = torch.randn(128)
+        layer.quant_bias = torch.randint(-128, 127, (256,))
+        layer.params_dtype = torch.float16
+
+        x = torch.randint(-128, 127, (32, 128), dtype=torch.int8)
+
+        expected_y_output = torch.randn(32, 256)
+        mock_matmul_compress_dequant.return_value = expected_y_output
+
+        output = self.method.apply(layer, x, tp_rank=0)
+
+        mock_quantize.assert_not_called()
+        mock_matmul_compress_dequant.assert_called_with(x, layer.weight, layer.index, layer.quant_bias, layer.deq_scale)
+        self.assertTrue(torch.equal(output, expected_y_output))

tests/ut/_310p/sample/test_sampler_310.py

@@ -0,0 +1,283 @@
+import sys
+import unittest
+from contextlib import nullcontext
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import MagicMock, patch
+
+import torch
+
+PROJECT_ROOT = Path(__file__).resolve().parents[4]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+if "vllm" not in sys.modules:
+    vllm_module = ModuleType("vllm")
+    vllm_envs_module = ModuleType("vllm.envs")
+    vllm_envs_module.VLLM_BATCH_INVARIANT = False  # type: ignore[attr-defined]
+    vllm_module.envs = vllm_envs_module  # type: ignore[attr-defined]
+    sys.modules["vllm"] = vllm_module
+    sys.modules["vllm.envs"] = vllm_envs_module
+
+if "vllm_ascend.sample.sampler" not in sys.modules:
+    sample_sampler_module = ModuleType("vllm_ascend.sample.sampler")
+    sample_sampler_module.DEFAULT_LOGPROBS_MODE = "raw_logprobs"  # type: ignore[attr-defined]
+    sample_sampler_module.AscendSampler = type("AscendSampler", (), {})  # type: ignore[attr-defined]
+    sample_sampler_module.AscendTopKTopPSampler = type("AscendTopKTopPSampler", (), {})  # type: ignore[attr-defined]
+    sys.modules["vllm_ascend.sample.sampler"] = sample_sampler_module
+
+if "vllm_ascend.utils" not in sys.modules:
+    utils_module = ModuleType("vllm_ascend.utils")
+    utils_module.global_stream = lambda: MagicMock()  # type: ignore[attr-defined]
+    utils_module.npu_stream_switch = lambda _: nullcontext()  # type: ignore[attr-defined]
+    sys.modules["vllm_ascend.utils"] = utils_module
+
+from vllm_ascend._310p.sample import sampler as sampler_310p  # noqa: E402
+
+
+class _FakeRow:
+    def __init__(self):
+        self.generators = []
+
+    def exponential_(self, generator=None):
+        self.generators.append(generator)
+        return self
+
+
+class _FakeQ:
+    def __init__(self, batch_size):
+        self.shape = (batch_size, 4)
+        self.default_exponential_called = False
+        self.rows = {idx: _FakeRow() for idx in range(batch_size)}
+
+    def cpu(self):
+        return self
+
+    def npu(self):
+        return self
+
+    def exponential_(self, generator=None):
+        if generator is None:
+            self.default_exponential_called = True
+        return self
+
+    def __getitem__(self, idx):
+        return self.rows[idx]
+
+    def __setitem__(self, idx, value):
+        self.rows[idx] = value
+
+
+def _empty_like_side_effect(q_instances, template):
+    if isinstance(template, _FakeRow):
+        return _FakeRow()
+    return next(q_instances)
+
+
+class _FakeCPUGenerator:
+    def __init__(self, device=None):
+        self.device = device
+        self.state = None
+        self.seed = None
+
+    def set_state(self, state):
+        self.state = state
+
+    def manual_seed(self, seed):
+        self.seed = seed
+
+
+class TestSampler310pStandalone(unittest.TestCase):
+    def tearDown(self):
+        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
+
+    def test_random_sample_310p_reuse_cpu_generator_cache(self):
+        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
+        probs = MagicMock()
+        probs.div_.return_value = probs
+        probs.argmax.return_value = probs
+        probs.view.return_value = torch.tensor([0])
+
+        fake_q_first = _FakeQ(batch_size=2)
+        fake_q_second = _FakeQ(batch_size=2)
+        q_instances = iter([fake_q_first, fake_q_second])
+
+        npu_stream = MagicMock()
+        generator = MagicMock()
+        generator.get_state.return_value = b"state"
+        generator.initial_seed.return_value = 7
+        generators = {1: generator}
+
+        with (
+            patch.object(sampler_310p, "npu_stream_switch", return_value=nullcontext()),
+            patch.object(sampler_310p, "global_stream", return_value=MagicMock()),
+            patch.object(
+                sampler_310p.torch,
+                "empty_like",
+                side_effect=lambda template: _empty_like_side_effect(q_instances, template),
+            ),
+            patch.object(sampler_310p.torch, "Generator", side_effect=_FakeCPUGenerator) as gen_ctor,
+            patch.object(
+                sampler_310p.torch,
+                "npu",
+                ModuleType("torch.npu"),
+                create=True,
+            ),
+        ):
+            sampler_310p.torch.npu.current_stream = MagicMock(return_value=npu_stream)
+            sampler_310p._random_sample_310p(probs, generators)
+            sampler_310p._random_sample_310p(probs, generators)
+
+        self.assertEqual(gen_ctor.call_count, 1)
+        self.assertIn(1, sampler_310p._CPU_GENERATOR_CACHE_310P)
+        cached_cpu_generator, source_generator_id = sampler_310p._CPU_GENERATOR_CACHE_310P[1]
+        self.assertIs(fake_q_first.rows[1].generators[0], cached_cpu_generator)
+        self.assertIs(fake_q_second.rows[1].generators[0], cached_cpu_generator)
+        self.assertEqual(source_generator_id, id(generator))
+        self.assertEqual(cached_cpu_generator.state, b"state")
+        self.assertIsNone(cached_cpu_generator.seed)
+        self.assertEqual(npu_stream.wait_stream.call_count, 2)
+
+    def test_random_sample_310p_fallback_to_initial_seed_when_set_state_failed(self):
+        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
+        probs = MagicMock()
+        probs.div_.return_value = probs
+        probs.argmax.return_value = probs
+        probs.view.return_value = torch.tensor([1])
+
+        fake_q = _FakeQ(batch_size=1)
+        q_instances = iter([fake_q])
+        npu_stream = MagicMock()
+        generator = MagicMock()
+        generator.get_state.side_effect = RuntimeError("state read failed")
+        generator.initial_seed.return_value = 1234
+        generators = {0: generator}
+
+        class _FailSetStateCPUGenerator(_FakeCPUGenerator):
+            def set_state(self, state):
+                raise RuntimeError("state set failed")
+
+        with (
+            patch.object(sampler_310p, "npu_stream_switch", return_value=nullcontext()),
+            patch.object(sampler_310p, "global_stream", return_value=MagicMock()),
+            patch.object(
+                sampler_310p.torch,
+                "empty_like",
+                side_effect=lambda template: _empty_like_side_effect(q_instances, template),
+            ),
+            patch.object(sampler_310p.torch, "Generator", side_effect=_FailSetStateCPUGenerator),
+            patch.object(
+                sampler_310p.torch,
+                "npu",
+                ModuleType("torch.npu"),
+                create=True,
+            ),
+        ):
+            sampler_310p.torch.npu.current_stream = MagicMock(return_value=npu_stream)
+            sampler_310p._random_sample_310p(probs, generators)
+
+        cached_cpu_generator, source_generator_id = sampler_310p._CPU_GENERATOR_CACHE_310P[0]
+        self.assertEqual(source_generator_id, id(generator))
+        self.assertEqual(cached_cpu_generator.seed, 1234)
+        self.assertIs(fake_q.rows[0].generators[0], cached_cpu_generator)
+        self.assertEqual(npu_stream.wait_stream.call_count, 1)
+
+    def test_random_sample_310p_rebuild_cache_when_generator_identity_changes(self):
+        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
+        probs = MagicMock()
+        probs.div_.return_value = probs
+        probs.argmax.return_value = probs
+        probs.view.return_value = torch.tensor([0])
+
+        fake_q_first = _FakeQ(batch_size=1)
+        fake_q_second = _FakeQ(batch_size=1)
+        q_instances = iter([fake_q_first, fake_q_second])
+        npu_stream = MagicMock()
+
+        generator_first = MagicMock()
+        generator_first.get_state.return_value = b"state-1"
+        generator_first.initial_seed.return_value = 11
+
+        generator_second = MagicMock()
+        generator_second.get_state.return_value = b"state-2"
+        generator_second.initial_seed.return_value = 22
+
+        with (
+            patch.object(sampler_310p, "npu_stream_switch", return_value=nullcontext()),
+            patch.object(sampler_310p, "global_stream", return_value=MagicMock()),
+            patch.object(
+                sampler_310p.torch,
+                "empty_like",
+                side_effect=lambda template: _empty_like_side_effect(q_instances, template),
+            ),
+            patch.object(sampler_310p.torch, "Generator", side_effect=_FakeCPUGenerator) as gen_ctor,
+            patch.object(
+                sampler_310p.torch,
+                "npu",
+                ModuleType("torch.npu"),
+                create=True,
+            ),
+        ):
+            sampler_310p.torch.npu.current_stream = MagicMock(return_value=npu_stream)
+            sampler_310p._random_sample_310p(probs, {0: generator_first})
+            sampler_310p._random_sample_310p(probs, {0: generator_second})
+
+        self.assertEqual(gen_ctor.call_count, 2)
+        first_cpu_generator = fake_q_first.rows[0].generators[0]
+        second_cpu_generator = fake_q_second.rows[0].generators[0]
+        self.assertIsNot(first_cpu_generator, second_cpu_generator)
+        self.assertEqual(first_cpu_generator.state, b"state-1")
+        self.assertEqual(second_cpu_generator.state, b"state-2")
+        cached_cpu_generator, source_generator_id = sampler_310p._CPU_GENERATOR_CACHE_310P[0]
+        self.assertIs(cached_cpu_generator, second_cpu_generator)
+        self.assertEqual(source_generator_id, id(generator_second))
+
+    def test_fill_cpu_exponential_310p_moves_has_draft_mask_to_cpu(self):
+        """Regression: NPU has_draft_mask must be moved to CPU before torch.where."""
+        sampler_310p._CPU_GENERATOR_CACHE_310P.clear()
+
+        q_cpu = torch.full((2, 4), 7.0)
+        cpu_mask = torch.tensor([True, False])
+        has_draft_mask = MagicMock()
+        has_draft_mask.cpu.return_value = cpu_mask
+
+        def _make_source_generator(seed: int):
+            source_generator = MagicMock()
+            seed_generator = torch.Generator(device="cpu")
+            seed_generator.manual_seed(seed)
+            source_generator.get_state.return_value = seed_generator.get_state()
+            source_generator.initial_seed.return_value = seed
+            return source_generator
+
+        where_conditions = []
+        real_where = torch.where
+
+        def where_spy(condition, x, y):
+            where_conditions.append(condition.detach().clone())
+            self.assertEqual(condition.device.type, "cpu")
+            self.assertEqual(x.device.type, "cpu")
+            self.assertEqual(y.device.type, "cpu")
+            return real_where(condition, x, y)
+
+        with patch.object(sampler_310p.torch, "where", side_effect=where_spy):
+            sampler_310p._fill_cpu_exponential_310p(
+                q_cpu,
+                {
+                    0: _make_source_generator(42),
+                    1: _make_source_generator(43),
+                },
+                has_draft_mask,
+            )
+
+        has_draft_mask.cpu.assert_called_once()
+        self.assertEqual(len(where_conditions), 2)
+        self.assertTrue(bool(where_conditions[0]))
+        self.assertFalse(bool(where_conditions[1]))
+        # Row 0 (masked): overwritten by seeded exponential via torch.where.
+        self.assertFalse(torch.equal(q_cpu[0], torch.full((4,), 7.0)))
+        # Row 1 (unmasked): also overwritten by the default exponential_ prefill.
+        self.assertFalse(torch.equal(q_cpu[1], torch.full((4,), 7.0)))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/_310p/spec_decode/test_llm_base_proposer_310.py

@@ -0,0 +1,83 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, writing
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.ops.rotary_embedding import AscendRotaryEmbedding310
+from vllm_ascend._310p.spec_decode.llm_base_proposer_310 import AscendSpecDecodeBaseProposer310
+from vllm_ascend.spec_decode.llm_base_proposer import AscendSpecDecodeBaseProposer
+
+
+class TestAscendSpecDecodeBaseProposer310(TestBase):
+    def test_run_merged_draft_sets_rope_flag_before_call(self):
+        flag_states = []
+
+        def mock_original(
+            self,
+            num_input_tokens,
+            batch_size,
+            token_indices_to_sample,
+            target_positions,
+            inputs_embeds,
+            multi_steps_attn_metadata,
+            num_tokens,
+            is_prefill=None,
+        ):
+            flag_states.append(AscendRotaryEmbedding310._is_drafting_update_enabled)
+            return torch.zeros(num_tokens, dtype=torch.long)
+
+        with (
+            patch.object(AscendSpecDecodeBaseProposer, "_run_merged_draft", mock_original),
+            patch("vllm_ascend._310p.spec_decode.llm_base_proposer_310._original_run_merged_draft", mock_original),
+        ):
+            proposer = object.__new__(AscendSpecDecodeBaseProposer310)
+            proposer._run_merged_draft(
+                num_input_tokens=4,
+                batch_size=2,
+                token_indices_to_sample=torch.tensor([0, 1]),
+                target_positions=torch.tensor([0, 1, 2, 3]),
+                inputs_embeds=torch.zeros(4, 128),
+                multi_steps_attn_metadata=None,
+                num_tokens=4,
+            )
+
+        self.assertEqual(len(flag_states), 1)
+        self.assertTrue(flag_states[0])
+        self.assertFalse(AscendRotaryEmbedding310._is_drafting_update_enabled)
+
+    def test_run_merged_draft_restores_rope_flag_after_exception(self):
+        def mock_original(*args, **kwargs):
+            raise RuntimeError("Test exception")
+
+        with (
+            patch.object(AscendSpecDecodeBaseProposer, "_run_merged_draft", mock_original),
+            patch("vllm_ascend._310p.spec_decode.llm_base_proposer_310._original_run_merged_draft", mock_original),
+        ):
+            proposer = object.__new__(AscendSpecDecodeBaseProposer310)
+            with self.assertRaises(RuntimeError):
+                proposer._run_merged_draft(
+                    num_input_tokens=4,
+                    batch_size=2,
+                    token_indices_to_sample=torch.tensor([0, 1]),
+                    target_positions=torch.tensor([0, 1, 2, 3]),
+                    inputs_embeds=torch.zeros(4, 128),
+                    multi_steps_attn_metadata=None,
+                    num_tokens=4,
+                )
+
+        self.assertFalse(AscendRotaryEmbedding310._is_drafting_update_enabled)

tests/ut/_310p/test_block_table_310p.py

@@ -0,0 +1,237 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import torch
+from vllm.distributed.parallel_state import GroupCoordinator
+
+from tests.ut.base import TestBase
+
+
+class TestBlockTable310(TestBase):
+    def setUp(self):
+        self.block_size = 128
+        self.max_num_reqs = 4
+        self.max_num_blocks_per_req = 128
+        self.max_num_batched_tokens = 512
+        self.pin_memory = False
+        self.device = torch.device("cpu")
+        self.kernel_sizes = [128]
+
+    def _create_block_table(self, dcp_world_size, dcp_rank, pcp_world_size, pcp_rank, cp_kv_cache_interleave_size):
+        with (
+            patch("vllm_ascend.worker.block_table.get_dcp_group") as mock_get_dcp_group,
+            patch("vllm_ascend.worker.block_table.get_pcp_group") as mock_get_pcp_group,
+        ):
+            mock_dcp_group = MagicMock(spec=GroupCoordinator)
+            mock_dcp_group.world_size = dcp_world_size
+            mock_dcp_group.rank_in_group = dcp_rank
+            mock_get_dcp_group.return_value = mock_dcp_group
+
+            mock_pcp_group = MagicMock(spec=GroupCoordinator)
+            mock_pcp_group.world_size = pcp_world_size
+            mock_pcp_group.rank_in_group = pcp_rank
+            mock_get_pcp_group.return_value = mock_pcp_group
+
+            from vllm_ascend._310p.block_table import BlockTable
+
+            return BlockTable(
+                block_size=self.block_size,
+                max_num_reqs=self.max_num_reqs,
+                max_num_blocks_per_req=self.max_num_blocks_per_req,
+                max_num_batched_tokens=self.max_num_batched_tokens,
+                pin_memory=self.pin_memory,
+                device=self.device,
+                kernel_sizes=self.kernel_sizes,
+                cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
+                num_speculative_tokens=0,
+            )
+
+    def _create_multi_group_block_table(
+        self,
+        dcp_world_size,
+        dcp_rank,
+        pcp_world_size,
+        pcp_rank,
+        cp_kv_cache_interleave_size,
+        block_sizes=None,
+        max_num_blocks=None,
+        kernel_sizes=None,
+    ):
+        block_sizes = block_sizes or [self.block_size]
+        max_num_blocks = max_num_blocks or [self.max_num_blocks_per_req] * len(block_sizes)
+        kernel_sizes = kernel_sizes or [[self.block_size]] * len(block_sizes)
+
+        with (
+            patch("vllm_ascend.worker.block_table.get_dcp_group") as mock_get_dcp_group,
+            patch("vllm_ascend.worker.block_table.get_pcp_group") as mock_get_pcp_group,
+        ):
+            mock_dcp_group = MagicMock(spec=GroupCoordinator)
+            mock_dcp_group.world_size = dcp_world_size
+            mock_dcp_group.rank_in_group = dcp_rank
+            mock_get_dcp_group.return_value = mock_dcp_group
+
+            mock_pcp_group = MagicMock(spec=GroupCoordinator)
+            mock_pcp_group.world_size = pcp_world_size
+            mock_pcp_group.rank_in_group = pcp_rank
+            mock_get_pcp_group.return_value = mock_pcp_group
+
+            from vllm_ascend._310p.block_table import MultiGroupBlockTable
+
+            return MultiGroupBlockTable(
+                max_num_reqs=self.max_num_reqs,
+                max_model_len=self.block_size * self.max_num_blocks_per_req,
+                max_num_batched_tokens=self.max_num_batched_tokens,
+                pin_memory=self.pin_memory,
+                device=self.device,
+                block_sizes=block_sizes,
+                max_num_blocks=max_num_blocks,
+                kernel_sizes=kernel_sizes,
+                cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
+            )
+
+    @staticmethod
+    def _setup_block_table_data(block_table, num_reqs=2):
+        for i in range(num_reqs):
+            block_ids = list(range(i * 4, (i + 1) * 4))
+            block_table.add_row(block_ids, i)
+
+    def test_compute_slot_mapping_with_query_start_loc_signature(self):
+        block_table = self._create_block_table(
+            dcp_world_size=1,
+            dcp_rank=0,
+            pcp_world_size=1,
+            pcp_rank=0,
+            cp_kv_cache_interleave_size=1,
+        )
+        self._setup_block_table_data(block_table, num_reqs=2)
+
+        query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
+        positions = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
+
+        block_table.compute_slot_mapping(2, query_start_loc, positions)
+
+        expected = np.array([0, 1, 512, 513], dtype=np.int32)
+        np.testing.assert_array_equal(block_table.slot_mapping.np[:4], expected)
+        np.testing.assert_array_equal(block_table.slot_mapping.gpu[:4].cpu().numpy(), expected)
+
+    def test_multi_group_compute_slot_mapping_accepts_none_compressed_args(self):
+        multi_group_block_table = self._create_multi_group_block_table(
+            dcp_world_size=1,
+            dcp_rank=0,
+            pcp_world_size=1,
+            pcp_rank=0,
+            cp_kv_cache_interleave_size=1,
+        )
+        self._setup_block_table_data(multi_group_block_table[0], num_reqs=2)
+
+        query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
+        positions = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
+
+        multi_group_block_table.compute_slot_mapping(2, query_start_loc, positions, None, None)
+
+        expected = np.array([0, 1, 512, 513], dtype=np.int32)
+        np.testing.assert_array_equal(multi_group_block_table[0].slot_mapping.np[:4], expected)
+        np.testing.assert_array_equal(multi_group_block_table[0].slot_mapping.gpu[:4].cpu().numpy(), expected)
+
+    def test_multi_group_compute_slot_mapping_uses_compressed_inputs_per_group(self):
+        multi_group_block_table = self._create_multi_group_block_table(
+            dcp_world_size=1,
+            dcp_rank=0,
+            pcp_world_size=1,
+            pcp_rank=0,
+            cp_kv_cache_interleave_size=1,
+            block_sizes=[self.block_size, self.block_size],
+            max_num_blocks=[self.max_num_blocks_per_req, self.max_num_blocks_per_req],
+            kernel_sizes=[[self.block_size], [self.block_size]],
+        )
+        for block_table in multi_group_block_table.block_tables:
+            self._setup_block_table_data(block_table, num_reqs=2)
+
+        query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
+        positions = torch.tensor([0, 1, 0, 1], dtype=torch.int64)
+        positions_compressed_list = [
+            np.array([0, 1], dtype=np.int64),
+            np.array([0], dtype=np.int64),
+        ]
+        req_indices_compressed_list = [
+            np.array([0, 0], dtype=np.int64),
+            np.array([1], dtype=np.int64),
+        ]
+
+        multi_group_block_table.compute_slot_mapping(
+            2,
+            query_start_loc,
+            positions,
+            positions_compressed_list,
+            req_indices_compressed_list,
+        )
+
+        np.testing.assert_array_equal(
+            multi_group_block_table[0].slot_mapping.np[:2],
+            np.array([0, 1], dtype=np.int32),
+        )
+        np.testing.assert_array_equal(
+            multi_group_block_table[1].slot_mapping.np[:1],
+            np.array([512], dtype=np.int32),
+        )
+
+    def test_compute_slot_mapping_with_req_indices_signature(self):
+        block_table = self._create_block_table(
+            dcp_world_size=4,
+            dcp_rank=0,
+            pcp_world_size=2,
+            pcp_rank=0,
+            cp_kv_cache_interleave_size=1,
+        )
+        self._setup_block_table_data(block_table, num_reqs=1)
+
+        req_indices = np.zeros(16, dtype=np.int32)
+        positions = np.arange(16, dtype=np.int32)
+
+        block_table.compute_slot_mapping(req_indices, positions)
+
+        expected = np.array([0, -1, -1, -1, -1, -1, -1, -1, 1, -1, -1, -1, -1, -1, -1, -1], dtype=np.int32)
+        np.testing.assert_array_equal(block_table.slot_mapping.np[:16], expected)
+        np.testing.assert_array_equal(block_table.slot_mapping.gpu[:16].cpu().numpy(), expected)
+
+    def test_compute_slot_mapping_rejects_device_tensor_inputs(self):
+        block_table = self._create_block_table(
+            dcp_world_size=1,
+            dcp_rank=0,
+            pcp_world_size=1,
+            pcp_rank=0,
+            cp_kv_cache_interleave_size=1,
+        )
+        self._setup_block_table_data(block_table, num_reqs=2)
+
+        req_indices = np.array([0, 0, 1, 1], dtype=np.int64)
+        device_positions = torch.empty(4, dtype=torch.int64, device="meta")
+
+        with self.assertRaisesRegex(TypeError, "D2H"):
+            block_table.compute_slot_mapping(req_indices, device_positions)
+
+        device_query_start_loc = torch.empty(3, dtype=torch.int32, device="meta")
+        positions = torch.arange(4, dtype=torch.int64)
+
+        with self.assertRaisesRegex(TypeError, "D2H"):
+            block_table.compute_slot_mapping(2, device_query_start_loc, positions)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/_310p/test_kv_block_zeroer_310p.py

@@ -0,0 +1,73 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from types import SimpleNamespace
+
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.kv_block_zeroer import AscendKVBlockZeroer310
+
+
+class TestAscendKVBlockZeroer310(TestBase):
+    def setUp(self):
+        self.zeroer = AscendKVBlockZeroer310(torch.device("cpu"), pin_memory=False)
+
+    def test_zero_block_ids_noop_when_empty(self):
+        kv = torch.ones(4, 2, 3)
+        self.zeroer._kv_tensors = [kv]
+        self.zeroer._logical_page_ratio = 1
+
+        self.zeroer.zero_block_ids([])
+        self.assertTrue(torch.all(kv == 1))
+
+    def test_zero_block_ids_zeros_target_slices(self):
+        kv = torch.ones(6, 2, 3)
+        self.zeroer._kv_tensors = [kv]
+        self.zeroer._logical_page_ratio = 2
+
+        self.zeroer.zero_block_ids([1])
+
+        self.assertTrue(torch.all(kv[:2] == 1))
+        self.assertTrue(torch.all(kv[2:4] == 0))
+        self.assertTrue(torch.all(kv[4:] == 1))
+
+    def test_init_meta_deduplicates_kv_pointers(self):
+        k_cache = torch.zeros(4, 2, 3)
+        v_cache = k_cache
+        layer_context = SimpleNamespace(kv_cache=(k_cache, v_cache))
+        spec = FullAttentionSpec(
+            block_size=128,
+            num_kv_heads=2,
+            head_size=64,
+            dtype=torch.float16,
+        )
+        group = SimpleNamespace(
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+            layer_names=["layer_0"],
+        )
+
+        self.zeroer.init_meta(
+            attn_groups_iter=[group],
+            kernel_block_sizes=[[64]],
+            cache_dtype="float16",
+            runner_only_attn_layers=set(),
+            static_forward_context={"layer_0": layer_context},
+        )
+
+        self.assertEqual(len(self.zeroer._kv_tensors), 1)
+        self.assertEqual(self.zeroer._logical_page_ratio, 2)

tests/ut/_310p/test_mamba_align_fallback_310p_source.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Source-level regressions for the 310P Mamba align fallback.
+
+The fallback is only active on 310P and depends on runtime NPU/vLLM state.
+Keep these checks import-free so they can run in lightweight DT environments
+while still guarding the important upstream semantic contract.
+"""
+
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[3]
+PATCH_MAMBA_UTILS = ROOT / "vllm_ascend" / "patch" / "worker" / "patch_mamba_utils.py"
+
+
+def _func(path: Path, name: str) -> ast.FunctionDef:
+    for node in ast.parse(path.read_text()).body:
+        if isinstance(node, ast.FunctionDef) and node.name == name:
+            return node
+    raise AssertionError(f"function {name} not found in {path}")
+
+
+def _src(node: ast.AST) -> str:
+    return ast.unparse(node)
+
+
+def test_310p_postprocess_fallback_preserves_upstream_metadata_semantics() -> None:
+    src = _src(_func(PATCH_MAMBA_UTILS, "_postprocess_mamba_align_gpu_cpu_fallback"))
+
+    assert "num_accepted_tokens_gpu" in src
+    assert "num_accepted_tokens_cpu_tensor[:num_reqs].copy_(num_accepted_tokens_gpu[:num_reqs])" in src
+    assert "num_tokens_running_state = num_computed_tokens[i] + num_scheduled_tokens[i] - num_draft_tokens[i]" in src
+    assert "new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens[i] - 1" in src
+    assert "aligned_new_computed_tokens = new_num_computed_tokens // block_size * block_size" in src
+    assert "if aligned_new_computed_tokens < num_tokens_running_state:" in src
+    assert "if src_block_idx == dest_block_idx:" in src
+    assert "num_accepted_tokens_cpu_tensor[i] = 1" in src
+
+
+def test_310p_postprocess_fallback_mirrors_state_copy_without_triton() -> None:
+    src = _src(_func(PATCH_MAMBA_UTILS, "_postprocess_mamba_align_gpu_cpu_fallback"))
+
+    assert "run_fused_postprocess" not in src
+    assert "postprocess_mamba_fused_kernel" not in src
+    assert "accept_token_bias = aligned_new_computed_tokens - num_tokens_running_state" in src
+    assert "if accept_token_bias == 0:" in src
+    assert "continue" in src
+    assert "for mamba_group_id in ctx.mamba_group_ids:" in src
+    assert "get_numpy_array()" in src
+    assert "copy_spec = state_copy_func(state, block_ids, src_block_idx, accept_token_bias + 1)" in src
+    assert "_tensor_view_from_data_ptr(state, copy_spec.start_addr, copy_spec.num_elements)" in src
+    assert "dst_state.copy_(src_state.clone())" in src

tests/ut/_310p/test_model_runner_310p.py

@@ -0,0 +1,143 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.config import CUDAGraphMode
+from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.model_runner_310p import NPUModelRunner310
+
+
+def _prepare_inputs_source() -> str:
+    source_path = Path(__file__).resolve().parents[3] / "vllm_ascend" / "_310p" / "model_runner_310p.py"
+    source = source_path.read_text(encoding="utf-8")
+    start = source.index("    def _prepare_inputs(")
+    end = source.index("    @torch.inference_mode()", start)
+    return source[start:end]
+
+
+def test_prepare_inputs_keeps_aclgraph_metadata_on_cpu() -> None:
+    source = _prepare_inputs_source()
+
+    assert "block_table.compute_slot_mapping(" in source
+    assert "req_indices," in source
+    assert "positions_np[:total_num_scheduled_tokens]" in source
+
+    assert "self.input_batch.block_table.compute_slot_mapping(" not in source
+    assert "query_start_loc.gpu[: num_reqs + 1]" not in source
+    assert "req_indices_gpu" not in source
+    assert "self.num_computed_tokens[req_indices_gpu]" not in source
+
+    assert "self.positions[:total_num_scheduled_tokens].copy_(" in source
+    assert "self._positions_cpu_buf[:total_num_scheduled_tokens]" in source
+    assert "self.seq_lens[:num_reqs].copy_(" in source
+    assert "self.optimistic_seq_lens_cpu[:num_reqs]" in source
+
+
+def test_model_forward_updates_mtp_full_graph_params_before_replay() -> None:
+    runner = object.__new__(NPUModelRunner310)
+    runner.uses_mrope = False
+    runner.enable_enpu = False
+    runner.speculative_config = SimpleNamespace(method="mtp")
+    runner.update_stream = MagicMock()
+    runner._all_gather_hidden_states_and_aux = MagicMock()
+
+    calls = []
+
+    def fake_update(*args):
+        calls.append("update")
+
+    def fake_model(**kwargs):
+        calls.append("model")
+        return torch.ones(1)
+
+    runner.model = fake_model
+    runner._update_full_graph_params_if_needed = fake_update
+    forward_context = SimpleNamespace(
+        cudagraph_runtime_mode=CUDAGraphMode.FULL,
+        capturing=False,
+        flash_comm_v1_enabled=False,
+    )
+
+    with patch(
+        "vllm_ascend._310p.model_runner_310p.get_forward_context",
+        return_value=forward_context,
+    ):
+        hidden_states = runner._model_forward(
+            8,
+            input_ids=torch.tensor([1]),
+            positions=torch.tensor([0]),
+        )
+
+    assert calls == ["update", "model"]
+    torch.testing.assert_close(hidden_states, torch.ones(1))
+
+
+class TestNPUModelRunner310(TestBase):
+    def test_may_reinitialize_input_batch_expands_prefix_mamba_block_table(self):
+        runner = object.__new__(NPUModelRunner310)
+        runner.max_num_reqs = 8
+        runner.max_model_len = 512
+        runner.max_encoder_len = 0
+        runner.max_num_tokens = 1024
+        runner.device = torch.device("cpu")
+        runner.pin_memory = False
+        runner.is_pooling_model = False
+        runner.model_config = SimpleNamespace(max_model_len=512, get_vocab_size=lambda: 32000)
+        runner.cache_config = SimpleNamespace(block_size=128, enable_prefix_caching=True)
+        runner.parallel_config = SimpleNamespace(cp_kv_cache_interleave_size=4)
+        runner.vllm_config = SimpleNamespace(speculative_config=None)
+        runner.offload_config = SimpleNamespace(uva=SimpleNamespace(cpu_offload_gb=0))
+        runner.input_batch = SimpleNamespace(logitsprocs=MagicMock())
+        attention_backend = SimpleNamespace(get_supported_kernel_block_sizes=lambda: [128, 64])
+        runner.attn_groups = [[SimpleNamespace(backend=attention_backend)]]
+
+        attention_spec = AttentionSpec(
+            block_size=128,
+            num_kv_heads=2,
+            head_size=64,
+            dtype=torch.float16,
+        )
+        mamba_spec = MambaSpec(
+            block_size=128,
+            shapes=((16,),),
+            dtypes=(torch.float16,),
+            mamba_cache_mode="align",
+            num_speculative_blocks=2,
+        )
+        kv_cache_config = SimpleNamespace(
+            kv_cache_groups=[
+                SimpleNamespace(kv_cache_spec=attention_spec),
+                SimpleNamespace(kv_cache_spec=mamba_spec),
+            ]
+        )
+
+        with (
+            patch("vllm_ascend._310p.model_runner_310p.NPUInputBatch") as mock_input_batch,
+            patch("vllm_ascend._310p.model_runner_310p.get_total_cp_world_size", return_value=1),
+        ):
+            runner.may_reinitialize_input_batch(kv_cache_config)
+
+        kwargs = mock_input_batch.call_args.kwargs
+        self.assertEqual(kwargs["block_sizes"], [128, 128])
+        self.assertEqual(kwargs["kernel_block_sizes"], [[128, 64], [0]])
+        self.assertEqual(kwargs["max_num_blocks_per_req"], [4, 6])
+        self.assertIs(kwargs["kv_cache_groups"], kv_cache_config.kv_cache_groups)
+        self.assertEqual(kwargs["cp_kv_cache_interleave_size"], 4)

tests/ut/_310p/test_sharded_state_loader_310p.py

@@ -0,0 +1,134 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend._310p.sharded_state_loader_310p import ShardedStateLoader310
+
+
+class MockQuantConfig:
+    """Mock quantization config for testing."""
+
+    def __init__(self, quant_type: str = "FLOAT"):
+        self.quant_description = {"model_quant_type": quant_type}
+
+
+class MockModel(torch.nn.Module):
+    """Mock model for testing."""
+
+    def __init__(self, quant_config=None, with_int_weights: bool = False):
+        super().__init__()
+        self.quant_config = quant_config
+        self.with_int_weights = with_int_weights
+        if with_int_weights:
+            self.linear = torch.nn.Linear(10, 10)
+            self.linear.weight = torch.nn.Parameter(
+                torch.randint(-127, 127, (10, 10), dtype=torch.int8), requires_grad=False
+            )
+            self.linear.bias = torch.nn.Parameter(torch.zeros(10, dtype=torch.int32), requires_grad=False)
+        else:
+            self.linear = torch.nn.Linear(10, 10)
+
+
+class TestShardedStateLoader310(TestBase):
+    """Test cases for ShardedStateLoader310."""
+
+    @patch("vllm.model_executor.model_loader.ShardedStateLoader._filter_subtensors")
+    @patch("vllm.distributed.get_tensor_model_parallel_rank")
+    @patch("safetensors.torch.save_file")
+    def test_save_model_with_nd_format_310(self, mock_save_file, mock_get_rank, mock_filter):
+        """Test save_model with ND format tensors (no conversion needed)."""
+        mock_get_rank.return_value = 0
+        mock_filter.side_effect = lambda x: x
+        mock_tensor = MagicMock(spec=torch.Tensor)
+
+        model = MockModel()
+        with (
+            patch.object(model, "state_dict", return_value={"linear.weight": mock_tensor}),
+            tempfile.TemporaryDirectory() as tmpdir,
+        ):
+            ShardedStateLoader310.save_model(model, tmpdir)
+
+            mock_save_file.assert_called_once()
+
+    @patch("vllm.model_executor.model_loader.ShardedStateLoader._filter_subtensors")
+    def test_generate_quant_description_float_model_310(self, mock_filter):
+        """Test generate_quant_description for float model."""
+        mock_filter.side_effect = lambda x: x
+        quant_config = MockQuantConfig(quant_type="FLOAT")
+        model = MockModel(quant_config=quant_config, with_int_weights=False)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ShardedStateLoader310.generate_quant_description(model, tmpdir, quant_config)
+
+            json_path = Path(tmpdir) / "parameters_type_map.json"
+            self.assertTrue(json_path.exists())
+
+            with open(json_path, encoding="utf-8") as f:
+                quant_description = json.load(f)
+
+            self.assertEqual(quant_description["model_quant_type"], "FLOAT")
+            self.assertEqual(quant_description["version"], "1.0.0")
+            self.assertIn("linear.weight", quant_description)
+            self.assertEqual(quant_description["linear.weight"], "FLOAT")
+            self.assertIn("linear.bias", quant_description)
+            self.assertEqual(quant_description["linear.bias"], "FLOAT")
+
+    @patch("vllm.model_executor.model_loader.ShardedStateLoader._filter_subtensors")
+    def test_generate_quant_description_no_quant_config_310(self, mock_filter):
+        """When quant_config is None, treat model as FLOAT."""
+        mock_filter.side_effect = lambda x: x
+        model = MockModel(quant_config=None, with_int_weights=False)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ShardedStateLoader310.generate_quant_description(model, tmpdir, None)
+
+            json_path = Path(tmpdir) / "parameters_type_map.json"
+            self.assertTrue(json_path.exists())
+
+            with open(json_path, encoding="utf-8") as f:
+                quant_description = json.load(f)
+
+            self.assertEqual(quant_description["model_quant_type"], "FLOAT")
+            self.assertEqual(quant_description["linear.weight"], "FLOAT")
+
+    @patch("vllm.model_executor.model_loader.ShardedStateLoader._filter_subtensors")
+    def test_generate_quant_description_int_model_310(self, mock_filter):
+        """Test generate_quant_description for int8 quantized model."""
+        mock_filter.side_effect = lambda x: x
+        quant_config = MockQuantConfig(quant_type="W8A8")
+        model = MockModel(quant_config=quant_config, with_int_weights=True)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            ShardedStateLoader310.generate_quant_description(model, tmpdir, quant_config)
+
+            json_path = Path(tmpdir) / "parameters_type_map.json"
+            self.assertTrue(json_path.exists())
+
+            with open(json_path, encoding="utf-8") as f:
+                quant_description = json.load(f)
+
+            self.assertEqual(quant_description["model_quant_type"], "W8A8")
+            self.assertEqual(quant_description["version"], "1.0.0")
+            self.assertIn("linear.weight", quant_description)
+            self.assertEqual(quant_description["linear.weight"], "W8A8")
+            self.assertIn("linear.bias", quant_description)
+            self.assertEqual(quant_description["linear.bias"], "W8A8")

tests/ut/_fake_weight/config.json

@@ -0,0 +1,28 @@
+{
+  "_name_or_path": "facebook/opt-125m",
+  "activation_dropout": 0.0,
+  "activation_function": "relu",
+  "architectures": [
+    "OPTForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "do_layer_norm_before": true,
+  "dropout": 0.1,
+  "eos_token_id": 2,
+  "ffn_dim": 3072,
+  "hidden_size": 768,
+  "init_std": 0.02,
+  "layerdrop": 0.0,
+  "max_position_embeddings": 2048,
+  "model_type": "opt",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "prefix": "</s>",
+  "torch_dtype": "float16",
+  "transformers_version": "4.21.0.dev0",
+  "use_cache": true,
+  "vocab_size": 50272,
+  "word_embed_proj_dim": 768
+}

tests/ut/_tools/test_ai_qos_tool.py

@@ -0,0 +1,300 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import importlib.util
+import json
+import sys
+from io import StringIO
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+TOOL_PATH = REPO_ROOT / "tools" / "ai_qos.py"
+MODULE_NAME = "vllm_ascend_tools_ai_qos"
+
+MASTER_IDS = (11, 12, 13, 7)
+
+
+def _load_ai_qos_tool(mock_ai: MagicMock | None = None):
+    if mock_ai is None:
+        mock_ai = MagicMock()
+
+        def get_qos_fn(device_id, master_id):
+            return (0, master_id, 42, 0, 0, 0)
+
+        mock_ai.get_qos.side_effect = get_qos_fn
+        mock_ai.get_bw.return_value = (0, 1, 2, 0)
+        mock_ai.get_fuse_mode.return_value = (0, 1, 1, 0)
+        mock_ai.set_bw.return_value = 0
+        mock_ai.set_qos.return_value = 0
+        mock_ai.set_fuse_gbl_config.return_value = 0
+
+    ascend = MagicMock()
+    ascend.ai_qos = mock_ai
+
+    sys.modules.pop(MODULE_NAME, None)
+    with patch.dict(
+        sys.modules,
+        {
+            "vllm_ascend": ascend,
+        },
+    ):
+        spec = importlib.util.spec_from_file_location(MODULE_NAME, TOOL_PATH)
+        assert spec and spec.loader
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
+    return mod, mock_ai
+
+
+def test_device_list_uses_all_visible_devices_when_env_unset(monkeypatch):
+    monkeypatch.delenv("ASCEND_RT_VISIBLE_DEVICES", raising=False)
+    mod, _ = _load_ai_qos_tool()
+    mock_torch = MagicMock()
+    mock_torch.npu.device_count.return_value = 4
+    with patch.dict(sys.modules, {"torch": mock_torch}):
+        assert mod._device_list() == [0, 1, 2, 3]
+
+
+def test_device_list_exits_when_env_unset_and_torch_query_fails(monkeypatch):
+    monkeypatch.delenv("ASCEND_RT_VISIBLE_DEVICES", raising=False)
+    mod, _ = _load_ai_qos_tool()
+    mock_torch = MagicMock()
+    mock_torch.npu.device_count.side_effect = RuntimeError("query failed")
+    with patch.dict(sys.modules, {"torch": mock_torch}), pytest.raises(SystemExit) as e:
+        mod._device_list()
+    assert e.value.code == 1
+
+
+def test_device_list_parses_visible_devices(monkeypatch):
+    mod, _ = _load_ai_qos_tool()
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0,2")
+    assert mod._device_list() == [0, 2]
+
+
+def test_device_list_parses_single_id(monkeypatch):
+    mod, _ = _load_ai_qos_tool()
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "3")
+    assert mod._device_list() == [3]
+
+
+def test_print_config_block(capsys):
+    mod, _ = _load_ai_qos_tool()
+    mod._print_config_block(["line a", "line b"])
+    out = capsys.readouterr().out
+    assert "system-view" in out and "line a" in out and "line b" in out
+    assert out.strip().endswith("commit")
+
+
+def test_load_first_apply_baseline_no_file(tmp_path):
+    mod, _ = _load_ai_qos_tool()
+    p = tmp_path / "missing.json"
+    assert mod._load_first_apply_baseline(p) is None
+
+
+def test_load_first_apply_baseline_malformed_json(tmp_path):
+    mod, _ = _load_ai_qos_tool()
+    p = tmp_path / "x.json"
+    p.write_text("{", encoding="utf-8")
+    assert mod._load_first_apply_baseline(p) is None
+
+
+def test_load_first_apply_baseline_invalid_original_qos(tmp_path):
+    mod, _ = _load_ai_qos_tool()
+    p = tmp_path / "x.json"
+    p.write_text(json.dumps({"original_qos": "bad"}), encoding="utf-8")
+    assert mod._load_first_apply_baseline(p) is None
+
+
+def test_load_first_apply_baseline_success(tmp_path):
+    mod, _ = _load_ai_qos_tool()
+    p = tmp_path / "x.json"
+    body = {
+        "original_qos": {"0": {"7": [7, 0, 0, 0, 0]}},
+        "original_sdma_mata": {"0": [0, 1, 2, 0]},
+        "original_fuse": {"0": [1, 1, 0]},
+    }
+    p.write_text(json.dumps(body), encoding="utf-8")
+    b = mod._load_first_apply_baseline(p)
+    assert b is not None
+    oq, osm, ofu = b
+    assert oq == body["original_qos"]
+    assert osm == body["original_sdma_mata"]
+    assert ofu == body["original_fuse"]
+
+
+def test_run_unset_exits_without_state_file(capsys):
+    mod, _ = _load_ai_qos_tool()
+    with pytest.raises(SystemExit) as e:
+        mod.run_unset(Path("/nonexistent/ai_qos_state.json"))
+    assert e.value.code == 1
+    err = capsys.readouterr().err
+    assert "No state file" in err
+
+
+def test_run_unset_parse_failed_bad_json_deletes_file(tmp_path, capsys):
+    mod, _ = _load_ai_qos_tool()
+    state = tmp_path / "ai_qos_state.json"
+    state.write_text("{", encoding="utf-8")
+    with pytest.raises(SystemExit) as e:
+        mod.run_unset(state)
+    assert e.value.code == 1
+    assert not state.is_file()
+    err = capsys.readouterr().err
+    assert "Failed to parse the state file." in err
+
+
+def test_run_unset_parse_failed_invalid_structure_deletes_file(tmp_path, capsys):
+    mod, _ = _load_ai_qos_tool()
+    state = tmp_path / "ai_qos_state.json"
+    state.write_text(
+        json.dumps(
+            {
+                "original_qos": {},
+                "printed_commands": [123],
+                "original_sdma_mata": {},
+                "original_fuse": {},
+            }
+        ),
+        encoding="utf-8",
+    )
+    with pytest.raises(SystemExit) as e:
+        mod.run_unset(state)
+    assert e.value.code == 1
+    assert not state.is_file()
+    assert "Failed to parse the state file." in capsys.readouterr().err
+
+
+def test_run_unset_restores_and_deletes_file(tmp_path):
+    mod, mock_ai = _load_ai_qos_tool()
+    state = tmp_path / "ai_qos_state.json"
+    data = {
+        "original_qos": {
+            "0": {
+                str(MASTER_IDS[0]): [11, 1, 0, 0, 0],
+            }
+        },
+        "original_sdma_mata": {"0": [0, 1, 2, 0]},
+        "original_fuse": {"0": [1, 0, 0]},
+        "printed_commands": ["hccs qos remap 1 0 0"],
+    }
+    state.write_text(json.dumps(data), encoding="utf-8")
+    out_buf = StringIO()
+    with patch("sys.stdout", out_buf):
+        mod.run_unset(state)
+    assert not state.is_file()
+    assert mock_ai.set_bw.called
+    assert mock_ai.set_qos.called
+    assert mock_ai.set_fuse_gbl_config.called
+    uo = out_buf.getvalue()
+    assert "system-view" in uo
+    assert "undo hccs qos remap 1 0 0" in uo
+    assert "commit" in uo
+
+
+def test_AiqosConfig_set_qos_captures_and_writes_state(tmp_path, monkeypatch):
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0")
+    mod, mock_ai = _load_ai_qos_tool()
+    cfg = {
+        "mode": "auto",
+        "aiqos_priority": {
+            "AIV_D2D": "high",
+            "AIV_H2D": "high",
+            "SDMA_D2D": "high",
+            "SDMA_H2D": "low",
+            "PCIEDMA_H2D": "high",
+        },
+    }
+    out_buf = StringIO()
+    with patch("sys.stdout", out_buf):
+        mod.AiqosConfig(cfg).set_qos(tmp_path / "state.json")
+    state = tmp_path / "state.json"
+    assert state.is_file()
+    j = json.loads(state.read_text(encoding="utf-8"))
+    assert "original_qos" in j and "printed_commands" in j
+    assert mock_ai.get_qos.call_count == 9
+    assert mock_ai.set_fuse_gbl_config.called
+    assert mock_ai.get_fuse_mode.called
+    assert mock_ai.get_bw.called
+    assert mock_ai.set_bw.called
+
+
+def test_AiqosConfig_second_apply_reuses_baseline_fewer_capture_qos(tmp_path, monkeypatch):
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0")
+    mod, mock_ai = _load_ai_qos_tool()
+    body = {
+        "original_qos": {
+            "0": {str(m): [m, 1, 0, 0, 0] for m in MASTER_IDS},
+        },
+        "original_sdma_mata": {"0": [0, 1, 1, 0]},
+        "original_fuse": {"0": [1, 1, 0]},
+        "printed_commands": [],
+    }
+    p = tmp_path / "s.json"
+    p.write_text(json.dumps(body), encoding="utf-8")
+    mock_ai.reset_mock()
+    cfg = {
+        "mode": "auto",
+        "aiqos_priority": {
+            "AIV_D2D": "high",
+            "AIV_H2D": "high",
+            "SDMA_D2D": "high",
+            "SDMA_H2D": "low",
+            "PCIEDMA_H2D": "high",
+        },
+    }
+    with patch("sys.stdout", StringIO()):
+        mod.AiqosConfig(cfg).set_qos(p)
+    n_with_baseline = mock_ai.get_qos.call_count
+    assert n_with_baseline == 4, "apply loop only: 4 masters, no capture get_qos"
+    p.unlink()
+    mock_ai.reset_mock()
+    with patch("sys.stdout", StringIO()):
+        mod.AiqosConfig(cfg).set_qos(p)
+    n_cold = mock_ai.get_qos.call_count
+    assert n_cold == 9, "1 dev cold: capture 4 + SDMA 1 + apply 4 = 9"
+
+
+def test_AiqosConfig_merges_baseline_when_device_list_grows(tmp_path, monkeypatch):
+    """Second apply with more NPU ids than the first-apply state must save baseline for new ids."""
+    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "0,1")
+    mod, _ = _load_ai_qos_tool()
+    body = {
+        "original_qos": {
+            "0": {str(m): [m, 1, 0, 0, 0] for m in MASTER_IDS},
+        },
+        "original_sdma_mata": {"0": [0, 1, 1, 0]},
+        "original_fuse": {"0": [1, 1, 0]},
+        "printed_commands": [],
+    }
+    p = tmp_path / "state.json"
+    p.write_text(json.dumps(body), encoding="utf-8")
+    cfg = {
+        "mode": "auto",
+        "aiqos_priority": {
+            "AIV_D2D": "high",
+            "AIV_H2D": "high",
+            "SDMA_D2D": "high",
+            "SDMA_H2D": "low",
+            "PCIEDMA_H2D": "high",
+        },
+    }
+    with patch("sys.stdout", StringIO()):
+        mod.AiqosConfig(cfg).set_qos(p)
+    j = json.loads(p.read_text(encoding="utf-8"))
+    assert "0" in j["original_qos"] and "1" in j["original_qos"]
+    assert "0" in j["original_sdma_mata"] and "1" in j["original_sdma_mata"]
+    assert "0" in j["original_fuse"] and "1" in j["original_fuse"]

tests/ut/_tools/test_docs_codegen.py

@@ -0,0 +1,584 @@
+from __future__ import annotations
+
+from io import StringIO
+from pathlib import Path
+from textwrap import dedent
+
+import pytest
+
+from tools.docs_codegen.cli import main
+from tools.docs_codegen.converters import RUN_DP_TEMPLATE_POSITIONALS
+from tools.docs_codegen.errors import DocsCodegenError
+from tools.docs_codegen.generator import GeneratorService
+from tools.docs_codegen.scanner import BlockScanner
+from tools.docs_codegen.utils import substitute_template_positionals
+
+
+def _write_text(path: Path, content: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(dedent(content).lstrip(), encoding="utf-8")
+
+
+def _write_single_node_yaml(tmp_path: Path) -> Path:
+    case_path = tmp_path / "cases" / "single_node.yaml"
+    _write_text(
+        case_path,
+        """
+        test_cases:
+          - name: default-case
+            model: default/model
+            envs:
+              SERVER_PORT: 8123
+            server_cmd: []
+          - name: selected-case
+            model: "Qwen/Test Model"
+            envs:
+              HCCL_BUFFSIZE: 1024
+              PROMPT: "hello world"
+              SERVER_PORT: DEFAULT_PORT
+            server_cmd:
+              - "--tensor-parallel-size"
+              - 2
+              - "--kv-transfer-config"
+              - '{"foo": "bar", "enabled": true}'
+            server_cmd_extra: "--trust-remote-code --enable-expert-parallel"
+        """,
+    )
+    return case_path.relative_to(tmp_path)
+
+
+def _write_multi_node_yaml(tmp_path: Path, *, invalid_command: bool = False) -> Path:
+    case_path = tmp_path / "cases" / "multi_node.yaml"
+    second_command = (
+        '"vllm serve multi-node/model extra-positional"'
+        if invalid_command
+        else """
+              - vllm
+              - serve
+              - multi-node/model
+              - "--headless"
+              - "--port"
+              - "$SERVER_PORT"
+        """
+    )
+    _write_text(
+        case_path,
+        f"""
+        deployment:
+          - envs:
+              LOCAL_IP: 127.0.0.1
+            server_cmd: "vllm serve first-host --port 8000"
+          - envs:
+              MASTER_IP: 10.0.0.1
+              SERVER_PORT: 9000
+            server_cmd: {second_command.rstrip()}
+        """,
+    )
+    return case_path.relative_to(tmp_path)
+
+
+def _write_model_code_doc(tmp_path: Path, content: str, *, name: str = "Demo.md") -> Path:
+    doc_path = tmp_path / "docs" / "models" / name
+    _write_text(doc_path, content)
+    return doc_path.relative_to(tmp_path)
+
+
+def _generate_block(tmp_path: Path, doc_path: Path, block_name: str) -> str:
+    service = GeneratorService(artifact_root="artifacts")
+    return service.generate_block(doc_path, block_name, dry_run=True)[1].content
+
+
+# Explicit CPU smart-UT routing keeps this guard out of the --run-all-cpu bucket.
+def test_block_scanner_parses_metadata_and_trims_raw_block(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        # Demo
+
+        ```{{model-code}}
+        :block_name: single
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        :case_index: 1
+
+        set -eux
+        {{{{ generated }}}}
+
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    blocks = BlockScanner().scan_document_blocks(doc_path)
+
+    assert len(blocks) == 1
+    block = blocks[0]
+    assert block.doc_path == doc_path
+    assert block.block_name == "single"
+    assert block.converter_tag == "single_node"
+    assert block.test_case_path == single_yaml.as_posix()
+    assert block.extra_options == (("case_index", "1"),)
+    assert block.directive_line == 3
+    assert block.raw_block_lines == ("set -eux", "{{ generated }}")
+
+
+def test_block_scanner_rejects_duplicate_block_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: serve
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        ```
+
+        ```{{model-code}}
+        :block_name: serve
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        ```
+        """,
+        name="Duplicate.md",
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        BlockScanner().scan_document_blocks(doc_path)
+
+    error_message = str(exc_info.value)
+    assert "docs/models/Duplicate.md:7: model-code generation error" in error_message
+    assert "block_name: serve" in error_message
+    assert "duplicated block_name 'serve'" in error_message
+    assert "previous declaration is on line 1" in error_message
+
+
+def test_block_scanner_rejects_unsupported_metadata(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: serve
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        :unknown_option: value
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        BlockScanner().scan_document_blocks(doc_path)
+
+    assert "unsupported metadata: unknown_option" in str(exc_info.value)
+
+
+def test_single_node_converter_uses_case_index_defaults_and_extra_args(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: single
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        :case_index: 1
+
+        set -eux
+        {{{{ generated }}}}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    script = _generate_block(tmp_path, doc_path, "single")
+
+    assert script.startswith("set -eux\nexport HCCL_BUFFSIZE=1024")
+    assert 'export PROMPT="hello world"' in script
+    assert "export SERVER_PORT=8000" in script
+    assert "ignored/model" not in script
+    assert "vllm serve 'Qwen/Test Model' \\" in script
+    assert "--tensor-parallel-size 2 \\" in script
+    assert '"foo": "bar"' in script
+    assert '"enabled": true' in script
+    assert "--trust-remote-code \\" in script
+    assert "--enable-expert-parallel" in script
+
+
+def test_single_node_converter_defaults_to_first_case_and_preserves_port(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: default_case
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    script = _generate_block(tmp_path, doc_path, "default_case")
+
+    assert script == "export SERVER_PORT=8123\n\nvllm serve default/model\n"
+
+
+def test_single_node_converter_reports_invalid_case_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: missing_case
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        :case_index: 3
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        _generate_block(tmp_path, doc_path, "missing_case")
+
+    assert "case_index 3 is out of range for 'test_cases' with 2 items" in str(exc_info.value)
+
+
+def test_multi_node_converter_uses_host_index_and_token_list(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    multi_yaml = _write_multi_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: worker
+        :converter_tag: multi_node
+        :test_case_path: {multi_yaml}
+        :host_index: 1
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    script = _generate_block(tmp_path, doc_path, "worker")
+
+    assert script.startswith("export MASTER_IP=10.0.0.1\nexport SERVER_PORT=9000")
+    assert "vllm serve multi-node/model \\" in script
+    assert "--headless \\" in script
+    assert "--port $SERVER_PORT" in script
+    assert "first-host" not in script
+
+
+def test_multi_node_converter_requires_host_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    multi_yaml = _write_multi_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: worker
+        :converter_tag: multi_node
+        :test_case_path: {multi_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        _generate_block(tmp_path, doc_path, "worker")
+
+    assert "converter_tag 'multi_node' requires host_index" in str(exc_info.value)
+
+
+def test_multi_node_converter_rejects_extra_positional_args(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    multi_yaml = _write_multi_node_yaml(tmp_path, invalid_command=True)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: worker
+        :converter_tag: multi_node
+        :test_case_path: {multi_yaml}
+        :host_index: 1
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        _generate_block(tmp_path, doc_path, "worker")
+
+    assert "unsupported positional argument 'extra-positional'" in str(exc_info.value)
+
+
+def test_generator_service_writes_selected_artifact(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: default_case
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    service = GeneratorService(artifact_root="artifacts")
+    output_path, generated_script = service.generate_block(doc_path, "default_case", dry_run=False)
+
+    assert output_path == Path("artifacts/Demo/default_case.sh")
+    assert output_path.read_text(encoding="utf-8") == generated_script.content
+    assert generated_script.content == "export SERVER_PORT=8123\n\nvllm serve default/model\n"
+
+
+def test_cli_generates_block_to_stdout(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    single_yaml = _write_single_node_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: default_case
+        :converter_tag: single_node
+        :test_case_path: {single_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+    stdout = StringIO()
+    stderr = StringIO()
+
+    exit_code = main(["--block", f"{doc_path}::default_case", "--dry-run", "--stdout"], stdout=stdout, stderr=stderr)
+
+    assert exit_code == 0
+    assert stderr.getvalue() == ""
+    assert stdout.getvalue() == (
+        "docs/_build/doc_codegen/Demo/default_case.sh\nexport SERVER_PORT=8123\n\nvllm serve default/model\n"
+    )
+
+
+def test_cli_rejects_invalid_block_reference():
+    stdout = StringIO()
+    stderr = StringIO()
+
+    exit_code = main(["--block", "docs/model.md"], stdout=stdout, stderr=stderr)
+
+    assert exit_code == 1
+    assert stdout.getvalue() == ""
+    assert "block reference must use '<doc_path>::<block_name>'" in stderr.getvalue()
+
+
+def _write_external_dp_yaml(tmp_path: Path, *, routing_type: str = "disaggregated_prefill") -> Path:
+    case_path = tmp_path / "cases" / "external_dp.yaml"
+    _write_text(
+        case_path,
+        f"""
+        model: "Eco-Tech/GLM-Test"
+        num_nodes: 2
+
+        routing:
+          type: "{routing_type}"
+          groups:
+            prefiller: [0]
+            decoder: [1]
+
+        config:
+          - node_index: 0
+            port_start: 7100
+            dp_rpc_port: 12321
+            dp_size: 4
+            dp_size_local: 2
+            dp_rank_start: 0
+            tp_size: 8
+            dp_address: "${{NODE_0_IP}}"
+          - node_index: 1
+            port_start: 7200
+            dp_rpc_port: 12321
+            dp_size: 8
+            dp_size_local: 4
+            dp_rank_start: 0
+            tp_size: 4
+            dp_address: "${{NODE_1_IP}}"
+
+        env_common: &env_common
+          HCCL_BUFFSIZE: "1024"
+          OMP_PROC_BIND: "false"
+
+        templates:
+          - node_index: 0
+            envs:
+              <<: *env_common
+              ASCEND_RT_VISIBLE_DEVICES: "${{VISIBLE_DEVICES}}"
+            server_cmd_template:
+              - --host
+              - "0.0.0.0"
+              - --port
+              - ${{PORT}}
+              - --data-parallel-size
+              - ${{DP_SIZE}}
+              - --data-parallel-rank
+              - ${{DP_RANK}}
+              - --tensor-parallel-size
+              - ${{TP_SIZE}}
+              - --profiler-config
+              - '{{"profiler":"torch","with_stack":false}}'
+              - --kv-transfer-config
+              - '{{"kv_connector": "MooncakeConnectorV1", "kv_role": "kv_producer", "kv_port": "30000"}}'
+          - node_index: 1
+            envs:
+              <<: *env_common
+              ASCEND_RT_VISIBLE_DEVICES: "${{VISIBLE_DEVICES}}"
+            server_cmd_template:
+              - --host
+              - "0.0.0.0"
+              - --port
+              - ${{PORT}}
+              - --data-parallel-size
+              - ${{DP_SIZE}}
+              - --tensor-parallel-size
+              - ${{TP_SIZE}}
+        """,
+    )
+    return case_path.relative_to(tmp_path)
+
+
+def test_substitute_template_positionals():
+    positionals = RUN_DP_TEMPLATE_POSITIONALS
+    assert substitute_template_positionals("${DP_SIZE}", positionals=positionals) == "$3"
+    assert substitute_template_positionals("--port ${PORT}", positionals=positionals) == "--port $2"
+    # Unknown braced variables and unbraced refs are left untouched.
+    assert substitute_template_positionals("${UNKNOWN}", positionals=positionals) == "${UNKNOWN}"
+    assert substitute_template_positionals("$SERVER_PORT", positionals=positionals) == "$SERVER_PORT"
+
+
+def test_external_dp_template_converter_maps_positionals(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    external_yaml = _write_external_dp_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: prefill_n0
+        :converter_tag: external_dp_template
+        :test_case_path: {external_yaml}
+        :host_index: 0
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    script = _generate_block(tmp_path, doc_path, "prefill_n0")
+
+    assert "export HCCL_BUFFSIZE=1024" in script
+    assert "export ASCEND_RT_VISIBLE_DEVICES=$1" in script
+    assert "SERVER_PORT" not in script
+    assert "vllm serve Eco-Tech/GLM-Test \\" in script
+    assert "--port $2 \\" in script
+    assert "--data-parallel-size $3 \\" in script
+    assert "--data-parallel-rank $4 \\" in script
+    assert "--tensor-parallel-size $7 \\" in script
+    # Space-free JSON values are quoted (not just whitespace-containing ones).
+    assert '--profiler-config \'{"profiler":"torch","with_stack":false}\' \\' in script
+    assert '"kv_role": "kv_producer"' in script
+    assert "${DP_SIZE}" not in script
+
+
+def test_external_dp_template_converter_requires_host_index(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    external_yaml = _write_external_dp_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: prefill_n0
+        :converter_tag: external_dp_template
+        :test_case_path: {external_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        _generate_block(tmp_path, doc_path, "prefill_n0")
+
+    assert "converter_tag 'external_dp_template' requires host_index" in str(exc_info.value)
+
+
+def test_external_dp_launch_converter_combines_all_nodes(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    external_yaml = _write_external_dp_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: launch
+        :converter_tag: external_dp_launch
+        :test_case_path: {external_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    script = _generate_block(tmp_path, doc_path, "launch")
+
+    assert script == (
+        "python launch_online_dp.py --dp-size 4 --tp-size 8 --dp-size-local 2 --dp-rank-start 0 "
+        "--dp-address ${NODE_0_IP} --dp-rpc-port 12321 --vllm-start-port 7100\n\n"
+        "python launch_online_dp.py --dp-size 8 --tp-size 4 --dp-size-local 4 --dp-rank-start 0 "
+        "--dp-address ${NODE_1_IP} --dp-rpc-port 12321 --vllm-start-port 7200\n"
+    )
+
+
+def test_external_dp_proxy_converter_expands_groups(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    external_yaml = _write_external_dp_yaml(tmp_path)
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: proxy
+        :converter_tag: external_dp_proxy
+        :test_case_path: {external_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    script = _generate_block(tmp_path, doc_path, "proxy")
+
+    assert script.startswith("python load_balance_proxy_server_example.py \\")
+    # Single-value flags stay inline; multi-value flags expand one value per line.
+    assert "  --host ${NODE_0_IP} \\" in script
+    assert "  --port 1999 \\" in script
+    assert "  --prefiller-hosts \\\n    ${NODE_0_IP} \\\n    ${NODE_0_IP} \\" in script
+    assert "  --prefiller-ports \\\n    7100 \\\n    7101 \\" in script
+    assert (
+        "  --decoder-hosts \\\n    ${NODE_1_IP} \\\n    ${NODE_1_IP} \\\n    ${NODE_1_IP} \\\n    ${NODE_1_IP} \\"
+        in script
+    )
+    assert "  --decoder-ports \\\n    7200 \\\n    7201 \\\n    7202 \\\n    7203" in script
+    assert script.rstrip().endswith("    7203")
+
+
+def test_external_dp_proxy_converter_rejects_unsupported_routing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    external_yaml = _write_external_dp_yaml(tmp_path, routing_type="generic_dp")
+    doc_path = _write_model_code_doc(
+        tmp_path,
+        f"""
+        ```{{model-code}}
+        :block_name: proxy
+        :converter_tag: external_dp_proxy
+        :test_case_path: {external_yaml}
+        ```
+        """,
+    )
+    monkeypatch.chdir(tmp_path)
+
+    with pytest.raises(DocsCodegenError) as exc_info:
+        _generate_block(tmp_path, doc_path, "proxy")
+
+    assert "only supports routing.type 'disaggregated_prefill'" in str(exc_info.value)

tests/ut/attention/a2/test_attention_cp.py

@@ -0,0 +1,623 @@
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from tests.ut.attention.utils import patch_distributed_groups
+from tests.ut.base import TestBase
+from vllm_ascend.attention import utils as attention_utils
+from vllm_ascend.attention.attention_v1 import AscendMetadata
+from vllm_ascend.attention.context_parallel.attention_cp import AscendAttentionCPImpl
+from vllm_ascend.attention.context_parallel.common_cp import (
+    AscendMetadataForPrefill,
+    AscendPCPMetadata,
+)
+
+
+class TestAscendAttentionCPImpl(TestBase):
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def setUp(self):
+        self.layer = MagicMock()
+        self.layer.layer_name = "test_layer"
+        self.layer._k_scale_float = 1.0
+        self.layer._v_scale_float = 1.0
+
+        self.attention_type = MagicMock()
+        self.attention_type.DECODER = "decoder"
+        self.attention_type.ENCODER = "encoder"
+
+        self.attn_metadata = MagicMock()
+        self.attn_metadata.return_value = "1"
+
+        self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
+        self.layer_no_quant.layer_name = "test_layer"
+        self.layer_no_quant._k_scale_float = 1.0
+        self.layer_no_quant._v_scale_float = 1.0
+        self.mock_vllm_config = MagicMock()
+        self.mock_vllm_config.speculative_config = None
+        self.config_patcher = patch(
+            "vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.utils_config_patcher = patch(
+            "vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.config_patcher.start()
+        self.utils_config_patcher.start()
+        attention_utils.needs_layer_aware_fia_graph_replay.cache_clear()
+        self.addCleanup(attention_utils.needs_layer_aware_fia_graph_replay.cache_clear)
+        self.addCleanup(self.utils_config_patcher.stop)
+        self.addCleanup(self.config_patcher.stop)
+
+        self.impl = AscendAttentionCPImpl(
+            num_heads=8,
+            head_size=64,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+    def test_init(self):
+        self.assertEqual(self.impl.pcp_size, 2)
+        self.assertEqual(self.impl.pcp_rank, 0)
+        self.assertEqual(self.impl.dcp_size, 2)
+        self.assertEqual(self.impl.dcp_rank, 0)
+
+    def test_forward_prefill_cp(self):
+        query = torch.randn(2, 4, 128)
+        key = torch.randn(4, 1, 128)
+        value = torch.randn(4, 1, 128)
+
+        def mock_attention_with_nomask_and_mask(q, k_mask, **kwargs):
+            mock_output = torch.randn_like(q)
+            mock_lse = torch.randn_like(k_mask)
+            return mock_output, mock_lse
+
+        self.impl._attention_with_nomask_and_mask = MagicMock()
+        self.impl._attention_with_nomask_and_mask.side_effect = mock_attention_with_nomask_and_mask
+
+        attn_metadata = MagicMock()
+        attn_metadata.prefill = MagicMock()
+        attn_metadata.prefill.pcp_metadata.q_head_idx = torch.tensor([0])
+        attn_metadata.prefill.pcp_metadata.q_tail_idx = torch.tensor([1])
+        attn_metadata.prefill.pcp_metadata.q_full_idx = torch.tensor([0, 1])
+        attn_metadata.prefill.pcp_metadata.kv_with_q_head_mask_idx = torch.tensor([0])
+        attn_metadata.prefill.pcp_metadata.kv_with_q_tail_nomask_idx = torch.tensor([0])
+        attn_metadata.prefill.pcp_metadata.kv_with_q_tail_mask_idx = torch.tensor([0])
+        attn_metadata.prefill.pcp_metadata.pcp_fa_query_idx = torch.tensor([0, 1])
+        attn_metadata.prefill.pcp_metadata.pcp_use_hybrid_attn = False
+
+        output, attn_lse = self.impl._forward_prefill_cp(query, key, value, attn_metadata)
+
+        self.assertEqual(output.shape[0], 2)
+        self.assertEqual(output.shape[1], 4)
+        self.assertEqual(output.shape[2], 128)
+
+    @patch("torch_npu.npu_attention_update")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2)
+    def test_forward_decode_pcp_dcp(
+        self,
+        mock_all2all,
+        mock_dcp,
+        mock_pcp,
+        mock_get_forward_context,
+        mock_npu_fused_infer_attention_score,
+        mock_npu_attention_update,
+    ):
+        query = torch.randn(2, 4, 64)
+        self.impl.key_cache = torch.randn(100, 64, 1, 64)
+        self.impl.value_cache = torch.randn(100, 64, 1, 64)
+
+        # Mock output
+        mock_npu_attention_update.return_value = (torch.randn(2 * 4, 64), None)
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        def mock_npu_fused_infer_attention_score_func(query, k_nope, value, **common_kwargs):
+            mock_output = torch.randn_like(query)
+            mock_lse = torch.randn(query.shape[0], query.shape[1], 1)
+            return mock_output, mock_lse
+
+        mock_npu_fused_infer_attention_score.side_effect = mock_npu_fused_infer_attention_score_func
+
+        attn_metadata = MagicMock()
+        attn_metadata.decode_meta = MagicMock()
+        attn_metadata.num_decodes_flatten = 5
+        attn_metadata.decode_meta.batch_seq_mask = torch.tensor([1, 0], dtype=torch.bool)
+        output = self.impl._forward_decode_pcp_dcp(query, attn_metadata)
+
+        self.assertEqual(output.shape[0], 2)
+        self.assertEqual(output.shape[1], 4)
+        self.assertEqual(output.shape[2], 64)
+
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_prefill_query_all_gather(self):
+        query = torch.randn(2, 4, 128)
+
+        attn_metadata = MagicMock()
+        attn_metadata.prefill = MagicMock()
+        attn_metadata.prefill.chunked_context = MagicMock()
+        attn_metadata.prefill.chunked_context.cp_kv_recover_idx_for_chunk = torch.tensor([1, 2, 3, 0])
+        output = self.impl._prefill_query_all_gather(attn_metadata, query)
+
+        self.assertEqual(output.shape[0], 4)
+        self.assertEqual(output.shape[1], 8)
+        self.assertEqual(output.shape[2], 128)
+
+    @patch("torch.ops.npu.npu_fused_infer_attention_score")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_compute_prefill_context(self, mock_npu_attention):
+        block_num = 100
+        block_size = 128
+        kv_num_heads = 1
+        head_size = 128
+        kv_cache = (
+            torch.randn(block_num, block_size, kv_num_heads, head_size),
+            torch.randn(block_num, block_size, kv_num_heads, head_size),
+        )
+
+        batch_size = 1024
+        self.impl.head_size = head_size
+        self.impl.num_heads = 4
+        num_heads = self.impl.num_heads * self.impl.dcp_size
+        query = torch.randn(batch_size, num_heads, head_size)
+
+        attn_metadata = MagicMock()
+        attn_metadata.prefill = MagicMock()
+        attn_metadata.prefill.chunked_context = MagicMock()
+        local_context_lens_allranks = torch.tensor([[[256, 256], [256, 256]]])
+        attn_metadata.prefill.chunked_context.local_context_lens_allranks = local_context_lens_allranks
+        attn_metadata.prefill.chunked_context.local_total_toks = local_context_lens_allranks[:, 0, 0].sum()
+
+        def mock_load_kv_for_chunk(attn_metadata, kv_cache, local_chunked_kv_lens_rank, query, total_toks):
+            return torch.randn(total_toks, kv_num_heads, head_size), torch.randn(total_toks, kv_num_heads, head_size)
+
+        self.impl._load_kv_for_chunk = MagicMock()
+        self.impl._load_kv_for_chunk.side_effect = mock_load_kv_for_chunk
+
+        mock_npu_attention.return_value = (
+            torch.randn(batch_size, num_heads, head_size),
+            torch.randn(batch_size, num_heads, 1),
+        )
+
+        context_output = self.impl._compute_prefill_context(query, kv_cache, attn_metadata)
+        local_context_output = torch.cat(context_output, dim=-1).permute([1, 2, 0]).contiguous()
+        global_context_output = self.impl._gather_global_context_output(local_context_output)
+        global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
+        result_output, result_lse = self.impl._update_global_context_output(global_context_output)
+
+        self.assertEqual(result_output.shape[0], batch_size)
+        self.assertEqual(result_output.shape[1], self.impl.num_heads)
+        self.assertEqual(result_output.shape[2], head_size)
+        self.assertEqual(result_lse.shape[0], batch_size)
+        self.assertEqual(result_lse.shape[1], self.impl.num_heads)
+        self.assertEqual(result_lse.shape[2], 1)
+
+    @patch("torch_npu.atb.npu_paged_cache_load")
+    def test_load_kv_for_chunk(self, mock_npu_paged_cache_load):
+        block_num = 100
+        block_size = 128
+        num_heads = 1
+        head_size = 128
+
+        kv_cache = (
+            torch.randn(block_num, block_size, num_heads, head_size),
+            torch.randn(block_num, block_size, num_heads, head_size),
+        )
+        query = torch.randn(4, 8, 128)
+        total_toks = 256
+        local_chunked_kv_lens_rank = torch.randn(total_toks)
+
+        attn_metadata = MagicMock()
+
+        key, value = self.impl._load_kv_for_chunk(
+            attn_metadata, kv_cache, local_chunked_kv_lens_rank, query, total_toks
+        )
+
+        self.assertEqual(key.shape[0], total_toks)
+        self.assertEqual(key.shape[1], num_heads)
+        self.assertEqual(key.shape[2], head_size)
+        self.assertEqual(value.shape[0], total_toks)
+        self.assertEqual(value.shape[1], num_heads)
+        self.assertEqual(value.shape[2], head_size)
+
+    @patch("torch_npu.Event", create=True)
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_reshape_and_cache(self, mock_event_class, mock_npu_reshape_and_cache):
+        num_tokens = 4
+        block_num = 100
+        block_size = 128
+        num_heads = 1
+        head_size = 128
+        self.impl.head_size = head_size
+        self.impl.is_kv_producer = False
+
+        kv_cache = (
+            torch.randn(block_num, block_size, num_heads, head_size),
+            torch.randn(block_num, block_size, num_heads, head_size),
+        )
+
+        attn_metadata = MagicMock()
+        attn_metadata.num_decode_tokens = 1
+        attn_metadata.num_decodes = 1
+        attn_metadata.num_prefills = 1
+        attn_metadata.slot_mapping = torch.randn(2)
+        attn_metadata.num_actual_tokens_pcp_padded = num_tokens * self.impl.pcp_size
+        attn_metadata.prefill = MagicMock()
+        attn_metadata.prefill.pcp_metadata.pcp_allgather_restore_idx = torch.tensor([0, 3, 1, 2, 0, 0, 0, 0])
+        attn_metadata.prefill.pcp_metadata.pcp_use_hybrid_attn = False
+        attn_metadata.prefill.pcp_metadata.pcp_padded_tokens_fla = 0
+        attn_metadata.prefill.pcp_metadata.pcp_enter_fa_restore_idx = torch.arange(num_tokens * 3 * self.impl.pcp_size)
+        attn_metadata.prefill.pcp_metadata.pcp_unpad_mask = torch.tensor(
+            [True, False, True, True, True, True, True, True]
+        )
+
+        query = torch.rand(num_tokens, num_heads, head_size)
+        key = torch.randn(num_tokens, num_heads, head_size)
+        value = torch.randn(num_tokens, num_heads, head_size)
+        output = torch.rand(num_tokens, num_heads * head_size)
+
+        query, key, value, output = self.impl.reshape_and_cache(query, key, value, kv_cache, attn_metadata, output)
+        self.assertEqual(key.shape[0], num_tokens * self.impl.pcp_size)
+        self.assertEqual(key.shape[1], num_heads)
+        self.assertEqual(key.shape[2], head_size)
+        self.assertEqual(value.shape[0], num_tokens * self.impl.pcp_size)
+        self.assertEqual(value.shape[1], num_heads)
+        self.assertEqual(value.shape[2], head_size)
+
+
+class TestUpdateNpuAttnOutLse(TestBase):
+    @patch_distributed_groups(needs_mocks=False)
+    def setUp(self):
+        self.layer = MagicMock()
+        self.layer.layer_name = "test_layer"
+        self.layer._k_scale_float = 1.0
+        self.layer._v_scale_float = 1.0
+
+        self.attention_type = MagicMock()
+        self.attention_type.DECODER = "decoder"
+        self.attention_type.ENCODER = "encoder"
+
+        self.attn_metadata = MagicMock()
+        self.attn_metadata.return_value = "1"
+
+        self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
+        self.layer_no_quant.layer_name = "test_layer"
+        self.layer_no_quant._k_scale_float = 1.0
+        self.layer_no_quant._v_scale_float = 1.0
+        self.mock_vllm_config = MagicMock()
+        self.mock_vllm_config.speculative_config = None
+        self.config_patcher = patch(
+            "vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.utils_config_patcher = patch(
+            "vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.config_patcher.start()
+        self.utils_config_patcher.start()
+        attention_utils.needs_layer_aware_fia_graph_replay.cache_clear()
+        self.addCleanup(attention_utils.needs_layer_aware_fia_graph_replay.cache_clear)
+        self.addCleanup(self.utils_config_patcher.stop)
+        self.addCleanup(self.config_patcher.stop)
+
+        self.impl = AscendAttentionCPImpl(
+            num_heads=8,
+            head_size=64,
+            scale=0.125,
+            num_kv_heads=2,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+        self.impl.pcp_size = 1
+        self.impl.dcp_size = 1
+
+        self.batch_size = 2
+
+        # sequence length per batch
+        self.q_lens_per_batch = [32, 64]
+        self.kv_lens_nomask_per_batch = [32, 64]
+        self.kv_lens_mask_per_batch = [32, 64]
+
+        # TND layout requires cumulative sum computation.
+        self.q_seqlens_cumsum = self._cumsum(self.q_lens_per_batch)  # [32, 96]
+        self.kv_seqlens_nomask_cumsum = self._cumsum(self.kv_lens_nomask_per_batch)  # [32, 96]
+        self.kv_seqlens_mask_cumsum = self._cumsum(self.kv_lens_mask_per_batch)  # [32, 96]
+
+        # Compute T value in TND layout
+        self.q_total_tokens = self.q_seqlens_cumsum[-1]
+        self.kv_total_nomask = self.kv_seqlens_nomask_cumsum[-1]  #
+        self.kv_total_mask = self.kv_seqlens_mask_cumsum[-1]
+
+    def _cumsum(self, arr: list[int]) -> list[int]:
+        result = []
+        total = 0
+        for val in arr:
+            total += val
+            result.append(total)
+        return result
+
+    def _build_attn_metadata(self, with_chunked_context=False):
+        attn_metadata = AscendMetadata()
+        attn_metadata.num_prefills = self.batch_size
+        attn_metadata.num_decodes = 0
+        attn_metadata.num_actual_tokens = self.q_total_tokens
+
+        prefill_metadata = AscendMetadataForPrefill()
+        pcp_metadata = AscendPCPMetadata()
+        pcp_metadata.attn_mask_seqlens = self.kv_seqlens_mask_cumsum
+        pcp_metadata.head_attn_nomask_seqlens = self.kv_seqlens_nomask_cumsum
+        pcp_metadata.tail_attn_nomask_seqlens = self.kv_seqlens_nomask_cumsum
+        prefill_metadata.pcp_metadata = pcp_metadata
+
+        prefill_metadata.actual_seq_lengths_q = torch.tensor(self.q_seqlens_cumsum)
+
+        if with_chunked_context:
+            chunked_context = AscendMetadataForPrefill.ChunkedContextMetadata(
+                actual_chunk_seq_lengths=self.kv_seqlens_mask_cumsum,
+                actual_seq_lengths_kv=self.kv_seqlens_mask_cumsum,
+                starts=None,
+                chunk_seq_mask_filtered_indices=None,
+            )
+            prefill_metadata.chunked_context = chunked_context
+        else:
+            prefill_metadata.chunked_context = None
+
+        attn_metadata.prefill = prefill_metadata
+        attn_metadata.decode_meta = None
+        return attn_metadata
+
+    @patch("torch.ops.npu.npu_fused_infer_attention_score")
+    def test_attention_with_nomask_none(self, mock_npu_attention):
+        # Mock input data
+        q = torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size)
+        q_seqlens = self.q_seqlens_cumsum
+        k_nomask = None
+        v_nomask = None
+        kv_seqlens_nomask = self.kv_seqlens_nomask_cumsum
+        k_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
+        v_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
+        kv_seqlens_mask = self.kv_seqlens_mask_cumsum
+        mask = torch.randn(self.q_total_tokens, self.kv_total_mask)
+        attn_metadata = self._build_attn_metadata(with_chunked_context=False)
+        # Mock output
+        mock_npu_attention.return_value = torch.randn(96, 8, 64), torch.randn(96, 8, 1)
+
+        # Call the method under test
+        output, attn_lse = self.impl._attention_with_nomask_and_mask(
+            q, q_seqlens, k_nomask, v_nomask, kv_seqlens_nomask, k_mask, v_mask, kv_seqlens_mask, mask, attn_metadata
+        )
+
+        # Verify only mask attention was invoked
+        mock_npu_attention.assert_called_with(
+            q,
+            k_mask,
+            v_mask,
+            num_heads=self.impl.num_heads,
+            num_key_value_heads=self.impl.num_kv_heads,
+            input_layout="TND",
+            atten_mask=mask,
+            scale=self.impl.scale,
+            sparse_mode=3,
+            antiquant_mode=0,
+            antiquant_scale=None,
+            softmax_lse_flag=True,
+            actual_seq_lengths_kv=kv_seqlens_mask,
+            actual_seq_lengths=q_seqlens,
+        )
+        # Assert the method call
+        self.assertEqual(mock_npu_attention.call_count, 1)
+        self.assertIsInstance(output, torch.Tensor)
+        self.assertIsInstance(attn_lse, torch.Tensor)
+        self.assertEqual(output.shape, (96, 8, 64))
+        self.assertEqual(attn_lse.shape, (96, 8, 1))
+
+    @patch("torch.ops.npu.npu_fused_infer_attention_score")
+    @patch("vllm_ascend.attention.context_parallel.attention_cp._update_out_and_lse")
+    def test_attention_with_nomask_and_mask_chunk(self, mock_update_out_and_lse, mock_npu_fused_infer_attention_score):
+        # Mock input data
+        q = torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size)
+        k_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
+        v_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
+        k_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
+        v_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
+
+        mask = torch.randn(self.q_total_tokens, self.kv_total_mask)
+        attn_metadata = self._build_attn_metadata(with_chunked_context=True)
+
+        # Mock output
+        mock_npu_fused_infer_attention_score.return_value = (
+            torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size),
+            torch.randn(self.q_total_tokens, self.impl.num_heads, 1),
+        )
+        mock_update_out_and_lse.return_value = (
+            torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size),
+            torch.randn(self.q_total_tokens, self.impl.num_heads, 1),
+        )
+        # Call the method under test
+        output, attn_lse = self.impl._attention_with_nomask_and_mask(
+            q=q,
+            q_seqlens=self.q_seqlens_cumsum,
+            k_nomask=k_nomask,
+            v_nomask=v_nomask,
+            kv_seqlens_nomask=self.kv_seqlens_nomask_cumsum,
+            k_mask=k_mask,
+            v_mask=v_mask,
+            kv_seqlens_mask=self.kv_seqlens_mask_cumsum,
+            mask=mask,
+            attn_metadata=attn_metadata,
+        )
+        # Assert the method call
+        self.assertEqual(mock_npu_fused_infer_attention_score.call_count, 2)
+        self.assertIsNotNone(output)
+        self.assertIsNotNone(attn_lse)
+
+    @patch("torch.ops.npu.npu_fused_infer_attention_score")
+    @patch("vllm_ascend.attention.context_parallel.attention_cp._npu_attn_out_lse_update")
+    def test_attention_with_nomask_and_mask_nochunk(
+        self, mock_npu_attn_out_lse_update, mock_npu_fused_infer_attention_score
+    ):
+        # Mock input data
+        q = torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size)
+        k_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
+        v_nomask = torch.randn(self.kv_total_nomask, self.impl.num_kv_heads, self.impl.head_size)
+        k_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
+        v_mask = torch.randn(self.kv_total_mask, self.impl.num_kv_heads, self.impl.head_size)
+        mask = torch.randn(self.q_total_tokens, self.kv_total_mask)
+
+        attn_metadata = self._build_attn_metadata(with_chunked_context=True)
+        attn_metadata.prefill.chunked_context = None
+
+        # Mock output
+        mock_npu_fused_infer_attention_score.return_value = (
+            torch.randn(self.q_total_tokens, self.impl.num_heads, self.impl.head_size),
+            torch.randn(self.q_total_tokens, self.impl.num_heads, 1),
+        )
+        mock_npu_attn_out_lse_update.return_value = torch.randn(
+            self.q_total_tokens, self.impl.num_heads, self.impl.head_size
+        )
+
+        # Call the method under test
+        output, attn_lse = self.impl._attention_with_nomask_and_mask(
+            q=q,
+            q_seqlens=self.q_seqlens_cumsum,
+            k_nomask=k_nomask,
+            v_nomask=v_nomask,
+            kv_seqlens_nomask=self.kv_seqlens_nomask_cumsum,
+            k_mask=k_mask,
+            v_mask=v_mask,
+            kv_seqlens_mask=self.kv_seqlens_mask_cumsum,
+            mask=mask,
+            attn_metadata=attn_metadata,
+        )
+        # Assert the method call
+        mock_npu_attn_out_lse_update.assert_called_once()
+        self.assertEqual(mock_npu_fused_infer_attention_score.call_count, 2)
+        self.assertIsNotNone(output)
+        self.assertEqual(attn_lse, None)
+
+    @patch("vllm_ascend.attention.context_parallel.attention_cp._npu_attn_out_lse_update")
+    def test_update_chunk_attn_out_lse_with_current_attn_out_lse(self, mock_npu_attn_out_lse_update):
+        # Mock input data
+        current_attn_output_prefill = torch.randn(32764, 8, 128)
+        current_attn_lse_prefill = torch.randn(32764, 8, 1)
+        attn_output_full_chunk = torch.randn(65528, 8, 128)
+        attn_lse_full_chunk = torch.randn(65528, 8, 1)
+        prefill_query = torch.randn(32764, 8, 128)
+        # mock attn_metadata
+        attn_metadata = self._build_attn_metadata(with_chunked_context=True)
+        attn_metadata.prefill.chunked_context.chunk_seq_mask_filtered_indices = torch.arange(32764, dtype=torch.int32)
+        attn_metadata.prefill.chunked_context.kv_inverse_idx_for_chunk = torch.arange(32764, dtype=torch.int32)
+        # Mock output
+        mock_npu_attn_out_lse_update.return_value = torch.randn(32764, 8, 128)
+        # test pcp_size > 1
+        self.impl.pcp_size = 2
+        self.impl.pcp_rank = 0
+        self.impl.dcp_group = None
+        self.impl.pcp_group = None
+        # Call the method under test
+        self.impl._update_chunk_attn_out_lse_with_current_attn_out_lse(
+            current_attn_output_prefill,
+            current_attn_lse_prefill,
+            attn_output_full_chunk,
+            attn_lse_full_chunk,
+            prefill_query,
+            attn_metadata,
+        )
+        # Assert the method call
+        mock_npu_attn_out_lse_update.assert_called_once()
+        # test pcp_size = 1
+        self.impl.pcp_size = 1
+        self.impl._update_chunk_attn_out_lse_with_current_attn_out_lse(
+            current_attn_output_prefill,
+            current_attn_lse_prefill,
+            attn_output_full_chunk,
+            attn_lse_full_chunk,
+            prefill_query,
+            attn_metadata,
+        )
+        self.assertEqual(mock_npu_attn_out_lse_update.call_count, 2)
+
+    @patch_distributed_groups(dcp_size=2, pcp_size=3)
+    def test_update_chunk_attn_out_lse_dcp2_pcp3(self, mock_all_to_all_single, mock_dcp, mock_pcp):
+        # Mock input data
+        prefix_chunk_output = torch.randn(2, 4, 8)
+        prefix_chunk_lse = torch.randn(2, 4, 1)
+        self.impl.dcp_size = 2
+        self.impl.pcp_size = 3
+        self.impl.head_size = 8
+
+        # Call the method under test
+        chunk_data = torch.cat([prefix_chunk_output, prefix_chunk_lse], dim=-1).permute([1, 2, 0]).contiguous()
+        global_context_output = self.impl._gather_global_context_output(chunk_data)
+        global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
+        output, lse = self.impl._update_global_context_output(global_context_output)
+
+        # Assert the method call
+        self.assertIsInstance(output, torch.Tensor)
+        self.assertIsInstance(lse, torch.Tensor)
+        self.assertEqual(output.shape, (2, 2, 8))
+        self.assertEqual(lse.shape, (2, 2, 1))
+
+        mock_all_to_all_single.assert_called_once()
+        mock_pcp.all_gather.assert_called_once()
+
+    @patch_distributed_groups(dcp_size=2)
+    def test_update_chunk_attn_out_lse_dcp2_pcp1(self, mock_all_to_all_single, mock_dcp, mock_pcp):
+        # Mock input data
+        prefix_chunk_output = torch.randn(2, 4, 8)
+        prefix_chunk_lse = torch.randn(2, 4, 1)
+
+        self.impl.dcp_size = 2
+        self.impl.pcp_size = 1
+        self.impl.head_size = 8
+
+        # Call the method under test
+        chunk_data = torch.cat([prefix_chunk_output, prefix_chunk_lse], dim=-1).permute([1, 2, 0]).contiguous()
+        global_context_output = self.impl._gather_global_context_output(chunk_data)
+        global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
+        output, lse = self.impl._update_global_context_output(global_context_output)
+
+        # Assert the method call
+        self.assertIsInstance(output, torch.Tensor)
+        self.assertIsInstance(lse, torch.Tensor)
+        self.assertEqual(output.shape, (2, 2, 8))
+        self.assertEqual(lse.shape, (2, 2, 1))
+
+        mock_all_to_all_single.assert_called_once()
+        mock_pcp.all_gather.assert_not_called()
+
+    @patch_distributed_groups(pcp_size=2)
+    def test_update_chunk_attn_out_lse_dcp1_pcp2(self, mock_all_to_all_single, mock_dcp, mock_pcp):
+        # Mock input data
+        prefix_chunk_output = torch.randn(2, 4, 8)
+        prefix_chunk_lse = torch.randn(2, 4, 1)
+
+        self.impl.dcp_size = 1
+        self.impl.pcp_size = 2
+        self.impl.head_size = 8
+
+        # Call the method under test
+        chunk_data = torch.cat([prefix_chunk_output, prefix_chunk_lse], dim=-1).permute([1, 2, 0]).contiguous()
+        global_context_output = self.impl._gather_global_context_output(chunk_data)
+        global_context_output = global_context_output.permute([2, 0, 1]).contiguous()
+        output, lse = self.impl._update_global_context_output(global_context_output)
+
+        # Assert the method call
+        self.assertIsInstance(output, torch.Tensor)
+        self.assertIsInstance(lse, torch.Tensor)
+        self.assertEqual(output.shape, (2, 4, 8))
+        self.assertEqual(lse.shape, (2, 4, 1))
+
+        mock_all_to_all_single.assert_not_called()
+        mock_pcp.all_gather.assert_called_once()

tests/ut/attention/a2/test_attention_cp_precision.py

@@ -0,0 +1,993 @@
+from functools import partial
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.config import set_current_vllm_config
+from vllm.forward_context import set_forward_context
+from vllm.utils.torch_utils import set_random_seed
+
+from tests.ut.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+    patch_distributed_groups,
+)
+from vllm_ascend.attention.attention_mask import AttentionMaskBuilder
+from vllm_ascend.attention.attention_v1 import AscendMetadata
+from vllm_ascend.attention.context_parallel.attention_cp import (
+    AscendAttentionCPImpl,
+)
+from vllm_ascend.attention.context_parallel.common_cp import (
+    AscendMetadataForDecode,
+    AscendMetadataForPrefill,
+    AscendPCPMetadata,
+)
+
+BATCH_SPECS = {
+    "single_prefill": BatchSpec(seq_lens=[128], query_lens=[128]),
+    "small_prefill": BatchSpec(seq_lens=[32, 48], query_lens=[32, 48]),
+    "medium_prefill": BatchSpec(seq_lens=[256, 512], query_lens=[256, 512]),
+    "large_prefill": BatchSpec(seq_lens=[1024, 2048], query_lens=[1024, 2048]),
+    "single_decode": BatchSpec(seq_lens=[32], query_lens=[1]),
+    "small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
+    "medium_decode": BatchSpec(seq_lens=[128, 256, 512, 1024], query_lens=[1, 1, 1, 1]),
+    "mixed_small": BatchSpec(seq_lens=[32, 40, 5, 5], query_lens=[1, 1, 5, 5]),
+    "mixed_medium": BatchSpec(seq_lens=[256, 512, 7, 7], query_lens=[1, 1, 7, 7]),
+    "mixed_large": BatchSpec(seq_lens=[1024, 2048, 16, 16], query_lens=[1, 1, 16, 16]),
+    "mtp_1_plus_3_small": BatchSpec(seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4]),
+    "mtp_1_plus_3_medium": BatchSpec(seq_lens=[1024, 2048, 3072, 4096], query_lens=[4, 4, 4, 4]),
+    "mtp_1_plus_3_tiny": BatchSpec(seq_lens=[64, 128], query_lens=[4, 4]),
+}
+
+MODELS = [
+    "Qwen/Qwen3-8B",
+]
+
+
+class MockAttentionLayer:
+    def __init__(self, device: torch.device):
+        self._q_scale = torch.tensor(1.0, device=device)
+        self._k_scale = torch.tensor(1.0, device=device)
+        self._v_scale = torch.tensor(1.0, device=device)
+        self._q_scale_float = 1.0
+        self._k_scale_float = 1.0
+        self._v_scale_float = 1.0
+        self.layer_name = "model.layers.0"
+
+
+def compute_sdpa_reference(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    batch_spec: BatchSpec,
+    scale: float,
+    num_q_heads: int,
+    num_kv_heads: int,
+) -> torch.Tensor:
+    """Compute reference attention output using PyTorch SDPA.
+
+    Iterates over each sequence in the batch and computes causal attention
+    matching the FIA sparse_mode=3 behavior with splitfuse layout.
+    Used for pure prefill tests where seq_lens == query_lens.
+    """
+    enable_gqa = num_q_heads != num_kv_heads
+    all_sdpa_outputs = []
+    q_offset = 0
+    kv_offset = 0
+
+    for i in range(batch_spec.batch_size):
+        q_len = batch_spec.query_lens[i]
+        kv_len = batch_spec.seq_lens[i]
+
+        q_i = q[q_offset : q_offset + q_len]
+        k_i = k[kv_offset : kv_offset + kv_len]
+        v_i = v[kv_offset : kv_offset + kv_len]
+
+        q_sdpa = q_i.unsqueeze(0).transpose(1, 2)
+        k_sdpa = k_i.unsqueeze(0).transpose(1, 2)
+        v_sdpa = v_i.unsqueeze(0).transpose(1, 2)
+
+        context_len = kv_len - q_len
+        if context_len > 0:
+            attn_mask = torch.ones(q_len, kv_len, dtype=torch.bool, device=q.device)
+            causal_mask = torch.tril(torch.ones(q_len, q_len, device=q.device))
+            attn_mask[:, context_len:] = causal_mask
+            sdpa_out = torch.nn.functional.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=attn_mask,
+                is_causal=False,
+                enable_gqa=enable_gqa,
+                scale=scale,
+            )
+        else:
+            sdpa_out = torch.nn.functional.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                is_causal=True,
+                enable_gqa=enable_gqa,
+                scale=scale,
+            )
+
+        all_sdpa_outputs.append(sdpa_out.transpose(1, 2).squeeze(0))
+        q_offset += q_len
+        kv_offset += kv_len
+
+    return torch.cat(all_sdpa_outputs, dim=0)
+
+
+def compute_mixed_sdpa_reference(
+    full_q: torch.Tensor,
+    full_k: torch.Tensor,
+    full_v: torch.Tensor,
+    batch_spec: BatchSpec,
+    scale: float,
+    num_q_heads: int,
+    num_kv_heads: int,
+) -> torch.Tensor:
+    """Compute per-sequence SDPA reference for mixed decode+prefill.
+
+    Each sequence gets its own Q/K/V with causal masking
+    (context tokens are visible to new tokens).
+    """
+    enable_gqa = num_q_heads != num_kv_heads
+    all_outputs = []
+    q_offset = 0
+    kv_offset = 0
+
+    for i in range(batch_spec.batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        q_i = full_q[q_offset : q_offset + q_len]
+        k_i = full_k[kv_offset : kv_offset + s_len]
+        v_i = full_v[kv_offset : kv_offset + s_len]
+
+        q_sdpa = q_i.unsqueeze(0).transpose(1, 2)
+        k_sdpa = k_i.unsqueeze(0).transpose(1, 2)
+        v_sdpa = v_i.unsqueeze(0).transpose(1, 2)
+
+        if context_len > 0:
+            attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=full_q.device)
+            causal_mask = torch.tril(torch.ones(q_len, q_len, device=full_q.device))
+            attn_mask[:, context_len:] = causal_mask
+            sdpa_out = torch.nn.functional.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=attn_mask,
+                is_causal=False,
+                enable_gqa=enable_gqa,
+                scale=scale,
+            )
+        else:
+            sdpa_out = torch.nn.functional.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                is_causal=True,
+                enable_gqa=enable_gqa,
+                scale=scale,
+            )
+
+        all_outputs.append(sdpa_out.transpose(1, 2).squeeze(0))
+        q_offset += q_len
+        kv_offset += s_len
+
+    return torch.cat(all_outputs, dim=0)
+
+
+def _make_kv_cache_for_decode(
+    batch_spec: BatchSpec,
+    num_kv_heads: int,
+    head_size: int,
+    block_size: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    block_table: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pre-populate KV cache with ALL tokens for decode-only tests.
+
+    Places context + new tokens sequentially in cache blocks
+    so that FIA paged attention can read them via block_table.
+    """
+    num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
+    num_blocks = max(num_blocks, 64)
+    k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
+    v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
+
+    kv_offset = 0
+    block_start = 0
+    for i in range(batch_spec.batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        k_full = key[kv_offset : kv_offset + s_len]
+        v_full = value[kv_offset : kv_offset + s_len]
+
+        context_k = k_full[:context_len].contiguous()
+        context_v = v_full[:context_len].contiguous()
+        new_k = k_full[context_len:].contiguous()
+        new_v = v_full[context_len:].contiguous()
+
+        num_blocks_for_seq = (s_len + block_size - 1) // block_size
+        block_table[i, :num_blocks_for_seq] = torch.arange(
+            block_start,
+            block_start + num_blocks_for_seq,
+            dtype=torch.int32,
+            device=device,
+        )
+
+        for t_idx in range(context_len):
+            blk = block_start + t_idx // block_size
+            pos = t_idx % block_size
+            k_cache[blk, pos] = context_k[t_idx]
+            v_cache[blk, pos] = context_v[t_idx]
+
+        for t_idx in range(q_len):
+            blk = block_start + (context_len + t_idx) // block_size
+            pos = (context_len + t_idx) % block_size
+            k_cache[blk, pos] = new_k[t_idx]
+            v_cache[blk, pos] = new_v[t_idx]
+
+        kv_offset += s_len
+        block_start += num_blocks_for_seq
+
+    return k_cache, v_cache
+
+
+def _make_kv_cache_for_mixed(
+    batch_spec: BatchSpec,
+    num_kv_heads: int,
+    head_size: int,
+    block_size: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    block_table: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pre-populate KV cache for mixed decode+prefill tests.
+
+    Only decode sequences have context tokens placed in the cache.
+    Prefill sequences (seq_lens == query_lens) have no context;
+    their KV goes through the direct FIA prefill path.
+    """
+    num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
+    num_blocks = max(num_blocks, 64)
+    k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
+    v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
+
+    kv_offset = 0
+    block_start = 0
+    for i in range(batch_spec.batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+        is_decode = q_len == 1
+
+        k_full = key[kv_offset : kv_offset + s_len]
+        v_full = value[kv_offset : kv_offset + s_len]
+
+        if is_decode and context_len > 0:
+            num_blocks_for_seq = (s_len + block_size - 1) // block_size
+            block_table[i, :num_blocks_for_seq] = torch.arange(
+                block_start,
+                block_start + num_blocks_for_seq,
+                dtype=torch.int32,
+                device=device,
+            )
+
+            for t_idx in range(context_len):
+                blk = block_start + t_idx // block_size
+                pos = t_idx % block_size
+                k_cache[blk, pos] = k_full[t_idx]
+                v_cache[blk, pos] = v_full[t_idx]
+
+            for t_idx in range(q_len):
+                blk = block_start + (context_len + t_idx) // block_size
+                pos = (context_len + t_idx) % block_size
+                k_cache[blk, pos] = k_full[context_len + t_idx]
+                v_cache[blk, pos] = v_full[context_len + t_idx]
+
+            block_start += num_blocks_for_seq
+
+        kv_offset += s_len
+
+    return k_cache, v_cache
+
+
+def build_cp_attn_metadata(
+    batch_spec: BatchSpec,
+    vllm_config,
+    device: torch.device,
+    pcp_size: int = 1,
+    pcp_rank: int = 0,
+    kv_cache_prepopulated: bool = False,
+    decode_threshold: int = 1,
+) -> AscendMetadata:
+    common_attn_metadata = create_common_attn_metadata(batch_spec, vllm_config.cache_config.block_size, device)
+
+    num_reqs = common_attn_metadata.num_reqs
+    num_actual_tokens = common_attn_metadata.num_actual_tokens
+
+    query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1]
+
+    num_decodes = sum(1 for ql in batch_spec.query_lens if ql <= decode_threshold)
+    num_prefills = batch_spec.batch_size - num_decodes
+    num_decode_tokens = sum(ql for ql in batch_spec.query_lens if ql <= decode_threshold)
+    num_prefill_tokens = num_actual_tokens - num_decode_tokens
+
+    block_table = common_attn_metadata.block_table_tensor
+    slot_mapping = common_attn_metadata.slot_mapping
+    query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
+    num_decodes_flatten = query_lens[:num_decodes].sum().item()
+
+    seq_lens_cpu = common_attn_metadata.seq_lens_cpu[:num_reqs]
+
+    num_actual_tokens_pcp_padded = num_actual_tokens * pcp_size
+
+    query_start_loc = query_start_loc_cpu.to(device, non_blocking=True)
+
+    attn_mask_builder = AttentionMaskBuilder(device)
+    attn_mask = attn_mask_builder.get_attention_mask(common_attn_metadata.causal, vllm_config.model_config)
+
+    prefill_metadata = None
+    if num_prefills > 0:
+        prefill_query_lens = query_lens[num_decodes:]
+        attn_mask_seqlens = torch.cumsum(prefill_query_lens, dim=0).tolist()
+        head_attn_nomask_seqlens = attn_mask_seqlens if pcp_rank > 0 else []
+        tail_attn_nomask_seqlens = attn_mask_seqlens
+
+        total_prefill_tokens = num_prefill_tokens
+        prefill_tokens_offset = num_decode_tokens
+
+        if pcp_size > 1:
+            chunk_size = total_prefill_tokens // pcp_size
+            rank_start = prefill_tokens_offset + pcp_rank * chunk_size
+            rank_end = rank_start + chunk_size
+
+            q_head_idx = torch.arange(rank_start, rank_end, device=device, dtype=torch.long)
+            q_tail_idx = torch.tensor([], device=device, dtype=torch.long)
+
+            kv_total = total_prefill_tokens * pcp_size
+            kv_with_q_head_mask_idx = torch.arange(
+                prefill_tokens_offset,
+                prefill_tokens_offset + kv_total,
+                device=device,
+                dtype=torch.long,
+            )
+            kv_with_q_head_nomask_idx = (
+                torch.arange(prefill_tokens_offset, prefill_tokens_offset + kv_total, device=device, dtype=torch.long)
+                if pcp_rank > 0
+                else torch.tensor([], device=device, dtype=torch.long)
+            )
+            q_full_idx = torch.arange(chunk_size, device=device, dtype=torch.long)
+            kv_with_q_tail_nomask_idx = torch.tensor([], device=device, dtype=torch.long)
+            kv_with_q_tail_mask_idx = torch.tensor([], device=device, dtype=torch.long)
+            pcp_allgather_restore_idx = list(range(total_prefill_tokens * pcp_size))
+        else:
+            q_head_idx = torch.arange(
+                prefill_tokens_offset,
+                prefill_tokens_offset + total_prefill_tokens,
+                device=device,
+                dtype=torch.long,
+            )
+            q_tail_idx = torch.tensor([], device=device, dtype=torch.long)
+            kv_with_q_head_mask_idx = torch.arange(
+                prefill_tokens_offset,
+                prefill_tokens_offset + total_prefill_tokens,
+                device=device,
+                dtype=torch.long,
+            )
+            kv_with_q_head_nomask_idx = torch.tensor([], device=device, dtype=torch.long)
+            kv_with_q_tail_nomask_idx = torch.tensor([], device=device, dtype=torch.long)
+            kv_with_q_tail_mask_idx = torch.tensor([], device=device, dtype=torch.long)
+            q_full_idx = torch.arange(total_prefill_tokens, device=device, dtype=torch.long)
+            pcp_allgather_restore_idx = None
+
+        pcp_metadata = AscendPCPMetadata(
+            q_head_idx=q_head_idx,
+            q_tail_idx=q_tail_idx,
+            kv_with_q_head_nomask_idx=kv_with_q_head_nomask_idx,
+            kv_with_q_head_mask_idx=kv_with_q_head_mask_idx,
+            kv_with_q_tail_nomask_idx=kv_with_q_tail_nomask_idx,
+            kv_with_q_tail_mask_idx=kv_with_q_tail_mask_idx,
+            attn_mask_seqlens=attn_mask_seqlens,
+            head_attn_nomask_seqlens=head_attn_nomask_seqlens,
+            tail_attn_nomask_seqlens=tail_attn_nomask_seqlens,
+            q_full_idx=q_full_idx,
+            pcp_use_hybrid_attn=False,
+            pcp_allgather_restore_idx=pcp_allgather_restore_idx,
+        )
+
+        prefill_cumsum_q = torch.cumsum(query_lens[num_decodes:], dim=0).to(device)
+
+        prefill_metadata = AscendMetadataForPrefill(
+            pcp_metadata=pcp_metadata,
+            pcp_exit_fa_scatter_idx=None,
+            chunked_context=None,
+            block_tables=block_table[num_decodes_flatten:, ...],
+            actual_seq_lengths_q=prefill_cumsum_q,
+        )
+
+    decode_metadata = None
+    if num_decodes > 0:
+        decode_query_lens = query_lens[:num_decodes].tolist()
+        decode_seq_lens = seq_lens_cpu[:num_decodes].tolist()
+
+        if kv_cache_prepopulated:
+            num_computed_tokens_arr = np.zeros((num_decodes_flatten, pcp_size, 1), dtype=np.int32)
+            flat_idx = 0
+            for i in range(num_decodes):
+                s_len = int(decode_seq_lens[i])
+                q_len = decode_query_lens[i]
+                context_len = s_len - q_len
+                for t in range(q_len):
+                    num_computed_tokens_arr[flat_idx, pcp_rank, 0] = context_len + t + 1
+                    flat_idx += 1
+        else:
+            num_computed_tokens_arr = np.zeros((num_decodes_flatten, pcp_size, 1), dtype=np.int32)
+
+        # Tile block_table for MTP: each decode request may have multiple tokens
+        if num_decodes_flatten > num_decodes:
+            tiled_rows = []
+            for i in range(num_decodes):
+                q_len = decode_query_lens[i]
+                row = block_table[i : i + 1]
+                tiled_rows.append(row.repeat(q_len, 1))
+            decode_block_tables = torch.cat(tiled_rows, dim=0)
+        else:
+            decode_block_tables = block_table[:num_decodes_flatten]
+
+        decode_metadata = AscendMetadataForDecode(
+            num_computed_tokens_of_pcp_dcp=num_computed_tokens_arr,
+            block_tables=decode_block_tables,
+        )
+
+    actual_seq_lengths_q = (
+        torch.arange(num_decodes_flatten, device=device) + 1
+        if num_decodes_flatten > 0
+        else torch.tensor([], device=device)
+    ).tolist() + torch.cumsum(query_lens[num_decodes:], dim=0).tolist()
+
+    attn_metadata = AscendMetadata(
+        num_actual_tokens=num_actual_tokens,
+        num_decode_tokens=num_decode_tokens,
+        num_actual_tokens_pcp_padded=num_actual_tokens_pcp_padded,
+        num_decodes_flatten=num_decodes_flatten,
+        block_tables=block_table,
+        query_start_loc=query_start_loc,
+        seq_lens=common_attn_metadata.seq_lens[:num_reqs],
+        seq_lens_cpu=seq_lens_cpu,
+        seq_lens_list=seq_lens_cpu.tolist(),
+        max_query_len=common_attn_metadata.max_query_len,
+        actual_seq_lengths_q=actual_seq_lengths_q,
+        slot_mapping=slot_mapping,
+        attn_mask=attn_mask,
+        attn_state=common_attn_metadata.attn_state,
+        num_prefills=num_prefills,
+        num_decodes=num_decodes,
+        prefill=prefill_metadata,
+        decode_meta=decode_metadata,
+    )
+    return attn_metadata
+
+
+def run_cp_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv_cache: tuple[torch.Tensor, torch.Tensor],
+    attn_metadata: AscendMetadata,
+    impl: AscendAttentionCPImpl,
+    device: torch.device,
+    vllm_config,
+) -> torch.Tensor:
+    """Run CP attention forward pass with proper setup.
+
+    Mocks reshape_and_cache (no-op) since KV cache is pre-populated
+    for decode tests and not needed for prefill tests.
+    """
+    mock_layer_entry = MagicMock()
+    for layer_name in ["placeholder"]:
+        vllm_config.compilation_config.static_forward_context[layer_name] = mock_layer_entry
+
+    num_tokens = query.shape[0]
+    num_q_heads = query.shape[1]
+    head_size = query.shape[2]
+    output = torch.empty(num_tokens, num_q_heads, head_size, dtype=query.dtype, device=device)
+
+    with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
+        from vllm.forward_context import get_forward_context
+
+        forward_ctx = get_forward_context()
+        forward_ctx.num_tokens = num_tokens
+        forward_ctx.is_draft_model = False
+        forward_ctx.is_draft_model_prefill = False
+        forward_ctx.capturing = False
+        forward_ctx.flash_comm_v1_enabled = False
+        forward_ctx.flashcomm_v2_enabled = False
+
+        mock_layer = MockAttentionLayer(device)
+
+        import vllm_ascend.device.device_op as device_op_module
+
+        original_reshape_and_cache = device_op_module.DeviceOperator.reshape_and_cache
+        original_kv_cache_load = device_op_module.DeviceOperator.kv_cache_load
+
+        def _mock_reshape_and_cache(key, value, key_cache, value_cache, slot_mapping):
+            return
+
+        def _mock_kv_cache_load(key_cache, value_cache, block_tables, seq_lens_kv, starts, key, value):
+            return
+
+        device_op_module.DeviceOperator.reshape_and_cache = staticmethod(_mock_reshape_and_cache)
+        device_op_module.DeviceOperator.kv_cache_load = staticmethod(_mock_kv_cache_load)
+
+        try:
+            output = impl.forward(mock_layer, query, key, value, kv_cache, attn_metadata, output=output)
+        finally:
+            device_op_module.DeviceOperator.reshape_and_cache = original_reshape_and_cache
+            device_op_module.DeviceOperator.kv_cache_load = original_kv_cache_load
+
+    return output
+
+
+@pytest.fixture(autouse=True)
+def default_mock_config():
+    mock_config = MagicMock()
+    mock_config.compilation_config = MagicMock()
+    mock_config.compilation_config.custom_ops = ["all"]
+    mock_config.compilation_config.static_forward_context = {}
+    mock_config.parallel_config = MagicMock()
+    mock_config.parallel_config.prefill_context_parallel_size = 1
+    mock_config.parallel_config.decode_context_parallel_size = 1
+    mock_config.parallel_config.tensor_parallel_size = 1
+    mock_config.model_config = MagicMock()
+    mock_config.model_config.dtype = torch.float16
+    mock_config.speculative_config = None
+    mock_config.cache_config = MagicMock()
+    mock_config.cache_config.block_size = 128
+    mock_config.kv_transfer_config = None
+
+    with set_current_vllm_config(mock_config):
+        yield mock_config
+
+
+@pytest.fixture(autouse=True)
+def mock_graph_params():
+    with patch("vllm_ascend.compilation.acl_graph.get_graph_params") as mock_get_graph:
+        graph_params = MagicMock()
+        graph_params.workspaces = {}
+        graph_params.handles = {}
+        graph_params.attn_params = {}
+        graph_params.events = {}
+        mock_get_graph.return_value = graph_params
+        with patch("vllm_ascend.compilation.acl_graph.get_draft_graph_params", return_value=graph_params):
+            yield
+
+
+def _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale):
+    return AscendAttentionCPImpl(
+        num_heads=num_q_heads,
+        head_size=head_size,
+        scale=scale,
+        num_kv_heads=num_kv_heads,
+        alibi_slopes=None,
+        sliding_window=None,
+        kv_cache_dtype="auto",
+        logits_soft_cap=None,
+        attn_type="decoder",
+        kv_sharing_target_layer_name=None,
+    )
+
+
+def _assert_close(output, reference, rtol, atol, backend_name):
+    assert output.shape == reference.shape, (
+        f"[{backend_name}] shape {output.shape} != reference shape {reference.shape}"
+    )
+    assert output.dtype == reference.dtype, (
+        f"[{backend_name}] dtype {output.dtype} != reference dtype {reference.dtype}"
+    )
+    assert torch.isfinite(output).all(), f"[{backend_name}] produced non-finite values"
+
+    def error_msg(msg: str, name: str):
+        return f"[{name}] output differs from SDPA baseline. {msg}"
+
+    torch.testing.assert_close(
+        output,
+        reference,
+        rtol=rtol,
+        atol=atol,
+        msg=partial(error_msg, name=backend_name),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Pure prefill (seq_lens == query_lens, no context)
+# ---------------------------------------------------------------------------
+def _test_cp_prefill_precision_no_cp(
+    batch_spec: BatchSpec,
+    model: str,
+    *,
+    block_size: int = 128,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+):
+    set_random_seed(42)
+
+    vllm_config = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=1,
+        max_model_len=max(batch_spec.seq_lens) + block_size,
+        block_size=block_size,
+        num_gpu_blocks=8192,
+    )
+    device = torch.device("npu")
+
+    num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+    num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+    head_size = vllm_config.model_config.get_head_size()
+    dtype = torch.bfloat16
+    scale = 1.0 / (head_size**0.5)
+
+    num_tokens = batch_spec.compute_num_tokens()
+    total_kv = sum(batch_spec.seq_lens)
+
+    query_vllm = torch.randn(num_tokens, num_q_heads, head_size, dtype=dtype, device=device)
+    key_vllm = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
+    value_vllm = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
+
+    sdpa_output = compute_sdpa_reference(
+        query_vllm,
+        key_vllm,
+        value_vllm,
+        batch_spec,
+        scale,
+        num_q_heads,
+        num_kv_heads,
+    )
+
+    attn_metadata = build_cp_attn_metadata(batch_spec, vllm_config, device, pcp_size=1, pcp_rank=0)
+
+    num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
+    num_blocks = max(num_blocks, 64)
+    kv_cache = (
+        torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device),
+        torch.zeros(num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device),
+    )
+
+    impl = _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale)
+    output = run_cp_attention(query_vllm, key_vllm, value_vllm, kv_cache, attn_metadata, impl, device, vllm_config)
+
+    _assert_close(output, sdpa_output, rtol, atol, "CP_Prefill")
+
+
+# ---------------------------------------------------------------------------
+# Pure decode (query_lens == 1, context in KV cache)
+# ---------------------------------------------------------------------------
+def _test_cp_decode_precision_no_cp(
+    batch_spec: BatchSpec,
+    model: str,
+    *,
+    block_size: int = 128,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+    decode_threshold: int = 1,
+):
+    set_random_seed(42)
+
+    vllm_config = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=1,
+        max_model_len=max(batch_spec.seq_lens) + block_size,
+        block_size=block_size,
+        num_gpu_blocks=8192,
+    )
+    device = torch.device("npu")
+
+    num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+    num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+    head_size = vllm_config.model_config.get_head_size()
+    dtype = torch.bfloat16
+    scale = 1.0 / (head_size**0.5)
+
+    total_kv = sum(batch_spec.seq_lens)
+
+    query_full = torch.randn(total_kv, num_q_heads, head_size, dtype=dtype, device=device)
+    key_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
+    value_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
+
+    sdpa_output = compute_mixed_sdpa_reference(
+        query_full,
+        key_full,
+        value_full,
+        batch_spec,
+        scale,
+        num_q_heads,
+        num_kv_heads,
+    )
+
+    # Backend inputs: only the new (decode) tokens
+    all_q, all_k, all_v = [], [], []
+    q_offset = 0
+    kv_offset = 0
+    for i in range(batch_spec.batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        q_i = query_full[q_offset : q_offset + q_len].contiguous()
+        k_full_i = key_full[kv_offset : kv_offset + s_len]
+        v_full_i = value_full[kv_offset : kv_offset + s_len]
+
+        all_q.append(q_i)
+        all_k.append(k_full_i[context_len:])
+        all_v.append(v_full_i[context_len:])
+
+        q_offset += q_len
+        kv_offset += s_len
+
+    query_vllm = torch.cat(all_q, dim=0)
+    key_vllm = torch.cat(all_k, dim=0)
+    value_vllm = torch.cat(all_v, dim=0)
+
+    attn_metadata = build_cp_attn_metadata(
+        batch_spec,
+        vllm_config,
+        device,
+        pcp_size=1,
+        pcp_rank=0,
+        kv_cache_prepopulated=True,
+        decode_threshold=decode_threshold,
+    )
+
+    k_cache, v_cache = _make_kv_cache_for_decode(
+        batch_spec,
+        num_kv_heads,
+        head_size,
+        block_size,
+        dtype,
+        device,
+        key_full,
+        value_full,
+        attn_metadata.block_tables,
+    )
+
+    # Re-tile decode block tables after _make_kv_cache_for_decode updates
+    # block_table in-place. Without this, the tiled decode_block_tables
+    # (created during build_cp_attn_metadata) still contains the original
+    # sequential block indices, causing FIA to read from empty cache blocks.
+    if attn_metadata.num_decodes_flatten > attn_metadata.num_decodes:
+        tiled_rows = []
+        for i in range(attn_metadata.num_decodes):
+            q_len = batch_spec.query_lens[i]
+            row = attn_metadata.block_tables[i : i + 1]
+            tiled_rows.append(row.repeat(q_len, 1))
+        attn_metadata.decode_meta.block_tables = torch.cat(tiled_rows, dim=0)
+
+    kv_cache = (k_cache, v_cache)
+
+    impl = _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale)
+    output = run_cp_attention(query_vllm, key_vllm, value_vllm, kv_cache, attn_metadata, impl, device, vllm_config)
+
+    _assert_close(output, sdpa_output, rtol, atol, "CP_Decode")
+
+
+# ---------------------------------------------------------------------------
+# Mixed decode + prefill (prefill sequences have seq_lens == query_lens)
+#
+# The CP attention non-chunked prefill path calls FIA directly with only the
+# new tokens as KV. Context tokens in the KV cache are NOT loaded for prefill
+# without chunked context. Therefore prefill sequences in mixed mode must
+# have seq_lens == query_lens (no context).
+# ---------------------------------------------------------------------------
+def _test_cp_mixed_precision_no_cp(
+    batch_spec: BatchSpec,
+    model: str,
+    *,
+    block_size: int = 128,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+):
+    set_random_seed(42)
+
+    vllm_config = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=1,
+        max_model_len=max(batch_spec.seq_lens) + block_size,
+        block_size=block_size,
+        num_gpu_blocks=8192,
+    )
+    device = torch.device("npu")
+
+    num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+    num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+    head_size = vllm_config.model_config.get_head_size()
+    dtype = torch.bfloat16
+    scale = 1.0 / (head_size**0.5)
+
+    total_kv = sum(batch_spec.seq_lens)
+
+    query_full = torch.randn(total_kv, num_q_heads, head_size, dtype=dtype, device=device)
+    key_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
+    value_full = torch.randn(total_kv, num_kv_heads, head_size, dtype=dtype, device=device)
+
+    sdpa_output = compute_mixed_sdpa_reference(
+        query_full,
+        key_full,
+        value_full,
+        batch_spec,
+        scale,
+        num_q_heads,
+        num_kv_heads,
+    )
+
+    # Backend inputs: only the new tokens
+    all_q, all_k, all_v = [], [], []
+    q_offset = 0
+    kv_offset = 0
+    for i in range(batch_spec.batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        q_i = query_full[q_offset : q_offset + q_len].contiguous()
+        k_full_i = key_full[kv_offset : kv_offset + s_len]
+        v_full_i = value_full[kv_offset : kv_offset + s_len]
+
+        all_q.append(q_i)
+        all_k.append(k_full_i[context_len:])
+        all_v.append(v_full_i[context_len:])
+
+        q_offset += q_len
+        kv_offset += s_len
+
+    query_vllm = torch.cat(all_q, dim=0)
+    key_vllm = torch.cat(all_k, dim=0)
+    value_vllm = torch.cat(all_v, dim=0)
+
+    attn_metadata = build_cp_attn_metadata(
+        batch_spec,
+        vllm_config,
+        device,
+        pcp_size=1,
+        pcp_rank=0,
+        kv_cache_prepopulated=True,
+    )
+
+    k_cache, v_cache = _make_kv_cache_for_mixed(
+        batch_spec,
+        num_kv_heads,
+        head_size,
+        block_size,
+        dtype,
+        device,
+        key_full,
+        value_full,
+        attn_metadata.block_tables,
+    )
+    kv_cache = (k_cache, v_cache)
+
+    impl = _create_cp_impl(vllm_config, device, num_q_heads, num_kv_heads, head_size, scale)
+    output = run_cp_attention(query_vllm, key_vllm, value_vllm, kv_cache, attn_metadata, impl, device, vllm_config)
+
+    _assert_close(output, sdpa_output, rtol, atol, "CP_Mixed")
+
+
+class TestCPAttentionPrecision:
+    """Precision tests for AscendAttentionCPImpl.
+
+    Validates that CP attention produces results matching
+    PyTorch SDPA within 1e-2 tolerance.
+
+    Test scenarios:
+    - Pure prefill (seq_lens == query_lens), PCP=1, DCP=1
+    - Pure decode, PCP=1, DCP=1 (context in KV cache)
+    - Mixed decode+prefill (prefill has seq_lens == query_lens), PCP=1, DCP=1
+    - MTP (Multi-Token Prediction) decode, PCP=1, DCP=1
+    """
+
+    @pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
+    @pytest.mark.parametrize(
+        "batch_spec_name",
+        [
+            "single_prefill",
+            "small_prefill",
+            "medium_prefill",
+            "large_prefill",
+        ],
+    )
+    @pytest.mark.parametrize("model", MODELS)
+    @patch_distributed_groups(dcp_size=1, pcp_size=1)
+    def test_cp_prefill_precision(
+        self,
+        mock_all2all,
+        mock_dcp,
+        mock_pcp,
+        batch_spec_name,
+        model,
+    ):
+        batch_spec = BATCH_SPECS[batch_spec_name]
+        _test_cp_prefill_precision_no_cp(batch_spec, model)
+
+    @pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
+    @pytest.mark.parametrize(
+        "batch_spec_name",
+        [
+            "single_decode",
+            "small_decode",
+            "medium_decode",
+        ],
+    )
+    @pytest.mark.parametrize("model", MODELS)
+    @patch_distributed_groups(dcp_size=1, pcp_size=1)
+    def test_cp_decode_precision(
+        self,
+        mock_all2all,
+        mock_dcp,
+        mock_pcp,
+        batch_spec_name,
+        model,
+    ):
+        batch_spec = BATCH_SPECS[batch_spec_name]
+        _test_cp_decode_precision_no_cp(batch_spec, model)
+
+    @pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
+    @pytest.mark.parametrize(
+        "batch_spec_name",
+        [
+            "mixed_small",
+            "mixed_medium",
+            "mixed_large",
+        ],
+    )
+    @pytest.mark.parametrize("model", MODELS)
+    @patch_distributed_groups(dcp_size=1, pcp_size=1)
+    def test_cp_mixed_precision(
+        self,
+        mock_all2all,
+        mock_dcp,
+        mock_pcp,
+        batch_spec_name,
+        model,
+    ):
+        batch_spec = BATCH_SPECS[batch_spec_name]
+        _test_cp_mixed_precision_no_cp(batch_spec, model)
+
+    @pytest.mark.skip(reason="Waiting for rebuild with irregular mask")
+    @pytest.mark.parametrize(
+        "batch_spec_name",
+        [
+            "mtp_1_plus_3_tiny",
+            "mtp_1_plus_3_small",
+            "mtp_1_plus_3_medium",
+        ],
+    )
+    @pytest.mark.parametrize("model", MODELS)
+    @patch_distributed_groups(dcp_size=1, pcp_size=1)
+    def test_cp_mtp_decode_precision(
+        self,
+        mock_all2all,
+        mock_dcp,
+        mock_pcp,
+        batch_spec_name,
+        model,
+    ):
+        """MTP decode: each request produces 1 target + 3 speculative tokens.
+
+        All tokens (context + new) are pre-populated in the KV cache.
+        FIA paged attention with actual_seq_lengths_kv enforces causal
+        masking per token within each request.
+        """
+        batch_spec = BATCH_SPECS[batch_spec_name]
+        _test_cp_decode_precision_no_cp(batch_spec, model, decode_threshold=4)

tests/ut/attention/a2/test_attention_v1.py

@@ -0,0 +1,775 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+
+import vllm_ascend.attention.attention_v1 as attn_module
+from tests.ut.base import TestBase
+from vllm_ascend.attention.attention_v1 import (
+    AscendAttentionBackend,
+    AscendAttentionBackendImpl,
+    AscendAttentionMetadataBuilder,
+    AscendAttentionState,
+)
+from vllm_ascend.attention.kvcomp_attn.attention_utils import get_kvcomp_decode_params, reshape_and_cache_kvcomp
+from vllm_ascend.attention.utils import (
+    AscendCommonAttentionMetadata,
+    cache_graph_workspace,
+    needs_layer_aware_fia_graph_replay,
+    using_paged_attention,
+)
+from vllm_ascend.device.device_op import A5DeviceAdaptor
+from vllm_ascend.device.utils import FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE
+from vllm_ascend.utils import AscendDeviceType
+
+LARGE_HEAD_PREFILL_PATH = "vllm_ascend.device.utils.npu_large_head_prefill_attention"
+
+
+class TestAttentionGraphHelpers(TestBase):
+    def test_cache_graph_workspace_keeps_first_workspace_by_default(self):
+        graph_params = SimpleNamespace(workspaces={1: torch.empty(4)})
+        candidate_workspace = torch.empty(8)
+
+        result = cache_graph_workspace(graph_params, 1, candidate_workspace, use_max_workspace=False)
+
+        self.assertEqual(result.numel(), 4)
+        self.assertEqual(graph_params.workspaces[1].numel(), 4)
+
+    def test_cache_graph_workspace_updates_to_larger_workspace(self):
+        graph_params = SimpleNamespace(workspaces={1: torch.empty(4)})
+        candidate_workspace = torch.empty(8)
+
+        result = cache_graph_workspace(graph_params, 1, candidate_workspace, use_max_workspace=True)
+
+        self.assertEqual(result.numel(), 8)
+        self.assertEqual(graph_params.workspaces[1].numel(), 8)
+
+    def test_large_head_uses_paged_attention_on_a2(self):
+        vllm_config = MagicMock()
+        vllm_config.speculative_config = None
+        with patch("vllm_ascend.attention.utils.get_ascend_device_type", return_value=AscendDeviceType.A2):
+            self.assertTrue(using_paged_attention(1, vllm_config, head_size=FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE))
+
+
+class TestAscendAttentionBackend(TestBase):
+    def setUp(self):
+        self.mock_config = MagicMock()
+
+        mock_parallel_config = MagicMock()
+        mock_parallel_config.prefill_context_parallel_size = 1
+        mock_parallel_config.decode_context_parallel_size = 1
+
+        self.mock_config.parallel_config = mock_parallel_config
+
+        self.utils_patcher = patch("vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_config)
+        self.utils_patcher.start()
+
+        from vllm_ascend.attention.utils import enable_cp
+
+        enable_cp.cache_clear()
+
+    def test_get_name(self):
+        self.assertEqual(AscendAttentionBackend.get_name(), "CUSTOM")
+
+    def test_get_impl_cls(self):
+        self.assertEqual(AscendAttentionBackend.get_impl_cls(), AscendAttentionBackendImpl)
+
+    def test_get_builder_cls(self):
+        self.assertEqual(AscendAttentionBackend.get_builder_cls(), AscendAttentionMetadataBuilder)
+
+    def test_get_kv_cache_shape_not(self):
+        result = AscendAttentionBackend.get_kv_cache_shape(10, 20, 30, 40)
+        self.assertEqual(result, (2, 10, 20, 30, 40))
+
+    def test_swap_blocks(self):
+        src_kv_cache = [torch.zeros((10, 20)), torch.zeros((10, 20))]
+        dst_kv_cache = [torch.zeros((10, 20)), torch.zeros((10, 20))]
+        src_to_dst = torch.tensor([[0, 1], [2, 3]])
+        AscendAttentionBackend.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
+        self.assertTrue(torch.all(dst_kv_cache[0][1] == src_kv_cache[0][0]))
+        self.assertTrue(torch.all(dst_kv_cache[1][3] == src_kv_cache[1][2]))
+
+    def test_copy_blocks(self):
+        kv_caches = [torch.zeros((10, 20)), torch.zeros((10, 20))]
+        src_to_dists = torch.tensor([[0, 1], [2, 3]])
+        AscendAttentionBackend.copy_blocks(kv_caches, src_to_dists)
+        self.assertTrue(torch.all(kv_caches[0][1] == kv_caches[0][0]))
+        self.assertTrue(torch.all(kv_caches[1][3] == kv_caches[1][2]))
+
+
+class TestAscendAttentionMetadataBuilder(TestBase):
+    def setUp(self):
+        self.mock_vllm_config = MagicMock()
+        self.mock_vllm_config.speculative_config = None
+        self.mock_vllm_config.model_config.max_model_len = 640
+        self.mock_vllm_config.model_config.hf_text_config.sliding_window = None
+        self.mock_vllm_config.cache_config.block_size = 64
+        self.mock_vllm_config.compilation_config.cudagraph_mode = None
+        self.mock_vllm_config.scheduler_config.max_num_seqs = 10
+        self.mock_vllm_config.scheduler_config.chunked_prefill_enabled = False
+        self.mock_device = "cpu:0"
+        torch.Tensor.pin_memory = lambda x: x  # noqa
+        self.builder = AscendAttentionMetadataBuilder(None, None, self.mock_vllm_config, self.mock_device)
+
+    def test_reorder_batch(self):
+        mock_input_batch = MagicMock()
+        mock_scheduler_output = MagicMock()
+
+        result = self.builder.reorder_batch(mock_input_batch, mock_scheduler_output)
+
+        self.assertFalse(result)
+
+    def test_unpadded_preserves_internal_seq_lens_cpu(self):
+        internal_seq_lens_cpu = torch.tensor([4, 5, 6], dtype=torch.int32)
+        common_attn_metadata = AscendCommonAttentionMetadata(
+            query_start_loc=torch.tensor([0, 2, 5, 9]),
+            query_start_loc_cpu=torch.tensor([0, 2, 5, 9]),
+            seq_lens=torch.tensor([4, 5, 6], dtype=torch.int32),
+            _seq_lens_cpu=internal_seq_lens_cpu,
+            seq_lens_cpu=None,
+            num_computed_tokens_cpu=None,
+            num_reqs=3,
+            num_actual_tokens=9,
+            max_query_len=4,
+            block_table_tensor=torch.zeros((3, 1), dtype=torch.int32),
+            slot_mapping=torch.arange(9, dtype=torch.int32),
+            causal=True,
+            actual_seq_lengths_q=[2, 3, 4],
+            positions=torch.arange(9),
+            attn_state=AscendAttentionState.ChunkedPrefill,
+            max_seq_len=6,
+        )
+
+        unpadded_metadata = common_attn_metadata.unpadded(num_actual_tokens=5, num_actual_reqs=2)
+
+        self.assertTrue(torch.equal(unpadded_metadata._seq_lens_cpu, internal_seq_lens_cpu[:2]))
+        self.assertIsNone(unpadded_metadata.seq_lens_cpu)
+
+    @patch("vllm_ascend.attention.attention_v1.AscendMetadata")
+    def test_build(self, mock_ascend_metadata):
+        common_attn_metadata = AscendCommonAttentionMetadata(
+            query_start_loc=torch.tensor([0, 2, 5, 9]),
+            query_start_loc_cpu=torch.tensor([0, 2, 5, 9]),
+            seq_lens_cpu=torch.tensor([4, 5, 6]),
+            num_reqs=3,
+            num_actual_tokens=15,
+            max_query_len=6,
+            decode_token_per_req=torch.tensor([1, 1, 1]),
+            block_table_tensor=torch.zeros((10, 10)),
+            slot_mapping=torch.tensor(range(20)),
+            actual_seq_lengths_q=torch.tensor([0, 1, 2]),
+            positions=torch.tensor([10, 10]),
+            attn_state=AscendAttentionState.ChunkedPrefill,
+            num_computed_tokens_cpu=None,
+            seq_lens=None,
+            max_seq_len=6,
+        )
+        mock_model = MagicMock()
+
+        self.builder.build(1, common_attn_metadata, mock_model)
+
+
+class TestAscendAttentionBackendImpl(TestBase):
+    def setUp(self):
+        self.mock_event = MagicMock()
+        self.mock_event.record.return_value = None
+        self.mock_event.wait.return_value = None
+
+        self.mock_stream = MagicMock()
+        self.event_patcher = patch("torch_npu.npu.Event", return_value=self.mock_event)
+        self.stream_patcher = patch("torch_npu.npu.current_stream", return_value=self.mock_stream)
+
+        self.event_patcher.start()
+        self.stream_patcher.start()
+
+        self.layer = MagicMock()
+        self.layer.layer_name = "test_layer"
+        self.layer._k_scale_float = 1.0
+        self.layer._v_scale_float = 1.0
+        self.attention_type = MagicMock()
+        self.attention_type.DECODER = "decoder"
+        self.attention_type.ENCODER = "encoder"
+        self.attn_metadata = MagicMock()
+        self.attn_metadata.return_value = "1"
+        self.layer_no_quant = MagicMock(spec=["layer_name", "_k_scale_float", "_v_scale_float"])
+        self.layer_no_quant.layer_name = "test_layer"
+        self.layer_no_quant._k_scale_float = 1.0
+        self.layer_no_quant._v_scale_float = 1.0
+        self.mock_vllm_config = MagicMock()
+        self.config_patcher = patch(
+            "vllm_ascend.attention.attention_v1.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.utils_config_patcher = patch(
+            "vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_vllm_config
+        )
+        self.config_patcher.start()
+        self.utils_config_patcher.start()
+        needs_layer_aware_fia_graph_replay.cache_clear()
+        self.addCleanup(needs_layer_aware_fia_graph_replay.cache_clear)
+        self.addCleanup(self.utils_config_patcher.stop)
+        self.addCleanup(self.config_patcher.stop)
+
+        self.impl = AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=64,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+        self.impl_192 = AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=192,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+        self.impl_error = AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=192,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=None,
+            kv_sharing_target_layer_name=None,
+        )
+
+        self.impl_swa = AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=64,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=1024,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+        self.impl_swa_sink = AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=64,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=1024,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+            sinks=torch.tensor([-3.4062], dtype=torch.bfloat16),
+        )
+
+        self.impl_large_head = AscendAttentionBackendImpl(
+            num_heads=8,
+            head_size=FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE,
+            scale=1.0,
+            num_kv_heads=8,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype="float16",
+            logits_soft_cap=None,
+            attn_type=self.attention_type.DECODER,
+            kv_sharing_target_layer_name=None,
+        )
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_large_head_prefill_uses_device_operator_fallback(self, mock_get_forward_context):
+        query = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        key = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        value = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.PrefillNoCache
+        metadata.actual_seq_lengths_q = [2]
+        metadata.causal = True
+        metadata.attn_mask = None
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        with patch(LARGE_HEAD_PREFILL_PATH, return_value=(torch.ones_like(query), None)) as mock_forward:
+            result = self.impl_large_head.forward_impl(query, key, value, (), metadata, output)
+
+        mock_forward.assert_called_once()
+        self.assertIs(result, output)
+        self.assertTrue(torch.equal(result, torch.ones_like(query)))
+
+    def test_supported_head_prefill_uses_fia(self):
+        query = torch.randn(2, 8, 64)
+        key = torch.randn(2, 8, 64)
+        value = torch.randn(2, 8, 64)
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.PrefillNoCache
+        metadata.actual_seq_lengths_q = [2]
+
+        self.impl.forward_fused_infer_attention = MagicMock(return_value=output)
+        with patch(LARGE_HEAD_PREFILL_PATH, return_value=(torch.empty_like(query), None)) as mock_forward:
+            result = self.impl.forward_impl(query, key, value, (), metadata, output)
+
+        mock_forward.assert_not_called()
+        self.impl.forward_fused_infer_attention.assert_called_once()
+        self.assertIs(result, output)
+
+    @patch("vllm_ascend.attention.attention_v1.using_paged_attention", return_value=True)
+    def test_decode_uses_paged_attention(self, mock_using_pa):
+        query = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+
+        self.impl_large_head.forward_paged_attention = MagicMock(return_value=output)
+        self.impl_large_head.forward_fused_infer_attention = MagicMock(return_value=output)
+
+        result = self.impl_large_head.forward_impl(query, None, None, (), metadata, output)
+
+        self.impl_large_head.forward_paged_attention.assert_called_once()
+        self.impl_large_head.forward_fused_infer_attention.assert_not_called()
+        self.assertIs(result, output)
+        mock_using_pa.assert_called_once()
+
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    def test_a5_device_operator_uses_fia_for_large_head(self, mock_fia):
+        query = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        key = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        value = torch.randn(2, 8, FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.PrefillNoCache
+        metadata.actual_seq_lengths_q = [2]
+
+        mock_fia.return_value = (torch.ones_like(query), None)
+        with patch(LARGE_HEAD_PREFILL_PATH, return_value=(torch.empty_like(query), None)) as mock_forward:
+            result = A5DeviceAdaptor.npu_fused_infer_attention_score(
+                query=query,
+                key=key,
+                value=value,
+                attn_metadata=metadata,
+                key_cache=None,
+                value_cache=None,
+                current_key=key,
+                current_value=value,
+                num_heads=8,
+                num_key_value_heads=8,
+                head_size=FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE,
+                scale=1.0,
+                is_prefill_no_cache=True,
+                block_table=None,
+                input_layout="TND",
+                block_size=128,
+                actual_seq_lengths=[2],
+                actual_seq_lengths_kv=[2],
+                sparse_mode=3,
+            )
+
+        mock_forward.assert_not_called()
+        mock_fia.assert_called_once()
+        self.assertEqual(result[0].shape, query.shape)
+
+    def test_forward_no_attn_metadata(self):
+        """Test forward pass when attn_metadata is None"""
+        query = torch.randn(10, 8 * 64)
+        key = torch.randn(10, 8 * 64)
+        value = torch.randn(10, 8 * 64)
+        kv_cache = torch.empty(2, 0, 0, 8, 64)
+        layer = self.layer_no_quant
+        output = torch.empty_like(query)
+
+        output = self.impl.forward(layer, query, key, value, kv_cache, None, output)
+
+        assert output.shape == (10, 8 * 64)
+
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_forward_fused_infer_attention(
+        self, mock_get_forward_context, mock_npu_fused_infer_attention_score, mock_npu_reshape_and_cache
+    ):
+        """Test forward pass in PrefillCacheHit state"""
+        query = torch.randn(10, 8, 64)
+        key = torch.randn(10, 8, 64)
+        value = torch.randn(10, 8, 64)
+        kv_cache = torch.empty(2, 5, 128, 8, 64)
+        output = torch.empty_like(query)
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.PrefillCacheHit
+        metadata.attn_mask = torch.randn(1, 1, 10, 10)
+        metadata.query_lens = torch.tensor([10])
+        metadata.seq_lens = torch.tensor([10])
+        metadata.actual_seq_lengths_q = [10]
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 10
+        metadata.num_decode_tokens = 0
+        metadata.num_decodes = 0
+        metadata.num_prefills = 10
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+        layer = self.layer_no_quant
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+        mock_npu_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), torch.ones(10, 8, 64))
+        output = self.impl.forward(layer, query, key, value, kv_cache, metadata, output)
+
+        mock_npu_fused_infer_attention_score.assert_called_once()
+        assert output.shape == (10, 8, 64)
+
+    @patch("vllm_ascend.attention.attention_v1.using_paged_attention")
+    @patch("torch_npu._npu_paged_attention")
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    def test_forward_paged_attention(
+        self, mock_get_forward_context, mock_npu_reshape_and_cache, mock_paged_attention, mock_using_paged_attention
+    ):
+        """Test forward pass in DecodeOnly state"""
+        query = torch.randn(4, 8 * 64)
+        key = torch.randn(4, 8 * 64)
+        value = torch.randn(4, 8 * 64)
+        kv_cache = torch.empty(2, 5, 128, 8, 64)
+        output = torch.empty_like(query)
+
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.seq_lens = torch.tensor([4])
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 4
+        metadata.slot_mapping = torch.zeros(4, dtype=torch.long)
+        metadata.num_decodes = 4
+        metadata.num_prefills = 0
+        layer = self.layer_no_quant
+        mock_using_paged_attention.return_value = True
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        output = self.impl.forward(layer, query, key, value, kv_cache, metadata, output)
+
+        mock_paged_attention.assert_called_once()
+        assert output.shape == (4, 8 * 64)
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("torch_npu._npu_reshape_and_cache")
+    def test_forward_decode_only_swa(
+        self, mock_npu_reshape_and_cache, mock_fused_infer_attention_score, mock_get_forward_context
+    ):
+        """Test forward pass in DecodeOnly state"""
+        query = torch.randn(10, 8 * 64)
+        key = torch.randn(10, 8 * 64)
+        value = torch.randn(10, 8 * 64)
+        kv_cache = torch.empty(2, 5, 128, 8, 64)
+        output = torch.empty(10, 8, 64)
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.seq_lens = torch.tensor([10] * 10)
+        metadata.actual_seq_lengths_q = [10]
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 100
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+        metadata.num_decodes = 10
+        metadata.num_prefills = 0
+        layer = self.layer_no_quant
+        mock_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), 1)
+        output = self.impl_swa.forward(layer, query, key, value, kv_cache, metadata, output)
+        print(output.shape)
+        mock_fused_infer_attention_score.assert_called_once()
+        assert output.shape == (10, 8, 64)
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("torch_npu.npu_fused_infer_attention_score_v2")
+    @patch("torch_npu._npu_reshape_and_cache")
+    def test_forward_decode_only_swa_sink(
+        self, mock_npu_reshape_and_cache, mock_fused_infer_attention_score, mock_get_forward_context
+    ):
+        """Test forward pass in DecodeOnly state"""
+        query = torch.randn(10, 8 * 64)
+        key = torch.randn(10, 8 * 64)
+        value = torch.randn(10, 8 * 64)
+        kv_cache = torch.empty(2, 5, 128, 8, 64)
+        output = torch.empty(10, 8, 64)
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.seq_lens = torch.tensor([10] * 10)
+        metadata.attn_mask = torch.randn(1, 1, 10, 10)
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 100
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+        metadata.num_decodes = 10
+        metadata.num_prefills = 0
+        layer = self.layer_no_quant
+        mock_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), 1)
+        output = self.impl_swa_sink.forward(layer, query, key, value, kv_cache, metadata, output)
+        print(output.shape)
+        mock_fused_infer_attention_score.assert_called_once()
+        assert output.shape == (10, 8, 64)
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("torch_npu._npu_paged_attention")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("torch_npu._npu_reshape_and_cache")
+    def test_forward_decode_only_swa_seq_len_mismatch(
+        self,
+        mock_npu_reshape_and_cache,
+        mock_fused_infer_attention_score,
+        mock_paged_attention,
+        mock_get_forward_context,
+    ):
+        """Test forward pass in DecodeOnly state when seq)len_mismatch"""
+        query = torch.randn(10, 8, 64)
+        key = torch.randn(10, 8, 64)
+        value = torch.randn(10, 8, 64)
+        kv_cache = torch.empty(2, 5, 128, 8, 64)
+        output = torch.empty_like(query)
+
+        metadata = self.attn_metadata
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.seq_lens = torch.tensor([10])  # len == 1 != query.size(0)==10
+        metadata.block_tables = torch.zeros(1, 5, dtype=torch.long)
+        metadata.num_actual_tokens = 10
+        metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
+        layer = self.layer_no_quant
+        metadata.num_decodes = 10
+        metadata.num_prefills = 0
+        metadata.actual_seq_lengths_q = [10]
+
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        mock_fused_infer_attention_score.return_value = (torch.ones(10, 8, 64), torch.ones(10, 8, 64))
+
+        output = self.impl_swa.forward(layer, query, key, value, kv_cache, metadata, output)
+
+        mock_paged_attention.assert_not_called()
+        mock_fused_infer_attention_score.assert_called_once()
+
+        assert output.shape == (10, 8, 64)
+
+    def test_get_kvcomp_params_early_exit(self):
+        """
+        Test that get_kvcomp_decode_params returns original values
+        when kvcomp is disabled or hashk_cache is missing.
+        """
+        query = torch.randn(10, 8, 64)
+        key = torch.randn(10, 8, 64)
+        block_table = torch.zeros(1, 5, dtype=torch.long)
+        actual_seq_lengths_kv = [10]
+
+        metadata = MagicMock()
+        # Mocking the case where hashk_caches is not properly initialized
+        kvcomp_metadata = MagicMock()
+        kvcomp_metadata.hashk_caches = [None]
+        metadata.kvcomp_metadata = kvcomp_metadata
+
+        self.impl.enable_hamming_sparse = True
+        self.impl.layerIndex = 0
+
+        res_bt, res_sl = get_kvcomp_decode_params(0, kvcomp_metadata, query, key, block_table, actual_seq_lengths_kv)
+
+        self.assertIs(res_bt, block_table)
+        self.assertEqual(res_sl, actual_seq_lengths_kv)
+
+    def test_get_kvcomp_params_reuse(self):
+        """
+        Test that in DecodeOnly state, if the current layer is a skip layer,
+        it correctly reuses the Hamming results from a previous layer.
+        """
+        query = torch.randn(10, 8, 64)
+        key = torch.randn(10, 8, 64)
+        block_table = torch.zeros(1, 5, dtype=torch.long)
+        actual_seq_lengths_kv = [10]
+
+        self.impl.enable_hamming_sparse = True
+        self.impl.layerIndex = 1
+
+        metadata = MagicMock()
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        expected_bt = torch.ones(1, 5)
+        expected_sl = torch.tensor([5])
+
+        # Construct kvcomp_metadata
+        kvcomp_metadata = MagicMock()
+        kvcomp_metadata.hashk_caches = [MagicMock(), MagicMock()]
+        kvcomp_metadata.hamming_output = expected_bt
+        kvcomp_metadata.seq_lens_from_hamming = expected_sl
+
+        kvcomp_config = MagicMock()
+        kvcomp_config.vllm_hash_attention_skip_layers = [False, True]
+        kvcomp_config.top_k_index_reuse = [0, 0]
+        kvcomp_metadata.kvcomp_config = kvcomp_config
+        metadata.kvcomp_metadata = kvcomp_metadata
+
+        metadata.hamming_output_records = [{"new_block_table": expected_bt, "new_seq_lens_list": expected_sl}, None]
+
+        res_bt, res_sl = get_kvcomp_decode_params(1, kvcomp_metadata, query, key, block_table, actual_seq_lengths_kv)
+
+        self.assertTrue(torch.equal(res_bt, expected_bt))
+        self.assertTrue(torch.equal(res_sl, expected_sl))
+
+    @patch("torch.ops._C_ascend.npu_reshape_and_cache_bnsd", create=True)
+    def test_get_kvcomp_params_prefill(self, mock_reshape_and_cache):
+        """
+        Test that in non-DecodeOnly state (e.g., Prefill), only Hash compute
+        and Cache update are performed, and original params are returned.
+        """
+        key = torch.randn(2, 8, 64)
+
+        self.impl.enable_hamming_sparse = True
+        self.impl.layerIndex = 0
+
+        metadata = MagicMock()
+        metadata.attn_state = AscendAttentionState.PrefillCacheHit
+        metadata.slot_mapping = torch.zeros(2)
+        metadata.actual_seq_lengths_q_device = torch.tensor([1, 1])
+        metadata.num_actual_tokens = 2
+        metadata.actual_query_lens = torch.tensor([1, 1], dtype=torch.int32)
+        metadata.query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32)
+
+        kvcomp_metadata = MagicMock()
+        kvcomp_metadata.hashk_caches = [MagicMock()]
+        kvcomp_config = MagicMock()
+        kvcomp_config.vllm_hash_attention_skip_layers = [False]
+        kvcomp_metadata.kvcomp_config = kvcomp_config
+
+        # Mock HashEncoder
+        hash_encoder = MagicMock()
+        hash_encoder.compute_hash.return_value = torch.ones(2, 8, 8)
+        kvcomp_metadata.hash_encoder = hash_encoder
+
+        metadata.kvcomp_metadata = kvcomp_metadata
+
+        reshape_and_cache_kvcomp(kvcomp_metadata, 0, key)
+
+        # Ensure cache update was called but Hamming was bypassed
+        self.assertTrue(mock_reshape_and_cache.called)
+
+    @patch("torch.ops._C_ascend.npu_reshape_and_cache_bnsd", create=True)
+    @patch("torch.ops._C_ascend.npu_hamming_dist_top_k", create=True)
+    def test_get_kvcomp_params_decode_hamming(self, mock_hamming, mock_reshape):
+        """
+        Test that in DecodeOnly state, the full flow including Hash computation
+        and Hamming Distance Top-K operation is executed.
+        """
+        query = torch.randn(2, 8, 64)
+        key = torch.randn(2, 8, 64)
+        block_table = torch.zeros(2, 5, dtype=torch.long)
+        actual_seq_lengths_kv = [10, 10]
+
+        self.impl.enable_hamming_sparse = True
+        self.impl.layerIndex = 0
+
+        metadata = MagicMock()
+        metadata.attn_state = AscendAttentionState.DecodeOnly
+        metadata.seq_lens = torch.tensor([10, 10])
+        metadata.actual_seq_lengths_q_device = torch.tensor([1, 1])
+        metadata.slot_mapping = torch.zeros(2)
+        metadata.block_tables = block_table
+        metadata.hamming_output_records = [None]
+        metadata.num_actual_tokens = 2
+        metadata.actual_query_lens = torch.tensor([1, 1], dtype=torch.int32)
+        metadata.query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32)
+        metadata.chunk_sizes_for_hamming = torch.tensor([64, 64], dtype=torch.int32)
+        metadata.max_seq_len_for_hamming = 1024
+        metadata.block_tables_for_hamming = torch.zeros(2, 10, dtype=torch.int32)
+        metadata.new_seq_lens_list = torch.tensor([5, 5], dtype=torch.int32)
+
+        kvcomp_metadata = MagicMock()
+        kvcomp_metadata.hashk_caches = [MagicMock()]
+
+        kvcomp_config = MagicMock()
+        kvcomp_config.vllm_hash_attention_skip_layers = [False]
+        kvcomp_config.chunk_size = 64
+        kvcomp_metadata.kvcomp_config = kvcomp_config
+
+        # Mock necessary Hamming parameters
+        kvcomp_metadata.chunk_sizes_for_hamming_full = torch.tensor([1, 1])
+        kvcomp_metadata.topk_for_hamming_full = torch.tensor([1, 1])
+        kvcomp_metadata.topk_for_hamming_full_cpu = torch.tensor([1, 1])
+        kvcomp_metadata.hamming_output = torch.zeros(2, 1)
+
+        # Mock HashEncoder
+        hash_encoder = MagicMock()
+        hash_encoder.compute_hash.return_value = torch.ones(2, 8, 8)
+        kvcomp_metadata.hash_encoder = hash_encoder
+
+        metadata.kvcomp_metadata = kvcomp_metadata
+
+        # Mock npu_hamming_dist_top_k output; note the squeeze(1) in implementation
+        mock_hamming.return_value = torch.ones(2, 1, 5)
+
+        res_bt, res_sl = get_kvcomp_decode_params(0, kvcomp_metadata, query, key, block_table, actual_seq_lengths_kv)
+
+        self.assertTrue(mock_reshape.called)
+        self.assertTrue(mock_hamming.called)
+        # Verify shape after squeeze(1) becomes (2, 5)
+        self.assertEqual(res_bt.shape, (2, 5))
+        self.assertTrue(torch.equal(res_bt, torch.ones(2, 5)))
+        # Verify the result is recorded in hamming_output_records
+        self.assertTrue(torch.equal(kvcomp_metadata.hamming_output, torch.ones(2, 5)))
+        self.assertIsNotNone(kvcomp_metadata.seq_lens_for_hamming)
+
+    @patch("torch.npu.stream")
+    @patch("torch.npu.graph_task_update_begin")
+    @patch("torch.npu.graph_task_update_end")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("vllm_ascend.attention.attention_v1.get_graph_params")
+    @patch("vllm_ascend.attention.attention_v1._EXTRA_CTX")
+    @patch("vllm_ascend.attention.attention_v1.using_paged_attention", return_value=False)
+    @patch("vllm_ascend.attention.attention_v1.needs_layer_aware_fia_graph_replay", return_value=False)
+    @patch("vllm_ascend.attention.attention_v1._ATTN_KEYS_BUFFER", new=[])
+    def test_update_graph_params(
+        self,
+        mock_needs_layer_aware_fia_graph_replay,
+        mock_using_paged_attention,
+        mock_EXTRA_CTX,
+        mock_get_graph_params,
+        mock_fia,
+        mock_graph_task_update_end,
+        mock_graph_task_update_begin,
+        mock_stream,
+    ):
+        """Test behavior when _ATTN_KEYS_BUFFER is [] after dummy_run."""
+
+        mock_EXTRA_CTX.sinks = False
+        mock_EXTRA_CTX.is_draft_model = False
+
+        param: list[MagicMock | None] = [MagicMock()] * 22
+        param[16] = None  # sliding_window
+        param[17] = None  # c8_k_aq_scale
+        param[21] = None  # layer_name
+
+        mock_get_graph_params.return_value.attn_params = {1: [tuple(param)] * 3}
+        mock_get_graph_params.return_value.handles = {1: [MagicMock()] * 3}
+        mock_get_graph_params.return_value.events = {1: [MagicMock()] * 3}
+
+        attn_metadata_keys = [
+            "model.layers.10.self_attn.attn",
+            "model.layers.2.self_attn.attn",
+            "model.layers.5.self_attn.attn",
+        ]
+        forward_context = MagicMock()
+        forward_context.attn_metadata = {key: MagicMock() for key in attn_metadata_keys}
+        # breakpoint()
+        self.impl.update_graph_params(self.mock_stream, forward_context, 1, self.mock_vllm_config)
+
+        expected = [
+            "model.layers.2.self_attn.attn",
+            "model.layers.5.self_attn.attn",
+            "model.layers.10.self_attn.attn",
+        ]
+        self.assertEqual(attn_module._ATTN_KEYS_BUFFER, expected)
+        self.assertEqual(mock_fia.out.call_count, 3)

tests/ut/attention/a2/test_attention_v1_precision.py

@@ -0,0 +1,397 @@
+from functools import partial
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.config import set_current_vllm_config
+from vllm.forward_context import set_forward_context
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.selector import get_attn_backend
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+from tests.ut.attention.utils import (
+    BatchSpec,
+    create_and_prepopulate_kv_cache,
+    create_common_attn_metadata,
+    create_standard_kv_cache_spec,
+    create_vllm_config,
+)
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
+
+
+@pytest.fixture(autouse=True)
+def default_vllm_config():
+    mock_config = MagicMock()
+    mock_config.compilation_config = MagicMock()
+    mock_config.compilation_config.custom_ops = ["all"]
+    mock_config.parallel_config = MagicMock()
+    mock_config.parallel_config.prefill_context_parallel_size = 1
+    mock_config.parallel_config.decode_context_parallel_size = 1
+    mock_config.parallel_config.tensor_parallel_size = 1
+    mock_config.model_config = MagicMock()
+    mock_config.model_config.dtype = torch.float16
+    mock_config.speculative_config = None
+    mock_config.cache_config = MagicMock()
+    mock_config.cache_config.block_size = 128
+    mock_config.kv_transfer_config = None
+    mock_config.additional_config = None
+    mock_config.quant_config = None
+
+    with set_current_vllm_config(mock_config):
+        yield mock_config
+
+
+BATCH_SPECS = {
+    "small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
+    "small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
+    "mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
+    "medium_decode": BatchSpec(
+        seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
+        query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
+    ),
+    "medium_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]),
+    "mixed_medium": BatchSpec(seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]),
+    "large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
+    "large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
+    "mixed_large": BatchSpec(seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]),
+    "single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
+    "single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
+    "small_encoder_prefill": BatchSpec(seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]),
+    "medium_encoder_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]),
+    "mtp_1_plus_3": BatchSpec(seq_lens=[256, 512, 1024, 1536], query_lens=[4, 4, 4, 4]),
+    "mtp_1_plus_7": BatchSpec(seq_lens=[512, 1024, 2048, 3072], query_lens=[8, 8, 8, 8]),
+    "mtp_small": BatchSpec(seq_lens=[64, 128, 256], query_lens=[4, 4, 4]),
+}
+
+
+class MockAttentionLayer:
+    """A mock attention layer for testing."""
+
+    def __init__(self, device: torch.device):
+        self._q_scale = torch.tensor(1.0, device=device)
+        self._k_scale = torch.tensor(1.0, device=device)
+        self._v_scale = torch.tensor(1.0, device=device)
+        self._q_scale_float = 1.0
+        self._k_scale_float = 1.0
+        self._v_scale_float = 1.0
+        self.layer_name = "model.layers.0"
+
+
+def run_attention_backend(
+    kv_cache_spec: FullAttentionSpec,
+    layer_names: list[str],
+    vllm_config,
+    device: torch.device,
+    common_attn_metadata: AscendCommonAttentionMetadata,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    kv_cache: torch.Tensor,
+    dtype: torch.dtype,
+    attn_type: AttentionType = AttentionType.DECODER,
+    sliding_window: int | None = None,
+) -> torch.Tensor:
+    """Run attention computation using the specified backend's AttentionImpl."""
+
+    num_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+    num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+    head_size = vllm_config.model_config.get_head_size()
+    scale = 1.0 / (head_size**0.5)
+
+    backend = get_attn_backend(head_size, dtype, None, use_mla=False, use_sparse=False, use_mm_prefix=False)
+    impl_cls = backend.get_impl_cls()
+    builder_cls = backend.get_builder_cls()
+
+    builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
+    attn_metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=common_attn_metadata,
+    )
+
+    impl = impl_cls(
+        num_heads=num_heads,
+        head_size=head_size,
+        scale=scale,
+        num_kv_heads=num_kv_heads,
+        alibi_slopes=None,
+        sliding_window=sliding_window,
+        attn_type=attn_type.value if hasattr(attn_type, "value") else attn_type,
+        kv_cache_dtype="auto",
+        logits_soft_cap=None,
+        kv_sharing_target_layer_name=None,
+    )
+
+    mock_layer = MockAttentionLayer(device)
+    output = torch.empty_like(query)
+    output = impl.forward(mock_layer, query, key, value, kv_cache, attn_metadata, output=output)
+
+    return output
+
+
+def compute_sdpa_reference(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    k_contexts: list[torch.Tensor],
+    v_contexts: list[torch.Tensor],
+    batch_spec: BatchSpec,
+    num_q_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    scale: float,
+    device: torch.device,
+    attn_type: AttentionType = AttentionType.DECODER,
+) -> torch.Tensor:
+    """Compute reference attention output using SDPA as golden baseline."""
+
+    all_sdpa_outputs = []
+    query_offset = 0
+    kv_offset = 0
+
+    for i in range(batch_spec.batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        q_i = query[query_offset : query_offset + q_len]
+        k_new_i = key[kv_offset : kv_offset + q_len]
+        v_new_i = value[kv_offset : kv_offset + q_len]
+
+        k_full_i = torch.cat([k_contexts[i], k_new_i], dim=0)
+        v_full_i = torch.cat([v_contexts[i], v_new_i], dim=0)
+
+        q_sdpa_in = q_i.unsqueeze(0).transpose(1, 2)
+        k_sdpa_in = k_full_i.unsqueeze(0).transpose(1, 2)
+        v_sdpa_in = v_full_i.unsqueeze(0).transpose(1, 2)
+
+        if attn_type == AttentionType.ENCODER_ONLY:
+            attn_mask = None
+        else:
+            attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
+            causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
+            attn_mask[:, context_len:] = causal_mask
+
+        sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
+            q_sdpa_in,
+            k_sdpa_in,
+            v_sdpa_in,
+            attn_mask=attn_mask,
+            is_causal=False,
+            enable_gqa=(num_q_heads != num_kv_heads),
+            scale=scale,
+        )
+
+        all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))
+        query_offset += q_len
+        kv_offset += q_len
+
+    return torch.cat(all_sdpa_outputs, dim=0)
+
+
+def _test_npu_attention_correctness(
+    batch_spec: BatchSpec,
+    model: str,
+    *,
+    attn_type: AttentionType = AttentionType.DECODER,
+    block_size: int = 128,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+    tensor_parallel_size: int = 1,
+):
+    """Test attention backend correctness with SDPA as reference."""
+
+    set_random_seed(42)
+
+    hf_config_override = None
+    if tensor_parallel_size > 1:
+        from vllm.config import ModelConfig
+
+        temp_config = ModelConfig(model=model, max_model_len=1)
+        original_num_heads = temp_config.hf_text_config.num_attention_heads
+        original_num_kv_heads = getattr(temp_config.hf_text_config, "num_key_value_heads", None)
+        hf_config_override = {
+            "num_attention_heads": original_num_heads // tensor_parallel_size,
+        }
+        if original_num_kv_heads is not None:
+            hf_config_override["num_key_value_heads"] = max(1, original_num_kv_heads // tensor_parallel_size)
+
+    vllm_config = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=1,
+        max_model_len=max(batch_spec.seq_lens),
+        block_size=block_size,
+        num_gpu_blocks=8192,
+        hf_config_override=hf_config_override,
+    )
+    device = torch.device("npu")
+
+    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+
+    batch_size = batch_spec.batch_size
+    num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+    num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+    head_size = vllm_config.model_config.get_head_size()
+    sliding_window = vllm_config.model_config.get_sliding_window()
+    dtype = torch.bfloat16
+    scale = 1.0 / (head_size**0.5)
+
+    k_contexts, v_contexts = [], []
+    all_q, all_k, all_v = [], [], []
+
+    for i in range(batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        q = torch.randn(q_len, num_q_heads, head_size, dtype=dtype, device=device)
+        k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
+        v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
+
+        all_q.append(q)
+        all_k.append(k_full[context_len:])
+        all_v.append(v_full[context_len:])
+
+        k_contexts.append(k_full[:context_len])
+        v_contexts.append(v_full[:context_len])
+
+    query = torch.cat(all_q, dim=0)
+    key = torch.cat(all_k, dim=0)
+    value = torch.cat(all_v, dim=0)
+
+    common_attn_metadata = create_common_attn_metadata(batch_spec, vllm_config.cache_config.block_size, device)
+    if attn_type == AttentionType.ENCODER_ONLY:
+        common_attn_metadata.causal = False
+
+    kv_cache = create_and_prepopulate_kv_cache(
+        k_contexts=k_contexts,
+        v_contexts=v_contexts,
+        block_size=block_size,
+        num_kv_heads=num_kv_heads,
+        head_size=head_size,
+        dtype=dtype,
+        device=device,
+        num_blocks=8192,
+        common_attn_metadata=common_attn_metadata,
+        randomize_blocks=False,
+    )
+
+    with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
+        from vllm.forward_context import get_forward_context
+
+        forward_ctx = get_forward_context()
+        forward_ctx.num_tokens = query.shape[0]
+        forward_ctx.is_draft_model = False
+        forward_ctx.is_draft_model_prefill = False
+        forward_ctx.capturing = False
+        forward_ctx.flash_comm_v1_enabled = False
+        forward_ctx.flashcomm_v2_enabled = False
+
+        backend_output = run_attention_backend(
+            kv_cache_spec,
+            ["placeholder"],
+            vllm_config,
+            device,
+            common_attn_metadata,
+            query,
+            key,
+            value,
+            kv_cache,
+            dtype,
+            sliding_window=sliding_window,
+            attn_type=attn_type,
+        )
+
+    sdpa_output = compute_sdpa_reference(
+        query,
+        key,
+        value,
+        k_contexts,
+        v_contexts,
+        batch_spec,
+        num_q_heads,
+        num_kv_heads,
+        head_size,
+        scale,
+        device,
+        attn_type=attn_type,
+    )
+
+    name = "GQA"
+    assert backend_output.shape == sdpa_output.shape, (
+        f"[{name}] shape {backend_output.shape} != SDPA shape {sdpa_output.shape}"
+    )
+    assert backend_output.dtype == sdpa_output.dtype, (
+        f"[{name}] dtype {backend_output.dtype} != SDPA dtype {sdpa_output.dtype}"
+    )
+
+    assert torch.isfinite(backend_output).all(), f"[{name}] produced non-finite values"
+
+    # Calculate and print differences for debugging
+    diff = torch.abs(backend_output - sdpa_output)
+    max_diff = diff.max().item()
+    mean_diff = diff.mean().item()
+    print(f"\n[{name}] Max difference: {max_diff:.6f}, Mean difference: {mean_diff:.6f}")
+    print(f"[{name}] Backend output range: [{backend_output.min().item():.6f}, {backend_output.max().item():.6f}]")
+    print(f"[{name}] SDPA output range: [{sdpa_output.min().item():.6f}, {sdpa_output.max().item():.6f}]")
+
+    def error_msg(msg: str, backend_name: str):
+        return f"[{backend_name}] output differs from SDPA baseline. {msg}"
+
+    torch.testing.assert_close(
+        backend_output,
+        sdpa_output,
+        rtol=rtol,
+        atol=atol,
+        msg=partial(error_msg, backend_name="GQA"),
+    )
+
+
+@pytest.mark.parametrize(
+    "batch_spec_name",
+    [
+        "small_decode",
+        "small_prefill",
+        "mixed_small",
+        "medium_decode",
+        "medium_prefill",
+        "mixed_medium",
+        "large_decode",
+        "large_prefill",
+        "single_decode",
+        "single_prefill",
+        "mtp_1_plus_3",
+        "mtp_1_plus_7",
+        "mtp_small",
+    ],
+)
+@pytest.mark.parametrize("model", ["Qwen/Qwen3-8B"])
+@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
+def test_causal_backend_correctness(default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int):
+    """Test backend's correctness with causal attention."""
+    batch_spec = BATCH_SPECS[batch_spec_name]
+
+    _test_npu_attention_correctness(
+        batch_spec,
+        model,
+        tensor_parallel_size=tensor_parallel_size,
+    )
+
+
+@pytest.mark.parametrize(
+    "batch_spec_name",
+    [
+        "small_encoder_prefill",
+        "medium_encoder_prefill",
+    ],
+)
+@pytest.mark.parametrize("model", ["Qwen/Qwen3-8B"])
+def test_encoder_only_backend_correctness(default_vllm_config, batch_spec_name: str, model: str):
+    """Test backend's correctness with encoder-only attention."""
+    batch_spec = BATCH_SPECS[batch_spec_name]
+
+    _test_npu_attention_correctness(
+        batch_spec,
+        model,
+        attn_type=AttentionType.ENCODER_ONLY,
+    )

tests/ut/attention/a2/test_common_cp.py

@@ -0,0 +1,138 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from vllm_ascend.attention.context_parallel.common_cp import (
+    _npu_attention_update,
+    _npu_attn_out_lse_update,
+    _update_out_and_lse,
+)
+
+
+class TestCommonCP(unittest.TestCase):
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_decode_context_model_parallel_world_size")
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_dcp_group")
+    @patch("torch.distributed.all_to_all_single")
+    def test_process_attn_out_lse_complex(
+        self, mock_all2all, mock_get_dcp_group, mock_get_dcp_size, mock_get_pcp_group
+    ):
+        from vllm_ascend.attention.context_parallel.common_cp import _process_attn_out_lse
+
+        pcp_size = 2
+        dcp_size = 2
+        mock_get_pcp_group.return_value.world_size = pcp_size
+        mock_get_dcp_size.return_value = dcp_size
+
+        mock_group = MagicMock()
+        mock_get_dcp_group.return_value.device_group = mock_group
+
+        bs, num_heads, head_dim = 4, 8, 64
+        attn_output = torch.randn(bs, num_heads, head_dim, dtype=torch.float16)
+        softmax_lse = torch.randn(bs, num_heads, 1, dtype=torch.float16)
+
+        def fake_all_gather(tensor, dim=0):
+            return torch.cat([tensor, tensor], dim=dim)
+
+        mock_get_pcp_group.return_value.all_gather = fake_all_gather
+
+        output = _process_attn_out_lse(attn_output, softmax_lse)
+
+        self.assertEqual(output.dtype, torch.float32)
+
+        # [4, 8, 64] + [4, 8, 1] -> [4, 8, 65] (Cat)
+        # DCP All2All -> [4, 8, 65]
+        # PCP AllGather (dim=0, size=2) -> [8, 8, 65]
+        expected_shape = (bs * pcp_size, num_heads, head_dim + 1)
+        self.assertEqual(output.shape, expected_shape)
+
+        mock_all2all.assert_called_once()
+        called_args = mock_all2all.call_args
+        self.assertEqual(called_args.kwargs["group"], mock_group)
+
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_decode_context_model_parallel_world_size")
+    def test_process_attn_out_lse_simple(self, mock_get_dcp_size, mock_get_pcp_group):
+        from vllm_ascend.attention.context_parallel.common_cp import _process_attn_out_lse
+
+        mock_get_pcp_group.return_value.world_size = 1
+        mock_get_dcp_size.return_value = 1
+
+        attn_output = torch.randn(2, 4, 16)
+        softmax_lse = torch.randn(2, 4, 1)
+
+        output = _process_attn_out_lse(attn_output, softmax_lse)
+
+        # concat: [2, 4, 16+1]
+        self.assertEqual(output.shape, (2, 4, 17))
+
+    @patch("torch_npu.npu_attention_update")
+    def test_npu_attn_out_lse_update(self, mock_npu_attention_update):
+        # Mock input data
+        attn_lse_mask = torch.randn(8, 128, 1)
+        attn_lse_nomask = torch.randn(8, 128, 1)
+        attn_out_mask = torch.randn(8, 128, 128)
+        attn_out_nomask = torch.randn(8, 128, 128)
+
+        mock_npu_attention_update.return_value = (torch.randn(8 * 128, 128), None)
+
+        # Call the method under test
+        output = _npu_attn_out_lse_update(attn_lse_mask, attn_lse_nomask, attn_out_mask, attn_out_nomask)
+
+        # Assertions
+        self.assertIsInstance(output, torch.Tensor)
+        self.assertEqual(output.shape, (8, 128, 128))
+        mock_npu_attention_update.assert_called_once()
+
+    def test_update_out_and_lse(self):
+        # Mock input data: [N, batch, heads, head_size/1]
+        out_list = torch.randn(3, 2, 4, 8)
+        lse_list = torch.randn(3, 2, 4, 1)
+
+        # Call the method under test
+        out_final, lse_final = _update_out_and_lse(out_list, lse_list)
+
+        # Assert shapes
+        self.assertEqual(out_final.shape, (2, 4, 8))
+        self.assertEqual(lse_final.shape, (2, 4, 1))
+        self.assertIsInstance(out_final, torch.Tensor)
+        self.assertIsInstance(lse_final, torch.Tensor)
+
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_decode_context_model_parallel_world_size")
+    @patch("torch_npu.npu_attention_update")
+    def test_npu_attention_update(self, mock_npu_update, mock_get_dcp, mock_get_pcp):
+        pcp_size = 2
+        dcp_size = 2
+        mock_get_pcp.return_value.world_size = pcp_size
+        mock_get_dcp.return_value = dcp_size
+
+        head_size = 64
+        S, H = 4, 8  # Sequence and Heads per segment
+
+        attn_out_lse = torch.randn(pcp_size * S, dcp_size * H, head_size + 1)
+
+        # mock NPU op return (N=PCP*DCP=4, S*H=32) -> [4*32, 64]
+        mock_npu_update.return_value = (torch.randn(pcp_size * dcp_size * S * H, head_size), None)
+
+        # test
+        result = _npu_attention_update(head_size, attn_out_lse)
+
+        self.assertEqual(result.shape, (pcp_size * dcp_size * S, H, head_size))
+        mock_npu_update.assert_called_once()
+
+    def test_out_lse_reshape(self):
+        # Mock input data
+        out_list = torch.randn(3, 2, 4, 8)  # [N, batch_size, num_heads, head_size]
+        lse_list = torch.randn(3, 2, 4, 1)  # [N, batch_size, num_heads, 1]
+
+        # Call the method under test
+        out_final, lse_final = _update_out_and_lse(out_list, lse_list)
+
+        # Assert the method call
+        self.assertEqual(out_final.shape, (2, 4, 8))  # [batch_size, num_heads, head_size]
+        self.assertEqual(lse_final.shape, (2, 4, 1))  # [batch_size, num_heads, 1]
+
+        self.assertIsInstance(out_final, torch.Tensor)
+        self.assertIsInstance(lse_final, torch.Tensor)

tests/ut/attention/a2/test_mla_cp.py

@@ -0,0 +1,899 @@
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.distributed.parallel_state import GroupCoordinator
+
+from tests.ut.attention.utils import patch_distributed_groups
+from tests.ut.base import TestBase
+from vllm_ascend.ascend_config import init_ascend_config
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+from vllm_ascend.attention.context_parallel.common_cp import (
+    CPChunkedContextMetadata,
+    _npu_attention_update,
+    _process_attn_out_lse,
+)
+from vllm_ascend.attention.context_parallel.mla_cp import AscendMlaCPImpl
+from vllm_ascend.attention.mla_v1 import ChunkedContextMetadata
+
+
+def get_pcp_split_info(pcp_rank, pcp_size, seq_lens):
+    q_head_idx, q_tail_idx = [], []
+    kv_with_q_head_nomask_idx, kv_with_q_head_mask_idx = [], []
+    kv_with_q_tail_nomask_idx, kv_with_q_tail_mask_idx = [], []
+    kv_tail_proj_idx: list[int] = []
+    kv_with_q_head_attn_idx_in_tail = []
+    kv_with_q_tail_attn_idx_in_tail = []
+    chunk_seqlens = []
+    kv_with_q_head_nomask_seqlens, kv_with_q_tail_nomask_seqlens = [], []
+    head_actual_seq_lengths_kv = []
+    tail_actual_seq_lengths_kv = []
+    q_req_offset = 0
+    kv_req_offset = 0
+    q_head_chunk_id = pcp_rank
+    q_tail_chunk_id = pcp_size * 2 - 1 - pcp_rank
+    for i, seq_len in enumerate(seq_lens):
+        chunk_len = seq_len // 2
+        chunk_seqlens.append(chunk_len)
+        q_head_idx.extend(list(range(q_req_offset, q_req_offset + chunk_len)))
+        kv_with_q_head_nomask_idx.extend(list(range(kv_req_offset, kv_req_offset + chunk_len * q_head_chunk_id)))
+        kv_with_q_head_mask_idx.extend(
+            list(range(kv_req_offset + chunk_len * q_head_chunk_id, kv_req_offset + chunk_len * (q_head_chunk_id + 1)))
+        )
+        kv_with_q_head_nomask_seqlens.append(chunk_len * q_head_chunk_id)
+
+        q_tail_idx.extend(list(range(q_req_offset + chunk_len, q_req_offset + chunk_len * 2)))
+        kv_with_q_tail_nomask_idx.extend(list(range(kv_req_offset, kv_req_offset + chunk_len * q_tail_chunk_id)))
+        kv_with_q_tail_mask_idx.extend(
+            list(range(kv_req_offset + chunk_len * q_tail_chunk_id, kv_req_offset + chunk_len * (q_tail_chunk_id + 1)))
+        )
+        kv_with_q_tail_nomask_seqlens.append(chunk_len * q_tail_chunk_id)
+
+        tail_proj_offset = len(kv_tail_proj_idx)
+        tail_proj_len = chunk_len * (q_tail_chunk_id + 1)
+        kv_tail_proj_idx.extend(list(range(kv_req_offset, kv_req_offset + tail_proj_len)))
+        kv_with_q_head_attn_idx_in_tail.extend(
+            list(range(tail_proj_offset, tail_proj_offset + chunk_len * (q_head_chunk_id + 1)))
+        )
+        kv_with_q_tail_attn_idx_in_tail.extend(list(range(tail_proj_offset, tail_proj_offset + tail_proj_len)))
+        head_actual_seq_lengths_kv.append(len(kv_with_q_head_attn_idx_in_tail))
+        tail_actual_seq_lengths_kv.append(len(kv_with_q_tail_attn_idx_in_tail))
+
+        q_req_offset += seq_len
+        kv_req_offset += seq_len * pcp_size
+    return (
+        torch.tensor(q_head_idx),
+        torch.tensor(q_tail_idx),
+        torch.tensor(kv_with_q_head_nomask_idx),
+        torch.tensor(kv_with_q_head_mask_idx),
+        torch.tensor(kv_with_q_tail_nomask_idx),
+        torch.tensor(kv_with_q_tail_mask_idx),
+        chunk_seqlens,
+        kv_with_q_head_nomask_seqlens,
+        kv_with_q_tail_nomask_seqlens,
+        torch.tensor(kv_tail_proj_idx),
+        torch.tensor(kv_with_q_head_attn_idx_in_tail),
+        torch.tensor(kv_with_q_tail_attn_idx_in_tail),
+        head_actual_seq_lengths_kv,
+        tail_actual_seq_lengths_kv,
+    )
+
+
+def get_chunk_metadata(
+    pcp_size,
+    dcp_size,
+    num_prefills,
+    num_decodes,
+    block_size,
+    num_computed_tokens_cpu,
+    num_reqs,
+    chunked_prefill_workspace_size,
+    num_computed_tokens_of_pcp_dcp,
+    cp_local_block_size,
+):
+    reqs_start = num_decodes
+    context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
+    max_context_len_cpu = context_lens_cpu.max().item()
+    num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
+    max_context_chunk = chunked_prefill_workspace_size // num_prefills_with_context_cpu
+    max_context_chunk = max_context_chunk // block_size * block_size
+
+    assert max_context_chunk > 0
+    num_chunks = (max_context_len_cpu + max_context_chunk - 1) // max_context_chunk
+    chunk_starts = torch.arange(num_chunks, dtype=torch.int32).unsqueeze(1).expand(-1, num_prefills) * max_context_chunk
+    chunk_ends = torch.min(context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk)
+    chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
+    cu_seq_lens_cpu = torch.zeros(num_chunks, num_prefills + 1, dtype=torch.int32)
+    torch.cumsum(chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32)
+
+    def cdiv(a, b):
+        return (a + b - 1) // b
+
+    if dcp_size * pcp_size > 1:
+        if num_computed_tokens_of_pcp_dcp is not None:
+            local_context_lens_allranks = torch.tensor(num_computed_tokens_of_pcp_dcp[reqs_start:num_reqs]).reshape(
+                -1, dcp_size * pcp_size
+            )
+        # Note(qcs): The max local context lengths
+        # padded to `cp_local_block_size`.
+        padded_local_context_lens_cpu = (
+            cdiv(
+                context_lens_cpu,
+                cp_local_block_size * pcp_size * dcp_size,
+            )
+            * cp_local_block_size
+        )
+        padded_local_max_context_chunk_across_ranks = (
+            cdiv(
+                max_context_chunk,
+                cp_local_block_size * pcp_size * dcp_size,
+            )
+            * cp_local_block_size
+        )
+        local_chunk_starts = (
+            torch.arange(num_chunks, dtype=torch.int32).unsqueeze(1).expand(-1, num_prefills)
+            * padded_local_max_context_chunk_across_ranks
+        )
+        local_chunk_ends = torch.min(
+            padded_local_context_lens_cpu.unsqueeze(0),
+            local_chunk_starts + padded_local_max_context_chunk_across_ranks,
+        )
+        padded_local_chunk_seq_lens = (local_chunk_ends - local_chunk_starts).clamp(min=0)
+        padded_local_cu_chunk_seq_lens_cpu = torch.zeros(num_chunks, num_prefills + 1, dtype=torch.int32)
+        torch.cumsum(
+            padded_local_chunk_seq_lens,
+            dim=1,
+            out=padded_local_cu_chunk_seq_lens_cpu[:, 1:],
+            dtype=torch.int32,
+        )
+        chunk_actual_seq_lengths_kv_list = [torch.cumsum(chunk_seq_lens[i], dim=0).tolist() for i in range(num_chunks)]
+        chunked_context_metadata = CPChunkedContextMetadata(
+            cu_seq_lens=cu_seq_lens_cpu.to(non_blocking=True),
+            starts=local_chunk_starts.to(non_blocking=True),
+            seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),
+            max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
+            chunk_seq_lens=chunk_seq_lens,
+            chunk_seq_lens_npu=chunk_seq_lens,
+            chunk_actual_seq_lengths_kv_list=chunk_actual_seq_lengths_kv_list,
+            workspace=None,
+            padded_chunk_seq_lens_npu=padded_local_chunk_seq_lens,
+            padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
+            local_context_lens_allranks=local_context_lens_allranks.tolist(),
+            padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to(non_blocking=True),
+            cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
+            chunk_size=padded_local_max_context_chunk_across_ranks,
+        )
+    else:
+        chunked_context_metadata = ChunkedContextMetadata(
+            cu_seq_lens=cu_seq_lens_cpu.to(non_blocking=True),
+            starts=chunk_starts.to(non_blocking=True),
+            seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
+            max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
+            chunk_seq_lens=chunk_seq_lens,
+            chunk_seq_lens_npu=chunk_seq_lens,
+            workspace=None,
+        )
+    return chunked_context_metadata
+
+
+class TestAscendMLAImpl(TestBase):
+    @patch("vllm.distributed.parallel_state._TP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    @patch("vllm.distributed.get_tensor_model_parallel_world_size", return_value=2)
+    @patch("vllm_ascend.attention.mla_v1.get_current_vllm_config")
+    @patch("vllm_ascend.attention.mla_v1.get_ascend_config")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def setUp(self, ascend_config, get_current_vllm_config, mock_get_tp_size, mock_tp):
+        mock_tp.world_size = 2
+        mock_tp.rank_in_group = MagicMock()
+        mock_tp.device_group = MagicMock()
+        vllm_config = MagicMock()
+        speculative_config = MagicMock()
+        model_config = MagicMock()
+        parallel_config = MagicMock()
+        parallel_config.prefill_context_parallel_size = 1
+        parallel_config.tensor_parallel_size = 2
+        speculative_config.num_speculative_tokens = 4
+        vllm_config.speculative_config = speculative_config
+        model_config.dtype = torch.float16
+        vllm_config.model_config = model_config
+        get_current_vllm_config.return_value = vllm_config
+        vllm_config.additional_config = {"refresh": True}
+        vllm_config.parallel_config = parallel_config
+        init_ascend_config(vllm_config)
+
+        num_heads = 256
+        head_size = 1024
+        scale = 0.1
+        num_kv_heads = 8
+        kv_cache_dtype = "auto"
+
+        kv_a_layernorm = MagicMock()
+        kv_a_layernorm.weight = torch.randn(96)
+        kv_a_layernorm.variance_epsilon = 1e-6
+        kwargs = {
+            "kv_lora_rank": 32,
+            "qk_nope_head_dim": 64,
+            "qk_rope_head_dim": 32,
+            "qk_head_dim": 96,
+            "v_head_dim": 128,
+            "q_lora_rank": 64,
+            "q_proj": MagicMock(),
+            "q_b_proj": MagicMock(),
+            "kv_b_proj": MagicMock(),
+            "o_proj": MagicMock(),
+            "kv_a_proj_with_mqa": MagicMock(),
+            "fused_qkv_a_proj": MagicMock(),
+            "kv_a_layernorm": kv_a_layernorm,
+            "rotary_emb": MagicMock(),
+        }
+
+        self.impl = AscendMlaCPImpl(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+            num_kv_heads=num_kv_heads,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype=kv_cache_dtype,
+            blocksparse_params=None,
+            logits_soft_cap=None,
+            attn_type=None,
+            kv_sharing_target_layer_name=None,
+            **kwargs,
+        )
+
+    def test_init(self):
+        self.assertEqual(self.impl.num_heads, 256)
+        self.assertEqual(self.impl.head_size, 1024)
+        self.assertEqual(self.impl.scale, 0.1)
+        self.assertEqual(self.impl.num_kv_heads, 8)
+        self.assertEqual(self.impl.kv_cache_dtype, "auto")
+        self.assertEqual(self.impl.kv_lora_rank, 32)
+        self.assertEqual(self.impl.qk_nope_head_dim, 64)
+        self.assertEqual(self.impl.qk_rope_head_dim, 32)
+        self.assertEqual(self.impl.qk_head_dim, 96)
+        self.assertEqual(self.impl.v_head_dim, 128)
+        self.assertIsNotNone(self.impl.q_proj)
+        self.assertIsNotNone(self.impl.kv_b_proj)
+        self.assertIsNotNone(self.impl.o_proj)
+        self.assertIsNotNone(self.impl.kv_a_proj_with_mqa)
+        self.assertIsNotNone(self.impl.kv_a_layernorm)
+        self.assertEqual(self.impl.num_queries_per_kv, 32)
+        self.assertEqual(self.impl.pcp_size, 2)
+        self.assertEqual(self.impl.dcp_size, 2)
+
+    @patch("torch.ops.vllm.maybe_all_gather_and_maybe_unpad")
+    @patch("vllm_ascend.attention.mla_v1.get_weight_prefetch_method", return_value=MagicMock())
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_mla_preprocess_dcp(self, mock_get_weight_prefetch_method, mock_maybe_all_gather_and_maybe_unpad):
+        self.impl.num_kv_heads = 1
+        self.impl.num_heads = 16
+        self.impl.qk_rope_head_dim = 64
+        self.impl.kv_lora_rank = 512
+        self.impl.q_lora_rank = 1536
+        self.impl.dcp_size = 2
+        self.impl.pcp_size = 2
+        block_num = 10
+        block_size = 128
+        batch_size = 2
+        hidden_size = 1024
+        hidden_states = torch.randn(batch_size, hidden_size)
+
+        kv_cache0 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.kv_lora_rank)
+        kv_cache1 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.qk_rope_head_dim)
+        kv_cache = (kv_cache0, kv_cache1)
+
+        attn_metadata = MagicMock()
+        attn_metadata.num_decodes = 2
+        attn_metadata.num_prefills = 0
+        attn_metadata.num_prefill_tokens = 0
+        attn_metadata.num_decode_tokens = 2
+        attn_metadata.num_actual_tokens = 2
+        attn_metadata.slot_mapping = torch.arange(4)
+        attn_metadata.decode.cos = torch.randn(2, 64)
+        attn_metadata.decode.sin = torch.randn(2, 64)
+
+        self.impl.q_a_layernorm = MagicMock()
+        self.impl.q_a_layernorm.return_value = torch.randn(attn_metadata.num_actual_tokens, self.impl.q_lora_rank)
+        self.impl.kv_a_proj_with_mqa = MagicMock()
+        self.impl.kv_a_proj_with_mqa.return_value = [
+            torch.randn(batch_size, self.impl.num_heads, self.impl.qk_rope_head_dim + self.impl.kv_lora_rank)
+        ]
+        self.impl.fused_qkv_a_proj = MagicMock()
+        self.impl.fused_qkv_a_proj.return_value = [
+            torch.randn(
+                attn_metadata.num_actual_tokens,
+                self.impl.qk_rope_head_dim + self.impl.kv_lora_rank + self.impl.q_lora_rank,
+            )
+        ]
+
+        self.impl.rope_single = MagicMock(side_effect=lambda x, cos, sin: x)
+        self.impl.exec_kv_decode = MagicMock()
+        self.impl.exec_kv_decode.return_value = [MagicMock(), MagicMock()]
+
+        self.impl._q_proj_and_k_up_proj = MagicMock()
+        self.impl._q_proj_and_k_up_proj.return_value = [
+            torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.kv_lora_rank),
+            torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.qk_rope_head_dim),
+        ]
+
+        mock_maybe_all_gather_and_maybe_unpad.side_effect = lambda x, label: x
+
+        decode_res, prefill_res = self.impl._mla_preprocess(
+            "mock_layer", hidden_states, kv_cache, attn_metadata, need_gather_q_kv=False
+        )
+
+        self.assertIsNotNone(decode_res)
+        self.assertIsNone(prefill_res)
+
+    @patch("torch_npu._npu_reshape_and_cache")
+    @patch("torch.ops.vllm.maybe_all_gather_and_maybe_unpad")
+    @patch("vllm_ascend.attention.mla_v1.get_weight_prefetch_method", return_value=MagicMock())
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_mla_preprocess_pcp(
+        self, mock_get_weight_prefetch_method, mock_maybe_all_gather_and_maybe_unpad, mock_npu_reshape_and_cache
+    ):
+        self.impl.num_kv_heads = 1
+        self.impl.num_heads = 16
+        self.impl.qk_rope_head_dim = 64
+        self.impl.kv_lora_rank = 512
+        self.impl.q_lora_rank = 1536
+        self.impl.dcp_size = 2
+        self.impl.pcp_size = 2
+        block_num = 10
+        block_size = 128
+        batch_size = 2
+        hidden_size = 1024
+        hidden_states = torch.randn(batch_size, hidden_size)
+
+        kv_cache0 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.kv_lora_rank)
+        kv_cache1 = torch.randn(block_num, block_size, self.impl.num_kv_heads, self.impl.qk_rope_head_dim)
+        kv_cache = (kv_cache0, kv_cache1)
+
+        attn_metadata = MagicMock()
+        attn_metadata.num_decodes = 0
+        attn_metadata.num_prefills = 2
+        attn_metadata.num_prefill_tokens = 2
+        attn_metadata.num_decode_tokens = 0
+        attn_metadata.num_actual_tokens = 2
+        attn_metadata.num_actual_tokens_pcp_padded = 4
+        attn_metadata.prefill.pcp_metadata = MagicMock()
+        attn_metadata.prefill.pcp_metadata.pcp_allgather_restore_idx = torch.arange(4)
+        tail_projection_len = 3
+        attn_metadata.prefill.pcp_metadata.kv_tail_proj_idx = torch.arange(tail_projection_len, dtype=torch.int64)
+        attn_metadata.slot_mapping = torch.arange(4)
+        attn_metadata.prefill.cos = torch.randn(2, 64)
+        attn_metadata.prefill.sin = torch.randn(2, 64)
+
+        self.impl.q_a_layernorm = MagicMock()
+        self.impl.q_a_layernorm.return_value = torch.randn(attn_metadata.num_actual_tokens, self.impl.q_lora_rank)
+        self.impl.kv_a_proj_with_mqa = MagicMock()
+        self.impl.kv_a_proj_with_mqa.return_value = [
+            torch.randn(batch_size, self.impl.num_heads, self.impl.qk_rope_head_dim + self.impl.kv_lora_rank)
+        ]
+        self.impl.fused_qkv_a_proj = MagicMock()
+        self.impl.fused_qkv_a_proj.return_value = [
+            torch.randn(
+                attn_metadata.num_actual_tokens,
+                self.impl.qk_rope_head_dim + self.impl.kv_lora_rank + self.impl.q_lora_rank,
+            )
+        ]
+
+        self.impl.rope_single = MagicMock(side_effect=lambda x, cos, sin: x)
+        self.impl.exec_kv_decode = MagicMock()
+        self.impl.exec_kv_decode.return_value = [MagicMock(), MagicMock()]
+
+        self.impl._q_proj_and_k_up_proj = MagicMock()
+        self.impl._q_proj_and_k_up_proj.return_value = [
+            torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.kv_lora_rank),
+            torch.randn(attn_metadata.num_decodes, self.impl.num_heads, self.impl.qk_rope_head_dim),
+        ]
+
+        mock_maybe_all_gather_and_maybe_unpad.side_effect = lambda x, label: x
+
+        self.impl.kv_a_layernorm = MagicMock()
+        self.impl.kv_a_layernorm.return_value = torch.randn(
+            attn_metadata.num_prefill_tokens, self.impl.num_kv_heads, self.impl.kv_lora_rank
+        )
+
+        self.impl.q_proj = MagicMock()
+        self.impl.q_proj.return_value = [
+            torch.randn(attn_metadata.num_prefill_tokens, self.impl.num_heads, self.impl.qk_head_dim)
+        ]
+        self.impl.kv_b_proj = MagicMock()
+        self.impl.kv_b_proj.return_value = [
+            torch.randn(
+                tail_projection_len,
+                self.impl.num_heads,
+                self.impl.v_head_dim + self.impl.qk_nope_head_dim,
+            )
+        ]
+        self.impl.rope_single = MagicMock(side_effect=lambda x, cos, sin: x)
+        self.impl.exec_kv_decode = MagicMock()
+        self.impl.exec_kv_decode.return_value = [MagicMock(), MagicMock()]
+        self.impl.exec_kv_prefill = MagicMock()
+        self.impl.exec_kv_prefill.return_value = [
+            torch.randn(attn_metadata.num_prefill_tokens, self.impl.num_heads, self.impl.qk_rope_head_dim),
+            torch.randn(attn_metadata.num_prefill_tokens, self.impl.num_heads, self.impl.kv_lora_rank),
+        ]
+
+        decode_res, prefill_res = self.impl._mla_preprocess(
+            "mock_layer", hidden_states, kv_cache, attn_metadata, need_gather_q_kv=False
+        )
+        self.assertIsNone(decode_res)
+        self.assertIsNotNone(prefill_res)
+        self.impl.kv_b_proj.assert_called_once()
+        self.assertEqual(self.impl.kv_b_proj.call_args.args[0].shape[0], tail_projection_len)
+
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_process_attn_out_lse(self):
+        self.impl.dcp_size = 2
+        self.impl.pcp_size = 2
+
+        B = 2
+        N = self.impl.num_heads
+        self.impl.kv_lora_rank = 512
+
+        attn_output = torch.randn(B, N, self.impl.kv_lora_rank)
+        softmax_lse = torch.randn(B, N, 1)
+
+        decode_metadata = MagicMock()
+        decode_metadata.actual_seq_lengths_q = MagicMock()
+        decode_metadata.seq_lens_list = MagicMock()
+        result = _process_attn_out_lse(attn_output, softmax_lse)
+
+        self.assertEqual(result.shape[0], B * self.impl.pcp_size)
+        self.assertEqual(result.shape[1], N)
+        self.assertEqual(result.shape[2], self.impl.kv_lora_rank + 1)
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch("torch_npu.npu_attention_update")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_forward_decode_pcp_dcp(
+        self, mock_npu_attention_update, mock_npu_fused_infer_attention_score, mock_get_forward_context
+    ):
+        self.impl.dcp_size = 2
+        self.impl.pcp_size = 2
+        self.impl.num_kv_heads = 1
+        self.impl.num_heads = 16
+        self.impl.kv_lora_rank = 64
+        self.impl.qk_nope_head_dim = 64
+        self.impl.spec_token_num = 1
+        B = 2
+        N = self.impl.num_heads * self.impl.dcp_size
+        BS = 128
+        NB = 100
+
+        q_nope = torch.randn(B, N, self.impl.qk_nope_head_dim)
+        q_pe = torch.randn(B, N, self.impl.qk_rope_head_dim)
+        k_nope = torch.randn(NB, 1, BS, self.impl.kv_lora_rank)
+        k_pe = torch.randn(NB, 1, BS, self.impl.qk_rope_head_dim)
+
+        attn_metadata = MagicMock()
+        attn_metadata.attn_state = AscendAttentionState.SpecDecoding
+        attn_metadata.decode = MagicMock()
+        attn_metadata.decode.actual_seq_lengths_q = MagicMock()
+        attn_metadata.decode.seq_lens_list = MagicMock()
+
+        self.impl.enable_kv_nz = True
+        self.impl.speculative_config = None
+        mock_npu_attention_update.return_value = (torch.randn(B, self.impl.num_heads, self.impl.kv_lora_rank), None)
+        mock_npu_fused_infer_attention_score.return_value = [
+            torch.randn(B, N, 1, self.impl.kv_lora_rank),
+            torch.randn(B, N, 1, 1),
+        ]
+        mock_get_forward_context.return_value = MagicMock(capturing=False)
+
+        self.impl._v_up_proj = MagicMock()
+        self.impl._v_up_proj.return_value = torch.randn(B, self.impl.v_head_dim)
+
+        result = self.impl._forward_decode(q_nope, q_pe, k_nope, k_pe, BS, attn_metadata)
+
+        self.assertEqual(result.shape[0], B)
+        self.assertEqual(result.shape[1], self.impl.v_head_dim)
+
+    @patch("torch_npu.atb.npu_paged_cache_load")
+    @patch("torch_npu.npu_attention_update")
+    @patch("torch_npu.npu_fused_infer_attention_score")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2)
+    def test_compute_prefill_context_with_dcp_pcp(
+        self, mock_all2all, mock_dcp, mock_pcp, mock_fia, mock_update, mock_load
+    ):
+        def mock_fia_attn(*args, **kwargs):
+            q = args[0]
+            v = args[2]
+            return (
+                torch.randn(q.shape[0], v.shape[1], v.shape[2], dtype=torch.float16),
+                torch.randn(v.shape[1], q.shape[0], dtype=torch.float16),
+            )
+
+        mock_fia.side_effect = mock_fia_attn
+
+        def mock_kv_b_proj(kv_c_normed):
+            return (
+                torch.randn(
+                    kv_c_normed.shape[0],
+                    self.impl.num_heads,
+                    self.impl.v_head_dim + self.impl.qk_nope_head_dim,
+                    dtype=torch.float16,
+                ),
+            )
+
+        def mock_reorg_kvcache(
+            allgatered_kv_c_normed: torch.Tensor,
+            allgatered_k_pe: torch.Tensor,
+            chunked_context: CPChunkedContextMetadata,
+            chunk_idx: int,
+            toks: int,
+        ):
+            return torch.randn(
+                chunked_context.cu_seq_lens_lst[chunk_idx][-1],
+                allgatered_kv_c_normed.shape[1],
+                allgatered_kv_c_normed.shape[2],
+            ), torch.randn(
+                chunked_context.cu_seq_lens_lst[chunk_idx][-1], allgatered_k_pe.shape[1], allgatered_k_pe.shape[2]
+            )
+
+        # mock proj
+        self.impl.kv_b_proj.side_effect = mock_kv_b_proj
+
+        def mock_update_fn(lse_list, out_list, mode):
+            total_len = out_list[0].shape[0]
+            D = out_list[0].shape[1]
+            return (torch.randn(total_len, D, dtype=torch.float32), None)
+
+        mock_update.side_effect = mock_update_fn
+        NUM_BLOCKS, BLOCK_SIZE = 10, 32  # fixed
+        USED_BLOCKS = 3
+        # pcp_size, dcp_size, nums_tokens_per_rank, nums_all_rank_context, num_prefills, num_decodes,
+        # num_seqs, cp_local_block_size, num_computed_tokens, num_computed_tokens_of_pcp_dcp
+        test_cases = [
+            (2, 2, [4], [128], 1, 0, 1, 1, [[[32, 32], [32, 32]]]),
+            (1, 2, [4], [128], 1, 0, 1, 1, [[[64, 64]]]),
+            (2, 1, [4], [128], 1, 0, 1, 1, [[[64], [64]]]),
+            (2, 2, [4, 7], [128, 128], 2, 0, 2, 1, [[[32, 32], [32, 32]], [[32, 32], [32, 32]]]),
+        ]
+        # kv cache tensor
+        kv_cache_0 = torch.randn(
+            NUM_BLOCKS, BLOCK_SIZE, self.impl.num_heads, self.impl.kv_lora_rank, dtype=torch.float16
+        )
+        kv_cache_1 = torch.randn(NUM_BLOCKS, BLOCK_SIZE, self.impl.num_heads, self.impl.v_head_dim, dtype=torch.float16)
+        kv_cache = [kv_cache_0, kv_cache_1]
+        max_model_len = 4096
+        max_num_seqs = 25
+        # create chunk context
+        chunked_prefill_workspace_size = min(max(8 * max_model_len, 4 * max_num_seqs * BLOCK_SIZE), 128 * 1024)
+        self.impl.prefill_mask = torch.triu(torch.ones(10, 10, dtype=torch.float16), 1)
+        for test_case in test_cases:
+            (
+                pcp_size,
+                dcp_size,
+                nums_tokens_per_rank,
+                nums_all_rank_context,
+                num_prefills,
+                num_decodes,
+                num_seqs,
+                cp_local_block_size,
+                num_computed_tokens_of_pcp_dcp,
+            ) = test_case
+            mock_dcp.world_size = dcp_size
+            mock_pcp.world_size = pcp_size
+            assert len(nums_tokens_per_rank) == len(nums_all_rank_context)
+            nums_context_per_rank = []
+            for num_all_rank_context in nums_all_rank_context:
+                assert num_all_rank_context % (pcp_size * dcp_size) == 0
+                nums_context_per_rank.append(num_all_rank_context // (pcp_size * dcp_size))
+            self.impl.dcp_size = dcp_size
+            self.impl.pcp_size = pcp_size
+            # create input
+            query = torch.randn(
+                sum(nums_tokens_per_rank), self.impl.num_heads, self.impl.qk_head_dim, dtype=torch.float16
+            )
+            q_nope = query[..., : self.impl.qk_nope_head_dim]
+            q_pe = query[..., self.impl.qk_nope_head_dim :]
+            prefix_out = torch.randn(
+                sum(nums_tokens_per_rank), self.impl.num_heads, self.impl.v_head_dim, dtype=torch.float16
+            )
+            prefix_lse = torch.randn(self.impl.num_heads, sum(nums_tokens_per_rank), dtype=torch.float16)
+            chunk_ctx = get_chunk_metadata(
+                pcp_size,
+                dcp_size,
+                num_prefills=num_prefills,
+                num_decodes=num_decodes,
+                block_size=BLOCK_SIZE,
+                num_computed_tokens_cpu=torch.tensor(nums_all_rank_context),
+                num_reqs=num_seqs,
+                chunked_prefill_workspace_size=chunked_prefill_workspace_size,
+                num_computed_tokens_of_pcp_dcp=num_computed_tokens_of_pcp_dcp,
+                cp_local_block_size=cp_local_block_size,
+            )
+            meta = MagicMock()
+            prefill_meta = MagicMock()
+            prefill_meta.query_lens = torch.tensor(nums_tokens_per_rank)
+            prefill_meta.block_table = torch.randint(0, USED_BLOCKS, (1, 64))  # (batch, max_blocks)
+            prefill_meta.chunked_context = chunk_ctx
+            meta.prefill = prefill_meta
+
+            with patch.object(self.impl, "_reorg_kvcache") as mock_reorg:
+                mock_reorg.side_effect = mock_reorg_kvcache
+
+                out, lse = self.impl._compute_prefill_context(
+                    q_nope, q_pe, kv_cache, self.impl.qk_rope_head_dim, meta, prefix_out, prefix_lse
+                )
+
+            iters = len(chunk_ctx.seq_tot)
+            self.impl.dcp_size = 1
+            self.impl.pcp_size = 1
+            self.assertEqual(mock_reorg.call_count, iters * (1 if dcp_size * pcp_size > 1 else 0))
+            self.assertEqual(mock_load.call_count, iters)
+            self.assertEqual(mock_fia.call_count, iters)
+            mock_reorg.reset_mock()
+            mock_load.reset_mock()
+            mock_fia.reset_mock()
+            mock_update.reset_mock()
+            mock_dcp.reset_mock()
+            mock_pcp.reset_mock()
+            self.assertEqual(out.shape, prefix_out.shape)
+
+    @patch_distributed_groups(dcp_size=2, pcp_size=2)
+    def test_reorg_kvcache_with_dcp_pcp(self, mock_all2all, mock_dcp, mock_pcp):
+        BLOCK_SIZE = 128  # fixed
+        max_model_len = 4096
+        max_num_seqs = 25
+        test_cases = [
+            (2, 2, [4], [128], 1, 0, 1, 1, [[[32, 32], [32, 32]]]),
+            (1, 2, [4], [128], 1, 0, 1, 1, [[[64, 64]]]),
+            (2, 1, [4], [128], 1, 0, 1, 1, [[[64], [64]]]),
+            (2, 2, [4, 7], [128, 128], 2, 0, 2, 1, [[[32, 32], [32, 32]], [[32, 32], [32, 32]]]),
+        ]
+        for test_case in test_cases:
+            (
+                pcp_size,
+                dcp_size,
+                nums_tokens_per_rank,
+                nums_all_rank_context,
+                num_prefills,
+                num_decodes,
+                num_seqs,
+                cp_local_block_size,
+                num_computed_tokens_of_pcp_dcp,
+            ) = test_case
+            if pcp_size * dcp_size == 1:
+                continue
+            self.impl.dcp_size = dcp_size
+            mock_dcp.world_size = dcp_size
+            mock_dcp.all_gather.reset_mock()
+            self.impl.pcp_size = pcp_size
+            mock_pcp.world_size = pcp_size
+            mock_pcp.all_gather.reset_mock()
+
+            chunked_prefill_workspace_size = min(max(8 * max_model_len, 4 * max_num_seqs * BLOCK_SIZE), 128 * 1024)
+            chunked_context = get_chunk_metadata(
+                pcp_size,
+                dcp_size,
+                num_prefills=num_prefills,
+                num_decodes=num_decodes,
+                block_size=BLOCK_SIZE,
+                num_computed_tokens_cpu=torch.tensor(nums_all_rank_context),
+                num_reqs=num_seqs,
+                chunked_prefill_workspace_size=chunked_prefill_workspace_size,
+                num_computed_tokens_of_pcp_dcp=num_computed_tokens_of_pcp_dcp,
+                cp_local_block_size=cp_local_block_size,
+            )
+
+            for i in range(len(chunked_context.seq_tot)):
+                allgatered_kv_c_normed = torch.randn(
+                    chunked_context.seq_tot[i], self.impl.num_heads, self.impl.kv_lora_rank
+                )
+                allgatered_k_pe = torch.randn(
+                    chunked_context.seq_tot[i], self.impl.num_heads, self.impl.qk_rope_head_dim
+                )
+                result_kv, result_k_pe = self.impl._reorg_kvcache(
+                    allgatered_kv_c_normed,
+                    allgatered_k_pe,
+                    chunked_context,
+                    chunk_idx=i,
+                    toks=chunked_context.seq_tot[i],
+                )
+                self.assertEqual(
+                    result_kv.shape,
+                    (chunked_context.cu_seq_lens_lst[i][-1], self.impl.num_heads, self.impl.kv_lora_rank),
+                )
+                self.assertEqual(
+                    result_k_pe.shape,
+                    (chunked_context.cu_seq_lens_lst[i][-1], self.impl.num_heads, self.impl.qk_rope_head_dim),
+                )
+
+                self.assertEqual(result_kv.shape[0], chunked_context.cu_seq_lens_lst[i][-1])
+                self.assertEqual(result_k_pe.shape[0], chunked_context.cu_seq_lens_lst[i][-1])
+
+                self.assertEqual(mock_dcp.all_gather.call_count, (1 if dcp_size > 1 else 0))
+                self.assertEqual(mock_pcp.all_gather.call_count, (1 if pcp_size > 1 else 0))
+
+    def test_out_lse_reshape(self):
+        test_cases = [10, 1, 128, 512]
+        for test_case in test_cases:
+            num_tokens = test_case
+            num_heads, head_dim = self.impl.num_heads, self.impl.v_head_dim
+            attn_out = torch.randn(num_tokens, num_heads, head_dim)
+            attn_lse = torch.randn(num_tokens, num_heads, 1)
+
+            out, lse = self.impl._out_lse_reshape(attn_out, attn_lse)
+
+            assert out.shape == (num_tokens * num_heads, head_dim)
+            assert out.is_contiguous()
+
+            assert lse.shape == (num_tokens * num_heads,)
+            assert lse.is_contiguous()
+
+            expected_out = attn_out.contiguous().view(-1, head_dim)
+            expected_lse = attn_lse.contiguous().view(-1)
+
+            assert torch.allclose(out, expected_out)
+            assert torch.allclose(lse, expected_lse)
+
+    @patch("torch_npu.npu_attention_update")
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_pcp_group")
+    @patch("vllm.distributed.parallel_state._PCP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    @patch("vllm_ascend.attention.context_parallel.common_cp.get_dcp_group")
+    @patch("vllm.distributed.parallel_state._DCP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    def test_npu_attention_update_with_dcp_pcp(
+        self, mock_dcp, mock_get_dcp_group, mock_pcp, mock_get_pcp_group, mock_npu_attention_update
+    ):
+        NUM_TOKENS = 10  # fixed
+        test_cases = [(1, 1), (1, 2), (2, 1), (2, 2), (2, 3)]
+        for test_case in test_cases:
+            self.impl.dcp_size = test_case[0]
+            self.impl.pcp_size = test_case[1]
+            num_heads, head_dim = self.impl.num_heads, self.impl.kv_lora_rank + 1
+
+            def mock_out_lse_reshape(attn_out, attn_lse):
+                attn_out = attn_out.contiguous().view(attn_out.shape[0] * attn_out.shape[1], attn_out.shape[2])
+                attn_lse = attn_lse.contiguous().view(attn_lse.shape[0] * attn_lse.shape[1] * attn_lse.shape[2])
+                return attn_out, attn_lse
+
+            self.impl._out_lse_reshape = MagicMock()
+            self.impl._out_lse_reshape.side_effect = mock_out_lse_reshape
+
+            def mock_npu_attention_update_effect(attn_lse_split_cp, attn_out_split_cp, update_type):
+                return torch.randn_like(attn_out_split_cp[0]), torch.randn_like(attn_lse_split_cp[0])
+
+            mock_npu_attention_update.side_effect = mock_npu_attention_update_effect
+
+            mock_pcp_group = MagicMock()
+            mock_pcp_group.world_size = self.impl.pcp_size
+            mock_get_pcp_group.return_value = mock_pcp_group
+
+            mock_dcp.world_size = self.impl.dcp_size
+            mock_dcp_group = MagicMock()
+            mock_get_dcp_group.return_value = mock_dcp_group
+            attn_out_lse = torch.randn(self.impl.pcp_size * NUM_TOKENS, self.impl.dcp_size * num_heads, head_dim)
+            out = _npu_attention_update(self.impl.kv_lora_rank, attn_out_lse)
+            self.impl.dcp_size = 1
+            self.impl.pcp_size = 1
+            assert out.shape == (NUM_TOKENS, num_heads, self.impl.kv_lora_rank)
+
+    @patch("torch.ops.npu.npu_fused_infer_attention_score")
+    def test_attention_with_optional_kv_select_with_dcp_pcp(self, mock_npu_fia):
+        num_heads = self.impl.num_heads
+        v_head_dim = self.impl.v_head_dim
+        qk_nope_head_dim = self.impl.qk_nope_head_dim
+        qk_rope_head_dim = self.impl.qk_rope_head_dim
+
+        def mock_npu_fia_effect(*args, **kwargs):
+            q_nope = args[0]
+            q_tokens = q_nope.shape[0]
+            out = torch.randn(q_tokens, num_heads, v_head_dim, dtype=torch.float16)
+            lse = torch.randn(q_tokens, num_heads, 1, dtype=torch.float32)
+            return out, lse
+
+        mock_npu_fia.side_effect = mock_npu_fia_effect
+
+        test_cases = [([8], 2, 2), ([8, 12], 2, 2)]
+        for test_case in test_cases:
+            scheduled_tokens, pcp_size, dcp_size = test_case
+            nums_tokens_per_rank = [num // pcp_size for num in scheduled_tokens]
+            seq_len_q, seq_len_k = sum(nums_tokens_per_rank), sum(scheduled_tokens)
+
+            q_nope = torch.randn(seq_len_q, num_heads, qk_nope_head_dim, dtype=torch.float16)
+            q_pe = torch.randn(seq_len_q, num_heads, qk_rope_head_dim, dtype=torch.float16)
+            k_nope = torch.randn(seq_len_k, num_heads, qk_nope_head_dim, dtype=torch.float16)
+            k_pe = torch.randn(seq_len_k, num_heads, qk_rope_head_dim, dtype=torch.float16)
+            value = torch.randn(seq_len_k, num_heads, v_head_dim, dtype=torch.float16)
+            mask = torch.triu(torch.ones(10, 10, dtype=torch.float16), 1)
+
+            attn_metadata = MagicMock()
+            attn_metadata.prefill = MagicMock()
+            attn_metadata.prefill.chunked_context = None
+
+            for rank in range(pcp_size):
+                info = get_pcp_split_info(rank, pcp_size, nums_tokens_per_rank)
+                q_head_idx = info[0]
+                chunk_seqlens = info[6]
+                kv_tail_proj_idx = info[9]
+                kv_with_q_head_attn_idx_in_tail = info[10]
+                head_actual_seq_lengths_kv = info[12]
+                attn_mask_seqlens = torch.cumsum(torch.tensor(chunk_seqlens, dtype=torch.int32), dim=0).tolist()
+
+                output_head, lse_head = self.impl._attention_with_optional_kv_select(
+                    q_nope=torch.index_select(q_nope, 0, q_head_idx),
+                    q_pe=torch.index_select(q_pe, 0, q_head_idx),
+                    k_nope=torch.index_select(k_nope, 0, kv_tail_proj_idx),
+                    k_pe=torch.index_select(k_pe, 0, kv_tail_proj_idx),
+                    value=torch.index_select(value, 0, kv_tail_proj_idx),
+                    kv_attn_idx=kv_with_q_head_attn_idx_in_tail,
+                    attn_mask_seqlens=attn_mask_seqlens,
+                    actual_seq_lengths_kv=head_actual_seq_lengths_kv,
+                    mask=mask,
+                    attn_metadata=attn_metadata,
+                )
+
+                self.assertEqual(output_head.shape, (q_head_idx.shape[0], num_heads, v_head_dim))
+                if lse_head is not None:
+                    self.assertEqual(lse_head.shape, (q_head_idx.shape[0], num_heads, 1))
+                else:
+                    self.assertIsNone(lse_head)
+                self.assertEqual(mock_npu_fia.call_count, 1)
+                self.assertEqual(mock_npu_fia.call_args.kwargs["sparse_mode"], 3)
+
+                mock_npu_fia.reset_mock()
+
+    @patch("torch.ops.npu.npu_fused_infer_attention_score")
+    def test_attention_with_optional_kv_select_trigger_chunked(self, mock_npu_fia):
+        num_heads = self.impl.num_heads
+        v_head_dim = self.impl.v_head_dim
+        qk_nope_head_dim = self.impl.qk_nope_head_dim
+        qk_rope_head_dim = self.impl.qk_rope_head_dim
+
+        def mock_npu_fia_effect(*args, **kwargs):
+            q_tokens = args[0].shape[0]
+            out = torch.randn(q_tokens, num_heads, v_head_dim, dtype=torch.float16)
+            lse = torch.randn(q_tokens, num_heads, 1, dtype=torch.float32)
+            return out, lse
+
+        mock_npu_fia.side_effect = mock_npu_fia_effect
+
+        test_cases = [([8], 2, 2)]
+        for test_case in test_cases:
+            scheduled_tokens, pcp_size, dcp_size = test_case
+            nums_tokens_per_rank = [num // pcp_size for num in scheduled_tokens]
+
+            q_nope = torch.randn(sum(nums_tokens_per_rank), num_heads, qk_nope_head_dim, dtype=torch.float16)
+            q_pe = torch.randn(sum(nums_tokens_per_rank), num_heads, qk_rope_head_dim, dtype=torch.float16)
+            k_nope = torch.randn(sum(scheduled_tokens), num_heads, qk_nope_head_dim, dtype=torch.float16)
+            k_pe = torch.randn(sum(scheduled_tokens), num_heads, qk_rope_head_dim, dtype=torch.float16)
+            value = torch.randn(sum(scheduled_tokens), num_heads, v_head_dim, dtype=torch.float16)
+            mask = torch.ones(10, 10, dtype=torch.float16)
+
+            attn_metadata = MagicMock()
+            attn_metadata.prefill = MagicMock()
+            attn_metadata.prefill.chunked_context = "TRIGGER_CHUNKED"
+
+            for rank in range(pcp_size):
+                info = get_pcp_split_info(rank, pcp_size, nums_tokens_per_rank)
+                q_head_idx = info[0]
+                chunk_seqlens = info[6]
+                kv_tail_proj_idx = info[9]
+                kv_with_q_head_attn_idx_in_tail = info[10]
+                head_actual_seq_lengths_kv = info[12]
+                attn_mask_seqlens = torch.cumsum(torch.tensor(chunk_seqlens, dtype=torch.int32), dim=0).tolist()
+
+                output_head, lse_head = self.impl._attention_with_optional_kv_select(
+                    q_nope=torch.index_select(q_nope, 0, q_head_idx),
+                    q_pe=torch.index_select(q_pe, 0, q_head_idx),
+                    k_nope=torch.index_select(k_nope, 0, kv_tail_proj_idx),
+                    k_pe=torch.index_select(k_pe, 0, kv_tail_proj_idx),
+                    value=torch.index_select(value, 0, kv_tail_proj_idx),
+                    kv_attn_idx=kv_with_q_head_attn_idx_in_tail,
+                    attn_mask_seqlens=attn_mask_seqlens,
+                    actual_seq_lengths_kv=head_actual_seq_lengths_kv,
+                    mask=mask,
+                    attn_metadata=attn_metadata,
+                )
+
+                self.assertEqual(mock_npu_fia.call_count, 1)
+                self.assertEqual(mock_npu_fia.call_args.kwargs["sparse_mode"], 3)
+                self.assertIsNotNone(lse_head)
+                self.assertEqual(output_head.shape, (q_head_idx.shape[0], num_heads, v_head_dim))
+
+                mock_npu_fia.reset_mock()

tests/ut/attention/a2/test_mla_precision.py

@@ -0,0 +1,681 @@
+import math
+from functools import partial
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from vllm.config import set_current_vllm_config
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.utils.torch_utils import set_random_seed
+from vllm.v1.attention.backend import AttentionType
+from vllm.v1.attention.selector import get_attn_backend
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+from tests.ut.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
+
+
+@pytest.fixture(autouse=True)
+def default_vllm_config():
+    mock_config = MagicMock()
+    mock_config.compilation_config = MagicMock()
+    mock_config.compilation_config.custom_ops = ["all"]
+    mock_config.parallel_config = MagicMock()
+    mock_config.parallel_config.prefill_context_parallel_size = 1
+    mock_config.parallel_config.decode_context_parallel_size = 1
+    mock_config.parallel_config.tensor_parallel_size = 1
+    mock_config.model_config = MagicMock()
+    mock_config.model_config.dtype = torch.float16
+    mock_config.speculative_config = None
+    mock_config.cache_config = MagicMock()
+    mock_config.cache_config.block_size = 128
+    mock_config.kv_transfer_config = None
+
+    with set_current_vllm_config(mock_config):
+        yield mock_config
+
+
+BATCH_SPECS = {
+    "small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
+    "small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
+    "mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
+    "medium_decode": BatchSpec(
+        seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
+        query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
+    ),
+    "medium_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]),
+    "mixed_medium": BatchSpec(seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]),
+    "large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
+    "large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
+    "mixed_large": BatchSpec(seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]),
+    "single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
+    "single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
+    # encoder-only
+    "small_encoder_prefill": BatchSpec(seq_lens=[32, 64, 128, 256], query_lens=[32, 64, 128, 256]),
+    "medium_encoder_prefill": BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[256, 512, 1024, 2048]),
+    "mtp_1_plus_3": BatchSpec(seq_lens=[256, 512, 1024, 1536], query_lens=[4, 4, 4, 4]),
+}
+
+
+class MockLinear:
+    def __init__(self, out_features=128, in_features=128, device=None, dtype=torch.bfloat16):
+        self.weight = torch.randn(out_features, in_features, dtype=dtype, device=device) / math.sqrt(in_features)
+        self.quant_method = UnquantizedLinearMethod()
+
+    def __call__(self, x, **kwargs):
+        if x.size(-1) != self.weight.size(-1):
+            self.weight = torch.randn(
+                self.weight.size(0), x.size(-1), dtype=self.weight.dtype, device=x.device
+            ) / math.sqrt(x.size(-1))
+
+        return (x @ self.weight.T, None)
+
+
+class MockLayerNorm:
+    def __init__(self, normalized_shape, device=None, dtype=torch.bfloat16):
+        self.weight = torch.ones(normalized_shape, dtype=dtype, device=device)
+        self.variance_epsilon = 1e-6
+
+    def __call__(self, x):
+        return x
+
+
+class MockRotary:
+    def __init__(self):
+        pass
+
+
+def create_mla_kv_cache(
+    k_nope_contexts: list[torch.Tensor],
+    k_pe_contexts: list[torch.Tensor],
+    block_size: int,
+    num_kv_heads: int,
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    num_blocks: int,
+    common_attn_metadata: AscendCommonAttentionMetadata,
+):
+    """Create MLA KV cache with two separate tensors.
+
+    MLA KV cache layout:
+    - k_cache (nope): (num_blocks, block_size, num_kv_heads, kv_lora_rank)
+    - v_cache (rope): (num_blocks, block_size, num_kv_heads, qk_rope_head_dim)
+    """
+    k_cache = torch.zeros(num_blocks, block_size, num_kv_heads, kv_lora_rank, dtype=dtype, device=device)
+    v_cache = torch.zeros(num_blocks, block_size, num_kv_heads, qk_rope_head_dim, dtype=dtype, device=device)
+
+    seq_lens = common_attn_metadata.seq_lens.cpu()
+    query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
+    query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
+    batch_size = len(k_nope_contexts)
+    block_table = common_attn_metadata.block_table_tensor
+
+    start_block_idx = 0
+    for i in range(batch_size):
+        k_nope_ctx = k_nope_contexts[i]
+        k_pe_ctx = k_pe_contexts[i]
+        context_len = int(seq_lens[i]) - int(query_lens[i])
+
+        num_blocks_for_seq = (int(seq_lens[i]) + block_size - 1) // block_size
+        block_table[i, :num_blocks_for_seq] = torch.arange(
+            start_block_idx, start_block_idx + num_blocks_for_seq, dtype=torch.int32
+        )
+
+        k_cache_flat = k_cache[start_block_idx : start_block_idx + num_blocks_for_seq].view(
+            -1, num_kv_heads, kv_lora_rank
+        )
+        v_cache_flat = v_cache[start_block_idx : start_block_idx + num_blocks_for_seq].view(
+            -1, num_kv_heads, qk_rope_head_dim
+        )
+        k_cache_flat[:context_len] = k_nope_ctx[:context_len]
+        v_cache_flat[:context_len] = k_pe_ctx[:context_len]
+
+        start_block_idx += num_blocks_for_seq
+
+    slot_mapping = common_attn_metadata.slot_mapping
+    for i in range(batch_size):
+        context_len_i = int(seq_lens[i]) - int(query_lens[i])
+        token_offsets = torch.arange(int(query_lens[i])) + context_len_i
+        block_indices = token_offsets // block_size
+        token_inter_block_offsets = token_offsets % block_size
+        start = int(query_start_loc_cpu[i])
+        end = int(query_start_loc_cpu[i + 1])
+        slot_mapping[start:end] = block_table[i, block_indices] * block_size + token_inter_block_offsets.to(device).to(
+            torch.int32
+        )
+
+    return (k_cache, v_cache)
+
+
+def run_mla_attention_backend(
+    kv_cache_spec: MLAAttentionSpec,
+    vllm_config,
+    device: torch.device,
+    common_attn_metadata: AscendCommonAttentionMetadata,
+    hidden_states: torch.Tensor,
+    kv_cache: tuple[torch.Tensor, torch.Tensor],
+    dtype: torch.bfloat16,
+    attn_type: AttentionType = AttentionType.DECODER,
+):
+    from vllm_ascend.ascend_config import init_ascend_config
+
+    init_ascend_config(vllm_config)
+
+    from vllm_ascend.ops import rotary_embedding
+
+    hf_config = vllm_config.model_config.hf_text_config
+    qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
+
+    rotary_embedding._cos_cache = torch.ones(8192, qk_rope_head_dim, dtype=dtype, device=device)
+    rotary_embedding._sin_cache = torch.zeros(8192, qk_rope_head_dim, dtype=dtype, device=device)
+    rotary_embedding._cos_mla = torch.ones(8192, 1, 1, qk_rope_head_dim, dtype=dtype, device=device)
+    rotary_embedding._sin_mla = torch.zeros(8192, 1, 1, qk_rope_head_dim, dtype=dtype, device=device)
+
+    from vllm.distributed.parallel_state import GroupCoordinator
+
+    mock_tp_group = MagicMock(spec=GroupCoordinator)
+    mock_tp_group.world_size = 1
+    mock_tp_group.rank = 0
+
+    mock_weight_prefetch = MagicMock()
+    mock_weight_prefetch.maybe_prefetch_mla_or_sla_weight_in_current_stream = MagicMock()
+
+    import vllm_ascend.utils as utils_module
+
+    original_weight_prefetch = utils_module._WEIGHT_PREFETCH_METHOD
+    utils_module._WEIGHT_PREFETCH_METHOD = mock_weight_prefetch
+    try:
+        with patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group):
+            num_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+            num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+            head_size = vllm_config.model_config.get_head_size()
+
+            kv_lora_rank = getattr(hf_config, "kv_lora_rank", 512)
+            q_lora_rank = getattr(hf_config, "q_lora_rank", 1536)
+            qk_nope_head_dim = getattr(hf_config, "qk_nope_head_dim", 128)
+            qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
+            qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+            v_head_dim = getattr(hf_config, "v_head_dim", 128)
+
+            backend = get_attn_backend(head_size, dtype, None, use_mla=True, use_sparse=False, use_mm_prefix=False)
+            impl_cls = backend.get_impl_cls()
+            builder_cls = backend.get_builder_cls()
+
+            mock_layer_entry = MagicMock()
+            for layer_name in ["placeholder"]:
+                vllm_config.compilation_config.static_forward_context[layer_name] = mock_layer_entry
+
+            builder = builder_cls(
+                kv_cache_spec,
+                ["placeholder"],
+                vllm_config,
+                device,
+            )
+            attn_metadata = builder.build(
+                common_prefix_len=0,
+                common_attn_metadata=common_attn_metadata,
+            )
+
+            hidden_size = num_heads * head_size
+            q_proj_out_dim = num_heads * qk_head_dim
+            q_b_proj_out_dim = num_heads * qk_head_dim
+            kv_b_proj_out_dim = num_heads * (qk_nope_head_dim + v_head_dim)
+            kv_b_proj_in_dim = kv_lora_rank
+            o_proj_out_dim = head_size * num_heads
+            o_proj_in_dim = num_heads * v_head_dim
+
+            impl = impl_cls(
+                num_heads=num_heads,
+                head_size=head_size,
+                scale=1.0 / (head_size**0.5),
+                num_kv_heads=num_kv_heads,
+                alibi_slopes=None,
+                sliding_window=None,
+                attn_type=attn_type.value if hasattr(attn_type, "value") else attn_type,
+                kv_cache_dtype="auto",
+                logits_soft_cap=None,
+                kv_sharing_target_layer_name=None,
+                q_lora_rank=q_lora_rank,
+                kv_lora_rank=kv_lora_rank,
+                qk_nope_head_dim=qk_nope_head_dim,
+                qk_rope_head_dim=qk_rope_head_dim,
+                qk_head_dim=qk_head_dim,
+                v_head_dim=v_head_dim,
+                q_proj=MockLinear(out_features=q_proj_out_dim, in_features=hidden_size, device=device, dtype=dtype),
+                q_b_proj=MockLinear(out_features=q_b_proj_out_dim, in_features=q_lora_rank, device=device, dtype=dtype),
+                kv_b_proj=MockLinear(
+                    out_features=kv_b_proj_out_dim, in_features=kv_b_proj_in_dim, device=device, dtype=dtype
+                ),
+                o_proj=MockLinear(out_features=o_proj_out_dim, in_features=o_proj_in_dim, device=device, dtype=dtype),
+                kv_a_layernorm=MockLayerNorm(normalized_shape=kv_lora_rank, device=device, dtype=dtype),
+                q_a_layernorm=MockLayerNorm(normalized_shape=q_lora_rank, device=device, dtype=dtype),
+                rotary_emb=MockRotary(),
+                fused_qkv_a_proj=None,
+                kv_a_proj_with_mqa=MockLinear(
+                    out_features=num_kv_heads * (kv_lora_rank + qk_rope_head_dim),
+                    in_features=hidden_size,
+                    device=device,
+                    dtype=dtype,
+                ),
+            )
+
+            impl.fa_quant_layer = False
+            impl.enable_mlapo = False
+            impl.process_weights_after_loading(dtype)
+
+            output = torch.empty_like(hidden_states)
+            output = impl.forward("layer_0", hidden_states, kv_cache, attn_metadata, output=output)
+    finally:
+        utils_module._WEIGHT_PREFETCH_METHOD = original_weight_prefetch
+
+    return output, impl
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+    x_normed = x * torch.rsqrt(variance + eps)
+    return (x_normed * weight).to(x.dtype)
+
+
+def npu_interleave_rope_simple(
+    x: torch.Tensor, cos: torch.Tensor | None = None, sin: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Simulate Ascend npu_interleave_rope with default cos=1, sin=0.
+
+    Interleave the last dimension: [x0, x1, x2, x3, ...] -> [x0, x2, ..., x1, x3, ...]
+    With cos=1, sin=0 (rope disabled), the function just returns the interleaved result.
+    """
+    even = x[..., 0::2]
+    odd = x[..., 1::2]
+    return torch.cat([even, odd], dim=-1).contiguous()
+
+
+def prefill_sdpa(
+    q_nope: torch.Tensor,
+    q_pe_raw: torch.Tensor,
+    k_pe_c: torch.Tensor,
+    k_nope_c: torch.Tensor,
+    v: torch.Tensor,
+    impl,
+    scale: float,
+    is_causal: bool = True,
+    context_len: int = 0,
+) -> torch.Tensor:
+    q_pe = npu_interleave_rope_simple(q_pe_raw)
+    q_full = torch.cat([q_nope, q_pe], dim=-1)
+    k_full = torch.cat([k_nope_c, k_pe_c], dim=-1)
+
+    q_sdpa = q_full.unsqueeze(0).transpose(1, 2)
+    k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
+    v_sdpa = v.unsqueeze(0).transpose(1, 2)
+
+    if context_len > 0 and is_causal:
+        q_len = q_full.shape[0]
+        kv_len = k_full.shape[0]
+        mask = torch.tril(
+            torch.ones(q_len, kv_len, device=q_full.device, dtype=torch.bool),
+            diagonal=context_len,
+        )
+        attn_out = torch.nn.functional.scaled_dot_product_attention(
+            q_sdpa,
+            k_sdpa,
+            v_sdpa,
+            attn_mask=mask,
+            enable_gqa=False,
+            scale=scale,
+        )
+    else:
+        attn_out = torch.nn.functional.scaled_dot_product_attention(
+            q_sdpa,
+            k_sdpa,
+            v_sdpa,
+            is_causal=is_causal,
+            enable_gqa=False,
+            scale=scale,
+        )
+    return attn_out.transpose(1, 2).squeeze(0)
+
+
+def decode_sdpa(
+    ql_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    k_pe: torch.Tensor,
+    k_nope: torch.Tensor,
+    impl,
+    scale: float,
+) -> torch.Tensor:
+    """Compute SDPA for decode path.
+
+    Decode path: q is W_UK projected (ql_nope), k/v are raw latent.
+    """
+    q_full = torch.cat([ql_nope, q_pe], dim=-1)
+    k_full = torch.cat([k_nope, k_pe], dim=-1)
+
+    q_sdpa = q_full.unsqueeze(0).transpose(1, 2)
+    k_sdpa = k_full.unsqueeze(0).transpose(1, 2)
+    v_sdpa = k_nope.unsqueeze(0).transpose(1, 2)
+
+    attn_out = torch.nn.functional.scaled_dot_product_attention(
+        q_sdpa,
+        k_sdpa,
+        v_sdpa,
+        is_causal=False,
+        enable_gqa=(impl.num_heads != impl.num_kv_heads),
+        scale=scale,
+    )
+    return attn_out.transpose(1, 2).squeeze(0)
+
+
+def compute_mla_sdpa_reference(
+    hidden_states: torch.Tensor,
+    k_nope_contexts: list[torch.Tensor],
+    k_pe_contexts: list[torch.Tensor],
+    impl,
+    batch_spec: BatchSpec,
+    scale: float,
+) -> torch.Tensor:
+    """Compute MLA reference using SDPA as golden baseline.
+
+    Handles three modes:
+    - Decode only: q via W_UK, k/v raw latent, _v_up_proj
+    - Prefill only: q raw, k/v via kv_b_proj, no _v_up_proj
+    - Mixed: decode first, then prefill, combined before o_proj
+    """
+    rms_eps = impl.kv_a_layernorm.variance_epsilon
+    rms_w = impl.kv_a_layernorm.weight
+
+    # --- Identify decode vs prefill tokens ---
+    num_decode_tokens = 0
+    for ql in batch_spec.query_lens:
+        if ql == 1:
+            num_decode_tokens += 1
+
+    # --- KV projection and normalization for all tokens ---
+    kv_no_split = impl.kv_a_proj_with_mqa(hidden_states)[0]
+    k_nope_all = kv_no_split[:, : impl.kv_lora_rank]
+    k_pe_all = kv_no_split[:, impl.kv_lora_rank :]
+    k_nope_normed = rms_norm(k_nope_all, rms_w, rms_eps)
+
+    # --- Q projection for all tokens ---
+    q_b = impl.q_proj(hidden_states)[0]
+    q_nope_all, q_pe_all = q_b.view(-1, impl.num_heads, impl.qk_head_dim).split(
+        [impl.qk_nope_head_dim, impl.qk_rope_head_dim], dim=-1
+    )
+
+    # Decode: q_nope is W_UK projected; Prefill: q_nope stays raw
+    q_nope_decode = q_nope_all[:num_decode_tokens]
+    q_nope_decode_t = q_nope_decode.transpose(0, 1).float()
+    ql_nope_decode = torch.bmm(q_nope_decode_t, impl.W_UK_T.float()).to(q_nope_all.dtype)
+    ql_nope_decode = ql_nope_decode.transpose(0, 1)
+
+    q_pe_decode = npu_interleave_rope_simple(q_pe_all[:num_decode_tokens])
+
+    q_nope_prefill = q_nope_all[num_decode_tokens:]
+    q_pe_prefill_raw = q_pe_all[num_decode_tokens:]
+
+    # --- Process each sequence ---
+    decode_outputs = []
+    prefill_outputs = []
+
+    token_offset = 0
+    for i in range(len(batch_spec.seq_lens)):
+        s_len_i = batch_spec.seq_lens[i]
+        q_len_i = batch_spec.query_lens[i]
+        context_len_i = s_len_i - q_len_i
+        is_decode = q_len_i == 1
+
+        if is_decode:
+            ql_nope_i = ql_nope_decode[token_offset : token_offset + q_len_i]
+            q_pe_i = q_pe_decode[token_offset : token_offset + q_len_i]
+
+            k_nope_dec = k_nope_normed[token_offset : token_offset + q_len_i].view(
+                q_len_i, impl.num_kv_heads, impl.kv_lora_rank
+            )
+            k_pe_dec = k_pe_all[token_offset : token_offset + q_len_i].view(
+                q_len_i, impl.num_kv_heads, impl.qk_rope_head_dim
+            )
+            k_pe_dec = npu_interleave_rope_simple(k_pe_dec)
+
+            k_nope_full_i = torch.cat([k_nope_contexts[i], k_nope_dec], dim=0)
+            k_pe_full_i = torch.cat([k_pe_contexts[i], k_pe_dec], dim=0)
+
+            sdpa_out = decode_sdpa(
+                ql_nope_i,
+                q_pe_i,
+                k_pe_full_i,
+                k_nope_full_i,
+                impl,
+                scale,
+            )
+            # _v_up_proj
+            sdpa_out = sdpa_out.transpose(0, 1).contiguous()
+            sdpa_out = torch.bmm(sdpa_out.float(), impl.W_UV.float()).to(sdpa_out.dtype)
+            sdpa_out = sdpa_out.permute(1, 0, 2)
+            sdpa_out = sdpa_out.reshape(-1, impl.num_heads * impl.v_head_dim)
+            decode_outputs.append(sdpa_out)
+
+        else:
+            q_nope_i = q_nope_prefill[token_offset - num_decode_tokens : token_offset - num_decode_tokens + q_len_i]
+            q_pe_raw_i = q_pe_prefill_raw[token_offset - num_decode_tokens : token_offset - num_decode_tokens + q_len_i]
+
+            k_nope_new = k_nope_normed[token_offset : token_offset + q_len_i].view(
+                q_len_i, impl.num_kv_heads, impl.kv_lora_rank
+            )
+            k_pe_raw_new = k_pe_all[token_offset : token_offset + q_len_i].view(
+                q_len_i, impl.num_kv_heads, impl.qk_rope_head_dim
+            )
+
+            # kv_b_proj on RMS-normed new tokens: k_nope_proj + v_proj
+            k_nope_new_flat = k_nope_new.view(-1, impl.kv_lora_rank)
+            kv_b_new = impl.kv_b_proj(k_nope_new_flat)[0].view(
+                q_len_i, impl.num_heads, impl.qk_nope_head_dim + impl.v_head_dim
+            )
+            k_nope_proj_new, v_proj_new = kv_b_new.split([impl.qk_nope_head_dim, impl.v_head_dim], dim=-1)
+
+            # kv_b_proj on raw context (cache stores raw data)
+            k_nope_ctx_raw = (
+                k_nope_contexts[i]
+                .view(context_len_i, impl.num_kv_heads, impl.kv_lora_rank)
+                .reshape(-1, impl.kv_lora_rank)
+            )
+            kv_b_ctx = impl.kv_b_proj(k_nope_ctx_raw)[0].view(
+                context_len_i, impl.num_heads, impl.qk_nope_head_dim + impl.v_head_dim
+            )
+            k_nope_proj_ctx, v_proj_ctx = kv_b_ctx.split([impl.qk_nope_head_dim, impl.v_head_dim], dim=-1)
+
+            k_nope_proj = torch.cat([k_nope_proj_ctx, k_nope_proj_new], dim=0)
+            v_proj = torch.cat([v_proj_ctx, v_proj_new], dim=0)
+
+            # k_pe: context raw (as stored in cache), new interleaved
+            k_pe_raw_ctx = k_pe_contexts[i].view(context_len_i, impl.num_kv_heads, impl.qk_rope_head_dim)
+            k_pe_new_interleaved = npu_interleave_rope_simple(k_pe_raw_new)
+            k_pe_cat = torch.cat([k_pe_raw_ctx, k_pe_new_interleaved], dim=0)
+            k_pe_expanded = k_pe_cat.expand(*k_nope_proj.shape[:-1], -1)
+
+            sdpa_out = prefill_sdpa(
+                q_nope_i,
+                q_pe_raw_i,
+                k_pe_expanded,
+                k_nope_proj,
+                v_proj,
+                impl,
+                scale,
+                is_causal=True,
+                context_len=context_len_i,
+            )
+            sdpa_out = sdpa_out.reshape(q_len_i, impl.num_heads * impl.v_head_dim)
+            prefill_outputs.append(sdpa_out)
+
+        token_offset += q_len_i
+
+    # --- Combine outputs in token order ---
+    all_outputs = decode_outputs + prefill_outputs
+    attn_output = torch.cat(all_outputs, dim=0)
+
+    # --- Output projection ---
+    sdpa_output = impl.o_proj(attn_output)[0]
+    return sdpa_output
+
+
+def _test_mla_attention_correctness(
+    batch_spec: BatchSpec,
+    model: str,
+    *,
+    attn_type: AttentionType = AttentionType.DECODER,
+    block_size: int = 128,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+    tensor_parallel_size: int = 1,
+):
+    set_random_seed(42)
+
+    vllm_config = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=1,
+        max_model_len=max(batch_spec.seq_lens),
+        block_size=block_size,
+        num_gpu_blocks=8192,
+    )
+    device = torch.device("npu")
+
+    hf_config = vllm_config.model_config.hf_text_config
+    kv_lora_rank = getattr(hf_config, "kv_lora_rank", 512)
+    qk_rope_head_dim = getattr(hf_config, "qk_rope_head_dim", 64)
+
+    batch_size = batch_spec.batch_size
+    num_q_heads = vllm_config.model_config.get_num_attention_heads(vllm_config.parallel_config)
+    num_kv_heads = vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config)
+    head_size = vllm_config.model_config.get_head_size()
+    dtype = torch.bfloat16
+    scale = 1.0 / (head_size**0.5)
+
+    kv_cache_spec = MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=num_kv_heads,
+        head_size=kv_lora_rank,
+        dtype=dtype,
+    )
+
+    k_nope_contexts, k_pe_contexts = [], []
+    for i in range(batch_size):
+        s_len = batch_spec.seq_lens[i]
+        q_len = batch_spec.query_lens[i]
+        context_len = s_len - q_len
+
+        k_nope_full = torch.randn(s_len, num_kv_heads, kv_lora_rank, dtype=dtype, device=device)
+        k_pe_full = torch.randn(s_len, num_kv_heads, qk_rope_head_dim, dtype=dtype, device=device)
+
+        k_nope_contexts.append(k_nope_full[:context_len])
+        k_pe_contexts.append(k_pe_full[:context_len])
+
+    common_attn_metadata = create_common_attn_metadata(batch_spec, vllm_config.cache_config.block_size, device)
+    if attn_type == AttentionType.ENCODER_ONLY:
+        common_attn_metadata.causal = False
+
+    num_blocks = sum((s + block_size - 1) // block_size for s in batch_spec.seq_lens)
+    num_blocks = max(16, num_blocks)
+
+    kv_cache = create_mla_kv_cache(
+        k_nope_contexts=k_nope_contexts,
+        k_pe_contexts=k_pe_contexts,
+        block_size=block_size,
+        num_kv_heads=num_kv_heads,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        dtype=dtype,
+        device=device,
+        num_blocks=num_blocks,
+        common_attn_metadata=common_attn_metadata,
+    )
+
+    num_tokens = common_attn_metadata.num_actual_tokens
+    hidden_size = num_q_heads * head_size
+    hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+
+    with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
+        from vllm.forward_context import get_forward_context
+
+        forward_ctx = get_forward_context()
+        forward_ctx.num_tokens = num_tokens
+        forward_ctx.is_draft_model = False
+        forward_ctx.is_draft_model_prefill = False
+        forward_ctx.capturing = False
+        forward_ctx.flash_comm_v1_enabled = False
+        forward_ctx.flashcomm_v2_enabled = False
+
+        backend_output, impl = run_mla_attention_backend(
+            kv_cache_spec,
+            vllm_config,
+            device,
+            common_attn_metadata,
+            hidden_states,
+            kv_cache,
+            dtype,
+            attn_type=attn_type,
+        )
+
+    sdpa_output = compute_mla_sdpa_reference(
+        hidden_states,
+        k_nope_contexts,
+        k_pe_contexts,
+        impl,
+        batch_spec,
+        scale,
+    )
+
+    # Compare (same pattern as test_gqa.py)
+    name = "MLA"
+    assert backend_output.shape == sdpa_output.shape, (
+        f"[{name}] shape {backend_output.shape} != SDPA shape {sdpa_output.shape}"
+    )
+    assert backend_output.dtype == sdpa_output.dtype, (
+        f"[{name}] dtype {backend_output.dtype} != SDPA dtype {sdpa_output.dtype}"
+    )
+
+    assert torch.isfinite(backend_output).all(), f"[{name}] produced non-finite values"
+
+    def error_msg(msg: str, backend_name: str):
+        return f"[{backend_name}] output differs from SDPA baseline. {msg}"
+
+    torch.testing.assert_close(
+        backend_output,
+        sdpa_output,
+        rtol=rtol,
+        atol=atol,
+        msg=partial(error_msg, backend_name="MLA"),
+    )
+
+
+@pytest.mark.parametrize(
+    "batch_spec_name",
+    [
+        "small_decode",
+        "small_prefill",
+        "mixed_small",
+        "medium_decode",
+        "medium_prefill",
+        "mixed_medium",
+        "large_decode",
+        "large_prefill",
+        "single_decode",
+        "single_prefill",
+        "mtp_1_plus_3",
+    ],
+)
+@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V2"])
+@pytest.mark.parametrize("tensor_parallel_size", [1])
+def test_mla_backend_correctness(default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int):
+    batch_spec = BATCH_SPECS[batch_spec_name]
+
+    _test_mla_attention_correctness(
+        batch_spec,
+        model,
+        tensor_parallel_size=tensor_parallel_size,
+    )

tests/ut/attention/a2/test_sfa_cp_precision.py

@@ -0,0 +1,844 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import math
+import sys
+from collections.abc import Callable
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+if "torch_npu._inductor" not in sys.modules:
+    sys.modules["torch_npu._inductor"] = MagicMock()
+
+from vllm.config import VllmConfig  # noqa: E402
+from vllm.forward_context import set_forward_context  # noqa: E402
+
+from tests.ut.attention.utils import (  # noqa: E402
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm_ascend.attention.context_parallel.sfa_cp import AscendSFACPImpl  # noqa: E402
+from vllm_ascend.utils import enable_custom_op
+
+enable_custom_op()
+
+SPARSE_COUNT = 2048
+_BLOCK_SIZE = 128
+_TEST_NUM_HEADS = 8
+DEFAULT_RTOL = 1e-2
+DEFAULT_ATOL = 1e-2
+
+_MAX_SIG_REL_ERR = 1e-2  # max |out-ref| / peak |ref|
+_MAX_MEAN_SIG_ERR = 5e-3  # mean |out-ref| / mean |ref|
+_MAX_REL_ERR = 1e-2  # max per-element rel err where |ref| >= floor
+_SIG_FLOOR_FRAC = 5e-1  # floor = this fraction of peak |ref|
+
+
+def _validate_spec(spec: BatchSpec) -> None:
+    """Require ``seq_len <= SPARSE_COUNT`` so sparse matches dense reference."""
+    for s, q in zip(spec.seq_lens, spec.query_lens):
+        assert q <= s, f"query_len ({q}) must not exceed seq_len ({s})"
+        assert s <= SPARSE_COUNT, (
+            f"seq_len ({s}) must be <= SPARSE_COUNT ({SPARSE_COUNT}) so the "
+            "sparse attention degenerates into dense attention for the "
+            "reference comparison."
+        )
+
+
+_VLLM_CONFIG_CACHE: dict = {}
+
+
+def _get_vllm_config(
+    model: str,
+    dtype: torch.dtype,
+    *,
+    max_model_len: int = 4096,
+    tensor_parallel_size: int = 1,
+) -> VllmConfig:
+    key = (model, dtype, tensor_parallel_size)
+    cfg = _VLLM_CONFIG_CACHE.get(key)
+    if cfg is not None:
+        return cfg
+    dtype_str = "bfloat16" if dtype == torch.bfloat16 else "float16"
+    sim_num_heads = max(1, _TEST_NUM_HEADS // tensor_parallel_size)
+    cfg = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=1,  # always TP=1; head split is simulated
+        max_model_len=max_model_len,
+        dtype=dtype_str,
+        block_size=_BLOCK_SIZE,
+        num_gpu_blocks=4096,
+        max_num_seqs=64,
+        max_num_batched_tokens=max(8192, max_model_len * 2),
+        enable_chunked_prefill=True,
+        hf_overrides={"quantization_config": None},
+        hf_config_override={
+            "num_attention_heads": sim_num_heads,
+            "num_key_value_heads": 1,
+        },
+    )
+    _VLLM_CONFIG_CACHE[key] = cfg
+    return cfg
+
+
+# Spec name prefixes drive the SFA-CP branch under test.
+BATCH_SPECS: dict[str, BatchSpec] = {
+    "decode_single": BatchSpec(seq_lens=[1024], query_lens=[1], name="decode_single"),
+    "decode_small_batch": BatchSpec(
+        seq_lens=[512, 1024, 1536, 2048], query_lens=[1, 1, 1, 1], name="decode_small_batch"
+    ),
+    "decode_large_batch": BatchSpec(seq_lens=[2048] * 8, query_lens=[1] * 8, name="decode_large_batch"),
+    "mtp_1_plus_1": BatchSpec(seq_lens=[512, 1024, 1536], query_lens=[2, 2, 2], name="mtp_1_plus_1"),
+    "mtp_1_plus_3": BatchSpec(seq_lens=[1024, 1536, 2048, 2048], query_lens=[4, 4, 4, 4], name="mtp_1_plus_3"),
+    "mtp_1_plus_7": BatchSpec(seq_lens=[1024, 1536, 2048], query_lens=[8, 8, 8], name="mtp_1_plus_7"),
+    "prefill_single": BatchSpec(seq_lens=[256], query_lens=[256], name="prefill_single"),
+    "prefill_small_batch": BatchSpec(seq_lens=[256, 512, 384], query_lens=[256, 512, 384], name="prefill_small_batch"),
+    "prefill_with_context": BatchSpec(seq_lens=[512, 1024], query_lens=[128, 256], name="prefill_with_context"),
+    "mixed_small": BatchSpec(seq_lens=[512, 1024, 256, 512], query_lens=[1, 1, 64, 128], name="mixed_small"),
+    "mixed_medium": BatchSpec(
+        seq_lens=[1024, 1536, 2048, 256, 512], query_lens=[1, 1, 1, 64, 128], name="mixed_medium"
+    ),
+}
+
+
+def _infer_mode(spec: BatchSpec) -> str:
+    """Return one of: ``decode``, ``mtp``, ``prefill``, ``mixed``."""
+    name = spec.name
+    for prefix in ("decode_", "mtp_", "prefill_", "mixed_"):
+        if name.startswith(prefix):
+            return prefix.rstrip("_")
+    raise ValueError(
+        f"BatchSpec name {name!r} does not start with a known mode prefix ('decode_', 'mtp_', 'prefill_', 'mixed_')"
+    )
+
+
+def _build_topk_indices(
+    seq_lens: list[int],
+    query_lens: list[int],
+    sparse_count: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Build causal topk indices with INVALID_IDX (-1) padding."""
+    num_tokens = sum(query_lens)
+    topk = torch.full((num_tokens, 1, sparse_count), -1, dtype=torch.int32, device=device)
+
+    cum_q = 0
+    for b, s_len in enumerate(seq_lens):
+        q_len = query_lens[b]
+        ctx_len = s_len - q_len
+        for j in range(q_len):
+            valid_end = ctx_len + j + 1
+            topk[cum_q + j, 0, :valid_end] = torch.arange(valid_end, dtype=torch.int32, device=device)
+        cum_q += q_len
+
+    return topk
+
+
+def _reference_sparse_attention(
+    ql_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    k_nope_full_per_req: list[torch.Tensor],
+    k_rope_full_per_req: list[torch.Tensor],
+    seq_lens: list[int],
+    query_lens: list[int],
+    scale: float,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Pure-PyTorch dense MQA baseline in fp32."""
+    outputs: list[torch.Tensor] = []
+
+    cum_q = 0
+    for b, s_len in enumerate(seq_lens):
+        q_len = query_lens[b]
+        ctx_len = s_len - q_len
+
+        K_nope = k_nope_full_per_req[b][:s_len].float()
+        K_rope = k_rope_full_per_req[b][:s_len].float()
+        K = torch.cat([K_nope, K_rope], dim=-1)
+        V = K_nope
+
+        for j in range(q_len):
+            t = cum_q + j
+            valid_end = ctx_len + j + 1
+
+            q_n = ql_nope[t].float()
+            q_p = q_pe[t].float()
+            Q = torch.cat([q_n, q_p], dim=-1)
+
+            K_b = K[:valid_end]
+            V_b = V[:valid_end]
+
+            scores = (Q @ K_b.transpose(0, 1)) * scale
+            attn = torch.softmax(scores, dim=-1)
+            out = attn @ V_b
+            outputs.append(out.to(out_dtype))
+
+        cum_q += q_len
+
+    return torch.stack(outputs, dim=0)
+
+
+def _build_cp_paged_kv_cache(
+    seq_lens: list[int],
+    k_nope_contexts: list[torch.Tensor],
+    k_rope_contexts: list[torch.Tensor],
+    block_size: int,
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+    cp_size: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Allocate a paged KV cache that matches the *decode* CP gather/block_table layout.
+
+    Each request is padded so its block count is exactly ``cp_size * L``
+    (i.e. divisible by the total CP world size). The local rank is treated
+    as "rank 0" and owns the first ``L`` of each request's logical blocks
+    (the actual values in ``local_block_table`` only matter via shape because
+    ``gather_kv_cross_cp`` is mocked).
+
+    The ``gathered_*`` tensors mirror what the impl would see after both
+    DCP and PCP all-gathers; layout matches the indexing produced by
+    ``gather_block_table``:
+        ``gathered[i*N + r*L + b]`` holds K data for logical block
+        ``b*cp_size + i`` of request ``r`` (where ``N == batch_size * L``).
+    """
+    batch_size = len(seq_lens)
+
+    raw_blocks_per_req = [(s + block_size - 1) // block_size for s in seq_lens]
+    max_raw = max(raw_blocks_per_req)
+    total_blocks_per_req = ((max_raw + cp_size - 1) // cp_size) * cp_size
+    L = total_blocks_per_req // cp_size
+
+    total_blocks = batch_size * total_blocks_per_req + 1  # +1 reserves block 0
+
+    full_k_nope_cache = torch.zeros(total_blocks, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
+    full_k_rope_cache = torch.zeros(total_blocks, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
+    full_block_table = torch.zeros(batch_size, total_blocks_per_req, dtype=torch.int32, device=device)
+
+    next_block = 1
+    for r in range(batch_size):
+        s_len = seq_lens[r]
+        for p in range(total_blocks_per_req):
+            full_block_table[r, p] = next_block
+            tok_start = p * block_size
+            tok_end = min(tok_start + block_size, s_len)
+            length = max(tok_end - tok_start, 0)
+            if length > 0:
+                full_k_nope_cache[next_block, :length, 0, :] = k_nope_contexts[r][tok_start:tok_end]
+                full_k_rope_cache[next_block, :length, 0, :] = k_rope_contexts[r][tok_start:tok_end]
+            next_block += 1
+
+    local_block_table = full_block_table[:, :L].contiguous()
+
+    N = batch_size * L
+    gathered_k_nope = torch.zeros(cp_size * N, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
+    gathered_k_rope = torch.zeros(cp_size * N, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
+    for i in range(cp_size):
+        for r in range(batch_size):
+            for b in range(L):
+                p_logical = b * cp_size + i
+                phys_block = int(full_block_table[r, p_logical].item())
+                dst = i * N + r * L + b
+                gathered_k_nope[dst] = full_k_nope_cache[phys_block]
+                gathered_k_rope[dst] = full_k_rope_cache[phys_block]
+
+    return (
+        full_k_nope_cache,
+        full_k_rope_cache,
+        full_block_table,
+        local_block_table,
+        gathered_k_nope,
+        gathered_k_rope,
+        L,
+    )
+
+
+def _build_cp_prefill_compact_metadata(
+    prefill_full_block_table: torch.Tensor,
+    full_k_nope_cache: torch.Tensor,
+    full_k_rope_cache: torch.Tensor,
+    cp_size: int,
+    block_size: int,
+    seq_lens: list[int],
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Build the prefill compact metadata + matching gathered KV view.
+
+    Reproduces the production
+    ``AscendSFACPMetadataBuilder.build_prefill_compact_block_metadata``
+    formula: ``block_table_cp[r, b*cp_size + i] = new_block_table[r, b] + i*M``
+    where ``M = num_unique_local_blocks``. Then materialises the
+    ``gathered_compact`` tensor that ``gather_kv_cross_cp_compact`` would
+    produce so that the kernel, when walking ``block_table_cp[r, :]``,
+    sees the contiguous K context of request ``r`` block-by-block.
+
+    Inputs
+    ------
+    prefill_full_block_table:
+        Shape ``(num_prefill_reqs, L * cp_size)``. The "full" (logical)
+        block table for prefill requests as built by
+        ``_build_cp_paged_kv_cache``. Entry ``[r, p]`` holds the physical
+        block id that owns logical position ``p`` of request ``r``.
+
+    The "local" (rank-0) prefill block table -- which the production code
+    actually receives in ``attn_metadata.block_table[num_decodes:]`` -- is
+    derived as ``prefill_full_block_table[:, :L]``. This matches the
+    convention used by ``_build_cp_paged_kv_cache`` for the decode path.
+    """
+    num_prefill_reqs, total_blocks_per_req = prefill_full_block_table.shape
+    assert total_blocks_per_req % cp_size == 0, (
+        f"prefill_full_block_table has {total_blocks_per_req} columns, which is not divisible by cp_size={cp_size}"
+    )
+    L = total_blocks_per_req // cp_size
+
+    prefill_local_block_table = prefill_full_block_table[:, :L].contiguous()
+
+    block_arange = torch.arange(cp_size, dtype=prefill_local_block_table.dtype, device=device)
+    valid_block_ids, new_block_table_flat = prefill_local_block_table.flatten().unique(return_inverse=True)
+    num_blocks = valid_block_ids.shape[0]
+    block_table_cp = (
+        new_block_table_flat.unsqueeze(-1).to(prefill_local_block_table)
+        + (block_arange * num_blocks).view(1, 1, -1).to(prefill_local_block_table)
+    ).reshape(prefill_local_block_table.shape[0], -1)
+
+    new_block_table_2d = new_block_table_flat.view(prefill_local_block_table.shape)
+
+    M = int(num_blocks)
+    gathered_compact_nope = torch.zeros(cp_size * M, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
+    gathered_compact_rope = torch.zeros(cp_size * M, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
+
+    for r in range(num_prefill_reqs):
+        s_len = seq_lens[r]
+        for b_local in range(L):
+            for i in range(cp_size):
+                p_logical = b_local * cp_size + i
+                tok_start = p_logical * block_size
+                tok_end = min(tok_start + block_size, s_len)
+                length = max(tok_end - tok_start, 0)
+                if length <= 0:
+                    continue
+                src_block = int(prefill_full_block_table[r, p_logical].item())
+                dst = int(new_block_table_2d[r, b_local].item()) + i * M
+                gathered_compact_nope[dst, :length, 0, :] = full_k_nope_cache[src_block, :length, 0, :]
+                gathered_compact_rope[dst, :length, 0, :] = full_k_rope_cache[src_block, :length, 0, :]
+
+    return valid_block_ids, block_table_cp, gathered_compact_nope, gathered_compact_rope
+
+
+def _make_fake_self(
+    *,
+    scale: float,
+    pcp_size: int,
+    dcp_size: int,
+    gather_kv_cross_cp_fn: Callable | None = None,
+    gather_kv_cross_cp_compact_fn: Callable | None = None,
+) -> MagicMock:
+    """Construct a ``MagicMock`` matching the slice of ``AscendSFACPImpl``
+    that ``_execute_sparse_flash_attention_process`` reads.
+
+    Bound methods that don't depend on init-time state
+    (``gather_block_table``, ``_execute_sparse_flash_attention``) are
+    delegated to the real (unbound) implementations so the production
+    kernel call stays under test. The collective gathers
+    (``gather_kv_cross_cp`` / ``gather_kv_cross_cp_compact``) are mocked
+    via ``side_effect`` because they require an initialised PCP / DCP
+    ``ProcessGroupHCCL``, which can't be created on a single rank.
+    """
+    fake_self = MagicMock()
+    fake_self.scale = scale
+    fake_self.pcp_size = pcp_size
+    fake_self.dcp_size = dcp_size
+
+    if gather_kv_cross_cp_fn is not None:
+        fake_self.gather_kv_cross_cp = MagicMock(side_effect=gather_kv_cross_cp_fn)
+    if gather_kv_cross_cp_compact_fn is not None:
+        fake_self.gather_kv_cross_cp_compact = MagicMock(side_effect=gather_kv_cross_cp_compact_fn)
+
+    fake_self.gather_block_table = lambda block_num, block_tables, block_arange: AscendSFACPImpl.gather_block_table(
+        fake_self, block_num, block_tables, block_arange
+    )
+    fake_self._execute_sparse_flash_attention = lambda *args, **kwargs: AscendSFACPImpl._execute_sparse_flash_attention(
+        fake_self, *args, **kwargs
+    )
+    fake_self._align_to_graph_bucket_tokens = lambda x, m: x
+    return fake_self
+
+
+def _run_sfa_cp_kernel(
+    *,
+    ql_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    full_k_nope_cache: torch.Tensor,
+    full_k_rope_cache: torch.Tensor,
+    local_block_table: torch.Tensor,
+    topk_indices: torch.Tensor,
+    cum_query_lens: torch.Tensor,
+    seq_lens_tensor: torch.Tensor,
+    scale: float,
+    pcp_size: int,
+    dcp_size: int,
+    num_decodes: int,
+    num_decode_tokens: int,
+    num_prefills: int,
+    device: torch.device,
+    gathered_k_nope: torch.Tensor | None = None,
+    gathered_k_rope: torch.Tensor | None = None,
+    valid_block_ids: torch.Tensor | None = None,
+    block_table_cp: torch.Tensor | None = None,
+    gathered_compact_nope: torch.Tensor | None = None,
+    gathered_compact_rope: torch.Tensor | None = None,
+    prefill_q_cum_seqlens: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run the SFA-CP kernel with mocked CP collectives."""
+    cp_size = pcp_size * dcp_size
+
+    gather_kv_cross_cp_fn = None
+    gather_kv_cross_cp_compact_fn = None
+
+    if num_decode_tokens > 0:
+        assert gathered_k_nope is not None and gathered_k_rope is not None, (
+            "decode branch requires gathered_k_nope/gathered_k_rope"
+        )
+        gathered_lookup = {
+            id(full_k_nope_cache): gathered_k_nope,
+            id(full_k_rope_cache): gathered_k_rope,
+        }
+
+        def gather_kv_cross_cp_fn(kv: torch.Tensor, block_tables: torch.Tensor):
+            gathered = gathered_lookup.get(id(kv))
+            assert gathered is not None, "gather_kv_cross_cp called with an unexpected kv tensor"
+
+            return gathered, block_tables.numel()
+
+    if num_prefills > 0:
+        assert (
+            valid_block_ids is not None
+            and block_table_cp is not None
+            and gathered_compact_nope is not None
+            and gathered_compact_rope is not None
+            and prefill_q_cum_seqlens is not None
+        ), (
+            "prefill branch requires valid_block_ids / block_table_cp / "
+            "gathered_compact_{nope,rope} / prefill_q_cum_seqlens"
+        )
+        gathered_compact_lookup = {
+            id(full_k_nope_cache): gathered_compact_nope,
+            id(full_k_rope_cache): gathered_compact_rope,
+        }
+
+        def gather_kv_cross_cp_compact_fn(kv: torch.Tensor, vbid: torch.Tensor):
+            gathered = gathered_compact_lookup.get(id(kv))
+            assert gathered is not None, "gather_kv_cross_cp_compact called with an unexpected kv tensor"
+
+            return gathered
+
+    fake_self = _make_fake_self(
+        scale=scale,
+        pcp_size=pcp_size,
+        dcp_size=dcp_size,
+        gather_kv_cross_cp_fn=gather_kv_cross_cp_fn,
+        gather_kv_cross_cp_compact_fn=gather_kv_cross_cp_compact_fn,
+    )
+
+    fake_attn_metadata = MagicMock()
+    fake_attn_metadata.block_table = local_block_table
+    fake_attn_metadata.num_decodes = num_decodes
+    fake_attn_metadata.num_decode_tokens = num_decode_tokens
+    fake_attn_metadata.num_prefills = num_prefills
+
+    fake_sfa_cp_metadata = MagicMock()
+    fake_sfa_cp_metadata.block_arange = torch.arange(cp_size, dtype=torch.int32, device=device)
+    if num_prefills > 0:
+        fake_sfa_cp_metadata.valid_block_ids = valid_block_ids
+        fake_sfa_cp_metadata.block_table_cp = block_table_cp
+        fake_sfa_cp_metadata.prefill_q_cum_seqlens = prefill_q_cum_seqlens
+    fake_attn_metadata.sfa_cp_metadata = fake_sfa_cp_metadata
+
+    cum_lens_arg = cum_query_lens if num_decode_tokens > 0 else prefill_q_cum_seqlens
+
+    return AscendSFACPImpl._execute_sparse_flash_attention_process(
+        fake_self,
+        ql_nope,
+        q_pe,
+        (full_k_nope_cache, full_k_rope_cache),
+        topk_indices,
+        fake_attn_metadata,
+        cum_lens_arg,
+        seq_lens_tensor,
+    )
+
+
+def _record_and_assert(
+    backend_output: torch.Tensor,
+    reference_output: torch.Tensor,
+    tag: str,
+    *,
+    dtype: torch.dtype,
+    atol: float = DEFAULT_ATOL,
+    rtol: float = DEFAULT_RTOL,
+) -> tuple[float, float]:
+    """Assert numerical closeness and record signal-relative metrics."""
+    assert backend_output.shape == reference_output.shape, (
+        f"[{tag}] backend shape {tuple(backend_output.shape)} != reference shape {tuple(reference_output.shape)}"
+    )
+    assert backend_output.dtype == reference_output.dtype, (
+        f"[{tag}] backend dtype {backend_output.dtype} != reference dtype {reference_output.dtype}"
+    )
+    assert torch.isfinite(backend_output).all(), f"[{tag}] sparse flash attention produced non-finite values"
+
+    torch.testing.assert_close(
+        backend_output,
+        reference_output,
+        rtol=rtol,
+        atol=atol,
+        msg=lambda m: f"[SFA-CP:{tag}] kernel output diverges from baseline. {m}",
+    )
+
+    ref_f32 = reference_output.float()
+    out_f32 = backend_output.float()
+    diff = (out_f32 - ref_f32).abs()
+    ref_abs = ref_f32.abs()
+    peak = float(ref_abs.max())
+    mean_ref_abs = float(ref_abs.mean())
+    sig_floor = peak * _SIG_FLOOR_FRAC
+
+    max_abs_err = float(diff.max())
+    mean_abs_err = float(diff.mean())
+    max_sig_rel_err = max_abs_err / peak if peak > 0 else 0.0
+    mean_sig_rel_err = mean_abs_err / mean_ref_abs if mean_ref_abs > 0 else 0.0
+
+    significant_mask = ref_abs >= sig_floor
+    if significant_mask.any():
+        per_elem_rel = diff[significant_mask] / ref_abs[significant_mask]
+        max_rel_err_sig = float(per_elem_rel.max())
+    else:
+        max_rel_err_sig = 0.0
+
+    assert max_sig_rel_err < _MAX_SIG_REL_ERR, (
+        f"[SFA-CP:{tag}] dtype={dtype} signal-relative max error "
+        f"{max_sig_rel_err * 100:.4f}% exceeds 1% budget "
+        f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})"
+    )
+    assert mean_sig_rel_err < _MAX_MEAN_SIG_ERR, (
+        f"[SFA-CP:{tag}] dtype={dtype} signal-relative mean error "
+        f"{mean_sig_rel_err * 100:.4f}% exceeds 0.5% drift budget "
+        f"(mean_ref_abs={mean_ref_abs:.4e}, mean_abs_err={mean_abs_err:.4e})"
+    )
+    assert max_rel_err_sig < _MAX_REL_ERR, (
+        f"[SFA-CP:{tag}] dtype={dtype} per-element relative error on "
+        f">={int(_SIG_FLOOR_FRAC * 100)}%-of-peak elements "
+        f"{max_rel_err_sig * 100:.4f}% exceeds 1% budget "
+        f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})"
+    )
+
+    return max_abs_err, max_rel_err_sig
+
+
+def _make_synthetic_kv_contexts(
+    seq_lens: list[int],
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+    """Per-request K context generated independently so values across
+    requests are uncorrelated."""
+    k_nope = [torch.randn(s, kv_lora_rank, dtype=dtype, device=device) * 0.1 for s in seq_lens]
+    k_rope = [torch.randn(s, qk_rope_head_dim, dtype=dtype, device=device) * 0.1 for s in seq_lens]
+    return k_nope, k_rope
+
+
+def _test_sfa_cp_correctness(
+    batch_spec: BatchSpec,
+    model: str,
+    *,
+    pcp_size: int,
+    dcp_size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    atol: float = 1e-2,
+    rtol: float = 1e-2,
+    tensor_parallel_size: int = 1,
+) -> None:
+    """Test ``AscendSFACPImpl`` against a fp32 dense MQA reference."""
+    mode = _infer_mode(batch_spec)
+    assert not (pcp_size > 1 and mode in ("prefill", "mixed")), (
+        f"PCP>1 {mode} is out of scope; the parametrize whitelist should not have generated this combination"
+    )
+
+    torch.manual_seed(2026)
+    _validate_spec(batch_spec)
+
+    vllm_config = _get_vllm_config(
+        model,
+        dtype,
+        tensor_parallel_size=tensor_parallel_size,
+    )
+    device = torch.device("npu")
+
+    seq_lens = list(batch_spec.seq_lens)
+    query_lens = list(batch_spec.query_lens)
+    batch_size = batch_spec.batch_size
+    num_tokens = batch_spec.compute_num_tokens()
+    cp_size = pcp_size * dcp_size
+
+    cache_config = vllm_config.cache_config
+    hf_text = vllm_config.model_config.hf_text_config
+    block_size = cache_config.block_size
+    kv_lora_rank = hf_text.kv_lora_rank
+    qk_rope_head_dim = hf_text.qk_rope_head_dim
+    num_heads = hf_text.num_attention_heads
+
+    head_dim = kv_lora_rank + qk_rope_head_dim
+    scale = 1.0 / math.sqrt(head_dim)
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        block_size=block_size,
+        device=device,
+    )
+
+    k_nope_contexts, k_rope_contexts = _make_synthetic_kv_contexts(
+        seq_lens,
+        kv_lora_rank,
+        qk_rope_head_dim,
+        dtype,
+        device,
+    )
+    (
+        full_k_nope_cache,
+        full_k_rope_cache,
+        full_block_table,
+        local_block_table,
+        gathered_k_nope,
+        gathered_k_rope,
+        _L,
+    ) = _build_cp_paged_kv_cache(
+        seq_lens=seq_lens,
+        k_nope_contexts=k_nope_contexts,
+        k_rope_contexts=k_rope_contexts,
+        block_size=block_size,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        cp_size=cp_size,
+        dtype=dtype,
+        device=device,
+    )
+
+    ql_nope = (
+        torch.randn(
+            num_tokens,
+            num_heads,
+            kv_lora_rank,
+            dtype=dtype,
+            device=device,
+        )
+        * 0.1
+    )
+    q_pe = (
+        torch.randn(
+            num_tokens,
+            num_heads,
+            qk_rope_head_dim,
+            dtype=dtype,
+            device=device,
+        )
+        * 0.1
+    )
+    topk_indices = _build_topk_indices(seq_lens, query_lens, SPARSE_COUNT, device)
+
+    cum_query_lens = common_attn_metadata.query_start_loc[1:].to(torch.int32)
+    seq_lens_tensor = common_attn_metadata.seq_lens.to(torch.int32)
+
+    decode_gathered_k_nope: torch.Tensor | None = None
+    decode_gathered_k_rope: torch.Tensor | None = None
+    prefill_full_block_table: torch.Tensor | None = None
+
+    if mode in ("decode", "mtp"):
+        num_decodes = batch_size
+        num_decode_tokens = num_tokens
+        num_prefills = 0
+        decode_gathered_k_nope = gathered_k_nope
+        decode_gathered_k_rope = gathered_k_rope
+        n_decode_reqs = batch_size
+    elif mode == "prefill":
+        num_decodes = 0
+        num_decode_tokens = 0
+        num_prefills = batch_size
+        prefill_full_block_table = full_block_table
+        n_decode_reqs = 0
+    else:  # mixed (chunked-prefill)
+        n_decode_reqs = sum(1 for q in query_lens if q == 1)
+        assert all(q == 1 for q in query_lens[:n_decode_reqs]), (
+            f"mixed spec {batch_spec.name}: decode requests must come first"
+        )
+        assert all(q > 1 for q in query_lens[n_decode_reqs:]), (
+            f"mixed spec {batch_spec.name}: prefill requests must come after decode"
+        )
+        num_decodes = n_decode_reqs
+        num_decode_tokens = n_decode_reqs
+        num_prefills = batch_size - n_decode_reqs
+        prefill_full_block_table = full_block_table[n_decode_reqs:]
+
+        L_decode = local_block_table[:n_decode_reqs].shape[1]
+        N_decode = n_decode_reqs * L_decode
+        decode_gathered_k_nope = torch.zeros(
+            cp_size * N_decode,
+            block_size,
+            1,
+            kv_lora_rank,
+            dtype=dtype,
+            device=device,
+        )
+        decode_gathered_k_rope = torch.zeros(
+            cp_size * N_decode,
+            block_size,
+            1,
+            qk_rope_head_dim,
+            dtype=dtype,
+            device=device,
+        )
+        for i in range(cp_size):
+            for r in range(n_decode_reqs):
+                for b in range(L_decode):
+                    p_logical = b * cp_size + i
+                    phys_block = int(full_block_table[r, p_logical].item())
+                    dst = i * N_decode + r * L_decode + b
+                    decode_gathered_k_nope[dst] = full_k_nope_cache[phys_block]
+                    decode_gathered_k_rope[dst] = full_k_rope_cache[phys_block]
+
+    valid_block_ids = None
+    block_table_cp = None
+    gathered_compact_nope = None
+    gathered_compact_rope = None
+    prefill_q_cum_seqlens = None
+    if num_prefills > 0:
+        valid_block_ids, block_table_cp, gathered_compact_nope, gathered_compact_rope = (
+            _build_cp_prefill_compact_metadata(
+                prefill_full_block_table=prefill_full_block_table,
+                full_k_nope_cache=full_k_nope_cache,
+                full_k_rope_cache=full_k_rope_cache,
+                cp_size=cp_size,
+                block_size=block_size,
+                seq_lens=seq_lens[n_decode_reqs:],
+                kv_lora_rank=kv_lora_rank,
+                qk_rope_head_dim=qk_rope_head_dim,
+                dtype=dtype,
+                device=device,
+            )
+        )
+        if n_decode_reqs > 0:
+            prefill_q_cum_seqlens = cum_query_lens[n_decode_reqs:] - cum_query_lens[n_decode_reqs - 1]
+        else:
+            prefill_q_cum_seqlens = cum_query_lens
+
+    with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
+        backend_output = _run_sfa_cp_kernel(
+            ql_nope=ql_nope,
+            q_pe=q_pe,
+            full_k_nope_cache=full_k_nope_cache,
+            full_k_rope_cache=full_k_rope_cache,
+            local_block_table=local_block_table,
+            topk_indices=topk_indices,
+            cum_query_lens=cum_query_lens,
+            seq_lens_tensor=seq_lens_tensor,
+            scale=scale,
+            pcp_size=pcp_size,
+            dcp_size=dcp_size,
+            num_decodes=num_decodes,
+            num_decode_tokens=num_decode_tokens,
+            num_prefills=num_prefills,
+            device=device,
+            gathered_k_nope=decode_gathered_k_nope,
+            gathered_k_rope=decode_gathered_k_rope,
+            valid_block_ids=valid_block_ids,
+            block_table_cp=block_table_cp,
+            gathered_compact_nope=gathered_compact_nope,
+            gathered_compact_rope=gathered_compact_rope,
+            prefill_q_cum_seqlens=prefill_q_cum_seqlens,
+        )
+    reference_output = _reference_sparse_attention(
+        ql_nope=ql_nope,
+        q_pe=q_pe,
+        k_nope_full_per_req=k_nope_contexts,
+        k_rope_full_per_req=k_rope_contexts,
+        seq_lens=seq_lens,
+        query_lens=query_lens,
+        scale=scale,
+        out_dtype=dtype,
+    )
+
+    dt = "bf16" if dtype == torch.bfloat16 else "fp16"
+    tag = f"{mode}|{batch_spec.name}|pcp={pcp_size}|dcp={dcp_size}|tp={tensor_parallel_size}|{dt}"
+
+    _record_and_assert(
+        backend_output,
+        reference_output,
+        tag,
+        dtype=dtype,
+        atol=atol,
+        rtol=rtol,
+    )
+
+
+# Whitelist of ``(batch_spec_name, pcp_size, dcp_size)`` combinations that are
+# in scope for the single-rank precision suite. ``PCP > 1`` prefill / mixed
+# requires per-rank PCP scheduler metadata (``q_head_idx``, ``q_tail_idx``,
+# ``q_full_idx``, ...) that only a multi-rank job can produce faithfully -- so
+# those combos are intentionally absent from the matrix (covered structurally
+# by ``tests/ut/attention/test_sfa_cp.py``).
+_TOPOLOGIES_ALL = [(2, 2), (2, 1), (1, 2), (1, 4)]
+_TOPOLOGIES_PCP1 = [(1, 2), (1, 4)]
+_TEST_CASES: list[tuple[str, int, int]] = [
+    (name, pcp, dcp)
+    for name in BATCH_SPECS
+    for (pcp, dcp) in (_TOPOLOGIES_ALL if _infer_mode(BATCH_SPECS[name]) in ("decode", "mtp") else _TOPOLOGIES_PCP1)
+]
+
+
+@pytest.mark.parametrize(
+    "batch_spec_name,pcp_size,dcp_size",
+    _TEST_CASES,
+    ids=[f"{n}-pcp{p}-dcp{d}" for (n, p, d) in _TEST_CASES],
+)
+@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V3.2-Exp"])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
+def test_sfa_cp_correctness(
+    batch_spec_name: str,
+    model: str,
+    pcp_size: int,
+    dcp_size: int,
+    dtype: torch.dtype,
+    tensor_parallel_size: int,
+) -> None:
+    """Test SFA-CP correctness across workload, topology, dtype, and TP size."""
+    _test_sfa_cp_correctness(
+        BATCH_SPECS[batch_spec_name],
+        model,
+        pcp_size=pcp_size,
+        dcp_size=dcp_size,
+        dtype=dtype,
+        tensor_parallel_size=tensor_parallel_size,
+    )

tests/ut/attention/a2/test_sfa_v1.py

@@ -0,0 +1,742 @@
+import sys
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.config import set_current_vllm_config
+from vllm.distributed.parallel_state import GroupCoordinator
+
+from tests.ut.attention.utils import patch_distributed_groups
+from tests.ut.base import TestBase
+from vllm_ascend.ascend_config import init_ascend_config
+from vllm_ascend.attention.attention_v1 import AscendAttentionState
+
+if "torch_npu._inductor" not in sys.modules:
+    sys.modules["torch_npu._inductor"] = MagicMock()
+
+from vllm_ascend.attention.sfa_v1 import (
+    AscendSFABackend,
+    AscendSFAImpl,
+    AscendSFAMetadata,
+    AscendSFAMetadataBuilder,
+    custom_kv_rmsnorm_rope,
+)
+from vllm_ascend.attention.utils import get_sfa_qsfa_packed_head_dim
+from vllm_ascend.device.device_op import DeviceOperator
+from vllm_ascend.utils import enable_dsa_cp
+
+
+class TestAscendSFABackend(TestBase):
+    def setUp(self):
+        self.mock_config = MagicMock()
+        mock_parallel_config = MagicMock()
+        mock_parallel_config.prefill_context_parallel_size = 1
+        mock_parallel_config.decode_context_parallel_size = 1
+        self.mock_config.parallel_config = mock_parallel_config
+        self.mock_config.model_config = MagicMock(spec=[])
+        self.config_context = set_current_vllm_config(self.mock_config)
+        self.config_context.__enter__()
+
+        self.utils_patcher = patch("vllm_ascend.attention.utils.get_current_vllm_config", return_value=self.mock_config)
+        self.utils_patcher.start()
+
+        from vllm_ascend.attention.utils import enable_cp
+
+        enable_cp.cache_clear()
+
+    def tearDown(self):
+        self.utils_patcher.stop()
+        self.config_context.__exit__(None, None, None)
+
+    def test_get_name(self):
+        self.assertEqual(AscendSFABackend.get_name(), "ASCEND_SFA")
+
+    def test_get_builder_cls(self):
+        self.assertEqual(AscendSFABackend.get_builder_cls(), AscendSFAMetadataBuilder)
+
+    def test_get_kv_cache_shape(self):
+        result = AscendSFABackend.get_kv_cache_shape(2, 4, 8, 128)
+        self.assertEqual(result, (2, 4, 8, 128))
+
+    def test_get_impl_cls(self):
+        result = AscendSFABackend.get_impl_cls()
+        self.assertEqual(result, AscendSFAImpl)
+
+    @patch("vllm_ascend.attention.sfa_v1.enable_cp")
+    def test_get_builder_cls_with_cp(self, mock_enable_cp):
+        mock_enable_cp.return_value = True
+        builder_cls = AscendSFABackend.get_builder_cls()
+        self.assertIsNotNone(builder_cls)
+
+    @patch("vllm_ascend.attention.sfa_v1.enable_cp")
+    def test_get_impl_cls_with_cp(self, mock_enable_cp):
+        mock_enable_cp.return_value = True
+        impl_cls = AscendSFABackend.get_impl_cls()
+        self.assertIsNotNone(impl_cls)
+
+
+class TestAscendSFABytePackedGather(TestBase):
+    @patch("vllm_ascend.attention.sfa_v1.get_tp_group")
+    def test_byte_packed_gather_preserves_mixed_dtype_tensors(self, mock_get_tp_group):
+        mock_get_tp_group.return_value = SimpleNamespace(world_size=1)
+        sfa_kv = torch.arange(12, dtype=torch.float16).view(2, 6)
+        k_li = torch.arange(16, dtype=torch.int8).view(2, 1, 8)
+        k_li_scale = torch.arange(2, dtype=torch.float32).view(2, 1)
+
+        gathered, handle, metadata = AscendSFAImpl._all_gather_byte_packed_async(
+            [
+                ("sfa_kv", sfa_kv),
+                ("k_li", k_li),
+                ("k_li_scale", k_li_scale),
+            ],
+            async_op=True,
+        )
+
+        self.assertIsNone(handle)
+        self.assertEqual(gathered.dtype, torch.int8)
+        restored = AscendSFAImpl._restore_byte_gathered_tensors(gathered, metadata)
+        for name, expected in (("sfa_kv", sfa_kv), ("k_li", k_li), ("k_li_scale", k_li_scale)):
+            self.assertEqual(restored[name].shape, expected.shape)
+            self.assertEqual(restored[name].dtype, expected.dtype)
+            self.assertTrue(torch.equal(restored[name], expected))
+
+    def test_byte_packed_gather_rejects_mismatched_token_counts(self):
+        with self.assertRaisesRegex(RuntimeError, "different token counts"):
+            AscendSFAImpl._all_gather_byte_packed_async(
+                [
+                    ("sfa_kv", torch.zeros(2, 6, dtype=torch.float16)),
+                    ("k_li", torch.zeros(3, 8, dtype=torch.int8)),
+                ],
+                async_op=True,
+            )
+
+
+class TestAscendSFADeviceOperator(TestBase):
+    def _make_common_inputs(self):
+        ql_nope = torch.randn(3, 4, 8)
+        q_pe = torch.randn(3, 4, 2)
+        topk_indices = torch.zeros(3, 1, dtype=torch.int32)
+        attn_metadata = MagicMock()
+        attn_metadata.block_table = torch.zeros(1, 4, dtype=torch.int32)
+        actual_seq_lengths_query = torch.tensor([3], dtype=torch.int32)
+        actual_seq_lengths_key = torch.tensor([3], dtype=torch.int32)
+        impl = MagicMock()
+        impl.scale = 0.125
+        impl.qk_rope_head_dim = 2
+        impl.sfa_qsfa_tile_size = 128
+        return (
+            impl,
+            ql_nope,
+            q_pe,
+            topk_indices,
+            attn_metadata,
+            actual_seq_lengths_query,
+            actual_seq_lengths_key,
+        )
+
+    def test_execute_sparse_flash_attention_returns_lse(self):
+        (
+            impl,
+            ql_nope,
+            q_pe,
+            topk_indices,
+            attn_metadata,
+            actual_seq_lengths_query,
+            actual_seq_lengths_key,
+        ) = self._make_common_inputs()
+        kv_cache = (
+            torch.randn(4, 1, 1, 8),
+            torch.randn(4, 1, 1, 2),
+        )
+        attn_output = torch.randn(3, 4, 8)
+        softmax_max = torch.zeros(1, 3, 4)
+        softmax_sum = torch.full((1, 3, 4), 2.0)
+
+        with patch.object(
+            torch.ops._C_ascend,
+            "npu_sparse_flash_attention",
+            create=True,
+            return_value=(attn_output, softmax_max, softmax_sum),
+        ) as mock_sfa:
+            output, softmax_lse = DeviceOperator.execute_sparse_flash_attention_process(
+                impl,
+                ql_nope,
+                q_pe,
+                kv_cache,
+                topk_indices,
+                attn_metadata,
+                actual_seq_lengths_query,
+                actual_seq_lengths_key,
+                return_lse=True,
+            )
+
+        self.assertIs(output, attn_output)
+        self.assertEqual(softmax_lse.shape, (3, 4, 1))
+        expected_lse = torch.full((3, 4, 1), torch.log(torch.tensor(2.0)).item())
+        self.assertTrue(torch.allclose(softmax_lse, expected_lse))
+        self.assertTrue(mock_sfa.call_args.kwargs["return_softmax_lse"])
+
+    def test_execute_sparse_flash_attention_c8_returns_lse(self):
+        (
+            impl,
+            ql_nope,
+            q_pe,
+            topk_indices,
+            attn_metadata,
+            actual_seq_lengths_query,
+            actual_seq_lengths_key,
+        ) = self._make_common_inputs()
+        packed_kv_cache = (torch.empty(4, 1, 1, 12, dtype=torch.int8),)
+        attn_output = torch.randn(3, 4, 8)
+        softmax_max = torch.ones(1, 3, 4)
+        softmax_sum = torch.full((1, 3, 4), 3.0)
+
+        with (
+            patch.object(
+                torch.ops._C_ascend,
+                "npu_kv_quant_sparse_flash_attention",
+                create=True,
+                return_value=(attn_output, softmax_max, softmax_sum),
+            ) as mock_qsfa,
+            patch(
+                "vllm_ascend.device.device_op.torch_npu.npu_kv_quant_sparse_flash_attention",
+                create=True,
+                side_effect=AssertionError("C8 SFA with LSE must use the custom op"),
+            ),
+        ):
+            output, softmax_lse = DeviceOperator.execute_sparse_flash_attention_process(
+                impl,
+                ql_nope,
+                q_pe,
+                packed_kv_cache,
+                topk_indices,
+                attn_metadata,
+                actual_seq_lengths_query,
+                actual_seq_lengths_key,
+                sparse_mode=0,
+                return_lse=True,
+            )
+
+        self.assertIs(output, attn_output)
+        expected_lse = torch.full((3, 4, 1), 1.0 + torch.log(torch.tensor(3.0)).item())
+        self.assertTrue(torch.allclose(softmax_lse, expected_lse))
+        call_kwargs = mock_qsfa.call_args.kwargs
+        self.assertIs(call_kwargs["key"], packed_kv_cache[0])
+        self.assertIs(call_kwargs["value"], packed_kv_cache[0])
+        self.assertEqual(call_kwargs["query"].shape, (3, 4, 10))
+        self.assertEqual(call_kwargs["sparse_mode"], 0)
+        self.assertTrue(call_kwargs["return_softmax_lse"])
+
+
+class TestAscendSFAKVQuantSparseAttention(TestBase):
+    @patch("vllm_ascend.attention.sfa_v1.torch_npu.npu_dynamic_block_quant")
+    @patch("vllm_ascend.attention.sfa_v1.torch_npu.npu_interleave_rope")
+    @patch("vllm_ascend.attention.sfa_v1.torch_npu.npu_rms_norm")
+    def test_pack_prefill_kv_cache(self, mock_rms_norm, mock_rope, mock_block_quant):
+        k_nope = torch.randn(2, 1, 1, 256, dtype=torch.bfloat16)
+        k_pe = torch.randn(2, 1, 1, 16, dtype=torch.bfloat16)
+        quantized = torch.randint(-128, 127, (2, 1, 256), dtype=torch.int8)
+        scales = torch.arange(1, 5, dtype=torch.float32).view(2, 1, 2)
+        mock_rms_norm.return_value = k_nope, None
+        mock_rope.return_value = k_pe
+        mock_block_quant.return_value = quantized, scales
+
+        actual_k_pe, actual_k_nope, actual_scales = custom_kv_rmsnorm_rope(
+            torch.randn(2, 1, 1, 272, dtype=torch.bfloat16),
+            torch.ones(256, dtype=torch.bfloat16),
+            torch.randn(2, 1, 1, 16),
+            torch.randn(2, 1, 1, 16),
+            256,
+            16,
+            dst_type=1,
+            tile_size=128,
+        )
+        packed_kv = torch.cat([actual_k_nope, actual_k_pe, actual_scales], dim=-1)
+
+        self.assertEqual(mock_block_quant.call_args.kwargs["dst_type"], 1)
+        self.assertEqual(mock_block_quant.call_args.kwargs["row_block_size"], 1)
+        self.assertEqual(mock_block_quant.call_args.kwargs["col_block_size"], 128)
+        self.assertEqual(packed_kv.shape, (2, 1, 1, 296))
+        self.assertTrue(torch.equal(packed_kv[..., :256], quantized.view_as(k_nope)))
+        self.assertTrue(torch.equal(packed_kv[..., 256:288], k_pe.contiguous().view(torch.int8)))
+        self.assertTrue(torch.equal(packed_kv[..., 288:], scales.view(2, 1, 1, 2).view(torch.int8)))
+
+    def test_execute_kv_quant_sparse_flash_attention(self):
+        impl = AscendSFAImpl.__new__(AscendSFAImpl)
+        impl.enable_sparse_sfa_c8 = True
+        impl.scale = 0.125
+        impl.sfa_qsfa_tile_size = 128
+        impl.qk_rope_head_dim = 16
+        ql_nope = torch.randn(3, 2, 32)
+        q_pe = torch.randn(3, 2, 16)
+        kv_cache = (torch.empty(4, 16, 1, 80, dtype=torch.int8),)
+        topk_indices = torch.zeros(3, 1, dtype=torch.int32)
+        attn_metadata = SimpleNamespace(block_table=torch.zeros(1, 4, dtype=torch.int32))
+        actual_seq_lengths = torch.tensor([3], dtype=torch.int32)
+        expected = torch.randn(3, 2, 32)
+
+        with (
+            patch.object(
+                torch.ops._C_ascend,
+                "npu_kv_quant_sparse_flash_attention",
+                create=True,
+                return_value=(expected, torch.empty(0), torch.empty(0)),
+            ) as mock_qsfa,
+            patch(
+                "vllm_ascend.device.device_op.torch_npu.npu_kv_quant_sparse_flash_attention",
+                create=True,
+                side_effect=AssertionError("Base must use _C_ascend custom op"),
+            ),
+        ):
+            result = impl._execute_sparse_flash_attention_process(
+                ql_nope,
+                q_pe,
+                kv_cache,
+                topk_indices,
+                attn_metadata,
+                actual_seq_lengths,
+                actual_seq_lengths,
+            )
+
+        self.assertIs(result, expected)
+        call_kwargs = mock_qsfa.call_args.kwargs
+        self.assertIs(call_kwargs["key"], kv_cache[0])
+        self.assertEqual(call_kwargs["query"].shape, (3, 2, 48))
+        self.assertEqual(call_kwargs["key_quant_mode"], 2)
+        self.assertEqual(call_kwargs["tile_size"], 128)
+        self.assertEqual(call_kwargs["return_softmax_lse"], False)
+
+    def test_prolog_v3_enables_packed_int8_kv_cache(self):
+        impl = AscendSFAImpl.__new__(AscendSFAImpl)
+        impl.enable_sparse_sfa_c8 = True
+        impl.has_indexer = True
+        impl.sfa_qsfa_tile_size = 128
+        impl.sfa_qsfa_k_nope_clip_alpha = torch.ones(1)
+        impl.sfa_qsfa_kr_cache_dummy = torch.empty(0, dtype=torch.bfloat16)
+        impl.local_num_heads = 2
+        impl.kv_lora_rank = 128
+        impl.qk_rope_head_dim = 16
+        impl.q_lora_rank = 8
+        impl.q_a_layernorm = SimpleNamespace(weight=SimpleNamespace(data=torch.ones(8)), variance_epsilon=1e-5)
+        impl.kv_a_layernorm = SimpleNamespace(weight=SimpleNamespace(data=torch.ones(128)), variance_epsilon=1e-5)
+        impl.weight_dq = torch.empty(1)
+        impl.weight_uq_qr = torch.empty(1)
+        impl.W_UK_T = torch.empty(1)
+        impl.weight_dkv_kr = torch.empty(1)
+        impl.dequant_scale_w_dq = torch.empty(1)
+        impl.dequant_scale_w_uq_qr = torch.empty(1)
+        impl.dequant_scale_w_dkv_kr = torch.empty(1)
+        k_cache = torch.empty(4, 16, 1, get_sfa_qsfa_packed_head_dim(128, 16), dtype=torch.int8)
+        dsa_k_cache = torch.empty(4, 16, 1, 128, dtype=torch.bfloat16)
+
+        with (
+            patch(
+                "vllm_ascend.device.device_op.torch_npu.npu_dynamic_quant",
+                return_value=(torch.empty(2, 8, dtype=torch.int8), torch.ones(2, 1)),
+            ),
+            patch(
+                "vllm_ascend.device.device_op.torch_npu.npu_mla_prolog_v3",
+                create=True,
+                return_value=(torch.randn(2, 2, 128), torch.randn(2, 2, 16), None, torch.randn(2, 8), None),
+            ) as mock_prolog,
+        ):
+            impl._sfa_preprocess_with_prolog_v3(
+                hidden_states=torch.randn(2, 8),
+                kv_cache=(k_cache, dsa_k_cache),
+                cos=torch.randn(2, 1, 1, 16),
+                sin=torch.randn(2, 1, 1, 16),
+                slot_mapping=torch.arange(2),
+                cache_mode="PA_BSND",
+            )
+
+        call_kwargs = mock_prolog.call_args.kwargs
+        self.assertIs(call_kwargs["kv_cache"], k_cache)
+        self.assertIs(call_kwargs["kr_cache"], impl.sfa_qsfa_kr_cache_dummy)
+        self.assertEqual(call_kwargs["kv_cache_quant_mode"], 3)
+        self.assertEqual(call_kwargs["ckvkr_repo_mode"], 1)
+        self.assertEqual(call_kwargs["quant_scale_repo_mode"], 1)
+
+
+class TestAscendSFAMetadata(TestBase):
+    def test_ascend_sfa_metadata_default(self):
+        num_actual_tokens = 100
+        slot_mapping = torch.randn(100, 4, 1024)
+        seq_lens = torch.tensor([30, 50])
+        cum_query_lens = torch.tensor([0, 30, 80])
+        block_table = torch.randint(0, 100, (100, 4))
+
+        rope_dim = 32
+        max_seq_len = int(seq_lens.max().item())
+        sin = torch.randn(max_seq_len, rope_dim)
+        cos = torch.randn(max_seq_len, rope_dim)
+
+        num_input_tokens = 2
+        head_dim = None
+        attn_mask = None
+        attn_state = AscendAttentionState.ChunkedPrefill
+
+        metadata = AscendSFAMetadata(
+            num_actual_tokens=num_actual_tokens,
+            slot_mapping=slot_mapping,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens,
+            cum_query_lens=cum_query_lens,
+            block_table=block_table,
+            sin=sin,
+            cos=cos,
+            num_input_tokens=num_input_tokens,
+            head_dim=head_dim,
+            attn_mask=attn_mask,
+            attn_state=attn_state,
+        )
+
+        self.assertEqual(metadata.num_actual_tokens, num_actual_tokens)
+        self.assertIs(metadata.slot_mapping, slot_mapping)
+        self.assertTrue(torch.equal(metadata.seq_lens, seq_lens))
+        self.assertTrue(torch.equal(metadata.cum_query_lens, cum_query_lens))
+        self.assertIs(metadata.block_table, block_table)
+        self.assertIs(metadata.sin, sin)
+        self.assertIs(metadata.cos, cos)
+        self.assertEqual(metadata.num_input_tokens, num_input_tokens)
+        self.assertIs(metadata.head_dim, head_dim)
+        self.assertIs(metadata.attn_mask, attn_mask)
+        self.assertEqual(metadata.attn_state, attn_state)
+
+
+class TestAscendSFAMetadataBuilder(TestBase):
+    @patch("vllm.distributed.parallel_state._TP", new_callable=lambda: MagicMock(spec=GroupCoordinator))
+    def setUp(self, mock_tp):
+        mock_tp.world_size = 2
+        mock_tp.rank_in_group = MagicMock()
+        mock_tp.device_group = MagicMock()
+
+        self.mock_cfg = MagicMock()
+
+        self.mock_cfg.parallel_config = MagicMock()
+        self.mock_cfg.parallel_config.tensor_parallel_size = 1
+        self.mock_cfg.parallel_config.prefill_context_parallel_size = 1
+        self.mock_cfg.parallel_config.decode_context_parallel_size = 1
+
+        self.mock_cfg.compilation_config = MagicMock()
+        self.mock_cfg.compilation_config.pass_config = MagicMock()
+        self.mock_cfg.compilation_config.pass_config.enable_sp = False
+
+        self.mock_cfg.speculative_config.num_speculative_tokens = 0
+
+        self.mock_cfg.additional_config = {"refresh": True}
+        init_ascend_config(self.mock_cfg)
+
+        self.patcher = patch("vllm.config.get_current_vllm_config", return_value=self.mock_cfg)
+        self.patcher.start()
+
+        mock_ascend_config = MagicMock()
+        mock_ascend_config.c8_enable_reshape_optim = False
+        mock_ascend_config.enable_mlapo = True
+        mock_ascend_config.enable_shared_expert_dp = False
+        mock_ascend_config.layer_sharding = None
+        self.ascend_config_patcher = patch(
+            "vllm_ascend.attention.sfa_v1.get_ascend_config",
+            return_value=mock_ascend_config,
+        )
+        self.ascend_config_patcher.start()
+
+        # Mock parent class __init__ to avoid complex initialization,
+        # but still set the essential attributes that child class needs
+        def mock_parent_init(
+            self, kv_cache_spec, layer_names, vllm_config, device, metadata_cls, supports_dcp_with_varlen
+        ):
+            self.metadata_cls = metadata_cls
+            self.kv_cache_spec = kv_cache_spec
+            self.model_config = vllm_config.model_config
+            self.vllm_config = vllm_config
+            self.device = device
+            self.chunked_prefill_workspace_size = 128 * 1024
+            self.chunked_prefill_workspace = torch.empty(
+                (self.chunked_prefill_workspace_size, vllm_config.model_config.get_head_size()),
+                dtype=vllm_config.model_config.dtype,
+                device=device,
+            )
+
+        self.parent_init_patcher = patch(
+            "vllm.model_executor.layers.attention.mla_attention.MLACommonMetadataBuilder.__init__", mock_parent_init
+        )
+        self.parent_init_patcher.start()
+
+        if hasattr(enable_dsa_cp, "cache_clear"):
+            enable_dsa_cp.cache_clear()
+
+    def tearDown(self):
+        self.patcher.stop()
+        self.ascend_config_patcher.stop()
+        self.parent_init_patcher.stop()
+
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_ascend_sfa_metadata_builder_default(self):
+        kv_cache_spec = MagicMock()
+        kv_cache_spec.block_size = 128
+        layer_names = ["layer1", "layer2"]
+        vllm_config = MagicMock()
+        vllm_config.cache_config.block_size = 16
+        vllm_config.scheduler_config.max_num_seqs = 16
+        vllm_config.model_config.max_model_len = 1024
+        vllm_config.model_config.get_head_size.return_value = 64
+        vllm_config.model_config.dtype = torch.float16
+        vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
+        speculative_config = MagicMock()
+        speculative_config.num_speculative_tokens = 4
+        vllm_config.speculative_config = speculative_config
+        device = torch.device("cpu")
+
+        builder = AscendSFAMetadataBuilder(
+            kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
+        )
+
+        assert builder.device == device
+        assert builder.vllm_config == vllm_config
+
+    @patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config")
+    @patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
+    @patch("vllm_ascend.attention.sfa_v1.enable_dsa_cp")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_ascend_sfa_metadata_builder_build(
+        self,
+        mock_enable_dsa_cp,
+        mock_get_cos_and_sin_mla,
+        mock_get_current_vllm_config,
+    ):
+        mock_enable_dsa_cp.return_value = False
+
+        cfg = MagicMock()
+        cfg.model_config = MagicMock()
+        cfg.model_config.hf_text_config = MagicMock()
+
+        mock_get_current_vllm_config.return_value = cfg
+        kv_cache_spec = MagicMock()
+        kv_cache_spec.block_size = 128
+        layer_names = ["layer1", "layer2"]
+        vllm_config = MagicMock()
+        vllm_config.cache_config.block_size = 16
+        vllm_config.scheduler_config.max_num_seqs = 16
+        vllm_config.model_config.max_model_len = 1024
+        vllm_config.model_config.get_head_size.return_value = 64
+        vllm_config.model_config.dtype = torch.float16
+        vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
+        speculative_config = MagicMock()
+        speculative_config.num_speculative_tokens = 4
+        vllm_config.speculative_config = speculative_config
+        device = torch.device("cpu")
+
+        builder = AscendSFAMetadataBuilder(
+            kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
+        )
+
+        common_attn_metadata = MagicMock()
+        common_attn_metadata.num_reqs = 10
+        common_attn_metadata.num_actual_tokens = 100
+        common_attn_metadata.query_start_loc = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
+        common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
+        common_attn_metadata.slot_mapping = torch.randn(100, 4, 1024)
+        common_attn_metadata.seq_lens_cpu = torch.tensor([2] * 10)
+        common_attn_metadata.positions = torch.randn(100)
+        common_attn_metadata.attn_mask = None
+        common_attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
+        common_attn_metadata.block_table_tensor = torch.randn(100, 4)
+        common_attn_metadata.cos = None
+        common_attn_metadata.sin = None
+        common_attn_metadata.num_input_tokens = 100
+
+        mock_get_cos_and_sin_mla.return_value = (torch.randn(100), torch.randn(100))
+
+        metadata = builder.build(
+            common_prefix_len=10,
+            common_attn_metadata=common_attn_metadata,
+        )
+
+        assert isinstance(metadata, AscendSFAMetadata)
+        assert metadata.num_actual_tokens == common_attn_metadata.num_actual_tokens
+        assert metadata.slot_mapping.shape == (100, 4, 1024)
+
+    @patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
+    @patch("vllm_ascend.attention.sfa_v1.get_tp_group")
+    def test_dsa_cp_metadata_builder_masks_graph_padding(
+        self,
+        mock_get_tp_group,
+        mock_get_cos_and_sin_mla,
+    ):
+        # TP8, graph size 80 and MTP3 produce 20 four-token request slots. With
+        # nine real requests, rank 6 splits a padded slot at its local boundary.
+        tp_group = MagicMock()
+        tp_group.world_size = 8
+        tp_group.rank_in_group = 6
+        mock_get_tp_group.return_value = tp_group
+        mock_get_cos_and_sin_mla.return_value = (
+            torch.zeros(80, 1, 1, 64),
+            torch.zeros(80, 1, 1, 64),
+        )
+
+        builder = AscendSFAMetadataBuilder.__new__(AscendSFAMetadataBuilder)
+        builder.kernel_block_size = 128
+        builder.model_config = MagicMock()
+        builder.model_config.get_head_size.return_value = 64
+        builder.attn_mask_builder = MagicMock()
+        builder.enable_dsa_cp = True
+        builder.actual_seq_lengths_query = torch.zeros(21, dtype=torch.int32)
+        builder.actual_seq_lengths_key = torch.zeros(21, dtype=torch.int32)
+        builder.spec_actual_seq_lengths_query = None
+        builder.spec_actual_seq_lengths_key = None
+        builder.metadata_cls = AscendSFAMetadata
+
+        common_attn_metadata = MagicMock()
+        common_attn_metadata.num_reqs = 20
+        common_attn_metadata.num_actual_tokens = 36
+        common_attn_metadata.num_input_tokens = 80
+        common_attn_metadata.query_start_loc = torch.arange(0, 81, 4, dtype=torch.int32)
+        common_attn_metadata.seq_lens = torch.zeros(20, dtype=torch.int32)
+        common_attn_metadata.seq_lens[:9] = torch.arange(128, 137, dtype=torch.int32)
+        common_attn_metadata._seq_lens_cpu = common_attn_metadata.seq_lens.clone()
+        common_attn_metadata.seq_lens_cpu = common_attn_metadata.seq_lens.clone()
+        common_attn_metadata.block_table_tensor = torch.zeros(20, 1, dtype=torch.int32)
+        common_attn_metadata.slot_mapping = torch.arange(80, dtype=torch.int64)
+        common_attn_metadata.positions = torch.arange(80, dtype=torch.int64)
+        common_attn_metadata.attn_state = AscendAttentionState.DecodeOnly
+        common_attn_metadata.causal = True
+
+        metadata = builder._build(common_attn_metadata)
+
+        local_seq_lens = metadata.dsa_cp_context.actual_seq_lengths_key
+        assert local_seq_lens[17].item() == 0
+        assert torch.all(local_seq_lens >= 0)
+
+    @patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config")
+    @patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
+    @patch("vllm_ascend.attention.sfa_v1.enable_dsa_cp", return_value=False)
+    @patch("vllm.distributed.parallel_state.get_tp_group")
+    @patch_distributed_groups(dcp_size=2, pcp_size=2, needs_mocks=False)
+    def test_ascend_sfa_metadata_builder_build_for_graph_capture(
+        self, mock_get_tp_group, mock_enable_dsa_cp, mock_get_cos_and_sin_mla, mock_get_current_vllm_config
+    ):
+        cfg = MagicMock()
+        cfg.model_config = MagicMock()
+        cfg.model_config.hf_text_config = MagicMock()
+
+        mock_get_current_vllm_config.return_value = cfg
+
+        kv_cache_spec = MagicMock()
+        kv_cache_spec.block_size = 128
+        layer_names = ["layer1", "layer2"]
+        vllm_config = MagicMock()
+        vllm_config.cache_config.block_size = 16
+        vllm_config.scheduler_config.max_num_seqs = 16
+        vllm_config.model_config.max_model_len = 1024
+        vllm_config.model_config.get_head_size.return_value = 64
+        vllm_config.model_config.dtype = torch.float16
+        vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
+        speculative_config = MagicMock()
+        speculative_config.num_speculative_tokens = 4
+        vllm_config.speculative_config = speculative_config
+        device = torch.device("cpu")
+
+        builder = AscendSFAMetadataBuilder(
+            kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
+        )
+
+        common_attn_metadata = MagicMock()
+        common_attn_metadata.num_reqs = 10
+        common_attn_metadata.num_actual_tokens = 100
+        common_attn_metadata.query_start_loc = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
+        common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
+        common_attn_metadata.slot_mapping = torch.randn(100, 4, 1024)
+        common_attn_metadata.seq_lens_cpu = torch.tensor([2] * 10)
+        common_attn_metadata.positions = torch.randn(100)
+        common_attn_metadata.attn_mask = None
+        common_attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
+        common_attn_metadata.block_table_tensor = torch.randn(100, 4)
+        common_attn_metadata.cos = None
+        common_attn_metadata.sin = None
+        common_attn_metadata.num_input_tokens = 100
+
+        mock_get_cos_and_sin_mla.return_value = (torch.randn(100), torch.randn(100))
+
+        attn_metadata = builder.build_for_graph_capture(
+            common_attn_metadata=common_attn_metadata,
+            attn_state=AscendAttentionState.DecodeOnly,
+        )
+
+        assert isinstance(attn_metadata, AscendSFAMetadata)
+        assert attn_metadata.attn_state == AscendAttentionState.DecodeOnly
+
+    @patch("vllm_ascend.attention.sfa_v1.get_current_vllm_config")
+    @patch("vllm_ascend.attention.sfa_v1.get_cos_and_sin_mla")
+    @patch("vllm_ascend.attention.sfa_v1.enable_dsa_cp", return_value=False)
+    @patch("torch.ops._C_ascend.store_kv_block_metadata", create=True)
+    def test_ascend_sfa_metadata_builder_build_with_c8_reshape_optim(
+        self,
+        store_kv_block_metadata,
+        mock_enable_dsa_cp,
+        mock_get_cos_and_sin_mla,
+        mock_get_current_vllm_config,
+    ):
+        cfg = MagicMock()
+        cfg.model_config = MagicMock()
+        cfg.model_config.hf_text_config = MagicMock()
+
+        mock_get_current_vllm_config.return_value = cfg
+        kv_cache_spec = MagicMock()
+        kv_cache_spec.block_size = 128
+        layer_names = ["layer1", "layer2"]
+        vllm_config = MagicMock()
+        vllm_config.cache_config.block_size = 16
+        vllm_config.scheduler_config.max_num_seqs = 16
+        vllm_config.model_config.max_model_len = 1024
+        vllm_config.model_config.get_head_size.return_value = 64
+        vllm_config.model_config.dtype = torch.float16
+        vllm_config.model_config.hf_text_config.qk_rope_head_dim = 64
+        speculative_config = MagicMock()
+        speculative_config.num_speculative_tokens = 4
+        vllm_config.speculative_config = speculative_config
+        device = torch.device("cpu")
+
+        builder = AscendSFAMetadataBuilder(
+            kv_cache_spec=kv_cache_spec, layer_names=layer_names, vllm_config=vllm_config, device=device
+        )
+
+        common_attn_metadata = MagicMock()
+        common_attn_metadata.num_reqs = 10
+        common_attn_metadata.num_actual_tokens = 100
+        common_attn_metadata.query_start_loc = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
+        common_attn_metadata.query_start_loc_cpu = torch.tensor([0, 10, 20, 30, 40, 50, 60, 70, 80, 90])
+        common_attn_metadata.slot_mapping = torch.randn(100, 4, 1024)
+        common_attn_metadata.seq_lens_cpu = torch.tensor([2] * 10)
+        common_attn_metadata.positions = torch.randn(100)
+        common_attn_metadata.attn_mask = None
+        common_attn_metadata.attn_state = AscendAttentionState.ChunkedPrefill
+        common_attn_metadata.block_table_tensor = torch.randn(100, 4)
+        common_attn_metadata.cos = None
+        common_attn_metadata.sin = None
+        common_attn_metadata.num_input_tokens = 100
+
+        mock_get_cos_and_sin_mla.return_value = (torch.randn(100), torch.randn(100))
+
+        with patch("vllm_ascend.attention.sfa_v1.get_ascend_config") as mock_get_ascend_config:
+            mock_ascend_config = MagicMock()
+            mock_ascend_config.c8_enable_reshape_optim = True
+            mock_get_ascend_config.return_value = mock_ascend_config
+
+            metadata = builder.build(
+                common_prefix_len=10,
+                common_attn_metadata=common_attn_metadata,
+            )
+
+        assert isinstance(metadata, AscendSFAMetadata)
+        assert metadata.num_actual_tokens == common_attn_metadata.num_actual_tokens
+        assert metadata.slot_mapping.shape == (100, 4, 1024)
+
+        store_kv_block_metadata.assert_called_once()
+        actual_args, _ = store_kv_block_metadata.call_args
+        assert torch.equal(actual_args[0], common_attn_metadata.slot_mapping)
+        assert actual_args[4] == 128
+
+        assert metadata.block_size == 128
+        assert metadata.group_len is actual_args[1]
+        assert metadata.group_key_idx is actual_args[2]
+        assert metadata.group_key_cache_idx is actual_args[3]

tests/ut/attention/a2/test_sfa_v1_precision.py

@@ -0,0 +1,487 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2024 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import math
+import os
+import sys
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm_ascend.utils import enable_custom_op
+
+enable_custom_op()
+
+# Metrics log: survives vLLM stdout/stderr redirection; path via env.
+_METRICS_LOG_PATH = os.environ.get(
+    "SFA_V1_PRECISION_METRICS_LOG",
+    "/tmp/sfa_v1_precision_metrics.log",
+)
+with open(_METRICS_LOG_PATH, "w", encoding="utf-8"):
+    pass
+
+
+def _emit_metric(line: str) -> None:
+    """Append ``line`` to the metrics log and echo it to stderr."""
+    with open(_METRICS_LOG_PATH, "a", encoding="utf-8") as f:
+        f.write(line + "\n")
+    try:
+        sys.stderr.write(line + "\n")
+        sys.stderr.flush()
+    except Exception:
+        pass
+
+
+if "torch_npu._inductor" not in sys.modules:
+    sys.modules["torch_npu._inductor"] = MagicMock()
+
+from vllm.forward_context import set_forward_context  # noqa: E402
+
+from tests.ut.attention.utils import (  # noqa: E402
+    BatchSpec,
+    create_common_attn_metadata,
+    create_vllm_config,
+)
+from vllm_ascend.attention.sfa_v1 import AscendSFAImpl  # noqa: E402
+
+SPARSE_COUNT = 2048  # indexer_select_post_process (sfa_v1)
+
+DEFAULT_RTOL = 1e-2
+DEFAULT_ATOL = 1e-2
+
+# Signal-relative checks (per-element |err|/|ref| near zero ref is unstable).
+_MAX_SIG_REL_ERR = 1e-2  # max |out-ref| / peak |ref|
+_MAX_MEAN_SIG_ERR = 5e-3  # mean |out-ref| / mean |ref|
+_MAX_REL_ERR = 1e-2  # max per-element rel err where |ref| >= floor
+_SIG_FLOOR_FRAC = 5e-1  # floor = this fraction of peak |ref|
+
+_BLOCK_SIZE = 128
+_TEST_NUM_HEADS = 8
+
+BATCH_SPECS: dict[str, BatchSpec] = {
+    "pure_decode_single": BatchSpec(
+        seq_lens=[1024],
+        query_lens=[1],
+        name="pure_decode_single",
+    ),
+    "pure_decode_small_batch": BatchSpec(
+        seq_lens=[64, 128, 256, 512],
+        query_lens=[1, 1, 1, 1],
+        name="pure_decode_small_batch",
+    ),
+    "pure_decode_large_batch": BatchSpec(
+        seq_lens=[2048] * 16,
+        query_lens=[1] * 16,
+        name="pure_decode_large_batch",
+    ),
+    "pure_prefill_single": BatchSpec(
+        seq_lens=[256],
+        query_lens=[256],
+        name="pure_prefill_single",
+    ),
+    "pure_prefill_small_batch": BatchSpec(
+        seq_lens=[128, 256, 384],
+        query_lens=[128, 256, 384],
+        name="pure_prefill_small_batch",
+    ),
+    "pure_prefill_with_context": BatchSpec(
+        seq_lens=[512, 1024],
+        query_lens=[128, 256],
+        name="pure_prefill_with_context",
+    ),
+    "mixed_small": BatchSpec(
+        seq_lens=[64, 128, 256, 512],
+        query_lens=[1, 1, 64, 128],
+        name="mixed_small",
+    ),
+    "mixed_medium": BatchSpec(
+        seq_lens=[1024, 1536, 2048, 256, 512],
+        query_lens=[1, 1, 1, 64, 128],
+        name="mixed_medium",
+    ),
+    "mtp_1_plus_1": BatchSpec(
+        seq_lens=[256, 512, 1024],
+        query_lens=[2, 2, 2],
+        name="mtp_1_plus_1",
+    ),
+    "mtp_1_plus_3": BatchSpec(
+        seq_lens=[256, 512, 1024, 1536],
+        query_lens=[4, 4, 4, 4],
+        name="mtp_1_plus_3",
+    ),
+    "mtp_1_plus_7": BatchSpec(
+        seq_lens=[512, 1024, 2048],
+        query_lens=[8, 8, 8],
+        name="mtp_1_plus_7",
+    ),
+}
+
+
+def _validate_spec(spec: BatchSpec) -> None:
+    """Require seq_len <= SPARSE_COUNT so sparse matches dense reference."""
+    for s, q in zip(spec.seq_lens, spec.query_lens):
+        assert q <= s, f"query_len ({q}) must not exceed seq_len ({s})"
+        assert s <= SPARSE_COUNT, (
+            f"seq_len ({s}) must be <= SPARSE_COUNT ({SPARSE_COUNT}) so the "
+            "sparse attention degenerates into dense attention for the "
+            "reference comparison."
+        )
+
+
+_VLLM_CONFIG_CACHE: dict = {}
+
+
+def _get_vllm_config(
+    model: str,
+    dtype: torch.dtype,
+    *,
+    max_model_len: int = 4096,
+    tensor_parallel_size: int = 1,
+):
+    """Cached ``VllmConfig`` for DSA/SFA (fp8 quant stripped; heads capped for UT)."""
+    key = (model, dtype, tensor_parallel_size)
+    cfg = _VLLM_CONFIG_CACHE.get(key)
+    if cfg is not None:
+        return cfg
+    dtype_str = "bfloat16" if dtype == torch.bfloat16 else "float16"
+    cfg = create_vllm_config(
+        model_name=model,
+        tensor_parallel_size=tensor_parallel_size,
+        max_model_len=max_model_len,
+        dtype=dtype_str,
+        block_size=_BLOCK_SIZE,
+        num_gpu_blocks=4096,
+        max_num_seqs=64,
+        max_num_batched_tokens=max(8192, max_model_len * 2),
+        enable_chunked_prefill=True,
+        hf_overrides={"quantization_config": None},
+        hf_config_override={
+            "num_attention_heads": _TEST_NUM_HEADS,
+            "num_key_value_heads": 1,
+        },
+    )
+    _VLLM_CONFIG_CACHE[key] = cfg
+    return cfg
+
+
+def _build_paged_kv_cache_from_metadata(
+    common_attn_metadata,
+    seq_lens: list[int],
+    k_nope_contexts: list[torch.Tensor],
+    k_rope_contexts: list[torch.Tensor],
+    block_size: int,
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Paged k_nope / k_rope caches from metadata block layout."""
+    blocks_per_seq = [(s + block_size - 1) // block_size for s in seq_lens]
+    total_blocks = sum(blocks_per_seq) + 1
+
+    k_nope_cache = torch.zeros(total_blocks, block_size, 1, kv_lora_rank, dtype=dtype, device=device)
+    k_rope_cache = torch.zeros(total_blocks, block_size, 1, qk_rope_head_dim, dtype=dtype, device=device)
+
+    block_table = common_attn_metadata.block_table_tensor
+    block_table.zero_()
+    next_block_id = 1
+    for b, s_len in enumerate(seq_lens):
+        n_blocks = blocks_per_seq[b]
+        for i in range(n_blocks):
+            block_id = next_block_id
+            block_table[b, i] = block_id
+            tok_start = i * block_size
+            tok_end = min(tok_start + block_size, s_len)
+            length = tok_end - tok_start
+            k_nope_cache[block_id, :length, 0, :] = k_nope_contexts[b][tok_start:tok_end]
+            k_rope_cache[block_id, :length, 0, :] = k_rope_contexts[b][tok_start:tok_end]
+            next_block_id += 1
+
+    return k_nope_cache, k_rope_cache, block_table
+
+
+def _build_topk_indices(
+    seq_lens: list[int],
+    query_lens: list[int],
+    sparse_count: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Causal top-k indices; shape ``(T, 1, sparse_count)`` int32, -1 pad."""
+    num_tokens = sum(query_lens)
+    topk = torch.full((num_tokens, 1, sparse_count), -1, dtype=torch.int32, device=device)
+
+    cum_q = 0
+    for b, s_len in enumerate(seq_lens):
+        q_len = query_lens[b]
+        ctx_len = s_len - q_len
+        for j in range(q_len):
+            valid_end = ctx_len + j + 1
+            topk[cum_q + j, 0, :valid_end] = torch.arange(valid_end, dtype=torch.int32, device=device)
+        cum_q += q_len
+
+    return topk
+
+
+def _reference_sparse_attention(
+    ql_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    k_nope_cache: torch.Tensor,
+    k_rope_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: list[int],
+    query_lens: list[int],
+    scale: float,
+    block_size: int,
+    out_dtype: torch.dtype,
+) -> torch.Tensor:
+    """Fp32 dense MQA softmax baseline over causal prefix."""
+    batch_size = len(seq_lens)
+    outputs: list[torch.Tensor] = []
+
+    cum_q = 0
+    for b in range(batch_size):
+        s_len = seq_lens[b]
+        q_len = query_lens[b]
+        ctx_len = s_len - q_len
+
+        n_blocks = (s_len + block_size - 1) // block_size
+        block_ids = block_table[b, :n_blocks].long()
+        k_blocks = k_nope_cache[block_ids]
+        k_rope_blocks = k_rope_cache[block_ids]
+
+        k_full = k_blocks.reshape(n_blocks * block_size, -1)[:s_len]
+        k_rope_full = k_rope_blocks.reshape(n_blocks * block_size, -1)[:s_len]
+
+        K = torch.cat([k_full, k_rope_full], dim=-1).float()
+        V = k_full.float()
+
+        for j in range(q_len):
+            t = cum_q + j
+            valid_end = ctx_len + j + 1
+
+            q_n = ql_nope[t].float()
+            q_p = q_pe[t].float()
+            Q = torch.cat([q_n, q_p], dim=-1)
+
+            K_b = K[:valid_end]
+            V_b = V[:valid_end]
+
+            scores = (Q @ K_b.transpose(0, 1)) * scale
+            attn = torch.softmax(scores, dim=-1)
+            out = attn @ V_b
+            outputs.append(out.to(out_dtype))
+
+        cum_q += q_len
+
+    return torch.stack(outputs, dim=0)
+
+
+def _run_sfa_kernel(
+    ql_nope: torch.Tensor,
+    q_pe: torch.Tensor,
+    k_nope_cache: torch.Tensor,
+    k_rope_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    topk_indices: torch.Tensor,
+    cum_query_lens: torch.Tensor,
+    seq_lens_tensor: torch.Tensor,
+    scale: float,
+) -> torch.Tensor:
+    """Call kernel via MagicMock self (only ``scale`` needed)."""
+    fake_self = MagicMock()
+    fake_self.scale = scale
+
+    fake_attn_metadata = MagicMock()
+    fake_attn_metadata.block_table = block_table
+
+    return AscendSFAImpl._execute_sparse_flash_attention_process(
+        fake_self,
+        ql_nope,
+        q_pe,
+        (k_nope_cache, k_rope_cache),
+        topk_indices,
+        fake_attn_metadata,
+        cum_query_lens,
+        seq_lens_tensor,
+    )
+
+
+def _run_precision_check(
+    spec: BatchSpec,
+    dtype: torch.dtype,
+    vllm_config,
+    *,
+    tensor_parallel_size: int,
+) -> None:
+    torch.manual_seed(2026)
+    _validate_spec(spec)
+
+    device = torch.device("npu")
+    seq_lens = list(spec.seq_lens)
+    query_lens = list(spec.query_lens)
+    batch_size = spec.batch_size
+    num_tokens = spec.compute_num_tokens()
+
+    cache_config = vllm_config.cache_config
+    hf_text = vllm_config.model_config.hf_text_config
+    block_size = cache_config.block_size
+    qk_rope_head_dim = hf_text.qk_rope_head_dim
+    kv_lora_rank = hf_text.kv_lora_rank
+    num_heads = hf_text.num_attention_heads
+
+    head_dim = kv_lora_rank + qk_rope_head_dim
+    scale = 1.0 / math.sqrt(head_dim)
+
+    common_attn_metadata = create_common_attn_metadata(spec, block_size=block_size, device=device)
+
+    k_nope_contexts = [torch.randn(s, kv_lora_rank, dtype=dtype, device=device) for s in seq_lens]
+    k_rope_contexts = [torch.randn(s, qk_rope_head_dim, dtype=dtype, device=device) for s in seq_lens]
+
+    k_nope_cache, k_rope_cache, block_table = _build_paged_kv_cache_from_metadata(
+        common_attn_metadata=common_attn_metadata,
+        seq_lens=seq_lens,
+        k_nope_contexts=k_nope_contexts,
+        k_rope_contexts=k_rope_contexts,
+        block_size=block_size,
+        kv_lora_rank=kv_lora_rank,
+        qk_rope_head_dim=qk_rope_head_dim,
+        dtype=dtype,
+        device=device,
+    )
+
+    ql_nope = torch.randn(num_tokens, num_heads, kv_lora_rank, dtype=dtype, device=device)
+    q_pe = torch.randn(num_tokens, num_heads, qk_rope_head_dim, dtype=dtype, device=device)
+
+    topk_indices = _build_topk_indices(seq_lens, query_lens, SPARSE_COUNT, device)
+
+    cum_query_lens = torch.tensor(
+        [sum(query_lens[: i + 1]) for i in range(batch_size)],
+        dtype=torch.int32,
+        device=device,
+    )
+    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+
+    with set_forward_context(attn_metadata=None, vllm_config=vllm_config):
+        backend_output = _run_sfa_kernel(
+            ql_nope=ql_nope,
+            q_pe=q_pe,
+            k_nope_cache=k_nope_cache,
+            k_rope_cache=k_rope_cache,
+            block_table=block_table,
+            topk_indices=topk_indices,
+            cum_query_lens=cum_query_lens,
+            seq_lens_tensor=seq_lens_tensor,
+            scale=scale,
+        )
+    reference_output = _reference_sparse_attention(
+        ql_nope=ql_nope,
+        q_pe=q_pe,
+        k_nope_cache=k_nope_cache,
+        k_rope_cache=k_rope_cache,
+        block_table=block_table,
+        seq_lens=seq_lens,
+        query_lens=query_lens,
+        scale=scale,
+        block_size=block_size,
+        out_dtype=dtype,
+    )
+
+    tag = f"{spec.name},tp={tensor_parallel_size}"
+    assert backend_output.shape == reference_output.shape, (
+        f"[{tag}] backend shape {tuple(backend_output.shape)} != reference shape {tuple(reference_output.shape)}"
+    )
+    assert backend_output.dtype == reference_output.dtype, (
+        f"[{tag}] backend dtype {backend_output.dtype} != reference dtype {reference_output.dtype}"
+    )
+    assert torch.isfinite(backend_output).all(), f"[{tag}] sparse flash attention produced non-finite values"
+
+    torch.testing.assert_close(
+        backend_output,
+        reference_output,
+        rtol=DEFAULT_RTOL,
+        atol=DEFAULT_ATOL,
+        msg=lambda m: f"[SFA:{tag}] kernel output diverges from baseline. {m}",
+    )
+
+    ref_f32 = reference_output.float()
+    out_f32 = backend_output.float()
+    diff = (out_f32 - ref_f32).abs()
+    ref_abs = ref_f32.abs()
+    peak = float(ref_abs.max())
+    mean_ref_abs = float(ref_abs.mean())
+    sig_floor = peak * _SIG_FLOOR_FRAC
+
+    max_abs_err = float(diff.max())
+    mean_abs_err = float(diff.mean())
+    max_sig_rel_err = max_abs_err / peak if peak > 0 else 0.0
+    mean_sig_rel_err = mean_abs_err / mean_ref_abs if mean_ref_abs > 0 else 0.0
+
+    significant_mask = ref_abs >= sig_floor
+    if significant_mask.any():
+        per_elem_rel = diff[significant_mask] / ref_abs[significant_mask]
+        max_rel_err_sig = float(per_elem_rel.max())
+    else:
+        max_rel_err_sig = 0.0
+
+    _emit_metric(
+        f"[SFA:{spec.name}] tp={tensor_parallel_size} dtype={dtype} "
+        f"peak={peak:.4e} "
+        f"max_abs_err={max_abs_err:.4e} "
+        f"max_sig_rel_err={max_sig_rel_err * 100:.4f}% "
+        f"mean_sig_rel_err={mean_sig_rel_err * 100:.4f}% "
+        f"max_rel_err_sig(>={int(_SIG_FLOOR_FRAC * 100)}%peak)="
+        f"{max_rel_err_sig * 100:.4f}%"
+    )
+
+    assert max_sig_rel_err < _MAX_SIG_REL_ERR, (
+        f"[SFA:{tag}] dtype={dtype} signal-relative max error "
+        f"{max_sig_rel_err * 100:.4f}% exceeds 1% budget "
+        f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})"
+    )
+    assert mean_sig_rel_err < _MAX_MEAN_SIG_ERR, (
+        f"[SFA:{tag}] dtype={dtype} signal-relative mean error "
+        f"{mean_sig_rel_err * 100:.4f}% exceeds 0.5% drift budget "
+        f"(mean_ref_abs={mean_ref_abs:.4e}, mean_abs_err={mean_abs_err:.4e})"
+    )
+    assert max_rel_err_sig < _MAX_REL_ERR, (
+        f"[SFA:{tag}] dtype={dtype} per-element relative error on "
+        f">={int(_SIG_FLOOR_FRAC * 100)}%-of-peak elements "
+        f"{max_rel_err_sig * 100:.4f}% exceeds 1% budget "
+        f"(peak={peak:.4e}, max_abs_err={max_abs_err:.4e})"
+    )
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V3.2-Exp"])
+@pytest.mark.parametrize("batch_spec_name", list(BATCH_SPECS.keys()))
+@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
+def test_sfa_sparse_flash_attention_precision(
+    batch_spec_name: str,
+    model: str,
+    dtype: torch.dtype,
+    tensor_parallel_size: int,
+) -> None:
+    """SFA kernel vs fp32 dense MQA reference (decode, prefill, mixed, MTP)."""
+    vllm_config = _get_vllm_config(model, dtype, tensor_parallel_size=tensor_parallel_size)
+    _run_precision_check(
+        BATCH_SPECS[batch_spec_name],
+        dtype,
+        vllm_config,
+        tensor_parallel_size=tensor_parallel_size,
+    )

tests/ut/attention/test_attention_fa3.py

@@ -0,0 +1,205 @@
+from importlib import import_module, util
+
+import numpy as np
+import pytest
+import torch
+import torch_npu
+
+
+def _fa3_available() -> bool:
+    try:
+        if util.find_spec("flash_attn_npu_v3") is None:
+            return False
+        mod = import_module("flash_attn_npu_v3")
+        return hasattr(mod, "flash_attn_with_kvcache")
+    except ImportError:
+        return False
+
+
+def ref_fused_infer_attention(
+    query,
+    key,
+    value,
+    block_table,
+    block_size,
+    actual_seq_lengths_q,
+    actual_seq_lengths_kv,
+    num_heads,
+    num_kv_heads,
+    head_size,
+    scale,
+    attn_mask,
+    causal,
+):
+    if not causal:
+        attn_output, _ = torch_npu.npu_fused_infer_attention_score(
+            query=query,
+            key=key,
+            value=value,
+            block_table=block_table,
+            input_layout="TND",
+            block_size=block_size,
+            actual_seq_lengths=actual_seq_lengths_q,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
+            num_key_value_heads=num_kv_heads,
+            num_heads=num_heads,
+            scale=scale,
+            sparse_mode=0,
+        )
+    else:
+        attn_output, _ = torch_npu.npu_fused_infer_attention_score(
+            query=query,
+            key=key,
+            value=value,
+            atten_mask=attn_mask,
+            block_table=block_table,
+            input_layout="TND",
+            block_size=block_size,
+            actual_seq_lengths=actual_seq_lengths_q,
+            actual_seq_lengths_kv=actual_seq_lengths_kv,
+            num_key_value_heads=num_kv_heads,
+            num_heads=num_heads,
+            scale=scale,
+            sparse_mode=3,
+        )
+
+    attn_output = attn_output.view(-1, num_heads, head_size)
+    return attn_output
+
+
+test_cases = [
+    # (data_type, batch_size, num_heads, kv_heads, q_seqlen, kv_seqlen, head_size, block_size, is_causal)
+    (torch.bfloat16, 1, 1, 1, 1024, 1024, 128, 128, False),
+    (torch.bfloat16, 5, 4, 1, 1024, 1024, 128, 128, True),
+    (torch.float16, 7, 16, 8, 512, 512, 128, 128, False),
+]
+
+
+@pytest.mark.skipif(not _fa3_available(), reason="flash_attn_npu_v3 is not installed")
+@pytest.mark.parametrize(
+    "data_type, batch_size, num_heads, kv_heads, q_seqlen, kv_seqlen, head_size, block_size, is_causal", test_cases
+)
+def test_fa_custom_ops_tnd(
+    data_type, batch_size, num_heads, kv_heads, q_seqlen, kv_seqlen, head_size, block_size, is_causal
+):
+    q_min_range = -1.0
+    q_max_range = 1.0
+    kv_min_range = -1.0
+    kv_max_range = 1.0
+    block_size = 128
+    num_blocks = 64
+
+    q_sequences = sorted(
+        torch.randint(low=1, high=q_seqlen + 1, size=(batch_size,)).tolist(), reverse=False
+    )  # actual_seq_lengths in fia need in ascending order
+    kv_sequences = [torch.randint(low=q, high=kv_seqlen + 1, size=(1,)).item() for q in q_sequences]
+
+    t_q_sum = sum(q_sequences)
+
+    query = (q_min_range + (q_max_range - q_min_range) * torch.rand(t_q_sum, num_heads, head_size)).to(data_type).npu()
+
+    key_cache = None
+    value_cache = None
+    block_tables = []
+
+    key_cache = (
+        (kv_min_range + (kv_max_range - kv_min_range) * torch.rand(num_blocks, block_size, kv_heads, head_size))
+        .to(data_type)
+        .npu()
+    )
+    value_cache = (
+        (kv_min_range + (kv_max_range - kv_min_range) * torch.rand(num_blocks, block_size, kv_heads, head_size))
+        .to(data_type)
+        .npu()
+    )
+    max_num_blocks_per_seq = (kv_seqlen + block_size - 1) // block_size
+    for i in range(batch_size):
+        block_table = [max_num_blocks_per_seq * i + j for j in range(max_num_blocks_per_seq)]
+        block_tables.append(block_table)
+    block_tables = torch.tensor(block_tables, dtype=torch.int32).npu()
+
+    q_seqlen_list = q_sequences
+    kv_seqlen_list = kv_sequences
+
+    scale = 1.0 / (head_size**0.5)
+    window_size_left = -1
+    window_size_right = -1
+    is_rotary_interleaved = False
+    num_splits = 0
+    kv_seqlen_list = torch.tensor(kv_seqlen_list, dtype=torch.int32).npu()
+    rotary_cos = None
+    rotary_sin = None
+    cache_batch_idx = None
+    leftpad_k = None
+    new_q_seqlen_list = None
+
+    new_q_seqlen_list = [0]
+    pre_seq_sum = 0
+    for i in range(batch_size):
+        pre_seq_sum += q_seqlen_list[i]
+        new_q_seqlen_list.append(pre_seq_sum)
+    new_q_seqlen_list = torch.tensor(new_q_seqlen_list, dtype=torch.int32).npu()
+
+    from flash_attn_npu_v3 import flash_attn_with_kvcache  # type: ignore[import-not-found]
+
+    out_out = flash_attn_with_kvcache(
+        query,
+        key_cache,
+        value_cache,
+        None,
+        None,
+        None,
+        rotary_cos=rotary_cos,
+        rotary_sin=rotary_sin,
+        cache_seqlens=kv_seqlen_list,
+        cache_batch_idx=cache_batch_idx,
+        cache_leftpad=leftpad_k,
+        page_table=block_tables,
+        cu_seqlens_q=new_q_seqlen_list,
+        cu_seqlens_k_new=None,
+        max_seqlen_q=q_seqlen,
+        rotary_seqlens=None,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
+        softmax_scale=None,
+        causal=is_causal,
+        window_size=[window_size_left, window_size_right],
+        attention_chunk=0,
+        softcap=0.0,
+        rotary_interleaved=is_rotary_interleaved,
+        scheduler_metadata=None,
+        num_splits=num_splits,
+        pack_gqa=None,
+        sm_margin=0,
+        return_softmax_lse=False,
+    )
+
+    ref_out = torch.empty((t_q_sum, num_heads, head_size), dtype=data_type)
+
+    if is_causal:
+        attn_mask = torch.triu(torch.ones(2048, 2048), diagonal=1).to(torch.int8).to(query.device)
+    else:
+        attn_mask = None
+
+    q_cumsum = torch.tensor(np.cumsum(q_sequences), dtype=torch.int32, device=query.device)
+
+    ref_out = ref_fused_infer_attention(
+        query,
+        key_cache.permute(0, 2, 1, 3).contiguous(),  # [num_blocks, num_kv_heads, block_size, head_size]
+        value_cache.permute(0, 2, 1, 3).contiguous(),
+        block_tables,
+        block_size,
+        q_cumsum,
+        kv_seqlen_list,
+        num_heads,
+        kv_heads,
+        head_size,
+        scale,
+        attn_mask,
+        is_causal,
+    )
+
+    rtol = 1e-2
+    atol = 1e-2
+    torch.testing.assert_close(out_out.cpu(), ref_out.cpu(), rtol=rtol, atol=atol)

tests/ut/attention/test_attention_mask.py

@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,114 +20,29 @@ from vllm_ascend.attention.attention_mask import AttentionMaskBuilder
 
 
 class TestAttentionMaskBuilder(TestBase):
-
-    def test_init_attention_mask_builder(self):
-        # generate attention_mask_builder with float16
-        attention_mask_builder = AttentionMaskBuilder(max_seq_len=1024,
-                                                      dtype=torch.float16)
-        self.assertEqual(attention_mask_builder._seq_len_cached, 1024)
-        self.assertEqual(attention_mask_builder.attn_mask_cache.dtype,
-                         torch.float16)
-        self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
-                         (1024, 1024))
-        self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
-                         torch.tensor(float("-inf"), dtype=torch.float16))
-
-        # generate attention_mask_builder with bfloat16
-        attention_mask_builder = AttentionMaskBuilder(max_seq_len=2048,
-                                                      dtype=torch.bfloat16)
-        self.assertEqual(attention_mask_builder._seq_len_cached, 2048)
-        self.assertEqual(attention_mask_builder.attn_mask_cache.dtype,
-                         torch.bfloat16)
-        self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
-                         (2048, 2048))
-        self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
-                         torch.tensor(1, dtype=torch.bfloat16))
-
-    def test_get_mask_scale_factor(self):
-        # supported data types
-        self.assertEqual(
-            AttentionMaskBuilder.get_mask_scale_factor(torch.float16), 1)
-        self.assertEqual(
-            AttentionMaskBuilder.get_mask_scale_factor(torch.bfloat16), -10000)
-        # mask_scale_factor now only supports data types: torch.float16 and torch.bfloat16
-        # Otherwise raise ValueError
-        with self.assertRaises(ValueError):
-            AttentionMaskBuilder.get_mask_scale_factor(torch.int8)
-
     def test_get_attn_mask(self):
         # if the len is less than max_seq_len, the attn_mask_cache will not be updated
-        attention_mask_builder = AttentionMaskBuilder(max_seq_len=1024,
-                                                      dtype=torch.float16)
-        attn_mask = attention_mask_builder.get_attn_mask(
-            max_seq_len=512, dtype=torch.float16, device=torch.device("cpu"))
+        attention_mask_builder = AttentionMaskBuilder(torch.device("cpu"))
+        attn_mask = attention_mask_builder.get_attn_mask(max_seq_len=512, dtype=torch.float16)
         self.assertEqual(attn_mask.shape, (512, 512))
-        self.assertEqual(attn_mask[0][-1],
-                         torch.tensor(float("-inf"), dtype=torch.float16))
-        self.assertEqual(attention_mask_builder._seq_len_cached, 1024)
-        self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
-                         (1024, 1024))
-        self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
-                         torch.tensor(float("-inf"), dtype=torch.float16))
+        self.assertEqual(attn_mask[0][-1], torch.tensor(float("-inf"), dtype=torch.float16))
+        self.assertEqual(attention_mask_builder._seq_len_cached, 512)
+        self.assertEqual(attention_mask_builder.attn_mask_cache.shape, (512, 512))
+        self.assertEqual(
+            attention_mask_builder.attn_mask_cache[0][-1], torch.tensor(float("-inf"), dtype=torch.float16)
+        )
 
         # if the len is greater than max_seq_len, the attn_mask_cache will be updated
-        attn_mask = attention_mask_builder.get_attn_mask(
-            max_seq_len=2048, dtype=torch.float16, device=torch.device("cpu"))
+        attn_mask = attention_mask_builder.get_attn_mask(max_seq_len=2048, dtype=torch.float16)
         self.assertEqual(attn_mask.shape, (2048, 2048))
-        self.assertEqual(attn_mask[0][-1],
-                         torch.tensor(float("-inf"), dtype=torch.float16))
+        self.assertEqual(attn_mask[0][-1], torch.tensor(float("-inf"), dtype=torch.float16))
         self.assertEqual(attention_mask_builder._seq_len_cached, 2048)
-        self.assertEqual(attention_mask_builder.attn_mask_cache.shape,
-                         (2048, 2048))
-        self.assertEqual(attention_mask_builder.attn_mask_cache[0][-1],
-                         torch.tensor(float("-inf"), dtype=torch.float16))
+        self.assertEqual(attention_mask_builder.attn_mask_cache.shape, (2048, 2048))
+        self.assertEqual(
+            attention_mask_builder.attn_mask_cache[0][-1], torch.tensor(float("-inf"), dtype=torch.float16)
+        )
 
     def test_get_splitfuse_attn_mask(self):
-        attention_mask_builder = AttentionMaskBuilder(max_seq_len=1024,
-                                                      dtype=torch.float16)
-        attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
-            seq_lens=torch.tensor([10, 20, 100]),
-            position=torch.tensor([7, 8, 9, 18, 19, 99]),
-            dtype=torch.float16,
-            device=torch.device("cpu"),
-        )
-        self.assertEqual(attn_mask.shape, (6, 100))
-        self.assertEqual(attention_mask_builder._seq_len_cached, 1024)
-
-        attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
-            seq_lens=torch.tensor([10, 3000, 2000]),
-            position=torch.tensor([7, 8, 9, 2999, 1999]),
-            dtype=torch.float16,
-            device=torch.device("cpu"),
-        )
-        self.assertEqual(attn_mask.shape, (5, 3000))
-        self.assertEqual(attention_mask_builder._seq_len_cached, 3000)
-
-        # splitfuse_attn_mask now only supports data types: torch.float16 and torch.bfloat16
-        # otherwise raise ValueError
-        with self.assertRaises(ValueError):
-            attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
-                seq_lens=torch.tensor([10, 20, 100]),
-                position=torch.tensor([7, 8, 9, 18, 19, 99]),
-                dtype=torch.int8,
-                device=torch.device("cpu"),
-            )
-
-    def test_mask_value_cleanliness(self):
-        attention_mask_builder = AttentionMaskBuilder(max_seq_len=6,
-                                                      dtype=torch.bfloat16)
-        self.assertEqual(attention_mask_builder.attn_mask_cache[-2][-1],
-                         torch.tensor(1, dtype=torch.bfloat16))
-
-        attn_mask = attention_mask_builder.get_splitfuse_attn_mask(
-            seq_lens=torch.tensor([6]),
-            position=torch.tensor([3, 4, 5]),
-            dtype=torch.bfloat16,
-            device=torch.device("cpu"),
-        )
-        self.assertEqual(
-            attn_mask[-2][-1],
-            torch.tensor(-10000, dtype=torch.bfloat16,
-                         device=attn_mask.device))
-        self.assertEqual(attention_mask_builder.attn_mask_cache[-2][-1],
-                         torch.tensor(1, dtype=torch.bfloat16))
+        attention_mask_builder = AttentionMaskBuilder(torch.device("cpu"))
+        attn_mask = attention_mask_builder.get_splitfuse_attn_mask()
+        self.assertEqual(attn_mask.shape, (2048, 2048))

tests/ut/attention/test_sfa_o_proj_tp.py

@@ -0,0 +1,103 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+import sys
+from unittest.mock import MagicMock
+
+import torch
+
+from tests.ut.base import TestBase
+
+if "torch_npu._inductor" not in sys.modules:
+    sys.modules["torch_npu._inductor"] = MagicMock()
+
+from vllm_ascend.attention.sfa_v1 import AscendSFAImpl
+
+
+class TestAscendSFAOProjTPParams(TestBase):
+    class _OProj(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.randn(4, 3), requires_grad=False)
+            self.aclnn_input_scale = torch.nn.Parameter(torch.randn(3), requires_grad=False)
+            self.weight_scale_second = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
+            self.weight_scale_second.input_dim = 1
+            self.weight_offset_second = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
+            self.weight_offset_second.input_dim = 1
+            self.extra_input_scale = torch.nn.Parameter(torch.randn(4, 2), requires_grad=False)
+            self.extra_input_scale.input_dim = 1
+            self.weight_scale = torch.nn.Parameter(torch.randn(4), requires_grad=False)
+
+    def setUp(self):
+        AscendSFAImpl.o_proj_full_pools.clear()
+
+    def _make_impl(self):
+        impl = AscendSFAImpl.__new__(AscendSFAImpl)
+        impl.tp_size = 2
+        impl.o_proj = self._OProj()
+        impl._is_o_proj_unquantized = lambda: False
+        return impl
+
+    def test_o_proj_tp_params_alias_original_storage(self):
+        impl = self._make_impl()
+        o_proj = impl.o_proj
+
+        impl._init_o_proj_tp_full_params()
+
+        self.assertEqual(impl.o_proj_tp_weight.data_ptr(), o_proj.weight.data_ptr())
+        self.assertEqual(
+            impl.o_proj_tp_aclnn_input_params["aclnn_input_scale"].data_ptr(),
+            o_proj.aclnn_input_scale.data_ptr(),
+        )
+        self.assertEqual(
+            impl.o_proj_tp_input_sharded_quant_params["weight_scale_second"].data_ptr(),
+            o_proj.weight_scale_second.data_ptr(),
+        )
+        self.assertEqual(
+            impl.o_proj_tp_input_sharded_quant_params["weight_offset_second"].data_ptr(),
+            o_proj.weight_offset_second.data_ptr(),
+        )
+        self.assertEqual(
+            impl.o_proj_tp_input_sharded_quant_params["extra_input_scale"].data_ptr(),
+            o_proj.extra_input_scale.data_ptr(),
+        )
+        self.assertNotIn("weight_scale", impl.o_proj_tp_input_sharded_quant_params)
+
+    def test_o_proj_full_weight_forward_restores_tp_storage(self):
+        impl = self._make_impl()
+        impl._init_o_proj_tp_full_params()
+        original_weight_ptr = impl.o_proj.weight.data_ptr()
+        original_scale_ptr = impl.o_proj.weight_scale_second.data_ptr()
+        full_weight_ptr = impl.o_proj_full_pool.data_ptr()
+        full_scale_ptr = impl.o_proj_full_input_sharded_quant_params["weight_scale_second"].data_ptr()
+
+        def _apply_with_full_weight(_attn_output):
+            self.assertEqual(impl.o_proj.weight.data_ptr(), full_weight_ptr)
+            self.assertEqual(impl.o_proj.weight_scale_second.data_ptr(), full_scale_ptr)
+            return torch.ones(2, 4)
+
+        impl._apply_o_proj_full_weight = MagicMock(side_effect=_apply_with_full_weight)
+
+        output, require_o_proj_forward = impl._handle_o_proj_weight_switch_and_forward(
+            attn_output=torch.randn(2, 3),
+            output=torch.empty(2, 4),
+            o_proj_full_handle=None,
+            o_proj_full_param_handles=[],
+            should_shard_weight=True,
+        )
+
+        self.assertEqual(impl.o_proj.weight.data_ptr(), original_weight_ptr)
+        self.assertEqual(impl.o_proj.weight_scale_second.data_ptr(), original_scale_ptr)
+        self.assertFalse(require_o_proj_forward)
+        self.assertTrue(torch.equal(output, torch.ones(2, 4)))

tests/ut/attention/utils.py

@@ -0,0 +1,341 @@
+from dataclasses import dataclass
+from functools import wraps
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.config import (
+    CacheConfig,
+    CompilationConfig,
+    DeviceConfig,
+    LoadConfig,
+    ModelConfig,
+    ParallelConfig,
+    SchedulerConfig,
+    VllmConfig,
+)
+from vllm.config.model import ModelDType
+from vllm.distributed.parallel_state import all_gather_fake
+from vllm.utils.math_utils import cdiv
+from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
+
+
+def patch_distributed_groups(dcp_size=1, dcp_rank=0, pcp_size=1, pcp_rank=0, needs_mocks=True):
+    """
+    Decorator to patch common distributed group mocks with configuration
+
+    Args:
+        dcp_size: DCP world size (default: 1)
+        dcp_rank: DCP rank (default: 0)
+        pcp_size: PCP world size (default: 1)
+        pcp_rank: PCP rank (default: 0)
+        needs_mocks: Whether to pass mock objects as the first arguments
+             after 'self' to the decorated function.
+             If True, the decorated function receives:
+                 func(self, mock_all_to_all_single, mock_dcp, mock_pcp, *args, **kwargs)
+             If False, mocks are not passed and function receives:
+                 func(self, *args, **kwargs)
+             (default: True)
+    """
+
+    def decorator(func):
+        @wraps(func)
+        @patch("torch.distributed.all_to_all_single")
+        @patch("vllm.distributed.parallel_state._PCP")
+        @patch("vllm.distributed.parallel_state._DCP")
+        def wrapper(self, mock_dcp, mock_pcp, mock_all_to_all_single, *args, **kwargs):
+            mock_dcp.rank_in_group = dcp_rank
+            mock_dcp.world_size = dcp_size
+            mock_dcp.device_group = MagicMock()
+
+            mock_dcp.all_gather = MagicMock()
+            mock_dcp.all_gather.side_effect = lambda input_, dim: all_gather_fake(
+                input_, dim, mock_dcp.world_size, "mock_dcp_group"
+            )
+
+            mock_pcp.rank_in_group = pcp_rank
+            mock_pcp.world_size = pcp_size
+            mock_pcp.device_group = MagicMock()
+
+            mock_pcp.all_gather = MagicMock()
+            mock_pcp.all_gather.side_effect = lambda input_, dim: all_gather_fake(
+                input_, dim, mock_pcp.world_size, "mock_pcp_group"
+            )
+
+            mock_all_to_all_single.side_effect = lambda output, input, *a, **kw: output.copy_(input)
+
+            if needs_mocks:
+                return func(self, mock_all_to_all_single, mock_dcp, mock_pcp, *args, **kwargs)
+            else:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+@dataclass
+class BatchSpec:
+    """Specification for a batch configuration (workload shape only)."""
+
+    seq_lens: list[int]
+    query_lens: list[int]
+
+    name: str = "unnamed"
+
+    @property
+    def batch_size(self):
+        return len(self.seq_lens)
+
+    def __post_init__(self):
+        assert len(self.seq_lens) == len(self.query_lens)
+
+    def compute_num_tokens(self):
+        return sum(self.query_lens)
+
+
+def create_common_attn_metadata(
+    batch_spec: BatchSpec,
+    block_size: int,
+    device: torch.device,
+    max_block_idx: int = 1000,
+    arange_block_indices: bool = True,
+) -> AscendCommonAttentionMetadata:
+    """Create CommonAttentionMetadata from a BatchSpec and ModelParams."""
+    # Create query start locations
+    query_start_loc = torch.zeros(batch_spec.batch_size + 1, dtype=torch.int32, device=device)
+    query_start_loc[1:] = torch.tensor(batch_spec.query_lens, dtype=torch.int32, device=device).cumsum(0)
+    query_start_loc_cpu = query_start_loc.cpu()
+    num_tokens = batch_spec.compute_num_tokens()
+    # Create sequence lengths
+    seq_lens = torch.tensor(batch_spec.seq_lens, dtype=torch.int32, device=device)
+    seq_lens_cpu = seq_lens.cpu()
+    max_seq_len = int(seq_lens_cpu.max())
+
+    # Create computed tokens (context length for each sequence)
+    context_lens = [batch_spec.seq_lens[i] - batch_spec.query_lens[i] for i in range(batch_spec.batch_size)]
+    num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
+
+    # Create block table and slot mapping
+    max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
+    num_blocks = batch_spec.batch_size * max_blocks
+    block_table_tensor = torch.arange(num_blocks, dtype=torch.int32, device=device).view(
+        batch_spec.batch_size, max_blocks
+    )
+    slot_mapping = torch.arange(num_tokens, dtype=torch.int32, device=device).view(num_tokens)
+
+    # Calculate max query length
+    max_query_len = max(batch_spec.query_lens)
+
+    # Create positions tensor
+    positions = torch.arange(num_tokens, dtype=torch.int32, device=device)
+
+    return AscendCommonAttentionMetadata(
+        query_start_loc=query_start_loc,
+        query_start_loc_cpu=query_start_loc_cpu,
+        seq_lens=seq_lens,
+        seq_lens_cpu=seq_lens_cpu,
+        _num_computed_tokens_cpu=num_computed_tokens_cpu,
+        num_reqs=batch_spec.batch_size,
+        num_actual_tokens=num_tokens,
+        max_query_len=max_query_len,
+        max_seq_len=max_seq_len,
+        block_table_tensor=block_table_tensor,
+        slot_mapping=slot_mapping,
+        causal=True,
+        positions=positions,
+    )
+
+
+def create_vllm_config(
+    model_name: str = "meta-llama/Meta-Llama-3-8B",
+    tensor_parallel_size: int = 1,
+    max_model_len: int = 1024,
+    dtype: ModelDType | torch.dtype = "auto",
+    num_gpu_blocks: int = 1000,
+    block_size: int = 16,
+    max_num_seqs: int = 256,
+    max_num_batched_tokens: int = 8192,
+    enable_chunked_prefill: bool = True,
+    add_mock_model_methods: bool = True,
+    hf_config_override: dict | None = None,
+    hf_overrides: dict | None = None,
+) -> VllmConfig:
+    """Create a VllmConfig for testing with reasonable defaults.
+
+    ``hf_overrides`` is forwarded to ``ModelConfig.__init__`` (vLLM-native
+    mechanism); this is required for fp8-quantized DSA models such as
+    ``deepseek-ai/DeepSeek-V3.2-Exp`` whose ``quantization_config`` would
+    otherwise be rejected by Ascend's ``ModelConfig`` validator.
+    """
+
+    model_config = ModelConfig(
+        model=model_name,
+        tokenizer=model_name,
+        trust_remote_code=False,
+        dtype=dtype,
+        seed=0,
+        max_model_len=max_model_len,
+        hf_overrides=hf_overrides or {},
+    )
+
+    cache_config = CacheConfig(
+        block_size=block_size,
+        cache_dtype="auto",
+    )
+    # Set cache blocks for testing
+    #   (these may be set during initialization normally)
+    cache_config.num_gpu_blocks = num_gpu_blocks
+    cache_config.num_cpu_blocks = 0
+
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=tensor_parallel_size,
+    )
+
+    scheduler_config = SchedulerConfig(
+        max_num_seqs=max_num_seqs,
+        max_num_batched_tokens=max_num_batched_tokens,
+        enable_chunked_prefill=enable_chunked_prefill,
+        max_model_len=model_config.max_model_len,
+        is_encoder_decoder=model_config.is_encoder_decoder,
+    )
+
+    device_config = DeviceConfig()
+    load_config = LoadConfig()
+    compilation_config = CompilationConfig()
+
+    if add_mock_model_methods:
+        # Add mock methods to satisfy backends that need them
+        # This is a workaround because tests don't build full, real models,
+        # but some backends expect to query the model for layer-specific
+        # parameters
+        import types
+
+        model_config.get_num_layers = types.MethodType(lambda self: 1, model_config)
+        model_config.get_sliding_window_for_layer = types.MethodType(lambda self, i: None, model_config)
+        model_config.get_logits_soft_cap_for_layer = types.MethodType(lambda self, i: 0.0, model_config)
+        model_config.get_sm_scale_for_layer = types.MethodType(
+            lambda self, i: 1.0 / model_config.get_head_size() ** 0.5, model_config
+        )
+
+    if hf_config_override:
+        # Apply the override to BOTH ``hf_config`` and ``hf_text_config`` so
+        # the attribute is visible to backends regardless of which view they
+        # reach for. For text-only models these usually point to the same
+        # object; for multimodal models ``hf_text_config`` may be different.
+        for k, v in hf_config_override.items():
+            setattr(model_config.hf_config, k, v)
+            if model_config.hf_text_config is not model_config.hf_config:
+                setattr(model_config.hf_text_config, k, v)
+
+    return VllmConfig(
+        model_config=model_config,
+        cache_config=cache_config,
+        parallel_config=parallel_config,
+        scheduler_config=scheduler_config,
+        device_config=device_config,
+        load_config=load_config,
+        compilation_config=compilation_config,
+    )
+
+
+def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
+    """Create a FullAttentionSpec from ModelParams only."""
+    return FullAttentionSpec(
+        block_size=vllm_config.cache_config.block_size,
+        num_kv_heads=vllm_config.model_config.get_num_kv_heads(vllm_config.parallel_config),
+        head_size=vllm_config.model_config.get_head_size(),
+        dtype=vllm_config.model_config.dtype,
+        sliding_window=vllm_config.model_config.get_sliding_window(),
+    )
+
+
+def create_and_prepopulate_kv_cache(
+    k_contexts: list[torch.Tensor],
+    v_contexts: list[torch.Tensor],
+    block_size: int,
+    num_kv_heads: int,
+    head_size: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    num_blocks: int,
+    common_attn_metadata: AscendCommonAttentionMetadata,
+    randomize_blocks: bool = True,
+) -> torch.Tensor:
+    """Create and prepopulate a KV cache with context data.
+
+    Args:
+        k_contexts: List of key context tensors for each sequence
+        v_contexts: List of value context tensors for each sequence
+        seq_lens: List of sequence lengths
+        block_size: Size of each block
+        num_kv_heads: Number of KV heads
+        head_size: Size of each head
+        dtype: Data type for the cache
+        device: Device to create the cache on
+        num_blocks: Total number of blocks in the cache
+        block_table: Block table tensor to populate
+        randomize_blocks: Whether to randomly permute blocks
+                          or use sequential order
+
+    Returns:
+        Tuple of (kv_cache, updated_block_table)
+    """
+    batch_size = len(k_contexts)
+    seq_lens = common_attn_metadata.seq_lens.cpu()
+    query_lens = common_attn_metadata.query_start_loc_cpu[1:] - common_attn_metadata.query_start_loc_cpu[:-1]
+    context_lens = seq_lens - query_lens
+    block_table = common_attn_metadata.block_table_tensor
+    slot_mapping = common_attn_metadata.slot_mapping
+    # Create KV cache
+    kv_cache = torch.zeros(2, num_blocks, block_size, num_kv_heads, head_size, dtype=dtype, device=device)
+    kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)
+
+    # Populate the cache with the context tokens
+    # Start from block_id=0
+    start_block_idx = 0
+    for i in range(batch_size):
+        k_context, v_context = k_contexts[i], v_contexts[i]
+        start = start_block_idx * block_size
+        end = start + k_context.shape[0]
+        kv_cache_flat[0, start:end, ...] = k_context
+        kv_cache_flat[1, start:end, ...] = v_context
+
+        # Stay block aligned and allocate enough blocks for the new tokens
+        start_block_idx += cdiv(int(seq_lens[i]), block_size)
+
+    blocks_end = start_block_idx
+
+    # Permute the context blocks
+    if randomize_blocks:
+        # Random permutation starting from block 0
+        perm = torch.randperm(blocks_end)
+    else:
+        # Sequential order starting from block 0
+        perm = torch.arange(blocks_end)
+    inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
+    inv_perm = torch.argsort(perm)
+    kv_cache[:, :blocks_end, ...] = kv_cache[:, perm, ...]
+    # Construct the right block table
+    # Start from block_id=0
+    start_block_idx = 0
+    for i in range(batch_size):
+        num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
+        start = start_block_idx
+        end = start + num_blocks_for_seq
+        block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
+        start_block_idx += num_blocks_for_seq
+
+        # Create a realistic slot mapping that corresponds to the block table
+    for i in range(batch_size):
+        token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
+        block_indices = token_offsets // block_size
+        token_inter_block_offsets = token_offsets % block_size
+        start = common_attn_metadata.query_start_loc_cpu[i]
+        end = common_attn_metadata.query_start_loc_cpu[i + 1]
+        slot_mapping[start:end] = block_table[i, block_indices] * block_size + token_inter_block_offsets.to(device).to(
+            torch.int32
+        )
+
+    return kv_cache

tests/ut/base.py

@@ -21,14 +21,13 @@ from vllm_ascend.utils import adapt_patch, register_ascend_customop
 
 
 class TestBase(unittest.TestCase):
-
     def __init__(self, *args, **kwargs):
         # adapt patch by default.
         adapt_patch(True)
         adapt_patch()
         register_ascend_customop()
         super().setUp()
-        super(TestBase, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
 
 
 class PytestBase:

tests/ut/compilation/a2/test_acl_graph.py

@@ -0,0 +1,985 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+import weakref
+from unittest.mock import MagicMock, Mock, patch
+
+import numpy as np
+import torch
+from vllm.compilation.cuda_graph import CUDAGraphOptions
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.forward_context import BatchDescriptor, ForwardContext
+
+from tests.ut.base import TestBase
+from vllm_ascend.attention.attention_v1 import AscendMetadata, AscendMetadataForDecode
+from vllm_ascend.attention.context_parallel.attention_cp import AscendAttentionCPImpl
+from vllm_ascend.attention.context_parallel.mla_cp import AscendMlaCPImpl
+from vllm_ascend.attention.mla_v1 import AscendMLADecodeMetadata, AscendMLAMetadata
+from vllm_ascend.compilation import acl_graph
+from vllm_ascend.compilation.acl_graph import (
+    ACLGraphEntry,
+    ACLGraphWrapper,
+    GraphParams,
+    get_draft_graph_params,
+    get_graph_params,
+    set_draft_graph_params,
+    set_graph_params,
+    update_draft_graph_params_workspaces,
+)
+from vllm_ascend.device_allocator.sleep_mem_optimized import AclGraphSleepWakeupManager
+
+
+class TestACLGraphEntry(TestBase):
+    def test_aclgraph_entry_initialization(self):
+        """Test ACLGraphEntry initialization with default values"""
+        batch_descriptor = BatchDescriptor(
+            num_tokens=30,
+            uniform=False,
+        )
+
+        entry = ACLGraphEntry(batch_descriptor=batch_descriptor)
+
+        self.assertEqual(entry.batch_descriptor, batch_descriptor)
+        self.assertIsNone(entry.aclgraph)
+        self.assertIsNone(entry.output)
+        self.assertIsNone(entry.input_addresses)
+
+    def test_aclgraph_entry_with_values(self):
+        """Test ACLGraphEntry initialization with specified values"""
+        batch_descriptor = BatchDescriptor(
+            num_tokens=30,
+            uniform=False,
+        )
+
+        mock_graph = MagicMock()
+        mock_output = MagicMock()
+        input_addresses = [12345, 67890]
+
+        entry = ACLGraphEntry(
+            batch_descriptor=batch_descriptor, aclgraph=mock_graph, output=mock_output, input_addresses=input_addresses
+        )
+
+        self.assertEqual(entry.batch_descriptor, batch_descriptor)
+        self.assertEqual(entry.aclgraph, mock_graph)
+        self.assertEqual(entry.output, mock_output)
+        self.assertEqual(entry.input_addresses, input_addresses)
+
+
+class TestACLGraphWrapper(TestBase):
+    def setUp(self):
+        """Set up test fixtures"""
+        super().setUp()
+
+        # Mock VllmConfig
+        self.mock_vllm_config = MagicMock(spec=VllmConfig)
+        self.mock_vllm_config.compilation_config = MagicMock()
+
+        # Mock runnable function
+        self.mock_runnable = MagicMock(return_value="test_output")
+
+        # Mock graph pool
+        self.mock_graph_pool = MagicMock()
+
+        # Mock CUDAGraphOptions
+        self.mock_cudagraph_options = MagicMock(spec=CUDAGraphOptions)
+        self.mock_cudagraph_options.debug_log_enable = False
+        self.mock_cudagraph_options.gc_disable = False
+        self.mock_cudagraph_options.weak_ref_output = False
+
+        # Mock BatchDescriptor
+        self.mock_batch_descriptor = BatchDescriptor(
+            num_tokens=30,
+            uniform=False,
+        )
+
+        # Mock ForwardContext
+        self.mock_forward_context = MagicMock(spec=ForwardContext)
+        self.mock_forward_context.batch_descriptor = self.mock_batch_descriptor
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    def test_initialization_with_default_options(self, mock_envs, mock_current_platform):
+        """Test ACLGraphWrapper initialization with default CUDAGraphOptions"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable, vllm_config=self.mock_vllm_config, runtime_mode=CUDAGraphMode.FULL
+        )
+
+        self.assertEqual(wrapper.runnable, self.mock_runnable)
+        self.assertEqual(wrapper.vllm_config, self.mock_vllm_config)
+        self.assertEqual(wrapper.graph_pool, self.mock_graph_pool)
+        self.assertEqual(wrapper.runtime_mode, CUDAGraphMode.FULL)
+        self.assertFalse(wrapper.is_debugging_mode)
+        self.assertIsInstance(wrapper.aclgraph_options, CUDAGraphOptions)
+        self.assertEqual(wrapper.concrete_aclgraph_entries, {})
+
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    def test_initialization_with_custom_options(self, mock_envs, mock_current_platform):
+        """Test ACLGraphWrapper initialization with custom CUDAGraphOptions"""
+        mock_envs.VLLM_LOGGING_LEVEL = "DEBUG"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        self.assertEqual(wrapper.runnable, self.mock_runnable)
+        self.assertEqual(wrapper.vllm_config, self.mock_vllm_config)
+        self.assertEqual(wrapper.graph_pool, self.mock_graph_pool)
+        self.assertEqual(wrapper.runtime_mode, CUDAGraphMode.FULL)
+        self.assertTrue(wrapper.is_debugging_mode)
+        self.assertEqual(wrapper.aclgraph_options, self.mock_cudagraph_options)
+        self.assertEqual(wrapper.concrete_aclgraph_entries, {})
+
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    def test_initialization_assertion_error(self, mock_envs, mock_current_platform):
+        """Test ACLGraphWrapper initialization raises AssertionError for NONE mode"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+
+        with self.assertRaises(AssertionError):
+            ACLGraphWrapper(
+                runnable=self.mock_runnable, vllm_config=self.mock_vllm_config, runtime_mode=CUDAGraphMode.NONE
+            )
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    def test_call_with_none_runtime_mode(
+        self, mock_envs, mock_current_platform, mock_get_forward_context, mock_get_forward_context_2
+    ):
+        """Test __call__ method when runtime mode is NONE"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.NONE
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        result = wrapper("arg1", "arg2")
+
+        # Should call the runnable directly without graph capture
+        self.mock_runnable.assert_called_once_with("arg1", "arg2")
+        self.assertEqual(result, "test_output")
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    def test_call_with_mismatched_runtime_mode(
+        self, mock_envs, mock_current_platform, mock_get_forward_context, mock_get_forward_context_2
+    ):
+        """Test __call__ method when runtime mode doesn't match wrapper mode"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE  # Different from FULL
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        result = wrapper("arg1", "arg2")
+
+        # Should call the runnable directly without graph capture
+        self.mock_runnable.assert_called_once_with("arg1", "arg2")
+        self.assertEqual(result, "test_output")
+
+    @patch("vllm_ascend.compilation.acl_graph.torch")
+    @patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.compilation_counter")
+    @patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors")
+    def test_call_capture_graph_first_time(
+        self,
+        mock_weak_ref_tensors,
+        mock_compilation_counter,
+        mock_envs,
+        mock_current_platform,
+        mock_get_forward_context,
+        mock_get_forward_context_2,
+        mock_validate_cudagraph_capturing_enabled,
+        mock_torch,
+    ):
+        """Test __call__ method captures graph for the first time"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+
+        # Mock torch.npu.NPUGraph
+        mock_npu_graph = MagicMock()
+        mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+        # Mock torch.npu.graph context manager
+        mock_graph_context = MagicMock()
+        mock_torch.npu.graph.return_value = mock_graph_context
+        mock_graph_context.__enter__ = Mock(return_value=None)
+        mock_graph_context.__exit__ = Mock(return_value=None)
+
+        # Mock weak_ref_tensors to return the same output
+        mock_weak_ref_tensors.return_value = "weak_ref_output"
+
+        # Ensure torch.Tensor can be correctly identified by isinstance
+        mock_torch.Tensor = torch.Tensor
+
+        # Set up the compilation counter mock
+        mock_compilation_counter.num_cudagraph_captured = 0
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # Create a real torch tensor for the test, not a mock
+        test_tensor = torch.tensor([1, 2, 3])
+
+        # Call the wrapper
+        result = wrapper(test_tensor, "arg2")
+
+        # Verify graph capture happened
+        mock_validate_cudagraph_capturing_enabled.assert_called_once()
+        mock_torch.npu.NPUGraph.assert_called_once()
+        mock_torch.npu.graph.assert_called_once_with(mock_npu_graph, pool=self.mock_graph_pool)
+        self.mock_runnable.assert_called_once_with(test_tensor, "arg2")
+
+        # Verify the entry was created and updated
+        self.assertIn(self.mock_batch_descriptor, wrapper.concrete_aclgraph_entries)
+        entry = wrapper.concrete_aclgraph_entries[self.mock_batch_descriptor]
+        self.assertEqual(entry.aclgraph, mock_npu_graph)
+        self.assertEqual(entry.output, "weak_ref_output")
+
+        # Verify compilation counter was incremented
+        self.assertEqual(mock_compilation_counter.num_cudagraph_captured, 1)
+
+        # Should return the original output (not weak ref)
+        self.assertEqual(result, "test_output")
+
+    @patch("vllm_ascend.compilation.acl_graph.torch")
+    @patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.compilation_counter")
+    @patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors")
+    def test_call_replay_graph(
+        self,
+        mock_weak_ref_tensors,
+        mock_compilation_counter,
+        mock_envs,
+        mock_current_platform,
+        mock_get_forward_context,
+        mock_get_forward_context_2,
+        mock_validate_cudagraph_capturing_enabled,
+        mock_torch,
+    ):
+        """Test __call__ method replays graph when already captured"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+        self.mock_forward_context.is_draft_model = False
+
+        # Mock torch.npu.NPUGraph
+        mock_npu_graph = MagicMock()
+        mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+        # Mock torch.npu.graph context manager
+        mock_graph_context = MagicMock()
+        mock_torch.npu.graph.return_value = mock_graph_context
+        mock_graph_context.__enter__ = Mock(return_value=None)
+        mock_graph_context.__exit__ = Mock(return_value=None)
+
+        # Mock weak_ref_tensors to return the same output
+        mock_weak_ref_tensors.return_value = "weak_ref_output"
+
+        # Ensure torch.Tensor can be correctly identified by isinstance
+        mock_torch.Tensor = torch.Tensor
+
+        # Set up the compilation counter mock
+        mock_compilation_counter.num_cudagraph_captured = 0
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # Create a real torch tensor for the test, not a mock
+        test_tensor = torch.tensor([1, 2, 3])
+
+        # First call to capture the graph
+        first_result = wrapper(test_tensor, "arg2")
+
+        # Verify graph capture happened during first call
+        mock_validate_cudagraph_capturing_enabled.assert_called_once()
+        mock_torch.npu.NPUGraph.assert_called_once()
+        mock_torch.npu.graph.assert_called_once()
+
+        # Reset mock to track second call
+        self.mock_runnable.reset_mock()
+        mock_npu_graph.reset_mock()
+
+        # Second call should replay the graph
+        second_result = wrapper(test_tensor, "arg2")
+
+        # Verify runnable was called only during capture (not during replay)
+        self.mock_runnable.assert_not_called()
+
+        # Verify graph replay happened
+        mock_npu_graph.replay.assert_called_once()
+
+        # Both calls should return the weak ref output
+        self.assertEqual(first_result, "test_output")  # Original output
+        self.assertEqual(second_result, "weak_ref_output")  # Weak ref output
+
+    @patch("vllm_ascend.compilation.acl_graph.torch")
+    @patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors")
+    def test_call_with_debug_mode_input_address_check(
+        self,
+        mock_weak_ref_tensors,
+        mock_envs,
+        mock_current_platform,
+        mock_get_forward_context,
+        mock_get_forward_context_2,
+        mock_validate_cudagraph_capturing_enabled,
+        mock_torch,
+    ):
+        """Test __call__ method with debug mode input address checking"""
+        mock_envs.VLLM_LOGGING_LEVEL = "DEBUG"  # Enable debug mode
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+        self.mock_forward_context.is_draft_model = False
+
+        # Mock torch.npu.NPUGraph
+        mock_npu_graph = MagicMock()
+        mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+        # Mock torch.npu.graph context manager
+        mock_graph_context = MagicMock()
+        mock_torch.npu.graph.return_value = mock_graph_context
+        mock_graph_context.__enter__ = Mock(return_value=None)
+        mock_graph_context.__exit__ = Mock(return_value=None)
+
+        # Mock weak_ref_tensors
+        mock_weak_ref_tensors.return_value = "weak_ref_output"
+
+        # Ensure torch.Tensor can be correctly identified by isinstance
+        mock_torch.Tensor = torch.Tensor
+
+        # Create a mock tensor as the output of runnable
+        mock_output_tensor = torch.tensor([4, 5, 6])
+        self.mock_runnable.return_value = mock_output_tensor
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # First call to capture the graph
+        tensor = torch.tensor([1, 2, 3])  # Create tensor once
+        _ = wrapper(tensor, "arg2")
+
+        # Second call with same tensor addresses should work
+        _ = wrapper(tensor, "arg2")  # Use the same tensor object
+
+        # Should not raise AssertionError
+        self.assertTrue(True)
+
+    @patch("vllm_ascend.compilation.acl_graph.torch")
+    @patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors")
+    def test_call_with_debug_mode_input_address_mismatch(
+        self,
+        mock_weak_ref_tensors,
+        mock_envs,
+        mock_current_platform,
+        mock_get_forward_context,
+        mock_get_forward_context_2,
+        mock_validate_cudagraph_capturing_enabled,
+        mock_torch,
+    ):
+        """Test __call__ method with debug mode input address mismatch raises AssertionError"""
+        mock_envs.VLLM_LOGGING_LEVEL = "DEBUG"  # Enable debug mode
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+
+        # Mock torch.npu.NPUGraph
+        mock_npu_graph = MagicMock()
+        mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+        # Mock torch.npu.graph context manager
+        mock_graph_context = MagicMock()
+        mock_torch.npu.graph.return_value = mock_graph_context
+        mock_graph_context.__enter__ = Mock(return_value=None)
+        mock_graph_context.__exit__ = Mock(return_value=None)
+
+        # Mock weak_ref_tensors
+        mock_weak_ref_tensors.return_value = "weak_ref_output"
+
+        # Ensure torch.Tensor can be correctly identified by isinstance
+        mock_torch.Tensor = torch.Tensor
+
+        # Create a mock tensor as the output of runnable
+        mock_output_tensor = torch.tensor([4, 5, 6])
+        self.mock_runnable.return_value = mock_output_tensor
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # First call to capture the graph
+        tensor1 = torch.tensor([1, 2, 3])
+        _ = wrapper(tensor1, "arg2")
+
+        # Second call with different tensor addresses should raise AssertionError
+        tensor2 = torch.tensor([4, 5, 6])  # Different values, different address
+
+        with self.assertRaises(AssertionError) as context:
+            wrapper(tensor2, "arg2")
+
+        self.assertIn("Input addresses for aclgraphs are different", str(context.exception))
+
+    @patch("vllm_ascend.compilation.acl_graph.torch")
+    @patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.compilation_counter")
+    @patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors")
+    @patch("vllm_ascend.compilation.acl_graph.patch")
+    def test_call_capture_graph_with_gc_disable(
+        self,
+        mock_patch,
+        mock_weak_ref_tensors,
+        mock_compilation_counter,
+        mock_envs,
+        mock_current_platform,
+        mock_get_forward_context,
+        mock_get_forward_context_2,
+        mock_validate_cudagraph_capturing_enabled,
+        mock_torch,
+    ):
+        """Test __call__ method captures graph with gc_disable option enabled"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+
+        # Enable gc_disable option
+        self.mock_cudagraph_options.gc_disable = True
+        # weak_ref_output is not enabled by default
+
+        # Mock torch.npu.NPUGraph
+        mock_npu_graph = MagicMock()
+        mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+        # Mock torch.npu.graph context manager
+        mock_graph_context = MagicMock()
+        mock_torch.npu.graph.return_value = mock_graph_context
+        mock_graph_context.__enter__ = Mock(return_value=None)
+        mock_graph_context.__exit__ = Mock(return_value=None)
+
+        # Mock patch context manager
+        mock_exit_stack = MagicMock()
+        mock_patch.return_value = mock_exit_stack
+        mock_exit_stack.enter_context = Mock()
+
+        # Mock weak_ref_tensors to simulate the actual behavior:
+        # 1. First call (inside the graph context) should return "inner_output"
+        # 2. Second call (for entry.output) should return "weak_ref_output"
+        mock_weak_ref_tensors.side_effect = ["inner_output", "weak_ref_output"]
+
+        # Ensure torch.Tensor can be correctly identified by isinstance
+        mock_torch.Tensor = torch.Tensor
+
+        # Set up the compilation counter mock
+        mock_compilation_counter.num_cudagraph_captured = 0
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # Create a real torch tensor for the test, not a mock
+        test_tensor = torch.tensor([1, 2, 3])
+
+        # Call the wrapper
+        result = wrapper(test_tensor, "arg2")
+
+        # Verify patch was called to disable gc
+        self.assertTrue(mock_patch.called)
+
+        # Verify graph capture happened
+        mock_validate_cudagraph_capturing_enabled.assert_called_once()
+        mock_torch.npu.NPUGraph.assert_called_once()
+        mock_torch.npu.graph.assert_called_once_with(mock_npu_graph, pool=self.mock_graph_pool)
+
+        # Should return the original output (not weak ref) since weak_ref_output is not enabled
+        self.assertEqual(result, "test_output")
+
+    @patch("vllm_ascend.compilation.acl_graph.torch")
+    @patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled")
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.compilation_counter")
+    @patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors")
+    def test_call_capture_graph_with_weak_ref_output(
+        self,
+        mock_weak_ref_tensors,
+        mock_compilation_counter,
+        mock_envs,
+        mock_current_platform,
+        mock_get_forward_context,
+        mock_get_forward_context_2,
+        mock_validate_cudagraph_capturing_enabled,
+        mock_torch,
+    ):
+        """Test __call__ method captures graph with weak_ref_output option enabled"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+
+        # Enable weak_ref_output option
+        self.mock_cudagraph_options.weak_ref_output = True
+
+        # Mock torch.npu.NPUGraph
+        mock_npu_graph = MagicMock()
+        mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+        # Mock torch.npu.graph context manager
+        mock_graph_context = MagicMock()
+        mock_torch.npu.graph.return_value = mock_graph_context
+        mock_graph_context.__enter__ = Mock(return_value=None)
+        mock_graph_context.__exit__ = Mock(return_value=None)
+
+        # Mock weak_ref_tensors to simulate the actual behavior:
+        # 1. First call (inside the graph context with weak_ref_output=True) should return "weak_ref_output"
+        # 2. Second call (for entry.output) should return "weak_ref_output"
+        mock_weak_ref_tensors.side_effect = ["weak_ref_output", "weak_ref_output"]
+
+        # Ensure torch.Tensor can be correctly identified by isinstance
+        mock_torch.Tensor = torch.Tensor
+
+        # Set up the compilation counter mock
+        mock_compilation_counter.num_cudagraph_captured = 0
+
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # Create a real torch tensor for the test, not a mock
+        test_tensor = torch.tensor([1, 2, 3])
+
+        # Call the wrapper
+        result = wrapper(test_tensor, "arg2")
+
+        # Verify weak_ref_tensors was called twice (once for inner output, once for final output)
+        self.assertEqual(mock_weak_ref_tensors.call_count, 2)
+
+        # Verify graph capture happened
+        mock_validate_cudagraph_capturing_enabled.assert_called_once()
+        mock_torch.npu.NPUGraph.assert_called_once()
+        mock_torch.npu.graph.assert_called_once_with(mock_npu_graph, pool=self.mock_graph_pool)
+
+        # Should return the weak ref output when weak_ref_output option is enabled
+        self.assertEqual(result, "weak_ref_output")
+
+    @patch("vllm_ascend.compilation.acl_graph.get_forward_context")
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch("vllm_ascend.compilation.acl_graph.current_platform")
+    @patch("vllm_ascend.compilation.acl_graph.envs")
+    @patch("vllm_ascend.compilation.acl_graph.logger")
+    def test_call_capture_graph_with_debug_log(
+        self, mock_logger, mock_envs, mock_current_platform, mock_get_forward_context, mock_get_forward_context_2
+    ):
+        """Test __call__ method captures graph with debug logging enabled"""
+        mock_envs.VLLM_LOGGING_LEVEL = "INFO"
+        mock_current_platform.get_global_graph_pool.return_value = self.mock_graph_pool
+        mock_get_forward_context.return_value = self.mock_forward_context
+        mock_get_forward_context_2.return_value = self.mock_forward_context
+        self.mock_forward_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
+
+        # Enable debug logging
+        self.mock_cudagraph_options.debug_log_enable = True
+        # weak_ref_output is not enabled by default
+
+        # Mock torch
+        with patch("vllm_ascend.compilation.acl_graph.torch") as mock_torch:
+            # Mock torch.npu.NPUGraph
+            mock_npu_graph = MagicMock()
+            mock_torch.npu.NPUGraph.return_value = mock_npu_graph
+
+            # Mock torch.npu.graph context manager
+            mock_graph_context = MagicMock()
+            mock_torch.npu.graph.return_value = mock_graph_context
+            mock_graph_context.__enter__ = Mock(return_value=None)
+            mock_graph_context.__exit__ = Mock(return_value=None)
+
+            # Ensure torch.Tensor can be correctly identified by isinstance
+            mock_torch.Tensor = torch.Tensor
+
+            # Mock weak_ref_tensors
+            with patch("vllm_ascend.compilation.acl_graph.weak_ref_tensors") as mock_weak_ref_tensors:
+                # Mock weak_ref_tensors to simulate the actual behavior:
+                # 1. First call (inside the graph context) should return "inner_output"
+                # 2. Second call (for entry.output) should return "weak_ref_output"
+                mock_weak_ref_tensors.side_effect = ["inner_output", "weak_ref_output"]
+
+                # Mock validate_cudagraph_capturing_enabled
+                with patch("vllm_ascend.compilation.acl_graph.validate_cudagraph_capturing_enabled"):
+                    wrapper = ACLGraphWrapper(
+                        runnable=self.mock_runnable,
+                        vllm_config=self.mock_vllm_config,
+                        runtime_mode=CUDAGraphMode.FULL,
+                        cudagraph_options=self.mock_cudagraph_options,
+                    )
+
+                    # Create a real torch tensor for the test, not a mock
+                    test_tensor = torch.tensor([1, 2, 3])
+
+                    # Call the wrapper
+                    _ = wrapper(test_tensor, "arg2")
+
+                    # Verify debug log was called
+                    mock_logger.debug.assert_called_once()
+
+    def test_getattr_access_runnable_attributes(self):
+        """Test __getattr__ method accesses runnable attributes"""
+        mock_runnable = MagicMock()
+        mock_runnable.test_attr = "test_value"
+
+        wrapper = ACLGraphWrapper(
+            runnable=mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # Should be able to access attributes of the runnable
+        self.assertEqual(wrapper.test_attr, "test_value")
+
+    def test_getattr_attribute_not_exists(self):
+        """Test __getattr__ method raises AttributeError for non-existent attributes"""
+
+        # Create a simple object without any attributes
+        class EmptyRunnable:
+            pass
+
+        mock_runnable = EmptyRunnable()
+
+        wrapper = ACLGraphWrapper(
+            runnable=mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        # Should raise AttributeError for non-existent attributes
+        with self.assertRaises(AttributeError) as context:
+            _ = wrapper.non_existent_attr
+
+        self.assertIn("Attribute non_existent_attr not found", str(context.exception))
+
+    def test_unwrap_method(self):
+        """Test unwrap method returns the original runnable"""
+        wrapper = ACLGraphWrapper(
+            runnable=self.mock_runnable,
+            vllm_config=self.mock_vllm_config,
+            runtime_mode=CUDAGraphMode.FULL,
+            cudagraph_options=self.mock_cudagraph_options,
+        )
+
+        unwrapped = wrapper.unwrap()
+        self.assertEqual(unwrapped, self.mock_runnable)
+
+    def test_acl_graph_wrappers_use_weak_refs(self):
+        self.assertIsInstance(acl_graph._acl_graph_wrappers, weakref.WeakSet)
+
+
+class TestSleepGraphParams(TestBase):
+    def test_clear_attention_workspaces_preserves_capture_size_keys(self):
+        graph_params = GraphParams(
+            events={4: [], 8: []},
+            workspaces={4: torch.empty(1), 8: torch.empty(2)},
+            handles={4: [], 8: []},
+            attn_params={4: [], 8: []},
+        )
+
+        with (
+            patch("vllm_ascend.compilation.acl_graph._graph_params", graph_params),
+            patch("vllm_ascend.compilation.acl_graph._draft_graph_params", None),
+            patch("vllm_ascend.compilation.acl_graph._draft_graph_prefill_params", None),
+        ):
+            AclGraphSleepWakeupManager.clear_all_attention_workspaces()
+
+        self.assertEqual(set(graph_params.workspaces), {4, 8})
+        self.assertIsNone(graph_params.workspaces[4])
+        self.assertIsNone(graph_params.workspaces[8])
+
+    def test_reset_graph_params_for_sleep_clears_registered_wrappers(self):
+        wrapper = MagicMock()
+        wrapper.concrete_aclgraph_entries = {"entry": object()}
+        wrapper.first_run_finished = True
+        empty_params = GraphParams(
+            events={},
+            workspaces={},
+            handles={},
+            attn_params={},
+        )
+        with (
+            patch("vllm_ascend.compilation.acl_graph._graph_params", empty_params),
+            patch("vllm_ascend.compilation.acl_graph._draft_graph_params", empty_params),
+            patch("vllm_ascend.compilation.acl_graph._draft_graph_prefill_params", empty_params),
+            patch("vllm_ascend.compilation.acl_graph._acl_graph_wrappers", [wrapper]),
+        ):
+            AclGraphSleepWakeupManager.reset_all_graph_params()
+
+        self.assertEqual(wrapper.concrete_aclgraph_entries, {})
+        self.assertFalse(wrapper.first_run_finished)
+
+
+class TestDraftGraphParams(TestBase):
+    def test_set_draft_graph_params(self):
+        with patch("vllm_ascend.compilation.acl_graph._draft_graph_params", new=None):
+            set_draft_graph_params([4])
+            from vllm_ascend.compilation.acl_graph import _draft_graph_params
+
+            self.assertIsNotNone(_draft_graph_params)
+
+    @patch("vllm_ascend.compilation.acl_graph._draft_graph_params")
+    def test_update_draft_graph_params_workspaces(self, draft_graph_params_mock):
+        draft_graph_params_mock.workspaces = {4: 5}
+        update_draft_graph_params_workspaces(4, 6)
+        self.assertEqual(draft_graph_params_mock.workspaces[4], 6)
+
+    @patch("vllm_ascend.compilation.acl_graph._draft_graph_params")
+    def test_get_draft_graph_params(self, draft_graph_params_mock):
+        graph_params = get_draft_graph_params()
+        self.assertIs(draft_graph_params_mock, graph_params)
+
+
+class TestPCPDCPGraphParams(TestBase):
+    def setUp(self):
+        self.update_stream = MagicMock(name="FakeStream")
+        import vllm_ascend.compilation.acl_graph as acl_graph_mod
+
+        self._prev_graph_params = acl_graph_mod._graph_params
+        acl_graph_mod._graph_params = None
+        graph_params = get_graph_params()
+        if graph_params is None:
+            set_graph_params(set([4]))
+            self.graph_params = get_graph_params()
+        else:
+            self.graph_params = graph_params
+        mock_event = MagicMock()
+        mock_event.record = MagicMock()
+        self.graph_params.events[4] = []
+        self.graph_params.handles[4] = []
+        self.graph_params.events[4].append(mock_event)
+        self.graph_params.handles[4].append(MagicMock())
+
+    def tearDown(self):
+        import vllm_ascend.compilation.acl_graph as acl_graph_mod
+
+        acl_graph_mod._graph_params = self._prev_graph_params
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch(
+        "torch.npu.graph_task_update_end",
+    )
+    @patch("torch.npu.graph_task_update_begin", MagicMock())
+    @patch("torch_npu.npu_fused_infer_attention_score.out", MagicMock())
+    def test_update_mla_dcp_pcp_params(self, _mock_graph_task_end, mock_context):
+        input_positions = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8])
+        block_table = torch.zeros(2, 5, dtype=torch.long)
+        seq_lens = torch.tensor([4, 4])
+        cp_seq_len = torch.tensor([2, 2])
+        max_seq_lens = 4
+        seq_lens_list = [4, 4]
+        slot_mapping = torch.zeros(8, dtype=torch.long)
+        query_start_loc = torch.tensor([0, 4])
+        block_tables = torch.zeros(2, 5, dtype=torch.long)
+
+        decode = AscendMLADecodeMetadata(
+            input_positions, block_table, seq_lens, max_seq_lens, seq_lens_list, cp_seq_len=cp_seq_len
+        )
+        metadata = AscendMLAMetadata(
+            8, 8, slot_mapping, query_start_loc, seq_lens, seq_lens, block_tables, 4, 4, 0, decode=decode
+        )
+        forward_context = MagicMock()
+        forward_context.attn_metadata = {"attn_layer_0": metadata}
+        forward_context.is_draft_model = False
+        mock_context.return_value = forward_context
+
+        num_heads = 256
+        scale = 0.1
+        num_kv_heads = 8
+        qk_head_dim = 96
+        qk_rope_head_dim = 32
+        qk_nope_head_dim = 64
+        query = torch.randn(4, num_heads, qk_head_dim)
+
+        q_nope = query[..., :qk_nope_head_dim]
+        q_pe = query[..., qk_rope_head_dim:]
+        k_nope = torch.randn(4, num_heads, qk_nope_head_dim)
+        k_pe = torch.randn(4, num_heads, qk_rope_head_dim)
+        input_layout = "BNSD"
+        actual_seq_lengths_kv = [1, 1]
+        out = torch.randn(2, 16, 128)
+        lse = torch.randn(2, 16, 8)
+        self.graph_params.attn_params[4] = []
+        self.graph_params.attn_params[4].append(
+            (
+                q_nope,
+                k_nope,
+                q_pe,
+                k_pe,
+                num_heads,
+                num_kv_heads,
+                input_layout,
+                None,
+                0,
+                scale,
+                block_table,
+                128,
+                None,
+                actual_seq_lengths_kv,
+                out,
+                lse,
+            )
+        )
+
+        with patch("torch_npu._C._npu_setStream", return_value=None):
+            AscendMlaCPImpl.update_graph_params(self.update_stream, forward_context, 4)
+
+        _mock_graph_task_end.assert_called_once()
+
+    @patch("vllm_ascend.ascend_forward_context.get_forward_context")
+    @patch(
+        "torch.npu.graph_task_update_end",
+    )
+    @patch("torch.npu.graph_task_update_begin", MagicMock())
+    @patch("torch_npu.npu_fused_infer_attention_score.out", MagicMock())
+    def test_update_attn_dcp_pcp_params(self, _mock_graph_task_end, mock_context):
+        block_table = torch.zeros(2, 5, dtype=torch.long)
+        num_heads = 256
+        scale = 0.1
+        num_kv_heads = 8
+        qk_head_dim = 96
+        qk_nope_head_dim = 64
+        query = torch.randn(4, num_heads, qk_head_dim)
+        q_nope = query[..., :qk_nope_head_dim]
+        k_nope = torch.randn(4, num_heads, qk_nope_head_dim)
+        actual_seq_lengths_kv = [1, 1]
+        actual_seq_lengths_q = np.array([1, 1])
+        out = torch.randn(2, 16, 128)
+        lse = torch.randn(2, 16, 8)
+
+        num_computed_tokens_of_pcp_dcp = np.array([[[1, 1], [1, 1]], [[1, 1], [1, 1]]])
+        decode = AscendMetadataForDecode(num_computed_tokens_of_pcp_dcp)
+        metadata = AscendMetadata(
+            num_actual_tokens_pcp_padded=[1, 1],
+            actual_seq_lengths_q=actual_seq_lengths_q,
+            num_decode_tokens=1,
+            decode_meta=decode,
+        )
+        forward_context = MagicMock()
+        forward_context.attn_metadata = {"attn_layer_0": metadata}
+        forward_context.is_draft_model = False
+        mock_context.return_value = forward_context
+
+        self.graph_params.attn_params[4] = []
+        self.graph_params.attn_params[4].append(
+            (
+                q_nope,
+                k_nope,
+                k_nope,
+                num_heads,
+                num_kv_heads,
+                scale,
+                block_table,
+                128,
+                actual_seq_lengths_kv,
+                actual_seq_lengths_q,
+                out,
+                lse,
+                2,
+                0,
+                0,
+                None,
+            )
+        )
+
+        with patch("torch_npu._C._npu_setStream", return_value=None):
+            AscendAttentionCPImpl.update_graph_params(self.update_stream, forward_context, 4, None)
+
+        _mock_graph_task_end.assert_called_once()

tests/ut/compilation/test_npugraph_ex_utils_check.py

@@ -0,0 +1,47 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from vllm_ascend.compilation.passes.utils.npugraph_ex_utils_check import extra_stream_scope_check
+
+
+def test_extra_stream_scope_check_logic():
+    """
+    Test the extra_stream_scope_check logic used in both fusion patterns.
+    This is a pure function test (copied logic for testability).
+    """
+
+    class MockNode:
+        def __init__(self, stream_label=None):
+            self.op = "call_function"
+            self.meta = {"stream_label": stream_label}
+
+    class MockMatch:
+        def __init__(self, nodes):
+            self.nodes = nodes
+
+    # Test 1: all default → OK
+    assert extra_stream_scope_check(MockMatch([MockNode(None), MockNode(None)])) is True
+
+    # Test 2: same non-default → OK
+    assert extra_stream_scope_check(MockMatch([MockNode("s1"), MockNode("s1")])) is True
+
+    # Test 3: mixed non-default → FAIL
+    assert extra_stream_scope_check(MockMatch([MockNode("s1"), MockNode("s2")])) is False
+
+    # Test 4: default + non-default → FAIL
+    assert extra_stream_scope_check(MockMatch([MockNode(None), MockNode("s1")])) is False
+
+    # Test 5: empty → OK
+    assert extra_stream_scope_check(MockMatch([])) is True

tests/ut/conftest.py

@@ -15,12 +15,143 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
+"""Shared UT setup.
 
-from vllm_ascend.utils import adapt_patch  # noqa E402
-from vllm_ascend.utils import register_ascend_customop
+NPU vs CPU routing is determined by directory convention, not decorators.
+See ``.github/workflows/scripts/select_tests.py`` and
+``.github/workflows/scripts/test_config.yaml`` for the routing rules.
+
+Conventions for UT directories:
+    tests/ut/<module>/            -> CPU runner (default)
+    tests/ut/<module>/a2/         -> A2 NPU x1
+    tests/ut/<module>/a2_2/       -> A2 NPU x2
+    tests/ut/<module>/a3_2/       -> A3 NPU x2
+    tests/ut/<module>/a3_4/       -> A3 NPU x4
+    tests/ut/<module>/310p/       -> 310P NPU x1
+"""
+
+import importlib.util
+import subprocess
+import sys
+import types
+from unittest.mock import MagicMock
+
+try:
+    # Note: do not import torch here for cpu env, which will lead to circle import error.
+    subprocess.run(["npu-smi", "info"], capture_output=True, check=True)
+    _npu_available = True
+except (subprocess.CalledProcessError, FileNotFoundError):
+    _npu_available = False
+
+if not _npu_available:
+    triton_runtime = MagicMock()
+    triton_runtime.driver.active.utils.get_device_properties.return_value = {
+        "num_aic": 8,
+        "num_vectorcore": 8,
+    }
+    sys.modules["triton.runtime"] = triton_runtime
+    torch_npu = types.ModuleType("torch_npu")
+    torch_npu.__spec__ = importlib.util.spec_from_loader("torch_npu", loader=None)
+    torch_npu.__path__ = []
+    torch_npu.npu = MagicMock()  # type: ignore[attr-defined]
+    torch_npu.profiler = MagicMock()  # type: ignore[attr-defined]
+    torch_npu.npu_fusion_attention = MagicMock()  # type: ignore[attr-defined]
+    torch_npu.npu_format_cast = MagicMock(side_effect=lambda weight, fmt: weight)  # type: ignore[attr-defined]
+    torch_npu._C = MagicMock()  # type: ignore[attr-defined]
+    torch_npu._C._NPUTaskGroupHandle = MagicMock
+    sys.modules["torch_npu"] = torch_npu
+    sys.modules["torch_npu._C"] = torch_npu._C
+    sys.modules["torch_npu._C._distributed_c10d"] = torch_npu._C._distributed_c10d
+    acl_rt = types.ModuleType("acl.rt")
+    acl_rt.__spec__ = importlib.util.spec_from_loader("acl.rt", loader=None)
+    acl_rt.memcpy = MagicMock()  # type: ignore[attr-defined]
+    acl_mod = types.ModuleType("acl")
+    acl_mod.__spec__ = importlib.util.spec_from_loader("acl", loader=None)
+    acl_mod.rt = acl_rt  # type: ignore[attr-defined]
+    sys.modules["acl"] = acl_mod
+    sys.modules["acl.rt"] = acl_rt
+    mooncake_engine = types.ModuleType("mooncake.engine")
+    mooncake_engine.__spec__ = importlib.util.spec_from_loader("mooncake.engine", loader=None)
+    mooncake_engine.TransferEngine = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["mooncake.engine"] = mooncake_engine
+    import torch
+
+    try:  # noqa: SIM105
+        torch.utils.rename_privateuse1_backend("npu")
+    except RuntimeError:
+        pass
+    torch.npu = MagicMock()
+    torch.npu.Stream = MagicMock
+    torch.version.cann = None
+    torch.distributed.is_hccl_available = MagicMock(return_value=True)
+
+import pytest
+
+mooncake_engine = types.ModuleType("mooncake.engine")
+mooncake_engine.__spec__ = importlib.util.spec_from_loader("mooncake.engine", loader=None)
+mooncake_engine.TransferEngine = MagicMock()  # type: ignore[attr-defined]
+sys.modules.setdefault("mooncake.engine", mooncake_engine)
+
+from vllm_ascend.utils import (  # noqa: E402
+    adapt_patch,
+    clear_enable_sp,
+    register_ascend_customop,
+)
+
+# Mock torch_npu AFTER vllm_ascend import to avoid circular import in accelerate
+if not _npu_available:
+    sys.modules["torch_npu"].npu.current_device = MagicMock(return_value=0)
+    sys.modules["torch_npu._inductor"] = MagicMock()
+    sys.modules["torch_npu"]._npu_flash_attention = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"]._npu_paged_attention_splitfuse = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"]._npu_reshape_and_cache = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"].npu_fused_infer_attention_score = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"]._npu_fused_infer_attention_score_get_max_workspace = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"].npu_moe_gating_top_k_softmax = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"].npu_quant_matmul = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"].npu_rms_norm = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"].npu_swiglu = MagicMock()  # type: ignore[attr-defined]
+    sys.modules["torch_npu"].npu_convert_weight_to_int4pack = MagicMock()  # type: ignore[attr-defined]
 
 adapt_patch()
 adapt_patch(True)
 
 # register Ascend CustomOp here because uts will use this
 register_ascend_customop()
+
+# Clean up any stale mock modules that may have been installed by
+# other test files (e.g., ascend_store/_mock_deps.py) which replace
+# real subpackages with MagicMock, breaking later imports.
+_stale_modules = [
+    k
+    for k in sys.modules
+    if k.startswith("vllm_ascend.distributed.kv_transfer.") and not isinstance(sys.modules[k], types.ModuleType)
+]
+for _m in _stale_modules:
+    del sys.modules[_m]
+
+
+@pytest.fixture(autouse=True)
+def _clear_enable_sp_before_test():
+    clear_enable_sp()
+    yield
+
+
+@pytest.fixture(autouse=True)
+def _mock_ascend_store_deps(request):
+    # ascend_store code imports vllm_ascend helpers (AttentionComputeStartGate,
+    # get/reset_attention_compute_start_gate, ...) which _mock_deps.py no longer
+    # mocks globally (mutating the real modules leaked into other UTs). Mock them
+    # per-test, scoped to the ascend_store tests only.
+    if "distributed/ascend_store/" not in request.node.nodeid:
+        yield
+        return
+    from unittest.mock import patch
+
+    _pfx = "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store"
+    with (
+        patch(f"{_pfx}.pool_worker.get_attention_compute_start_gate"),
+        patch(f"{_pfx}.pool_worker.reset_attention_compute_start_gate"),
+        patch(f"{_pfx}.config_data.AttentionComputeStartGate", type("AttentionComputeStartGate", (), {})),
+    ):
+        yield

tests/ut/core/test_profiling_chunk.py

@@ -0,0 +1,424 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from unittest.mock import MagicMock, patch
+
+import torch
+from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
+from vllm.sampling_params import SamplingParams
+from vllm.utils.hashing import sha256
+from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.request import Request
+from vllm.v1.structured_output import StructuredOutputManager
+
+from tests.ut.base import TestBase
+from vllm_ascend.ascend_config import ProfilingChunkConfig, clear_ascend_config, init_ascend_config
+from vllm_ascend.core.profiling_chunk_predictor import ChunkSizePredictor, ProfilingChunkManager
+from vllm_ascend.core.scheduler_profiling_chunk import ProfilingChunkScheduler
+
+MODEL = "Qwen/Qwen3-0.6B"
+BLOCK_SIZE = 16
+MAX_NUM_BATCHED_TOKENS = 8192
+MAX_NUM_SEQS = 16
+
+
+def create_requests(num_requests, num_tokens=10, max_tokens=16):
+    init_none_hash(sha256)
+    sampling_params = SamplingParams(ignore_eos=False, max_tokens=max_tokens)
+    requests = []
+    for i in range(num_requests):
+        request = Request(
+            request_id=f"{i}",
+            prompt_token_ids=[i] * num_tokens,
+            sampling_params=sampling_params,
+            pooling_params=None,
+            block_hasher=get_request_block_hasher(BLOCK_SIZE, sha256),
+        )
+        requests.append(request)
+    return requests
+
+
+def make_output(scheduler):
+    req_ids = [req.request_id for req in scheduler.running]
+    req_id_to_index = {req.request_id: i for i, req in enumerate(scheduler.running)}
+    sampled_token_ids = [[1000]] * len(scheduler.running)
+    return ModelRunnerOutput(
+        req_ids=req_ids,
+        req_id_to_index=req_id_to_index,
+        sampled_token_ids=sampled_token_ids,
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+
+
+# ===================================================================
+# ProfilingChunkConfig
+# ===================================================================
+
+
+class TestProfilingChunkConfig(TestBase):
+    def test_default_values(self):
+        cfg = ProfilingChunkConfig()
+        self.assertFalse(cfg.enabled)
+        self.assertAlmostEqual(cfg.smooth_factor, 1.0)
+        self.assertEqual(cfg.min_chunk, 4096)
+
+    def test_invalid_smooth_factor_raises(self):
+        with self.assertRaises(ValueError):
+            ProfilingChunkConfig({"smooth_factor": 0.0})
+        with self.assertRaises(ValueError):
+            ProfilingChunkConfig({"smooth_factor": 1.5})
+
+    def test_invalid_min_chunk_raises(self):
+        with self.assertRaises(ValueError):
+            ProfilingChunkConfig({"min_chunk": 0})
+
+    @patch("vllm.config.VllmConfig.__post_init__", MagicMock())
+    @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock())
+    @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config")
+    def test_enabled_without_pp_raises(self, _mock):
+        clear_ascend_config()
+        vllm_config = VllmConfig()
+        vllm_config.model_config = MagicMock()
+        vllm_config.additional_config = {
+            "profiling_chunk_config": {"enabled": True},
+            "refresh": True,
+        }
+        vllm_config.parallel_config.pipeline_parallel_size = 1
+        with self.assertRaises(ValueError) as ctx:
+            init_ascend_config(vllm_config)
+        self.assertIn("pipeline parallelism", str(ctx.exception))
+        clear_ascend_config()
+
+    @patch("vllm.config.VllmConfig.__post_init__", MagicMock())
+    @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock())
+    @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config")
+    def test_enabled_with_pp_ok(self, _mock):
+        clear_ascend_config()
+        vllm_config = VllmConfig()
+        vllm_config.model_config = MagicMock()
+        vllm_config.additional_config = {
+            "profiling_chunk_config": {"enabled": True},
+            "refresh": True,
+        }
+        vllm_config.parallel_config.pipeline_parallel_size = 2
+        ascend_config = init_ascend_config(vllm_config)
+        self.assertTrue(ascend_config.profiling_chunk_config.enabled)
+        clear_ascend_config()
+
+    @patch("vllm.config.VllmConfig.__post_init__", MagicMock())
+    @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock())
+    @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config")
+    def test_disabled_without_pp_ok(self, _mock):
+        clear_ascend_config()
+        vllm_config = VllmConfig()
+        vllm_config.model_config = MagicMock()
+        vllm_config.additional_config = {"refresh": True}
+        ascend_config = init_ascend_config(vllm_config)
+        self.assertFalse(ascend_config.profiling_chunk_config.enabled)
+        clear_ascend_config()
+
+
+# ===================================================================
+# ChunkSizePredictor
+# ===================================================================
+
+
+class TestChunkSizePredictor(TestBase):
+    @staticmethod
+    def _make_data(a, b, c, seq_lens):
+        return [a * seq_len * seq_len + b * seq_len + c for seq_len in seq_lens]
+
+    def test_fit_and_predict(self):
+        predictor = ChunkSizePredictor()
+        seq_lens = list(range(64, 8256, 128))
+        latencies = self._make_data(1e-6, 0.01, 1.0, seq_lens)
+
+        self.assertTrue(predictor.fit(seq_lens, latencies))
+        predictor.set_target_latency(8192)
+        predictor.is_ready = True
+
+        chunk = predictor.predict(num_computed_tokens=0, base_chunk_size=8192, page_size=128)
+        self.assertIsNotNone(chunk)
+        self.assertEqual(chunk % 128, 0)
+
+    def test_predict_decreases_with_history(self):
+        predictor = ChunkSizePredictor()
+        seq_lens = list(range(64, 8256, 128))
+        latencies = self._make_data(1e-6, 0.01, 1.0, seq_lens)
+        predictor.fit(seq_lens, latencies)
+        predictor.set_target_latency(8192)
+        predictor.is_ready = True
+
+        c0 = predictor.predict(0, 8192, 128)
+        c1 = predictor.predict(4096, 8192, 128)
+        c2 = predictor.predict(16384, 8192, 128)
+        self.assertGreaterEqual(c0, c1)
+        self.assertGreaterEqual(c1, c2)
+
+    def test_predict_not_ready_returns_none(self):
+        predictor = ChunkSizePredictor()
+        self.assertIsNone(predictor.predict(0, 8192, 128))
+
+    def test_fit_chunk_and_predict_with_history(self):
+        predictor = ChunkSizePredictor()
+        predictor.is_ready = True
+        predictor.target_latency = 50.0
+
+        data = []
+        for i in range(10):
+            c, h = 1000 + i * 100, i * 500
+            data.append([(c + h) * c, c + h, 1, 1e-9 * (c + h) * c + 0.001 * (c + h) + 0.5])
+        self.assertTrue(predictor.fit_chunk(data))
+        predictor.with_history_ready = True
+
+        result = predictor.predict_with_history(1000, 8192, 128)
+        self.assertIsNotNone(result)
+        self.assertEqual(result % 128, 0)
+
+
+# ===================================================================
+# ProfilingChunkManager
+# ===================================================================
+
+
+class TestProfilingChunkManager(TestBase):
+    def test_not_ready_before_profiling(self):
+        mgr = ProfilingChunkManager(base_chunk_size=8192, page_size=128)
+        self.assertFalse(mgr.is_ready)
+        self.assertIsNone(mgr.predict_chunk_size(0, 1.0))
+
+    def test_run_profiling_success(self):
+        mgr = ProfilingChunkManager(base_chunk_size=8192, page_size=128)
+        seq_lens = list(range(64, 8256, 128))
+        latencies = [1e-6 * seq_len * seq_len + 0.01 * seq_len + 1.0 for seq_len in seq_lens]
+        self.assertTrue(mgr.predictor.fit(seq_lens, latencies))
+        mgr.predictor.set_target_latency(8192)
+        mgr.predictor.is_ready = True
+        mgr._profiling_done = True
+
+        self.assertTrue(mgr.is_ready)
+        self.assertIsNotNone(mgr.predict_chunk_size(0, 1.0))
+
+    def test_run_profiling_all_fail(self):
+        mgr = ProfilingChunkManager(base_chunk_size=8192, page_size=128)
+        too_few_seq_lens = [64, 128, 256]
+        too_few_latencies = [1.0, 2.0, 3.0]
+        self.assertFalse(mgr.predictor.fit(too_few_seq_lens, too_few_latencies))
+        self.assertFalse(mgr.is_ready)
+        self.assertIsNone(mgr.predict_chunk_size(0, 1.0))
+
+    def test_record_batch_refines_model(self):
+        mgr = ProfilingChunkManager(base_chunk_size=8192, page_size=128)
+        seq_lens = list(range(64, 8256, 128))
+        latencies = [1e-6 * seq_len * seq_len + 0.01 * seq_len + 1.0 for seq_len in seq_lens]
+        mgr.predictor.fit(seq_lens, latencies)
+        mgr.predictor.set_target_latency(8192)
+        mgr.predictor.is_ready = True
+        mgr._profiling_done = True
+
+        for i in range(10):
+            mgr.record_batch_execution_time([(4096 - i * 100, i * 500)], 0.05 + i * 0.01)
+        self.assertGreaterEqual(len(mgr.chunked_fit_data), 10)
+        self.assertTrue(mgr.history_ready)
+
+
+# ===================================================================
+# ProfilingChunkScheduler
+# ===================================================================
+
+
+class TestProfilingChunkScheduler(TestBase):
+    @patch("vllm_ascend.ascend_config.AscendConfig.__init__", MagicMock(return_value=None))
+    @patch("vllm_ascend.ascend_config.get_ascend_config")
+    @patch("vllm.config.ModelConfig.__post_init__", MagicMock())
+    @patch("vllm.config.VllmConfig.__post_init__", MagicMock())
+    @patch("vllm.config.device.DeviceConfig.__post_init__", MagicMock())
+    def create_scheduler(self, mock_get_ascend_config):
+        profiling_cfg = MagicMock()
+        profiling_cfg.enabled = True
+        profiling_cfg.smooth_factor = 0.8
+        profiling_cfg.min_chunk = 256
+        mock_get_ascend_config.return_value = MagicMock(profiling_chunk_config=profiling_cfg)
+
+        mock_hf_config = MagicMock()
+        mock_hf_config.model_type = "qwen3"
+        mock_hf_config.is_encoder_decoder = False
+        mock_hf_config.architectures = ["Qwen3ForCausalLM"]
+        model_config = ModelConfig(
+            model=MODEL,
+            tokenizer=MODEL,
+            trust_remote_code=True,
+            dtype="float16",
+            seed=42,
+            max_model_len=MAX_NUM_BATCHED_TOKENS,
+        )
+        model_config.hf_config = mock_hf_config
+        model_config.hf_text_config = MagicMock()
+        model_config.hf_text_config.is_encoder_decoder = False
+        model_config.runner_type = "generate"
+
+        scheduler_config = SchedulerConfig(
+            max_num_seqs=MAX_NUM_SEQS,
+            max_model_len=MAX_NUM_BATCHED_TOKENS,
+            long_prefill_token_threshold=0,
+            disable_chunked_mm_input=False,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS,
+            is_encoder_decoder=False,
+        )
+        scheduler_config.max_num_encoder_input_tokens = 10000
+        scheduler_config.encoder_cache_size = 10000
+        scheduler_config.chunked_prefill_enabled = True
+
+        cache_config = CacheConfig(
+            block_size=BLOCK_SIZE,
+            gpu_memory_utilization=0.9,
+            cache_dtype="auto",
+        )
+
+        vllm_config = VllmConfig(
+            scheduler_config=scheduler_config,
+            model_config=model_config,
+            cache_config=cache_config,
+        )
+        vllm_config.parallel_config.pipeline_parallel_size = 2
+        from unittest.mock import PropertyMock
+
+        type(model_config).is_encoder_decoder = PropertyMock(return_value=False)
+        vllm_config.model_config.hf_config.is_encoder_decoder = False
+
+        kv_cache_config = KVCacheConfig(
+            num_blocks=10000,
+            kv_cache_tensors=[],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    ["layer"],
+                    FullAttentionSpec(block_size=BLOCK_SIZE, num_kv_heads=1, head_size=1, dtype=torch.float32),
+                )
+            ],
+        )
+        kv_cache_config.hash_block_size = BLOCK_SIZE
+        cache_config.num_gpu_blocks = 10000
+
+        scheduler = ProfilingChunkScheduler(
+            vllm_config=vllm_config,
+            kv_cache_config=kv_cache_config,
+            block_size=BLOCK_SIZE,
+            log_stats=True,
+            structured_output_manager=MagicMock(spec=StructuredOutputManager),
+        )
+
+        should_advance = MagicMock()
+        should_advance.return_value = False
+        scheduler.structured_output_manager.should_advance = should_advance
+
+        return scheduler
+
+    def test_scheduler_init(self):
+        scheduler = self.create_scheduler()
+        self.assertIsNotNone(scheduler.profiling_chunk_manager)
+        self.assertFalse(scheduler._profiling_initialized)
+
+    def test_run_profiling_chunk_init_success(self):
+        scheduler = self.create_scheduler()
+        mock_executor = MagicMock()
+        mock_executor.collective_rpc.return_value = [10.0]
+
+        scheduler.run_profiling_chunk_init(mock_executor)
+
+        self.assertTrue(scheduler._profiling_initialized)
+        self.assertTrue(scheduler.profiling_chunk_manager.is_ready)
+
+    def test_run_profiling_chunk_init_skips_second_call(self):
+        scheduler = self.create_scheduler()
+        mock_executor = MagicMock()
+        mock_executor.collective_rpc.return_value = [10.0]
+
+        scheduler.run_profiling_chunk_init(mock_executor)
+        call_count = mock_executor.collective_rpc.call_count
+
+        scheduler.run_profiling_chunk_init(mock_executor)
+        self.assertEqual(mock_executor.collective_rpc.call_count, call_count)
+
+    def test_run_profiling_chunk_init_none_executor(self):
+        scheduler = self.create_scheduler()
+        scheduler.run_profiling_chunk_init(None)
+        self.assertTrue(scheduler._profiling_initialized)
+        self.assertFalse(scheduler.profiling_chunk_manager.is_ready)
+
+    def test_schedule_new_requests(self):
+        scheduler = self.create_scheduler()
+        requests = create_requests(num_requests=5)
+        for req in requests:
+            scheduler.add_request(req)
+
+        output = scheduler.schedule()
+        self.assertEqual(len(output.scheduled_new_reqs), 5)
+        self.assertEqual(len(scheduler.waiting), 0)
+        self.assertEqual(len(scheduler.running), 5)
+
+    def test_schedule_with_profiling_ready(self):
+        """After profiling is ready, schedule() should still work correctly."""
+        scheduler = self.create_scheduler()
+        mock_executor = MagicMock()
+        mock_executor.collective_rpc.return_value = [10.0]
+        scheduler.run_profiling_chunk_init(mock_executor)
+        self.assertTrue(scheduler.profiling_chunk_manager.is_ready)
+
+        requests = create_requests(num_requests=3, num_tokens=100)
+        for req in requests:
+            scheduler.add_request(req)
+
+        output = scheduler.schedule()
+        self.assertGreater(len(output.scheduled_new_reqs), 0)
+        total = sum(output.num_scheduled_tokens.values())
+        self.assertGreater(total, 0)
+
+    def test_schedule_chunked_prefill_running(self):
+        """Running requests with num_computed_tokens > 0 get dynamic chunk."""
+        scheduler = self.create_scheduler()
+        mock_executor = MagicMock()
+        mock_executor.collective_rpc.return_value = [10.0]
+        scheduler.run_profiling_chunk_init(mock_executor)
+
+        requests = create_requests(num_requests=1, num_tokens=2000, max_tokens=16)
+        for req in requests:
+            scheduler.add_request(req)
+
+        output1 = scheduler.schedule()
+        self.assertEqual(len(output1.scheduled_new_reqs), 1)
+
+        model_output = make_output(scheduler)
+        scheduler.update_from_output(output1, model_output)
+
+        output2 = scheduler.schedule()
+        self.assertGreater(output2.total_num_scheduled_tokens, 0)
+
+    def test_update_from_output(self):
+        scheduler = self.create_scheduler()
+        requests = create_requests(num_requests=3)
+        for req in requests:
+            scheduler.add_request(req)
+
+        output = scheduler.schedule()
+        model_output = make_output(scheduler)
+        scheduler.update_from_output(output, model_output)
+
+        self.assertEqual(len(scheduler.running), 3)

tests/ut/core/test_recompute_scheduler.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import MethodType, SimpleNamespace
+from unittest.mock import MagicMock
+
+from vllm.sampling_params import SamplingParams
+from vllm.v1.request import Request
+from vllm.v1.sample.rejection_sampler import PLACEHOLDER_TOKEN_ID
+
+from vllm_ascend.core.recompute_scheduler import RecomputeScheduler
+
+
+def test_pd_consumer_first_step_injects_placeholder_spec_tokens():
+    scheduler = RecomputeScheduler.__new__(RecomputeScheduler)
+    scheduler.requests = {}
+    scheduler.is_kv_producer = False
+    scheduler.is_hybrid_model = False
+    scheduler.is_mtp_kv_consumer = True
+    scheduler.num_spec_tokens = 1
+    scheduler.max_model_len = 1024
+    scheduler.log_stats = False
+    scheduler.connector = None
+
+    enqueued_requests = []
+
+    def enqueue_waiting_request(self, request):
+        enqueued_requests.append(request)
+
+    scheduler._enqueue_waiting_request = MethodType(enqueue_waiting_request, scheduler)
+
+    request = Request(
+        request_id="pd-consumer-first-step",
+        prompt_token_ids=[1, 2, 3, 4],
+        sampling_params=SamplingParams(max_tokens=8),
+        pooling_params=None,
+    )
+
+    scheduler.add_request(request)
+
+    assert enqueued_requests == [request]
+    assert scheduler.requests[request.request_id] is request
+    assert request.spec_token_ids == [PLACEHOLDER_TOKEN_ID]
+    assert request.num_tokens_with_spec == request.num_tokens + 1
+
+
+def test_update_from_output_settles_finished_request_in_flight_tokens():
+    scheduler = RecomputeScheduler.__new__(RecomputeScheduler)
+    request = SimpleNamespace(
+        num_in_flight_tokens=1,
+        is_finished=lambda: True,
+    )
+    scheduler.requests = {"request": request}
+    scheduler.perf_metrics = None
+    scheduler.connector = None
+    scheduler.enable_return_routed_experts = False
+    scheduler.kv_cache_manager = MagicMock()
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.finished_req_ids_dict = {}
+    scheduler.make_stats = MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"request": 1},
+        recomputed_reqs=None,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        routed_experts=None,
+    )
+
+    assert scheduler.update_from_output(scheduler_output, model_runner_output) == {}
+    assert request.num_in_flight_tokens == 0

tests/ut/device/test_device_op.py

@@ -0,0 +1,119 @@
+from unittest import mock
+
+import pytest
+import torch
+
+from vllm_ascend.device.device_op import A5DeviceAdaptor, BaseDeviceAdaptor
+
+
+def test_npu_flash_attention_uses_fusion_attention_for_fp32():
+    query = torch.randn(5, 4, 64, dtype=torch.float32)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    seq_lens_cpu = torch.tensor([2, 3], dtype=torch.int32)
+    expected = torch.randn_like(query)
+
+    with (
+        mock.patch(
+            "vllm_ascend.device.device_op.torch_npu.npu_fusion_attention",
+            return_value=(expected,),
+        ) as mock_fusion_attention,
+        mock.patch(
+            "vllm_ascend.device.device_op.torch_npu._npu_flash_attention_unpad",
+            create=True,
+        ) as mock_flash_attention,
+    ):
+        output = BaseDeviceAdaptor.npu_flash_attention(
+            query=query,
+            key=key,
+            value=value,
+            seq_lens_cpu=seq_lens_cpu,
+            head_num=4,
+            scale_value=0.125,
+            num_kv_heads=4,
+        )
+
+    assert output is expected
+    mock_flash_attention.assert_not_called()
+    mock_fusion_attention.assert_called_once()
+    call_kwargs = mock_fusion_attention.call_args.kwargs
+    assert call_kwargs["query"] is query
+    assert call_kwargs["key"] is key
+    assert call_kwargs["value"] is value
+    assert call_kwargs["actual_seq_qlen"] == [2, 5]
+    assert all(isinstance(seq_len, int) for seq_len in call_kwargs["actual_seq_qlen"])
+    assert call_kwargs["actual_seq_kvlen"] is call_kwargs["actual_seq_qlen"]
+    assert call_kwargs["head_num"] == 4
+    assert call_kwargs["scale"] == 0.125
+    assert call_kwargs["input_layout"] == "TND"
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_npu_flash_attention_uses_unpad_attention_for_low_precision(dtype):
+    query = torch.randn(5, 4, 64, dtype=dtype)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    seq_lens_cpu = torch.tensor([2, 3], dtype=torch.int32)
+
+    def fake_flash_attention(*, query, key, value, seq_len, scale_value, num_heads, num_kv_heads, out):
+        out.copy_(query + 1)
+
+    with (
+        mock.patch(
+            "vllm_ascend.device.device_op.torch_npu.npu_fusion_attention",
+        ) as mock_fusion_attention,
+        mock.patch(
+            "vllm_ascend.device.device_op.torch_npu._npu_flash_attention_unpad",
+            side_effect=fake_flash_attention,
+            create=True,
+        ) as mock_flash_attention,
+    ):
+        output = BaseDeviceAdaptor.npu_flash_attention(
+            query=query,
+            key=key,
+            value=value,
+            seq_lens_cpu=seq_lens_cpu,
+            head_num=4,
+            scale_value=0.125,
+            num_kv_heads=4,
+        )
+
+    mock_fusion_attention.assert_not_called()
+    mock_flash_attention.assert_called_once()
+    call_kwargs = mock_flash_attention.call_args.kwargs
+    assert call_kwargs["query"] is query
+    assert call_kwargs["key"] is key
+    assert call_kwargs["value"] is value
+    assert call_kwargs["seq_len"] is seq_lens_cpu
+    assert call_kwargs["num_heads"] == 4
+    assert call_kwargs["num_kv_heads"] == 4
+    assert call_kwargs["scale_value"] == 0.125
+    torch.testing.assert_close(output, query + 1)
+
+
+def test_a5_npu_flash_attention_uses_python_sequence_lengths():
+    query = torch.randn(5, 4, 64, dtype=torch.float16)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    seq_lens_cpu = torch.tensor([2, 3], dtype=torch.int32)
+    expected = torch.randn_like(query)
+
+    with mock.patch(
+        "vllm_ascend.device.device_op.torch_npu.npu_fusion_attention",
+        return_value=(expected,),
+    ) as mock_fusion_attention:
+        output = A5DeviceAdaptor.npu_flash_attention(
+            query=query,
+            key=key,
+            value=value,
+            seq_lens_cpu=seq_lens_cpu,
+            head_num=4,
+            scale_value=0.125,
+            num_kv_heads=4,
+        )
+
+    assert output is expected
+    call_kwargs = mock_fusion_attention.call_args.kwargs
+    assert call_kwargs["actual_seq_qlen"] == [2, 5]
+    assert all(isinstance(seq_len, int) for seq_len in call_kwargs["actual_seq_qlen"])
+    assert call_kwargs["actual_seq_kvlen"] is call_kwargs["actual_seq_qlen"]

tests/ut/device_allocator/a2/__init__.py

@@ -0,0 +1,14 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#

tests/ut/device_allocator/a2/test_find_loaded_library.py

@@ -0,0 +1,28 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from tests.ut.base import PytestBase
+from vllm_ascend.device_allocator.camem import find_loaded_library
+
+
+class TestFindLoadedLibrary(PytestBase):
+    def test_find_loaded_library_success_and_not_found(self):
+        path = find_loaded_library("libc")
+        assert path is not None, "Expected to find libc library"
+        assert path.endswith(".so.6") or ".so" in path
+        assert "libc" in path
+
+        path = find_loaded_library("non_existent_library")
+        assert path is None, "Expected to not find non-existent library"

tests/ut/device_allocator/test_camem.py

@@ -19,11 +19,13 @@ import pytest
 import torch
 
 from tests.ut.base import PytestBase
-from vllm_ascend.device_allocator.camem import (AllocationData, CaMemAllocator,
-                                                create_and_map,
-                                                find_loaded_library,
-                                                get_pluggable_allocator,
-                                                unmap_and_release)
+from vllm_ascend.device_allocator.camem import (
+    AllocationData,
+    CaMemAllocator,
+    create_and_map,
+    get_pluggable_allocator,
+    unmap_and_release,
+)
 
 
 def dummy_malloc(args):
@@ -35,44 +37,34 @@ def dummy_free(ptr):
 
 
 class TestCaMem(PytestBase):
-
-    def test_find_loaded_library_success_and_not_found(self):
-        path = find_loaded_library("libc")
-        assert path is not None, "Expected to find libc library"
-        assert path.endswith(".so.6") or ".so" in path
-        assert "libc" in path
-
-        path = find_loaded_library("non_existent_library")
-        assert path is None, "Expected to not find non-existent library"
-
-    @pytest.mark.parametrize("handle", [
-        (1, 2, 3),
-        ("device", 99),
-        (None, ),
-    ])
+    @pytest.mark.parametrize(
+        "handle",
+        [
+            (1, 2, 3),
+            ("device", 99),
+            (None,),
+        ],
+    )
     def test_create_and_map_calls_python_create_and_map(self, handle):
-        with patch("vllm_ascend.device_allocator.camem.python_create_and_map"
-                   ) as mock_create:
+        with patch("vllm_ascend.device_allocator.camem.python_create_and_map") as mock_create:
             create_and_map(handle)
             mock_create.assert_called_once_with(*handle)
 
-    @pytest.mark.parametrize("handle", [
-        (42, "bar"),
-        ("foo", ),
-    ])
+    @pytest.mark.parametrize(
+        "handle",
+        [
+            (42, "bar"),
+            ("foo",),
+        ],
+    )
     def test_unmap_and_release_calls_python_unmap_and_release(self, handle):
-        with patch(
-                "vllm_ascend.device_allocator.camem.python_unmap_and_release"
-        ) as mock_release:
+        with patch("vllm_ascend.device_allocator.camem.python_unmap_and_release") as mock_release:
             unmap_and_release(handle)
             mock_release.assert_called_once_with(*handle)
 
     @patch("vllm_ascend.device_allocator.camem.init_module")
-    @patch(
-        "vllm_ascend.device_allocator.camem.torch.npu.memory.NPUPluggableAllocator"
-    )
-    def test_get_pluggable_allocator(self, mock_allocator_class,
-                                     mock_init_module):
+    @patch("vllm_ascend.device_allocator.camem.torch.npu.memory.NPUPluggableAllocator")
+    def test_get_pluggable_allocator(self, mock_allocator_class, mock_init_module):
         mock_allocator_instance = MagicMock()
         mock_allocator_class.return_value = mock_allocator_instance
 
@@ -128,10 +120,16 @@ class TestCaMem(PytestBase):
             2000: data2,
         }
 
-        # mock is_pin_memory_available, return False as some machine only has cpu
-        with patch(
-                "vllm_ascend.device_allocator.camem.NPUPlatform.is_pin_memory_available",
-                return_value=False):
+        # Mock torch.empty to force pin_memory=False
+        original_torch_empty = torch.empty
+
+        def mock_torch_empty(*args, **kwargs):
+            # If pin_memory was explicitly set to True, change it to False
+            if "pin_memory" in kwargs and kwargs["pin_memory"] is True:
+                kwargs["pin_memory"] = False
+            return original_torch_empty(*args, **kwargs)
+
+        with patch("vllm_ascend.device_allocator.camem.torch.empty", side_effect=mock_torch_empty):
             allocator.sleep(offload_tags="tag1")
 
         # only offload tag1, other tag2 call unmap_and_release
@@ -144,8 +142,7 @@ class TestCaMem(PytestBase):
 
     @patch("vllm_ascend.device_allocator.camem.create_and_map")
     @patch("vllm_ascend.device_allocator.camem.memcpy")
-    def test_wake_up_loads_and_clears_cpu_backup(self, mock_memcpy,
-                                                 mock_create_and_map):
+    def test_wake_up_loads_and_clears_cpu_backup(self, mock_memcpy, mock_create_and_map):
         allocator = CaMemAllocator.get_instance()
 
         handle = (1, 10, 1000, 0)
@@ -168,9 +165,7 @@ class TestCaMem(PytestBase):
         mock_ctx.__enter__.return_value = "data"
         mock_ctx.__exit__.return_value = None
 
-        with patch(
-                "vllm_ascend.device_allocator.camem.use_memory_pool_with_allocator",
-                return_value=mock_ctx):
+        with patch("vllm_ascend.device_allocator.camem.use_memory_pool_with_allocator", return_value=mock_ctx):
             with allocator.use_memory_pool(tag="my_tag"):
                 assert allocator.current_tag == "my_tag"
             # restore old tag after context manager exits

tests/ut/device_allocator/test_sleep_mem_optimized.py

@@ -0,0 +1,107 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+from contextlib import nullcontext
+from dataclasses import dataclass
+from unittest.mock import MagicMock, patch
+
+from vllm_ascend.device_allocator.sleep_mem_optimized import (
+    AclGraphSleepWakeupManager,
+    HcclSleepWakeupManager,
+    SleepWakeupManager,
+)
+
+
+@dataclass
+class DummyGraphParams:
+    events: dict[int, list]
+    workspaces: dict[int, object]
+    extra_handles: dict[int, list]
+    metadata: dict[int, tuple]
+
+
+def test_acl_graph_reset_graph_params_clears_list_values_only():
+    workspace = object()
+    params = DummyGraphParams(
+        events={1: ["event"]},
+        workspaces={1: workspace},
+        extra_handles={1: ["handle"]},
+        metadata={1: ("keep",)},
+    )
+
+    AclGraphSleepWakeupManager.reset_graph_params(params)
+
+    assert params.events == {1: []}
+    assert params.extra_handles == {1: []}
+    assert params.workspaces == {1: workspace}
+    assert params.metadata == {1: ("keep",)}
+
+
+def test_acl_graph_wakeup_waits_for_kv_cache_tag():
+    model_runner = MagicMock()
+    manager = AclGraphSleepWakeupManager(MagicMock(), lambda: model_runner)
+
+    manager.wakeup(tags=["weights"])
+    model_runner.capture_model.assert_not_called()
+
+    manager.wakeup(tags=["kv_cache"])
+    model_runner.capture_model.assert_called_once_with()
+
+
+def test_sleep_wakeup_manager_skips_acl_sleep_when_aclgraph_disabled():
+    model_runner = MagicMock()
+    model_runner.use_aclgraph = False
+    manager = SleepWakeupManager(MagicMock(), MagicMock(), lambda: model_runner)
+    manager.acl_graph.sleep = MagicMock()
+    manager.hccl.sleep = MagicMock()
+    with patch(
+        "vllm_ascend.device_allocator.sleep_mem_optimized.torch.npu.mem_get_info",
+        side_effect=[(10, 20), (12, 20)],
+    ):
+        manager.sleep()
+
+    manager.acl_graph.sleep.assert_not_called()
+    manager.hccl.sleep.assert_called_once_with()
+
+
+def test_sleep_wakeup_manager_cleans_acl_before_hccl_when_aclgraph_enabled():
+    model_runner = MagicMock()
+    model_runner.use_aclgraph = True
+    manager = SleepWakeupManager(MagicMock(), MagicMock(), lambda: model_runner)
+    calls = []
+    manager.acl_graph.sleep = MagicMock(side_effect=lambda: calls.append("acl"))
+    manager.hccl.sleep = MagicMock(side_effect=lambda: calls.append("hccl"))
+
+    mem_info = [(10, 20), (12, 20), (12, 20), (13, 20)]
+    with patch("vllm_ascend.device_allocator.sleep_mem_optimized.torch.npu.mem_get_info", side_effect=mem_info):
+        manager.sleep()
+
+    assert calls == ["acl", "hccl"]
+
+
+def test_hccl_wakeup_restores_and_refreshes_moe_groups():
+    manager = HcclSleepWakeupManager(MagicMock(), MagicMock())
+
+    with (
+        patch("vllm_ascend.device_allocator.sleep_mem_optimized.set_current_vllm_config", return_value=nullcontext()),
+        patch.object(manager, "restore_hccl", return_value=2) as mock_restore,
+        patch.object(manager, "refresh_moe_hccl_groups") as mock_refresh,
+    ):
+        manager.wakeup()
+
+    mock_restore.assert_called_once_with()
+    mock_refresh.assert_called_once_with()

tests/ut/distributed/ascend_store/_mock_deps.py

@@ -0,0 +1,463 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+"""Mock heavy dependencies (torch, vllm, etc.) for ascend_store unit tests.
+
+IMPORTANT: This module MUST be imported before any vllm_ascend or vllm
+imports in each test file.
+
+Usage at the top of each test file:
+    import tests.ut.distributed.ascend_store._mock_deps  # noqa: F401, E402
+"""
+
+import importlib.util
+import logging
+import os
+import sys
+import types
+from typing import Any
+from unittest.mock import MagicMock
+
+# ---------------------------------------------------------------------------
+# Mock torch / torch_npu
+# ---------------------------------------------------------------------------
+if "torch" not in sys.modules and importlib.util.find_spec("torch") is None:
+    _torch = types.ModuleType("torch")
+    _torch.Tensor = MagicMock  # type: ignore[attr-defined]
+    _torch.bool = "bool"  # type: ignore[attr-defined]
+    _torch.float16 = "float16"  # type: ignore[attr-defined]
+    _torch.float32 = "float32"  # type: ignore[attr-defined]
+    _torch.zeros = MagicMock(return_value=MagicMock())  # type: ignore[attr-defined]
+    _torch.sum = MagicMock(return_value=0)  # type: ignore[attr-defined]
+    _torch.device = MagicMock()  # type: ignore[attr-defined]
+    _torch.distributed = MagicMock()  # type: ignore[attr-defined]
+    _npu = MagicMock()
+    _npu.Event = MagicMock
+    _npu.current_device = MagicMock(return_value=0)
+    _npu.set_device = MagicMock()
+    _torch.npu = _npu  # type: ignore[attr-defined]
+    sys.modules["torch"] = _torch
+    sys.modules["torch.distributed"] = _torch.distributed  # type: ignore[attr-defined]
+
+if "torch_npu" not in sys.modules:
+    sys.modules["torch_npu"] = MagicMock()
+    sys.modules["torch_npu._inductor"] = MagicMock()
+
+# ---------------------------------------------------------------------------
+# Mock vllm modules
+# ---------------------------------------------------------------------------
+_MOCK_VLLM_DEPS = importlib.util.find_spec("vllm") is None
+_vllm_mock_modules = [
+    "vllm",
+    "vllm.config",
+    "vllm.distributed",
+    "vllm.distributed.kv_events",
+    "vllm.distributed.kv_transfer",
+    "vllm.distributed.kv_transfer.kv_connector",
+    "vllm.distributed.kv_transfer.kv_connector.factory",
+    "vllm.distributed.kv_transfer.kv_connector.v1",
+    "vllm.distributed.kv_transfer.kv_connector.v1.base",
+    "vllm.distributed.parallel_state",
+    "vllm.envs",
+    "vllm.forward_context",
+    "vllm.logger",
+    "vllm.model_executor",
+    "vllm.model_executor.layers",
+    "vllm.model_executor.layers.linear",
+    "vllm.model_executor.layers.quantization",
+    "vllm.platforms",
+    "vllm.utils",
+    "vllm.utils.hashing",
+    "vllm.utils.math_utils",
+    "vllm.utils.network_utils",
+    "vllm.v1",
+    "vllm.v1.attention",
+    "vllm.v1.attention.backend",
+    "vllm.v1.core",
+    "vllm.v1.core.block_pool",
+    "vllm.v1.core.kv_cache_manager",
+    "vllm.v1.core.kv_cache_utils",
+    "vllm.v1.core.sched",
+    "vllm.v1.core.sched.output",
+    "vllm.v1.core.single_type_kv_cache_manager",
+    "vllm.v1.kv_cache_interface",
+    "vllm.v1.kv_cache_spec_registry",
+    "vllm.v1.outputs",
+    "vllm.v1.request",
+    "vllm.v1.serial_utils",
+]
+if _MOCK_VLLM_DEPS:
+    for _mod_name in _vllm_mock_modules:
+        if _mod_name not in sys.modules:
+            sys.modules[_mod_name] = MagicMock()
+
+if _MOCK_VLLM_DEPS:
+    sys.modules["vllm.utils.math_utils"].cdiv = lambda a, b: -(-a // b)  # type: ignore[attr-defined]
+    sys.modules["vllm.logger"].logger = logging.getLogger("vllm")  # type: ignore[attr-defined]
+
+_base_mod: Any = (
+    sys.modules["vllm.distributed.kv_transfer.kv_connector.v1.base"] if _MOCK_VLLM_DEPS else types.SimpleNamespace()
+)
+_base_mod.KVConnectorBase_V1 = type("KVConnectorBase_V1", (), {"__init__": lambda self, **kw: None})  # type: ignore[attr-defined]
+_base_mod.KVConnectorMetadata = type("KVConnectorMetadata", (), {})  # type: ignore[attr-defined]
+_base_mod.KVConnectorWorkerMetadata = type("KVConnectorWorkerMetadata", (), {})  # type: ignore[attr-defined]
+_base_mod.KVConnectorRole = MagicMock()  # type: ignore[attr-defined]
+_base_mod.KVConnectorRole.SCHEDULER = "SCHEDULER"
+_base_mod.KVConnectorRole.WORKER = "WORKER"
+_base_mod.SupportsHMA = type("SupportsHMA", (), {})  # type: ignore[attr-defined]
+
+_events_mod: Any = sys.modules["vllm.distributed.kv_events"] if _MOCK_VLLM_DEPS else types.SimpleNamespace()
+_events_mod.KVCacheEvent = type("KVCacheEvent", (), {})  # type: ignore[attr-defined]
+_events_mod.KVConnectorKVEvents = type("KVConnectorKVEvents", (), {})  # type: ignore[attr-defined]
+
+
+class _FakeAggregator:
+    def __init__(self, *args, **kwargs):
+        self._mock = MagicMock()
+
+    def __getattr__(self, name):
+        return getattr(self._mock, name)
+
+
+_events_mod.KVEventAggregator = _FakeAggregator  # type: ignore[attr-defined]
+_events_mod.BlockStored = type(  # type: ignore[attr-defined]
+    "BlockStored",
+    (),
+    {"__init__": lambda self, **kwargs: self.__dict__.update(kwargs)},
+)
+
+_kv_cache_utils_mod: Any = sys.modules["vllm.v1.core.kv_cache_utils"] if _MOCK_VLLM_DEPS else types.SimpleNamespace()
+_kv_cache_utils_mod.BlockHash = bytes  # type: ignore[attr-defined]
+_kv_cache_utils_mod.maybe_convert_block_hash = lambda x: x  # type: ignore[attr-defined]
+
+
+class _FakeKVCacheBlock:
+    def __init__(self, block_id=0, **kwargs):
+        self.block_id = block_id
+        self.__dict__.update(kwargs)
+
+
+class _FakeKVCacheSpec:
+    def __init__(self, block_size=16, **kwargs):
+        self.block_size = block_size
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def __eq__(self, other):
+        return type(self) is type(other) and self.__dict__ == getattr(other, "__dict__", {})
+
+    def copy_with_new_block_size(self, block_size):
+        kwargs = self.__dict__.copy()
+        kwargs["block_size"] = block_size
+        return type(self)(**kwargs)
+
+    @property
+    def page_size_bytes(self):
+        num_kv_heads = getattr(self, "num_kv_heads", 1)
+        head_size = getattr(self, "head_size", 1)
+        dtype = getattr(self, "dtype", None)
+        dtype_size = getattr(dtype, "itemsize", None)
+        if dtype_size is None and dtype is not None and hasattr(dtype, "element_size"):
+            dtype_size = dtype.element_size()
+        return self.block_size * num_kv_heads * head_size * int(dtype_size or 1) * 2
+
+
+class _FakeFullAttentionSpec(_FakeKVCacheSpec):
+    pass
+
+
+class _FakeSlidingWindowSpec(_FakeKVCacheSpec):
+    def __init__(self, block_size=16, sliding_window=32, **kwargs):
+        super().__init__(block_size=block_size, sliding_window=sliding_window, **kwargs)
+
+
+class _FakeMambaSpec(_FakeKVCacheSpec):
+    def __init__(self, block_size=16, **kwargs):
+        super().__init__(block_size=block_size, **kwargs)
+        self.num_speculative_blocks = getattr(self, "num_speculative_blocks", 0)
+
+
+class _FakeUniformTypeKVCacheSpecs(_FakeKVCacheSpec):
+    def __init__(self, block_size=16, kv_cache_specs=None, **kwargs):
+        super().__init__(block_size=block_size, **kwargs)
+        self.kv_cache_specs = kv_cache_specs or {}
+
+    @classmethod
+    def from_specs(cls, kv_cache_specs):
+        if not kv_cache_specs:
+            return None
+        first_spec = next(iter(kv_cache_specs.values()))
+        return cls(
+            block_size=getattr(first_spec, "block_size", 16),
+            kv_cache_specs=kv_cache_specs,
+        )
+
+
+class _FakeKVCacheGroupSpec:
+    def __init__(self, layer_names=None, kv_cache_spec=None, is_eagle_group=False):
+        self.layer_names = layer_names or []
+        self.kv_cache_spec = kv_cache_spec or _FakeFullAttentionSpec()
+        self.is_eagle_group = is_eagle_group
+
+
+class _FakeKVCacheConfig:
+    def __init__(self, num_blocks=1, kv_cache_tensors=None, kv_cache_groups=None):
+        self.num_blocks = num_blocks
+        self.kv_cache_tensors = kv_cache_tensors or []
+        self.kv_cache_groups = kv_cache_groups or []
+
+
+_kv_cache_utils_mod.KVCacheBlock = _FakeKVCacheBlock  # type: ignore[attr-defined]
+_kv_cache_utils_mod.BlockHashList = list  # type: ignore[attr-defined]
+
+
+class _FakeBlockPool:
+    def __init__(self, *args, **kwargs):
+        self.null_block = _FakeKVCacheBlock(block_id=0)
+        self._next_block_id = 1
+
+    def get_new_blocks(self, num_blocks):
+        blocks = []
+        for _ in range(num_blocks):
+            blocks.append(_FakeKVCacheBlock(block_id=self._next_block_id))
+            self._next_block_id += 1
+        return blocks
+
+
+if _MOCK_VLLM_DEPS:
+    sys.modules["vllm.v1.core.block_pool"].BlockPool = _FakeBlockPool  # type: ignore[attr-defined]
+
+
+class _FakeSingleTypeKVCacheManager:
+    def __init__(self, *args, **kwargs):
+        self._mock = MagicMock()
+
+    def __getattr__(self, name):
+        return getattr(self._mock, name)
+
+    @classmethod
+    def reachable_block_mask(
+        cls,
+        start_block,
+        end_block,
+        alignment_tokens,
+        kv_cache_spec,
+        use_eagle,
+        retention_interval=None,
+        num_prompt_tokens=None,
+    ):
+        return None
+
+    @classmethod
+    def find_longest_cache_hit(
+        cls,
+        block_hashes,
+        max_length,
+        kv_cache_group_ids,
+        block_pool,
+        kv_cache_spec,
+        drop_eagle_block=False,
+        alignment_tokens=16,
+        dcp_world_size=1,
+        pcp_world_size=1,
+    ):
+        computed: tuple[list[object], ...] = tuple([] for _ in kv_cache_group_ids)
+        max_blocks = max_length // kv_cache_spec.block_size
+        for block_hash in list(block_hashes)[:max_blocks]:
+            cached = block_pool.get_cached_block(block_hash, kv_cache_group_ids)
+            if not cached:
+                break
+            for blocks, block in zip(computed, cached):
+                blocks.append(block)
+        if drop_eagle_block and computed and computed[0]:
+            for blocks in computed:
+                blocks.pop()
+        return computed
+
+
+class _FakeSlidingWindowManager(_FakeSingleTypeKVCacheManager):
+    @classmethod
+    def reachable_block_mask(
+        cls,
+        start_block,
+        end_block,
+        alignment_tokens,
+        kv_cache_spec,
+        use_eagle,
+        retention_interval=None,
+        num_prompt_tokens=None,
+    ):
+        if alignment_tokens is None:
+            return None
+        per_segment = max(alignment_tokens // kv_cache_spec.block_size, 1)
+        return [(idx + 1) % per_segment == 0 for idx in range(start_block, end_block)]
+
+
+_single_type_mod: Any = (
+    sys.modules["vllm.v1.core.single_type_kv_cache_manager"] if _MOCK_VLLM_DEPS else types.SimpleNamespace()
+)
+_single_type_mod.SingleTypeKVCacheManager = _FakeSingleTypeKVCacheManager  # type: ignore[attr-defined]
+_single_type_mod.FullAttentionManager = _FakeSingleTypeKVCacheManager  # type: ignore[attr-defined]
+_single_type_mod.SlidingWindowManager = _FakeSlidingWindowManager  # type: ignore[attr-defined]
+_single_type_mod.MambaManager = _FakeSingleTypeKVCacheManager  # type: ignore[attr-defined]
+_single_type_mod.spec_manager_map = {  # type: ignore[attr-defined]
+    _FakeFullAttentionSpec: _FakeSingleTypeKVCacheManager,
+    _FakeSlidingWindowSpec: _FakeSlidingWindowManager,
+    _FakeMambaSpec: _FakeSingleTypeKVCacheManager,
+}
+
+_kv_interface_mod: Any = sys.modules["vllm.v1.kv_cache_interface"] if _MOCK_VLLM_DEPS else types.SimpleNamespace()
+_kv_interface_mod.KVCacheSpec = _FakeKVCacheSpec  # type: ignore[attr-defined]
+_kv_interface_mod.FullAttentionSpec = _FakeFullAttentionSpec  # type: ignore[attr-defined]
+_kv_interface_mod.SlidingWindowSpec = _FakeSlidingWindowSpec  # type: ignore[attr-defined]
+_kv_interface_mod.MambaSpec = _FakeMambaSpec  # type: ignore[attr-defined]
+_kv_interface_mod.UniformTypeKVCacheSpecs = _FakeUniformTypeKVCacheSpecs  # type: ignore[attr-defined]
+_kv_interface_mod.KVCacheGroupSpec = _FakeKVCacheGroupSpec  # type: ignore[attr-defined]
+_kv_interface_mod.KVCacheConfig = _FakeKVCacheConfig  # type: ignore[attr-defined]
+
+
+class _FakeKVCacheSpecRegistry:
+    @classmethod
+    def get_manager_class(cls, kv_cache_spec):
+        if isinstance(kv_cache_spec, _FakeSlidingWindowSpec):
+            return _FakeSlidingWindowManager
+        return _FakeSingleTypeKVCacheManager
+
+
+if _MOCK_VLLM_DEPS:
+    sys.modules["vllm.v1.kv_cache_spec_registry"].KVCacheSpecRegistry = _FakeKVCacheSpecRegistry  # type: ignore[attr-defined]
+
+_sched_output_mod: Any = sys.modules["vllm.v1.core.sched.output"] if _MOCK_VLLM_DEPS else types.SimpleNamespace()
+_sched_output_mod.NewRequestData = MagicMock  # type: ignore[attr-defined]
+
+if _MOCK_VLLM_DEPS:
+    sys.modules["vllm.envs"].VLLM_RPC_BASE_PATH = "/tmp/vllm_rpc"  # type: ignore[attr-defined]
+
+# ---------------------------------------------------------------------------
+# Mock external backends
+# ---------------------------------------------------------------------------
+for _mod_name in [
+    "mooncake",
+    "mooncake.engine",
+    "mooncake.store",
+    "memcache_hybrid",
+    "yr",
+    "yr.datasystem",
+    "yr.datasystem.hetero_client",
+    "yr.datasystem.kv_client",
+    "yr.datasystem.object_client",
+    "zmq",
+]:
+    if _mod_name not in sys.modules:
+        sys.modules[_mod_name] = MagicMock()
+
+# ---------------------------------------------------------------------------
+# Mock vllm_ascend transitive imports
+# ---------------------------------------------------------------------------
+
+
+def _make_pkg(name, path=""):
+    mod = types.ModuleType(name)
+    mod.__path__ = [path]  # type: ignore[attr-defined]
+    mod.__package__ = name  # type: ignore[attr-defined]
+    return mod
+
+
+_vllm_ascend_real_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "vllm_ascend"))
+_vllm_ascend_package_paths = {
+    "vllm_ascend": _vllm_ascend_real_path,
+    "vllm_ascend.distributed": os.path.join(_vllm_ascend_real_path, "distributed"),
+}
+for _pkg, _path in _vllm_ascend_package_paths.items():
+    if _pkg not in sys.modules:
+        sys.modules[_pkg] = _make_pkg(_pkg, _path)
+
+_distributed_utils = types.ModuleType("vllm_ascend.distributed.utils")
+_distributed_utils.get_decode_context_model_parallel_rank = MagicMock(  # type: ignore[attr-defined]
+    return_value=0
+)
+_distributed_utils.get_decode_context_model_parallel_world_size = MagicMock(  # type: ignore[attr-defined]
+    return_value=1
+)
+sys.modules["vllm_ascend.distributed.utils"] = _distributed_utils
+
+_kv_transfer_init = _make_pkg("vllm_ascend.distributed.kv_transfer")
+_kv_transfer_init.register_connector = MagicMock()  # type: ignore[attr-defined]
+sys.modules["vllm_ascend.distributed.kv_transfer"] = _kv_transfer_init
+
+_kv_utils_pkg = _make_pkg("vllm_ascend.distributed.kv_transfer.utils")
+sys.modules["vllm_ascend.distributed.kv_transfer.utils"] = _kv_utils_pkg
+sys.modules["vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine"] = MagicMock()
+
+_kv_pool_pkg = _make_pkg("vllm_ascend.distributed.kv_transfer.kv_pool")
+sys.modules["vllm_ascend.distributed.kv_transfer.kv_pool"] = _kv_pool_pkg
+
+_ascend_store_real_path = os.path.join(
+    os.path.dirname(__file__),
+    "..",
+    "..",
+    "..",
+    "..",
+    "vllm_ascend",
+    "distributed",
+    "kv_transfer",
+    "kv_pool",
+    "ascend_store",
+)
+_ascend_store_pkg = _make_pkg(
+    "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store",
+    os.path.abspath(_ascend_store_real_path),
+)
+sys.modules["vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store"] = _ascend_store_pkg
+
+_backend_pkg = _make_pkg(
+    "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend",
+    os.path.join(os.path.abspath(_ascend_store_real_path), "backend"),
+)
+sys.modules["vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend"] = _backend_pkg
+
+# Mirror the real backend/__init__.py entry points. The scheduler/worker resolve
+# the backend class dynamically via ``importlib.import_module(path)``; tests that
+# exercise those paths patch ``<module>.importlib`` locally (see
+# test_pool_scheduler.py / test_pool_worker.py) so the backend resolves to a
+# MagicMock. Do NOT register the backends in sys.modules or globally wrap
+# import_module here: test_backend.py imports the real backend classes and also
+# relies on ``mock.patch`` (which itself calls importlib.import_module) resolving
+# those real modules.
+_backend_module_paths = {
+    "mooncake": "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend",
+    "memcache": "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.memcache_backend",
+    "yuanrong": "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.yuanrong_backend",
+}
+_backend_pkg.backend_map = {  # type: ignore[attr-defined]
+    "mooncake": {"name": "MooncakeBackend", "path": _backend_module_paths["mooncake"]},
+    "memcache": {"name": "MemcacheBackend", "path": _backend_module_paths["memcache"]},
+    "yuanrong": {"name": "YuanrongBackend", "path": _backend_module_paths["yuanrong"]},
+}
+
+if "vllm_ascend.utils" not in sys.modules or not hasattr(sys.modules["vllm_ascend.utils"], "AscendDeviceType"):
+    _ascend_utils = MagicMock()
+    _ascend_utils.AscendDeviceType = MagicMock()
+    _ascend_utils.get_ascend_device_type = MagicMock()
+    sys.modules["vllm_ascend.utils"] = _ascend_utils
+
+# NOTE: vllm_ascend.{ascend_config, memcache_comm_fence} and their helpers
+# (get_ascend_config, AttentionComputeStartGate, ...) are intentionally NOT
+# mocked here. Doing so by mutating these real modules leaks into every other
+# UT in the same pytest session (breaking test_ascend_config / test_platform,
+# which collect after ascend_store and bind the polluted symbols at import).
+# These helpers are mocked per-test, scoped to the ascend_store tests only,
+# via the autouse fixture in tests/ut/conftest.py.

tests/ut/distributed/ascend_store/test_ascend_store_connector.py

@@ -0,0 +1,459 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import types
+import unittest
+from unittest.mock import MagicMock, patch
+
+# isort: off
+import tests.ut.distributed.ascend_store._mock_deps  # noqa: F401, E402
+from vllm.distributed.kv_events import KVCacheEvent
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector import (
+    AscendStoreConnector,
+    AscendStoreKVEvents,
+)
+
+# isort: on
+
+
+def _mock_events(num_workers=1):
+    events = AscendStoreKVEvents(num_workers=num_workers)
+    events._aggregator = MagicMock()
+    return events
+
+
+class TestAscendStoreKVEvents(unittest.TestCase):
+    def _make_events(self, num_workers=1):
+        return _mock_events(num_workers=num_workers)
+
+    def test_add_and_get_events(self):
+        ev = self._make_events()
+        mock_events = [MagicMock(spec=KVCacheEvent), MagicMock(spec=KVCacheEvent)]
+        ev.add_events(mock_events)
+        ev._aggregator.get_all_events.return_value = mock_events
+        result = ev.get_all_events()
+        self.assertEqual(result, mock_events)
+
+    def test_aggregate(self):
+        ev = self._make_events()
+        common = [MagicMock()]
+        ev._aggregator.get_common_events.return_value = common
+        result = ev.aggregate()
+        self.assertIs(result, ev)
+        ev._aggregator.clear_events.assert_called()
+        ev._aggregator.add_events.assert_called_with(common)
+        ev._aggregator.reset_workers.assert_called()
+
+    def test_increment_workers(self):
+        ev = self._make_events()
+        ev.increment_workers(3)
+        ev._aggregator.increment_workers.assert_called_with(3)
+
+    def test_get_number_of_workers(self):
+        ev = self._make_events()
+        ev._aggregator.get_number_of_workers.return_value = 5
+        self.assertEqual(ev.get_number_of_workers(), 5)
+
+    def test_clear_events(self):
+        ev = self._make_events()
+        ev.clear_events()
+        ev._aggregator.clear_events.assert_called()
+        ev._aggregator.reset_workers.assert_called()
+
+    def test_repr(self):
+        ev = self._make_events()
+        ev._aggregator.get_all_events.return_value = []
+        s = repr(ev)
+        self.assertIn("AscendStoreKVEvents", s)
+
+
+class TestAscendStoreConnector(unittest.TestCase):
+    def _make_vllm_config(self, kv_role="kv_producer", extra_config=None):
+        config = MagicMock()
+        config.kv_transfer_config.kv_role = kv_role
+        config.kv_transfer_config.kv_connector = "AscendStoreConnector"
+        config.kv_transfer_config.kv_connector_extra_config = extra_config or {}
+        config.parallel_config.rank = 0
+        return config
+
+    def test_pp_handshake_metadata_is_ignored(self):
+        connector = AscendStoreConnector.__new__(AscendStoreConnector)
+        metadata = {
+            (0, 0): MagicMock(),
+            (1, 0): MagicMock(),
+        }
+        original_metadata = metadata.copy()
+
+        result = connector.set_xfer_handshake_metadata_pp_aware(metadata)
+
+        self.assertIsNone(result)
+        self.assertEqual(metadata, original_metadata)
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolScheduler")
+    def test_init_scheduler_role(self, mock_scheduler_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        _connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.SCHEDULER,
+            kv_cache_config=MagicMock(),
+        )
+        mock_scheduler_cls.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_init_worker_role(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        _connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        mock_worker_cls.assert_called_once()
+        mock_lookup_cls.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolScheduler")
+    def test_scheduler_methods_delegate(self, mock_scheduler_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.SCHEDULER,
+            kv_cache_config=MagicMock(),
+        )
+        mock_sched = mock_scheduler_cls.return_value
+
+        # get_num_new_matched_tokens
+        mock_sched.get_num_new_matched_tokens.return_value = (10, False)
+        result = connector.get_num_new_matched_tokens(MagicMock(), 5)
+        self.assertEqual(result, (10, False))
+
+        # update_state_after_alloc
+        connector.update_state_after_alloc(MagicMock(), MagicMock(), 10)
+        mock_sched.update_state_after_alloc.assert_called_once()
+
+        # build_connector_meta
+        connector.build_connector_meta(MagicMock())
+        mock_sched.build_connector_meta.assert_called_once()
+
+        # request_finished
+        mock_sched.request_finished.return_value = (True, None)
+        result = connector.request_finished(MagicMock(), [1, 2])
+        self.assertEqual(result, (True, None))
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolScheduler")
+    def test_update_connector_output_no_events(self, mock_scheduler_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.SCHEDULER,
+            kv_cache_config=MagicMock(),
+        )
+        output = MagicMock()
+        output.kv_cache_events = None
+        connector.update_connector_output(output)
+        self.assertIsNone(connector._kv_cache_events)
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolScheduler")
+    def test_update_connector_output_with_events(self, mock_scheduler_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.SCHEDULER,
+            kv_cache_config=MagicMock(),
+        )
+        events = _mock_events(num_workers=1)
+        mock_kv_events = [MagicMock()]
+        events._aggregator.get_all_events.return_value = mock_kv_events
+        events._aggregator.get_number_of_workers.return_value = 1
+
+        output = MagicMock()
+        output.kv_cache_events = events
+        connector.update_connector_output(output)
+        self.assertIsNotNone(connector._kv_cache_events)
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolScheduler")
+    def test_update_connector_output_accumulate(self, mock_scheduler_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.SCHEDULER,
+            kv_cache_config=MagicMock(),
+        )
+        # First update
+        events1 = _mock_events(num_workers=1)
+        events1._aggregator.get_all_events.return_value = [MagicMock()]
+        events1._aggregator.get_number_of_workers.return_value = 1
+        output1 = MagicMock()
+        output1.kv_cache_events = events1
+        connector.update_connector_output(output1)
+
+        # Second update
+        events2 = _mock_events(num_workers=1)
+        events2._aggregator.get_all_events.return_value = [MagicMock()]
+        events2._aggregator.get_number_of_workers.return_value = 1
+        output2 = MagicMock()
+        output2.kv_cache_events = events2
+        connector.update_connector_output(output2)
+        self.assertIsNotNone(connector._kv_cache_events)
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolScheduler")
+    def test_take_events(self, mock_scheduler_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.SCHEDULER,
+            kv_cache_config=MagicMock(),
+        )
+        # No events
+        result = list(connector.take_events())
+        self.assertEqual(result, [])
+
+        # With events
+        events = _mock_events(num_workers=1)
+        mock_event = MagicMock()
+        events._aggregator.get_common_events.return_value = [mock_event]
+        events._aggregator.get_all_events.return_value = [mock_event]
+        connector._kv_cache_events = events
+        result = list(connector.take_events())
+        self.assertEqual(len(result), 1)
+        self.assertIsNone(connector._kv_cache_events)
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_worker_methods(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        mock_worker = mock_worker_cls.return_value
+
+        # register_kv_caches
+        connector.register_kv_caches({"layer1": MagicMock()})
+        mock_worker.register_kv_caches.assert_called_once()
+
+        # start_load_kv
+        connector._get_connector_metadata = MagicMock(return_value=MagicMock())
+        connector.start_load_kv(MagicMock())
+        mock_worker.start_load_kv.assert_called_once()
+
+        # wait_for_save (non-consumer)
+        connector.kv_role = "kv_producer"
+        connector.use_layerwise = False
+        connector.wait_for_save()
+        mock_worker.wait_for_save.assert_called_once()
+
+        # get_finished
+        mock_worker.get_finished.return_value = ({"r1"}, {"r2"})
+        done_s, done_r = connector.get_finished({"r1"})
+        self.assertEqual(done_s, {"r1"})
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_wait_for_layer_load_not_layerwise(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config(extra_config={"use_layerwise": False})
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        # Should return immediately without calling worker
+        connector.wait_for_layer_load("layer_0")
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_save_kv_layer_not_layerwise(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config(extra_config={"use_layerwise": False})
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        connector.save_kv_layer("layer_0", MagicMock(), MagicMock())
+        # Should return immediately
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_save_kv_layer_consumer(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config(kv_role="kv_consumer", extra_config={"use_layerwise": True})
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        connector.save_kv_layer("layer_0", MagicMock(), MagicMock())
+        # Consumer should not save
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_wait_for_save_consumer(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config(kv_role="kv_consumer")
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        connector.wait_for_save()
+        mock_worker_cls.return_value.wait_for_save.assert_not_called()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_get_kv_connector_kv_cache_events_empty(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        mock_worker_cls.return_value.get_kv_events.return_value = []
+        result = connector.get_kv_connector_kv_cache_events()
+        self.assertIsNone(result)
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.LookupKeyServer")
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector.KVPoolWorker")
+    def test_get_kv_connector_kv_cache_events_with_events(self, mock_worker_cls, mock_lookup_cls):
+        config = self._make_vllm_config()
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        connector = AscendStoreConnector(
+            vllm_config=config,
+            role=KVConnectorRole.WORKER,
+            kv_cache_config=None,
+        )
+        mock_worker_cls.return_value.get_kv_events.return_value = [MagicMock()]
+        result = connector.get_kv_connector_kv_cache_events()
+        self.assertIsNotNone(result)
+        self.assertIsInstance(result, AscendStoreKVEvents)
+
+
+class TestAscendStoreConnectorLayerwise(unittest.TestCase):
+    """Test connector methods that are specific to layerwise mode."""
+
+    connector_mod: types.ModuleType
+
+    @classmethod
+    def setUpClass(cls):
+        from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store import ascend_store_connector
+
+        cls.connector_mod = ascend_store_connector
+
+    def test_requires_piecewise_for_cudagraph_enabled(self):
+        self.assertTrue(
+            self.connector_mod.AscendStoreConnector.requires_piecewise_for_cudagraph({"use_layerwise": True})
+        )
+
+    def test_requires_piecewise_for_cudagraph_disabled(self):
+        self.assertFalse(
+            self.connector_mod.AscendStoreConnector.requires_piecewise_for_cudagraph({"use_layerwise": False})
+        )
+
+    def test_requires_piecewise_for_cudagraph_missing(self):
+        self.assertFalse(self.connector_mod.AscendStoreConnector.requires_piecewise_for_cudagraph({}))
+
+    def test_wait_for_save_layerwise_returns_early(self):
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        with (
+            patch.object(self.connector_mod, "KVPoolWorker") as mock_worker_cls,
+            patch.object(self.connector_mod, "LookupKeyServer") as _mock_lookup_cls,
+        ):
+            config = MagicMock()
+            config.kv_transfer_config.kv_role = "kv_producer"
+            config.kv_transfer_config.kv_connector = "AscendStoreConnector"
+            config.kv_transfer_config.kv_connector_extra_config = {"use_layerwise": True}
+            config.parallel_config.rank = 0
+
+            connector = self.connector_mod.AscendStoreConnector(
+                vllm_config=config,
+                role=KVConnectorRole.WORKER,
+                kv_cache_config=None,
+            )
+            connector.wait_for_save()
+            mock_worker_cls.return_value.wait_for_save.assert_not_called()
+
+    def test_save_kv_layer_layerwise_producer(self):
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        with (
+            patch.object(self.connector_mod, "KVPoolWorker") as mock_worker_cls,
+            patch.object(self.connector_mod, "LookupKeyServer") as _mock_lookup_cls,
+        ):
+            config = MagicMock()
+            config.kv_transfer_config.kv_role = "kv_producer"
+            config.kv_transfer_config.kv_connector = "AscendStoreConnector"
+            config.kv_transfer_config.kv_connector_extra_config = {"use_layerwise": True}
+            config.parallel_config.rank = 0
+
+            connector = self.connector_mod.AscendStoreConnector(
+                vllm_config=config,
+                role=KVConnectorRole.WORKER,
+                kv_cache_config=None,
+            )
+            connector._get_connector_metadata = MagicMock(return_value=MagicMock())
+            connector.save_kv_layer("layer_0", MagicMock(), MagicMock())
+            mock_worker_cls.return_value.save_kv_layer.assert_called_once()
+
+    def test_wait_for_layer_load_layerwise(self):
+        from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+
+        with (
+            patch.object(self.connector_mod, "KVPoolWorker") as mock_worker_cls,
+            patch.object(self.connector_mod, "LookupKeyServer") as _mock_lookup_cls,
+        ):
+            config = MagicMock()
+            config.kv_transfer_config.kv_role = "kv_consumer"
+            config.kv_transfer_config.kv_connector = "AscendStoreConnector"
+            config.kv_transfer_config.kv_connector_extra_config = {"use_layerwise": True}
+            config.parallel_config.rank = 0
+
+            connector = self.connector_mod.AscendStoreConnector(
+                vllm_config=config,
+                role=KVConnectorRole.WORKER,
+                kv_cache_config=None,
+            )
+            connector.wait_for_layer_load("layer_0")
+            mock_worker_cls.return_value.wait_for_layer_load.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/ascend_store/test_backend.py

@@ -0,0 +1,665 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import json
+import os
+import tempfile
+import unittest
+from unittest.mock import MagicMock, patch
+
+import tests.ut.distributed.ascend_store._mock_deps  # noqa: F401, E402
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.backend import Backend
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend import (
+    MooncakeStoreConfig,
+    _convert_to_bytes,
+    _parse_global_segment_size,
+    _ssd_setup_kwargs,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.yuanrong_backend import (
+    YuanrongConfig,
+    YuanrongHelper,
+)
+
+
+def _format_log_call(call):
+    args = call.args
+    return args[0] % args[1:]
+
+
+# =========================================================================
+# Backend ABC
+# =========================================================================
+class TestBackendABC(unittest.TestCase):
+    def test_cannot_instantiate(self):
+        with self.assertRaises(TypeError):
+            Backend(MagicMock())  # type: ignore[abstract]
+
+
+def _make_mooncake_store_config(**overrides) -> MooncakeStoreConfig:
+    """Build MooncakeStoreConfig via from_file(); inherits from_file() defaults."""
+    config = dict(overrides)
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        json.dump(config, f)
+        f.flush()
+        path = f.name
+    try:
+        return MooncakeStoreConfig.from_file(path)
+    finally:
+        os.unlink(path)
+
+
+# =========================================================================
+# MooncakeStoreConfig
+# =========================================================================
+class TestMooncakeStoreConfig(unittest.TestCase):
+    def test_from_file(self):
+        config = {
+            "metadata_server": "127.0.0.1:2379",
+            "global_segment_size": "2GB",
+            "local_buffer_size": "1GB",
+            "protocol": "ascend",
+            "device_name": "npu0",
+            "master_server_address": "127.0.0.1:8080",
+        }
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+            json.dump(config, f)
+            f.flush()
+            path = f.name
+
+        try:
+            cfg = MooncakeStoreConfig.from_file(path)
+            self.assertEqual(cfg.metadata_server, "127.0.0.1:2379")
+            self.assertEqual(cfg.global_segment_size, 2 * 1024**3)
+            self.assertEqual(cfg.local_buffer_size, 1 * 1024**3)
+            self.assertEqual(cfg.protocol, "ascend")
+            self.assertEqual(cfg.device_name, "npu0")
+        finally:
+            os.unlink(path)
+
+    def test_from_file_defaults(self):
+        config = {
+            "metadata_server": "localhost:2379",
+            "master_server_address": "localhost:8080",
+        }
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+            json.dump(config, f)
+            f.flush()
+            path = f.name
+
+        try:
+            cfg = MooncakeStoreConfig.from_file(path)
+            self.assertEqual(cfg.protocol, "ascend")
+            self.assertEqual(cfg.device_name, "")
+            self.assertFalse(cfg.enable_ssd_offload)
+            self.assertEqual(cfg.ssd_offload_path, "")
+        finally:
+            os.unlink(path)
+
+    def test_from_file_ssd_offload(self):
+        ssd_path = TestMooncakeStoreConfig._writable_ssd_path()
+        self.addCleanup(lambda: os.rmdir(ssd_path))
+        cfg = _make_mooncake_store_config(
+            enable_ssd_offload=True,
+            ssd_offload_path=ssd_path,
+        )
+        self.assertTrue(cfg.enable_ssd_offload)
+        self.assertEqual(cfg.ssd_offload_path, ssd_path)
+
+    def test_ssd_offload_requires_absolute_path(self):
+        with self.assertRaises(ValueError):
+            _make_mooncake_store_config(
+                enable_ssd_offload=True,
+                ssd_offload_path="relative/path",
+            )
+
+    def test_ssd_offload_requires_path_in_json(self):
+        with self.assertRaises(ValueError):
+            _make_mooncake_store_config(enable_ssd_offload=True)
+
+    @staticmethod
+    def _writable_ssd_path() -> str:
+        return tempfile.mkdtemp(prefix="mooncake_ssd_ut_")
+
+    @patch(
+        "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend."
+        "mooncake_backend._mooncake_setup_supports_ssd_offload",
+        return_value=False,
+    )
+    def test_ssd_setup_kwargs_off_when_disabled(self, _mock_supports):
+        cfg = _make_mooncake_store_config()
+        self.assertEqual(_ssd_setup_kwargs(cfg), {})
+
+    @patch(
+        "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend."
+        "mooncake_backend._mooncake_setup_supports_ssd_offload",
+        return_value=False,
+    )
+    def test_ssd_setup_kwargs_raises_on_old_mooncake(self, _mock_supports):
+        ssd_path = TestMooncakeStoreConfig._writable_ssd_path()
+        self.addCleanup(lambda: os.rmdir(ssd_path))
+        cfg = _make_mooncake_store_config(
+            enable_ssd_offload=True,
+            ssd_offload_path=ssd_path,
+        )
+        with self.assertRaises(RuntimeError):
+            _ssd_setup_kwargs(cfg)
+
+    @patch(
+        "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend."
+        "mooncake_backend._mooncake_setup_supports_ssd_offload",
+        return_value=True,
+    )
+    def test_ssd_setup_kwargs_when_supported(self, _mock_supports):
+        ssd_path = TestMooncakeStoreConfig._writable_ssd_path()
+        self.addCleanup(lambda: os.rmdir(ssd_path))
+        cfg = _make_mooncake_store_config(
+            enable_ssd_offload=True,
+            ssd_offload_path=ssd_path,
+        )
+        self.assertEqual(
+            _ssd_setup_kwargs(cfg),
+            {
+                "enable_ssd_offload": cfg.enable_ssd_offload,
+                "ssd_offload_path": cfg.ssd_offload_path,
+            },
+        )
+
+    def test_load_from_env_missing(self):
+        with patch.dict(os.environ, {}, clear=True):
+            os.environ.pop("MOONCAKE_CONFIG_PATH", None)
+            with self.assertRaises(ValueError):
+                MooncakeStoreConfig.load_from_env()
+
+    def test_load_from_env(self):
+        config = {
+            "metadata_server": "host:1234",
+            "master_server_address": "host:5678",
+        }
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+            json.dump(config, f)
+            f.flush()
+            path = f.name
+
+        try:
+            with patch.dict(os.environ, {"MOONCAKE_CONFIG_PATH": path}):
+                cfg = MooncakeStoreConfig.load_from_env()
+                self.assertEqual(cfg.metadata_server, "host:1234")
+        finally:
+            os.unlink(path)
+
+
+class TestParseGlobalSegmentSize(unittest.TestCase):
+    def test_int(self):
+        self.assertEqual(_parse_global_segment_size(1024), 1024)
+
+    def test_gb(self):
+        self.assertEqual(_parse_global_segment_size("2GB"), 2 * 1024**3)
+
+    def test_mb(self):
+        self.assertEqual(_parse_global_segment_size("512MB"), 512 * 1024**2)
+
+    def test_kb(self):
+        self.assertEqual(_parse_global_segment_size("256KB"), 256 * 1024)
+
+    def test_b(self):
+        self.assertEqual(_parse_global_segment_size("4096B"), 4096)
+
+    def test_no_unit(self):
+        self.assertEqual(_parse_global_segment_size("2048"), 2048)
+
+    def test_float_input(self):
+        self.assertEqual(_parse_global_segment_size(2048.0), 2048)
+
+    def test_empty_string(self):
+        with self.assertRaises(ValueError):
+            _parse_global_segment_size("")
+
+    def test_invalid_format(self):
+        with self.assertRaises(ValueError):
+            _parse_global_segment_size("abcGB")
+
+    def test_unsupported_type(self):
+        with self.assertRaises(TypeError):
+            _parse_global_segment_size(None)  # type: ignore[arg-type]
+
+
+class TestConvertToBytes(unittest.TestCase):
+    def test_valid(self):
+        self.assertEqual(_convert_to_bytes("10", 1, "10"), 10)
+        self.assertEqual(_convert_to_bytes("1.5", 1024, "1.5KB"), int(1.5 * 1024))
+
+    def test_invalid_number(self):
+        with self.assertRaises(ValueError):
+            _convert_to_bytes("abc", 1, "abc")
+
+
+# =========================================================================
+# YuanrongConfig
+# =========================================================================
+class TestYuanrongConfig(unittest.TestCase):
+    def test_load_from_env(self):
+        with patch.dict(
+            os.environ,
+            {
+                "DS_WORKER_ADDR": "host:1234",
+                "DS_ENABLE_EXCLUSIVE_CONNECTION": "1",
+                "DS_ENABLE_REMOTE_H2D": "0",
+            },
+        ):
+            cfg = YuanrongConfig.load_from_env()
+            self.assertEqual(cfg.worker_addr, "host:1234")
+            self.assertTrue(cfg.enable_exclusive_connection)
+            self.assertFalse(cfg.enable_remote_h2d)
+
+    def test_load_from_env_missing(self):
+        with patch.dict(os.environ, {}, clear=True):
+            os.environ.pop("DS_WORKER_ADDR", None)
+            with self.assertRaises(ValueError):
+                YuanrongConfig.load_from_env()
+
+    def test_load_from_env_defaults(self):
+        with patch.dict(os.environ, {"DS_WORKER_ADDR": "h:1"}):
+            cfg = YuanrongConfig.load_from_env()
+            self.assertFalse(cfg.enable_exclusive_connection)
+            self.assertFalse(cfg.enable_remote_h2d)
+
+
+# =========================================================================
+# YuanrongHelper
+# =========================================================================
+class TestYuanrongHelper(unittest.TestCase):
+    def setUp(self):
+        self.blob_cls = MagicMock()
+        self.blob_list_cls = MagicMock()
+        self.helper = YuanrongHelper(self.blob_cls, self.blob_list_cls)
+
+    def test_normalize_keys_short_valid(self):
+        keys = ["abc-123", "key_2"]
+        result = self.helper.normalize_keys(keys)
+        self.assertEqual(result, keys)
+
+    def test_normalize_keys_with_invalid_chars(self):
+        keys = ["key with spaces/and.dots"]
+        result = self.helper.normalize_keys(keys)
+        self.assertEqual(len(result), 1)
+        # Should not contain the original invalid chars
+        self.assertNotIn(" ", result[0])
+        self.assertNotIn("/", result[0])
+        # Should have hash suffix
+        self.assertIn("__", result[0])
+
+    def test_normalize_keys_at_max_length(self):
+        max_length_key = "a" * 1024
+        result = self.helper.normalize_keys([max_length_key])
+        self.assertEqual(result, [max_length_key])
+
+    def test_normalize_keys_over_max_length(self):
+        long_key = "a" * 1025
+        result = self.helper.normalize_keys([long_key])
+        self.assertEqual(len(result), 1)
+        self.assertEqual(len(result[0]), 1024)
+        self.assertIn("__", result[0])
+
+    def test_make_blob_lists(self):
+        self.helper._device_id = 0
+        addrs = [[100, 200], [300, 400]]
+        sizes = [[10, 20], [30, 40]]
+        result = self.helper.make_blob_lists(addrs, sizes)
+        self.assertEqual(len(result), 2)
+        self.assertEqual(self.blob_cls.call_count, 4)
+
+    def test_make_blob_lists_length_mismatch(self):
+        self.helper._device_id = 0
+        with self.assertRaises(ValueError):
+            self.helper.make_blob_lists([[1]], [[1, 2], [3, 4]])
+
+    def test_make_blob_lists_inner_length_mismatch(self):
+        self.helper._device_id = 0
+        with self.assertRaises(ValueError):
+            self.helper.make_blob_lists([[1, 2]], [[1]])
+
+    def test_make_blob_lists_no_device(self):
+        self.helper._device_id = None
+        with self.assertRaises(RuntimeError):
+            self.helper.make_blob_lists([[1]], [[1]])
+
+
+# =========================================================================
+# MooncakeBackend (mocked store)
+# =========================================================================
+class TestMooncakeBackendMethods(unittest.TestCase):
+    def _make_backend(self):
+        from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend import MooncakeBackend
+
+        with (
+            patch.dict(os.environ, {"MOONCAKE_CONFIG_PATH": "/dev/null"}),
+            patch.object(MooncakeBackend, "__init__", lambda self, pc: None),
+        ):
+            backend = MooncakeBackend.__new__(MooncakeBackend)
+            backend.store = MagicMock()
+            backend.config = MagicMock()
+            backend.local_seg = "127.0.0.1:1234"
+            backend._lazy_init = False
+            backend._store_initialized = True
+            backend._use_fabric_mem = False
+            backend._store_init_lock = MagicMock()
+            backend.local_seg = None
+            return backend
+
+    def test_exists(self):
+        b = self._make_backend()
+        b.store.batch_is_exist.return_value = [1, 0]
+        result = b.exists(["k1", "k2"])
+        self.assertEqual(result, [1, 0])
+
+    def test_put(self):
+        b = self._make_backend()
+        b.store.batch_put_from_multi_buffers.return_value = [0, 0]
+        b.put(["k1"], [[100]], [[10]])
+        b.store.batch_put_from_multi_buffers.assert_called_once()
+
+    def test_put_error(self):
+        b = self._make_backend()
+        b.store.batch_put_from_multi_buffers.return_value = [-1]
+        b.put(["k1"], [[100]], [[10]])  # Should log error but not raise
+
+    def test_put_exception(self):
+        b = self._make_backend()
+        b.store.batch_put_from_multi_buffers.side_effect = RuntimeError("backend fail")
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend.logger"
+        ) as mock_logger:
+            b.put(["k1"], [[100]], [[10]])  # Should log error but not raise
+        error_log = _format_log_call(mock_logger.error.call_args)
+        self.assertIn("RuntimeError", error_log)
+        self.assertIn("backend fail", error_log)
+
+    def test_get(self):
+        b = self._make_backend()
+        b.store.batch_get_into_multi_buffers.return_value = [0]
+        b.get(["k1"], [[100]], [[10]])
+        b.store.batch_get_into_multi_buffers.assert_called_once()
+
+    def test_get_error(self):
+        b = self._make_backend()
+        b.store.batch_get_into_multi_buffers.return_value = [-1]
+        b.get(["k1"], [[100]], [[10]])
+
+    def test_get_exception(self):
+        b = self._make_backend()
+        b.store.batch_get_into_multi_buffers.side_effect = RuntimeError("backend fail")
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend.logger"
+        ) as mock_logger:
+            b.get(["k1"], [[100]], [[10]])
+        error_log = _format_log_call(mock_logger.error.call_args)
+        self.assertIn("RuntimeError", error_log)
+        self.assertIn("backend fail", error_log)
+
+    def test_register_buffer(self):
+        b = self._make_backend()
+        with (
+            patch(
+                "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend.global_te"
+            ) as mock_te,
+            patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend.get_ip"),
+        ):
+            b.register_buffer([100], [200])
+            mock_te.register_buffer.assert_called_once()
+
+
+# =========================================================================
+# YuanrongBackend (mocked store)
+# =========================================================================
+class TestYuanrongBackendMethods(unittest.TestCase):
+    def _make_backend(self):
+        from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.yuanrong_backend import YuanrongBackend
+
+        with patch.object(YuanrongBackend, "__init__", lambda self, pc: None):
+            backend = YuanrongBackend.__new__(YuanrongBackend)
+            backend._helper = MagicMock()
+            backend._helper._device_id = 0
+            backend._helper.normalize_keys = lambda keys: keys
+            backend._helper.make_blob_lists = lambda a, s: [MagicMock() for _ in a]
+            backend._hetero_client = MagicMock()
+            backend._ds_set_param = MagicMock()
+            backend._is_a2 = False
+            backend._registered_buffers = None
+            backend._buffers_registered = False
+            backend.config = YuanrongConfig(
+                worker_addr="127.0.0.1:0",
+                enable_exclusive_connection=False,
+                enable_remote_h2d=False,
+            )
+            backend.rank = 0
+            return backend
+
+    def test_exists_empty(self):
+        b = self._make_backend()
+        result = b.exists([])
+        self.assertEqual(result, [])
+
+    def test_exists(self):
+        b = self._make_backend()
+        b._hetero_client.exist.return_value = [True, False]
+        result = b.exists(["k1", "k2"])
+        self.assertEqual(result, [1, 0])
+
+    def test_exists_exception(self):
+        b = self._make_backend()
+        b._hetero_client.exist.side_effect = Exception("fail")
+        result = b.exists(["k1"])
+        self.assertEqual(result, [0])
+
+    def test_get_empty(self):
+        b = self._make_backend()
+        result = b.get([], [], [])
+        self.assertEqual(result, [])
+        b._hetero_client.mget_h2d.assert_not_called()
+
+    def test_get(self):
+        b = self._make_backend()
+        b._hetero_client.mget_h2d.return_value = []
+        result = b.get(["k1"], [[100]], [[10]])
+        self.assertEqual(result, [0])
+        b._hetero_client.mget_h2d.assert_called_once()
+
+    def test_get_partial_failure(self):
+        b = self._make_backend()
+        b._hetero_client.mget_h2d.return_value = ["k2"]
+        result = b.get(["k1", "k2", "k3"], [[100], [200], [300]], [[10], [20], [30]])
+        self.assertEqual(result, [0, 1, 0])
+
+    def test_get_failed_keys(self):
+        b = self._make_backend()
+        b._hetero_client.mget_h2d.return_value = ["k1"]
+        result = b.get(["k1"], [[100]], [[10]])  # Should log error
+        self.assertEqual(result, [1])
+
+    def test_get_exception(self):
+        b = self._make_backend()
+        b._hetero_client.mget_h2d.side_effect = RuntimeError("backend fail")
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.yuanrong_backend.logger"
+        ) as mock_logger:
+            result = b.get(["k1"], [[100]], [[10]])
+        error_log = _format_log_call(mock_logger.error.call_args)
+        self.assertIsNone(result)
+        self.assertIn("RuntimeError", error_log)
+        self.assertIn("backend fail", error_log)
+
+    def test_put_empty(self):
+        b = self._make_backend()
+        b.put([], [], [])
+        b._hetero_client.mset_d2h.assert_not_called()
+
+    def test_put(self):
+        b = self._make_backend()
+        b.put(["k1"], [[100]], [[10]])
+        b._hetero_client.mset_d2h.assert_called_once()
+
+    def test_put_exception(self):
+        b = self._make_backend()
+        b._hetero_client.mset_d2h.side_effect = RuntimeError("backend fail")
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.yuanrong_backend.logger"
+        ) as mock_logger:
+            b.put(["k1"], [[100]], [[10]])
+        error_log = _format_log_call(mock_logger.error.call_args)
+        self.assertIn("RuntimeError", error_log)
+        self.assertIn("backend fail", error_log)
+
+    def test_register_buffer_noop_when_remote_h2d_disabled(self):
+        b = self._make_backend()
+        b.register_buffer([100], [200])
+        b._hetero_client.pre_register_device_memory.assert_not_called()
+
+    def test_register_buffer_when_remote_h2d_enabled(self):
+        b = self._make_backend()
+        b.config.enable_remote_h2d = True
+        b.register_buffer([100], [200])
+        b._hetero_client.pre_register_device_memory.assert_called_once_with([100], [200])
+
+    def test_register_buffer_noop_on_a2(self):
+        # A2 must not register (opposite of memcache_backend's _is_a2 gating).
+        b = self._make_backend()
+        b._is_a2 = True
+        b.config.enable_remote_h2d = True
+        b.register_buffer([100], [200])
+        b._hetero_client.pre_register_device_memory.assert_not_called()
+
+    def test_register_buffer_idempotent(self):
+        b = self._make_backend()
+        b.config.enable_remote_h2d = True
+        b.register_buffer([100], [200])
+        b.register_buffer([300], [400])
+        b._hetero_client.pre_register_device_memory.assert_called_once_with([100], [200])
+
+    def test_register_buffers_if_needed_no_buffers(self):
+        b = self._make_backend()
+        b.config.enable_remote_h2d = True
+        b._registered_buffers = None
+        b._register_buffers_if_needed()
+        b._hetero_client.pre_register_device_memory.assert_not_called()
+
+    def test_register_buffers_if_needed_already_registered(self):
+        b = self._make_backend()
+        b.config.enable_remote_h2d = True
+        b._registered_buffers = ([100], [200])
+        b._buffers_registered = True
+        b._register_buffers_if_needed()
+        b._hetero_client.pre_register_device_memory.assert_not_called()
+
+    def test_register_buffers_if_needed_disabled(self):
+        b = self._make_backend()
+        b.config.enable_remote_h2d = False
+        b._registered_buffers = ([100], [200])
+        b._register_buffers_if_needed()
+        b._hetero_client.pre_register_device_memory.assert_not_called()
+
+    def test_ensure_device_ready(self):
+        b = self._make_backend()
+        b._helper._device_id = None
+        b.set_device = MagicMock()
+        b._ensure_device_ready()
+        b.set_device.assert_called_once()
+
+    def test_ensure_device_ready_already_set(self):
+        b = self._make_backend()
+        b._helper._device_id = 0
+        b.set_device = MagicMock()
+        b._ensure_device_ready()
+        b.set_device.assert_not_called()
+
+
+# =========================================================================
+# MemcacheBackend (mocked store)
+# =========================================================================
+class TestMemcacheBackendMethods(unittest.TestCase):
+    def _make_backend(self):
+        from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.memcache_backend import MemcacheBackend
+
+        with patch.object(MemcacheBackend, "__init__", lambda self, pc: None):
+            backend = MemcacheBackend.__new__(MemcacheBackend)
+            backend.store = MagicMock()
+            backend.local_rank = 0
+            # Set internal state to avoid lazy init logic during tests
+            backend._lazy_init = False
+            backend._store_initialized = True
+            backend._is_a2 = False
+            backend._registered_buffers = None
+            backend._buffers_registered = False
+            return backend
+
+    def test_exists(self):
+        b = self._make_backend()
+        b.store.batch_is_exist.return_value = [1]
+        self.assertEqual(b.exists(["k1"]), [1])
+
+    def test_register_buffer(self):
+        b = self._make_backend()
+        b._is_a2 = True
+        b.register_buffer([100], [200])
+        b.store.register_buffer.assert_called_once()
+
+    def test_get(self):
+        b = self._make_backend()
+        b.store.batch_get_into_layers.return_value = [0]
+        b.get(["k1"], [[100]], [[10]])
+        b.store.batch_get_into_layers.assert_called_once()
+
+    def test_get_error(self):
+        b = self._make_backend()
+        b.store.batch_get_into_layers.return_value = [1]  # non-zero = error
+        b.get(["k1"], [[100]], [[10]])
+
+    def test_get_exception(self):
+        b = self._make_backend()
+        b.store.batch_get_into_layers.side_effect = RuntimeError("backend fail")
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.memcache_backend.logger"
+        ) as mock_logger:
+            b.get(["k1"], [[100]], [[10]])
+        error_log = _format_log_call(mock_logger.error.call_args)
+        self.assertIn("RuntimeError", error_log)
+        self.assertIn("backend fail", error_log)
+
+    def test_put(self):
+        b = self._make_backend()
+        b.store.batch_put_from_layers.return_value = [0]
+        b.put(["k1"], [[100]], [[10]])
+        b.store.batch_put_from_layers.assert_called_once()
+
+    def test_put_error(self):
+        b = self._make_backend()
+        b.store.batch_put_from_layers.return_value = [1]
+        b.put(["k1"], [[100]], [[10]])
+
+    def test_put_exception(self):
+        b = self._make_backend()
+        b.store.batch_put_from_layers.side_effect = RuntimeError("backend fail")
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.memcache_backend.logger"
+        ) as mock_logger:
+            b.put(["k1"], [[100]], [[10]])
+        error_log = _format_log_call(mock_logger.error.call_args)
+        self.assertIn("RuntimeError", error_log)
+        self.assertIn("backend fail", error_log)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/ascend_store/test_config_data.py

@@ -0,0 +1,581 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import hashlib
+import unittest
+from unittest.mock import MagicMock
+
+import tests.ut.distributed.ascend_store._mock_deps  # noqa: F401, E402
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
+    AscendConnectorMetadata,
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    LayerMultiBlockReqMeta,
+    LayerPoolKey,
+    LoadSpec,
+    PoolKey,
+    ReqMeta,
+    RequestTracker,
+    get_block_hashes,
+)
+
+_GROUPED_BLOCK_HASH_DOMAIN = b"vllm-ascend-grouped-block-hash-v1\0"
+_GROUPED_BLOCK_HASH_LENGTH_PREFIX_BYTES = 4
+
+
+def _expected_grouped_hash(*block_hashes):
+    hasher = hashlib.sha256()
+    hasher.update(_GROUPED_BLOCK_HASH_DOMAIN)
+    hasher.update(len(block_hashes).to_bytes(_GROUPED_BLOCK_HASH_LENGTH_PREFIX_BYTES, "big"))
+    for block_hash in block_hashes:
+        hash_bytes = block_hash.encode("utf-8") if isinstance(block_hash, str) else bytes(block_hash)
+        hasher.update(len(hash_bytes).to_bytes(_GROUPED_BLOCK_HASH_LENGTH_PREFIX_BYTES, "big"))
+        hasher.update(hash_bytes)
+    return hasher.digest()
+
+
+class TestKeyMetadata(unittest.TestCase):
+    def test_fields(self):
+        meta = KeyMetadata(
+            model_name="llama",
+            head_or_tp_rank=0,
+            pcp_rank=0,
+            dcp_rank=0,
+            pp_rank=0,
+        )
+        self.assertEqual(meta.model_name, "llama")
+        self.assertEqual(meta.head_or_tp_rank, 0)
+        self.assertEqual(meta.pcp_rank, 0)
+        self.assertEqual(meta.dcp_rank, 0)
+        self.assertEqual(meta.pp_rank, 0)
+
+
+class TestPoolKey(unittest.TestCase):
+    def setUp(self):
+        self.meta = KeyMetadata("llama", 1, 2, 3, 0)
+
+    def test_hash_equal(self):
+        k1 = PoolKey(self.meta, "abc123")
+        k2 = PoolKey(self.meta, "abc123")
+        self.assertEqual(hash(k1), hash(k2))
+
+    def test_hash_diff(self):
+        k1 = PoolKey(self.meta, "abc123")
+        k2 = PoolKey(self.meta, "def456")
+        self.assertNotEqual(hash(k1), hash(k2))
+
+    def test_to_string(self):
+        k = PoolKey(self.meta, "hash1")
+        s = k.to_string()
+        self.assertIn("llama", s)
+        self.assertIn("@pcp2", s)
+        self.assertIn("@dcp3", s)
+        self.assertIn("@head_or_tp_rank:1", s)
+        self.assertIn("@pp_rank:0", s)
+        self.assertIn("hash1", s)
+
+    def test_pp_ranks_use_distinct_keys(self):
+        other_pp_meta = KeyMetadata("llama", 1, 2, 3, 1)
+        pp0_key = PoolKey(self.meta, "hash1")
+        pp1_key = PoolKey(other_pp_meta, "hash1")
+
+        self.assertNotEqual(pp0_key.to_string(), pp1_key.to_string())
+        self.assertIn("@pp_rank:0", pp0_key.to_string())
+        self.assertIn("@pp_rank:1", pp1_key.to_string())
+
+    def test_split_layers(self):
+        k = PoolKey(self.meta, "hash1")
+        layers = k.split_layers(3)
+        self.assertEqual(len(layers), 3)
+        for i, lk in enumerate(layers):
+            self.assertIsInstance(lk, LayerPoolKey)
+            self.assertEqual(lk.layer_id, i)
+            self.assertEqual(lk.chunk_hash, "hash1")
+
+
+class TestLayerPoolKey(unittest.TestCase):
+    def test_hash(self):
+        meta = KeyMetadata("model", 0, 0, 0, 0)
+        k1 = LayerPoolKey(meta, "h1", 0)
+        k2 = LayerPoolKey(meta, "h1", 1)
+        self.assertNotEqual(hash(k1), hash(k2))
+
+    def test_to_string_contains_layer_id(self):
+        meta = KeyMetadata("model", 0, 0, 0, 0)
+        k = LayerPoolKey(meta, "h1", 5)
+        s = k.to_string()
+        self.assertIn("@layer_id:5", s)
+        self.assertIn("model", s)
+        self.assertTrue(s.endswith("@h1"))
+
+
+class TestChunkedTokenDatabase(unittest.TestCase):
+    def setUp(self):
+        self.meta = KeyMetadata("llama", 0, 0, 0, 0)
+        self.db = ChunkedTokenDatabase([self.meta], block_size=[16], partitions=None)
+        self.db.set_group_buffers({0: [1000, 2000]}, {0: [160, 320]}, group_num_layers={0: 1})
+
+    def test_make_key_by_hash(self):
+        key = self.db._make_key_by_hash("abc")
+        self.assertIsInstance(key, PoolKey)
+        self.assertEqual(key.chunk_hash, "abc")
+
+    def test_process_tokens_empty(self):
+        result = list(self.db.process_tokens(32, []))
+        self.assertEqual(result, [])
+
+    def test_process_tokens_with_str_hashes(self):
+        hashes = ["aaa", "bbb"]
+        result = list(self.db.process_tokens(32, hashes))
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0][0], 0)  # start
+        self.assertEqual(result[0][1], 16)  # end
+        self.assertEqual(result[1][0], 16)
+        self.assertEqual(result[1][1], 32)
+
+    def test_process_tokens_with_bytes_hashes(self):
+        hashes = [b"\xaa\xbb", b"\xcc\xdd"]
+        result = list(self.db.process_tokens(32, hashes))
+        self.assertEqual(len(result), 2)
+
+    def test_process_tokens_with_mask(self):
+        hashes = ["a", "b", "c"]
+        result = list(self.db.process_tokens(48, hashes, mask_num=16))
+        # first chunk (start=0 < mask_num=16) should be skipped
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0][0], 16)
+
+    def test_process_tokens_with_tail_clipped_block_ids_maps_tail_chunks(self):
+        db = ChunkedTokenDatabase([self.meta], block_size=[128], partitions=None)
+        hashes = [bytes([idx % 251]) * 32 for idx in range(128)]
+
+        result = list(
+            db.process_token_key_strings_with_block_ids(
+                128 * 128,
+                hashes,
+                [1000, 1001, 1002, 1003],
+            )
+        )
+
+        self.assertEqual(
+            [start for start, _, _, _, _ in result],
+            [124 * 128, 125 * 128, 126 * 128, 127 * 128],
+        )
+        self.assertEqual(
+            [block_id for _, _, _, _, block_id in result],
+            [1000, 1001, 1002, 1003],
+        )
+
+    def test_process_tokens_token_len_shorter_than_all_blocks(self):
+        hashes = ["a", "b", "c", "d"]
+        # token_len=32 means only first 2 blocks valid
+        result = list(self.db.process_tokens(32, hashes))
+        self.assertEqual(len(result), 2)
+
+    def test_process_tokens_rehashes_grouped_hashes(self):
+        db = ChunkedTokenDatabase([self.meta], block_size=[16], partitions=None, hash_block_size=8)
+        result = list(db.process_tokens(32, ["a", "b", "c", "d"]))
+        self.assertEqual(len(result), 2)
+        self.assertEqual(result[0][2].chunk_hash, _expected_grouped_hash("a", "b").hex())
+        self.assertEqual(len(result[0][2].chunk_hash), 64)
+
+    def test_key_strings_match_pool_keys(self):
+        hashes = ["aaa", "bbb", "ccc"]
+        pool_keys = list(self.db.process_tokens(40, hashes))
+        self.assertEqual(
+            list(self.db.process_token_key_strings(40, hashes)),
+            [
+                (start, end, key.to_string(), hash_val)
+                for (start, end, key), hash_val in zip(pool_keys, hashes, strict=True)
+            ],
+        )
+
+        block_ids = [5, 6]
+        self.assertEqual(
+            list(self.db.process_token_key_strings_with_block_ids(32, hashes, block_ids)),
+            [
+                (start, end, key.to_string(), hash_val, block_id)
+                for (start, end, key), hash_val, block_id in zip(pool_keys[:2], hashes[:2], block_ids, strict=True)
+            ],
+        )
+
+    def test_direct_keys_preserve_multigroup_layerwise_key_semantics(self):
+        group_metadata = [
+            KeyMetadata("llama", 0, 0, 0, 0),
+            KeyMetadata("llama", 1, 0, 0, 0),
+        ]
+        db = ChunkedTokenDatabase(group_metadata, block_size=[16, 32], partitions=None, hash_block_size=16)
+        db.set_group_buffers(
+            {0: [1000], 1: [2000]},
+            {0: [160], 1: [320]},
+            group_cache_families={0: "c1", 1: "c2"},
+            group_num_layers={0: 2, 1: 2},
+        )
+        hashes = ["a", "b", "c", "d"]
+
+        pool_key_result = list(db.process_tokens(64, hashes, kv_cache_group_id=1))
+        direct_key_result = list(db.process_token_key_strings(64, hashes, kv_cache_group_id=1))
+
+        self.assertEqual(len(pool_key_result), 1)
+        self.assertEqual(
+            direct_key_result[0][:3],
+            (pool_key_result[0][0], pool_key_result[0][1], pool_key_result[0][2].to_string()),
+        )
+        layer_key = pool_key_result[0][2].split_layers(2)[1]
+        self.assertIn("@group:1@cache_role:kv@cache_family:c2@layer_id:1", layer_key.to_string())
+
+    def test_key_strings_pre_shard_after_filtering(self):
+        hashes = ["a", "b", "c", "d"]
+        store_mask = [True, False, True, True]
+        result = list(
+            self.db.process_token_key_strings_with_block_ids(
+                64,
+                hashes,
+                [10, 11, 12, 13],
+                chunk_filter=lambda start: store_mask[start // 16],
+                shard_rank=1,
+                shard_size=2,
+            )
+        )
+        self.assertEqual([(start, end, block_id) for start, end, _, _, block_id in result], [(32, 48, 12)])
+
+    def test_get_block_hashes_rehashes_groups(self):
+        for hashes in (["a", "b", "c", "d"], [b"a", b"b", b"c", b"d"]):
+            with self.subTest(hash_type=type(hashes[0])):
+                result = get_block_hashes(hashes, group_block_size=32, hash_block_size=16)
+                expected = [
+                    _expected_grouped_hash(hashes[0], hashes[1]),
+                    _expected_grouped_hash(hashes[2], hashes[3]),
+                ]
+                self.assertEqual(list(result), expected)
+
+    def test_prepare_value(self):
+        addr, size, block_id = self.db.prepare_value(0, 16, [5, 6, 7])
+        self.assertEqual(block_id, 5)
+        self.assertEqual(len(addr), 2)
+        self.assertEqual(addr[0], 1000 + 5 * 160)
+        self.assertEqual(addr[1], 2000 + 5 * 320)
+        self.assertEqual(size[0], 160)
+        self.assertEqual(size[1], 320)
+
+    def test_prepare_value_partial_block(self):
+        addr, size, block_id = self.db.prepare_value(0, 8, [5])
+        self.assertEqual(size[0], 80)  # 160/16*8
+        self.assertEqual(size[1], 160)  # 320/16*8
+
+    def test_prepare_value_uses_block_id_override(self):
+        addr, size, block_id = self.db.prepare_value(64, 80, [5], block_id=99)
+        self.assertEqual(block_id, 99)
+        self.assertEqual(addr[0], 1000 + 99 * 160)
+        self.assertEqual(addr[1], 2000 + 99 * 320)
+        self.assertEqual(size[0], 160)
+        self.assertEqual(size[1], 320)
+
+    def test_prepare_value_layer(self):
+        addr, size, block_id = self.db.prepare_value_layer(0, 16, [5, 6], layer_id=0)
+        self.assertEqual(block_id, 5)
+        self.assertEqual(len(addr), 2)
+        # layer_id=0, entries_per_layers=2 => group_addrs[0] and group_addrs[1]
+        self.assertEqual(addr[0], 1000 + 5 * 160)
+        self.assertEqual(addr[1], 2000 + 5 * 320)
+
+    def test_decode_adaptor_prefill_pp_no_partitions(self):
+        key, addr, size = self.db.decode_adaptor_prefill_pp(["k1"], [[1, 2]], [[10, 20]])
+        self.assertEqual(key, ["k1"])
+
+    def test_decode_adaptor_prefill_pp_single_partition(self):
+        db = ChunkedTokenDatabase([self.meta], [16], partitions=[4])
+        key, addr, size = db.decode_adaptor_prefill_pp(["k1"], [[1, 2]], [[10, 20]])
+        self.assertEqual(key, ["k1"])
+
+    def test_decode_adaptor_prefill_pp_multi_partition(self):
+        db = ChunkedTokenDatabase([self.meta], [16], partitions=[2, 2])
+        db.set_group_buffers({0: [1000, 2000]}, {0: [160, 320]})
+        keys = ["k1@pp_rank:0"]
+        addrs = [[1, 2, 3, 4, 5, 6, 7, 8]]
+        sizes = [[10, 20, 30, 40, 50, 60, 70, 80]]
+        new_keys, new_addrs, new_sizes = db.decode_adaptor_prefill_pp(keys, addrs, sizes)
+        self.assertEqual(len(new_keys), 2)
+        self.assertIn("@pp_rank:0", new_keys[0])
+        self.assertIn("@pp_rank:1", new_keys[1])
+
+
+class TestLoadSpec(unittest.TestCase):
+    def test_fields(self):
+        spec = LoadSpec(vllm_cached_tokens=10, kvpool_cached_tokens=20, can_load=True)
+        self.assertEqual(spec.vllm_cached_tokens, 10)
+        self.assertEqual(spec.kvpool_cached_tokens, 20)
+        self.assertTrue(spec.can_load)
+        self.assertEqual(spec.token_len, 0)
+
+    def test_token_len_default(self):
+        spec = LoadSpec(0, 0, False, token_len=128)
+        self.assertEqual(spec.token_len, 128)
+
+
+class TestRequestTracker(unittest.TestCase):
+    def test_from_new_request(self):
+        new_req = MagicMock()
+        new_req.req_id = "req-1"
+        new_req.block_ids = [10, 20, 30]
+        new_req.prompt_token_ids = list(range(100))
+
+        tracker = RequestTracker.from_new_request(new_req, num_tokens_to_compute=48)
+        self.assertEqual(tracker.req_id, "req-1")
+        self.assertEqual(tracker.token_len, 48)
+        self.assertEqual(tracker.allocated_block_ids, [10, 20, 30])
+        self.assertEqual(len(tracker.token_ids), 48)
+        self.assertEqual(tracker.num_saved_tokens, 0)
+
+    def test_from_new_request_nested_block_ids(self):
+        new_req = MagicMock()
+        new_req.req_id = "req-2"
+        new_req.block_ids = [[10, 20], [30, 40]]
+        new_req.prompt_token_ids = list(range(32))
+
+        tracker = RequestTracker.from_new_request(new_req, num_tokens_to_compute=32)
+        self.assertEqual(tracker.allocated_block_ids, [10, 20])
+
+    def test_update_with_list(self):
+        tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1, 2])
+        tracker.update([3, 4])
+        self.assertEqual(tracker.allocated_block_ids, [1, 2, 3, 4])
+
+    def test_update_with_tuple(self):
+        tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1])
+        tracker.update(([5, 6], [7, 8]))
+        self.assertEqual(tracker.allocated_block_ids, [1, 5, 6])
+
+    def test_update_with_empty(self):
+        tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1])
+        tracker.update([])
+        self.assertEqual(tracker.allocated_block_ids, [1])
+
+    def test_update_invalid_type(self):
+        tracker = RequestTracker(req_id="r1", token_len=16, allocated_block_ids=[1])
+        with self.assertRaises(ValueError):
+            tracker.update("invalid")  # type: ignore[arg-type]
+
+    def test_update_mamba_with_tuple(self):
+        tracker = RequestTracker(
+            req_id="r1", token_len=16, allocated_block_ids_by_group=[[1], [2], [3], [4]], block_sizes=[16] * 4
+        )
+        tracker.update(([5, 6], [0, 7], [0, 8], [0, 9]))
+        self.assertEqual(tracker.allocated_block_ids_by_group[0], [1, 5, 6])
+        self.assertEqual(tracker.allocated_block_ids_by_group[1], [2, 0, 7])
+        self.assertEqual(tracker.allocated_block_ids_by_group[2], [3, 0, 8])
+        self.assertEqual(tracker.allocated_block_ids_by_group[3], [4, 0, 9])
+
+    def test_update_mamba_mtp_with_tuple_chunk2(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids_by_group=[
+                [1, 2],
+                [0, 3, 4, 5, 6],
+                [0, 7, 8, 9, 10],
+                [0, 11, 12, 13, 14],
+            ],
+            mamba_group_ids=[1, 2, 3],
+            num_speculative_blocks=3,
+            block_sizes=[16] * 4,
+        )
+
+        tracker.update(([15, 16], [4, 17], [8, 18], [12, 19]), 32)
+        self.assertEqual(tracker.allocated_block_ids_by_group[0], [1, 2, 15, 16])
+        self.assertEqual(tracker.allocated_block_ids_by_group[1], [0, 3, 0, 5, 6, 4, 17])
+        self.assertEqual(tracker.allocated_block_ids_by_group[2], [0, 7, 0, 9, 10, 8, 18])
+        self.assertEqual(tracker.allocated_block_ids_by_group[3], [0, 11, 0, 13, 14, 12, 19])
+
+    def test_update_mamba_mtp_with_tuple_chunk8(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=128,
+            allocated_block_ids_by_group=[
+                [1, 2, 3, 4, 5, 6, 7, 8],
+                [0, 0, 0, 0, 0, 0, 0, 9, 10, 11, 12],
+                [0, 0, 0, 0, 0, 0, 0, 13, 14, 15, 16],
+                [0, 0, 0, 0, 0, 0, 0, 17, 18, 19, 20],
+            ],
+            mamba_group_ids=[1, 2, 3],
+            num_speculative_blocks=3,
+            block_sizes=[16] * 4,
+        )
+
+        tracker.update(
+            (
+                [21, 22, 23, 24, 25, 26, 27, 28],
+                [0, 0, 0, 0, 10, 11, 12, 29],
+                [0, 0, 0, 0, 14, 15, 16, 30],
+                [0, 0, 0, 0, 18, 19, 20, 31],
+            ),
+            128,
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[0], [1, 2, 3, 4, 5, 6, 7, 8, 21, 22, 23, 24, 25, 26, 27, 28]
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[1], [0, 0, 0, 0, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0, 10, 11, 12, 29]
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[2], [0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 15, 16, 30]
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[3], [0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 19, 20, 31]
+        )
+
+
+class TestReqMeta(unittest.TestCase):
+    def test_from_request_tracker_basic_save(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=0,
+            token_ids=list(range(32)),
+        )
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, block_hashes=[b"h1", b"h2"])
+        self.assertIsNotNone(meta)
+        self.assertEqual(meta.req_id, "r1")
+        self.assertTrue(meta.can_save)
+        self.assertEqual(meta.token_len_chunk, 32)
+        self.assertIsNone(meta.load_spec)
+
+    def test_from_request_tracker_skip_save(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=0,
+        )
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, skip_save=True)
+        self.assertIsNone(meta)
+
+    def test_from_request_tracker_with_load_spec(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=0,
+        )
+        load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=32, can_load=True)
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, load_spec=load_spec, skip_save=True)
+        self.assertIsNotNone(meta)
+        self.assertIsNotNone(meta.load_spec)
+
+    def test_from_request_tracker_load_spec_cannot_load(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=32,
+        )
+        load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=32, can_load=False)
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, load_spec=load_spec, skip_save=True)
+        # can_load=False => load_spec set to None in meta,
+        # but skip_save+load_spec input is not None, so meta is still created
+        self.assertIsNotNone(meta)
+        self.assertIsNone(meta.load_spec)
+        self.assertFalse(meta.can_save)
+
+    def test_from_request_tracker_partial_tokens_discarded(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=20,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=0,
+        )
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, discard_partial_chunks=True)
+        self.assertIsNotNone(meta)
+        self.assertEqual(meta.token_len_chunk, 16)
+
+    def test_from_request_tracker_no_discard(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=20,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=0,
+        )
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16, discard_partial_chunks=False)
+        self.assertIsNotNone(meta)
+        self.assertEqual(meta.token_len_chunk, 20)
+
+    def test_from_request_tracker_already_saved(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=32,
+        )
+        meta = ReqMeta.from_request_tracker(tracker, cache_transfer_granularity=16)
+
+        # num_saved_tokens=32, chunk_boundary=ceil(33/16)*16=48 > 32
+        # so skip_save, and no load_spec => None
+        self.assertIsNone(meta)
+
+    def test_from_request_tracker_with_original_block_size(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids=[0, 1],
+            num_saved_tokens=0,
+        )
+        # Provide block_hashes (2 full blocks for token_len=32 / granularity=16)
+        # so the boundary_without_hash short-circuit does not zero out the save
+        # length and skip; this exercises the original_block_size propagation.
+        meta = ReqMeta.from_request_tracker(
+            tracker,
+            cache_transfer_granularity=16,
+            original_block_size=8,
+            block_hashes=[b"h0", b"h1"],
+        )
+        self.assertIsNotNone(meta)
+        self.assertEqual(meta.original_block_size, 8)
+
+
+class TestAscendConnectorMetadata(unittest.TestCase):
+    def test_add_request(self):
+        meta = AscendConnectorMetadata(unfinished_request_ids=set(), preempted_req_ids=set())
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=16,
+            block_ids=[0],
+            block_hashes=[],
+        )
+        meta.add_request(req)
+        self.assertEqual(len(meta.requests), 1)
+        self.assertEqual(meta.requests[0].req_id, "r1")
+
+
+class TestLayerMultiBlockReqMeta(unittest.TestCase):
+    def test_fields(self):
+        meta = LayerMultiBlockReqMeta(
+            req_id="r1",
+            keys=[],
+            starts=[0, 16],
+            ends=[16, 32],
+            block_ids=[0, 1],
+            layer_id=2,
+        )
+        self.assertEqual(meta.req_id, "r1")
+        self.assertEqual(meta.layer_id, 2)
+        self.assertTrue(meta.is_last_chunk)
+        self.assertIsNone(meta.current_event)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/ascend_store/test_coordinator.py

@@ -0,0 +1,291 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import unittest
+from dataclasses import dataclass, replace
+from unittest.mock import patch
+
+# isort: off
+import tests.ut.distributed.ascend_store._mock_deps  # noqa: F401, E402
+import torch
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec, SlidingWindowSpec
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store import config_data
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    get_block_hashes,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.coordinator import (
+    AscendStoreCoordinator,
+    ExternalCachedBlockPool,
+)
+
+# isort: on
+
+
+def _hashes(num_blocks: int) -> list[bytes]:
+    return [bytes([idx % 251]) * 32 for idx in range(num_blocks)]
+
+
+def _full_spec(block_size: int) -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+    )
+
+
+def _sliding_spec(block_size: int, sliding_window: int) -> SlidingWindowSpec:
+    return SlidingWindowSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        sliding_window=sliding_window,
+    )
+
+
+@dataclass(frozen=True)
+class _FakeCompressedSpec:
+    block_size: int
+    compress_ratio: int
+
+    def copy_with_new_block_size(self, block_size):
+        return replace(self, block_size=block_size)
+
+
+class _FakeCompressedManager:
+    @classmethod
+    def find_longest_cache_hit(
+        cls,
+        block_hashes,
+        max_length,
+        kv_cache_group_ids,
+        block_pool,
+        kv_cache_spec,
+        drop_eagle_block=False,
+        alignment_tokens=16,
+        **kwargs,
+    ):
+        computed: tuple[list[object], ...] = tuple([] for _ in kv_cache_group_ids)
+        logical_block_size = kv_cache_spec.block_size * kv_cache_spec.compress_ratio
+        max_blocks = max_length // logical_block_size
+        for block_hash in list(block_hashes)[:max_blocks]:
+            cached = block_pool.get_cached_block(block_hash, kv_cache_group_ids)
+            if not cached:
+                break
+            for blocks, block in zip(computed, cached):
+                blocks.append(block)
+        return computed
+
+
+class TestAscendStoreCoordinator(unittest.TestCase):
+    def test_load_mask_grouped_hashes_are_reused_by_key_build(self):
+        block_hashes = _hashes(4)
+        coord = AscendStoreCoordinator(
+            [KVCacheGroupSpec(["layer.0"], _full_spec(16))],
+            scheduler_block_size=16,
+            hash_block_size=8,
+            group_block_sizes=[16],
+            group_cache_families=["c1"],
+        )
+        db = ChunkedTokenDatabase(
+            [KeyMetadata("model", 0, 0, 0, 0)],
+            block_size=[16],
+            partitions=None,
+            hash_block_size=8,
+        )
+        db.cache_coordinator = coord
+        grouped_hash_cache: config_data.GroupedBlockHashCache = {}
+
+        with patch.object(
+            config_data,
+            "_rehash_block_hash_group",
+            wraps=config_data._rehash_block_hash_group,
+        ) as rehash:
+            self.assertEqual(
+                db.load_mask(
+                    block_hashes,
+                    32,
+                    grouped_hash_cache=grouped_hash_cache,
+                ),
+                ([True, True],),
+            )
+            self.assertEqual(rehash.call_count, 2)
+
+            keys = list(
+                db.process_token_key_strings_with_block_ids(
+                    32,
+                    block_hashes,
+                    [10, 11],
+                    grouped_hash_cache=grouped_hash_cache,
+                )
+            )
+
+        self.assertEqual(len(keys), 2)
+        self.assertEqual(rehash.call_count, 2)
+
+    def test_compressed_group_hits_on_effective_granularity(self):
+        block_hashes = _hashes(128)
+        grouped_hash = get_block_hashes(block_hashes, group_block_size=128 * 128, hash_block_size=128)[0]
+        coord = AscendStoreCoordinator(
+            [KVCacheGroupSpec(["layer.0"], _full_spec(128))],
+            scheduler_block_size=128 * 128,
+            hash_block_size=128,
+            group_block_sizes=[128],
+            group_cache_families=["c128"],
+        )
+
+        _, hit_length = coord.find_longest_cache_hit(
+            block_hashes,
+            128 * 128,
+            ExternalCachedBlockPool({(0, bytes(grouped_hash))}),
+        )
+
+        self.assertEqual(hit_length, 128 * 128)
+
+    def test_compressed_spec_does_not_apply_ratio_twice(self):
+        block_hashes = _hashes(128)
+        grouped_hash = get_block_hashes(block_hashes, group_block_size=128 * 128, hash_block_size=128)[0]
+
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.coordinator._get_manager_class",
+            return_value=_FakeCompressedManager,
+        ):
+            coord = AscendStoreCoordinator(
+                [KVCacheGroupSpec(["layer.0"], _FakeCompressedSpec(block_size=128, compress_ratio=128))],
+                scheduler_block_size=128 * 128,
+                hash_block_size=128,
+                group_block_sizes=[128],
+                group_cache_families=["c128"],
+            )
+
+            _, hit_length = coord.find_longest_cache_hit(
+                block_hashes,
+                128 * 128,
+                ExternalCachedBlockPool({(0, bytes(grouped_hash))}),
+            )
+
+        self.assertEqual(coord.group_effective_specs[0].compress_ratio, 1)
+        self.assertEqual(hit_length, 128 * 128)
+
+    def test_missing_required_group_returns_zero(self):
+        block_hashes = _hashes(128)
+        c1_exists = {(0, block_hash) for block_hash in block_hashes}
+        coord = AscendStoreCoordinator(
+            [
+                KVCacheGroupSpec(["layer.0"], _full_spec(128)),
+                KVCacheGroupSpec(["layer.1"], _full_spec(128)),
+            ],
+            scheduler_block_size=128 * 128,
+            hash_block_size=128,
+            group_block_sizes=[128, 128],
+            group_cache_families=["c1", "c128"],
+        )
+
+        _, hit_length = coord.find_longest_cache_hit(
+            block_hashes,
+            128 * 128,
+            ExternalCachedBlockPool(c1_exists),
+        )
+
+        self.assertEqual(hit_length, 0)
+
+    def test_store_mask_uses_manager_reachability(self):
+        coord = AscendStoreCoordinator(
+            [KVCacheGroupSpec(["layer.0"], _sliding_spec(block_size=128, sliding_window=256))],
+            scheduler_block_size=512,
+            hash_block_size=128,
+            group_block_sizes=[128],
+            group_cache_families=["c1"],
+        )
+
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.coordinator._reachable_block_mask",
+            return_value=[False, False, False, True],
+        ):
+            masks = coord.store_mask(512)
+
+        self.assertEqual(masks, ([False, False, False, True],))
+
+    def test_lookup_mask_uses_reachability_without_retention(self):
+        coord = AscendStoreCoordinator(
+            [KVCacheGroupSpec(["layer.0"], _sliding_spec(block_size=128, sliding_window=256))],
+            scheduler_block_size=512,
+            hash_block_size=128,
+            group_block_sizes=[128],
+            group_cache_families=["c1"],
+            retention_interval=256,
+        )
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.coordinator._reachable_block_mask",
+            return_value=[False, False, False, True],
+        ) as reachable:
+            masks = coord.lookup_mask(512)
+
+        self.assertEqual(masks, ([False, False, False, True],))
+        self.assertIsNone(reachable.call_args.kwargs["retention_interval"])
+
+    def test_store_mask_propagates_eagle_to_same_spec_siblings(self):
+        calls = []
+
+        def fake_reachable_block_mask(*args, **kwargs):
+            calls.append(kwargs["use_eagle"])
+            return [True, False, True, False]
+
+        shared_spec = _sliding_spec(block_size=128, sliding_window=256)
+        coord = AscendStoreCoordinator(
+            [
+                KVCacheGroupSpec(["layer.0"], shared_spec),
+                KVCacheGroupSpec(["layer.mtp"], shared_spec, is_eagle_group=True),
+            ],
+            scheduler_block_size=512,
+            hash_block_size=128,
+            group_block_sizes=[128, 128],
+            group_cache_families=["c1", "c1"],
+        )
+
+        with patch(
+            "vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.coordinator._reachable_block_mask",
+            side_effect=fake_reachable_block_mask,
+        ):
+            masks = coord.store_mask(512)
+
+        self.assertEqual(calls, [True, True])
+        self.assertEqual(masks, ([True, False, True, False], [True, False, True, False]))
+
+    def test_compressed_masks_stay_unmasked(self):
+        coord = AscendStoreCoordinator(
+            [KVCacheGroupSpec(["layer.0"], _sliding_spec(block_size=128, sliding_window=512))],
+            scheduler_block_size=2048,
+            hash_block_size=128,
+            group_block_sizes=[128],
+            group_cache_families=["c4"],
+        )
+
+        self.assertEqual(coord.store_mask(2048, num_prompt_tokens=2048), ([True] * 4,))
+        with patch.object(
+            coord,
+            "find_longest_cache_hit",
+            return_value=(([False, False, False, True],), 2048),
+        ):
+            self.assertEqual(coord.load_mask(_hashes(16), 2048), ([True] * 4,))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/ascend_store/test_kv_transfer.py

@@ -0,0 +1,664 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
+import threading
+import unittest
+from unittest.mock import MagicMock, patch
+
+# isort: off
+import tests.ut.distributed.ascend_store._mock_deps  # noqa: F401, E402
+from vllm.distributed.kv_events import BlockStored
+from vllm.v1.core.kv_cache_utils import maybe_convert_block_hash
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store import config_data
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    LayerMultiBlockReqMeta,
+    LayerPoolKey,
+    LoadSpec,
+    ReqMeta,
+)
+
+# isort: on
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer import (
+    KVCacheStoreLayerRecvingThread,
+    KVCacheStoreLayerSendingThread,
+    KVCacheStoreRecvingThread,
+    KVCacheStoreSendingThread,
+    KVTransferThread,
+)
+
+
+class FakeStore:
+    def __init__(self, exists_result=None):
+        self.exists_result = exists_result or []
+        self.put_calls = []
+        self.get_calls = []
+
+    def set_device(self):
+        pass
+
+    def exists(self, keys):
+        return self.exists_result[: len(keys)]
+
+    def put(self, keys, addrs, sizes):
+        self.put_calls.append((list(keys), list(addrs), list(sizes)))
+
+    def get(self, keys, addrs, sizes):
+        self.get_calls.append((list(keys), list(addrs), list(sizes)))
+
+
+class FakeTokenDatabase(ChunkedTokenDatabase):
+    def __init__(self, block_size=16):
+        super().__init__([KeyMetadata("m", 0, 0, 0, 0)], [block_size], None)
+        self.set_group_buffers({0: [1000]}, {0: [block_size]}, {0: [1]}, group_num_layers={0: 1})
+
+
+class MaskedFakeTokenDatabase(FakeTokenDatabase):
+    def __init__(self, block_size=16, masks=([True],)):
+        super().__init__(block_size)
+        self.masks = masks
+
+    def store_mask(self, token_len, num_prompt_tokens=None):
+        return self.masks
+
+    def load_mask(self, block_hashes, token_len, grouped_hash_cache=None):
+        return self.masks
+
+    def mask_allows_chunk(self, masks, kv_cache_group_id, start):
+        if masks is None:
+            return True
+        block_idx = start // self.get_block_size(kv_cache_group_id)
+        return block_idx < len(masks[kv_cache_group_id]) and masks[kv_cache_group_id][block_idx]
+
+
+class TestKVTransferThread(unittest.TestCase):
+    def _make_thread(self, exists_result=None):
+        store = FakeStore(exists_result or [])
+        db = FakeTokenDatabase()
+        t = KVTransferThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            ready_event=threading.Event(),
+            name="test",
+        )
+        return t, store
+
+    def test_add_request(self):
+        t, _ = self._make_thread()
+        req = MagicMock()
+        t.add_request(req)
+        self.assertFalse(t.request_queue.empty())
+
+    def test_get_and_clear_finished_requests(self):
+        t, _ = self._make_thread()
+        t.set_finished_request("r1")
+        t.set_finished_request("r2")
+        finished = t.get_and_clear_finished_requests()
+        self.assertEqual(finished, {"r1", "r2"})
+        self.assertEqual(t.get_and_clear_finished_requests(), set())
+
+    def test_lookup_all_exist(self):
+        t, _ = self._make_thread([1, 1, 1])
+        result = t.lookup(["k1", "k2", "k3"])
+        self.assertEqual(result, [True, True, True])
+
+    def test_lookup_partial(self):
+        t, _ = self._make_thread([1, 0, 1])
+        result = t.lookup(["k1", "k2", "k3"])
+        self.assertEqual(result, [True, False, True])
+
+    def test_lookup_exception(self):
+        t, store = self._make_thread()
+        store.exists = MagicMock(side_effect=Exception("conn fail"))
+        result = t.lookup(["k1"])
+        self.assertEqual(result, [False])
+
+    def test_update_and_get_kv_events(self):
+        t, _ = self._make_thread()
+        event1 = BlockStored(
+            block_hashes=["h1"],
+            parent_block_hash=None,
+            token_ids=[1, 2, 3],
+            block_size=16,
+            lora_id=None,
+            medium="cpu",
+            lora_name=None,
+        )
+        event2 = BlockStored(
+            block_hashes=["h2"],
+            parent_block_hash="h1",
+            token_ids=[4, 5, 6],
+            block_size=16,
+            lora_id=None,
+            medium="cpu",
+            lora_name=None,
+        )
+        t.update_kv_event([event1, event2])
+        events = t.get_kv_events()
+        self.assertEqual(len(events), 2)
+        # After get, events should be cleared
+        self.assertEqual(len(t.get_kv_events()), 0)
+
+    def test_handle_request_base_noop(self):
+        t, _ = self._make_thread()
+        # Base class _handle_request does nothing
+        t._handle_request(MagicMock())
+
+
+class TestKVCacheStoreSendingThread(unittest.TestCase):
+    def _make_thread(self, exists_result=None, kv_role="kv_producer", enable_kv_event=False):
+        store = FakeStore(exists_result or [0, 0, 0, 0])
+        db = FakeTokenDatabase()
+        t = KVCacheStoreSendingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            put_step=1,
+            kv_role=kv_role,
+            ready_event=threading.Event(),
+            group_uses_align_state=[False],
+            enable_kv_event=enable_kv_event,
+        )
+        return t, store
+
+    def test_handle_request_puts_missing_keys(self):
+        t, store = self._make_thread([1, 0, 1, 0])
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=64,
+            block_ids=[0, 1, 2, 3],
+            block_hashes=[b"h0", b"h1", b"h2", b"h3"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.put_calls), 1)
+        keys, _, _ = store.put_calls[0]
+        self.assertEqual(len(keys), 2)
+
+    def test_handle_request_all_exist_no_put(self):
+        t, store = self._make_thread([1, 1])
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=32,
+            block_ids=[0, 1],
+            block_hashes=[b"h0", b"h1"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.put_calls), 0)
+
+    def test_handle_request_not_in_stored(self):
+        t, store = self._make_thread([0])
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=16,
+            block_ids=[0],
+            block_hashes=[b"h0"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.put_calls), 0)
+
+    def test_handle_request_with_kv_event(self):
+        t, store = self._make_thread([1, 0, 1], enable_kv_event=True)
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=48,
+            block_ids=[0, 1, 2],
+            block_hashes=[b"h0", b"h1", b"h2"],  # type: ignore[arg-type]
+            current_event=None,
+            token_ids=list(range(48)),
+            original_block_size=16,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        events = t.get_kv_events()
+        self.assertEqual(len(events), 1)
+        self.assertEqual(events[0].block_hashes, [maybe_convert_block_hash(b"h1")])
+        self.assertEqual(events[0].parent_block_hash, maybe_convert_block_hash(b"h0"))
+
+    def test_save_reuses_grouped_hashes_for_kv_events(self):
+        store = FakeStore([0, 0])
+        db = ChunkedTokenDatabase(
+            [KeyMetadata("m", 0, 0, 0, 0)],
+            [16],
+            None,
+            hash_block_size=8,
+        )
+        db.set_group_buffers({0: [1000]}, {0: [16]}, {0: [1]}, group_num_layers={0: 1})
+        thread = KVCacheStoreSendingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            put_step=1,
+            kv_role="kv_producer",
+            ready_event=threading.Event(),
+            group_uses_align_state=[False],
+            enable_kv_event=True,
+        )
+        request = ReqMeta(
+            req_id="r1",
+            token_len_chunk=32,
+            block_ids=[0, 1],
+            block_hashes=[b"h0", b"h1", b"h2", b"h3"],  # type: ignore[arg-type]
+            current_event=None,
+            token_ids=list(range(32)),
+            original_block_size=16,
+        )
+        thread.add_stored_request("r1")
+        thread.request_queue.put(request)
+
+        with patch.object(
+            config_data,
+            "_rehash_block_hash_group",
+            wraps=config_data._rehash_block_hash_group,
+        ) as rehash:
+            thread._handle_request(request)
+
+        self.assertEqual(rehash.call_count, 2)
+        self.assertEqual(len(thread.get_kv_events()), 2)
+
+    def test_handle_request_consumer_role(self):
+        t, store = self._make_thread([0], kv_role="kv_consumer")
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=16,
+            block_ids=[0],
+            block_hashes=[b"h0"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.put_calls), 1)
+
+    def test_add_dec_delete_stored_request(self):
+        t, _ = self._make_thread()
+        t.add_stored_request("r1")
+        t.add_stored_request("r1")
+        self.assertEqual(t.stored_requests["r1"], 2)
+        t.dec_stored_request("r1")
+        self.assertEqual(t.stored_requests["r1"], 1)
+        t.delete_finished_stored_request("r1")
+        self.assertNotIn("r1", t.stored_requests)
+
+    def test_dec_nonexistent_request(self):
+        t, _ = self._make_thread()
+        t.dec_stored_request("nonexist")  # should not raise
+
+    def test_delete_nonexistent_request(self):
+        t, _ = self._make_thread()
+        t.delete_finished_stored_request("nonexist")  # should not raise
+
+    def test_handle_request_with_current_event(self):
+        t, store = self._make_thread([0])
+        event = MagicMock()
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=16,
+            block_ids=[0],
+            block_hashes=[b"h0"],  # type: ignore[arg-type]
+            current_event=event,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        event.synchronize.assert_called_once()
+
+    def test_handle_request_dcp_size_gt_1(self):
+        store = FakeStore([0, 0])
+        db = FakeTokenDatabase()
+        t = KVCacheStoreSendingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=2,
+            put_step=1,
+            kv_role="kv_producer",
+            ready_event=threading.Event(),
+            group_uses_align_state=[False],
+        )
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=32,
+            block_ids=[0, 1],
+            block_hashes=[b"h0", b"h1"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        # dcp_size > 1 means no slicing
+        self.assertEqual(len(store.put_calls), 1)
+
+    def test_handle_request_applies_store_mask(self):
+        store = FakeStore([0, 0])
+        db = MaskedFakeTokenDatabase(masks=([True, False],))
+        t = KVCacheStoreSendingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            put_step=1,
+            kv_role="kv_producer",
+            ready_event=threading.Event(),
+            group_uses_align_state=[False],
+        )
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=32,
+            block_ids=[0, 1],
+            block_hashes=[b"h0", b"h1"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        keys, _, _ = store.put_calls[0]
+        self.assertEqual(len(keys), 1)
+
+    def test_handle_request_skips_compressed_hit_in_raw_token_domain(self):
+        t, store = self._make_thread([0, 0])
+        t.token_database.group_cache_families["kv"][0] = "c4"
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=128,
+            block_ids=[0, 1],
+            block_hashes=[f"h{i}" for i in range(8)],
+            load_spec=LoadSpec(
+                vllm_cached_tokens=0,
+                kvpool_cached_tokens=63,
+                kvpool_store_skip_tokens=64,
+                can_load=True,
+            ),
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        keys, addrs, _ = store.put_calls[0]
+        self.assertEqual(len(keys), 1)
+        self.assertEqual(addrs, [[1001]])
+
+    def test_save_exception_cleans_queue_lifecycle(self):
+        t, store = self._make_thread([0])
+        store.put = MagicMock(side_effect=RuntimeError("put failed"))
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=16,
+            block_ids=[0],
+            block_hashes=[b"h0"],  # type: ignore[arg-type]
+        )
+        t.add_stored_request("r1")
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(t.request_queue.unfinished_tasks, 0)
+        self.assertNotIn("r1", t.stored_requests)
+
+
+class TestKVCacheStoreRecvingThread(unittest.TestCase):
+    def test_handle_request(self):
+        store = FakeStore()
+        db = FakeTokenDatabase()
+        t = KVCacheStoreRecvingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            ready_event=threading.Event(),
+            invalid_block_ids=set(),
+            invalid_block_ids_lock=threading.Lock(),
+        )
+        load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=32, can_load=True, token_len=32)
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=32,
+            block_ids=[0, 1],
+            block_hashes=[b"h0", b"h1"],  # type: ignore[arg-type]
+            load_spec=load_spec,
+        )
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.get_calls), 1)
+        finished = t.get_and_clear_finished_requests()
+        self.assertIn("r1", finished)
+
+    def test_handle_request_applies_load_mask(self):
+        store = FakeStore()
+        db = MaskedFakeTokenDatabase(masks=([True, False],))
+        t = KVCacheStoreRecvingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            ready_event=threading.Event(),
+            invalid_block_ids=set(),
+            invalid_block_ids_lock=threading.Lock(),
+        )
+        load_spec = LoadSpec(vllm_cached_tokens=0, kvpool_cached_tokens=32, can_load=True, token_len=32)
+        req = ReqMeta(
+            req_id="r1",
+            token_len_chunk=32,
+            block_ids=[0, 1],
+            block_hashes=[b"h0", b"h1"],  # type: ignore[arg-type]
+            load_spec=load_spec,
+        )
+        t.request_queue.put(req)
+        t._handle_request(req)
+        keys, _, _ = store.get_calls[0]
+        self.assertEqual(len(keys), 1)
+
+
+@unittest.skip("LayerMultiBlockReqMeta API is deprecated, tests need update for LayerTransferTask")
+class TestKVCacheStoreLayerSendingThread(unittest.TestCase):
+    def _make_thread(self, exists_result=None, num_layers=2):
+        store = FakeStore(exists_result or [0, 0])
+        db = FakeTokenDatabase()
+        t = KVCacheStoreLayerSendingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            tp_size=1,
+            dcp_size=1,
+            put_step=1,
+            my_key_index=0,
+            num_ranks_per_layer=1,
+            page_size_bytes=32,
+            ready_event=threading.Event(),
+            num_layers=num_layers,
+            layer_save_finished_events=[threading.Event() for _ in range(num_layers)],
+            sync_save_events=[],
+        )
+        return t, store
+
+    def _make_layer_req(self, layer_id=0, is_last_chunk=False, num_keys=2):
+        meta = KeyMetadata("m", 0, 0, 0, 0)
+        keys = [LayerPoolKey(meta, f"h{i}", layer_id) for i in range(num_keys)]
+        return LayerMultiBlockReqMeta(
+            req_id="r1",
+            keys=keys,
+            starts=[i * 16 for i in range(num_keys)],
+            ends=[(i + 1) * 16 for i in range(num_keys)],
+            block_ids=list(range(num_keys)),
+            layer_id=layer_id,
+            is_last_chunk=is_last_chunk,
+            current_event=None,
+            token_ids=list(range(num_keys * 16)),
+            original_block_size=16,
+            block_hashes=[f"h{i}".encode() for i in range(num_keys)],
+        )
+
+    def test_handle_request_puts_missing(self):
+        t, store = self._make_thread([1, 0])
+        req = self._make_layer_req(layer_id=0)
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.put_calls), 1)
+        keys, _, _ = store.put_calls[0]
+        self.assertEqual(len(keys), 1)
+
+    def test_handle_request_all_exist_not_last(self):
+        t, store = self._make_thread([1, 1])
+        req = self._make_layer_req(layer_id=0, is_last_chunk=False)
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.put_calls), 0)
+
+    def test_handle_request_all_exist_last_chunk_final_layer(self):
+        t, store = self._make_thread([1, 1], num_layers=2)
+        req = self._make_layer_req(layer_id=1, is_last_chunk=True)
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        finished = t.get_and_clear_finished_requests()
+        self.assertIn("r1", finished)
+
+    def test_handle_request_empty_keys(self):
+        t, store = self._make_thread()
+        _meta = KeyMetadata("m", 0, 0, 0, 0)
+        req = LayerMultiBlockReqMeta(
+            req_id="r1",
+            keys=[],
+            starts=[],
+            ends=[],
+            block_ids=[],
+            layer_id=0,
+            is_last_chunk=True,
+        )
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        finished = t.get_and_clear_finished_requests()
+        self.assertNotIn("r1", finished)
+
+    def test_handle_request_with_current_event(self):
+        t, store = self._make_thread([0])
+        event = MagicMock()
+        meta = KeyMetadata("m", 0, 0, 0, 0)
+        req = LayerMultiBlockReqMeta(
+            req_id="r1",
+            keys=[LayerPoolKey(meta, "h0", 0)],
+            starts=[0],
+            ends=[16],
+            block_ids=[0],
+            layer_id=0,
+            is_last_chunk=False,
+            current_event=event,
+        )
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        event.synchronize.assert_called_once()
+
+    def test_handle_request_last_chunk_final_layer_with_missing(self):
+        t, store = self._make_thread([0], num_layers=2)
+        req = self._make_layer_req(layer_id=1, is_last_chunk=True, num_keys=1)
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        finished = t.get_and_clear_finished_requests()
+        self.assertIn("r1", finished)
+
+    def test_layerwise_kv_event_published_on_final_layer(self):
+        t, store = self._make_thread([0], num_layers=2)
+        req = self._make_layer_req(layer_id=1, is_last_chunk=True, num_keys=1)
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        events = t.get_kv_events()
+        self.assertEqual(len(events), 1)
+        self.assertEqual(events[0].block_hashes, [maybe_convert_block_hash(b"h0")])
+        self.assertEqual(events[0].token_ids, list(range(16)))
+        self.assertEqual(events[0].block_size, 16)
+
+    def test_layerwise_kv_event_not_published_before_final_layer(self):
+        t, store = self._make_thread([0], num_layers=2)
+        req = self._make_layer_req(layer_id=0, is_last_chunk=False, num_keys=1)
+        t.add_stored_request(req.req_id)
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(t.get_kv_events(), [])
+
+    def test_layerwise_kv_event_uses_missing_blocks_from_previous_layers(self):
+        t, store = self._make_thread([0], num_layers=2)
+        first_layer_req = self._make_layer_req(layer_id=0, is_last_chunk=True, num_keys=1)
+        t.add_stored_request(first_layer_req.req_id)
+        t.request_queue.put(first_layer_req)
+        t._handle_request(first_layer_req)
+        t.m_store.exists_result = [1]
+        final_layer_req = self._make_layer_req(layer_id=1, is_last_chunk=True, num_keys=1)
+        t.request_queue.put(final_layer_req)
+        t._handle_request(final_layer_req)
+        events = t.get_kv_events()
+        self.assertEqual(len(events), 1)
+        self.assertEqual(events[0].block_hashes, [maybe_convert_block_hash(b"h0")])
+
+
+@unittest.skip("LayerMultiBlockReqMeta API is deprecated, tests need update for LayerTransferTask")
+class TestKVCacheStoreLayerRecvingThread(unittest.TestCase):
+    def test_handle_request(self):
+        store = FakeStore()
+        db = FakeTokenDatabase()
+        get_event = threading.Event()
+        t = KVCacheStoreLayerRecvingThread(
+            m_store=store,
+            token_database=db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            ready_event=threading.Event(),
+            get_event=get_event,
+            invalid_block_ids=set(),
+            invalid_block_ids_lock=threading.Lock(),
+        )
+        meta = KeyMetadata("m", 0, 0, 0, 0)
+        req = LayerMultiBlockReqMeta(
+            req_id="r1",
+            keys=[LayerPoolKey(meta, "h0", 0)],
+            starts=[0],
+            ends=[16],
+            block_ids=[0],
+            layer_id=0,
+        )
+        t.request_queue.put(req)
+        t._handle_request(req)
+        self.assertEqual(len(store.get_calls), 1)
+        self.assertTrue(get_event.is_set())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/device_communicators/test_pyhccl.py

@@ -4,8 +4,7 @@ from unittest.mock import MagicMock, patch
 from vllm.distributed.utils import StatelessProcessGroup
 
 from tests.ut.base import TestBase
-from vllm_ascend.distributed.device_communicators.pyhccl import \
-    PyHcclCommunicator
+from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator
 
 
 class MockHcclLib:
@@ -17,7 +16,6 @@ class MockUniqueId:
 
 
 class TestPyHcclCommunicator(TestBase):
-
     @patch.dict(os.environ, {"RANK": "0", "WORLD_SIZE": "1"})
     def test_world_size_1_return_early(self):
         comm = PyHcclCommunicator(
@@ -29,26 +27,17 @@ class TestPyHcclCommunicator(TestBase):
 
     @patch.dict(os.environ, {"RANK": "0", "WORLD_SIZE": "2"})
     def test_load_hccl_fail(self):
-        comm = PyHcclCommunicator(group=StatelessProcessGroup(
-            0, 2, None, None),
-                                  device="npu:0",
-                                  library_path="/not/exist/path/libhccl.so")
+        comm = PyHcclCommunicator(
+            group=StatelessProcessGroup(0, 2, None, None), device="npu:0", library_path="/not/exist/path/libhccl.so"
+        )
         self.assertTrue(comm.disabled)
 
-    @patch(
-        "vllm_ascend.distributed.device_communicators.pyhccl_wrapper.HCCLLibrary",
-        MockHcclLib)
-    @patch(
-        "vllm_ascend.distributed.device_communicators.pyhccl_wrapper.hcclUniqueId",
-        MockUniqueId)
+    @patch("vllm_ascend.distributed.device_communicators.pyhccl_wrapper.HCCLLibrary", MockHcclLib)
+    @patch("vllm_ascend.distributed.device_communicators.pyhccl_wrapper.hcclUniqueId", MockUniqueId)
     @patch("torch.npu.device")
-    @patch("vllm_ascend.utils.current_stream",
-           return_value=MagicMock(npu_stream=5678))
+    @patch("vllm_ascend.utils.current_stream", return_value=MagicMock(npu_stream=5678))
     def test_stateless_group(self, *_):
-        group = StatelessProcessGroup(rank=3,
-                                      world_size=4,
-                                      store=None,
-                                      socket=None)
+        group = StatelessProcessGroup(rank=3, world_size=4, store=None)
 
         comm = PyHcclCommunicator(group=group, device=3)
 
@@ -56,21 +45,17 @@ class TestPyHcclCommunicator(TestBase):
         self.assertEqual(comm.world_size, 4)
 
     @patch.dict(os.environ, {"RANK": "1", "WORLD_SIZE": "2"})
-    @patch(
-        "vllm_ascend.distributed.device_communicators.pyhccl_wrapper.HCCLLibrary",
-        MockHcclLib)
-    @patch(
-        "vllm_ascend.distributed.device_communicators.pyhccl_wrapper.hcclUniqueId",
-        MockUniqueId)
+    @patch("vllm_ascend.distributed.device_communicators.pyhccl_wrapper.HCCLLibrary", MockHcclLib)
+    @patch("vllm_ascend.distributed.device_communicators.pyhccl_wrapper.hcclUniqueId", MockUniqueId)
     @patch("torch.distributed.is_initialized", return_value=True)
     @patch("torch.distributed.get_backend", return_value="nccl")
+    @patch("torch.distributed.Backend.HCCL", "hccl", create=True)
     @patch("torch.distributed.get_rank", return_value=1)
     @patch("torch.distributed.get_world_size", return_value=2)
     @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1])
     @patch("torch.distributed.broadcast")
     @patch("torch.npu.device")
-    @patch("vllm_ascend.utils.current_stream",
-           return_value=MagicMock(npu_stream=1234))
+    @patch("vllm_ascend.utils.current_stream", return_value=MagicMock(npu_stream=1234))
     def test_multi_gpu_pg_torch(
         self,
         *_,

tests/ut/distributed/device_communicators/test_pyhccl_wrapper.py

@@ -5,13 +5,21 @@ from torch.distributed import ReduceOp
 
 from tests.ut.base import TestBase
 from vllm_ascend.distributed.device_communicators.pyhccl_wrapper import (
-    Function, HCCLLibrary, aclrtStream_t, buffer_type, hcclComm_t,
-    hcclDataType_t, hcclDataTypeEnum, hcclRedOp_t, hcclRedOpTypeEnum,
-    hcclResult_t, hcclUniqueId)
+    Function,
+    HCCLLibrary,
+    aclrtStream_t,
+    buffer_type,
+    hcclComm_t,
+    hcclDataType_t,
+    hcclDataTypeEnum,
+    hcclRedOp_t,
+    hcclRedOpTypeEnum,
+    hcclResult_t,
+    hcclUniqueId,
+)
 
 
 class TestHcclUniqueId(TestBase):
-
     def test_construct(self):
         uid = hcclUniqueId()
         uid.internal[0] = 12
@@ -20,7 +28,6 @@ class TestHcclUniqueId(TestBase):
 
 
 class TestHcclDataTypeEnum(TestBase):
-
     def test_torch_dtype_mapping(self):
         expected = {
             torch.int8: hcclDataTypeEnum.hcclInt8,
@@ -35,8 +42,7 @@ class TestHcclDataTypeEnum(TestBase):
 
         for torch_dtype, expected_enum in expected.items():
             with self.subTest(torch_dtype=torch_dtype):
-                self.assertEqual(hcclDataTypeEnum.from_torch(torch_dtype),
-                                 expected_enum)
+                self.assertEqual(hcclDataTypeEnum.from_torch(torch_dtype), expected_enum)
 
     def test_unsupported_dtype_raises(self):
         with self.assertRaises(ValueError):
@@ -44,7 +50,6 @@ class TestHcclDataTypeEnum(TestBase):
 
 
 class TestHcclRedOpTypeEnum(TestBase):
-
     def test_torch_reduce_op_mapping(self):
         expected = {
             ReduceOp.SUM: hcclRedOpTypeEnum.hcclSum,
@@ -55,8 +60,7 @@ class TestHcclRedOpTypeEnum(TestBase):
 
         for torch_op, expected_enum in expected.items():
             with self.subTest(torch_op=torch_op):
-                self.assertEqual(hcclRedOpTypeEnum.from_torch(torch_op),
-                                 expected_enum)
+                self.assertEqual(hcclRedOpTypeEnum.from_torch(torch_op), expected_enum)
 
     def test_unsupported_op_raises(self):
         unsupported_op = "NOT_EXIST"
@@ -65,7 +69,6 @@ class TestHcclRedOpTypeEnum(TestBase):
 
 
 class TestFunction(TestBase):
-
     def test_construct_with_valid_args(self):
         func = Function(name="foo", restype=int, argtypes=[int, str, float])
         self.assertEqual(func.name, "foo")
@@ -74,7 +77,6 @@ class TestFunction(TestBase):
 
 
 class TestHCLLLibrary(TestBase):
-
     def test_init_with_nonexistent_so(self):
         fake_path = "/definitely/not/exist/libhccl.so"
         with self.assertRaises(OSError):
@@ -127,7 +129,6 @@ class TestHCLLLibrary(TestBase):
 
     @patch.object(HCCLLibrary, "HCCL_CHECK")
     def test_hccl_all_reduce(self, mock_hccl_check):
-
         lib = HCCLLibrary.__new__(HCCLLibrary)
         lib._funcs = {"HcclAllReduce": MagicMock(return_value=0)}
         sendbuff = buffer_type()
@@ -138,16 +139,13 @@ class TestHCLLLibrary(TestBase):
         comm = hcclComm_t()
         stream = aclrtStream_t()
 
-        lib.hcclAllReduce(sendbuff, recvbuff, count, datatype, op, comm,
-                          stream)
+        lib.hcclAllReduce(sendbuff, recvbuff, count, datatype, op, comm, stream)
 
-        lib._funcs["HcclAllReduce"].assert_called_once_with(
-            sendbuff, recvbuff, count, datatype, op, comm, stream)
+        lib._funcs["HcclAllReduce"].assert_called_once_with(sendbuff, recvbuff, count, datatype, op, comm, stream)
         mock_hccl_check.assert_called_once_with(0)
 
     @patch.object(HCCLLibrary, "HCCL_CHECK")
     def test_hccl_broad_cast(self, mock_hccl_check):
-
         lib = HCCLLibrary.__new__(HCCLLibrary)
         lib._funcs = {"HcclBroadcast": MagicMock(return_value=0)}
         buff = buffer_type()
@@ -159,8 +157,7 @@ class TestHCLLLibrary(TestBase):
 
         lib.hcclBroadcast(buff, count, datatype, root, comm, stream)
 
-        lib._funcs["HcclBroadcast"].assert_called_once_with(
-            buff, count, datatype, root, comm, stream)
+        lib._funcs["HcclBroadcast"].assert_called_once_with(buff, count, datatype, root, comm, stream)
         mock_hccl_check.assert_called_once_with(0)
 
     @patch.object(HCCLLibrary, "HCCL_CHECK")

tests/ut/distributed/kv_transfer/test_kv_transfer_failures.py

@@ -0,0 +1,240 @@
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+"""Unit tests for KV transfer failure handling in ascend_store.
+
+This module tests the record_failed_blocks function which handles KV transfer
+failures by recording which blocks failed to load during the transfer process.
+"""
+
+import types
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+if not hasattr(torch, "npu"):
+    torch.npu = types.SimpleNamespace(Event=type("Event", (), {}))  # type: ignore[attr-defined]
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.ascend_store_connector import AscendStoreConnector
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer import record_failed_blocks
+
+
+class TestRecordFailedBlocks(unittest.TestCase):
+    """Test cases for the record_failed_blocks function.
+
+    The record_failed_blocks function takes a list of block IDs and their corresponding
+    return codes from a KV transfer operation, and returns a set of block IDs that failed
+    (i.e., those with non-zero return codes).
+    """
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_all_blocks_succeed(self, mock_logger: MagicMock):
+        """Test when all blocks are transferred successfully (all return codes are 0)."""
+        block_ids: list[int] = [1, 2, 3, 4, 5]
+        ret_codes: list[int] = [0, 0, 0, 0, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, set())
+        self.assertEqual(len(result), 0)
+        mock_logger.error.assert_not_called()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_all_blocks_fail(self, mock_logger: MagicMock):
+        """Test when all blocks fail to transfer (all return codes are non-zero)."""
+        block_ids: list[int] = [1, 2, 3, 4, 5]
+        ret_codes: list[int] = [1, 2, 3, 4, 5]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {1, 2, 3, 4, 5})
+        self.assertEqual(len(result), 5)
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_partial_blocks_fail(self, mock_logger: MagicMock):
+        """Test when some blocks fail and some succeed."""
+        block_ids: list[int] = [1, 2, 3, 4, 5]
+        ret_codes: list[int] = [0, 1, 0, 2, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {2, 4})
+        self.assertEqual(len(result), 2)
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_empty_lists(self, mock_logger: MagicMock):
+        """Test with empty block_ids and ret_codes."""
+        block_ids: list[int] = []
+        ret_codes: list[int] = []
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, set())
+        mock_logger.error.assert_not_called()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_single_block_succeed(self, mock_logger: MagicMock):
+        """Test with a single block that succeeds."""
+        block_ids: list[int] = [42]
+        ret_codes: list[int] = [0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, set())
+        mock_logger.error.assert_not_called()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_single_block_fail(self, mock_logger: MagicMock):
+        """Test with a single block that fails."""
+        block_ids: list[int] = [42]
+        ret_codes: list[int] = [1]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {42})
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_negative_return_codes(self, mock_logger: MagicMock):
+        """Test with negative return codes (error conditions)."""
+        block_ids: list[int] = [1, 2, 3]
+        ret_codes: list[int] = [0, -1, -2]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {2, 3})
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_large_block_ids(self, mock_logger: MagicMock):
+        """Test with large block ID values."""
+        block_ids: list[int] = [1000000, 2000000, 3000000]
+        ret_codes: list[int] = [0, 1, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {2000000})
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_mixed_error_codes(self, mock_logger: MagicMock):
+        """Test with various non-zero error codes."""
+        block_ids: list[int] = [10, 20, 30, 40, 50]
+        ret_codes: list[int] = [0, -1, 100, 0, 999]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {20, 30, 50})
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_logs_failed_blocks(self, mock_logger: MagicMock):
+        """Test that failed blocks are logged."""
+        block_ids: list[int] = [1, 2, 3]
+        ret_codes: list[int] = [0, 1, 2]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {2, 3})
+        mock_logger.error.assert_called_once()
+        call_args = mock_logger.error.call_args[0]
+        log_msg = call_args[0]
+        self.assertIn("Failed to load blocks", log_msg)
+        # The last argument is the failed blocks set
+        self.assertEqual(call_args[-1], {2, 3})
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_no_log_when_all_succeed(self, mock_logger: MagicMock):
+        """Test that no error is logged when all blocks succeed."""
+        block_ids: list[int] = [1, 2, 3]
+        ret_codes: list[int] = [0, 0, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, set())
+        mock_logger.error.assert_not_called()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_non_hybrid_single_block_semantics(self, mock_logger: MagicMock):
+        """Test non-hybrid callers still map one return code to one block."""
+        block_ids: list[int] = [10, 11, 12]
+        ret_codes: list[int] = [0, 1, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {11})
+        mock_logger.error.assert_called_once()
+
+
+class TestRecordFailedBlocksEdgeCases(unittest.TestCase):
+    """Additional edge case tests for record_failed_blocks."""
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_duplicate_block_ids_all_fail(self, mock_logger: MagicMock):
+        """Test with duplicate block IDs that all fail."""
+        # Note: This tests the behavior with duplicates
+        # The set will deduplicate, but all should be marked as failed
+        block_ids: list[int] = [1, 1, 2, 2]
+        ret_codes: list[int] = [1, 1, 2, 2]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        # Set deduplicates, so we get unique failed block IDs
+        self.assertEqual(result, {1, 2})
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_zero_block_id_with_failure(self, mock_logger: MagicMock):
+        """Test with block ID 0 failing."""
+        block_ids: list[int] = [0, 1, 2]
+        ret_codes: list[int] = [1, 0, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {0})
+        mock_logger.error.assert_called_once()
+
+    @patch("vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer.logger")
+    def test_consecutive_failures(self, mock_logger: MagicMock):
+        """Test with consecutive block failures."""
+        block_ids: list[int] = [100, 101, 102, 103, 104]
+        ret_codes: list[int] = [1, 1, 1, 0, 0]
+
+        result = record_failed_blocks(block_ids, ret_codes)
+
+        self.assertEqual(result, {100, 101, 102})
+        mock_logger.error.assert_called_once()
+
+
+class TestAscendStoreConnector(unittest.TestCase):
+    """Regression tests for connector-level load failure reporting."""
+
+    def test_get_block_ids_with_load_errors_forwards_to_worker(self):
+        connector = AscendStoreConnector.__new__(AscendStoreConnector)
+        connector.connector_worker = MagicMock()
+        connector.connector_worker.get_block_ids_with_load_errors.return_value = {3, 7}
+
+        result = connector.get_block_ids_with_load_errors()
+
+        self.assertEqual(result, {3, 7})
+        connector.connector_worker.get_block_ids_with_load_errors.assert_called_once_with()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/mooncake/test_mooncake_config_data.py

@@ -0,0 +1,74 @@
+import sys
+import types
+import unittest
+from unittest.mock import MagicMock
+
+fake_engine = types.ModuleType("mooncake.engine")
+fake_engine.TransferEngine = MagicMock()  # type: ignore[attr-defined]
+sys.modules["mooncake.engine"] = fake_engine
+fake_store = types.ModuleType("mooncake.store")
+fake_store.ReplicateConfig = MagicMock()  # type: ignore[attr-defined]
+sys.modules["mooncake.store"] = fake_store
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.backend.mooncake_backend import (  # noqa: E402
+    _convert_to_bytes,
+    _parse_global_segment_size,
+)
+
+
+class TestParseGlobalSegmentSize(unittest.TestCase):
+    def test_int_input(self):
+        self.assertEqual(_parse_global_segment_size(1024), 1024)
+        self.assertEqual(_parse_global_segment_size(0), 0)
+
+    def test_gb_unit(self):
+        self.assertEqual(_parse_global_segment_size("2GB"), 2 * 1024**3)
+        self.assertEqual(_parse_global_segment_size("1.5GB"), int(1.5 * 1024**3))
+        self.assertEqual(_parse_global_segment_size(" 2 GB "), 2 * 1024**3)
+
+    def test_gb_unit_edge_cases(self):
+        with self.assertRaises(ValueError):
+            _parse_global_segment_size("GB")
+        with self.assertRaises(ValueError):
+            _parse_global_segment_size("abcGB")
+
+    def test_mb_unit(self):
+        self.assertEqual(_parse_global_segment_size("512MB"), 512 * 1024**2)
+        self.assertEqual(_parse_global_segment_size("0.5MB"), int(0.5 * 1024**2))
+        self.assertEqual(_parse_global_segment_size("1024MB"), 1024 * 1024**2)
+
+    def test_kb_unit(self):
+        self.assertEqual(_parse_global_segment_size("256KB"), 256 * 1024)
+        self.assertEqual(_parse_global_segment_size("1.25KB"), int(1.25 * 1024))
+
+    def test_b_unit(self):
+        self.assertEqual(_parse_global_segment_size("4096B"), 4096)
+        self.assertEqual(_parse_global_segment_size("1024b"), 1024)
+
+    def test_no_unit(self):
+        self.assertEqual(_parse_global_segment_size("2048"), 2048)
+        self.assertEqual(_parse_global_segment_size("0"), 0)
+
+    def test_non_string_non_int_input(self):
+        self.assertEqual(_parse_global_segment_size(2048.0), 2048)
+        self.assertEqual(_parse_global_segment_size(True), 1)
+
+        with self.assertRaises(TypeError):
+            _parse_global_segment_size(None)
+
+        with self.assertRaises(TypeError):
+            _parse_global_segment_size({"size": 1024})
+
+
+class TestConvertToBytes(unittest.TestCase):
+    def test_valid_conversion(self):
+        self.assertEqual(_convert_to_bytes("10", 1, "10"), 10)
+        self.assertEqual(_convert_to_bytes("1.5", 1024, "1.5KB"), int(1.5 * 1024))
+        self.assertEqual(_convert_to_bytes("0", 1024**3, "0GB"), 0)
+
+    def test_invalid_numbers(self):
+        with self.assertRaises(ValueError):
+            _convert_to_bytes("abc", 1, "abc")
+
+        with self.assertRaises(ValueError):
+            _convert_to_bytes("1.2.3", 1024, "1.2.3KB")

tests/ut/distributed/mooncake/test_mooncake_kv_transfer.py

@@ -0,0 +1,73 @@
+import threading
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+if not hasattr(torch, "npu"):
+    torch.npu = SimpleNamespace(Event=object)  # type: ignore[attr-defined]
+
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.config_data import (
+    ChunkedTokenDatabase,
+    KeyMetadata,
+    ReqMeta,
+)
+from vllm_ascend.distributed.kv_transfer.kv_pool.ascend_store.kv_transfer import (
+    KVCacheStoreSendingThread,
+)
+
+
+class _FakeStore:
+    def __init__(self, exists_result: list[int]):
+        self.exists_result = exists_result
+        self.put_calls: list[tuple[list[str], list[list[int]], list[list[int]]]] = []
+
+    def set_device(self):
+        return None
+
+    def exists(self, keys: list[str]) -> list[int]:
+        # Return exact number of states for requested keys.
+        return self.exists_result[: len(keys)]
+
+    def put(self, keys, addrs, sizes):
+        self.put_calls.append((list(keys), list(addrs), list(sizes)))
+
+
+class TestKVTransferMissingKeyPut(unittest.TestCase):
+    def test_sending_thread_only_puts_missing_keys(self):
+        store = _FakeStore(exists_result=[1, 0, 1, 0])
+        token_db = ChunkedTokenDatabase([KeyMetadata("m", 0, 0, 0, 0)], [16], None)
+        token_db.set_group_buffers({0: [1000]}, {0: [16]}, {0: [1]})
+        thread = KVCacheStoreSendingThread(
+            m_store=store,
+            token_database=token_db,
+            block_size=16,
+            tp_rank=0,
+            dcp_size=1,
+            put_step=1,
+            kv_role="kv_producer",
+            ready_event=threading.Event(),
+            group_uses_align_state=[False],
+            enable_kv_event=False,
+        )
+
+        req_meta = ReqMeta(
+            req_id="req-1",
+            token_len_chunk=64,
+            block_ids=[0, 1, 2, 3],
+            block_hashes=[b"h0", b"h1", b"h2", b"h3"],  # type: ignore[arg-type]
+            current_event=None,
+        )
+        thread.add_stored_request("req-1")
+        thread.request_queue.put(req_meta)
+        thread._handle_request(req_meta)
+
+        self.assertEqual(len(store.put_calls), 1)
+        put_keys, put_addrs, put_sizes = store.put_calls[0]
+        self.assertEqual(len(put_keys), 2)
+        self.assertEqual(put_addrs, [[1001], [1003]])
+        self.assertEqual(put_sizes, [[16], [16]])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/distributed/test_communicator.py

@@ -4,19 +4,14 @@ from unittest.mock import MagicMock, patch
 import torch
 import torch.distributed as dist
 
-from vllm_ascend.distributed.communicator import NPUCommunicator
+from vllm_ascend.distributed.device_communicators.npu_communicator import NPUCommunicator
 
 
 class TestNPUCommunicator(unittest.TestCase):
-
     @patch("vllm.config.get_current_vllm_config", return_value=None)
     @patch("torch.npu.current_device", return_value=MagicMock())
     @patch("torch.npu.set_device", return_value=MagicMock())
-    @patch("torch.distributed.get_process_group_ranks",
-           return_value={
-               0: 0,
-               1: 1
-           })
+    @patch("torch.distributed.get_process_group_ranks", return_value={0: 0, 1: 1})
     @patch("torch.distributed.get_group_rank", return_value={0: 0, 1: 1})
     @patch("torch.distributed.is_initialized", return_value=True)
     @patch("torch.distributed.get_rank", return_value=1)
@@ -27,15 +22,8 @@ class TestNPUCommunicator(unittest.TestCase):
     @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1])
     @patch("torch.npu.device")
     def test_all_to_all_with_sizes(self, *_):
-
-        def patched_all_to_all(output_tensor_list,
-                               input_tensor_list,
-                               group=None,
-                               async_op=False):
-            output_tensor_list[:] = ([
-                torch.tensor([10, 20]),
-                torch.tensor([50, 60])
-            ])
+        def patched_all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False):
+            output_tensor_list[:] = [torch.tensor([10, 20]), torch.tensor([50, 60])]
 
         torch.distributed.all_to_all = patched_all_to_all
 
@@ -43,22 +31,17 @@ class TestNPUCommunicator(unittest.TestCase):
         gather_sizes = [2, 2]
         input_ = torch.tensor([10, 20, 30, 40])
 
-        comm = NPUCommunicator(cpu_group=dist.group.WORLD)
+        with patch.dict(dist.distributed_c10d._world.pg_map, {dist.group.WORLD: MagicMock()}, clear=False):
+            comm = NPUCommunicator(cpu_group=dist.group.WORLD)
 
-        output = comm.all_to_all(input_,
-                                 scatter_sizes=scatter_sizes,
-                                 gather_sizes=gather_sizes)
+        output = comm.all_to_all(input_, scatter_sizes=scatter_sizes, gather_sizes=gather_sizes)
 
         assert output.tolist() == [10, 20, 50, 60]
 
     @patch("vllm.config.get_current_vllm_config", return_value=None)
     @patch("torch.npu.current_device", return_value=MagicMock())
     @patch("torch.npu.set_device", return_value=MagicMock())
-    @patch("torch.distributed.get_process_group_ranks",
-           return_value={
-               0: 0,
-               1: 1
-           })
+    @patch("torch.distributed.get_process_group_ranks", return_value={0: 0, 1: 1})
     @patch("torch.distributed.get_group_rank", return_value={0: 0, 1: 1})
     @patch("torch.distributed.is_initialized", return_value=True)
     @patch("torch.distributed.get_rank", return_value=1)
@@ -69,21 +52,15 @@ class TestNPUCommunicator(unittest.TestCase):
     @patch("torch.distributed.get_process_group_ranks", return_value=[0, 1])
     @patch("torch.npu.device")
     def test_all_to_all_without_sizes(self, *_):
-
-        def patched_all_to_all(output_tensor_list,
-                               input_tensor_list,
-                               group=None,
-                               async_op=False):
-            output_tensor_list[:] = ([
-                torch.tensor([[10, 20]]),
-                torch.tensor([[50, 60]])
-            ])
+        def patched_all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False):
+            output_tensor_list[:] = [torch.tensor([[10, 20]]), torch.tensor([[50, 60]])]
 
         torch.distributed.all_to_all = patched_all_to_all
 
         input_ = torch.tensor([[10, 20], [30, 40]])
 
-        comm = NPUCommunicator(cpu_group=dist.group.WORLD)
-        output = comm.all_to_all(input_, scatter_dim=0, gather_dim=0)
+        with patch.dict(dist.distributed_c10d._world.pg_map, {dist.group.WORLD: MagicMock()}, clear=False):
+            comm = NPUCommunicator(cpu_group=dist.group.WORLD)
+            output = comm.all_to_all(input_, scatter_dim=0, gather_dim=0)
 
         assert output.tolist() == [[10, 20], [50, 60]]

tests/ut/distributed/test_parallel_state.py

@@ -1,48 +1,157 @@
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 from vllm.config import ParallelConfig
 
 from vllm_ascend.distributed.parallel_state import (
-    _LMTP, _MC2, _OTP, destroy_ascend_model_parallel, get_lmhead_tp_group,
-    get_mc2_group, get_otp_group, init_ascend_model_parallel)
+    _FLASHCOMM2_ODP,
+    _FLASHCOMM2_OTP,
+    _LMTP,
+    _MC2,
+    _OTP,
+    _P_TP,
+    destroy_ascend_model_parallel,
+    get_flashcomm2_odp_group,
+    get_flashcomm2_otp_group,
+    get_global_rank,
+    get_lmhead_tp_group,
+    get_mc2_group,
+    get_otp_group,
+    get_p_tp_group,
+    init_ascend_model_parallel,
+)
 
 
 @pytest.fixture
 def parallel_config():
-    return ParallelConfig(data_parallel_size=2,
-                          tensor_parallel_size=2,
-                          pipeline_parallel_size=2)
+    return ParallelConfig(
+        data_parallel_size=2,
+        tensor_parallel_size=4,
+        pipeline_parallel_size=2,
+    )
 
 
 @pytest.fixture
 def mock_distributed():
-    with patch('torch.distributed.is_initialized', return_value=True), \
-         patch('torch.distributed.get_world_size', return_value=8), \
-         patch('torch.distributed.get_backend', return_value='nccl'), \
-         patch('vllm_ascend.distributed.parallel_state.get_world_group') as mock_group:
+    with (
+        patch("torch.distributed.is_initialized", return_value=True),
+        patch("torch.distributed.get_world_size", return_value=16),
+        patch("torch.distributed.get_backend", return_value="nccl"),
+        patch("vllm_ascend.distributed.parallel_state.get_world_group") as mock_group,
+        patch("vllm_ascend.distributed.parallel_state.get_tp_group") as mock_tp_group,
+    ):
         mock_group.return_value.local_rank = 0
         mock_group.return_value.device_group = MagicMock()
+        mock_tp_group.return_value.world_size = 4
         yield
 
 
 def test_init_ascend_model_parallel(mock_distributed, parallel_config):
     mock_ascend_config = MagicMock()
-    mock_ascend_config.lmhead_tensor_parallel_size = 2
-    mock_ascend_config.oproj_tensor_parallel_size = 2
-    with patch('vllm_ascend.distributed.parallel_state.model_parallel_initialized', return_value=False), \
-         patch('vllm_ascend.distributed.parallel_state.init_model_parallel_group'), \
-         patch('vllm_ascend.distributed.parallel_state.get_ascend_config', return_value=mock_ascend_config):
+    mock_ascend_config.finegrained_tp_config.lmhead_tensor_parallel_size = 2
+    mock_ascend_config.finegrained_tp_config.oproj_tensor_parallel_size = 2
+    mock_ascend_config.finegrained_tp_config.embedding_tensor_parallel_size = 2
+    mock_ascend_config.finegrained_tp_config.mlp_tensor_parallel_size = 2
+    mock_ascend_config.flashcomm2_oproj_tensor_parallel_size = 2
+    mock_ascend_config.pd_tp_ratio = 2
+    mock_ascend_config.num_head_replica = 0
+    mock_ascend_config.pd_head_ratio = 2
+    mock_ascend_config.enable_flashcomm2_parallel_size = 2
+    mock_ascend_config.enable_context_parallel = False
+    mock_vllm_config = MagicMock()
+    mock_vllm_config.kv_transfer_config.is_kv_producer = True
+    with (
+        patch("vllm_ascend.distributed.parallel_state.model_parallel_initialized", return_value=False),
+        patch("vllm_ascend.distributed.parallel_state.init_model_parallel_group"),
+        patch("vllm_ascend.distributed.parallel_state.get_current_vllm_config", return_value=mock_vllm_config),
+        patch("vllm_ascend.distributed.parallel_state.get_ascend_config", return_value=mock_ascend_config),
+        patch("vllm_ascend.utils.get_ascend_config", return_value=mock_ascend_config),
+    ):
         init_ascend_model_parallel(parallel_config)
 
         mc2_group = get_mc2_group()
         lmheadtp_group = get_lmhead_tp_group()
         otp_group = get_otp_group()
+        flashcomm2_otp_group = get_flashcomm2_otp_group()
+        flashcomm2_odp_group = get_flashcomm2_odp_group()
+        p_tp_group = get_p_tp_group()
         assert mc2_group is not None
         assert otp_group is not None
+        assert flashcomm2_otp_group is not None
+        assert flashcomm2_odp_group is not None
         assert lmheadtp_group is not None
+        assert p_tp_group is not None
 
         destroy_ascend_model_parallel()
         assert _MC2 is None
         assert _LMTP is None
         assert _OTP is None
+        assert _FLASHCOMM2_OTP is None
+        assert _FLASHCOMM2_ODP is None
+        assert _P_TP is None
+
+
+def _build_parallel_config(
+    tensor_parallel_size=1,
+    pipeline_parallel_size=1,
+    prefill_context_parallel_size=1,
+    data_parallel_index=0,
+):
+    return SimpleNamespace(
+        tensor_parallel_size=tensor_parallel_size,
+        pipeline_parallel_size=pipeline_parallel_size,
+        prefill_context_parallel_size=prefill_context_parallel_size,
+        data_parallel_index=data_parallel_index,
+    )
+
+
+@pytest.mark.parametrize(
+    "parallel_config_kwargs, rank_in_group, expected",
+    [
+        # No parallelism at all (single card): replica_size == 1.
+        (dict(tensor_parallel_size=1), 0, 0),
+        # TP only: rank_in_group is the local rank within the single replica.
+        (dict(tensor_parallel_size=4), 0, 0),
+        (dict(tensor_parallel_size=4), 3, 3),
+        # Dense DP: world group spans one replica, rank_in_group is local and
+        # data_parallel_index supplies the DP offset.
+        (dict(tensor_parallel_size=4, data_parallel_index=0), 2, 2),
+        (dict(tensor_parallel_size=4, data_parallel_index=1), 2, 6),
+        # MoE DP / external_launcher: world group spans all DP ranks, so
+        # rank_in_group is already global; the modulo strips the DP offset and
+        # data_parallel_index re-adds it (result equals rank_in_group).
+        (dict(tensor_parallel_size=4, data_parallel_index=1), 6, 6),
+        (dict(tensor_parallel_size=4, data_parallel_index=1), 7, 7),
+        # TP * PP * prefill-CP all contribute to replica_size; DCP/EP do not.
+        (dict(tensor_parallel_size=2, pipeline_parallel_size=2, data_parallel_index=1), 1, 5),
+        (
+            dict(
+                tensor_parallel_size=2, pipeline_parallel_size=2, prefill_context_parallel_size=2, data_parallel_index=1
+            ),
+            3,
+            11,
+        ),
+    ],
+)
+def test_get_global_rank(parallel_config_kwargs, rank_in_group, expected):
+    parallel_config = _build_parallel_config(**parallel_config_kwargs)
+    with patch("vllm_ascend.distributed.parallel_state.get_world_group") as mock_group:
+        mock_group.return_value.rank_in_group = rank_in_group
+        assert get_global_rank(parallel_config) == expected
+
+
+def test_get_global_rank_defaults_to_current_config():
+    parallel_config = _build_parallel_config(tensor_parallel_size=4, data_parallel_index=1)
+    mock_vllm_config = MagicMock()
+    mock_vllm_config.parallel_config = parallel_config
+    with (
+        patch(
+            "vllm_ascend.distributed.parallel_state.get_current_vllm_config",
+            return_value=mock_vllm_config,
+        ),
+        patch("vllm_ascend.distributed.parallel_state.get_world_group") as mock_group,
+    ):
+        mock_group.return_value.rank_in_group = 3
+        # data_parallel_index(1) * replica_size(4) + 3 == 7
+        assert get_global_rank() == 7

tests/ut/distributed/weight_transfer/test_npu_ipc_engine.py

@@ -0,0 +1,159 @@
+#
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+"""Regression tests for the NPU IPC weight transfer engine.
+
+These cover two bugs that broke ``examples/rl/rlhf_http_npu_ipc.py``:
+
+1. ``NPUIPCWeightTransferEngine.__init__`` did not accept the ``model``
+   argument that ``WeightTransferEngineFactory.create_engine`` passes,
+   raising ``TypeError: __init__() takes 3 positional arguments but 4
+   were given`` at engine construction.
+2. ``receive_weights`` / ``packed_npu_ipc_consumer`` unpacked the stored
+   IPC handle as ``func, args`` even though the producer stored only the
+   ``reduce_tensor`` *args*, raising ``ValueError: too many values to
+   unpack (expected 2)``. Aligned with upstream vLLM's CUDA IPC engine:
+   the producer stores args only and the consumer rebuilds with the
+   well-known ``rebuild_npu_tensor``.
+"""
+
+import inspect
+import sys
+import types
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from vllm_ascend.distributed.weight_transfer import npu_ipc_engine
+from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import (
+    NPUIPCWeightTransferEngine,
+)
+
+_MODULE = "vllm_ascend.distributed.weight_transfer.npu_ipc_engine"
+
+
+def _patch_rebuild_npu_tensor(rebuild_func):
+    """Install a fake ``torch_npu.multiprocessing.reductions`` module.
+
+    The engine imports ``rebuild_npu_tensor`` lazily from ``torch_npu``,
+    which is only a stub on CPU CI runners, so provide a fake submodule.
+    """
+    fake_mod = types.ModuleType("torch_npu.multiprocessing.reductions")
+    fake_mod.rebuild_npu_tensor = rebuild_func  # type: ignore[attr-defined]
+    return patch.dict(
+        sys.modules,
+        {
+            "torch_npu.multiprocessing": types.ModuleType("torch_npu.multiprocessing"),
+            "torch_npu.multiprocessing.reductions": fake_mod,
+        },
+    )
+
+
+def test_init_accepts_model_argument():
+    """Bug 1: __init__ must accept the optional ``model`` argument."""
+    params = inspect.signature(NPUIPCWeightTransferEngine.__init__).parameters
+    assert "model" in params
+
+
+def test_init_passes_model_to_super():
+    """Bug 1: the ``model`` argument must be forwarded to the base engine."""
+    captured = {}
+
+    def fake_init(self, config, parallel_config, model=None):
+        captured["args"] = (config, parallel_config, model)
+
+    with patch.object(npu_ipc_engine.WeightTransferEngine, "__init__", fake_init):
+        NPUIPCWeightTransferEngine("config", "parallel_config", "model")
+
+    assert captured["args"] == ("config", "parallel_config", "model")
+
+
+def test_unpacked_send_stores_reduce_tensor_args_only():
+    """Bug 2 (producer): the handle stores only the ``reduce_tensor`` args.
+
+    This matches upstream vLLM's CUDA IPC engine, which drops the rebuild
+    func and relies on the consumer using the well-known rebuild function.
+    """
+    npu_uuid = "node-0"
+
+    rebuild_args = (None, None, None, None, None, None, 999, None)
+    fake_reduce = MagicMock(return_value=("rebuild_func_sentinel", rebuild_args))
+
+    captured = {}
+
+    def send_mode(update_info):
+        captured["update_info"] = update_info
+
+    trainer_args = MagicMock()
+    trainer_args.send_mode = send_mode
+    trainer_args.packed = False
+
+    iterator = iter([("model.weight", torch.zeros(3))])
+
+    with patch(f"{_MODULE}.reduce_tensor", fake_reduce):
+        NPUIPCWeightTransferEngine._send_unpacked(iterator, trainer_args, npu_uuid)
+
+    update_info = captured["update_info"]
+    assert isinstance(update_info.ipc_handles, list)
+    stored = update_info.ipc_handles[0][npu_uuid]
+    # Only the args tuple is stored, not a (func, args) pair.
+    assert stored == rebuild_args
+
+
+def test_receive_weights_rebuilds_with_rebuild_npu_tensor():
+    """Bug 2 (consumer): receive_weights rebuilds via ``rebuild_npu_tensor``.
+
+    Verifies the args-only handle is consumed without unpacking errors and
+    that the receiver's device index is written into the rebuild args.
+    """
+    npu_uuid = "node-0"
+    device_index = 0
+
+    rebuilt_weight = torch.tensor([1.0, 2.0, 3.0])
+    seen = {}
+
+    def fake_rebuild(*args):
+        seen["args"] = args
+        return rebuilt_weight
+
+    # Sender stores 999 at index 6; the receiver must overwrite it.
+    rebuild_args = (None, None, None, None, None, None, 999, None)
+
+    update_info = NPUIPCWeightTransferEngine.update_info_cls(
+        names=["model.weight"],
+        dtype_names=["float32"],
+        shapes=[[3]],
+        ipc_handles=[{npu_uuid: rebuild_args}],
+        packed=False,
+    )
+
+    engine = object.__new__(NPUIPCWeightTransferEngine)
+    received = {}
+
+    def load_weights(weights):
+        received["weights"] = weights
+
+    with (
+        _patch_rebuild_npu_tensor(fake_rebuild),
+        patch(f"{_MODULE}.npu_generate_uuid", return_value=npu_uuid),
+        patch("torch.accelerator.current_device_index", return_value=device_index),
+    ):
+        engine.receive_weights(update_info, load_weights)
+
+    assert received["weights"][0][0] == "model.weight"
+    assert torch.equal(received["weights"][0][1], rebuilt_weight)
+    # Index 6 (device index) overwritten with the receiver's device.
+    assert seen["args"][6] == device_index

tests/ut/eplb/adaptor/test_vllm_adaptor.py

@@ -0,0 +1,194 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+from transformers import DeepseekV2Config
+
+from vllm_ascend.eplb.adaptor.vllm_adaptor import EPLB_EXPERT_WEIGHT_NAMES, VllmEplbAdaptor
+from vllm_ascend.quantization.quant_type import QuantType
+
+
+class TestVllmAdaptor(unittest.TestCase):
+    def setUp(self):
+        VllmEplbAdaptor._registered_moe_layers = []
+
+        n_routed_experts = 256
+        self.mock_layer = MagicMock()
+        self.mock_layer.local_num_experts = n_routed_experts
+        self.mock_layer.ep_rank = 0
+        self.mock_layer.quant_type = QuantType.W8A8
+        self.mock_layer.w13_weight_list = [torch.randn(256, 128) for _ in range(n_routed_experts)]
+        self.mock_layer.w2_weight_list = [torch.randn(128, 256) for _ in range(n_routed_experts)]
+        self.mock_layer.w13_weight_scale_fp32_list = [torch.tensor([1.0]) for _ in range(n_routed_experts)]
+        self.mock_layer.w2_weight_scale_list = [torch.tensor([1.0]) for _ in range(n_routed_experts)]
+        self.mock_layer.w13_weight = torch.randn(n_routed_experts, 256, 128)
+        self.mock_layer.w2_weight = torch.randn(n_routed_experts, 128, 256)
+        self.mock_layer.moe_load = torch.randn(n_routed_experts)
+        self.mock_layer.global_expert_map = torch.arange(n_routed_experts * 4).reshape(n_routed_experts, 4)
+        self.mock_layer.get_log2phy_map.return_value = torch.arange(4)
+        self.mock_layer.clear_moe_load = MagicMock()
+        VllmEplbAdaptor.register_layer(self.mock_layer)
+
+        mock_model = MagicMock()
+        mock_model.model.named_parameters.return_value = dict()
+        config = DeepseekV2Config(n_routed_experts=n_routed_experts)
+        mock_model.config = config
+        del mock_model.language_model
+        self.model = mock_model
+        num_dense_layers = getattr(config, "first_k_dense_replace", 0)
+        self.model.model.layers[num_dense_layers].mlp.experts.quant_type = QuantType.W8A8
+
+        self.mock_rank = patch("vllm_ascend.eplb.adaptor.vllm_adaptor.dist.get_rank", return_value=0).start()
+        self.mock_size = patch("vllm_ascend.eplb.adaptor.vllm_adaptor.dist.get_world_size", return_value=4).start()
+
+    @patch("torch.empty_like", return_value=torch.zeros(16, 32))
+    @patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
+    def test_init_fp16(self, mock_get_config, mock_func):
+        mock_config = MagicMock()
+        mock_config.enable_fused_mc2 = 1
+        mock_get_config.return_value = mock_config
+        self.model.quant_config = None
+        adaptor = VllmEplbAdaptor(self.model)
+        self.assertEqual(adaptor.expert_weight_key_per_layer[0], (QuantType.NONE, True))
+        self.assertIs(adaptor.expert_param_per_layer[0][0][0], self.mock_layer.w13_weight_list[0])
+        self.assertIs(adaptor.expert_param_per_layer[0][0][1], self.mock_layer.w2_weight_list[0])
+
+    @patch("torch.empty_like", return_value=torch.zeros(16, 32))
+    @patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
+    def test_init_w8a8(self, mock_get_config, mock_func):
+        mock_config = MagicMock()
+        mock_config.enable_fused_mc2 = 0
+        mock_get_config.return_value = mock_config
+        VllmEplbAdaptor(self.model)
+
+    @patch("torch.empty_like", return_value=torch.zeros(16, 32))
+    @patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
+    def test_language_model_w8a8(self, mock_get_config, mock_func):
+        mock_config = MagicMock()
+        mock_config.enable_fused_mc2 = 0
+        mock_get_config.return_value = mock_config
+        model = MagicMock()
+        model.language_model = self.model
+        model.config.text_config = self.model.config
+        VllmEplbAdaptor(model)
+
+    def test_pp_eplb_adaptor_init_with_registered_layer(self):
+        """PP+EPLB: adaptor picks up MoE layers registered via register_layer."""
+        VllmEplbAdaptor._registered_moe_layers = []
+        layer = MagicMock()
+        layer.local_num_experts = 4
+        layer.ep_rank = 0
+        layer.quant_type = QuantType.W8A8
+        layer.w13_weight_list = [torch.randn(256, 128) for _ in range(4)]
+        layer.w2_weight_list = [torch.randn(128, 256) for _ in range(4)]
+        layer.w13_weight_scale_fp32_list = [torch.tensor([1.0]) for _ in range(4)]
+        layer.w2_weight_scale_list = [torch.tensor([1.0]) for _ in range(4)]
+        layer.moe_load = torch.randn(4)
+        layer.global_expert_map = torch.arange(16).reshape(4, 4)
+        layer.get_log2phy_map.return_value = torch.arange(4)
+        VllmEplbAdaptor.register_layer(layer)
+
+        with patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config") as mock_get_config:
+            mock_config = MagicMock()
+            mock_config.enable_fused_mc2 = 0
+            mock_get_config.return_value = mock_config
+            model = MagicMock()
+            model.quant_config = MagicMock()
+            model.config.first_k_dense_replace = 0
+            del model.language_model
+            adaptor = VllmEplbAdaptor(model)
+
+        self.assertEqual(adaptor.num_moe_layers, 1)
+        self.assertEqual(adaptor.num_local_experts, 4)
+        self.assertEqual(adaptor.ep_rank, 0)
+
+    @patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
+    def test_init_mixed_quant_type_per_layer(self, mock_get_config):
+        mock_config = MagicMock()
+        mock_config.enable_fused_mc2 = 1
+        mock_get_config.return_value = mock_config
+
+        VllmEplbAdaptor._registered_moe_layers = []
+        num_local_experts = 2
+        w8a8_layer = MagicMock()
+        w8a8_layer.local_num_experts = num_local_experts
+        w8a8_layer.ep_rank = 0
+        w8a8_layer.quant_type = QuantType.W8A8
+        w8a8_layer.w13_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
+        w8a8_layer.w2_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
+        w8a8_layer.w13_weight_scale_fp32_list = [torch.randn(1) for _ in range(num_local_experts)]
+        w8a8_layer.w2_weight_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
+        w8a8_layer.fused_w1_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
+        w8a8_layer.fused_w2_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
+        w8a8_layer.moe_load = torch.zeros(num_local_experts)
+        w8a8_layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
+        w8a8_layer.get_log2phy_map.return_value = torch.arange(4)
+
+        mxfp8_layer = MagicMock()
+        mxfp8_layer.local_num_experts = num_local_experts
+        mxfp8_layer.ep_rank = 0
+        mxfp8_layer.quant_type = QuantType.MXFP8
+        mxfp8_layer.w13_weight = torch.randn(num_local_experts, 2, 2)
+        mxfp8_layer.w2_weight = torch.randn(num_local_experts, 2, 2)
+        mxfp8_layer.w13_weight_scale = torch.randn(num_local_experts, 1)
+        mxfp8_layer.w2_weight_scale = torch.randn(num_local_experts, 1)
+        mxfp8_layer.moe_load = torch.zeros(num_local_experts)
+        mxfp8_layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
+        mxfp8_layer.get_log2phy_map.return_value = torch.arange(4)
+
+        VllmEplbAdaptor.register_layer(w8a8_layer)
+        VllmEplbAdaptor.register_layer(mxfp8_layer)
+
+        model = MagicMock()
+        model.quant_config = MagicMock()
+        model.config.first_k_dense_replace = 0
+        del model.language_model
+        adaptor = VllmEplbAdaptor(model)
+
+        w8a8_key = (QuantType.W8A8, True)
+        mxfp8_key = (QuantType.MXFP8, True)
+        self.assertEqual(adaptor.expert_weight_key_per_layer[0], w8a8_key)
+        self.assertEqual(adaptor.expert_weight_key_per_layer[1], mxfp8_key)
+        self.assertEqual(len(adaptor.buffer_tensor_list[w8a8_key][0]), len(EPLB_EXPERT_WEIGHT_NAMES[w8a8_key]))
+        self.assertEqual(len(adaptor.buffer_tensor_list[mxfp8_key][0]), len(EPLB_EXPERT_WEIGHT_NAMES[mxfp8_key]))
+        self.assertEqual(len(adaptor.expert_param_per_layer[0][0]), len(EPLB_EXPERT_WEIGHT_NAMES[w8a8_key]))
+        self.assertEqual(len(adaptor.expert_param_per_layer[1][0]), len(EPLB_EXPERT_WEIGHT_NAMES[mxfp8_key]))
+
+    @patch("vllm_ascend.eplb.adaptor.vllm_adaptor.get_ascend_config")
+    def test_reused_buffer_requires_same_expert_weight_shape(self, mock_get_config):
+        mock_config = MagicMock()
+        mock_config.enable_fused_mc2 = 0
+        mock_get_config.return_value = mock_config
+
+        VllmEplbAdaptor._registered_moe_layers = []
+        num_local_experts = 2
+        for weight_shape in [(2, 2), (3, 2)]:
+            layer = MagicMock()
+            layer.local_num_experts = num_local_experts
+            layer.ep_rank = 0
+            layer.quant_type = QuantType.W8A8
+            layer.w13_weight_list = [torch.randn(*weight_shape) for _ in range(num_local_experts)]
+            layer.w2_weight_list = [torch.randn(2, 2) for _ in range(num_local_experts)]
+            layer.w13_weight_scale_fp32_list = [torch.randn(1) for _ in range(num_local_experts)]
+            layer.w2_weight_scale_list = [torch.randn(1) for _ in range(num_local_experts)]
+            layer.moe_load = torch.zeros(num_local_experts)
+            layer.global_expert_map = torch.arange(num_local_experts * 4).reshape(num_local_experts, 4)
+            layer.get_log2phy_map.return_value = torch.arange(4)
+            VllmEplbAdaptor.register_layer(layer)
+
+        model = MagicMock()
+        model.quant_config = MagicMock()
+        model.config.first_k_dense_replace = 0
+        del model.language_model
+
+        with self.assertRaisesRegex(AssertionError, "EPLB expert weight shapes mismatch"):
+            VllmEplbAdaptor(model)
+
+    def tearDown(self):
+        self.mock_rank.stop()
+        self.mock_size.stop()
+        VllmEplbAdaptor._registered_moe_layers = []
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ut/eplb/core/a2/expert_map.json

@@ -0,0 +1,17 @@
+{
+    "moe_layer_count":
+    1,
+    "layer_list": [{
+        "layer_id":
+        0,
+        "device_count":
+        2,
+        "device_list": [{
+            "device_id": 0,
+            "device_expert": [7, 2, 0, 3, 5]
+        }, {
+            "device_id": 1,
+            "device_expert": [6, 1, 4, 7, 2]
+        }]
+    }]
+}

tests/ut/eplb/core/a2/test_eplb_utils.py

@@ -0,0 +1,117 @@
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+# isort: off
+import torch
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig, FusedMoEParallelConfig
+
+from vllm_ascend.ascend_config import init_ascend_config
+from vllm_ascend.eplb.core.eplb_utils import generate_log2phy_map, init_eplb_config
+from vllm_ascend.utils import vllm_version_is
+# isort: on
+
+
+class TestAscendConfig(unittest.TestCase):
+    @patch("vllm.config.VllmConfig.__post_init__", MagicMock())
+    @patch("vllm_ascend.platform.NPUPlatform._fix_incompatible_config")
+    def setUp(self, mock_fix_incompatible_config):
+        vllm_config = VllmConfig()
+        vllm_config.model_config = MagicMock()
+        vllm_config.additional_config = {
+            "refresh": True,
+            "eplb_config": {"dynamic_eplb": True, "num_redundant_experts": 2},
+        }
+        from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+
+        moe_parallel_config = FusedMoEParallelConfig(2, 0, 1, 2, 1, 1, 1, 1, 1, True, "hccl", enable_eplb=True)
+        if vllm_version_is("0.23.0"):
+            moe_config = FusedMoEConfig(
+                num_experts=8,
+                experts_per_token=8,
+                hidden_dim=8192,
+                intermediate_size_per_partition=5,
+                num_local_experts=8,
+                num_logical_experts=8,
+                activation="silu",
+                device="npu",
+                routing_method=RoutingMethodType.Simulated,
+                moe_parallel_config=moe_parallel_config,
+                in_dtype=torch.float16,
+            )
+        else:
+            from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+
+            moe_config = FusedMoEConfig(
+                num_experts=8,
+                experts_per_token=8,
+                hidden_dim=8192,
+                intermediate_size=10,
+                num_local_experts=8,
+                num_logical_experts=8,
+                activation=MoEActivation.SILU,
+                device="npu",
+                routing_method=RoutingMethodType.Simulated,
+                moe_parallel_config=moe_parallel_config,
+                in_dtype=torch.float16,
+            )
+        moe_config.supports_eplb = True
+        self.vllm_config = vllm_config
+        self.moe_config = moe_config
+        self.mock_npu_patcher = patch("torch.Tensor.npu", new=lambda self: self)
+        self.mock_npu_patcher.start()
+        os.environ["DYNAMIC_EPLB"] = "true"
+
+    def tearDown(self):
+        self.mock_npu_patcher.stop()
+        os.environ.pop("DYNAMIC_EPLB", None)
+
+    def test_init_eplb_config_with_eplb(self):
+        eplb_config = init_ascend_config(self.vllm_config).eplb_config
+        _, expert_map, log2phy, redundant_experts = init_eplb_config(eplb_config, 0, self.moe_config)
+        gt_expert_map = torch.tensor([4, -1, -1, -1, 0, 1, 2, 3])
+        gt_log2phy = torch.tensor([9, 1, 2, 3, 5, 6, 7, 8])
+        self.assertTrue(torch.equal(expert_map, gt_expert_map))
+        self.assertTrue(torch.equal(log2phy, gt_log2phy))
+        self.assertEqual(redundant_experts, 2)
+
+    def test_init_eplb_config_with_eplb_withmap(self):
+        _TEST_DIR = os.path.dirname(__file__)
+        self.vllm_config.additional_config["eplb_config"]["expert_map_path"] = _TEST_DIR + "/expert_map.json"
+        eplb_config = init_ascend_config(self.vllm_config).eplb_config
+        _, expert_map, log2phy, redundant_experts = init_eplb_config(eplb_config, 0, self.moe_config)
+        gt_expert_map = torch.tensor([-1, 1, 4, -1, 2, -1, 0, 3])
+        gt_log2phy = torch.tensor([2, 6, 9, 3, 7, 4, 5, 8])
+        self.assertTrue(torch.equal(expert_map, gt_expert_map))
+        self.assertTrue(torch.equal(log2phy, gt_log2phy))
+        self.assertEqual(redundant_experts, 2)
+
+    def test_generate_log2phy_map_rotates_tail_tp_rank_with_tp_size(self):
+        global_expert_map = [
+            torch.tensor([0, -1], dtype=torch.int32),
+            torch.tensor([0, -1], dtype=torch.int32),
+            torch.tensor([0, -1], dtype=torch.int32),
+            torch.tensor([0, -1], dtype=torch.int32),
+            torch.tensor([-1, 0], dtype=torch.int32),
+            torch.tensor([-1, 0], dtype=torch.int32),
+            torch.tensor([-1, 0], dtype=torch.int32),
+            torch.tensor([-1, 0], dtype=torch.int32),
+        ]
+
+        fallback_tail_dp1 = generate_log2phy_map(global_expert_map, ep_rank=7)
+        rotated_tail_dp0 = generate_log2phy_map(global_expert_map, ep_rank=3, tp_size=4)
+        rotated_tail_dp1 = generate_log2phy_map(global_expert_map, ep_rank=7, tp_size=4)
+
+        self.assertTrue(torch.equal(fallback_tail_dp1, torch.tensor([3, 7], dtype=torch.int32)))
+        self.assertTrue(torch.equal(rotated_tail_dp0, torch.tensor([3, 4], dtype=torch.int32)))
+        self.assertTrue(torch.equal(rotated_tail_dp1, torch.tensor([0, 5], dtype=torch.int32)))
+
+    def test_init_eplb_config_without_eplb(self):
+        self.vllm_config.additional_config = {"refresh": True}
+        eplb_config = init_ascend_config(self.vllm_config).eplb_config
+        _, expert_map, log2phy, redundant_experts = init_eplb_config(eplb_config, 0, self.moe_config)
+        gt_expert_map = torch.tensor([-1, -1, -1, -1, 0, 1, 2, 3])
+        self.assertIsNone(log2phy)
+        self.assertTrue(torch.equal(expert_map, gt_expert_map))
+        self.assertEqual(redundant_experts, 0)

tests/ut/eplb/core/policy/test_policy_factory.py

@@ -0,0 +1,36 @@
+import unittest
+
+import torch
+
+from vllm_ascend.eplb.core.eplb_worker import EplbWorker
+from vllm_ascend.eplb.core.policy.policy_factory import PolicyFactory
+from vllm_ascend.eplb.core.policy.policy_flashlb import generate_layered_experts
+
+
+class TestEplbRebalancePolicies(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(42)
+        self.current_expert_table = generate_layered_experts()
+        x = torch.rand(100, 58, 32, 9)
+        x = x**10
+        self.expert_workload = (x * 999 + 1).long()
+        self.hotness = EplbWorker._calculate_hotness(self.current_expert_table, self.expert_workload.sum(0))
+
+    @unittest.mock.patch("torch.npu.device_count", return_value=16)
+    def test_swift_balance_rebalance_experts(self, mock_count):
+        swift_policy = PolicyFactory.generate_policy(2)
+        _, _, new_placement = swift_policy.rebalance_experts(self.current_expert_table, self.expert_workload.sum(0))
+        update_mean, _ = EplbWorker._compute_imbalance(new_placement, self.hotness)
+
+        self.assertLessEqual(update_mean, 1.08)
+
+    def test_flashlb_rebalance_experts(self):
+        flashlb_policy = PolicyFactory.generate_policy(3)
+        _, _, new_placement = flashlb_policy.rebalance_experts(self.current_expert_table, self.expert_workload)
+        update_mean, _ = EplbWorker._compute_imbalance(new_placement, self.hotness)
+
+        self.assertLessEqual(update_mean, 1.1)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)