[feat] group gemm

ex_engine/python/moe_dispatch.py

@@ -0,0 +1,172 @@
+"""moe_dispatch.py — Load ix_moe_bridge.so and dispatch MoE forward.
+
+3-level fallback:
+  Tier 0: ix_moe_bridge.fused_moe_forward (C++ fused 7-step pipeline)
+  Tier 1: ix_moe_bridge individual ops (topk + expand + gemm + silu + gemm + combine)
+  Tier 2: Pure PyTorch fallback (F.linear loop)
+
+Used by: patch_moe_hot_path.py → replaces Qwen3_5MoE.forward()
+
+Reference: ex_engine/python/corex_moe.py (237L)
+"""
+import os
+import sys
+import logging
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger("moe_dispatch")
+
+# --- Load bridge .so ---
+_bridge = None
+_tier = 2  # default: PyTorch fallback
+
+
+def _try_load_bridge():
+    global _bridge, _tier
+
+    # Try 1: prebuilt .so
+    search_paths = [
+        os.path.join(os.path.dirname(__file__), "ix_moe_bridge.so"),
+        os.path.join(os.path.dirname(__file__), "..", "prebuilt", "ix_moe_bridge.so"),
+        os.path.join(os.path.dirname(__file__), "..", "ix_moe_bridge.so"),
+    ]
+    for p in search_paths:
+        if os.path.isfile(p):
+            try:
+                import importlib.util
+                spec = importlib.util.spec_from_file_location("ix_moe_bridge", p)
+                mod = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(mod)
+                _bridge = mod
+                logger.info(f"[moe_dispatch] ✓ Loaded bridge from {p}")
+                break
+            except Exception as e:
+                logger.warning(f"[moe_dispatch] Failed to load {p}: {e}")
+
+    # Try 2: torch JIT compiled module
+    if _bridge is None:
+        try:
+            import ix_moe_bridge
+            _bridge = ix_moe_bridge
+            logger.info("[moe_dispatch] ✓ Loaded bridge via import")
+        except ImportError:
+            pass
+
+    if _bridge is None:
+        logger.warning("[moe_dispatch] Bridge not available, using PyTorch fallback")
+        _tier = 2
+        return
+
+    # Check what functions are available
+    try:
+        if hasattr(_bridge, 'fused_moe_forward'):
+            _tier = 0
+            logger.info("[moe_dispatch] Tier 0: fused pipeline available")
+        elif hasattr(_bridge, 'topk_softmax') and hasattr(_bridge, 'group_gemm'):
+            _tier = 1
+            logger.info("[moe_dispatch] Tier 1: individual ops available")
+        else:
+            _tier = 2
+            logger.warning("[moe_dispatch] Bridge loaded but missing functions")
+    except Exception as e:
+        logger.warning(f"[moe_dispatch] Function check failed: {e}")
+        _tier = 2
+
+
+_try_load_bridge()
+
+
+# ============================================================================
+# Tier 2: Pure PyTorch fallback (identical to base vllm behavior)
+# ============================================================================
+
+def _pytorch_moe_forward(hidden_states, router_logits, w13, w2,
+                          topk, num_experts, renormalize):
+    """Python fallback: softmax → topk → loop over experts with F.linear."""
+    gating = torch.softmax(router_logits.float(), dim=-1)
+    topk_weights, topk_ids = torch.topk(gating, topk, dim=-1)
+    if renormalize:
+        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)
+    topk_weights = topk_weights.to(hidden_states.dtype)
+
+    # Per-expert loop
+    final_output = torch.zeros_like(hidden_states)
+    for k in range(topk):
+        expert_ids = topk_ids[:, k]  # [T]
+        weights_k = topk_weights[:, k].unsqueeze(-1)  # [T, 1]
+        for e in range(num_experts):
+            mask = (expert_ids == e)
+            if not mask.any():
+                continue
+            expert_input = hidden_states[mask]
+            # gate_up = expert_input @ w13[e].T  → [n, 2*inter]
+            gate_up = F.linear(expert_input, w13[e])
+            inter = gate_up.shape[-1] // 2
+            gate = torch.sigmoid(gate_up[:, :inter])
+            up = gate_up[:, inter:]
+            activated = gate * up  # SiLU approximated as sigmoid * x (should be silu_and_mul)
+            # down = activated @ w2[e].T → [n, hidden]
+            down = F.linear(activated, w2[e])
+            final_output[mask] += weights_k[mask] * down
+
+    return final_output
+
+
+# ============================================================================
+# Tier 1: Individual bridge ops
+# ============================================================================
+
+def _bridge_individual_moe_forward(hidden_states, router_logits, w13, w2,
+                                    topk, num_experts, renormalize):
+    """Use individual bridge ops: topk → gen_idx → expand → gemm → silu → gemm → combine."""
+    topk_weights, topk_ids, _ = _bridge.topk_softmax(router_logits, topk, False)
+    if renormalize:
+        topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)
+
+    idx_results = _bridge.moe_gen_idx(topk_ids.view(-1).to(torch.int32), num_experts)
+    src_dst, dst_src, expert_sizes = idx_results[0], idx_results[1], idx_results[2]
+
+    expanded = _bridge.moe_expand_input(hidden_states, src_dst, dst_src, topk)
+
+    gate_up = _bridge.group_gemm(expanded, w13, expert_sizes, w13.size(1))
+    activated = _bridge.silu_and_mul(gate_up)
+    down = _bridge.group_gemm(activated, w2, expert_sizes, w2.size(1))
+    output = _bridge.moe_combine_result(down, topk_weights)
+
+    return output
+
+
+# ============================================================================
+# Public API
+# ============================================================================
+
+def moe_forward(hidden_states, router_logits, w13, w2,
+                topk, num_experts, renormalize=True):
+    """Dispatch MoE forward to best available implementation."""
+    if _tier == 0:
+        try:
+            return _bridge.fused_moe_forward(
+                hidden_states, router_logits, w13, w2,
+                topk, num_experts, renormalize)
+        except Exception as e:
+            logger.warning(f"[moe_dispatch] Tier 0 failed: {e}, falling to Tier 1")
+            pass
+
+    if _tier <= 1 and _bridge is not None:
+        try:
+            return _bridge_individual_moe_forward(
+                hidden_states, router_logits, w13, w2,
+                topk, num_experts, renormalize)
+        except Exception as e:
+            logger.warning(f"[moe_dispatch] Tier 1 failed: {e}, falling to Tier 2")
+            pass
+
+    return _pytorch_moe_forward(
+        hidden_states, router_logits, w13, w2,
+        topk, num_experts, renormalize)
+
+
+def get_tier():
+    """Return current dispatch tier (0=fused, 1=individual, 2=pytorch)."""
+    return _tier

