Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions tensorrt_llm/_torch/attention_backend/trtllm.py
Original file line number Diff line number Diff line change
Expand Up @@ -300,9 +300,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)
Expand Down
Loading