kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-04 02:02:01 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] convert to CUDA tensor before calling _resmooth_kernel. (#10770)

Browse Source 
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
This commit is contained in:
Yuxian Qiu 2026-01-17 16:18:34 +08:00 committed by GitHub
parent b65560fc32
commit cef67b4f8d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorrt_llm/quantization/utils/fp8_utils.py
View File

@ -103,6 +103,9 @@ def resmooth_to_fp8_e8m0(
assert weight.dtype == torch.float8_e4m3fn
assert weight_scale.dtype == torch.float32
weight = weight.cuda()
weight_scale = weight_scale.cuda()
orig_shape = weight.shape
M, K = orig_shape[-2:]
w_view = weight.view(-1, M, K)