From f4df481dbb4f26c156779af92214877ac8d17334 Mon Sep 17 00:00:00 2001
From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Date: Fri, 17 Jul 2026 13:12:55 +0800
Subject: [PATCH 1/9] [TRTLLM-14054][perf] Reduce per-step host preparation
 overhead in the PyTorch executor decode path (#16313)

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
---
 .../sparse/deepseek_v4/cache_manager.py       |   7 +-
 .../_torch/attention_backend/trtllm.py        |  44 +++-
 .../_torch/modules/mamba/mamba2_metadata.py   |  11 +-
 .../_torch/pyexecutor/kv_cache_manager_v2.py  |   3 +
 .../_torch/pyexecutor/mamba_cache_manager.py  |  13 +-
 .../_torch/pyexecutor/model_engine.py         | 215 ++++++++++++++++++
 .../_torch/pyexecutor/resource_manager.py     |  40 +++-
 .../test_kv_block_offset_overlap_race.py      |  47 ++++
 .../_torch/speculative/test_eagle3.py         |  75 ++++++
 9 files changed, 426 insertions(+), 29 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
index c7ca1891b642..b46d8478ffd6 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
@@ -1301,8 +1301,13 @@ def copy_batch_block_offsets(
         beam_width: int,
         num_contexts: int,
         num_seqs: int,
+        max_blocks: Optional[int] = None,
     ) -> None:
-        """For compatibility with AttentionOp, copy only the SWA block offsets."""
+        """For compatibility with AttentionOp, copy only the SWA block offsets.
+
+        max_blocks is accepted for signature parity with KVCacheManager; the
+        copy below is already bounded by the precomputed SWA table width.
+        """
         assert beam_width == 1, "DSV4 only supports beam width 1 now"
         assert dst_tensor.is_cuda, "copy_batch_block_offsets expects a CUDA destination"
         dst_tensor.fill_(BAD_PAGE_INDEX)
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 931830253cba..40d15970398b 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -31,6 +31,7 @@
 from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned
 from tensorrt_llm.bindings.internal import thop
 from tensorrt_llm.functional import AttentionMaskType
+from tensorrt_llm.math_utils import ceil_div
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
 from ..utils import (compute_swizzled_sf_shape, get_global_attrs,
@@ -586,24 +587,47 @@ def prepare(self) -> None:
         # kv block offsets
         assert self.request_ids is not None
         if self.kv_cache_manager is not None:
-            self.kv_cache_manager.copy_batch_block_offsets(
-                self.kv_cache_block_offsets, self.request_ids, self.beam_width,
-                self.num_contexts, self.num_seqs)
-
-            error_message = (
-                f"The max KV cache length of input sequences ({self.kv_lens[:self.num_seqs].max()}) "
+            max_kv_len = int(self.kv_lens[:self.num_seqs].max())
+            assert max_kv_len <= self.kv_cache_manager.max_seq_len, (
+                f"The max KV cache length of input sequences ({max_kv_len}) "
                 f"exceeds the KV cache manager's maximum supported length "
                 f"({self.kv_cache_manager.max_seq_len}).")
 
-            assert self.kv_lens[:self.num_seqs].max(
-            ) <= self.kv_cache_manager.max_seq_len, error_message
+            # On the non-speculative path the host kv_lens snapshot bounds
+            # every block-table access, so the staged/H2D width can be capped
+            # at the batch's maximum instead of max_seq_len's worth of
+            # columns. Speculative decoding must stage the full width:
+            # draft/tree sub-steps and the overlap scheduler advance
+            # kv_lens_cuda on device past the host snapshot, and their
+            # kernels dereference block columns a host-derived cap would
+            # leave unstaged (uninitialized in this buffer).
+            spec_active = (self.draft_kv_cache_manager is not None
+                           or self.is_spec_decoding_enabled
+                           or bool(self.kv_cache_params.num_extra_kv_tokens) or
+                           (self.runtime_features is not None and
+                            self.runtime_features.has_speculative_draft_tokens))
+            max_blocks = None
+            if not spec_active and self.kv_cache_manager.tokens_per_block:
+                max_blocks = ceil_div(max_kv_len,
+                                      self.kv_cache_manager.tokens_per_block)
+            self.kv_cache_manager.copy_batch_block_offsets(
+                self.kv_cache_block_offsets,
+                self.request_ids,
+                self.beam_width,
+                self.num_contexts,
+                self.num_seqs,
+                max_blocks=max_blocks)
 
             # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists
             if self.draft_kv_cache_manager is not None:
                 # Use the wrapper method which works for both V1 and V2
                 self.draft_kv_cache_manager.copy_batch_block_offsets(
-                    self.draft_kv_cache_block_offsets, self.request_ids,
-                    self.beam_width, self.num_contexts, self.num_seqs)
+                    self.draft_kv_cache_block_offsets,
+                    self.request_ids,
+                    self.beam_width,
+                    self.num_contexts,
+                    self.num_seqs,
+                    max_blocks=max_blocks)
 
         # Don't pass self.kv_lens as kv_lens here because it includes extra
         # tokens. Use the actual KV length (without extra tokens) for
diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py
index 5cab7283f033..9c323799a77a 100644
--- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py
+++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py
@@ -369,9 +369,14 @@ def prepare(self, attn_metadata: AttentionMetadata):
                     self.state_indices_cpu[:batch_size], non_blocking=True)
             else:
                 # indices is a Python sequence (e.g. List[int]); data
-                # already lives on host, CPU staging is fine.
-                for i, idx in enumerate(indices):
-                    self.state_indices_cpu[i] = idx
+                # already lives on host, CPU staging is fine. One bulk
+                # conversion instead of a per-element tensor write.
+                assert len(indices) == batch_size, (
+                    f"get_state_indices() returned {len(indices)} entries for "
+                    f"a batch of {batch_size} requests.")
+                self.state_indices_cpu[:batch_size].copy_(
+                    torch.as_tensor(indices,
+                                    dtype=self.state_indices_cpu.dtype))
                 self.state_indices[:batch_size].copy_(
                     self.state_indices_cpu[:batch_size], non_blocking=True)
 
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index 46d3c55a0771..ffe7b6f5f83e 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -3194,7 +3194,10 @@ def copy_batch_block_offsets(
         beam_width: int,
         num_contexts: int,
         num_seqs: int,
+        max_blocks: Optional[int] = None,
     ):
+        # max_blocks is accepted for signature parity with KVCacheManager; the
+        # device-side copy op here already scales with allocated blocks only.
         assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2"
 
         copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width)
diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
index df137d96de6d..d02e19d5b95c 100644
--- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
@@ -1984,15 +1984,18 @@ def _setup_state_indices(self, requests=None) -> None:
 
         self.cuda_state_indices.copy_(self._host_state_indices,
                                       non_blocking=True)
+        is_dummy = [req.is_dummy for req in requests]
         self._refresh_dummy_request_mask(
-            [req.is_dummy for req in self.requests])
+            is_dummy if requests is
+            self.requests else [req.is_dummy for req in self.requests])
 
         # Build request_id → pool block offset mapping so that
         # get_state_indices can return indices in arbitrary request order.
