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

Rewrite pipeline test: use raw checkpoint weights, compare runner vs dynamic-gs reference

Browse Source
This commit is contained in:
biondizzle
2026-05-17 18:10:05 +00:00
parent e51eafe288
commit 5e63a0d8a3
1 changed files with 183 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

308
tests/test_pipeline_real_weights.py
View File

@@ -7,7 +7,6 @@ import torch
import sys
import os
import glob
import math
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/..')
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/../vllm')
@@ -22,47 +21,106 @@ LAYER_IDX = 0
NUM_EXPERTS = 48
HIDDEN_SIZE = 7168
INTERMEDIATE_SIZE = 18432
# Note: gate and up each have INTERMEDIATE_SIZE outputs
# L1 GEMM output = 2 * INTERMEDIATE_SIZE
NUM_TOKENS = 64
TOP_K = 6
SWIGLU_LIMIT = 10.0
DEVICE = "cuda"
def load_expert_weights(layer_idx, num_experts):
"""Load NVFP4 weights for one layer from the checkpoint."""
"""Load NVFP4 weights for one layer from the checkpoint.
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")))
l1_fp4 = []
l1_sf = []
l1_gs = []
l2_fp4 = []
l2_sf = []
l2_gs = []
experts = []
for shard_path in shards:
with safe_open(shard_path, framework="pt", device="cpu") as f:
for e in range(num_experts):
w13_key = f"model.layers.{layer_idx}.mlp.experts.{e}.w13_weight"
sf13_key = f"model.layers.{layer_idx}.mlp.experts.{e}.w13_weight_scale"
gs13_key = f"model.layers.{layer_idx}.mlp.experts.{e}.w13_weight_scale_2"
w2_key = f"model.layers.{layer_idx}.mlp.experts.{e}.w2_weight"
sf2_key = f"model.layers.{layer_idx}.mlp.experts.{e}.w2_weight_scale"
gs2_key = f"model.layers.{layer_idx}.mlp.experts.{e}.w2_weight_scale_2"
if len(experts) > e:
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
if w13_key in f.keys() and len(l1_fp4) <= e:
l1_fp4.append(f.get_tensor(w13_key).to(DEVICE))
l1_sf.append(f.get_tensor(sf13_key).to(DEVICE))
l1_gs.append(f.get_tensor(gs13_key).to(DEVICE))
l2_fp4.append(f.get_tensor(w2_key).to(DEVICE))
l2_sf.append(f.get_tensor(sf2_key).to(DEVICE))
l2_gs.append(f.get_tensor(gs2_key).to(DEVICE))
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(l1_fp4) >= num_experts:
if len(experts) >= num_experts:
break
if len(l1_fp4) != num_experts:
raise RuntimeError(f"Only loaded {len(l1_fp4)}/{num_experts} experts")
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
"""
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))
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))
return {
'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs,
@@ -70,95 +128,6 @@ def load_expert_weights(layer_idx, num_experts):
}
def run_reference(hidden_states, topk_weights, topk_ids, weights, swiglu_limit=None):
"""Reference MoE: per-expert processing with dynamic gs (quantize_to_nvfp4)."""
from cutedsl.quantize import quantize_to_nvfp4
from cutedsl.gemm import run_nvfp4_grouped_gemm
num_tokens = hidden_states.shape[0]
top_k = topk_ids.shape[1]
num_experts = len(weights['l1_fp4'])
# Sort tokens by expert
flat_ids = topk_ids.reshape(-1).cpu().numpy()
flat_weights = topk_weights.reshape(-1)
token_indices = torch.arange(num_tokens).unsqueeze(1).expand(-1, top_k).reshape(-1)
sort_idx = torch.argsort(topk_ids.reshape(-1), stable=True)
sorted_ids = topk_ids.reshape(-1)[sort_idx]
sorted_weights = topk_weights.reshape(-1)[sort_idx]
sorted_token_ids = token_indices[sort_idx]
# Compute expert offsets
expert_id_range = torch.arange(num_experts)
tokens_per_expert = (sorted_ids.unsqueeze(1).cpu() == expert_id_range.unsqueeze(0)).sum(dim=0)
expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32)
for e in range(num_experts):
expert_offsets[e + 1] = expert_offsets[e] + tokens_per_expert[e].item()
num_slots = num_tokens * top_k
slot_hidden = hidden_states[sorted_token_ids]
# Stack weights for GEMM
l1_mat_b, l1_scale_b, l1_gsb = _stack_weights(weights['l1_fp4'], weights['l1_sf'], weights['l1_gs'])
l2_mat_b, l2_scale_b, l2_gsb = _stack_weights(weights['l2_fp4'], weights['l2_sf'], weights['l2_gs'])
# L1 with dynamic gs
x_fp4, x_sf, gs_val = quantize_to_nvfp4(slot_hidden)
l1_gsa = torch.full((num_experts,), gs_val, dtype=torch.float32, device=DEVICE)
l1_scale_a = _assemble_scales_ref(x_sf, expert_offsets, num_experts)
l1_out = run_nvfp4_grouped_gemm(
mat_a=x_fp4, mat_b=l1_mat_b,
scale_a=l1_scale_a, scale_b=l1_scale_b,
expert_offsets=expert_offsets[1:],
global_scale_a=l1_gsa, global_scale_b=l1_gsb,
)
# SiLU(gate) * up
gate = l1_out[:, :INTERMEDIATE_SIZE]
up = l1_out[:, INTERMEDIATE_SIZE:]
gate_silu = torch.nn.functional.silu(gate)
if swiglu_limit is not None:
gate_silu = gate_silu.clamp(max=swiglu_limit)
up = up.clamp(min=-swiglu_limit, max=swiglu_limit)
activated = gate_silu * up
# L2 with dynamic gs
l2_x_fp4, l2_x_sf, l2_gs_val = quantize_to_nvfp4(activated)
l2_gsa = torch.full((num_experts,), l2_gs_val, dtype=torch.float32, device=DEVICE)
l2_scale_a = _assemble_scales_ref(l2_x_sf, expert_offsets, num_experts)
l2_out = run_nvfp4_grouped_gemm(
mat_a=l2_x_fp4, mat_b=l2_mat_b,
scale_a=l2_scale_a, scale_b=l2_scale_b,
expert_offsets=expert_offsets[1:],
global_scale_a=l2_gsa, global_scale_b=l2_gsb,
)
# Scatter-add
y = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)
weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype)
y.scatter_add_(0, sorted_token_ids.unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted_out)
return y
def _stack_weights(fp4_list, sf_list, gs_list):
"""Stack expert weights into GEMM format."""
mat_b = torch.stack(fp4_list)
scale_b = torch.stack(sf_list)
gsb = torch.stack(gs_list)
return mat_b, scale_b, gsb
def _assemble_scales_ref(x_sf, expert_offsets, num_experts):
"""Reference scale assembly using assemble_scales_3d_side from bridge."""
from cutedsl.bridge import assemble_scales_3d_side
return assemble_scales_3d_side(x_sf, expert_offsets, num_experts)
def main():
torch.cuda.set_device(0)
torch.manual_seed(42)
@@ -167,7 +136,8 @@ def main():
print(f" swiglu_limit={SWIGLU_LIMIT}")
print("\nLoading weights from checkpoint...")
weights = load_expert_weights(LAYER_IDX, NUM_EXPERTS)
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} "
@@ -179,19 +149,12 @@ def main():
# Realistic top-k: uneven distribution
topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE)
for i in range(NUM_TOKENS):
experts = torch.randperm(NUM_EXPERTS)[:TOP_K]
topk_ids[i] = experts
experts_perm = torch.randperm(NUM_EXPERTS)[:TOP_K]
topk_ids[i] = experts_perm
topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K
# ---- Reference ----
print("\n--- Reference (dynamic gs, per-expert scale assembly) ---")
with torch.no_grad():
ref_out = run_reference(hidden_states, topk_weights, topk_ids, weights, swiglu_limit=SWIGLU_LIMIT)
print(f"Reference: amax={ref_out.amax().item():.4f} mean={ref_out.mean().item():.4f}")
print(f" NaN: {torch.isnan(ref_out).any().item()} Inf: {torch.isinf(ref_out).any().item()}")
# ---- Runner ----
print("\n--- CuTeDSL Runner (warmup gs, full-buffer swizzle) ---")
print("\n--- CuTeDSL Runner (warmup gs, full-buffer swizzle, swiglu_limit) ---")
runner = CuTeDSLMoERunner(
num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE,
intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS,
@@ -206,12 +169,102 @@ def main():
runner.set_swiglu_limit(SWIGLU_LIMIT)
with torch.no_grad():
# Compute warmup gs
runner.compute_activation_global_scales(hidden_states, topk_weights, topk_ids)
print(f"Warmup gs: L1={runner._l1_activation_global_scale:.6f} L2={runner._l2_activation_global_scale:.6f}")
l1_gs_val = runner._l1_activation_global_scale
l2_gs_val = runner._l2_activation_global_scale
print(f"Warmup gs: L1={l1_gs_val:.6f} L2={l2_gs_val:.6f}")
# Run
runner_out = runner.run(hidden_states, topk_weights, topk_ids)
print(f"Runner: amax={runner_out.amax().item():.4f} mean={runner_out.mean().item():.4f}")
print(f" NaN: {torch.isnan(runner_out).any().item()} Inf: {torch.isinf(runner_out).any().item()}")
# ---- Reference: same runner but with dynamic gs (quantize_to_nvfp4) ----
print("\n--- Reference (dynamic gs via quantize_to_nvfp4) ---")
# We'll use the same runner infrastructure but manually call the reference path
from cutedsl.quantize import quantize_to_nvfp4
from cutedsl.gemm import run_nvfp4_grouped_gemm
from cutedsl.bridge import assemble_scales_3d_side, make_b_k_major, assemble_scales_3d_side
with torch.no_grad():
# Stack weights for GEMM
l1_mat_b = torch.stack(weights['l1_fp4'])
l1_scale_b = torch.stack(weights['l1_sf'])
l1_gsb = torch.stack(weights['l1_gs'])
l2_mat_b = torch.stack(weights['l2_fp4'])
l2_scale_b = torch.stack(weights['l2_sf'])
l2_gsb = torch.stack(weights['l2_gs'])
# Make B-K major (required by GEMM)
l1_mat_b = make_b_k_major(l1_mat_b)
l1_scale_b = assemble_scales_3d_side(l1_scale_b)
l2_mat_b = make_b_k_major(l2_mat_b)
l2_scale_b = assemble_scales_3d_side(l2_scale_b)
# Sort tokens by expert
flat_ids = topk_ids.reshape(-1)
flat_weights = topk_weights.reshape(-1)
sort_idx = flat_ids.argsort(stable=True)
sorted_ids = flat_ids[sort_idx]
sorted_weights = flat_weights[sort_idx]
token_indices = torch.arange(NUM_TOKENS, device=DEVICE).unsqueeze(1).expand(-1, TOP_K).reshape(-1)
sorted_token_ids = token_indices[sort_idx]
# Expert offsets
expert_id_range = torch.arange(NUM_EXPERTS, device=DEVICE)
tokens_per_expert = (sorted_ids.unsqueeze(1) == expert_id_range.unsqueeze(0)).sum(dim=0).int()
expert_offsets = torch.zeros(NUM_EXPERTS + 1, dtype=torch.int32, device=DEVICE)
expert_offsets[1:] = tokens_per_expert.cumsum(0)
slot_hidden = hidden_states[sorted_token_ids]
# L1: dynamic gs
x_fp4, x_sf, l1_gs_dyn = quantize_to_nvfp4(slot_hidden)
l1_gsa = torch.full((NUM_EXPERTS,), l1_gs_dyn, dtype=torch.float32, device=DEVICE)
l1_scale_a = assemble_scales_3d_side(x_sf, expert_offsets[:NUM_EXPERTS+1], NUM_EXPERTS)
l1_out = run_nvfp4_grouped_gemm(
mat_a=x_fp4, mat_b=l1_mat_b,
scale_a=l1_scale_a, scale_b=l1_scale_b,
expert_offsets=expert_offsets[1:],
global_scale_a=l1_gsa, global_scale_b=l1_gsb,
)
print(f" L1 gs (dynamic): {l1_gs_dyn:.6f}")
print(f" L1 out: amax={l1_out.amax().item():.4f}")
# SiLU(gate) * up with swiglu_limit
gate = l1_out[:, :INTERMEDIATE_SIZE]
up = l1_out[:, INTERMEDIATE_SIZE:]
gate_silu = torch.nn.functional.silu(gate)
if SWIGLU_LIMIT is not None:
gate_silu = gate_silu.clamp(max=SWIGLU_LIMIT)
up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT)
activated = gate_silu * up
print(f" activated: amax={activated.amax().item():.4f}")
# L2: dynamic gs
l2_x_fp4, l2_x_sf, l2_gs_dyn = quantize_to_nvfp4(activated)
l2_gsa = torch.full((NUM_EXPERTS,), l2_gs_dyn, dtype=torch.float32, device=DEVICE)
l2_scale_a = assemble_scales_3d_side(l2_x_sf, expert_offsets[:NUM_EXPERTS+1], NUM_EXPERTS)
l2_out = run_nvfp4_grouped_gemm(
mat_a=l2_x_fp4, mat_b=l2_mat_b,
scale_a=l2_scale_a, scale_b=l2_scale_b,
expert_offsets=expert_offsets[1:],
global_scale_a=l2_gsa, global_scale_b=l2_gsb,
)
print(f" L2 gs (dynamic): {l2_gs_dyn:.6f}")
# Scatter-add
ref_out = torch.zeros(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE)
weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype)
ref_out.scatter_add_(0, sorted_token_ids.unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted_out)
print(f"Reference: amax={ref_out.amax().item():.4f} mean={ref_out.mean().item():.4f}")
print(f" NaN: {torch.isnan(ref_out).any().item()} Inf: {torch.isinf(ref_out).any().item()}")
# ---- Comparison ----
print("\n--- Comparison ---")
cos = torch.nn.functional.cosine_similarity(
@@ -239,6 +292,11 @@ def main():
print(f"\n⚠️ MARGINAL: cosine {cos:.6f}")
else:
print(f"\n❌ FAIL: cosine {cos:.6f}")
# Print gs comparison
print(f"\n--- GS Comparison ---")
print(f" L1: dynamic={l1_gs_dyn:.6f} warmup={l1_gs_val:.6f} ratio={l1_gs_val/l1_gs_dyn:.4f}")
print(f" L2: dynamic={l2_gs_dyn:.6f} warmup={l2_gs_val:.6f} ratio={l2_gs_val/l2_gs_dyn:.4f}")
if __name__ == "__main__":