[None][fix] Runtime fixes found while bringing up Mistral Large 4

3rdparty/fetch_content.json

@@ -62,7 +62,8 @@
       "git_repository": "https://github.com/deepseek-ai/FlashMLA.git",
       "git_tag": "15f13e5030374295491c5ce31b02d7e63a7772c6",
       "git_submodules_recurse": true,
-      "source_subdir": "dont-add-this-project-with-add-subdirectory"
+      "source_subdir": "dont-add-this-project-with-add-subdirectory",
+      "patch_file": "patches/flashmla_release_gil.patch"
     },
     {
       "name": "fmt",

3rdparty/patches/flashmla_release_gil.patch

@@ -0,0 +1,17 @@
+--- a/csrc/api/api.cpp
++++ b/csrc/api/api.cpp
+@@ -7,9 +7,9 @@
+
+ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+     m.doc() = "FlashMLA";
+-    m.def("sparse_decode_fwd", &sparse_attn_decode_interface);
+-    m.def("dense_decode_fwd", &dense_attn_decode_interface);
+-    m.def("sparse_prefill_fwd", &sparse_attn_prefill_interface);
+-    m.def("dense_prefill_fwd", &FMHACutlassSM100FwdRun);
+-    m.def("dense_prefill_bwd", &FMHACutlassSM100BwdRun);
++    m.def("sparse_decode_fwd", &sparse_attn_decode_interface, pybind11::call_guard<pybind11::gil_scoped_release>());
++    m.def("dense_decode_fwd", &dense_attn_decode_interface, pybind11::call_guard<pybind11::gil_scoped_release>());
++    m.def("sparse_prefill_fwd", &sparse_attn_prefill_interface, pybind11::call_guard<pybind11::gil_scoped_release>());
++    m.def("dense_prefill_fwd", &FMHACutlassSM100FwdRun, pybind11::call_guard<pybind11::gil_scoped_release>());
++    m.def("dense_prefill_bwd", &FMHACutlassSM100BwdRun, pybind11::call_guard<pybind11::gil_scoped_release>());
+ }

tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py

@@ -12,11 +12,12 @@
     IterationStats,
     KvCacheStats,
     RequestStats,
+    SpecDecodingStats,
 )
 from tensorrt_llm.logger import logger
 
 from .scheduler.adp_router import RankIterStatsPayload, RankState
-from .scheduler.rank_state import RankKvCacheStats
+from .scheduler.rank_state import RankKvCacheStats, RankSpecDecodingStats
 
 _ITERATION_STATS_SCALAR_FIELDS = (
     "timestamp",
@@ -33,10 +34,7 @@
     "pinned_mem_usage",
 )
 
-_ITERATION_STATS_OPTIONAL_FIELDS = (
-    "static_batching_stats",
-    "specdec_stats",
-)
+_ITERATION_STATS_OPTIONAL_FIELDS = ("static_batching_stats",)
 
 
 def _snapshot_kv_cache_stats(stats: KvCacheStats | None) -> RankKvCacheStats | None:
@@ -78,6 +76,51 @@ def _restore_kv_cache_stats(snapshot: RankKvCacheStats) -> KvCacheStats:
     return stats
 
 