-        for i, req in enumerate(requests):
-            self._request_id_to_state_index[
-                req.py_request_id] = self._host_state_indices[i].item()
-            self._request_id_to_is_dummy[req.py_request_id] = req.is_dummy
+        # Bulk tolist avoids a per-request tensor-index + .item() round-trip.
+        state_values = self._host_state_indices[:n].tolist()
+        for req, value, dummy in zip(requests, state_values, is_dummy):
+            self._request_id_to_state_index[req.py_request_id] = value
+            self._request_id_to_is_dummy[req.py_request_id] = dummy
 
     def get_state_indices(self,
                           request_ids: Optional[List[int]] = None,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 62a442972830..dff5b770704a 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -675,6 +675,18 @@ def __init__(
         self.position_ids_cuda = torch.empty((self.max_num_tokens, ),
                                              dtype=torch.int,
                                              device='cuda')
+        # Steady-state generation-only prepare cache (non-speculative overlap
+        # decode). Holds the per-request lists that are invariant while the
+        # scheduled generation batch keeps the same composition, plus a pinned
+        # cached-token counter advanced by one per step (host-side bookkeeping
+        # only; the device position buffer is advanced in place and this
+        # buffer is never the source of an async H2D). Invalidated (set to
+        # None) by every full _prepare_tp_inputs pass.
+        self._steady_gen_cache: Optional[Dict[str, Any]] = None
+        self._steady_gen_positions_pinned = torch.empty(
+            (self.max_num_tokens, ),
+            dtype=torch.int,
+            pin_memory=prefer_pinned())
         if self.use_mrope:
             self.mrope_position_ids_cuda = torch.empty(
                 (3, 1, self.max_num_tokens), dtype=torch.int, device='cuda')
@@ -3269,6 +3281,143 @@ def _apply_incremental_update_target(
 
         return inputs, self.gather_ids_cuda[:num_generation_tokens]
 
+    def _can_use_steady_gen_fast_prepare(
+            self, scheduled_requests: ScheduledRequests,
+            new_tokens_device: Optional[torch.Tensor],
+            next_draft_tokens_device: Optional[torch.Tensor],
+            spec_metadata: Optional[SpecMetadata]) -> bool:
+        """Check whether the cached steady-state generation prepare applies.
+
+        The cache is only recorded by a full _prepare_tp_inputs pass whose
+        batch consisted purely of non-dummy generation requests that all had
+        a previous overlap-scheduler tensor (see the recording site), so the
+        per-step check only needs to confirm the dynamic conditions: still a
+        generation-only batch with the exact same requests in the same order.
+        """
+        cache = self._steady_gen_cache
+        if cache is None or self.is_warmup:
+            return False
+        if new_tokens_device is None or next_draft_tokens_device is not None \
+                or spec_metadata is not None:
+            return False
+        if scheduled_requests.num_context_requests > 0:
+            return False
+        generation_requests = scheduled_requests.generation_requests
+        if len(generation_requests) != cache['num_requests']:
+            return False
+        return cache['request_ids'] == [
+            request.py_request_id for request in generation_requests
+        ]
+
+    @nvtx_range("_apply_steady_gen_fast_prepare")
+    def _apply_steady_gen_fast_prepare(
+            self, kv_cache_manager: Union[KVCacheManager, KVCacheManagerV2],
+            attn_metadata: AttentionMetadata,
+            new_tensors_device: SampleStateTensors,
+            resource_manager: Optional[ResourceManager]):
+        """Prepare inputs for an unchanged generation-only batch.
+
+        Every request advanced by exactly one committed token since the last
+        prepare, so instead of re-walking the batch in Python this advances
+        the cached positions in place (device position buffer plus a pinned
+        host counter), reuses the seq-slot buffer already on device, and
+        refreshes only the per-step metadata. For mrope models (recorded only
+        for batches with no actual mrope work) the (3,1,N) broadcast buffer
+        the model reads is the one advanced.
+        """
+        cache = self._steady_gen_cache
+        num_requests = cache['num_requests']
+
+        # Positions and cached-token counts are the same values in this
+        # regime; advance both by one. The device-side position buffer is
+        # advanced in place: it still holds the previous step's positions
+        # because only _prepare_tp_inputs writes it and the cache validity
+        # invariant guarantees the previous pass wrote these same rows. This
+        # avoids reusing a mutated pinned buffer as the source of an async
+        # H2D whose previous-step copy may still be pending under the overlap
+        # scheduler (the nvbug 6293536 hazard class; see
+        # KVCacheManager._stage_block_offsets_for_copy). The pinned buffer is
+        # host-side bookkeeping only.
+        use_mrope = cache['use_mrope']
+        positions = self._steady_gen_positions_pinned[:num_requests]
+        positions.add_(1)
+        if use_mrope:
+            # Text-only batch on an mrope model: the recording pass broadcast
+            # the scalar positions onto all three axes of the (3,1,N) buffer,
+            # which is what the model (and any captured CUDA graph) reads, so
+            # advance it in place. position_ids_cuda is reseeded by the next
+            # full pass.
+            self.mrope_position_ids_cuda[:, :, :num_requests].add_(1)
+        else:
+            self.position_ids_cuda[:num_requests].add_(1)
+        num_cached_tokens_per_seq = positions.tolist()
+
+        # Gather this step's input tokens from the previous iteration's device
+        # sample buffer; the seq-slot indices in previous_batch_indices_cuda
+        # are unchanged since the last full pass.
+        previous_slots = self.previous_batch_indices_cuda[:num_requests]
+        new_tokens = new_tensors_device.new_tokens[:1, previous_slots, :self.
+                                                   max_beam_width]
+        self.input_ids_cuda[:num_requests * self.max_beam_width].copy_(
+            new_tokens.flatten(), non_blocking=True)
+
+        if not attn_metadata.is_cuda_graph:
+            attn_metadata.seq_lens = cache['seq_lens_ones']
+        attn_metadata.beam_width = 1
+        attn_metadata.request_ids = cache['request_ids']
+        attn_metadata.prompt_lens = cache['prompt_lens']
+        attn_metadata.num_contexts = 0
+        attn_metadata.num_chunked_ctx_requests = 0
+        attn_metadata.kv_cache_params = KVCacheParams(
+            use_cache=True,
+            num_cached_tokens_per_seq=num_cached_tokens_per_seq,
+            num_extra_kv_tokens=get_num_extra_kv_tokens(None))
+        attn_metadata.kv_cache_manager = kv_cache_manager
+        if hasattr(self.model.model_config.pretrained_config, 'chunk_size'):
+            attn_metadata.mamba_chunk_size = \
+                self.model.model_config.pretrained_config.chunk_size
+        with nvtx_range("steady_gen_metadata_prepare"):
+            attn_metadata.prepare()
+
+        attn_all_rank_num_tokens = self._get_all_rank_num_tokens(attn_metadata)
+        padded_num_tokens, can_run_piecewise_cuda_graph, attn_all_rank_num_tokens = \
+            self._get_padding_params(num_requests, 0, attn_all_rank_num_tokens)
+        set_per_request_piecewise_cuda_graph_flag(can_run_piecewise_cuda_graph)
+        attn_metadata.padded_num_tokens = (
+            padded_num_tokens if padded_num_tokens != num_requests else None)
+        virtual_num_tokens = num_requests
+        if attn_metadata.padded_num_tokens is not None:
+            self.input_ids_cuda[num_requests:padded_num_tokens].fill_(0)
+            # Zero-fill the padding tail of whichever position layout the
+            # model consumes, matching the full pass.
+            if use_mrope:
+                self.mrope_position_ids_cuda[:, :, num_requests:
+                                             padded_num_tokens].fill_(0)
+            else:
+                self.position_ids_cuda[num_requests:padded_num_tokens].fill_(0)
+            virtual_num_tokens = padded_num_tokens
+
+        self.iter_states['num_ctx_requests'] = 0
+        self.iter_states['num_ctx_tokens'] = 0
+        self.iter_states['num_generation_tokens'] = num_requests
+        self.iter_states['cached_kv_tokens'] = sum(num_cached_tokens_per_seq)
+
+        if use_mrope:
+            final_position_ids = \
+                self.mrope_position_ids_cuda[:, :, :virtual_num_tokens]
+        else:
+            final_position_ids = \
+                self.position_ids_cuda[:virtual_num_tokens].unsqueeze(0)
+        inputs = {
+            'attn_metadata': attn_metadata,
+            'input_ids': self.input_ids_cuda[:virtual_num_tokens],
+            'position_ids': final_position_ids,
+            'inputs_embeds': None,
+            'multimodal_params': [],
+            'resource_manager': resource_manager,
+        }
+        return inputs, None
+
     def _prepare_tp_inputs(
             self,
             scheduled_requests: ScheduledRequests,
@@ -3305,12 +3454,28 @@ def _prepare_tp_inputs(
         if self._can_use_incremental_update(scheduled_requests,
                                             new_tokens_device,
                                             next_draft_tokens_device):
+            # Spec engines never record the steady-gen cache, but invalidate
+            # defensively so the two fast paths can never interleave if the
+            # gates ever evolve.
+            self._steady_gen_cache = None
             return self._apply_incremental_update(
                 scheduled_requests, kv_cache_manager, attn_metadata,
                 spec_metadata, new_tensors_device, cache_indirection_buffer,
                 num_accepted_tokens_device, req_id_to_old_request,
                 resource_manager)
 
+        if self._can_use_steady_gen_fast_prepare(scheduled_requests,
+                                                 new_tokens_device,
+                                                 next_draft_tokens_device,
+                                                 spec_metadata):
+            return self._apply_steady_gen_fast_prepare(kv_cache_manager,
+                                                       attn_metadata,
+                                                       new_tensors_device,
+                                                       resource_manager)
+        # Any full pass invalidates the steady-state cache; it is re-recorded
+        # at the end of this pass when the batch qualifies.
+        self._steady_gen_cache = None
+
         # Hoist self.use_mrope to a function-scope local so the per-request /
         # per-context-request mrope branches use LOAD_FAST instead of LOAD_ATTR.
         _use_mrope = self.use_mrope
@@ -4226,6 +4391,12 @@ def previous_seq_slots_device():
 
         if hasattr(self.model.model_config.pretrained_config, 'chunk_size'):
             attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size
+        # Some sparse backends (RocketKV) clamp
+        # kv_cache_params.num_cached_tokens_per_seq in place during prepare(),
+        # and KVCacheParams holds the list by reference. Snapshot the true
+        # pre-prepare counts so the steady-gen recording below stores values
+        # that the per-step prepare() can re-clamp from scratch.
+        num_cached_tokens_snapshot = list(num_cached_tokens_per_seq)
         attn_metadata.prepare()
         cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs(
             cross_encoder_hidden_states,
@@ -4354,6 +4525,50 @@ def previous_seq_slots_device():
             self.previous_request_ids = all_gen_request_ids
             self.has_previous_device_draft = next_draft_tokens_device is not None
 
+            # Record the steady-state generation cache when this pass handled
+            # purely non-dummy generation requests that all carried a previous
+            # overlap-scheduler tensor (previous_batch_len == _n_gen implies
+            # every request took that branch and none appended input_ids).
+            # While the batch composition holds, the next passes only need to
+            # advance positions by one and refresh per-step metadata.
+            # MRoPE models are supported only for batches with no actual mrope
+            # work (text-only requests, empty mrope lists below): the full
+            # pass routes use_mrope models through the (3,1,N)
+            # mrope_position_ids_cuda layout even then (to keep torch.compile
+            # guards stable), with all three axes equal to the scalar
+            # positions, so the fast path advances that buffer in place and
+            # returns the same layout (see _apply_steady_gen_fast_prepare).
+            if (self.spec_config is None and not self.is_draft_model
+                    and spec_metadata is None and new_tokens_device is not None
+                    and self.guided_decoder is None
+                    and not self.enable_attention_dp and not mrope_position_ids
+                    and not mrope_delta_write_seq_slots
+                    and not mrope_delta_read_seq_slots
+                    and not self.use_beam_search and self.max_beam_width == 1
+                    and not is_enc_dec and not _has_cp_helix
+                    and num_ctx_requests == 0 and not extend_requests
+                    and not first_draft_requests and _n_gen > 0
+                    and previous_batch_len == _n_gen and num_tokens == 0
+                    and not _has_any_multimodal_request
+                    and not multimodal_params_list and not lora_params
+                    and attn_metadata.padded_num_tokens is None
+                    and self._get_position_id_offset() == 0):
+                self._steady_gen_positions_pinned[:_n_gen].copy_(
+                    torch.as_tensor(num_cached_tokens_snapshot,
+                                    dtype=torch.int))
+                self._steady_gen_cache = {
+                    'num_requests':
+                    _n_gen,
+                    'request_ids':
+                    all_gen_request_ids,
+                    'prompt_lens':
+                    prompt_lengths,
+                    'seq_lens_ones':
+                    maybe_pin_memory(torch.ones(_n_gen, dtype=torch.int)),
+                    'use_mrope':
+                    _use_mrope,
+                }
+
         return inputs, self.gather_ids_cuda[:len(
             gather_ids)] if self.enable_spec_decode else None
 
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 5bac35ab06d0..9bcbf9cd57bb 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -2173,9 +2173,13 @@ def _validate_and_adjust_attention_windows(
     def pin_blocks(self, request_id: int):
         self.impl.pin_blocks(request_id)
 
-    def copy_batch_block_offsets(self, dst_tensor: torch.Tensor,
-                                 request_ids: List[int], beam_width: int,
-                                 num_context: int, num_seqs: int):
+    def copy_batch_block_offsets(self,
+                                 dst_tensor: torch.Tensor,
+                                 request_ids: List[int],
+                                 beam_width: int,
+                                 num_context: int,
+                                 num_seqs: int,
+                                 max_blocks: Optional[int] = None):
         # Fill the persistent host buffer in place, exactly as before. CPU-side
         # consumers read self.host_kv_cache_block_offsets directly and depend on
         # its persistent, max_batch-sized layout: DSA sparse attention, the
@@ -2241,23 +2245,39 @@ def copy_batch_block_offsets(self, dst_tensor: torch.Tensor,
         # matching the already-safe kv_lens / block_ids_per_seq staging. The
         # persistent buffer above is untouched by this and stays valid for the
         # synchronous CPU readers.
-        host_block_offsets = self._stage_block_offsets_for_copy(num_seqs)
+        host_block_offsets = self._stage_block_offsets_for_copy(
+            num_seqs, max_blocks)
+        width = host_block_offsets.shape[-1]
         for pool_idx in range(self.num_pools):
-            dst_tensor[pool_idx, :num_seqs].copy_(host_block_offsets[pool_idx],
-                                                  non_blocking=True)
+            dst_tensor[pool_idx, :num_seqs, :, :width].copy_(
+                host_block_offsets[pool_idx], non_blocking=True)
 
-    def _stage_block_offsets_for_copy(self, num_rows: int) -> torch.Tensor:
+    def _stage_block_offsets_for_copy(
+            self,
+            num_rows: int,
+            max_blocks: Optional[int] = None) -> torch.Tensor:
         """Snapshot the first ``num_rows`` rows of the persistent host block
         offset buffer into a fresh pinned buffer, to serve as the private source
-        of an asynchronous H2D copy (nvbug 6293536)."""
+        of an asynchronous H2D copy (nvbug 6293536).
+
+        ``max_blocks`` bounds the copied block width. The buffer is laid out
+        for max_seq_len (max_blocks_per_seq columns) but consumers only read
+        each sequence's allocated block prefix, so a caller that knows the
+        batch's maximum KV length can skip the unused tail — with a large
+        max_seq_len the tail dominates the copy cost."""
+        if max_blocks is None:
+            width = self.max_blocks_per_seq
+        else:
+            width = min(max(max_blocks, 1), self.max_blocks_per_seq)
         host_block_offsets = torch.empty(self.num_pools,
                                          num_rows,
                                          2,
-                                         self.max_blocks_per_seq,
+                                         width,
                                          dtype=torch.int32,
                                          pin_memory=prefer_pinned(),
                                          device='cpu')
-        host_block_offsets.copy_(self.host_kv_cache_block_offsets[:, :num_rows])
+        host_block_offsets.copy_(
+            self.host_kv_cache_block_offsets[:, :num_rows, :, :width])
         return host_block_offsets
 
     def truncate_blocks(self, target_tokens: List[int],
diff --git a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py
index 92aba91ead1b..ff5743430938 100644
--- a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py
+++ b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py
@@ -102,6 +102,53 @@ def _reference_offsets(mgr, ids):
     return dst[:, : len(ids)].clone()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+def test_copy_batch_block_offsets_max_blocks_staging_width():
+    """``max_blocks`` bounds the staged H2D width; ``None`` stages the full
+    allocated width.
+
+    The bounded copy must leave destination columns beyond the cap untouched
+    (callers only pass a cap when nothing reads past it), and the unbounded
+    copy must ship every allocated block: speculative decoding allocates
+    blocks past the host kv_lens snapshot and its kernels dereference those
+    columns, so capping by a host-derived block count corrupted the device
+    block table (EAGLE3 warmup illegal memory access).
+    """
+    mgr = _build_manager()
+
+    ids = list(range(1, 1 + _BATCH))
+    toks = [_TOKENS_PER_SEQ] * _BATCH  # 5 allocated blocks per sequence
+    mgr.add_dummy_requests(request_ids=ids, token_nums=toks, prepare_resource=True)
+    allocated_blocks = _TOKENS_PER_SEQ // _TOKENS_PER_BLOCK
+
+    ref = _reference_offsets(mgr, ids)
+
+    sentinel = -12345
+    capped_width = 2
+    assert capped_width < allocated_blocks <= mgr.max_blocks_per_seq
+
+    dst = torch.full(
+        (mgr.num_pools, 2 * _BATCH, 2, mgr.max_blocks_per_seq),
+        sentinel,
+        dtype=torch.int32,
+        device="cuda",
+    )
+    mgr.copy_batch_block_offsets(dst, ids, 1, _BATCH, _BATCH, max_blocks=capped_width)
+    torch.cuda.synchronize()
+    assert torch.equal(dst[:, :_BATCH, :, :capped_width], ref[..., :capped_width])
+    assert (dst[:, :_BATCH, :, capped_width:] == sentinel).all(), (
+        "bounded staging must not write past the requested block width"
+    )
+
+    dst.fill_(sentinel)
+    mgr.copy_batch_block_offsets(dst, ids, 1, _BATCH, _BATCH, max_blocks=None)
+    torch.cuda.synchronize()
+    assert torch.equal(dst[:, :_BATCH, :, :allocated_blocks], ref[..., :allocated_blocks]), (
+        "max_blocks=None must stage every allocated block, including blocks "
+        "past the batch's current kv length (speculative decoding reads them)"
+    )
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
 def test_copy_batch_block_offsets_survives_overlap_overwrite():
     mgr = _build_manager()
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index 6298565205eb..4e7faecbbea2 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -207,6 +207,81 @@ def test_kv_lens_runtime_with_eagle3_one_model():
         f"kv_lens should be {expected_kv_lens_with_extra.tolist()}, but got {kv_lens_internal.tolist()}"
 
 
+def _make_mock_kv_cache_manager(num_seqs: int) -> MagicMock:
+    mock_kv_cache_manager = MagicMock()
+    mock_kv_cache_manager.tokens_per_block = 32
+    mock_kv_cache_manager.num_pools = 1
+    mock_kv_cache_manager.num_attention_op_pools = 1
+    mock_kv_cache_manager.max_blocks_per_seq = 16
+    mock_kv_cache_manager.max_batch_size = num_seqs
+    mock_kv_cache_manager.max_seq_len = 512
+    mock_kv_cache_manager.copy_batch_block_offsets = MagicMock()
+    return mock_kv_cache_manager
+
+
+@pytest.mark.parametrize("spec_signal", [
+    None, "num_extra_kv_tokens", "is_spec_decoding_enabled",
+    "has_speculative_draft_tokens", "draft_kv_cache_manager"
+])
+def test_block_offsets_staging_width_spec_gate(spec_signal):
+    """prepare() caps the staged block-table width by the batch's max KV
+    length only on the non-speculative path.
+
+    Any speculative-decoding signal must disable the cap (max_blocks=None):
+    spec kernels address block columns past the host kv_lens snapshot
+    (device-side kv_lens advances in draft/tree sub-steps, draft-token blocks
+    are allocated ahead), so a host-derived cap leaves columns they
+    dereference unstaged. Regression test for the EAGLE3 warmup illegal
+    memory access.
+    """
+    num_seqs = 3
+    prompt_lens = [50, 100, 75]
+    seq_lens_q = [1, 1, 1]
+    num_cached_tokens_per_seq = [
+        prompt_lens[i] - seq_lens_q[i] for i in range(num_seqs)
+    ]
+
+    mock_kv_cache_manager = _make_mock_kv_cache_manager(num_seqs)
+    metadata_kwargs = dict(
+        max_num_requests=num_seqs,
+        max_num_tokens=sum(seq_lens_q),
+        kv_cache_manager=mock_kv_cache_manager,
+    )
+    mock_draft_manager = None
+    if spec_signal == "draft_kv_cache_manager":
+        mock_draft_manager = _make_mock_kv_cache_manager(num_seqs)
+        metadata_kwargs["draft_kv_cache_manager"] = mock_draft_manager
+
+    attn_metadata = TrtllmAttentionMetadata(**metadata_kwargs)
+    if spec_signal == "is_spec_decoding_enabled":
+        attn_metadata.is_spec_decoding_enabled = True
+    elif spec_signal == "has_speculative_draft_tokens":
+        attn_metadata.runtime_features.has_speculative_draft_tokens = True
+
+    attn_metadata.request_ids = list(range(1, num_seqs + 1))
+    attn_metadata.prompt_lens = prompt_lens
+    attn_metadata._seq_lens = torch.tensor(seq_lens_q, dtype=torch.int32)
+    attn_metadata._seq_lens_kv = torch.tensor(seq_lens_q, dtype=torch.int32)
+    attn_metadata.kv_cache_params = KVCacheParams(
+        use_cache=True,
+        num_cached_tokens_per_seq=num_cached_tokens_per_seq,
+        num_extra_kv_tokens=(7 if spec_signal == "num_extra_kv_tokens" else 0))
+
+    attn_metadata.prepare()
+
+    if spec_signal is None:
+        # Non-speculative: capped at ceil(max kv len / tokens_per_block).
+        expected_max_blocks = -(-max(prompt_lens) //
+                                mock_kv_cache_manager.tokens_per_block)
+    else:
+        expected_max_blocks = None
+    call_kwargs = mock_kv_cache_manager.copy_batch_block_offsets.call_args.kwargs
+    assert call_kwargs["max_blocks"] == expected_max_blocks
+    if mock_draft_manager is not None:
+        draft_kwargs = mock_draft_manager.copy_batch_block_offsets.call_args.kwargs
+        assert draft_kwargs["max_blocks"] is None
+
+
 @pytest.mark.parametrize(
     "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp,use_hf_speculative_model",
     [

From 459839cc271f33693d666f6fe586e37dbfe2c2a1 Mon Sep 17 00:00:00 2001
From: dongfengy <99041270+dongfengy@users.noreply.github.com>
Date: Thu, 16 Jul 2026 22:31:39 -0700
Subject: [PATCH 2/9] [https://nvbugs/6404567][fix] Clamp piecewise cudagraph
 captures to the reachable ceiling instead of force-appending it (#16256)

Signed-off-by: Dongfeng Yu <dongfengy@nvidia.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
---
 .../_torch/pyexecutor/model_engine.py         | 38 +++++----
 tests/unittest/llmapi/test_llm_args.py        | 85 ++++++++++++++-----
 2 files changed, 84 insertions(+), 39 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index dff5b770704a..d83f51a1bbe3 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -125,29 +125,35 @@ def _filter_piecewise_capture_num_tokens(
 ) -> Tuple[list[int], list[int]]:
     """Cap piecewise CUDA graph capture candidates at the engine's reachable
     `num_tokens` ceiling `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)`
-    and ensure the ceiling itself is captured.
+    clamping user-requested sizes above it down to the ceiling.
 
     Each in-flight request must leave room for at least one decode token,
     so the ceiling is the largest forward-pass `num_tokens` the warmup
-    builder can construct. Including it in the capture set closes the
-    runtime padding gap between the next-largest candidate and the ceiling
-    (otherwise ISLs in that gap have no graph >= them and fall back to
-    eager).
-
-    Returns `(kept, unrecordable)` where `kept` is sorted ascending,
-    deduped, and contains the ceiling whenever it is positive.
+    builder can construct. Candidates above the ceiling cannot be
+    recorded; clamping them down to the ceiling preserves the user's
+    intent (a requested 128 becomes 127 when only 127 is recordable)
+    without inventing capture sizes the user never asked
+    for. Appending sizes beyond the user's list is harmful: runtime
+    padding rounds iterations up to the nearest captured size, so a far
+    appended ceiling (e.g. 65536 over a list topping at 13914) would
+    make every iteration in the gap execute the full ceiling shape.
+
+    Returns `(kept, unrecordable)` where `kept` is sorted ascending and
+    deduped, with above-ceiling candidates clamped to the ceiling.
     `unrecordable` is the sorted unique set of input entries above the
-    ceiling but within `max_num_tokens`.
+    ceiling but within `max_num_tokens` (the clamped ones, reported so
+    the caller's warning fires).
     """
     max_capturable_num_tokens = max(
         0, max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps))
     piecewise_capacity_limit = min(max_num_tokens, max_capturable_num_tokens)
-    kept = sorted(
-        {i
-         for i in candidate_num_tokens if 0 < i <= piecewise_capacity_limit})
-    if piecewise_capacity_limit > 0 and (not kept or kept[-1]
-                                         < piecewise_capacity_limit):
-        kept.append(piecewise_capacity_limit)
+    if piecewise_capacity_limit > 0:
+        kept = sorted({
+            min(i, piecewise_capacity_limit)
+            for i in candidate_num_tokens if 0 < i <= max_num_tokens
+        })
+    else:
+        kept = []
     unrecordable = sorted({
         i
         for i in candidate_num_tokens
@@ -524,7 +530,7 @@ def __init__(
                 f"{unrecordable}: exceeds reachable ceiling "
                 f"max_batch_size*(max_seq_len-1-num_extra_decoding_steps)="
                 f"{max(0, self.batch_size * (self.max_seq_len - 1 - num_extra_decoding_steps))}. "
-                f"Capturing the ceiling itself; raise max_seq_len for larger graphs."
+                f"Clamping them to the ceiling; raise max_seq_len for larger graphs."
             )
 
         try:
diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py
index c77fa26d7fd9..b76bdea0810b 100644
--- a/tests/unittest/llmapi/test_llm_args.py
+++ b/tests/unittest/llmapi/test_llm_args.py
@@ -1368,18 +1368,19 @@ class TestPiecewiseCudaGraphCaptureDefaults:
        powers-of-2 + 256-stride list when `enable_piecewise_cuda_graph`
        is True (and stays `None` otherwise). The fixed list keeps the
        capture set small to bound startup time and CUDA graph memory;
-       the model-engine filter (invariants 2 and 3) ensures the largest
-       reachable size is always captured even when it is not in this
-       default list.
+       the model-engine filter (invariants 2 and 3) clamps out-of-range
+       entries to the reachable ceiling and never invents sizes beyond
+       this list.
     2. `_filter_piecewise_capture_num_tokens` caps the candidate list at
        `max_batch_size * (max_seq_len - 1 - num_extra_decoding_steps)` --
        the largest forward-pass `num_tokens` the warmup builder can
        construct, since every in-flight request must leave room for at
        least one decode token.
-    3. The reachable ceiling itself is always present in the returned
-       capture set (when positive), so runtime ISLs in the gap between
-       the next-largest candidate and the ceiling get a graph rather
-       than falling back to eager.
+    3. Candidates above the reachable ceiling are clamped down to the
+       ceiling (a requested 128 becomes 127), and no size beyond the
+       user's list is ever invented (an appended far ceiling would make
+       runtime padding execute the full ceiling shape for every
+       iteration in the gap).
     """
 
     _EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS = [2**i for i in range(8)] + list(
@@ -1436,14 +1437,50 @@ def test_torch_llm_args_capture_num_tokens_default_when_piecewise_enabled(
         )
         assert args.torch_compile_config.capture_num_tokens == self._EXPECTED_DEFAULT_CAPTURE_NUM_TOKENS
 
-    def test_piecewise_filter_drops_entries_above_reachable_ceiling(self):
-        """Drop candidates above `max_batch_size * (max_seq_len - 1)`.
+    def test_piecewise_filter_never_invents_far_ceiling(self):
+        """A ceiling far above the largest candidate is NOT added.
 
-        Without the cap, the warmup loop would silently skip these entries
-        and the outer padding logic would pad to a target with no captured
-        graph. They must be removed from `kept` and surfaced in
-        `unrecordable` so the warning fires. The ceiling itself is then
-        appended so ISLs in the gap still get a graph.
+        Runtime padding rounds each iteration up to the nearest captured
+        size, so an invented far ceiling (e.g. 65536 over a list topping
+        out at 13914) would make every iteration in the gap execute the
+        full ceiling shape. The filter must never invent sizes the user
+        did not request.
+        """
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
+
+        candidates = [512, 1024, 2048, 4096, 8192, 13914]
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            candidates,
+            max_num_tokens=65536,
+            max_batch_size=896,
+            max_seq_len=32768,
+        )
+        assert kept == candidates
+        assert unrecordable == []
+
+    def test_piecewise_filter_clamps_multiple_oversized_candidates(self):
+        """All above-ceiling candidates collapse to one ceiling entry."""
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            _filter_piecewise_capture_num_tokens
+
+        kept, unrecordable = _filter_piecewise_capture_num_tokens(
+            [64, 120, 128, 200, 256],
+            max_num_tokens=256,
+            max_batch_size=1,
+            max_seq_len=128,
+        )
+        # Ceiling: 1 * (128 - 1) = 127; 128/200/256 clamp to 127, deduped.
+        assert kept == [64, 120, 127]
+        assert unrecordable == [128, 200, 256]
+
+    def test_piecewise_filter_clamps_entries_above_reachable_ceiling(self):
+        """Clamp candidates above `max_batch_size * (max_seq_len - 1)`.
+
+        Entries above the ceiling cannot be recorded by the warmup loop;
+        they are clamped down to the ceiling and surfaced in
+        `unrecordable` so the warning fires. ISLs in the gap still get a
+        graph at the nearest recordable size.
         """
         from tensorrt_llm._torch.pyexecutor.model_engine import \
             _filter_piecewise_capture_num_tokens
@@ -1501,9 +1538,8 @@ def test_piecewise_filter_subtracts_extra_decoding_steps(self):
 
         Drafting loops consume extra decode steps; the filter must mirror
         the `max_seq_len - 1 - num_extra_decoding_steps` constraint
-        applied when warmup requests are built. The ceiling is appended
-        whenever it is strictly greater than the largest surviving
-        candidate.
+        applied when warmup requests are built. Candidates above the
+        reduced ceiling are clamped down to it; nothing is appended.
         """
         from tensorrt_llm._torch.pyexecutor.model_engine import \
             _filter_piecewise_capture_num_tokens
@@ -1517,7 +1553,7 @@ def test_piecewise_filter_subtracts_extra_decoding_steps(self):
             max_seq_len=128,
             num_extra_decoding_steps=5,
         )
-        assert kept[-1] == 122
+        assert kept[-1] == 120  # nothing above the 122 ceiling to clamp
         assert 120 in kept
         assert unrecordable == []
         # Same setup with 9 extra decoding steps -> ceiling 118; 120 drops.
@@ -1565,9 +1601,12 @@ def test_piecewise_filter_returns_empty_when_ceiling_is_zero(self):
         assert kept == []
         assert unrecordable == [1, 2, 4]
 
-    def test_piecewise_filter_appends_ceiling_when_only_smaller_candidates(
-            self):
-        """No candidate near the ceiling -> ceiling still appended."""
+    def test_piecewise_filter_keeps_small_candidates_unchanged(self):
+        """No candidate above the ceiling -> the list is used as-is.
+
+        The ceiling (1016 here) is not appended; iterations above the
+        largest candidate run eagerly at their true size.
+        """
         from tensorrt_llm._torch.pyexecutor.model_engine import \
             _filter_piecewise_capture_num_tokens
 
@@ -1577,8 +1616,8 @@ def test_piecewise_filter_appends_ceiling_when_only_smaller_candidates(
             max_batch_size=8,
             max_seq_len=128,
         )
-        # Ceiling: 8 * (128 - 1) = 1016.
-        assert kept == [1, 2, 4, 8, 1016]
+        # Ceiling: 8 * (128 - 1) = 1016 -- far above max candidate 8.
+        assert kept == [1, 2, 4, 8]
 
 
 class TestTorchLlmArgs:

From ed1a0b9bfad2074b49452aa0597ee9a476cb47a5 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Fri, 17 Jul 2026 14:04:04 +0800
Subject: [PATCH 3/9] [https://nvbugs/6272673][chore] Unwaive DeepSeek V3 Lite
 NVFP4 test (#16398)

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 tests/integration/test_lists/waives.txt | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index ab58cd72d4e4..87872ed30939 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -83,7 +83,6 @@ accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backe
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6384625)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp2pp2-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445472)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6445456)
-accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6272673)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6422432)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-pp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6245394)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=2-tp2pp2-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=False] SKIP (https://nvbugs/6384625)

From a2ab5934caee528b73c7275d18546ea1fc87635e Mon Sep 17 00:00:00 2001
From: Iman Tabrizian <10105175+Tabrizian@users.noreply.github.com>
Date: Fri, 17 Jul 2026 01:12:20 -0700
Subject: [PATCH 4/9] [None][fix] DSA CuteDSL paged-MQA-logits: persistent
 output buffer (CUDA-graph stale-pointer IMA) (#16307)

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
---
 .../_torch/custom_ops/cute_dsl_custom_ops.py  | 24 ++++--
 .../attention/sparse/dsa/test_dsa_indexer.py  | 85 +++++++++++++++++++
 2 files changed, 101 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
index 0f674cb20e67..95ac81158e8d 100644
--- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
+++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
@@ -7241,10 +7241,14 @@ def forward(
             SPLIT_KV = compute_block_kv * 2  # NUM_MATH_WG = 2
             aligned_max_ctx = (
                 (max_context_len + SPLIT_KV - 1) // SPLIT_KV) * SPLIT_KV
-            logits = torch.empty(
-                (B * next_n, aligned_max_ctx),
-                device=q.device,
-                dtype=output_dtype,
+            # Use a persistent arena buffer instead of a per-forward torch.empty
+            # so the output address stays stable across CUDA-graph replays.
+            _reserve = torch.cuda.is_current_stream_capturing()
+            logits = get_memory_buffers().get_buffer(
+                [B * next_n, aligned_max_ctx],
+                output_dtype,
+                buffer_name="cute_dsl_mqa_logits",
+                reserve_buffer=_reserve,
             )
             logits = logits[:, :max_context_len]
 
@@ -8089,10 +8093,14 @@ def forward(
             SPLIT_KV = compute_block_kv * 2  # NUM_MATH_WG = 2
             aligned_max_ctx = (
                 (max_context_len + SPLIT_KV - 1) // SPLIT_KV) * SPLIT_KV
-            logits = torch.empty(
-                (B * next_n, aligned_max_ctx),
-                device=q.device,
-                dtype=output_dtype,
+            # Use a persistent arena buffer instead of a per-forward torch.empty
+            # so the output address stays stable across CUDA-graph replays.
+            _reserve = torch.cuda.is_current_stream_capturing()
+            logits = get_memory_buffers().get_buffer(
+                [B * next_n, aligned_max_ctx],
+                output_dtype,
+                buffer_name="cute_dsl_mqa_logits",
+                reserve_buffer=_reserve,
             )
             logits = logits[:, :max_context_len]
 
diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
index d6929e0506d9..98247a7f1c79 100644
--- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
+++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py
@@ -3165,6 +3165,91 @@ def test_prepare_swaps_and_restore_recovers(self):
         torch.testing.assert_close(meta.host_kv_cache_block_offsets, original_host_offsets)
 
 
+@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
+@skip_pre_blackwell
+def test_cutedsl_mqa_logits_output_buffer_persistent():
+    """Regression: the CuteDSL paged-MQA-logits output must have a STABLE address
+    across calls, not be a per-forward ``torch.empty``.
+
+    At long context the ``[B*next_n, kv_len]`` output is large; as a churning
+    transient it goes stale under CUDA-graph replay when another subsystem
+    co-captured in the same graph (e.g. the MTP / one-model spec sampler)
+    perturbs the shared pool. It must instead be drawn from the reserved
+    ``get_memory_buffers`` arena (like the CuteDSL topk runner), so its address is
+    identical across calls.
+
+    Fails before the fix (two distinct ``torch.empty`` addresses while the first
+    output is kept alive); passes after (same reserved-arena address).
+    """
+    from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers
+
+    batch_size, next_n = 4, 1
+    head_dim, block_size, index_topk = 128, 64, 2048
+    heads = 32
+    kv_len = 4096
+
+    cache_manager, sparse_attn_config = create_dsa_cache_manager(
+        batch_size=batch_size,
+        head_dim=head_dim,
+        tokens_per_block=block_size,
+        max_seq_len=kv_len,
+        num_layers=1,
+        index_topk=index_topk,
+    )
+    create_indexer(sparse_attn_config, layer_idx=0)
+
+    request_ids = list(range(batch_size))
+    kv_lens = torch.full((batch_size,), kv_len, dtype=torch.int32)
+    cache_manager.add_dummy_requests(
+        request_ids=request_ids, token_nums=kv_lens.tolist(), is_gen=False, prepare_resource=True
+    )
+
+    metadata = _create_mock_metadata(
+        request_ids,
+        batch_size,
+        num_contexts=0,
+        num_generations=batch_size,
+        seq_lens=torch.full((batch_size,), next_n, dtype=torch.int32),
+        kv_lens=kv_lens.clone(),
+        num_cached_tokens=[kv_len - next_n] * batch_size,
+        cache_manager=cache_manager,
+        num_ctx_tokens=0,
+        num_tokens=batch_size * next_n,
+        max_draft_tokens=next_n - 1,
+        index_topk=index_topk,
+        use_cute_dsl_paged_mqa_logits=True,
+    )
+    Indexer.prepare(metadata)
+
+    kv_cache = cache_manager.get_indexer_k_cache_buffers(0)
+    q = torch.randn((batch_size, next_n, heads, head_dim), device="cuda", dtype=torch.bfloat16).to(
+        torch.float8_e4m3fn
+    )
+    weights = torch.randn((batch_size * next_n, heads), device="cuda", dtype=torch.float32)
+    context_lens = metadata.gen_indexer_kv_lens_cuda_runtime
+    block_table = metadata.indexer_k_cache_block_offsets[0:batch_size]
+    sched = metadata.scheduler_metadata_buffer
+
+    def _mqa():
+        return torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits(
+            q, kv_cache, weights, context_lens, block_table, sched, kv_len
+        )
+
+    out1 = _mqa()
+    ptr1 = out1.data_ptr()
+    out2 = _mqa()  # out1 kept alive: a fresh torch.empty would land elsewhere
+    ptr2 = out2.data_ptr()
+
+    assert ptr1 == ptr2, (
+        f"CuteDSL mqa-logits output address changed across calls "
+        f"({ptr1:#x} -> {ptr2:#x}); it must be a persistent reserved-arena buffer "
+        f"to avoid stale-pointer IMA under CUDA-graph replay"
+    )
+    assert "cute_dsl_mqa_logits" in get_memory_buffers().buffers, (
+        "CuteDSL mqa-logits output must be drawn from the get_memory_buffers arena"
+    )
+
+
 @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available")
 @skip_pre_hopper
 def test_topk_indices_buffer_cuda_graph():

From 29f4552bfc67502329b668e7afa76164429b0833 Mon Sep 17 00:00:00 2001
From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com>
Date: Fri, 17 Jul 2026 16:25:11 +0800
Subject: [PATCH 5/9] [None][fix] Fix unfused RoPE for yarn models: double
 rotation and GPT-OSS pairing (#16378)

Signed-off-by: tianruih <tianruih@nvidia.com>
---
 .../_torch/models/modeling_gpt_oss.py         |  5 +-
 tensorrt_llm/functional.py                    |  3 +-
 .../_torch/modules/test_rotary_embedding.py   | 80 +++++++++++++++++++
 3 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py
index 00b5c77c8951..252f04ccc517 100644
--- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py
+++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py
@@ -60,7 +60,10 @@ def __init__(
                 beta_fast=pretrained_config.rope_scaling['beta_fast'],
                 beta_slow=pretrained_config.rope_scaling['beta_slow'],
                 duplicate_data=False),
-            is_neox=False,
+            # GPT-OSS applies NeoX-style (rotate-half) RoPE, matching the HF
+            # reference. The fused kernel ignores this flag for yarn (which
+            # masked the wrong value); the unfused path honors it.
+            is_neox=True,
         )
 
         super().__init__(
diff --git a/tensorrt_llm/functional.py b/tensorrt_llm/functional.py
index f4b97746f615..f39b6eac7a03 100644
--- a/tensorrt_llm/functional.py
+++ b/tensorrt_llm/functional.py
@@ -64,7 +64,8 @@ class PositionEmbeddingType(IntEnum):
 
     def is_rope(self) -> bool:
         return self in [
-            self.rope_gptj, self.rope_gpt_neox, self.long_rope, self.mrope
+            self.rope_gptj, self.rope_gpt_neox, self.long_rope, self.mrope,
+            self.yarn
         ]
 
     def is_mrope(self) -> bool:
diff --git a/tests/unittest/_torch/modules/test_rotary_embedding.py b/tests/unittest/_torch/modules/test_rotary_embedding.py
index cc20dee61a45..13a861c0a494 100644
--- a/tests/unittest/_torch/modules/test_rotary_embedding.py
+++ b/tests/unittest/_torch/modules/test_rotary_embedding.py
@@ -245,3 +245,83 @@ def test_with_qk_rope_head_dim(self):
         rp = RopeParams.from_config(self._make_config(qk_rope_head_dim=64))
         assert rp.duplicate_data is True
         assert rp.dim == 64
+
+
+class TestUnfusedRopeOwnership:
+    """With rope_fusion=False the Python rotary module owns RoPE; the backend
+    must receive no position-embedding params. yarn is not listed in
+    PositionEmbeddingType.is_rope(), which used to leak the params through and
+    made the attention kernel rotate a second time (double RoPE)."""
+
+    def test_unfused_yarn_rope_is_applied_exactly_once(self):
+        from tensorrt_llm._torch.attention_backend.interface import \
+            PositionalEmbeddingParams
+        from tensorrt_llm._torch.model_config import ModelConfig
+        from tensorrt_llm._torch.modules.attention import Attention
+        from tensorrt_llm.functional import PositionEmbeddingType
+
+        yarn_params = PositionalEmbeddingParams(
+            type=PositionEmbeddingType.yarn,
+            rope=RopeParams(
+                dim=32,
+                theta=150000,
+                scale_type=RotaryScalingType.yarn,
+                scale=32.0,
+                max_positions=1024,
+                original_max_positions=256,
+                beta_fast=32,
+                beta_slow=1,
+                duplicate_data=False,
+            ),
+            is_neox=True,
+        )
+        attn = Attention(
+            hidden_size=256,
+            num_attention_heads=8,
+            num_key_value_heads=8,
+            max_position_embeddings=1024,
+            bias=False,
+            pos_embd_params=yarn_params,
+            rope_fusion=False,
+            layer_idx=0,
+            dtype=torch.bfloat16,
+            config=ModelConfig(),
+        )
+
+        assert attn.rotary_emb is not None
+        # 0 means the kernel side received no position embedding.
+        assert attn.attn.position_embedding_type == 0
+
+    def test_unfused_yarn_rotation_matches_fused_kernel_convention(self):
+        """The unfused (Python) yarn rotation must match the NeoX rotate-half
+        convention the fused kernel applies with the same cos/sin table."""
+        head_dim = 64
+        num_pos, num_heads = 64, 4
+        rope_params = RopeParams(
+            dim=head_dim,
+            theta=150000,
+            scale_type=RotaryScalingType.yarn,
+            scale=32.0,
+            max_positions=1024,
+            original_max_positions=256,
+            beta_fast=32,
+            beta_slow=1,
+            duplicate_data=False,
+        )
+        emb = RotaryEmbedding(rope_params, head_dim=head_dim, is_neox=True)
+        torch.manual_seed(0)
+        q = torch.randn(num_pos, num_heads * head_dim)
+        positions = torch.arange(num_pos).cuda()
+        # Single-target call takes the pure-torch path.
+        q_unfused = emb(positions, [q.cuda()])[0].cpu()
+
+        # Reference: NeoX rotate-half with the exact table the kernel reads.
+        table = emb.rotary_cos_sin.cpu()  # (max_pos, 2, head_dim/2)
+        cos = table[:num_pos, 0, :].unsqueeze(1)
+        sin = table[:num_pos, 1, :].unsqueeze(1)
+        qh = q.view(num_pos, num_heads, head_dim)
+        q1, q2 = qh[..., :head_dim // 2], qh[..., head_dim // 2:]
+        q_ref = torch.cat((q1 * cos - q2 * sin, q2 * cos + q1 * sin),
+                          dim=-1).reshape(num_pos, -1)
+
+        torch.testing.assert_close(q_unfused, q_ref, rtol=1e-5, atol=1e-5)

From 9fdfd10c17610b9033b6a8807e33d9a6b87a7389 Mon Sep 17 00:00:00 2001
From: yufeiwu-nv <230315618+yufeiwu-nv@users.noreply.github.com>
Date: Fri, 17 Jul 2026 16:56:29 +0800
Subject: [PATCH 6/9] [None][test] Update LLM performance test cases to use
 increased input/output lengths (#16538)

Signed-off-by: yufeiwu-nv <230315618+yufeiwu-nv@users.noreply.github.com>
---
 .../test_lists/qa/llm_perf_core.yml           | 128 +++++++++---------
 1 file changed, 63 insertions(+), 65 deletions(-)

diff --git a/tests/integration/test_lists/qa/llm_perf_core.yml b/tests/integration/test_lists/qa/llm_perf_core.yml
index 7a5fd834aefb..792d4fafdff4 100644
--- a/tests/integration/test_lists/qa/llm_perf_core.yml
+++ b/tests/integration/test_lists/qa/llm_perf_core.yml
@@ -28,14 +28,12 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:128,128]
   - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:500,2000]
   - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:2000,500]
-  - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:1000,1000]
   - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:1000,2000]
-  - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_9b-bench-pytorch-bfloat16-input_output_len:8000,1000]
   - perf/test_perf.py::test_perf[qwen3_4b_eagle3-bench-pytorch-streaming-bfloat16-maxbs:4-kv_frac:0.6-input_output_len:500,100-reqs:200-con:4]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_nano_8b_fp8-bench-pytorch-float8-maxnt:5000-input_output_len:5000,500-reqs:8-con:1]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_nano_8b_fp8-bench-pytorch-float8-input_output_len:500,2000-reqs:8-con:1]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_nano_8b_fp8-bench-pytorch-float8-input_output_len:1000,1000-reqs:8-con:1]
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_nano_8b_fp8-bench-pytorch-float8-input_output_len:8000,1000-reqs:8-con:1]
    # test overlap scheduler
    # con:1 paired with a small model is an intentional design choice—it amplifies host-side overhead and simplifies execution timelines to the maximum extent.
   - perf/test_perf.py::test_perf[qwen3_0.6b-bench-pytorch-bfloat16-maxnt:2048-input_output_len:8000,1000-reqs:256-con:1-pp:4-gpus:4]
@@ -53,7 +51,7 @@ llm_perf_core:
   #nemotron_nano_12b_v2
   - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-bfloat16-maxbs:1-input_output_len:128,128-reqs:10-con:1] #min_latency
   #qwen3.5_27b (dense BF16 52G, 2-GPU)
-  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:2-gpus:2] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:2-gpus:2] #min_latency
   - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct-bench-pytorch-streaming-bfloat16-input_output_len:128,128-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct-bench-pytorch-bfloat16-input_output_len:128,128-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct_fp8-bench-pytorch-streaming-float8-input_output_len:2000,200-gpus:8]
@@ -62,9 +60,9 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:128,128-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:500,2000-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:2000,500-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:1000,1000-tp:4-gpus:4]
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:8000,1000-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:1000,2000-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250-tp:4-gpus:4] #max_throughput
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:8000,1000-con:250-tp:4-gpus:4] #max_throughput
 
 
 # 3: H100, H20 test cases
