[#20567][perf] DFlash: snapshot the prefill context page counts at prepare time

tensorrt_llm/_torch/speculative/dflash.py

@@ -38,6 +38,16 @@
 )
 from .interface import SpecMetadata, SpecWorkerBase
 
+# TLLM_DFLASH_PREFILL_CTX_SNAPSHOT=1 (read at import): _store_prefill_context
+# reads the context rows' draft page counts from a host snapshot taken in
+# DFlashSpecMetadata.prepare() -- after attn_metadata.prepare() filled
+# draft_kv_cache_block_offsets, before the target forward is enqueued --
+# instead of a blocking .tolist() of a device tensor at the end of the forward.
+# That read made every context forward wait for its own GPU work, so the
+# executor's host work between chunks could not overlap the next chunk. Same
+# formula on the same tensor, so the values are identical.
+_PREFILL_CTX_SNAPSHOT = os.environ.get("TLLM_DFLASH_PREFILL_CTX_SNAPSHOT", "0") == "1"
+
 # Backends that keep the draft's private context K/V in a paged HND pool.
 # They share the pool layout, page table and paged append op
 _PAGED_ATTENTION_BACKENDS = ("TRTLLM", "FA4")
@@ -629,6 +639,8 @@ def prepare(self):
                 pin_memory=prefer_pinned(),
             )
             worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True)
+            if _PREFILL_CTX_SNAPSHOT:
+                worker._snapshot_ctx_block_counts(num_contexts)
 
     def is_layer_capture(self, layer_id: int) -> bool:
         return layer_id in self._capture_layer_set
@@ -750,6 +762,10 @@ def __init__(
         # Pages the manager handed this request, per batch row. Saturates at
         # the table width on the default backend -- dflash_allocated_ctx_limit.
         self._ctx_block_counts = None
+        # TLLM_DFLASH_PREFILL_CTX_SNAPSHOT: the offsets tensor of the last
+        # refresh and the prepare-time host snapshot of the context rows.
+        self._ctx_offsets_src = None
+        self._ctx_counts_snapshot = None
         self._ctx_pool_idx = 0
         self._ctx_block_divisor = 1  # encoded offset -> raw pool block index
         self._ctx_kv_indptr = None
@@ -1023,6 +1039,7 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool:
         # pool, and the drafter addresses blocks, not planes.
         # clone(): .to() is a no-op for an int64 source, and the in-place ops
         # below would then edit attn_metadata's own tensor.
+        self._ctx_offsets_src = src
         encoded = src[self._ctx_pool_idx, :num_seqs, 0].to(torch.int64).clone()
         # Before the clamp erases the placeholders. Informative only on the
         # python V2 backend; the cpp one already mapped BAD_PAGE_INDEX to 0.
@@ -1031,6 +1048,58 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool:
         self._ctx_block_tables[:num_seqs].copy_(decoded.to(torch.int32))
         return True
 
+    def _snapshot_ctx_block_counts(self, num_contexts: int) -> None:
+        """Stream-ordered host copy of the context rows' draft page counts.
+
+        Called from DFlashSpecMetadata.prepare(), after attn_metadata.prepare()
+        refilled draft_kv_cache_block_offsets on this stream and before the
+        target forward is enqueued. The counts are computed as
+        _refresh_ctx_block_tables does (rows [0, num_contexts) of the same
+        tensor, which nothing writes between here and the refresh) and copied
+        non-blocking into pinned memory behind an event. The offsets tensor is
+        the persistent buffer seen by the previous refresh;
+        _take_ctx_block_counts only accepts the snapshot if the forward's
+        attn_metadata carries that same tensor.
+        """
+        self._ctx_counts_snapshot = None
+        src = self._ctx_offsets_src
+        if (
+            num_contexts <= 0
+            or src is None
+            or self._ctx_block_tables is None
+            or not prefer_pinned()
+            or torch.cuda.is_current_stream_capturing()
+        ):
+            return
+        encoded = src[self._ctx_pool_idx, :num_contexts, 0]
+        counts = (encoded >= 0).sum(dim=1)
+        counts_host = torch.empty(counts.shape, dtype=counts.dtype, pin_memory=prefer_pinned())
+        counts_host.copy_(counts, non_blocking=True)
+        event = torch.cuda.Event()
+        event.record()
+        binding = (self._ctx_kv_manager, self._ctx_pool_idx)
+        self._ctx_counts_snapshot = (src, num_contexts, binding, counts_host, event)
+
+    def _take_ctx_block_counts(self, attn_metadata, num_contexts: int) -> Optional[torch.Tensor]:
+        """The host page counts snapshotted for this forward, or None."""
+        snapshot = self._ctx_counts_snapshot
+        self._ctx_counts_snapshot = None
+        if snapshot is None:
+            return None
+        src, n, binding, counts_host, event = snapshot
+        if (
+            n != num_contexts
+            or src is not getattr(attn_metadata, "draft_kv_cache_block_offsets", None)
+            or src is not self._ctx_offsets_src
+            or binding != (self._ctx_kv_manager, self._ctx_pool_idx)
+            or self._ctx_block_tables is None
+        ):
+            return None
+        # Completes once the GPU has started this forward (the copy sits in
+        # front of it), not after it.
+        event.synchronize()
+        return counts_host
+
     def _lazy_init_ctx_buffers(
         self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None
     ):
@@ -1486,11 +1555,17 @@ def _store_prefill_context(
         # it must fit that row's allocation as well as the arena's length.
         # One tolist() for the whole loop: the lengths themselves come from
         # _ctx_len_host, so the loop body below stays sync-free.
-        ctx_alloc = (
-            (self._ctx_block_counts[:num_contexts] * self._ctx_page_size).tolist()
-            if self._ctx_block_tables is not None
+        counts = (
+            self._take_ctx_block_counts(attn_metadata, num_contexts)
+            if _PREFILL_CTX_SNAPSHOT
             else None
         )
+        if self._ctx_block_tables is None:
+            ctx_alloc = None
+        elif counts is not None:
+            ctx_alloc = (counts * self._ctx_page_size).tolist()
+        else:
+            ctx_alloc = (self._ctx_block_counts[:num_contexts] * self._ctx_page_size).tolist()
         for i in range(num_contexts):
             req_id = spec_metadata.request_ids[i]
             slen = int(attn_metadata._seq_lens[i])

tests/integration/test_lists/test-db/l0_b200.yml

@@ -251,6 +251,7 @@ l0_b200:
   # stream/event, no SM gating), covered on A30 pre-merge; drop that whole
   # function's leaves on Blackwell, keep the rest of the sampler dir.
   - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor"
+  - unittest/_torch/speculative/test_dflash_prefill_ctx_snapshot.py
   - unittest/_torch/speculative/test_eagle3.py
   - unittest/_torch/speculative/test_lilicorr_loading.py
   - unittest/_torch/speculative/test_dspark_cute_dsl_attention.py

tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py

@@ -27,6 +27,7 @@
 import pytest
 import torch
 
+from tensorrt_llm._torch.speculative import dflash
 from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker
 
 pytestmark = pytest.mark.skipif(
@@ -48,8 +49,13 @@ def _worker(lengths):
         _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"),
         _ctx_len_host=list(lengths),
         _batch_to_slot=torch.zeros(8, dtype=torch.long, device="cuda"),
+        # TLLM_DFLASH_PREFILL_CTX_SNAPSHOT state: not bound to the draft pool.
+        _ctx_offsets_src=None,
+        _ctx_counts_snapshot=None,
+        _ctx_block_tables=None,
     )
     w._write_ctx_len = partial(DFlashWorker._write_ctx_len, w)
+    w._snapshot_ctx_block_counts = partial(DFlashWorker._snapshot_ctx_block_counts, w)
     w._advance_ctx_len = partial(DFlashWorker._advance_ctx_len, w)
     w._assign_slot = lambda rid, *a, **k: None
     return w
@@ -70,7 +76,14 @@ def _cuda(values):
     return torch.tensor(values, dtype=torch.long, device="cuda")
 
 
-def test_padded_steps_keep_the_dummy_slot_empty():
+@pytest.fixture(params=[False, True], ids=["ctx-snapshot-off", "ctx-snapshot-on"])
+def prefill_ctx_snapshot(request, monkeypatch):
+    """prepare() with and without TLLM_DFLASH_PREFILL_CTX_SNAPSHOT."""
+    monkeypatch.setattr(dflash, "_PREFILL_CTX_SNAPSHOT", request.param)
+    return request.param
+
+
+def test_padded_steps_keep_the_dummy_slot_empty(prefill_ctx_snapshot):
     worker = _worker([40, 300, 7, 0, 0])
     # Three requests in a graph of five.
     _prepare(worker, [11, 12, 13, FLOOR + 1, FLOOR + 1], num_generations=5)
@@ -102,7 +115,7 @@ def test_dummy_slot_reset_replays_in_a_cuda_graph():
         assert worker._ctx_len.tolist() == [40 + 2 * step, 300 + 3 * step, 7, 0, 0], step
 
 
-def test_prepare_resets_only_evicted_slots():
+def test_prepare_resets_only_evicted_slots(prefill_ctx_snapshot):
     worker = _worker([40, 300, 7, 0, 96])
     written = []
     write = worker._write_ctx_len
@@ -117,3 +130,5 @@ def record(updates):
     assert written == [{}, {1: 0}]
     assert worker._ctx_len.tolist() == [40, 0, 7, 0, 96]
     assert 1 in worker._free_slots and 12 not in worker._req_to_slot
+    # Generation-only steps take no snapshot.
+    assert worker._ctx_counts_snapshot is None

tests/unittest/_torch/speculative/test_dflash_prefill_ctx_snapshot.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""DFlash prefill: the prepare-time host snapshot of the draft page counts."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.speculative.dflash import DFlashWorker
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+NUM_SLOTS, MAX_BLOCKS = 9, 16
+
+
+def _worker(divisor: int = 2) -> DFlashWorker:
+    worker = object.__new__(DFlashWorker)
+    worker._ctx_block_tables = torch.zeros(NUM_SLOTS, MAX_BLOCKS, dtype=torch.int32, device="cuda")
+    worker._ctx_block_counts = torch.zeros(NUM_SLOTS, dtype=torch.long, device="cuda")
+    worker._ctx_pool_idx = 1
+    worker._ctx_block_divisor = divisor
+    worker._ctx_kv_manager = object()
+    worker._ctx_offsets_src = None
+    worker._ctx_counts_snapshot = None
+    return worker
+
+
+def _fill(offsets: torch.Tensor, counts: list[int], divisor: int) -> None:
+    """Rows of pool 1 with counts[i] encoded pages, the rest placeholders (-1)."""
+    offsets.fill_(-1)
+    for row, count in enumerate(counts):
+        pages = torch.arange(count, dtype=torch.int32, device="cuda") * divisor + 7 * row
+        offsets[1, row, :, :count] = pages
+
+
+def test_snapshot_matches_refreshed_counts() -> None:
+    worker = _worker()
+    offsets = torch.empty(2, NUM_SLOTS, 2, MAX_BLOCKS, dtype=torch.int32, device="cuda")
+    meta = SimpleNamespace(draft_kv_cache_block_offsets=offsets)
+    # No refresh seen yet: nothing to snapshot.
+    worker._snapshot_ctx_block_counts(3)
+    assert worker._ctx_counts_snapshot is None
+    _fill(offsets, [1, 2, 3, 4], 2)
+    assert worker._refresh_ctx_block_tables(meta, 4)
+    for step, counts in enumerate(([5, 0, 16], [16, 3, 1, 9, 2], [2])):
+        # The next iteration's attn_metadata.prepare() refills the same tensor.
+        _fill(offsets, counts, 2)
+        n = len(counts)
+        worker._snapshot_ctx_block_counts(n)
+        # The forward runs, then refreshes the table from the same tensor.
+        torch.cuda._sleep(1_000_000)
+        assert worker._refresh_ctx_block_tables(meta, n)
+        snap = worker._take_ctx_block_counts(meta, n)
+        assert snap is not None and not snap.is_cuda, step
+        assert snap.tolist() == counts == worker._ctx_block_counts[:n].tolist()
+        # Consumed once.
+        assert worker._take_ctx_block_counts(meta, n) is None
+
+
+@pytest.mark.parametrize("change", ["num_contexts", "tensor", "manager", "pool"])
+def test_snapshot_is_rejected_when_the_binding_changed(change: str) -> None:
+    worker = _worker()
+    offsets = torch.empty(2, NUM_SLOTS, 2, MAX_BLOCKS, dtype=torch.int32, device="cuda")
+    meta = SimpleNamespace(draft_kv_cache_block_offsets=offsets)
+    _fill(offsets, [3, 1], 2)
+    worker._refresh_ctx_block_tables(meta, 2)
+    worker._snapshot_ctx_block_counts(2)
+    n = 2
+    if change == "num_contexts":
+        n = 1
+    elif change == "tensor":
+        meta = SimpleNamespace(draft_kv_cache_block_offsets=offsets.clone())
+    elif change == "manager":
+        worker._ctx_kv_manager = object()
+    else:
+        worker._ctx_pool_idx = 0
+    assert worker._take_ctx_block_counts(meta, n) is None
+
+
+def test_no_snapshot_under_graph_capture() -> None:
+    worker = _worker()
+    offsets = torch.empty(2, NUM_SLOTS, 2, MAX_BLOCKS, dtype=torch.int32, device="cuda")
+    meta = SimpleNamespace(draft_kv_cache_block_offsets=offsets)
+    _fill(offsets, [3, 1], 2)
+    worker._refresh_ctx_block_tables(meta, 2)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        worker._snapshot_ctx_block_counts(2)
+    assert worker._ctx_counts_snapshot is None