diff --git a/single_shot_inference.py b/single_shot_inference.py index 339dca6d..a4cf7966 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -350,7 +350,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, # 1. Q: q_a (NVFP4 GEMM) → q_a_norm → q_b (NVFP4 GEMM) → q_b_norm q_a = prod_lin['q_a'](x_normed) - if li < 3: + if VERBOSE >= 2 and li < 3: # Compare q_a with PyTorch reference q_a_ref = do_nvfp4_linear_ref(x_normed, w, pfx, 'q_a_proj') if q_a_ref is not None: @@ -399,7 +399,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, # 6. Production FMHA attn_out = _run_production_fmha(q_heads, all_kv, n_h, hd, T, seq_len, scale, dev, li, w, pfx) - if li < 3: + if VERBOSE >= 2 and li < 3: # Compare with PyTorch reference k_exp = all_kv.unsqueeze(0).expand(n_h, -1, -1).contiguous() v_exp = k_exp.clone() @@ -419,7 +419,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, a_grp = a_flat.reshape(T, o_groups, gid); oa_3d = oa_bf.reshape(o_groups, o_rank, gid) g_out = torch.bmm(a_grp.permute(1, 0, 2), oa_3d.transpose(1, 2)) g_flat = g_out.permute(1, 0, 2).reshape(T, o_groups * o_rank) - if li < 3: + if VERBOSE >= 2 and li < 3: print(f" L{li} wo_a: |g_flat|={g_flat.abs().max().item():.6f} shape={g_flat.shape}", flush=True) F_attn = prod_lin['o_b'](g_flat) else: @@ -430,7 +430,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, else: log.warning(f"L{li}: No o_a_proj weight, zero attention output") F_attn = torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev) - if li < 3: + if VERBOSE >= 2 and li < 3: print(f" L{li} F_attn: |F_attn|={F_attn.abs().max().item():.6f}", flush=True) return F_attn, q_a @@ -444,13 +444,13 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id): torch.cuda.synchronize(x.device) if topk_ids.max().item() >= 384 or topk_ids.min().item() < 0: print(f" L{li} BAD topk_ids: min={topk_ids.min().item()} max={topk_ids.max().item()}", flush=True) - if li < 3: + if VERBOSE >= 2 and li < 3: print(f" L{li} MoE input: |x|={x.abs().max().item():.4f} has_nan={torch.isnan(x).any().item()}", flush=True) routed_out = moe_runner.run(x, topk_w, topk_ids) - if li < 3: + if VERBOSE >= 2 and li < 3: print(f" L{li} MoE routed: |out|={routed_out.abs().max().item():.4f} has_nan={torch.isnan(routed_out).any().item()}", flush=True) shared_out = se_runner.run(x) - if li < 3: + if VERBOSE >= 2 and li < 3: has_nan = torch.isnan(shared_out).any().item() out_max = shared_out.abs().max().item() if not has_nan else float('nan') print(f" L{li} MoE shared: |out|={out_max:.4f} has_nan={has_nan}", flush=True) @@ -794,7 +794,7 @@ def main(): err = torch.cuda.current_stream(gpu).query() print(f" CRASH at token {pi} layer {li} gpu {gpu}: {e}", flush=True) raise - if pi == 0 and li < 3: + if VERBOSE >= 2 and pi == 0 and li < 3: torch.cuda.synchronize(gpu) print(f" Token {pi} L{li}: OK |X|={X.abs().max().item():.1f}", flush=True) X = X.to('cuda:0'); torch.cuda.set_device(0) @@ -829,7 +829,7 @@ def main(): next_id = torch.argmax(logits, -1).item(); all_tokens.append(next_id) dt = time.time() - t1 has_nan = torch.isnan(logits.float()).any().item() - if step % 5 == 0 or has_nan: + if step % 1 == 0 or has_nan: tv, ti = torch.topk(logits[0], 5) top5 = ' '.join(f'{tokenizer.decode([t.item()])}({v.item():.1f})' for t, v in zip(ti[:5], tv[:5])) print(f" Step {step}: {next_id} '{tokenizer.decode([next_id])}' ({dt:.2f}s) "