135 lines
5.2 KiB
Python
135 lines
5.2 KiB
Python
"""
|
|
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].tolist() # (top_k,) → CPU list, ONE sync
|
|
ws = topk_weights[0] # (top_k,) stays on GPU
|
|
|
|
for i in range(top_k):
|
|
eid = eids[i]
|
|
|
|
# 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
|