[Feature]: FlexAttention + encoder_decoder support

[JakubCerven]

🚀 The feature, motivation and pitch

Currently on compute capability < 8.0 (Turing, Volta) encoder_decoder raises NotImplementedError: FlexAttention does not support encoder_decoder attention
Tested on openai/whisper-large-v3 + T4 + vllm/vllm-openai:nightly-0f7acdd73ca6316c8ae0474c0a9c4fc264e87a7b

cc @russellb @DarkLight1337

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[russellb]

Thanks for filing an issue. I started looking at this, but it doesn't seem as trivial as it was to update flashinfer and flash_attn to support this.

[Isotr0py]

it doesn't seem as trivial as it was to update flashinfer and flash_attn to support this

FlexAttention has supported both bidirectional attention (encoder) and casual attention (decoder), so I think the only problem is just how to prepare the attention mask effectively.

Perhaps I can look into this feature as well this weekend :)

[russellb]

FWIW, here's the first change that I know is needed:

commit 8ea0e33fdecfd2f5dd47556790896e6be566d772 (HEAD -> whisper-flex-attn)
Author: Russell Bryant <rbryant@redhat.com>
Date:   Thu Sep 11 17:59:55 2025 +0000

    flex_attention: allow ENCODER_DECODER attention
    
    Signed-off-by: Russell Bryant <rbryant@redhat.com>

diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py
index 662d39845..f2574bc78 100644
--- a/vllm/v1/attention/backends/flex_attention.py
+++ b/vllm/v1/attention/backends/flex_attention.py
@@ -631,8 +631,8 @@ class FlexAttentionImpl(AttentionImpl):
         self.num_kv_heads = num_kv_heads
         self.attn_type = attn_type
 
-        if attn_type not in (AttentionType.ENCODER_ONLY,
-                             AttentionType.DECODER):
+        if attn_type not in (AttentionType.ENCODER_ONLY, AttentionType.DECODER,
+                             AttentionType.ENCODER_DECODER):
             raise NotImplementedError(
                 f"FlexAttention does not support {attn_type} attention")
 
@@ -712,8 +712,7 @@ class FlexAttentionImpl(AttentionImpl):
 
         num_actual_tokens = attn_metadata.num_actual_tokens
 
-        if not attn_metadata.causal:
-            assert self.attn_type == AttentionType.ENCODER_ONLY
+        if self.attn_type == AttentionType.ENCODER_ONLY:
 
             query, key_tensor, value_tensor = map(
                 lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
@@ -730,19 +729,24 @@ class FlexAttentionImpl(AttentionImpl):
                 value_tensor = value_tensor[:, :, :num_actual_tokens, :]
 
         else:
-            assert self.attn_type == AttentionType.DECODER
+            assert self.attn_type in (AttentionType.DECODER,
+                                      AttentionType.ENCODER_DECODER)
             key_cache, value_cache = kv_cache.unbind(0)
 
-            torch.ops._C_cache_ops.reshape_and_cache_flash(
-                key,
-                value,
-                key_cache,
-                value_cache,
-                attn_metadata.slot_mapping,
-                self.kv_cache_dtype,
-                layer._k_scale,
-                layer._v_scale,
-            )
+            # key and value may be None in the case of cross attention. They are
+            # calculated once based on the output from the encoder and then
+            # cached in KV cache.
+            if key is not None and value is not None:
+                torch.ops._C_cache_ops.reshape_and_cache_flash(
+                    key,
+                    value,
+                    key_cache,
+                    value_cache,
+                    attn_metadata.slot_mapping,
+                    self.kv_cache_dtype,
+                    layer._k_scale,
+                    layer._v_scale,
+                )
 
             # View out the block_size dim
             key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)

but here's how it breaks:

