Fix full layer test: use correct checkpoint key names
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).
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user