fix: restore proper quantize_weight_to_nvfp4 — K is the packed dim, not N

quantize_to_nvfp4() only packs the last dimension, but for weight
matrices (K, N), K is the packed dimension. The weight quantizer
reshapes (k_blocks, block_size, N) and computes block scales along
the K block dimension. This was accidentally replaced with a simple
delegation to quantize_to_nvfp4, producing wrong tensor shapes.

cutedsl/bridge.py

@@ -179,17 +179,50 @@ def quantize_weight_to_nvfp4(w_bf16, block_size=SF_VEC_SIZE):
     """Quantize BF16 weight matrix to NVFP4.
     
     The weight is (K, N) where K is the input dim (packed dimension).
-
+    Block scales are computed along K (dim 0).
+    
+    NOTE: NOT cudagraph-safe — uses .max() for global scale.
+    
     Args:
         w_bf16: (K, N) BF16 weight matrix
-        block_size: NVFP4 block size
     
     Returns:
-        w_fp4: (K//2, N) float4_e2m1fn_x2 — native PyTorch FP4
-        w_sf: (K//16, N) float8_e4m3fn — block scales
+        w_fp4: (K//2, N) float4_e2m1fn_x2 — K is the packed dim
+        w_sf: (K//16, N) float8_e4m3fn — block scales along K
         global_scale: float32 scalar
     """
-    return quantize_to_nvfp4(w_bf16, block_size)
+    K, N = w_bf16.shape
+    w_f32 = w_bf16.float()
+    amax = w_f32.abs().max().clamp(min=1e-8).float()
+    global_scale = amax / (6.0 * 448.0)
+    w_norm = w_f32 / global_scale
+
+    k_blocks = ceil_div(K, block_size)
+    if K % block_size != 0:
+        w_norm = torch.nn.functional.pad(w_norm, (0, 0, 0, k_blocks * block_size - K))
+
+    w_reshaped = w_norm.reshape(k_blocks, block_size, N)
+    w_block_amax = w_reshaped.abs().amax(dim=1).clamp(min=1e-8)
+    block_scale = (w_block_amax / 6.0).to(torch.float8_e4m3fn)
+
+    block_sf_expanded = block_scale.float().unsqueeze(1)
+    w_scaled = w_reshaped / block_sf_expanded.clamp(min=1e-8)
+
+    # Nearest E2M1 — memory-efficient clamp approach
+    signs = torch.sign(w_scaled)
+    abs_scaled = w_scaled.abs().clamp(max=6.0)
+    
+    half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8)
+    step_to_idx = _get_step_to_idx_lut(w_bf16.device)
+    indices = step_to_idx[half_steps.long()]
+
+    nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)
+    even = nibbles[:, ::2, :]
+    odd = nibbles[:, 1::2, :]
+    packed = (odd << 4) | even
+
+    w_fp4 = packed.reshape(K // 2, N).view(torch.float4_e2m1fn_x2)
+    return w_fp4, block_scale, global_scale
 
 
 # ── Tensor Layout Conversion ───────────────────────────────────────────