feat: port NaiveBatchedExperts from ds_vllm — view transpose + cublas transB

Source: 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

New files (ported from ds_vllm, adapted for BI-V100):
  ex_engine/moe/__init__.py
  ex_engine/moe/activation.py
    - MoEActivation enum + apply_moe_activation
    - torch.ops._C.silu_and_mul replaced with F.silu(gate)*up fallback
  ex_engine/moe/naive_batched_experts.py
    - naive_batched_moe_forward()
    - Decode: per-expert loop, w13[eid].transpose(0,1) is VIEW (zero copy)
    - @ operator → cublas passes transB=CUBLAS_OP_T internally
    - Prefill: group tokens by expert, batch @ per expert

Modified:
  qwen3_6_scripts/qwen3_5.py
    - Import naive_batched_moe_forward
    - Tier 0.5: after ix_fused_moe, before corex point-optimized loop
    - Uses existing topk routing (xllm/corex/pytorch)

Key difference from previous approach:
  - NO physical transpose (was 22ms overhead)
  - NO weight gather into contiguous buffer
  - View transpose is O(0), cublas handles transB

ex_engine/moe/__init__.py

@@ -0,0 +1,10 @@
+"""
+ex_engine.moe — MoE expert computation for BI-V100
+
+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 ex_engine.moe.naive_batched_experts import naive_batched_moe_forward
+from ex_engine.moe.activation import MoEActivation, apply_moe_activation

ex_engine/moe/activation.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MoE activation function enum and utilities."""
+
+from enum import Enum
+
+import torch
+import torch.nn.functional as F
+
+
+class MoEActivation(Enum):
+    """Activation functions for MoE layers."""
+
+    # Gated activations (gate * activation(up)) expect input of shape [..., 2*d]
+    # and produce output of shape [..., d]
+    SILU = "silu"
+    GELU = "gelu"
+    GELU_TANH = "gelu_tanh"
+    RELU2 = "relu2"
+    SWIGLUOAI = "swigluoai"
+    SWIGLUSTEP = "swiglustep"
+
+    # Non-gated activations (no mul with gate) expect input of shape [..., d]
+    # and produce output of shape [..., d].
+    # NOTE: Non-gated activations require the "_no_mul" suffix to be present.
+    SILU_NO_MUL = "silu_no_mul"
+    GELU_NO_MUL = "gelu_no_mul"
+    GELU_TANH_NO_MUL = "gelu_tanh_no_mul"
+    RELU2_NO_MUL = "relu2_no_mul"
+
+    @property
+    def is_gated(self) -> bool:
+        """Returns True if activation expects gate*activation(up) pattern.
+
+        Gated activations expect input tensor with 2x the output size,
+        where the first half is the gate and second half is the up projection.
+        """
+        return not self.value.endswith("_no_mul")
+
+    @property
+    def custom_op_name(self) -> str:
+        """Maps to the CustomOp name of activations
+        in vllm/model_executor/layers/activation.py."""
+        return _CUSTOM_OP_NAMES[self]
+
+    def without_mul(self) -> "MoEActivation":
+        """Get the non-gated variant of this activation.
+
+        For activations that have a _no_mul variant, returns that variant.
+        For activations without a _no_mul variant (or already _no_mul),
+        returns self.
+        """
+        return _WITHOUT_MUL.get(self, self)
+
+    @classmethod
+    def from_str(cls, s: str) -> "MoEActivation":
+        """Parse from string for backward compatibility."""
+        s = _STR_ALIASES.get(s, s)
+        for member in cls:
+            if member.value == s:
+                return member
+        valid = [m.value for m in cls]
+        raise ValueError(f"Unknown MoE activation: {s!r}. Valid activations: {valid}")
+
+
+# Module-level lookup tables used by MoEActivation functions.
+_STR_ALIASES: dict[str, str] = {
+    "gelu_pytorch_tanh": "gelu_tanh",
+}
+
+_CUSTOM_OP_NAMES: dict[MoEActivation, str] = {
+    MoEActivation.SILU: "silu_and_mul",
+    MoEActivation.GELU: "gelu_and_mul",
+    MoEActivation.GELU_TANH: "gelu_tanh_and_mul",
+    MoEActivation.SWIGLUOAI: "swigluoai_and_mul",
+    MoEActivation.SWIGLUSTEP: "swiglustep_and_mul",
+    MoEActivation.RELU2: "relu2",
+    MoEActivation.SILU_NO_MUL: "silu_and_mul",
+    MoEActivation.GELU_NO_MUL: "gelu_and_mul",
+    MoEActivation.GELU_TANH_NO_MUL: "gelu_tanh_and_mul",
+    MoEActivation.RELU2_NO_MUL: "relu2",
+}
+
+_WITHOUT_MUL: dict[MoEActivation, MoEActivation] = {
+    MoEActivation.SILU: MoEActivation.SILU_NO_MUL,
+    MoEActivation.GELU: MoEActivation.GELU_NO_MUL,
+    MoEActivation.GELU_TANH: MoEActivation.GELU_TANH_NO_MUL,
+    MoEActivation.RELU2: MoEActivation.RELU2_NO_MUL,
+}
+
+
+def activation_without_mul(activation: str) -> str:
+    """Get the non-gated variant of an activation function.
+
+    Args:
+        activation: The activation function name (e.g., "silu", "gelu")
+
+    Returns:
+        The non-gated activation name (e.g., "silu_no_mul", "gelu_no_mul")
+    """
+    return MoEActivation.from_str(activation).without_mul().value
+
+
+def apply_moe_activation(
+    activation: MoEActivation,
+    output: torch.Tensor,
+    input: torch.Tensor,
+) -> torch.Tensor:
+    """Apply MoE activation function."""
+    assert input.dim() == 2, "Input must be 2D"
+    assert output.dim() == 2, "Output must be 2D"
+    if activation.is_gated:
+        assert output.size(-1) * 2 == input.size(-1), (
+            f"{activation.value} expects 2x ratio: "
+            f"{output.size(-1) * 2} vs {input.size(-1)}"
+        )
+    else:
+        assert output.size(-1) == input.size(-1), (
+            f"{activation.value} expects equal sizes: "
+            f"{output.size(-1)} vs {input.size(-1)}"
+        )
+
+    # 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)
+    elif activation == MoEActivation.GELU:
+        d = output.size(-1)
+        gate = input[..., :d]
+        up = input[..., d:]
+        output.copy_(F.gelu(gate) * up)
+    elif activation == MoEActivation.GELU_TANH:
+        d = output.size(-1)
+        gate = input[..., :d]
+        up = input[..., d:]
+        output.copy_(F.gelu(gate, approximate="tanh") * up)
+    elif activation == MoEActivation.SWIGLUOAI:
+        d = output.size(-1)
+        gate = input[..., :d]
+        up = input[..., d:]
+        output.copy_(F.silu(gate) * up)
+    elif activation == MoEActivation.SWIGLUSTEP:
+        d = output.size(-1)
+        gate = input[..., :d]
+        up = input[..., d:]
+        output.copy_(F.silu(gate) * up)
+
+    # Activations without gated multiplication
+    elif activation == MoEActivation.SILU_NO_MUL:
+        output.copy_(F.silu(input))
+    elif activation == MoEActivation.GELU_NO_MUL:
+        output.copy_(F.gelu(input))
+    elif activation == MoEActivation.GELU_TANH_NO_MUL:
+        output.copy_(F.gelu(input, approximate="tanh"))
+    elif activation == MoEActivation.RELU2_NO_MUL:
+        F.relu(input, inplace=True)
+        torch.square(input, out=output)
+    else:
+        raise ValueError(f"Unsupported FusedMoe activation: {activation}")
+
+    return output

