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

jenkins/L0_MergeRequest.groovy

@@ -1345,6 +1345,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars)
         "tests/integration/test_lists/test-db/l0_gb300.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_gpus_perf_sanity.yml",
+        "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu2_gen1_node2_gpu8.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu2_gen1_node8_gpu32.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8.yml",

jenkins/L0_Test.groovy

@@ -1731,6 +1731,10 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG
     def jobWorkspace = "/home/svc_tensorrt/bloom/scripts/${jobUID}"
     def disaggMultiNodeMode = stageName.contains("Disagg-PerfSanity")
     def aggMultiNodeMode = !disaggMultiNodeMode && nodeCount > 1 && stageName.contains("PerfSanity")
+    // Forces --segment=${nodeCount} so all nodes land in one NVLink domain. Aggregated
+    // multi-node jobs need this: the runtime only takes the MNNVL path when every TP rank
+    // shares a fabric, and a cross-rack allocation silently degrades instead of failing.
+    def singleNvlinkDomainMode = stageName.contains("1_NVLink_Domain")
     def infraDryRun = isInfraDryRun()
     if (infraDryRun) {
         testList = INFRA_DRY_RUN_TEST_CONTEXT
@@ -1976,7 +1980,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG
                         .replace("${ARTIFACTORY_DOCKER_HOST}/", "${ARTIFACTORY_DOCKER_HOST}#")
                 }
                 def mounts = getMountListForSlurmTest(cluster, true).join(",")
-                String[] taskArgs = getNodeArgs(nodeCount, gpuCount, disaggMultiNodeMode)
+                String[] taskArgs = getNodeArgs(nodeCount, gpuCount, disaggMultiNodeMode || singleNvlinkDomainMode)
                 if (taskArgs == null) {
                     error "Invalid Slurm test stage name is set"
                 }
@@ -6335,13 +6339,18 @@ def launchTestJobs(pipeline, testFilter, globalVars)
     fullSet += SBSASlurmTestConfigs.keySet()
 
     multiNodesSBSAConfigs = [
-        // Each testcase uses 8 GPUs and 2 nodes.
+        // Each GB200 testcase below uses 8 GPUs and 2 nodes.
         // https://nvbugs/5598863 (uncorrectable NVLink error detected during the execution) may not exist in OCI machines.
         "GB200-8_GPUs-2_Nodes-PyTorch-1": ["auto:gb200-flex", "l0_gb200_multi_nodes", 1, 2, 8, 2],
         "GB200-8_GPUs-2_Nodes-PyTorch-2": ["auto:gb200-flex", "l0_gb200_multi_nodes", 2, 2, 8, 2],
         "GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-1": ["auto:gb200-flex", "l0_gb200_multi_nodes", 1, 3, 8, 2],
         "GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-2": ["auto:gb200-flex", "l0_gb200_multi_nodes", 2, 3, 8, 2],
         "GB200-8_GPUs-2_Nodes-PyTorch-Post-Merge-3": ["auto:gb200-flex", "l0_gb200_multi_nodes", 3, 3, 8, 2],
+        // GB300 accuracy post-merge aggregated (4 GPUs per node). One test list per topology,
+        // spelled out here rather than via buildStageConfigs: test_to_stage_mapping.py resolves
+        // stage <-> test by list name with a line-based parser, so a shared list or a helper's
+        // output breaks the mapping. For 1_NVLink_Domain see singleNvlinkDomainMode.
+        "GB300-8_GPUs-2_Nodes-1_NVLink_Domain-PyTorch-Post-Merge-1": ["auto:gb300-flex", "l0_gb300_multi_nodes_node2_gpu8", 1, 1, 8, 2],
     ]
     // PerfSanity post-merge aggregated
     // 2 Nodes

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

@@ -1438,7 +1438,23 @@ 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
+        is_external_drafter = (
+            self._speculative_config.spec_dec_mode.is_external_drafter())
+        if self._mapping.enable_attention_dp and not is_external_drafter:
+            # Attention-DP uses the target KV cache manager for MTP draft layers
+            # so draft cache state follows the DP-sharded target requests. That
+            # works only because MTP draft layers are target-shaped and can be
+            # appended to the target pool (get_pp_layers). An external drafter
+            # has its own architecture and layer count, so there is nothing to
+            # append and the bail would strand it on the private context arena
+            # -- dense in max_seq_len, and with KV heads unsharded under
+            # attention DP that is the allocation that OOMs at long ISL. Its
+            # own manager follows the DP sharding regardless: same mapping (so
+            # `tp_size = 1 if enable_attention_dp` applies to it too) and block
+            # tables copied per rank from attn_metadata.request_ids.
             logger.info(
                 "Attention DP is enabled, separate draft KV cache is not supported."
             )

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -2717,6 +2717,43 @@ 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. The caller must treat
+        that as fatal: the request stays in the scheduled batch, so skipping it
+        only defers the failure to copy_batch_block_offsets(), which asserts in
+        C++ on the unmapped request ID. Slots are released early on a
+        context-only worker (release_index_slot), so saturation here means the
+        draft mirror is out of sync with the target, not merely busy.
+        """
+        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.
 
@@ -2727,23 +2764,12 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests):
         """
         with request_context(True, 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,
+                    raise RuntimeError(
+                        f"Draft KV cache mirror exhausted its IndexMapper on "
+                        f"context request {req.py_request_id}"
                     )
-                    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()
                 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}"
@@ -2762,10 +2788,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(
@@ -3607,6 +3634,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

@@ -591,6 +591,10 @@ def __init__(
         # kv cache events
         self.kv_cache_manager = self.resource_manager.resource_managers.get(
             ResourceManagerType.KV_CACHE_MANAGER)
+        # A separate draft manager mirrors the target's slot lifecycle, so it
+        # needs the same early index-slot release on a context-only worker.
+        self.draft_kv_cache_manager = self.resource_manager.resource_managers.get(
+            ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
         # V2 owns KV allocation, suspend, resume, and context finalization.
         # The executor skips the V1 terminate/pause paths and finalizes V2
         # context resources before transfer or response handling can terminate
@@ -7514,8 +7518,9 @@ 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)
+                for mgr in (self.kv_cache_manager, self.draft_kv_cache_manager):
+                    if hasattr(mgr, 'release_index_slot'):
+                        mgr.release_index_slot(req.py_request_id)
                 # Order is important here: we need to start the transfer before responding
                 # to make sure the blocks are stored for reuse before they are sent.
                 self.async_transfer_manager.start_transfer(req)

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -507,10 +507,17 @@ 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 crashes in disaggregated mode
+        # when the draft layers get their own KV cache manager (RMSNorm invalid
+        # argument). The root cause was never found and the bug closed
+        # will-not-fix, so the workaround stands -- but only for the mode it was
+        # reported on. It was originally a blanket disable, which also stranded
+        # external drafters on their private context arena; that arena is dense
+        # in max_seq_len and, with KV heads unsharded under attention DP, is
+        # what makes long-ISL disaggregated serving unservable.
+        is_eagle3 = (spec_config.spec_dec_mode.is_eagle3()
+                     or spec_config.spec_dec_mode.is_eagle3_one_model())
+        if cache_transceiver_config is not None and is_eagle3:
             spec_config._allow_separate_draft_kv_cache = False
 
     # chunk_unit_size may be changed to 64 when using flash mla