#!/usr/bin/env python3 """ DeepSeek-V4 MoE NaN Reproduction Test Finds where NaN originates in the MoE forward pass. Tests individual experts (gate+up+down) with the CuTeDSL NVFP4 linear runner. Then tests the grouped GEMM MoE runner with stacked weights. Key insight: DeepSeek-V4 is a MegaMoE with 384 experts. The NaN might come from: 1. Weight loading / quantization 2. Activation quantization (quantize_activation_nvfp4) 3. The grouped GEMM kernel 4. The combine/scatter step Usage (on B200): cd /root/nvfp4-megamoe-kernel PYTHONPATH=/root/nvfp4-megamoe-kernel tests/venv/bin/python tests/test_moe_nan_b200.py """ import sys, os, json, torch, torch.nn.functional as F from safetensors import safe_open REPO = "/root/nvfp4-megamoe-kernel" sys.path.insert(0, REPO) MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" DEV = "cuda:0" H = 7168 INTERMEDIATE = 3072 # DeepSeek-V4 MoE intermediate NUM_EXPERTS = 384 TOPK = 6 EPS = 1e-6 E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) _cache = {} def P(k, wm, md): if k in _cache: return _cache[k] with safe_open(os.path.join(md, wm[k]), framework="pt") as f: t = f.get_tensor(k) _cache[k] = t return t def rms(x, w, eps=1e-6): v = x.float().pow(2).mean(-1, keepdim=True) return (w.float() * (x * torch.rsqrt(v+eps)).float()).to(x.dtype) def make_runner(w, sf, gs_t, inf, outf): from dsv4.layers.linear import Nvfp4Linear fp4 = w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous() s = sf.to(torch.float8_e4m3fn) if sf.dtype != torch.float8_e4m3fn else sf s = s.permute(1,0).contiguous() gs = gs_t.max().item() if gs_t.numel() > 1 else gs_t.item() r = Nvfp4Linear(in_features=inf, out_features=outf, max_num_tokens=8192, device=str(w.device)) r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] r.finalize_weights(); r._ensure_initialized() return r def test_single_expert(layer_id=2, expert_id=0): """Test a single expert's gate+up+down with CuTeDSL NVFP4 linear.""" torch.cuda.set_device(0) torch.cuda.empty_cache() _cache.clear() with open(os.path.join(MODEL, "model.safetensors.index.json")) as f: wm = json.load(f)["weight_map"] G = lambda k: P(k, wm, MODEL).to(DEV) p = f"model.layers.{layer_id}" m = f"{p}.mlp" e = f"{m}.experts.{expert_id}" emb = G("model.embed_tokens.weight") fnorm = G(f"{p}.post_attention_layernorm.weight") # Load expert weights gate_w = G(f"{e}.gate_proj.weight"); gate_sf = G(f"{e}.gate_proj.weight_scale"); gate_gs = G(f"{e}.gate_proj.weight_scale_2") up_w = G(f"{e}.up_proj.weight"); up_sf = G(f"{e}.up_proj.weight_scale"); up_gs = G(f"{e}.up_proj.weight_scale_2") down_w = G(f"{e}.down_proj.weight"); down_sf = G(f"{e}.down_proj.weight_scale"); down_gs = G(f"{e}.down_proj.weight_scale_2") print(f" Expert {expert_id}:") print(f" gate: shape={gate_w.shape} dtype={gate_w.dtype} sf_shape={gate_sf.shape} gs={gate_gs.tolist()}") print(f" up: shape={up_w.shape} dtype={up_w.dtype}") print(f" down: shape={down_w.shape} dtype={down_w.dtype}") print(f" gate NaN: {torch.isnan(gate_w.float()).any()}") print(f" gate_gs NaN: {torch.isnan(gate_gs).any()}") print(f" gate input_scale: exists={f'{e}.gate_proj.input_scale' in wm}") # Check for zero or extreme gs values for name, gs in [("gate", gate_gs), ("up", up_gs), ("down", down_gs)]: if gs.numel() > 0: print(f" {name} gs: min={gs.min().item():.6f} max={gs.max().item():.6f}") if gs.min().item() == 0: print(f" WARNING: {name} gs has zero value — will cause division by zero!") r_gate = make_runner(gate_w, gate_sf, gate_gs, H, gate_w.shape[0]) r_up = make_runner(up_w, up_sf, up_gs, H, up_w.shape[0]) r_down = make_runner(down_w, down_sf, down_gs, INTERMEDIATE, down_w.shape[0]) # Test with various token counts for num_tokens in [1, 4, 8, 16]: token_ids = torch.randint(1, 1000, (num_tokens,), dtype=torch.long, device=DEV) hidden = emb[token_ids] normed = rms(hidden, fnorm, EPS) with torch.no_grad(): gate_out = r_gate.run(normed) up_out = r_up.run(normed) # Check gate and up gate_nan = torch.isnan(gate_out).any().item() up_nan = torch.isnan(up_out).any().item() if gate_nan or up_nan: print(f" {num_tokens} tokens: gate NaN={gate_nan} up NaN={up_nan}") # Find which row has NaN gate_nan_rows = torch.isnan(gate_out).any(dim=1).nonzero().flatten().tolist() print(f" Gate NaN rows: {gate_nan_rows}") continue # SiLU activation activated = F.silu(gate_out) * up_out act_nan = torch.isnan(activated).any().item() if act_nan: print(f" {num_tokens} tokens: NaN after SiLU activation!") continue down_out = r_down.run(activated) down_nan = torch.isnan(down_out).any().item() if down_nan: print(f" {num_tokens} tokens: down NaN={down_nan}") continue print(f" {num_tokens} tokens: amax={down_out.amax():.4f} OK") del r_gate, r_up, r_down torch.cuda.empty_cache() def test_quantize_activation(): """Test the activation quantization used by the MoE grouped GEMM.""" from cutedsl.nvfp4_linear import quantize_activation_nvfp4 torch.cuda.set_device(0) for num_tokens in [1, 4, 8, 16]: # Create realistic input (after SiLU * up) x = torch.randn(num_tokens, INTERMEDIATE, dtype=torch.bfloat16, device=DEV) # quantize_activation_nvfp4 returns (x_sf, x_gs) or similar # The grouped GEMM needs quantized activation as input try: result = quantize_activation_nvfp4(x, num_tokens) if isinstance(result, tuple): for i, r in enumerate(result): if r is not None and r.is_floating_point(): print(f" {num_tokens} tokens: quantize result[{i}] NaN={torch.isnan(r).any()}") else: print(f" {num_tokens} tokens: quantize result NaN={torch.isnan(result).any()}") except Exception as e: print(f" {num_tokens} tokens: quantize failed: {e}") print() def test_grouped_gemm_shapes(): """Test the CuTeDSL grouped GEMM with MegaMoE-like shapes.""" from cutedsl.moe import run_nvfp4_grouped_gemm torch.cuda.set_device(0) # Create simple test: 4 experts, 8 tokens, top-2 num_experts = 4 num_tokens = 8 hidden_size = 512 # Small for testing intermediate_size = 256 # Allocate weight tensors (random) # L1: (num_experts, 2*intermediate_size, hidden_size//2) fp4 # L2: (num_experts, hidden_size, intermediate_size//2) fp4 l1_shape = (num_experts, 2 * intermediate_size, hidden_size // 2) l2_shape = (num_experts, hidden_size, intermediate_size // 2) print(f" Testing grouped GEMM with:") print(f" num_experts={num_experts}, num_tokens={num_tokens}") print(f" l1_shape={l1_shape}, l2_shape={l2_shape}") # This test just checks if the kernel can handle various expert distributions # without NaN. The actual weight values don't matter for NaN detection. print(f" (Skipping — requires proper weight packing from vLLM model loader)") print(f" The CuTeDSL grouped GEMM needs weights in a specific packed format") print(f" that the vLLM model loader creates during model initialization.") print() def main(): print("=" * 70) print(" DeepSeek-V4 MoE NaN Reproduction Test") print(" Finds where NaN originates in the MoE forward pass") print("=" * 70) print("\n=== Test 1: Single expert gate+up+down ===") for expert_id in [0, 1, 100, 383]: test_single_expert(layer_id=2, expert_id=expert_id) _cache.clear() print("\n=== Test 2: Activation quantization ===") test_quantize_activation() print("\n=== Test 3: Grouped GEMM shapes ===") test_grouped_gemm_shapes() print(f"\n{'='*70}") print(f" Summary: If single experts produce NaN, the issue is in weight") print(f" loading or the CuTeDSL NVFP4 linear kernel. If they're fine,") print(f" the NaN comes from the grouped GEMM or the combine step.") print(f"{'='*70}") if __name__ == "__main__": main()