+def _snapshot_specdec_stats(stats: SpecDecodingStats | None) -> RankSpecDecodingStats | None:
+    """Capture a rank's speculative-decoding counters.
+
+    Args:
+        stats: Local speculative-decoding stats, or None when they are unavailable.
+
+    Returns:
+        Immutable counter snapshot, or None when unavailable.
+    """
+    if stats is None:
+        return None
+    return RankSpecDecodingStats(
+        num_draft_tokens=stats.num_draft_tokens,
+        num_accepted_tokens=stats.num_accepted_tokens,
+        num_requests_with_draft_tokens=stats.num_requests_with_draft_tokens,
+    )
+
+
+def _restore_specdec_stats(
+    snapshot: RankSpecDecodingStats, rank0_stats: SpecDecodingStats | None
+) -> SpecDecodingStats:
+    """Create a native speculative-decoding stats object for one rank's row.
+
+    The counters and the acceptance length derived from them are the rank's
+    own. The draft timings keep rank 0's values, like the row's other timings.
+
+    Args:
+        snapshot: Counters captured with the rank's iteration payload.
+        rank0_stats: Rank 0's speculative-decoding stats for the same iteration.
+
+    Returns:
+        A fresh native stats object.
+    """
+    stats = SpecDecodingStats()
+    for name, value in zip(snapshot._fields, snapshot):
+        setattr(stats, name, value)
+    num_requests = snapshot.num_requests_with_draft_tokens
+    stats.acceptance_length = (
+        (snapshot.num_accepted_tokens + num_requests) / num_requests if num_requests else 0.0
+    )
+    stats.iter_latency_ms = rank0_stats.iter_latency_ms if rank0_stats is not None else 0.0
+    stats.draft_overhead = rank0_stats.draft_overhead if rank0_stats is not None else 0.0
+    return stats
+
+
 @dataclass
 class ADPIterStatsRecord:
     """Append-ready stats row produced by Attention-DP fanout."""
@@ -141,6 +184,7 @@ def make_payload(stats: IterationStats) -> RankIterStatsPayload:
             num_paused_kv_tokens=ifb.num_paused_kv_tokens,
             kv_cache_stats=_snapshot_kv_cache_stats(stats.kv_cache_stats),
             cross_kv_cache_stats=_snapshot_kv_cache_stats(stats.cross_kv_cache_stats),
