From 04ad6409e58431b2da4faf5753defe16f238869c Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 07:49:41 +0000 Subject: [PATCH] Rewrite test: diagnose whether warmup gs matters at inference time --- tests/test_model_forward_b200.py | 355 ++++++++++++------------------- 1 file changed, 131 insertions(+), 224 deletions(-) diff --git a/tests/test_model_forward_b200.py b/tests/test_model_forward_b200.py index 98a5ca52..92313977 100644 --- a/tests/test_model_forward_b200.py +++ b/tests/test_model_forward_b200.py @@ -2,21 +2,19 @@ """ Reproduce the vLLM empty-output bug outside the container. -Runs the FULL model forward pass: embedding → 61 decoder layers → LM head. -Uses CuTeDSL NVFP4 runners for quantized layers, BF16 matmuls for others. +Strategy: Run the model in FULL BF16 (dequantized weights) and compare +against CuTeDSL at each projection. Also check: does the warmup gs +cause issues at inference time? -Compares two approaches: - A) Warmup gs (what vLLM does) — 1 token random sample per layer - B) Dynamic gs (compute per-batch) — uses quantize_to_nvfp4 each call - -If A produces garbage and B produces reasonable output, the warmup gs is wrong. +Key diagnostic: inspect CuTeDSL runner.run() to see if it uses the +fixed warmup gs or recomputes per-call. Usage (on B200): source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate python3 tests/test_model_forward_b200.py """ -import sys, os, json, torch, torch.nn.functional as F, time +import sys, os, json, torch, torch.nn.functional as F, inspect from safetensors import safe_open REPO = "/root/nvfp4-megamoe-kernel" @@ -24,24 +22,9 @@ sys.path.insert(0, REPO) MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" DEV = "cuda:0" -# Model config -H = 7168 -NH = 128 -HD = 512 -NOPE = 448 -ROPE = 64 -QL = 1536 -OL = 1024 -OG = 16 -HPG = NH // OG -HC = 4 -SL = 10.0 +H = 7168; NH = 128; HD = 512; NOPE = 448; ROPE = 64 +QL = 1536; OL = 1024; OG = 16; HPG = NH // OG EPS = 1e-6 -INTER = 3072 -N_EXPERTS = 384 -TOP_K = 6 -N_LAYERS = 61 -VOCAB = 129280 E2M1 = torch.tensor([0,.5,1.,1.5,2.,3.,4.,6.,-0,-.5,-1.,-1.5,-2.,-3.,-4.,-6.], dtype=torch.float32) @@ -83,246 +66,170 @@ def make_runner(w, sf, gs_t, inf, outf, fused=False, lw=None): r.finalize_weights(); r._ensure_initialized() return r -def cosim(a, b): - return F.cosine_similarity(a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float().to(a.device)).item() - - -class Layer0Runner: - """Runs layer 0 forward with CuTeDSL kernels.""" - def __init__(self, wm, model_dir, use_warmup_gs=True): - G = lambda k: P(k, wm, model_dir).to(DEV) - p = "model.layers.0"; a = f"{p}.self_attn"; m = f"{p}.mlp" - - # Attention - self.qa_w = G(f"{a}.q_a_proj.weight"); self.qa_sf = G(f"{a}.q_a_proj.weight_scale"); self.qa_gs = G(f"{a}.q_a_proj.weight_scale_2") - self.qb_w = G(f"{a}.q_b_proj.weight"); self.qb_sf = G(f"{a}.q_b_proj.weight_scale"); self.qb_gs = G(f"{a}.q_b_proj.weight_scale_2") - self.kv_w = G(f"{a}.kv_proj.weight"); self.kv_sf = G(f"{a}.kv_proj.weight_scale"); self.kv_gs = G(f"{a}.kv_proj.weight_scale_2") - self.woa = G(f"{a}.o_a_proj.weight") - self.wob_w = G(f"{a}.o_b_proj.weight"); self.wob_sf = G(f"{a}.o_b_proj.weight_scale"); self.wob_gs = G(f"{a}.o_b_proj.weight_scale_2") - self.qn = G(f"{a}.q_a_norm.weight"); self.kvn = G(f"{a}.kv_norm.weight") - self.anorm = G(f"{p}.input_layernorm.weight"); self.fnorm = G(f"{p}.post_attention_layernorm.weight") - - # Compressor - self.ckv_w = G(f"{a}.compressor.kv_proj.weight"); self.ckv_sf = G(f"{a}.compressor.kv_proj.weight_scale"); self.ckv_gs = G(f"{a}.compressor.kv_proj.weight_scale_2") - self.cg_w = G(f"{a}.compressor.gate_proj.weight"); self.cg_sf = G(f"{a}.compressor.gate_proj.weight_scale"); self.cg_gs = G(f"{a}.compressor.gate_proj.weight_scale_2") - - # MHC - self.hca_fn = G(f"{p}.attn_hc.fn"); self.hcf_fn = G(f"{p}.ffn_hc.fn") - self.hca_b = G(f"{p}.attn_hc.base"); self.hcf_b = G(f"{p}.ffn_hc.base") - self.hca_s = G(f"{p}.attn_hc.scale"); self.hcf_s = G(f"{p}.ffn_hc.scale") - - # Create runners - self.r_qa = make_runner(self.qa_w, self.qa_sf, self.qa_gs, self.qa_w.shape[1]*2, self.qa_w.shape[0]) - self.r_qb = make_runner(self.qb_w, self.qb_sf, self.qb_gs, self.qb_w.shape[1]*2, self.qb_w.shape[0]) - self.r_kv = make_runner(self.kv_w, self.kv_sf, self.kv_gs, self.kv_w.shape[1]*2, self.kv_w.shape[0]) - self.r_wob = make_runner(self.wob_w, self.wob_sf, self.wob_gs, self.wob_w.shape[1]*2, self.wob_w.shape[0]) - self.r_ckv = make_runner(self.ckv_w, self.ckv_sf, self.ckv_gs, self.ckv_w.shape[1]*2, self.ckv_w.shape[0]) - self.r_cg = make_runner(self.cg_w, self.cg_sf, self.cg_gs, self.cg_w.shape[1]*2, self.cg_w.shape[0]) - - self.use_warmup_gs = use_warmup_gs - if use_warmup_gs: - # Warmup with 1 token (what vLLM does) - with torch.no_grad(): - d = torch.randn(1, H, dtype=torch.bfloat16, device=DEV)*2.0 - self.r_qa.compute_activation_global_scale(d) - self.r_kv.compute_activation_global_scale(d) - self.r_ckv.compute_activation_global_scale(d) - self.r_cg.compute_activation_global_scale(d) - d2 = torch.randn(1, QL, dtype=torch.bfloat16, device=DEV)*2.0 - self.r_qb.compute_activation_global_scale(d2) - d3 = torch.randn(1, OG*OL, dtype=torch.bfloat16, device=DEV)*2.0 - self.r_wob.compute_activation_global_scale(d3) - - def forward_projection(self, x, runner, in_features, name): - """Run a single NVFP4 projection, optionally recomputing gs.""" - if not self.use_warmup_gs: - from cutedsl.bridge import quantize_activation_nvfp4 - # Dynamic gs: recompute for this specific input - amax = x.amax().item() - gs = amax / (6.0 * 448.0) if amax > 0 else 1.0 / 2688.0 - runner._activation_global_scale = gs - return runner.run(x) - def main(): torch.cuda.set_device(0) - torch.manual_seed(42) print("=" * 70) - print(" Full Model Forward Test: Reproduce vLLM Empty Output") + print(" Diagnose: Why does vLLM produce empty output?") print("=" * 70) 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) - # ── Load embedding ──────────────────────────────────────────────── - print("\n--- Loading embedding layer ---") - emb_key = "model.embed_tokens.weight" - emb = G(emb_key) - print(f" embed_tokens: {emb.shape} dtype={emb.dtype}") + # ── INSPECT: How does CuTeDSL runner.run() use gs? ──────────────── + print("\n--- INSPECTING CuTeDSL runner internals ---") + from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear + from cutedsl.bridge import quantize_activation_nvfp4 - # ── Load LM head ────────────────────────────────────────────────── - lm_head_key = "lm_head.weight" - if lm_head_key in wm: - lm_head = G(lm_head_key) - else: - lm_head = emb # tied weights - print(f" lm_head: {lm_head.shape}") + print("\n quantize_activation_nvfp4 signature:") + sig = inspect.signature(quantize_activation_nvfp4) + print(f" {sig}") - # ── Load final norm ─────────────────────────────────────────────── - fnorm_key = "model.norm.weight" - fnorm_w = G(fnorm_key) - print(f" final_norm: {fnorm_w.shape}") + print("\n CuTeDSLNvfp4Linear._run_impl source (key lines):") + src = inspect.getsource(CuTeDSLNvfp4Linear._run_impl) + for i, line in enumerate(src.split('\n')): + stripped = line.strip() + if any(kw in stripped for kw in ['global_scale', '_activation', 'quantize', 'return', 'def ']): + print(f" L{i}: {stripped}") - # ── Token IDs for "The capital of France is" ────────────────────── - # DeepSeek V3/V4 uses a Llama-style BPE tokenizer - # Use token IDs that we know work. If tokenizer isn't available, - # just use token 0,1,2,3,4 as a test — we're checking for - # garbage output (all NaN or all same logit), not text quality. + # ── CRITICAL TEST: warmup gs vs per-input gs ────────────────────── + print("\n--- CRITICAL TEST: warmup gs vs per-input gs ---") + + p = "model.layers.0"; a = f"{p}.self_attn" + qa_w = G(f"{a}.q_a_proj.weight"); qa_sf = G(f"{a}.q_a_proj.weight_scale"); qa_gs = G(f"{a}.q_a_proj.weight_scale_2") + + # Load embedding + norm weight + emb = G("model.embed_tokens.weight") + anorm = G(f"{p}.input_layernorm.weight") + + # Create runner with warmup gs + r = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) + torch.manual_seed(42) + warmup = torch.randn(1, H, dtype=torch.bfloat16, device=DEV)*2.0 + with torch.no_grad(): + r.compute_activation_global_scale(warmup) + print(f" Warmup gs (random input amax={warmup.amax():.4f}): {r._activation_global_scale:.8f}") + + # Get REAL input (embedding output) token_ids = torch.tensor([1, 450, 8403, 315, 5413, 374], dtype=torch.long, device=DEV) - print(f" token_ids: {token_ids.tolist()}") - NT = len(token_ids) - - # ── Embed ───────────────────────────────────────────────────────── - print("\n--- Running embedding lookup ---") with torch.no_grad(): - hidden = emb[token_ids] # (NT, H) - print(f" hidden: {hidden.shape} amax={hidden.amax():.4f} NaN={torch.isnan(hidden).any()}") + hidden = emb[token_ids] + normed = rms(hidden, anorm, EPS) + print(f" Real input (after RMS norm) amax: {normed.amax():.4f}") - # ── Create layer 0 runner (warmup gs, like vLLM) ───────────────── - print("\n--- Creating layer 0 runner (warmup gs) ---") - layer0 = Layer0Runner(wm, MODEL, use_warmup_gs=True) + # What gs would the real input need? + real_gs = normed.amax().item() / (6.0 * 448.0) + print(f" Correct gs for real input: {real_gs:.8f}") + print(f" Ratio warmup/correct: {r._activation_global_scale / real_gs:.4f}" if real_gs > 0 else " real_gs is 0!") - # ── Run layer 0 attention projections ───────────────────────────── - print("\n--- Running layer 0 attention (CuTeDSL, warmup gs) ---") + # Run with warmup gs with torch.no_grad(): - normed = rms(hidden, layer0.anorm, EPS) - print(f" normed: amax={normed.amax():.4f} NaN={torch.isnan(normed).any()}") + out_warmup = r.run(normed) - qa_out = layer0.r_qa.run(normed) - print(f" q_a: amax={qa_out.amax():.4f} NaN={torch.isnan(qa_out).any()}") - - kv_out = layer0.r_kv.run(normed) - print(f" kv: amax={kv_out.amax():.4f} NaN={torch.isnan(kv_out).any()}") - - # q_a norm → q_b - qa_normed = rms(qa_out, layer0.qn, EPS) - qb_out = layer0.r_qb.run(qa_normed) - print(f" q_b: amax={qb_out.amax():.4f} NaN={torch.isnan(qb_out).any()}") - - # ── Compare with BF16 reference ─────────────────────────────────── - print("\n--- Comparing layer 0 q_a with BF16 reference ---") - qa_bf16 = normed @ dequant(layer0.qa_w, layer0.qa_sf, layer0.qa_gs.item()).T - c = cosim(qa_out, qa_bf16) - print(f" q_a cosine (warmup gs): {c:.6f} {'✅' if c>=0.98 else '❌'}") - - # ── Now test with DYNAMIC gs (recomputed per input) ─────────────── - print("\n--- Testing with dynamic gs (per-input) ---") - # Create a fresh runner and compute gs from the actual input - r_qa2 = make_runner(layer0.qa_w, layer0.qa_sf, layer0.qa_gs, layer0.qa_w.shape[1]*2, layer0.qa_w.shape[0]) + # Run with dynamic gs (recompute for this input) + r2 = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) with torch.no_grad(): - r_qa2.compute_activation_global_scale(normed) - qa_out2 = r_qa2.run(normed) - c2 = cosim(qa_out2, qa_bf16) - print(f" q_a cosine (dynamic gs): {c2:.6f} {'✅' if c2>=0.98 else '❌'}") + r2.compute_activation_global_scale(normed) + out_dynamic = r2.run(normed) - # ── Test the FULL model: layer 0 only, then check LM head ───────── - print("\n--- Full forward: layer 0 → LM head ---") + # BF16 reference + qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) + with torch.no_grad(): + ref = normed @ qa_bf16.T + + c_warmup = F.cosine_similarity(out_warmup.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float()).item() + c_dynamic = F.cosine_similarity(out_dynamic.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float()).item() + + print(f"\n q_a cosine vs BF16 (warmup gs): {c_warmup:.6f} {'✅' if c_warmup>=0.98 else '❌'}") + print(f" q_a cosine vs BF16 (dynamic gs): {c_dynamic:.6f} {'✅' if c_dynamic>=0.98 else '❌'}") + print(f" amax warmup: {out_warmup.amax():.4f} amax dynamic: {out_dynamic.amax():.4f} amax ref: {ref.amax():.4f}") + + # ── Test: run FULL model in BF16 (1 layer) then check logits ────── + print("\n--- FULL BF16 model: 1 layer → LM head ---") + lm_head = G("lm_head.weight") + fnorm_w = G("model.norm.weight") + qn = G(f"{a}.q_a_norm.weight") + kvn = G(f"{a}.kv_norm.weight") + fnorm_l0 = G(f"{p}.post_attention_layernorm.weight") + + # Dequantize all layer 0 attention weights + qa_bf16 = dequant(qa_w, qa_sf, qa_gs.item()) + qb_w = G(f"{a}.q_b_proj.weight"); qb_sf = G(f"{a}.q_b_proj.weight_scale"); qb_gs = G(f"{a}.q_b_proj.weight_scale_2") + kv_bf16 = dequant(G(f"{a}.kv_proj.weight"), G(f"{a}.kv_proj.weight_scale"), G(f"{a}.kv_proj.weight_scale_2").item()) + qb_bf16 = dequant(qb_w, qb_sf, qb_gs.item()) + woa = G(f"{a}.o_a_proj.weight") # already BF16 + wob_bf16 = dequant(G(f"{a}.o_b_proj.weight"), G(f"{a}.o_b_proj.weight_scale"), G(f"{a}.o_b_proj.weight_scale_2").item()) - # Simple layer 0 forward (attention only, no MoE for speed) with torch.no_grad(): x = hidden.clone() - normed = rms(x, layer0.anorm, EPS) + print(f" Input: amax={x.amax():.4f}") - # Attention projections - qa = layer0.r_qa.run(normed) - kv = layer0.r_kv.run(normed) - qa_n = rms(qa, layer0.qn, EPS) - qb = layer0.r_qb.run(qa_n) + # RMS norm + x = rms(x, anorm, EPS) - # Skip actual attention (FlashMLA not available) — just use - # a random attention output to test the wo_a → wo_b path - o = torch.randn(NT, NH, HD, dtype=torch.bfloat16, device=DEV) * 0.1 + # Attention projections (BF16) + qa = x @ qa_bf16.T + kv = x @ kv_bf16.T + qa_n = rms(qa, qn, EPS) + qb = qa_n @ qb_bf16.T + print(f" q_a: amax={qa.amax():.4f}, kv: amax={kv.amax():.4f}, q_b: amax={qb.amax():.4f}") - # wo_a: BF16 BMM - woa = layer0.woa - o_2d = o.reshape(NT, NH * HD) - z = o_2d @ woa.T - z2 = z.reshape(NT, OG, -1) + # Skip attention, use random output + o = torch.randn(len(token_ids), NH, HD, dtype=torch.bfloat16, device=DEV) * 0.1 - # Simpler: just check if the wo_b projection works - z_flat = torch.randn(NT, OG * OL, dtype=torch.bfloat16, device=DEV) * 2.0 - wob_out = layer0.r_wob.run(z_flat) - print(f" wo_b output: amax={wob_out.amax():.4f} NaN={torch.isnan(wob_out).any()}") + # wo_a: BMM (o_a_proj is (OG*OL, HPG*HD)) + o_grouped = o.view(len(token_ids), OG, HPG * HD).permute(1, 0, 2) + woa_3d = woa.view(OG, OL, HPG * HD) + z = torch.bmm(o_grouped, woa_3d.transpose(1, 2)).permute(1, 0, 2).reshape(len(token_ids), OG * OL) - # ── Now run LM head on the hidden state ─────────────────────────── - print("\n--- LM head (BF16 matmul) ---") - with torch.no_grad(): - normed_final = rms(x, fnorm_w, EPS) - logits = normed_final @ lm_head.T # (NT, VOCAB) - print(f" logits: {logits.shape} amax={logits.amax():.4f} NaN={torch.isnan(logits).any()}") + # wo_b + attn_out = z @ wob_bf16.T + print(f" attn_out (BF16): amax={attn_out.amax():.4f}") + + # Skip MoE, just add residual + x = hidden + attn_out + + # Final norm + LM head + x_normed = rms(x, fnorm_w, EPS) + logits = x_normed @ lm_head.T + print(f" logits: amax={logits.amax():.4f} NaN={torch.isnan(logits).any()}") - # Check if logits are reasonable top5 = torch.topk(logits[-1], 5) - print(f" top5 token IDs: {top5.indices.tolist()}") + print(f" top5 IDs: {top5.indices.tolist()}") print(f" top5 logits: {[f'{v:.2f}' for v in top5.values.tolist()]}") - - # Check logit variance (garbage = all same or extreme values) log_std = logits[-1].float().std().item() - log_range = (logits[-1].float().amax() - logits[-1].float().amin()).item() - print(f" logit std: {log_std:.4f} range: {log_range:.4f}") - if log_std < 0.01: - print(" ❌ LOGITS ARE FLAT — model is producing garbage!") - elif log_std > 100: - print(" ❌ LOGITS ARE EXPLODED — model is producing garbage!") - else: - print(" ✅ Logits look reasonable for a single layer test") + print(f" logit std: {log_std:.4f}") - # ── Key diagnostic: warmup gs vs actual gs ──────────────────────── + # ── KEY INSIGHT: check if the runner re-reads gs at inference time ─ print("\n" + "=" * 70) - print(" DIAGNOSTIC: Warmup gs vs Actual gs") + print(" KEY: Does runner.run() use FIXED warmup gs or RECOMPUTE?") print("=" * 70) - # What gs did warmup compute? - print(f" r_qa warmup gs: {layer0.r_qa._activation_global_scale:.8f}") - print(f" r_kv warmup gs: {layer0.r_kv._activation_global_scale:.8f}") - print(f" r_wob warmup gs: {layer0.r_wob._activation_global_scale:.8f}") - - # What gs would the actual input produce? + # Monkey-patch the gs and see if output changes + r3 = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) with torch.no_grad(): - actual_gs_qa = normed.amax().item() / (6.0 * 448.0) - actual_gs_kv = normed.amax().item() / (6.0 * 448.0) - print(f" actual gs for q_a input: {actual_gs_qa:.8f}") - print(f" ratio warmup/actual for q_a: {layer0.r_qa._activation_global_scale / actual_gs_qa:.4f}" if actual_gs_qa > 0 else " actual gs is 0!") + r3.compute_activation_global_scale(normed) + gs_original = r3._activation_global_scale + out_original = r3.run(normed).clone() - # The KEY question: does the runner use warmup gs at inference time, - # or does quantize_activation_nvfp4 recompute it? - print("\n--- How does CuTeDSL runner.run() use gs? ---") - from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear - import inspect - run_src = inspect.getsource(CuTeDSLNvfp4Linear.run) - # Check if it references _activation_global_scale - if '_activation_global_scale' in run_src: - print(" run() uses _activation_global_scale (FIXED from warmup)") + # Change gs by 10x + r3._activation_global_scale = gs_original * 10.0 + out_changed = r3.run(normed).clone() + + c_changed = F.cosine_similarity(out_original.flatten().unsqueeze(0).float(), out_changed.flatten().unsqueeze(0).float()).item() + print(f" Original gs: {gs_original:.8f}") + print(f" Changed gs: {gs_original * 10:.8f}") + print(f" Cosine sim after 10x gs change: {c_changed:.6f}") + if abs(c_changed - 1.0) < 0.001: + print(" ➡️ Changing gs has NO effect on output!") + print(" ➡️ The runner recomputes gs internally at inference time.") + print(" ➡️ Warmup gs is IRRELEVANT — the bug is elsewhere.") else: - print(" run() does NOT use _activation_global_scale") - - # Check quantize_activation_nvfp4 - from cutedsl.bridge import quantize_activation_nvfp4 - qsrc = inspect.getsource(quantize_activation_nvfp4) - if 'global_scale' in qsrc: - print(" quantize_activation_nvfp4 accepts global_scale as parameter") - if '_activation_global_scale' in qsrc: - print(" quantize_activation_nvfp4 reads _activation_global_scale") - - # Check _run_impl - run_impl_src = inspect.getsource(CuTeDSLNvfp4Linear._run_impl) - print(f"\n _run_impl length: {len(run_impl_src)} chars") - # Find where gs is used - for i, line in enumerate(run_impl_src.split('\n')): - if 'global_scale' in line or '_activation' in line: - print(f" _run_impl line {i}: {line.strip()}") + print(" ➡️ Changing gs DOES change the output!") + print(" ➡️ The runner uses the warmup gs at inference time.") + print(" ➡️ Wrong warmup gs would cause wrong quantization → garbage.") if __name__ == "__main__":