feat: integrate moe_compute_index kernel into MoE prefill path

Verified on real BI-V100:
  moe_compute_index: 11.48x speedup (0.035ms vs 0.397ms)
  moe_combine_result: 2.66x speedup (0.022ms vs 0.059ms)

Integration:
- qwen3_5.py: import corex_moe_index_combine, use in prefill path
  with _USE_COREX_MOE_INDEX_COMBINE flag (env BI100_MOE_COREX_INDEX_COMBINE)
  Falls back to PyTorch argsort+bincount if .so unavailable
- patch_ops.sh: compile corex_moe_index_combine.cu during docker build

qwen3_6_scripts/patch_ops.sh

@@ -246,6 +246,14 @@ if source != installed:
     raise SystemExit("runtime api_server overlay identity mismatch")
 PY
 
+build_stage "compiling CoreX MoE index+combine kernel"
+if [[ -x /usr/local/corex-3.2.3/bin/clang++ ]]; then
+    bash ./build_corex_moe_index_combine.sh "${VLLM_ROOT}" || \
+        echo "[WARN] moe_index_combine build failed — will use PyTorch fallback"
+else
+    echo "[WARN] corex clang++ not found — skipping moe_index_combine build"
+fi
+
 build_stage "compiling submission Python sources"
 find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
 build_stage "patch script completed"

qwen3_6_scripts/qwen3_5.py

@@ -143,6 +143,11 @@ try:
 except ImportError:
     _corex_moe_topk_softmax = None
 
+try:
+    from vllm import corex_moe_index_combine as _corex_moe_index_combine
+except ImportError:
+    _corex_moe_index_combine = None
+
 from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA,
                                                    SupportsMultiModal)
 
@@ -187,6 +192,9 @@ _USE_COREX_MOE_DIRECT_ROUTED = (
 _USE_COREX_MOE_TOPK_SOFTMAX = (
     _corex_moe_topk_softmax is not None
     and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True))
+_USE_COREX_MOE_INDEX_COMBINE = (
+    _corex_moe_index_combine is not None
+    and env_bool("BI100_MOE_COREX_INDEX_COMBINE", True))
 _USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
 
 
@@ -1701,17 +1709,28 @@ class Qwen3_5MoeSparseBlock(nn.Module):
                 out = (expert_out * ws.unsqueeze(-1)).sum(
                     0, keepdim=True).to(hidden_states.dtype)   # (1, H)
         else:
-            # General path (prefill / multi-seq): group assignments once. The
-            # previous implementation scanned the full (T, top_k) routing
-            # matrix and ran nonzero() for every active expert.
+            # General path (prefill / multi-seq): group assignments once.
             out = torch.zeros_like(hidden_states)
             flat_eids = topk_ids.reshape(-1)
-            order = torch.argsort(flat_eids, stable=True)
-            sorted_tok_ids = torch.arange(
-                T, device=topk_ids.device).repeat_interleave(self.top_k)[order]
-            sorted_weights = topk_weights.reshape(-1)[order]
-            expert_counts = torch.bincount(
-                flat_eids, minlength=w13.shape[0]).tolist()
+
+            if _USE_COREX_MOE_INDEX_COMBINE:
+                # Fused CUDA: histogram + prefix_sum + place (11.5x faster)
+                src_dst, dst_src, expert_sizes = \
+                    _corex_moe_index_combine.moe_compute_index(
+                        flat_eids, w13.shape[0])
+                sorted_tok_ids = torch.arange(
+                    T, device=topk_ids.device
+                ).repeat_interleave(self.top_k)[dst_src.long()]
+                sorted_weights = topk_weights.reshape(-1)[dst_src.long()]
+                expert_counts = expert_sizes.tolist()
+            else:
+                order = torch.argsort(flat_eids, stable=True)
+                sorted_tok_ids = torch.arange(
+                    T, device=topk_ids.device
+                ).repeat_interleave(self.top_k)[order]
+                sorted_weights = topk_weights.reshape(-1)[order]
+                expert_counts = torch.bincount(
+                    flat_eids, minlength=w13.shape[0]).tolist()
 
             start = 0
             for eid, count in enumerate(expert_counts):