fix(PROVEN): _moe_C compiles and runs on real BI-V100 hardware
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
26
Dockerfile
26
Dockerfile
@@ -8,31 +8,17 @@ COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
|
||||
COPY ./computility-run.yaml /workspace/computility-run.yaml
|
||||
COPY ./ex_engine /workspace/ex_engine
|
||||
|
||||
# Step 1: Build EX Engine .so libraries
|
||||
RUN chmod +x /workspace/ex_engine/build.sh && \
|
||||
bash /workspace/ex_engine/build.sh --corex 2>&1 | tee /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] ex_engine build exit code: $?"
|
||||
# Step 1: Compile _moe_C (CUB-based topk_softmax + moe_align_block_size)
|
||||
# Proven on real BI-V100: WARP_SIZE=64, -cl-fast-relaxed-math, cub/block/block_reduce.cuh
|
||||
RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] _moe_C precompile exit code: $?"
|
||||
|
||||
# Step 2: Precompile MoE CUDA kernels
|
||||
RUN python3 /workspace/ex_engine/precompile_moe_topk.py 2>&1 | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] moe_topk precompile exit code: $?"
|
||||
|
||||
# Step 3: Precompile vllm v0.5.5 MoE kernels
|
||||
RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] moe_v055 precompile exit code: $?"
|
||||
|
||||
# Step 4: Deploy patches (serving + engine fixes)
|
||||
# Step 2: Deploy patches (serving + engine fixes)
|
||||
RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
|
||||
bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
|
||||
echo "[Dockerfile] patch_ops exit code: $?"
|
||||
|
||||
# Step 5: Precompile ix_moe_bridge.cpp → links to ixformer::infer::topk_softmax()
|
||||
# This is the C++ pybind bridge that makes ixformer SDK callable from Python.
|
||||
# Source: upstream_ref/xllm/xllm/core/kernels/ilu/fused_moe.cpp call pattern
|
||||
RUN python3 /workspace/ex_engine/precompile_ix_bridge.py 2>&1 | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] ix_bridge precompile exit code: $?"
|
||||
|
||||
# Step 6: Precompile GDN kernel (needs vllm in path, so after patch_ops)
|
||||
# Step 3: Precompile GDN kernel (needs vllm in path, so after patch_ops)
|
||||
RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
|
||||
/workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/ex_build.log ; \
|
||||
echo "[Dockerfile] gdn precompile exit code: $?"
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#define WARP_SIZE 32
|
||||
#define WARP_SIZE 64
|
||||
#else
|
||||
#define WARP_SIZE warpSize
|
||||
#endif
|
||||
|
||||
@@ -23,7 +23,7 @@
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#include <cub/util_type.cuh>
|
||||
#include <cub/cub.cuh>
|
||||
#include <cub/block/block_reduce.cuh>
|
||||
#else
|
||||
#include <hipcub/util_type.hpp>
|
||||
#include <hipcub/hipcub.hpp>
|
||||
|
||||
@@ -1,95 +1,57 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
precompile_moe_kernels.py — JIT compile vllm v0.5.5 MoE CUDA kernels for BI-V100.
|
||||
Precompile _moe_C extension: topk_softmax + moe_align_block_size.
|
||||
|
||||
Produces: moe_kernels.so with:
|
||||
- topk_softmax(topk_weights, topk_indices, token_expert_indices, gating_output)
|
||||
- moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad)
|
||||
|
||||
Usage:
|
||||
python3 precompile_moe_kernels.py # JIT compile
|
||||
python3 precompile_moe_kernels.py --test # compile + smoke test
|
||||
Proven on real BI-V100 hardware:
|
||||
- WARP_SIZE=64 (not 32)
|
||||
- cub/block/block_reduce.cuh (not cub/cub.cuh which pulls radix_sort)
|
||||
- -cl-fast-relaxed-math (not --use_fast_math which is nvcc-only)
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import os, sys, logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger("precompile_moe")
|
||||
|
||||
def compile_moe_kernels():
|
||||
"""JIT compile MoE CUDA kernels via torch.utils.cpp_extension."""
|
||||
def main():
|
||||
import torch
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
moe_dir = os.path.join(script_dir, 'csrc', 'moe_v055')
|
||||
base = os.path.dirname(os.path.abspath(__file__))
|
||||
v055 = os.path.join(base, "csrc", "moe_v055")
|
||||
|
||||
sources = [
|
||||
os.path.join(moe_dir, 'moe_pybind.cpp'),
|
||||
os.path.join(moe_dir, 'topk_softmax_kernels.cu'),
|
||||
os.path.join(moe_dir, 'moe_align_block_size_kernels.cu'),
|
||||
os.path.join(v055, "topk_softmax_kernels.cu"),
|
||||
os.path.join(v055, "moe_align_block_size_kernels.cu"),
|
||||
os.path.join(v055, "moe_pybind.cpp"),
|
||||
]
|
||||
for s in sources:
|
||||
if not os.path.exists(s):
|
||||
logger.error("MISSING: %s", s)
|
||||
sys.exit(1)
|
||||
|
||||
include_paths = [
|
||||
v055,
|
||||
os.path.join(base, "csrc", "moe"),
|
||||
os.path.join(base, "csrc"),
|
||||
"/usr/local/corex/include",
|
||||
]
|
||||
|
||||
for s in sources:
|
||||
if not os.path.isfile(s):
|
||||
raise FileNotFoundError(f"Missing: {s}")
|
||||
logger.info("Sources: %s", sources)
|
||||
logger.info("Compiling _moe_C...")
|
||||
|
||||
print(f"[moe_kernels] Compiling from {moe_dir}")
|
||||
t0 = time.time()
|
||||
try:
|
||||
mod = load(
|
||||
name="_moe_C",
|
||||
sources=sources,
|
||||
extra_include_paths=include_paths,
|
||||
extra_cuda_cflags=["-O3", "-cl-fast-relaxed-math"],
|
||||
extra_cflags=["-O2", "-std=c++17"],
|
||||
verbose=True,
|
||||
)
|
||||
fns = [x for x in dir(mod) if not x.startswith("_")]
|
||||
logger.info("SUCCESS: _moe_C functions: %s", fns)
|
||||
except Exception as e:
|
||||
logger.error("FAILED: %s", e)
|
||||
sys.exit(1)
|
||||
|
||||
mod = load(
|
||||
name='moe_kernels',
|
||||
sources=sources,
|
||||
extra_include_paths=[moe_dir],
|
||||
extra_cflags=['-O2', '-std=c++17'],
|
||||
extra_cuda_cflags=['-O2', '--expt-relaxed-constexpr'],
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
dt = time.time() - t0
|
||||
funcs = [x for x in dir(mod) if not x.startswith('_')]
|
||||
print(f"[moe_kernels] Compiled in {dt:.1f}s — functions: {funcs}")
|
||||
return mod
|
||||
|
||||
|
||||
def smoke_test(mod):
|
||||
"""Quick functional test of compiled kernels."""
|
||||
import torch
|
||||
|
||||
print("\n=== Smoke test ===")
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
if device == 'cpu':
|
||||
print(" SKIP: no CUDA device")
|
||||
return
|
||||
|
||||
# Test topk_softmax
|
||||
num_tokens, num_experts, topk = 4, 8, 2
|
||||
gating = torch.randn(num_tokens, num_experts, device=device, dtype=torch.float32)
|
||||
topk_weights = torch.empty(num_tokens, topk, device=device, dtype=torch.float32)
|
||||
topk_indices = torch.empty(num_tokens, topk, device=device, dtype=torch.int32)
|
||||
token_expert_indices = torch.empty(num_tokens, topk, device=device, dtype=torch.int32)
|
||||
|
||||
mod.topk_softmax(topk_weights, topk_indices, token_expert_indices, gating)
|
||||
|
||||
print(f" topk_softmax: weights={topk_weights.shape}, NaN={topk_weights.isnan().any()}")
|
||||
print(f" weights[0] = {topk_weights[0].tolist()}")
|
||||
print(f" indices[0] = {topk_indices[0].tolist()}")
|
||||
|
||||
# Test moe_align_block_size
|
||||
block_size = 4
|
||||
max_num_tokens_padded = (num_tokens * topk + num_experts * block_size)
|
||||
sorted_ids = torch.empty(max_num_tokens_padded, device=device, dtype=torch.int32)
|
||||
expert_ids = torch.empty(max_num_tokens_padded // block_size, device=device, dtype=torch.int32)
|
||||
num_tokens_post_pad = torch.empty(1, device=device, dtype=torch.int32)
|
||||
|
||||
mod.moe_align_block_size(topk_indices, num_experts, block_size,
|
||||
sorted_ids, expert_ids, num_tokens_post_pad)
|
||||
|
||||
print(f" moe_align: sorted_ids[:8]={sorted_ids[:8].tolist()}, "
|
||||
f"num_post_pad={num_tokens_post_pad.item()}")
|
||||
|
||||
print("\n ✓ All smoke tests passed")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
mod = compile_moe_kernels()
|
||||
if '--test' in sys.argv:
|
||||
smoke_test(mod)
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user