diff --git a/tests/layertest.py b/tests/layertest.py index 716e8291..29749831 100644 --- a/tests/layertest.py +++ b/tests/layertest.py @@ -1,13 +1,10 @@ #!/usr/bin/env python3 """ -Layer 0 kernel comparison test: CuTeDSL NVFP4 kernel vs BF16 reference. +Layer 0 full MoE pipeline test: CuTeDSL NVFP4 vs BF16 reference. -No vLLM, no Docker, no tensor parallelism. Just raw weights + CuTeDSL kernel. -If cosine < 0.99, the test exits with error. - -Uses the bridge layer in cutedsl/bridge.py for tensor layout conversion. +Tests the complete pipeline: L1→SiLU→L2→scatter +If cosine < 0.99, exits with error. """ - import os import sys import json @@ -15,42 +12,25 @@ import glob import torch from safetensors import safe_open -# Add repo root to path 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, - quantize_weight_to_nvfp4, - assemble_scales_2d_side, - assemble_scales_3d_side, - make_b_k_major, - compute_expert_offsets, - run_nvfp4_grouped_gemm, -) - from cutedsl.moe_pipeline import ( - stage_activation, prepare_nvfp4_moe_weights, run_nvfp4_moe, ) -# ── Constants ────────────────────────────────────────────────────────── - NVFP4_MODEL_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" LAYER_IDX = 0 DEVICE = "cuda" COSINE_THRESHOLD = 0.99 -# E2M1 FP4 lookup table (for BF16 dequant reference) 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) -# ── Checkpoint loading ───────────────────────────────────────────────── - def find_shards(model_dir): index_path = os.path.join(model_dir, "model.safetensors.index.json") key_to_shard = {} @@ -68,7 +48,6 @@ def find_shards(model_dir): def load_layer_tensors(model_dir, layer_idx): - """Load all tensors for a specific layer. Keys normalized (no 'model.' prefix).""" key_to_shard = find_shards(model_dir) layer_prefix = f"layers.{layer_idx}." shard_to_keys = {} @@ -85,10 +64,7 @@ def load_layer_tensors(model_dir, layer_idx): return tensors -# ── NVFP4 Dequantization (BF16 reference) ───────────────────────────── - def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): - """Dequantize NVFP4 (E2M1 + E4M3 + global) to BF16.""" device = packed_uint8.device lut = E2M1_LUT.to(device) lower = lut[(packed_uint8 & 0x0F).long()] @@ -104,7 +80,6 @@ def dequantize_nvfp4_weight(packed_uint8, scale_e4m3, global_scale): def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices): - """Dequantize expert weights from NVFP4 checkpoint → BF16.""" experts = {} for e in expert_indices: expert = {} @@ -124,10 +99,7 @@ def dequantize_nvfp4_experts(nvfp4_tensors, layer_idx, expert_indices): return experts -# ── BF16 MoE Forward ─────────────────────────────────────────────────── - def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights): - """Run MoE forward pass in pure BF16 (torch.matmul).""" num_tokens, hidden_size = hidden_states.shape top_k = expert_ids.shape[1] output = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) @@ -149,81 +121,6 @@ def moe_forward_bf16(hidden_states, experts, expert_ids, expert_weights): return output -# ── CuTeDSL NVFP4 Kernel MoE Forward ────────────────────────────────── - -def moe_forward_nvfp4_l1_only(slot_hidden, nvfp4_tensors, layer_idx, expert_indices, tokens_per_expert): - """Run L1 (gate+up) GEMM using CuTeDSL. - - slot_hidden is already laid out slot-major: [expert0_tokens | expert1_tokens | ...] - """ - num_slots, hidden_size = slot_hidden.shape - num_experts = len(expert_indices) - - # Quantize activation - x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden) - - # Load and quantize weights - w_fp4_list = [] - w_sf_list = [] - w_gs_list = [] - - for e in expert_indices: - gate_w_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight" - gate_sf_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale" - gate_gs_key = f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight_scale_2" - up_w_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight" - up_sf_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale" - up_gs_key = f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight_scale_2" - - gate_w_bf16 = dequantize_nvfp4_weight( - nvfp4_tensors[gate_w_key].to(DEVICE), - nvfp4_tensors[gate_sf_key].to(DEVICE), - nvfp4_tensors[gate_gs_key].item(), - ) - up_w_bf16 = dequantize_nvfp4_weight( - nvfp4_tensors[up_w_key].to(DEVICE), - nvfp4_tensors[up_sf_key].to(DEVICE), - nvfp4_tensors[up_gs_key].item(), - ) - - # Fuse gate + up, transpose to (K=hidden, N=6144) - fused = torch.cat([gate_w_bf16, up_w_bf16], dim=0) # (6144, 7168) - l1_w_bf16 = fused.T # (7168, 6144) - l1_w_fp4, l1_w_sf, l1_w_gs = quantize_weight_to_nvfp4(l1_w_bf16) - - w_fp4_list.append(l1_w_fp4) - w_sf_list.append(l1_w_sf) - w_gs_list.append(l1_w_gs) - - # Stack and convert to K-major - mat_b = make_b_k_major(torch.stack(w_fp4_list)) - - # Assemble scale factors - # scale_a: per-expert activation scales, split by expert offsets - x_sf_parts = [] - offset = 0 - for tpe in tokens_per_expert: - x_sf_parts.append(x_sf[offset:offset+tpe]) - offset += tpe - scale_a = assemble_scales_2d_side(x_sf_parts) - scale_b = assemble_scales_3d_side(w_sf_list) - - # Expert offsets - expert_offsets = compute_expert_offsets(tokens_per_expert, num_experts) - - # Global scales - global_scale_a = torch.tensor([x_igs] * num_experts, dtype=torch.float32, device=DEVICE) - global_scale_b = torch.tensor(w_gs_list, dtype=torch.float32, device=DEVICE) - - # Run kernel - out = run_nvfp4_grouped_gemm( - mat_a=x_fp4, mat_b=mat_b, - scale_a=scale_a, scale_b=scale_b, - expert_offsets=expert_offsets, - global_scale_a=global_scale_a, global_scale_b=global_scale_b, - ) - return out - def main(): torch.manual_seed(42) expert_indices = [0, 1, 2] @@ -231,70 +128,60 @@ def main(): num_tokens = 4 hidden_size = 7168 - # ── Load NVFP4 checkpoint ── print("=" * 70) print(" Loading NVFP4 checkpoint layer 0") print("=" * 70) nvfp4_tensors = load_layer_tensors(NVFP4_MODEL_DIR, LAYER_IDX) - expert_keys = [k for k in sorted(nvfp4_tensors.keys()) if 'experts.0.' in k] print(f" {len(nvfp4_tensors)} tensors loaded") - for key in expert_keys[:3]: - t = nvfp4_tensors[key] - print(f" {key}: dtype={t.dtype} shape={tuple(t.shape)}") - # ── Prepare NVFP4 weights ── - print(" - Preparing NVFP4 weights (dequant → re-quant)...") + # Prepare weights + print("\n Preparing NVFP4 weights...") weights = prepare_nvfp4_moe_weights(nvfp4_tensors, LAYER_IDX, expert_indices) print(f" L1: {len(weights['l1_fp4'])} experts, shape {weights['l1_fp4'][0].shape}") print(f" L2: {len(weights['l2_fp4'])} experts, shape {weights['l2_fp4'][0].shape}") - # ── Dequantize → BF16 reference ── - print(" - Dequantizing NVFP4 → BF16 reference...") + # Dequantize for BF16 reference + print("\n Dequantizing NVFP4 -> BF16 reference...") nvfp4_experts_bf16 = dequantize_nvfp4_experts(nvfp4_tensors, LAYER_IDX, expert_indices) - # ── Create test input ── + # Test input 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) - # ── BF16 full MoE reference ── - print(" - Running BF16 MoE reference...") + # BF16 reference + print("\n Running BF16 MoE reference...") ref_output = moe_forward_bf16(hidden_states, nvfp4_experts_bf16, expert_ids, expert_weights) print(f" BF16 ref: amax={ref_output.abs().max():.4f} mean={ref_output.float().mean():.6f}") del nvfp4_experts_bf16 torch.cuda.empty_cache() - # ── CuTeDSL NVFP4 full MoE pipeline ── - print(" - Running CuTeDSL NVFP4 MoE pipeline (first run compiles, ~1-2 min)...") + # CuTeDSL NVFP4 pipeline + print("\n Running CuTeDSL NVFP4 MoE pipeline (first run compiles)...") kernel_output = run_nvfp4_moe( hidden_states, expert_ids, expert_weights, weights, expert_indices, ) print(f" Kernel: amax={kernel_output.abs().max():.4f} mean={kernel_output.float().mean():.6f}") - # ── Compare ── + # Compare cosine = torch.nn.functional.cosine_similarity( kernel_output.flatten().unsqueeze(0).float(), ref_output.flatten().unsqueeze(0).float(), ).item() mse = (kernel_output.float() - ref_output.float()).pow(2).mean().item() - print(f" -{'=' * 70}") + print(f"\n{'=' * 70}") print(f" RESULT: cosine={cosine:.6f} MSE={mse:.6e}") print(f"{'=' * 70}") if cosine < COSINE_THRESHOLD: - print(f" ❌ FAIL: cosine {cosine:.6f} < {COSINE_THRESHOLD}") + print(f" FAIL: cosine {cosine:.6f} < {COSINE_THRESHOLD}") sys.exit(1) else: - print(f" ✅ PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}") + print(f" PASS: cosine {cosine:.6f} >= {COSINE_THRESHOLD}") if __name__ == "__main__":