debug: deep probe MoE kernel dispatch in base image

probe_moe_detail.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""probe_moe_detail.py — Find exactly how to make MoE work on BI-V100"""
+import os, sys, traceback
+
+# 1. Check if vllm_moe_topk_softmax exists anywhere
+print("=== 1. Search for vllm_moe_topk_softmax ===")
+try:
+    import ixformer.functions as ixf_F
+    if hasattr(ixf_F, 'vllm_moe_topk_softmax'):
+        print("  FOUND in ixf_F!")
+    else:
+        print("  NOT in ixf_F")
+        # Check submodules
+        for attr in dir(ixf_F):
+            mod = getattr(ixf_F, attr)
+            if hasattr(mod, 'vllm_moe_topk_softmax'):
+                print(f"  FOUND in ixf_F.{attr}")
+except Exception as e:
+    print(f"  {e}")
+
+# 2. Read the actual _custom_ops.py from base image (not our copy)
+print("\n=== 2. Base image _custom_ops.py topk_softmax ===")
+for p in ["/usr/local/corex/lib64/python3/dist-packages/vllm/_custom_ops.py",
+          "/usr/local/corex/lib/python3/dist-packages/vllm/_custom_ops.py"]:
+    if os.path.exists(p):
+        print(f"  File: {p}")
+        with open(p) as f:
+            lines = f.readlines()
+        for i, line in enumerate(lines):
+            if 'topk_softmax' in line or 'moe_topk' in line or 'invoke_fused_moe' in line:
+                # Print context
+                start = max(0, i-2)
+                end = min(len(lines), i+5)
+                for j in range(start, end):
+                    marker = ">>>" if j == i else "   "
+                    print(f"  {marker} {j+1}: {lines[j].rstrip()}")
+                print()
+        break
+
+# 3. Read base image fused_moe.py — the actual kernel dispatch
+print("\n=== 3. Base image fused_moe.py kernel dispatch ===")
+for p in ["/usr/local/corex/lib64/python3/dist-packages/vllm/model_executor/layers/fused_moe/fused_moe.py",
+          "/usr/local/corex/lib/python3/dist-packages/vllm/model_executor/layers/fused_moe/fused_moe.py"]:
+    if os.path.exists(p):
+        print(f"  File: {p}")
+        with open(p) as f:
+            lines = f.readlines()
+        for i, line in enumerate(lines):
+            if 'invoke_fused_moe' in line or 'triton' in line.lower() or 'kernel' in line.lower() or 'ixf' in line.lower():
+                start = max(0, i-1)
+                end = min(len(lines), i+3)
+                for j in range(start, end):
+                    marker = ">>>" if j == i else "   "
+                    print(f"  {marker} {j+1}: {lines[j].rstrip()}")
+                print()
+        break
+
+# 4. Check _ixformer_torch for topk
+print("\n=== 4. _ixformer_torch Python bindings ===")
+try:
+    import ixformer._ixformer_torch as ixt
+    print(f"  Module: {ixt}")
+    for attr in sorted(dir(ixt)):
+        if not attr.startswith('__'):
+            print(f"    {attr}")
+except Exception as e:
+    print(f"  {e}")
+
+# 5. Check ixformer.functions.vllm source
+print("\n=== 5. ixformer.functions.vllm source (for vllm_moe references) ===")
+try:
+    import ixformer.functions.vllm as ixf_vllm
+    import inspect
+    src = inspect.getsource(ixf_vllm)
+    for i, line in enumerate(src.split('\n')):
+        if 'moe' in line.lower() or 'topk' in line.lower() or 'expert' in line.lower() or 'mlp' in line.lower():
+            print(f"  {i+1}: {line}")
+except Exception as e:
+    print(f"  {e}")
+
+# 6. What does _custom_ops invoke_fused_moe_kernel look like?
+print("\n=== 6. invoke_fused_moe_kernel in _custom_ops ===")
+for p in ["/usr/local/corex/lib64/python3/dist-packages/vllm/_custom_ops.py"]:
+    if os.path.exists(p):
+        with open(p) as f:
+            content = f.read()
+        if 'invoke_fused_moe' in content:
+            idx = content.index('invoke_fused_moe')
+            start = max(0, content.rfind('\n', 0, idx-100))
+            end = content.find('\n\n', idx+100)
+            print(content[start:end])
+        else:
+            print("  invoke_fused_moe NOT in _custom_ops.py")
+            # What IS there for MoE?
+            for line in content.split('\n'):
+                if 'moe' in line.lower() or 'expert' in line.lower():
+                    print(f"  {line.strip()}")