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

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:
-  - 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()