#!/usr/bin/env python3 """Production FMHA layer comparison test — real model weights, real pipeline. Strategy: 1. Run the full production pipeline (single_shot_inference.py forward_layer) for all prefill tokens through layers 0-4. 2. On the LAST prefill token, for each layer, ALSO run the reference FMHA (dequantize KV to BF16, run PyTorch SDPA) on the SAME gathered KV that the production kernel saw. 3. Compare raw FMHA output (before inverse RoPE, before output projection). This isolates the FMHA kernel's accuracy from the rest of the pipeline. Production values: HD=512, NOPE=448, ROPE=64, H=128, 61 layers, 8 GPUs. """ import os, sys, json, math, time import torch import torch.nn.functional as F CHECKPOINT_DIR = os.environ.get( "CHECKPOINT_DIR", "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4") NUM_GPUS = int(os.environ.get("NUM_GPUS", "8")) DEVICE = "cuda:0" def cosine(a, b): return F.cosine_similarity(a.flatten().float(), b.flatten().float(), dim=0).item() def main(): torch.manual_seed(42) print("=" * 70) print("PRODUCTION FMHA LAYER COMPARISON TEST") print("=" * 70) with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f: cfg = json.load(f) n_layers = cfg["num_hidden_layers"] H = cfg["hidden_size"] hd = cfg["head_dim"] n_h = cfg["num_attention_heads"] rd = cfg.get("qk_rope_head_dim", 64) nope_dim = hd - rd cr = cfg.get("compress_ratios", [128] * n_layers) print(f"Model: {n_layers} layers, {n_h} heads, hd={hd}, rope_dim={rd}") from single_shot_inference import ( load_all_weights, make_nvfp4_linear, get_nvfp4_weight, rmsnorm, unweighted_rmsnorm, _apply_rope, build_rope_cache, KVCache, Compressor, Indexer, forward_layer, moe_forward, _load_moe_weights_stacked, _load_shared_expert_weights, _cache_layer_weights_no_experts, ) from dsv4.layers.mhc import mHCLayer, mHCContext from dsv4.layers.router import Router from dsv4.layers.moe import Nvfp4MoE from dsv4.layers.shared_expert import Nvfp4SharedExpert from dsv4.layers.grouped_linear import Nvfp4GroupedLinear from dsv4.layers.linear import Nvfp4Linear from dsv4.ops.quantize import ( rmsnorm_quantize_nvfp4, mhc_rmsnorm_quantize_nvfp4, dequantize_nvfp4, quantize_to_nvfp4, ) print("Loading weights...") all_w = load_all_weights(CHECKPOINT_DIR) TEST_LAYERS = 5 o_groups = cfg.get("o_groups", 16) o_rank = cfg.get("o_lora_rank", 1024) n_ih = cfg.get("index_n_heads", 64) ihd = cfg.get("index_head_dim", 128) itk = cfg.get("index_topk", 1024) rope_caches = {g: build_rope_cache(65536, rd, f"cuda:{g}", 10000., "yarn", 16., 4096, 32, 1) for g in range(NUM_GPUS)} # Build all production components (same as single_shot main()) prod_lins, attn_mhcs, ffn_mhcs = {}, {}, {} attn_norms, ffn_norms = {}, {} compressors, indexers, kv_caches = {}, {}, {} routers, moe_runners, se_runners = {}, {}, {} for li in range(TEST_LAYERS): gpu = li % NUM_GPUS dev = f"cuda:{gpu}" torch.cuda.set_device(gpu) pfx = f"model.layers.{li}.self_attn" mlp_pfx = f"model.layers.{li}.mlp" ratio = cr[li] if li < len(cr) else 128 # Attention linears pl = {} pl['q_a'] = make_nvfp4_linear(H, 1536, dev, all_w, pfx, 'q_a_proj') pl['q_b'] = make_nvfp4_linear(1536, H * hd, dev, all_w, pfx, 'q_b_proj') pl['kv'] = make_nvfp4_linear(H, hd, dev, all_w, pfx, 'kv_proj') hpg = n_h // o_groups wo_a = Nvfp4GroupedLinear(n_local_groups=o_groups, heads_per_group=hpg, head_dim=hd, o_lora_rank=o_rank, max_num_tokens=8192, device=dev) oa_w, oa_ws, oa_ws2, oa_isc = get_nvfp4_weight(all_w, pfx, 'o_a_proj') if oa_w is not None and oa_ws is not None: wo_a.load_nvfp4_weight(oa_w.to(dev), oa_ws.to(dev), oa_ws2.to(dev) if oa_ws2 is not None else None, oa_isc.to(dev) if oa_isc is not None else None) else: oa_bf = all_w.get(f"{pfx}.o_a_proj.weight") if oa_bf is not None: wo_a.set_bf16_weight(oa_bf.bfloat16().to(dev)) pl['o_a'] = wo_a; wo_a._use_runtime_gsa = True pl['o_b'] = make_nvfp4_linear(o_groups * o_rank, H, dev, all_w, pfx, 'o_b_proj') prod_lins[li] = pl # mHC for tag, blocks, fn_s, base_s, scale_s in [ ("attn", attn_mhcs, f"model.layers.{li}.attn_hc.fn", f"model.layers.{li}.attn_hc.base", f"model.layers.{li}.attn_hc.scale"), ("ffn", ffn_mhcs, f"model.layers.{li}.ffn_hc.fn", f"model.layers.{li}.ffn_hc.base", f"model.layers.{li}.ffn_hc.scale"), ]: fn, base, scale = all_w.get(fn_s), all_w.get(base_s), all_w.get(scale_s) if fn is not None and base is not None and scale is not None: m = mHCLayer(hidden_dim=H, n_hc=4, t_max_sinkhorn=20, device=dev) n = 4 m.load_weights( W_pre=fn[0:n].to(dev, torch.float32), W_post=fn[n:2*n].to(dev, torch.float32), W_comb=fn[2*n:].to(dev, torch.float32), S_pre=base[0:n].reshape(1, n).to(dev, torch.float32), S_post=base[n:2*n].reshape(n, 1).to(dev, torch.float32), S_comb=base[2*n:].reshape(n, n).to(dev, torch.float32), alpha_pre=scale[0].item(), alpha_post=scale[1].item(), alpha_comb=scale[2].item()) blocks[li] = m an_k = f"model.layers.{li}.input_layernorm.weight" if an_k in all_w: attn_norms[li] = all_w[an_k].to(dev, torch.float32) fn_k = f"model.layers.{li}.post_attention_layernorm.weight" if fn_k in all_w: ffn_norms[li] = all_w[fn_k].to(dev, torch.float32) max_comp = (8192 + ratio - 1) // ratio if ratio > 0 else 0 kv_caches[li] = KVCache(hd, cfg.get("sliding_window", 128), max_comp=max_comp, device=dev, indexer_key_dim=ihd, compress_ratio=ratio, indexer_top_k=itk, rope_dim=rd) if ratio > 0: compressors[li] = Compressor(ratio, hd, H, dev) if ratio == 4: indexers[li] = Indexer(n_ih, ihd, itk, dev) # Router is_hash = (li < cfg.get("num_hash_layers", 3)) and (f"{mlp_pfx}.gate.tid2eid" in all_w) router = Router(hidden_size=H, num_experts=cfg["n_routed_experts"], top_k=cfg.get("num_experts_per_tok", 6), routed_scaling_factor=cfg.get("routed_scaling_factor", 2.5), mode="hash" if is_hash else "dense", vocab_size=cfg.get("vocab_size", 128000) if is_hash else None, device=dev) if is_hash: router.load_weights(hash_lut=all_w[f"{mlp_pfx}.gate.tid2eid"].to(dev, torch.int32)) else: eb = all_w.get(f"{mlp_pfx}.gate.e_score_correction_bias") gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, mlp_pfx, 'gate') E = cfg["n_routed_experts"] if gate_w is not None and gate_ws is not None: gate_lin = Nvfp4Linear(in_features=H, out_features=E, device=dev) gate_lin.fp4 = [gate_w.to(dev).view(torch.float4_e2m1fn_x2) if gate_w.dtype == torch.uint8 else gate_w.to(dev)] gate_lin.sf = [gate_ws.to(dev)] ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0 isc_v = gate_isc.float().item() if gate_isc is not None else 1.0/(6.0*448.0) gate_lin.gs = [1.0] gate_lin.ws2 = [torch.tensor([ws2_v], device=dev, dtype=torch.float32)] gate_lin._activation_global_scale = isc_v gate_lin._use_runtime_gsa = True gate_lin.finalize_weights() router.load_nvfp4_gate(gate_lin) router.load_weights(e_bias=eb.to(dev, torch.float32)) else: # BF16 gate weight — quantize to NVFP4 gw = all_w.get(f"{mlp_pfx}.gate.weight") if gw is not None: g_bf16 = gw if gw.shape == (E, H) else gw.T.contiguous() g_bf16 = g_bf16.bfloat16().to(dev) g_fp4, g_sf, g_gs = quantize_to_nvfp4(g_bf16) gate_lin = Nvfp4Linear(in_features=H, out_features=E, device=dev) gate_lin.fp4 = [g_fp4] gate_lin.sf = [g_sf] gate_lin.gs = [g_gs] gate_lin.ws2 = [torch.tensor([g_gs], device=dev, dtype=torch.float32)] gate_lin._activation_global_scale = 1.0 / (6.0 * 448.0) gate_lin._use_runtime_gsa = True gate_lin.finalize_weights() router.load_nvfp4_gate(gate_lin) router.load_weights(e_bias=eb.to(dev, torch.float32)) router.finalize_weights(); routers[li] = router moe = Nvfp4MoE(num_experts=cfg["n_routed_experts"], hidden_size=H, intermediate_size=cfg.get("moe_intermediate_size", 3072), top_k=cfg.get("num_experts_per_tok", 6), device=dev) moe.set_swiglu_limit(cfg.get("swiglu_limit", 10.0)); moe.set_fused_swiglu(True) _load_moe_weights_stacked(all_w, li, mlp_pfx, dev, moe, cfg) moe._ensure_stacked(); moe._use_runtime_gsa = True; moe_runners[li] = moe se = Nvfp4SharedExpert(hidden_size=H, intermediate_size=cfg.get("moe_intermediate_size", 3072), device=dev, swiglu_limit=cfg.get("swiglu_limit", 10.0)) se.set_fused_swiglu(True) _load_shared_expert_weights(all_w, li, mlp_pfx, dev, se, cfg) se._ensure_initialized(); se._use_runtime_gsa = True; se_runners[li] = se torch.cuda.empty_cache() for li in range(TEST_LAYERS): pfx = f"model.layers.{li}.self_attn.compressor" dev = f"cuda:{li % NUM_GPUS}" if li in compressors: compressors[li].load(all_w, pfx, dev=dev) if li in indexers: indexers[li].load(all_w, f"{pfx}.indexer", dev=dev) print("Components built") # Embedding + tokenizer from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR) bos = tokenizer.bos_token_id or 0 USER_TOKEN, ASSISTANT_TOKEN, THINK_START = 128803, 128804, 128821 input_ids = [bos, USER_TOKEN] input_ids += tokenizer.encode('\n\nThe capital of France is', add_special_tokens=False) input_ids.append(ASSISTANT_TOKEN); input_ids.append(THINK_START) print(f"Input: {len(input_ids)} tokens") torch.cuda.set_device(0) embed_w = all_w.get("model.embed_tokens.weight") embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to(DEVICE)) devs_list = [f"cuda:{g}" for g in range(NUM_GPUS)] layer_w = _cache_layer_weights_no_experts(all_w, TEST_LAYERS, devs_list) del all_w; import gc; gc.collect() for g in range(NUM_GPUS): torch.cuda.set_device(g); torch.cuda.empty_cache() torch.cuda.set_device(0) # ================================================================ # PHASE 1: Run full production pipeline to populate KV caches # ================================================================ print(f"\nPhase 1: Populating KV caches...") for pi, tid_val in enumerate(input_ids): t1 = time.time() tid = torch.tensor([tid_val], dtype=torch.long, device=DEVICE) pos = torch.tensor([pi], dtype=torch.long, device=DEVICE) tid32 = torch.tensor([tid_val], dtype=torch.int32, device=DEVICE) X = mHCLayer.init_state(embed(tid)) for li in range(TEST_LAYERS): gpu = li % NUM_GPUS if X.device != torch.device(f"cuda:{gpu}"): X = X.to(f"cuda:{gpu}") torch.cuda.set_device(gpu) X = forward_layer(X, layer_w[li], li, cfg, *rope_caches[gpu], attn_mhcs.get(li), ffn_mhcs.get(li), attn_norms.get(li), ffn_norms.get(li), kv_caches[li], pos, tid32, compressors.get(li), indexers.get(li), moe_runners.get(li), se_runners.get(li), routers.get(li), prod_lin=prod_lins.get(li), _use_fused_rmsnorm_quantize=True) if pi % 5 == 0: print(f" Token {pi}/{len(input_ids)}: {time.time()-t1:.2f}s", flush=True) # ================================================================ # PHASE 2: For each layer, gather KV, run production FMHA, compare vs SDPA # ================================================================ print(f"\nPhase 2: FMHA comparison per layer...") results = {} for li in range(TEST_LAYERS): gpu = li % NUM_GPUS dev = f"cuda:{gpu}" torch.cuda.set_device(gpu) ratio = cr[li] if li < len(cr) else 128 k_cache = kv_caches[li] # Gather KV in mixed format (same as production path) if k_cache.n_comp > 0: if ratio > 4: kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = k_cache.gather_mixed_all() else: kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = k_cache.gather_mixed_swa_only() else: kv_nope_fp8, kv_nope_scale, kv_rope_bf16 = k_cache.gather_mixed_swa_only() seq_len = kv_nope_scale.shape[0] if seq_len == 0: print(f" L{li}: SKIPPED (seq_len=0)") continue # Generate a test Q (random, on this GPU) q_bf16 = torch.randn(1, n_h, 1, hd, dtype=torch.bfloat16, device=dev) * 0.5 # 1. Run production mixed FP8 FMHA from dsv4.kernels.attention.fmha_mixed_fp8_op import fmha_mixed_fp8_decode_raw scale_val = 1.0 / math.sqrt(hd) try: o_prod, lse_prod = fmha_mixed_fp8_decode_raw( q_bf16, kv_nope_fp8, kv_nope_scale, kv_rope_bf16, scale_val, rope_dim=rd) except Exception as e: print(f" L{li}: PROD FMHA FAILED: {e}") results[li] = {'cos': -1.0, 'error': str(e)} continue # 2. Reference: dequantize KV, run SDPA nope_dequant = kv_nope_fp8.view(torch.float8_e4m3fn).float() * kv_nope_scale.unsqueeze(-1).float() kv_full = torch.cat([nope_dequant.bfloat16(), kv_rope_bf16], dim=-1) # (N, hd) k_4d = kv_full.unsqueeze(0).unsqueeze(0).expand(1, 1, -1, -1) # (1, 1, N, hd) v_4d = k_4d.clone() o_ref = F.scaled_dot_product_attention(q_bf16, k_4d, v_4d, scale=scale_val) # (1, H, 1, hd) # 3. Compare cos_val = cosine(o_prod, o_ref) mag_prod = o_prod.float().abs().max().item() mag_ref = o_ref.float().abs().max().item() # Per-head cosine o_prod_h = o_prod.float().squeeze(2) # (1, H, hd) → (H, hd) after squeeze o_ref_h = o_ref.float().squeeze(2) if o_prod_h.dim() == 3: o_prod_h = o_prod_h.squeeze(0) if o_ref_h.dim() == 3: o_ref_h = o_ref_h.squeeze(0) per_head_cos = F.cosine_similarity(o_prod_h, o_ref_h, dim=-1) min_head = per_head_cos.min().item() mean_head = per_head_cos.mean().item() results[li] = { 'cos': cos_val, 'mag_prod': mag_prod, 'mag_ref': mag_ref, 'seq_len': seq_len, 'ratio': ratio, 'min_head_cos': min_head, 'mean_head_cos': mean_head, } status = "PASS" if cos_val >= 0.999 else "FAIL" print(f" L{li}: {status} cos={cos_val:.6f} min_head={min_head:.6f} mean_head={mean_head:.6f} " f"|prod|={mag_prod:.4f} |ref|={mag_ref:.4f} seq_len={seq_len} ratio={ratio}") if cos_val < 0.999: worst = per_head_cos.argsort()[:5] print(f" Worst heads: {worst.tolist()} cos={[f'{c:.4f}' for c in per_head_cos[worst].tolist()]}") # Summary print("\n" + "=" * 70) print("SUMMARY") print("=" * 70) all_pass = True for li in sorted(results.keys()): r = results[li] c = r.get('cos', -1.0) status = "PASS" if c >= 0.999 else "FAIL" if c < 0.999: all_pass = False print(f" L{li}: {status} cos={c:.6f} seq={r.get('seq_len','?')} ratio={r.get('ratio','?')}") print() if all_pass: print("ALL PASSED (cos >= 0.999)") else: print("SOME FAILED — see per-layer output above") return 0 if all_pass else 1 if __name__ == "__main__": sys.exit(main())