init v0.23.0

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

tests/ut/ops/a2/test_gdn_chunk_meta.py

@@ -0,0 +1,439 @@
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm-ascend project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from vllm_ascend.ops.triton.fla import chunk, chunk_o, chunk_o_update
+from vllm_ascend.utils import enable_custom_op
+
+enable_custom_op()
+
+
+class _FakeKernel:
+    def __init__(self):
+        self.grid = None
+        self.grid_result = None
+        self.launch_kwargs: dict[str, object] | None = None
+
+    def __getitem__(self, grid):
+        self.grid = grid
+        self.grid_result = grid({"BV": 128})
+
+        def launch(**kwargs):
+            self.launch_kwargs = kwargs
+
+        return launch
+
+
+class _DummyTensor:
+    def __init__(self, name: str):
+        self.name = name
+        self.shape = (1,)
+        self.dtype = torch.float32
+
+    def unsqueeze(self, dim: int):
+        return self
+
+    def new_empty(self, *shape):
+        return _DummyTensor(f"{self.name}.new_empty")
+
+    def __getitem__(self, item):
+        return self
+
+    def __setitem__(self, item, value):
+        return None
+
+    def __add__(self, other):
+        return self
+
+    def __sub__(self, other):
+        return self
+
+    def transpose(self, dim0, dim1):
+        return self
+
+    def contiguous(self):
+        return self
+
+    def to(self, *args, **kwargs):
+        return self
+
+
+class _GatherResult:
+    def __init__(self, items):
+        self.items = items
+
+    def __getitem__(self, item):
+        if isinstance(item, tuple):
+            item = item[0]
+        return self.items[item]
+
+
+def _patch_missing_cdiv(monkeypatch: pytest.MonkeyPatch, module) -> None:
+    if hasattr(module.triton, "cdiv"):
+        return
+    monkeypatch.setattr(
+        module.triton,
+        "cdiv",
+        lambda x, y: (x + y - 1) // y,
+        raising=False,
+    )
+
+
+@pytest.mark.parametrize("target", ["chunk_o", "chunk_o_update"])
+def test_chunk_leaf_wrappers_use_prebuilt_chunk_offsets(
+    monkeypatch: pytest.MonkeyPatch,
+    target: str,
+):
+    fake_kernel = _FakeKernel()
+    sentinel = torch.tensor([0, 2, 5], dtype=torch.int32)
+    cu_seqlens = torch.tensor([0, 4, 7], dtype=torch.int32)
+
+    if target == "chunk_o":
+        _patch_missing_cdiv(monkeypatch, chunk_o)
+        monkeypatch.setattr(chunk_o, "chunk_fwd_kernel_o", fake_kernel)
+        monkeypatch.setattr(
+            chunk_o,
+            "prepare_chunk_offsets",
+            lambda *args, **kwargs: pytest.fail("prepare_chunk_offsets should not be called"),
+        )
+        chunk_o.chunk_fwd_o(
+            q=torch.zeros((2, 4, 1, 8), dtype=torch.float32),
+            k=torch.zeros((2, 4, 1, 8), dtype=torch.float32),
+            v=torch.zeros((2, 4, 1, 16), dtype=torch.float32),
+            h=torch.zeros((4, 1, 8, 16), dtype=torch.float32),
+            g=torch.zeros((2, 4, 1), dtype=torch.float32),
+            cu_seqlens=cu_seqlens,
+            chunk_offsets=sentinel,
+        )
+    else:
+        _patch_missing_cdiv(monkeypatch, chunk_o_update)
+        monkeypatch.setattr(chunk_o_update, "chunk_fwd_kernel_o_update", fake_kernel)
+        monkeypatch.setattr(
+            chunk_o_update,
+            "prepare_chunk_offsets",
+            lambda *args, **kwargs: pytest.fail("prepare_chunk_offsets should not be called"),
+        )
+        chunk_o_update.chunk_fwd_o_update(
+            q=torch.zeros((2, 4, 1, 8), dtype=torch.float32),
+            v=torch.zeros((2, 4, 1, 16), dtype=torch.float32),
+            h=torch.zeros((4, 1, 8, 16), dtype=torch.float32),
+            h_update=torch.zeros((5, 1, 8, 8), dtype=torch.float32),
+            updated_h_state=torch.zeros((1, 8, 16), dtype=torch.float32),
+            cu_seqlens=cu_seqlens,
+            chunk_offsets=sentinel,
+        )
+
+    assert fake_kernel.launch_kwargs is not None
+    assert fake_kernel.launch_kwargs["chunk_offsets"] is sentinel
+
+
+def test_chunk_gated_delta_rule_fwd_threads_prebuilt_chunk_offsets(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    chunk_offsets = torch.tensor([0, 2, 5], dtype=torch.int32)
+    update_chunk_offsets = torch.tensor([0, 3, 7], dtype=torch.int32)
+    final_chunk_indices = torch.tensor([1, 3], dtype=torch.int32)
+    prebuilt_meta = type(
+        "PrebuiltMeta",
+        (),
+        {
+            "block_indices_cumsum": None,
+            "cu_seqlens_host": (0, 4, 7),
+            "chunk_indices_chunk64_host": (0, 0, 1, 0),
+            "chunk_indices_chunk64": None,
+            "chunk_offsets_chunk64": chunk_offsets,
+            "update_chunk_offsets_chunk64": update_chunk_offsets,
+            "final_chunk_indices_chunk64": final_chunk_indices,
+            "chunk_indices_large_block": None,
+            "keep_meta": None,
+            "cu_seqlens_kern": None,
+        },
+    )()
+
+    q = _DummyTensor("q")
+    k = _DummyTensor("k")
+    v = _DummyTensor("v")
+    g = _DummyTensor("g")
+    beta = _DummyTensor("beta")
+    initial_state = _DummyTensor("initial_state")
+
+    non_pcp_calls: list[tuple[str, object]] = []
+    pcp_calls: list[tuple[str, object]] = []
+
+    def run_case(world_size: int, calls: list[tuple[str, object]]):
+        group = type(
+            "Group",
+            (),
+            {
+                "world_size": world_size,
+                "rank_in_group": 0,
+                "all_gather": lambda self, value, dim: _GatherResult([_DummyTensor("g0"), _DummyTensor("g1")]),
+            },
+        )()
+
+        monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
+        monkeypatch.setattr(chunk, "get_pcp_group", lambda: group)
+        monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *args, **kwargs: _DummyTensor("g_cumsum"))
+        monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *args, **kwargs: _DummyTensor("A"))
+        monkeypatch.setattr(chunk, "solve_tril", lambda *args, **kwargs: _DummyTensor("A_solved"))
+        monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *args, **kwargs: (_DummyTensor("w"), _DummyTensor("u")))
+        monkeypatch.setattr(
+            chunk,
+            "chunk_gated_delta_rule_fwd_h",
+            lambda *args, **kwargs: (_DummyTensor("h"), _DummyTensor("v_new"), _DummyTensor("final_state")),
+        )
+        monkeypatch.setattr(
+            chunk,
+            "chunk_gated_delta_rule_fwd_hupdate",
+            lambda *args, **kwargs: _DummyTensor("h_update"),
+        )
+        monkeypatch.setattr(
+            chunk.torch,
+            "matmul",
+            lambda *args, **kwargs: _DummyTensor("matmul"),
+            raising=False,
+        )
+        monkeypatch.setattr(
+            chunk.torch,
+            "zeros_like",
+            lambda *args, **kwargs: _DummyTensor("zeros_like"),
+            raising=False,
+        )
+        monkeypatch.setattr(
+            torch.ops._C_ascend,
+            "chunk_gated_delta_rule_fwd_h",
+            lambda *args, **kwargs: (_DummyTensor("h"), _DummyTensor("v_new"), _DummyTensor("final_state")),
+            raising=False,
+        )
+        monkeypatch.setattr(
+            torch.ops._C_ascend,
+            "chunk_fwd_o",
+            lambda *args, **kwargs: _DummyTensor("o_ascend"),
+            raising=False,
+        )
+
+        def fake_chunk_fwd_o(*args, **kwargs):
+            calls.append(("o", kwargs["chunk_offsets"]))
+            return _DummyTensor("o")
+
+        def fake_chunk_fwd_o_update(*args, **kwargs):
+            calls.append(("o_update", kwargs["chunk_offsets"]))
+            return _DummyTensor("h_updated")
+
+        monkeypatch.setattr(chunk, "chunk_fwd_o", fake_chunk_fwd_o)
+        if world_size > 1:
+            monkeypatch.setattr(chunk, "chunk_gated_delta_rule_fwd_hupdate", fake_chunk_fwd_o_update)
+
+        chunk.chunk_gated_delta_rule_fwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            scale=1.0,
+            initial_state=initial_state,
+            output_final_state=False,
+            cu_seqlens=torch.tensor([0, 4, 7], dtype=torch.int32),
+            prebuilt_meta=prebuilt_meta,
+        )
+
+    run_case(1, non_pcp_calls)
+    assert non_pcp_calls == []
+
+    run_case(2, pcp_calls)
+    assert pcp_calls == [("o_update", chunk_offsets)]
+
+
+def test_chunk_gated_delta_rule_fwd_uses_prebuilt_metadata_without_runtime_tolist(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    prebuilt_meta = type(
+        "PrebuiltMeta",
+        (),
+        {
+            "block_indices_cumsum": None,
+            "cu_seqlens_host": (0, 4, 7),
+            "chunk_indices_chunk64_host": (0, 0, 1, 0),
+            "chunk_indices_chunk64": torch.tensor([[0, 0], [1, 0]], dtype=torch.int32),
+            "chunk_offsets_chunk64": torch.tensor([0, 1, 2], dtype=torch.int32),
+            "update_chunk_offsets_chunk64": torch.tensor([0, 2, 4], dtype=torch.int32),
+            "final_chunk_indices_chunk64": torch.tensor([1, 3], dtype=torch.int32),
+            "chunk_indices_large_block": None,
+            "keep_meta": None,
+            "cu_seqlens_kern": None,
+        },
+    )()
+
+    q = _DummyTensor("q")
+    k = _DummyTensor("k")
+    v = _DummyTensor("v")
+    g = _DummyTensor("g")
+    beta = _DummyTensor("beta")
+    initial_state = _DummyTensor("initial_state")
+
+    captured: dict[str, tuple[int, ...] | None] = {}
+
+    monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
+    monkeypatch.setattr(
+        chunk,
+        "get_pcp_group",
+        lambda: type("Group", (), {"world_size": 1, "rank_in_group": 0})(),
+    )
+    monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *args, **kwargs: _DummyTensor("g_cumsum"))
+    monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *args, **kwargs: _DummyTensor("A"))
+    monkeypatch.setattr(chunk, "solve_tril", lambda *args, **kwargs: _DummyTensor("A_solved"))
+    monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *args, **kwargs: (_DummyTensor("w"), _DummyTensor("u")))
+    monkeypatch.setattr(
+        torch.ops._C_ascend,
+        "chunk_gated_delta_rule_fwd_h",
+        lambda *args, **kwargs: (
+            captured.update(
+                {
+                    "cu_seqlens": kwargs["cu_seqlens"],
+                    "chunk_indices": kwargs["chunk_indices"],
+                }
+            )
+            or (_DummyTensor("h"), _DummyTensor("v_new"), _DummyTensor("final_state"))
+        ),
+        raising=False,
+    )
+    monkeypatch.setattr(
+        torch.ops._C_ascend,
+        "chunk_fwd_o",
+        lambda *args, **kwargs: _DummyTensor("o_ascend"),
+        raising=False,
+    )
+    monkeypatch.setattr(
+        torch.Tensor,
+        "tolist",
+        lambda self: pytest.fail("runtime should not convert device tensors to host tuples"),
+    )
+
+    chunk.chunk_gated_delta_rule_fwd(
+        q=q,
+        k=k,
+        v=v,
+        g=g,
+        beta=beta,
+        scale=1.0,
+        initial_state=initial_state,
+        output_final_state=False,
+        cu_seqlens=torch.tensor([0, 4, 7], dtype=torch.int32),
+        prebuilt_meta=prebuilt_meta,
+    )
+
+    assert captured["cu_seqlens"] == prebuilt_meta.cu_seqlens_host
+    assert captured["chunk_indices"] == prebuilt_meta.chunk_indices_chunk64_host
+
+
+def test_chunk_gated_delta_rule_fwd_pcp_chaining_subtracts_initial_state(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """PCP chaining uses (updated_state[i-1] - initial_state), not updated_state[i-1].
+
+    With s0 != 0 (subsequent prefill chunk), the fix subtracts s0 to avoid
+    double-counting Φ_i·s0. Verified by checking the returned final_state
+    matches the sequential result Φ_1·(Φ_0·s0+p_0)+p_1.
+    """
+    torch.manual_seed(42)
+    N, H, K, V = 1, 2, 4, 4
+    s0 = torch.randn(N, H, K, V)
+    phi_0 = torch.randn(N, H, K, K)
+    phi_1 = torch.randn(N, H, K, K)
+    p_0 = torch.randn(N, H, K, V)
+    p_1 = torch.randn(N, H, K, V)
+
+    # Each rank computes final_state = Φ_i · s0 + p_i (from shared s0)
+    rank0_fs = torch.matmul(phi_0, s0) + p_0
+    rank1_fs = torch.matmul(phi_1, s0) + p_1
+    # h_update shape [1, N, H, K, K]; after [:, [0], :, :, :] → [1, N, H, K, K]
+    h_update_tensor = phi_0.unsqueeze(0)
+
+    prebuilt_meta = type(
+        "PrebuiltMeta",
+        (),
+        {
+            "block_indices_cumsum": None,
+            "cu_seqlens_host": (0, N),
+            "chunk_indices_chunk64_host": (0, 0),
+            "chunk_indices_chunk64": None,
+            "chunk_offsets_chunk64": torch.tensor([0, 1], dtype=torch.int32),
+            "update_chunk_offsets_chunk64": torch.tensor([0, 2], dtype=torch.int32),
+            "final_chunk_indices_chunk64": torch.tensor([0], dtype=torch.int32),
+            "chunk_indices_large_block": None,
+            "num_decodes": 0,
+            "keep_meta": None,
+            "cu_seqlens_kern": None,
+        },
+    )()
+
+    all_gather_returns = [
+        torch.stack([rank0_fs, rank1_fs]),  # all_final_state: [2, N, H, K, V]
+        torch.stack([phi_0, phi_1]),  # all_final_h_update: [2, N, H, K, K]
+    ]
+
+    group = type(
+        "Group",
+        (),
+        {
+            "world_size": 2,
+            "rank_in_group": 0,
+            "all_gather": lambda self, value, dim: all_gather_returns.pop(0),
+        },
+    )()
+
+    monkeypatch.setattr(chunk, "get_forward_context", lambda: type("Ctx", (), {"attn_metadata": None})())
+    monkeypatch.setattr(chunk, "get_pcp_group", lambda: group)
+    monkeypatch.setattr(chunk, "chunk_local_cumsum", lambda *a, **kw: _DummyTensor("g_cumsum"))
+    monkeypatch.setattr(chunk, "chunk_scaled_dot_kkt_fwd", lambda *a, **kw: _DummyTensor("A"))
+    monkeypatch.setattr(chunk, "solve_tril", lambda *a, **kw: _DummyTensor("A_solved"))
+    monkeypatch.setattr(chunk, "recompute_w_u_fwd", lambda *a, **kw: (_DummyTensor("w"), _DummyTensor("u")))
+    monkeypatch.setattr(
+        torch.ops._C_ascend,
+        "chunk_gated_delta_rule_fwd_h",
+        lambda *a, **kw: (_DummyTensor("h"), _DummyTensor("v_new"), rank0_fs),
+        raising=False,
+    )
+    monkeypatch.setattr(
+        chunk,
+        "chunk_gated_delta_rule_fwd_hupdate",
+        lambda *a, **kw: h_update_tensor,
+    )
+    monkeypatch.setattr(
+        torch.ops._C_ascend,
+        "chunk_fwd_o",
+        lambda *a, **kw: _DummyTensor("o_ascendc"),
+        raising=False,
+    )
+
+    result = chunk.chunk_gated_delta_rule_fwd(
+        q=_DummyTensor("q"),
+        k=_DummyTensor("k"),
+        v=_DummyTensor("v"),
+        g=_DummyTensor("g"),
+        beta=_DummyTensor("beta"),
+        scale=1.0,
+        initial_state=s0,
+        output_final_state=False,
+        cu_seqlens=torch.tensor([0, N], dtype=torch.int32),
+        prebuilt_meta=prebuilt_meta,
+    )
+
+    final_state = result[3]
+    # Sequential: Φ_1·(Φ_0·s0 + p_0) + p_1
+    expected = torch.matmul(phi_1, torch.matmul(phi_0, s0) + p_0) + p_1
+    torch.testing.assert_close(final_state, expected, rtol=1e-4, atol=1e-4)

tests/ut/ops/a2/test_token_dispatcher.py

@@ -0,0 +1,827 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+
+from unittest.mock import MagicMock, PropertyMock, patch
+
+import numpy as np
+import pytest
+import torch
+
+from tests.ut.base import TestBase
+from vllm_ascend.ops.fused_moe.moe_runtime_args import (
+    MoEAllGatherCombineMetadata,
+    MoEAllToAllCombineMetadata,
+    MoEMC2CombineMetadata,
+    MoEQuantParams,
+    MoERoutingParams,
+    MoETokenDispatchInput,
+)
+
+from vllm_ascend.ops.fused_moe.token_dispatcher import (  # isort: skip
+    AscendDeviceType,
+    EXPERT_TOKEN_NUMS_TYPE_COUNT,
+    EXPERT_TOKEN_NUMS_TYPE_CUMSUM,
+    TokenDispatcherWithAll2AllV,
+    TokenDispatcherWithAllGather,
+    TokenDispatcherWithMC2,
+)
+from vllm_ascend.ops.fused_moe.moe_stage_params import MoEMxfpParams
+from vllm_ascend.quantization.quant_type import QuantType
+
+MXFP4_TEST_DTYPE = getattr(torch, "float4_e2m1fn_x2", torch.float16)
+
+
+def build_token_dispatch_input_fixture(
+    *,
+    hidden_states: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    expert_map: torch.Tensor | None = None,
+    global_redundant_expert_num: int = 0,
+    apply_router_weight_on_input: bool = False,
+    pertoken_scale: torch.Tensor | None = None,
+    quant_type: QuantType = QuantType.NONE,
+    comm_quant_mode: int | None = None,
+    act_quant_type: torch.dtype | None = None,
+    is_per_channel_weight: bool = False,
+    mc2_mask: torch.Tensor | None = None,
+) -> MoETokenDispatchInput:
+    mxfp_spec = None
+    if quant_type in (QuantType.MXFP8, QuantType.MXFP4):
+        mxfp_spec = MoEMxfpParams(act_quant_type=act_quant_type)
+    return MoETokenDispatchInput(
+        hidden_states=hidden_states,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        routing=MoERoutingParams(
+            expert_map=expert_map,
+            global_redundant_expert_num=global_redundant_expert_num,
+            mc2_mask=mc2_mask,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+            pertoken_scale=pertoken_scale,
+        ),
+        quant=MoEQuantParams(
+            quant_type=quant_type,
+            comm_quant_mode=comm_quant_mode,
+            mxfp=mxfp_spec,
+            is_per_channel_weight=is_per_channel_weight,
+        ),
+    )
+
+
+class TestTokenDispatcherWithMC2(TestBase):
+    def setUp(self):
+        self.config_patcher = patch("vllm_ascend.ops.fused_moe.token_dispatcher.get_current_vllm_config")
+        self.mock_get_config = self.config_patcher.start()
+
+        mock_config = MagicMock()
+
+        mock_config.scheduler_config.max_num_seqs = 256
+
+        mock_config.compilation_config.custom_ops = ["all"]
+
+        mock_config.speculative_config = None
+
+        mock_config.parallel_config.tensor_parallel_size = 1
+
+        self.mock_get_config.return_value = mock_config
+        self.mc2_tokens_capacity = 128
+        self.mc2_capacity_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.get_mc2_tokens_capacity",
+            return_value=self.mc2_tokens_capacity,
+        )
+        self.mock_get_mc2_tokens_capacity = self.mc2_capacity_patch.start()
+
+        self.mc2_group = MagicMock()
+        self.mc2_group.device_group.return_value._get_backend.return_value.get_hccl_comm_name.return_value = "hccl_123"
+        self.mc2_group.rank_in_group = 0
+        self.mc2_group.world_size = 8
+        self.mc2_group_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.get_mc2_group", return_value=self.mc2_group
+        )
+        self.mc2_group_patch.start()
+
+        self.rank_group_patch = patch("torch.distributed.get_rank", return_value=0)
+        self.rank_group_patch.start()
+
+        # Mock get_forward_context().mc2_mask
+        self.forward_context = MagicMock()
+        self.forward_context.mc2_mask = torch.tensor([1, 0, 1])
+        self.forward_context_patch = patch(
+            "vllm.forward_context.get_forward_context", return_value=self.forward_context
+        )
+        self.forward_context_patch.start()
+
+        # Mock get_ascend_device_type()
+        self.ascend_soc_version_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.get_ascend_device_type", return_value=AscendDeviceType.A3
+        )
+        self.ascend_soc_version_patch.start()
+
+        # Mock get_ascend_config() and is_hierarchical_communication_enabled()
+        mock_ascend_config = MagicMock()
+        mock_ascend_config.enable_mc2_hierarchy_comm = False
+        mock_ascend_config.eplb_config = MagicMock()
+        mock_ascend_config.eplb_config.dynamic_eplb = False
+        self.ascend_config_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.get_ascend_config", return_value=mock_ascend_config
+        )
+        self.ascend_config_patch.start()
+        self.ascend_config_utils_patch = patch("vllm_ascend.utils.get_ascend_config", return_value=mock_ascend_config)
+        self.ascend_config_utils_patch.start()
+        self.hier_comm_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.is_hierarchical_communication_enabled", return_value=False
+        )
+        self.hier_comm_patch.start()
+        self.skip_allreduce_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.should_skip_allreduce_across_dp_group", return_value=False
+        )
+        self.mock_skip_allreduce = self.skip_allreduce_patch.start()
+
+        kwargs = {"with_quant": False, "top_k": 8, "num_experts": 128}
+        self.dispatcher = TokenDispatcherWithMC2(**kwargs)
+
+    def tearDown(self):
+        self.config_patcher.stop()
+        self.mc2_capacity_patch.stop()
+        self.mc2_group_patch.stop()
+        self.rank_group_patch.stop()
+        self.forward_context_patch.stop()
+        self.ascend_soc_version_patch.stop()
+        self.ascend_config_patch.stop()
+        self.ascend_config_utils_patch.stop()
+        self.hier_comm_patch.stop()
+        self.skip_allreduce_patch.stop()
+
+    def test_init(self):
+        self.assertEqual(self.dispatcher.ep_rank_id, 0)
+        self.assertEqual(self.dispatcher.ep_world_size, 8)
+        self.assertTrue(self.dispatcher.enable_dispatch_v2)
+        self.assertTrue(self.dispatcher.need_extra_args)
+        self.assertEqual(self.dispatcher.global_bs, 0)
+
+    def test_init_uses_mc2_capacity_for_non_uniform_global_bs(self):
+        self.mock_get_config.return_value.parallel_config.tensor_parallel_size = 4
+        self.mock_skip_allreduce.return_value = True
+
+        dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128)
+
+        self.assertEqual(dispatcher.global_bs, 256)
+
+    def test_get_dispatch_mc2_kwargs_with_skip_allreduce_omits_mc2_mask(self):
+        self.mock_get_config.return_value.parallel_config.tensor_parallel_size = 4
+        self.mock_skip_allreduce.return_value = True
+        dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128)
+
+        hidden_states = torch.randn(10, 128)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        topk_weights = torch.randn(10, 1)
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        mc2_mask = torch.tensor([True, False, True, False])
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            mc2_mask=mc2_mask,
+        )
+
+        kwargs = dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+
+        self.assertEqual(kwargs["global_bs"], 256)
+        self.assertNotIn("x_active_mask", kwargs)
+
+    def test_get_dispatch_mc2_kwargs_without_skip_allreduce_keeps_mc2_mask(self):
+        hidden_states = torch.randn(10, 128)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        topk_weights = torch.randn(10, 1)
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        mc2_mask = torch.tensor([True, False, True, False])
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            mc2_mask=mc2_mask,
+        )
+
+        kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+
+        self.assertEqual(kwargs["global_bs"], 0)
+        self.assertIs(kwargs["x_active_mask"], mc2_mask)
+
+    def test_get_dispatch_mc2_kwargs_without_quant(self):
+        hidden_states = torch.randn(10, 128)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        topk_weights = torch.randn(10, 1)
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            global_redundant_expert_num=0,
+            apply_router_weight_on_input=False,
+            pertoken_scale=None,
+        )
+        kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+        self.assertIn("x", kwargs)
+        self.assertIn("expert_ids", kwargs)
+        self.assertEqual(kwargs["moe_expert_num"], 8)
+
+    def test_token_permutation_dispatch(self):
+        hidden_states = torch.randn(10, 128)
+        topk_weights = torch.randn(10, 1)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+
+        with patch(
+            "torch_npu.npu_moe_distribute_dispatch_v2", return_value=(torch.randn(10, 128),) * 5 + (None, None)
+        ) as mock_dispatch:
+            token_dispatch_input = build_token_dispatch_input_fixture(
+                hidden_states=hidden_states,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                expert_map=expert_map,
+            )
+            output = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+            mock_dispatch.assert_called_once()
+            self.assertEqual(output.group_list_type, 0)  # group_list_type == 0
+            self.assertIsInstance(output.combine_metadata, MoEMC2CombineMetadata)
+
+    def test_w4a8_per_channel_dispatch_uses_count_group_list(self):
+        hidden_states = torch.randn(10, 128)
+        topk_weights = torch.randn(10, 1)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        self.dispatcher.enable_dispatch_v2 = True
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            quant_type=QuantType.W4A8,
+            is_per_channel_weight=True,
+        )
+        with patch(
+            "torch_npu.npu_moe_distribute_dispatch_v2", return_value=(torch.randn(10, 128),) * 5 + (None, None)
+        ) as mock_dispatch:
+            output = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        mock_dispatch.assert_called_once()
+        self.assertEqual(mock_dispatch.call_args.kwargs["expert_token_nums_type"], EXPERT_TOKEN_NUMS_TYPE_COUNT)
+        self.assertEqual(output.group_list_type, EXPERT_TOKEN_NUMS_TYPE_COUNT)
+
+    def test_w4a8_group_dispatch_keeps_prefix_sum_group_list(self):
+        hidden_states = torch.randn(10, 128)
+        topk_weights = torch.randn(10, 1)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            quant_type=QuantType.W4A8,
+            is_per_channel_weight=False,
+        )
+        kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+
+        self.assertEqual(kwargs["expert_token_nums_type"], EXPERT_TOKEN_NUMS_TYPE_CUMSUM)
+
+    def test_get_combine_mc_kwargs_with_quant(self):
+        hidden_states = torch.randn(10, 128)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        topk_weights = torch.randn(10, 1)
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        ep_recv_counts = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        tp_recv_counts = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+        assist_info_for_combine = torch.arange(10)
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+        )
+
+        combine_metadata = MoEMC2CombineMetadata(
+            topk_ids=topk_ids,
+            topk_weights=topk_weights,
+            expert_map=expert_map,
+            ep_recv_counts=ep_recv_counts,
+            tp_recv_counts=tp_recv_counts,
+            assist_info_for_combine=assist_info_for_combine,
+            expand_scales=None,
+            quant=token_dispatch_input.quant,
+        )
+
+        self.dispatcher.need_extra_args = True
+        self.dispatcher.enable_dispatch_v2 = True
+        self.dispatcher.moe_expert_num = len(expert_map)
+        kwargs = self.dispatcher.get_combine_mc_kwargs(hidden_states, combine_metadata)
+        self.assertIn("tp_send_counts", kwargs)
+
+    def test_get_dispatch_mc2_kwargs_with_mxfp8_quant(self):
+        hidden_states = torch.randn(10, 128)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        topk_weights = torch.randn(10, 1)
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+
+        self.dispatcher.a5_need_extra_args = True
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            quant_type=QuantType.MXFP8,
+            act_quant_type=torch.float8_e4m3fn,
+        )
+
+        kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+
+        self.assertTrue(token_dispatch_input.quant.dispatch_with_quant)
+        self.assertEqual(kwargs["quant_mode"], 4)
+        self.assertEqual(kwargs["y_dtype"], torch.float8_e4m3fn)
+
+    def test_get_dispatch_mc2_kwargs_with_mxfp4_quant(self):
+        hidden_states = torch.randn(10, 128)
+        topk_weights = torch.randn(10, 1)
+        topk_ids = torch.randint(0, 8, (10, 1))
+        expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7])
+
+        self.dispatcher.a5_need_extra_args = True
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            quant_type=QuantType.MXFP4,
+            act_quant_type=MXFP4_TEST_DTYPE,
+        )
+        kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+
+        self.assertTrue(token_dispatch_input.quant.dispatch_with_quant)
+        self.assertEqual(kwargs["quant_mode"], 4)
+        self.assertIn("y_dtype", kwargs)
+        self.assertNotEqual(kwargs["y_dtype"], torch.float8_e4m3fn)
+
+        with patch(
+            "torch_npu.npu_moe_distribute_dispatch_v2",
+            return_value=(
+                torch.randn(10, 128),
+                torch.randn(10, 1),
+                torch.arange(10, dtype=torch.int32),
+                torch.tensor([10], dtype=torch.int64),
+                torch.tensor([10], dtype=torch.int64),
+                torch.tensor([10], dtype=torch.int64),
+                torch.randn(10, 1),
+            ),
+        ) as mock_dispatch:
+            output = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        mock_dispatch.assert_called_once()
+        self.assertIsNotNone(output.dynamic_scale)
+        self.assertTrue(output.combine_metadata.quant.dispatch_with_quant)
+
+
+def test_allgather_token_dispatch_quant_mode_without_dynamic_scale():
+    dispatcher = TokenDispatcherWithAllGather(top_k=2, num_experts=128)
+    hidden_states = torch.randn(3, 128)
+    topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]])
+    topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int32)
+    init_routing_output = (
+        torch.randn(6, 128),
+        torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32),
+        torch.tensor([2, 2, 2], dtype=torch.int32),
+        torch.randn(6, 4),
+    )
+
+    cases = [
+        {
+            "quant_type": QuantType.MXFP8,
+            "act_quant_type": torch.float8_e4m3fn,
+            "expected_quant_mode": 3,
+            "expected_act_quant_type": torch.float8_e4m3fn,
+            "expect_dynamic_scale": True,
+        },
+        {
+            "quant_type": QuantType.MXFP4,
+            "act_quant_type": MXFP4_TEST_DTYPE,
+            "expected_quant_mode": -1,
+            "expected_act_quant_type": None,
+            "expect_dynamic_scale": False,
+        },
+    ]
+    for case in cases:
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            quant_type=case["quant_type"],
+            act_quant_type=case["act_quant_type"],
+        )
+
+        with patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.DeviceOperator.npu_moe_init_routing",
+            return_value=init_routing_output,
+        ) as mock_init_routing:
+            output = dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        init_kwargs = mock_init_routing.call_args.kwargs
+        assert init_kwargs["quant_mode"] == case["expected_quant_mode"]
+        assert init_kwargs["act_quant_type"] == case["expected_act_quant_type"]
+        assert (output.dynamic_scale is not None) == case["expect_dynamic_scale"]
+
+
+def test_allgather_token_dispatch_mxfp4_keeps_prequantized_scale():
+    dispatcher = TokenDispatcherWithAllGather(top_k=2, num_experts=128)
+    hidden_states = torch.randn(3, 128)
+    topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]])
+    topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int32)
+    pertoken_scale = torch.randn(3, 4)
+    returned_scale = torch.randn(6, 4)
+    init_routing_output = (
+        torch.randn(6, 128),
+        torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32),
+        torch.tensor([2, 2, 2], dtype=torch.int32),
+        returned_scale,
+    )
+
+    token_dispatch_input = build_token_dispatch_input_fixture(
+        hidden_states=hidden_states,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        pertoken_scale=pertoken_scale,
+        quant_type=QuantType.MXFP4,
+        act_quant_type=MXFP4_TEST_DTYPE,
+    )
+
+    with patch(
+        "vllm_ascend.ops.fused_moe.token_dispatcher.DeviceOperator.npu_moe_init_routing",
+        return_value=init_routing_output,
+    ) as mock_init_routing:
+        output = dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+    init_kwargs = mock_init_routing.call_args.kwargs
+    assert init_kwargs["scale"] is pertoken_scale
+    assert init_kwargs["quant_mode"] == -1
+    assert init_kwargs["act_quant_type"] == MXFP4_TEST_DTYPE
+    assert output.dynamic_scale is returned_scale
+
+
+class TestTokenDispatcherWithAllGather(TestBase):
+    def setUp(self):
+        # Mock dependencies
+        kwargs = {
+            "apply_router_weight_on_input": False,
+            "top_k": 2,
+            "max_num_tokens": 100,
+            "ep_size": 2,
+            "num_experts": 128,
+            "with_quant": False,
+        }
+        self.dispatcher = TokenDispatcherWithAllGather(**kwargs)
+
+        # Mock NPU functions
+        self.patcher_npu_moe_init_routing_custom = patch("torch.ops._C_ascend.npu_moe_init_routing_custom")
+        self.mock_npu_moe_init_routing_custom = self.patcher_npu_moe_init_routing_custom.start()
+        self.mock_npu_moe_init_routing_custom.return_value = (
+            torch.randn(6, 128),  # sorted_hidden_states
+            torch.tensor([0, 1, 2, 3, 4, 5]),  # expanded_row_idx
+            torch.tensor([0, 1, 0, 1, 0, 1]),  # expanded_expert_idx
+            torch.tensor([0, 1, 0, 1, 0, 1]),
+        )
+        self.patcher_npu_moe_token_unpermute = patch("torch_npu.npu_moe_token_unpermute")
+        self.mock_npu_moe_token_unpermute = self.patcher_npu_moe_token_unpermute.start()
+        self.mock_npu_moe_token_unpermute.return_value = torch.randn(6, 128)
+
+    def tearDown(self):
+        self.patcher_npu_moe_init_routing_custom.stop()
+        self.patcher_npu_moe_token_unpermute.stop()
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_without_expert_map(self):
+        hidden_states = torch.randn(3, 128)
+        topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]])
+        topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]])
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+        )
+        results = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        # Verify npu_moe_init_routing is called
+        self.mock_npu_moe_init_routing_custom.assert_called_once()
+        args, kwargs = self.mock_npu_moe_init_routing_custom.call_args
+
+        self.assertEqual(results.group_list_type, 1)
+        self.assertIsInstance(results.combine_metadata, MoEAllGatherCombineMetadata)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_with_expert_map(self):
+        self.dispatcher.expert_map = torch.tensor([0, 1, 2, 3])
+        hidden_states = torch.randn(3, 128)
+        topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]])
+        topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]])
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+        )
+        results = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        # Verify npu_moe_init_routing is called
+        self.mock_npu_moe_init_routing_custom.assert_called_once()
+        args, kwargs = self.mock_npu_moe_init_routing_custom.call_args
+
+        self.assertEqual(results.group_list_type, 1)
+        self.assertIsInstance(results.combine_metadata, MoEAllGatherCombineMetadata)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_without_quant(self):
+        kwargs = {
+            "apply_router_weight_on_input": False,
+            "top_k": 2,
+            "max_num_tokens": 100,
+            "ep_size": 2,
+            "num_experts": 128,
+        }
+        self.dispatcher_quant = TokenDispatcherWithAllGather(**kwargs)
+
+        hidden_states = torch.randn(3, 128)
+        topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]])
+        topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]])
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+        )
+        results = self.dispatcher_quant.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        self.assertEqual(results.group_list_type, 1)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_with_quant(self):
+        kwargs = {
+            "apply_router_weight_on_input": False,
+            "top_k": 2,
+            "max_num_tokens": 100,
+            "ep_size": 2,
+            "num_experts": 128,
+        }
+        self.dispatcher_quant = TokenDispatcherWithAllGather(**kwargs)
+
+        hidden_states = torch.randn(3, 128)
+        topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]])
+        topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]])
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            quant_type=QuantType.W8A8,
+        )
+        results = self.dispatcher_quant.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        self.assertIsNotNone(results.hidden_states)
+        self.assertIsNotNone(results.group_list)
+        self.assertIsNotNone(results.dynamic_scale)
+        self.assertEqual(results.group_list_type, 1)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_combine_with_expert_map(self):
+        hidden_states = torch.randn(6, 128)
+        combine_metadata = MoEAllGatherCombineMetadata(
+            expanded_row_idx=torch.tensor([0, 1, 1, 1, 1, 1]),
+            topk_weights=torch.tensor([0.5, 0.5, 0.5, 0.5, 0.5, 0.5]),
+            restore_shape=torch.Size([6, 128]),
+        )
+        final_hidden_states = self.dispatcher.token_combine(hidden_states, combine_metadata)
+        self.assertEqual(final_hidden_states.shape, (6, 128))
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_combine_without_expert_map(self):
+        hidden_states = torch.randn(6, 128)
+        combine_metadata = MoEAllGatherCombineMetadata(
+            expanded_row_idx=torch.tensor([0, 1, 1, 1, 1, 1]),
+            topk_weights=torch.tensor([0.5, 0.5, 0.5, 0.5, 0.5, 0.5]),
+            restore_shape=torch.Size([6, 128]),
+        )
+        final_hidden_states = self.dispatcher.token_combine(hidden_states, combine_metadata)
+        self.mock_npu_moe_token_unpermute.assert_called_once()
+        self.assertEqual(final_hidden_states.shape, (6, 128))
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_with_router_weight(self):
+        hidden_states = torch.randn(3, 128)
+        topk_weights = torch.tensor([[0.7], [0.6], [0.5]])  # topk=1
+        topk_ids = torch.tensor([[0], [1], [2]])
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            apply_router_weight_on_input=True,
+        )
+        results = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+        self.assertEqual(results.hidden_states.shape, (6, 128))
+        self.assertIsInstance(results.combine_metadata, MoEAllGatherCombineMetadata)
+
+
+class TestTokenDispatcherWithAll2AllV(TestBase):
+    def setUp(self):
+        # Patch properties
+        patcher1 = patch.object(
+            TokenDispatcherWithAll2AllV, "ep_group", new_callable=PropertyMock, return_value=MagicMock()
+        )
+        patcher2 = patch.object(TokenDispatcherWithAll2AllV, "ep_rank", new_callable=PropertyMock, return_value=0)
+        patcher3 = patch.object(TokenDispatcherWithAll2AllV, "ep_size", new_callable=PropertyMock, return_value=2)
+
+        self.addCleanup(patcher1.stop)
+        self.addCleanup(patcher2.stop)
+        self.addCleanup(patcher3.stop)
+
+        self.mock_ep_group_prop = patcher1.start()
+        self.mock_ep_rank_prop = patcher2.start()
+        self.mock_ep_size_prop = patcher3.start()
+
+        # Mock torch_npu.npu_moe_token_permute
+        patcher4 = patch("torch_npu.npu_moe_token_permute")
+        self.mock_npu_moe_token_permute = patcher4.start()
+        self.addCleanup(patcher4.stop)
+        self.mock_npu_moe_token_permute.return_value = (torch.randn(16, 16), torch.arange(16))
+
+        # Mock torch_npu.npu_moe_token_unpermute
+        patcher5 = patch("torch_npu.npu_moe_token_unpermute")
+        self.mock_npu_moe_token_unpermute = patcher5.start()
+        self.addCleanup(patcher5.stop)
+        self.mock_npu_moe_token_unpermute.return_value = torch.randn(8, 16)
+
+        # Mock async_all_to_all
+        patcher6 = patch("vllm_ascend.ops.fused_moe.comm_utils.async_all_to_all")
+        self.mock_async_all_to_all = patcher6.start()
+        self.addCleanup(patcher6.stop)
+        self.mock_async_all_to_all.return_value = (None, torch.randn(16, 16), MagicMock())
+
+        # Mock gather_from_sequence_parallel_region
+        patcher7 = patch("vllm_ascend.ops.fused_moe.token_dispatcher.gather_from_sequence_parallel_region")
+        self.mock_gather_from_sequence_parallel_region = patcher7.start()
+        self.addCleanup(patcher7.stop)
+        self.mock_gather_from_sequence_parallel_region.return_value = torch.tensor(
+            [[2, 2, 2, 2], [2, 2, 2, 2]], dtype=torch.int64
+        )
+
+        # Mock torch.histc
+        patcher8 = patch("torch.histc")
+        self.mock_histc = patcher8.start()
+        self.addCleanup(patcher8.stop)
+        self.mock_histc.return_value = torch.tensor([2, 2, 2, 2], dtype=torch.int64)
+
+        # Mock torch.npu.current_device
+        patcher9 = patch("torch.npu.current_device")
+        self.mock_current_device = patcher9.start()
+        self.addCleanup(patcher9.stop)
+        self.mock_current_device.return_value = "cpu"
+
+        # Mock torch_npu.npu_dynamic_quant
+        patcher10 = patch("torch_npu.npu_dynamic_quant")
+        self.mock_npu_dynamic_quant = patcher10.start()
+        self.addCleanup(patcher10.stop)
+        self.mock_npu_dynamic_quant.return_value = (torch.randn(16, 16), torch.randn(16))
+
+        # Mock torch.ops._C_ascend.npu_moe_init_routing_custom
+        patcher11 = patch("torch.ops._C_ascend.npu_moe_init_routing_custom")
+        self.mock_npu_moe_init_routing_custom = patcher11.start()
+        self.addCleanup(patcher11.stop)
+        self.mock_npu_moe_init_routing_custom.return_value = (
+            torch.randn(16, 16),
+            torch.arange(16),
+            None,
+            torch.randn(16),
+        )
+
+        # Mock torch.repeat_interleave
+        patcher12 = patch("torch.repeat_interleave")
+        self.mock_repeat_interleave = patcher12.start()
+        self.addCleanup(patcher12.stop)
+        self.mock_repeat_interleave.return_value = torch.arange(16)
+
+        self.dispatcher = TokenDispatcherWithAll2AllV(top_k=2, num_experts=4, num_local_experts=2, with_quant=False)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch(self):
+        hidden_states = torch.randn(8, 16)
+        topk_weights = torch.rand(8, 4)
+        topk_ids = torch.randint(0, 4, (8, 2)).long()
+        expert_map = torch.tensor([0, 1, 2, 3])
+
+        self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32)
+        self.dispatcher.local_expert_indices = [0, 1]
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+        )
+        result = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        self.assertIsNotNone(result.hidden_states)
+        self.assertIsNotNone(result.group_list)
+        self.assertEqual(result.group_list_type, 1)
+        self.assertIsInstance(result.combine_metadata, MoEAllToAllCombineMetadata)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_combine(self):
+        hidden_states = torch.randn(16, 16)
+        combine_metadata = MoEAllToAllCombineMetadata(
+            input_splits=np.array([4, 4]),
+            output_splits=np.array([4, 4]),
+            topk_weights=torch.rand(8, 4),
+            reversed_local_input_permutation_mapping=torch.arange(8),
+            reversed_global_input_permutation_mapping=torch.arange(16),
+            hidden_shape=torch.Size([8, 16]),
+            hidden_shape_before_permute=torch.Size([8, 16]),
+        )
+        self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32)
+        self.dispatcher.local_expert_indices = [0, 1]
+
+        output = self.dispatcher.token_combine(hidden_states, combine_metadata)
+        self.assertIsNotNone(output)
+        self.assertEqual(output.shape, (8, 16))
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_with_quant(self):
+        self.dispatcher = TokenDispatcherWithAll2AllV(top_k=2, num_experts=4, num_local_experts=2)
+
+        hidden_states = torch.randn(8, 16)
+        topk_weights = torch.rand(8, 4)
+        topk_ids = torch.randint(0, 4, (8, 2)).long()
+        expert_map = torch.tensor([0, 1, 2, 3])
+
+        self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32)
+        self.dispatcher.local_expert_indices = [0, 1]
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            quant_type=QuantType.W8A8,
+        )
+        result = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        self.assertIsNotNone(result.hidden_states)
+        self.assertIsNotNone(result.group_list)
+        self.assertIsNotNone(result.dynamic_scale)
+        self.assertEqual(result.group_list_type, 1)
+        self.assertIsInstance(result.combine_metadata, MoEAllToAllCombineMetadata)
+
+    @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.")
+    def test_token_dispatch_with_quant_no_active_tokens(self):
+        self.dispatcher = TokenDispatcherWithAll2AllV(top_k=2, num_experts=4, num_local_experts=2)
+
+        self.mock_repeat_interleave.return_value = torch.tensor([], dtype=torch.long)
+
+        hidden_states = torch.randn(8, 16)
+        topk_weights = torch.rand(8, 4)
+        topk_ids = torch.randint(0, 4, (8, 2)).long()
+        expert_map = torch.tensor([0, 1, 2, 3])
+
+        self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32)
+        self.dispatcher.local_expert_indices = [0, 1]
+
+        token_dispatch_input = build_token_dispatch_input_fixture(
+            hidden_states=hidden_states,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            expert_map=expert_map,
+            quant_type=QuantType.W8A8,
+        )
+        result = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input)
+
+        self.assertIsNotNone(result.hidden_states)
+        self.assertIsNotNone(result.group_list)
+        self.assertIsNotNone(result.dynamic_scale)
+        self.assertEqual(result.group_list_type, 1)

