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 14, 05:30 UTC
fp8 kv cache support fix for torch.compile (#22758)
Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com>
Aleksandr Malyshev committed
Sep 16, 2025 at 14:27 UTC
3053a22b330cd7170dce6f33f3a2043c64a99599
2 files changed
+4
-2
vllm/model_executor/layers/quantization/kv_cache.py
+3
-1
@@ -125,7 +125,9 @@ class BaseKVCacheMethod(QuantizeMethodBase):
125
126
# These are used in the final Attention.forward()
127
layer._q_scale.copy_(q_scale)
128
- layer._q_scale_float = q_scale
128
+ layer._q_scale_float = q_scale.item() if isinstance(
129
+ q_scale, torch.Tensor) else q_scale
130
+
131
layer._prob_scale.copy_(prob_scale)
132
if layer.kv_cache_dtype == "fp8" and (q_scale == 1.0
133
or prob_scale == 1.0):
vllm/v1/attention/backends/triton_attn.py
+1
-1
@@ -361,7 +361,7 @@ class TritonAttentionImpl(AttentionImpl):
361
key_cache = key_cache.view(self.fp8_dtype)
362
value_cache = value_cache.view(self.fp8_dtype)
363
num_tokens, num_heads, head_size = query.shape
364
- assert layer._q_scale == 1.0, \
364
+ assert layer._q_scale_float == 1.0, \
365
"A non 1.0 q_scale is not currently supported."
366
if current_platform.is_cuda():
367
# Skip Q quantization on ROCm and XPU, enable this on cuda