@@ -89,17 +87,17 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:128,128-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:500,2000-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:2000,500-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:1000,1000-tp:2-gpus:2]
+  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:8000,1000-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:1000,2000-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250-tp:2-gpus:2] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:8000,1000-con:250-tp:2-gpus:2] #max_throughput
   #llama_v3.3_nemotron_super_49b (nemotron-nas BF16 94G, 2-GPU)
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:128,128-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:500,2000-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:2000,500-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:1000,1000-tp:2-gpus:2]
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:8000,1000-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:1000,2000-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:2-gpus:2] #min_latency
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250-tp:2-gpus:2] #max_throughput
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:2-gpus:2] #min_latency
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b-bench-pytorch-bfloat16-input_output_len:8000,1000-con:250-tp:2-gpus:2] #max_throughput
 
 
 # 4: H100, H20, GB200, B200, B300, GB300, RTX6000-Server test cases
@@ -115,40 +113,40 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[qwen3_235b_a22b_fp8-bench-pytorch-float8-input_output_len:1000,2000-con:256-ep:8-gpus:8]
   #nemotron_nano_12b_v2
   - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-bfloat16-input_output_len:512,512]
-  - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1] #min_latency
-  - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250] #max_throughput
+  - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-bfloat16-maxbs:1-input_output_len:8000,1000-reqs:10-con:1] #min_latency
+  - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-bfloat16-input_output_len:8000,1000-con:250] #max_throughput
   - perf/test_perf.py::test_perf[nemotron_nano_12b_v2-bench-pytorch-streaming-bfloat16-input_output_len:500,2000-con:250] #max_throughput streaming
   #qwen3.5_27b (dense BF16 52G, 1-GPU)
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:128,128]
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:500,2000]
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:2000,500]
-  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:1000,1000]
+  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:8000,1000]
   - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:1000,2000]
