[None][feat] DSpark Unified KVCache standalone

docs/source/features/speculative-decoding.md

@@ -156,7 +156,7 @@ Reference: [DFlash: Distilled Flash Speculative Decoding](https://arxiv.org/pdf/
 * `speculative_model`: Path or HuggingFace model ID for the DFlash draft model.
 * `mask_token_id`: Token ID used as the mask token for parallel prediction. If not set, it is read from the draft model config.
 * `target_layer_ids`: List of target model layer indices whose hidden states are captured for cross-attention in the draft model. If not set, read from the draft model config.
-* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"VANILLA"` (the default) uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a private paged context K/V cache and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only.
+* `attention_backend`: Cross-attention backend for the draft model, independent of the drafter's standard attention modules. `"VANILLA"` (the default) uses FlashAttention with a contiguous context K/V cache and runs anywhere. `"TRTLLM"` uses TRTLLM-Gen FMHA (via FlashInfer) over a paged context K/V cache held by the KV cache manager (see [Draft KV cache](#dflash-draft-kv-cache)) and supports SM100/SM103 only. `"FA4"` uses the flash-attn CuTe DSL kernels on the same paged cache and supports SM90 only.
 
 ```python
 from tensorrt_llm.llmapi import DFlashDecodingConfig
@@ -170,6 +170,37 @@ speculative_config = DFlashDecodingConfig(
 llm = LLM("/path/to/target_model", speculative_config=speculative_config)
 ```
 
+#### DFlash draft KV cache
+
+With `attention_backend="TRTLLM"`, TRTLLM target attention, and the default
+`KVCacheManagerV2`, the drafter's layers are allocated in the target's KV cache
+manager. Target and draft layers share one page table, so one KV budget covers
+both and each request's draft pages are allocated, padded and freed with its
+target pages. The draft layers keep the drafter's own dtype, so an FP8 target KV
+cache (`kv_cache_config.dtype: fp8`) is paired with a BF16 draft cache.
+
+In disaggregated serving, direct worker launches require
+`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` on both the context and generation workers
+to enable shared draft-KV transfer. The `disaggregated_mpi_worker` launcher
+enables this setting automatically when both servers enable compatible
+speculation and use the same `tensor_parallel_size` and
+`enable_attention_dp`.
+
+When sharing is enabled, the context server's draft pages are transferred to
+the generation server with the target pages, so the drafter starts generation
+with the prompt context instead of an empty one. This requires the draft layers
+to have the target's rank-local KV-head count; the draft dtype and head
+dimension may differ, except that an NVFP4 target KV cache requires an
+identical draft layout. Both servers must use the same layout.
+
+The drafter keeps a separate draft KV cache manager instead when any of the
+following applies: context parallelism larger than 1, a per-layer KV-head count
+on the target or drafter, a specialized KV cache manager (for example, hybrid Mamba
+models), Gemma 4 hybrid attention, or, in disaggregated serving, a draft layout
+that cannot be transferred with the target's. When the layers are shared, the
+worker log reports `target and draft KV layers share one KV cache manager` at
+startup.
+
 [DFlash 2](https://inco.ai/blog/dflash2/) is also supported. The same `DFlashDecodingConfig` can be used for DFlash 2; no extra arguments are required.
 
 [LiLiCorr](https://arxiv.org/abs/2608.20530) checkpoints use the same

tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py

@@ -604,6 +604,12 @@ def _run_mla_decode(
         params: FmhaParams,
         kernel_dtype: torch.dtype,
     ) -> None:
+        from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import (
+            KVCacheManagerV2,
+            Role,
+        )
+        from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX
+
         attn = params.attn
         meta = params.meta
 
@@ -680,6 +686,21 @@ def _run_mla_decode(
         page_table = self._generation_page_table(
             page_table_layer, meta.num_contexts, batch_size, layer_in_pool
         )
+        if (
+            layers_in_pool == 1
+            and isinstance(meta.kv_cache_manager, KVCacheManagerV2)
+            and meta.kv_cache_manager._use_per_layer_page_tables
+        ):
+            # V2 per-layer indices include the layer offset, but get_buffers()
+            # already starts at this layer's KEY buffer. Apply the offset once.
+            converter = meta.kv_cache_manager.impl.get_page_index_converter(
+                local_layer_idx, Role.KEY
+            )
+            layer_offset = int(converter.layer_offset)
+            if layer_offset:
+                page_table = torch.where(
+                    page_table == BAD_PAGE_INDEX, page_table, page_table - layer_offset
+                )
 
         # KVCacheManager exposes NHD pages as [num_pages, 1, page_size, 1, head_dim].
         # The CuTe DSL kernel consumes a paged [page_size, dim, num_pages] view

tensorrt_llm/_torch/disaggregation/native/peer.py

@@ -180,6 +180,29 @@ 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; "
+                "set TRTLLM_DISAGG_SHARED_DRAFT_KV identically on both workers."
+            )
+        if self._ri.num_draft_layers:
+            if self._ri.tp_size != peer_ri.tp_size:
+                raise ValueError(
+                    "PeerRegistrar: a shared draft KV cache needs equal tp_size "
+                    f"(local={self._ri.tp_size}, peer={peer_ri.tp_size}); "
+                    "unset TRTLLM_DISAGG_SHARED_DRAFT_KV or use the same TP on both workers."
+                )
+            if self._ri.attention.enable_attention_dp != peer_ri.attention.enable_attention_dp:
+                raise ValueError(
+                    "PeerRegistrar: a shared draft KV cache needs the same attention DP "
+                    f"setting (local={self._ri.attention.enable_attention_dp}, "
+                    f"peer={peer_ri.attention.enable_attention_dp}); unset "
+                    "TRTLLM_DISAGG_SHARED_DRAFT_KV or match it on both workers."
+                )
 
         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