fix: clamp block_amax to E4M3 max (448) in quantize_activation_nvfp4 — prevents NaN from overflow

2026-06-01 04:59:06 +00:00
parent eb9c46f8cb
commit 6e53e3007c
1 changed files with 1 additions and 1 deletions
--- a/dsv4/ops/quantize.py
+++ b/dsv4/ops/quantize.py
@@ -145,7 +145,7 @@ def quantize_activation_nvfp4(x_bf16, global_scale, block_size=SF_VEC_SIZE):
    zero_block = block_amax < (6.0 * 2.0 ** -9)
    x_reshaped = torch.where(zero_block.unsqueeze(-1),
                              torch.zeros_like(x_reshaped), x_reshaped)
-    block_amax = block_amax.clamp(min=1e-8)
+    block_amax = block_amax.clamp(min=1e-8, max=6.0 * 448.0)  # E4M3 max = 448
    block_scale = (block_amax / 6.0).to(torch.float8_e4m3fn)
    block_scale = torch.where(zero_block, torch.zeros_like(block_scale), block_scale)