data: port complete MoE + xllm layer call chains from upstream repos

MoE call chain from ds_vllm (vllm-project/vllm latest):
  ex_engine/moe/ — 20 files, 8736 lines
  - modular_kernel.py (1630 lines) — base classes for modular MoE
  - experts/fused_batched_moe.py (972 lines) — NaiveBatchedExperts
  - prepare_finalize/batched.py (171 lines) — token grouping by expert
  - topk_weight_and_reduce.py (176 lines) — scatter-add finalize
  - fused_moe.py (1740 lines) — main fused_moe dispatch
  - config.py (1407 lines) — FusedMoEQuantConfig
  - activation.py, utils.py, layer.py, etc.

xllm layer code (jd-opensource/xllm):
  ex_engine/xllm_layers/ — 39 files, 5859 lines
  - ilu/fused_moe.cpp (797 lines) — production ixformer 7-step MoE pipeline
  - ilu/attention.cpp (189 lines) — paged_attention + flash_attn bridge
  - npu_torch/qwen3_gated_delta_net_base.cpp (576 lines) — GDN reference
  - common/rms_norm.cpp, rotary_embedding.cpp, activation.cpp, dense_mlp.cpp

xllm ILU kernels — synced 10 files to upstream (diffs from prior edits)

These are reference implementations, NOT hand-written.
Source repos: vllm-project/vllm, jd-opensource/xllm

ex_engine/moe/__init__.py

@@ -1,10 +1,160 @@
-"""
-ex_engine.moe — MoE expert computation for BI-V100
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-Ported from:
-  upstream_ref/ds_vllm/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
-  upstream_ref/ds_vllm/vllm/model_executor/layers/fused_moe/activation.py
-"""
+from contextlib import contextmanager
+from typing import Any
 
-from ex_engine.moe.naive_batched_experts import naive_batched_moe_forward
-from ex_engine.moe.activation import MoEActivation, apply_moe_activation
+from vllm.model_executor.layers.fused_moe.activation import (
+    MoEActivation,
+    activation_without_mul,
+    apply_moe_activation,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+    RoutingMethodType,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.layer import (
+    FusedMoE,
+    fused_moe_make_expert_params_mapping,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEActivationFormat,
+    FusedMoEExpertsModular,
+    FusedMoEPrepareAndFinalizeModular,
+)
+from vllm.model_executor.layers.fused_moe.routed_experts import (
+    FusedMoeWeightScaleSupported,
+    RoutedExperts,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
+    FusedMoERouter,
+)
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import (
+    MoERunner,
+)
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+)
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+    UnquantizedFusedMoEMethod,
+)
+from vllm.triton_utils import HAS_TRITON
+
+_config: dict[str, Any] | None = None
+
+
+@contextmanager
+def override_config(config):
+    global _config
+    old_config = _config
+    _config = config
+    yield
+    _config = old_config
+
+
+def get_config() -> dict[str, Any] | None:
+    return _config
+
+
+__all__ = [
+    "FusedMoE",
+    "FusedMoERouter",
+    "FusedMoEConfig",
+    "FusedMoEQuantConfig",
+    "FusedMoEParallelConfig",
+    "FusedMoEMethodBase",
+    "MoEActivation",
+    "UnquantizedFusedMoEMethod",
+    "FusedMoeWeightScaleSupported",
+    "FusedMoEExpertsModular",
+    "FusedMoEActivationFormat",
+    "FusedMoEPrepareAndFinalizeModular",
+    "GateLinear",
+    "MoERunner",
+    "RoutingMethodType",
+    "RoutedExperts",
+    "SharedExperts",
+    "activation_without_mul",
+    "apply_moe_activation",
+    "fused_moe_make_expert_params_mapping",
+    "override_config",
+    "get_config",
+]
+
+if HAS_TRITON:
+    # import to register the custom ops
+    from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
+        BatchedDeepGemmExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+        CutlassBatchedExpertsFp8,
+        CutlassExpertsFp8,
+        CutlassExpertsW4A8Fp8,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
+        DeepGemmExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
+        BatchedTritonExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
+        AiterExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
+        TritonOrDeepGemmExperts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
+        TritonExperts,
+        TritonWNA16Experts,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
+        XPUExperts,
+        XPUExpertsFp8,
+        XPUExpertsMxFp4,
+    )
+    from vllm.model_executor.layers.fused_moe.fused_moe import (
+        fused_experts,
+        get_config_file_name,
+    )
+    from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
+        fused_topk,
+    )
+    from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
+        GroupedTopk,
+    )
+
+    __all__ += [
+        "AiterExperts",
+        "fused_topk",
+        "fused_experts",
+        "get_config_file_name",
+        "GroupedTopk",
+        "CutlassExpertsFp8",
+        "CutlassBatchedExpertsFp8",
+        "CutlassExpertsW4A8Fp8",
+        "TritonExperts",
+        "TritonWNA16Experts",
+        "BatchedTritonExperts",
+        "DeepGemmExperts",
+        "BatchedDeepGemmExperts",
+        "TritonOrDeepGemmExperts",
+        "XPUExperts",
+        "XPUExpertsFp8",
+        "XPUExpertsBlockFp8",
+        "XPUExpertsMxFp8",
+        "XPUExpertsMxFp4",
+    ]
+else:
+    # Some model classes directly use the custom ops. Add placeholders
+    # to avoid import errors.
+    def _raise_exception(method: str):
+        raise NotImplementedError(f"{method} is not implemented as lack of triton.")
+
+    fused_topk = lambda *args, **kwargs: _raise_exception("fused_topk")
+    fused_experts = lambda *args, **kwargs: _raise_exception("fused_experts")

ex_engine/moe/activation.py

