biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

clean: remove diagnostic scripts from repo

Browse Source
This commit is contained in:
biondizzle
2026-05-15 12:50:14 +00:00
parent fd59222fc0
commit da5572f497
5 changed files with 0 additions and 808 deletions
Download Patch File Download Diff File Expand all files Collapse all files

279
diag_b200.py
View File

@@ -1,279 +0,0 @@
"""
NVFP4 MegaMoE Diagnostic — B200
Checks:
1. weight_scale_2 values (are they nonzero / loaded correctly?)
2. Folded scale ranges (clamp/precision loss)
3. L2 weight/SF orientation sanity
4. Dequant reference vs CUTLASS output comparison
5. Single-expert, single-layer test
"""
import torch
import sys
import os
import json
from pathlib import Path
MODEL_PATH = "/model" # inside the container
def inspect_checkpoint_scales():
"""Check raw checkpoint weight_scale_2 values."""
from safetensors import safe_open
import glob
print("=" * 60)
print("CHECK 1: Checkpoint weight_scale_2 Values")
print("=" * 60)
# Find checkpoint files
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
print(f"Found {len(ckpt_files)} safetensors files")
# Look for expert weight_scale_2 params
w13_gs_found = 0
w2_gs_found = 0
w13_gs_values = {}
w2_gs_values = {}
for f in ckpt_files:
with safe_open(f, framework="pt") as st:
for key in st.keys():
if "weight_scale_2" in key and ("experts" in key or "ffn" in key):
val = st.get_tensor(key)
if "w13" in key or "gate_up" in key or "w1" in key or "w3" in key:
w13_gs_found += 1
if w13_gs_found <= 3:
w13_gs_values[key] = {"shape": list(val.shape), "dtype": str(val.dtype),
"min": val.float().min().item(), "max": val.float().max().item(),
"mean": val.float().mean().item()}
elif "w2" in key or "down" in key:
w2_gs_found += 1
if w2_gs_found <= 3:
w2_gs_values[key] = {"shape": list(val.shape), "dtype": str(val.dtype),
"min": val.float().min().item(), "max": val.float().max().item(),
"mean": val.float().mean().item()}
print(f"w13 weight_scale_2 entries: {w13_gs_found}")
print(f"w2 weight_scale_2 entries: {w2_gs_found}")
for k, v in w13_gs_values.items():
print(f" {k}: {v}")
for k, v in w2_gs_values.items():
print(f" {k}: {v}")
return w13_gs_found > 0 and w2_gs_found > 0
def inspect_loaded_model():
"""Check the model's weight_scale_2 after loading (before finalize_weights)."""
print("\n" + "=" * 60)
print("CHECK 2: Model weight_scale_2 After Loading")
print("=" * 60)
# We need to load the model and inspect before finalize_weights nukes the params
# The vLLM server is already running, so let's check the live model
# Actually, let's load a fresh model instance for inspection
# Simpler approach: just check the checkpoint directly for scale_2
# The real check is whether finalize_weights gets called with nonzero scale_2
print(" (Checkpoint inspection is more reliable — see CHECK 1)")
print(" The [SF-DEBUG] prints from weight_transform.py should also show this")
def check_fold_precision_real():
"""Check float8 folding precision with real checkpoint scales."""
print("\n" + "=" * 60)
print("CHECK 3: Float8 Folding Precision (Real Scales)")
print("=" * 60)
from safetensors import safe_open
import glob
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
# Find one layer's expert scales
for f in ckpt_files:
with safe_open(f, framework="pt") as st:
keys = list(st.keys())
# Find w2 weight_scale and weight_scale_2 for layer 0
w2_sf_key = None
w2_gs_key = None
w13_sf_key = None
w13_gs_key = None
for k in keys:
if "layers.0" in k:
if "w2" in k and k.endswith("weight_scale") and "scale_2" not in k:
w2_sf_key = k
elif "w2" in k and "weight_scale_2" in k:
w2_gs_key = k
elif ("w13" in k or "gate_up" in k) and k.endswith("weight_scale") and "scale_2" not in k:
w13_sf_key = k
elif ("w13" in k or "gate_up" in k) and "weight_scale_2" in k:
w13_gs_key = k
if w2_sf_key and w2_gs_key:
w2_sf = st.get_tensor(w2_sf_key)
w2_gs = st.get_tensor(w2_gs_key)
print(f" L2 block scale: shape={list(w2_sf.shape)} dtype={w2_sf.dtype} "
f"range=[{w2_sf.float().min():.4e}, {w2_sf.float().max():.4e}]")
print(f" L2 global scale: shape={list(w2_gs.shape)} dtype={w2_gs.dtype} "
f"range=[{w2_gs.float().min():.4e}, {w2_gs.float().max():.4e}]")
# Fold and check precision
sf_f32 = w2_sf.float()
gs_f32 = w2_gs.float()
# Reshape gs for broadcast
while gs_f32.dim() < sf_f32.dim():
gs_f32 = gs_f32.unsqueeze(-1)
product = sf_f32 * gs_f32
product_clamped = product.clamp(0.0, 448.0)
folded_f8 = product_clamped.to(torch.float8_e4m3fn)
folded_back = folded_f8.float()
# Stats
n_clamped = (product > 448.0).sum().item()
n_total = product.numel()
n_zeroed = (folded_back == 0.0).sum().item()
rel_err = (folded_back - product).abs() / product.clamp(min=1e-10)
print(f"\n L2 Fold results:")
print(f" Clamped to 448: {n_clamped}/{n_total} ({100*n_clamped/n_total:.1f}%)")
print(f" Zeroed (subnormal): {n_zeroed}/{n_total} ({100*n_zeroed/n_total:.1f}%)")
print(f" Rel error: max={rel_err.max():.4f} mean={rel_err.mean():.4f} p99={rel_err.quantile(0.99):.4f}")
# Show distribution of folded values
fb_hist = torch.histc(folded_back, bins=10, min=0, max=448)
print(f" Folded value histogram (0-448, 10 bins): {fb_hist.int().tolist()}")
# CRITICAL CHECK: is the product range within float8?
print(f" Product range: [{product.min():.4e}, {product.max():.4e}]")
if n_clamped > 0:
print(f" ⚠️ {n_clamped} values clamped — this IS precision loss!")
if w13_sf_key and w13_gs_key:
w13_sf = st.get_tensor(w13_sf_key)
w13_gs = st.get_tensor(w13_gs_key)
print(f"\n L1 block scale: shape={list(w13_sf.shape)} dtype={w13_sf.dtype} "
f"range=[{w13_sf.float().min():.4e}, {w13_sf.float().max():.4e}]")
print(f" L1 global scale: shape={list(w13_gs.shape)} dtype={w13_gs.dtype} "
f"range=[{w13_gs.float().min():.4e}, {w13_gs.float().max():.4e}]")
break # Just check one file that has layer 0
def check_l2_weight_semantics():
"""Verify L2 weight layout by dequantizing and checking against reference."""
print("\n" + "=" * 60)
print("CHECK 4: L2 Weight Dequantization Sanity")
print("=" * 60)
from safetensors import safe_open
import glob
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
for f in ckpt_files:
with safe_open(f, framework="pt") as st:
keys = list(st.keys())
# Find layer 0 w2 weight, weight_scale, weight_scale_2
w2_w = w2_sf = w2_gs = None
for k in keys:
if "layers.0" in k:
if "w2" in k and k.endswith(".weight") and "scale" not in k:
w2_w = st.get_tensor(k)
elif "w2" in k and "weight_scale" == k.split(".")[-1]:
w2_sf = st.get_tensor(k)
elif "w2" in k and "weight_scale_2" in k:
w2_gs = st.get_tensor(k)
if w2_w is not None and w2_sf is not None and w2_gs is not None:
print(f" w2_weight: shape={list(w2_w.shape)} dtype={w2_w.dtype}")
print(f" w2_weight_scale: shape={list(w2_sf.shape)} dtype={w2_sf.dtype}")
print(f" w2_weight_scale_2: shape={list(w2_gs.shape)} dtype={w2_gs.dtype}")
# Dequantize a small patch
# w2 is down_proj: (hidden, intermediate) in BF16, or (hidden, inter//2) uint8 for NVFP4
if w2_w.dtype == torch.uint8:
# Unpack E2M1
FP4_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6,
-0, -0.5, -1, -1.5, -2, -3, -4, -6],
dtype=torch.float32, device=w2_w.device)
lower = FP4_LUT[(w2_w[:4, :8] & 0x0F).long()]
upper = FP4_LUT[((w2_w[:4, :8] >> 4) & 0x0F).long()]
unpacked = torch.empty(4, 16, dtype=torch.float32)
unpacked[:, 0::2] = lower
unpacked[:, 1::2] = upper
# Apply scales
sf_slice = w2_sf[:4, :1].float() # (4, 1)
gs = w2_gs.float()
print(f" Dequantized w2[:4, :16] with sf[:4,:1]={sf_slice.flatten().tolist()}")
print(f" global_scale_2 = {gs.item() if gs.numel() == 1 else gs[:4].flatten().tolist()}")
dequant = unpacked * sf_slice * gs.float()
print(f" Dequantized range: [{dequant.min():.4f}, {dequant.max():.4f}]")
print(f" Dequantized[:2, :8]: {dequant[:2, :8].tolist()}")
else:
print(f" w2_weight is {w2_w.dtype}, not uint8 — may be BF16 checkpoint")
print(f" w2[:4, :8] = {w2_w[:4, :8].tolist()}")
break
def check_ep_reduce_contract():
"""Verify the EP all-reduce contract with a synthetic test."""
print("\n" + "=" * 60)
print("CHECK 5: EP Reduce Contract (Synthetic)")
print("=" * 60)
# Simulate 2 ranks
M, HIDDEN = 4, 8
# Rank 0: experts 0,1 — tokens routed to expert 0 (slot_weight=0.7) and 1 (slot_weight=0.3)
y0 = torch.zeros(M, HIDDEN, dtype=torch.bfloat16)
slot_token_0 = torch.tensor([0, 0, 1, 2, 3]) # which tokens
slot_weight_0 = torch.tensor([0.7, 0.3, 0.5, 0.6, 0.4], dtype=torch.bfloat16)
l2_slots_0 = torch.randn(5, HIDDEN, dtype=torch.bfloat16)
y0.index_add_(0, slot_token_0, l2_slots_0 * slot_weight_0.unsqueeze(1))
# Rank 1: experts 2,3 — token 0 also routed to expert 2
y1 = torch.zeros(M, HIDDEN, dtype=torch.bfloat16)
slot_token_1 = torch.tensor([0, 1])
slot_weight_1 = torch.tensor([0.2, 0.5], dtype=torch.bfloat16)
l2_slots_1 = torch.randn(2, HIDDEN, dtype=torch.bfloat16)
y1.index_add_(0, slot_token_1, l2_slots_1 * slot_weight_1.unsqueeze(1))
# All-reduce (sum)
y_final = y0 + y1 # simulated all-reduce
# Verify: token 0 should have contributions from rank0 (experts 0,1) and rank1 (expert 2)
expected_0 = (0.7 * l2_slots_0[0] + 0.3 * l2_slots_0[1] + 0.2 * l2_slots_1[0]).bfloat16()
actual_0 = y_final[0].bfloat16()
diff = (expected_0 - actual_0).abs().max().item()
print(f" Token 0: expected vs actual diff = {diff:.6f} ✓" if diff < 0.01 else f" Token 0: MISMATCH diff = {diff}")
print(f" EP reduce contract is correct — sum of partial rank outputs gives full result")
if __name__ == "__main__":
print("NVFP4 MegaMoE Diagnostic — B200")
print(f"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}")
print(f"GPUs: {torch.cuda.device_count()}")
print()
try:
inspect_checkpoint_scales()
except Exception as e:
print(f"CHECK 1 FAILED: {e}")
try:
check_fold_precision_real()
except Exception as e:
print(f"CHECK 3 FAILED: {e}")
try:
check_l2_weight_semantics()
except Exception as e:
print(f"CHECK 4 FAILED: {e}")
try:
check_ep_reduce_contract()
except Exception as e:
print(f"CHECK 5 FAILED: {e}")

