[TRTLLM-16722][feat] Dflash Unified KVCache

tensorrt_llm/_torch/disaggregation/native/peer.py

@@ -180,6 +180,14 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool:
             peer_ri.page_table,
         )
         AttentionPolicy.validate_peer_compatible(self_page_table, peer_ri.page_table)
+        if self._ri.num_draft_layers != peer_ri.num_draft_layers:
+            # The layer-count tolerance below would transfer only the target
+            # layers and leave the drafter's prompt KV unset on the receiver.
+            raise ValueError(
+                "PeerRegistrar: unified draft KV layer count differs "
+                f"(local={self._ri.num_draft_layers}, peer={peer_ri.num_draft_layers}); "
+                "both peers must use the same speculative decoding KV cache layout."
+            )
 
         self_layers = sum(self._ri.layer_num_per_pp)
         peer_layers = sum(peer_ri.layer_num_per_pp)

tensorrt_llm/_torch/disaggregation/native/rank_info.py

@@ -49,6 +49,9 @@ class RankInfo:
     aux_meta: Optional[AuxBufferMeta] = None
     page_table: Optional[KVCachePageTable] = None
     bounded_replay_on_generation: bool = False
+    # Layers a unified KV cache manager owns for the drafter (layer specs with
+    # a DRAFT owner). They ride the target's pages, so both peers must agree.
+    num_draft_layers: int = 0
 
     @property
     def tp_size_per_dp_group(self) -> int:
@@ -114,6 +117,9 @@ def from_kv_cache_manager(
             aux_meta=aux_buffer_meta,
             page_table=build_page_table_from_manager(kvm),
             bounded_replay_on_generation=getattr(kvm, "bounded_replay_on_generation", False),
+            num_draft_layers=(
+                len(kvm.draft_layer_ids) if getattr(kvm, "layer_specs", None) is not None else 0
+            ),
         )
 
     @classmethod

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py