+            specdec_stats=_snapshot_specdec_stats(stats.specdec_stats),
         )
 
     def queue(
@@ -266,6 +310,11 @@ def _make_rank_iter_stats(
             if nested_stats is not None:
                 setattr(stats, attr, nested_stats)
 
+        if payload.specdec_stats is not None:
+            stats.specdec_stats = _restore_specdec_stats(
+                payload.specdec_stats, rank0_stats.specdec_stats
+            )
+
         ifb = InflightBatchingStats()
         ifb.num_context_requests = payload.num_context_requests
         ifb.num_ctx_tokens = payload.num_ctx_tokens

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py

@@ -404,6 +404,17 @@ def _warmup_scheduled(
             and self.guided_decoder is None
             and not isinstance(kv_cache_manager, MambaHybridCacheManager)
         )
+        # A capturable guided decoder rules out the general warmup's generation
+        # requests only: the autotuner below runs its context step with the
+        # decoder installed. Run the context requests anyway, so that the
+        # autotuner's max-token step is not the engine's first forward (with
+        # the buffers it allocates lazily) and runs after the release below,
+        # as it does without guided decoding.
+        can_run_context_warmup = can_run_general_warmup or self._agreement.agree_flag(
+            self._context_warmups_supported
+            and not self.mapping.has_cp_helix()
+            and not isinstance(kv_cache_manager, MambaHybridCacheManager)
+        )
 
         log_mem_snapshot("warmup/before_warmup")
         # Compile the DSv4 indexer-Q CuTe DSL kernels before the first
@@ -425,14 +436,19 @@ def _warmup_scheduled(
             ):
                 self._run_attention_warmup(resource_manager, can_run_general_warmup)
 
-        if can_run_general_warmup:
+        if can_run_context_warmup:
             # Specialize torch.compile graphs across the key input shapes before CUDA graph capture.
             with self._warmup_timer.phase(
                 "general", metrics=self._metrics, metric_name="general_warmup_seconds"
             ):
-                warmup_requests_configs = self._agreement.agree_shapes(
-                    self._get_full_general_warmup_requests(resource_manager)
-                )
+                warmup_requests_configs = self._get_full_general_warmup_requests(resource_manager)
+                if not can_run_general_warmup:
+                    warmup_requests_configs = [
+                        (num_tokens, num_gen_tokens)
+                        for num_tokens, num_gen_tokens in warmup_requests_configs
+                        if num_gen_tokens == 0
+                    ]
+                warmup_requests_configs = self._agreement.agree_shapes(warmup_requests_configs)
                 # Currently graph has not been captured, disable cuda graph for this warmup.
                 with self._graphs.no_cuda_graph():
                     self._general_warmup(resource_manager, warmup_requests_configs)
@@ -900,6 +916,10 @@ def _general_warmup_impl(
                     # Peers are inside the same forward's collectives and
                     # cannot follow a rank-local skip.
                     raise
+                if self.guided_decoder is not None:
+                    # Input preparation enqueues the batch before model forward.
+                    # An OOM may leave it unconsumed, or leave callbacks pending.
+                    self.guided_decoder.reset_warmup_state()
                 logger.warning(
                     f"OOM during general warmup with {num_tokens} tokens, "
                     f"{num_gen_tokens} generation tokens. Skipping."

tensorrt_llm/_torch/pyexecutor/guided_decoder.py

@@ -581,6 +581,20 @@ def _drain_host_functions(self) -> None:
         if not torch.cuda.is_current_stream_capturing():
             self.bitmask_event.synchronize()
 
+    def reset_warmup_state(self) -> None:
+        """Discard abandoned dummy batches after an eager warmup failure.
+
+        Only call during warmup, with no live requests or graph capture. Wait
+        on the stream itself: a failed forward may have enqueued callbacks
+        without reaching the bitmask-event record. Once they finish, no
+        callback can consume or overwrite the cleared request snapshots.
+        Warmup dummies have no grammar matchers to roll back.
+        """
+        self.stream.synchronize()
+        self.queue = Queue()
+        self.requests = None
+        self.requests_hostfunc = None
+
     @nvtx_range("GuidedDecoder.add_batch")
     def add_batch(self,
                   scheduled_requests: ScheduledRequests,

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -6366,6 +6366,17 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests):
                 self._after_shrink(req.py_request_id, kv_cache)
 
             for req in scheduled_batch.generation_requests:
+                # A paired pool's generation dummy has a fixed shape, and
+                # add_dummy_requests reserves its full forward capacity in both
+                # pools, even when it joins the batch after scheduler admission.
+                # An unpaired pool is not handed to add_dummy_requests: the
+                # mirror the scheduler admits for its dummy is empty and grows
+                # below like any other.
+                preallocated_dummy = (
+                    self.enable_joint_kv_cache_reuse
+                    and req.is_dummy_request
+                    and req.py_request_id in self.kv_cache_map
+                )
                 kv_cache = self._mirror_draft_kv_cache(req)
                 if kv_cache is None:
                     raise RuntimeError(
@@ -6382,7 +6393,7 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests):
                     )
                 # A paired scheduler already grew both pools atomically; the
                 # recorded width says so, so do not charge this step twice.
-                if req.py_request_id in self._allocated_draft_lens:
+                if preallocated_dummy or req.py_request_id in self._allocated_draft_lens:
                     continue
                 new_cap = self._required_gen_capacity(req, kv_cache.capacity)
                 if not kv_cache.resize(new_cap):

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -5235,13 +5235,14 @@ def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None:
         if spec_resource_manager is not None:
             spec_resource_manager.free_resources(dummy_request)
         self.kv_cache_manager.free_resources(dummy_request)
-        # The V2 scheduler also admits the draft pool's mirror while it
-        # schedules the dummy (KVCacheV2Scheduler._admit_unpaired_draft), so
-        # that pool holds an empty cache under the dummy ID too. Left behind,
-        # it outlives this dummy, and the deferred previous batch -- which
-        # still holds the earlier, finished dummy with the same ID -- resizes
-        # it to that dummy's history in update_resources ("History length
-        # cannot exceed capacity").
+        # The draft pool holds a cache under the dummy ID too: a paired pool
+        # gets it with the dummy (add_dummy_requests), and the V2 scheduler
+        # admits an unpaired pool's mirror while it schedules the dummy
+        # (KVCacheV2Scheduler._admit_unpaired_draft). Left behind, it outlives
+        # this dummy, and the deferred previous batch -- which still holds the
+        # earlier, finished dummy with the same ID -- resizes it to that
+        # dummy's history in update_resources ("History length cannot exceed
+        # capacity").
         draft_kv_cache_manager = getattr(self, "draft_kv_cache_manager", None)
         if (getattr(self, "_is_kv_manager_v2", False)
                 and draft_kv_cache_manager is not None):
@@ -9572,6 +9573,11 @@ def _pad_attention_dp_dummy_request(self):
                 key="attention_dp_dummy_insufficient_kv_capacity")
             return
 
+        # A paired draft pool admits every scheduled request together with the
+        # target pool, so the dummy needs a cache in both.
+        draft_kv_cache_manager = (self.draft_kv_cache_manager
+                                  if self.enable_joint_kv_cache_reuse else None)
+
         if (not self._enable_adp_dummy_fixes
                 or self.kv_cache_transceiver is None):
             try:
@@ -9580,6 +9586,7 @@ def _pad_attention_dp_dummy_request(self):
                     token_nums=token_nums,
                     is_gen=self._adp_dummy_is_gen,
                     max_num_draft_tokens=self.max_total_draft_tokens,
+                    draft_kv_cache_manager=draft_kv_cache_manager,
                 )
             except OutOfPagesError:
                 dummy_requests = None
@@ -9598,7 +9605,13 @@ def _pad_attention_dp_dummy_request(self):
             spec_resource_manager = self.resource_manager.get_resource_manager(
                 ResourceManagerType.SPEC_RESOURCE_MANAGER)
             if spec_resource_manager is not None:
-                spec_resource_manager.add_dummy_requests(dummy_request_ids)
+                try:
+                    spec_resource_manager.add_dummy_requests(dummy_request_ids)
+                except NoFreeSlotsError:
+                    self.kv_cache_manager.free_resources(llm_request)
+                    if draft_kv_cache_manager is not None:
+                        draft_kv_cache_manager.free_resources(llm_request)
+                    return
             self.active_requests.append(llm_request)
             return
 
@@ -9615,6 +9628,7 @@ def _pad_attention_dp_dummy_request(self):
                 token_nums=token_nums,
                 is_gen=self._adp_dummy_is_gen,
                 max_num_draft_tokens=self.max_total_draft_tokens,
+                draft_kv_cache_manager=draft_kv_cache_manager,
             )
         except OutOfPagesError:
             dummy_requests = None
