improve quant script with help of claude

nvfp4gguf/quant.py

@@ -518,13 +518,18 @@ def optimize_block_scales_torch(
     for delta in (-4, -3, -2, -1, 1, 2, 3, 4):
         candidate_bit_list.append(torch.clamp(base_bits.to(torch.int16) + delta, 1, 126).to(torch.uint8))
 
+    # Max-derived candidates (the natural block scale from the largest element)
+    for delta in (-1, 0, 1):
+        candidate_bit_list.append(torch.clamp(b_max.to(torch.int16) + delta, 1, 126).to(torch.uint8))
+
     # Outlier-clipped scales (2nd and 3rd max)
     for delta in (-1, 0, 1):
         candidate_bit_list.append(torch.clamp(b_sub.to(torch.int16) + delta, 1, 126).to(torch.uint8))
     candidate_bit_list.append(torch.clamp(b_3rd, 1, 126))
 
-    # Outlier dampening (85% max)
+    # Outlier dampening (85% and 75% of max — clips isolated spikes)
     candidate_bit_list.append(torch.clamp(b_damp1, 1, 126))
+    candidate_bit_list.append(torch.clamp(b_damp2, 1, 126))
 
     # Baseline MSE from original (codes, scale) so we never regress
     if bake:
@@ -591,10 +596,15 @@ def coordinate_descent_refine_torch(
     c_t = codes.to(device=target_device, dtype=torch.uint8)
     bs_t = block_scale.to(device=target_device, dtype=torch.float32)
 
+    # Initial code refinement against the post-UE4M3 runtime scale before
+    # entering the scale optimization loop.  The learned rounding chose codes
+    # against the full-precision scale, but the runtime scale is quantized to
+    # UE4M3.  Snapping codes first gives the scale optimizer a better baseline.
+    c_t = refine_codes_for_scale_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device)
+
     total_rounds = max(int(rounds), 1)
     for r in range(total_rounds):
         prev_bs = bs_t
-        prev_c = c_t
         c_t, bs_t = optimize_block_scales_torch(w_t, c_t, bs_t, per_tensor, bake, device=target_device)
         scale_diff = int((bs_t != prev_bs).sum().item())
         if r > 0 and scale_diff == 0:
@@ -639,7 +649,7 @@ def sweep_per_tensor_scale_torch(
     best_bs = bs_t
     best_codes = c_t
 
-    multipliers = [0.95, 0.97, 0.98, 0.99, 1.01, 1.02, 1.03, 1.05]
+    multipliers = [0.93, 0.95, 0.97, 0.98, 0.99, 1.01, 1.02, 1.03, 1.05, 1.07]
 
     for m in multipliers:
         cand_pt = per_tensor * m
@@ -661,6 +671,30 @@ def sweep_per_tensor_scale_torch(
             best_bs = cand_bs
             best_codes = opt_codes.reshape(rows, cols)
 
+    # Fine refinement around the best multiplier from the coarse pass
+    if best_pt != per_tensor:
+        coarse_m = best_pt / per_tensor
+        for fine_delta in (-0.005, -0.003, -0.001, 0.001, 0.003, 0.005):
+            m = coarse_m + fine_delta
+            cand_pt = per_tensor * m
+            cand_bs = bs_t / m
+            cand_bits = ue4m3_bits_torch(cand_bs)
+            cand_rt = lut[cand_bits.long()] * float(cand_pt)
+            cand_rs_3d = cand_rt.unsqueeze(-1)
+
+            safe_rs = cand_rs_3d.clamp_min(1e-30)
+            scaled = w_blocks / safe_rs
+            opt_codes = snap_to_fp4_codes_torch(scaled, fp4_levels_sorted, fp4_level_map)
+
+            recon = fp4_table[opt_codes.long()] * cand_rs_3d
+            cand_mse = float((w_blocks - recon).square().mean().item())
+
+            if cand_mse < best_total_mse:
+                best_total_mse = cand_mse
+                best_pt = cand_pt
+                best_bs = cand_bs
+                best_codes = opt_codes.reshape(rows, cols)
+
     if best_pt != per_tensor:
         orig_mse_val = float((w_blocks - orig_recon).square().mean().item())
         log.info(