[ROCm][Bugfix] Ignore descales for unquantized AITER caches (#56726)
Signed-off-by: donglu <donglu@cohere.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
dongluw committed
Sep 16, 2026 at 21:29 UTC
19b6ff62f24d477a2760602b9e94eb6e4dd7dcdd
2 files changed
+85
-22
tests/v1/attention/test_attention_backends.py
+59
-7
@@ -243,14 +243,19 @@ def create_and_prepopulate_kv_cache(
243
class MockAttentionLayer:
244
"""A mock attention layer for testing."""
245
246
- def __init__(self, device: torch.device):
246
+ def __init__(
247
+ self,
248
+ device: torch.device,
249
+ k_scale: float = 1.0,
250
+ v_scale: float = 1.0,
251
+ ):
252
self._q_scale = torch.tensor(1.0, device=device)
248
- self._k_scale = torch.tensor(1.0, device=device)
249
- self._v_scale = torch.tensor(1.0, device=device)
253
+ self._k_scale = torch.tensor(k_scale, device=device)
254
+ self._v_scale = torch.tensor(v_scale, device=device)
255
# Add float versions for flashinfer
256
self._q_scale_float = 1.0
252
- self._k_scale_float = 1.0
253
- self._v_scale_float = 1.0
257
+ self._k_scale_float = k_scale
258
+ self._v_scale_float = v_scale
259
260
261
def _clone_kv_cache_in_layout(
@@ -285,6 +290,8 @@ def run_attention_backend(
290
kv_cache_dtype: str = "auto",
291
sinks: torch.Tensor | None = None,
292
use_cuda_graph: bool = False,
293
+ layer_k_scale: float = 1.0,
294
+ layer_v_scale: float = 1.0,
295
) -> torch.Tensor:
296
"""Run attention computation using the specified backend's AttentionImpl."""
297
use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
@@ -356,7 +363,7 @@ def run_attention_backend(
363
)
364
365
# Create mock layer and output buffer
359
- mock_layer = MockAttentionLayer(device)
366
+ mock_layer = MockAttentionLayer(device, layer_k_scale, layer_v_scale)
367
output = torch.empty_like(query)
368
369
if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
@@ -386,9 +393,11 @@ def run_attention_backend(
393
)
394
graph.replay()
395
else:
389
- output = impl.forward(
396
+ backend_output = impl.forward(
397
mock_layer, query, key, value, kv_cache, attn_metadata, output=output
398
)
399
+ if backend_output is not None:
400
+ output = backend_output
401
402
return output
403
@@ -413,6 +422,8 @@ def _test_backend_correctness(
422
model_dtype: torch.dtype | None = None,
423
max_num_seqs: int | None = None,
424
max_num_batched_tokens: int | None = None,
425
+ layer_k_scale: float = 1.0,
426
+ layer_v_scale: float = 1.0,
427
):
428
"""Test that all backends produce similar outputs to a reference implementation
429
using FlexAttention or an explicit attention-sink reference.
@@ -687,6 +698,8 @@ def _test_backend_correctness(
698
kv_cache_dtype=kv_cache_dtype,
699
sinks=sinks,
700
use_cuda_graph=use_cuda_graph,
701
+ layer_k_scale=layer_k_scale,
702
+ layer_v_scale=layer_v_scale,
703
)
704
705
# Check shape and dtype consistency
@@ -1317,6 +1330,45 @@ def test_sliding_window_backend_correctness(
1330
)
1331
1332
1333
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific test")
1334
+def test_rocm_aiter_fa_unquantized_cache_ignores_kv_scales(
1335
+ default_vllm_config,
1336
+):
1337
+ """Stale scale tensors must not affect an unquantized KV cache."""
1338
+ from vllm._aiter_ops import is_aiter_found_and_supported
1339
+
1340
+ if not is_aiter_found_and_supported():
1341
+ pytest.skip("AITER is required")
1342
+
1343
+ def sliding_window_mask_mod(
1344
+ b: torch.Tensor,
1345
+ h: torch.Tensor,
1346
+ q_idx: torch.Tensor,
1347
+ kv_idx: torch.Tensor,
1348
+ *,
1349
+ context_len: int,
1350
+ sliding_window: int,
1351
+ ):
1352
+ causal_mask = q_idx + context_len >= kv_idx
1353
+ window_mask = q_idx + context_len - kv_idx < sliding_window
1354
+ return causal_mask & window_mask
1355
+
1356
+ model = "microsoft/Phi-tiny-MoE-instruct"
1357
+ batch_spec = BATCH_SPECS["small_decode"]
1358
+ model_config = ModelConfig(model=model, max_model_len=max(batch_spec.seq_lens))
1359
+ sliding_window = model_config.get_sliding_window()
1360
+ assert sliding_window is not None
1361
+
1362
+ _test_backend_correctness(
1363
+ batch_spec,
1364
+ model,
1365
+ [AttentionBackendEnum.ROCM_AITER_FA],
1366
+ partial(sliding_window_mask_mod, sliding_window=sliding_window),
1367
+ layer_k_scale=0.25,
1368
+ layer_v_scale=0.5,
1369
+ )
1370
+
1371
+
1372
@pytest.mark.parametrize(
1373
"batch_spec_name",
1374
[
vllm/v1/attention/backends/rocm_aiter_fa.py
+26
-15
@@ -989,6 +989,20 @@ class AiterFlashAttentionImpl(AttentionImpl):
989
"cross-attention."
990
)
991
992
+ def _get_kv_cache_descales(
993
+ self,
994
+ layer: AttentionLayer,
995
+ num_decodes: int,
996
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
997
+ # AITER treats non-null descales as an instruction to dequantize.
998
+ if not is_quantized_kv_cache(self.kv_cache_dtype):
999
+ return None, None
1000
+ descale_shape = (num_decodes, self.num_kv_heads)
1001
+ return (
1002
+ layer._k_scale.expand(descale_shape),
1003
+ layer._v_scale.expand(descale_shape),
1004
+ )
1005
+
1006
def extend_for_sliding_window(
1007
self,
1008
attn_metadata: AiterFlashAttentionMetadata,
@@ -1367,6 +1381,9 @@ class AiterFlashAttentionImpl(AttentionImpl):
1381
or decode_max_query_len > 1
1382
or self.sinks is not None
1383
):
1384
+ k_descale, v_descale = self._get_kv_cache_descales(
1385
+ layer, num_decodes
1386
+ )
1387
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1388
"Shuffle KV cache layout is not supported with sliding "
1389
"window, sinks, or speculative decoding (multi-token decode)."
@@ -1376,7 +1393,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
1393
flash_attn_with_kvcache,
1394
)
1395
1379
- descale_shape = (num_decodes, key_cache.shape[2])
1396
decode_query = query[:num_decode_tokens].reshape(
1397
num_decodes,
1398
decode_max_query_len,
@@ -1393,8 +1409,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1409
window_size=self.sliding_window,
1410
softcap=self.logits_soft_cap,
1411
q_descale=None,
1396
- k_descale=layer._k_scale.expand(descale_shape),
1397
- v_descale=layer._v_scale.expand(descale_shape),
1412
+ k_descale=k_descale,
1413
+ v_descale=v_descale,
1414
page_table=attn_metadata.block_table[:num_decodes],
1415
)
1416
output[:num_decode_tokens].copy_(
@@ -1412,10 +1428,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
1428
unified_attention,
1429
)
1430
1415
- descale_shape = (
1416
- num_decodes,
1417
- key_cache.shape[2],
1418
- )
1431
unified_attention(
1432
q=query[:num_decode_tokens],
1433
k=key_cache,
@@ -1434,8 +1446,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1446
block_table=attn_metadata.block_table[:num_decodes],
1447
softcap=self.logits_soft_cap,
1448
q_descale=None,
1437
- k_descale=layer._k_scale.expand(descale_shape),
1438
- v_descale=layer._v_scale.expand(descale_shape),
1449
+ k_descale=k_descale,
1450
+ v_descale=v_descale,
1451
sinks=self.sinks,
1452
)
1453
return
@@ -1449,6 +1461,9 @@ class AiterFlashAttentionImpl(AttentionImpl):
1461
use_unified_attention = self.head_size < _MIN_HEAD_SIZE_FOR_LL4MI
1462
1463
if use_unified_attention:
1464
+ k_descale, v_descale = self._get_kv_cache_descales(
1465
+ layer, num_decodes
1466
+ )
1467
assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (
1468
"unified_attention fallback with shuffle layout "
1469
"is not supported yet."
@@ -1460,10 +1475,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
1475
decode_cu_seqlens_q = attn_metadata.query_start_loc[
1476
: num_decodes + 1
1477
]
1463
- descale_shape = (
1464
- num_decodes,
1465
- key_cache.shape[2],
1466
- )
1478
unified_attention(
1479
q=query[:num_decode_tokens],
1480
k=key_cache,
@@ -1480,8 +1491,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1491
block_table=attn_metadata.block_table[:num_decodes],
1492
softcap=self.logits_soft_cap,
1493
q_descale=None,
1483
- k_descale=layer._k_scale.expand(descale_shape),
1484
- v_descale=layer._v_scale.expand(descale_shape),
1494
+ k_descale=k_descale,
1495
+ v_descale=v_descale,
1496
)
1497
elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1498
_, num_heads, head_size = query.shape