feat(moe): EP experts forward收纳到kernel底下实现多impls管理，并且添加 NPU grouped matmul的forward分支。

src/twinkle/kernel/ops/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Backend-agnostic operator interfaces with per-platform implementations.
+
+Each op family defines an abstract base class (e.g. ``EpExpertsGmm``) plus a
+dispatcher that picks the first eligible backend implementation from a
+lazily-built registry. Callers invoke the dispatcher only; backend details
+(NPU, GPU, ...) stay hidden behind the interface.
+"""
+from .ep import EpExpertsGmm, ep_forward
+
+__all__ = [
+    'EpExpertsGmm',
+    'ep_forward',
+]

src/twinkle/kernel/ops/ep/__init__.py

@@ -0,0 +1,112 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Grouped-matmul interface for expert-parallel (EP) experts compute.
+
+``EpExpertsGmm`` is the abstract backend interface; each platform provides a
+subclass (e.g. ``NpuEpExpertsGmm`` in ``ep_experts_gmm_npu.py``) that is
+registered lazily. The generic per-expert loop (``LoopEpExpertsGmm`` in
+``ep_experts_gmm_loop.py``) is registered last as the always-eligible
+fallback. ``ep_forward`` is the single entry point used by the EP
+forward: it dispatches to the first eligible backend, so it always returns a
+result.
+"""
+from __future__ import annotations
+
+import torch
+from abc import ABC, abstractmethod
+from torch import nn
+
+from twinkle import get_logger
+
+logger = get_logger()
+
+
+class EpExpertsGmm(ABC):
+    """Backend interface for batched EP experts compute via grouped matmul.
+
+    A backend computes all local experts in one shot from the permuted token
+    buffer (tokens sorted by expert, contiguous chunks per local expert),
+    replacing the per-expert Python loop.
+    """
+
+    name: str
+    # Fallback backends (e.g. the per-expert loop) are always eligible and are
+    # registered last; hitting one means every accelerated backend declined.
+    fallback: bool = False
+
+    @abstractmethod
+    def ineligible_reason(self, experts_mod: nn.Module) -> str | None:
+        """Return why this backend cannot handle ``experts_mod``, or None if it can."""
+
+    @abstractmethod
+    def forward(
+        self,
+        experts_mod: nn.Module,
+        permuted_tokens: torch.Tensor,
+        num_global_sum_tokens_per_local_expert: torch.Tensor,
+        experts_per_rank: int,
+    ) -> torch.Tensor:
+        """Compute local experts on the permuted token buffer."""
+
+
+_IMPLS: list[EpExpertsGmm] | None = None
+_PATH_LOGGED = False
+_WARN_LOGGED = False
+
+
+def _get_impls() -> list[EpExpertsGmm]:
+    """Build the backend registry on first use.
+
+    Each backend module is imported defensively: platforms lacking its
+    dependencies (e.g. no torch_npu) simply skip that backend.
+    """
+    global _IMPLS
+    if _IMPLS is None:
+        _IMPLS = []
+        try:
+            from .npu import NpuEpExpertsGmm
+            _IMPLS.append(NpuEpExpertsGmm())
+        except ImportError:
+            pass
+        # The loop fallback has no optional dependencies and is always last:
+        # it guarantees the dispatcher never runs out of backends.
+        from .loop import LoopEpExpertsGmm
+        _IMPLS.append(LoopEpExpertsGmm())
+    return _IMPLS
+
+
+def ep_forward(
+    experts_mod: nn.Module,
+    permuted_tokens: torch.Tensor,
+    num_global_sum_tokens_per_local_expert: torch.Tensor,
+    experts_per_rank: int,
+) -> torch.Tensor:
+    """Dispatch to the first eligible backend (the loop fallback guarantees a hit).
+
+    One-time logging: INFO on the first accelerated dispatch (process-wide);
+    WARNING with the ineligibility reasons when falling back to the loop
+    (once per experts instance, as reasons may differ per MoE block).
+    """
+    if permuted_tokens.numel() == 0:
+        # Preserve the autograd edge to token_pre_all2all. Returning a new
+        # empty tensor can make this rank skip the matching backward
+        # all-to-all, causing EP collective order divergence.
+        return permuted_tokens
+
+    global _PATH_LOGGED, _WARN_LOGGED
+    ineligible: list[str] = []
+    for impl in _get_impls():
+        reason = impl.ineligible_reason(experts_mod)
+        if reason is not None:
+            ineligible.append(f'{impl.name}: {reason}')
+            continue
+        if impl.fallback:
+            if ineligible and not _WARN_LOGGED:
+                detail = '; '.join(ineligible)
+                logger.warning(f'EP experts compute: grouped matmul disabled ({detail}); '
+                               f'falling back to {impl.name}.')
+                _WARN_LOGGED = True
+        elif not _PATH_LOGGED:
+            logger.info(f'EP experts compute: using {impl.name} grouped matmul (per-expert loop replaced).')
+            _PATH_LOGGED = True
+        return impl.forward(experts_mod, permuted_tokens, num_global_sum_tokens_per_local_expert, experts_per_rank)
+    raise RuntimeError('no EP experts backend registered (loop fallback missing)')

