From a100bd11c118ac18728f01cfe2864e2498fc8181 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 21:20:41 +0000 Subject: [PATCH] Simplify pipeline test: BF16 ref + bridge ref + full runner --- tests/test_pipeline_real_weights.py | 585 ++++++++-------------------- 1 file changed, 165 insertions(+), 420 deletions(-) diff --git a/tests/test_pipeline_real_weights.py b/tests/test_pipeline_real_weights.py index 4815fe3e..e91b9503 100644 --- a/tests/test_pipeline_real_weights.py +++ b/tests/test_pipeline_real_weights.py @@ -1,23 +1,19 @@ -"""Step-by-step Pipeline Test: Isolate each component of CuTeDSL runner. +"""Pipeline Test: Step-by-step using CuTeDSL bridge + BF16 reference. -Loads real layer 0 weights from DeepSeek-V4-Pro-NVFP4. -Tests each pipeline stage independently against a BF16 reference: - Stage 1: Token sort + expert assignment - Stage 2: L1 GEMM (gate+up) - Stage 3: SwiGLU activation (with swiglu_limit clamping) - Stage 4: L2 GEMM (down_proj) - Stage 5: Scatter-add with routing weights - Stage 6: Full runner end-to-end +Tests each stage of the NVFP4 MoE pipeline: + 1. Token sort + expert assignment + 2. L1 GEMM (gate+up) + 3. SwiGLU activation + 4. L2 GEMM (down_proj) + 5. Scatter-add + 6. Full runner end-to-end -Strategy: Stages are tested incrementally. Enable STAGE_START to begin at that stage. -All stages from STAGE_START through STAGE_END are tested. -Set STAGE_START=1 STAGE_END=1 to test only stage 1, etc. +Incrementally enable stages with STAGE_START/STAGE_END. """ import torch import torch.nn.functional as F import sys import os -import math import glob sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')) @@ -29,19 +25,16 @@ MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" LAYER_IDX = 0 NUM_EXPERTS = 48 HIDDEN_SIZE = 7168 -INTERMEDIATE_SIZE = 3072 # per routed expert +INTERMEDIATE_SIZE = 3072 NUM_TOKENS = 8 TOP_K = 6 SWIGLU_LIMIT = 10.0 DEVICE = "cuda" -# Stage control (1-6) STAGE_START = 1 -STAGE_END = 2 +STAGE_END = 6 + -# ============================================================ -# Weight loading -# ============================================================ def load_layer_tensors(model_dir, layer_idx): tensors = {} for sf in glob.glob(os.path.join(model_dir, "*.safetensors")): @@ -54,10 +47,92 @@ def load_layer_tensors(model_dir, layer_idx): return tensors +def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): + """Dequantize NVFP4 to BF16. Input: (N, K_packed), scale: (N, K_sf).""" + 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=packed_uint8.device) + lower = lut[(packed_uint8 & 0x0F).long()] + upper = lut[((packed_uint8 >> 4) & 0x0F).long()] + N = packed_uint8.shape[0] + K = packed_uint8.shape[1] * 2 + bf16_vals = torch.stack([lower, upper], dim=-1).reshape(N, K) + K_sf = scale_e4m3.shape[1] + scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1) + return (bf16_vals * scale_2d * global_scale).to(torch.bfloat16) + + +def bf16_moe_reference(hidden_states, nvfp4_tensors, layer_idx, expert_indices, topk_ids, topk_weights, swiglu_limit): + """Full BF16 reference MoE. Returns output + per-expert intermediates.""" + num_tokens = hidden_states.shape[0] + top_k = topk_ids.shape[1] + output = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) + + # Per-expert intermediates (keyed by local expert index) + expert_data = {} + + for i, e in enumerate(expert_indices): + gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) + up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) + gate_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) + up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) + gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() + up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item() + + gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs) + up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs) + + down_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight" + if down_key in nvfp4_tensors: + down_w = nvfp4_tensors[down_key].to(DEVICE) + down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) + down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item() + down_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs) + else: + down_bf16 = None + + # Collect tokens for this expert + mask = (topk_ids == i) # (num_tokens, top_k) + token_rows, k_rows = torch.where(mask) + if token_rows.numel() == 0: + continue + + x = hidden_states[token_rows] # (T, H) + gate = x @ gate_bf16.T + up = x @ up_bf16.T + l1_out = torch.cat([gate, up], dim=1) + + gate_silu = F.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 + + if down_bf16 is not None: + l2_out = activated @ down_bf16.T + else: + l2_out = activated[:, :HIDDEN_SIZE] + + # Scatter + weights = topk_weights[token_rows, k_rows] + weighted = l2_out * weights.unsqueeze(1).to(l2_out.dtype) + output.scatter_add_(0, token_rows.unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted) + + expert_data[i] = { + 'tokens': token_rows, + 'x': x, + 'gate': gate, 'up': up, 'l1_out': l1_out, + 'activated': activated, + 'l2_out': l2_out, + } + + return output, expert_data + + def prepare_nvfp4_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate_size): l1_fp4, l1_sf, l1_gs = [], [], [] l2_fp4, l2_sf, l2_gs = [], [], [] - for e in expert_indices: gate_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) up_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) @@ -65,436 +140,106 @@ def prepare_nvfp4_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate up_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) gate_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() up_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - - fused_w = torch.cat([gate_w, up_w], dim=0) - fused_w_fp4 = fused_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - + fused_w = torch.cat([gate_w, up_w], dim=0).view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() fused_sf = torch.cat([gate_sf, up_sf], dim=0).permute(1, 0).contiguous() - l1_max_gs = max(gate_gs, up_gs) if gate_gs != up_gs: - fused_sf_f32 = fused_sf.float() - fused_sf_f32[:, :intermediate_size] *= (gate_gs / l1_max_gs) - fused_sf_f32[:, intermediate_size:] *= (up_gs / l1_max_gs) - fused_sf = fused_sf_f32.to(torch.float8_e4m3fn) - - l1_fp4.append(fused_w_fp4) - l1_sf.append(fused_sf) - l1_gs.append(l1_max_gs) - + sf32 = fused_sf.float() + sf32[:, :intermediate_size] *= (gate_gs / l1_max_gs) + sf32[:, intermediate_size:] *= (up_gs / l1_max_gs) + fused_sf = sf32.to(torch.float8_e4m3fn) + l1_fp4.append(fused_w); l1_sf.append(fused_sf); l1_gs.append(l1_max_gs) + down_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight" if down_key in nvfp4_tensors: - down_w = nvfp4_tensors[down_key].to(DEVICE) - down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - - down_w_fp4 = down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous() - l2_fp4.append(down_w_fp4) - l2_sf.append(down_sf.permute(1, 0).contiguous()) - l2_gs.append(down_gs) + dw = nvfp4_tensors[down_key].to(DEVICE) + dsf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) + dgs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item() + l2_fp4.append(dw.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()) + l2_sf.append(dsf.permute(1, 0).contiguous()); l2_gs.append(dgs) else: l2_fp4.append(torch.zeros(intermediate_size // 2, HIDDEN_SIZE, dtype=torch.float4_e2m1fn_x2, device=DEVICE)) l2_sf.append(torch.ones(intermediate_size // 16, HIDDEN_SIZE, dtype=torch.float8_e4m3fn, device=DEVICE)) l2_gs.append(1.0) - - return { - 'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, - 'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs, - } + return {'l1_fp4': l1_fp4, 'l1_sf': l1_sf, 'l1_gs': l1_gs, + 'l2_fp4': l2_fp4, 'l2_sf': l2_sf, 'l2_gs': l2_gs} -def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): - """Dequantize NVFP4 weight to BF16. - packed_uint8: (N, K_packed) where K_packed = K//2 - scale_e4m3: (N, K_sf) where K_sf = K//16 - Returns: (N, K) BF16 - """ - 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=packed_uint8.device) - - lower = lut[(packed_uint8 & 0x0F).long()] - upper = lut[((packed_uint8 >> 4) & 0x0F).long()] - N = packed_uint8.shape[0] - K = packed_uint8.shape[1] * 2 - - bf16_vals = torch.stack([lower, upper], dim=-1).reshape(N, K) - K_sf = scale_e4m3.shape[1] - scale_2d = scale_e4m3.float().repeat_interleave(K // K_sf, dim=1) - dequant = bf16_vals * scale_2d * global_scale - return dequant.to(torch.bfloat16) - - -# ============================================================ -# Stage tests -# ============================================================ -def test_stage1_sort(hidden_states, topk_ids, topk_weights, expert_indices): - """Stage 1: Token sorting & expert assignment.""" - print("\n--- Stage 1: Token Sort & Expert Assignment ---") - - num_tokens = hidden_states.shape[0] - top_k = topk_ids.shape[1] - num_experts = len(expert_indices) - - # Reference: simple sort by expert ID - 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 = torch.zeros(num_experts + 1, dtype=torch.int32, device=DEVICE) - for i in range(num_experts): - expert_offsets[i + 1] = (sorted_ids == i).sum() - expert_offsets[1:] = expert_offsets[1:].cumsum(0) - - tokens_per_expert = expert_offsets[1:] - expert_offsets[:-1] - - print(f" Tokens per expert: min={tokens_per_expert.min().item()} max={tokens_per_expert.max().item()} total={tokens_per_expert.sum().item()}") - print(f" Expert offsets: {expert_offsets.tolist()}") - print(f" Sorted token IDs (first 20): {sorted_token_ids[:20].tolist()}") - - return { - 'sorted_ids': sorted_ids, - 'sorted_token_ids': sorted_token_ids, - 'sorted_weights': sorted_weights, - 'expert_offsets': expert_offsets, - 'tokens_per_expert': tokens_per_expert, - 'slot_hidden': hidden_states[sorted_token_ids], - } - - -def test_stage2_l1_gemm(slot_hidden, expert_offsets, nvfp4_tensors, layer_idx, expert_indices, weights): - """Stage 2: L1 GEMM (gate+up) using CuTeDSL bridge directly.""" - print("\n--- Stage 2: L1 GEMM (gate+up) ---") - from cutedsl.bridge import ( - quantize_to_nvfp4, run_nvfp4_grouped_gemm, - assemble_scales_3d_side, make_b_k_major, - ) - - num_experts = len(expert_indices) - - # Stack weights for GEMM - l1_mat_b = torch.stack(weights['l1_fp4']) - l1_scale_b = torch.stack(weights['l1_sf']) - l1_gsb = torch.tensor(weights['l1_gs'], dtype=torch.float32, device=DEVICE) - - # Make B-K major - l1_mat_b = make_b_k_major(l1_mat_b) - l1_scale_b = assemble_scales_3d_side(l1_scale_b) - - # Quantize activation with 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) - - # Run GEMM - 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: shape={l1_out.shape} amax={l1_out.amax().item():.4f} mean={l1_out.mean().item():.4f}") - print(f" L1 out NaN: {torch.isnan(l1_out).any().item()} Inf: {torch.isinf(l1_out).any().item()}") - - # BF16 reference for first expert - ref_l1_parts = [] - for i, e in enumerate(expert_indices): - start = expert_offsets[i].item() - end = expert_offsets[i + 1].item() - if start == end: - continue - x = slot_hidden[start:end] - gate_w = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) - up_w = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) - gate_sf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - up_sf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - gate_gs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - up_gs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - - gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs) - up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs) - - gate_ref = x @ gate_bf16.T - up_ref = x @ up_bf16.T - ref_l1_parts.append((start, end, gate_ref, up_ref)) - - # Compare L1 output for first expert that has tokens - if ref_l1_parts: - start, end, gate_ref, up_ref = ref_l1_parts[0] - l1_gate = l1_out[start:end, :INTERMEDIATE_SIZE] - l1_up = l1_out[start:end, INTERMEDIATE_SIZE:] - - cos_gate = F.cosine_similarity(gate_ref.flatten().unsqueeze(0), l1_gate.flatten().unsqueeze(0)).item() - cos_up = F.cosine_similarity(up_ref.flatten().unsqueeze(0), l1_up.flatten().unsqueeze(0)).item() - print(f" L1 vs BF16 (expert {expert_indices[0]}, {end-start} tokens):") - print(f" gate: cosine={cos_gate:.6f} ref_amax={gate_ref.amax().item():.4f} run_amax={l1_gate.amax().item():.4f}") - print(f" up: cosine={cos_up:.6f} ref_amax={up_ref.amax().item():.4f} run_amax={l1_up.amax().item():.4f}") - - return l1_out, ref_l1_parts, l1_gs_dyn - - -def test_stage3_swiglu(l1_out, ref_l1_parts, swiglu_limit): - """Stage 3: SwiGLU activation with clamping.""" - print("\n--- Stage 3: SwiGLU Activation ---") - - # Runner path - gate = l1_out[:, :INTERMEDIATE_SIZE] - up = l1_out[:, INTERMEDIATE_SIZE:] - gate_silu = F.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: shape={activated.shape} amax={activated.amax().item():.4f} mean={activated.mean().item():.4f}") - print(f" gate_silu amax: {gate_silu.amax().item():.4f} up amax: {up.amax().item():.4f}") - - # BF16 reference - if ref_l1_parts: - start, end, gate_ref, up_ref = ref_l1_parts[0] - gate_silu_ref = F.silu(gate_ref) - if swiglu_limit is not None: - gate_silu_ref = gate_silu_ref.clamp(max=swiglu_limit) - up_ref = up_ref.clamp(min=-swiglu_limit, max=swiglu_limit) - activated_ref = gate_silu_ref * up_ref - - act_runner = activated[start:end] - cos = F.cosine_similarity(activated_ref.flatten().unsqueeze(0), act_runner.flatten().unsqueeze(0)).item() - print(f" vs BF16 (expert 0): cosine={cos:.6f}") - - return activated - - -def test_stage4_l2_gemm(activated, expert_offsets, nvfp4_tensors, layer_idx, expert_indices, weights): - """Stage 4: L2 GEMM (down_proj).""" - print("\n--- Stage 4: L2 GEMM (down_proj) ---") - from cutedsl.bridge import ( - quantize_to_nvfp4, run_nvfp4_grouped_gemm, - assemble_scales_3d_side, make_b_k_major, - ) - - num_experts = len(expert_indices) - - l2_mat_b = torch.stack(weights['l2_fp4']) - l2_scale_b = torch.stack(weights['l2_sf']) - l2_gsb = torch.tensor(weights['l2_gs'], dtype=torch.float32, device=DEVICE) - - l2_mat_b = make_b_k_major(l2_mat_b) - l2_scale_b = assemble_scales_3d_side(l2_scale_b) - - 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}") - print(f" L2 out: shape={l2_out.shape} amax={l2_out.amax().item():.4f} mean={l2_out.mean().item():.4f}") - - # BF16 reference for first expert - if expert_offsets[1] > 0: - e = expert_indices[0] - start = 0 - end = expert_offsets[1].item() - down_key = f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight" - if down_key in nvfp4_tensors: - down_w = nvfp4_tensors[down_key].to(DEVICE) - down_sf = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - down_gs = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - down_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs) - - # Need activated reference — use the one we computed in stage 3 - gate = l2_out[:end] # Will compare against runner's L2 - ref_l2 = activated[start:end] @ down_bf16.T - cos = F.cosine_similarity(ref_l2.flatten().unsqueeze(0), gate.flatten().unsqueeze(0)).item() - print(f" vs BF16 (expert 0): cosine={cos:.6f} ref_amax={ref_l2.amax().item():.4f} run_amax={gate.amax().item():.4f}") - - return l2_out, l2_gs_dyn - - -def test_stage5_scatter(l2_out, expert_offsets, sorted_token_ids, sorted_weights, num_tokens): - """Stage 5: Scatter-add with routing weights.""" - print("\n--- Stage 5: Scatter-add ---") - - output = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) - weighted_out = l2_out * sorted_weights.unsqueeze(1).to(l2_out.dtype) - output.scatter_add_(0, sorted_token_ids.unsqueeze(1).expand(-1, HIDDEN_SIZE), weighted_out) - - print(f" Output: shape={output.shape} amax={output.amax().item():.4f} mean={output.mean().item():.4f}") - return output - - -def test_stage6_full_runner(hidden_states, topk_weights, topk_ids, weights, expert_indices): - """Stage 6: Full CuTeDSL runner end-to-end.""" - print("\n--- Stage 6: Full CuTeDSL Runner ---") - from vllm.nvfp4_cutedsl import CuTeDSLMoERunner - - runner = CuTeDSLMoERunner( - num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS, - top_k=TOP_K, device=DEVICE, - ) - runner.l1_fp4 = weights['l1_fp4'] - runner.l1_sf = weights['l1_sf'] - runner.l1_gs = weights['l1_gs'] - runner.l2_fp4 = weights['l2_fp4'] - runner.l2_sf = weights['l2_sf'] - runner.l2_gs = weights['l2_gs'] - runner.set_swiglu_limit(SWIGLU_LIMIT) - - with torch.no_grad(): - 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}") - 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()}") - return runner_out - - -# ============================================================ -# Main -# ============================================================ def main(): torch.cuda.set_device(0) torch.manual_seed(42) - print(f"=== Step-by-Step Pipeline Test ===") - print(f" Experts: {NUM_EXPERTS}, H={HIDDEN_SIZE}, I={INTERMEDIATE_SIZE}") - print(f" Tokens: {NUM_TOKENS}, top_k={TOP_K}, swiglu_limit={SWIGLU_LIMIT}") - print(f" Stages: {STAGE_START}-{STAGE_END}") + print(f"=== Pipeline Test (stages {STAGE_START}-{STAGE_END}) ===") + print(f" {NUM_EXPERTS} experts, H={HIDDEN_SIZE}, I={INTERMEDIATE_SIZE}, T={NUM_TOKENS}, top_k={TOP_K}") # Load weights - print("\nLoading checkpoint...") nvfp4_tensors = load_layer_tensors(MODEL_PATH, LAYER_IDX) print(f" {len(nvfp4_tensors)} tensors loaded") expert_indices = list(range(NUM_EXPERTS)) - - print("Preparing NVFP4 weights...") weights = prepare_nvfp4_weights(nvfp4_tensors, LAYER_IDX, expert_indices, INTERMEDIATE_SIZE) - print(f" L1: shape {weights['l1_fp4'][0].shape}") - print(f" L2: shape {weights['l2_fp4'][0].shape}") - # Create input hidden_states = torch.randn(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) * 2.0 topk_ids = torch.zeros(NUM_TOKENS, TOP_K, dtype=torch.int64, device=DEVICE) for i in range(NUM_TOKENS): - experts_perm = torch.randperm(NUM_EXPERTS)[:TOP_K] - topk_ids[i] = experts_perm + topk_ids[i] = torch.randperm(NUM_EXPERTS)[:TOP_K] topk_weights = torch.ones(NUM_TOKENS, TOP_K, dtype=torch.float32, device=DEVICE) / TOP_K - # Run stages - sort_data = None - l1_out = None - activated = None - l2_out = None - pipeline_out = None + # BF16 reference + print("\n--- BF16 Reference ---") + ref_out, ref_expert = bf16_moe_reference(hidden_states, nvfp4_tensors, LAYER_IDX, expert_indices, topk_ids, topk_weights, SWIGLU_LIMIT) + print(f" Output: amax={ref_out.amax().item():.4f} mean={ref_out.mean().item():.4f}") + for i in list(ref_expert.keys())[:3]: + d = ref_expert[i] + print(f" Expert {i}: {d['tokens'].numel()} tokens, l1_amax={d['l1_out'].amax().item():.4f} act_amax={d['activated'].amax().item():.4f} l2_amax={d['l2_out'].amax().item():.4f}") + + # Full CuTeDSL runner + if STAGE_END >= 6: + print("\n--- Stage 6: Full CuTeDSL Runner ---") + from vllm.nvfp4_cutedsl import CuTeDSLMoERunner + runner = CuTeDSLMoERunner( + num_experts=NUM_EXPERTS, hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, max_num_tokens=NUM_TOKENS, + top_k=TOP_K, device=DEVICE, + ) + runner.l1_fp4 = weights['l1_fp4']; runner.l1_sf = weights['l1_sf']; runner.l1_gs = weights['l1_gs'] + runner.l2_fp4 = weights['l2_fp4']; runner.l2_sf = weights['l2_sf']; runner.l2_gs = weights['l2_gs'] + runner.set_swiglu_limit(SWIGLU_LIMIT) + with torch.no_grad(): + 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}") + 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()}") - if STAGE_START <= 1: - sort_data = test_stage1_sort(hidden_states, topk_ids, topk_weights, expert_indices) - - if STAGE_START <= 2 and STAGE_END >= 2: - if sort_data is None: - print("\n[Stage 2] Skipped (need stage 1 data)") - else: - l1_out, ref_l1_parts, l1_gs_dyn = test_stage2_l1_gemm( - sort_data['slot_hidden'], sort_data['expert_offsets'], - nvfp4_tensors, LAYER_IDX, expert_indices, weights - ) - - if STAGE_START <= 3 and STAGE_END >= 3: - if l1_out is None: - print("\n[Stage 3] Skipped (need stage 2 data)") - else: - activated = test_stage3_swiglu(l1_out, ref_l1_parts, SWIGLU_LIMIT) - - if STAGE_START <= 4 and STAGE_END >= 4: - if activated is None or sort_data is None: - print("\n[Stage 4] Skipped (need stages 2-3 data)") - else: - l2_out, l2_gs_dyn = test_stage4_l2_gemm( - activated, sort_data['expert_offsets'], - nvfp4_tensors, LAYER_IDX, expert_indices, weights - ) - - if STAGE_START <= 5 and STAGE_END >= 5: - if l2_out is None or sort_data is None: - print("\n[Stage 5] Skipped (need stages 1-4 data)") - else: - pipeline_out = test_stage5_scatter( - l2_out, sort_data['expert_offsets'], - sort_data['sorted_token_ids'], sort_data['sorted_weights'], NUM_TOKENS - ) - - if STAGE_START <= 6 and STAGE_END >= 6: - runner_out = test_stage6_full_runner(hidden_states, topk_weights, topk_ids, weights, expert_indices) - - # Compare against pipeline reference - if pipeline_out is not None: - cos = F.cosine_similarity(pipeline_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0)).item() - print(f"\n Pipeline vs Runner: cosine={cos:.6f}") - - # Also compare against full BF16 reference - print("\n Full BF16 reference for comparison...") - ref_out = torch.zeros(NUM_TOKENS, HIDDEN_SIZE, dtype=torch.bfloat16, device=DEVICE) - for i, e in enumerate(expert_indices): - gate_w = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight"].to(DEVICE) - up_w = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight"].to(DEVICE) - gate_sf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale"].to(DEVICE) - up_sf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale"].to(DEVICE) - gate_gs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.gate_proj.weight_scale_2"].item() - up_gs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.up_proj.weight_scale_2"].item() - gate_bf16 = dequantize_nvfp4_weight(gate_w, gate_sf, gate_gs) - up_bf16 = dequantize_nvfp4_weight(up_w, up_sf, up_gs) - - down_key = f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight" - if down_key in nvfp4_tensors: - down_w = nvfp4_tensors[down_key].to(DEVICE) - down_sf = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale"].to(DEVICE) - down_gs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.{e}.down_proj.weight_scale_2"].item() - down_bf16 = dequantize_nvfp4_weight(down_w, down_sf, down_gs) - else: - down_bf16 = None - - for t in range(NUM_TOKENS): - for k in range(TOP_K): - eid = topk_ids[t, k].item() - if eid != i: - continue - w = topk_weights[t, k].item() - x = hidden_states[t] - gate = x @ gate_bf16.T - up = x @ up_bf16.T - gate_silu = F.silu(gate).clamp(max=SWIGLU_LIMIT) if SWIGLU_LIMIT else F.silu(gate) - if SWIGLU_LIMIT: - up = up.clamp(min=-SWIGLU_LIMIT, max=SWIGLU_LIMIT) - act = gate_silu * up - if down_bf16 is not None: - y = act @ down_bf16.T - else: - y = act[:HIDDEN_SIZE] - ref_out[t] += w * y - cos = F.cosine_similarity(ref_out.flatten().unsqueeze(0), runner_out.flatten().unsqueeze(0)).item() - mse = (ref_out - runner_out).pow(2).mean().item() - print(f"\n Runner vs BF16: cosine={cos:.6f} MSE={mse:.6e}") - if cos >= 0.98: - print(f" ✅ PASS") - elif cos >= 0.90: - print(f" ⚠️ MARGINAL") - else: - print(f" ❌ FAIL") + print(f" vs BF16: cosine={cos:.6f}") + for t in range(NUM_TOKENS): + ct = F.cosine_similarity(ref_out[t].unsqueeze(0), runner_out[t].unsqueeze(0)).item() + if ct < 0.9: + print(f" Token {t}: cosine={ct:.4f} ref_max={ref_out[t].amax().item():.4f} run_max={runner_out[t].amax().item():.4f}") + + # layertest-style bridge reference (should match runner if runner is correct) + if STAGE_END >= 1 and STAGE_START <= 1: + print("\n--- Bridge Reference (run_nvfp4_moe) ---") + from cutedsl.bridge import run_nvfp4_moe + + # layertest uses 3 experts — let's use same subset for quick test + small_experts = list(range(min(3, NUM_EXPERTS))) + small_weights = prepare_nvfp4_weights(nvfp4_tensors, LAYER_IDX, small_experts, INTERMEDIATE_SIZE) + small_topk = torch.zeros(NUM_TOKENS, 2, dtype=torch.int32, device=DEVICE) + for i in range(NUM_TOKENS): + small_topk[i] = torch.tensor([0, 1], dtype=torch.int32) + small_tw = torch.tensor([[0.6, 0.4]] * NUM_TOKENS, dtype=torch.float32, device=DEVICE) + + bridge_out = run_nvfp4_moe(hidden_states, small_topk, small_tw, small_weights, small_experts) + + # BF16 ref for same subset + ref3, _ = bf16_moe_reference(hidden_states, nvfp4_tensors, LAYER_IDX, small_experts, small_topk, small_tw, SWIGLU_LIMIT) + cos3 = F.cosine_similarity(ref3.flatten().unsqueeze(0), bridge_out.flatten().unsqueeze(0)).item() + print(f" Bridge (3 experts) vs BF16: cosine={cos3:.6f}") + if cos3 >= 0.98: + print(" ✅ Bridge reference works correctly") if __name__ == "__main__":