66
diag_fold.py
View File

@@ -1,66 +0,0 @@
"""
Diagnostic: Check global scale folding precision for NVFP4 weights.
The fold is: sf_f32 * gs → clamp(0, 448) → float8_e4m3fn
Question: how much precision is lost in the float8 round-trip?
"""
import torch
# Simulate typical NVFP4 scale distributions
# block_scale (float8_e4m3fn) range: roughly 0.06 to 448
# global_scale (float32) range: varies per expert
# Test 1: If global_scale >> 1, product can exceed 448 → clamp → loss
# Test 2: If global_scale << 1, product can go subnormal → loss
# Test 3: Quantization error from 3-bit mantissa
# Simulate a range of scale values
block_scales = torch.tensor([0.0625, 0.125, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 448.0], dtype=torch.float32)
global_scales = torch.tensor([0.001, 0.01, 0.1, 1.0, 10.0, 100.0, 1000.0], dtype=torch.float32)
print("=== Float8 Folding Precision Analysis ===\n")
print(f"block_scales: {block_scales.tolist()}")
print(f"global_scales: {global_scales.tolist()}\n")
total_clamped = 0
total_subnormal = 0
max_rel_error = 0.0
for gs in global_scales:
products = block_scales * gs
clamped = products.clamp(0.0, 448.0)
folded_f8 = clamped.to(torch.float8_e4m3fn)
roundtrip = folded_f8.to(torch.float32)
n_clamped = (products > 448.0).sum().item()
n_subnormal = (roundtrip > 0).logical_and(roundtrip < 0.0625).sum().item() # rough check
rel_errors = torch.where(roundtrip > 0, (roundtrip - clamped).abs() / clamped.clamp(min=1e-10), torch.zeros_like(clamped))
max_err = rel_errors.max().item()
total_clamped += n_clamped
total_subnormal += n_subnormal
max_rel_error = max(max_rel_error, max_err)
if n_clamped > 0 or max_err > 0.05:
print(f"gs={gs:.3f}: {n_clamped} clamped, max_rel_err={max_err:.4f}")
for i, (p, c, r) in enumerate(zip(products, clamped, roundtrip)):
if abs(r - c) / max(abs(c), 1e-10) > 0.01:
print(f" block={block_scales[i]:.4f} product={p:.4f} clamped={c:.4f} roundtrip={r:.4f} err={abs(r-c)/max(abs(c),1e-10):.4f}")
print(f"\nTotal clamped: {total_clamped}, Total subnormal: {total_subnormal}, Max relative error: {max_rel_error:.4f}")
# The real check: what's the float8_e4m3fn step size at various magnitudes?
print("\n=== Float8 E4M3 Step Sizes ===")
test_vals = [0.01, 0.1, 1.0, 10.0, 100.0, 448.0]
for v in test_vals:
f8 = torch.tensor(v, dtype=torch.float32).to(torch.float8_e4m3fn)
back = f8.to(torch.float32)
# Find next representable value
u8 = f8.view(torch.uint8)
next_u8 = u8 + 1
next_f8 = next_u8.view(torch.float8_e4m3fn)
next_val = next_f8.to(torch.float32)
step = next_val - back
rel_step = step / back if back > 0 else 0
print(f" value={v:.3f} → f8={back:.6f} → next={next_val:.6f} step={step:.6f} rel={rel_step:.4f}")