tests/ut/ops/a2/test_weight_prefetch.py

@@ -0,0 +1,458 @@
+#
+# 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 unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_ascend.ops.weight_prefetch import (
+    MAX_PREFETCH_WEIGHT_SIZE,
+    MOE_PREFETCH_TOKEN_THRESHOLD,
+    SUPPORTED_MODULES,
+    ModuleWeightPrefetchConfig,
+    WeightPrefetchMethod,
+    maybe_npu_prefetch,
+)
+
+
+class TestModuleWeightPrefetchConfig:
+    def test_init_with_valid_module_name(self):
+        for module_name in SUPPORTED_MODULES:
+            config = ModuleWeightPrefetchConfig(module_name=module_name)
+            assert config.module_name == module_name
+            assert config.enable is False
+            assert config.is_active_this_forward is False
+            assert config.prefetch_ratio == {}
+            assert config.linear_prefix_map == {}
+
+    def test_init_with_invalid_module_name(self):
+        with pytest.raises(AssertionError, match="Invalid module name"):
+            ModuleWeightPrefetchConfig(module_name="invalid_module")
+
+    def test_prefetch_ratio_filtering(self):
+        config = ModuleWeightPrefetchConfig(
+            module_name="attn",
+            prefetch_ratio={"qkv": 0.8, "o": 1.2, "invalid": -0.5},
+        )
+        assert "qkv" in config.prefetch_ratio
+        assert "o" not in config.prefetch_ratio
+        assert "invalid" not in config.prefetch_ratio
+
+    def test_enable_logic_with_prefetch_ratio(self):
+        config = ModuleWeightPrefetchConfig(
+            module_name="attn",
+            enable=True,
+            prefetch_ratio={"qkv": 0.8},
+        )
+        assert config.enable is True
+
+    def test_enable_logic_without_prefetch_ratio(self):
+        config = ModuleWeightPrefetchConfig(
+            module_name="attn",
+            enable=True,
+            prefetch_ratio={},
+        )
+        assert config.enable is False
+
+
+class TestWeightPrefetchMethod:
+    @pytest.fixture
+    def mock_weight_prefetch_config(self):
+        config = MagicMock()
+        config.enabled = True
+        config.prefetch_ratio = {
+            "attn": {"qkv": 0.8, "o": 0.8},
+            "moe": {"gate_up": 0.8},
+            "mlp": {"gate_up": 1.0, "down": 1.0},
+        }
+        return config
+
+    @pytest.fixture
+    def mock_vllm_config(self):
+        config = MagicMock()
+        config.model_config = MagicMock()
+        config.model_config.hf_config = MagicMock()
+        config.model_config.hf_config.model_type = "llama"
+        return config
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    def test_init_non_moe_model(
+        self,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        assert method.is_moe is False
+        assert method.attn.enable is True
+        assert method.moe.enable is False
+        assert method.mlp.enable is True
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=True)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    def test_init_moe_model(
+        self,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        assert method.is_moe is True
+        assert method.attn.enable is True
+        assert method.moe.enable is True
+        assert method.mlp.enable is False
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    def test_init_disabled_config(
+        self,
+        mock_get_config,
+        mock_is_moe,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        disabled_config = MagicMock()
+        disabled_config.enabled = False
+        disabled_config.prefetch_ratio = {}
+
+        method = WeightPrefetchMethod(disabled_config)
+
+        assert method.attn.enable is False
+        assert method.moe.enable is False
+        assert method.mlp.enable is False
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_preprocess")
+    def test_maybe_prefetch_attn_weight_preprocess_enabled(
+        self,
+        mock_prefetch,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        weight = torch.randn(1024, 1024)
+        start_flag = torch.tensor([1])
+
+        method.maybe_prefetch_attn_weight_preprocess(
+            layer_cls_name="AscendQKVParallelLinear",
+            weight=weight,
+            start_flag=start_flag,
+        )
+
+        mock_prefetch.assert_called_once()
+        call_kwargs = mock_prefetch.call_args[1]
+        assert call_kwargs["weight"] is weight
+        assert call_kwargs["start_flag"] is start_flag
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_preprocess")
+    def test_maybe_prefetch_attn_weight_preprocess_disabled(
+        self,
+        mock_prefetch,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        disabled_config = MagicMock()
+        disabled_config.enabled = False
+        disabled_config.prefetch_ratio = {}
+
+        method = WeightPrefetchMethod(disabled_config)
+
+        weight = torch.randn(1024, 1024)
+        start_flag = torch.tensor([1])
+
+        method.maybe_prefetch_attn_weight_preprocess(
+            layer_cls_name="AscendQKVParallelLinear",
+            weight=weight,
+            start_flag=start_flag,
+        )
+
+        mock_prefetch.assert_not_called()
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_postprocess")
+    def test_maybe_prefetch_attn_weight_postprocess_enabled(
+        self,
+        mock_postprocess,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        stop_flag = torch.tensor([1])
+        method.maybe_prefetch_attn_weight_postprocess(
+            layer_cls_name="AscendQKVParallelLinear",
+            stop_flag=stop_flag,
+        )
+
+        mock_postprocess.assert_called_once_with(stop_flag)
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_postprocess")
+    def test_maybe_prefetch_attn_weight_postprocess_disabled(
+        self,
+        mock_postprocess,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        disabled_config = MagicMock()
+        disabled_config.enabled = False
+        disabled_config.prefetch_ratio = {}
+
+        method = WeightPrefetchMethod(disabled_config)
+
+        stop_flag = torch.tensor([1])
+        method.maybe_prefetch_attn_weight_postprocess(
+            layer_cls_name="AscendQKVParallelLinear",
+            stop_flag=stop_flag,
+        )
+
+        mock_postprocess.assert_not_called()
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=True)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("vllm_ascend.ops.weight_prefetch.get_forward_context")
+    @patch("vllm_ascend.ops.weight_prefetch._EXTRA_CTX")
+    @patch("torch.ops.vllm.prefetch_preprocess")
+    def test_maybe_prefetch_moe_weight_preprocess_enabled(
+        self,
+        mock_prefetch,
+        mock_extra_ctx,
+        mock_get_forward_context,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        mock_get_forward_context.return_value = MagicMock()
+
+        mock_model_instance = MagicMock()
+        mock_layer = MagicMock()
+        mock_layer.mlp.experts.w13_weight = torch.randn(1024, 1024)
+        mock_model_instance.model.layers = [mock_layer]
+        mock_extra_ctx.model_instance = mock_model_instance
+        mock_extra_ctx.layer_idx = 1
+
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        hidden_states = torch.randn(MOE_PREFETCH_TOKEN_THRESHOLD + 10, 1024)
+        method.maybe_prefetch_moe_weight_preprocess(hidden_states, "gate_up")
+
+        assert method.moe.is_active_this_forward is True
+        mock_prefetch.assert_called_once()
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=True)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    def test_maybe_prefetch_moe_weight_preprocess_below_threshold(
+        self,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        hidden_states = torch.randn(MOE_PREFETCH_TOKEN_THRESHOLD - 10, 1024)
+        method.maybe_prefetch_moe_weight_preprocess(hidden_states, "gate_up")
+
+        assert method.moe.is_active_this_forward is False
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=True)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_postprocess")
+    def test_maybe_prefetch_moe_weight_postprocess_enabled(
+        self,
+        mock_postprocess,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+        method.moe.is_active_this_forward = True
+
+        stop_flag = torch.tensor([1])
+        method.maybe_prefetch_moe_weight_postprocess(stop_flag)
+
+        mock_postprocess.assert_called_once_with(stop_flag)
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_postprocess")
+    def test_maybe_prefetch_moe_weight_postprocess_disabled(
+        self,
+        mock_postprocess,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+        method.moe.is_active_this_forward = False
+
+        stop_flag = torch.tensor([1])
+        method.maybe_prefetch_moe_weight_postprocess(stop_flag)
+
+        mock_postprocess.assert_not_called()
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_preprocess")
+    def test_maybe_prefetch_mla_or_sla_weight_enabled(
+        self,
+        mock_prefetch,
+        mock_get_config,
+        mock_is_moe,
+        mock_weight_prefetch_config,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        method = WeightPrefetchMethod(mock_weight_prefetch_config)
+
+        inputs = torch.randn(1024, 1024)
+        dependency = torch.tensor([1])
+
+        method.maybe_prefetch_mla_or_sla_weight_in_current_stream(
+            inputs=inputs,
+            dependency=dependency,
+            max_size=1024,
+        )
+
+        mock_prefetch.assert_called_once()
+
+    @patch("vllm_ascend.ops.weight_prefetch.is_moe_model", return_value=False)
+    @patch("vllm_ascend.ops.weight_prefetch.get_current_vllm_config")
+    @patch("torch.ops.vllm.prefetch_preprocess")
+    def test_maybe_prefetch_mla_or_sla_weight_disabled(
+        self,
+        mock_prefetch,
+        mock_get_config,
+        mock_is_moe,
+        mock_vllm_config,
+    ):
+        mock_get_config.return_value = mock_vllm_config
+        disabled_config = MagicMock()
+        disabled_config.enabled = False
+        disabled_config.prefetch_ratio = {}
+
+        method = WeightPrefetchMethod(disabled_config)
+
+        inputs = torch.randn(1024, 1024)
+        dependency = torch.tensor([1])
+
+        method.maybe_prefetch_mla_or_sla_weight_in_current_stream(
+            inputs=inputs,
+            dependency=dependency,
+        )
+
+        mock_prefetch.assert_not_called()
+
+
+class TestMaybeNpuPrefetch:
+    @patch("torch_npu.npu_prefetch")
+    def test_maybe_npu_prefetch_enabled(self, mock_prefetch):
+        inputs = torch.randn(1024, 1024)
+        dependency = torch.tensor([1])
+
+        maybe_npu_prefetch(inputs, dependency, enabled=True)
+
+        mock_prefetch.assert_called_once()
+        call_args = mock_prefetch.call_args[0]
+        assert call_args[0] is inputs
+        assert call_args[1] is dependency
+
+    @patch("torch_npu.npu_prefetch")
+    def test_maybe_npu_prefetch_disabled(self, mock_prefetch):
+        inputs = torch.randn(1024, 1024)
+        dependency = torch.tensor([1])
+
+        maybe_npu_prefetch(inputs, dependency, enabled=False)
+
+        mock_prefetch.assert_not_called()
+
+    @patch("torch_npu.npu_prefetch")
+    def test_maybe_npu_prefetch_max_size_calculation(self, mock_prefetch):
+        inputs = torch.randn(100, 100, dtype=torch.float32)
+        dependency = torch.tensor([1])
+
+        maybe_npu_prefetch(inputs, dependency, max_size=0, enabled=True)
+
+        expected_size = inputs.element_size() * inputs.numel()
+        call_args = mock_prefetch.call_args[0]
+        assert call_args[2] == expected_size
+
+    @patch("torch_npu.npu_prefetch")
+    def test_maybe_npu_prefetch_with_custom_max_size(self, mock_prefetch):
+        inputs = torch.randn(100, 100, dtype=torch.float32)
+        dependency = torch.tensor([1])
+        custom_max_size = 1000
+
+        maybe_npu_prefetch(inputs, dependency, max_size=custom_max_size, enabled=True)
+
+        call_args = mock_prefetch.call_args[0]
+        assert call_args[2] == custom_max_size
+
+    @patch("torch_npu.npu_prefetch")
+    def test_maybe_npu_prefetch_with_offset(self, mock_prefetch):
+        inputs = torch.randn(1024, 1024)
+        dependency = torch.tensor([1])
+        offset = 100
+
+        maybe_npu_prefetch(inputs, dependency, offset=offset, enabled=True)
+
+        call_args = mock_prefetch.call_args[0]
+        assert call_args[3] == offset
+
+
+class TestConstants:
+    def test_supported_modules(self):
+        assert "attn" in SUPPORTED_MODULES
+        assert "mlp" in SUPPORTED_MODULES
+        assert "moe" in SUPPORTED_MODULES
+
+    def test_moe_prefetch_token_threshold(self):
+        assert MOE_PREFETCH_TOKEN_THRESHOLD == 96
+
+    def test_max_prefetch_weight_size(self):
+        assert MAX_PREFETCH_WEIGHT_SIZE == 18 * 1024 * 1024