fix cross attention #28346

fsx950223 · 2025-11-08T11:47:55Z

Purpose

Test Plan

pytest entrypoints/openai/test_transcription_validation.py::test_basic_audio[openai/whisper-large-v3-turbo] --maxfail=1

Test Result

=============================================================================== test session starts ===============================================================================
platform linux -- Python 3.12.11, pytest-8.4.1, pluggy-1.6.0
rootdir: /mnt/raid0/sixifang/vllm
configfile: pyproject.toml
plugins: asyncio-1.1.0, anyio-4.10.0
asyncio: mode=Mode.STRICT, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 1 item                                                                                                                                                                  

entrypoints/openai/test_transcription_validation.py .                                                                                                                       [100%]

================================================================================ warnings summary =================================================================================
<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

../../../../../usr/lib/python3/dist-packages/pyparsing.py:108
  /usr/lib/python3/dist-packages/pyparsing.py:108: DeprecationWarning: module 'sre_constants' is deprecated
    import sre_constants

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================================================== 1 passed, 3 warnings in 98.74s (0:01:38) =====================================================================
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Signed-off-by: fsx950223 <fsx950223@outlook.com>

gemini-code-assist

Code Review

This pull request aims to add cross-attention support to the Triton attention backend. While the intent is good, it introduces a critical issue where cross-attention is incorrectly treated as causal attention, which will lead to incorrect model outputs. The underlying Triton kernel, unified_attention, currently only supports causal attention and would need to be modified to correctly handle cross-attention. The other change to add ROCm support in a utility function appears correct.

gemini-code-assist · 2025-11-08T11:50:13Z

vllm/v1/attention/backends/triton_attn.py

+        if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
            raise NotImplementedError(
-                "Encoder self-attention and "
-                "encoder/decoder cross-attention "
-                "are not implemented for "
-                "TritonAttentionImpl"
+                "Encoder self-attention is not implemented for TritonAttentionImpl"
            )
-
+        self.attn_type = attn_type


This change enables AttentionType.ENCODER_DECODER, but the forward method of TritonAttentionImpl subsequently calls unified_attention with causal=True hardcoded. Cross-attention is non-causal, so this will result in incorrect attention calculations for encoder-decoder models.

Furthermore, the unified_attention function itself asserts that it only supports causal attention. To properly support cross-attention, unified_attention and its underlying Triton kernel must be updated to handle non-causal cases. Without this, enabling ENCODER_DECODER here introduces a bug.

Is this valid?

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2025-11-08T11:52:01Z

vllm/v1/attention/backends/triton_attn.py

        max_seqlen_k = attn_metadata.max_seq_len
        block_table = attn_metadata.block_table