96
diag_fold_real.py
View File

@@ -1,96 +0,0 @@
"""
Critical check: weight_scale_2 values are ~4.65e-05 (TINY).
When folded: block_sf * 4.65e-05 → most products near zero → float8 can't represent
This is likely THE bug: folding a float8 scale by a tiny global scale produces
subnormal/zero values in float8.
"""
from safetensors import safe_open
import glob
import os
import torch
MODEL_PATH = "/model"
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
# Get layer 0, expert 0 scales
for f in ckpt_files:
with safe_open(f, framework="pt") as st:
keys = list(st.keys())
if any("layers.0.mlp.experts.0.gate_proj.weight_scale" in k for k in keys):
# Gate
gate_sf = st.get_tensor("model.layers.0.mlp.experts.0.gate_proj.weight_scale")
gate_gs = st.get_tensor("model.layers.0.mlp.experts.0.gate_proj.weight_scale_2")
# Up
up_sf = st.get_tensor("model.layers.0.mlp.experts.0.up_proj.weight_scale")
up_gs = st.get_tensor("model.layers.0.mlp.experts.0.up_proj.weight_scale_2")
# Down
down_sf = st.get_tensor("model.layers.0.mlp.experts.0.down_proj.weight_scale")
down_gs = st.get_tensor("model.layers.0.mlp.experts.0.down_proj.weight_scale_2")
print("=" * 60)
print("LAYER 0, EXPERT 0 — Scale Analysis")
print("=" * 60)
for name, sf, gs in [("gate", gate_sf, gate_gs), ("up", up_sf, up_gs), ("down", down_sf, down_gs)]:
sf_f32 = sf.float()
gs_f32 = gs.float()
product = sf_f32 * gs_f32
product_clamped = product.clamp(0.0, 448.0)
folded_f8 = product_clamped.to(torch.float8_e4m3fn)
folded_back = folded_f8.float()
n_total = product.numel()
n_clamped = (product > 448.0).sum().item()
n_zeroed = (folded_back == 0.0).sum().item()
n_nonzero_orig = (sf_f32 > 0).sum().item()
n_nonzero_folded = (folded_back > 0).sum().item()
rel_err = (folded_back - product).abs() / product.clamp(min=1e-10)
print(f"\n {name}_proj:")
print(f" block_sf: shape={list(sf.shape)} range=[{sf_f32.min():.4e}, {sf_f32.max():.4e}] unique_u8={torch.unique(sf.view(torch.uint8)).numel()}")
print(f" global_sf: {gs_f32.item():.6e}")
print(f" product (sf*gs): range=[{product.min():.4e}, {product.max():.4e}]")
print(f" folded (float8): range=[{folded_back.min():.4e}, {folded_back.max():.4e}]")
print(f" Clamped to 448: {n_clamped}/{n_total} ({100*n_clamped/n_total:.1f}%)")
print(f" Became zero: {n_zeroed}/{n_total} ({100*n_zeroed/n_total:.1f}%)")
print(f" Was nonzero → became zero: {n_nonzero_orig - n_nonzero_folded}/{n_nonzero_orig}")
print(f" Rel error: max={rel_err.max():.4f} mean={rel_err.mean():.4f}")
# Show the float8 step size at the product magnitude
if product.max() > 0:
typical = product.median().item()
if typical > 0:
f8_typ = torch.tensor(typical, dtype=torch.float32).to(torch.float8_e4m3fn)
f8_back = f8_typ.float()
if f8_back > 0:
step = (f8_typ.view(torch.uint8) + 1).view(torch.float8_e4m3fn).float() - f8_back
print(f" Float8 step at median ({typical:.4e}): Δ={step.item():.4e} rel={step.item()/f8_back.item():.2%}")
break
# Now check: what if we DON'T fold, and instead pass global_scale as GEMM alpha?
print("\n" + "=" * 60)
print("ALTERNATIVE: Pass global_scale as GEMM alpha")
print("=" * 60)
print("""
The fold is lossy because float8 can't represent the product range.
But if we DON'T fold, the CUTLASS GEMM needs a separate global scale mechanism.
Option 1: Multiply the GEMM alpha by the weight's global_scale
- alpha already carries the activation global scale
- We could fold weight global scale into alpha: alpha_new = alpha * weight_gs
- BUT: alpha is a single scalar, weight_gs varies per-expert
- For grouped GEMM, each expert needs its own alpha
Option 2: Keep block scales as-is (no fold), multiply output by global_scale
- After GEMM: output *= weight_global_scale
- This is exact (float32 multiply on bf16 output)
- Requires passing global_scale to nvfp4_mega_moe_full
Option 3: Fold global_scale into the GEMM alpha per-expert
- In cutlass_grouped_nvfp4_gemm, each expert gets its own alpha
- alpha_expert = l1_global_scale * l1_weight_global_scale[expert_id]
- This is EXACT and doesn't lose precision
- The block scales stay at their original float8 values (no folding)
""")

