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
10
ex_engine/moe/__init__.py
Normal file
10
ex_engine/moe/__init__.py
Normal file
@@ -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
|
||||
165
ex_engine/moe/activation.py
Normal file
165
ex_engine/moe/activation.py
Normal 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
|
||||
134
ex_engine/moe/naive_batched_experts.py
Normal file
134
ex_engine/moe/naive_batched_experts.py
Normal file
@@ -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
|
||||
@@ -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)
|
||||
# ---------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user