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 8, 11:15 UTC
[ROCm] [aiter] Split KV cache update for AiterFlashAttention (#33681)
Signed-off-by: kliuae <kuanfu.liu@embeddedllm.com>
kliuae committed
Feb 12, 2026 at 00:26 UTC
64f570ab56cab7e8977c611b78f9a44a9a9f033c
1 file changed
+68
-40
vllm/v1/attention/backends/rocm_aiter_fa.py
+68
-40
@@ -11,6 +11,7 @@ from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.attention import Attention
14
+from vllm.model_executor.layers.attention.attention import get_attention_context
15
from vllm.platforms import current_platform
16
from vllm.utils.math_utils import cdiv
17
from vllm.utils.platform_utils import get_cu_count
@@ -687,6 +688,8 @@ class AiterFlashAttentionBackend(AttentionBackend):
688
def get_supported_head_sizes(cls) -> list[int]:
689
return [64, 128, 256]
690
691
+ forward_includes_kv_cache_update: bool = False
692
+
693
@staticmethod
694
def get_name() -> str:
695
return "FLASH_ATTN"
@@ -982,49 +985,10 @@ class AiterFlashAttentionImpl(AttentionImpl):
985
# performance to make sure it does not introduce any overhead.
986
num_actual_tokens = attn_metadata.num_actual_tokens
987
key_cache, value_cache = kv_cache.unbind(0)
985
- # key and value may be None in the case of cross attention. They are
986
- # calculated once based on the output from the encoder and then cached
987
- # in KV cache.
988
+
989
if self.kv_cache_dtype.startswith("fp8"):
990
key_cache = key_cache.view(current_platform.fp8_dtype())
991
value_cache = value_cache.view(current_platform.fp8_dtype())
991
- if (
992
- self.kv_sharing_target_layer_name is None
993
- and key is not None
994
- and value is not None
995
- ):
996
- # Reshape the input keys and values and store them in the cache.
997
- # Skip this if sharing KV cache with an earlier attention layer.
998
- # NOTE(woosuk): Here, key and value are padded while slot_mapping
999
- # is not padded. However, we don't need to do
1000
- # key[:num_actual_tokens] and value[:num_actual_tokens] because
1001
- # the reshape_and_cache_flash op uses the slot_mapping's shape
1002
- # to determine the number of actual tokens.
1003
- if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1004
- # We may calculate per token quant scale in
1005
- # reshape_and_cache_shuffle_triton which might differ from
1006
- # vllm's style when shuffle layout is used.
1007
- reshape_and_cache_shuffle_triton(
1008
- key,
1009
- value,
1010
- key_cache,
1011
- value_cache,
1012
- attn_metadata.slot_mapping,
1013
- self.kv_cache_dtype,
1014
- attn_metadata.k_scale,
1015
- attn_metadata.v_scale,
1016
- )
1017
- else:
1018
- torch.ops._C_cache_ops.reshape_and_cache_flash(
1019
- key,
1020
- value,
1021
- key_cache,
1022
- value_cache,
1023
- attn_metadata.slot_mapping,
1024
- self.kv_cache_dtype,
1025
- layer._k_scale,
1026
- layer._v_scale,
1027
- )
992
993
# decode:extend:prefill
994
query = query[:num_actual_tokens]
@@ -1215,3 +1179,67 @@ class AiterFlashAttentionImpl(AttentionImpl):
1179
)
1180
1181
return output
1182
+
1183
+ def do_kv_cache_update(
1184
+ self,
1185
+ layer: Attention,
1186
+ key: torch.Tensor,
1187
+ value: torch.Tensor,
1188
+ kv_cache: torch.Tensor,
1189
+ slot_mapping: torch.Tensor,
1190
+ ):
1191
+ attn_metadata, _, _ = get_attention_context(layer.layer_name)
1192
+ if attn_metadata is None:
1193
+ # Profiling run.
1194
+ return
1195
+
1196
+ key_cache, value_cache = kv_cache.unbind(0)
1197
+
1198
+ # key and value may be None in the case of cross attention. They are
1199
+ # calculated once based on the output from the encoder and then cached
1200
+ # in KV cache.
1201
+ if self.kv_cache_dtype.startswith("fp8"):
1202
+ key_cache = key_cache.view(current_platform.fp8_dtype())
1203
+ value_cache = value_cache.view(current_platform.fp8_dtype())
1204
+ if (
1205
+ self.kv_sharing_target_layer_name is None
1206
+ and key is not None
1207
+ and value is not None
1208
+ ):
1209
+ # Reshape the input keys and values and store them in the cache.
1210
+ # Skip this if sharing KV cache with an earlier attention layer.
1211
+ # NOTE(woosuk): Here, key and value are padded while slot_mapping
1212
+ # is not padded. However, we don't need to do
1213
+ # key[:num_actual_tokens] and value[:num_actual_tokens] because
1214
+ # the reshape_and_cache_flash op uses the slot_mapping's shape
1215
+ # to determine the number of actual tokens.
1216
+ if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1217
+ # We may calculate per token quant scale in
1218
+ # reshape_and_cache_shuffle_triton which might differ from
1219
+ # vllm's style when shuffle layout is used.
1220
+ k_scale = attn_metadata.k_scale
1221
+ v_scale = attn_metadata.v_scale
1222
+ assert k_scale is not None and v_scale is not None, (
1223
+ "k_scale and v_scale are required for shuffled update"
1224
+ )
1225
+ reshape_and_cache_shuffle_triton(
1226
+ key,
1227
+ value,
1228
+ key_cache,
1229
+ value_cache,
1230
+ slot_mapping,
1231
+ self.kv_cache_dtype,
1232
+ k_scale,
1233
+ v_scale,
1234
+ )
1235
+ else:
1236
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
1237
+ key,
1238
+ value,
1239
+ key_cache,
1240
+ value_cache,
1241
+ slot_mapping,
1242
+ self.kv_cache_dtype,
1243
+ layer._k_scale,
1244
+ layer._v_scale,
1245
+ )