ex_engine/python/patch_moe_hot_path.py

@@ -0,0 +1,109 @@
+"""patch_moe_hot_path.py — Replace Qwen3_5MoE.forward() with bridge dispatch.
+
+This is the key performance patch: replaces the Python expert-loop MoE
+with a single C++ call that does all 7 steps fused.
+
+Called by: patch_ops.sh during Docker build
+Target:    vllm.model_executor.models.qwen3_5.Qwen3_5MoE
+
+Reference: ex_engine/python/patch_vllm_hot_path.py (200L)
+"""
+import sys
+import logging
+import torch
+
+logger = logging.getLogger("patch_moe_hot_path")
+
+
+def apply_moe_patch():
+    """Monkey-patch Qwen3_5MoE.forward to use moe_dispatch."""
+    try:
+        from ex_engine.python.moe_dispatch import moe_forward, get_tier
+    except ImportError:
+        try:
+            from moe_dispatch import moe_forward, get_tier
+        except ImportError:
+            logger.warning("[moe_patch] moe_dispatch not available, skipping patch")
+            return False
+
+    tier = get_tier()
+    logger.info(f"[moe_patch] moe_dispatch tier={tier}")
+
+    # Find the MoE class
+    moe_cls = None
+    try:
+        from vllm.model_executor.models.qwen3_5 import Qwen3_5MoE
+        moe_cls = Qwen3_5MoE
+    except ImportError:
+        pass
+
+    if moe_cls is None:
+        # Try to find it in sys.modules (may be registered under different name)
+        for mod_name, mod in sys.modules.items():
+            if hasattr(mod, 'Qwen3_5MoE'):
+                moe_cls = getattr(mod, 'Qwen3_5MoE')
+                break
+
+    if moe_cls is None:
+        logger.warning("[moe_patch] Qwen3_5MoE class not found")
+        return False
+
+    # Save original forward
+    _original_forward = moe_cls.forward
+
+    def patched_forward(self, hidden_states, *args, **kwargs):
+        """Patched MoE forward using bridge dispatch."""
+        # Get router logits
+        # In Qwen3_5, the gate + shared_expert_gate are concatenated:
+        #   router_and_shared_gate = self.gate(hidden_states)
+        #   router_logits = router_and_shared_gate[..., :self.num_experts]
+        #   shared_gate = router_and_shared_gate[..., -1]
+        router_and_shared_gate = self.gate(hidden_states)
+        router_logits = router_and_shared_gate[..., :self.num_experts]
+
+        # Shared expert (if any) — run in parallel
+        shared_output = None
+        if hasattr(self, 'shared_expert') and self.shared_expert is not None:
+            if hasattr(self, 'shared_expert_gate'):
+                shared_gate = torch.sigmoid(
+                    router_and_shared_gate[..., -1].unsqueeze(-1))
+            else:
+                shared_gate = None
+
+        # Routed experts via bridge
+        try:
+            routed_output = moe_forward(
+                hidden_states.view(-1, hidden_states.shape[-1]),
+                router_logits.view(-1, router_logits.shape[-1]),
+                self.w13_weight if hasattr(self, 'w13_weight') else self.experts.w13_weight,
+                self.w2_weight if hasattr(self, 'w2_weight') else self.experts.w2_weight,
+                topk=self.top_k,
+                num_experts=self.num_experts,
+                renormalize=True,
+            )
+            routed_output = routed_output.view_as(hidden_states)
+        except Exception as e:
+            logger.warning(f"[moe_patch] Bridge failed ({e}), using original forward")
+            return _original_forward(self, hidden_states, *args, **kwargs)
+
+        # Add shared expert output
+        if hasattr(self, 'shared_expert') and self.shared_expert is not None:
+            shared_out = self.shared_expert(hidden_states)
+            if shared_gate is not None:
+                shared_out = shared_out * shared_gate
+            routed_output = routed_output + shared_out
+
+        return routed_output
+
+    # Only patch if we have a real bridge (not pure Python fallback)
+    if tier < 2:
+        moe_cls.forward = patched_forward
+        logger.info(f"[moe_patch] ✓ Patched Qwen3_5MoE.forward (tier={tier})")
+        return True
+    else:
+        logger.info("[moe_patch] Tier 2 (Python only), not patching")
+        return False
+
+
+if __name__ == "__main__":
+    apply_moe_patch()