ex_engine/moe/naive_batched_experts.py

@@ -0,0 +1,134 @@
+"""
+naive_batched_experts.py — MoE expert computation for BI-V100
+
+Ported from:
+  upstream_ref/ds_vllm/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+  class NaiveBatchedExperts.apply()
+
+Key design from upstream:
+  - w1[expert].transpose(0, 1) is a VIEW (zero copy)
+  - @ operator lets cublas pass transB=CUBLAS_OP_T internally
+  - No physical transpose, no gather of full weight matrices
+  - Per-expert loop with early exit on num_tokens == 0
+
+Adaptations for BI-V100:
+  - Removed modular_kernel / FusedMoEExpertsModular base class
+  - Removed triton kernels (BatchedTritonExperts)
+  - Removed quantization (FP8, INT8, INT4)
+  - Removed workspace_shapes / MoEActivation enum dependency
+  - activation uses F.silu directly (torch.ops._C.silu_and_mul not available)
+  - Standalone function, not a class — called from qwen3_5.py
+"""
+
+import torch
+import torch.nn.functional as F
+from typing import Optional
+
+
+def _resize_cache(x: torch.Tensor, v: tuple) -> torch.Tensor:
+    """Shrink tensor and reshape. From ds_vllm utils.py."""
+    from math import prod
+    assert prod(v) <= x.numel(), f"{v} ({prod(v)}) <= {x.shape} ({x.numel()})"
+    return x.flatten()[:prod(v)].view(*v)
+
+
+def naive_batched_moe_forward(
+    hidden_states: torch.Tensor,     # (T, H) or (1, H) for decode
+    w13: torch.Tensor,               # (E, 2*I, H) — gate+up fused weights
+    w2: torch.Tensor,                # (E, H, I)   — down weights
+    topk_ids: torch.Tensor,          # (T, top_k) — selected expert ids
+    topk_weights: torch.Tensor,      # (T, top_k) — routing weights
+    act_fn: Optional[object] = None, # SiluAndMul instance or None
+) -> torch.Tensor:
+    """
+    MoE expert forward — ported from NaiveBatchedExperts.apply().
+
+    For each selected expert:
+      1. FC1: input @ w1[expert].transpose(0, 1)  — view transpose, cublas transB
+      2. Activation: silu_and_mul (gated)
+      3. FC2: act @ w2[expert].transpose(0, 1)
+
+    Source: upstream_ref/ds_vllm/.../experts/fused_batched_moe.py lines 611-647
+    """
+    T = hidden_states.shape[0]
+    H = hidden_states.shape[1]
+    I = w2.shape[2]          # intermediate size (per partition)
+    top_k = topk_ids.shape[1]
+
+    # Output accumulator
+    out = torch.zeros(T, H, dtype=hidden_states.dtype, device=hidden_states.device)
+
+    if T == 1:
+        # === Decode path (single token) ===
+        # From NaiveBatchedExperts.apply():
+        #   input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
+        #
+        # For decode, each expert sees exactly 1 token.
+        # expert ids are in topk_ids[0] (shape: top_k,)
+        eids = topk_ids[0]        # (top_k,)
+        ws = topk_weights[0]      # (top_k,)
+
+        for i in range(top_k):
+            eid = eids[i].item()
+
+            # FC1: (1, H) @ (H, 2*I) → (1, 2*I)
+            # w13[eid] is (2*I, H), .transpose(0, 1) is (H, 2*I) — VIEW, zero copy
+            # @ lets cublas use transB=CUBLAS_OP_T
+            gate_up = hidden_states @ w13[eid].transpose(0, 1)  # (1, 2*I)
+
+            # Activation: silu_and_mul
+            # From upstream apply_moe_activation():
+            #   gate = input[..., :d], up = input[..., d:]
+            #   output = F.silu(gate) * up
+            if act_fn is not None:
+                act = act_fn(gate_up)     # SiluAndMul: (1, 2*I) → (1, I)
+            else:
+                gate = gate_up[..., :I]
+                up = gate_up[..., I:]
+                act = F.silu(gate) * up   # (1, I)
+
+            # FC2: (1, I) @ (I, H) → (1, H)
+            # w2[eid] is (H, I), .transpose(0, 1) is (I, H) — VIEW, zero copy
+            expert_out = act @ w2[eid].transpose(0, 1)  # (1, H)
+
+            # Weighted accumulate
+            out += ws[i] * expert_out
+
+    else:
+        # === Prefill path (multiple tokens) ===
+        # Group tokens by expert, then batch-process each expert.
+        # From NaiveBatchedExperts.apply() — the for-expert loop.
+        flat_eids = topk_ids.reshape(-1)                 # (T * top_k,)
+        flat_weights = topk_weights.reshape(-1)           # (T * top_k,)
+        flat_token_ids = torch.arange(
+            T, device=hidden_states.device
+        ).repeat_interleave(top_k)                        # (T * top_k,)
+
+        num_experts = w13.shape[0]
+        for expert in range(num_experts):
+            mask = (flat_eids == expert)
+            if not mask.any():
+                continue
+
+            token_ids = flat_token_ids[mask]               # tokens assigned to this expert
+            weights = flat_weights[mask]                   # their routing weights
+            expert_input = hidden_states[token_ids]        # (num, H)
+
+            # FC1: (num, H) @ (H, 2*I) → (num, 2*I)
+            gate_up = expert_input @ w13[expert].transpose(0, 1)
+
+            # Activation
+            if act_fn is not None:
+                act = act_fn(gate_up)
+            else:
+                gate = gate_up[..., :I]
+                up = gate_up[..., I:]
+                act = F.silu(gate) * up
+
+            # FC2: (num, I) @ (I, H) → (num, H)
+            expert_out = act @ w2[expert].transpose(0, 1)
+
+            # Weighted scatter-add back
+            out.index_add_(0, token_ids, expert_out * weights.unsqueeze(1))
+
+    return out

