98 lines
3.9 KiB
Python
98 lines
3.9 KiB
Python
#!/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()}")
|