src/twinkle/kernel/ops/ep/loop.py

@@ -0,0 +1,66 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Generic per-expert Python loop backend for the EP experts interface.
+
+This is the always-eligible fallback: it runs on any platform (pure PyTorch)
+and is registered last, so the dispatcher can never run out of backends.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from . import EpExpertsGmm
+
+
+class LoopEpExpertsGmm(EpExpertsGmm):
+    """Per-expert Python loop over the permuted token buffer (generic fallback)."""
+
+    name = 'per-expert loop'
+    fallback = True
+
+    def ineligible_reason(self, experts_mod: nn.Module) -> str | None:
+        # The loop is the reference implementation: always eligible.
+        return None
+
+    def forward(
+        self,
+        experts_mod: nn.Module,
+        permuted_tokens: torch.Tensor,
+        num_global_sum_tokens_per_local_expert: torch.Tensor,
+        experts_per_rank: int,
+    ) -> torch.Tensor:
+        input_dtype = permuted_tokens.dtype
+
+        cumsum = torch.zeros(experts_per_rank + 1, dtype=torch.long)
+        for i in range(experts_per_rank):
+            cumsum[i + 1] = cumsum[i] + int(num_global_sum_tokens_per_local_expert[i].item())
+
+        output_chunks = []
+        for i in range(experts_per_rank):
+            start = int(cumsum[i].item())
+            end = int(cumsum[i + 1].item())
+            expert_in = permuted_tokens[start:end]
+            if expert_in.numel() == 0:
+                output_chunks.append(expert_in)
+                continue
+
+            gate_up = experts_mod.gate_up_proj[i]
+            down = experts_mod.down_proj[i]
+            compute_dtype = gate_up.dtype
+            if expert_in.dtype != compute_dtype:
+                expert_in = expert_in.to(compute_dtype)
+            gate_up_out = F.linear(expert_in, gate_up)
+            if hasattr(experts_mod, '_apply_gate'):
+                out = experts_mod._apply_gate(gate_up_out)
+            else:
+                gate, up = gate_up_out.chunk(2, dim=-1)
+                out = experts_mod.act_fn(gate) * up
+            out = F.linear(out, down)
+
+            if out.dtype != input_dtype:
+                out = out.to(input_dtype)
+            output_chunks.append(out)
+
+        return torch.cat(
+            output_chunks, dim=0) if output_chunks else permuted_tokens.new_empty(0, permuted_tokens.size(-1))

src/twinkle/kernel/ops/ep/npu.py