qwen3_6_scripts/qwen3_5.py

@@ -245,6 +245,24 @@ if _USE_IX_FUSED_MOE:
 else:
     logger.info("ix_fused_moe unavailable — using point-optimized Python MoE")
 
+# naive_batched_moe_forward: ported from ds_vllm NaiveBatchedExperts
+# Uses view transpose + @ operator (cublas transB), no physical transpose
+try:
+    from ex_engine.moe.naive_batched_experts import naive_batched_moe_forward
+    _HAS_NAIVE_BATCHED_MOE = True
+except ImportError:
+    try:
+        import sys, os
+        sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+        from ex_engine.moe.naive_batched_experts import naive_batched_moe_forward
+        _HAS_NAIVE_BATCHED_MOE = True
+    except ImportError:
+        _HAS_NAIVE_BATCHED_MOE = False
+        naive_batched_moe_forward = None
+_USE_NAIVE_BATCHED_MOE = (
+    _HAS_NAIVE_BATCHED_MOE
+    and env_bool("BI100_MOE_NAIVE_BATCHED", True))
+
 
 # ---------------------------------------------------------------------------
 # Qwen3.6 vision tower and vLLM 0.6 multimodal input integration
@@ -1686,6 +1704,38 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                 self.top_k, w13.shape[0],
                 True)  # renormalize
 
