From e18ece8f3ab9e66041cd67c0aa5c6663ea3ae1d9 Mon Sep 17 00:00:00 2001 From: dylan Date: Sat, 15 Aug 2026 13:05:45 +0000 Subject: [PATCH] =?UTF-8?q?feat:=20port=20NaiveBatchedExperts=20from=20ds?= =?UTF-8?q?=5Fvllm=20=E2=80=94=20view=20transpose=20+=20cublas=20transB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 10 ++ ex_engine/moe/activation.py | 165 +++++++++++++++++++++++++ ex_engine/moe/naive_batched_experts.py | 134 ++++++++++++++++++++ qwen3_6_scripts/qwen3_5.py | 50 ++++++++ 4 files changed, 359 insertions(+) create mode 100644 ex_engine/moe/__init__.py create mode 100644 ex_engine/moe/activation.py create mode 100644 ex_engine/moe/naive_batched_experts.py diff --git a/ex_engine/moe/__init__.py b/ex_engine/moe/__init__.py new file mode 100644 index 00000000..66c220e8 --- /dev/null +++ b/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 diff --git a/ex_engine/moe/activation.py b/ex_engine/moe/activation.py new file mode 100644 index 00000000..7216b5a0 --- /dev/null +++ b/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 diff --git a/ex_engine/moe/naive_batched_experts.py b/ex_engine/moe/naive_batched_experts.py new file mode 100644 index 00000000..ecc4b11e --- /dev/null +++ b/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 diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 5fc77d96..a05c80c6 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/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) # ---------------------------------------------------------------