[https://nvbugs/6571220][fix] Correct unfused attention context workspace sizing

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -799,7 +799,10 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3
 
     auto const batch_size = static_cast<size_t>(max_num_seq);
     auto const kv_seq_length = (isCrossAttention() ? cross_kv_length : input_seq_length);
-    size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * max_num_tokens * kv_seq_length;
+    // The unfused-MHA buffers below must upper-bound the enqueueContext carve, which sizes them by
+    // batch_size * input_seq_length (not num_tokens): with padding removal the actual token count can be
+    // smaller than batch_size * max(context q length), so sizing by max_num_tokens underestimates.
+    size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * kv_seq_length;
     size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1);
     size_t const rotary_inv_freq_size = sizeof(float) * batch_size * mRotaryEmbeddingDim / 2;
 
@@ -819,7 +822,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3
     size_t const v_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * kv_seq_length * local_hidden_units_kv;
     size_t const qk_buf_size
         = mEnableContextFMHA ? 0 : size * batch_size * mNumHeads * input_seq_length * kv_seq_length;
-    size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * max_num_tokens * local_hidden_units_qo;
+    size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_length * local_hidden_units_qo;
     size_t const qk_buf_float_size
         = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_length * kv_seq_length;
     int dim_q_per_head = (mMLAParams.qk_rope_head_dim + mMLAParams.qk_nope_head_dim);
@@ -894,8 +897,8 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3
         ? sizeof(float) * tc::divUp(local_hidden_units_kv, std::max(1, mSageAttnNumEltsPerBlkV))
         : 0;
 
-    size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
-    size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * max_num_tokens;
+    size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_length;
+    size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * cross_kv_length;
     // Each token holds (batch_idx, token_idx_in_seq) int2.
     size_t const tokens_info_size = sizeof(int2) * max_num_tokens;
     size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0;

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -345,7 +345,8 @@ class RunnerBase
     virtual ~RunnerBase() = default;
     virtual void prepare(AttentionOp& op) const = 0;
     virtual int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size,
-        int const num_gen_tokens, int const max_blocks_per_sequence, int const ctx_total_kv_len = 0) const
+        int const num_gen_tokens, int const max_blocks_per_sequence, int const ctx_total_kv_len = 0,
+        int const maxCrossKvLength = 0) const
         = 0;
     // typically, we use single qkv input, but for context MLA, we use separate qkv inputs
     virtual void run(AttentionOp& op, bool const is_context, int32_t const seq_offset, int32_t const num_seqs,
@@ -410,10 +411,11 @@ class Runner : public RunnerBase
     }
 
     int64_t getWorkspaceSize(AttentionOp const& op, int const num_tokens, int const max_attention_window_size,
-        int const num_gen_tokens, int const max_blocks_per_sequence, int const ctx_total_kv_len = 0) const override
+        int const num_gen_tokens, int const max_blocks_per_sequence, int const ctx_total_kv_len = 0,
+        int const maxCrossKvLength = 0) const override
     {
         size_t const context_workspace_size = op.getWorkspaceSizeForContext(
-            op.mType, max_num_requests, op.mMaxContextLength, 0, num_tokens, ctx_total_kv_len);
+            op.mType, max_num_requests, op.mMaxContextLength, maxCrossKvLength, num_tokens, ctx_total_kv_len);
         size_t const generation_workspace_size = op.getWorkspaceSizeForGeneration(
             op.mType, max_num_requests, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence);
 
@@ -895,6 +897,7 @@ class Runner : public RunnerBase
                 auto const& cross_kv_tensor = cross_kv.value();
                 enqueue_params.cross_kv = static_cast<T const*>(cross_kv_tensor.data_ptr());
                 enqueue_params.num_encoder_tokens = static_cast<int32_t>(cross_kv_tensor.size(0));
+                // Kept in step with maxCrossKvLength in attention(), which sizes the workspace carved here.
                 enqueue_params.cross_kv_length
                     = host_past_key_value_lengths.slice(0, seq_offset, seq_offset + num_seqs).max().item<int32_t>();
             }
