[https://nvbugs/5580099][fix] Cherry pick IMA issue fix from release/1.1

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -577,6 +577,7 @@ def is_nvfp4_output_kernel_available(
 @dataclass(kw_only=True)
 class TrtllmAttentionMetadata(AttentionMetadata):
     workspace: Optional[torch.Tensor] = None
+    cuda_graph_workspace: Optional[torch.Tensor] = None
 
     # TrtllmAttention needs to know the beam width to access to the cache indirection buffer,
     # when beam search is enabled.
@@ -692,6 +693,14 @@ def _post_init_with_buffers(self, buffers) -> None:
                 device='cuda',
                 dtype=torch.int8,
             )
+
+        if self.cuda_graph_workspace is None:
+            self.cuda_graph_workspace = torch.empty(
+                (0, ),
+                device='cuda',
+                dtype=torch.int8,
+            )
+
         if self.kv_cache_manager is not None:
             self.kv_cache_block_offsets = self.get_empty(
                 buffers,
@@ -1335,8 +1344,9 @@ def forward(
             host_kv_cache_pool_pointers=metadata.host_kv_cache_pool_pointers,
             host_kv_cache_pool_mapping=metadata.host_kv_cache_pool_mapping,
             block_ids_per_seq=metadata.block_ids_per_seq,
-            workspace=metadata.
-            workspace,  # re-enable it, if pass None to it, fp8 mla will encounter invalid cuda free issue.
+            # re-enable it, if pass None to it, fp8 mla will encounter invalid cuda free issue.
+            workspace=metadata.workspace
+            if not metadata.is_cuda_graph else metadata.cuda_graph_workspace,
             cache_indirection=metadata.cache_indirection,
             kv_scale_orig_quant=self.kv_scale_orig_quant,
             kv_scale_quant_orig=self.kv_scale_quant_orig,

tests/integration/defs/test_e2e.py

@@ -148,7 +148,12 @@ def _check_mem_usage(file, mem_info, ranks_num=1):
         f"Running memory information: peak mem {peak}, model mem {model_size}, kv mem {kv_mem_size}, extra {extra}, total {min_total}, activation {activation_memory}, tmp_kv {tmp_kv}, fraction  {fraction}, none-torch memory at starttime {start_time_mem}"
     )
 
-    assert peak - tmp_kv <= e_peak + start_time_mem + delta, f"peak memory {peak} is larger than expected {e_peak}"
+    increased_peak_mem = peak - tmp_kv - e_peak - start_time_mem - delta
+    assert increased_peak_mem <= 0, (
+        f"increased peak memory {increased_peak_mem} is larger than 0,"
+        f" which is calculated as peak ({peak}) - tmp_kv ({tmp_kv}) -"
+        f" e_peak ({e_peak}) - start_time_mem ({start_time_mem}) - delta ({delta})."
+    )
     assert kv_mem_size >= e_kv_mem_size - delta, f"kv memory size {kv_mem_size} is smaller than expected {e_kv_mem_size}"
     # assert model_size <= e_model_size + delta, f"model memory {model_size} is larger than expected {e_model_size}"
     # assert max(extra) <= e_extra + delta, f"extra memory size {extra} is larger than expected {e_extra}"
@@ -1950,7 +1955,7 @@ def test_ptp_quickstart_advanced_mtp(llm_root, llm_venv, model_name,
                 "--use_one_model",
             ],
             stdout=running_log)
-        _check_mem_usage(running_log, [54.60, 0, 0, 0])
+        _check_mem_usage(running_log, [54.90, 0, 0, 0])
 
 
 @pytest.mark.parametrize("model_name,model_path", [