Add full layer 0 B200 test: CuTeDSL vs BF16 reference
Tests each attention/FFN projection individually against BF16 dequantized reference, then runs full layer forward. Identifies exactly where garbage enters the pipeline. Key finding: checkpoint uses different names than vLLM: - q_a_proj, q_b_proj, kv_proj (not fused_wqa_wkv) - q_a_norm (not q_norm) - compressor.* (C4A layers only) - sinks (attn_sink)
This commit is contained in:
340
tests/test_full_layer_b200.py
Normal file
340
tests/test_full_layer_b200.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user