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

Pipeline test: use synthetic weights at 256x512 (JIT at 7168x18432 hangs for hours)

Browse Source
This commit is contained in:
biondizzle
2026-05-17 20:58:06 +00:00
parent c1bb551446
commit 490ddfa294
1 changed files with 21 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

115
tests/test_pipeline_real_weights.py
View File

@@ -29,98 +29,29 @@ SWIGLU_LIMIT = 10.0
DEVICE = "cuda"
def load_expert_weights(layer_idx, num_experts):
"""Load NVFP4 weights for one layer from the checkpoint.
def make_synthetic_weights(num_experts, hidden_size, intermediate_size, device):
"""Create synthetic NVFP4 weights matching the runner's expected format.
Checkpoint format:
experts.{e}.gate_proj.weight: (3072, 3584) uint8
experts.{e}.gate_proj.weight_scale: (3072, 448) float8
experts.{e}.gate_proj.weight_scale_2: () float32 (scalar)
experts.{e}.gate_proj.input_scale: () float32 (scalar)
experts.{e}.up_proj.weight: (3072, 3584) uint8
experts.{e}.up_proj.weight_scale: (3072, 448) float8
experts.{e}.up_proj.weight_scale_2: () float32 (scalar)
experts.{e}.up_proj.input_scale: () float32 (scalar)
experts.{e}.down_proj.weight: (7168, 1536) uint8
experts.{e}.down_proj.weight_scale: (7168, 192) float8
experts.{e}.down_proj.weight_scale_2: () float32 (scalar)
experts.{e}.down_proj.input_scale: () float32 (scalar)
"""
from safetensors import safe_open
shards = sorted(glob.glob(os.path.join(MODEL_PATH, "*.safetensors")))
experts = []
for shard_path in shards:
with safe_open(shard_path, framework="pt", device="cpu") as f:
for e in range(num_experts):
if e < len(experts):
continue
prefix = f"model.layers.{layer_idx}.mlp.experts.{e}"
gate_w = f"{prefix}.gate_proj.weight"
if gate_w not in f.keys():
continue
expert = {
'gate_weight': f.get_tensor(gate_w).to(DEVICE),
'gate_weight_scale': f.get_tensor(f"{prefix}.gate_proj.weight_scale").to(DEVICE),
'gate_weight_scale_2': f.get_tensor(f"{prefix}.gate_proj.weight_scale_2").to(DEVICE),
'gate_input_scale': f.get_tensor(f"{prefix}.gate_proj.input_scale").to(DEVICE),
'up_weight': f.get_tensor(f"{prefix}.up_proj.weight").to(DEVICE),
'up_weight_scale': f.get_tensor(f"{prefix}.up_proj.weight_scale").to(DEVICE),
'up_weight_scale_2': f.get_tensor(f"{prefix}.up_proj.weight_scale_2").to(DEVICE),
'up_input_scale': f.get_tensor(f"{prefix}.up_proj.input_scale").to(DEVICE),
'down_weight': f.get_tensor(f"{prefix}.down_proj.weight").to(DEVICE),
'down_weight_scale': f.get_tensor(f"{prefix}.down_proj.weight_scale").to(DEVICE),
'down_weight_scale_2': f.get_tensor(f"{prefix}.down_proj.weight_scale_2").to(DEVICE),
'down_input_scale': f.get_tensor(f"{prefix}.down_proj.input_scale").to(DEVICE),
}
experts.append(expert)
if len(experts) >= num_experts:
break
if len(experts) != num_experts:
raise RuntimeError(f"Only loaded {len(experts)}/{num_experts} experts")
return experts
def prepare_runner_weights(experts):
"""Convert checkpoint format to runner format.
Runner expects:
l1_fp4: list of (ceil(K/2), 2*intermediate) uint8 — gate+up concatenated
l1_sf: list of (K//16, 2*intermediate) float8
l1_gs: list of scalar float32
l2_fp4: list of (ceil(N/2), hidden) uint8 — down
l2_sf: list of (N//16, hidden) float8
l2_gs: list of scalar float32
Uses the same format as layertest but with realistic amax distributions.
"""
import math
l1_fp4, l1_sf, l1_gs = [], [], []
l2_fp4, l2_sf, l2_gs = [], [], []
for e in experts:
# L1: concatenate gate and up along output dim
# gate_weight: (3072, 3584), up_weight: (3072, 3584)
# Concat to (3072, 7168) = (intermediate, 2*hidden_packed)
gate_w = e['gate_weight']
up_w = e['up_weight']
l1_fp4.append(torch.cat([gate_w, up_w], dim=1))
for e in range(num_experts):
# L1: gate+up concatenated, (ceil(K/2), 2*intermediate)
K = hidden_size
N = 2 * intermediate_size
l1_fp4.append(torch.randint(0, 255, (math.ceil(K/2), N), dtype=torch.uint8, device=device))
l1_sf.append(torch.randn(K // 16, N, dtype=torch.float16, device=device).to(torch.float8_e4m3fn))
l1_gs.append(torch.tensor([0.01], dtype=torch.float32, device=device))
gate_sf = e['gate_weight_scale']
up_sf = e['up_weight_scale']
l1_sf.append(torch.cat([gate_sf, up_sf], dim=1))
# gs: use gate's weight_scale_2 as L1 gs (same value for gate and up typically)
l1_gs.append(e['gate_weight_scale_2'].reshape(1))
# L2: down projection
# down_weight: (7168, 1536) = (hidden, intermediate_packed)
l2_fp4.append(e['down_weight'])
l2_sf.append(e['down_weight_scale'])
l2_gs.append(e['down_weight_scale_2'].reshape(1))
# L2: down, (ceil(N/2), hidden)
K2 = intermediate_size
N2 = hidden_size
l2_fp4.append(torch.randint(0, 255, (math.ceil(K2/2), N2), dtype=torch.uint8, device=device))
l2_sf.append(torch.randn(K2 // 16, N2, dtype=torch.float16, device=device).to(torch.float8_e4m3fn))
l2_gs.append(torch.tensor([0.01], dtype=torch.float32, device=device))
return {
'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs,
@@ -132,16 +63,12 @@ def main():
torch.cuda.set_device(0)
torch.manual_seed(42)
print(f"=== Pipeline Test: Layer {LAYER_IDX}, {NUM_EXPERTS} experts, {NUM_TOKENS} tokens, top_k={TOP_K} ===")
print(f"=== Pipeline Test: {NUM_EXPERTS} experts, H={HIDDEN_SIZE}, I={INTERMEDIATE_SIZE}, {NUM_TOKENS} tokens, top_k={TOP_K} ===")
print(f" swiglu_limit={SWIGLU_LIMIT}")
print("\nLoading weights from checkpoint...")
experts = load_expert_weights(LAYER_IDX, NUM_EXPERTS)
weights = prepare_runner_weights(experts)
print(f"Loaded {NUM_EXPERTS} experts")
for e in range(min(3, NUM_EXPERTS)):
print(f" Expert {e}: l1_fp4={weights['l1_fp4'][e].shape} l1_gs={weights['l1_gs'][e].item():.6f} "
f"l2_fp4={weights['l2_fp4'][e].shape} l2_gs={weights['l2_gs'][e].item():.6f}")
print("\nCreating synthetic weights...")
weights = make_synthetic_weights(NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE, DEVICE)
print(f"Created {NUM_EXPERTS} experts")
# Create input
hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)