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
This commit is contained in:
dylan
2026-08-15 13:05:45 +00:00
parent 6f1904aa8c
commit e18ece8f3a
4 changed files with 359 additions and 0 deletions

165
ex_engine/moe/activation.py Normal file
View File

@@ -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