From 5e63a0d8a367bd56f06c0a14b4f19a4ab3648382 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 18:10:05 +0000 Subject: [PATCH] Rewrite pipeline test: use raw checkpoint weights, compare runner vs dynamic-gs reference --- tests/test_pipeline_real_weights.py | 308 +++++++++++++++++----------- 1 file changed, 183 insertions(+), 125 deletions(-) diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py index 784f8f72..0fa05b11 100644 --- a/tests/test_pipeline_real_weights.py +++ b/tests/test_pipeline_real_weights.py @@ -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__":