@@ -1018,9 +1018,10 @@ def _try_schedule_disagg_gen_init(
         Returns ``(action, tokens)``.  *tokens* is 0 because disagg requests
         don't participate in the forward pass token budget.
         """
-        # Cache-transceiver mode disables the separate one-model draft manager,
-        # so disagg generation init has no paired-reuse path. Supporting one
-        # would also require draft KV transfer and history_length=prompt_len.
+        # Only the target manager receives prompt KV. Drafters that need it
+        # under disagg keep their layers in that manager (MTP/Eagle fold them
+        # in; DFlash uses explicit draft layer specs). A separate draft manager
+        # gets no transfer and no paired-reuse path here.
         if not self.kv_cache_manager.prepare_disagg_gen_init(req):
             logger.debug("prepare_disagg_gen_init failed for request %s", req.py_request_id)
             return ScheduleAction.SKIP, 0

tensorrt_llm/_torch/speculative/dflash.py

@@ -31,6 +31,7 @@
 
 from ..attention.backends import AttentionMetadata
 from ..modules.logits_processor import LogitsProcessor
+from ..pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role
 from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager
 from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID
 from ..pyexecutor.resource_manager import BaseResourceManager
@@ -615,19 +616,30 @@ def prepare(self):
             # prefilling, so _store_prefill_context -- the only other assigner --
             # never runs and every request would land on the dummy slot.
             num_contexts = max(0, len(self.request_ids) - self.num_generations)
-            for rid in self.request_ids[num_contexts:]:
+            assigned = {}
+            for request_idx, rid in enumerate(self.request_ids[num_contexts:], num_contexts):
                 if (
                     rid != ATTENTION_DP_DUMMY_REQUEST_ID
                     and rid < worker._graph_dummy_id_floor
                     and rid not in worker._req_to_slot
                 ):
-                    if worker._assign_slot(rid) is None:
+                    slot = worker._assign_slot(rid, updates=assigned)
+                    if slot is None:
                         # Not fatal -- spec dec is lossless, so the dummy-slot
                         # fallback only costs acceptance. Warn so it is not silent.
                         logger.warning(
                             f"DFlash: no free drafter slot for generation request {rid}; "
                             "it shares the dummy slot and will draft poorly."
                         )
+                    elif worker._ctx_unified_cache:
+                        # New requests have no in-flight predecessor. Subsequent
+                        # accepted tokens advance the device length during replay.
+                        token_offset = sum(self.seq_lens[:request_idx])
+                        assigned[slot] = min(
+                            int(self.host_position_ids[token_offset]), worker._max_ctx
+                        )
+            if assigned:
+                worker._write_ctx_len(assigned)
 
             # Route unknown request IDs (cuda graph padding or warmup dummies)
             # to dummy slot to avoid corrupting real request's context
@@ -794,6 +806,9 @@ def __init__(
         # the table width on the default backend -- dflash_allocated_ctx_limit.
         self._ctx_block_counts = None
         self._ctx_pool_idx = 0
+        self._ctx_layer_ids = None
+        self._ctx_unified_cache = False
+        self._ctx_block_offset = 0
         self._ctx_block_divisor = 1  # encoded offset -> raw pool block index
         self._ctx_kv_indptr = None
         self._ctx_kv_last_page_len = None
@@ -1083,7 +1098,18 @@ def _managed_ctx_pool(self, draft_kv_cache_manager, L, nkv, hd, dtype, kv_factor
             # it: _util.py skips that bail for an external drafter, which is
             # what DSpark is, so this drafter does get a manager under DP.
             return None
-        layers = [draft_kv_cache_manager.get_buffers(i, kv_layout="HND") for i in range(L)]
+        draft_layer_ids = tuple(getattr(draft_kv_cache_manager, "draft_layer_ids", ()))
+        layer_ids = draft_layer_ids if draft_layer_ids else tuple(range(L))
+        if len(layer_ids) != L:
+            logger.warning(
+                f"DFlash: KV cache manager exposes {len(layer_ids)} draft layers "
+                f"but the drafter has {L}; using the private ctx arena."
+            )
+            return None
+        self._ctx_layer_ids = layer_ids
+        layers = [
+            draft_kv_cache_manager.get_buffers(layer_id, kv_layout="HND") for layer_id in layer_ids
+        ]
         base = layers[0]
         # kv_factor 1 is the MLA drafter's SELFKONLY pool: one latent, no V.
         expected = (kv_factor, nkv, base.size(-2), hd)
@@ -1116,10 +1142,40 @@ def _init_ctx_block_tables(
         """
         mgr = draft_kv_cache_manager
         pool_mapping = mgr.kv_cache_pool_mapping
-        layer_offset = mgr.layer_offsets[0]
+        first_layer_id = self._ctx_layer_ids[0] if self._ctx_layer_ids is not None else 0
+        layer_offset = mgr.layer_offsets[first_layer_id]
         self._ctx_pool_idx = int(pool_mapping[layer_offset, 0])
+        per_layer_tables = bool(getattr(mgr, "_use_per_layer_page_tables", False))
         num_pool_layers = int((pool_mapping[:, 0] == self._ctx_pool_idx).sum())
-        if num_pool_layers != L:
+        self._ctx_block_offset = 0
+        if isinstance(mgr, KVCacheManagerV2):
+            converter = mgr.impl.get_page_index_converter(layer_offset, Role.KEY)
+            layer_group_id = mgr.layer_to_pool_mapping_dict[layer_offset]
+            if mgr.enable_swa_scratch_reuse:
+                raise NotImplementedError(
+                    "DFlash managed context does not support SWA scratch pages"
+                )
+            for layer_id in self._ctx_layer_ids:
+                local_id = mgr.layer_offsets[layer_id]
+                if (
+                    mgr.layer_to_pool_mapping_dict[local_id] != layer_group_id
+                    or mgr.get_layer_page_index_scale(layer_id) != converter.scale
+                ):
+                    if self._ctx_unified_cache:
+                        raise NotImplementedError(
+                            "DFlash's shared context block table requires draft layers "
+                            "with the same lifecycle and page-index scale"
+                        )
+                    logger.warning(
+                        "DFlash: draft KV layers have different lifecycles or page-index "
+                        "scales; using the private ctx arena."
+                    )
+                    return False
+            if per_layer_tables:
+                # Attention tables address the physical pool origin, whereas
+                # get_buffers() starts at this layer's own K buffer.
+                self._ctx_block_offset = int(converter.layer_offset)
+        elif num_pool_layers != L:
             logger.warning(
                 f"DFlash: draft KV pool {self._ctx_pool_idx} spans {num_pool_layers} layers but "
                 f"the drafter has {L}; using the private ctx arena."
@@ -1128,12 +1184,23 @@ def _init_ctx_block_tables(
         # V1 slices a page-major pool, so stride(0) spans every layer; V2 wraps
         # the layer base in a dense view over interleaved pages, so stride(0) is
         # one layer slot. Guessing wrong is silent: acceptance quietly drops.
-        layer_block_elems = mgr.kv_factor * nkv * page_size * hd
+        manager_kv_factor = (
+            mgr.get_layer_kv_factor(first_layer_id)
+            if hasattr(mgr, "get_layer_kv_factor")
+            else mgr.kv_factor
+        )
+        if manager_kv_factor != getattr(self, "_ctx_kv_factor", manager_kv_factor):
+            logger.warning(
+                "DFlash: draft KV layer factor does not match the drafter; "
+                "using the private ctx arena."
+            )
+            return False
+        layer_block_elems = manager_kv_factor * nkv * page_size * hd
         stride0 = pool[0].stride(0)
         if stride0 == layer_block_elems:
-            self._ctx_block_divisor = mgr.kv_factor
+            self._ctx_block_divisor = manager_kv_factor
         elif stride0 == num_pool_layers * layer_block_elems:
-            self._ctx_block_divisor = num_pool_layers * mgr.kv_factor
+            self._ctx_block_divisor = num_pool_layers * manager_kv_factor
         else:
             logger.warning(
                 f"DFlash: draft KV pool layer stride {stride0} matches neither V1's "
@@ -1166,14 +1233,17 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool:
         """
         if self._ctx_block_tables is None or num_seqs <= 0:
             return False
-        src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None)
+        offsets_name = (
+            "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets"
+        )
+        src = getattr(attn_metadata, offsets_name, None)
         if src is None:
             # The table is bound but unfillable. Leaving it at zeros would send
             # every request's every layer to pool block 0 -- silently.
             raise RuntimeError(
                 "DFlash: bound to the draft KV pool but "
                 f"{type(attn_metadata).__name__} carries no "
-                "draft_kv_cache_block_offsets this iteration."
+                f"{offsets_name} this iteration."
             )
         # Field 0 of the K/V pair: both name the same block for a kv_factor=2
         # pool, and the drafter addresses blocks, not planes. The block count
@@ -1185,19 +1255,26 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool:
             self._ctx_block_divisor,
             self._ctx_block_tables,
             self._ctx_block_counts,
+            block_offset=self._ctx_block_offset,
         )
         return True
 
     def _lazy_init_ctx_buffers(
         self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None
     ):
         self._validate_draft_attention_backend(draft_model)
-        # Only TrtllmAttentionMetadata builds draft_kv_cache_block_offsets.
-        # Without it the table stays all-zero and every request would read and
-        # write pool block 0 -- silently, at acceptance 1.0.
-        if draft_kv_cache_manager is not None and not hasattr(
-            attn_metadata, "draft_kv_cache_block_offsets"
-        ):
+        self._ctx_unified_cache = (
+            draft_kv_cache_manager is not None
+            and draft_kv_cache_manager is getattr(attn_metadata, "kv_cache_manager", None)
+        )
+        offsets_name = (
+            "kv_cache_block_offsets" if self._ctx_unified_cache else "draft_kv_cache_block_offsets"
+        )
+        if draft_kv_cache_manager is not None and not hasattr(attn_metadata, offsets_name):
+            if self._ctx_unified_cache:
+                raise NotImplementedError(
+                    f"DFlash unified KV requires {offsets_name} in attention metadata"
+                )
             logger.warning(
                 f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; "
                 "using the private ctx arena."
@@ -1306,6 +1383,7 @@ def _lazy_init_ctx_buffers(
         nkv = draft_model._num_kv_heads
         hd = draft_model._head_dim
         kv_factor = getattr(draft_model, "_kv_factor", 2)
+        self._ctx_kv_factor = kv_factor
         capacity = self._max_ctx + self._compute_block_size
         # Published before any drafter table is built, so a drafter may read it
         # instead of its config cap. _max_ctx is attn_metadata.max_seq_len, which
@@ -1361,6 +1439,8 @@ def _lazy_init_ctx_buffers(
             ):
                 pool = None
             if pool is None:
+                if self._ctx_unified_cache:
+                    raise ValueError("DFlash could not bind the unified manager's draft KV layers")
                 # Per-slot reservation is an arena-only notion; the pool path
                 # takes its pages from the manager one iteration at a time.
                 self._ctx_pages_per_slot = (capacity + page_size - 1) // page_size
@@ -1399,7 +1479,7 @@ def _lazy_init_ctx_buffers(
                 # Only a manager that publishes and matches its own blocks can
                 # hand back a reused prefix's drafter K/V; the private arena and
                 # an unpaired draft pool both start every request at 0.
-                self._ctx_reuse_addressable = bool(
+                self._ctx_reuse_addressable = self._ctx_unified_cache or bool(
                     getattr(draft_kv_cache_manager, "enable_joint_kv_cache_reuse", False)
                 )
             self._ctx_kv_last_page_len = torch.full(
@@ -1819,10 +1899,18 @@ def _forward_impl(
 
         # Lazy init buffers and attach worker reference for prepare()
         draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager)
+        previous_ctx_len = self._ctx_len
         self._lazy_init_ctx_buffers(
             draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager
         )
         spec_metadata._dflash_worker = self
+        if (
+            self._ctx_unified_cache
+            and self._ctx_len is not previous_ctx_len
+            and spec_metadata.request_ids
+        ):
+            # New buffers clear request slots; seed them before generation.
+            spec_metadata.prepare()
         # Before any store: prefill and decode both address pages through it.
         # Returning False here means an empty batch -- the missing-offsets case
         # raises inside, with the metadata type in the message.
@@ -2342,7 +2430,7 @@ def prepare_1st_drafter_inputs(
                     projected_rows = projected_to_store.reshape(num_gens, K_plus_1, -1)
                     lilicorr_anchor = last_accepted_hidden(projected_rows, gen_num_accepted)
                 gen_num_accepted_long = gen_num_accepted.long()
-                col_idx = self._ctx_len[slots].unsqueeze(1) + offsets_kp1.unsqueeze(0)
+                col_idx = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0)
                 write_mask = offsets_kp1.unsqueeze(0) < gen_num_accepted_long.unsqueeze(1)
                 if self._ctx_block_tables is not None:
                     # Bound by what the manager handed THIS request: the masked

tensorrt_llm/_torch/speculative/interface.py

@@ -179,7 +179,8 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata,
     Call restore_attn_metadata_after_draft_replay in a finally block.
     Returns saved state or None if no-op.
     """
-    if draft_kv_cache_manager is None:
+    if (draft_kv_cache_manager is None
+            or draft_kv_cache_manager is attn_metadata.kv_cache_manager):
         return None
     if not isinstance(attn_metadata, TrtllmAttentionMetadata):
         return None
@@ -3239,8 +3240,10 @@ def get_draft_kv_cache_manager(self, resource_manager):
         Get the draft KV cache manager if using separate KV cache layouts.
         """
         if self.use_separate_draft_kv_cache and resource_manager is not None:
-            return resource_manager.get_resource_manager(
+            draft = resource_manager.get_resource_manager(
                 ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
+            if draft is not None:
+                return draft
         if resource_manager is not None:
             target = resource_manager.get_resource_manager(
                 ResourceManagerType.KV_CACHE_MANAGER)
@@ -3261,7 +3264,8 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager):
         """
 
         # draft_kv_cache_manager is None if using two-engine speculative decoding or not enabling separate draft KV cache.
-        if draft_kv_cache_manager is None:
+        if (draft_kv_cache_manager is None
+                or draft_kv_cache_manager is attn_metadata.kv_cache_manager):
             yield attn_metadata
             return
 

tensorrt_llm/_torch/speculative/spec_tail_kernels.py

@@ -134,6 +134,7 @@ def _decode_block_table_kernel(
     stride_src_col,
     n_blocks,
     divisor,
+    block_offset,
     table_ptr,
     stride_table,
     counts_ptr,
@@ -148,21 +149,32 @@ def _decode_block_table_kernel(
         cols = start + offs
         mask = cols < n_blocks
         encoded = tl.load(src_row + cols * stride_src_col, mask=mask, other=-1).to(tl.int64)
+        encoded -= block_offset
         count += (encoded >= 0).to(tl.int64)
         tl.store(table_row + cols, (tl.maximum(encoded, 0) // divisor).to(tl.int32), mask=mask)
     tl.store(counts_ptr + row, tl.sum(count, axis=0))
 
 
 def decode_draft_block_table(
-    src: torch.Tensor, divisor: int, table: torch.Tensor, counts: torch.Tensor
+    src: torch.Tensor,
+    divisor: int,
+    table: torch.Tensor,
+    counts: torch.Tensor,
+    *,
+    block_offset: int = 0,
 ) -> None:
-    """``counts[i] = (src[i] >= 0).sum()``; ``table[i] = src[i].clamp(min=0) // divisor``.
+    """Decode layer-relative block offsets without modifying ``src``.
+
+    For ``encoded = src - block_offset``, computes
+    ``counts[i] = (encoded[i] >= 0).sum()`` and
+    ``table[i] = encoded[i].clamp(min=0) // divisor``.
 
     Args:
         src: ``[rows, blocks]`` encoded block offsets (any strides).
         divisor: encoded offset -> pool block index divisor.
         table: ``[>= rows, blocks]`` int32 output, unit inner stride.
         counts: ``[>= rows]`` int64 output.
+        block_offset: this layer's K-buffer offset within each physical slot.
     """
     rows, n_blocks = src.shape
     assert table.shape[-1] == n_blocks and table.stride(-1) == 1
@@ -174,6 +186,7 @@ def decode_draft_block_table(
         src.stride(1),
         n_blocks,
         int(divisor),
+        int(block_offset),
         table,
         table.stride(0),
         counts,

tensorrt_llm/_torch/speculative/utils.py

@@ -804,7 +804,9 @@ def get_draft_kv_cache_manager(spec_config, resource_manager):
     target = resource_manager.get_resource_manager(
         ResourceManagerType.KV_CACHE_MANAGER)
     get_view = getattr(target, "get_draft_kv_cache_view", None)
-    return get_view() if get_view is not None else None
+    view = get_view() if get_view is not None else None
+    # A unified manager's request resources are already handled through target.
+    return None if view is target else view
 
 
 def update_spec_config_from_model_config(spec_config,