-  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-maxbs:1-input_output_len:8000,1000-reqs:10-con:1] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_27b-bench-pytorch-bfloat16-input_output_len:8000,1000-con:250] #max_throughput
   #qwen3.5_35b_a3b_fp8 (MoE FP8 36G, 1-GPU)
   - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-input_output_len:128,128]
   - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-input_output_len:500,2000]
   - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-input_output_len:2000,500]
-  - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-input_output_len:1000,1000]
+  - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-input_output_len:8000,1000]
   - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-input_output_len:1000,2000]
-  - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-maxbs:512-input_output_len:1000,1000-con:256] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_35b_a3b_fp8-bench-pytorch-float8-maxbs:512-input_output_len:8000,1000-con:256] #max_throughput
   #llama_v3.3_nemotron_super_49b_fp8 (nemotron-nas FP8 49G, 2-GPU for safety)
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:128,128-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:500,2000-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:2000,500-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:1000,1000-tp:2-gpus:2]
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:8000,1000-tp:2-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:1000,2000-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:2-gpus:2] #min_latency
-  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:1000,1000-con:250-tp:2-gpus:2] #max_throughput
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:2-gpus:2] #min_latency
+  - perf/test_perf.py::test_perf[llama_v3.3_nemotron_super_49b_fp8-bench-pytorch-float8-input_output_len:8000,1000-con:250-tp:2-gpus:2] #max_throughput
     #qwen3.5_122b_a10b (MoE BF16 234G, 4-GPU)
   - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:500,2000-ep:4-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:2000,500-ep:4-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:1000,1000-ep:4-tp:4-gpus:4]