(EngineCore_DP0 pid=1345647) Traceback (most recent call last):
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
(EngineCore_DP0 pid=1345647)     self.run()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/.local/share/uv/python/cpython-3.12.11-linux-x86_64-gnu/lib/python3.12/multiprocessing/process.py", line 108, in run
(EngineCore_DP0 pid=1345647)     self._target(*self._args, **self._kwargs)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 716, in run_engine_core
(EngineCore_DP0 pid=1345647)     raise e
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 705, in run_engine_core
(EngineCore_DP0 pid=1345647)     engine_core.run_busy_loop()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 732, in run_busy_loop
(EngineCore_DP0 pid=1345647)     self._process_engine_step()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 758, in _process_engine_step
(EngineCore_DP0 pid=1345647)     outputs, model_executed = self.step_fn()
(EngineCore_DP0 pid=1345647)                               ^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 288, in step
(EngineCore_DP0 pid=1345647)     model_output = self.execute_model_with_error_logging(
(EngineCore_DP0 pid=1345647)                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 274, in execute_model_with_error_logging
(EngineCore_DP0 pid=1345647)     raise err
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/engine/core.py", line 265, in execute_model_with_error_logging
(EngineCore_DP0 pid=1345647)     return model_fn(scheduler_output)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/executor/abstract.py", line 103, in execute_model
(EngineCore_DP0 pid=1345647)     output = self.collective_rpc("execute_model",
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/executor/uniproc_executor.py", line 83, in collective_rpc
(EngineCore_DP0 pid=1345647)     return [run_method(self.driver_worker, method, args, kwargs)]
(EngineCore_DP0 pid=1345647)             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/utils/__init__.py", line 3067, in run_method
(EngineCore_DP0 pid=1345647)     return func(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
(EngineCore_DP0 pid=1345647)     return func(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/worker/gpu_worker.py", line 446, in execute_model
(EngineCore_DP0 pid=1345647)     output = self.model_runner.execute_model(scheduler_output,
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
(EngineCore_DP0 pid=1345647)     return func(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/worker/gpu_model_runner.py", line 2101, in execute_model
(EngineCore_DP0 pid=1345647)     model_output = self.model(
(EngineCore_DP0 pid=1345647)                    ^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
(EngineCore_DP0 pid=1345647)     return self._call_impl(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
(EngineCore_DP0 pid=1345647)     return forward_call(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/model_executor/models/whisper.py", line 901, in forward
(EngineCore_DP0 pid=1345647)     decoder_outputs = self.model(
(EngineCore_DP0 pid=1345647)                       ^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
(EngineCore_DP0 pid=1345647)     return self._call_impl(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
(EngineCore_DP0 pid=1345647)     return forward_call(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/model_executor/models/whisper.py", line 606, in forward
(EngineCore_DP0 pid=1345647)     decoder_outputs = self.decoder(
(EngineCore_DP0 pid=1345647)                       ^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
(EngineCore_DP0 pid=1345647)     return self._call_impl(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
(EngineCore_DP0 pid=1345647)     return forward_call(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/model_executor/models/whisper.py", line 575, in forward
(EngineCore_DP0 pid=1345647)     hidden_states = decoder_layer(
(EngineCore_DP0 pid=1345647)                     ^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
(EngineCore_DP0 pid=1345647)     return self._call_impl(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
(EngineCore_DP0 pid=1345647)     return forward_call(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/model_executor/models/whisper.py", line 458, in forward
(EngineCore_DP0 pid=1345647)     hidden_states = self.self_attn(hidden_states=hidden_states)
(EngineCore_DP0 pid=1345647)                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
(EngineCore_DP0 pid=1345647)     return self._call_impl(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
(EngineCore_DP0 pid=1345647)     return forward_call(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/model_executor/models/whisper.py", line 265, in forward
(EngineCore_DP0 pid=1345647)     attn_output = self.attn(q, k, v)
(EngineCore_DP0 pid=1345647)                   ^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
(EngineCore_DP0 pid=1345647)     return self._call_impl(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
(EngineCore_DP0 pid=1345647)     return forward_call(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/attention/layer.py", line 306, in forward
(EngineCore_DP0 pid=1345647)     torch.ops.vllm.unified_attention_with_output(
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_ops.py", line 1243, in __call__
(EngineCore_DP0 pid=1345647)     return self._op(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/attention/layer.py", line 597, in unified_attention_with_output
(EngineCore_DP0 pid=1345647)     self.impl.forward(self,
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/vllm/v1/attention/backends/flex_attention.py", line 770, in forward
(EngineCore_DP0 pid=1345647)     out = flex_attention_compiled(
(EngineCore_DP0 pid=1345647)           ^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 736, in compile_wrapper
(EngineCore_DP0 pid=1345647)     return fn(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1495, in __call__
(EngineCore_DP0 pid=1345647)     return self._torchdynamo_orig_callable(
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 629, in __call__
(EngineCore_DP0 pid=1345647)     return _compile(
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1164, in _compile
(EngineCore_DP0 pid=1345647)     raise InternalTorchDynamoError(
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1111, in _compile
(EngineCore_DP0 pid=1345647)     guarded_code = compile_inner(code, one_graph, hooks, transform)
(EngineCore_DP0 pid=1345647)                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_utils_internal.py", line 97, in wrapper_function
(EngineCore_DP0 pid=1345647)     return function(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 793, in compile_inner
(EngineCore_DP0 pid=1345647)     return _compile_inner(code, one_graph, hooks, transform)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 832, in _compile_inner
(EngineCore_DP0 pid=1345647)     out_code = transform_code_object(code, transform)
(EngineCore_DP0 pid=1345647)                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/bytecode_transformation.py", line 1424, in transform_code_object
(EngineCore_DP0 pid=1345647)     transformations(instructions, code_options)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 267, in _fn
(EngineCore_DP0 pid=1345647)     return fn(*args, **kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 753, in transform
(EngineCore_DP0 pid=1345647)     tracer.run()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 3497, in run
(EngineCore_DP0 pid=1345647)     super().run()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1363, in run
(EngineCore_DP0 pid=1345647)     while self.step():
(EngineCore_DP0 pid=1345647)           ^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1267, in step
(EngineCore_DP0 pid=1345647)     self.dispatch_table[inst.opcode](self, inst)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2294, in LOAD_ATTR
(EngineCore_DP0 pid=1345647)     self._load_attr(inst)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2282, in _load_attr
(EngineCore_DP0 pid=1345647)     result = BuiltinVariable(getattr).call_function(
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 1189, in call_function
(EngineCore_DP0 pid=1345647)     return handler(tx, args, kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 1009, in builtin_dispatch
(EngineCore_DP0 pid=1345647)     rv = fn(tx, args, kwargs)
(EngineCore_DP0 pid=1345647)          ^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 887, in call_self_handler
(EngineCore_DP0 pid=1345647)     result = self_handler(tx, *args, **kwargs)
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 2099, in call_getattr
(EngineCore_DP0 pid=1345647)     return obj.var_getattr(tx, name)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/user_defined.py", line 1250, in var_getattr
(EngineCore_DP0 pid=1345647)     ).call_function(tx, [], {})
(EngineCore_DP0 pid=1345647)       ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/functions.py", line 1055, in call_function
(EngineCore_DP0 pid=1345647)     return super().call_function(tx, args, kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/functions.py", line 529, in call_function
(EngineCore_DP0 pid=1345647)     return super().call_function(tx, args, kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/functions.py", line 293, in call_function
(EngineCore_DP0 pid=1345647)     return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1210, in inline_user_function_return
(EngineCore_DP0 pid=1345647)     return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 3698, in inline_call
(EngineCore_DP0 pid=1345647)     return tracer.inline_call_()
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 3901, in inline_call_
(EngineCore_DP0 pid=1345647)     self.run()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1363, in run
(EngineCore_DP0 pid=1345647)     while self.step():
(EngineCore_DP0 pid=1345647)           ^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 1267, in step
(EngineCore_DP0 pid=1345647)     self.dispatch_table[inst.opcode](self, inst)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2294, in LOAD_ATTR
(EngineCore_DP0 pid=1345647)     self._load_attr(inst)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2282, in _load_attr
(EngineCore_DP0 pid=1345647)     result = BuiltinVariable(getattr).call_function(
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 1189, in call_function
(EngineCore_DP0 pid=1345647)     return handler(tx, args, kwargs)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builtin.py", line 822, in <lambda>
(EngineCore_DP0 pid=1345647)     tx, [v.realize() for v in args], kwargs
(EngineCore_DP0 pid=1345647)          ^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/lazy.py", line 67, in realize
(EngineCore_DP0 pid=1345647)     self._cache.realize()
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/lazy.py", line 32, in realize
(EngineCore_DP0 pid=1345647)     self.vt = builder.VariableBuilder(tx, self.source)(self.value)
(EngineCore_DP0 pid=1345647)               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builder.py", line 440, in __call__
(EngineCore_DP0 pid=1345647)     vt = self._wrap(value)
(EngineCore_DP0 pid=1345647)          ^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builder.py", line 652, in _wrap
(EngineCore_DP0 pid=1345647)     return type_dispatch(self, value)
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builder.py", line 2031, in wrap_tensor
(EngineCore_DP0 pid=1345647)     example_value = wrap_to_fake_tensor_and_record(
(EngineCore_DP0 pid=1345647)                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builder.py", line 3381, in wrap_to_fake_tensor_and_record
(EngineCore_DP0 pid=1345647)     fake_e = wrap_fake_exception(
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/utils.py", line 2753, in wrap_fake_exception
(EngineCore_DP0 pid=1345647)     return fn()
(EngineCore_DP0 pid=1345647)            ^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_dynamo/variables/builder.py", line 3382, in <lambda>
(EngineCore_DP0 pid=1345647)     lambda: tx.fake_mode.from_tensor(
(EngineCore_DP0 pid=1345647)             ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.py", line 2948, in from_tensor
(EngineCore_DP0 pid=1345647)     return self.fake_tensor_converter.from_real_tensor(
(EngineCore_DP0 pid=1345647)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.py", line 404, in from_real_tensor
(EngineCore_DP0 pid=1345647)     out = self.meta_converter(
(EngineCore_DP0 pid=1345647)           ^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_subclasses/meta_utils.py", line 1896, in __call__
(EngineCore_DP0 pid=1345647)     t_desc = self.describer.describe_tensor(t, trace=trace)
(EngineCore_DP0 pid=1345647)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/_subclasses/meta_utils.py", line 312, in describe_tensor
(EngineCore_DP0 pid=1345647)     storage = self.describe_storage(t.untyped_storage(), trace=trace)
(EngineCore_DP0 pid=1345647)                                     ^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=1345647) torch._dynamo.exc.InternalTorchDynamoError: RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. Stack trace: File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 903, in create_block_mask
(EngineCore_DP0 pid=1345647)     block_mask = _create_sparse_block_from_block_mask(
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 764, in _create_sparse_block_from_block_mask
(EngineCore_DP0 pid=1345647)     partial_bm = _dense_to_ordered(partial_blocks)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 182, in _dense_to_ordered
(EngineCore_DP0 pid=1345647)     col_indices.to(torch.int32, memory_format=torch.contiguous_format),. To prevent overwriting, clone the tensor outside of torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.
(EngineCore_DP0 pid=1345647) 
(EngineCore_DP0 pid=1345647) from user code:
(EngineCore_DP0 pid=1345647)    File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 1361, in flex_attention
(EngineCore_DP0 pid=1345647)     block_mask_q_len = block_mask.shape[-2]
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 415, in shape
(EngineCore_DP0 pid=1345647)     *batch_dims, _, _ = self.kv_indices.shape
(EngineCore_DP0 pid=1345647) 
(EngineCore_DP0 pid=1345647) Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"
(EngineCore_DP0 pid=1345647) 

[russellb]

and this is a quick way to test it:

VLLM_ATTENTION_BACKEND=FLEX_ATTENTION python examples/offline_inference/audio_language.py --model whisper

[Isotr0py]

(EngineCore_DP0 pid=1345647) torch._dynamo.exc.InternalTorchDynamoError: RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run. Stack trace: File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 903, in create_block_mask
(EngineCore_DP0 pid=1345647)     block_mask = _create_sparse_block_from_block_mask(
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 764, in _create_sparse_block_from_block_mask
(EngineCore_DP0 pid=1345647)     partial_bm = _dense_to_ordered(partial_blocks)
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 182, in _dense_to_ordered
(EngineCore_DP0 pid=1345647)     col_indices.to(torch.int32, memory_format=torch.contiguous_format),. To prevent overwriting, clone the tensor outside of torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.
(EngineCore_DP0 pid=1345647) 
(EngineCore_DP0 pid=1345647) from user code:
(EngineCore_DP0 pid=1345647)    File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 1361, in flex_attention
(EngineCore_DP0 pid=1345647)     block_mask_q_len = block_mask.shape[-2]
(EngineCore_DP0 pid=1345647)   File "/home/rbryant/vllm/venv/lib/python3.12/site-packages/torch/nn/attention/flex_attention.py", line 415, in shape
(EngineCore_DP0 pid=1345647)     *batch_dims, _, _ = self.kv_indices.shape
(EngineCore_DP0 pid=1345647) 
(EngineCore_DP0 pid=1345647) Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

I have met this issue before when adding sliding windows support to flex attention (#24089 (comment)). And seems the root issue is that FlexAttention will try to access tenosrs outside CUDAGraphs when enabling hybrid allocator.

Perhaps we should make FlexAttention compatible with hybrid allocator first. 😿