diff --git a/probe_moe_detail.py b/probe_moe_detail.py new file mode 100644 index 00000000..86215c66 --- /dev/null +++ b/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()}")