e18ece8f3a
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
11 lines
386 B
Python
11 lines
386 B
Python
"""
|
|
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
|