[None][feat] Page the DSpark drafter context through the draft KV cache manager

tensorrt_llm/_torch/models/modeling_dflash.py

@@ -873,7 +873,10 @@ def dflash_forward(
             )
             block_tables = ctx_page_table.index_select(0, cache_batch_idx_i32.long())
             pages_per_slot = block_tables.size(1)
-            page_size = ctx_kv_cache.size(-2)
+            # Index the layer first: ctx_kv_cache is an [L, ...] tensor for the
+            # private arena but a per-layer list when bound to the draft KV
+            # cache manager's pool. page_size sits at -2 either way.
+            page_size = ctx_kv_cache[0].size(-2)
             kv_indices = block_tables.flatten()
             kv_indptr = torch.arange(
                 0,

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -1500,7 +1500,18 @@ def _should_create_separate_draft_kv_cache(self) -> bool:
         in the target model and don't produce a separate ModelConfig. We fall
         back to the target model's config via _get_effective_draft_config().
         """
-        if self._mapping.enable_attention_dp:
+        if self._speculative_config is None:
+            # No drafter at all, so there is nothing to give a manager to.
+            return False
+        # Narrower than is_external_drafter(): PARD and DRAFT_TARGET_ONE_MODEL
+        # never reach the arena this carve-out exists for.
+        spec_dec_mode = self._speculative_config.spec_dec_mode
+        is_standalone_drafter = (spec_dec_mode.is_dflash()
+                                 or spec_dec_mode.is_dspark())
+        if self._mapping.enable_attention_dp and not is_standalone_drafter:
+            # This bail suits MTP, whose draft layers are target-shaped and
+            # appendable to the target pool. A standalone drafter has nothing to
+            # append, so it would be stranded on the private arena instead.
             logger.info(
                 "Attention DP is enabled, separate draft KV cache is not supported."
             )

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -2950,6 +2950,41 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests):
             self._prepare_draft_resources(scheduled_batch)
             return
 
+    def _mirror_draft_kv_cache(self, req: LlmRequest):
+        """The draft manager's entry for ``req``, created on first sight.
+
+        Mirrors what the target manager does in ``_prepare_context_impl``:
+        create when missing rather than fail, passing ``is_dummy`` through. The
+        draft mirror never looks up block reuse (``tokens=None``) and never
+        commits, so this cannot share the target's implementation.
+
+        Reached from both the context and the generation loop, because not
+        every generation request passed through context on this worker:
+        attention DP injects its idle placeholder straight as a generation
+        request when a rank has no real work (ATTENTION_DP_DUMMY_REQUEST_ID,
+        see py_executor._adp_dummy_is_gen), and a disaggregated generation
+        worker receives prompt KV rather than prefilling it.
+
+        Returns None when the IndexMapper is saturated. A context request can
+        retry next iteration; a generation request cannot, and skipping it only
+        defers the failure to copy_batch_block_offsets(), which asserts in C++
+        on the unmapped request ID.
+        """
+        kv_cache = self.kv_cache_map.get(req.py_request_id)
+        if kv_cache is not None:
+            return kv_cache
+        kv_cache = self._create_kv_cache(
+            req.py_request_id,
+            req.lora_task_id,
+            None,
+            cache_salt=req.cache_salt,
+            is_dummy=req.is_dummy,
+        )
+        if kv_cache is None:
+            return None
+        kv_cache.stop_committing()
+        return kv_cache
+
     def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests):
         """Create/resize KV caches in the draft V2 manager for scheduled requests.
 
@@ -2962,23 +2997,16 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests):
         # separate draft engine still uses the legacy draft request view.
         with request_context(not self.enable_joint_kv_cache_reuse, scheduled_batch):
             for req in scheduled_batch.context_requests:
-                kv_cache = self.kv_cache_map.get(req.py_request_id)
+                kv_cache = self._mirror_draft_kv_cache(req)
                 if kv_cache is None:
