diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 53dcfbc5af2a..09ba91d52c4e 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -1743,14 +1743,13 @@ def is_chunked_prefill_for_mla_context(
                 and metadata.num_ctx_cached_tokens > 0
                 and metadata.runtime_features.chunked_prefill)
 
-    def is_chunked_prefill_mla_context_for_warmup(
+    def has_cached_kv_for_mla_context_warmup(
         self,
         metadata: TrtllmAttentionMetadata,
     ) -> bool:
-        """Chunked prefill MLA context check for warmup; does not check num_ctx_cached_tokens."""
+        """KV cache reuse / chunked prefill MLA context check for warmup, do not check num_ctx_cached_tokens."""
         return (self.is_mla_enable and metadata.kv_cache_manager is not None
-                and metadata.enable_context_mla_with_cached_kv
-                and metadata.runtime_features.chunked_prefill)
+                and metadata.enable_context_mla_with_cached_kv)
 
     def load_paged_kv_cache_for_mla(
         self,
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 616724550497..4d2cd971b4af 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/tensorrt_llm/_torch/modules/attention.py
@@ -2359,9 +2359,11 @@ def forward_context_with_chunked_prefill(
 
     @staticmethod
     @functools.cache
-    def cached_warmup_forward_context_with_chunked_prefill(
-            num_heads_tp, qk_nope_head_dim, qk_rope_head_dim, kv_lora_rank,
-            v_head_dim, dtype, device):
+    def cached_warmup_forward_context_with_cached_kv(num_heads_tp,
+                                                     qk_nope_head_dim,
+                                                     qk_rope_head_dim,
+                                                     kv_lora_rank, v_head_dim,
+                                                     dtype, device):
         """Warmup torch.compile for cat operations with different tensor layouts.
 
         Tensors are marked with torch._dynamo.maybe_mark_dynamic(..., 0) on the
@@ -2430,9 +2432,13 @@ def forward_context(
         if isinstance(self.mha, TrtllmAttention):
             assert isinstance(attn_metadata, TrtllmAttentionMetadata)
             trtllm_attention = cast(TrtllmAttention, self.mha)
-            if trtllm_attention.is_chunked_prefill_mla_context_for_warmup(
+            # Warm up maybe_compiled_cat for both the chunked-prefill path and
+            # the cached-kv path; without this, the cached-kv prefill
+            # (block-reuse without chunked_prefill) recompiles per shape and
+            # may stall inside inductor's compile worker.
+            if trtllm_attention.has_cached_kv_for_mla_context_warmup(
                     attn_metadata):
-                self.cached_warmup_forward_context_with_chunked_prefill(
+                self.cached_warmup_forward_context_with_cached_kv(
                     self.num_heads_tp, self.qk_nope_head_dim,
                     self.qk_rope_head_dim, self.kv_lora_rank, self.v_head_dim,
                     q.dtype, q.device)