-        descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
+        descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])

        unified_attention(
            q=query[:num_actual_tokens],


Cross-attention still forced through causal mask

The constructor now allows AttentionType.ENCODER_DECODER, but the execution path still unconditionally calls unified_attention(..., causal=True, …). Cross-attention builders explicitly set attn_metadata.causal = False so that queries can see the full encoder sequence. For batches where multiple decoder queries are processed at once (e.g. teacher-forced decoding or speculative decoding of several tokens), keeping causal=True restricts each query to only the prefix of encoder keys and produces incorrect attention weights. If cross-attention is intended to work for Triton, this should respect attn_metadata.causal (and the kernel would need to support the non-causal case); otherwise the change to accept ENCODER_DECODER silently yields wrong results.

Useful? React with 👍 / 👎.

maleksan85

LGTM! Thanks!

mergify · 2025-11-11T12:53:36Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fsx950223.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: fsx950223 <fsx950223@outlook.com>

tjtanaa · 2025-11-12T20:19:39Z

vllm/v1/worker/utils.py

            # in the same decoder block.
-            if current_platform.is_cuda() or current_platform.is_xpu():
+            if (
+                current_platform.is_cuda()


NITS: use current_platform.is_cuda_alike() and remove current_platform.is_rocm()

tjtanaa · 2025-11-12T20:20:24Z

/gemini review

gemini-code-assist

Code Review

This pull request aims to add cross-attention support to the Triton attention backend. While the changes in vllm/v1/worker/utils.py and the robustness improvement in triton_attn.py are correct, there is a critical issue in the main logic for handling cross-attention. The current implementation incorrectly attempts to cache encoder keys and values into the decoder's paged KV cache, which will lead to memory corruption or a crash. I've provided a detailed comment and a suggested fix to prevent this bug.

vllm/v1/attention/backends/triton_attn.py

tjtanaa · 2025-11-12T23:28:53Z

@SageMoore @tdoublep @bringlein Could you help to take a look? Thank you.

Signed-off-by: fsx950223 <fsx950223@outlook.com>

vllm/v1/attention/backends/triton_attn.py

tjtanaa · 2025-11-19T14:33:57Z

@SageMoore @tdoublep @bringlein Could you help to take a look? Thank you.

SageMoore

Looks reasonable. Just one question.

SageMoore · 2025-11-19T15:54:44Z

vllm/v1/attention/backends/triton_attn.py

@@ -346,7 +347,7 @@ def forward(
        max_seqlen_k = attn_metadata.max_seq_len
        block_table = attn_metadata.block_table

-        descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
+        descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])


This change is a noop right? We are just getting the num_kv_heads from a different spot? Is there some subtle difference that I'm missing?

Can we just leave the original code then?

No, key may be None.

sammysun0711 · 2025-11-21T05:23:03Z

@fsx950223, could you please rebase?

DarkLight1337 · 2025-11-21T12:55:37Z

H100 test is failing in latest nightly: https://buildkite.com/vllm/ci/builds/40068/steps/canvas?jid=019aa55f-e011-41eb-bf7f-b02500cc76d4#019aa55f-e011-41eb-bf7f-b02500cc76d4

So we can merge this

Signed-off-by: fsx950223 <fsx950223@outlook.com>

Signed-off-by: fsx950223 <fsx950223@outlook.com> Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu>

Signed-off-by: fsx950223 <fsx950223@outlook.com>

fix cross attention

232b99f

Signed-off-by: fsx950223 <fsx950223@outlook.com>

fsx950223 requested a review from tdoublep as a code owner November 8, 2025 11:47

mergify bot added the v1 label Nov 8, 2025

gemini-code-assist bot reviewed Nov 8, 2025

View reviewed changes

chatgpt-codex-connector bot reviewed Nov 8, 2025

View reviewed changes

coderfeli mentioned this pull request Nov 8, 2025

[CI Failure][AMD] Encoder-Decoder Models Fail on AMD CI #27442

Closed

3 tasks

zhewenl requested review from DarkLight1337, LucasWilkinson, markmc, sarckk and ywang96 November 8, 2025 17:28

zhewenl mentioned this pull request Nov 8, 2025

☂️[AMD][CI Failure]: AMD CI Issues Master #27680

Open

14 tasks

maleksan85 approved these changes Nov 8, 2025

View reviewed changes

ywang96 assigned russellb Nov 9, 2025

apinge mentioned this pull request Nov 10, 2025

[ROCm] Support for Whisper v1 with Aiter Unified Attention and Aiter Flash Attention #28376

Merged

5 tasks

zhewenl mentioned this pull request Nov 10, 2025

[CI/Build] Skip encoder-decoder models on AMD #28156

Draft

mergify bot added the needs-rebase label Nov 11, 2025

Merge remote-tracking branch 'upstream/main' into fix_cross_attention

a5bd54a

Signed-off-by: fsx950223 <fsx950223@outlook.com>

mergify bot removed the needs-rebase label Nov 12, 2025

vllmellm mentioned this pull request Nov 12, 2025

[Bug]: encoder_decoder models (e.g. Whisper) is not working in vLLM 0.11 with ROCm #28184

Open

1 task

tjtanaa reviewed Nov 12, 2025

View reviewed changes

gemini-code-assist bot reviewed Nov 12, 2025

View reviewed changes

vllm/v1/attention/backends/triton_attn.py Show resolved Hide resolved

change api

9a17437

Signed-off-by: fsx950223 <fsx950223@outlook.com>

apinge reviewed Nov 13, 2025

View reviewed changes

vllm/v1/attention/backends/triton_attn.py Outdated Show resolved Hide resolved

fsx950223 force-pushed the fix_cross_attention branch from 2b244f5 to 9a17437 Compare November 14, 2025 03:42

SageMoore reviewed Nov 19, 2025

View reviewed changes

SageMoore approved these changes Nov 20, 2025

View reviewed changes

tjtanaa approved these changes Nov 21, 2025

View reviewed changes

DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 21, 2025

fsx950223 added 2 commits November 21, 2025 13:26

Merge branch 'main' into fix_cross_attention

086ebdf

Merge branch 'main' into fix_cross_attention

f5c18a9

vllm-bot merged commit fc9f821 into vllm-project:main Nov 21, 2025
42 of 45 checks passed

ywang96 pushed a commit to ywang96/vllm that referenced this pull request Nov 23, 2025

fix cross attention (vllm-project#28346)

467e997

Signed-off-by: fsx950223 <fsx950223@outlook.com>

lpapavassiliou pushed a commit to lpapavassiliou/vllm that referenced this pull request Nov 24, 2025

fix cross attention (vllm-project#28346)

7fb4a2c

Signed-off-by: fsx950223 <fsx950223@outlook.com>

RunkaiTao pushed a commit to RunkaiTao/vllm that referenced this pull request Nov 24, 2025

fix cross attention (vllm-project#28346)

529e200

Signed-off-by: fsx950223 <fsx950223@outlook.com> Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu>

bringlein pushed a commit to bringlein/vllm that referenced this pull request Nov 26, 2025

fix cross attention (vllm-project#28346)

7791794

Signed-off-by: fsx950223 <fsx950223@outlook.com>

devpatelio pushed a commit to SumanthRH/vllm that referenced this pull request Nov 29, 2025

fix cross attention (vllm-project#28346)

52f2032

Signed-off-by: fsx950223 <fsx950223@outlook.com>

kitaekatt pushed a commit to kitaekatt/vllm that referenced this pull request Dec 1, 2025

fix cross attention (vllm-project#28346)

cd3e6d6

Signed-off-by: fsx950223 <fsx950223@outlook.com>

Uh oh!

fix cross attention #28346

fix cross attention #28346

Conversation

fsx950223 commented Nov 8, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 8, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector bot Nov 8, 2025

Choose a reason for hiding this comment

Uh oh!

maleksan85 left a comment

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Nov 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tjtanaa commented Nov 12, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

tjtanaa commented Nov 12, 2025

Uh oh!

Uh oh!

tjtanaa commented Nov 19, 2025

Uh oh!

SageMoore left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

fsx950223 Nov 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sammysun0711 commented Nov 21, 2025

Uh oh!

DarkLight1337 commented Nov 21, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

fsx950223 commented Nov 8, 2025 •

edited by github-actions bot

Loading

fsx950223 Nov 20, 2025 •

edited

Loading