@@ -0,0 +1,103 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""NPU backend for the EP experts grouped-matmul interface."""
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from . import EpExpertsGmm
+
+
+class NpuEpExpertsGmm(EpExpertsGmm):
+    """Batched EP experts compute via ``torch_npu.npu_grouped_matmul``."""
+
+    name = 'NPU'
+
+    def ineligible_reason(self, experts_mod: nn.Module) -> str | None:
+        """Check whether the tensor-experts EP compute can use NPU grouped matmul.
+
+        Falls back to the per-expert Python loop when:
+        - torch_npu / NPU is unavailable;
+        - there is no way to compute the gated activation: no ``_apply_gate``
+          hook and a non-SiLU ``act_fn`` (``npu_swiglu`` == silu(gate) * up);
+        - the packed weights are not plain 3D floating tensors (e.g. DTensor
+          that was not unsharded, quantized weights, etc.).
+
+        Note on ``_apply_gate``: transformers injects a default implementation
+        (``_default_apply_gate``: chunk + act_fn(gate) * up) onto many MoE
+        expert classes. With SiLU activation that is exactly ``npu_swiglu``, so
+        it does not block the fused path. A *custom* gate hook is still
+        supported — it is called eagerly on the grouped-matmul output.
+        """
+        apply_gate = getattr(experts_mod, '_apply_gate', None)
+        is_default_gate = getattr(apply_gate, '__name__', None) == '_default_apply_gate'
+        if apply_gate is None or is_default_gate:
+            # Activation will be computed by npu_swiglu == silu(gate) * up.
+            act_fn = getattr(experts_mod, 'act_fn', None)
+            if act_fn is None:
+                return 'no act_fn attribute'
+            # nn.SiLU and transformers' SiLUActivation both compute F.silu
+            # exactly; either is equivalent to the silu inside npu_swiglu.
+            if act_fn.__class__.__name__ not in ('SiLU', 'SiLUActivation'):
+                return f'act_fn is {act_fn.__class__.__name__}, not SiLU'
+        # else: custom gate hook is called eagerly in the GMM path — allowed.
+        try:
+            import torch_npu  # noqa: F401
+        except ImportError:
+            return 'torch_npu not importable'
+        if not torch.npu.is_available():
+            return 'NPU not available'
+        for name in ('gate_up_proj', 'down_proj'):
+            weight = getattr(experts_mod, name, None)
+            # Note: accept plain tensors and nn.Parameter. Under FSDP2 the
+            # pre-forward hook has already unsharded params by the time
+            # ep_forward runs; under PEFT target_parameters the parametrized
+            # property returns a merged plain tensor. Anything else (DTensor,
+            # quantized types) falls back to the loop.
+            if not isinstance(weight, torch.Tensor) or weight.__class__.__name__ == 'DTensor':
+                return f'{name} is {type(weight).__name__}, not a plain tensor'
+            if weight.ndim != 3 or not weight.dtype.is_floating_point:
+                return f'{name} has shape {tuple(weight.shape)} dtype {weight.dtype}'
+        return None
+
+    def forward(
+        self,
+        experts_mod: nn.Module,
+        permuted_tokens: torch.Tensor,
+        num_global_sum_tokens_per_local_expert: torch.Tensor,
+        experts_per_rank: int,
+    ) -> torch.Tensor:
+        """Compute local experts with one grouped matmul instead of a Python loop.
+
+        Mathematically equivalent to the per-expert loop: grouped matmul
+        applies each expert's weight to its contiguous token chunk, and
+        ``npu_swiglu(x) == silu(gate) * up`` for gate/up = x.chunk(2, -1).
+        The group list stays on device, eliminating the per-expert ``.item()``
+        host synchronizations of the loop implementation.
+        """
+        import torch_npu
+
+        from twinkle.kernel.npu_impls.moe import GmmFunction, _get_cached_expert_weights
+
+        input_dtype = permuted_tokens.dtype
+        hidden_dim = permuted_tokens.size(-1)
+        # counts arrives as a CPU tensor from preprocess(); grouped matmul needs
+        # the group list on device.
+        group_list = num_global_sum_tokens_per_local_expert.to(device=permuted_tokens.device, dtype=torch.int64)
+        gate_up_weight, down_weight = _get_cached_expert_weights(experts_mod, input_dtype, hidden_dim)
+
+        expert_in = permuted_tokens
+        if expert_in.dtype != gate_up_weight.dtype:
+            expert_in = expert_in.to(gate_up_weight.dtype)
+        intermediate = GmmFunction.apply(expert_in, group_list, gate_up_weight)
+        apply_gate = getattr(experts_mod, '_apply_gate', None)
+        if apply_gate is not None and getattr(apply_gate, '__name__', '') != '_default_apply_gate':
+            # Custom gate hook: honor it eagerly instead of the fused swiglu.
+            activated = apply_gate(intermediate)
+        else:
+            # No hook, or transformers' default gate: silu(gate) * up == npu_swiglu.
+            activated = torch_npu.npu_swiglu(intermediate, dim=-1)
+        out = GmmFunction.apply(activated, group_list, down_weight)
+        if out.dtype != input_dtype:
+            out = out.to(input_dtype)
+        return out

src/twinkle/model/transformers/moe/ep_utils.py

@@ -8,6 +8,8 @@
 import torch.distributed as dist
 from typing import Optional
 
+from twinkle import torch_util
+
 
 # ========================== comm ==========================
 class _AllToAll(torch.autograd.Function):
@@ -113,7 +115,11 @@ def permute(tokens: torch.Tensor, expert_mask: torch.Tensor):
     sorted_indices = token_indices.masked_select(expert_mask)
 
     # use the mapping to permute the tokens
-    permuted_input = tokens.index_select(0, sorted_indices)
+    # NOTE: use advanced indexing instead of index_select — index_select's
+    # backward (index_add) is broken on some torch_npu/CANN versions
+    # (aclnnIndexAdd fails for any dtype), while x[idx] backward (index_put)
+    # works. Mathematically identical.
+    permuted_input = tokens[sorted_indices]
 
     return permuted_input, sorted_indices
 
@@ -203,6 +209,12 @@ def preprocess(
     )
     dist.all_gather_into_tensor(num_global_tokens_per_expert, num_local_tokens_per_expert, group=ep_group)
 
+    # The collective may run on a separate stream (e.g. HCCL on NPU) whose
+    # completion is not always ordered with the host reads (.tolist()/.item())
+    # below — observed in practice as garbage split sizes. Force a device sync
+    # before reading gathered results back on the host.
+    torch_util.synchronize()
+
     # [ep_size, num_local_experts]
     start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts
     num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous()
@@ -211,11 +223,13 @@ def preprocess(
     output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist()
 
     # [num_local_expert]
-    num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(
-        torch.device('cpu'), non_blocking=True)
+    # NOTE: keep these copies synchronous. With non_blocking=True the CPU
+    # tensors are read (tolist()/item()) before the async D2H copy lands,
+    # producing garbage split sizes (observed on NPU).
+    num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(torch.device('cpu'))
 
     num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to(
-        torch.device('cpu'), non_blocking=True)
+        torch.device('cpu'))
 
     return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert