diff --git a/tensorrt_llm/_torch/models/modeling_dflash.py b/tensorrt_llm/_torch/models/modeling_dflash.py
index b9693465b6fe..778530fb8056 100644
--- a/tensorrt_llm/_torch/models/modeling_dflash.py
+++ b/tensorrt_llm/_torch/models/modeling_dflash.py
@@ -1838,6 +1838,11 @@ def dflash_forward(
                     cache_batch_idx=cache_batch_idx_i32,
                     causal=causal,
                     window_size=window_size,
+                    # FA2's default heuristic sizes splits from k_cache.size(1),
+                    # the max_seq_len-dense arena, not the live context, so long
+                    # contexts run on split 0 with a few CTAs. A fixed count is
+                    # CUDA-graph safe; splits past the live length exit early.
+                    num_splits=128,
                 )
                 if pack_gqa:
                     # Undo the packing: [B, group*blk, h_kv, d] -> [B, blk, h_q, d].
diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py
index ae37141bfb3e..84c8194b0f92 100644
--- a/tensorrt_llm/_torch/speculative/dflash.py
+++ b/tensorrt_llm/_torch/speculative/dflash.py
@@ -844,8 +844,8 @@ def _prepare_managed_history(self, request_ids: list[int]) -> list[int]:
             torch.tensor(
                 [self._req_to_slot.get(request_id, self._dummy_slot) for request_id in request_ids],
                 dtype=torch.long,
-                device=self._batch_to_slot.device,
-            )
+                pin_memory=prefer_pinned(),
+            ).to(self._batch_to_slot.device, non_blocking=True)
         )
         self._managed_request_ids = tuple(
             request_id for request_id in request_ids if request_id in real_request_ids
@@ -891,8 +891,8 @@ def snapshot_managed_draft_history(self) -> DraftHistoryUpdate | None:
         slots = torch.tensor(
             [self._req_to_slot[request_id] for request_id in self._managed_request_ids],
             dtype=torch.long,
-            device=self._ctx_len.device,
-        )
+            pin_memory=prefer_pinned(),
+        ).to(self._ctx_len.device, non_blocking=True)
         lengths = self._ctx_len[slots]
         positions = lengths + self._ctx_position_offset[slots]
         return DraftHistoryUpdate.capture(
@@ -1091,8 +1091,8 @@ def _refresh_ctx_block_tables(
                         for rid in request_ids[:num_seqs]
                     ],
                     dtype=torch.long,
-                    device=self._ctx_block_counts.device,
-                )
+                    pin_memory=prefer_pinned(),
+                ).to(self._ctx_block_counts.device, non_blocking=True)
             )
             return True
         src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None)