@@ -122,32 +122,17 @@ def apply_moe_activation(
 
     # Activations with gated multiplication (gate × activation(up))
     if activation == MoEActivation.SILU:
-        # BI-V100: torch.ops._C.silu_and_mul not available
-        # Use corex_attn_head_rms_norm pattern: try C++ first, fallback to PyTorch
-        d = output.size(-1)
-        gate = input[..., :d]
-        up = input[..., d:]
-        output.copy_(F.silu(gate) * up)
+        torch.ops._C.silu_and_mul(output, input)
     elif activation == MoEActivation.GELU:
-        d = output.size(-1)
-        gate = input[..., :d]
-        up = input[..., d:]
-        output.copy_(F.gelu(gate) * up)
+        torch.ops._C.gelu_and_mul(output, input)
     elif activation == MoEActivation.GELU_TANH:
-        d = output.size(-1)
-        gate = input[..., :d]
-        up = input[..., d:]
-        output.copy_(F.gelu(gate, approximate="tanh") * up)
+        torch.ops._C.gelu_tanh_and_mul(output, input)
     elif activation == MoEActivation.SWIGLUOAI:
-        d = output.size(-1)
-        gate = input[..., :d]
-        up = input[..., d:]
-        output.copy_(F.silu(gate) * up)
+        torch.ops._C.swigluoai_and_mul(output, input)
     elif activation == MoEActivation.SWIGLUSTEP:
-        d = output.size(-1)
-        gate = input[..., :d]
-        up = input[..., d:]
-        output.copy_(F.silu(gate) * up)
+        from vllm.model_executor.layers.activation import swiglustep_and_mul_triton
+
+        swiglustep_and_mul_triton(output, input)
 
     # Activations without gated multiplication
     elif activation == MoEActivation.SILU_NO_MUL:

ex_engine/moe/experts/fallback.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
+
+
+class FallbackExperts(mk.FusedMoEExpertsModular, ABC):
+    """Base class for runtime dispatching of expert implementations."""
+
+    def __init__(
+        self,
+        experts: mk.FusedMoEExpertsModular,
+        fallback_experts: mk.FusedMoEExpertsModular,
+    ):
+        super().__init__(
+            moe_config=experts.moe_config, quant_config=experts.quant_config
+        )
+        self.fallback_experts = fallback_experts
+        self.experts = experts
+
+    @staticmethod
+    def get_clses() -> tuple[
+        type[mk.FusedMoEExpertsModular],
+        type[mk.FusedMoEExpertsModular],
+    ]:
+        """
+        Get the cls for the experts and fallback experts.
+
+        Subclasses should implement this method, so that
+        we have a consistent way to call the _supports_*
+        class methods below.
+        """
+        raise NotImplementedError(
+            "Subclasses must return the cls for the experts and fallback experts."
+        )
+
+    @classmethod
+    def activation_format(
+        cls: type["FallbackExperts"],
+    ) -> mk.FusedMoEActivationFormat:
+        experts_cls, fallback_cls = cls.get_clses()
+        assert experts_cls.activation_format() == fallback_cls.activation_format()
+        return experts_cls.activation_format()
+
+    @classmethod
+    def _supports_current_device(cls) -> bool:
+        experts_cls, fallback_cls = cls.get_clses()
+        return (
+            experts_cls._supports_current_device()
+            and fallback_cls._supports_current_device()
+        )
+
+    @classmethod
+    def _supports_no_act_and_mul(cls) -> bool:
+        experts_cls, fallback_cls = cls.get_clses()
+        return (
+            experts_cls._supports_no_act_and_mul()
+            and fallback_cls._supports_no_act_and_mul()
+        )
+
+    @classmethod
+    def _supports_quant_scheme(
+        cls,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        experts_cls, fallback_cls = cls.get_clses()
+        return experts_cls._supports_quant_scheme(
+            weight_key, activation_key
+        ) and fallback_cls._supports_quant_scheme(weight_key, activation_key)
+
+    @classmethod
+    def _supports_activation(cls, activation: MoEActivation) -> bool:
+        experts_cls, fallback_cls = cls.get_clses()
+        return experts_cls._supports_activation(
+            activation
+        ) and fallback_cls._supports_activation(activation)
+
+    @classmethod
+    def _supports_parallel_config(
+        cls, moe_parallel_config: FusedMoEParallelConfig
+    ) -> bool:
+        experts_cls, fallback_cls = cls.get_clses()
+        return experts_cls._supports_parallel_config(
+            moe_parallel_config
+        ) and fallback_cls._supports_parallel_config(moe_parallel_config)
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        e_war = self.experts.finalize_weight_and_reduce_impl()
+        fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl()
+        is_dge_war = e_war is not None
+        is_fbe_war = fbe_war is not None
+
+        if is_dge_war and is_fbe_war:
+            assert e_war == fbe_war, (
+                "Both implementations should agree on WeightAndReduce impls. "
+                f"Got e_war: {e_war}, and fbe_war: {fbe_war}"
+            )
+
+        if e_war is not None:
+            return e_war
+        assert fbe_war is not None
+        return fbe_war
+
+    @abstractmethod
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _select_experts_impl(
+        self,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+    ) -> mk.FusedMoEExpertsModular:
+        raise NotImplementedError
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        experts = self._select_experts_impl(hidden_states, w1, w2)
+        experts.apply(
+            output,
+            hidden_states,
+            w1,
+            w2,
+            topk_weights,
+            topk_ids,
+            activation,
+            global_num_experts,
+            expert_map,
+            a1q_scale,
+            a2_scale,
+            workspace13,
+            workspace2,
+            expert_tokens_meta,
+            apply_router_weight_on_input,
+        )

ex_engine/moe/experts/fused_batched_moe.py

@@ -0,0 +1,972 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Fused batched MoE kernel."""
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceDelegate,
+)
+from vllm.model_executor.layers.fused_moe.utils import (
+    _resize_cache,
+    moe_kernel_quantize_input,
+    normalize_batched_scales_shape,
+    swiglu_limit_func,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    group_broadcast,
+    kFp8Dynamic128Sym,
+    kFp8DynamicTensorSym,
+    kFp8DynamicTokenSym,
+    kFp8Static128BlockSym,
+    kFp8StaticChannelSym,
+    kFp8StaticTensorSym,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def moe_mmk(
+    a_ptrs,
+    b_ptrs,
+    K,
+    expert_id,
+    a_scale_ptr,
+    b_scale_ptr,
+    # The stride variables represent how much to increase the ptr by when
+    # moving by 1 element in a particular dimension. E.g. `stride_am` is
+    # how much to increase `a_ptr` by to get the element one row down
+    # (A has M rows).
+    stride_ak: tl.int64,
+    stride_bk: tl.int64,
+    stride_ase: tl.int64,
+    stride_asm: tl.int64,
+    stride_ask: tl.int64,
+    stride_bse: tl.int64,
+    stride_bsk: tl.int64,
+    stride_bsn: tl.int64,
+    # Offsets and masks
+    offs_m,
+    offs_n,
+    offs_bn,
+    mask_m,
+    # Block size for block-wise quantization
+    group_n: tl.constexpr,
+    group_k: tl.constexpr,
+    # Meta-parameters
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    compute_type: tl.constexpr,
+    use_w8a8: tl.constexpr,
+    use_w8a16: tl.constexpr,
+    per_act_token_quant: tl.constexpr,
+):
+    offs_k = tl.arange(0, BLOCK_K)
+
+    if use_w8a16:
+        b_scale_ptrs = (
+            b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn
+        )
+        b_scale = tl.load(b_scale_ptrs)
+
+    if use_w8a8:
+        # block-wise
+        if group_k > 0 and group_n > 0:
+            a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
+            offs_bsn = offs_bn // group_n
+            b_scale_ptrs = b_scale_ptr + offs_bsn * stride_bsn
+
+        # per act token
+        elif per_act_token_quant:
+            # Load per-token scale for activations
+            a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
+            a_scale = tl.load(a_scale_ptrs, mask=mask_m, other=0.0)[:, None]
+
+            b_scale_ptrs = b_scale_ptr + offs_bn[None, :] * stride_bsn
+            b_scale = tl.load(b_scale_ptrs)
+
+        # tensor-wise
+        else:
+            a_scale = tl.load(a_scale_ptr)
+            b_scale = tl.load(b_scale_ptr)
+
+    # -----------------------------------------------------------
+    # Iterate to compute a block of the C matrix.
+    # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
+    # of fp32 values for higher accuracy.
+    # `accumulator` will be converted back to fp16 after the loop.
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        # Load the next block of A and B, generate a mask by checking the
+        # K dimension.
+        a = tl.load(
+            a_ptrs,
+            mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
+        # We accumulate along the K dimension.
+        if use_w8a16:
+            accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
+        elif use_w8a8:
+            if group_k > 0 and group_n > 0:
+                k_start = k * BLOCK_K
+                offs_ks = k_start // group_k
+                a_scale = tl.load(
+                    a_scale_ptrs + offs_ks * stride_ask, mask=mask_m, other=0.0
+                )
+                b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
+
+                accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
+            else:
+                # acc used to enable fp8_fast_accum
+                accumulator = tl.dot(a, b, acc=accumulator)
+        else:
+            accumulator += tl.dot(a, b)
+
+        # Advance the ptrs to the next K block.
+        a_ptrs += BLOCK_K * stride_ak
+        b_ptrs += BLOCK_K * stride_bk
+
+    if use_w8a16:
+        accumulator = (accumulator * b_scale).to(compute_type)
+    elif use_w8a8:
+        if group_k > 0 and group_n > 0:
+            accumulator = accumulator.to(compute_type)
+        else:
+            accumulator = (accumulator * a_scale * b_scale).to(compute_type)
+    else:
+        accumulator = accumulator.to(compute_type)
+
+    return accumulator
+
+
+@triton.jit
+def expert_triton_kernel(
+    a_ptr,  # [max_tokens, K]
+    b_ptr,  # [K, N]
+    c_ptr,  # [max_tokens, N]
+    expert_id,
+    compute_type: tl.constexpr,
+    # Dimensions
+    M,
+    N,
+    K,
+    # Quantization data
+    a_scale_ptr,
+    b_scale_ptr,
+    b_zp_ptr,
+    # strides
+    stride_am: tl.int64,
+    stride_ak: tl.int64,
+    stride_bk: tl.int64,
+    stride_bn: tl.int64,
+    stride_cm: tl.int64,
+    stride_cn: tl.int64,
+    stride_ase: tl.int64,
+    stride_asm: tl.int64,
+    stride_ask: tl.int64,
+    stride_bse: tl.int64,
+    stride_bsk: tl.int64,
+    stride_bsn: tl.int64,
+    # offsets
+    offs_bn,
+    # Blockwise quantization data
+    group_n,
+    group_k,
+    # Quantization schemes
+    use_fp8_w8a8: tl.constexpr,
+    use_int8_w8a16: tl.constexpr,
+    per_act_token_quant: tl.constexpr,
+    # Kernel config
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    offs_m = tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N) % N
+    offs_k = tl.arange(0, BLOCK_K)
+    mask_m = offs_m < M
+
+    # Make grids of a + b pointers
+    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
+
+    accumulator = moe_mmk(
+        a_ptrs,
+        b_ptrs,
+        K,
+        expert_id,
+        a_scale_ptr,
+        b_scale_ptr,
+        # The stride variables represent how much to increase the ptr by when
+        # moving by 1 element in a particular dimension. E.g. `stride_am` is
+        # how much to increase `a_ptr` by to get the element one row down
+        # (A has M rows).
+        stride_ak,
+        stride_bk,
+        stride_ase,
+        stride_asm,
+        stride_ask,
+        stride_bse,
+        stride_bsk,
+        stride_bsn,
+        # Offsets and masks
+        offs_m,
+        offs_n,
+        offs_bn,
+        mask_m,
+        # Block size for block-wise quantization
+        group_n,
+        group_k,
+        # Meta-parameters
+        BLOCK_M,
+        BLOCK_N,
+        BLOCK_K,
+        compute_type,
+        use_fp8_w8a8,
+        use_int8_w8a16,
+        per_act_token_quant,
+    )
+
+    # store in C
+    offs_cn = tl.arange(0, BLOCK_N)
+    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_cn[None, :] * stride_cn
+    c_mask = mask_m[:, None] & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, accumulator, mask=c_mask)
+
+
+@triton.jit
+def batched_triton_kernel(
+    a_ptr,  # [E, max_num_tokens, K]
+    b_ptr,  # [E, K, N]
+    c_ptr,  # [E, max_num_tokens, N]
+    expert_num_tokens,  # [E]
+    compute_type: tl.constexpr,
+    # Dimensions
+    max_num_tokens,
+    K,
+    N,
+    # Quantization data
+    a_scale_ptr,
+    b_scale_ptr,
+    b_zp_ptr,
+    # The stride variables represent how much to increase the ptr by when
+    # moving by 1 element in a particular dimension. E.g. `stride_am` is
+    # how much to increase `a_ptr` by to get the element one row down
+    # (A has M rows).
+    stride_ae: tl.int64,
+    stride_am: tl.int64,
+    stride_ak: tl.int64,
+    stride_be: tl.int64,
+    stride_bk: tl.int64,
+    stride_bn: tl.int64,
+    stride_ce: tl.int64,
+    stride_cm: tl.int64,
+    stride_cn: tl.int64,
+    stride_ase: tl.int64,
+    stride_asm: tl.int64,
+    stride_ask: tl.int64,
+    stride_bse: tl.int64,
+    stride_bsk: tl.int64,
+    stride_bsn: tl.int64,
+    # Blockwise quantization data
+    group_n: tl.constexpr,
+    group_k: tl.constexpr,
+    # Quantization schemes
+    use_fp8_w8a8: tl.constexpr,
+    use_int8_w8a16: tl.constexpr,
+    per_act_token_quant: tl.constexpr,
+    # Kernel config
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    expert_id = tl.program_id(axis=0)
+    e_num_tokens = tl.load(expert_num_tokens + expert_id)
+    if e_num_tokens == 0:
+        # Early exit
+        return
+
+    # axis 1 is M_blocks * N_blocks
+    pid_mn = tl.program_id(axis=1)
+    # num_pid_m = tl.cdiv(max_num_tokens, BLOCK_M)
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    pid_m = pid_mn // num_pid_n
+    pid_n = pid_mn % num_pid_n
+
+    cta_m_start = pid_m * BLOCK_M
+    cta_n_start = pid_n * BLOCK_N
+    if cta_m_start >= e_num_tokens:
+        # Early exit
+        return
+
+    cta_m_size = min(BLOCK_M, e_num_tokens - cta_m_start)
+    cta_n_size = min(BLOCK_N, N - cta_n_start)
+
+    a_ptr = a_ptr + expert_id * stride_ae + cta_m_start * stride_am
+    b_ptr = b_ptr + expert_id * stride_be + cta_n_start * stride_bn
+    c_ptr = (
+        c_ptr
+        + expert_id * stride_ce
+        + cta_m_start * stride_cm
+        + cta_n_start * stride_cn
+    )
+
+    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)) % N
+
+    if use_fp8_w8a8:
+        a_scale_ptr = a_scale_ptr + expert_id * stride_ase
+        b_scale_ptr = b_scale_ptr + expert_id * stride_bse
+
+        # block-wise
+        if group_k > 0 and group_n > 0 or per_act_token_quant:
+            a_scale_ptr = a_scale_ptr + cta_m_start * stride_asm
+
+    expert_triton_kernel(
+        a_ptr,
+        b_ptr,
+        c_ptr,
+        expert_id,
+        compute_type,
+        cta_m_size,  # M
+        cta_n_size,  # N
+        K,  # K
+        a_scale_ptr,
+        b_scale_ptr,
+        b_zp_ptr,
+        # Strides
+        stride_am,
+        stride_ak,
+        stride_bk,
+        stride_bn,
+        stride_cm,
+        stride_cn,
+        stride_ase,
+        stride_asm,
+        stride_ask,
+        stride_bse,
+        stride_bsk,
+        stride_bsn,
+        # offsets
+        offs_bn,
+        # Blockwise quantization data
+        group_n,
+        group_k,
+        # Quantization schemes
+        use_fp8_w8a8,
+        use_int8_w8a16,
+        per_act_token_quant,
+        # Kernel config
+        BLOCK_M,
+        BLOCK_N,
+        BLOCK_K,
+    )
+
+
+def invoke_moe_batched_triton_kernel(
+    A: torch.Tensor,  # [E, max_tokens, K]
+    B: torch.Tensor,  # [E, N, K]
+    C: torch.Tensor,  # [E, max_tokens, N]
+    expert_num_tokens: torch.Tensor,  # [E]
+    compute_type: tl.dtype,
+    # Quantization data
+    A_scale: torch.Tensor | None,
+    B_scale: torch.Tensor | None,
+    B_zp: torch.Tensor,
+    # Quantization schemes
+    use_fp8_w8a8: bool,
+    use_int8_w8a16: bool,
+    use_int4_w4a16: bool,
+    config: dict[str, int],
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+):
+    assert not use_int4_w4a16
+    max_num_tokens = A.size(1)
+    K = A.size(2)
+    N = C.size(2)
+
+    BLOCK_M = config["BLOCK_SIZE_M"]
+    BLOCK_N = config["BLOCK_SIZE_N"]
+    BLOCK_K = config["BLOCK_SIZE_K"]
+
+    grid = (
+        expert_num_tokens.size(0),
+        triton.cdiv(max_num_tokens, BLOCK_M) * triton.cdiv(B.size(1), BLOCK_N),
+    )
+
+    A_scale = normalize_batched_scales_shape(A_scale, expert_num_tokens.shape[0])
+
+    if B_scale is not None and B_scale.ndim == 1:
+        assert B_scale.numel() == expert_num_tokens.shape[0]
+        B_scale = B_scale.view(-1, 1, 1)
+
+    assert A_scale is None or A_scale.ndim == 3, (
+        f"{0 if A_scale is None else A_scale.shape}"
+    )
+    assert B_scale is None or B_scale.ndim == 1 or B_scale.ndim == 3, (
+        f"{0 if B_scale is None else B_scale.shape}"
+    )
+
+    if B_scale is not None:
+        if B_scale.ndim == 1:
+            stride_bse = 1
+            stride_bsk = 0
+            stride_bsn = 0
+        else:
+            stride_bse = B_scale.stride(0)
+            stride_bsk = B_scale.stride(2)
+            stride_bsn = B_scale.stride(1)
+
+    else:
+        stride_bse = 0
+        stride_bsk = 0
+        stride_bsn = 0
+
+    if A_scale is not None:
+        stride_ase = A_scale.stride(0)
+        stride_asm = A_scale.stride(1)
+        stride_ask = A_scale.stride(2)
+    else:
+        stride_ase = 0
+        stride_asm = 0
+        stride_ask = 0
+
+    batched_triton_kernel[grid](
+        A,
+        B,
+        C,
+        expert_num_tokens,
+        compute_type,
+        # Dimensions
+        max_num_tokens,
+        K,
+        N,
+        # Quantization data
+        A_scale,
+        B_scale,
+        B_zp,
+        # Strides
+        A.stride(0),
+        A.stride(1),
+        A.stride(2),
+        B.stride(0),
+        B.stride(2),
+        B.stride(1),
+        C.stride(0),
+        C.stride(1),
+        C.stride(2),
+        stride_ase,
+        stride_asm,
+        stride_ask,
+        stride_bse,
+        stride_bsk,
+        stride_bsn,
+        # Blockwise quantization data
+        0 if block_shape is None else block_shape[0],
+        0 if block_shape is None else block_shape[1],
+        # Quantization schemes
+        use_fp8_w8a8,
+        use_int8_w8a16,
+        per_act_token_quant,
+        # Kernel config
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+    )
+
+
+class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
+    """
+    A reference MoE expert class that operates on expert batched format,
+    i.e. E x max_num_tokens x K.  This is the format that the batched
+    dispatch/combine kernels use.
+    """
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int,
+        num_dispatchers: int,
+    ):
+        super().__init__(
+            moe_config=moe_config,
+            quant_config=quant_config,
+            max_num_tokens=max_num_tokens,
+            num_dispatchers=num_dispatchers,
+        )
+        assert not self.quant_config.use_int8_w8a8, "NYI"
+        assert not self.quant_config.use_int8_w8a16, "NYI"
+        assert not self.quant_config.use_int4_w4a16, "NYI"
+        assert self.quant_config.ocp_mx_scheme is None, "NYI"
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.BatchedExperts
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        raise NotImplementedError(
+            "NaiveBatchedExperts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        raise NotImplementedError(
+            "NaiveBatchedExperts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        raise NotImplementedError(
+            "NaiveBatchedExperts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        raise NotImplementedError(
+            "NaiveBatchedExperts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        raise NotImplementedError(
+            "NaiveBatchedExperts is not yet used by an Oracle. "
+            "This method should not be called."
+        )
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        # Let PrepareAndFinalize::finalize() decide the impl.
+        return TopKWeightAndReduceDelegate()
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        assert self.num_dispatchers is not None
+        assert self.max_num_tokens is not None
+        num_dp = self.num_dispatchers
+        num_experts = local_num_experts
+        workspace13 = (num_experts, self.max_num_tokens * num_dp, K)
+        workspace2 = (self.max_num_tokens * num_dp, N)
+        output = workspace13
+        return (workspace13, workspace2, output)
+
+    def dequant(self, t: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+        assert self.quant_config.is_quantized
+        f32 = torch.float32
+        if self.quant_config.is_per_act_token or self.quant_config.is_per_tensor:
+            return t.to(f32) * scale
+        else:
+            return t.to(f32) * group_broadcast(scale, t.shape)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        assert hidden_states.dim() == 3
+        assert expert_tokens_meta is not None
+        expert_num_tokens = expert_tokens_meta.expert_num_tokens
+
+        num_local_experts = w1.size(0)
+        assert num_local_experts == w1.size(0), f"{num_local_experts} == {w1.size(0)}"
+
+        N = w1.size(1) // 2
+
+        for expert in range(num_local_experts):
+            # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor
+            if (
+                torch.compiler.is_compiling()
+                or torch.cuda.is_current_stream_capturing()
+            ):
+                num = hidden_states.shape[1]
+            else:
+                num = int(expert_num_tokens[expert].item())
+
+            if num == 0:
+                continue
+
+            tmp = _resize_cache(workspace2, (num, N))
+
+            if self.quant_config.is_quantized:
+                assert a1q_scale is not None and self.w1_scale is not None
+                input = self.dequant(hidden_states[expert, :, :], a1q_scale[expert])
+                w1_dq = self.dequant(w1[expert], self.w1_scale[expert])
+                input = input[:num] @ w1_dq.transpose(0, 1)
+            else:
+                input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
+
+            self.activation(activation, tmp, input.to(tmp.dtype))
+
+            if self.quant_config.is_quantized:
+                assert self.w2_scale is not None
+                w2_dq = self.dequant(w2[expert], self.w2_scale[expert])
+            else:
+                w2_dq = w2[expert]
+
+            output[expert, :num, :] = tmp @ w2_dq.transpose(0, 1).to(tmp.dtype)
+
+
+def batched_moe_kernel_quantize_input(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    num_tokens: int,
+    E: int,
+    N: int,
+    expert_num_tokens: torch.Tensor,
+    qtype: torch.dtype | None,
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing():
+        # Note: this does a bunch of extra work because expert_num_tokens is
+        # ignored but it does support torch.compile + cudagraphs.
+        hidden_dim = A.size(-1)
+        assert A_scale is None or A_scale.ndim <= 2, (
+            f"{A_scale.shape if A_scale is not None else None}"
+        )
+        A_q, A_q_scale = moe_kernel_quantize_input(
+            A.view(-1, hidden_dim), A_scale, qtype, per_act_token_quant, block_shape
+        )
+        A_q = A_q.view(E, -1, hidden_dim)
+        A_q_scale = normalize_batched_scales_shape(A_q_scale, E)
+
+        return A_q, A_q_scale
+    elif qtype is None:
+        return A, normalize_batched_scales_shape(A_scale, E)
+    else:
+        A_q = torch.empty_like(A, dtype=qtype)
+
+        if per_act_token_quant:
+            assert block_shape is None
+            scale_shape = (E, num_tokens, 1)
+        elif block_shape is not None:
+            _, block_k = block_shape
+            k_tiles = (A.shape[-1] + block_k - 1) // block_k
+            scale_shape = (E, num_tokens, k_tiles)
+        else:
+            scale_shape = (E, 1, 1)
+
+        A_q_scale = torch.zeros(scale_shape, dtype=torch.float32, device=A.device)
+
+        num_experts = expert_num_tokens.numel()
+
+        A_scale = normalize_batched_scales_shape(A_scale, num_experts)
+
+        for e in range(E):
+            num_tokens = int(expert_num_tokens[e].item())
+            if num_tokens > 0:
+                if A_scale is not None:
+                    scales = A_scale[e, : min(num_tokens, A_scale.shape[1])]
+                else:
+                    scales = None
+                A_q[e, :num_tokens], tmp_scale = moe_kernel_quantize_input(
+                    A[e, :num_tokens],
+                    scales,
+                    qtype,
+                    per_act_token_quant,
+                    block_shape,
+                )
+                assert tmp_scale is not None
+                A_q_scale[e, : tmp_scale.shape[0]] = tmp_scale
+
+        return A_q, A_q_scale
+
+
+class BatchedTritonExperts(mk.FusedMoEExpertsModular):
+    """
+    A Triton based MoE expert class that operates on expert batched format,
+    i.e. E x max_num_tokens x K.  This is the format that the batched
+    dispatch/combine kernels use.
+    """
+
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int,
+        num_dispatchers: int,
+    ):
+        super().__init__(
+            moe_config=moe_config,
+            quant_config=quant_config,
+            max_num_tokens=max_num_tokens,
+            num_dispatchers=num_dispatchers,
+        )
+        assert not self.quant_config.use_int8_w8a8, "NYI"
+        assert not self.quant_config.use_int8_w8a16, "NYI"
+        assert not self.quant_config.use_int4_w4a16, "NYI"
+        assert self.quant_config.ocp_mx_scheme is None, "NYI"
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.BatchedExperts
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        return current_platform.is_cuda_alike()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return True
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        p = current_platform
+        if p.is_rocm():
+            from vllm.platforms.rocm import on_gfx9
+
+            is_rocm_on_gfx9 = on_gfx9()
+        else:
+            is_rocm_on_gfx9 = False
+
+        device_supports_fp8 = is_rocm_on_gfx9 or (
+            p.is_cuda() and p.has_device_capability((8, 9))
+        )
+
+        supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
+        if device_supports_fp8:
+            supported += [
+                (kFp8Static128BlockSym, kFp8Dynamic128Sym),
+                (kFp8StaticChannelSym, kFp8DynamicTokenSym),
+                (kFp8StaticTensorSym, kFp8DynamicTokenSym),
+                (kFp8StaticTensorSym, kFp8StaticTensorSym),
+                (kFp8StaticTensorSym, kFp8DynamicTensorSym),
+            ]
+        return (weight_key, activation_key) in supported
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in [
+            MoEActivation.SILU,
+            MoEActivation.GELU,
+            MoEActivation.GELU_TANH,
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SILU_NO_MUL,
+            MoEActivation.GELU_NO_MUL,
+            MoEActivation.GELU_TANH_NO_MUL,
+            MoEActivation.RELU2_NO_MUL,
+        ]
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return True
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        # Let PrepareAndFinalize::finalize() decide the impl.
+        return TopKWeightAndReduceDelegate()
+
+    def activation(
+        self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
+    ) -> None:
+        gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
+        if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
+            swiglu_limit_func(output, input, float(gemm1_clamp_limit))
+            return
+
+        super().activation(activation, output, input)
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        assert self.num_dispatchers is not None
+        assert self.max_num_tokens is not None
+        num_dp = self.num_dispatchers
+        num_experts = local_num_experts
+        max_num_tokens = self.max_num_tokens
+        activation_out_dim = self.adjust_N_for_activation(N, activation)
+        workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
+        workspace2 = (num_experts, max_num_tokens * num_dp, activation_out_dim)
+        output = (num_experts, max_num_tokens * num_dp, K)
+        return (workspace13, workspace2, output)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor,
+        workspace2: torch.Tensor,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        # Check constraints.
+        if self.quant_config.use_int4_w4a16:
+            assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
+        else:
+            assert hidden_states.size(-1) == w1.size(2), (
+                f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
+            )
+
+        assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+        assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
+        assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
+        assert hidden_states.dtype in [
+            torch.float32,
+            torch.float16,
+            torch.bfloat16,
+            torch.float8_e4m3fn,
+            torch.float8_e4m3fnuz,
+        ]
+        assert expert_tokens_meta is not None
+
+        expert_num_tokens = expert_tokens_meta.expert_num_tokens
+
+        E, max_num_tokens, N, K, top_k_num = self.moe_problem_size(
+            hidden_states, w1, w2, topk_ids
+        )
+
+        assert w1.size(0) == E
+        assert w2.size(0) == E
+
+        config_dtype = self.quant_config.config_name(hidden_states.dtype)
+
+        config = try_get_optimal_moe_config(
+            w1.size(),
+            w2.size(),
+            top_k_num,
+            config_dtype,
+            max_num_tokens,
+            block_shape=self.block_shape,
+        )
+
+        if hidden_states.dtype == torch.bfloat16:
+            compute_type = tl.bfloat16
+        elif hidden_states.dtype == torch.float16:
+            compute_type = tl.float16
+        elif hidden_states.dtype == torch.float32:
+            compute_type = tl.float32
+        elif hidden_states.dtype == current_platform.fp8_dtype():
+            compute_type = tl.bfloat16
+        else:
+            raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
+
+        # We can reuse the memory between these because by the time we need
+        # cache3, we're done with cache1
+        intermediate_cache1 = _resize_cache(workspace13, (E, max_num_tokens, N))
+        activation_out_dim = self.adjust_N_for_activation(N, activation)
+        intermediate_cache2 = _resize_cache(
+            workspace2, (E, max_num_tokens, activation_out_dim)
+        )
+
+        # TODO(bnell): should this be done for any quantized type?
+        if self.quant_config.use_fp8_w8a8:
+            intermediate_cache1.fill_(0)
+
+        a1q_scale = normalize_batched_scales_shape(a1q_scale, E)
+
+        # MM1
+        invoke_moe_batched_triton_kernel(
+            A=hidden_states,
+            B=w1,
+            C=intermediate_cache1,
+            expert_num_tokens=expert_num_tokens,
+            compute_type=compute_type,
+            A_scale=a1q_scale,
+            B_scale=self.w1_scale,
+            B_zp=self.w1_zp,
+            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
+            use_int8_w8a16=self.quant_config.use_int8_w8a16,
+            use_int4_w4a16=self.quant_config.use_int4_w4a16,
+            config=config,
+            per_act_token_quant=self.per_act_token_quant,
+            block_shape=self.block_shape,
+        )
+
+        intermediate_cache2.fill_(0)
+
+        # TODO (bnell): use triton utility from batched deep gemm.
+        self.activation(
+            activation,
+            intermediate_cache2.view(-1, activation_out_dim),
+            intermediate_cache1.view(-1, N),
+        )
+
+        qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
+            intermediate_cache2,
+            a2_scale,
+            max_num_tokens,
+            E,
+            N,
+            expert_num_tokens,
+            self.quant_dtype,
+            self.per_act_token_quant,
+            self.block_shape,
+        )
+
+        invoke_moe_batched_triton_kernel(
+            A=qintermediate_cache2,
+            B=w2,
+            C=output,
+            expert_num_tokens=expert_num_tokens,
+            compute_type=compute_type,
+            A_scale=a2q_scale,
+            B_scale=self.w2_scale,
+            B_zp=self.w2_zp,
+            use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
+            use_int8_w8a16=self.quant_config.use_int8_w8a16,
+            use_int4_w4a16=self.quant_config.use_int4_w4a16,
+            config=config,
+            per_act_token_quant=self.per_act_token_quant,
+            block_shape=self.block_shape,
+        )

ex_engine/moe/fused_moe_method_base.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEExpertsModular,
+    FusedMoEPrepareAndFinalizeModular,
+)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizeMethodBase,
+)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+    from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts
+
+logger = init_logger(__name__)
+
+
+class FusedMoEMethodBase(QuantizeMethodBase):
+    def __init__(self, moe: FusedMoEConfig):
+        super().__init__()
+        self.moe: FusedMoEConfig = moe
+        self.moe_quant_config: FusedMoEQuantConfig | None = None
+        self.moe_kernel: mk.FusedMoEKernel | None = None
+
+    @property
+    def supports_internal_mk(self) -> bool:
+        # NOTE(rob): temporary attribute to indicate support for
+        # completed migration to the new internal MK interface.
+        return self.moe_kernel is not None
+
+    @property
+    def mk_can_overlap_shared_experts(self) -> bool:
+        # NOTE(rob): temporary attribute to indicate support for
+        # completed migration to the new internal MK interface.
+        return (
+            self.moe_kernel is not None and self.moe_kernel.can_overlap_shared_experts
+        )
+
+    @abstractmethod
+    def create_weights(
+        self,
+        layer: "RoutedExperts",
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        raise NotImplementedError
+
+    def uses_weight_scale_2_pattern(self) -> bool:
+        """
+        Returns True if this quantization method uses 'weight_scale_2' pattern
+        for per-tensor weight scales (e.g., FP4 variants), False otherwise.
+
+        This method should be overridden by subclasses that use the
+        'weight_scale_2' pattern instead of the standard 'weight_scale' pattern.
+        """
+        return False
+
+    def maybe_roundup_sizes(
+        self,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        act_dtype: torch.dtype,
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> tuple[int, int]:
+        """
+        Given layer hidden size and intermediate size per partition and MoE
+        configurations, round up hidden_size and intermediate_size_per_partition
+        if necessary.
+
+        Args:
+            hidden_size: Layer hidden-size
+            intermediate_size_per_partition: Intermediate size per partition for
+                the layer.
+            act_dtype: Data type of the layer activations.
+            moe_parallel_config: Fused MoE parallelization strategy configuration.
+
+        Return:
+            A tuple of (rounded_hidden_size, rounded_intermediate_size_per_partition),
+            where:
+                - rounded_hidden_size is the possibly rounded up hidden size.
+                - rounded_intermediate_size_per_partition is the possibly rounded
+                  up intermediate size per partition.
+        """
+        from .all2all_utils import maybe_roundup_layer_hidden_size
+
+        return maybe_roundup_layer_hidden_size(
+            hidden_size, act_dtype, moe_parallel_config
+        ), intermediate_size_per_partition
+
+    def maybe_make_prepare_finalize(
+        self,
+        routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+    ) -> FusedMoEPrepareAndFinalizeModular | None:
+        from .all2all_utils import maybe_make_prepare_finalize
+
+        pf = maybe_make_prepare_finalize(
+            self.moe, self.moe_quant_config, routing_tables
+        )
+        assert pf is None or isinstance(pf, FusedMoEPrepareAndFinalizeModular)
+        return pf
+
+    def select_gemm_impl(
+        self,
+        prepare_finalize: FusedMoEPrepareAndFinalizeModular,
+        layer: "RoutedExperts",
+    ) -> FusedMoEExpertsModular:
+        # based on the all2all implementation, select the appropriate
+        # gemm implementation
+        raise ValueError(
+            f"{self.__class__.__name__} uses the new modular kernel initialization "
+            "logic. This function should not be called."
+        )
+
+    @abstractmethod
+    def get_fused_moe_quant_config(
+        self, layer: "RoutedExperts"
+    ) -> FusedMoEQuantConfig | None:
+        raise NotImplementedError
+
+    @property
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        if self.moe_kernel is not None:
+            return self.moe_kernel.prepare_finalize.topk_indices_dtype()
+        return None
+
+    @property
+    def skip_forward_padding(self) -> bool:
+        """Whether to skip the padding in the forward before applying the moe method."""
+        return False
+
+    @property
+    def has_unpadded_output(self) -> bool:
+        """
+        Indicates that the hidden_states output might be the unpadded
+        hidden_states shape rather than the full padded shape.
+        """
+        return False
+
+    @property
+    def supports_eplb(self) -> bool:
+        return False
+
+    @property
+    def method_name(self) -> str:
+        return self.__class__.__name__
+
+    @property
+    def is_monolithic(self) -> bool:
+        if self.moe_kernel is None:
+            if hasattr(self, "experts_cls"):
+                return self.experts_cls.is_monolithic()
+            else:
+                return False
+        return self.moe_kernel.is_monolithic
+
+    def apply(
+        self,
+        layer: "RoutedExperts",
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts: "SharedExperts | None",
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """
+        Apply the MoE operation using modular kernels.
+
+        Args:
+            layer: RoutedExperts instance containing weight parameters
+            x: Input tensor
+            topk_weights: Expert weights from router
+            topk_ids: Selected expert IDs from router
+            shared_experts_input: Input for shared experts (if any)
+
+        Returns:
+            Output tensor from routed experts
+        """
+        raise NotImplementedError
+
+    def apply_monolithic(
+        self,
+        layer: "RoutedExperts",
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Apply the MoE operation using monolithic kernels.
+
+        Args:
+            layer: RoutedExperts instance containing weight parameters
+            x: Input tensor
+            router_logits: Router logits (routing done internally)
+
+        Returns:
+            Output tensor from routed experts
+        """
+        raise NotImplementedError

ex_engine/moe/fused_moe_modular_method.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEKernel,
+    FusedMoEPrepareAndFinalizeModular,
+)
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.fused_moe.routed_experts import (
+        RoutedExperts,
+    )
+
+logger = init_logger(__name__)
+
+
+# --8<-- [start:modular_fused_moe]
+@CustomOp.register("modular_fused_moe")
+class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
+    # --8<-- [end:modular_fused_moe]
+
+    def __init__(
+        self, old_quant_method: FusedMoEMethodBase, moe_kernel: FusedMoEKernel
+    ):
+        super().__init__(moe_kernel.moe_config)
+        self.moe_quant_config = old_quant_method.moe_quant_config
+        self.moe_kernel = moe_kernel
+        self.old_quant_method = old_quant_method
+        logger.debug("Swapping out %s", self.old_quant_method.__class__.__name__)
+
+    @property
+    def wraps_legacy_quant_method(self) -> bool:
+        return not self.old_quant_method.supports_internal_mk
+
+    @staticmethod
+    def make(
+        routed_experts: "RoutedExperts",
+        old_quant_method: FusedMoEMethodBase,
+        prepare_finalize: FusedMoEPrepareAndFinalizeModular,
+    ) -> "FusedMoEModularMethod":
+        return FusedMoEModularMethod(
+            old_quant_method,
+            FusedMoEKernel(
+                prepare_finalize,
+                old_quant_method.select_gemm_impl(prepare_finalize, routed_experts),
+            ),
+        )
+
+    @property
+    def skip_forward_padding(self) -> bool:
+        return self.old_quant_method.skip_forward_padding
+
+    @property
+    def has_unpadded_output(self) -> bool:
+        return self.old_quant_method.has_unpadded_output
+
+    @property
+    def supports_eplb(self) -> bool:
+        return self.old_quant_method.supports_eplb
+
+    @property
+    def method_name(self) -> str:
+        return self.old_quant_method.method_name
+
+    def create_weights(
+        self,
+        layer: "RoutedExperts",
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        raise NotImplementedError
+
+    def get_fused_moe_quant_config(
+        self, layer: "RoutedExperts"
+    ) -> FusedMoEQuantConfig | None:
+        return self.moe_quant_config
+
+    def apply(
+        self,
+        layer: "RoutedExperts",
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts: SharedExperts | None,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            hidden_states=x,
+            w1=layer.w13_weight,
+            w2=layer.w2_weight,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            expert_map=layer.expert_map,
+            shared_experts=shared_experts,
+            shared_experts_input=shared_experts_input,
+        )

ex_engine/moe/layer.py

@@ -0,0 +1,406 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config import ParallelConfig, get_current_vllm_config
+from vllm.distributed import (
+    get_dp_group,
+    get_pcp_group,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+)
+from vllm.model_executor.layers.fused_moe.expert_map_manager import (
+    ExpertMapManager,
+)
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
+    FusedMoERouter,
+)
+from vllm.model_executor.layers.fused_moe.router.router_factory import (
+    create_fused_moe_router,
+)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import (
+    MoERunner,
+)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+)
+
+logger = init_logger(__name__)
+
+
+def make_parallel_config(
+    tp_size: int | None,
+    dp_size: int | None,
+    pcp_size: int | None,
+    is_sequence_parallel: bool,
+    parallel_config: ParallelConfig,
+) -> FusedMoEParallelConfig:
+    tp_size_ = (
+        tp_size if tp_size is not None else get_tensor_model_parallel_world_size()
+    )
+    dp_size_ = dp_size if dp_size is not None else get_dp_group().world_size
+    pcp_size_ = pcp_size if pcp_size is not None else get_pcp_group().world_size
+    sp_size = tp_size_ if is_sequence_parallel else 1
+
+    moe_parallel_config = FusedMoEParallelConfig.make(
+        tp_size_=tp_size_,
+        pcp_size_=pcp_size_,
+        dp_size_=dp_size_,
+        sp_size_=sp_size,
+        vllm_parallel_config=parallel_config,
+    )
+
+    assert moe_parallel_config.is_sequence_parallel == is_sequence_parallel
+
+    logger.debug("FusedMoEParallelConfig = %s", str(moe_parallel_config))
+
+    return moe_parallel_config
+
+
+def determine_expert_counts(
+    num_experts: int,
+    num_redundant_experts: int,
+    n_shared_experts: int | None,
+    is_act_and_mul: bool,
+) -> tuple[int, int, int]:
+    global_num_experts = num_experts + num_redundant_experts
+    logical_num_experts = num_experts
+    # ROCm aiter shared experts fusion
+    # AITER only supports gated activations (silu/gelu), so disable it
+    # for non-gated MoE (is_act_and_mul=False)
+    # rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul
+    aiter_fmoe_shared_expert_enabled = (
+        rocm_aiter_ops.is_fusion_moe_shared_experts_enabled() and is_act_and_mul
+    )
+
+    num_fused_shared_experts = (
+        n_shared_experts
+        if n_shared_experts is not None and aiter_fmoe_shared_expert_enabled
+        else 0
+    )
+    if not aiter_fmoe_shared_expert_enabled and num_fused_shared_experts != 0:
+        raise ValueError(
+            "n_shared_experts is only supported on ROCm aiter when "
+            "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is enabled"
+        )
+
+    return global_num_experts, logical_num_experts, num_fused_shared_experts
+
+
+# TODO: rename this
+def FusedMoE(
+    num_experts: int,  # Global number of experts
+    top_k: int,
+    hidden_size: int,
+    intermediate_size: int,
+    params_dtype: torch.dtype | None = None,
+    renormalize: bool = True,
+    use_grouped_topk: bool = False,
+    num_expert_group: int | None = None,
+    topk_group: int | None = None,
+    quant_config: QuantizationConfig | None = None,
+    tp_size: int | None = None,
+    dp_size: int | None = None,
+    pcp_size: int | None = None,
+    prefix: str = "",
+    custom_routing_function: Callable | None = None,
+    router: FusedMoERouter | None = None,
+    scoring_func: str = "softmax",
+    routed_scaling_factor: float = 1.0,
+    swiglu_limit: float | None = None,
+    e_score_correction_bias: torch.Tensor | None = None,
+    apply_router_weight_on_input: bool = False,
+    activation: str = "silu",
+    enable_eplb: bool = False,
+    num_redundant_experts: int = 0,
+    has_bias: bool = False,
+    is_sequence_parallel: bool = False,
+    expert_mapping: list[tuple[str, str, int, str]] | None = None,
+    n_shared_experts: int | None = None,
+    router_logits_dtype: torch.dtype | None = None,
+    gate: torch.nn.Module | None = None,
+    shared_experts: torch.nn.Module | None = None,
+    shared_expert_gate: torch.nn.Module | None = None,
+    routed_input_transform: torch.nn.Module | None = None,
+    routed_output_transform: torch.nn.Module | None = None,
+    apply_routed_scale_to_output: bool = False,
+    zero_expert_type: str | None = None,
+    hash_indices_table: torch.Tensor | None = None,
+    runner_cls: type[MoERunner] | None = None,
+    runner_args: dict[str, Any] | None = None,
+    routed_experts_cls: type[RoutedExperts] | None = None,
+    routed_experts_args: dict[str, Any] | None = None,
+) -> MoERunner:
+    """Factory function for creating MoE execution pipeline.
+
+    Creates and configures a complete MoE execution pipeline including:
+    - Router (for token-to-expert assignment)
+    - RoutedExperts (containing expert weight parameters)
+    - MoERunner (orchestrates the complete forward pass)
+
+    The experts contain both MergedColumnParallel weights (gate_up_proj/w13)
+    and RowParallelLinear weights (down_proj/w2).
+
+    Note: Mixtral uses w1, w2, and w3 for gate, up, and down_proj. We
+    copy that naming convention here and handle any remapping in the
+    load_weights function in each model implementation.
+
+    Args:
+        num_experts: Number of experts in the model (global count)
+        top_k: Number of experts selected for each token
+        hidden_size: Input hidden state size of the transformer
+        intermediate_size: Intermediate size of the experts
+        params_dtype: Data type for the parameters
+        renormalize: Whether to renormalize the logits in the router
+        use_grouped_topk: Whether to use grouped top-k routing
+        num_expert_group: Number of expert groups for grouped top-k
+        topk_group: Top-k value per group for grouped top-k
+        quant_config: Quantization configuration
+        tp_size: Tensor parallelism size (None = use global default)
+        dp_size: Data parallelism size (None = use global default)
+        pcp_size: Pipeline context parallelism size (None = use global default)
+        prefix: Layer name prefix for weight loading
+        custom_routing_function: Custom routing function override
+        router: Pre-configured router instance (None = create default)
+        scoring_func: Scoring function for routing ("softmax" or others)
+        routed_scaling_factor: Scaling factor applied to topk_weights or output
+        swiglu_limit: SwiGLU activation limit
+        e_score_correction_bias: Expert score correction bias tensor
+        apply_router_weight_on_input: Whether to apply router weights on input
+        activation: Activation function name ("silu", "gelu", etc.)
+        enable_eplb: Whether to enable expert parallelism load balancer
+        num_redundant_experts: Number of redundant experts for EPLB
+        has_bias: Whether expert layers have bias terms
+        is_sequence_parallel: Whether sequence parallelism is enabled
+        expert_mapping: Expert parameter mapping for weight loading
+        n_shared_experts: Number of shared experts (ROCm aiter only)
+        router_logits_dtype: Data type for router logits buffers
+        gate: Pre-configured gate module
+        shared_experts: Pre-configured shared experts module
+        shared_expert_gate: Pre-configured shared expert gate module
+        routed_input_transform: Input transformation module
+        routed_output_transform: Output transformation module
+        apply_routed_scale_to_output: Whether to apply routed_scaling_factor to
+                                      output instead of topk_weights
+        zero_expert_type: Type of zero expert handling
+        hash_indices_table: Hash table for expert indices
+        runner_cls: Custom MoERunner class (None = use default MoERunner)
+        runner_args: Additional arguments for runner constructor
+        routed_experts_cls: Custom RoutedExperts class (None = use default)
+        routed_experts_args: Additional arguments for routed_experts constructor
+
+    Returns:
+        MoERunner: Configured MoE execution pipeline ready for forward passes
+    """
+    vllm_config = get_current_vllm_config()
+
+    layer_name = prefix
+
+    moe_activation = MoEActivation.from_str(activation)
+    is_act_and_mul = moe_activation.is_gated
+
+    moe_parallel_config = make_parallel_config(
+        tp_size=tp_size,
+        dp_size=dp_size,
+        pcp_size=pcp_size,
+        is_sequence_parallel=is_sequence_parallel,
+        parallel_config=vllm_config.parallel_config,
+    )
+
+    global_num_experts, logical_num_experts, num_fused_shared_experts = (
+        determine_expert_counts(
+            num_experts,
+            num_redundant_experts,
+            n_shared_experts,
+            is_act_and_mul,
+        )
+    )
+
+    # Initialize EPLB manager (or None?)
+    eplb_state: EplbLayerState | None = None
+    if enable_eplb:
+        use_ep = moe_parallel_config.use_ep
+        ep_size = moe_parallel_config.ep_size
+        if use_ep and global_num_experts % ep_size != 0:
+            raise ValueError(
+                f"EPLB currently only supports even distribution of "
+                f"experts across ranks. Got {global_num_experts} experts "
+                f"and {ep_size} EP ranks."
+            )
+        eplb_state = EplbLayerState()
+    else:
+        assert num_redundant_experts == 0, (
+            "Redundant experts are only supported with EPLB."
+        )
+
+    max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+
+    # Create ExpertMapManager to handle expert mapping and placement for EP.
+    # See ExpertMapManager for a detailed description of what it does and when
+    # it is required.
+    expert_map_manager = ExpertMapManager(
+        max_num_batched_tokens=max_num_batched_tokens,
+        top_k=top_k,
+        global_num_experts=global_num_experts,
+        num_redundant_experts=num_redundant_experts,
+        num_expert_group=num_expert_group,
+        moe_parallel_config=moe_parallel_config,
+        placement_strategy=vllm_config.parallel_config.expert_placement_strategy,
+        enable_eplb=eplb_state is not None,
+        num_fused_shared_experts=num_fused_shared_experts,
+        rocm_aiter_enabled=rocm_aiter_ops.is_fused_moe_enabled() and is_act_and_mul,
+    )
+
+    # TODO(bnell): we should not have to create a router if the kernel is
+    # monolithic.
+    if router is None:
+        router = create_fused_moe_router(
+            top_k=top_k,
+            global_num_experts=global_num_experts,
+            eplb_state=eplb_state,
+            renormalize=renormalize,
+            use_grouped_topk=use_grouped_topk,
+            num_expert_group=num_expert_group,
+            topk_group=topk_group,
+            custom_routing_function=custom_routing_function,
+            scoring_func=scoring_func,
+            # When apply_routed_scale_to_output is True, we set the scaling factor
+            # to 1.0 so it ends up being a nop. Applying the scale will be handled
+            # by the runner in this case.
+            # The member variable must be set in the same way as the router since
+            # some quantization methods can access it.
+            routed_scaling_factor=routed_scaling_factor
+            if not apply_routed_scale_to_output
+            else 1.0,
+            e_score_correction_bias=e_score_correction_bias,
+            num_fused_shared_experts=num_fused_shared_experts,
+            zero_expert_type=zero_expert_type,
+            num_logical_experts=logical_num_experts,
+            hash_indices_table=hash_indices_table,
+        )
+
+    if params_dtype is None:
+        params_dtype = torch.get_default_dtype()
+
+    # FIXME (varun): We should have a better way of inferring the activation
+    # datatype. This works for now as the tensor datatype entering the MoE
+    # operation is typically unquantized (i.e. float16/bfloat16).
+    if vllm_config.model_config is not None:
+        moe_in_dtype = vllm_config.model_config.dtype
+    else:
+        # TODO (bnell): This is a hack to get test_mixtral_moe to work
+        # since model_config is not set in the pytest test.
+        moe_in_dtype = params_dtype
+
+    moe_config = FusedMoEConfig(
+        num_experts=global_num_experts,
+        experts_per_token=top_k,
+        hidden_dim=hidden_size,
+        intermediate_size=intermediate_size,
+        num_local_experts=expert_map_manager.local_num_experts,
+        num_logical_experts=logical_num_experts,
+        moe_parallel_config=moe_parallel_config,
+        in_dtype=moe_in_dtype,
+        moe_backend=vllm_config.kernel_config.moe_backend,
+        router_logits_dtype=router_logits_dtype,
+        max_num_tokens=max_num_batched_tokens,
+        has_bias=has_bias,
+        is_lora_enabled=vllm_config.lora_config is not None,
+        activation=moe_activation,
+        device=vllm_config.device_config.device,
+        routing_method=router.routing_method_type,  # Not ideal
+        swiglu_limit=swiglu_limit,
+        max_capture_size=vllm_config.compilation_config.max_cudagraph_capture_size,
+    )
+
+    logger.debug("FusedMoEConfig = %s", moe_config)
+
+    # Create RoutedExperts instance BEFORE create_weights()
+    # This will hold all expert weight parameters
+    if routed_experts_cls is None:
+        routed_experts_cls = RoutedExperts
+
+    assert params_dtype is not None
+    routed_experts = routed_experts_cls(
+        layer_name,
+        params_dtype,
+        moe_config,
+        quant_config,
+        expert_map_manager=expert_map_manager,
+        expert_mapping=expert_mapping,
+        # Extra params that are needed by quant_methods, pass along for now
+        # Prefer getting these from other sources, e.g. moe_config or
+        # router object
+        renormalize=renormalize,
+        use_grouped_topk=use_grouped_topk,
+        num_expert_group=num_expert_group,
+        topk_group=topk_group,
+        custom_routing_function=custom_routing_function,
+        scoring_func=scoring_func,
+        routed_scaling_factor=routed_scaling_factor
+        if not apply_routed_scale_to_output
+        else 1.0,
+        swiglu_limit=swiglu_limit,
+        # TODO get from router? needs to be truncated?
+        e_score_correction_bias=e_score_correction_bias,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        **routed_experts_args if routed_experts_args is not None else {},
+    )
+
+    if runner_cls is None:
+        runner_cls = MoERunner
+
+    runner = runner_cls(
+        layer_name=layer_name,
+        moe_config=moe_config,
+        router=router,
+        routed_experts=routed_experts,
+        enable_dbo=vllm_config.parallel_config.enable_dbo,
+        gate=gate,
+        shared_expert_gate=shared_expert_gate,
+        shared_experts=shared_experts,
+        routed_input_transform=routed_input_transform,
+        routed_output_transform=routed_output_transform,
+        # When apply_routed_scale_to_output is True, we allow
+        # the scaling factor to be passed to the runner, otherwise
+        # we pass 1.0 so it ends up being a nop.
+        routed_scaling_factor=routed_scaling_factor
+        if apply_routed_scale_to_output
+        else 1.0,
+        **runner_args if runner_args is not None else {},
+    )
+
+    return runner
+
+
+def fused_moe_make_expert_params_mapping(
+    model: torch.nn.Module,
+    ckpt_gate_proj_name: str,
+    ckpt_down_proj_name: str,
+    ckpt_up_proj_name: str,
+    num_experts: int,
+    num_redundant_experts: int = 0,
+    routed_experts_prefix: str = "routed_experts",
+) -> list[tuple[str, str, int, str]]:
+    """Delegate to EPLB manager."""
+    return RoutedExperts.make_expert_params_mapping(
+        model,
+        ckpt_gate_proj_name,
+        ckpt_down_proj_name,
+        ckpt_up_proj_name,
+        num_experts,
+        num_redundant_experts,
+        routed_experts_prefix,
+    )