@@ -1359,8 +1362,18 @@ void attention(torch::Tensor q, std::optional<torch::Tensor> k, std::optional<to
         = beam_width == 1 ? attention_window_size : cache_indirection.value().size(2);
     int32_t const max_blocks_per_sequence
         = use_kv_cache && kv_cache_block_offsets.has_value() ? kv_cache_block_offsets.value().size(-1) : 0;
-    int64_t const workspace_size = runner->getWorkspaceSize(
-        *op, num_tokens, max_attention_window_size, num_gen_tokens, max_blocks_per_sequence, ctx_total_kv_len);
+    // For cross-attention, several unfused-path context buffers scale with the encoder KV length.
+    // Mirror the context-stage enqueue, which uses the max past-KV length over the context sequences
+    // as cross_kv_length; sizing with 0 here under-allocates the workspace and the carved views in
+    // enqueueContext land past the end of the allocation. The enqueue also gates on cross_kv.has_value(),
+    //  so this can over-allocate relative to the carve; that is safe.
+    int32_t maxCrossKvLength = 0;
+    if (op->isCrossAttention() && num_contexts > 0)
+    {
+        maxCrossKvLength = host_past_key_value_lengths.slice(0, 0, num_contexts).max().item<int32_t>();
+    }
+    int64_t const workspace_size = runner->getWorkspaceSize(*op, num_tokens, max_attention_window_size, num_gen_tokens,
+        max_blocks_per_sequence, ctx_total_kv_len, maxCrossKvLength);
     TLLM_LOG_TRACE("Expected workspace size is %ld bytes", workspace_size);
 
     torch::Tensor workspace;

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -532,21 +532,6 @@ def __init__(
                 "disabled.")
             self.cuda_graph_config = None
 
-        if (self.cuda_graph_config is not None and self.dtype == torch.float32
-                and self._is_encoder_decoder_model()):
-            # fp32 enc-dec runs unfused cross-attention, whose thop workspace
-            # size query hardcodes cross_kv_length=0 (attentionOp.cpp,
-            # Runner::getWorkspaceSize) and undersizes the workspace. The
-            # graph-capture warmup runs cross_attn in isolation, so the carve
-            # overruns the allocation (surfaces as cublas EXECUTION_FAILED).
-            # Keep eager until the upstream size query is fixed.
-            logger.warning(
-                "Decoder CUDA graphs are not supported for float32 "
-                "encoder-decoder models. Decoder CUDA graphs will be disabled; "
-                "use a half-precision checkpoint or "
-                "model_kwargs={'torch_dtype': ...} to enable them.")
-            self.cuda_graph_config = None
-
         cuda_graph_batch_sizes = self.cuda_graph_config.batch_sizes if self.cuda_graph_config else CudaGraphConfig.model_fields[
             'batch_sizes'].default
         cuda_graph_padding_enabled = self.cuda_graph_config.enable_padding if self.cuda_graph_config else CudaGraphConfig.model_fields[

tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py

@@ -506,6 +506,15 @@ def _mixed_batch_test_case(
         exact_match=True,
         feature_id="bf16-kv-v1-decoder-cuda-graph-on-greedy-batch2",
     ),
+    _mixed_batch_test_case(
+        model_name="t5-small",
+        torch_dtype="float32",
+        use_kv_cache_manager_v2=False,
+        num_beams=1,
+        num_return_sequences=1,
+        exact_match=True,
+        feature_id="fp32-kv-v1-decoder-cuda-graph-on-greedy-batch2",
+    ),
     _mixed_batch_test_case(
         model_name="t5-small",
         torch_dtype="bfloat16",
@@ -566,6 +575,19 @@ def _decoder_cuda_graph_config(
     )
 
 
+def _assert_decoder_cuda_graphs_captured(llm: LLM) -> None:
+    """Introspect the in-process engine (single-process mode only).
+
+    Guards against the engine silently declining decoder graphs: without it a
+    workspace-sizing regression that disables capture would still pass the
+    output checks. The enc-dec encoder step stays eager.
+    """
+    model_engine = llm._executor.engine.model_engine
+    assert not model_engine.encoder_cuda_graph_runner.enabled
+    assert model_engine.cuda_graph_runner.enabled
+    assert model_engine.cuda_graph_runner.graphs
+
+
 class _SleepLogitsProcessor:
     def __init__(self, delay_seconds: float) -> None:
         self.delay_seconds = delay_seconds
@@ -751,6 +773,9 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch(
     exact_match: bool,
 ) -> None:
     monkeypatch.setenv("TRTLLM_SKIP_KV_CACHE_ESTIMATION", "1")
+    # Single-process worker so _assert_decoder_cuda_graphs_captured can reach the engine; the
+    # default proxy executor runs it in another process.
+    monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
 
     model_path = _get_t5_model_path(model_name)
     tokenizer = AutoTokenizer.from_pretrained(model_path)
@@ -790,6 +815,7 @@ def test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch(
         )
 
         assert len(responses) == len(_MIXED_ENCODER_SOURCE_TEXTS)
+        _assert_decoder_cuda_graphs_captured(llm)
 
         for request_idx, response in enumerate(responses):
             expected_token_ids = (

tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py

@@ -105,9 +105,8 @@ def _make_llm(
     cuda_graph_batch_sizes: list[int] | None = None,
     tensor_parallel_size: int = 1,
 ) -> LLM:
-    # CudaGraphConfig captures the decode step only; fp32 enc-dec declines
-    # graphs at engine init (workspace-sizing guard), so requesting them must
-    # still work for every dtype.
+    # CudaGraphConfig captures the decode step only; requesting graphs must
+    # work for every dtype.
     cuda_graph_config = (
         CudaGraphConfig(batch_sizes=cuda_graph_batch_sizes, enable_padding=True)
         if cuda_graph_batch_sizes is not None
@@ -259,10 +258,10 @@ def _assert_decoder_cuda_graph_state(llm: LLM, captured: bool) -> None:
 # (torch_dtype override or None for checkpoint fp32, kv manager v2, decoder
 # cuda-graph batch sizes, graphs must capture, TP size). KVCacheManagerV2
 # requires beam width 1, so v2 rides greedy; the fp32+graphs-requested case
-# asserts the engine declines graphs (fp32 enc-dec guard) yet stays exact.
+# covers fp32 enc-dec capturing decoder graphs.
 _FEATURE_COMBINATION_CASES = [
     pytest.param(None, True, None, False, 1, id="fp32-kv-v2-graphs-off-greedy"),
-    pytest.param(None, False, [1, 2], False, 1, id="fp32-kv-v1-graphs-requested-greedy"),
+    pytest.param(None, False, [1, 2], True, 1, id="fp32-kv-v1-graphs-requested-greedy"),
     pytest.param("bfloat16", False, [1, 2], True, 1, id="bf16-kv-v1-decoder-graphs-on-greedy"),
     pytest.param("bfloat16", True, [1, 2], True, 1, id="bf16-kv-v2-decoder-graphs-on-greedy"),
     pytest.param("float16", False, None, False, 1, id="fp16-kv-v1-graphs-off-greedy"),

tests/integration/test_lists/test-db/l0_h100.yml

@@ -341,6 +341,7 @@ l0_h100:
   - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2-t5-small]
   - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v1-decoder-cuda-graph-on-beam2-batch2-flan-t5-small]
   - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[bf16-kv-v2-decoder-cuda-graph-on-greedy-batch2-t5-small]
+  - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch[fp32-kv-v1-decoder-cuda-graph-on-greedy-batch2-t5-small]
   - llmapi/test_llm_api_pytorch_t5.py::test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_beam_search[fp32-kv-v1-graphs-off-beam2]
   - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy]