From 7b95e76723af85061927c2dee9d589dd31ff2f75 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 17 May 2026 07:33:20 +0000 Subject: [PATCH] test: runner vs pipeline comparison + scale assembly comparison --- tests/test_runner_vs_pipeline.py | 202 +++++++++++++++++++++++++++++++ tests/test_scale_assembly.py | 108 +++++++++++++++++ 2 files changed, 310 insertions(+) create mode 100644 tests/test_runner_vs_pipeline.py create mode 100644 tests/test_scale_assembly.py diff --git a/tests/test_runner_vs_pipeline.py b/tests/test_runner_vs_pipeline.py new file mode 100644 index 00000000..8ae00ad6 --- /dev/null +++ b/tests/test_runner_vs_pipeline.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +""" +Test A: Compare moe_pipeline output vs CuTeDSLMoERunner output. + +Uses the same weights and inputs. If they differ, the runner is broken. +Runs on the B200 host (not inside Docker): + source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate + python3 tests/test_runner_vs_pipeline.py +""" +import os, sys, json, torch +from safetensors import safe_open + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + +from cutedsl.moe_pipeline import run_nvfp4_moe +from vllm.nvfp4_cutedsl import CuTeDSLMoERunner +from cutedsl.bridge import quantize_to_nvfp4, quantize_weight_to_nvfp4, make_b_k_major, assemble_scales_3d_side, compute_expert_offsets + +MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +DEVICE = "cuda" +LAYER_IDX = 0 +E2M1_LUT = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32) + + +def find_shards(model_dir): + index_path = os.path.join(model_dir, "model.safetensors.index.json") + key_to_shard = {} + if os.path.exists(index_path): + with open(index_path) as f: + index = json.load(f) + for key, shard in index["weight_map"].items(): + key_to_shard[key] = os.path.join(model_dir, shard) + return key_to_shard + + +def load_layer_tensors(model_dir, layer_idx): + key_to_shard = find_shards(model_dir) + layer_prefix = f"layers.{layer_idx}." + tensors = {} + for key, shard in key_to_shard.items(): + norm_key = key.removeprefix("model.") + if not norm_key.startswith(layer_prefix): + continue + with safe_open(shard, framework="pt") as f: + if key in f.keys(): + tensors[norm_key] = f.get_tensor(key) + return tensors + + +def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): + device = packed_uint8.device + lut = E2M1_LUT.to(device) + lower = lut[(packed_uint8 & 0x0F).long()] + upper = lut[((packed_uint8 >> 4) & 0x0F).long()] + out_features = packed_uint8.shape[0] + in_features = packed_uint8.shape[1] * 2 + unpacked = torch.empty(out_features, in_features, dtype=torch.float32, device=device) + unpacked[:, 0::2] = lower + unpacked[:, 1::2] = upper + block_scale = scale_e4m3.float() + block_expanded = block_scale.repeat_interleave(16, dim=1)[:, :in_features] + return (unpacked * block_expanded * global_scale).to(torch.bfloat16) + + +def prepare_direct_weights(nvfp4_tensors, layer_idx, expert_indices, intermediate_size): + """Direct view-cast path (same as layertest).""" + 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) + 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() + 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_sf = torch.cat([gate_sf, up_sf], dim=0).permute(1, 0).contiguous() + max_gs = max(gate_gs, up_gs) + if gate_gs != up_gs: + f32 = fused_sf.float() + f32[:, :intermediate_size] *= (gate_gs / max_gs) + f32[:, intermediate_size:] *= (up_gs / max_gs) + fused_sf = f32.to(torch.float8_e4m3fn) + l1_fp4.append(fused_w_fp4) + l1_sf.append(fused_sf) + l1_gs.append(max_gs) + down_w = nvfp4_tensors[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight"].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() + l2_fp4.append(down_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()) + l2_sf.append(down_sf.permute(1, 0).contiguous()) + l2_gs.append(down_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 main(): + torch.manual_seed(42) + expert_indices = [0, 1, 2] + num_experts = len(expert_indices) + hidden_size = 7168 + intermediate_size = 3072 + top_k = 2 + num_tokens = 4 + + print("=" * 70) + print(" Loading checkpoint") + print("=" * 70) + nvfp4_tensors = load_layer_tensors(MODEL_DIR, LAYER_IDX) + print(f" {len(nvfp4_tensors)} tensors loaded") + + weights = prepare_direct_weights(nvfp4_tensors, LAYER_IDX, expert_indices, intermediate_size) + + hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 + expert_ids = torch.tensor([[0, 1]] * num_tokens, dtype=torch.int32, device=DEVICE) + expert_weights = torch.tensor([[0.6, 0.4]] * num_tokens, dtype=torch.float32, device=DEVICE) + + # ── Path 1: moe_pipeline (reference, uses quantize_to_nvfp4) ── + print("\n Running moe_pipeline (dynamic gs)...") + pipeline_out = run_nvfp4_moe( + hidden_states.clone(), expert_ids.clone(), expert_weights.clone(), + weights, expert_indices, + ) + print(f" Pipeline: amax={pipeline_out.abs().max():.4f}, mean={pipeline_out.float().mean():.6f}") + + # ── Path 2: CuTeDSLMoERunner with checkpoint input_scale (what vLLM uses) ── + print("\n Running CuTeDSLMoERunner (checkpoint gs)...") + runner = CuTeDSLMoERunner(num_experts, hidden_size, intermediate_size, device=DEVICE) + runner.prepare_weights_direct( + [w.clone() for w in weights['l1_fp4']], + [w.clone() for w in weights['l1_sf']], + list(weights['l1_gs']), + [w.clone() for w in weights['l2_fp4']], + [w.clone() for w in weights['l2_sf']], + list(weights['l2_gs']), + ) + # Set checkpoint input_scale (what vLLM does in finalize_weights) + igs = nvfp4_tensors[f"layers.{LAYER_IDX}.mlp.experts.0.gate_proj.input_scale"].item() + runner._l1_activation_global_scale = igs + runner._l2_activation_global_scale = igs + print(f" Checkpoint input_scale: {igs:.10f}") + + # Build topk_weights and topk_ids in the format the runner expects + # runner.run expects topk_ids as expert indices (0-based within our expert set) + topk_weights = expert_weights + topk_ids = expert_ids + + runner_out = runner.run(hidden_states.clone(), topk_weights, topk_ids) + print(f" Runner (ckpt gs): amax={runner_out.abs().max():.4f}, mean={runner_out.float().mean():.6f}") + + cos_ckpt = torch.nn.functional.cosine_similarity( + runner_out.flatten().unsqueeze(0).float(), + pipeline_out.flatten().unsqueeze(0).float(), + ).item() + print(f" Cosine vs pipeline: {cos_ckpt:.6f}") + + # ── Path 3: CuTeDSLMoERunner with dynamic gs ── + print("\n Running CuTeDSLMoERunner (dynamic gs)...") + # We can't use quantize_to_nvfp4 in the runner (cudagraph), but we can + # compute the gs from the input and set it before calling run + x_igs = (hidden_states.abs().max().item()) / (6.0 * 448.0) + runner2 = CuTeDSLMoERunner(num_experts, hidden_size, intermediate_size, device=DEVICE) + runner2.prepare_weights_direct( + [w.clone() for w in weights['l1_fp4']], + [w.clone() for w in weights['l1_sf']], + list(weights['l1_gs']), + [w.clone() for w in weights['l2_fp4']], + [w.clone() for w in weights['l2_sf']], + list(weights['l2_gs']), + ) + runner2._l1_activation_global_scale = x_igs + runner2._l2_activation_global_scale = x_igs + print(f" Dynamic gs (from input amax): {x_igs:.10f}") + + runner2_out = runner2.run(hidden_states.clone(), topk_weights, topk_ids) + print(f" Runner (dynamic gs): amax={runner2_out.abs().max():.4f}, mean={runner2_out.float().mean():.6f}") + + cos_dyn = torch.nn.functional.cosine_similarity( + runner2_out.flatten().unsqueeze(0).float(), + pipeline_out.flatten().unsqueeze(0).float(), + ).item() + print(f" Cosine vs pipeline: {cos_dyn:.6f}") + + # ── Summary ── + print(f"\n{'=' * 70}") + print(f" RESULTS") + print(f"{'=' * 70}") + print(f" Runner with checkpoint gs vs pipeline: {cos_ckpt:.6f}") + print(f" Runner with dynamic gs vs pipeline: {cos_dyn:.6f}") + if cos_dyn > 0.95: + print(f" ✅ Dynamic gs fixes the problem — gs is the only bug") + elif cos_dyn < 0.5 and cos_ckpt < 0.5: + print(f" ❌ Both runner paths are broken — scale assembly is also wrong") + else: + print(f" ⚠️ Partial match — multiple issues") + + +if __name__ == "__main__": + main() diff --git a/tests/test_scale_assembly.py b/tests/test_scale_assembly.py new file mode 100644 index 00000000..30cf42be --- /dev/null +++ b/tests/test_scale_assembly.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +""" +Test B: Compare _assemble_scales_cudagraph_safe vs assemble_scales_2d_side. + +Both should produce identical output given the same x_sf and expert_offsets. +If they differ, the cudagraph-safe path has a bug. + +Runs on the B200 host: + source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate + python3 tests/test_scale_assembly.py +""" +import os, sys, torch + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + +from cutedsl.bridge import quantize_to_nvfp4, assemble_scales_2d_side +from cutedsl.kernel.moe.torch_scaled_grouped_mm import pad_and_swizzle_single, ceil_div +from vllm.nvfp4_cutedsl import CuTeDSLMoERunner + + +def test_scale_assembly(): + """Compare the two scale assembly methods with realistic data.""" + DEVICE = "cuda" + num_experts = 3 + hidden_size = 7168 + intermediate_size = 3072 + + # Create a runner just to use its _assemble_scales_cudagraph_safe + runner = CuTeDSLMoERunner(num_experts, hidden_size, intermediate_size, device=DEVICE) + # Trigger _ensure_stacked and buffer allocation with dummy weights + def rand_fp4(*shape): + return torch.randint(0, 256, shape, dtype=torch.uint8, device=DEVICE).view(torch.float4_e2m1fn_x2) + def rand_sf(*shape): + return torch.rand(shape, dtype=torch.float16, device=DEVICE).to(torch.float8_e4m3fn) + runner.prepare_weights_direct( + [rand_fp4(3584, intermediate_size * 2) for _ in range(num_experts)], + [rand_sf(3584 // 16, intermediate_size * 2) for _ in range(num_experts)], + [0.1] * num_experts, + [rand_fp4(1536, hidden_size) for _ in range(num_experts)], + [rand_sf(1536 // 16, hidden_size) for _ in range(num_experts)], + [0.1] * num_experts, + ) + runner._ensure_stacked() + + # Test with different token distributions + test_cases = [ + ("4 tokens, expert 0 gets 2, expert 1 gets 2, expert 2 gets 0", [2, 2, 0]), + ("8 tokens, expert 0 gets 4, expert 1 gets 3, expert 2 gets 1", [4, 3, 1]), + ("4 tokens, expert 0 gets 4, expert 1 gets 0, expert 2 gets 0", [4, 0, 0]), + ("3 tokens, expert 0 gets 1, expert 1 gets 1, expert 2 gets 1", [1, 1, 1]), + ] + + all_pass = True + for desc, tokens_per_expert in test_cases: + total_tokens = sum(tokens_per_expert) + + # Create input and quantize + x = torch.randn(total_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) * 2.0 + x_fp4, x_sf, x_igs = quantize_to_nvfp4(x) + + # Path 1: assemble_scales_2d_side (per-expert split) + x_sf_parts = [] + offset = 0 + for tpe in tokens_per_expert: + x_sf_parts.append(x_sf[offset:offset + tpe]) + offset += tpe + scale_a_ref = assemble_scales_2d_side(x_sf_parts) + + # Path 2: _assemble_scales_cudagraph_safe (GPU-only) + expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=DEVICE) + expert_offsets[1:] = torch.tensor(tokens_per_expert, dtype=torch.int32).cumsum(0) + scale_a_cudagraph = runner._assemble_scales_cudagraph_safe(x_sf, expert_offsets) + + # Compare + # Note: shapes may differ due to padding, but the data in the + # padded rows should match (up to the total number of rows used by the kernel) + if scale_a_ref.shape != scale_a_cudagraph.shape: + print(f" {desc}") + print(f" Shape mismatch: ref={scale_a_ref.shape}, cg={scale_a_cudagraph.shape}") + all_pass = False + continue + + match = torch.equal(scale_a_ref, scale_a_cudagraph) + if not match: + # Check how many bytes differ + diff = (scale_a_ref.view(torch.uint8) != scale_a_cudagraph.view(torch.uint8)).sum().item() + total = scale_a_ref.numel() + pct = diff / total * 100 + print(f" {desc}") + print(f" MISMATCH: {diff}/{total} bytes differ ({pct:.1f}%)") + print(f" ref range: [{scale_a_ref.view(torch.uint8).min()}, {scale_a_ref.view(torch.uint8).max()}]") + print(f" cg range: [{scale_a_cudagraph.view(torch.uint8).min()}, {scale_a_cudagraph.view(torch.uint8).max()}]") + all_pass = False + else: + print(f" {desc}: ✅ MATCH") + + print(f"\n{'=' * 70}") + if all_pass: + print(" ALL SCALE ASSEMBLY TESTS PASSED ✅") + else: + print(" SCALE ASSEMBLY TESTS FAILED ❌") + print(f"{'=' * 70}") + return all_pass + + +if __name__ == "__main__": + test_scale_assembly()