Fix NaN from stale FP4 scale padding: torch.empty → torch.zeros

Padding rows in the swizzled scale tensor were uninitialized (torch.empty),
containing stale NaN from prior GPU allocations. The TRT-LLM mm_fp4 kernel
with use_8x4_sf_layout=True reads padding scales and applies them to real
rows, contaminating output with NaN.

Zero-filling ensures padding scales contribute 0 * data = 0.

Fixes: https://github.com/flashinfer-ai/flashinfer/issues/2861

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
This commit is contained in:
Elvir Crncevic
2026-03-28 13:21:00 -04:00
committed by Tyler Michael Smith
co-authored by Claude Opus 4.6
parent eadf848ea0
commit d2afd40d6f
+2 -2
View File
@@ -56,11 +56,11 @@ def create_fp4_scale_tensor(
rounded_m = round_up(m, 128)
scale_n = n // block_size
rounded_n = round_up(scale_n, 4)
return torch.empty(
return torch.zeros(
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
)
else:
return torch.empty((m, n // block_size), device=device, dtype=torch.uint8)
return torch.zeros((m, n // block_size), device=device, dtype=torch.uint8)
def create_fp4_output_tensors(