-                    kv_cache = self._create_kv_cache(
-                        req.py_request_id,
-                        req.lora_task_id,
-                        None,
-                        cache_salt=req.cache_salt,
-                        is_dummy=req.is_dummy,
+                    # Retryable here, unlike the generation loop below: a
+                    # context request has not drafted yet, so the next
+                    # iteration can mirror it once slots free up.
+                    logger.warning(
+                        f"Draft KV cache mirror has no free IndexMapper slot for "
+                        f"context request {req.py_request_id}; retrying next iteration."
                     )
-                    if kv_cache is None:
-                        # Saturated IndexMapper (e.g. slots held by disagg
-                        # generation transfers in flight): skip mirroring this
-                        # request for now; it is retried next iteration once
-                        # slots free up, before the request runs any spec-dec
-                        # forward that needs the mirror.
-                        continue
-                    kv_cache.stop_committing()
+                    continue
                 if not self._resume_and_restore(req.py_request_id, kv_cache):
                     raise RuntimeError(
                         f"Failed to resume draft KV cache for request {req.py_request_id}"
@@ -2997,10 +3025,11 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests):
                     )
 
             for req in scheduled_batch.generation_requests:
-                kv_cache = self.kv_cache_map.get(req.py_request_id)
+                kv_cache = self._mirror_draft_kv_cache(req)
                 if kv_cache is None:
                     raise RuntimeError(
-                        f"Missing draft KV cache for generation request {req.py_request_id}"
+                        f"Draft KV cache mirror exhausted its IndexMapper on "
+                        f"generation request {req.py_request_id}"
                     )
                 if not self._resume_and_restore(req.py_request_id, kv_cache):
                     raise RuntimeError(
@@ -3849,6 +3878,10 @@ def release_index_slot(self, request_id: int) -> None:
         the KV cache blocks are still being transferred via NIXL/UCX.
         """
         kv_cache = self.kv_cache_map.get(request_id)
+        if self.is_draft and (kv_cache is None or request_id in self._early_freed_index_requests):
+            # The draft mirror only holds a slot for requests it actually
+            # mirrored, and the target may release the same request twice.
+            return
         if kv_cache is not None:
             for i in range(self.max_beam_width):
                 for pool_idx in range(self.num_pools):

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -7742,9 +7742,12 @@ def _send_disagg_ctx_kv_async(self,
                     # Forward is done for this request — release the
                     # IndexMapper slot so new requests can reuse it.
                     # KV blocks stay allocated for the upcoming transfer.
-                    if hasattr(self.kv_cache_manager, 'release_index_slot'):
-                        self.kv_cache_manager.release_index_slot(
-                            req.py_request_id)
+                    # getattr: executors built by unit tests skip __init__
+                    # and so never get the draft binding.
+                    for mgr in (self.kv_cache_manager,
+                                getattr(self, 'draft_kv_cache_manager', None)):
+                        if hasattr(mgr, 'release_index_slot'):
+                            mgr.release_index_slot(req.py_request_id)
                     # Order matters: start_transfer commits the request's blocks to the reuse
                     # tree and pins them, and must run before respond_and_send_async sends the
                     # final KV slice and (for the Python transceiver) transitions the request toward completion.

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -496,10 +496,12 @@ def create_py_executor(
         if hasattr(spec_config, '_max_batch_size'):
             spec_config._max_batch_size = max_batch_size
 
-        # WAR for https://nvbugs/5807902
-        # Disable separate draft KV cache in disaggregated mode
-        # Enable separate pool for None DI + Non-KVBM and Aggregated + KVBM
-        if cache_transceiver_config is not None:
+        # WAR for https://nvbugs/5807902 (Eagle3 disagg RMSNorm crash, closed
+        # will-not-fix). Keep the blanket disable; carve out only the standalone
+        # drafters, which it stranded on their private max_seq_len-dense arena.
+        is_standalone_drafter = (spec_config.spec_dec_mode.is_dflash()
+                                 or spec_config.spec_dec_mode.is_dspark())
+        if cache_transceiver_config is not None and not is_standalone_drafter:
             spec_config._allow_separate_draft_kv_cache = False
 
     # chunk_unit_size may be changed to 64 when using flash mla