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

fix(PROVEN): _moe_C compiles and runs on real BI-V100 hardware

Browse Source 
Tested on real machine (cc-b2042074, BI-V100, IX-ML 3.2.3):
  _moe_C.topk_softmax() → SUCCESS, correct output

Two fixes proven on hardware:
1. cuda_compat.h: WARP_SIZE=64 (BI-V100 warp is 64, not 32)
2. topk_softmax_kernels.cu: cub/block/block_reduce.cuh instead of cub/cub.cuh
   (cub.cuh pulls radix_sort which has WARP_SIZE conflict)

Key finding: ixformer SDK on this base image does NOT have topk_softmax.
The ixformer::infer namespace from xllm's ixformer.h is for newer SDK.
We MUST compile our own _moe_C kernel — which now works.

Build flags (clang 16, ivcore10):
  CUDA: -O3 -cl-fast-relaxed-math (NOT --use_fast_math)
  C++:  -O2 -std=c++17

Dockerfile simplified: 3 steps (was 6)
_custom_ops.py: _moe_C as Priority 0, in-place vllm API
This commit is contained in:
project6-dev
2026-08-11 01:50:42 +00:00
parent 1cd8ca0649
commit 0478628f17
5 changed files with 91 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
qwen3_6_scripts/_custom_ops.py
View File

@@ -1008,15 +1008,39 @@ def _init_ix_bridge():
def _init_moe_topk():
global _moe_topk_ext, _moe_topk_init_done
_moe_topk_init_done = True
# 0. Try ix_bridge first (calls ixformer C++ SDK directly)
# 0. Try _moe_C (CUB-based, proven on BI-V100 real hardware 2026-08-11)
try:
import _moe_C as ext
if hasattr(ext, 'topk_softmax'):
_moe_topk_ext = ext
logger.info("topk_softmax: loaded _moe_C (CUB BlockReduce, WARP_SIZE=64)")
return
except ImportError:
pass
# 0b. Try loading from torch cache
import glob as _glob
for pattern in [
"/root/.cache/torch_extensions/py310_cu102/_moe_C/_moe_C.so",
"/root/.cache/torch_extensions/*/_moe_C/*.so",
]:
for so_path in _glob.glob(pattern):
try:
torch.ops.load_library(so_path)
import _moe_C as ext
_moe_topk_ext = ext
logger.info("topk_softmax: loaded _moe_C from %s", so_path)
return
except Exception:
pass
# 0c. Try ix_bridge (calls ixformer C++ SDK if available)
_init_ix_bridge()
if _ix_bridge_mod:
return # ix_bridge loaded, no need for CUDA kernel
# 1. Try import precompiled module (torch cache from Docker build)
return
# 1. Try import old precompiled module (torch cache from Docker build)
try:
import moe_topk_softmax_v3 as ext
_moe_topk_ext = ext
logger.info("topk_softmax: loaded precompiled CUDA kernel")
logger.info("topk_softmax: loaded precompiled moe_topk_softmax_v3")
return
except ImportError:
pass
@@ -1088,15 +1112,21 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
except Exception as e:
logger.warning("topk_softmax ix_bridge failed (%s), trying CUDA kernel", e)
# Priority 1: Our CUDA kernel (fused warp-shuffle, ~5x faster than PyTorch)
# Priority 1: CUDA kernel (_moe_C or moe_topk_softmax_v3)
if _moe_topk_ext is not None:
try:
gating = gating_output if isinstance(gating_output, torch.Tensor) else gating_output
topk_k = topk_weights.shape[1]
results = _moe_topk_ext.moe_topk_softmax(gating, topk_k, False)
topk_weights.copy_(results[0].to(topk_weights.dtype))
topk_ids.copy_(results[1].to(topk_ids.dtype))
token_expert_indicies.copy_(results[2].to(token_expert_indicies.dtype))
if hasattr(_moe_topk_ext, 'topk_softmax'):
# _moe_C style: in-place (vllm standard API)
_moe_topk_ext.topk_softmax(topk_weights, topk_ids,
token_expert_indicies, gating.float())
elif hasattr(_moe_topk_ext, 'moe_topk_softmax'):
# old v3 style: returns tuple
topk_k = topk_weights.shape[1]
results = _moe_topk_ext.moe_topk_softmax(gating, topk_k, False)
topk_weights.copy_(results[0].to(topk_weights.dtype))
topk_ids.copy_(results[1].to(topk_ids.dtype))
token_expert_indicies.copy_(results[2].to(token_expert_indicies.dtype))
return
except Exception as e:
logger.warning("topk_softmax CUDA kernel failed (%s), falling back to PyTorch", e)