From 2a6f9a10b11bacba94e801be9b97453d2d6a88a1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 1 Jun 2026 22:05:22 +0000
Subject: [PATCH] lm_head: fall back to BF16 F.linear for stability
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

NVFP4 quantize_from_buffer produces CUDA error on large-magnitude
inputs (|X|>500 at L60 output). BF16 lm_head is correct and only
runs once per decode step — not a bottleneck.

TODO: debug the NVFP4 path for large activations and re-enable.
---
 single_shot_inference.py | 32 +++++++++-----------------------
 1 file changed, 9 insertions(+), 23 deletions(-)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index 024bac24..f09f4732 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -829,29 +829,12 @@ def main():
     torch.cuda.set_device(0)
     embed_w = all_w.get("model.embed_tokens.weight")
     embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16().to('cuda:0'))
-    # lm_head: quantize to NVFP4 for tensor-core acceleration
-    # Weight is (vocab_size, hidden_size) = (N, K) in BF16
-    # quantize_weight_to_nvfp4 expects (K, N), so transpose first
-    # But Nvfp4Linear expects (N_packed, K_packed) from checkpoint layout
-    # quantize_weight_to_nvfp4 returns (K//2, N) which IS (K_packed, N)
-    # So we need to transpose the weight, quantize as (K, N),
-    # then the result (K//2, N) needs to be transposed to (N, K//2) for Nvfp4Linear.
+    # lm_head: BF16 for now — NVFP4 path has CUDA error on large-magnitude inputs
+    # TODO: debug quantize_from_buffer for |X|>500 and re-enable NVFP4
     lm_w_raw = all_w.get("lm_head.weight", embed_w).bfloat16().to('cuda:0')
-    from dsv4.layers.linear import Nvfp4Linear
-    lm_head_lin = Nvfp4Linear(lm_w_raw.shape[1], lm_w_raw.shape[0], max_num_tokens=8192, device='cuda:0')
-    from dsv4.ops.quantize import quantize_weight_to_nvfp4
-    # quantize_weight_to_nvfp4 takes (K, N) → returns (K//2, N), (K//16, N), gs
-    lm_fp4, lm_sf, lm_gs = quantize_weight_to_nvfp4(lm_w_raw.T.contiguous())  # (K//2, N) = (3584, 128K)
-    # Nvfp4Linear expects fp4 in (N_packed, K_packed) layout, so transpose
-    lm_head_lin.fp4 = [lm_fp4.permute(1, 0).contiguous()]  # (N, K_packed) = (128K, 3584)
-    lm_head_lin.sf = [lm_sf.permute(1, 0).contiguous()]  # (N, K_sf) = (128K, 448)
-    lm_head_lin.gs = [lm_gs]  # global scale from weight quantization
-    lm_head_lin.ws2 = [None]  # no separate weight_scale_2
-    lm_head_lin._activation_global_scale = 1.0 / (6.0 * 448.0)  # placeholder
-    lm_head_lin._use_runtime_gsa = True
-    lm_head_lin.finalize_weights()
-    lm_w = None  # free BF16 weight
-    print("  lm_head: NVFP4 production GEMM")
+    lm_head_w = lm_w_raw  # keep as BF16 for F.linear
+    lm_head_lin = None  # signal: use BF16 path
+    print("  lm_head: BF16 F.linear (NVFP4 deferred)")
     final_norm_w = all_w.get("model.norm.weight")
     if final_norm_w is not None: final_norm_w = final_norm_w.to('cuda:0', torch.float32)
 
@@ -987,7 +970,10 @@ def main():
         X = X.to('cuda:0'); torch.cuda.set_device(0)
         x_out = hc_head.forward(X) if hc_head is not None else X[:, 0, :]
         if final_norm_w is not None: x_out = rmsnorm(x_out, final_norm_w)
-        logits = lm_head_lin(x_out)
+        if lm_head_lin is not None:
+            logits = lm_head_lin(x_out)
+        else:
+            logits = torch.nn.functional.linear(x_out, lm_head_w)
         # Validate logits before sampling
         if step == 0 or torch.isnan(logits.float()).any().item():
             print(f"  logits: shape={list(logits.shape)} dtype={logits.dtype} "