[ROCm][CI] Only require q_scale==1.0 for fp8 query in RocmAttention (#46148)
Signed-off-by: stefankoncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Jun 23, 2026 at 01:25 UTC
91ba720b75f0d74e2af83765f80078f2a551e881
1 file changed
+12
-3
vllm/v1/attention/backends/rocm_attn.py
+12
-3
@@ -421,9 +421,18 @@ class RocmAttentionImpl(AttentionImpl):
421
if is_quantized_kv_cache(self.kv_cache_dtype):
422
key_cache = key_cache.view(self.fp8_dtype)
423
value_cache = value_cache.view(self.fp8_dtype)
424
- assert layer._q_scale_float == 1.0, (
425
- "A non 1.0 q_scale is not currently supported."
426
- )
424
+ # chunked_prefill_paged_decode runs attention with a full-precision
425
+ # query (it does not quantize Q to fp8 and does not consume
426
+ # q_scale), so q_scale only matters when the query itself is fp8.
427
+ # For a non-fp8 query, q_scale is not applicable and is ignored
428
+ # (mirrors TritonAttentionImpl). This avoids spuriously failing on
429
+ # checkpoints that carry a non-1.0 q_scale while keeping the query
430
+ # in full precision.
431
+ if query.dtype == self.fp8_dtype and layer._q_scale_float != 1.0:
432
+ raise NotImplementedError(
433
+ "A non 1.0 q_scale with an fp8 query is not currently "
434
+ "supported by RocmAttentionImpl."
435
+ )
436
437
cu_seqlens_q = attn_metadata.query_start_loc
438
seqused_k = attn_metadata.seq_lens