328
diag_issues.py
View File

@@ -1,328 +0,0 @@
"""
Diagnostic script for NVFP4 mega_moe issues.
Run on the B200 server. Checks:
1. Global scale folding precision (float8 round-trip)
2. L2 weight/SF orientation (transpose correctness)
3. EP aggregation contract (local vs all-reduce)
4. Folded scale float8 precision loss
Usage: python diag_issues.py
"""
import torch
import sys
import os
# Try to import the model components
try:
from nvfp4_megamoe_kernel import (
transform_nvfp4_weights_for_mega_moe,
stage_activation,
nvfp4_mega_moe_full,
)
HAS_KERNEL = True
except ImportError:
HAS_KERNEL = False
print("WARNING: nvfp4_megamoe_kernel not importable, some checks will be skipped")
def check_fold_precision():
"""Check 1: Float8 folding precision.
The fold is: sf_f32 * gs → clamp(0, 448) → float8_e4m3fn
Question: are we silently destroying critical precision?
"""
print("=" * 60)
print("CHECK 1: Global Scale Folding Precision")
print("=" * 60)
# Simulate realistic scale distributions
# NVFP4 block scales (float8_e4m3fn) are typically in range [0.06, 448]
# Global scales are per-expert float32
# Test with realistic ranges
for gs_val in [0.001, 0.01, 0.1, 1.0, 10.0, 100.0]:
# Simulate 1000 block scales
sf = torch.rand(48, 64, 192) * 448 # Smaller for quantile perf
sf_f8 = sf.clamp(0.0, 448.0).to(torch.float8_e4m3fn)
sf_back = sf_f8.to(torch.float32)
# Fold: product then cast back
product = sf_back * gs_val
product_clamped = product.clamp(0.0, 448.0)
folded_f8 = product_clamped.to(torch.float8_e4m3fn)
folded_back = folded_f8.to(torch.float32)
# Compare against the "correct" product (sf_f32 * gs, no float8 intermediate)
correct_product = sf * gs_val
# Count how many values are lost to clamping or zero
n_clamped = (product > 448.0).sum().item()
n_zeroed = (folded_back == 0.0).sum().item() - (correct_product == 0.0).sum().item()
# Relative error
rel_err = (folded_back - correct_product).abs() / correct_product.clamp(min=1e-10)
max_rel = rel_err.max().item()
mean_rel = rel_err.mean().item()
p99_rel = rel_err.quantile(0.99).item()
print(f" gs={gs_val:>8.3f}: clamped={n_clamped:>8d} zeroed={n_zeroed:>8d} "
f"max_rel={max_rel:.4f} mean_rel={mean_rel:.4f} p99_rel={p99_rel:.4f}")
def check_l2_orientation():
"""Check 2: L2 weight/SF orientation.
The down_proj maps intermediate→hidden. In PyTorch, weight is (out, in) = (hidden, intermediate).
After NVFP4 packing: (hidden, intermediate//2).
After transpose for CUTLASS col-major B: (intermediate//2, hidden).
The CUTLASS GEMM computes: D = alpha * A @ B where A is (M, K) and B is (K, N).
K = intermediate (contraction dim), N = hidden (output dim).
Packed B is (K_half, N) in memory (column-major for CUTLASS).
Question: is the transpose correct for the CUTLASS B layout?
"""
print("\n" + "=" * 60)
print("CHECK 2: L2 Weight/SF Orientation")
print("=" * 60)
# Simulate L2 weight and SF
E, HIDDEN, INTER = 48, 7168, 3072
K_half = INTER // 2 # 1536
sf_K = INTER // 16 # 192
# Checkpoint shapes
w2_weight_shape = (E, HIDDEN, K_half) # (E, N_out, K_in//2)
w2_sf_shape = (E, HIDDEN, sf_K) # (E, N_out, sf_K)
# After transpose
w2_weight_transposed = (E, K_half, HIDDEN) # (E, K_half, N) — CUTLASS col-major B
w2_sf_transposed = (E, sf_K, HIDDEN) # (E, sf_K, N)
# CUTLASS expects for the grouped GEMM:
# weights: (E, K_half, N) ✓
# weight_sf: (E, sf_K, N) — but which is K_sf and which is N?
# The remap kernel gets: MN=N=HIDDEN, K_sf=INTER//16=192, col_major_src=true
# Source is (K_sf, MN) = (192, 7168) row-major ✓
print(f" Checkpoint w2_weight: {w2_weight_shape}")
print(f" Checkpoint w2_sf: {w2_sf_shape}")
print(f" After transpose w2_weight: {w2_weight_transposed}")
print(f" After transpose w2_sf: {w2_sf_transposed}")
print(f" CUTLASS expects B: (K_half={K_half}, N={HIDDEN})")
print(f" CUTLASS expects SFB: (K_sf={sf_K}, N={HIDDEN})")
print(f" ✓ Shapes match")
# BUT: check if the DATA is semantically correct
# The transpose swaps (N, K_half) → (K_half, N)
# For the weight, row i of (N, K_half) becomes column i of (K_half, N)
# In row-major, element [i,j] of (N, K_half) goes to offset i*K_half + j
# After transpose, it's at offset j*N + i in (K_half, N)
# CUTLASS column-major B reads logical (n,k) at offset n + k*N
# Where n is the output dim (hidden) and k is the contraction dim (intermediate)
# For packed FP4: k ranges 0..K_half-1 (2 values per byte)
# So logical (n, k_half) at offset n + k_half * N
# Our data: element at memory offset k_half * N + n (row-major (K_half, N))
# = k_half * N + n = n + k_half * N ← SAME ✓
print(f" ✓ CUTLASS column-major stride matches our row-major (K_half, N) layout")
def check_ep_aggregation():
"""Check 3: EP aggregation contract.
Each rank computes y = sum over local experts of (routing_weight * expert_output).
Then all-reduce sums across EP ranks.
The contract is: final_y = sum_ranks(y_rank)
Question: is the local y correctly computed such that the all-reduce gives the right answer?
"""
print("\n" + "=" * 60)
print("CHECK 3: EP Aggregation Contract")
print("=" * 60)
# Simulate: 2 EP ranks, 4 total experts, topk=2
# Rank 0 has experts 0,1; Rank 1 has experts 2,3
# Token is routed to experts 0 and 2 (one per rank)
# On Rank 0: slot for expert 0, slot_weight * l2_output → index_add to y
# On Rank 1: slot for expert 2, slot_weight * l2_output → index_add to y
# All-reduce: y_final = y_rank0 + y_rank1 ✓
# POTENTIAL ISSUE: what if the same token is routed to multiple experts
# on the same rank? index_add_ handles this correctly (sums in-place).
# POTENTIAL ISSUE: what if a token has NO experts on a rank?
# y stays at 0 for that token → correct, other ranks contribute.
# POTENTIAL ISSUE: is slot_weight correctly applied?
# In nvfp4_mega_moe_full:
# y.index_add_(0, slot_token, l2_slots * slot_weight.unsqueeze(1))
# l2_slots is (num_slots, HIDDEN) bf16
# slot_weight is (num_slots,) float32, unsqueezed to (num_slots, 1)
# So each slot output is scaled by its routing weight before accumulating.
# This is correct: final = sum_k(w_k * expert_k(x))
print(" ✓ Local index_add_ + all-reduce contract is correct")
print(" ✓ slot_weight applied before index_add (correct)")
print(" NOTE: This assumes all-reduce uses SUM (not AVG). Verify with torch.distributed.")
# Check the vllm code uses all_reduce (sum by default)
# torch.distributed.all_reduce defaults to ReduceOp.SUM ✓
def check_fold_vs_nofold():
"""Check 4: What happens if global scale is NOT folded?
If weight_scale_2 is not folded into the block scales, the weights are
effectively used without their global scaling factor. This would produce
finite but semantically garbage output — exactly the symptom.
"""
print("\n" + "=" * 60)
print("CHECK 4: Global Scale Folding Verification")
print("=" * 60)
# The fold happens in transform_nvfp4_weights_for_mega_moe:
# 1. sf_f32 = weight_scale.to(float32)
# 2. sf_f32 *= weight_scale_2 (global scale)
# 3. sf_out = sf_f32.clamp(0, 448).to(float8_e4m3fn)
# If weight_scale_2 is None (not provided), the fold is skipped
# and only block scales are used. This would be a bug.
# Check: is weight_scale_2 actually non-None when finalize_weights is called?
# From the code:
# transform_nvfp4_weights_for_mega_moe(
# ..., l1_weight_scale_2=self.w13_weight_scale_2.data.contiguous(), ...)
# self.w13_weight_scale_2 is initialized as nn.Parameter(torch.zeros(num_local_experts, 2))
# It's loaded from checkpoint in weight_loader (shard_id w1→[e,0], w3→[e,1])
# If the checkpoint doesn't contain weight_scale_2 for experts,
# the parameter stays at zeros. Folding with gs=0 → all scales become 0 → garbage.
print(" If weight_scale_2 is all zeros (not loaded from checkpoint):")
sf = torch.tensor([1.0, 2.0, 4.0, 8.0, 16.0])
gs_zero = 0.0
folded = (sf * gs_zero).clamp(0, 448).to(torch.float8_e4m3fn)
print(f" sf={sf.tolist()} * gs=0 → folded={folded.to(torch.float32).tolist()}")
print(" ALL SCALES GO TO ZERO → all outputs are zero → garbage")
print("\n If weight_scale_2 is correctly loaded (typical values):")
gs = torch.tensor([0.5, 1.0, 2.0, 5.0, 10.0])
for g in gs:
folded = (sf * g).clamp(0, 448).to(torch.float8_e4m3fn)
correct = sf * g
rel_err = ((folded.to(torch.float32) - correct).abs() / correct).mean()
print(f" gs={g:.1f}: mean_rel_err={rel_err:.4f}")
def check_l2_sf_transpose_semantics():
"""Check 5: After transposing L2 SF, is the data in the right layout?
The w2_weight_scale in checkpoint is (E, N, sf_K) = (E, hidden, inter//16).
This means: for each expert, for each output row (hidden dim), we have sf_K block scales
along the input dimension.
After transpose: (E, sf_K, N) = (E, inter//16, hidden).
This means: for each expert, for each block along the input dim, we have N=hidden scale values.
CUTLASS SFB is (K_sf, N) where K_sf is the contraction dim's scale groups.
K_sf = K // 16 = inter // 16. N = hidden.
The CUTLASS remap expects col_major_src=True, so it reads src[k_sf * N + m].
With N=hidden and K_sf=inter//16, this accesses the (E, inter//16, hidden) tensor correctly.
BUT WAIT: The CUTLASS SFB layout is defined for the B matrix which is ColumnMajor.
For ColumnMajor B with shape (N, K), the SFB layout might have a different
semantic mapping than what we're providing.
Let me check: does CUTLASS SFB index by (N_idx, K_sf_idx) or (K_sf_idx, N_idx)?
"""
print("\n" + "=" * 60)
print("CHECK 5: L2 SF Transpose Semantics (Deep Dive)")
print("=" * 60)
# The key question: after the transpose, does SFB[i, j] contain the right value?
#
# Original (checkpoint): weight_scale[E, hidden_row, sf_k_block]
# = the block scale for expert E, output row hidden_row, input block sf_k_block
#
# The GEMM operation: Y = X @ W where W is (K, N) = (inter, hidden)
# SFB should be: for each output column n and each input block k_sf,
# SFB[n, k_sf] = scale for column n, input block k_sf
# OR (depending on CUTLASS convention):
# SFB[k_sf, n] = scale for input block k_sf, output column n
#
# In CUTLASS NVFP4, SFB has the same (K, N) structure as B.
# B is ColumnMajor (N, K), so B[n, k] is at memory n + k * N.
# SFB should follow the same (N, K_sf) → ColumnMajor → (K_sf, N) row-major in memory.
#
# Our source (after transpose): (E, sf_K, N) = (E, K_sf, N) row-major
# Element [e, k_sf, n] = original [e, n, k_sf] = checkpoint scale for expert e, output n, input block k_sf
# The remap reads: src[k_sf * N + n] (col_major_src=true)
# = element [e, k_sf, n] = correct scale for (n, k_sf) in the B matrix
# ✓ This is correct!
print(" L2 SF transpose semantics are correct")
print(" After transpose: (E, K_sf, N) with col_major_src=True")
print(" remap reads src[k_sf * N + n] = original scale[e, n, k_sf] ✓")
def check_w13_gate_up_split():
"""Check 6: Is the gate/up split for w13 scale_2 folding aligned
with the actual weight layout after transpose?
w13_weight shape: (E, 2*INTER, HIDDEN//2)
w13_weight_scale shape: (E, 2*INTER, HIDDEN//16)
The fold splits: gate = first INTER rows, up = last INTER rows
Then applies gs[:,0] to gate, gs[:,1] to up
After transpose:
w13_weight: (E, HIDDEN//2, 2*INTER)
w13_sf: (E, HIDDEN//16, 2*INTER)
The gate/up split is now along the LAST dim (N), not the middle.
But the fold happens BEFORE the transpose, so the split is correct.
After transpose, the gate portion is columns 0..INTER-1 and up is INTER..2*INTER-1.
This is still semantically correct for the CUTLASS GEMM.
"""
print("\n" + "=" * 60)
print("CHECK 6: w13 Gate/Up Split Alignment")
print("=" * 60)
print(" Fold happens before transpose → gate/up split is on dim 1 (N)")
print(" After transpose, split is on dim 2 (N) — last dimension")
print(" CUTLASS GEMM sees N=2*INTER with gate first, up second ✓")
print(" The folded scales correctly reflect gate_gs and up_gs ✓")
if __name__ == "__main__":
if not torch.cuda.is_available():
print("WARNING: No CUDA — some checks will be approximate")
check_fold_precision()
check_l2_orientation()
check_ep_aggregation()
check_fold_vs_nofold()
check_l2_sf_transpose_semantics()
check_w13_gate_up_split()
print("\n" + "=" * 60)
print("SUMMARY")
print("=" * 60)
print("""
Most likely suspects for "finite but garbage" output:
1. weight_scale_2 not loaded → all-zero global scales → folded sf = 0
CHECK: Print w13_weight_scale_2 and w2_weight_scale_2 after loading
2. Float8 folding precision: 12-95% relative error for small global scales
This is a QUALITY issue, not a garbage issue
BUT: if global scales are very small (<<1), entire scale groups zero out
3. L2 weight/SF: shapes and semantics look correct after analysis
The transpose + CUTLASS col-major + SFB remap are consistent
4. EP aggregation: contract looks correct (local sum + all_reduce)
ACTION ITEMS:
a) Run the model with debug prints showing weight_scale_2 values
b) Check if any folded scales clamp to 0 or 448 (precision ceiling)
c) Compare folded sf values against reference (unfolded) computation
d) Test with a single expert to isolate EP issues
""")

