zheyuf · Oct 27, 2026 · Oct 21, 2026 · Oct 21, 2026 · Oct 21, 2026 · Oct 21, 2026
@@ -107,6 +107,10 @@ struct BufferConfig
     std::optional<int> tokensPerBlockOverride;
 
     bool isSparse = false; //!< Whether history uses the sparse attention lifecycle.
+
+    //! Buffers may coalesce only when their group and expanded size match.
+    //! This does not change their lifecycle or shared slot allocation.
+    int coalescingGroup = 0;
 };
 
 // ---------------------------------------------------------------------------

@@ -154,13 +154,13 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config)
     LifeCycleRegistry registry{config};
     int tokensPerBlock = config.tokensPerBlock;
 
-    // Map: lifeCycleId → (bufferSize → list of BufferId).
-    // Outer map key is LifeCycleId; inner map key is expanded buffer size.
+    // Map: lifeCycleId → ((coalescingGroup, expanded buffer size) → list of BufferId).
     // NOTE: Python uses insertion-ordered dict/defaultdict here. The std::map
     // grouping below intentionally remains sorted, which can make observable
     // pool-group indices/layout order differ from Python. Later lookups use
     // StorageConfig mappings, so this is not expected to affect correctness.
-    std::map<LifeCycleId, std::map<size_t, std::vector<BufferId>>> bufferGroups;
+    using CoalescingKey = std::pair<int, size_t>;
+    std::map<LifeCycleId, std::map<CoalescingKey, std::vector<BufferId>>> bufferGroups;
     std::unordered_map<BufferId, int, BufferIdHash> expansionMap;
 
     for (auto const& layer : config.layers)
@@ -179,7 +179,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config)
 
                     BufferId bid{cfg.layerId, buf.role};
                     expansionMap[bid] = exp;
-                    bufferGroups[lcId][expandedSize].push_back(bid);
+                    bufferGroups[lcId][{buf.coalescingGroup, expandedSize}].push_back(bid);
                 }
             },
             layer);
@@ -189,14 +189,14 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config)
     // Each variant: coalesced buffers sorted by size descending.
     std::vector<SlotDescVariant> slotGroups;
     slotGroups.reserve(bufferGroups.size());
-    for (auto const& [lcId, sizeMap] : bufferGroups)
+    for (auto const& [lcId, coalescingMap] : bufferGroups)
     {
         SlotDescVariant var;
         var.lifeCycleId = lcId;
-        for (auto const& [sz, bufIds] : sizeMap)
+        for (auto const& [key, bufIds] : coalescingMap)
         {
             CoalescedBuffer cb;
-            cb.singleBufferSize = sz;
+            cb.singleBufferSize = key.second;
             cb.bufferIds = bufIds;
             var.coalescedBuffers.push_back(std::move(cb));
         }

@@ -1597,12 +1597,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
         .def("assert_valid", &kv::DiskCacheTierConfig::assertValid) DEF_COPY(kv::DiskCacheTierConfig);
 
     nb::class_<kv::BufferConfig>(m, "BufferConfig")
-        .def(nb::init<kv::DataRole, size_t, std::optional<int>, bool>(), nb::arg("role"), nb::arg("size"),
-            nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false)
+        .def(nb::init<kv::DataRole, size_t, std::optional<int>, bool, int>(), nb::arg("role"), nb::arg("size"),
+            nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false,
+            nb::arg("coalescing_group") = 0)
         .def_rw("role", &kv::BufferConfig::role)
         .def_rw("size", &kv::BufferConfig::size)
         .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride)
-        .def_rw("is_sparse", &kv::BufferConfig::isSparse) DEF_COPY(kv::BufferConfig);
+        .def_rw("is_sparse", &kv::BufferConfig::isSparse)
+        .def_rw("coalescing_group", &kv::BufferConfig::coalescingGroup) DEF_COPY(kv::BufferConfig);
 
     nb::class_<kv::AttentionLayerConfig>(m, "AttentionLayerConfig")
         .def(nb::init<kv::LayerId, std::vector<kv::BufferConfig>, std::optional<int>, std::optional<int>>(),

@@ -1659,14 +1659,17 @@ KvCachePoolPointers buildKvCachePoolPointers(at::Tensor const& hostKvCachePoolPo
     }
     else
     {
-        TORCH_CHECK(hostKvCachePoolPointers.dim() == 2);
+        TORCH_CHECK(hostKvCachePoolPointers.dim() == 2 || hostKvCachePoolPointers.dim() == 3);
+        // Non-FP4 layers use only the data pointers in a mixed-format pool table.
+        auto const dataPointers
+            = hostKvCachePoolPointers.dim() == 3 ? hostKvCachePoolPointers.select(2, 0) : hostKvCachePoolPointers;
         pointers.primaryPoolPtr = reinterpret_cast<void*>(
             reinterpret_cast<char*>(
-                readHostTensor2D<int64_t>(hostKvCachePoolPointers, poolIndex, 0, "host_kv_cache_pool_pointers"))
+                readHostTensor2D<int64_t>(dataPointers, poolIndex, 0, "host_kv_cache_pool_pointers"))
             + intraPoolOffset);
         pointers.secondaryPoolPtr = reinterpret_cast<void*>(
             reinterpret_cast<char*>(
-                readHostTensor2D<int64_t>(hostKvCachePoolPointers, poolIndex, 1, "host_kv_cache_pool_pointers"))
+                readHostTensor2D<int64_t>(dataPointers, poolIndex, 1, "host_kv_cache_pool_pointers"))
             + intraPoolOffset);
     }
     return pointers;

