@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a94a699c3f
Fork 0

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, 19:45 UTC

[ROCm][FP8] Fix for adjustments needed only for fnuz (#14689)

Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>

Gregory Shtrasberg committed Mar 12, 2025 at 23:14 UTC a94a699c3ff9bfc23a35f147150a826b753bbf6a
1 file changed +2 -2
vllm/model_executor/layers/quantization/kv_cache.py
+2 -2
@@ -50,7 +50,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):
50 # We prefer to use separate k_scale and v_scale if present
51 k_scale = layer.k_scale.to("cpu").tolist()
52 v_scale = layer.v_scale.to("cpu").tolist()
53 - if current_platform.is_rocm():
53 + if current_platform.is_fp8_fnuz():
54 k_scale *= 2
55 v_scale *= 2
56 elif layer.k_scale < 0.0 and layer.v_scale < 0.0:
@@ -66,7 +66,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):
66 scale_to_duplicate = max(layer.k_scale, layer.v_scale)
67 k_scale = scale_to_duplicate.to("cpu").tolist()
68 v_scale = scale_to_duplicate.to("cpu").tolist()
69 - if current_platform.is_rocm():
69 + if current_platform.is_fp8_fnuz():
70 k_scale *= 2
71 v_scale *= 2
72