39
diag_keys.py
View File

@@ -1,39 +0,0 @@
"""Find ALL weight_scale_2 keys in the checkpoint for layer 0 experts."""
from safetensors import safe_open
import glob
import os
MODEL_PATH = "/model"
ckpt_files = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
# Collect ALL keys that mention layer 0 experts and scale
scale_keys = []
for f in ckpt_files:
with safe_open(f, framework="pt") as st:
for key in st.keys():
if "layers.0" in key and "experts.0" in key and "scale" in key.lower():
val = st.get_tensor(key)
scale_keys.append((key, list(val.shape), str(val.dtype), val.float().min().item(), val.float().max().item()))
scale_keys.sort()
for k, s, d, mn, mx in scale_keys:
print(f" {k} shape={s} dtype={d} range=[{mn:.4e}, {mx:.4e}]")
print(f"\nTotal: {len(scale_keys)} scale keys for layer 0 expert 0")
# Also find gate_proj and up_proj weight_scale_2 keys
print("\n--- All weight_scale_2 keys with gate/up/down for layer 0 ---")
ws2_keys = []
for f in ckpt_files:
with safe_open(f, framework="pt") as st:
for key in st.keys():
if "layers.0" in key and "weight_scale_2" in key:
val = st.get_tensor(key)
ws2_keys.append((key, list(val.shape), str(val.dtype), val.float().min().item(), val.float().max().item()))
ws2_keys.sort()
for k, s, d, mn, mx in ws2_keys[:10]:
print(f" {k} shape={s} dtype={d} range=[{mn:.4e}, {mx:.4e}]")
if len(ws2_keys) > 10:
print(f" ... and {len(ws2_keys)-10} more")
print(f"Total: {len(ws2_keys)} weight_scale_2 keys for layer 0")