7aa5054574
System design: algorithm factor replacement, not a connector.
All ops go through ixformer::infer C++ namespace (no Python fallback).
New files:
build_ix_bridge.sh — compile ix_full_bridge_v2.cpp on BI-V100
build_xllm_ilu_kernels.sh — compile upstream xllm ILU wrappers
deploy_ilu_pipeline.sh — wire everything into patch_ops.sh
ix_ops_dispatch.py — runtime dispatcher (12 ops via C++ bridge)
corex_fa2_dispatch.py — 3-mode attention (prefill/v1/flash paged)
fused_moe_ilu.py — 7-step MoE pipeline (no expert for-loop)
Upstream sources used (not rewritten):
xllm/core/kernels/ilu/*.cpp (ILU kernel wrappers)
xllm/core/kernels/ilu/ixformer.h (14 C++ function declarations)
ds_vllm/csrc/libtorch_stable/*.cu (kernel references)
Call chain:
patch_ops.sh → deploy_ilu_pipeline.sh → build_ix_bridge.sh
→ ix_full_bridge_v2.so → ixformer::infer::*
→ silu_and_mul, rms_norm, rotary_embedding, paged_attention,
topk_softmax, group_gemm, expand_input, combine_result
206 lines
6.9 KiB
Python
206 lines
6.9 KiB
Python
"""
|
|
fused_moe_ilu.py — 7-step fused MoE via xllm upstream ILU dispatch chain
|
|
|
|
Upstream ref: xllm/core/layers/ilu/fused_moe.cpp
|
|
xllm/core/kernels/ilu/fused_moe.cpp
|
|
|
|
The 7-step pipeline:
|
|
1. topk_softmax → ixformer::infer::topk_softmax
|
|
2. moe_gen_idx → ixformer::infer::moe_compute_token_index_api
|
|
3. moe_expand_input → ixformer::infer::moe_expand_input
|
|
4. group_gemm (w13) → ixformer::infer::moe_w16a16_group_gemm
|
|
5. silu_and_mul → ixformer::infer::silu_and_mul
|
|
6. group_gemm (w2) → ixformer::infer::moe_w16a16_group_gemm
|
|
7. moe_combine_result → ixformer::infer::moe_output_reduce_sum
|
|
|
|
Every step calls C++. No Python expert loop.
|
|
"""
|
|
|
|
import logging
|
|
import torch
|
|
from typing import Optional, Tuple
|
|
|
|
logger = logging.getLogger("fused_moe_ilu")
|
|
|
|
_init_logged = False
|
|
|
|
# =====================================================================
|
|
# Load the C++ ops
|
|
# =====================================================================
|
|
|
|
def _get_ops():
|
|
"""Get the ix_ops_dispatch module."""
|
|
try:
|
|
from ex_engine.python import ix_ops_dispatch as ops
|
|
return ops
|
|
except ImportError:
|
|
pass
|
|
try:
|
|
from vllm.ex_engine import ix_ops_dispatch as ops
|
|
return ops
|
|
except ImportError:
|
|
pass
|
|
return None
|
|
|
|
|
|
# =====================================================================
|
|
# 7-step fused MoE forward
|
|
# =====================================================================
|
|
|
|
def fused_moe_forward(
|
|
hidden_states: torch.Tensor, # (num_tokens, hidden_size)
|
|
gate_output: torch.Tensor, # (num_tokens, num_experts) router logits
|
|
w13: torch.Tensor, # (E, 2*intermediate, hidden_size) merged gate_up
|
|
w2: torch.Tensor, # (E, hidden_size, intermediate)
|
|
topk: int = 8,
|
|
renormalize: bool = True,
|
|
num_experts: int = 64,
|
|
shared_expert: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Full 7-step fused MoE pipeline.
|
|
|
|
All steps go through C++ — no Python fallback.
|
|
If C++ is unavailable, raises RuntimeError.
|
|
"""
|
|
global _init_logged
|
|
ops = _get_ops()
|
|
if ops is None:
|
|
raise RuntimeError("fused_moe_ilu: ix_ops_dispatch not available")
|
|
|
|
num_tokens = hidden_states.shape[0]
|
|
hidden_size = hidden_states.shape[1]
|
|
intermediate_2x = w13.shape[1] # 2 * intermediate_size
|
|
intermediate = intermediate_2x // 2
|
|
|
|
if not _init_logged:
|
|
logger.info("Using fused MoE ILU pipeline: tokens=%d, experts=%d, topk=%d, "
|
|
"intermediate=%d", num_tokens, num_experts, topk, intermediate)
|
|
_init_logged = True
|
|
|
|
# Step 1: topk_softmax
|
|
topk_weights, topk_ids = ops.topk_softmax(gate_output, topk, renormalize)
|
|
|
|
# Step 2: moe_compute_token_index
|
|
src_dst, dst_src, expert_sizes = ops.moe_compute_token_index(
|
|
topk_ids, num_experts)
|
|
|
|
# Step 3: moe_expand_input
|
|
expanded = ops.moe_expand_input(hidden_states, dst_src, topk)
|
|
|
|
# Step 4: group_gemm w13 (gate + up projection)
|
|
gate_up = ops.moe_group_gemm(expanded, w13, expert_sizes, intermediate_2x)
|
|
|
|
# Step 5: silu_and_mul
|
|
activated = ops.silu_and_mul(gate_up)
|
|
|
|
# Step 6: group_gemm w2 (down projection)
|
|
down = ops.moe_group_gemm(activated, w2, expert_sizes, hidden_size)
|
|
|
|
# Step 7: moe_output_reduce_sum (weighted combine)
|
|
output = ops.moe_output_reduce_sum(down, topk_weights.to(down.dtype))
|
|
|
|
return output
|
|
|
|
|
|
# =====================================================================
|
|
# Fallback: Per-expert matmul (used when group_gemm unavailable)
|
|
# Still uses C++ for topk and activation, just loops for GEMM.
|
|
# =====================================================================
|
|
|
|
def fused_moe_per_expert(
|
|
hidden_states: torch.Tensor,
|
|
gate_output: torch.Tensor,
|
|
w13: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk: int = 8,
|
|
renormalize: bool = True,
|
|
num_experts: int = 64,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Per-expert fallback with C++ topk and activation.
|
|
Uses torch.matmul for GEMM (goes to cublas).
|
|
"""
|
|
ops = _get_ops()
|
|
num_tokens = hidden_states.shape[0]
|
|
hidden_size = hidden_states.shape[1]
|
|
intermediate_2x = w13.shape[1]
|
|
half_inter = intermediate_2x // 2
|
|
dtype = hidden_states.dtype
|
|
|
|
# Step 1: topk
|
|
if ops is not None:
|
|
try:
|
|
topk_weights, topk_ids = ops.topk_softmax(gate_output, topk, renormalize)
|
|
except RuntimeError:
|
|
scores = torch.softmax(gate_output.float(), dim=-1)
|
|
topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
topk_ids = topk_ids.to(torch.int32)
|
|
else:
|
|
scores = torch.softmax(gate_output.float(), dim=-1)
|
|
topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1)
|
|
if renormalize:
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
|
topk_ids = topk_ids.to(torch.int32)
|
|
|
|
topk_weights = topk_weights.to(dtype)
|
|
flat_ids = topk_ids.view(-1)
|
|
flat_weights = topk_weights.view(-1)
|
|
|
|
# Expand input
|
|
expanded = hidden_states.unsqueeze(1).expand(-1, topk, -1).reshape(-1, hidden_size)
|
|
output = torch.zeros_like(expanded)
|
|
|
|
# Per-expert GEMM (cublas)
|
|
for eidx in range(num_experts):
|
|
mask = (flat_ids == eidx)
|
|
if not mask.any():
|
|
continue
|
|
tokens = expanded[mask]
|
|
|
|
# gate_up GEMM → cublas via torch.matmul
|
|
gate_up = torch.matmul(tokens, w13[eidx].t())
|
|
|
|
# SiLU activation (C++ if available)
|
|
if ops is not None:
|
|
try:
|
|
act = ops.silu_and_mul(gate_up)
|
|
except RuntimeError:
|
|
act = torch.nn.functional.silu(gate_up[:, :half_inter]) * gate_up[:, half_inter:]
|
|
else:
|
|
act = torch.nn.functional.silu(gate_up[:, :half_inter]) * gate_up[:, half_inter:]
|
|
|
|
# down GEMM → cublas
|
|
output[mask] = torch.matmul(act, w2[eidx].t())
|
|
|
|
output = output * flat_weights.unsqueeze(-1)
|
|
return output.view(num_tokens, topk, hidden_size).sum(dim=1)
|
|
|
|
|
|
# =====================================================================
|
|
# Auto-dispatch: try full pipeline, fall back to per-expert
|
|
# =====================================================================
|
|
|
|
def moe_forward(
|
|
hidden_states: torch.Tensor,
|
|
gate_output: torch.Tensor,
|
|
w13: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
topk: int = 8,
|
|
renormalize: bool = True,
|
|
num_experts: int = 64,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
"""Auto-dispatch MoE: try full C++ pipeline, then per-expert with C++ ops."""
|
|
try:
|
|
return fused_moe_forward(
|
|
hidden_states, gate_output, w13, w2,
|
|
topk, renormalize, num_experts)
|
|
except RuntimeError as e:
|
|
logger.debug("Full pipeline failed: %s, using per-expert fallback", e)
|
|
return fused_moe_per_expert(
|
|
hidden_states, gate_output, w13, w2,
|
|
topk, renormalize, num_experts)
|