diff --git a/tests/test_full_layer_b200.py b/tests/test_full_layer_b200.py new file mode 100644 index 00000000..f6446e05 --- /dev/null +++ b/tests/test_full_layer_b200.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python3""" +Full decoder layer 0 test: ALL components using CuTeDSL kernels, NO vLLM. + +Tests each attention + FFN projection individually (CuTeDSL vs BF16 ref), +then runs the full layer forward to identify where garbage enters. + +Usage (on B200): + source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate + export CUDA_TOOLKIT_PATH=/usr/local/cuda + python3 tests/test_full_layer_b200.py +""" + +import sys, os, json, math, torch, torch.nn.functional as F +from safetensors import safe_open + +REPO_ROOT = "/root/nvfp4-megamoe-kernel" +sys.path.insert(0, REPO_ROOT) + +MODEL_PATH = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +DEVICE = "cuda:0" +LAYER_IDX = 0 + +# Model config +HIDDEN = 7168 +N_HEADS = 128 +HEAD_DIM = 512 +NOPE = 448 +ROPE = 64 +Q_LORA = 1536 +O_LORA = 1024 +O_GROUPS = 16 +HPG = N_HEADS // O_GROUPS # 8 +HC_MULT = 4 +SWIGLU_LIM = 10.0 +RMS_EPS = 1e-6 +INTER = 3072 +N_SHARED = 1 +TOP_K = 6 +N_TOKENS = 4 + +E2M1_LUT = 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(key, wm, model_dir): + if key in _cache: + return _cache[key] + with safe_open(os.path.join(model_dir, wm[key]), framework="pt") as f: + t = f.get_tensor(key) + _cache[key] = t + return t + +def dequant(w, sf, gs): + dev = w.device + lut = E2M1_LUT.to(dev) + lo = lut[(w & 0xF).long()] + hi = lut[((w >> 4) & 0xF).long()] + O, I2 = w.shape + I = I2 * 2 + u = torch.empty(O, I, dtype=torch.float32, device=dev) + u[:, 0::2] = lo; u[:, 1::2] = hi + bs = sf.float().repeat_interleave(16, dim=1)[:O, :I] + return (u * bs * gs).to(torch.bfloat16) + +def rms_norm(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 build_cos_sin(max_pos=4096, rope_dim=ROPE): + half = rope_dim // 2 + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, half, dtype=torch.float32) / half)) + freqs = torch.outer(torch.arange(max_pos, dtype=torch.float32), inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], dim=-1) + +def inv_rope_bf16(o, pos, csc, nope=NOPE, rope=ROPE): + if rope == 0 or o.numel() == 0: return o + half = rope // 2 + cos = csc[pos, :half].unsqueeze(1).to(o.dtype) + sin = csc[pos, half:].unsqueeze(1).to(o.dtype) + r = o.clone() + or_ = o[:, :, nope:] + r[:, :, nope:][:, :, 0::2] = or_[:, :, 0::2] * cos + or_[:, :, 1::2] * sin + r[:, :, nope:][:, :, 1::2] = -or_[:, :, 0::2] * sin + or_[:, :, 1::2] * cos + return r + +def mhc_pre(residual, fn, scale, base, eps, pre_eps, sk_eps, post_mult, sk_rep): + hm = residual.shape[-2] + hs = residual.shape[-1] + os_ = residual.shape[:-2] + rf = residual.view(-1, hm, hs) + nt = rf.shape[0] + x = rf.view(nt, hm * hs).float() + mixes = x @ fn.t() + ss = x.square().sum(-1, keepdim=True) + mixes = mixes * torch.rsqrt(ss / (hm * hs) + eps) + pre = torch.sigmoid(mixes[:, :hm] * scale[0] + base[:hm]) + pre_eps + post = torch.sigmoid(mixes[:, hm:2*hm] * scale[1] + base[hm:2*hm]) * post_mult + comb = (mixes[:, 2*hm:].view(nt, hm, hm) * scale[2] + base[2*hm:].view(1, hm, hm)) + comb = torch.softmax(comb, -1) + sk_eps + comb = comb / (comb.sum(-2, keepdim=True) + sk_eps) + for _ in range(sk_rep - 1): + comb = comb / (comb.sum(-1, keepdim=True) + sk_eps) + comb = comb / (comb.sum(-2, keepdim=True) + sk_eps) + li = (pre.unsqueeze(-1) * rf.float()).sum(1).to(torch.bfloat16) + return (post.view(*os_, hm, 1), comb.view(*os_, hm, hm), li.view(*os_, hs)) + +def mhc_post(x, residual, post, comb): + mr = torch.einsum("...ij,...ih->...jh", comb.float(), residual.float()) + pt = post.float() * x.unsqueeze(-2).float() + return (mr + pt).to(residual.dtype) + +def make_runner(w, sf, gs_tensor, in_feat, out_feat, is_fused=False, lw=None): + from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear + dev = w.device + 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() + + if is_fused and gs_tensor.numel() == 2: + g1, g2 = gs_tensor[0].item(), gs_tensor[1].item() + gs = max(g1, g2) + if g1 != g2: + s32 = s.float() + sp = lw[0] if lw else out_feat // 2 + s32[:sp] *= g1 / gs; s32[sp:] *= g2 / gs + s = s32.to(torch.float8_e4m3fn) + else: + gs = gs_tensor.max().item() if gs_tensor.numel() > 1 else gs_tensor.item() + + r = CuTeDSLNvfp4Linear(in_features=in_feat, out_features=out_feat, max_num_tokens=8192, device=str(dev)) + r.fp4 = [fp4]; r.sf = [s]; r.gs = [gs] + r.finalize_weights(); r._ensure_initialized() + return r + +def cosine_sim(a, b): + return F.cosine_similarity(a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float()).item() + +def main(): + torch.cuda.set_device(0) + torch.manual_seed(42) + + print("=" * 70) + print(" Full Layer 0 Test: CuTeDSL NVFP4 vs BF16 Reference") + print("=" * 70) + + with open(os.path.join(MODEL_PATH, "model.safetensors.index.json")) as f: + wm = json.load(f)["weight_map"] + L = lambda k: P(k, wm, MODEL_PATH).to(DEVICE) + + pre = f"model.layers.{LAYER_IDX}" + ap = f"{pre}.self_attn" + mp = f"{pre}.mlp" + + # ── Load all weights ────────────────────────────────────────────── + print("\n--- Loading weights ---") + + # Attention + wqa_wkv_w = L(f"{ap}.fused_wqa_wkv.weight") + wqa_wkv_sf = L(f"{ap}.fused_wqa_wkv.weight_scale") + wqa_wkv_gs = L(f"{ap}.fused_wqa_wkv.weight_scale_2") + wq_b_w = L(f"{ap}.wq_b.weight") + wq_b_sf = L(f"{ap}.wq_b.weight_scale") + wq_b_gs = L(f"{ap}.wq_b.weight_scale_2") + wo_a_w = L(f"{ap}.o_a_proj.weight") # BF16 + wo_b_w = L(f"{ap}.o_b_proj.weight") + wo_b_sf = L(f"{ap}.o_b_proj.weight_scale") + wo_b_gs = L(f"{ap}.o_b_proj.weight_scale_2") + + # Norms + attn_norm = L(f"{pre}.input_layernorm.weight") + ffn_norm = L(f"{pre}.post_attention_layernorm.weight") + q_norm = L(f"{ap}.q_norm.weight") + kv_norm = L(f"{ap}.kv_norm.weight") + + # MHC + hc_attn_fn = L(f"{pre}.hc_attn_fn") + hc_ffn_fn = L(f"{pre}.hc_ffn_fn") + hc_attn_base = L(f"{pre}.hc_attn_base") + hc_ffn_base = L(f"{pre}.hc_ffn_base") + hc_attn_scale = L(f"{pre}.hc_attn_scale") + hc_ffn_scale = L(f"{pre}.hc_ffn_scale") + + for name, t in [ + ("fused_wqa_wkv", wqa_wkv_w), ("wq_b", wq_b_w), + ("wo_a", wo_a_w), ("wo_b", wo_b_w), + ("hc_attn_fn", hc_attn_fn), ("hc_ffn_fn", hc_ffn_fn), + ]: + print(f" {name}: shape={t.shape} dtype={t.dtype}") + + # ── Create CuTeDSL runners ──────────────────────────────────────── + print("\n--- Creating CuTeDSL runners ---") + + # fused_wqa_wkv: (q_a_rank + head_dim, hidden//2) = (2048, 3584) + r_wqa = make_runner(wqa_wkv_w, wqa_wkv_sf, wqa_wkv_gs, + HIDDEN, wqa_wkv_w.shape[0], is_fused=True, lw=[Q_LORA, HEAD_DIM]) + # wq_b: (num_heads*head_dim, q_lora_rank//2) = (65536, 768) + r_wqb = make_runner(wq_b_w, wq_b_sf, wq_b_gs, + wq_b_w.shape[1]*2, wq_b_w.shape[0]) + # wo_b: (hidden, o_groups*o_lora_rank//2) = (7168, 8192) + r_wob = make_runner(wo_b_w, wo_b_sf, wo_b_gs, + wo_b_w.shape[1]*2, wo_b_w.shape[0]) + + # Warmup all runners + print(" Warming up...") + d1 = torch.randn(N_TOKENS, HIDDEN, dtype=torch.bfloat16, device=DEVICE) * 2.0 + r_wqa.compute_activation_global_scale(d1) + d2 = torch.randn(N_TOKENS, Q_LORA, dtype=torch.bfloat16, device=DEVICE) * 2.0 + r_wqb.compute_activation_global_scale(d2) + d3 = torch.randn(N_TOKENS, O_GROUPS * O_LORA, dtype=torch.bfloat16, device=DEVICE) * 2.0 + r_wob.compute_activation_global_scale(d3) + print(" Done.") + + # ── Per-projection BF16 vs CuTeDSL comparison ──────────────────── + print("\n" + "=" * 70) + print(" PROJECTION-LEVEL: CuTeDSL vs BF16 Reference") + print("=" * 70) + + torch.manual_seed(123) + test_x = torch.randn(N_TOKENS, HIDDEN, dtype=torch.bfloat16, device=DEVICE) * 2.0 + + # 1. fused_wqa_wkv + with torch.no_grad(): + cutedsl_out = r_wqa.run(test_x) + wqa_bf16 = dequant(wqa_wkv_w, wqa_wkv_sf, wqa_wkv_gs.max().item() if wqa_wkv_gs.numel() > 1 else wqa_wkv_gs.item()) + # For fused: dequant each sub-proj separately + # Actually the global scales may differ — use per-sub dequant + if wqa_wkv_gs.numel() == 2: + gs1, gs2 = wqa_wkv_gs[0].item(), wqa_wkv_gs[1].item() + q_a_bf16 = dequant(wqa_wkv_w[:Q_LORA], wqa_wkv_sf[:Q_LORA], gs1) + kv_bf16 = dequant(wqa_wkv_w[Q_LORA:], wqa_wkv_sf[Q_LORA:], gs2) + ref_out = torch.cat([test_x @ q_a_bf16.T, test_x @ kv_bf16.T], dim=-1) + else: + ref_out = test_x @ dequant(wqa_wkv_w, wqa_wkv_sf, wqa_wkv_gs.item()).T + c = cosine_sim(cutedsl_out, ref_out) + print(f" fused_wqa_wkv: cosine={c:.6f} {'✅' if c >= 0.98 else '❌'} cutedsl_amax={cutedsl_out.amax():.4f} ref_amax={ref_out.amax():.4f}") + + # 2. wq_b + test_q = torch.randn(N_TOKENS, Q_LORA, dtype=torch.bfloat16, device=DEVICE) * 2.0 + with torch.no_grad(): + cutedsl_out = r_wqb.run(test_q) + ref_out = test_q @ dequant(wq_b_w, wq_b_sf, wq_b_gs.item()).T + c = cosine_sim(cutedsl_out, ref_out) + print(f" wq_b: cosine={c:.6f} {'✅' if c >= 0.98 else '❌'} cutedsl_amax={cutedsl_out.amax():.4f}") + + # 3. wo_b + test_z = torch.randn(N_TOKENS, O_GROUPS * O_LORA, dtype=torch.bfloat16, device=DEVICE) * 2.0 + with torch.no_grad(): + cutedsl_out = r_wob.run(test_z) + ref_out = test_z @ dequant(wo_b_w, wo_b_sf, wo_b_gs.item()).T + c = cosine_sim(cutedsl_out, ref_out) + print(f" wo_b: cosine={c:.6f} {'✅' if c >= 0.98 else '❌'} cutedsl_amax={cutedsl_out.amax():.4f}") + + # 4. wo_a (BF16, no CuTeDSL — just matmul) + test_o = torch.randn(N_TOKENS, N_HEADS, HEAD_DIM, dtype=torch.bfloat16, device=DEVICE) * 0.1 + csc = build_cos_sin().to(DEVICE) + pos = torch.arange(N_TOKENS, dtype=torch.int64, device=DEVICE) + o_inv = inv_rope_bf16(test_o, pos, csc) + o_g = o_inv.view(N_TOKENS, O_GROUPS, HPG * HEAD_DIM).permute(1, 0, 2) + wo_a_3d = wo_a_w.view(O_GROUPS, O_LORA, HPG * HEAD_DIM) + z_bmm = torch.bmm(o_g, wo_a_3d.transpose(1, 2)).permute(1, 0, 2).reshape(N_TOKENS, O_GROUPS * O_LORA) + # Reference: flatten matmul + z_ref = o_inv.reshape(N_TOKENS, N_HEADS * HEAD_DIM) @ wo_a_w.T + c = cosine_sim(z_bmm, z_ref) + print(f" wo_a (BMM): cosine={c:.6f} {'✅' if c >= 0.99 else '❌'} (should be ~1.0, BF16)") + + # ── 5. Shared expert (CuTeDSL vs BF16) ─────────────────────────── + print("\n--- Shared Expert: CuTeDSL vs BF16 ---") + from cutedsl.shared_expert_pipeline import CuTeDSLSharedExpertRunner + + se_gw = L(f"{mp}.shared_experts.gate_proj.weight") + se_gsf = L(f"{mp}.shared_experts.gate_proj.weight_scale") + se_ggs = L(f"{mp}.shared_experts.gate_proj.weight_scale_2").item() + se_uw = L(f"{mp}.shared_experts.up_proj.weight") + se_usf = L(f"{mp}.shared_experts.up_proj.weight_scale") + se_ugs = L(f"{mp}.shared_experts.up_proj.weight_scale_2").item() + se_dw = L(f"{mp}.shared_experts.down_proj.weight") + se_dsf = L(f"{mp}.shared_experts.down_proj.weight_scale") + se_dgs = L(f"{mp}.shared_experts.down_proj.weight_scale_2").item() + + se_inter = INTER * N_SHARED # 3072 + + se_gu_w = torch.cat([se_gw, se_uw], dim=0) + se_gu_sf = torch.cat([se_gsf, se_usf], dim=0) + se_mgs = max(se_ggs, se_ugs) + if se_ggs != se_ugs: + sf32 = se_gu_sf.float() + sf32[:se_inter] *= se_ggs / se_mgs + sf32[se_inter:] *= se_ugs / se_mgs + se_gu_sf = sf32.to(torch.float8_e4m3fn) + + ser = CuTeDSLSharedExpertRunner(hidden_size=HIDDEN, intermediate_size=se_inter, + max_num_tokens=8192, device=DEVICE, swiglu_limit=SWIGLU_LIM) + ser.l1_fp4 = [se_gu_w.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] + ser.l1_sf = [se_gu_sf.permute(1, 0).contiguous()] + ser.l1_gs = [se_mgs] + ser.l2_fp4 = [se_dw.view(torch.float4_e2m1fn_x2).permute(1, 0).contiguous()] + ser.l2_sf = [se_dsf.permute(1, 0).contiguous()] + ser.l2_gs = [se_dgs] + ser.finalize_weights() + ser._ensure_initialized() + + test_se = torch.randn(N_TOKENS, HIDDEN, dtype=torch.bfloat16, device=DEVICE) * 2.0 + ser.compute_activation_global_scales(test_se) + + with torch.no_grad(): + se_out = ser.run(test_se) + + # BF16 ref + g_bf16 = dequant(se_gw, se_gsf, se_ggs) + u_bf16 = dequant(se_uw, se_usf, se_ugs) + d_bf16 = dequant(se_dw, se_dsf, se_dgs) + with torch.no_grad(): + g = test_se @ g_bf16.T + u = test_se @ u_bf16.T + act = F.silu(g.clamp(max=SWIGLU_LIM)) * u.clamp(min=-SWIGLU_LIM, max=SWIGLU_LIM) + se_ref = act @ d_bf16.T + + c = cosine_sim(se_out, se_ref) + print(f" shared_expert: cosine={c:.6f} {'✅' if c >= 0.98 else '❌'} cutedsl_amax={se_out.amax():.4f} ref_amax={se_ref.amax():.4f}") + + # ── 6. MHC sanity check ────────────────────────────────────────── + print("\n--- MHC sanity check ---") + x_hc = torch.randn(N_TOKENS, HC_MULT, HIDDEN, dtype=torch.bfloat16, device=DEVICE) * 0.1 + post_m, res_m, li = mhc_pre(x_hc, hc_attn_fn, hc_attn_scale, hc_attn_base, + RMS_EPS, 1e-6, 1e-6, 2.0, 20) + print(f" mhc_pre output: amax={li.amax():.4f} NaN={torch.isnan(li).any()}") + new_res = mhc_post(li, x_hc, post_m, res_m) + print(f" mhc_post output: amax={new_res.amax():.4f} NaN={torch.isnan(new_res).any()}") + + # ── Summary ─────────────────────────────────────────────────────── + print("\n" + "=" * 70) + print(" If all projections pass (cosine >= 0.98), the CuTeDSL kernels") + print(" are correct and the bug is in vLLM's pipeline, not our kernels.") + print(" If any fail, that projection needs debugging.") + print("=" * 70) + + +if __name__ == "__main__": + main()