+        # ---------------------------------------------------------------
+        # Tier 0.5: NaiveBatchedExperts from ds_vllm
+        # Per-expert loop with view transpose + @ (cublas transB)
+        # No physical transpose, no weight gather copy
+        # Source: ds_vllm/vllm/.../experts/fused_batched_moe.py
+        # ---------------------------------------------------------------
+        if _USE_NAIVE_BATCHED_MOE:
+            w13 = self.experts.w13_weight  # (E, 2*I, H)
+            w2 = self.experts.w2_weight    # (E, H, I)
+
+            # topk routing (reuse existing corex/xllm/pytorch topk)
+            if _USE_XLLM_MOE:
+                topk_weights, topk_ids = _xllm_moe.moe_fused_topk(
+                    router_logits, self.top_k, True, None, "softmax")
+                topk_ids = topk_ids.to(torch.int64)
+                topk_weights = topk_weights.to(hidden_states.dtype)
+            elif _USE_COREX_MOE_TOPK_SOFTMAX:
+                topk_weights, topk_ids = _corex_moe_topk_softmax.moe_topk_softmax(
+                    router_logits.float(), self.top_k, True)
+                topk_ids = topk_ids.to(torch.int64)
+                topk_weights = topk_weights.to(hidden_states.dtype)
+            else:
+                topk_logits, topk_ids = torch.topk(
+                    router_logits.float(), self.top_k, dim=-1)
+                topk_weights = torch.softmax(topk_logits, dim=-1)
+                topk_weights = topk_weights.to(hidden_states.dtype)
+
+            return naive_batched_moe_forward(
+                hidden_states, w13, w2,
+                topk_ids, topk_weights,
+                act_fn=self.act_fn)
+
         # ---------------------------------------------------------------
         # Tier 1: Point-optimized Python loop (individual corex .so)
         # ---------------------------------------------------------------