Fix checkpoint keys: attn_hc.*, compressor.*, q_a_proj/q_b_proj/kv_proj
This commit is contained in:
@@ -2,8 +2,7 @@
|
||||
"""
|
||||
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.
|
||||
Tests each attention + FFN projection individually (CuTeDSL vs BF16 ref).
|
||||
|
||||
Usage (on B200):
|
||||
source /root/nvfp4-megamoe-kernel/tests/.venv/bin/activate
|
||||
@@ -17,9 +16,8 @@ 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
|
||||
# Model config (layer 0, compress_ratio=128 → C4A)
|
||||
H = 7168
|
||||
NH = 128
|
||||
HD = 512
|
||||
@@ -28,7 +26,7 @@ ROPE = 64
|
||||
QL = 1536
|
||||
OL = 1024
|
||||
OG = 16
|
||||
HPG = NH // OG # 8
|
||||
HPG = NH // OG
|
||||
HC = 4
|
||||
SL = 10.0
|
||||
EPS = 1e-6
|
||||
@@ -58,203 +56,203 @@ 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 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(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(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, 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_t, inf, outf, fused=False, lw=None):
|
||||
from cutedsl.nvfp4_linear import CuTeDSLNvfp4Linear
|
||||
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)
|
||||
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_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]
|
||||
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 cosim(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().to(a.device)).item()
|
||||
|
||||
def main():
|
||||
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(" Layer 0: 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, "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.{L}"; a=f"{p}.self_attn"; m=f"{p}.mlp"
|
||||
p = "model.layers.0"; a = f"{p}.self_attn"; m = f"{p}.mlp"
|
||||
|
||||
# Checkpoint key names (NOT vLLM names!)
|
||||
# ── Load attention weights (correct checkpoint key 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")
|
||||
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")
|
||||
qn = G(f"{a}.q_a_norm.weight"); kvn = G(f"{a}.kv_norm.weight")
|
||||
anorm = G(f"{p}.input_layernorm.weight"); fnorm = G(f"{p}.post_attention_layernorm.weight")
|
||||
|
||||
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)]:
|
||||
# Compressor (C4A path)
|
||||
ckv_w = G(f"{a}.compressor.kv_proj.weight"); ckv_sf = G(f"{a}.compressor.kv_proj.weight_scale"); ckv_gs = G(f"{a}.compressor.kv_proj.weight_scale_2")
|
||||
cg_w = G(f"{a}.compressor.gate_proj.weight"); cg_sf = G(f"{a}.compressor.gate_proj.weight_scale"); cg_gs = G(f"{a}.compressor.gate_proj.weight_scale_2")
|
||||
ckn = G(f"{a}.compressor.kv_norm.weight")
|
||||
cpb = G(f"{a}.compressor.position_bias")
|
||||
sinks = G(f"{a}.sinks")
|
||||
|
||||
# MHC
|
||||
hca_fn = G(f"{p}.attn_hc.fn"); hcf_fn = G(f"{p}.ffn_hc.fn")
|
||||
hca_b = G(f"{p}.attn_hc.base"); hcf_b = G(f"{p}.ffn_hc.base")
|
||||
hca_s = G(f"{p}.attn_hc.scale"); hcf_s = G(f"{p}.ffn_hc.scale")
|
||||
|
||||
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),
|
||||
("comp.kv_proj", ckv_w), ("comp.gate_proj", cg_w),
|
||||
("sinks", sinks), ("comp.position_bias", cpb),
|
||||
("attn_hc.fn", hca_fn)]:
|
||||
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()}")
|
||||
|
||||
# Create CuTeDSL runners (separate projections, not fused)
|
||||
# ── Create CuTeDSL runners ────────────────────────────────────────
|
||||
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])
|
||||
|
||||
# Compressor runners
|
||||
r_ckv = make_runner(ckv_w, ckv_sf, ckv_gs, ckv_w.shape[1]*2, ckv_w.shape[0])
|
||||
r_cg = make_runner(cg_w, cg_sf, cg_gs, cg_w.shape[1]*2, cg_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]}")
|
||||
print(f" comp.kv: in={ckv_w.shape[1]*2} out={ckv_w.shape[0]}")
|
||||
print(f" comp.gate: in={cg_w.shape[1]*2} out={cg_w.shape[0]}")
|
||||
|
||||
# Warmup
|
||||
print(" Warming up...")
|
||||
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
|
||||
d1 = torch.randn(NT, H, dtype=torch.bfloat16, device=DEV)*2.0
|
||||
for r in [r_qa, r_kv, r_ckv, r_cg]:
|
||||
r.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
|
||||
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+"\n PROJECTION-LEVEL: CuTeDSL vs BF16\n"+"="*70)
|
||||
# ── Per-projection BF16 vs CuTeDSL comparison ────────────────────
|
||||
print("\n" + "="*70)
|
||||
print(" PROJECTION-LEVEL: CuTeDSL vs BF16")
|
||||
print("="*70)
|
||||
torch.manual_seed(123)
|
||||
tx=torch.randn(NT,H,dtype=torch.bfloat16,device=DEV)*2.0
|
||||
tx = torch.randn(NT, H, dtype=torch.bfloat16, device=DEV)*2.0
|
||||
|
||||
results = {}
|
||||
|
||||
# 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}")
|
||||
with torch.no_grad(): co = r_qa.run(tx)
|
||||
ref = tx @ dequant(qa_w, qa_sf, qa_gs.item()).T
|
||||
c = cosim(co, ref); results['q_a_proj'] = c
|
||||
print(f" q_a_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.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}")
|
||||
with torch.no_grad(): co = r_kv.run(tx)
|
||||
ref = tx @ dequant(kv_w, kv_sf, kv_gs.item()).T
|
||||
c = cosim(co, ref); results['kv_proj'] = c
|
||||
print(f" kv_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.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}")
|
||||
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); results['q_b_proj'] = c
|
||||
print(f" q_b_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}")
|
||||
|
||||
# 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}")
|
||||
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); results['wo_b_proj'] = c
|
||||
print(f" wo_b_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}")
|
||||
|
||||
# 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 '❌'}")
|
||||
# compressor kv_proj
|
||||
with torch.no_grad(): co = r_ckv.run(tx)
|
||||
ref = tx @ dequant(ckv_w, ckv_sf, ckv_gs.item()).T
|
||||
c = cosim(co, ref); results['comp.kv_proj'] = c
|
||||
print(f" comp.kv_proj: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}")
|
||||
|
||||
# Shared expert
|
||||
# compressor gate_proj
|
||||
with torch.no_grad(): co = r_cg.run(tx)
|
||||
ref = tx @ dequant(cg_w, cg_sf, cg_gs.item()).T
|
||||
c = cosim(co, ref); results['comp.gate_proj'] = c
|
||||
print(f" comp.gate: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={co.amax():.4f} ref={ref.amax():.4f}")
|
||||
|
||||
# ── Shared expert ─────────────────────────────────────────────────
|
||||
print("\n--- Shared Expert: CuTeDSL vs BF16 ---")
|
||||
from cutedsl.shared_expert_pipeline import CuTeDSLSharedExpertRunner
|
||||
|
||||
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()
|
||||
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")
|
||||
sugs = 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()
|
||||
|
||||
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)
|
||||
si = INTER
|
||||
sgu_w = torch.cat([sgw, suw], 0); sgu_sf = torch.cat([sgsf, susf], 0)
|
||||
smgs = max(sggs, sugs)
|
||||
if sggs != sugs:
|
||||
s32 = sgu_sf.float(); s32[:si] *= sggs/smgs; s32[si:] *= sugs/smgs
|
||||
sgu_sf = s32.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 = 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()
|
||||
|
||||
tse=torch.randn(NT,H,dtype=torch.bfloat16,device=DEV)*2.0
|
||||
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)
|
||||
with torch.no_grad(): so = ser.run(tse)
|
||||
|
||||
# BF16 ref
|
||||
gb=dequant(sgw,sgsf,sggs); ub=dequant(suw,susf,sugest); db=dequant(sdw,sdsf,sdgs)
|
||||
gb = dequant(sgw, sgsf, sggs); ub = dequant(suw, susf, sugs); db = dequant(sdw, sdsf, sdgs)
|
||||
with torch.no_grad():
|
||||
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)
|
||||
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); results['shared_expert'] = c
|
||||
print(f" shared_expert: cosine={c:.6f} {'✅' if c>=0.98 else '❌'} amax={so.amax():.4f} ref={sref.amax():.4f}")
|
||||
|
||||
# MHC sanity
|
||||
print("\n--- MHC sanity check ---")
|
||||
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()}")
|
||||
# ── MHC sanity ────────────────────────────────────────────────────
|
||||
print("\n--- MHC weight shapes ---")
|
||||
print(f" attn_hc.fn: {hca_fn.shape} ffn_hc.fn: {hcf_fn.shape}")
|
||||
print(f" attn_hc.base: {hca_b.shape} ffn_hc.base: {hcf_b.shape}")
|
||||
print(f" attn_hc.scale: {hca_s.shape} ffn_hc.scale: {hcf_s.shape}")
|
||||
print(f" attn_hc.scale values: {hca_s.tolist()}")
|
||||
print(f" sinks: {sinks.shape}")
|
||||
|
||||
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.")
|
||||
# ── Summary ───────────────────────────────────────────────────────
|
||||
print("\n" + "="*70)
|
||||
print(" SUMMARY")
|
||||
print("="*70)
|
||||
all_pass = True
|
||||
for name, c in results.items():
|
||||
status = '✅' if c >= 0.98 else '❌'
|
||||
if c < 0.98: all_pass = False
|
||||
print(f" {name}: {c:.6f} {status}")
|
||||
if all_pass:
|
||||
print("\n All projections pass! CuTeDSL kernels match BF16 reference.")
|
||||
print(" The bug is in vLLM's pipeline, not our kernels.")
|
||||
else:
|
||||
print("\n Some projections FAIL. Need to debug those specific kernels.")
|
||||
|
||||
if __name__=="__main__":
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user