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 5, 17:15 UTC
[ROCm][perf] Shuffle KV cache to use paged_attention_common (#32914)
Signed-off-by: Samu Tamminen <stammine@amd.com> Co-authored-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
Samu Tamminen committed
Apr 1, 2026 at 06:30 UTC
c49497726ba4a8a2b86d7b0fccb9ced5f53d18a8
2 files changed
+89
-4
vllm/_aiter_ops.py
+51
@@ -2070,5 +2070,56 @@ class rocm_aiter_ops:
2070
out_=out_,
2071
)
2072
2073
+ @staticmethod
2074
+ def paged_attention_common(
2075
+ Q: torch.Tensor,
2076
+ K: torch.Tensor,
2077
+ V: torch.Tensor,
2078
+ tmp_out: torch.Tensor,
2079
+ max_logits: torch.Tensor,
2080
+ exp_sums: torch.Tensor,
2081
+ max_seq_len: int,
2082
+ block_tables: torch.Tensor,
2083
+ context_lens: torch.Tensor,
2084
+ block_tables_stride0: int,
2085
+ scale: float,
2086
+ K_QScale_hip: torch.Tensor,
2087
+ V_QScale_hip: torch.Tensor,
2088
+ K_QScale_asm: torch.Tensor,
2089
+ V_QScale_asm: torch.Tensor,
2090
+ out_: torch.Tensor,
2091
+ kv_cache_dtype: str,
2092
+ ):
2093
+ """
2094
+ Paged attention common function.
2095
+
2096
+ This function is NOT wrapped with @is_aiter_supported decorator
2097
+ to allow explicit backend selection via attention_config to work
2098
+ even when VLLM_ROCM_USE_AITER=0.
2099
+
2100
+ Note: This performs lazy import of aiter.paged_attention_common
2101
+ """
2102
+ from aiter import paged_attention_common
2103
+
2104
+ return paged_attention_common(
2105
+ Q=Q,
2106
+ K=K,
2107
+ V=V,
2108
+ tmp_out=tmp_out,
2109
+ max_logits=max_logits,
2110
+ exp_sums=exp_sums,
2111
+ max_seq_len=max_seq_len,
2112
+ block_tables=block_tables,
2113
+ context_lens=context_lens,
2114
+ block_tables_stride0=block_tables_stride0,
2115
+ scale=scale,
2116
+ K_QScale_hip=K_QScale_hip,
2117
+ V_QScale_hip=V_QScale_hip,
2118
+ K_QScale_asm=K_QScale_asm,
2119
+ V_QScale_asm=V_QScale_asm,
2120
+ out_=out_,
2121
+ kv_cache_dtype=kv_cache_dtype,
2122
+ )
2123
+
2124
2125
rocm_aiter_ops.register_ops_once()
vllm/v1/attention/backends/rocm_aiter_fa.py
+38
-4
@@ -1247,7 +1247,23 @@ class AiterFlashAttentionImpl(AttentionImpl):
1247
v_descale=layer._v_scale.expand(descale_shape),
1248
)
1249
elif rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1250
- num_blocks, block_size, num_kv_heads, head_size = key_cache.shape
1250
+ _, num_heads, head_size = query.shape
1251
+ num_seqs = attn_metadata.seq_lens.shape[0]
1252
+ max_num_partitions = (
1253
+ attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1
1254
+ ) // _PARTITION_SIZE_ROCM
1255
+ tmp_out = torch.empty(
1256
+ (num_seqs, num_heads, max_num_partitions, head_size),
1257
+ dtype=query.dtype,
1258
+ device=query.device,
1259
+ )
1260
+ exp_sums = torch.empty(
1261
+ (num_seqs, num_heads, max_num_partitions),
1262
+ dtype=torch.float32,
1263
+ device=query.device,
1264
+ )
1265
+ max_logits = torch.empty_like(exp_sums)
1266
+ num_blocks, block_size, num_kv_heads, _ = key_cache.shape
1267
x = 16 // key_cache.element_size()
1268
k_cache_template = torch.empty(
1269
[num_blocks, num_kv_heads, head_size // x, block_size, x],
@@ -1261,18 +1277,36 @@ class AiterFlashAttentionImpl(AttentionImpl):
1277
)
1278
new_key_cache = key_cache.view_as(k_cache_template)
1279
new_value_cache = value_cache.view_as(v_cache_template)
1264
- rocm_aiter_ops.pa_fwd_asm(
1280
+ k_qscale = (
1281
+ layer._k_scale
1282
+ if attn_metadata.k_scale is None
1283
+ else attn_metadata.k_scale
1284
+ )
1285
+ v_qscale = (
1286
+ layer._v_scale
1287
+ if attn_metadata.v_scale is None
1288
+ else attn_metadata.v_scale
1289
+ )
1290
+ rocm_aiter_ops.paged_attention_common(
1291
Q=query[:num_decode_tokens],
1292
K=new_key_cache,
1293
V=new_value_cache,
1294
+ tmp_out=tmp_out,
1295
+ max_logits=max_logits,
1296
+ exp_sums=exp_sums,
1297
+ max_seq_len=attn_metadata.max_seq_len,
1298
block_tables=attn_metadata.block_table[:num_decodes],
1299
context_lens=attn_metadata.seq_lens[:num_decodes],
1300
block_tables_stride0=attn_metadata.block_table[
1301
:num_decodes
1302
].stride(0),
1273
- K_QScale=attn_metadata.k_scale,
1274
- V_QScale=attn_metadata.v_scale,
1303
+ scale=self.scale,
1304
+ K_QScale_hip=k_qscale,
1305
+ V_QScale_hip=v_qscale,
1306
+ K_QScale_asm=k_qscale,
1307
+ V_QScale_asm=v_qscale,
1308
out_=output[:num_decode_tokens],
1309
+ kv_cache_dtype=self.kv_cache_dtype,
1310
)
1311
else:
1312
_, num_heads, head_size = query.shape