@@ -9631,6 +9645,8 @@ def _pad_attention_dp_dummy_request(self):
                 spec_resource_manager.add_dummy_requests(dummy_request_ids)
             except NoFreeSlotsError:
                 self.kv_cache_manager.free_resources(dummy_request)
+                if draft_kv_cache_manager is not None:
+                    draft_kv_cache_manager.free_resources(dummy_request)
                 return
 
         assert dummy_request is not None
@@ -9706,6 +9722,11 @@ def _pad_empty_attention_dp_batch(
             return
 
         dummy_request_ids = [ATTENTION_DP_DUMMY_REQUEST_ID]
+        # Allocate the paired draft cache with the dummy, where a shortage
+        # degrades to an empty batch, rather than leaving it to the draft
+        # manager's prepare_resources mirror, which raises on one.
+        draft_kv_cache_manager = (self.draft_kv_cache_manager
+                                  if self.enable_joint_kv_cache_reuse else None)
         try:
             # Degrade to an empty batch rather than propagate: the ranks have
             # yet to agree on can_queue, so a rank-local raise would strand the
@@ -9716,6 +9737,7 @@ def _pad_empty_attention_dp_batch(
                 is_gen=not context_dummy,
                 max_num_draft_tokens=(0 if context_dummy else
                                       self.max_total_draft_tokens),
+                draft_kv_cache_manager=draft_kv_cache_manager,
             )
         except (OutOfPagesError, NoFreeSlotsError):
             dummy_requests = None
@@ -9734,6 +9756,8 @@ def _pad_empty_attention_dp_batch(
                 spec_resource_manager.add_dummy_requests(dummy_request_ids)
             except NoFreeSlotsError:
                 self.kv_cache_manager.free_resources(dummy_request)
+                if draft_kv_cache_manager is not None:
+                    draft_kv_cache_manager.free_resources(dummy_request)
                 return
 
         dummy_request.is_attention_dp_dummy = True