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 9, 12:30 UTC
[ROCm] Enabling forward_includes_kv_cache on ROCm MHA backends (#33106)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Jan 28, 2026 at 00:36 UTC
22ad64950119db2a6fa3911cd87682f8f339bc3a
3 files changed
+127
-87
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+34
-22
@@ -10,7 +10,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
10
QuantKey,
11
kFp8StaticTensorSym,
12
)
13
-from vllm.v1.attention.backend import AttentionType
13
+from vllm.v1.attention.backend import AttentionLayer, AttentionType
14
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
15
from vllm.v1.attention.backends.rocm_attn import (
16
RocmAttentionBackend,
@@ -24,6 +24,8 @@ logger = init_logger(__name__)
24
class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
25
accept_output_buffer: bool = True
26
27
+ forward_includes_kv_cache_update: bool = False
28
+
29
@staticmethod
30
def get_name() -> str:
31
return "ROCM_AITER_UNIFIED_ATTN"
@@ -142,27 +144,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
144
145
key_cache, value_cache = kv_cache.unbind(0)
146
145
- # key and value may be None in the case of cross attention. They are
146
- # calculated once based on the output from the encoder and then cached
147
- # in KV cache.
148
- if (
149
- self.kv_sharing_target_layer_name is None
150
- and key is not None
151
- and value is not None
152
- ):
153
- # Reshape the input keys and values and store them in the cache.
154
- # Skip this if sharing KV cache with an earlier attention layer.
155
- ops.reshape_and_cache_flash(
156
- key,
157
- value,
158
- key_cache,
159
- value_cache,
160
- attn_metadata.slot_mapping,
161
- self.kv_cache_dtype,
162
- layer._k_scale,
163
- layer._v_scale,
164
- )
165
-
147
if self.kv_cache_dtype.startswith("fp8"):
148
key_cache = key_cache.view(self.fp8_dtype)
149
value_cache = value_cache.view(self.fp8_dtype)
@@ -204,3 +185,34 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
185
)
186
187
return output
188
+
189
+ def do_kv_cache_update(
190
+ self,
191
+ layer: AttentionLayer,
192
+ key: torch.Tensor,
193
+ value: torch.Tensor,
194
+ kv_cache: torch.Tensor,
195
+ slot_mapping: torch.Tensor,
196
+ ):
197
+ key_cache, value_cache = kv_cache.unbind(0)
198
+
199
+ # key and value may be None in the case of cross attention. They are
200
+ # calculated once based on the output from the encoder and then cached
201
+ # in KV cache.
202
+ if (
203
+ self.kv_sharing_target_layer_name is None
204
+ and key is not None
205
+ and value is not None
206
+ ):
207
+ # Reshape the input keys and values and store them in the cache.
208
+ # Skip this if sharing KV cache with an earlier attention layer.
209
+ ops.reshape_and_cache_flash(
210
+ key,
211
+ value,
212
+ key_cache,
213
+ value_cache,
214
+ slot_mapping,
215
+ self.kv_cache_dtype,
216
+ layer._k_scale,
217
+ layer._v_scale,
218
+ )
vllm/v1/attention/backends/rocm_attn.py
+55
-40
@@ -18,6 +18,7 @@ from vllm.v1.attention.backend import (
18
AttentionBackend,
19
AttentionCGSupport,
20
AttentionImpl,
21
+ AttentionLayer,
22
AttentionMetadataBuilder,
23
AttentionType,
24
CommonAttentionMetadata,
@@ -193,6 +194,8 @@ class RocmAttentionBackend(AttentionBackend):
194
"FlexAttention backend which supports all head sizes."
195
)
196
197
+ forward_includes_kv_cache_update: bool = False
198
+
199
@staticmethod
200
def get_name() -> str:
201
return "ROCM_ATTN"
@@ -330,6 +333,56 @@ class RocmAttentionImpl(AttentionImpl):
333
kv_cache, self.num_kv_heads, self.head_size
334
)
335
336
+ if self.kv_cache_dtype.startswith("fp8"):
337
+ key_cache = key_cache.view(self.fp8_dtype)
338
+ value_cache = value_cache.view(self.fp8_dtype)
339
+ assert layer._q_scale_float == 1.0, (
340
+ "A non 1.0 q_scale is not currently supported."
341
+ )
342
+
343
+ cu_seqlens_q = attn_metadata.query_start_loc
344
+ seqused_k = attn_metadata.seq_lens
345
+ max_seqlen_q = attn_metadata.max_query_len
346
+ max_seqlen_k = attn_metadata.max_seq_len
347
+ block_table = attn_metadata.block_table
348
+
349
+ # Compute attention and update output up to `num_actual_tokens`.
350
+ chunked_prefill_paged_decode(
351
+ query=query[:num_actual_tokens],
352
+ key=key[:num_actual_tokens] if key is not None else None,
353
+ value=value[:num_actual_tokens] if value is not None else None,
354
+ output=output[:num_actual_tokens],
355
+ kv_cache_dtype=self.kv_cache_dtype,
356
+ key_cache=key_cache,
357
+ value_cache=value_cache,
358
+ block_table=block_table,
359
+ query_start_loc=cu_seqlens_q,
360
+ seq_lens=seqused_k,
361
+ max_seq_len=max_seqlen_k,
362
+ max_query_len=max_seqlen_q,
363
+ k_scale=layer._k_scale,
364
+ v_scale=layer._v_scale,
365
+ alibi_slopes=self.alibi_slopes,
366
+ sliding_window=self.sliding_window[0],
367
+ sm_scale=self.scale,
368
+ output_scale=output_scale,
369
+ sinks=self.sinks,
370
+ )
371
+
372
+ return output
373
+
374
+ def do_kv_cache_update(
375
+ self,
376
+ layer: AttentionLayer,
377
+ key: torch.Tensor,
378
+ value: torch.Tensor,
379
+ kv_cache: torch.Tensor,
380
+ slot_mapping: torch.Tensor,
381
+ ):
382
+ key_cache, value_cache = PagedAttention.split_kv_cache(
383
+ kv_cache, self.num_kv_heads, self.head_size
384
+ )
385
+
386
# key and value may be None in the case of cross attention. They are
387
# calculated once based on the output from the encoder and then cached
388
# in KV cache.
@@ -354,7 +407,7 @@ class RocmAttentionImpl(AttentionImpl):
407
value,
408
key_cache,
409
value_cache,
357
- attn_metadata.slot_mapping,
410
+ slot_mapping,
411
self.kv_cache_dtype,
412
layer._k_scale,
413
layer._v_scale,
@@ -367,46 +420,8 @@ class RocmAttentionImpl(AttentionImpl):
420
value,
421
key_cache,
422
value_cache,
370
- attn_metadata.slot_mapping,
423
+ slot_mapping,
424
self.kv_cache_dtype,
425
layer._k_scale,
426
layer._v_scale,
427
)
375
-
376
- if self.kv_cache_dtype.startswith("fp8"):
377
- key_cache = key_cache.view(self.fp8_dtype)
378
- value_cache = value_cache.view(self.fp8_dtype)
379
- assert layer._q_scale_float == 1.0, (
380
- "A non 1.0 q_scale is not currently supported."
381
- )
382
-
383
- cu_seqlens_q = attn_metadata.query_start_loc
384
- seqused_k = attn_metadata.seq_lens
385
- max_seqlen_q = attn_metadata.max_query_len
386
- max_seqlen_k = attn_metadata.max_seq_len
387
- block_table = attn_metadata.block_table
388
-
389
- # Compute attention and update output up to `num_actual_tokens`.
390
- chunked_prefill_paged_decode(
391
- query=query[:num_actual_tokens],
392
- key=key[:num_actual_tokens] if key is not None else None,
393
- value=value[:num_actual_tokens] if value is not None else None,
394
- output=output[:num_actual_tokens],
395
- kv_cache_dtype=self.kv_cache_dtype,
396
- key_cache=key_cache,
397
- value_cache=value_cache,
398
- block_table=block_table,
399
- query_start_loc=cu_seqlens_q,
400
- seq_lens=seqused_k,
401
- max_seq_len=max_seqlen_k,
402
- max_query_len=max_seqlen_q,
403
- k_scale=layer._k_scale,
404
- v_scale=layer._v_scale,
405
- alibi_slopes=self.alibi_slopes,
406
- sliding_window=self.sliding_window[0],
407
- sm_scale=self.scale,
408
- output_scale=output_scale,
409
- sinks=self.sinks,
410
- )
411
-
412
- return output
vllm/v1/attention/backends/triton_attn.py
+38
-25
@@ -21,6 +21,7 @@ from vllm.v1.attention.backend import (
21
AttentionBackend,
22
AttentionCGSupport,
23
AttentionImpl,
24
+ AttentionLayer,
25
AttentionMetadataBuilder,
26
AttentionType,
27
CommonAttentionMetadata,
@@ -271,6 +272,8 @@ class TritonAttentionBackend(AttentionBackend):
272
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
273
return [MultipleOf(16)]
274
275
+ forward_includes_kv_cache_update: bool = False
276
+
277
@staticmethod
278
def get_name() -> str:
279
return "TRITON_ATTN"
@@ -461,31 +464,6 @@ class TritonAttentionImpl(AttentionImpl):
464
465
# For decoder and cross-attention, use KV cache as before
466
key_cache, value_cache = kv_cache.unbind(1)
464
-
465
- if (
466
- self.kv_sharing_target_layer_name is None
467
- and key is not None
468
- and value is not None
469
- ):
470
- # Reshape the input keys and values and store them in the cache.
471
- # Skip this if sharing KV cache with an earlier attention layer.
472
- if self.kv_cache_dtype.startswith("fp8"):
473
- key_cache = key_cache.view(self.fp8_dtype)
474
- value_cache = value_cache.view(self.fp8_dtype)
475
- # triton kernel does not support uint8 kv_cache
476
- # (because some explicit casts (e.g. float8_e4m3fnuz)
477
- # are not supported)
478
- triton_reshape_and_cache_flash(
479
- key,
480
- value,
481
- key_cache,
482
- value_cache,
483
- attn_metadata.slot_mapping,
484
- self.kv_cache_dtype,
485
- layer._k_scale,
486
- layer._v_scale,
487
- )
488
-
467
if self.kv_cache_dtype.startswith("fp8"):
468
if key_cache.dtype != self.fp8_dtype:
469
key_cache = key_cache.view(self.fp8_dtype)
@@ -585,3 +563,38 @@ class TritonAttentionImpl(AttentionImpl):
563
sliding_window_k=self.sliding_window[1],
564
)
565
return output
566
+
567
+ def do_kv_cache_update(
568
+ self,
569
+ layer: AttentionLayer,
570
+ key: torch.Tensor,
571
+ value: torch.Tensor,
572
+ kv_cache: torch.Tensor,
573
+ slot_mapping: torch.Tensor,
574
+ ):
575
+ # For decoder and cross-attention, use KV cache as before
576
+ key_cache, value_cache = kv_cache.unbind(1)
577
+
578
+ if (
579
+ self.kv_sharing_target_layer_name is None
580
+ and key is not None
581
+ and value is not None
582
+ ):
583
+ # Reshape the input keys and values and store them in the cache.
584
+ # Skip this if sharing KV cache with an earlier attention layer.
585
+ if self.kv_cache_dtype.startswith("fp8"):
586
+ key_cache = key_cache.view(self.fp8_dtype)
587
+ value_cache = value_cache.view(self.fp8_dtype)
588
+ # triton kernel does not support uint8 kv_cache
589
+ # (because some explicit casts (e.g. float8_e4m3fnuz)
590
+ # are not supported)
591
+ triton_reshape_and_cache_flash(
592
+ key,
593
+ value,
594
+ key_cache,
595
+ value_cache,
596
+ slot_mapping,
597
+ self.kv_cache_dtype,
598
+ layer._k_scale,
599
+ layer._v_scale,
600
+ )