+  - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:8000,1000-ep:4-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:1000,2000-ep:4-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:4-tp:4-gpus:4] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-maxbs:512-input_output_len:1000,1000-con:256-ep:4-tp:4-gpus:4] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:4-tp:4-gpus:4] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-maxbs:512-input_output_len:8000,1000-con:256-ep:4-tp:4-gpus:4] #max_throughput
 
 
 
@@ -180,7 +178,7 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:256-maxnt:1024-kv_frac:0.85-input_output_len:2000,2000-reqs:200-ep:4-tp:4-gpus:4] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:1000-maxnt:5000-kv_frac:0.85-input_output_len:5000,500-reqs:2000-ep:4-tp:4-gpus:4] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:32-maxnt:32768-input_output_len:8192,1024-reqs:20-con:1-ep:1-tp:4-gpus:4] TIMEOUT(120)
-  - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-kv_frac:0.85-input_output_len:1000,1000-reqs:2000-ep:4-tp:4-gpus:4] TIMEOUT(120)
+  - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-kv_frac:0.85-input_output_len:8000,1000-reqs:2000-ep:4-tp:4-gpus:4] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-kv_frac:0.85-input_output_len:1000,2000-reqs:3000-ep:4-tp:4-gpus:4] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:1000-maxnt:5000-kv_frac:0.85-input_output_len:5000,500-reqs:20000-ep:4-tp:4-gpus:4] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:512-input_output_len:128,128-ep:4-tp:4-gpus:4]
@@ -189,29 +187,29 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:128,128-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:500,2000-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:2000,500-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:1000,1000-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:4-gpus:4] #min_latency
-  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:512-input_output_len:1000,1000-con:256-tp:4-gpus:4] #max_throughput
+  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:8000,1000-tp:4-gpus:4]
+  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:4-gpus:4] #min_latency
+  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:512-input_output_len:8000,1000-con:256-tp:4-gpus:4] #max_throughput
   #deepseek_v3.2_fp4 (FP4 389G, 4-GPU)
   - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-input_output_len:128,128-ep:4-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:4-tp:4-gpus:4] #min_latency
-  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:512-kv_frac:0.85-input_output_len:1000,1000-con:512-ep:4-tp:4-gpus:4] TIMEOUT(120) #max_throughput
+  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:4-tp:4-gpus:4] #min_latency
+  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:512-kv_frac:0.85-input_output_len:8000,1000-con:512-ep:4-tp:4-gpus:4] TIMEOUT(120) #max_throughput
   #llama_v3.1_nemotron_ultra_253b_fp8 (nemotron-nas FP8 241G, 4-GPU)
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:128,128-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:500,2000-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:2000,500-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:1000,1000-tp:4-gpus:4]
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:8000,1000-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:1000,2000-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:4-gpus:4] #min_latency
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:1000,1000-con:250-tp:4-gpus:4] #max_throughput
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:4-gpus:4] #min_latency
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:8000,1000-con:250-tp:4-gpus:4] #max_throughput
   #qwen3.5_397b_a17b_fp4 (MoE FP4 234G, 4-GPU)
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:128,128-ep:4-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:500,2000-ep:4-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:2000,500-ep:4-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:1000,1000-ep:4-tp:4-gpus:4]
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:8000,1000-ep:4-tp:4-gpus:4]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:1000,2000-ep:4-tp:4-gpus:4]
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:4-tp:4-gpus:4] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:512-input_output_len:1000,1000-con:512-ep:4-tp:4-gpus:4] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:4-tp:4-gpus:4] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:512-input_output_len:8000,1000-con:512-ep:4-tp:4-gpus:4] #max_throughput
   #nemotron_3_super_120b_nvfp4 (Hybrid MoE+SSM+Attn FP4 76G, 4-GPU ep=4 tp=4, throughput config)
   #these test config come from docs/source/deployment-guide/deployment-guide-for-nemotron-3-on-trtllm.md
   - perf/test_perf.py::test_perf[nemotron_3_super_120b_nvfp4-serve-pytorch-float4-maxbs:512-maxnt:2048-kv_frac:0.8-input_output_len:1024,1024-reqs:5-con:1-ep:4-tp:4-gpus:4] #min_latency
@@ -243,30 +241,30 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[gpt_oss_120b_fp4-bench-pytorch-float4-maxbs:720-maxnt:16384-input_output_len:1024,1024-reqs:100-con:32-ep:8-tp:8-gpus:8]
 
   # deepseek_r1_0528_fp4
-  - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-kv_frac:0.85-input_output_len:1000,1000-reqs:20000-ep:8-tp:8-gpus:8] TIMEOUT(120)
+  - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-kv_frac:0.85-input_output_len:8000,1000-reqs:20000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-kv_frac:0.85-input_output_len:1000,2000-reqs:3000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   #kimi_k2.5_fp4 (multimodal MoE FP4 553G, 8-GPU ep=8)
   - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-input_output_len:128,128-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-input_output_len:500,2000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-input_output_len:2000,500-ep:8-tp:8-gpus:8] TIMEOUT(120)
-  - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-input_output_len:1000,1000-ep:8-tp:8-gpus:8] TIMEOUT(120)
+  - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-input_output_len:8000,1000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-input_output_len:1000,2000-ep:8-tp:8-gpus:8] TIMEOUT(120)
-  - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-maxbs:512-maxnt:2048-kv_frac:0.6-input_output_len:1000,1000-con:512-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
+  - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[kimi_k2.5_fp4-bench-pytorch-float4-maxbs:512-maxnt:2048-kv_frac:0.6-input_output_len:8000,1000-con:512-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
   #glm_5_fp8 (MoE FP8 708G, 8-GPU ep=8)
   - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-input_output_len:128,128-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-input_output_len:500,2000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-input_output_len:2000,500-ep:8-tp:8-gpus:8] TIMEOUT(120)
-  - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-input_output_len:1000,1000-ep:8-tp:8-gpus:8] TIMEOUT(120)
+  - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-input_output_len:8000,1000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-tp:8-gpus:8] TIMEOUT(120)
-  - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-kv_frac:0.6-input_output_len:1000,1000-con:512-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
+  - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[glm_5_fp8-bench-pytorch-float8-maxbs:512-maxnt:2048-kv_frac:0.6-input_output_len:8000,1000-con:512-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
   #deepseek_v3.2_fp4 (FP4 389G, 8-GPU ep=8)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp4-bench-pytorch-float4-maxbs:384-maxnt:1536-input_output_len:1000,2000-reqs:10000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max throughput test
   - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-input_output_len:128,128-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-input_output_len:1000,2000-ep:8-tp:8-gpus:8] TIMEOUT(120)
-  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:384-maxnt:1536-input_output_len:1000,1000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
+  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[deepseek_v3.2_fp4-bench-pytorch-float4-maxbs:384-maxnt:1536-input_output_len:8000,1000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
 
 
 # 8: H100, H20, B200, B300, RTX6000-Server test cases
@@ -285,26 +283,26 @@ llm_perf_core:
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:128,128-ep:8-gpus:8]
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:500,2000-ep:8-gpus:8]
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:2000,500-ep:8-gpus:8]
-  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:1000,1000-ep:8-gpus:8]
+  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:8000,1000-ep:8-gpus:8]
   - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-gpus:8]
-  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:512-input_output_len:1000,1000-con:512-ep:8-gpus:8] #max_throughput
+  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[minimax_m2.5_fp8-bench-pytorch-float8-maxbs:512-input_output_len:8000,1000-con:512-ep:8-gpus:8] #max_throughput
   #llama_v3.1_nemotron_ultra_253b_fp8 (nemotron-nas FP8 241G, 8-GPU)
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:128,128-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:500,2000-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:2000,500-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:1000,1000-tp:8-gpus:8]
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:8000,1000-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:1000,2000-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:1000,1000-con:250-tp:8-gpus:8] #max_throughput
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b_fp8-bench-pytorch-float8-input_output_len:8000,1000-con:250-tp:8-gpus:8] #max_throughput
   #qwen3.5_397b_a17b_fp8 (MoE FP8 380G, 8-GPU ep=8)
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:128,128-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:500,2000-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:2000,500-ep:8-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:1000,1000-ep:8-tp:8-gpus:8]
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:8000,1000-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-input_output_len:1000,2000-ep:8-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-maxbs:512-input_output_len:1000,1000-con:512-ep:8-tp:8-gpus:8] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp8-bench-pytorch-float8-maxbs:512-input_output_len:8000,1000-con:512-ep:8-tp:8-gpus:8] #max_throughput
   - perf/test_perf.py::test_perf[qwen3.5_122b_a10b-bench-pytorch-bfloat16-input_output_len:128,128-ep:4-tp:4-gpus:4]
 
 
@@ -320,27 +318,27 @@ llm_perf_core:
         gt: 90000
   tests:
   # deepseek_r1_0528
-  - perf/test_perf.py::test_perf[deepseek_r1_0528_fp8-bench-pytorch-float8-input_output_len:1000,1000-reqs:20000-ep:8-tp:8-gpus:8] TIMEOUT(120)
+  - perf/test_perf.py::test_perf[deepseek_r1_0528_fp8-bench-pytorch-float8-input_output_len:8000,1000-reqs:20000-ep:8-tp:8-gpus:8] TIMEOUT(120)
   - perf/test_perf.py::test_perf[deepseek_r1_0528_fp8-bench-pytorch-float8-input_output_len:1000,2000-reqs:3000-ep:8-tp:8-gpus:8] TIMEOUT(100)
   #deepseek_v3.2_fp8 (FP8 645G, 8-GPU ep=8)
   - perf/test_perf.py::test_perf[deepseek_v3.2_fp8-bench-pytorch-float8-input_output_len:128,128-ep:8-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[deepseek_v3.2_fp8-bench-pytorch-float8-maxbs:384-maxnt:1536-input_output_len:1000,1000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
+  - perf/test_perf.py::test_perf[deepseek_v3.2_fp8-bench-pytorch-float8-maxbs:384-maxnt:1536-input_output_len:8000,1000-con:3072-ep:8-tp:8-gpus:8] TIMEOUT(120) #max_throughput
   #qwen3.5_397b_a17b_fp4 (MoE FP4 234G, 8-GPU ep=8)
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:128,128-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:500,2000-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:2000,500-ep:8-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:1000,1000-ep:8-tp:8-gpus:8]
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:8000,1000-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-input_output_len:1000,2000-ep:8-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:512-input_output_len:1000,1000-con:512-ep:8-tp:8-gpus:8] #max_throughput
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-ep:8-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[qwen3.5_397b_a17b_fp4-bench-pytorch-float4-maxbs:512-input_output_len:8000,1000-con:512-ep:8-tp:8-gpus:8] #max_throughput
    #llama_v3.1_nemotron_ultra_253b (nemotron-nas BF16 474G, 8-GPU)
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:128,128-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:500,2000-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:2000,500-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:1000,1000-tp:8-gpus:8]
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:8000,1000-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:1000,2000-tp:8-gpus:8]
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-maxbs:1-input_output_len:1000,1000-reqs:10-con:1-tp:8-gpus:8] #min_latency
-  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:1000,1000-con:250-tp:8-gpus:8] #max_throughput
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-maxbs:1-input_output_len:8000,1000-reqs:10-con:1-tp:8-gpus:8] #min_latency
+  - perf/test_perf.py::test_perf[llama_v3.1_nemotron_ultra_253b-bench-pytorch-bfloat16-input_output_len:8000,1000-con:250-tp:8-gpus:8] #max_throughput
 
 # 10: RTX-6000 Server test cases
 - condition:
@@ -354,7 +352,7 @@ llm_perf_core:
   #llama_v3.3_70b
   - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct-bench-pytorch-bfloat16-maxbs:1-input_output_len:128,128-reqs:10-gpus:2]
   - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct_fp4-bench-pytorch-float4-input_output_len:128,128-tp:2-gpus:2]
-  - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct_fp4-bench-pytorch-float4-maxbs:1024-maxnt:4096-kv_frac:0.85-input_output_len:1000,1000-reqs:3000-tp:8-gpus:8]
+  - perf/test_perf.py::test_perf[llama_v3.3_70b_instruct_fp4-bench-pytorch-float4-maxbs:1024-maxnt:4096-kv_frac:0.85-input_output_len:8000,1000-reqs:3000-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3_235b_a22b_fp4-bench-pytorch-float4-input_output_len:1000,2000-con:8-ep:8-tp:8-gpus:8]
   - perf/test_perf.py::test_perf[qwen3_235b_a22b_fp4-bench-pytorch-float4-input_output_len:1000,2000-con:512-ep:8-tp:8-gpus:8]
   # deepseek_r1_0528

