A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[ROCm][Bugfix] Use platform specific FP8 dtype (#15717)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Apr 4, 2025 at 12:40 UTC
40a36ccfeb49df1351cb1d9250ebf29f6f6805ae
1 file changed
+1
-1
vllm/attention/ops/prefix_prefill.py
+1
-1
@@ -753,7 +753,7 @@ if triton.__version__ >= "2.1.0":
753
assert (v_cache.dtype == torch.uint8)
754
755
if kv_cache_dtype in ("fp8", "fp8_e4m3"):
756
- target_dtype = torch.float8_e4m3fn
756
+ target_dtype = current_platform.fp8_dtype()
757
elif kv_cache_dtype == "fp8_e5m2":
758
target_dtype = torch.float8_e5m2
759
else: