From 6731dda7db53a4da4d813f352b16c03c73fbe2ff Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 20 Nov 2025 09:13:03 +0000 Subject: [PATCH] [https://nvbugs/5676748][fix] Fix mismatched nvfp4 gemm sf shape. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- tensorrt_llm/_torch/attention_backend/trtllm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 423fb7beebb6..464396ab785f 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -329,9 +329,9 @@ def create_output(self, q: torch.Tensor, out_dtype: torch.dtype): size_per_token // num_nvfp4_elements_per_container), dtype=torch.uint8) # Create a sf (scaling factors) tensor for NVFP4 (use INT8 as the container dtype). - output_sf = q.new_empty(compute_swizzled_sf_shape( - num_tokens, size_per_token // scaling_vector_size), - dtype=torch.uint8) + padded_row, padded_col = compute_swizzled_sf_shape( + num_tokens, size_per_token // scaling_vector_size) + output_sf = q.new_empty(padded_row * padded_col, dtype=torch.uint8) else: output = q.new_empty((num_tokens, self.num_heads * v_head_size), dtype=out_dtype)