From 24325208d27c8686398bdbd059891235c0fbfd36 Mon Sep 17 00:00:00 2001 From: Elvir Crncevic Date: Wed, 25 Mar 2026 22:24:13 +0100 Subject: [PATCH] =?UTF-8?q?Fix=20NaN=20from=20stale=20FP4=20scale=20paddin?= =?UTF-8?q?g:=20torch.empty=20=E2=86=92=20torch.zeros?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Signed-off-by: Elvir Crncevic --- vllm/_custom_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 6c9ca07dba9a..da26075fe18f 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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(