From 6a8c3a83ff6080d65ae64073dceb1700b1aa649f Mon Sep 17 00:00:00 2001 From: Viraat Chandra Date: Sun, 11 Oct 2026 20:27:01 -0700 Subject: [PATCH 1/2] [None][perf] DFlash VANILLA draft attention: pass explicit FA2 num_splits The VANILLA DFlash drafter attends through flash_attn_with_kvcache over a dense per-slot context arena whose seqlen dimension is min(max_seq_len, max_position_embeddings). FA2's split-KV heuristic sizes num_splits from k_cache.size(1), i.e. the arena capacity, not the live context length. With a 262k arena and ~38k live tokens it picks ~4-6 splits of ~44-65k tokens, so only split 0 has work: ~70-100 CTAs on a 212-SM GPU, each streaming the whole context serially, in every drafter layer on every step. Pass num_splits=128 explicitly. It does not depend on the live length, so it is CUDA-graph safe; splits past the live length exit early and the combine kernel is a few microseconds. Measured on Kimi-K3 + DSpark drafter (draft 7), Rubin VR200 (sm107, 212 SMs), disaggregated 1P1D TP8/TP8 C16 with full-prompt drafter context (#19733): per-call drafter attention 612 us -> ~150 us, gen step -1.2 ms at batch 8 and -2.3 ms at batch 10-12, output tok/s/user +6.3%, acceptance length unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Viraat Chandra --- tensorrt_llm/_torch/models/modeling_dflash.py | 5 +++++ 1 file changed, 5 insertions(+) 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]. From adf62c6765ae3c02afe057186b21aece8c65d24c Mon Sep 17 00:00:00 2001 From: Viraat Chandra Date: Sun, 11 Oct 2026 20:34:47 -0700 Subject: [PATCH 2/2] [None][perf] DFlash unified draft cache: avoid per-step host syncs on index H2D snapshot_managed_draft_history and the managed-history / block-count refreshes build small index tensors with torch.tensor(list, device=cuda). That is a pageable H2D copy, which synchronizes the host with the stream on every step and defeats the overlap scheduler: the next step's prepare can no longer run while the current forward is on the GPU. Build them in pinned host memory and copy non_blocking, as the rest of dflash.py already does (prefer_pinned()). Measured on Kimi-K3 + DSpark drafter (draft 7), Rubin VR200, disaggregated 1P1D TP8/TP8 C16 with full-prompt drafter context, on top of the FA2 num_splits fix: nsys showed a 26 ms cudaStreamSynchronize per step (the whole forward) right after the graph launch, and GPU idle of 6.0 vs 1.5 ms/step without drafter context. With this change the gen step drops 2.2-2.5 ms at equal batch, output tok/s/user +15% and total tok/s/GPU +6%. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Viraat Chandra --- tensorrt_llm/_torch/speculative/dflash.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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)