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 15, 11:00 UTC
[Fix] Support CUDAGraph capture for encoder-decoder on ROCm (#18104)
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
Luka Govedič committed
May 13, 2025 at 22:31 UTC
176a95c670f676e88175c6d3a507ace0b1c35f3d
1 file changed
+8
-8
vllm/attention/backends/utils.py
+8
-8
@@ -345,10 +345,10 @@ class CommonAttentionState(AttentionState):
345
if is_encoder_decoder_model:
346
# The encoder decoder model works only with XFormers and
347
# Flash Attention backend. Assert the same.
348
- assert self.runner.attn_backend.get_name() in\
349
- ["XFORMERS", "FLASH_ATTN"], \
350
- f"Expected attn_backend name to be either 'XFORMERS' or " \
351
- f"'FLASH_ATTN', but "\
348
+ assert self.runner.attn_backend.get_name() in \
349
+ ["XFORMERS", "FLASH_ATTN", "ROCM_FLASH"], \
350
+ f"Expected attn_backend name to be either 'XFORMERS'," \
351
+ f"'ROCM_FLASH', or 'FLASH_ATTN', but " \
352
f"got '{self.runner.attn_backend.get_name()}'"
353
self._update_captured_metadata_for_enc_dec_model(
354
batch_size=batch_size, attn_metadata=attn_metadata)
@@ -367,10 +367,10 @@ class CommonAttentionState(AttentionState):
367
if is_encoder_decoder_model:
368
# The encoder decoder model works only with XFormers and
369
# Flash Attention backend. Assert the same.
370
- assert self.runner.attn_backend.get_name() in\
371
- ["XFORMERS", "FLASH_ATTN"], \
372
- f"Expected attn_backend name to be either 'XFORMERS' or "\
373
- f"'FLASH_ATTN', but "\
370
+ assert self.runner.attn_backend.get_name() in \
371
+ ["XFORMERS", "FLASH_ATTN", "ROCM_FLASH"], \
372
+ f"Expected attn_backend name to be either 'XFORMERS'," \
373
+ f"'ROCM_FLASH', or 'FLASH_ATTN', but " \
374
f"got '{self.runner.attn_backend.get_name()}'"
375
self._add_additonal_input_buffers_for_enc_dec_model(
376
attn_metadata=attn_metadata, input_buffers=input_buffers)