feat: CUTLASS Cu10 grouped GEMM — real device verified

BI-V100 real device results:
  moe_group_gemm: err=0.000015 PASS
  moe_decode_cutlass: NaN=False PASS
  cutlass grouped: 4.77ms vs torch.mm loop: 9.38ms → 1.97x speedup

Fix: gemm_grouped.cu ldb=K (not N) for ColumnMajor B view
Link: -lcuinfer from /usr/local/corex-3.2.3/lib64/libcuinfer.so.7

ex_engine/python/gemm_dispatch.py

@@ -0,0 +1,180 @@
+"""gemm_dispatch.py — Unified GEMM dispatch for MoE group matmul.
+
+AST Layer 2: selects best available GEMM backend on real device.
+
+Backend priority:
+  1. gemm_grouped.so  (cutlass Cu10 TensorOp, per-expert GEMM)
+  2. ix_moe_bridge.so (cuinferCustomGemm, per-expert loop)
+  3. corex_batched_gemm.so (cutlass batched, decode-only)
+  4. hgemm.so (blocktiling kernel from siboehm)
+  5. torch.mm loop (PyTorch fallback)
+
+Reference: ex_engine/python/ix_ops_dispatch.py (407L)
+"""
+import os
+import logging
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger("gemm_dispatch")
+
+# --- Backend loading ---
+_cutlass_grouped = None
+_moe_bridge = None
+_batched_gemm = None
+_hgemm = None
+_backend = "torch"
+
+
+def _try_load(name):
+    """Try to load a .so module by name."""
+    # Search paths
+    search = [
+        os.path.join(os.path.dirname(__file__), f"{name}.so"),
+        os.path.join(os.path.dirname(__file__), "..", "prebuilt", f"{name}.so"),
+        os.path.join(os.path.dirname(__file__), "..", f"{name}.so"),
+    ]
+    for p in search:
+        if os.path.isfile(p):
+            try:
+                import importlib.util
+                spec = importlib.util.spec_from_file_location(name, p)
+                mod = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(mod)
+                return mod
+            except Exception as e:
+                logger.debug(f"[gemm] Failed to load {p}: {e}")
+    # Try direct import
+    try:
+        import importlib
+        return importlib.import_module(name)
+    except ImportError:
+        return None
+
+
+def _init_backends():
+    global _cutlass_grouped, _moe_bridge, _batched_gemm, _hgemm, _backend
+
+    _cutlass_grouped = _try_load("gemm_grouped")
+    if _cutlass_grouped and hasattr(_cutlass_grouped, "moe_group_gemm"):
+        _backend = "cutlass_grouped"
+        logger.info("[gemm] Backend: cutlass_grouped (Cu10 TensorOp)")
+        return
+
+    _moe_bridge = _try_load("ix_moe_bridge")
+    if _moe_bridge and hasattr(_moe_bridge, "group_gemm"):
+        _backend = "cuinfer"
+        logger.info("[gemm] Backend: cuinfer (via ix_moe_bridge)")
+        return
+
+    _batched_gemm = _try_load("corex_batched_gemm")
+    if _batched_gemm and hasattr(_batched_gemm, "batched_gemm_fp16"):
+        _backend = "cutlass_batched"
+        logger.info("[gemm] Backend: cutlass_batched")
+        return
+
+    _hgemm = _try_load("hgemm")
+    if _hgemm and hasattr(_hgemm, "moe_expert_gemm"):
+        _backend = "hgemm"
+        logger.info("[gemm] Backend: hgemm (blocktiling)")
+        return
+
+    _backend = "torch"
+    logger.info("[gemm] Backend: torch (F.linear fallback)")
+
+
+_init_backends()
+
+
+# ============================================================================
+# Public API
+# ============================================================================
+
+def group_gemm(input_tokens, weights, expert_counts, output_dim):
+    """Per-expert GEMM: output[offset:offset+count] = input[offset:offset+count] @ W[e]^T
+
+    Args:
+        input_tokens: (total_tokens, K) fp16
+        weights:      (num_experts, N, K) fp16, TN layout
+        expert_counts: (num_experts,) int32
+        output_dim:   N (output dimension)
+
+    Returns:
+        (total_tokens, N) fp16
+    """
+    if _backend == "cutlass_grouped":
+        return _cutlass_grouped.moe_group_gemm(input_tokens, weights, expert_counts)
+
+    if _backend == "cuinfer":
+        return _moe_bridge.group_gemm(input_tokens, weights, expert_counts, output_dim)
+
+    if _backend == "hgemm":
+        return _hgemm.moe_expert_gemm(input_tokens, weights, expert_counts)
+
+    # torch fallback
+    return _torch_group_gemm(input_tokens, weights, expert_counts)
+
+
+def moe_decode_gemm(hidden, w13_sel, w2_sel, topk_weights):
+    """Single-token MoE decode: batched GEMM over topk experts.
+
+    Args:
+        hidden:       (1, H) fp16
+        w13_sel:      (topk, 2*I, H) fp16
+        w2_sel:       (topk, H, I) fp16
+        topk_weights: (topk,) float32
+
+    Returns:
+        (1, H) fp16
+    """
+    if _backend == "cutlass_grouped" and hasattr(_cutlass_grouped, "moe_decode_cutlass"):
+        return _cutlass_grouped.moe_decode_cutlass(hidden, w13_sel, w2_sel, topk_weights)
+
+    if _backend == "cutlass_batched" and _batched_gemm is not None:
+        return _batched_gemm.moe_decode_fused(hidden, w13_sel, w2_sel, topk_weights)
+
+    # torch fallback
+    return _torch_moe_decode(hidden, w13_sel, w2_sel, topk_weights)
+
+
+def get_backend():
+    return _backend
+
+
+# ============================================================================
+# Fallbacks
+# ============================================================================
+
+def _torch_group_gemm(input_tokens, weights, expert_counts):
+    """PyTorch fallback: per-expert F.linear loop."""
+    num_experts = weights.size(0)
+    N = weights.size(1)
+    output = torch.zeros(input_tokens.size(0), N,
+                          device=input_tokens.device, dtype=input_tokens.dtype)
+
+    counts_cpu = expert_counts.cpu().to(torch.int32)
+    offset = 0
+    for e in range(num_experts):
+        cnt = counts_cpu[e].item()
+        if cnt <= 0:
+            offset += cnt
+            continue
+        x = input_tokens[offset:offset+cnt]
+        w = weights[e]  # (N, K)
+        output[offset:offset+cnt] = F.linear(x, w)
+        offset += cnt
+
+    return output
+
+
+def _torch_moe_decode(hidden, w13_sel, w2_sel, topk_weights):
+    """PyTorch fallback for single-token MoE decode."""
+    topk = w13_sel.size(0)
+    results = []
+    for k in range(topk):
+        gate_up = F.linear(hidden, w13_sel[k])
+        inter = gate_up.shape[-1] // 2
+        act = torch.silu(gate_up[:, :inter]) * gate_up[:, inter:]
+        down = F.linear(act, w2_sel[k])
+        results.append(down * topk_weights[k].to(down.dtype))
+    return sum(results)