From fd7425741343d7584699042103446bc295a6e12b Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Fri, 17 Jul 2026 02:13:37 -0700
Subject: [PATCH 7/9] [https://nvbugs/6336747][fix] Honor GenerationResult
 timeout (#16453)

* Why?

While investigating NVBug 6336747, we found a cluster epilog could
remove PyTorch shared-memory files belonging to another same-account
job on the node.

The resulting worker abort exposed a separate API bug: when no terminal
response arrived, we could could still wait on the result indefinitely
because it ignored the timeout.

The cluster cleanup issue is being fixed independently. Honoring the
existing timeout contract gives callers a reliable bound when an
executor stops producing responses.

* What?

Apply timeout as a total budget across all response steps, measured
with a monotonic clock. Pass the remaining budget to both queue
implementations and raise `TimeoutError` when it expires. Preserve
nonblocking `timeout=0 `behavior for responses already queued, and add
focused coverage for these paths.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 tensorrt_llm/executor/result.py         | 24 +++++++--
 tensorrt_llm/llmapi/utils.py            | 16 ++++--
 tests/integration/test_lists/waives.txt |  4 --
 tests/unittest/llmapi/test_executor.py  | 71 +++++++++++++++++++++++++
 4 files changed, 104 insertions(+), 11 deletions(-)

diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index a3014da7072f..7870becd21f8 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -995,7 +995,11 @@ def _handle_ray_response(self, response: Any):
         return response
 
     def _result_step(self, timeout: Optional[float] = None):
-        response = self.queue.get()
+        # Honor `timeout`: a bounded `queue.get()` lets the caller regain control if the executor
+        # worker dies silently without pushing a terminal response, instead of blocking potentially
+        # indefinitely.
+        # Raises `queue.Empty` on timeout; `result()` turns that into a `TimeoutError`.
+        response = self.queue.get(timeout=timeout)
         # Fast-fail: when a worker dies, the proxy enqueues EngineDeadError onto
         # every pending result so this get() unblocks instead of hanging forever
         # on a queue whose producer is gone. Record it as the sticky terminal
@@ -1022,15 +1026,29 @@ def result(self, timeout: Optional[float] = None) -> "GenerationResult":
         """Wait for the completion of the request, and return the result.
 
         Args:
-            timeout (float, optional): Timeout. Defaults to None.
+            timeout (float, optional): The maximum number of seconds to wait for the request to
+                complete. `None` (default) waits indefinitely.
+                The timeout is a total budget across all streaming steps, not per-step.
 
         Returns:
             tensorrt_llm.executor.result.GenerationResult: generation result.
+
+        Raises:
+            TimeoutError: If the request does not complete within `timeout` seconds. Bounding the
+                wait prevents a silently-dead executor worker from hanging the caller forever.
         """
         if self._terminal_error is not None:
             raise self._terminal_error
+        deadline = None if timeout is None else time.monotonic() + timeout
         while not self._done:
-            self._result_step(timeout)
+            remaining = (None if deadline is None else max(
+                0.0, deadline - time.monotonic()))
+            try:
+                self._result_step(remaining)
+            except Empty:
+                raise TimeoutError(
+                    f"Request {self.request_id} did not complete within "
+                    f"{timeout} seconds.") from None
         return self
 
     async def aresult(self) -> "GenerationResult":
diff --git a/tensorrt_llm/llmapi/utils.py b/tensorrt_llm/llmapi/utils.py
index f1b49fc87af8..acfc2decd9d2 100644
--- a/tensorrt_llm/llmapi/utils.py
+++ b/tensorrt_llm/llmapi/utils.py
@@ -18,7 +18,7 @@
 from contextlib import nullcontext
 from functools import wraps
 from pathlib import Path
-from queue import Queue
+from queue import Empty, Queue
 from typing import (Any, Callable, ContextManager, Iterable, List, Optional,
                     Tuple, Type, get_type_hints)
 
@@ -534,12 +534,20 @@ def get(self, timeout=None):
 
         # We can't call asyncio.run_coroutine_threadsafe(self._aq.get(), self.loop) and wait the returned Future,
         # since we are in the same event loop, and we can't yield the thread while waiting result.
-        deadline = None if timeout is None else time.time() + timeout
-        while deadline is None or time.time() < deadline:
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
             try:
                 return self._aq.unsafe_get()
             except asyncio.QueueEmpty:
-                time.sleep(0.01)
+                if deadline is not None:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        # Match `queue.Queue.get()` semantics; a silent `None` return would be
+                        # mis-handled downstream as an unknown response type.
+                        raise Empty() from None
+                    time.sleep(min(0.01, remaining))
+                else:
+                    time.sleep(0.01)
 
 
 def get_numa_aware_cpu_affinity(device_id):
diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index 87872ed30939..41524c7704ac 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -23,8 +23,6 @@ accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen
 accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=False] SKIP (https://nvbugs/6427411)
 accuracy/test_disaggregated_serving.py::TestQwen3NextInstruct::test_auto_dtype[use_py_transceiver=True] SKIP (https://nvbugs/6402054)
 accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5748664)
-accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] SKIP (https://nvbugs/6336747)
-accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[qwen3vl_2b_instruct] SKIP (https://nvbugs/6422294)
 accuracy/test_llm_api_autodeploy.py::TestNemotronSuperV3::test_mtp[nvfp4_ws8_80gb-trtllm] SKIP (https://nvbugs/6450341)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp] SKIP (https://nvbugs/6428101)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[throughput_mtp_trtllm] SKIP (https://nvbugs/6426868)
@@ -132,8 +130,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] SKIP (https:/
 accuracy/test_llm_api_pytorch.py::TestStep3_7::test_fp8_block_scales[tp_size=4-ep_size=4-mtp_nextn=3] SKIP (https://nvbugs/6367805)
 accuracy/test_llm_api_pytorch.py::TestStep3_7::test_nvfp4[tp_size=4-ep_size=4-mtp_nextn=3] SKIP (https://nvbugs/6367805)
 accuracy/test_llm_api_pytorch_multimodal.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6248827)
-accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[fp8_mmmu_encoder_cuda_graph] SKIP (https://nvbugs/6336747)
-accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[nvfp4] SKIP (https://nvbugs/6336747)
 accuracy/test_llm_api_pytorch_multimodal.py::TestStep3_7::test_nvfp4[mtp_nextn=3] SKIP (https://nvbugs/6367805)
 accuracy/test_llm_api_pytorch_ray.py::TestLlama3_1_8BInstruct::test_pp2_ray SKIP (https://nvbugs/6427411)
 cpp/test_e2e.py::test_benchmarks[bart-90] SKIP (https://nvbugs/5550689)
diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py
index 1b11587f7e15..85bd42b50f63 100644
--- a/tests/unittest/llmapi/test_executor.py
+++ b/tests/unittest/llmapi/test_executor.py
@@ -2,8 +2,10 @@
 import datetime
 import tempfile
 import threading
+import time
 from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
+from queue import Empty
 
 import pytest
 import torch
@@ -17,6 +19,7 @@
                                    GenerationResultBase, PostprocWorker)
 from tensorrt_llm.executor.ipc import FusedIpcQueue, ZeroMqQueue
 from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer
+from tensorrt_llm.llmapi.utils import AsyncQueue
 from tensorrt_llm.sampling_params import SamplingParams
 
 # isort: off
@@ -115,6 +118,74 @@ def test_GenerationResult():
     assert result._done
 
 
+def test_result_timeout_raises():
+    request = GenerationRequest(prompt_token_ids=[12, 23, 34],
+                                sampling_params=SamplingParams(max_tokens=4))
+    result = GenerationResult(request)
+
+    # Queue stays empty (no worker pushing responses) -> must time out fast, not block indefinitely.
+    start = time.monotonic()
+    with pytest.raises(TimeoutError):
+        result.result(timeout=0.1)
+    elapsed = time.monotonic() - start
+    assert elapsed < 2.0, f"result() did not honor timeout (took {elapsed:.2f}s)"
+    assert not result._done
+
+
+def test_result_timeout_budget_across_steps():
+    request = GenerationRequest(prompt_token_ids=[12, 23, 34],
+                                sampling_params=SamplingParams(max_tokens=4))
+    result = GenerationResult(request)
+
+    # A single non-final response is available, then the queue goes empty and the request never
+    # completes.
+    result.queue.put(create_rsp(33, finished=False))
+
+    start = time.monotonic()
+    with pytest.raises(TimeoutError):
+        result.result(timeout=0.1)
+    elapsed = time.monotonic() - start
+    assert elapsed < 2.0, f"result() did not honor timeout (took {elapsed:.2f}s)"
+    assert not result._done
+
+
+def test_result_zero_timeout_completes_with_queued_responses():
+    request = GenerationRequest(prompt_token_ids=[12, 23, 34],
+                                sampling_params=SamplingParams(max_tokens=4))
+    result = GenerationResult(request)
+
+    result.queue.put(create_rsp(33, finished=False))
+    result.queue.put(create_rsp(44, finished=True))
+
+    assert result.result(timeout=0) is result
+    assert result._done
+    assert len(result.outputs[0].token_ids) == 2
+
+
+def test_sync_queue_zero_timeout_checks_for_queued_item():
+    queue = AsyncQueue()
+    queue.put("ready")
+
+    with pytest.warns(UserWarning):
+        assert queue.sync_q.get(timeout=0) == "ready"
+    with pytest.warns(UserWarning), pytest.raises(Empty):
+        queue.sync_q.get(timeout=0)
+
+
+def test_result_completes_within_timeout():
+    request = GenerationRequest(prompt_token_ids=[12, 23, 34],
+                                sampling_params=SamplingParams(max_tokens=4))
+    result = GenerationResult(request)
+
+    result.queue.put(create_rsp(33, finished=False))
+    result.queue.put(create_rsp(44, finished=True))
+
+    ret = result.result(timeout=30.0)
+    assert ret is result
+    assert result._done
+    assert len(result.outputs[0].token_ids) == 2
+
+
 def test_DetokenizedGenerationResultBase():
     sampling_params = SamplingParams(max_tokens=4)
     model_path = llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0"

From 28be4aebd73b7c59b184336ea3cbdfbcced4fa1b Mon Sep 17 00:00:00 2001
From: TensorRT LLM AI Agent <296075020+trtllm-agent@users.noreply.github.com>
Date: Fri, 17 Jul 2026 17:17:36 +0800
Subject: [PATCH 8/9] =?UTF-8?q?[https://nvbugs/6460072][fix]=20Split=20the?=
 =?UTF-8?q?=20caller=20=E2=80=94=20for=20`use=5Flm=5Fhead=5Ftp=5Fin=5Fadp?=
 =?UTF-8?q?=20and=20is=5Fall=5Fgreedy=5Fsample`,=20keep=E2=80=A6=20(#16440?=
 =?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
---
 tensorrt_llm/_torch/speculative/eagle3.py    | 21 ++++-----
 tensorrt_llm/_torch/speculative/interface.py | 47 +++++++++++++++-----
 2 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index d9b5be1ecc1e..731d0236f049 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -894,18 +894,17 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                                                                 self._d2t,
                                                                 draft_step=i)
 
-                # When ADP+LM-head-TP pads logits to max_num_requests, the
-                # padded rows are zero-filled placeholders only required so
-                # every TP rank produces logits of identical shape for the
-                # LM-head-TP all-gather. Drop them *before* sampling: the
-                # per-request sampling params (temperatures/top_k/top_p) are
-                # sized to token_count (== batch_size), so the padded logits
-                # would otherwise fail to broadcast in apply_temperature. This
-                # also keeps next_draft_tokens and the draft_probs buffer
-                # token_count-sized without a post-hoc trim.
+                # ADP+LM-head-TP logits are the LM-head-TP group's row-stacked
+                # batch (each rank's rows padded to max_num_requests, then
+                # all-gathered along dim 0) with the vocab sharded across the
+                # group. Rows [:token_count] would be group rank 0's requests,
+                # not this rank's, and a per-rank argmax would return a
+                # shard-local index -- so keep the full stacked logits and let
+                # greedy_sample_draft_with_tp_gather combine the group's vocab
+                # shards and slice this rank's own row segment; only then trim
+                # the max_num_requests padding down to token_count.
                 mapping_lm_head_tp = None
                 if use_lm_head_tp_in_adp:
-                    logits = logits[:token_count]
                     # The MTP head built this per-forward mapping when producing
                     # the vocab-sharded logits; the sampler needs it to gather.
                     mapping_lm_head_tp = getattr(
@@ -917,6 +916,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata,
                     batch_size,
                     draft_step=i,
                     mapping_lm_head_tp=mapping_lm_head_tp)
+                if use_lm_head_tp_in_adp:
+                    new_draft_token = new_draft_token[:token_count]
                 next_draft_tokens.append(new_draft_token)
 
                 # Update hidden states for the next iteration.
diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py
index 308fc86ad6b9..48ae77f716e8 100644
--- a/tensorrt_llm/_torch/speculative/interface.py
+++ b/tensorrt_llm/_torch/speculative/interface.py
@@ -1378,12 +1378,15 @@ def maybe_gather_sharded_draft_logits(self,
         (see ``_draft_logits_are_sharded``); replicated full-vocab logits are
         returned unchanged.
 
-        Plain TP gathers vocab shards over ``self.mapping``. Under ADP + LM-head
-        TP the worker has already trimmed the LM-head-TP padding rows so each rank
-        holds ``[token_count, vocab_shard]`` for its own tokens; a vocab-dim
-        all-gather over ``mapping_lm_head_tp`` restores full vocab (no token
-        re-slice is needed after the trim).
-        """
+        Plain TP gathers vocab shards over ``self.mapping``. ADP + LM-head TP
+        never reaches this path: rejection sampling (the only consumer of
+        advanced draft sampling) is config-gated off under attention DP, and
+        the group-stacked sharded logits it produces are handled by the greedy
+        path in ``greedy_sample_draft_with_tp_gather``.
+        """
+        assert mapping_lm_head_tp is None, (
+            "Advanced draft sampling is not supported under ADP + LM-head TP "
+            "(rejection sampling is config-gated off with attention DP)")
         if (spec_metadata is None or spec_metadata.is_all_greedy_sample
                 or not self._draft_logits_are_sharded(logits, spec_metadata)):
             return logits
@@ -1750,7 +1753,26 @@ def greedy_sample_draft_with_tp_gather(self,
         vocab-sharded (see ``_draft_logits_are_sharded``) -- e.g. a borrowed or
         gathered full-vocab draft head. Returns tokens in draft-vocab space (the
         caller applies d2t). Expects 2D ``[num_tokens, vocab_shard]`` logits.
+
+        Under ADP + LM-head TP (``mapping_lm_head_tp`` given) the logits are the
+        LM-head-TP group's row-stacked batch (``tp_size`` segments of
+        ``max_num_requests`` padded rows, all-gathered along dim 0 by the MTP
+        shared head) with the vocab sharded across the group. The global argmax
+        must combine the group's vocab shards, and each rank must read its own
+        row segment at offset ``tp_rank * max_num_requests`` -- NOT rows
+        ``[:batch]``, which belong to group rank 0.
         """
+        if (mapping_lm_head_tp is not None
+                and getattr(mapping_lm_head_tp, "tp_size", 1) > 1):
+            from ..distributed.ops import allgather
+            combined = self._get_local_max_and_combined(logits,
+                                                        mapping_lm_head_tp)
+            gathered = allgather(combined, mapping_lm_head_tp, dim=-1)
+            group_size = mapping_lm_head_tp.tp_size
+            local_rows = logits.shape[0] // group_size
+            own_segment = gathered.view(group_size, local_rows,
+                                        -1)[mapping_lm_head_tp.tp_rank]
+            return self._get_draft_tokens_from_gathered(own_segment)
         mapping = self.mapping
         sharded = self._draft_logits_are_sharded(logits, spec_metadata)
         if (sharded and mapping is not None
@@ -1760,9 +1782,9 @@ def greedy_sample_draft_with_tp_gather(self,
             combined = self._get_local_max_and_combined(logits)
             gathered = allgather(combined, mapping, dim=-1)
             return self._get_draft_tokens_from_gathered(gathered)
-        # No TP gather for attention-DP (incl. ADP + LM-head TP): each rank owns
-        # its own requests, so a per-rank argmax is the correct proposal and a
-        # cross-rank gather here would desync the ranks (see
+        # No cross-rank gather for plain attention-DP: each rank owns its own
+        # requests with replicated full-vocab logits, so a per-rank argmax is
+        # the correct proposal and a gather would desync the ranks (see
         # _draft_logits_are_sharded). Plain argmax; caller applies d2t.
         return torch.argmax(logits, dim=-1).type(torch.int32)
 
@@ -1908,7 +1930,12 @@ def sample_draft_tokens(self,
             tokens = self.greedy_sample_draft_with_tp_gather(
                 logits.reshape(-1, logits.shape[-1]), spec_metadata,
                 mapping_lm_head_tp)
-            tokens = tokens.reshape(batch_shape)
+            if mapping_lm_head_tp is None:
+                tokens = tokens.reshape(batch_shape)
+            # else: ADP+LM-head-TP (2D step form only) -- the sampler returned
+            # this rank's own row segment, 1/tp_size of the stacked input rows,
+            # so the input batch shape no longer applies. Keep as-is; the
+            # caller trims the max_num_requests padding to token_count.
         else:
             # Advanced sampling gathers the vocab-sharded draft logits to full
             # vocab, then samples (scattering this step's proposal distribution

From 01f34634e6d657f3007d99fc67306af12fa1ad1a Mon Sep 17 00:00:00 2001
From: Yiqing Yan <yiqingy@nvidia.com>
Date: Fri, 17 Jul 2026 17:26:40 +0800
Subject: [PATCH 9/9] [None][infra] Add blossom-ci authorized users (#16541)

Signed-off-by: yiqingy <yiqingy@tensorrt-llm-infra-debug-vm-01.nvidia.com>
Co-authored-by: yiqingy <yiqingy@tensorrt-llm-infra-debug-vm-01.nvidia.com>
---
 .github/workflows/blossom-ci.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml
index d67685f216f2..e398b8591d3c 100644
--- a/.github/workflows/blossom-ci.yml
+++ b/.github/workflows/blossom-ci.yml
@@ -187,6 +187,7 @@ jobs:
         "JunyiXu-nv",
         "JyChang012",
         "kaiyux",
+        "Kambili",
         "kanghui0204",
         "karljang",
         "karthikvetrivel",