From 3adec0d505b13b67cb1b3b0abfae0f8b01937751 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Tue, 14 Jul 2026 09:40:07 -0700 Subject: [PATCH 1/2] [None][fix] Size the attention context workspace with the real cross-KV length Runner::getWorkspaceSize passed cross_kv_length=0, so the unfused-path buffers were sized at zero while enqueueContext carved them at the real length. The carved views overrun the allocation, surfacing as CUBLAS_STATUS_EXECUTION_FAILED at the QK^T GEMM. FP32 enc-dec hits it, since FP32 takes the unfused path. Pass the context sequences' max past-KV length through the size query. Every changed buffer is !mEnableContextFMHA-gated, so fused FP16/BF16 is unchanged. Also drops the model_engine guard that disabled fp32 enc-dec CUDA graphs pending this fix, and flips the Whisper case that pinned it. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 11 ++++++---- cpp/tensorrt_llm/thop/attentionOp.cpp | 21 ++++++++++++++----- .../_torch/pyexecutor/model_engine.py | 15 ------------- .../defs/llmapi/test_llm_api_pytorch_t5.py | 9 ++++++++ .../llmapi/test_llm_api_pytorch_whisper.py | 9 ++++---- .../test_lists/test-db/l0_h100.yml | 1 + 6 files changed, 37 insertions(+), 29 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 762ed7b54164..db09ff48e618 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -777,7 +777,10 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 auto const batch_size = static_cast(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; @@ -797,7 +800,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); @@ -865,8 +868,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; diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 811810a8d5a2..d767ea9ad62f 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -343,7 +343,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 max_cross_kv_length = 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, @@ -408,10 +409,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 max_cross_kv_length = 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, max_cross_kv_length, 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); @@ -1347,8 +1349,17 @@ void attention(torch::Tensor q, std::optional k, std::optionalgetWorkspaceSize( - *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. + int32_t max_cross_kv_length = 0; + if (op->isCrossAttention() && num_contexts > 0) + { + max_cross_kv_length = host_past_key_value_lengths.slice(0, 0, num_contexts).max().item(); + } + 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, max_cross_kv_length); TLLM_LOG_TRACE("Expected workspace size is %ld bytes", workspace_size); torch::Tensor workspace; diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5afc55544c8c..e3703f359571 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -466,21 +466,6 @@ def __init__( "decoder CUDA graphs. CUDA graphs will be 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( - "CUDA graphs are not supported for float32 encoder-decoder " - "models. 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[ diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py index 492b88ff3d26..3d83b9b0ac29 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -505,6 +505,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", diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py index bb72ccb7121a..533cfb269ca6 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_whisper.py +++ b/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"), diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index f4a3fcefef03..5d35069485ac 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -343,6 +343,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_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] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v1-graphs-requested-greedy] From 56cb39cf55c50d5828233b8c865169553b940651 Mon Sep 17 00:00:00 2001 From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:02:19 -0700 Subject: [PATCH 2/2] [https://nvbugs/6571220][fix] Address review on attention context workspace sizing Rename the new workspace-length parameter to lower camelCase per the C++ coding guidelines, and cross-reference the two sites that derive the context-stage cross-KV length so the size query and the enqueue carve stay in agreement. Assert the T5 mixed-encoder-length cases actually capture decoder CUDA graphs, which needs the single-process worker to reach the engine. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/attentionOp.cpp | 16 +++++++++------- .../defs/llmapi/test_llm_api_pytorch_t5.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 35697751b296..3e36b5566a8d 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -346,7 +346,7 @@ class RunnerBase 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, - int const max_cross_kv_length = 0) const + 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, @@ -412,10 +412,10 @@ 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, - int const max_cross_kv_length = 0) const override + int const maxCrossKvLength = 0) const override { size_t const context_workspace_size = op.getWorkspaceSizeForContext( - op.mType, max_num_requests, op.mMaxContextLength, max_cross_kv_length, 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); @@ -897,6 +897,7 @@ class Runner : public RunnerBase auto const& cross_kv_tensor = cross_kv.value(); enqueue_params.cross_kv = static_cast(cross_kv_tensor.data_ptr()); enqueue_params.num_encoder_tokens = static_cast(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(); } @@ -1365,14 +1366,15 @@ void attention(torch::Tensor q, std::optional k, std::optionalisCrossAttention() && num_contexts > 0) { - max_cross_kv_length = host_past_key_value_lengths.slice(0, 0, num_contexts).max().item(); + maxCrossKvLength = host_past_key_value_lengths.slice(0, 0, num_contexts).max().item(); } 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, max_cross_kv_length); + max_blocks_per_sequence, ctx_total_kv_len, maxCrossKvLength); TLLM_LOG_TRACE("Expected workspace size is %ld bytes", workspace_size); torch::Tensor workspace; diff --git a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py index 3d83b9b0ac29..22e141611308 100644 --- a/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py +++ b/tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py @@ -574,6 +574,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 @@ -759,6 +772,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) @@ -798,6 +814,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 = (