ex_engine/moe/moe_align_block_size.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.triton_utils import triton
+from vllm.utils.math_utils import round_up
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+    expert_map: torch.Tensor | None = None,
+    pad_sorted_ids: bool = False,
+    ignore_invalid_experts: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Aligns the token distribution across experts to be compatible with block
+    size for matrix multiplication.
+
+    Note: In the case of expert_parallel, moe_align_block_size initially
+    considers all experts as valid and aligns all tokens appropriately.
+    Before the function returns it marks the experts_ids that are not in
+    the current GPU rank as -1 so the MoE matmuls could skip those blocks.
+    This requires the num_experts input arg to be the num global experts.
+
+    Parameters:
+    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
+        top-k expert indices for each token.
+    - block_size: The block size used in block matrix multiplication.
+    - num_experts: The total number of experts.
+    - expert_map: A tensor of shape [num_experts] that maps the expert index
+        from the global space to the local index space of the current
+        expert parallel shard. If the expert is not in the current expert
+        parallel shard, the mapping is set to -1.
+    - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
+        should be padded to a multiple of block_size,
+    - ignore_invalid_experts: A flag indicating whether to ignore invalid
+        experts. When False, all expert_ids in topk_ids will participate in
+        counting and ranking, but invalid experts in expert_ids will be marked
+        as -1. When True, all invalid expert_ids in topk_ids will be ignored
+        and will not participate in counting or ranking, and there will be no
+        -1 in expert_ids.
+
+    Returns:
+    - sorted_token_ids: A tensor containing the sorted token indices according
+        to their allocated expert.
+    - expert_ids: A tensor indicating the assigned expert index for each block.
+    - num_tokens_post_padded: The total number of tokens after padding,
+        ensuring divisibility by block_size.
+
+    This function pads the number of tokens that each expert needs to process
+    so that it is divisible by block_size.
+    Padding ensures that during block matrix multiplication, the dimensions
+    align correctly.
+
+    Example:
+    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
+    block_size = 4, and num_experts = 4:
+    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
+        with each expert needing to process 3 tokens.
+    - As block_size is 4, we pad 1 token for each expert.
+    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
+    - Then append padding tokens [12, 12, 12, 12] for each block.
+    - After sorting by expert index, we obtain token_ids
+        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
+        Tokens 12 are non-existent (padding) and are ignored in
+        the subsequent matrix multiplication.
+    - The padding ensures that the total number of tokens is now divisible
+        by block_size for proper block matrix operations.
+    """
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    if pad_sorted_ids:
+        max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    if topk_ids.numel() < num_experts:
+        max_num_tokens_padded = min(
+            topk_ids.numel() * block_size, max_num_tokens_padded
+        )
+    sorted_ids = torch.empty(
+        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
+    )
+    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
+    expert_ids = torch.empty(
+        (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
+    )
+    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
+
+    ops.moe_align_block_size(
+        topk_ids,
+        num_experts,
+        block_size,
+        sorted_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        expert_map if ignore_invalid_experts else None,
+    )
+
+    if expert_map is not None and not ignore_invalid_experts:
+        expert_ids = expert_map[expert_ids]
+
+    return sorted_ids, expert_ids, num_tokens_post_pad
+
+
+def batched_moe_align_block_size(
+    max_tokens_per_batch: int, block_size: int, expert_num_tokens: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Given num_batches, max_tokens_per_batch, block_size and the number of
+    valid-tokens in each batch, prepare sorted_token_ids, expert_ids and
+    num_tokens_post_pad. sorted_token_ids, expert_ids and num_tokens_post_pad
+    have the same semantics as in moe_align_block_size.
+
+    This function is intended to be a drop in replacement for
+    moe_align_batch_size for the batched case.
+
+    Parameters:
+    - max_tokens_per_batch (int): Number of tokens in each batch (both
+        valid and invalid).
+    - block_size (int): block_size to align the data to.
+    - expert_num_tokens (torch.Tensor): expert_num_tokens[i], indicates
+        the number of valid tokens in batch i.
+
+    Returns:
+    - sorted_token_ids (torch.Tensor): Torch tensor of size
+        (num_batches * max_tokens_per_batch) indicating the token indices for
+        that block.
+    - expert_ids (torch.Tensor): Torch tensor of size
+        ceil((num_batches * max_tokens_per_batch) / block_size) indicating
+        what expert to use for each block.
+    - num_tokens_post_pad (torch.Tensor): Torch tensor of size 1
+        indicating the number of valid blocks with actual data to
+        process. This is represented in terms of num tokens.
+    Example:
+    Let num_batches=5, max_tokens_per_batch=8, block_size=4, and
+    expert_num_tokens=[2, 3, 0, 6, 8]. This expert_num_tokens tensor
+    indicates that,
+     - The first 2 tokens in the 0th batch are valid and the rest 6 are
+     invalid (i.e. in the 2D hidden_states tensor of shape,
+     [num_batches * max_tokens_per_batch, K], indices 0, 1 are valid)
+     - The first 3 tokens in the 1st batch are valid. i.e. indices 8, 9, 10
+     - 0 tokens in the 2nd batch are valid
+     - first 6 tokens in the  3rd batch are valid. i.e. indices,
+     24, 25, 26, 27, 28, 29
+     - so on ...
+
+     In this case,
+      sorted_token_ids will be [0, 1, 40, 40,
+                                8, 9, 10, 40,
+                                24, 25, 26, 27,
+                                28, 29, 40, 40,
+                                32, 33, 34, 35,
+                                36, 37, 38, 39,
+                                40, 40, 40, 40,
+                                (rest all 40, 40, 40, 40)
+                                ...]
+      Here, 40 represents an invalid index. as there is no token index 40.
+      The gemm kernel using this sorted_token_ids is expected to skip the
+      gemm computation when it encounters this invalid index.
+
+      expert_ids will be [0, 1, 3, 3, 4, 5, 5, -1, -1, (rest all -1) ...]
+      Here, -1 represents an invalid expert. The gemm kernel using this
+      expert_ids is expected to skip the gemm computation when it encounters
+      an expert of id -1.
+
+      num_tokens_post_pad will be 24 as sorted_token_ids has valid entries
+      until 24.
+    """
+
+    B = expert_num_tokens.size(0)
+    device = expert_num_tokens.device
+
+    # Round up so each batch can be split to blocks evenly.
+    max_num_tokens_padded = B * round_up(max_tokens_per_batch, block_size)
+
+    sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device=device)
+    assert max_num_tokens_padded % block_size == 0
+    max_num_m_blocks = max_num_tokens_padded // block_size
+    expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device=device)
+    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=device)
+
+    ops.batched_moe_align_block_size(
+        max_tokens_per_batch,
+        block_size,
+        expert_num_tokens,
+        sorted_ids,
+        expert_ids,
+        num_tokens_post_pad,
+    )
+
+    return sorted_ids, expert_ids, num_tokens_post_pad

