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

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
+# Step 1: Compile _moe_C (CUB-based topk_softmax + moe_align_block_size)
-RUN chmod +x /workspace/ex_engine/build.sh && \
+# Proven on real BI-V100: WARP_SIZE=64, -cl-fast-relaxed-math, cub/block/block_reduce.cuh
-    bash /workspace/ex_engine/build.sh --corex 2>&1 | tee /workspace/ex_build.log ; \
+RUN python3 /workspace/ex_engine/precompile_moe_kernels.py 2>&1 | tee /workspace/ex_build.log ; \
-    echo "[Dockerfile] ex_engine build exit code: $?"
+    echo "[Dockerfile] _moe_C precompile exit code: $?"
 
-# Step 2: Precompile MoE CUDA kernels
+# Step 2: Deploy patches (serving + engine fixes)
-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)
 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()
+# Step 3: Precompile GDN kernel (needs vllm in path, so after patch_ops)
-# 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)
 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: $?"

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

ex_engine/csrc/moe_v055/topk_softmax_kernels.cu

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

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:
+Proven on real BI-V100 hardware:
-  - topk_softmax(topk_weights, topk_indices, token_expert_indices, gating_output)
+- WARP_SIZE=64 (not 32)
-  - moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids, expert_ids, num_tokens_post_pad)
+- 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)
-Usage:
-  python3 precompile_moe_kernels.py           # JIT compile
-  python3 precompile_moe_kernels.py --test    # compile + smoke test
 """
-import os
+import os, sys, logging
-import sys
+logging.basicConfig(level=logging.INFO)
-import time
+logger = logging.getLogger("precompile_moe")
 
-def compile_moe_kernels():
+def main():
-    """JIT compile MoE CUDA kernels via torch.utils.cpp_extension."""
     import torch
     from torch.utils.cpp_extension import load
 
-    script_dir = os.path.dirname(os.path.abspath(__file__))
+    base = os.path.dirname(os.path.abspath(__file__))
-    moe_dir = os.path.join(script_dir, 'csrc', 'moe_v055')
+    v055 = os.path.join(base, "csrc", "moe_v055")
 
     sources = [
-        os.path.join(moe_dir, 'moe_pybind.cpp'),
+        os.path.join(v055, "topk_softmax_kernels.cu"),
-        os.path.join(moe_dir, 'topk_softmax_kernels.cu'),
+        os.path.join(v055, "moe_align_block_size_kernels.cu"),
-        os.path.join(moe_dir, '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:
+    logger.info("Sources: %s", sources)
-        if not os.path.isfile(s):
+    logger.info("Compiling _moe_C...")
-            raise FileNotFoundError(f"Missing: {s}")
 
-    print(f"[moe_kernels] Compiling from {moe_dir}")
+    try:
-    t0 = time.time()
+        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(
+if __name__ == "__main__":
-        name='moe_kernels',
+    main()
-        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)

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
+        return
-    # 1. Try import precompiled module (torch cache from Docker build)
+    # 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]
+            if hasattr(_moe_topk_ext, 'topk_softmax'):
-            results = _moe_topk_ext.moe_topk_softmax(gating, topk_k, False)
+                # _moe_C style: in-place (vllm standard API)
-            topk_weights.copy_(results[0].to(topk_weights.dtype))
+                _moe_topk_ext.topk_softmax(topk_weights, topk_ids,
-            topk_ids.copy_(results[1].to(topk_ids.dtype))
+                                           token_expert_indicies, gating.float())
-            token_expert_indicies.copy_(results[2].to(token_expert_indicies.dtype))
+            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)