Files
nvfp4-megamoe-kernel/probe_hf_indexer.py

65 lines
2.4 KiB
Python

#!/usr/bin/env python3
"""Probe the HF DeepSeekV4 indexer implementation to understand the correct architecture.
Specifically: what shape are the indexer compressed keys, and how does scoring work?
Run via: fire_b200_test probe_hf_indexer.py
"""
import sys, os
# Find the HF modeling file
candidates = [
"/root/dsv4-nvfp4-workspace/venv/lib/python3.12/site-packages/transformers/models/deepseek_v4/modeling_deepseek_v4.py",
"/root/dsv4-nvfp4-workspace/venv/lib/python*/site-packages/transformers/models/deepseek_v4/modeling_deepseek_v4.py",
]
# Also try to find it dynamically
import glob
matches = glob.glob("/root/dsv4-nvfp4-workspace/venv/lib/python*/site-packages/transformers/models/deepseek_v4/modeling_deepseek_v4.py")
if matches:
candidates = matches
found = None
for c in candidates:
if os.path.exists(c):
found = c
break
if found is None:
# Try pip show
import subprocess
result = subprocess.run(["find", "/root/dsv4-nvfp4-workspace/venv", "-name", "modeling_deepseek_v4.py"],
capture_output=True, text=True)
if result.stdout.strip():
found = result.stdout.strip().split('\n')[0]
if found:
print(f"Found: {found}")
# Read and print the indexer-related code
with open(found) as f:
lines = f.readlines()
# Find class definitions and indexer-related methods
in_relevant = False
indent = 0
for i, line in enumerate(lines):
# Look for indexer, compress, lightning, score keywords
lower = line.lower()
if any(kw in lower for kw in ['indexer', 'lightning', 'index_score', 'index_topk', 'compress_indexer', 'indexer_head']):
# Print surrounding context
start = max(0, i - 2)
end = min(len(lines), i + 20)
print(f"\n--- Line {i+1} ---")
for j in range(start, end):
marker = ">>>" if j == i else " "
print(f"{marker} {j+1}: {lines[j]}", end='')
else:
print("DeepSeek V4 modeling file not found. Checking what's available...")
result = subprocess.run(["find", "/root/dsv4-nvfp4-workspace/venv", "-name", "modeling_deepseek*.py"],
capture_output=True, text=True)
print(result.stdout[:2000] if result.stdout else "No deepseek modeling files found")
# Try pip
result2 = subprocess.run(["pip", "show", "transformers"], capture_output=True, text=True)
print(result2.stdout[:500])
print("\nDone.")