feat: integrate xllm_moe into qwen3_5.py MoE hot path

xllm_moe.so provides 3 fused CUDA kernels compiled for ivcore10:
  - moe_fused_topk: CUB topk + softmax (replaces corex_moe_topk_softmax)
  - moe_compute_index: histogram + prefix_sum + place (replaces corex_moe_index_combine)
  - moe_combine_result: reorder + weighted sum (available but not yet wired to output)

Dispatch priority in _pure_pytorch_experts():
  Tier 0: xllm_moe (if available)
  Tier 1: corex_moe_* individual .so
  Tier 2: PyTorch fallback

Integration points:
  1. Topk routing: xllm_moe.moe_fused_topk → corex_moe_topk_softmax → torch.topk
  2. Index computation: xllm_moe.moe_compute_index → corex_moe_index_combine → torch.argsort
  3. Expert loop: still Python F.linear (next target: batch GEMM)

patch_ops.sh already deploys all prebuilt/*.so including xllm_moe.so
This commit is contained in:
claude
2026-08-14 11:37:21 +00:00
parent e147c283e3
commit 865c18f852

View File

@@ -148,6 +148,14 @@ try:
except ImportError:
_corex_moe_index_combine = None
try:
from vllm import xllm_moe as _xllm_moe
except ImportError:
try:
import xllm_moe as _xllm_moe
except ImportError:
_xllm_moe = None
try:
from vllm import corex_gdn_chunk_recurrent as _corex_gdn_chunk_recurrent
except ImportError:
@@ -202,6 +210,11 @@ _USE_COREX_MOE_TOPK_SOFTMAX = (
_USE_COREX_MOE_INDEX_COMBINE = (
_corex_moe_index_combine is not None
and env_bool("BI100_MOE_COREX_INDEX_COMBINE", True))
_USE_XLLM_MOE = (
_xllm_moe is not None
and env_bool("BI100_MOE_XLLM", True))
if _USE_XLLM_MOE:
logger.info("xllm_moe ENABLED — fused_topk + compute_index + combine_result")
_USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
# ix_fused_moe: full 7-step fused MoE pipeline via ixformer C++ API
@@ -1666,7 +1679,12 @@ class Qwen3_5MoeSparseBlock(nn.Module):
# ---------------------------------------------------------------
# Fused topk+softmax: single CUB kernel vs 2 PyTorch ops.
# Source: xllm/core/kernels/cuda/moe/moe_topk_softmax_kernels.cuh
if _USE_COREX_MOE_TOPK_SOFTMAX:
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)
@@ -1762,7 +1780,16 @@ class Qwen3_5MoeSparseBlock(nn.Module):
out = torch.zeros_like(hidden_states)
flat_eids = topk_ids.reshape(-1)
if _USE_COREX_MOE_INDEX_COMBINE:
if _USE_XLLM_MOE:
# xllm CUDA: histogram + prefix_sum + place
src_dst, dst_src, expert_sizes = \
_xllm_moe.moe_compute_index(flat_eids, w13.shape[0])
sorted_tok_ids = torch.arange(
T, device=topk_ids.device
).repeat_interleave(self.top_k)[dst_src.long()]
sorted_weights = topk_weights.reshape(-1)[dst_src.long()]
expert_counts = expert_sizes.tolist()
elif _USE_COREX_MOE_INDEX_COMBINE:
# Fused CUDA: histogram + prefix_sum + place (11.5x faster)
src_dst, dst_src, expert_sizes = \
_corex_moe_index_combine.moe_compute_index(