ex_engine/moe/moe_fused_mul_sum.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+from torch._subclasses.fake_tensor import FakeTensor
+
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def moe_fused_mul_sum_kernel(
+    inputs_ptr,
+    topk_weights_ptr,
+    outputs_ptr,
+    top_ids_ptr,
+    expert_map_ptr,
+    num_tokens,
+    stride_m,
+    has_expert_map: tl.constexpr,
+    top_k: tl.constexpr,
+    size: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_k = tl.program_id(0)
+    pid_m = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
+
+    m_mask = offs_m < num_tokens
+    k_mask = offs_k < size
+    mask = m_mask[:, None] & k_mask[None, :]
+
+    a_base = inputs_ptr + (offs_m * stride_m)[:, None] + offs_k[None, :]
+    b_base = topk_weights_ptr + offs_m * top_k
+
+    acc = tl.zeros((BLOCK_M, BLOCK_K), dtype=tl.float32)
+
+    for n in tl.static_range(top_k):
+        b_val = tl.load(b_base + n, mask=m_mask, other=0.0).to(tl.float32)
+        if has_expert_map:
+            id_val = tl.load(top_ids_ptr + offs_m * top_k + n, mask=m_mask, other=0)
+            expert_mask = tl.load(expert_map_ptr + id_val) >= 0
+            a_vec = tl.load(
+                a_base + n * size,
+                mask=mask & expert_mask[:, None],
+                other=0.0,
+            ).to(tl.float32)
+        else:
+            a_vec = tl.load(
+                a_base + n * size,
+                mask=mask,
+                other=0.0,
+            ).to(tl.float32)
+        acc += a_vec * b_val[:, None]
+
+    out_ptrs = outputs_ptr + (offs_m * size)[:, None] + offs_k[None, :]
+    tl.store(
+        out_ptrs,
+        acc.to(outputs_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def _heuristic_config(
+    num_tokens: int,
+    top_k: int,
+    size: int,
+    element_size: int,
+):
+    is_fp32 = element_size > 2
+    is_sm90_plus = current_platform.has_device_capability(90)
+    is_sm80_before = not current_platform.has_device_capability(80)
+
+    if current_platform.has_device_capability(90):
+        # SM90/SM100+: prefer small tiles + many CTAs.
+        if is_fp32:
+            BLOCK_M = 1 if num_tokens <= 4 else 2
+        else:
+            if num_tokens <= 4:
+                BLOCK_M = 1
+            elif num_tokens <= 128:
+                BLOCK_M = 2
+            else:
+                BLOCK_M = 4
+    elif is_fp32:
+        if num_tokens <= 4:
+            BLOCK_M = 1
+        elif num_tokens <= 32:
+            BLOCK_M = 2
+        elif num_tokens <= 128:
+            BLOCK_M = 4
+        else:
+            BLOCK_M = 4
+    else:
+        if num_tokens <= 4:
+            BLOCK_M = 1
+        elif num_tokens <= 32:
+            BLOCK_M = 2
+        elif num_tokens <= 128:
+            BLOCK_M = 4
+        elif num_tokens <= 1024:
+            BLOCK_M = 16
+        else:
+            BLOCK_M = 8
+
+    if is_fp32:
+        max_block_k = 256
+    elif is_sm80_before or is_sm90_plus:
+        max_block_k = 512
+    else:
+        max_block_k = 1024
+    BLOCK_K = min(triton.next_power_of_2(size), max_block_k)
+    BLOCK_K = max(BLOCK_K, 256)
+
+    total = BLOCK_M * BLOCK_K
+    if is_fp32:
+        num_warps = max(8, min(16, total // 64))
+    else:
+        num_warps = max(4, min(16, total // 256))
+
+    if is_sm80_before:
+        num_warps = min(num_warps, 8)
+        num_stages = 2
+    elif is_sm90_plus:
+        num_warps = min(num_warps, 8)
+        num_stages = 4 if total <= 2048 else 2
+    else:
+        num_stages = 4 if total <= 2048 else 2
+
+    return BLOCK_M, BLOCK_K, num_warps, num_stages
+
+
+def moe_fused_mul_sum(
+    inputs: torch.Tensor,
+    topk_weights: torch.Tensor,
+    outputs: torch.Tensor | None = None,
+    topk_ids: torch.Tensor | None = None,
+    expert_map: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """
+    Fused kernel for MoE (Mixture of Experts) to perform weighted summation
+    of expert outputs.
+
+    Args:
+        inputs: The output from experts.
+            Shape: (num_tokens, top_k, hidden_size).
+        topk_weights: The weights assigned to each expert for each token.
+            Shape: (num_tokens, top_k).
+        outputs: Optional pre-allocated output tensor.
+            Shape: (num_tokens, hidden_size).
+        topk_ids: Optional indices of the top-k experts. Used when
+            `expert_map` is provided. Shape: (num_tokens, top_k).
+        expert_map: Optional mapping for Expert Parallelism. A value < 0
+            indicates an invalid token/expert pair that will be skipped.
+
+    Returns:
+        The fused weighted sum of expert outputs.
+        Shape: (num_tokens, hidden_size).
+    """
+    assert inputs.ndim == 3
+    assert topk_weights.ndim == 2
+    assert inputs.is_contiguous()
+    assert topk_weights.is_contiguous()
+    assert inputs.dtype in (torch.float32, torch.float16, torch.bfloat16)
+    assert topk_weights.dtype in (torch.float32, torch.float16, torch.bfloat16)
+
+    num_tokens, top_k, size = inputs.shape
+    output_shape = (num_tokens, size)
+    if outputs is None:
+        outputs = torch.empty(output_shape, dtype=inputs.dtype, device=inputs.device)
+
+    assert outputs.shape == output_shape
+    assert topk_weights.shape == (num_tokens, top_k)
+
+    if not isinstance(inputs, FakeTensor):
+        BLOCK_M, BLOCK_K, num_warps, num_stages = _heuristic_config(
+            num_tokens,
+            top_k,
+            size,
+            inputs.element_size(),
+        )
+        grid = (triton.cdiv(size, BLOCK_K), triton.cdiv(num_tokens, BLOCK_M))
+        moe_fused_mul_sum_kernel[grid](
+            inputs,
+            topk_weights,
+            outputs,
+            topk_ids,
+            expert_map,
+            num_tokens,
+            top_k * size,
+            expert_map is not None,
+            top_k,
+            size,
+            BLOCK_M,
+            BLOCK_K,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+
+    return outputs

ex_engine/moe/moe_permute_unpermute.py

@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass, field
+
+import torch
+
+
+@dataclass
+class MoEPermuteScratch:
+    # Reused metadata buffers for repeated grouped-MoE permutes.
+    max_num_tokens: int
+    topk: int
+    num_experts: int
+    num_local_experts: int
+    device: torch.device
+    hidden_size: int | None = None
+    hidden_dtype: torch.dtype | None = None
+    token_expert_indices: torch.Tensor = field(init=False)
+    expert_first_token_offset: torch.Tensor = field(init=False)
+    permuted_idx: torch.Tensor = field(init=False)
+    inv_permuted_idx: torch.Tensor = field(init=False)
+    permuted_hidden_states: torch.Tensor | None = field(init=False, default=None)
+    sort_workspace: torch.Tensor = field(init=False)
+    permuted_experts_id: torch.Tensor = field(init=False)
+    sorted_row_idx: torch.Tensor = field(init=False)
+    topk_ids_int32: torch.Tensor = field(init=False)
+    topk_ids_for_sort: torch.Tensor = field(init=False)
+    max_expanded_rows: int = field(init=False)
+
+    def __post_init__(self) -> None:
+        assert self.max_num_tokens > 0
+        assert self.topk > 0
+        assert self.num_experts > 0
+        assert self.num_local_experts > 0
+        if self.hidden_size is None:
+            assert self.hidden_dtype is None
+        else:
+            assert self.hidden_dtype is not None
+
+        self.max_expanded_rows = self.max_num_tokens * self.topk
+        self.token_expert_indices = torch.arange(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        self.expert_first_token_offset = torch.empty(
+            self.num_local_experts + 1, dtype=torch.int64, device=self.device
+        )
+        self.permuted_idx = torch.empty(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        self.inv_permuted_idx = torch.empty(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        if self.hidden_size is not None:
+            hidden_numel = self.max_expanded_rows * self.hidden_size
+            self.permuted_hidden_states = torch.empty(
+                hidden_numel, dtype=self.hidden_dtype, device=self.device
+            )
+        self.permuted_experts_id = torch.empty(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        self.sorted_row_idx = torch.empty(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        self.topk_ids_int32 = torch.empty(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        self.topk_ids_for_sort = torch.empty(
+            self.max_expanded_rows, dtype=torch.int32, device=self.device
+        )
+        sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size(
+            self.max_expanded_rows, self.num_experts
+        )
+        self.sort_workspace = torch.empty(
+            sorter_size, dtype=torch.int8, device=self.device
+        )
+        # torch.device("cuda") in config, after initialized,
+        # will be changed to cuda:{index}, so we need to refresh here.
+        self.device = self.token_expert_indices.device
+
+    def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None:
+        n_token, n_hidden = hidden_states.shape
+        assert hidden_states.device == self.device
+        assert topk_ids.device == self.device
+        assert n_token <= self.max_num_tokens
+        assert topk_ids.size(1) == self.topk
+        assert topk_ids.size(0) == n_token
+        if self.hidden_size is not None:
+            assert n_hidden == self.hidden_size
+            assert hidden_states.dtype == self.hidden_dtype
+            assert self.permuted_hidden_states is not None
+
+    def token_expert_indices_view(self, n_token: int) -> torch.Tensor:
+        return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk)
+
+    def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
+        if topk_ids.dtype == torch.int32:
+            return topk_ids
+        numel = topk_ids.numel()
+        topk_ids_int32 = self.topk_ids_int32[:numel].view_as(topk_ids)
+        topk_ids_int32.copy_(topk_ids)
+        return topk_ids_int32
+
+
+def moe_permute(
+    hidden_states: torch.Tensor,
+    a1q_scale: torch.Tensor | None,
+    topk_ids: torch.Tensor,
+    n_expert: int,
+    n_local_expert: int = -1,
+    expert_map: torch.Tensor | None = None,
+    permuted_hidden_states: torch.Tensor | None = None,
+    scratch: MoEPermuteScratch | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    This function expands and permutes activation to gather uncontinuous tokens
+      for each expert.
+    Parameters:
+    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
+    - a1q_scale (Optional[torch.Tensor]): quant scale for hidden_states
+    - topk_ids (torch.Tensor): topk expert route id for each token.
+    - n_expert (int): The number of expert.
+    - n_local_expert (int): The number of expert in current EP rank.
+    - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
+        from the global expert space to the local expert space of the expert
+        parallel shard.
+    - permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor.
+        If None, the output tensor will be created in this function.
+    Returns:
+    - permuted_hidden_states (torch.Tensor): permuted activation.
+    - a1q_scale (Optional[torch.Tensor]): permuted quant scale for hidden_states
+        if original scale not per-tensor scaling
+    - expert_first_token_offset (torch.Tensor): offset of the first token
+       of each expert for standard grouped gemm.
+    - inv_permuted_idx (torch.Tensor): idx map for moe_unpermute.
+    - permuted_idx (torch.Tensor): idx map from hidden to permuted_hidden.
+    """
+    n_token, n_hidden = hidden_states.size()
+    topk = topk_ids.size(1)
+    assert (n_hidden * hidden_states.element_size()) % 16 == 0, (
+        "permue kernel need hidden dim align to 16B"
+    )
+    permuted_row_size = n_token * topk
+    if n_local_expert == -1:
+        n_local_expert = n_expert
+    if permuted_hidden_states is None:
+        if scratch is None:
+            permuted_hidden_states = torch.empty(
+                (permuted_row_size, n_hidden),
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+        else:
+            scratch.validate(hidden_states, topk_ids)
+            hidden_numel = permuted_row_size * n_hidden
+            scratch_hidden_states = scratch.permuted_hidden_states
+            assert scratch_hidden_states is not None
+            permuted_hidden_states = scratch_hidden_states[:hidden_numel].view(
+                permuted_row_size, n_hidden
+            )
+    assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), (
+        f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}"
+        f" but got {permuted_hidden_states.size()}"
+    )
+
+    if scratch is None:
+        token_expert_indices = torch.arange(
+            0, n_token * topk, dtype=torch.int32, device=hidden_states.device
+        ).reshape((n_token, topk))
+
+        expert_first_token_offset = torch.empty(
+            n_local_expert + 1, dtype=torch.int64, device=hidden_states.device
+        )
+        permuted_idx = torch.full(
+            (permuted_row_size,),
+            n_token * topk,
+            dtype=torch.int32,
+            device=hidden_states.device,
+        )
+        inv_permuted_idx = torch.empty(
+            (n_token, topk), dtype=torch.int32, device=hidden_states.device
+        )
+        topk_ids_int32 = topk_ids.to(torch.int32)
+        torch.ops._moe_C.moe_permute(
+            hidden_states,
+            topk_ids_int32,
+            token_expert_indices,
+            expert_map,
+            n_expert,
+            n_local_expert,
+            topk,
+            permuted_hidden_states,
+            expert_first_token_offset,
+            inv_permuted_idx,
+            permuted_idx,
+        )
+    else:
+        scratch.validate(hidden_states, topk_ids)
+        assert n_expert == scratch.num_experts
+        assert n_local_expert == scratch.num_local_experts
+        token_expert_indices = scratch.token_expert_indices_view(n_token)
+        expert_first_token_offset = scratch.expert_first_token_offset
+        permuted_idx = scratch.permuted_idx[:permuted_row_size]
+        permuted_idx.fill_(permuted_row_size)
+        inv_permuted_idx = scratch.inv_permuted_idx[:permuted_row_size].view(
+            n_token, topk
+        )
+        permuted_experts_id = scratch.permuted_experts_id[:permuted_row_size].view(
+            n_token, topk
+        )
+        sorted_row_idx = scratch.sorted_row_idx[:permuted_row_size].view(n_token, topk)
+        topk_ids_for_sort = scratch.topk_ids_for_sort[:permuted_row_size].view(
+            n_token, topk
+        )
+        topk_ids_int32 = scratch.prepare_topk_ids(topk_ids)
+        torch.ops._moe_C.moe_permute_with_scratch(
+            hidden_states,
+            topk_ids_int32,
+            token_expert_indices,
+            expert_map,
+            n_expert,
+            n_local_expert,
+            topk,
+            permuted_hidden_states,
+            expert_first_token_offset,
+            inv_permuted_idx,
+            permuted_idx,
+            scratch.sort_workspace,
+            permuted_experts_id,
+            sorted_row_idx,
+            topk_ids_for_sort,
+        )
+
+    if a1q_scale is not None and a1q_scale.dim() > 1:
+        a1q_scale = a1q_scale[permuted_idx.clamp(max=n_token * topk - 1) // topk]
+    return (
+        permuted_hidden_states,
+        a1q_scale,
+        expert_first_token_offset,
+        inv_permuted_idx.flatten(),
+        permuted_idx,
+    )
+
+
+def moe_unpermute(
+    out: torch.Tensor,
+    permuted_hidden_states: torch.Tensor,
+    topk_weights: torch.Tensor,
+    inv_permuted_idx: torch.Tensor,
+    expert_first_token_offset: torch.Tensor | None = None,
+) -> None:
+    """
+    This function expands and permutes activation to gathering uncontinuous
+      tokens for each expert.
+    Parameters:
+    - out (torch.Tensor): output tensor
+    - permuted_hidden_states (torch.Tensor): permuted activation.
+    - topk_weights (torch.Tensor): topk expert route weight for each token.
+    - inv_permuted_idx (torch.Tensor): row idx map for moe_unpermute.
+    - expert_first_token_offset (Optional[torch.Tensor]): offset of the first
+      token of each expert for grouped gemm.
+    Returns:
+    - hidden_states (torch.Tensor): The reduced and unpermuted activation
+      tensor.
+    """
+    topk = topk_weights.size(1)
+    n_hidden = permuted_hidden_states.size(-1)
+    assert (n_hidden * permuted_hidden_states.element_size()) % 16 == 0, (
+        "unpermue kernel need hidden dim align to 16B"
+    )
+
+    torch.ops._moe_C.moe_unpermute(
+        permuted_hidden_states,
+        topk_weights,
+        inv_permuted_idx,
+        expert_first_token_offset,
+        topk,
+        out,
+    )
+
+
+def moe_permute_unpermute_supported():
+    return torch.ops._moe_C.moe_permute_unpermute_supported()

ex_engine/moe/prepare_finalize/__init__.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
+    BatchedPrepareAndFinalize,
+)
+from vllm.model_executor.layers.fused_moe.prepare_finalize.naive_dp_ep import (
+    MoEPrepareAndFinalizeNaiveDPEPModular,
+    MoEPrepareAndFinalizeNaiveDPEPMonolithic,
+    make_moe_prepare_and_finalize_naive_dp_ep,
+)
+from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import (
+    MoEPrepareAndFinalizeNoDPEPModular,
+    MoEPrepareAndFinalizeNoDPEPMonolithic,
+    make_moe_prepare_and_finalize_no_dp_ep,
+)
+
+__all__ = [
+    "BatchedPrepareAndFinalize",
+    "MoEPrepareAndFinalizeNaiveDPEPMonolithic",
+    "MoEPrepareAndFinalizeNaiveDPEPModular",
+    "make_moe_prepare_and_finalize_naive_dp_ep",
+    "MoEPrepareAndFinalizeNoDPEPMonolithic",
+    "MoEPrepareAndFinalizeNoDPEPModular",
+    "make_moe_prepare_and_finalize_no_dp_ep",
+    # deepep_ht, deepep_ll, and flashinfer_a2a are not
+    # imported here as they have optional dependencies (deep_ep, flashinfer).
+    # Import them directly from their modules as needed.
+]

ex_engine/moe/prepare_finalize/batched.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceDelegate,
+    TopKWeightAndReduceNaiveBatched,
+)
+from vllm.model_executor.layers.fused_moe.utils import (
+    moe_kernel_quantize_input,
+    normalize_scales_shape,
+)
+
+
+class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
+    """
+    A reference prepare/finalize class that reorganizes the tokens into
+    expert batched format, i.e. E x max_num_tokens x K.  This is the format
+    that the batched dispatch/combine kernels use.
+    """
+
+    def __init__(
+        self,
+        max_num_tokens: int,
+        num_local_experts: int,
+        num_dispatchers: int,
+        rank: int,
+    ):
+        super().__init__()
+        self.max_num_tokens = max_num_tokens
+        self.num_local_experts = num_local_experts
+        self.rank = rank
+        self.num_dispatchers_ = num_dispatchers
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.BatchedExperts
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return self.max_num_tokens
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return None
+
+    def num_dispatchers(self) -> int:
+        return self.num_dispatchers_
+
+    def output_is_reduced(self) -> bool:
+        return False
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareResultType:
+        if defer_input_quant:
+            raise NotImplementedError(
+                f"{self.__class__.__name__} does not support defer_input_quant=True. "
+                "Please select an MoE kernel that accepts quantized inputs."
+            )
+        assert a1.dim() == 2
+        assert topk_ids.dim() == 2
+        assert topk_ids.size(0) == a1.size(0)
+
+        if apply_router_weight_on_input:
+            topk = topk_ids.size(1)
+            # TODO: this only works for topK=1, will need to update for topK>1
+            assert topk == 1, (
+                "apply_router_weight_on_input is only implemented for topk=1"
+            )
+            a1.mul_(topk_weights.to(a1.dtype))
+
+        num_tokens, hidden_dim = a1.size()
+        topk = topk_ids.size(1)
+
+        tokens_per_expert = torch.zeros(num_experts, dtype=torch.int, device=a1.device)
+
+        num_local_experts = self.num_local_experts
+
+        if quant_config.quant_dtype is None:
+            b_type = a1.dtype
+        else:
+            b_type = quant_config.quant_dtype
+
+        b_a1 = torch.zeros(
+            (num_local_experts, self.max_num_tokens, hidden_dim),
+            dtype=b_type,
+            device=a1.device,
+        )
+
+        if quant_config.is_quantized:
+            scale_shape = quant_config.batched_scale_shape(
+                num_local_experts, self.max_num_tokens, hidden_dim
+            )
+
+            b_a1_scale = torch.empty(scale_shape, dtype=torch.float32, device=a1.device)
+        else:
+            assert quant_config.a1_scale is None
+            b_a1_scale = None
+
+        first_expert = num_local_experts * self.rank
+        last_expert = first_expert + num_local_experts
+
+        a1_scale = normalize_scales_shape(quant_config.a1_scale)
+
+        for expert_id in range(first_expert, last_expert):
+            topks = torch.any(topk_ids == expert_id, dim=1).flatten()
+            rows = torch.count_nonzero(topks.flatten())
+            if rows == 0:
+                continue
+            idx = expert_id - first_expert
+            tokens_per_expert[idx] = rows
+            rhs = a1[: topks.numel()][topks]
+            if quant_config.quant_dtype is not None:
+                if a1_scale is not None:
+                    if quant_config.is_per_act_token:
+                        rhs_a1_scale = a1_scale[: topks.numel()][topks]
+                    else:
+                        rhs_a1_scale = a1_scale
+                else:
+                    rhs_a1_scale = None
+                b_a1[idx, :rows, :], b_s = moe_kernel_quantize_input(
+                    rhs,
+                    rhs_a1_scale,
+                    quant_config.quant_dtype,
+                    quant_config.per_act_token_quant,
+                    quant_config.block_shape,
+                )
+                assert b_s is not None
+                if quant_config.is_per_act_token:
+                    b_a1_scale[idx, :rows] = b_s[:rows]
+                else:
+                    b_a1_scale[idx, : b_s.shape[0]] = b_s
+            else:
+                b_a1[idx, :rows, :] = rhs
+
+        assert b_a1_scale is None or b_a1_scale.ndim == 3
+
+        expert_tokens_meta = mk.ExpertTokensMetadata(
+            expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None
+        )
+
+        return b_a1, b_a1_scale, expert_tokens_meta, None, None
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
+            weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)
+        weight_and_reduce_impl.apply(
+            output=output,
+            fused_expert_output=fused_expert_output,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )

ex_engine/moe/prepare_finalize/no_dp_ep.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceContiguous,
+    TopKWeightAndReduceDelegate,
+)
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+
+
+def _quantize_input(
+    a1: torch.Tensor,
+    quant_config: FusedMoEQuantConfig,
+    defer_input_quant: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    # Defer input quant to moe kernel for backends (e.g. AITER, FI)
+    # which use a single kernel call for quant + experts.
+    if defer_input_quant:
+        return a1, None
+
+    input_sf = (
+        quant_config.a1_gscale if quant_config.use_nvfp4_w4a4 else quant_config.a1_scale
+    )
+    a1q, a1q_scale = moe_kernel_quantize_input(
+        a1,
+        input_sf,
+        quant_dtype=quant_config.quant_dtype,
+        per_act_token_quant=quant_config.per_act_token_quant,
+        block_shape=quant_config.block_shape,
+        is_scale_swizzled=quant_config.is_scale_swizzled,
+        mx_alignment=quant_config.mx_alignment,
+    )
+
+    return a1q, a1q_scale
+
+
+class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular):
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return None
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return None
+
+    def num_dispatchers(self) -> int:
+        return 1
+
+    def output_is_reduced(self) -> bool:
+        return False
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareResultType:
+        if apply_router_weight_on_input:
+            topk = topk_ids.size(1)
+            # TODO: this only works for topK=1, will need to update for topK>1
+            assert topk == 1, (
+                "apply_router_weight_on_input is only implemented for topk=1"
+            )
+            a1 = a1 * topk_weights.to(a1.dtype)
+
+        a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant)
+
+        return a1q, a1q_scale, None, None, None
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
+            weight_and_reduce_impl = TopKWeightAndReduceContiguous()
+        weight_and_reduce_impl.apply(
+            output=output,
+            fused_expert_output=fused_expert_output,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )
+
+
+class MoEPrepareAndFinalizeNoDPEPMonolithic(mk.FusedMoEPrepareAndFinalizeMonolithic):
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return None
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return None
+
+    def num_dispatchers(self) -> int:
+        return 1
+
+    def output_is_reduced(self) -> bool:
+        return False
+
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        router_logits: torch.Tensor,
+        quant_config: FusedMoEQuantConfig,
+        defer_input_quant: bool = False,
+    ) -> mk.PrepareMonolithicResultType:
+        a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant)
+        return a1q, a1q_scale, router_logits
+
+    def finalize(
+        self,
+        fused_expert_output: torch.Tensor,
+    ) -> torch.Tensor:
+        return fused_expert_output
+
+
+def make_moe_prepare_and_finalize_no_dp_ep(
+    use_monolithic: bool,
+) -> MoEPrepareAndFinalizeNoDPEPModular | MoEPrepareAndFinalizeNoDPEPMonolithic:
+    return (
+        MoEPrepareAndFinalizeNoDPEPMonolithic()
+        if use_monolithic
+        else MoEPrepareAndFinalizeNoDPEPModular()
+    )

ex_engine/moe/topk_weight_and_reduce.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+import vllm._custom_ops as ops
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+
+
+class TopKWeightAndReduceDelegate(mk.TopKWeightAndReduce):
+    """
+    Useful in the case when some FusedMoEExpertsModular
+    implementation does not perform weight application and reduction
+    but cannot address the needs of all the compatible PrepareAndFinalize
+    implementations.
+    For example, BatchedTritonExperts is compatible with both batched
+    PrepareAndFinalize implementations like DeepEPLLPrepareAndFinalize and
+    BatchedPrepareAndFinalize. Some PrepareAndFinalize implementations do
+    the weight-application + reduction as part of the combine kernel, while
+    BatchedPrepareAndFinalize needs an explicit implementation. To facilitate
+    this case, the BatchedTritonExperts could use TopKWeightAndReduceDelegate
+    so the PrepareAndFinalize implementations could choose how to
+    weight + reduce.
+    """
+
+    def __eq__(self, other):
+        return isinstance(other, TopKWeightAndReduceDelegate)
+
+    def apply(
+        self,
+        output: torch.Tensor | None,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+    ) -> torch.Tensor:
+        raise RuntimeError(
+            "The caller is expected to choose an appropriate "
+            "TopKWeightAndReduce implementation."
+        )
+
+
+class TopKWeightAndReduceNoOP(mk.TopKWeightAndReduce):
+    """
+    The fused_experts outputs have already been weight applied and reduced.
+    This implementation is a no-op.
+    """
+
+    def __eq__(self, other):
+        return isinstance(other, TopKWeightAndReduceNoOP)
+
+    def apply(
+        self,
+        output: torch.Tensor | None,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+    ) -> torch.Tensor:
+        # Weight application and reduction operations are already done.
+        if output is None:
+            return fused_expert_output
+
+        # Skip self-copy when caller aliased fused_out to output upstream.
+        if output is fused_expert_output:
+            return output
+
+        # MoEPrepareAndFinalizeNoDPEPModular needs the output to be in the `output`
+        # tensor.
+        assert output.size() == fused_expert_output.size(), (
+            "output shape is expected to match the fused_expert_output shape. "
+            f"But got output={output.size()}, "
+            f"used_expert_output={fused_expert_output.size()}"
+        )
+        output.copy_(fused_expert_output, non_blocking=True)
+        return output
+
+
+class TopKWeightAndReduceContiguous(mk.TopKWeightAndReduce):
+    """
+    TopKWeightAndReduce implementation for a fused_experts output
+    of shape (m, topk, K)
+    """
+
+    def __eq__(self, other):
+        return isinstance(other, TopKWeightAndReduceContiguous)
+
+    def apply(
+        self,
+        output: torch.Tensor | None,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+    ) -> torch.Tensor:
+        m, num_topk = topk_ids.size()
+        k = fused_expert_output.size(-1)
+        if fused_expert_output.ndim == 2:
+            fused_expert_output = fused_expert_output.view(m, num_topk, k)
+
+        assert fused_expert_output.size() == (m, num_topk, k), (
+            f"Expected fused_expert_output size {(m, num_topk, k)}. But got "
+            f"{fused_expert_output.size()}"
+        )
+
+        if not apply_router_weight_on_input:
+            fused_expert_output.mul_(topk_weights.view(m, -1, 1))
+
+        if output is None:
+            output = torch.empty(
+                (m, k),
+                device=fused_expert_output.device,
+                dtype=fused_expert_output.dtype,
+            )
+        assert output.size() == (m, k), (
+            f"Expected output size {(m, k)}. But got {output.size()}"
+        )
+
+        ops.moe_sum(fused_expert_output, output)
+        return output
+
+
+class TopKWeightAndReduceNaiveBatched(mk.TopKWeightAndReduce):
+    """
+    TopKWeightAndReduce implementation for a fused_experts output
+    of shape (num_experts, batch_size, K)
+    """
+
+    def __init__(self, rank: int):
+        self.rank = rank
+
+    def __eq__(self, other):
+        return isinstance(other, TopKWeightAndReduceNaiveBatched) and (
+            other.rank == self.rank
+        )
+
+    def apply(
+        self,
+        output: torch.Tensor | None,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+    ) -> torch.Tensor:
+        assert fused_expert_output.ndim == 3
+        num_tokens = topk_ids.size(0)
+        num_local_experts = fused_expert_output.size(0)
+        K = fused_expert_output.size(-1)
+
+        if output is None:
+            output = torch.zeros(
+                (num_tokens, K),
+                device=fused_expert_output.device,
+                dtype=fused_expert_output.dtype,
+            )
+        else:
+            output.fill_(0)
+
+        assert output.size() == (num_tokens, K), (
+            f"Expected output size {(num_tokens, K)}, but got {output.size()}"
+        )
+
+        first_expert = num_local_experts * self.rank
+        last_expert = first_expert + num_local_experts
+
+        for expert_id in range(first_expert, last_expert):
+            matching_tokens = topk_ids == expert_id
+            topks = torch.any(matching_tokens, dim=1).flatten()
+            rows = torch.count_nonzero(topks)
+            rhs = fused_expert_output[expert_id - first_expert, :rows, :]
+            if not apply_router_weight_on_input:
+                rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1))
+            output[topks] = output[topks] + rhs
+
+        return output

ex_engine/moe/utils.py

@@ -0,0 +1,441 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from math import prod
+
+import torch
+import torch.nn.functional as F
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    per_token_group_quant_fp8,
+)
+from vllm.model_executor.layers.quantization.utils.int8_utils import (
+    per_token_group_quant_int8,
+    per_token_quant_int8,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+    quant_dequant_mxfp4,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
+    quant_dequant_mxfp6,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    mxfp8_e4m3_quantize,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    ref_nvfp4_quant_dequant,
+)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    per_tensor_dequantize,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+
+
+@triton.jit
+def _count_expert_num_tokens(
+    topk_ids_ptr,
+    expert_num_tokens_ptr,
+    num_experts,
+    topk_numel,
+    expert_map,
+    HAS_EXPERT_MAP: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    curr_expert = tl.program_id(0)
+
+    offsets = tl.arange(0, BLOCK_SIZE)
+    topk_ids_ptrs = topk_ids_ptr + offsets
+
+    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
+    for x in range(tl.cdiv(topk_numel, BLOCK_SIZE)):
+        mask = offsets < (topk_numel - x * BLOCK_SIZE)
+        expert_ids = tl.load(topk_ids_ptrs, mask=mask, other=-1)
+        if HAS_EXPERT_MAP:
+            expert_map_ptrs = expert_map + expert_ids
+            expert_map_mask = expert_ids >= 0
+            expert_ids = tl.load(expert_map_ptrs, mask=expert_map_mask, other=-1)
+
+        has_curr_expert = tl.where(expert_ids == curr_expert, 1, 0)
+        acc = acc + has_curr_expert
+        topk_ids_ptrs += BLOCK_SIZE
+
+    if curr_expert < num_experts:
+        tl.store(expert_num_tokens_ptr + curr_expert, tl.sum(acc))
+
+
+def count_expert_num_tokens(
+    topk_ids: torch.Tensor, num_local_experts: int, expert_map: torch.Tensor | None
+) -> torch.Tensor:
+    """
+    Count the number to tokens assigned to each expert.
+
+    Parameters:
+    - topk_ids (torch.Tensor): Tensor mapping each token to its
+    list of experts.
+    - num_local_experts (int): Number of experts in this rank.
+    - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
+    from the global expert space to the local expert space of the expert
+    parallel shard.
+
+    Returns:
+    A tensor of size num_local_experts, where tensor[i] holds the number
+    of tokens assigned to the ith expert.
+    """
+    assert topk_ids.dtype.is_signed, "The kernel uses -1 to represent invalid topk_ids"
+    expert_num_tokens = torch.empty(
+        (num_local_experts), device=topk_ids.device, dtype=torch.int32
+    )
+
+    grid = num_local_experts
+    BLOCK_SIZE = min(topk_ids.numel(), 1024)
+    BLOCK_SIZE = triton.next_power_of_2(BLOCK_SIZE)
+
+    _count_expert_num_tokens[(grid,)](
+        topk_ids,
+        expert_num_tokens,
+        num_local_experts,
+        topk_ids.numel(),
+        expert_map,
+        HAS_EXPERT_MAP=expert_map is not None,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+    return expert_num_tokens
+
+
+def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
+    """
+    Shrink the given tensor and apply the given view to it.  This is
+    used to resize the intermediate fused_moe caches.
+    """
+    assert prod(v) <= x.numel(), (
+        f"{v} ({prod(v)}) <= {x.shape} ({x.numel()})"
+    )  # CUDAGRAPH unfriendly?
+    return x.flatten()[: prod(v)].view(*v)
+
+
+def _nvfp4_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    is_sf_swizzled_layout: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return ops.scaled_fp4_quant(A, A_scale, is_sf_swizzled_layout=is_sf_swizzled_layout)
+
+
+def _fp8_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    per_act_token: bool,
+    block_shape: list[int] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Perform fp8 quantization on the inputs.  If a block_shape
+    is provided, the output will be blocked.
+    """
+    if block_shape is None:
+        # TODO(luka): use QuantFP8 custom op
+        #  https://github.com/vllm-project/vllm/issues/20711
+        A, A_scale = ops.scaled_fp8_quant(
+            A, A_scale, use_per_token_if_dynamic=per_act_token
+        )
+    else:
+        assert not per_act_token
+        assert len(block_shape) == 2
+        _, block_k = block_shape[0], block_shape[1]
+        A, A_scale = per_token_group_quant_fp8(A, block_k)
+        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)
+
+    return A, A_scale
+
+
+def _int8_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    per_act_token: bool,
+    block_shape: list[int] | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Perform int8 quantization on the inputs.  If a block_shape
+    is provided, the output will be blocked.
+    """
+
+    # If weights are per-channel (per_channel_quant=True), then
+    # activations apply per-token quantization. Otherwise, assume
+    # activation tensor-wise fp8/int8 quantization, dynamic or static
+    if block_shape is None:
+        if per_act_token:
+            A, A_scale = per_token_quant_int8(A)
+        elif A_scale is not None:
+            # Static per-tensor: use the optimized CUDA kernel
+            A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale)
+        elif A_scale is None:
+            # Dynamic per-tensor: compute scale then quantize via kernel
+            A_scale = torch.clamp(A.abs().max() / 127.0, min=1e-10)
+            A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale)
+    else:
+        assert not per_act_token
+        assert len(block_shape) == 2
+        _, block_k = block_shape[0], block_shape[1]
+        A, A_scale = per_token_group_quant_int8(A, block_k)
+        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)
+
+    return A, A_scale
+
+
+def _mxfp4_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+) -> tuple[torch.Tensor, None]:
+    assert block_shape is None
+    # TODO: native mxfp4 is currently not integrated in vllm,
+    # so simulating even on devices supporting this data type natively.
+    # Once integrated, `current_platform.supports_mx()` should be used to
+    # control quantize+dequantize, or simply quantize here down to mxfp4.
+    A = quant_dequant_mxfp4(A)
+
+    return A, None
+
+
+def _mxfp8_e4m3_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+    is_sf_swizzled_layout: bool = False,
+    mx_alignment: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert A_scale is None
+    assert not per_act_token_quant
+    assert block_shape is None or block_shape == [1, 32]
+    return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout, mx_alignment)
+
+
+def _mxfp6_e3m2_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+) -> tuple[torch.Tensor, None]:
+    assert block_shape is None
+
+    # TODO: native mxfp6 is currently not integrated in vllm,
+    # so simulating even on devices supporting this data type natively.
+    # Eventually, there should be a check based on
+    # `current_platform.supports_mx()` here.
+    A = quant_dequant_mxfp6(A, quant_dtype="fp6_e3m2")
+
+    return A, None
+
+
+def _mxfp6_e2m3_quantize(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+) -> tuple[torch.Tensor, None]:
+    assert block_shape is None
+
+    # TODO: native mxfp6 is currently not integrated in vllm,
+    # so simulating even on devices supporting this data type natively.
+    # Eventually, there should be a check based on
+    # `current_platform.supports_mx()` here.
+    A = quant_dequant_mxfp6(A, quant_dtype="fp6_e2m3")
+
+    return A, None
+
+
+def moe_kernel_quantize_input(
+    A: torch.Tensor,
+    A_scale: torch.Tensor | None,
+    quant_dtype: None | torch.dtype | str,
+    per_act_token_quant: bool,
+    block_shape: list[int] | None = None,
+    is_scale_swizzled: bool = True,
+    ocp_mx_scheme: str | None = None,
+    quantization_emulation: bool = False,
+    mx_alignment: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation
+    if ocp_mx_scheme is not None:
+        if ocp_mx_scheme in {"w_mxfp4", "w_mxfp4_a_mxfp4"}:
+            pass  # No QDQ needed for these schemes
+        elif ocp_mx_scheme.endswith("a_fp8"):
+            # Perform QDQ (quantize and dequantize) on activation for emulation
+            # purpose, because there is no native kernel for weight in ocp_mx_scheme
+            # and activation in FP8. The implementation is based on existing
+            # non-emulation ops.
+            qA, qA_scale = ops.scaled_fp8_quant(
+                A, A_scale, use_per_token_if_dynamic=False
+            )
+            A = per_tensor_dequantize(qA, qA_scale).to(A.dtype)
+            # After QDQ, we don't need further quantization
+            return A, None
+        # else: For other schemes (e.g., *_a_mxfp6_e3m2, *_a_mxfp6_e2m3),
+        # weights are already dequantized, and we proceed with normal
+        # activation quantization below.
+
+    if quant_dtype == current_platform.fp8_dtype():
+        if quantization_emulation:
+            raise NotImplementedError(
+                f"moe_kernel_quantize_input does not support quant_dtype={quant_dtype}"
+                " MOE quantization emulation. Please open an issue."
+            )
+        return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
+    elif quant_dtype == torch.int8:
+        if quantization_emulation:
+            raise NotImplementedError(
+                "moe_kernel_quantize_input does not support quant_dtype=torch.int8"
+                " MOE quantization emulation. Please open an issue."
+            )
+        return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
+    elif quant_dtype == "nvfp4":
+        if not quantization_emulation:
+            return _nvfp4_quantize(A, A_scale, is_sf_swizzled_layout=is_scale_swizzled)
+        else:
+            A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
+            return A, None
+    elif quant_dtype == "mxfp4":
+        if not quantization_emulation:
+            raise NotImplementedError(
+                "moe_kernel_quantize_input should not be used for native"
+                " quant_dtype='mxfp4' MOE. Please open an issue."
+            )
+        return _mxfp4_quantize(A, A_scale, per_act_token_quant, block_shape)
+    elif quant_dtype == "mxfp8":
+        # TODO: `quant_dtype == "mxfp8"` is ambiguous,
+        # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`.
+        if quantization_emulation:
+            raise NotImplementedError(
+                "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE "
+                "quantization emulation. Please open an issue."
+            )
+        return _mxfp8_e4m3_quantize(
+            A,
+            A_scale,
+            per_act_token_quant,
+            block_shape,
+            is_sf_swizzled_layout=is_scale_swizzled,
+            mx_alignment=mx_alignment,
+        )
+    elif quant_dtype == "mxfp6_e3m2":
+        if not quantization_emulation:
+            raise NotImplementedError(
+                "moe_kernel_quantize_input should not be used for native "
+                " quant_dtype='mxfp6_e3m2'MOE. Please open an issue."
+            )
+
+        return _mxfp6_e3m2_quantize(A, A_scale, per_act_token_quant, block_shape)
+    elif quant_dtype == "mxfp6_e2m3":
+        if not quantization_emulation:
+            raise NotImplementedError(
+                "moe_kernel_quantize_input should not be used for native"
+                " quant_dtype='mxfp6_e2m3' MOE. Please open an issue."
+            )
+
+        return _mxfp6_e2m3_quantize(A, A_scale, per_act_token_quant, block_shape)
+    else:
+        return A, A_scale
+
+
+def normalize_scales_shape(scales: torch.Tensor | None) -> torch.Tensor | None:
+    if scales is not None:
+        if scales.numel() == 1:
+            scales = scales.view(1, 1)
+        else:
+            scales = scales.view(-1, scales.size(-1))
+    return scales
+
+
+def normalize_batched_scales_shape(
+    scales: torch.Tensor | None,
+    num_experts: int,
+) -> torch.Tensor | None:
+    if scales is not None and scales.ndim < 3:
+        if scales.numel() == 1:
+            scales = scales.view(1)
+            scales = torch.repeat_interleave(scales, num_experts, dim=0).view(
+                num_experts, 1, 1
+            )
+        else:
+            scales = scales.view(num_experts, -1, scales.size(-1))
+
+    return scales
+
+
+@triton.jit
+def _pack_topk_ids_weights_kernel(
+    topk_ids_ptr,
+    topk_weights_ptr,
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+    USE_GDC: tl.constexpr,
+    launch_pdl: tl.constexpr,  # triton metadata
+):
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+        tl.extra.cuda.gdc_wait()
+    expert_id = tl.load(topk_ids_ptr + offsets, mask=mask, other=0).to(tl.int32)
+    expert_id_shifted = expert_id << 16
+
+    weight = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0)
+    weight_bf16 = weight.to(tl.bfloat16)
+    weight_int16 = weight_bf16.to(tl.int16, bitcast=True)
+
+    weight_int32 = weight_int16.to(tl.int32) & 0xFFFF
+
+    packed = expert_id_shifted | weight_int32
+    tl.store(output_ptr + offsets, packed, mask=mask)
+
+
+def trtllm_moe_pack_topk_ids_weights(
+    topk_ids: torch.Tensor,
+    topk_weights: torch.Tensor,
+    block_size: int = 1024,
+) -> torch.Tensor:
+    assert topk_ids.shape == topk_weights.shape
+    assert topk_ids.is_contiguous() and topk_weights.is_contiguous()
+
+    original_shape = topk_ids.shape
+    ids_flat = topk_ids.reshape(-1)
+    weights_flat = topk_weights.reshape(-1)
+
+    n_elements = ids_flat.numel()
+    output = torch.empty(n_elements, dtype=torch.int32, device=topk_ids.device)
+
+    use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90)
+    grid = (triton.cdiv(n_elements, block_size),)
+    _pack_topk_ids_weights_kernel[grid](
+        ids_flat,
+        weights_flat,
+        output,
+        n_elements,
+        BLOCK_SIZE=block_size,
+        USE_GDC=use_gdc,
+        launch_pdl=use_gdc,
+    )
+    return output.reshape(original_shape)
+
+
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+def swiglu_limit_func(
+    output: torch.Tensor,
+    input: torch.Tensor,  # first half is gate, second half is up
+    swiglu_limit: float = 0.0,
+) -> None:
+    d = input.shape[1] // 2
+    gate = input[:, :d]
+    up = input[:, d:]
+
+    if swiglu_limit > 0:
+        gate = torch.clamp(gate, max=swiglu_limit)
+        up = torch.clamp(up, min=-swiglu_limit, max=swiglu_limit)
+
+    output.copy_(F.silu(gate) * up)

ex_engine/xllm_kernels/ilu/activation.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/attention.cpp

@@ -1,5 +1,5 @@
 
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/fused_moe.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/group_gemm.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2026 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/ilu_ops_api.h

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/ixformer.h

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/matmul.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/norm.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/rope.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_kernels/ilu/utils.h

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2025 The xLLM Authors. 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.

ex_engine/xllm_layers/common/activation.cpp

@@ -0,0 +1,38 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "activation.h"
+
+#include "kernels/ops_api.h"
+namespace xllm {
+namespace layer {
+
+ActivationImpl::ActivationImpl(const std::string& act_mode, bool is_gated)
+    : act_mode_(act_mode), is_gated_(is_gated) {}
+
+void ActivationImpl::forward(torch::Tensor& input, torch::Tensor& output) {
+  xllm::kernel::ActivationParams activation_params;
+  activation_params.input = input;
+  activation_params.output = output;
+  activation_params.act_mode = act_mode_;
+  activation_params.is_gated = is_gated_;
+  xllm::kernel::active(activation_params);
+  // Unified assignment: NPU returns new tensor, others modify in-place (no-op
+  // assignment)
+  output = activation_params.output;
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/activation.h

@@ -0,0 +1,38 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include <string>
+
+namespace xllm {
+namespace layer {
+
+class ActivationImpl : public torch::nn::Module {
+ public:
+  ActivationImpl(const std::string& act_mode, bool is_gated);
+
+  void forward(torch::Tensor& input, torch::Tensor& output);
+
+ private:
+  std::string act_mode_;
+  bool is_gated_;
+};
+TORCH_MODULE(Activation);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/dense_mlp.cpp

@@ -0,0 +1,141 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "dense_mlp.h"
+
+#include <glog/logging.h>
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+DenseMLPImpl::DenseMLPImpl(int64_t hidden_size,
+                           int64_t intermediate_size,
+                           bool is_gated,
+                           bool has_bias,
+                           const std::string& hidden_act,
+                           bool enable_result_reduction,
+                           const QuantArgs& quant_args,
+                           ProcessGroup* process_group,
+                           const torch::TensorOptions& options,
+                           const std::string& module_prefix)
+    : is_gated_(is_gated),
+      intermediate_size_(intermediate_size),
+      process_group_(process_group),
+      hidden_act_(hidden_act) {
+  // Check if using w8a8 smoothquant quantization
+  is_smoothquant_ = quant_args.quant_method() == kQuantMethodSmoothquant;
+
+  if (is_smoothquant_) {
+    // Safety check: only w8a8 smoothquant is supported
+    if (quant_args.bits() != 8 || !quant_args.activation_dynamic()) {
+      LOG(FATAL)
+          << "DenseMLP w8a8 mode only supports w8a8 smoothquant quantization. "
+          << "Got bits=" << quant_args.bits()
+          << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+  }
+
+  // Determine extra args based on quantization mode
+  LinearExtraArgs gate_up_proj_extra_args("none", false);
+  LinearExtraArgs down_proj_extra_args("none", false);
+  if (is_smoothquant_) {
+    // For per-token smoothquant, use specific args
+    down_proj_extra_args = LinearExtraArgs(hidden_act_, is_gated_);
+  }
+
+  // 1. gate + up
+  int64_t out_feature = is_gated_ ? intermediate_size_ * 2 : intermediate_size_;
+  gate_up_proj_ =
+      register_module("gate_up_proj",
+                      ColumnParallelLinear(hidden_size,
+                                           out_feature,
+                                           /*bias=*/has_bias,
+                                           /*gather_output=*/false,
+                                           quant_args,
+                                           process_group_,
+                                           options,
+                                           gate_up_proj_extra_args));
+
+  act_ = register_module("act", Activation(hidden_act_, is_gated_));
+
+  // 2. down
+  const auto down_proj_quant_args =
+      module_prefix.empty()
+          ? quant_args
+          : quant_args.for_module(module_prefix + ".down_proj");
+  down_proj_ = register_module("down_proj",
+                               RowParallelLinear(intermediate_size_,
+                                                 hidden_size,
+                                                 /*bias=*/has_bias,
+                                                 /*input_is_parallelized=*/true,
+                                                 enable_result_reduction,
+                                                 down_proj_quant_args,
+                                                 process_group_,
+                                                 options,
+                                                 down_proj_extra_args));
+}
+
+torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
+  // input shape: [num_tokens, hidden_size]
+  auto gate_up = gate_up_proj_->forward(hidden_states);
+
+  if (is_smoothquant_) {
+    // For w8a8 quantization, the active operation is fused with the down_proj
+    return down_proj_->forward(gate_up);
+  } else {
+    torch::Tensor output;
+    if (Device::type_str() != "npu") {
+      int64_t batch_size = gate_up.sizes()[0];
+      output = torch::empty(
+          {batch_size, intermediate_size_ / process_group_->world_size()},
+          gate_up.options());
+    }
+
+    act_->forward(gate_up, output);
+    return down_proj_->forward(output);
+  }
+}
+
+void DenseMLPImpl::load_state_dict(const StateDict& state_dict) {
+  gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."});
+  down_proj_->load_state_dict(state_dict.get_dict_with_prefix("down_proj."));
+}
+
+void DenseMLPImpl::load_state_dict(const StateDict& state_dict,
+                                   const std::vector<std::string>& gate_up_name,
+                                   const std::string& down_name) {
+  if (is_gated_) {
+    CHECK_EQ(gate_up_name.size(), 2);
+    gate_up_proj_->load_state_dict(state_dict, gate_up_name);
+  } else {
+    CHECK_EQ(gate_up_name.size(), 1);
+    gate_up_proj_->load_state_dict(
+        state_dict.get_dict_with_prefix(gate_up_name[0]));
+  }
+  down_proj_->load_state_dict(state_dict.get_dict_with_prefix(down_name));
+}
+
+std::optional<torch::Tensor> DenseMLPImpl::get_fp8_input_scale() const {
+  if (gate_up_proj_) {
+    return gate_up_proj_->get_input_scale();
+  }
+  return std::nullopt;
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/dense_mlp.h

@@ -0,0 +1,67 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include "activation.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "linear.h"
+
+namespace xllm {
+namespace layer {
+
+class DenseMLPImpl : public torch::nn::Module {
+ public:
+  DenseMLPImpl() = default;
+  DenseMLPImpl(int64_t hidden_size,
+               int64_t intermediate_size,
+               bool is_gated,
+               bool has_bias,
+               const std::string& hidden_act,
+               bool enable_result_reduction,
+               const QuantArgs& quant_args,
+               ProcessGroup* process_group,
+               const torch::TensorOptions& options,
+               const std::string& module_prefix = "");
+
+  torch::Tensor forward(const torch::Tensor& hidden_states);
+
+  void load_state_dict(const StateDict& state_dict);
+  void load_state_dict(const StateDict& state_dict,
+                       const std::vector<std::string>& gate_up_name,
+                       const std::string& down_name);
+
+  // Get FP8 input scale from gate_up_proj for fused RMSNorm+FP8 quantization
+  std::optional<torch::Tensor> get_fp8_input_scale() const;
+
+ private:
+  bool is_gated_;
+  int64_t intermediate_size_;
+  ProcessGroup* process_group_;
+  ColumnParallelLinear gate_up_proj_{nullptr};
+  RowParallelLinear down_proj_{nullptr};
+  Activation act_{nullptr};
+  bool is_smoothquant_;
+  std::string hidden_act_;
+};
+TORCH_MODULE(DenseMLP);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/fused_moe.cpp

@@ -0,0 +1,58 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "fused_moe.h"
+
+#include <glog/logging.h>
+
+namespace xllm {
+namespace layer {
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& /*model_args*/,
+                           const FusedMoEArgs& /*moe_args*/,
+                           const QuantArgs& /*quant_args*/,
+                           const ParallelArgs& /*parallel_args*/,
+                           const torch::TensorOptions& /*options*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+}
+
+torch::Tensor FusedMoEImpl::forward_experts(
+    const torch::Tensor& /*hidden_states*/,
+    const torch::Tensor& /*router_logits*/,
+    bool /*enable_all2all_communication*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+  return torch::Tensor();
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& /*hidden_states*/,
+                                    const ModelInputParams& /*input_params*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+  return torch::Tensor();
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& /*state_dict*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/fused_moe.h

@@ -0,0 +1,54 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include "dense_mlp.h"
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "fused_moe_base.h"
+#include "linear.h"
+
+namespace xllm {
+namespace layer {
+
+// FusedMoE common implementation - placeholder for unsupported backends
+// Actual implementations are in backend-specific fused_moe.h files.
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_experts(const torch::Tensor& hidden_states,
+                                const torch::Tensor& router_logits,
+                                bool enable_all2all_communication);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rms_norm.cpp

@@ -0,0 +1,144 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "rms_norm.h"
+
+#include <glog/logging.h>
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+const static std::string kLayerNormMode = "layernorm";
+const static std::string kRmsNormMode = "rmsnorm";
+
+RMSNormImpl::RMSNormImpl(int64_t dim,
+                         double eps,
+                         const torch::TensorOptions& options)
+    : norm_dim_(dim), eps_(eps), mode_(kRmsNormMode) {
+  weight_ = register_parameter("weight",
+                               torch::empty({dim}, options),
+                               /*requires_grad=*/false);
+}
+
+RMSNormImpl::RMSNormImpl(const ModelContext& context)
+    : RMSNormImpl(context.get_model_args().hidden_size(),
+                  context.get_model_args().rms_norm_eps(),
+                  context.get_tensor_options()) {}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> RMSNormImpl::forward(
+    torch::Tensor& input,
+    std::optional<torch::Tensor> residual,
+    std::optional<torch::Tensor> inplace_output) {
+  auto org_shape = input.sizes().vec();
+  input = input.reshape({-1, norm_dim_});
+
+  torch::Tensor output;
+  if (Device::type_str() != "npu") {
+    if (inplace_output.has_value()) {
+      output = inplace_output.value();
+      output = output.reshape({-1, norm_dim_});
+    } else {
+      output = torch::empty_like(input);
+    }
+  }
+
+  std::optional<torch::Tensor> residual_out;
+  if (residual.has_value()) {
+    residual.value() = residual.value().reshape({-1, norm_dim_});
+    if (Device::type_str() == "mlu" || Device::type_str() == "ilu") {
+      residual_out = residual.value();
+    }
+  }
+
+  xllm::kernel::FusedLayerNormParams fused_layernorm_params;
+  fused_layernorm_params.input = input;
+  fused_layernorm_params.residual = residual;
+  fused_layernorm_params.output = output;
+  fused_layernorm_params.residual_out = residual_out;
+  fused_layernorm_params.weight = weight_;
+  fused_layernorm_params.eps = eps_;
+  fused_layernorm_params.mode = mode_;
+  fused_layernorm_params.store_output_before_norm = residual_out.has_value();
+  if (bias_.defined()) {
+    fused_layernorm_params.beta = bias_;
+  }
+
+  xllm::kernel::fused_layernorm(fused_layernorm_params);
+
+  output = fused_layernorm_params.output;
+  residual_out = fused_layernorm_params.residual_out;
+
+  output = output.view(org_shape);
+  if (residual_out.has_value()) {
+    residual_out.value() = residual_out.value().view(org_shape);
+  }
+  return std::make_tuple(output, residual_out);
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
+RMSNormImpl::forward_fp8(torch::Tensor& input,
+                         const torch::Tensor& fp8_scale,
+                         std::optional<torch::Tensor> residual) {
+  // Only supported on CUDA for now
+  CHECK(Device::type_str() == "cuda")
+      << "forward_fp8 is only supported on CUDA";
+  CHECK(mode_ == kRmsNormMode)
+      << "forward_fp8 only supports RMSNorm mode, not LayerNorm";
+
+  if (residual.has_value()) {
+    // Fused Add + RMSNorm + FP8 Quantization
+    xllm::kernel::FusedAddRmsNormStaticFp8QuantParams params;
+    params.input = input;
+    params.residual = residual.value();
+    params.weight = weight_;
+    params.scale = fp8_scale;
+    params.epsilon = eps_;
+
+    auto [output, updated_residual] =
+        xllm::kernel::fused_add_rms_norm_static_fp8_quant(params);
+
+    return std::make_tuple(output, updated_residual);
+  } else {
+    // RMSNorm + FP8 Quantization (no residual)
+    xllm::kernel::RmsNormStaticFp8QuantParams params;
+    params.input = input;
+    params.weight = weight_;
+    params.scale = fp8_scale;
+    params.epsilon = eps_;
+
+    auto output = xllm::kernel::rms_norm_static_fp8_quant(params);
+
+    return std::make_tuple(output, std::nullopt);
+  }
+}
+
+void RMSNormImpl::load_state_dict(const StateDict& state_dict) {
+  LOAD_WEIGHT(weight);
+  if (bias_.defined()) {
+    LOAD_WEIGHT(bias);
+  }
+}
+
+void RMSNormImpl::set_layernorm_mode() {
+  mode_ = kLayerNormMode;
+  bias_ = register_parameter(
+      "bias", torch::empty({norm_dim_}, weight_.options()), false);
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rms_norm.h

@@ -0,0 +1,64 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include "core/framework/model_context.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+
+namespace xllm {
+namespace layer {
+
+class RMSNormImpl : public torch::nn::Module {
+ public:
+  RMSNormImpl(int64_t dim, double eps, const torch::TensorOptions& options);
+  RMSNormImpl(const ModelContext& context);
+
+  // Standard forward: returns (normalized_output, updated_residual)
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      torch::Tensor& input,
+      std::optional<torch::Tensor> residual = std::nullopt,
+      std::optional<torch::Tensor> inplace_output = std::nullopt);
+
+  // Fused forward with FP8 quantization output (for static quantization)
+  // Returns: (fp8_quantized_output, updated_residual)
+  // This combines RMSNorm + FP8 quantization to reduce memory bandwidth
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward_fp8(
+      torch::Tensor& input,
+      const torch::Tensor& fp8_scale,
+      std::optional<torch::Tensor> residual = std::nullopt);
+
+  void set_layernorm_mode();
+
+  void load_state_dict(const StateDict& state_dict);
+
+  torch::Tensor weight() const { return weight_; }
+  torch::Tensor bias() const { return bias_; }
+  double eps() const { return eps_; }
+
+ private:
+  DEFINE_WEIGHT(weight);
+  DEFINE_WEIGHT(bias);
+  int64_t norm_dim_;
+  double eps_;
+  std::string mode_;
+};
+TORCH_MODULE(RMSNorm);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rotary_embedding.cpp

@@ -0,0 +1,307 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "rotary_embedding.h"
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+RotaryEmbeddingImpl::RotaryEmbeddingImpl(const ModelContext& context) {
+  LOG(FATAL) << "Not implement currently.";
+}
+
+RotaryEmbeddingImpl::RotaryEmbeddingImpl(int64_t rotary_dim,
+                                         int64_t max_position_embeddings,
+                                         int64_t rope_theta,
+                                         bool interleaved,
+                                         const torch::TensorOptions& options)
+    : interleaved_(interleaved) {
+  auto inv_freq = rotary::compute_inv_freq(rotary_dim, rope_theta, options);
+  const auto cos_sin = rotary::compute_cos_sin_cache(
+      rotary_dim, max_position_embeddings, interleaved, inv_freq, options);
+  cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin);
+
+  auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1);
+  cos_ = cos_sin_vec[0].view({-1, rotary_dim});
+  sin_ = cos_sin_vec[1].view({-1, rotary_dim});
+
+  // Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels.
+  const auto dev = Device::type_str();
+  if (dev == "cuda" || dev == "ilu" || dev == "musa") {
+    auto chunks = cos_sin_cache_.chunk(4, -1);
+    precomputed_cos_sin_cache_ =
+        torch::cat({chunks[0], chunks[2]}, -1).contiguous();
+  }
+}
+
+void RotaryEmbeddingImpl::forward(torch::Tensor& q,
+                                  torch::Tensor& k,
+                                  const torch::Tensor& positions,
+                                  const torch::Tensor& cu_query_lens,
+                                  int64_t max_query_len,
+                                  bool is_prompt) {
+  bool discrete;
+  std::optional<torch::Tensor> position_ids;
+  if (is_prompt) {
+    discrete = false;
+    if (Device::type_str() == "cuda" || Device::type_str() == "npu" ||
+        Device::type_str() == "ilu" || Device::type_str() == "musa") {
+      position_ids = positions;
+    }
+  } else {
+    discrete = true;
+    position_ids = positions;
+  }
+
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = q;
+  rotary_params.k = k;
+  rotary_params.sin = sin_;
+  rotary_params.cos = cos_;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
+  rotary_params.position_ids = position_ids;
+  rotary_params.cu_query_lens = cu_query_lens;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = discrete;
+  rotary_params.max_query_len = max_query_len;
+  xllm::kernel::apply_rotary(rotary_params);
+
+  q = rotary_params.q;
+  k = rotary_params.k;
+}
+
+// Single tensor forward for MLA architecture
+void RotaryEmbeddingImpl::forward(torch::Tensor& input,
+                                  const torch::Tensor& positions,
+                                  const torch::Tensor& cu_query_lens,
+                                  int64_t max_query_len,
+                                  bool is_prompt) {
+  bool discrete;
+  std::optional<torch::Tensor> position_ids;
+  if (is_prompt) {
+    discrete = false;
+    if (Device::type_str() == "cuda" || Device::type_str() == "npu" ||
+        Device::type_str() == "ilu") {
+      position_ids = positions;
+    }
+  } else {
+    discrete = true;
+    position_ids = positions;
+  }
+
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = input;
+  rotary_params.sin = sin_;
+  rotary_params.cos = cos_;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.position_ids = position_ids;
+  rotary_params.cu_query_lens = cu_query_lens;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = discrete;
+  rotary_params.max_query_len = max_query_len;
+  xllm::kernel::apply_rotary(rotary_params);
+
+  input = rotary_params.q;
+}
+
+MRotaryEmbeddingImpl::MRotaryEmbeddingImpl(
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    int64_t rope_theta,
+    bool interleaved,
+    const std::vector<int64_t>& rope_scaling_mrope_section,
+    const torch::TensorOptions& options)
+    : RotaryEmbeddingImpl(rotary_dim,
+                          max_position_embeddings,
+                          rope_theta,
+                          interleaved,
+                          options),
+      mrope_section_(rope_scaling_mrope_section) {
+  mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device());
+}
+
+void MRotaryEmbeddingImpl::forward(torch::Tensor& q,
+                                   torch::Tensor& k,
+                                   const torch::Tensor& positions,
+                                   const AttentionMetadata& attn_metadata) {
+  bool only_prefill =
+      (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
+  if (!only_prefill || mrope_section_.empty()) {
+    torch::Tensor position_ids = positions;
+    if (positions.dim() == 2) {
+      position_ids = positions[0];
+    }
+    return RotaryEmbeddingImpl::forward(q,
+                                        k,
+                                        position_ids,
+                                        attn_metadata.q_cu_seq_lens,
+                                        attn_metadata.max_query_len,
+                                        attn_metadata.is_prefill);
+  }
+
+  int64_t num_tokens = positions.size(-1);
+  mrope_cu_seq_lens_[1] = num_tokens;
+  CHECK(attn_metadata.mrope_cos.defined() && attn_metadata.mrope_sin.defined());
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = q;
+  rotary_params.k = k;
+  rotary_params.sin = attn_metadata.mrope_sin;
+  rotary_params.cos = attn_metadata.mrope_cos;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
+  rotary_params.position_ids = std::nullopt;
+  rotary_params.cu_query_lens = mrope_cu_seq_lens_;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = false;
+  rotary_params.max_query_len = num_tokens;
+  xllm::kernel::apply_rotary(rotary_params);
+
+  q = rotary_params.q;
+  k = rotary_params.k;
+}
+
+DeepseekScalingRotaryEmbeddingImpl::DeepseekScalingRotaryEmbeddingImpl(
+    int64_t head_size,
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    int64_t rope_scaling_original_max_position_embeddings,
+    int64_t rope_theta,
+    bool interleaved,
+    float scaling_factor,
+    float extrapolation_factor,
+    float attn_factor,
+    float beta_fast,
+    float beta_slow,
+    float mscale,
+    float mscale_all_dim,
+    const torch::TensorOptions& options)
+    : head_size_(head_size),
+      rotary_dim_(rotary_dim),
+      interleaved_(interleaved) {
+  auto inv_freq = rotary::apply_deepseek_yarn_rope_scaling(
+      scaling_factor,
+      extrapolation_factor,
+      beta_fast,
+      beta_slow,
+      rotary_dim,
+      rope_theta,
+      rope_scaling_original_max_position_embeddings);
+  const auto cos_sin = rotary::compute_cos_sin_cache(rotary_dim,
+                                                     max_position_embeddings,
+                                                     interleaved,
+                                                     scaling_factor,
+                                                     attn_factor,
+                                                     mscale,
+                                                     mscale_all_dim,
+                                                     inv_freq,
+                                                     options);
+  cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin);
+
+  auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1);
+  cos_ = cos_sin_vec[0].view({-1, rotary_dim});
+  sin_ = cos_sin_vec[1].view({-1, rotary_dim});
+
+  // Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels.
+  const auto dev = Device::type_str();
+  if (dev == "cuda" || dev == "ilu" || dev == "musa") {
+    auto chunks = cos_sin_cache_.chunk(4, -1);
+    precomputed_cos_sin_cache_ =
+        torch::cat({chunks[0], chunks[2]}, -1).contiguous();
+  }
+}
+
+void DeepseekScalingRotaryEmbeddingImpl::forward(
+    torch::Tensor& input,
+    const torch::Tensor& positions,
+    const torch::Tensor& cu_query_lens,
+    int64_t max_query_len,
+    bool is_prompt) {
+  const int32_t dim = -1;
+  bool discrete;
+  std::optional<torch::Tensor> position_ids;
+  if (is_prompt) {
+    discrete = false;
+    position_ids = std::nullopt;
+  } else {
+    discrete = true;
+    position_ids = positions;
+    max_query_len = 1;
+  }
+  auto input_rot = input.slice(dim, 0, rotary_dim_);
+  torch::Tensor input_pass;
+  if (rotary_dim_ < head_size_) {
+    input_pass = input.slice(dim, rotary_dim_, head_size_);
+  }
+
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = input_rot;
+  rotary_params.sin = sin_;
+  rotary_params.cos = cos_;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
+  rotary_params.position_ids = position_ids;
+  rotary_params.cu_query_lens = cu_query_lens;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = discrete;
+  rotary_params.max_query_len = max_query_len;
+  xllm::kernel::apply_rotary(rotary_params);
+  input_rot = rotary_params.q;
+
+  if (rotary_dim_ < head_size_) {
+    input = torch::cat({input_rot, input_pass}, dim);
+  } else {
+    input = input_rot;
+  }
+}
+
+// Factory function: creates the appropriate RoPE type based on model args
+std::shared_ptr<RotaryEmbeddingBase> create_mla_rotary_embedding(
+    const ModelArgs& args,
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    bool interleaved,
+    const torch::TensorOptions& options) {
+  if (args.rope_scaling_rope_type() == "deepseek_yarn") {
+    return std::make_shared<DeepseekScalingRotaryEmbeddingImpl>(
+        rotary_dim,  // head_size (same as rotary_dim for MLA)
+        rotary_dim,
+        max_position_embeddings,
+        args.rope_scaling_original_max_position_embeddings(),
+        args.rope_theta(),
+        interleaved,
+        args.rope_scaling_factor(),
+        args.rope_extrapolation_factor(),
+        args.rope_scaling_attn_factor(),
+        args.rope_scaling_beta_fast(),
+        args.rope_scaling_beta_slow(),
+        args.rope_scaling_mscale(),
+        args.rope_scaling_mscale_all_dim(),
+        options);
+  } else {
+    // default rope type
+    return std::make_shared<RotaryEmbeddingImpl>(rotary_dim,
+                                                 max_position_embeddings,
+                                                 args.rope_theta(),
+                                                 interleaved,
+                                                 options);
+  }
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rotary_embedding.h

@@ -0,0 +1,158 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+#include <torch/types.h>
+
+#include <memory>
+
+#include "attention_metadata.h"
+#include "core/framework/model_context.h"
+#include "framework/model/model_args.h"
+#include "rotary_embedding_util.h"
+
+namespace xllm {
+namespace layer {
+
+class RotaryEmbeddingBase : public torch::nn::Module {
+ public:
+  ~RotaryEmbeddingBase() override = default;
+
+  virtual void forward(torch::Tensor& input,
+                       const torch::Tensor& positions,
+                       const torch::Tensor& cu_query_lens,
+                       int64_t max_query_len,
+                       bool is_prompt) = 0;
+  virtual const torch::Tensor& get_sin_cache() const = 0;
+  virtual const torch::Tensor& get_cos_cache() const = 0;
+  virtual const bool get_interleaved() const = 0;
+};
+
+class RotaryEmbeddingImpl : public RotaryEmbeddingBase {
+ public:
+  RotaryEmbeddingImpl(int64_t rotary_dim,
+                      int64_t max_position_embeddings,
+                      int64_t rope_theta,
+                      bool interleaved,
+                      const torch::TensorOptions& options);
+  RotaryEmbeddingImpl(const ModelContext& context);
+
+  void forward(torch::Tensor& q,
+               torch::Tensor& k,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt);
+  // Single tensor forward for MLA architecture
+  void forward(torch::Tensor& input,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt) override;
+
+  const torch::Tensor& precomputed_cos_sin_cache() {
+    return precomputed_cos_sin_cache_;
+  }
+
+  torch::Tensor get_cos_sin_cache() { return cos_sin_cache_; }
+  const torch::Tensor& get_sin_cache() const override { return sin_; }
+  const torch::Tensor& get_cos_cache() const override { return cos_; }
+  const bool get_interleaved() const override { return interleaved_; }
+
+ protected:
+  bool interleaved_;
+  torch::Tensor cos_sin_cache_;
+  // Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels.
+  // Avoids chunk/cat operations on every forward call.
+  torch::Tensor precomputed_cos_sin_cache_;
+
+ private:
+  torch::Tensor sin_;
+  torch::Tensor cos_;
+};
+TORCH_MODULE(RotaryEmbedding);
+
+class MRotaryEmbeddingImpl : public RotaryEmbeddingImpl {
+ public:
+  MRotaryEmbeddingImpl(int64_t rotary_dim,
+                       int64_t max_position_embeddings,
+                       int64_t rope_theta,
+                       bool interleaved,
+                       const std::vector<int64_t>& rope_scaling_mrope_section,
+                       const torch::TensorOptions& options);
+
+  void forward(torch::Tensor& q,
+               torch::Tensor& k,
+               const torch::Tensor& positions,
+               const AttentionMetadata& attn_metadata);
+
+ private:
+  std::vector<int64_t> mrope_section_;
+  torch::Tensor mrope_cu_seq_lens_;
+};
+TORCH_MODULE(MRotaryEmbedding);
+
+class DeepseekScalingRotaryEmbeddingImpl : public RotaryEmbeddingBase {
+ public:
+  DeepseekScalingRotaryEmbeddingImpl(
+      int64_t head_size,
+      int64_t rotary_dim,
+      int64_t max_position_embeddings,
+      int64_t rope_scaling_original_max_position_embeddings,
+      int64_t rope_theta,
+      bool interleaved,
+      float scaling_factor,
+      float extrapolation_factor,
+      float attn_factor,
+      float beta_fast,
+      float beta_slow,
+      float mscale,
+      float mscale_all_dim,
+      const torch::TensorOptions& options);
+
+  void forward(torch::Tensor& input,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt) override;
+  const torch::Tensor& get_sin_cache() const override { return sin_; }
+  const torch::Tensor& get_cos_cache() const override { return cos_; }
+  const bool get_interleaved() const override { return interleaved_; }
+
+ private:
+  int64_t head_size_;
+  int64_t rotary_dim_;
+  bool interleaved_;
+  torch::Tensor sin_;
+  torch::Tensor cos_;
+  torch::Tensor cos_sin_cache_;
+  // Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels.
+  // Avoids chunk/cat operations on every forward call.
+  torch::Tensor precomputed_cos_sin_cache_;
+};
+TORCH_MODULE(DeepseekScalingRotaryEmbedding);
+
+// Factory function: creates the appropriate RoPE type based on model args
+std::shared_ptr<RotaryEmbeddingBase> create_mla_rotary_embedding(
+    const ModelArgs& args,
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    bool interleaved,
+    const torch::TensorOptions& options);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/attention.cpp

@@ -0,0 +1,189 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "attention.h"
+
+#include "kernels/ilu/ilu_ops_api.h"
+#include "kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             float scale,
+                             int64_t num_kv_heads,
+                             int64_t sliding_window)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(head_size),
+      use_fused_mla_qkv_(false),
+      enable_lighting_indexer_(false),
+      enable_mla_(false),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             int64_t num_kv_heads,
+                             int64_t v_head_dim,
+                             int64_t sliding_window,
+                             float scale,
+                             bool use_fused_mla_qkv,
+                             bool enable_lighting_indexer,
+                             bool enable_mla)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(v_head_dim),
+      use_fused_mla_qkv_(use_fused_mla_qkv),
+      enable_lighting_indexer_(enable_lighting_indexer),
+      enable_mla_(enable_mla),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
+    const AttentionMetadata& attn_metadata,
+    torch::Tensor& query,
+    torch::Tensor& key,
+    torch::Tensor& value,
+    KVCache& kv_cache) {
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  torch::Tensor output;
+  if (enable_mla_) {
+    output = torch::empty({query.size(0), num_heads_ * v_head_dim_},
+                          query.options());
+  } else {
+    output = torch::empty_like(query);
+  }
+  if (attn_metadata.is_dummy) {
+    return std::make_tuple(output, output_lse);
+  }
+
+  bool only_prefill =
+      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
+  int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_;
+  torch::Tensor k_cache = kv_cache.get_k_cache();
+  std::optional<torch::Tensor> v_cache;
+  std::optional<torch::Tensor> v;
+  if (!enable_mla_) {
+    v = value.view({-1, num_kv_heads, head_size_});
+    v_cache = kv_cache.get_v_cache();
+  }
+
+  bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_);
+  if (!skip_process_cache) {
+    xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
+    reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_});
+    reshape_paged_cache_params.value = v;
+    reshape_paged_cache_params.k_cache = k_cache;
+    reshape_paged_cache_params.v_cache = v_cache;
+    reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping;
+    xllm::kernel::reshape_paged_cache(reshape_paged_cache_params);
+  }
+
+  if (enable_lighting_indexer_ || !only_prefill) {
+    decoder_forward(query, output, k_cache, v_cache, attn_metadata);
+  } else {
+    prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
+  }
+
+  int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_;
+  output = output.view({-1, num_heads_ * head_size});
+  return {output, output_lse};
+}
+
+void AttentionImpl::prefill_forward(torch::Tensor& query,
+                                    torch::Tensor& key,
+                                    torch::Tensor& value,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  query = query.view({-1, num_heads_, head_size_});
+  output = output.view({-1, num_heads_, head_size_v});
+  // torch::Tensor k_cache_ = k_cache;
+  // torch::Tensor v_cache_ = v_cache.value();
+  xllm::kernel::ilu::batch_prefill(query,
+                                   k_cache,
+                                   v_cache,
+                                   output,
+                                   output_lse,
+                                   attn_metadata.q_cu_seq_lens,
+                                   attn_metadata.kv_cu_seq_lens,
+                                   /*alibi_slope=*/std::nullopt,
+                                   /*attn_bias=*/std::nullopt,
+                                   /*q_quant_scale=*/std::nullopt,
+                                   /*k_quant_scale=*/std::nullopt,
+                                   /*v_quant_scale=*/std::nullopt,
+                                   attn_metadata.block_table,
+                                   attn_metadata.max_query_len,
+                                   attn_metadata.max_seq_len,
+                                   scale_,
+                                   attn_metadata.is_causal,
+                                   sliding_window_,
+                                   /*window_size_right=*/-1,
+                                   attn_metadata.compute_dtype,
+                                   /*return_lse=*/false);
+}
+
+void AttentionImpl::decoder_forward(torch::Tensor& query,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  query = query.view({-1, 1, num_heads_, head_size_});
+  output = output.view({-1, 1, num_heads_, head_size_v});
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+
+  int64_t block_aligned_max_seq_len =
+      attn_metadata.block_table.size(-1) * k_cache.size(2);
+
+  xllm::kernel::ilu::batch_decode(query,
+                                  k_cache,
+                                  output,
+                                  attn_metadata.block_table,
+                                  attn_metadata.kv_seq_lens,
+                                  v_cache,
+                                  output_lse,
+                                  /*q_quant_scale=*/std::nullopt,
+                                  /*k_quant_scale=*/std::nullopt,
+                                  /*v_quant_scale=*/std::nullopt,
+                                  /*out_quant_scale=*/std::nullopt,
+                                  /*alibi_slope=*/std::nullopt,
+                                  attn_metadata.attn_mask,
+                                  attn_metadata.compute_dtype,
+                                  block_aligned_max_seq_len,
+                                  sliding_window_,
+                                  /*window_size_right=*/-1,
+                                  scale_,
+                                  /*return_lse=*/false,
+                                  attn_metadata.is_causal,
+                                  /*kv_cache_quant_bit_size=*/-1);
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/attention.h

@@ -0,0 +1,82 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include <tuple>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_input_params.h"
+#include "layers/common/attention_metadata.h"
+
+namespace xllm {
+namespace layer {
+class AttentionImpl : public torch::nn::Module {
+ public:
+  AttentionImpl() = default;
+
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                float scale,
+                int64_t num_kv_heads,
+                int64_t sliding_window);
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                int64_t num_kv_heads,
+                int64_t v_head_dim,
+                int64_t sliding_window,
+                float scale,
+                bool use_fused_mla_qkv,
+                bool enable_lighting_indexer,
+                bool enable_mla);
+
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      const AttentionMetadata& attn_metadata,
+      torch::Tensor& query,
+      torch::Tensor& key,
+      torch::Tensor& value,
+      KVCache& kv_cache);
+
+  void prefill_forward(torch::Tensor& query,
+                       torch::Tensor& key,
+                       torch::Tensor& value,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+  void decoder_forward(torch::Tensor& query,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+ private:
+  int64_t num_heads_;
+  int64_t head_size_;
+  float scale_;
+  int64_t num_kv_heads_;
+  int64_t v_head_dim_;
+  bool use_fused_mla_qkv_;
+  bool enable_lighting_indexer_;
+  bool enable_mla_;
+  int64_t sliding_window_;
+};
+TORCH_MODULE(Attention);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/fused_moe.cpp

@@ -0,0 +1,797 @@
+/* Copyright 2026 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#include "fused_moe.h"
+
+#include <glog/logging.h>
+
+#include <iomanip>
+
+#include "common/global_flags.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "kernels/ops_api.h"
+#include "layers/common/dp_utils.h"
+#include "util/utils.h"
+
+namespace {
+
+int32_t get_dtype_size(torch::ScalarType dtype) {
+  return static_cast<int32_t>(torch::elementSize(dtype));
+}
+
+}  // namespace
+
+namespace xllm {
+namespace layer {
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
+                           const FusedMoEArgs& moe_args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options)
+    : num_total_experts_(static_cast<int64_t>(model_args.n_routed_experts())),
+      topk_(model_args.num_experts_per_tok()),
+      num_expert_group_(model_args.n_group()),
+      topk_group_(model_args.topk_group()),
+      route_scale_(model_args.routed_scaling_factor()),
+      hidden_size_(model_args.hidden_size()),
+      n_shared_experts_(model_args.n_shared_experts()),
+      is_gated_(moe_args.is_gated),
+      renormalize_(model_args.norm_topk_prob() ? 1 : 0),
+      hidden_act_(model_args.hidden_act()),
+      scoring_func_(model_args.scoring_func()),
+      quant_args_(quant_args),
+      parallel_args_(parallel_args),
+      options_(options),
+      device_(options.device()) {
+  const int64_t num_experts = num_total_experts_;
+  const int64_t intermediate_size =
+      static_cast<int64_t>(model_args.moe_intermediate_size());
+  const std::string& topk_method = model_args.topk_method();
+  int64_t ep_size = parallel_args.ep_size();
+  int64_t ep_rank = 0;
+  tp_pg_ = parallel_args.tp_group_;
+  if (ep_size > 1) {
+    ep_rank = parallel_args.moe_ep_group_->rank();
+    tp_pg_ = parallel_args.moe_tp_group_;
+  }
+
+  // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
+  // supported
+  if (!quant_args.quant_method().empty()) {
+    if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
+        !quant_args.activation_dynamic()) {
+      LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
+                    "quant_method is set. "
+                 << "Got quant_method=" << quant_args.quant_method()
+                 << ", bits=" << quant_args.bits()
+                 << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+    // If confirmed as smoothquant w8a8, set is_smoothquant_ to true
+    is_smoothquant_ = true;
+  } else {
+    is_smoothquant_ = false;
+  }
+
+  // Deep EP initialization check
+  enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1;
+  if (enable_deep_ep_) {
+    // for now, we only implement the deep ep for decode stage.
+    // so we will assume the max_token_num is limited to max_batch_size * (1+K)
+    // K is the number of speculative tokens.
+    int64_t dispatch_token_size;
+    if (quant_args.quant_method() == "smoothquant") {
+      // float32 is for the scale of the quantized input
+      dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) +
+                            get_dtype_size(torch::kFloat32);
+    } else {
+      dispatch_token_size =
+          hidden_size_ * get_dtype_size(options_.dtype().toScalarType());
+    }
+    torch::ScalarType combine_dtype = options_.dtype().toScalarType();
+    int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype);
+    // Ensure calculation base is at least ep_size
+    int64_t effective_seqs =
+        std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size);
+    // NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size,
+    // regardless of the dp size. To ensure robust scheduling and account
+    // for the worst-case scenario, we must guarantee that each rank is capable
+    // of handling the maximum possible number of tokens. Therefore, we define
+    // max_num_tokens_per_rank as the full maximum value, without dividing by
+    // either the rank count or the dp size.
+    int64_t max_num_tokens_per_rank =
+        (1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_;
+
+    // make sure that all layers share the same deep ep instance
+    //  so that the memory footprint is minimized
+    deep_ep_ = DeepEPManager::get_instance(dispatch_token_size,
+                                           combine_token_size,
+                                           max_num_tokens_per_rank,
+                                           num_experts,
+                                           parallel_args,
+                                           options_);
+
+    // obtain the buffer and parameters of deep ep
+    deep_ep_buffer_ = deep_ep_->get_buffer();
+    deep_ep_params_ = deep_ep_->get_params();
+
+    // intermediate buffer that can be initialized once
+    // we place these tensor here in order to speed up forward pass
+    int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv;
+    int64_t token_bytes = is_smoothquant_
+                              ? get_dtype_size(torch::kInt8)
+                              : get_dtype_size(options_.dtype().toScalarType());
+    token_bytes = token_bytes * hidden_size_;
+    int64_t head_size = n_tokens_recv * token_bytes;
+    dispatch_recv_token_tensor_head_ =
+        deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size)
+            .view({n_tokens_recv, token_bytes});
+    // input scale in smoothquant
+    if (is_smoothquant_) {
+      int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32);
+      dispatch_recv_token_tensor_tail_ =
+          deep_ep_buffer_.combine_send_token_tensor
+              .narrow(0, head_size, tail_size)
+              .view({n_tokens_recv, -1});
+    }
+  }
+
+  // calculate the number of experts per rank
+  num_experts_per_rank_ = num_experts / ep_size;
+  start_expert_id_ = ep_rank * num_experts_per_rank_;
+
+  if (topk_method == "noaux_tc") {
+    e_score_correction_bias_ = register_parameter(
+        "e_score_correction_bias", torch::empty({num_experts}, options), false);
+  }
+
+  gate_ = register_module(
+      "gate_proj",
+      ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
+  if (n_shared_experts_ > 0) {
+    ProcessGroup* shared_expert_pg;
+    if (parallel_args_.ep_size() > 1) {
+      // we use tp=1 for shared experts computation in deep ep mode
+      CHECK(parallel_args_.ep_size() == parallel_args_.world_size())
+          << "Models with shared experts only support ep_size equal to "
+             "world size for now.";
+      shared_expert_pg = parallel_args.moe_tp_group_;
+    } else {
+      shared_expert_pg = parallel_args.process_group_;
+    }
+    // The shared experts computation can proceed in parallel with the
+    // final communication step during the MoE computation, as long as it
+    // remains independent of any communication operations. For optimal
+    // performance, ensure that the shared experts layer on each rank always
+    // maintains its own unique weights.
+    shared_experts_ =
+        register_module("shared_experts",
+                        DenseMLP(hidden_size_,
+                                 intermediate_size * n_shared_experts_,
+                                 is_gated_,
+                                 false,
+                                 hidden_act_,
+                                 /*enable_result_reduction=*/true,
+                                 quant_args,
+                                 shared_expert_pg,
+                                 options));
+  }
+
+  // create weight buffer
+  const int64_t world_size = tp_pg_->world_size();
+  int64_t local_intermediate_size = intermediate_size / world_size;
+  if (is_smoothquant_) {
+    auto quant_option = options_.dtype(torch::kInt8);
+    auto fp_option = options_.dtype(torch::kFloat32);
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            quant_option),
+        false);
+    w13_scale_ = register_parameter(
+        "w13_scale",
+        torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
+                     fp_option),
+        false);
+    // Note: We do not check enable_deep_ep_ here, since smooth quantization
+    // information may be needed even when deep EP mode is disabled. This allows
+    // retrieving quantization parameters for any subset of experts as required.
+    input_smooth_ = register_parameter(
+        "input_smooth",
+        torch::empty({num_total_experts_, hidden_size_}, fp_option),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            quant_option),
+        false);
+    w2_scale_ = register_parameter(
+        "w2_scale",
+        torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
+        false);
+    act_smooth_ = register_parameter(
+        "act_smooth",
+        torch::empty({num_experts_per_rank_, local_intermediate_size},
+                     fp_option),
+        false);
+
+  } else {
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            options_),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            options_),
+        false);
+  }
+}
+
+torch::Tensor FusedMoEImpl::create_group_gemm_output(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& group_list,
+    torch::ScalarType dtype,
+    torch::Tensor& workspace) {
+  // unify shape logic: define the target shape once.
+  bool is_3d_weight = (b.dim() != 2);
+  int64_t num_tokens = a.size(0);
+  int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0);
+
+  std::vector<int64_t> output_shape;
+  int64_t required_elements = num_tokens * out_dim;
+
+  if (is_3d_weight) {
+    output_shape = {num_tokens, out_dim};
+  } else {
+    output_shape = {group_list.size(0), num_tokens, out_dim};
+    required_elements *= group_list.size(0);
+  }
+
+  auto options = a.options().dtype(dtype);
+
+  // non-smoothquant: direct allocation
+  if (!is_smoothquant_) {
+    return torch::empty(output_shape, options);
+  }
+
+  // smoothquant: managed workspace logic
+  if (!workspace.defined()) {
+    // Lazy initialization: allocate max buffer for the lifecycle
+    // Note: accessing class members w13_ and w2_ directly for context
+    int64_t max_width = std::max(w13_.size(1), w2_.size(1));
+    workspace = torch::empty({num_tokens * max_width}, options);
+  }
+
+  // view construction
+  CHECK(workspace.numel() >= required_elements)
+      << "FusedMoE Workspace too small! Alloc: " << workspace.numel()
+      << ", Req: " << required_elements;
+
+  // utilize the pre-calculated output_shape
+  return workspace.slice(0, 0, required_elements).view(output_shape);
+}
+
+torch::Tensor FusedMoEImpl::select_experts(
+    const torch::Tensor& hidden_states_2d,
+    const torch::Tensor& router_logits_2d,
+    SelectedExpertInfo& selected_expert_info,
+    bool enable_all2all_communication) {
+  // prepare the parameters for select_experts
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1: apply softmax topk or sigmoid topk / routing logic
+  torch::Tensor reduce_weight;
+  torch::Tensor expert_id;
+  {
+    xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
+    moe_active_topk_params.input = router_logits_2d;
+    moe_active_topk_params.topk = topk_;
+    moe_active_topk_params.num_expert_group = num_expert_group_;
+    moe_active_topk_params.topk_group = topk_group_;
+    moe_active_topk_params.normalize = renormalize_;
+    moe_active_topk_params.normed_by = "topk_logit";
+    moe_active_topk_params.scoring_func = scoring_func_;
+    moe_active_topk_params.route_scale = route_scale_;
+    moe_active_topk_params.e_score_correction_bias = e_score_correction_bias;
+    std::tie(reduce_weight, expert_id) =
+        xllm::kernel::moe_active_topk(moe_active_topk_params);
+  }
+
+  // Step 2: generate expert ids
+  torch::Tensor gather_idx;
+  torch::Tensor combine_idx;
+  torch::Tensor token_count;
+  std::optional<torch::Tensor> cusum_token_count;
+  {
+    xllm::kernel::MoeGenIdxParams moe_gen_idx_params;
+    moe_gen_idx_params.expert_id = expert_id;
+    moe_gen_idx_params.expert_num = num_total_experts_;
+    std::vector<torch::Tensor> output_vec =
+        xllm::kernel::moe_gen_idx(moe_gen_idx_params);
+    gather_idx = output_vec[0];
+    combine_idx = output_vec[1];
+    token_count = output_vec[2];
+    // during all2all communication, we do not need cusum_token_count in the
+    // following computation
+    if (enable_all2all_communication) {
+      cusum_token_count = std::nullopt;
+    } else {
+      cusum_token_count = output_vec[3];
+    }
+  }
+
+  // Step 3: expand and quantize input if needed
+  torch::Tensor expand_hidden_states;
+  torch::Tensor hidden_states_scale;
+  torch::Tensor token_count_slice;
+  // all2all related variables
+  torch::Tensor dispatch_send_token_tensor;
+  // in all2all, the input is scattered, so there is no need to slice the token
+  // count, and we can use the dispatch buffer directly
+  if (enable_all2all_communication) {
+    token_count_slice = token_count;
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    int64_t dispatch_bytes =
+        num_token_expand * deep_ep_params_.dispatch_token_size;
+    dispatch_send_token_tensor =
+        deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes)
+            .view({num_token_expand, deep_ep_params_.dispatch_token_size});
+  } else {
+    token_count_slice =
+        token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size);
+  }
+
+  if (is_smoothquant_) {
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = hidden_states_2d;
+    // use dispatch_send_token_tensor buffer for input
+    //  to reduce memory footprint
+    if (enable_all2all_communication) {
+      scaled_quantize_params.smooth = input_smooth_;
+      scaled_quantize_params.output =
+          dispatch_send_token_tensor.slice(1, 0, hidden_size_);
+    } else {
+      scaled_quantize_params.smooth = input_smooth_.slice(
+          0, start_expert_id_, start_expert_id_ + expert_size);
+      scaled_quantize_params.gather_index_start_position =
+          cusum_token_count.value().index({start_expert_id_}).unsqueeze(0);
+    }
+    scaled_quantize_params.token_count = token_count_slice;
+    scaled_quantize_params.gather_index = gather_idx;
+    scaled_quantize_params.act_mode = "none";
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = false;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(expand_hidden_states, hidden_states_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+    if (enable_all2all_communication) {
+      // since view_as_dtype has not supported stride yet,
+      //  we need to copy the scale output to the dispatch buffer
+      torch::Tensor dispatch_scale_slice =
+          dispatch_send_token_tensor.slice(1, hidden_size_);
+      torch::Tensor hidden_states_scale_bytes =
+          view_as_dtype(hidden_states_scale, torch::kInt8)
+              .view_as(dispatch_scale_slice);
+      dispatch_scale_slice.copy_(hidden_states_scale_bytes);
+    }
+  } else {
+    xllm::kernel::MoeExpandInputParams moe_expand_input_params;
+    moe_expand_input_params.input = hidden_states_2d;
+    moe_expand_input_params.gather_index = gather_idx;
+    moe_expand_input_params.combine_idx = combine_idx;
+    moe_expand_input_params.topk = topk_;
+    expand_hidden_states =
+        xllm::kernel::moe_expand_input(moe_expand_input_params);
+    if (enable_all2all_communication) {
+      // use copy to place the output inside the dispatch buffer
+      torch::Tensor dispatch_tensor =
+          view_as_dtype(expand_hidden_states, torch::kChar);
+      dispatch_send_token_tensor.copy_(dispatch_tensor);
+    }
+  }
+
+  // collect the selected tensor
+  selected_expert_info.reduce_weight = reduce_weight;
+  selected_expert_info.combine_idx = combine_idx;
+  selected_expert_info.token_count_slice = token_count_slice;
+  selected_expert_info.cusum_token_count = cusum_token_count;
+  if (is_smoothquant_) {
+    selected_expert_info.input_scale = hidden_states_scale;
+  }
+
+  return expand_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states,
+                                            const torch::Tensor& router_logits,
+                                            bool enable_all2all_communication) {
+  if (!stream_initialized_) {
+    // update device record
+    device_ = xllm::Device(hidden_states.device());
+
+    // acquire streams from the pool again
+    routed_stream_ = device_.get_stream_from_pool();
+    shared_stream_ = device_.get_stream_from_pool();
+    stream_initialized_ = true;
+  }
+
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+
+  // prepare the parameters for MoE computation
+  torch::Tensor shared_expert_output;
+  torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
+  torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
+  torch::Tensor hidden_states_2d =
+      hidden_states.reshape({-1, hidden_states.size(-1)});
+  torch::Tensor router_logits_2d =
+      router_logits.reshape({-1, router_logits.size(-1)});
+  int64_t group_gemm_max_dim = enable_all2all_communication
+                                   ? deep_ep_params_.max_num_tokens_recv / topk_
+                                   : hidden_states_2d.size(0);
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1-3: select experts
+  SelectedExpertInfo selected_expert_info;
+  torch::Tensor expand_hidden_states =
+      select_experts(hidden_states_2d,
+                     router_logits_2d,
+                     selected_expert_info,
+                     enable_all2all_communication);
+
+  // Communciation Step 1: Dipatch
+  // intermediate outputs that are used both in dispatch and combine
+  torch::Tensor gather_by_rank_index;
+  torch::Tensor token_sum;
+  if (enable_all2all_communication) {
+    int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_;
+
+    // 1. Dispatch Step: Generate layout and send data
+    deep_ep_->dispatch_step(dispatch_token_num,
+                            selected_expert_info.token_count_slice);
+
+    // 2. Process Result: Generate indices and unpack to computation buffer
+    // use the buffer during initialization for the output
+    expand_hidden_states = dispatch_recv_token_tensor_head_;
+    std::optional<torch::Tensor> output_tail = std::nullopt;
+    if (is_smoothquant_) {
+      output_tail = dispatch_recv_token_tensor_tail_;
+      // update selected_expert_info with the tail (input scale)
+      selected_expert_info.input_scale = output_tail;
+    }
+
+    DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result(
+        num_experts_per_rank_, expand_hidden_states, output_tail);
+
+    // Extract metadata for subsequent steps
+    gather_by_rank_index = deep_ep_meta.gather_rank_index;
+    selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice;
+    token_sum = deep_ep_meta.token_sum;
+  }
+
+  // common gemm workspace for reduce memory footprint
+  torch::Tensor gemm_workspace;
+
+  // Step 4: group gemm 1
+  torch::Tensor gemm1_out =
+      create_group_gemm_output(expand_hidden_states,
+                               w13_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    torch::ScalarType a_dtype =
+        is_smoothquant_ ? torch::kInt8 : hidden_states_dtype;
+    group_gemm_params.a =
+        view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_});
+    group_gemm_params.b = w13_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      torch::Tensor a_scale =
+          selected_expert_info.input_scale.value().flatten();
+      selected_expert_info.input_scale =
+          view_as_dtype(a_scale, torch::kFloat32);
+      group_gemm_params.a_scale = selected_expert_info.input_scale;
+      group_gemm_params.b_scale = w13_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm1_out;
+    group_gemm_params.combine_idx = std::nullopt;
+    gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Step 5: activation or scaled quantization(fused with activation)
+  torch::Tensor act_out;
+  torch::Tensor act_out_scale;
+  if (is_smoothquant_) {
+    int64_t slice_dim = gemm1_out.size(1);
+    if (is_gated_) slice_dim /= 2;
+    // slice operation is a view, does not take up extra memory, but points to
+    // the same memory
+    act_out = expand_hidden_states.slice(1, 0, slice_dim);
+    act_out_scale =
+        selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0));
+    // call scaled quantization kernel (also fused with activation)
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = gemm1_out;
+    scaled_quantize_params.smooth = act_smooth_;
+    scaled_quantize_params.token_count = selected_expert_info.token_count_slice;
+    scaled_quantize_params.output = act_out;
+    scaled_quantize_params.output_scale = act_out_scale;
+    scaled_quantize_params.act_mode = hidden_act_;
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = is_gated_;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(act_out, act_out_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+  } else {
+    act_out = is_gated_
+                  ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous()
+                  : gemm1_out;
+    // call activation kernel
+    xllm::kernel::ActivationParams activation_params;
+    activation_params.input = gemm1_out;
+    activation_params.output = act_out;
+    activation_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    activation_params.act_mode = hidden_act_;
+    activation_params.is_gated = is_gated_;
+    activation_params.start_expert_id = start_expert_id_;
+    activation_params.expert_size = expert_size;
+    xllm::kernel::active(activation_params);
+  }
+
+  // Step 6: group gemm 2
+  torch::Tensor gemm2_out =
+      create_group_gemm_output(act_out,
+                               w2_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    group_gemm_params.a = act_out;
+    group_gemm_params.b = w2_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      group_gemm_params.a_scale = act_out_scale;
+      group_gemm_params.b_scale = w2_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm2_out;
+    group_gemm_params.combine_idx = selected_expert_info.combine_idx;
+    gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Communciation Step 2: Combine
+  if (enable_all2all_communication) {
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    // Delegate pack, layout generation and combine to DeepEP
+    torch::Tensor combine_send_layout =
+        deep_ep_->combine_step_pack(gemm2_out,
+                                    gather_by_rank_index,
+                                    token_sum,
+                                    hidden_size_,
+                                    hidden_states_dtype);
+
+    // create a wait event for the current stream to finish computation
+    auto current_stream = device_.current_stream();
+    routed_stream_->wait_stream(*current_stream);
+    // pure communciation kernel: dispatch
+    {
+      torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+      gemm2_out = deep_ep_->combine_step_comm(combine_send_layout,
+                                              num_token_expand,
+                                              hidden_size_,
+                                              hidden_states_dtype);
+    }
+
+    // pure computation kernel: shared experts
+    if (n_shared_experts_ > 0) {
+      shared_stream_->wait_stream(*current_stream);
+      torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+      shared_expert_output = shared_experts_(hidden_states);
+    }
+
+    // join for parallelization
+    current_stream->wait_stream(*routed_stream_);
+    if (n_shared_experts_ > 0) {
+      current_stream->wait_stream(*shared_stream_);
+    }
+  }
+
+  // After group gemm is finished, some tensors are no
+  // longer needed. We must explicitly release the memory.
+  expand_hidden_states = torch::Tensor();
+  selected_expert_info.input_scale = std::nullopt;
+  act_out = torch::Tensor();
+
+  // Step 7: combine the intermediate results and get the final hidden states
+  torch::Tensor final_hidden_states;
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::MoeCombineResultParams moe_combine_result_params;
+    moe_combine_result_params.input = gemm2_out;
+    moe_combine_result_params.reduce_weight =
+        selected_expert_info.reduce_weight;
+    moe_combine_result_params.gather_ids = selected_expert_info.combine_idx;
+    moe_combine_result_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    moe_combine_result_params.start_expert_id = start_expert_id_;
+    moe_combine_result_params.expert_size = expert_size;
+    moe_combine_result_params.bias = std::nullopt;
+    // if all2all communication is enabled and shared output is provided,
+    //  we will fused the add up to combine result
+    if (enable_all2all_communication && n_shared_experts_ > 0) {
+      moe_combine_result_params.residual =
+          shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+    }
+
+    final_hidden_states =
+        xllm::kernel::moe_combine_result(moe_combine_result_params);
+  }
+
+  // reshape the final hidden states to the original shape
+  final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
+
+  if (enable_all2all_communication) {
+    return final_hidden_states;
+  }
+
+  // Communciation Step 3: AllReduce for non-all2all communication
+  // shared experts can be parallelized with the final communication step
+  // during moe computation.
+  auto current_stream = device_.current_stream();
+  routed_stream_->wait_stream(*current_stream);
+  {
+    torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+    if (tp_pg_->world_size() > 1) {
+      final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
+    }
+    if (parallel_args_.ep_size() > 1) {
+      final_hidden_states = parallel_state::reduce(
+          final_hidden_states, parallel_args_.moe_ep_group_);
+    }
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_stream_->wait_stream(*current_stream);
+    torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+    // for non all2all, we compute the shared experts parallelized with the
+    // final communication step
+    shared_expert_output = shared_experts_(hidden_states);
+    shared_expert_output =
+        shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+  }
+
+  // join for parallelization
+  current_stream->wait_stream(*routed_stream_);
+  if (n_shared_experts_ > 0) {
+    current_stream->wait_stream(*shared_stream_);
+    final_hidden_states += shared_expert_output;
+  }
+
+  return final_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
+                                    const ModelInputParams& input_params) {
+  // we only support all2all communication for decode stage for now
+  bool enable_all2all_communication =
+      enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(),
+                                     input_params.dp_is_decode.end(),
+                                     [](int32_t val) { return val == 1; });
+
+  bool is_dp_ep_parallel =
+      parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1;
+  // during all2all communication, the output has been
+  //  gathered and sliced by dispatch and combine steps,
+  //  so we do not need to gather input and slice output again
+  bool need_gather_and_slice =
+      is_dp_ep_parallel && !enable_all2all_communication;
+
+  auto input = hidden_states;
+  if (need_gather_and_slice) {
+    input = parallel_state::gather(input,
+                                   parallel_args_.dp_local_process_group_,
+                                   input_params.dp_global_token_nums);
+  }
+  // MoE Gate
+  auto router_logits = gate_(input);
+
+  // MoE Experts
+  auto output =
+      forward_experts(input, router_logits, enable_all2all_communication);
+
+  if (need_gather_and_slice) {
+    output = get_dp_local_slice(output, input_params, parallel_args_);
+  }
+
+  return output;
+}
+
+void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
+  if (e_score_correction_bias_.defined() &&
+      !e_score_correction_bias_is_loaded_) {
+    LOAD_WEIGHT(e_score_correction_bias);
+  }
+}
+
+void FusedMoEImpl::load_experts(const StateDict& state_dict) {
+  const int64_t rank = tp_pg_->rank();
+  const int64_t world_size = tp_pg_->world_size();
+  const int64_t start_expert_id = start_expert_id_;
+  const int64_t num_experts_per_rank = num_experts_per_rank_;
+  const int64_t num_total_experts = num_total_experts_;
+  std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
+  if (is_smoothquant_) {
+    LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
+    LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
+    // When supporting DeepEP All2All mode,
+    // we need to load the complete set of expert weights corresponding to
+    // "up_proj.smooth". Note that even if deep EP mode is not enabled, it
+    // remains possible to retrieve the smooth quantization information for a
+    // subset of experts. Therefore, we intentionally do not check whether
+    // deep_ep_ is enabled in this case.
+    LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
+    LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
+  } else {
+    LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
+    LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
+  }
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
+  if (state_dict.size() == 0) {
+    return;
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_experts_->load_state_dict(
+        state_dict.get_dict_with_prefix("shared_experts."));
+  }
+  gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
+  load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
+  load_experts(state_dict.get_dict_with_prefix("experts."));
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/fused_moe.h

@@ -0,0 +1,131 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/deep_ep.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/fused_moe_base.h"
+#include "layers/common/linear.h"
+#include "platform/device.h"
+#include "util/tensor_helper.h"
+
+namespace xllm {
+namespace layer {
+
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_experts(const torch::Tensor& hidden_states,
+                                const torch::Tensor& router_logits,
+                                bool enable_all2all_communication);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+
+ private:
+  // struct to store the selected expert info
+  struct SelectedExpertInfo {
+    torch::Tensor reduce_weight;
+    torch::Tensor combine_idx;
+    torch::Tensor token_count_slice;
+    std::optional<torch::Tensor> cusum_token_count;
+    std::optional<torch::Tensor> input_scale;
+  };
+
+  // initial steps for MoE computation, select the experts for each token
+  torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
+                               const torch::Tensor& router_logits_2d,
+                               SelectedExpertInfo& selected_expert_info,
+                               bool enable_all2all_communication);
+
+ private:
+  int64_t num_total_experts_;
+  int64_t topk_;
+  int64_t num_expert_group_;
+  int64_t topk_group_;
+  double route_scale_;
+  int64_t hidden_size_;
+  int64_t n_shared_experts_;
+  bool is_gated_;
+  int64_t renormalize_;
+  std::string hidden_act_;
+  std::string scoring_func_;
+  bool is_smoothquant_;
+
+  int64_t num_experts_per_rank_;
+  int64_t start_expert_id_;
+
+  // Deep EP related parameters
+  bool enable_deep_ep_;
+  DeepEPBuffer deep_ep_buffer_;
+  DeepEPParams deep_ep_params_;
+  torch::Tensor dispatch_recv_token_tensor_head_;
+  torch::Tensor dispatch_recv_token_tensor_tail_;
+
+  // steams for parallel shared experts
+  std::unique_ptr<Stream> shared_stream_;
+  std::unique_ptr<Stream> routed_stream_;
+  xllm::Device device_;
+  bool stream_initialized_ = false;
+
+  ReplicatedLinear gate_{nullptr};
+  DenseMLP shared_experts_{nullptr};
+  DeepEP deep_ep_{nullptr};
+
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_;
+  torch::TensorOptions options_;
+  ProcessGroup* tp_pg_;
+
+  DEFINE_WEIGHT(w13);
+  DEFINE_FUSED_WEIGHT(w1);
+  DEFINE_FUSED_WEIGHT(w3);
+  DEFINE_FUSED_WEIGHT(w2);
+  DEFINE_WEIGHT(e_score_correction_bias);
+  DEFINE_WEIGHT(w13_scale);
+  DEFINE_FUSED_WEIGHT(w1_scale);
+  DEFINE_FUSED_WEIGHT(w3_scale);
+  DEFINE_FUSED_WEIGHT(w2_scale);
+  DEFINE_FUSED_WEIGHT(input_smooth);
+  DEFINE_FUSED_WEIGHT(act_smooth);
+
+  void load_e_score_correction_bias(const StateDict& state_dict);
+  void load_experts(const StateDict& state_dict);
+  // create the group gemm output tensor with the workspace
+  torch::Tensor create_group_gemm_output(const torch::Tensor& a,
+                                         const torch::Tensor& b,
+                                         const torch::Tensor& group_list,
+                                         torch::ScalarType dtype,
+                                         torch::Tensor& workspace);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2026 The xLLM Authors. 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
@@ -123,53 +123,19 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
 }
 
 std::pair<torch::Tensor, torch::Tensor>
-Qwen3_5GatedDeltaNetImpl::project_decode_inputs(
-    const torch::Tensor& hidden_states) {
-  const auto reshape_projection = [](const torch::Tensor& projection) {
-    return projection.view({projection.size(0), -1, projection.size(-1)});
-  };
-  auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states));
-  auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states));
-  auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states));
-  auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states));
-  return {merge_qkvz_from_split_activations(qkv, z_proj),
-          merge_ba_from_split_activations(b_proj, a_proj)};
-}
-
-std::pair<torch::Tensor, torch::Tensor>
-Qwen3_5GatedDeltaNetImpl::project_flat_inputs(
-    const torch::Tensor& hidden_states) {
-  auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0);
-  auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0);
-  auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0);
-  auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0);
-  auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj);
-  auto ba = merge_ba_from_split_activations(b_proj, a_proj);
-  return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(),
-          ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()};
-}
-
-std::optional<
-    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
-Qwen3_5GatedDeltaNetImpl::project_split_inputs(
+Qwen3_5GatedDeltaNetImpl::project_padded_inputs(
     const torch::Tensor& hidden_states,
     const AttentionMetadata& attn_metadata) {
-  auto qkv = reshape_projected_tokens_with_pad(
-      attn_metadata, in_proj_qkv_->forward(hidden_states));
-  auto z_proj = reshape_projected_tokens_with_pad(
-      attn_metadata, in_proj_z_->forward(hidden_states));
-  auto b_proj = reshape_projected_tokens_with_pad(
-      attn_metadata, in_proj_b_->forward(hidden_states));
-  auto a_proj = reshape_projected_tokens_with_pad(
-      attn_metadata, in_proj_a_->forward(hidden_states));
-
-  const int64_t batch_size = qkv.size(0);
-  const int64_t seq_len = qkv.size(1);
-  auto z =
-      z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
-  auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
-  auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
-  return std::make_tuple(qkv, z, b, a);
+  auto qkv = reshape_qkvz_with_pad(attn_metadata,
+                                   in_proj_qkv_->forward(hidden_states));
+  auto z_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states));
+  auto b_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states));
+  auto a_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states));
+  return {merge_qkvz_from_split_activations(qkv, z_proj),
+          merge_ba_from_split_activations(b_proj, a_proj)};
 }
 
 void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(

ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2026 The xLLM Authors. 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.
@@ -17,9 +17,7 @@ limitations under the License.
 
 #include <torch/torch.h>
 
-#include <optional>
 #include <string>
-#include <tuple>
 #include <utility>
 
 #include "qwen3_next_gated_delta_net.h"
@@ -36,15 +34,9 @@ class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
                            const torch::TensorOptions& options);
 
  protected:
-  std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
-      const torch::Tensor& hidden_states) override;
-  std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
-      const torch::Tensor& hidden_states) override;
-  std::optional<
-      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
-  project_split_inputs(const torch::Tensor& hidden_states,
-                       const AttentionMetadata& attn_metadata) override;
-  bool use_fla_ssm_state_layout() const override { return true; }
+  std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
+      const torch::Tensor& hidden_states,
+      const AttentionMetadata& attn_metadata) override;
 
   void load_projection_state_dict(const StateDict& state_dict) override;
   void verify_projection_weights(const std::string& prefix) const override;

ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2026 The xLLM Authors. 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
@@ -15,12 +15,9 @@ limitations under the License.
 #include <glog/logging.h>
 #include <torch/torch.h>
 
-#include <optional>
 #include <tuple>
 
-#include "xllm/core/kernels/npu/npu_ops_api.h"
 #include "xllm/core/kernels/ops_api.h"
-#include "xllm/core/platform/npu/acl_graph_task_update_context.h"
 
 namespace xllm {
 namespace layer {
@@ -31,31 +28,6 @@ torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
   return x / norm;
 }
 
-torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor,
-                                  int64_t target_heads,
-                                  int64_t head_dim) {
-  const int64_t current_heads = tensor.size(head_dim);
-  if (current_heads == target_heads) {
-    return tensor;
-  }
-  CHECK_GT(current_heads, 0) << "current heads must be positive";
-  CHECK_EQ(target_heads % current_heads, 0)
-      << "target heads must be divisible by current heads, target_heads="
-      << target_heads << ", current_heads=" << current_heads;
-
-  const int64_t repeats = target_heads / current_heads;
-  std::vector<int64_t> view_shape = tensor.sizes().vec();
-  view_shape.insert(view_shape.begin() + head_dim + 1, 1);
-  std::vector<int64_t> expand_shape = view_shape;
-  expand_shape[head_dim + 1] = repeats;
-  std::vector<int64_t> output_shape = tensor.sizes().vec();
-  output_shape[head_dim] = target_heads;
-  return tensor.unsqueeze(head_dim + 1)
-      .expand(expand_shape)
-      .reshape(output_shape)
-      .contiguous();
-}
-
 std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
     torch::Tensor query,
     torch::Tensor key,
@@ -80,9 +52,6 @@ std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
   value = to_float32_and_transpose(value);
   beta = to_float32_and_transpose(beta);
   g = to_float32_and_transpose(g);
-  const int64_t value_num_heads = value.size(1);
-  query = repeat_tensor_heads(query, value_num_heads, 1);
-  key = repeat_tensor_heads(key, value_num_heads, 1);
 
   int64_t batch_size = key.size(0);
   int64_t num_heads = key.size(1);
@@ -150,15 +119,12 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
   value = to_float32(value);
   beta = to_float32(beta);
   g = to_float32(g);
-  const int64_t value_num_heads = value.size(1);
-  query = repeat_tensor_heads(query, value_num_heads, 1);
-  key = repeat_tensor_heads(key, value_num_heads, 1);
 
-  int64_t batch_size = query.size(0);
-  int64_t num_heads = query.size(1);
-  int64_t sequence_length = query.size(2);
-  int64_t k_head_dim = key.size(-1);
-  int64_t v_head_dim = value.size(-1);
+  auto batch_size = query.size(0);
+  auto num_heads = query.size(1);
+  auto sequence_length = query.size(2);
+  auto k_head_dim = key.size(-1);
+  auto v_head_dim = value.size(-1);
 
   int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
   query = torch::nn::functional::pad(
@@ -276,164 +242,6 @@ std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
   core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
   return std::make_tuple(core_attn_out, last_recurrent_state);
 }
-
-int64_t get_checkpoint_stride(const torch::Tensor& conv_cache,
-                              const torch::Tensor& ssm_cache) {
-  if (!conv_cache.defined() || !ssm_cache.defined() ||
-      conv_cache.numel() == 0 || ssm_cache.numel() == 0) {
-    return 1;
-  }
-  CHECK_GT(conv_cache.size(0), 0) << "conv cache must have positive batch dim";
-  CHECK_EQ(ssm_cache.size(0) % conv_cache.size(0), 0)
-      << "ssm cache checkpoint layout mismatch, ssm_rows=" << ssm_cache.size(0)
-      << ", conv_rows=" << conv_cache.size(0);
-  return ssm_cache.size(0) / conv_cache.size(0);
-}
-
-torch::Tensor build_linear_state_base_indices(
-    const torch::Tensor& logical_state_indices,
-    int64_t checkpoint_stride) {
-  if (checkpoint_stride == 1) {
-    return logical_state_indices;
-  }
-  return logical_state_indices * checkpoint_stride;
-}
-
-torch::Tensor expand_sequence_tensor_to_batch(const torch::Tensor& tensor,
-                                              int64_t target_batch,
-                                              const char* tensor_name) {
-  CHECK(tensor.defined()) << tensor_name << " must be defined";
-  CHECK_EQ(tensor.dim(), 1) << tensor_name << " must be a 1D tensor.";
-  const int64_t source_batch = tensor.size(0);
-  if (source_batch == target_batch) {
-    return tensor.contiguous();
-  }
-  CHECK_GT(source_batch, 0) << tensor_name << " must not be empty.";
-  CHECK_EQ(target_batch % source_batch, 0)
-      << tensor_name << " cannot be expanded from " << source_batch << " to "
-      << target_batch;
-  const int64_t repeat_count = target_batch / source_batch;
-  return tensor.unsqueeze(1)
-      .expand({source_batch, repeat_count})
-      .reshape({target_batch})
-      .contiguous();
-}
-
-torch::Tensor run_causal_conv1d_graph_update(
-    const std::shared_ptr<xllm::npu::AclGraphTaskUpdateContext>& graph_context,
-    const torch::Tensor& x,
-    const torch::Tensor& weight,
-    const torch::Tensor& conv_state,
-    const std::optional<torch::Tensor>& bias,
-    const std::vector<int64_t>& query_start_loc,
-    const std::vector<int64_t>& cache_indices,
-    const std::vector<int64_t>& num_accepted_tokens,
-    xllm::npu::CausalConv1dGraphBranch branch) {
-  CHECK(graph_context != nullptr && graph_context->capturing)
-      << "causal_conv1d graph update can only be registered during capture";
-
-  c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream();
-  auto event = std::make_shared<c10_npu::NPUEvent>(ACL_EVENT_EXTERNAL);
-  event->block(stream);
-  event->reset(stream);
-
-  torch::Tensor output;
-  c10_npu::graph_task_group_begin(stream);
-  const std::vector<int64_t> empty_host_args;
-  CHECK(!query_start_loc.empty())
-      << "query_start_loc must be populated for causal_conv1d graph update";
-  CHECK_EQ(query_start_loc.back(), x.size(0))
-      << "query_start_loc must be padded to x.shape[0] during graph capture";
-  CHECK_EQ(cache_indices.size() + 1, query_start_loc.size())
-      << "cache_indices must be sequence-scoped";
-  if (branch == xllm::npu::CausalConv1dGraphBranch::kSpecVerify) {
-    CHECK_EQ(num_accepted_tokens.size(), cache_indices.size())
-        << "num_accepted_tokens must be sequence-scoped for spec verify";
-  }
-
-  output = torch::empty_like(x);
-  xllm::kernel::causal_conv1d_out(output,
-                                  x,
-                                  weight,
-                                  conv_state,
-                                  bias,
-                                  torch::IntArrayRef(query_start_loc),
-                                  torch::IntArrayRef(cache_indices),
-                                  torch::IntArrayRef(empty_host_args),
-                                  torch::IntArrayRef(num_accepted_tokens),
-                                  xllm::npu::kCausalConv1dActivationSilu,
-                                  xllm::npu::kCausalConv1dGraphPadSlotId,
-                                  xllm::npu::kCausalConv1dRunModeUpdate);
-  c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream);
-
-  xllm::npu::CausalConv1dGraphTask task;
-  task.output = output;
-  task.x = x;
-  task.weight = weight;
-  task.conv_state = conv_state;
-  task.bias = bias;
-  task.activation_mode = xllm::npu::kCausalConv1dActivationSilu;
-  task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId;
-  task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate;
-  task.branch = branch;
-  task.handle = handle;
-  task.event = std::move(event);
-  graph_context->causal_conv1d_tasks.emplace_back(std::move(task));
-  return output;
-}
-
-torch::Tensor run_spec_verify_gated_delta_rule(
-    torch::Tensor query,
-    torch::Tensor key,
-    torch::Tensor value,
-    torch::Tensor g,
-    torch::Tensor beta,
-    torch::Tensor& ssm_cache,
-    const torch::Tensor& checkpoint_indices,
-    const torch::Tensor& num_accepted_tokens,
-    const torch::Tensor& cu_seq_lens,
-    const std::vector<int32_t>& q_seq_lens_vec,
-    double scale) {
-  const auto device = value.device();
-  const int64_t batch_size = value.size(0);
-  const int64_t seq_len = value.size(1);
-  const int64_t total_seq_len = batch_size * seq_len;
-  CHECK_EQ(cu_seq_lens.numel(), batch_size + 1)
-      << "GDN spec verify cu_seq_lens must be cumulative.";
-  CHECK_EQ(q_seq_lens_vec.size(), static_cast<size_t>(batch_size))
-      << "GDN spec verify q_seq_lens_vec must be per sequence.";
-  for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
-    CHECK_EQ(q_seq_lens_vec[batch_idx], seq_len)
-        << "Qwen3.5 spec verify fused recurrent path expects dense "
-           "same-length validate tokens.";
-  }
-
-  xllm::kernel::FusedRecurrentGatedDeltaRuleParams params;
-  params.q = query.reshape({1, total_seq_len, query.size(-2), query.size(-1)})
-                 .contiguous();
-  params.k =
-      key.reshape({1, total_seq_len, key.size(-2), key.size(-1)}).contiguous();
-  params.v = value.reshape({1, total_seq_len, value.size(-2), value.size(-1)})
-                 .contiguous();
-  params.g = g.to(torch::kFloat32)
-                 .reshape({1, total_seq_len, g.size(-1)})
-                 .contiguous();
-  params.beta = beta.reshape({1, total_seq_len, beta.size(-1)}).contiguous();
-  params.scale = static_cast<float>(scale);
-  params.initial_state = ssm_cache;
-  params.inplace_final_state = true;
-  params.cu_seqlens = cu_seq_lens.to(torch::kLong).contiguous();
-  params.ssm_state_indices = checkpoint_indices.contiguous();
-  params.num_accepted_tokens =
-      num_accepted_tokens.to(device, torch::kInt32).contiguous();
-  params.use_qk_l2norm_in_kernel = true;
-
-  auto output_and_state =
-      xllm::kernel::fused_recurrent_gated_delta_rule(params);
-  return output_and_state.first.view(
-      {batch_size, seq_len, value.size(-2), value.size(-1)});
-}
-
 }  // namespace
 
 Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl(
@@ -495,11 +303,7 @@ void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict(
 
   if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) {
     conv1d_->load_state_dict(
-        StateDict({{"weight", w.squeeze(1)}},
-                  static_cast<std::string>(state_dict.prefix()) + "conv1d."),
-        shard_tensor_count,
-        shard_sizes);
-    conv1d_->weight().set_(conv1d_->weight().transpose(0, 1).contiguous());
+        StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes);
   }
   o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
   if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
@@ -518,279 +322,87 @@ void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights(
                           << prefix << "A_log";
 }
 
-std::pair<torch::Tensor, torch::Tensor>
-Qwen3GatedDeltaNetBaseImpl::project_padded_inputs(
-    const torch::Tensor& hidden_states,
-    const AttentionMetadata& attn_metadata) {
-  if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) {
-    auto [qkvz_flat, ba_flat] = project_flat_inputs(hidden_states);
-    return {reshape_projected_tokens_with_pad(attn_metadata, qkvz_flat),
-            reshape_projected_tokens_with_pad(attn_metadata, ba_flat)};
-  }
-  return project_decode_inputs(hidden_states);
-}
-
 torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
     const torch::Tensor& hidden_states,
     const AttentionMetadata& attn_metadata,
     KVCache& kv_cache,
     const ModelInputParams& input_params) {
-  // Early-return on dummy shards. Under dp>1, an empty shard is padded with a
-  // fake token by worker_impl but its GDN state tensors (kv_cache_tokens_nums,
-  // linear_state_ids etc.) are left undefined. This mirrors the is_dummy
-  // early-return in Attention::forward (npu_torch/attention.cpp). Uses
-  // zeros_like rather than empty_like so downstream post-norm / mlp do not
-  // read uninitialized data. Placed before FlashComm1 sequence gather so
-  // dummy shards do not enter the collective and waste bandwidth.
-  if (attn_metadata.is_dummy) {
-    return torch::zeros_like(hidden_states);
-  }
-  const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
-  torch::Tensor h = hidden_states;
-  if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
-    h = gather_sequence(hidden_states, *fc1_ctx);
-  }
+  auto [qkvz_padded, ba_padded] =
+      project_padded_inputs(hidden_states, attn_metadata);
+  int64_t batch_size = qkvz_padded.size(0);
+  int64_t seq_len = qkvz_padded.size(1);
+
+  torch::Tensor qkvz_flat =
+      qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
+  torch::Tensor ba_flat =
+      ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
+  xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
+  fused_params.mixed_qkvz = qkvz_flat;
+  fused_params.mixed_ba = ba_flat;
+  fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
+  fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
+  fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
+  fused_params.head_v = static_cast<int32_t>(head_v_dim_);
 
-  // Save the gathered hidden-state size for potential padding later.
-  const int64_t original_num_tokens = h.size(0);
-  const bool use_spec_verify = input_params.is_spec_verify;
-  const bool is_any_prefill =
-      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
   torch::Tensor mixed_qkv, z, b, a;
-  torch::Tensor processed_q, processed_k, processed_v;
-  int64_t batch_size = 0;
-  int64_t seq_len = 0;
+  std::tie(mixed_qkv, z, b, a) =
+      xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
 
-  // Qwen3.5 stores qkv, z, b, and a as separate projection weights, so it can
-  // use their outputs directly in every forward mode. Qwen3Next stores qkvz
-  // and ba as packed weights and uses the fused-split fallback below.
-  auto split_inputs = project_split_inputs(h, attn_metadata);
-  if (split_inputs.has_value()) {
-    std::tie(mixed_qkv, z, b, a) = split_inputs.value();
-    batch_size = mixed_qkv.size(0);
-    seq_len = mixed_qkv.size(1);
-  } else {
-    auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata);
-    batch_size = qkvz_padded.size(0);
-    seq_len = qkvz_padded.size(1);
-
-    torch::Tensor qkvz_flat =
-        qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
-    torch::Tensor ba_flat =
-        ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
-    xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
-    fused_params.mixed_qkvz = qkvz_flat;
-    fused_params.mixed_ba = ba_flat;
-    fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
-    fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
-    fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
-    fused_params.head_v = static_cast<int32_t>(head_v_dim_);
-
-    std::tie(mixed_qkv, z, b, a) =
-        xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
-
-    mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
-    z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
-    b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
-    a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
-  }
-
-  const bool fla_ssm_state_layout = use_fla_ssm_state_layout();
-  const int64_t local_q_heads = num_k_heads_ / tp_size_;
-  const int64_t local_v_heads = num_v_heads_ / tp_size_;
-  const int64_t local_conv_dim =
-      2 * local_q_heads * head_k_dim_ + local_v_heads * head_v_dim_;
-  bool used_direct_prefill_qkv = false;
+  mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
+  z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
+  b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+  a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
 
   torch::Tensor conv_cache = kv_cache.get_conv_cache();
   torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
-  torch::Device device = mixed_qkv.device();
-  torch::Tensor conv_weight = conv1d_->weight();
-  torch::Tensor logical_state_indices =
-      get_linear_state_indices(input_params, device);
-  const int64_t checkpoint_stride =
-      get_checkpoint_stride(conv_cache, ssm_cache);
-  torch::Tensor linear_state_base_indices =
-      build_linear_state_base_indices(logical_state_indices, checkpoint_stride);
-  auto graph_context = input_params.graph.acl_graph_task_update_context;
-  const bool register_conv1d_graph_update =
-      graph_context != nullptr && graph_context->capturing;
+  torch::Tensor g, beta, core_attn_out, last_recurrent_state;
+  auto device = mixed_qkv.device();
+  auto conv_weight = conv1d_->weight();
+  auto linear_state_indices = get_linear_state_indices(input_params, device);
 
-  if (!use_spec_verify && is_any_prefill) {
-    torch::IntArrayRef num_accepted_tokens_opt;
-    std::vector<int64_t> linear_state_indices_vec(
-        input_params.embedding.linear_state_ids.begin(),
-        input_params.embedding.linear_state_ids.end());
-    torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv);
+  if (attn_metadata.is_prefill) {
+    mixed_qkv = mixed_qkv.transpose(1, 2);
+    torch::Tensor conv_state =
+        (seq_len < conv_kernel_size_ - 1)
+            ? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len})
+        : (seq_len > conv_kernel_size_ - 1)
+            ? mixed_qkv.narrow(
+                  -1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1)
+            : mixed_qkv;
+    conv_state = conv_state.transpose(1, 2).contiguous();
+    conv_cache.index_put_({linear_state_indices},
+                          conv_state.to(conv_cache.dtype()));
+    torch::Tensor bias;
+    auto conv_output =
+        torch::conv1d(mixed_qkv,
+                      conv_weight.unsqueeze(1).to(device),
+                      bias,
+                      /*stride=*/std::vector<int64_t>{1},
+                      /*padding=*/std::vector<int64_t>{3},
+                      /*dilation=*/std::vector<int64_t>{1},
+                      /*groups=*/static_cast<int64_t>(mixed_qkv.size(1)));
+    mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len));
 
-    const bool direct_qkv_model_supported =
-        fla_ssm_state_layout && num_k_heads_ % tp_size_ == 0 &&
-        num_v_heads_ % tp_size_ == 0 && local_q_heads > 0 &&
-        local_v_heads > 0 && head_k_dim_ == 128 && head_v_dim_ == 128;
-    const bool direct_qkv_metadata_available =
-        attn_metadata.q_seq_lens_vec.size() ==
-            static_cast<size_t>(batch_size) &&
-        input_params.parallel.query_start_loc.size() ==
-            static_cast<size_t>(batch_size + 1) &&
-        input_params.embedding.linear_state_ids.size() ==
-            static_cast<size_t>(batch_size) &&
-        input_params.linear_state_validity_mask.size() ==
-            static_cast<size_t>(batch_size);
-    int64_t total_valid_tokens = 0;
-    bool direct_qkv_lengths_valid = direct_qkv_metadata_available;
-    if (direct_qkv_metadata_available) {
-      for (const int32_t valid_len : attn_metadata.q_seq_lens_vec) {
-        direct_qkv_lengths_valid =
-            direct_qkv_lengths_valid && valid_len >= 0 && valid_len <= seq_len;
-        total_valid_tokens += valid_len;
-      }
-    }
-    const bool direct_qkv_sequence_supported =
-        direct_qkv_model_supported && direct_qkv_lengths_valid &&
-        conv_input.dim() == 2 && total_valid_tokens == conv_input.size(0);
-    const bool direct_qkv_shape_supported =
-        direct_qkv_sequence_supported && conv_input.size(1) == local_conv_dim &&
-        conv_weight.dim() == 2 && conv_weight.size(0) == 4 &&
-        conv_weight.size(1) == local_conv_dim && conv_cache.dim() == 3 &&
-        conv_cache.size(1) >= 3 && conv_cache.size(2) == local_conv_dim;
-    const bool direct_qkv_dtype_supported =
-        direct_qkv_shape_supported &&
-        conv_input.scalar_type() == torch::kBFloat16 &&
-        conv_weight.scalar_type() == torch::kBFloat16 &&
-        conv_cache.scalar_type() == torch::kBFloat16;
-    const bool use_direct_prefill_qkv =
-        direct_qkv_dtype_supported && conv_input.is_contiguous() &&
-        conv_weight.is_contiguous() && conv_cache.is_contiguous();
-    if (use_direct_prefill_qkv) {
-      std::tie(processed_q, processed_k, processed_v) =
-          xllm::kernel::npu::causal_conv1d_qkv(
-              conv_input,
-              conv_weight,
-              conv_cache,
-              torch::IntArrayRef(input_params.parallel.query_start_loc),
-              torch::IntArrayRef(linear_state_indices_vec),
-              torch::IntArrayRef(input_params.linear_state_validity_mask),
-              local_q_heads,
-              local_v_heads,
-              head_k_dim_,
-              head_v_dim_);
-      used_direct_prefill_qkv = true;
-    } else {
-      mixed_qkv = xllm::kernel::causal_conv1d(
-          conv_input,
-          conv_weight,
-          conv_cache,
-          std::optional<torch::Tensor>(),  // bias (no bias for qwen3)
-          torch::IntArrayRef(input_params.parallel.query_start_loc),
-          torch::IntArrayRef(linear_state_indices_vec),
-          torch::IntArrayRef(input_params.linear_state_validity_mask),
-          num_accepted_tokens_opt,
-          xllm::npu::kCausalConv1dActivationSilu,
-          xllm::npu::kCausalConv1dGraphPadSlotId,
-          xllm::npu::kCausalConv1dRunModeForward);
-
-      mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv);
-      mixed_qkv = mixed_qkv.transpose(1, 2);
-    }
   } else {
-    if (use_spec_verify) {
-      CHECK(input_params.num_accepted_tokens.defined())
-          << "num_accepted_tokens must be populated for Qwen3.5 spec verify";
-    }
-    torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv);
-    const auto& num_accepted = use_spec_verify
-                                   ? input_params.num_accepted_tokens_host
-                                   : std::vector<int64_t>();
-    const std::vector<int64_t> linear_state_indices_host(
-        input_params.embedding.linear_state_ids.begin(),
-        input_params.embedding.linear_state_ids.end());
-    if (register_conv1d_graph_update) {
-      if (use_spec_verify) {
-        const auto conv1d_branch =
-            xllm::npu::CausalConv1dGraphBranch::kSpecVerify;
-        mixed_qkv = run_causal_conv1d_graph_update(
-            graph_context,
-            conv_input,
-            conv_weight,
-            conv_cache,
-            std::optional<torch::Tensor>(),
-            input_params.parallel.query_start_loc,
-            linear_state_indices_host,
-            num_accepted,
-            conv1d_branch);
-      } else {
-        auto conv_input_2d = conv_input.dim() == 3
-                                 ? conv_input.reshape({-1, conv_input.size(-1)})
-                                 : conv_input;
-        xllm::kernel::CausalConv1dUpdateParams conv1d_params;
-        conv1d_params.x = conv_input_2d;
-        conv1d_params.conv_state = conv_cache;
-        conv1d_params.weight = conv_weight;
-        conv1d_params.conv_state_indices = logical_state_indices;
-        conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
-        conv1d_params.max_query_len = attn_metadata.max_query_len;
-        mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
-        if (conv_input.dim() == 3) {
-          mixed_qkv =
-              mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)});
-        }
-      }
-    } else {
-      if (use_spec_verify) {
-        torch::Tensor output = torch::empty_like(conv_input);
-        xllm::kernel::causal_conv1d_out(
-            output,
-            conv_input,
-            conv_weight,
-            conv_cache,
-            std::optional<torch::Tensor>(),
-            torch::IntArrayRef(input_params.parallel.query_start_loc),
-            torch::IntArrayRef(linear_state_indices_host),
-            torch::IntArrayRef(std::vector<int64_t>()),
-            torch::IntArrayRef(num_accepted),
-            xllm::npu::kCausalConv1dActivationSilu,
-            xllm::npu::kCausalConv1dGraphPadSlotId,
-            xllm::npu::kCausalConv1dRunModeUpdate);
-        mixed_qkv = output;
-      } else {
-        auto conv_input_2d = conv_input.dim() == 3
-                                 ? conv_input.reshape({-1, conv_input.size(-1)})
-                                 : conv_input;
-        xllm::kernel::CausalConv1dUpdateParams conv1d_params;
-        conv1d_params.x = conv_input_2d;
-        conv1d_params.conv_state = conv_cache;
-        conv1d_params.weight = conv_weight;
-        conv1d_params.conv_state_indices = logical_state_indices;
-        conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
-        conv1d_params.max_query_len = attn_metadata.max_query_len;
-        mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
-        if (conv_input.dim() == 3) {
-          mixed_qkv =
-              mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)});
-        }
-      }
-    }
-    mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv);
+    xllm::kernel::CausalConv1dUpdateParams conv1d_params;
+    conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)});
+    conv1d_params.conv_state = conv_cache;
+    conv1d_params.weight = conv_weight;
+    conv1d_params.conv_state_indices = linear_state_indices;
+    conv1d_params.block_idx_last_scheduled_token =
+        std::optional<torch::Tensor>();
+    conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
+    conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
+    conv1d_params.max_query_len = attn_metadata.max_query_len;
+    mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
+    // Reshape back to 3D [batch_size, dim, seq_len]
+    mixed_qkv =
+        mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous();
     mixed_qkv = mixed_qkv.transpose(1, 2);
   }
-  const bool use_fused_sigmoid_gdn_decode =
-      fla_ssm_state_layout && !use_spec_verify && !is_any_prefill &&
-      checkpoint_stride == 1;
-  torch::Tensor g;
-  torch::Tensor beta;
+
   // Compute gated delta net decay and beta terms.
-  if (use_spec_verify || attn_metadata.is_chunked_prefill ||
-      checkpoint_stride > 1) {
-    beta = torch::sigmoid(b);
-    torch::Tensor A_log_exp = A_log_.exp();
-    torch::Tensor a_float = a.to(torch::kFloat32);
-    torch::Tensor a_plus_dt = a_float + dt_bias_;
-    torch::Tensor softplus_out = torch::nn::functional::softplus(
-        a_plus_dt,
-        torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0));
-    g = -A_log_exp * softplus_out;
-    g = g.to(a.dtype()).contiguous();
-  } else if (attn_metadata.is_prefill) {
+  if (attn_metadata.is_prefill) {
     xllm::kernel::FusedGdnGatingParams gdn_params;
     gdn_params.A_log = A_log_;
     gdn_params.a = a.contiguous().view({-1, a.size(-1)});
@@ -801,7 +413,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
     std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
     g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)});
     beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)});
-  } else if (!use_fused_sigmoid_gdn_decode) {
+  } else {
     xllm::kernel::FusedGdnGatingParams gdn_params;
     gdn_params.A_log = A_log_;
     gdn_params.a = a.view({-1, a.size(-1)});
@@ -811,216 +423,57 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
     gdn_params.threshold = 20.0f;
     std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
   }
-  if (!used_direct_prefill_qkv) {
-    std::tie(processed_q, processed_k, processed_v) =
-        process_mixed_qkv(mixed_qkv);
-  }
-  torch::Tensor core_attn_out;
-  torch::Tensor last_recurrent_state;
+  auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
   // Apply chunked or recurrent gated-delta attention and update caches.
-  if (use_spec_verify) {
-    torch::Tensor spec_num_accepted_tokens = expand_sequence_tensor_to_batch(
-        input_params.num_accepted_tokens.to(device, torch::kInt32),
-        batch_size,
-        "num_accepted_tokens");
-    torch::Tensor spec_linear_state_base_indices =
-        expand_sequence_tensor_to_batch(
-            linear_state_base_indices, batch_size, "linear_state_base_indices");
-    torch::Tensor step_offsets =
-        torch::arange(seq_len,
-                      torch::TensorOptions()
-                          .dtype(spec_linear_state_base_indices.dtype())
-                          .device(device));
-    torch::Tensor checkpoint_indices =
-        spec_linear_state_base_indices.unsqueeze(1) + step_offsets;
-    double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
-    core_attn_out =
-        run_spec_verify_gated_delta_rule(processed_q,
-                                         processed_k,
-                                         processed_v,
-                                         g,
-                                         beta,
-                                         ssm_cache,
-                                         checkpoint_indices,
-                                         spec_num_accepted_tokens,
-                                         attn_metadata.q_cu_seq_lens,
-                                         attn_metadata.q_seq_lens_vec,
-                                         scale);
-  } else if (is_any_prefill) {
-    CHECK_GE(attn_metadata.q_seq_lens_vec.size(),
-             static_cast<size_t>(batch_size))
-        << "q_seq_lens_vec must be populated for Qwen3.5 prefill.";
-    const bool use_single_prefill_pack =
-        batch_size == 1 && attn_metadata.q_seq_lens_vec.size() == 1 &&
-        attn_metadata.q_seq_lens_vec[0] == seq_len;
-    torch::Tensor packed_processed_q;
-    torch::Tensor packed_processed_k;
-    torch::Tensor packed_processed_v;
-    torch::Tensor packed_g_tensor;
-    torch::Tensor packed_beta_tensor;
-    if (use_single_prefill_pack) {
-      packed_processed_q = processed_q;
-      packed_processed_k = processed_k;
-      packed_processed_v = processed_v;
-      packed_g_tensor = g;
-      packed_beta_tensor = beta;
-    } else {
-      std::vector<torch::Tensor> packed_q;
-      std::vector<torch::Tensor> packed_k;
-      std::vector<torch::Tensor> packed_v;
-      std::vector<torch::Tensor> packed_g;
-      std::vector<torch::Tensor> packed_beta;
-      packed_q.reserve(batch_size);
-      packed_k.reserve(batch_size);
-      packed_v.reserve(batch_size);
-      packed_g.reserve(batch_size);
-      packed_beta.reserve(batch_size);
-      for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
-        const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];
-        if (!used_direct_prefill_qkv) {
-          packed_q.emplace_back(processed_q[batch_idx].narrow(
-              /*dim=*/0, /*start=*/0, valid_len));
-          packed_k.emplace_back(processed_k[batch_idx].narrow(
-              /*dim=*/0, /*start=*/0, valid_len));
-          packed_v.emplace_back(processed_v[batch_idx].narrow(
-              /*dim=*/0, /*start=*/0, valid_len));
-        }
-        packed_g.emplace_back(
-            g[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len));
-        packed_beta.emplace_back(
-            beta[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len));
-      }
-      if (used_direct_prefill_qkv) {
-        packed_processed_q = processed_q;
-        packed_processed_k = processed_k;
-        packed_processed_v = processed_v;
-      } else {
-        packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0);
-        packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0);
-        packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0);
-      }
-      packed_g_tensor = torch::cat(packed_g, 0).unsqueeze(0);
-      packed_beta_tensor = torch::cat(packed_beta, 0).unsqueeze(0);
-    }
-
-    xllm::kernel::MegaChunkGdnParams mega_chunk_gdn_params;
-    mega_chunk_gdn_params.q = packed_processed_q;
-    mega_chunk_gdn_params.k = packed_processed_k;
-    mega_chunk_gdn_params.v = packed_processed_v;
-    mega_chunk_gdn_params.g = packed_g_tensor;
-    mega_chunk_gdn_params.beta = packed_beta_tensor;
+  if (attn_metadata.is_prefill) {
+    xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params;
+    chunk_gated_delta_params.q = processed_q;
+    chunk_gated_delta_params.k = processed_k;
+    chunk_gated_delta_params.v = processed_v;
+    chunk_gated_delta_params.g = g;
+    chunk_gated_delta_params.beta = beta;
     // Get initial state from ssm_cache for sequences with previous state
     // Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
     torch::Tensor initial_state_tensor =
-        torch::index_select(ssm_cache, 0, linear_state_base_indices);
-    CHECK_EQ(input_params.linear_state_validity_mask.size(),
-             input_params.embedding.linear_state_ids.size())
-        << "linear state validity mask must be sequence-scoped.";
-    for (size_t i = 0; i < input_params.linear_state_validity_mask.size();
-         ++i) {
-      if (input_params.linear_state_validity_mask[i] == 0) {
-        initial_state_tensor.select(0, static_cast<int64_t>(i)).fill_(0.0);
-      }
-    }
-    if (!fla_ssm_state_layout && attn_metadata.is_chunked_prefill) {
-      initial_state_tensor =
-          initial_state_tensor.transpose(-1, -2).contiguous();
-    }
-    mega_chunk_gdn_params.initial_state = initial_state_tensor;
-    mega_chunk_gdn_params.output_final_state = true;
-    mega_chunk_gdn_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
-    mega_chunk_gdn_params.q_seq_lens = c10::ArrayRef<int32_t>(
-        attn_metadata.q_seq_lens_vec.data(), static_cast<size_t>(batch_size));
-    mega_chunk_gdn_params.use_qk_l2norm_in_kernel = !used_direct_prefill_qkv;
-    torch::Tensor packed_core_attn_out;
-    std::tie(packed_core_attn_out, last_recurrent_state) =
-        xllm::kernel::mega_chunk_gdn(mega_chunk_gdn_params);
-    if (use_single_prefill_pack) {
-      core_attn_out = packed_core_attn_out;
-      if (core_attn_out.scalar_type() != processed_v.scalar_type()) {
-        core_attn_out = core_attn_out.to(processed_v.scalar_type());
-      }
-    } else {
-      core_attn_out =
-          used_direct_prefill_qkv
-              ? torch::zeros({batch_size, seq_len, local_v_heads, head_v_dim_},
-                             z.options())
-              : torch::zeros_like(processed_v);
-      int64_t packed_offset = 0;
-      for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) {
-        const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];
-        core_attn_out[batch_idx]
-            .narrow(/*dim=*/0, /*start=*/0, valid_len)
-            .copy_(packed_core_attn_out[0].narrow(
-                /*dim=*/0, packed_offset, valid_len));
-        packed_offset += valid_len;
-      }
-    }
-    torch::Tensor state_to_store = fla_ssm_state_layout
-                                       ? last_recurrent_state
-                                       : last_recurrent_state.transpose(-1, -2);
-    ssm_cache.index_put_({linear_state_base_indices},
-                         state_to_store.to(ssm_cache.dtype()));
-  } else if (checkpoint_stride > 1) {
-    auto ssm_state =
-        torch::index_select(ssm_cache, 0, linear_state_base_indices);
-    if (!fla_ssm_state_layout) {
-      ssm_state = ssm_state.transpose(-1, -2);
-    }
-    ssm_state = ssm_state.contiguous();
+        torch::index_select(ssm_cache, 0, linear_state_indices);
+    // Todo: chunked-prefill/prefix-cache use initial_state
+    initial_state_tensor.fill_(0.0);
+    chunk_gated_delta_params.initial_state = initial_state_tensor;
+    chunk_gated_delta_params.output_final_state = true;
+    chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
+    chunk_gated_delta_params.head_first = false;
+    chunk_gated_delta_params.use_qk_l2norm_in_kernel = true;
     std::tie(core_attn_out, last_recurrent_state) =
-        torch_recurrent_gated_delta_rule(
-            processed_q, processed_k, processed_v, g, beta, ssm_state);
-    torch::Tensor state_to_store = fla_ssm_state_layout
-                                       ? last_recurrent_state
-                                       : last_recurrent_state.transpose(-1, -2);
-    ssm_cache.index_put_({linear_state_base_indices},
-                         state_to_store.to(ssm_cache.dtype()));
+        xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
+    ssm_cache.index_put_(
+        {linear_state_indices},
+        last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
   } else {
+    processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
+    processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
+    auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
+    torch::Tensor actual_seq_lengths =
+        torch::cat({zero, attn_metadata.q_seq_lens}, 0);
     double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
-    if (fla_ssm_state_layout) {
-      xllm::kernel::FusedSigmoidGatingDeltaRuleUpdateParams params;
-      params.A_log = A_log_.contiguous();
-      params.a = a.contiguous();
-      params.dt_bias = dt_bias_.contiguous();
-      params.q = processed_q.contiguous();
-      params.k = processed_k.contiguous();
-      params.v = processed_v.contiguous();
-      params.b = b.contiguous();
-      params.initial_state_source = ssm_cache;
-      params.initial_state_indices = linear_state_base_indices.contiguous();
-      params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous();
-      params.scale = static_cast<float>(scale);
-      params.use_qk_l2norm_in_kernel = true;
-      params.softplus_beta = 1.0f;
-      params.softplus_threshold = 20.0f;
-      core_attn_out =
-          xllm::kernel::fused_sigmoid_gating_delta_rule_update(params);
-    } else {
-      processed_q = xllm::kernel::l2_norm(processed_q, /*eps=*/1e-6);
-      processed_k = xllm::kernel::l2_norm(processed_k, /*eps=*/1e-6);
-      auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
-      torch::Tensor actual_seq_lengths =
-          torch::cat({zero, attn_metadata.q_seq_lens}, 0);
-      core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
-                          processed_q.reshape(
-                              {-1, processed_q.size(-2), processed_q.size(-1)}),
-                          processed_k.reshape(
-                              {-1, processed_k.size(-2), processed_k.size(-1)}),
-                          processed_v.reshape(
-                              {-1, processed_v.size(-2), processed_v.size(-1)}),
-                          ssm_cache,
-                          beta.squeeze(0).contiguous(),
-                          scale,
-                          actual_seq_lengths,
-                          logical_state_indices,
-                          c10::nullopt,
-                          g.squeeze(0).contiguous(),
-                          c10::nullopt)
-                          .unsqueeze(0)
-                          .contiguous();
-    }
+    core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
+                        processed_q.reshape(
+                            {-1, processed_q.size(-2), processed_q.size(-1)}),
+                        processed_k.reshape(
+                            {-1, processed_k.size(-2), processed_k.size(-1)}),
+                        processed_v.reshape(
+                            {-1, processed_v.size(-2), processed_v.size(-1)}),
+                        ssm_cache,
+                        beta.squeeze(0).contiguous(),
+                        scale,
+                        actual_seq_lengths,
+                        linear_state_indices,
+                        c10::nullopt,
+                        g.squeeze(0).contiguous(),
+                        c10::nullopt)
+                        .unsqueeze(0)
+                        .contiguous();
   }
+
   auto z_reshaped = z.view({-1, z.size(-1)});
   auto core_attn_out_reshaped =
       core_attn_out.view({-1, core_attn_out.size(-1)});
@@ -1033,47 +486,25 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
   auto rearranged_norm =
       norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)});
   rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm);
-  // For chunked prefill or spec verify, reshape_projected_tokens_with_pad may
-  // pad each batch to max_len, causing output tokens > original_num_tokens. We
-  // need to slice back to original_num_tokens to match the residual shape.
-  if (rearranged_norm.size(0) > original_num_tokens) {
-    // Slice excess padding tokens
-    rearranged_norm =
-        rearranged_norm.slice(0, 0, original_num_tokens).contiguous();
-  }
-  if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
-    return o_proj_->forward(rearranged_norm,
-                            row_parallel_reduce_mode_for_fc1(*fc1_ctx));
-  }
-  return o_proj_->forward(rearranged_norm);
+  auto attn_output = o_proj_->forward(rearranged_norm);
+  return attn_output;
 }
 
 torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
     const AttentionMetadata& attn_metadata,
     const torch::Tensor& padded_qkvz) const {
-  const bool has_padded_queries =
-      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
-  if (!has_padded_queries) {
+  if (!attn_metadata.is_prefill) {
     return padded_qkvz;
   }
   std::vector<torch::Tensor> valid_batches;
-  const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
-  int64_t bs = has_host_lens
-                   ? static_cast<int64_t>(attn_metadata.q_seq_lens_vec.size())
-                   : attn_metadata.q_seq_lens.size(0);
-  valid_batches.reserve(bs);
+  int64_t bs = attn_metadata.q_seq_lens.size(0);
   int64_t max_len = attn_metadata.max_query_len;
   const auto& ori_seq_lens = attn_metadata.q_seq_lens;
   auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1});
   for (int64_t b = 0; b < bs; ++b) {
-    int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b]
-                                    : ori_seq_lens[b].template item<int64_t>();
-    torch::Tensor valid_batch =
-        reshaped_qkvz[b].slice(/*dim=*/0, /*start=*/0, ori_len);
-    valid_batches.emplace_back(valid_batch);
-  }
-  if (valid_batches.size() == 1) {
-    return valid_batches[0].contiguous();
+    int64_t ori_len = ori_seq_lens[b].template item<int64_t>();
+    torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len);
+    valid_batches.push_back(valid_batch);
   }
   return torch::cat(valid_batches, 0).contiguous();
 }
@@ -1081,60 +512,41 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
 torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices(
     const ModelInputParams& input_params,
     const torch::Device& device) const {
-  CHECK(!input_params.embedding.linear_state_ids.empty())
+  CHECK(!input_params.linear_state_ids.empty())
       << "linear_state_ids must be populated for gated delta net";
-  if (input_params.embedding.linear_state_indices.defined()) {
-    auto indices = input_params.embedding.linear_state_indices;
-    if (indices.device() != device || indices.scalar_type() != torch::kInt) {
-      indices =
-          indices.to(torch::TensorOptions().dtype(torch::kInt).device(device),
-                     /*non_blocking=*/true,
-                     /*copy=*/true);
-    }
-    return indices.contiguous();
+  if (input_params.linear_state_indices.defined()) {
+    return input_params.linear_state_indices;
   }
   return torch::tensor(
-      input_params.embedding.linear_state_ids,
+      input_params.linear_state_ids,
       torch::TensorOptions().dtype(torch::kInt).device(device));
 }
 
-torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad(
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
     const AttentionMetadata& attn_metadata,
-    const torch::Tensor& projected_tokens) const {
-  const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty();
-  int64_t bs = has_host_lens
-                   ? static_cast<int64_t>(attn_metadata.q_seq_lens_vec.size())
-                   : attn_metadata.q_seq_lens.size(0);
+    const torch::Tensor& qkvz) const {
+  int64_t bs = attn_metadata.q_seq_lens.size(0);
   int64_t max_len = attn_metadata.max_query_len;
   const auto& start_loc = attn_metadata.q_seq_lens;
-  const bool need_padding =
-      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
-  if (!need_padding) {
-    return projected_tokens.view({bs, -1, projected_tokens.size(-1)});
-  }
-  if (has_host_lens && bs == 1 && attn_metadata.q_seq_lens_vec[0] == max_len &&
-      projected_tokens.dim() == 2 && projected_tokens.size(0) == max_len) {
-    return projected_tokens.view({1, max_len, projected_tokens.size(-1)});
+  if (!attn_metadata.is_prefill) {
+    return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
   }
   std::vector<torch::Tensor> batches;
-  batches.reserve(bs);
   int64_t idx = 0;
   for (int64_t b = 0; b < bs; ++b) {
-    int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b]
-                                    : start_loc[b].template item<int64_t>();
-    torch::Tensor batch =
-        projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous();
+    int64_t cur_len = start_loc[b].template item<int64_t>();
+    torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
     idx = idx + cur_len;
     if (batch.size(0) != max_len) {
       batch = batch.size(0) > max_len
-                  ? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous()
+                  ? batch.slice(0, 0, max_len).contiguous()
                   : torch::nn::functional::pad(
                         batch,
                         torch::nn::functional::PadFuncOptions(
                             {0, 0, 0, max_len - batch.size(0)}))
                         .contiguous();
     }
-    batches.emplace_back(batch);
+    batches.push_back(batch);
   }
   auto ret = torch::stack(batches, 0).contiguous();
   return ret;

ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h

@@ -1,4 +1,4 @@
-/* Copyright 2025-2026 The xLLM Authors.
+/* Copyright 2026 The xLLM Authors. 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.
@@ -17,7 +17,6 @@ limitations under the License.
 
 #include <torch/torch.h>
 
-#include <optional>
 #include <string>
 #include <tuple>
 #include <utility>
@@ -52,40 +51,19 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
                         const ModelInputParams& input_params);
 
  protected:
-  virtual std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
-      const torch::Tensor& hidden_states) = 0;
-  virtual std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
-      const torch::Tensor& hidden_states) = 0;
-  // Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a
-  // weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns
-  // nullopt to select the fused-split fallback.
-  virtual std::optional<
-      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
-  project_split_inputs(const torch::Tensor& hidden_states,
-                       const AttentionMetadata& attn_metadata) {
-    return std::nullopt;
-  }
-  virtual bool use_fla_ssm_state_layout() const { return false; }
+  virtual std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
+      const torch::Tensor& hidden_states,
+      const AttentionMetadata& attn_metadata) = 0;
 
   void load_common_state_dict(const StateDict& state_dict);
   void verify_common_loaded_weights(const std::string& prefix) const;
 
-  torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
-                                         const torch::Device& device) const;
-
-  std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
-      const torch::Tensor& hidden_states,
-      const AttentionMetadata& attn_metadata);
-
+  torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata,
+                                      const torch::Tensor& qkvz) const;
   torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
                                    const torch::Tensor& padded_qkvz) const;
-
-  // Projection outputs are packed as [total_tokens, dim], while GDN kernels
-  // consume dense [batch, max_query_len, dim] tensors. Split the packed tokens
-  // by query length and pad each sequence before entering the kernels.
-  torch::Tensor reshape_projected_tokens_with_pad(
-      const AttentionMetadata& attn_metadata,
-      const torch::Tensor& projected_tokens) const;
+  torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
+                                         const torch::Device& device) const;
 
   std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
       torch::Tensor& mixed_qkv) const;