@@ -152,7 +152,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. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker and under `modeling_v2`. `"VANILLA"` 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. `"AUTO"` (the default) resolves to `"TRTLLM"` where it can run and every drafter layer is windowed (sliding-window drafters such as DFlash 2), and to `"VANILLA"` otherwise, including on a disaggregated worker with `TRTLLM_DISAGG_SHARED_DRAFT_KV=0` and under `modeling_v2`. `"VANILLA"` 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
@@ -166,6 +166,39 @@ 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, 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. Both servers therefore need
+the same `speculative_config`: a generation server whose context server keeps a
+different number of draft layers, for example one without a drafter, fails
+those requests at the first KV transfer. To draft on the generation server
+alone, set `attention_backend="VANILLA"` or `TRTLLM_DISAGG_SHARED_DRAFT_KV=0`
+there; it then drafts without the prompt context, at a lower acceptance length.
+The target and draft layers
+are resharded by their own KV-head counts, so the draft's KV-head count, head
+dimension and dtype may differ from the target's, and the context and
+generation servers may use different tensor-parallel sizes or attention DP
+settings. An NVFP4 target KV cache requires a draft with an identical dtype
+and KV-head count. The target's and draft's K-only (MLA) or K/V cache kind must
+match, and both servers must run the same target and drafter.
+
+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.
 
 ### User-provided drafting

@@ -329,6 +329,8 @@ def block_offsets(
             host is None
             or getattr(manager, "sparse_metadata_batch", None) is not None
             or getattr(manager, "_use_per_layer_page_tables", True)
+            # Shared attention-op tables are not the life cycles' SHARED-path tables.
+            or getattr(manager, "_attention_op_pool_groups", None) is not None
             or not dst.is_cuda
             or dst.dtype != torch.int32
             or host.dtype != torch.int32

@@ -1352,20 +1352,23 @@ def _post_init_with_buffers(self, buffers) -> None:
             self._max_num_blocks_per_seq = self.kv_cache_manager.max_blocks_per_seq
 
             # Layers may share one page-index list only when they are in the
-            # same pool AND have the same page-index scale: VSWA splits pools,
+            # same pool, page-index scale, and KV factor: VSWA splits pools,
             # and per-layer geometry (e.g. Gemma4 sliding/global head_dim)
             # splits scales even when all windows collapse to max_seq_len and
-            # is_vswa is False. Guarded on V2-specific attributes (V1 managers
-            # lack the per-pool infrastructure).
+            # is_vswa is False. Mixed K-only/KV layers can use different view
+            # indices even with the same physical scale. Guarded on V2-specific
+            # attributes (V1 managers lack the per-pool infrastructure).
             mgr = self.kv_cache_manager
             get_scale = getattr(mgr, 'get_layer_page_index_scale', None)
+            get_kv_factor = getattr(mgr, 'get_layer_kv_factor', None)
             layer_space: Dict[int, int] = {}
             if hasattr(mgr, 'layer_to_pool_mapping_dict') and get_scale:
                 space_ids = {}
                 for layer_idx in attention_layer_indices:
                     layer_offset = mgr.layer_offsets[layer_idx]
                     key = (mgr.layer_to_pool_mapping_dict[layer_offset],
-                           get_scale(layer_idx))
+                           get_scale(layer_idx),
+                           get_kv_factor(layer_idx) if get_kv_factor else None)
                     space_ids.setdefault(key, len(space_ids))
                     layer_space[layer_idx] = space_ids[key]
             if layer_space:

@@ -79,6 +79,11 @@ def get_paged_kv_storage_unsupported_reason(
     if isinstance(manager, KVCacheManagerV2):
         if manager.enable_swa_scratch_reuse:
             return "KVCacheManagerV2 SWA scratch reuse is not supported."
+        # get_kv_page_offset indexes the life cycles' kv_offset by attention-op pool.
+        if getattr(manager, "_use_per_layer_page_tables", False) or (
+            getattr(manager, "_attention_op_pool_groups", None) is not None
+        ):
+            return "KVCacheManagerV2 per-layer or shared attention-op tables are not supported."
     elif isinstance(manager, KVCacheManager):
         if manager.num_pools != 1:
             return "KVCacheManagerV1 with multiple memory pools is not supported."

@@ -877,6 +877,7 @@ def _get_batch_cache_indices_by_pool_id(
         is_kv_aggregate: bool = True,
         num_blocks_per_seq: Optional[Sequence[int]] = None,
         index_scale: Optional[int] = None,
+        kv_factor: Optional[int] = None,
     ):
         """Return page indices; padded entries remain ``BAD_PAGE_INDEX`` (-1).
 
@@ -905,6 +906,8 @@ def _get_batch_cache_indices_by_pool_id(
             index_scale: Kept for compatibility with the base virtual method;
                 M3 bypasses the V1 block-id conversion entirely, so any
                 caller-supplied scale is ignored alongside ``index_scales``.
+            kv_factor: Kept for compatibility with the base virtual method;
+                M3 returns raw slot IDs without KV aggregation.
         """
         res = []
         for req_idx, req_id in enumerate(request_ids):

@@ -235,10 +235,12 @@ def reserve(
         *,
         timeout: Optional[float] = None,
         extra_bytes: int = 0,
+        writers_split_layers: bool = True,
     ) -> bool:
         """Reserve a region and record its address for the senders. False falls back to per-fragment.
         ``extra_bytes`` sizes the sender's non-paged payload (mamba/KDA recurrent state) that rides
-        the same coalesced write as the KV blocks."""
+        the same coalesced write as the KV blocks. ``writers_split_layers`` says whether the fan-in
+        writers hold different layers (PP fan-in) rather than a head slice of every layer."""
 
     @abstractmethod
     def writer_base(self, rid_slice, writer_index: int) -> Optional[int]:

@@ -235,13 +235,15 @@ def reserve(
         *,
         timeout: Optional[float] = _RESERVE_TIMEOUT_S,
         extra_bytes: int = 0,
+        writers_split_layers: bool = True,
     ) -> bool:
         """Reserve a region and create its state, recording the address for the senders. Returns
         False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total
         must divide across the writers. ``extra_bytes`` is the non-paged payload the sender appends
         to the same coalesced write (mamba/KDA recurrent state, sized by the receiver via
         ``mamba_receiver_payload_bytes``); the region must cover it or the write would overrun into the
-        neighboring slot."""
+        neighboring slot. ``writers_split_layers`` is False when every fan-in writer sends a head
+        slice of every layer (TP fan-in), True when the writers hold different layers (PP fan-in)."""
         total = 0
         has_state_group = False
         # Block tables are positional with -1 holes (see Chunk); only the
@@ -311,10 +313,12 @@ def reserve(
                 f"({total % num_writers}B remainder); head-mismatch explosion NOT mitigated",
                 warn_key="kv-bounce-uneven-fanin",
             )
-        if num_writers > 1:
+        if num_writers > 1 and writers_split_layers:
             # Fan-in gives each writer an equal share of the region, which only matches where it
-            # writes when all writers send the same bytes. Equal layer count guarantees that only
-            # when the per-block sizes match, so require that here, else fall back.
+            # writes when all writers send the same bytes. A TP fan-in writer sends 1/num_writers
+            # of every block of every group, so its share is exact for any per-group block sizes.
+            # PP fan-in writers send different layers, and an equal layer count gives equal bytes
+            # only when the per-block sizes match, so require that here, else fall back.
             # Exclude None entries (STATE groups) — they are already guarded above.
             present_slot_bytes = {
                 self._block_bytes_per_group[g]
@@ -475,12 +479,23 @@ def _scatter_loop(self):
             ok = True
             try:
                 # Scatter each writer's fragments from its own source, never one global offset, so a
-                # missing or fallback writer cannot shift where the others are read from.
+                # missing or fallback writer cannot shift where the others are read from. All writers
+                # go into ONE batched copy: the copy helper reuses one pinned metadata buffer per
+                # stream and needs a sync between calls, and one call costs one host round trip
+                # instead of one per writer.
+                srcs, dsts, frag_sizes = [], [], []
                 for src_base, dst_ptrs, sizes in descs:
                     p = Plan(dst_ptrs, dst_ptrs, sizes, int(sizes.sum()))
-                    scatter_contiguous(
-                        src_base, p.dst_ptrs, p.sizes, p.offsets, stream=self._scatter_stream
-                    )
+                    srcs.append(np.int64(src_base) + p.offsets)
+                    dsts.append(p.dst_ptrs)
+                    frag_sizes.append(p.sizes)
+                scatter_contiguous(
+                    0,
+                    np.concatenate(dsts),
+                    np.concatenate(frag_sizes),
+                    np.concatenate(srcs),
+                    stream=self._scatter_stream,
+                )
                 CUASSERT(cudart.cudaStreamSynchronize(self._scatter_stream))
             except Exception as e:
                 # a scatter failure must not kill the worker nor be reported as success
@@ -524,6 +539,7 @@ def reserve(
         *,
         timeout: Optional[float] = _RESERVE_TIMEOUT_S,
         extra_bytes: int = 0,
+        writers_split_layers: bool = True,
     ) -> bool:
         return False