From 0478628f1721e0ed2d6f4ad897e2cbfd081cf9cd Mon Sep 17 00:00:00 2001 From: project6-dev Date: Tue, 11 Aug 2026 01:50:42 +0000 Subject: [PATCH] fix(PROVEN): _moe_C compiles and runs on real BI-V100 hardware MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- Dockerfile | 26 +--- ex_engine/csrc/moe_v055/cuda_compat.h | 2 +- .../csrc/moe_v055/topk_softmax_kernels.cu | 2 +- ex_engine/precompile_moe_kernels.py | 124 ++++++------------ qwen3_6_scripts/_custom_ops.py | 50 +++++-- 5 files changed, 91 insertions(+), 113 deletions(-) diff --git a/Dockerfile b/Dockerfile index e4af123f..0edf1acf 100644 --- a/Dockerfile +++ b/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: $?" diff --git a/ex_engine/csrc/moe_v055/cuda_compat.h b/ex_engine/csrc/moe_v055/cuda_compat.h index 82e55613..03daf40f 100644 --- a/ex_engine/csrc/moe_v055/cuda_compat.h +++ b/ex_engine/csrc/moe_v055/cuda_compat.h @@ -5,7 +5,7 @@ #endif #ifndef USE_ROCM - #define WARP_SIZE 32 + #define WARP_SIZE 64 #else #define WARP_SIZE warpSize #endif diff --git a/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu b/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu index 5273e0a5..2df0a94a 100644 --- a/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu +++ b/ex_engine/csrc/moe_v055/topk_softmax_kernels.cu @@ -23,7 +23,7 @@ #ifndef USE_ROCM #include - #include + #include #else #include #include diff --git a/ex_engine/precompile_moe_kernels.py b/ex_engine/precompile_moe_kernels.py index 54278243..ad6a908f 100644 --- a/ex_engine/precompile_moe_kernels.py +++ b/ex_engine/precompile_moe_kernels.py @@ -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() diff --git a/qwen3_6_scripts/_custom_ops.py b/qwen3_6_scripts/_custom_ops.py index eb7b6c23..9962a147 100644 --- a/qwen3_6_scripts/_custom_ops.py +++ b/qwen3_6_scripts/_custom_ops.py @@ -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)