dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: port NaiveBatchedExperts from ds_vllm — view transpose + cublas transB

Browse Source 
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
2026-08-15 13:05:45 +00:00
parent 6f1904aa8c
commit e18ece8f3a
4 changed files with 359 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
qwen3_6_scripts/qwen3_5.py
View File

@@ -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)
# ---------------------------------------------------------------