From f97762cc9ff6827dc0206fa64a1f64e16ce94058 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 19 May 2026 07:16:33 +0000 Subject: [PATCH] Fix full layer test: use correct checkpoint key names MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Checkpoint uses q_a_proj/q_b_proj/kv_proj/q_a_norm — NOT the vLLM fused names (fused_wqa_wkv, wq_b, q_norm). --- tests/test_full_layer_b200.py | 454 ++++++++++++++-------------------- 1 file changed, 187 insertions(+), 267 deletions(-) diff --git a/tests/test_full_layer_b200.py b/tests/test_full_layer_b200.py index f6446e05..37c8e46e 100644 --- a/tests/test_full_layer_b200.py +++ b/tests/test_full_layer_b200.py @@ -1,4 +1,5 @@ -#!/usr/bin/env python3""" +#!/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), @@ -6,335 +7,254 @@ 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 +REPO = "/root/nvfp4-megamoe-kernel" +sys.path.insert(0, REPO) +MODEL = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4" +DEV = "cuda:0" +L = 0 # layer index # Model config -HIDDEN = 7168 -N_HEADS = 128 -HEAD_DIM = 512 +H = 7168 +NH = 128 +HD = 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 +QL = 1536 +OL = 1024 +OG = 16 +HPG = NH // OG # 8 +HC = 4 +SL = 10.0 +EPS = 1e-6 INTER = 3072 -N_SHARED = 1 -TOP_K = 6 -N_TOKENS = 4 +NT = 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) +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(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 +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 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] + d = w.device; lut = E2M1.to(d) + 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=d) + 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): +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) + 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 csc(max_p=4096, rd=ROPE): + hf = rd//2 + inv = 1.0/(10000.0**(torch.arange(0,hf,dtype=torch.float32)/hf)) + fr = torch.outer(torch.arange(max_p,dtype=torch.float32), inv) + return torch.cat([fr.cos(), fr.sin()], -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 +def inv_rope(o, pos, cs, nope=NOPE, rope=ROPE): + if rope==0 or o.numel()==0: return o + hf=rope//2; c=cs[pos,:hf].unsqueeze(1).to(o.dtype); s=cs[pos,hf:].unsqueeze(1).to(o.dtype) + r=o.clone(); q=r[:,:,nope:] + r[:,:,nope:][:,:,0::2]=q[:,:,0::2]*c+q[:,:,1::2]*s + r[:,:,nope:][:,:,1::2]=-q[:,:,0::2]*s+q[:,:,1::2]*c 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_pre(res, fn, sc, bs, eps, pe, se, pm, sr): + hm=res.shape[-2]; hs=res.shape[-1]; os_=res.shape[:-2] + rf=res.view(-1,hm,hs); nt=rf.shape[0] + x=rf.view(nt,hm*hs).float(); mx=x@fn.t(); ss=x.square().sum(-1,keepdim=True) + mx=mx*torch.rsqrt(ss/(hm*hs)+eps) + pre=torch.sigmoid(mx[:,:hm]*sc[0]+bs[:hm])+pe + post=torch.sigmoid(mx[:,hm:2*hm]*sc[1]+bs[hm:2*hm])*pm + cb=mx[:,2*hm:].view(nt,hm,hm)*sc[2]+bs[2*hm:].view(1,hm,hm) + cb=torch.softmax(cb,-1)+se; cb=cb/(cb.sum(-2,keepdim=True)+se) + for _ in range(sr-1): + cb=cb/(cb.sum(-1,keepdim=True)+se); cb=cb/(cb.sum(-2,keepdim=True)+se) + li=(pre.unsqueeze(-1)*rf.float()).sum(1).to(torch.bfloat16) + return (post.view(*os_,hm,1),cb.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 mhc_post(x, res, post, comb): + mr=torch.einsum("...ij,...ih->...jh",comb.float(),res.float()) + pt=post.float()*x.unsqueeze(-2).float() + return (mr+pt).to(res.dtype) -def make_runner(w, sf, gs_tensor, in_feat, out_feat, is_fused=False, lw=None): +def make_runner(w, sf, gs_t, inf, outf, 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) + 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 fused and gs_t.numel()==2: + g1,g2=gs_t[0].item(),gs_t[1].item(); gs=max(g1,g2) + if g1!=g2: + s32=s.float(); sp=lw[0] if lw else outf//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] + gs=gs_t.max().item() if gs_t.numel()>1 else gs_t.item() + r=CuTeDSLNvfp4Linear(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 cosine_sim(a, b): - return F.cosine_similarity(a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float()).item() +def cosim(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) + torch.cuda.set_device(0); torch.manual_seed(42) + print("="*70+"\n Layer 0 Test: CuTeDSL NVFP4 vs BF16 Reference\n"+"="*70) - print("=" * 70) - print(" Full Layer 0 Test: CuTeDSL NVFP4 vs BF16 Reference") - 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) - 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) + p=f"model.layers.{L}"; a=f"{p}.self_attn"; m=f"{p}.mlp" - pre = f"model.layers.{LAYER_IDX}" - ap = f"{pre}.self_attn" - mp = f"{pre}.mlp" - - # ── Load all weights ────────────────────────────────────────────── + # Checkpoint key names (NOT vLLM names!) print("\n--- Loading weights ---") + 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") + 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_w=G(f"{a}.kv_proj.weight"); kv_sf=G(f"{a}.kv_proj.weight_scale"); kv_gs=G(f"{a}.kv_proj.weight_scale_2") + woa=G(f"{a}.o_a_proj.weight") # BF16 + wob_w=G(f"{a}.o_b_proj.weight"); wob_sf=G(f"{a}.o_b_proj.weight_scale"); wob_gs=G(f"{a}.o_b_proj.weight_scale_2") + anorm=G(f"{p}.input_layernorm.weight"); fnorm=G(f"{p}.post_attention_layernorm.weight") + qn=G(f"{a}.q_a_norm.weight"); kvn=G(f"{a}.kv_norm.weight") + hca_fn=G(f"{p}.hc_attn_fn"); hcf_fn=G(f"{p}.hc_ffn_fn") + hca_b=G(f"{p}.hc_attn_base"); hcf_b=G(f"{p}.hc_ffn_base") + hca_s=G(f"{p}.hc_attn_scale"); hcf_s=G(f"{p}.hc_ffn_scale") - # 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") + for nm,t in [("q_a_proj",qa_w),("q_b_proj",qb_w),("kv_proj",kv_w),("o_a_proj",woa),("o_b_proj",wob_w)]: + print(f" {nm}: shape={t.shape} dtype={t.dtype}") + print(f" q_a_proj gs: {qa_gs.tolist()}") + print(f" q_b_proj gs: {qb_gs.tolist()}") + print(f" kv_proj gs: {kv_gs.tolist()}") - # 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 ──────────────────────────────────────── + # Create CuTeDSL runners (separate projections, not fused) print("\n--- Creating CuTeDSL runners ---") + r_qa = make_runner(qa_w, qa_sf, qa_gs, qa_w.shape[1]*2, qa_w.shape[0]) + r_qb = make_runner(qb_w, qb_sf, qb_gs, qb_w.shape[1]*2, qb_w.shape[0]) + r_kv = make_runner(kv_w, kv_sf, kv_gs, kv_w.shape[1]*2, kv_w.shape[0]) + r_wob = make_runner(wob_w, wob_sf, wob_gs, wob_w.shape[1]*2, wob_w.shape[0]) + print(f" q_a: in={qa_w.shape[1]*2} out={qa_w.shape[0]}") + print(f" q_b: in={qb_w.shape[1]*2} out={qb_w.shape[0]}") + print(f" kv: in={kv_w.shape[1]*2} out={kv_w.shape[0]}") + print(f" wo_b: in={wob_w.shape[1]*2} out={wob_w.shape[0]}") - # 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 + # Warmup 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 + d1=torch.randn(NT,H,dtype=torch.bfloat16,device=DEV)*2.0 + r_qa.compute_activation_global_scale(d1); r_kv.compute_activation_global_scale(d1) + d2=torch.randn(NT,QL,dtype=torch.bfloat16,device=DEV)*2.0 + r_qb.compute_activation_global_scale(d2) + d3=torch.randn(NT,OG*OL,dtype=torch.bfloat16,device=DEV)*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) - + # Per-projection BF16 vs CuTeDSL comparison + print("\n"+"="*70+"\n PROJECTION-LEVEL: CuTeDSL vs BF16\n"+"="*70) torch.manual_seed(123) - test_x = torch.randn(N_TOKENS, HIDDEN, dtype=torch.bfloat16, device=DEVICE) * 2.0 + tx=torch.randn(NT,H,dtype=torch.bfloat16,device=DEV)*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}") + # q_a_proj + with torch.no_grad(): co=r_qa.run(tx) + ref=tx@dequant(qa_w,qa_sf,qa_gs.item()).T + c=cosim(co,ref) + print(f" q_a_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.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}") + # kv_proj + with torch.no_grad(): co=r_kv.run(tx) + ref=tx@dequant(kv_w,kv_sf,kv_gs.item()).T + c=cosim(co,ref) + print(f" kv_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.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}") + # q_b_proj + tq=torch.randn(NT,QL,dtype=torch.bfloat16,device=DEV)*2.0 + with torch.no_grad(): co=r_qb.run(tq) + ref=tq@dequant(qb_w,qb_sf,qb_gs.item()).T + c=cosim(co,ref) + print(f" q_b_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.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)") + # wo_b_proj + tz=torch.randn(NT,OG*OL,dtype=torch.bfloat16,device=DEV)*2.0 + with torch.no_grad(): co=r_wob.run(tz) + ref=tz@dequant(wob_w,wob_sf,wob_gs.item()).T + c=cosim(co,ref) + print(f" wo_b_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}") - # ── 5. Shared expert (CuTeDSL vs BF16) ─────────────────────────── + # wo_a (BF16 — just test BMM correctness) + to_=torch.randn(NT,NH,HD,dtype=torch.bfloat16,device=DEV)*0.1 + cs_=csc().to(DEV); pos=torch.arange(NT,dtype=torch.int64,device=DEV) + oi=inv_rope(to_,pos,cs_) + og=oi.view(NT,OG,HPG*HD).permute(1,0,2) + wa3=woa.view(OG,OL,HPG*HD) + z_bmm=torch.bmm(og,wa3.transpose(1,2)).permute(1,0,2).reshape(NT,OG*OL) + z_ref=oi.reshape(NT,NH*HD)@woa.T + c=cosim(z_bmm,z_ref) + print(f" wo_a (BF16 BMM): cosine={c:.6f} {'✅' if c>=0.99 else '❌'}") + + # Shared expert 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() + sgw=G(f"{m}.shared_experts.gate_proj.weight"); sgsf=G(f"{m}.shared_experts.gate_proj.weight_scale") + sggs=G(f"{m}.shared_experts.gate_proj.weight_scale_2").item() + suw=G(f"{m}.shared_experts.up_proj.weight"); susf=G(f"{m}.shared_experts.up_proj.weight_scale") + sugest=G(f"{m}.shared_experts.up_proj.weight_scale_2").item() + sdw=G(f"{m}.shared_experts.down_proj.weight"); sdsf=G(f"{m}.shared_experts.down_proj.weight_scale") + sdgs=G(f"{m}.shared_experts.down_proj.weight_scale_2").item() - se_inter = INTER * N_SHARED # 3072 + si=INTER # 3072 + sgu_w=torch.cat([sgw,suw],0); sgu_sf=torch.cat([sgsf,susf],0) + smgs=max(sggs,sugest) + if sggs!=sugest: + s32=sgu_sf.float(); s32[:si]*=sggs/smgs; s32[si:]*=sugest/smgs + sgu_sf=s32.to(torch.float8_e4m3fn) - 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=H,intermediate_size=si,max_num_tokens=8192, + device=DEV,swiglu_limit=SL) + ser.l1_fp4=[sgu_w.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()] + ser.l1_sf=[sgu_sf.permute(1,0).contiguous()]; ser.l1_gs=[smgs] + ser.l2_fp4=[sdw.view(torch.float4_e2m1fn_x2).permute(1,0).contiguous()] + ser.l2_sf=[sdsf.permute(1,0).contiguous()]; ser.l2_gs=[sdgs] + ser.finalize_weights(); ser._ensure_initialized() - 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) + tse=torch.randn(NT,H,dtype=torch.bfloat16,device=DEV)*2.0 + ser.compute_activation_global_scales(tse) + with torch.no_grad(): so=ser.run(tse) # 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) + gb=dequant(sgw,sgsf,sggs); ub=dequant(suw,susf,sugest); db=dequant(sdw,sdsf,sdgs) 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 + g_=tse@gb.T; u_=tse@ub.T + act=F.silu(g_.clamp(max=SL))*u_.clamp(min=-SL,max=SL) + sref=act@db.T + c=cosim(so,sref) + print(f" shared_expert: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={so.amax():.4f} ref={sref.amax():.4f}") - 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 ────────────────────────────────────────── + # MHC sanity 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()}") + xhc=torch.randn(NT,HC,H,dtype=torch.bfloat16,device=DEV)*0.1 + pm,rm,li=mhc_pre(xhc,hca_fn,hca_s,hca_b,EPS,1e-6,1e-6,2.0,20) + print(f" mhc_pre: li amax={li.amax():.4f} NaN={torch.isnan(li).any()}") + nr=mhc_post(li,xhc,pm,rm) + print(f" mhc_post: amax={nr.amax():.4f} NaN={torch.isnan(nr).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("\n"+"="*70) + print(" If all projections pass (cosine >= 0.98), CuTeDSL kernels are") + print(" correct and the bug is in vLLM's pipeline, not our kernels.") print(" If any fail, that projection needs debugging.") - print("=" * 70) + print("="*70) - -if __name__ == "__main__": +if __name__=="__main__": main()