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")