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 6, 19:15 UTC
Enable RoPE+KV cache fusion for ROCm AITER FA (non-shuffle layout) (#35786)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Mar 13, 2026 at 02:33 UTC
a4ad9db54169694baae152d6a86dd4050263148f
2 files changed
+47
-1
tests/compile/passes/test_rope_kvcache_fusion.py
+1
@@ -196,6 +196,7 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
196
AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
197
AttentionBackendEnum.TRITON_ATTN,
198
AttentionBackendEnum.ROCM_ATTN,
199
+ AttentionBackendEnum.ROCM_AITER_FA,
200
],
201
)
202
@pytest.mark.parametrize("enable_rope_custom_op", [True]) # [True, False])
vllm/v1/attention/backends/rocm_aiter_fa.py
+46
-1
@@ -20,6 +20,7 @@ from vllm.v1.attention.backend import (
20
AttentionBackend,
21
AttentionCGSupport,
22
AttentionImpl,
23
+ AttentionLayer,
24
AttentionMetadataBuilder,
25
AttentionType,
26
CommonAttentionMetadata,
@@ -1308,7 +1309,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1309
1310
def do_kv_cache_update(
1311
self,
1311
- layer: Attention,
1312
+ layer: AttentionLayer,
1313
key: torch.Tensor,
1314
value: torch.Tensor,
1315
kv_cache: torch.Tensor,
@@ -1359,3 +1360,47 @@ class AiterFlashAttentionImpl(AttentionImpl):
1360
layer._k_scale,
1361
layer._v_scale,
1362
)
1363
+
1364
+ def fused_rope_kvcache_supported(self):
1365
+ # Only support fusion when shuffle KV cache layout is not used;
1366
+ # shuffle layout uses a different cache update path.
1367
+ return (
1368
+ rocm_aiter_ops.is_enabled()
1369
+ and not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
1370
+ )
1371
+
1372
+ def do_rope_and_kv_cache_update(
1373
+ self,
1374
+ layer: AttentionLayer,
1375
+ query: torch.Tensor,
1376
+ key: torch.Tensor,
1377
+ value: torch.Tensor,
1378
+ positions: torch.Tensor,
1379
+ cos_sin_cache: torch.Tensor,
1380
+ is_neox: bool,
1381
+ kv_cache: torch.Tensor,
1382
+ layer_slot_mapping: torch.Tensor,
1383
+ ):
1384
+ key_cache, value_cache = kv_cache.unbind(0)
1385
+ flash_layout = True
1386
+
1387
+ is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
1388
+ if is_fp8_kv_cache:
1389
+ key_cache = key_cache.view(current_platform.fp8_dtype())
1390
+ value_cache = value_cache.view(current_platform.fp8_dtype())
1391
+
1392
+ rocm_aiter_ops.triton_rope_and_cache(
1393
+ query,
1394
+ key,
1395
+ value,
1396
+ positions,
1397
+ cos_sin_cache,
1398
+ is_neox,
1399
+ key_cache,
1400
+ value_cache,
1401
+ layer_slot_mapping,
1402
+ layer._k_scale,
1403
+ layer._v_scale,
1404
+ flash_layout,
1405
+ is_fp8_kv_cache,
1406
+ )