[pull] main from NVIDIA:main

jenkins/L0_Test.groovy

@@ -4674,12 +4674,12 @@ def launchTestJobs(pipeline, testFilter)
         "DGX_B200-8_GPUs-PyTorch-4": ["auto:dgx-b200-flex", "l0_dgx_b200", 4, 4, 8, 1, true],
         "DGX_B200-8_GPUs-AutoDeploy-Post-Merge-1": ["auto:dgx-b200-flex", "l0_dgx_b200", 1, 1, 8, 1, true],
         "DGX_B200-4_GPUs-Verl-Post-Merge-1": ["auto:dgx-b200-flex", "l0_verl", 1, 1, 4, 1, true],
-        "B300-PyTorch-1": ["auto:dgx-b300-flex", "l0_b300", 1, 2, 1, 1, true],
-        "B300-PyTorch-2": ["auto:dgx-b300-flex", "l0_b300", 2, 2, 1, 1, true],
-        "B300-PyTorch-Post-Merge-1": ["auto:dgx-b300-flex", "l0_b300", 1, 1, 1, 1, true],
-        "DGX_B300-4_GPUs-PyTorch-1": ["auto:dgx-b300-flex", "l0_dgx_b300", 1, 1, 4, 1, true],
-        "DGX_B300-4_GPUs-PyTorch-Post-Merge-1": ["auto:dgx-b300-flex", "l0_dgx_b300", 1, 2, 4, 1, true],
-        "DGX_B300-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b300-flex", "l0_dgx_b300", 2, 2, 4, 1, true],
+        // "B300-PyTorch-1": ["auto:dgx-b300-flex", "l0_b300", 1, 2, 1, 1, true],
+        // "B300-PyTorch-2": ["auto:dgx-b300-flex", "l0_b300", 2, 2, 1, 1, true],
+        // "B300-PyTorch-Post-Merge-1": ["auto:dgx-b300-flex", "l0_b300", 1, 1, 1, 1, true],
+        // "DGX_B300-4_GPUs-PyTorch-1": ["auto:dgx-b300-flex", "l0_dgx_b300", 1, 1, 4, 1, true],
+        // "DGX_B300-4_GPUs-PyTorch-Post-Merge-1": ["auto:dgx-b300-flex", "l0_dgx_b300", 1, 2, 4, 1, true],
+        // "DGX_B300-4_GPUs-PyTorch-Post-Merge-2": ["auto:dgx-b300-flex", "l0_dgx_b300", 2, 2, 4, 1, true],
         // VisualGen PerfSanity post-merge test
         "DGX_B200-8_GPUs-PyTorch-VisualGen-PerfSanity-Post-Merge-1": ["auto:dgx-b200-flex", "l0_b200_visual_gen_perf_sanity", 1, 1, 8, 1, true],
         // PerfSanity post-merge tests

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -2388,10 +2388,12 @@ def sparse_attn_indexer(
         has_prefill = num_contexts > 0
         num_gen_tokens = num_tokens - num_ctx_tokens
 
-        topk_indices_buffer = torch.empty(
+        topk_indices_buffer = metadata.get_empty(
+            metadata.cuda_graph_buffers,
             (hidden_states.shape[0], self.index_topk),
+            cache_name="indexer_topk_out_buffer",
             dtype=torch.int32,
-            device=hidden_states.device)
+            capture_graph=metadata.is_cuda_graph)
         if not use_custom_topk:
             topk_indices_buffer[:hidden_states.shape[0]] = -1
 

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -167,12 +167,7 @@ def __enter__(self):
     def __exit__(self, _exc_type, _exc_val, _exc_tb):
         self.shutdown()
 
-    def _create_kv_slice(
-        self,
-        req: LlmRequest,
-        token_range: Optional[TokenRange] = None,
-        is_last_slice: bool = True,
-    ) -> KVSlice:
+    def _create_kv_slice(self, req: LlmRequest) -> KVSlice:
         adapter = self._reuse_adapter
         tpb = adapter.tokens_per_block
         assert self._page_table is not None
@@ -185,7 +180,8 @@ def _create_kv_slice(
             else [0] * len(layer_groups)
         )
 
-        if token_range is None and req.prompt_len > 0:
+        token_range = None
+        if req.prompt_len > 0:
             # Align with KV cache allocation (prepare_disagg_gen_init /
             # _get_context_bytes), which reserves prompt_len +
             # num_extra_kv_tokens slots for speculative decoding methods
@@ -236,7 +232,7 @@ def _create_kv_slice(
             mamba_state_index = self._kv_cache_manager.mamba_cache_index[req.py_request_id]
 
         return KVSlice(
-            is_last_slice=is_last_slice,
+            is_last_slice=True,
             block_ids_per_layer_groups=groups,
             mamba_state_index=mamba_state_index,
             token_range=token_range,

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -1254,14 +1254,20 @@ def get_max_atten_window_upper_bound(self, blocks_in_primary_pool,
         assert max_atten_window_upper_bound > 0, "Impossible to fit in any sequence in kvCache"
         return max_atten_window_upper_bound
 
+    def _resolve_window_size(
+            self,
+            window_size: Optional[int],
+            error_msg: str = "window_size must be provided for VSWA") -> int:
+        if window_size is not None:
+            return window_size
+        if self.is_vswa:
+            raise ValueError(error_msg)
+        return self.max_attention_window_vec[0]
+
     def get_cache_indices(self,
                           request: LlmRequest,
                           window_size: Optional[int] = None) -> List[int]:
-        if window_size is None:
-            if len(self.max_attention_window_vec) > 1:
-                raise ValueError("window_size must be provided for VSWA")
-            window_size = self.max_attention_window_vec[0]
-
+        window_size = self._resolve_window_size(window_size)
         result = self.impl.get_cache_block_ids(request.py_request_id,
                                                window_size)
         assert len(result) == 1
@@ -1298,16 +1304,7 @@ def commit_and_get_block_hashes(
         hash matches what ``storeBlocks`` would later compute. Beam-width-1
         only; the connector enforces this at startup.
         """
-        if window_size is None:
-            # ``is_vswa`` (distinct window sizes) is the real VSWA signal; a
-            # uniform per-layer vector such as ``[4096, 4096, ...]`` has
-            # ``len > 1`` yet a single effective window, so keying off the
-            # length would spuriously reject it for connector callers that omit
-            # ``window_size``.
-            if self.is_vswa:
-                raise ValueError("window_size must be provided for VSWA")
-            window_size = self.max_attention_window_vec[0]
-
+        window_size = self._resolve_window_size(window_size)
         return list(
             self.impl.commit_and_get_block_hashes_for_request(
                 request, window_size))
@@ -1331,10 +1328,7 @@ def get_priority_by_block_id(self,
         Returns:
             The retention priority of the block (0-100), or default priority (35) if not found.
         """
-        if window_size is None:
-            if len(self.max_attention_window_vec) > 1:
-                raise ValueError("window_size must be provided for VSWA")
-            window_size = self.max_attention_window_vec[0]
+        window_size = self._resolve_window_size(window_size)
         return self.impl.get_priority_by_block_id(block_id, window_size)
 
     def get_batch_cache_indices(
@@ -1348,10 +1342,9 @@ def get_batch_cache_indices(
         beam_width = beam_width or 1
         if window_size is None:
             if layer_idx is None:
-                if len(self.max_attention_window_vec) > 1:
-                    raise ValueError(
-                        "layer_idx or window_size must be provided for VSWA")
-                window_size = self.max_attention_window_vec[0]
+                window_size = self._resolve_window_size(
+                    window_size,
+                    "layer_idx or window_size must be provided for VSWA")
             else:
                 layer_offset = self.layer_offsets[layer_idx]
                 # Explicit layer_offset -> window_size mapping (no modulo

tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py

@@ -3163,3 +3163,85 @@ def test_prepare_swaps_and_restore_recovers(self):
         assert meta.kv_cache_manager is original_kv_mgr
         torch.testing.assert_close(meta.kv_cache_block_offsets, original_offsets)
         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_hopper
+def test_topk_indices_buffer_cuda_graph():
+    from tensorrt_llm._torch.memory_buffer_utils import Buffers
+
+    batch_size, next_n = 4, 3
+    head_dim, block_size = 128, 64
+    index_topk = 2048
+    kv_len = 512
+    layer_idx = 0
+    num_tokens = batch_size * next_n
+
+    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,
+    )
+    indexer = create_indexer(sparse_attn_config, layer_idx=layer_idx)
+
+    request_ids = list(range(batch_size))
+    kv_lens = torch.tensor([kv_len] * batch_size, 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.tensor([next_n] * batch_size, 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=num_tokens,
+        max_draft_tokens=next_n - 1,
+        index_topk=index_topk,
+        enable_indexer_skip=True,
+    )
+    assert metadata.skip_indexer_for_gen_reqs, (
+        "test setup must hit the skip-indexer path (kv_len <= index_topk)"
+    )
+
+    metadata.is_cuda_graph = True
+    metadata.cuda_graph_buffers = Buffers()
+
+    hidden_states = torch.zeros((num_tokens, 1), device="cuda", dtype=torch.bfloat16)
+    dummy = torch.zeros((num_tokens, 1), device="cuda", dtype=torch.bfloat16)
+
+    def _run_indexer():
+        return indexer.sparse_attn_indexer(
+            metadata,
+            hidden_states,
+            q_fp8=dummy,
+            k_fp8=dummy,
+            k_scale=dummy,
+            weights=dummy,
+            update_k_cache=False,
+        )
+
+    out1 = _run_indexer()
+    ptr1 = out1.data_ptr()
+    out2 = _run_indexer()
+    ptr2 = out2.data_ptr()
+
+    assert ptr1 == ptr2, (
+        f"indexer topk-output buffer address changed across calls "
+        f"({ptr1:#x} -> {ptr2:#x}); under CUDA-graph capture it must be a "
+        f"persistent reserved-arena buffer to avoid stale-pointer IMA"
+    )
+    assert "indexer_topk_out_buffer" in metadata.cuda_graph_buffers.buffers, (
+        "indexer topk-output buffer must be drawn from the cuda_graph_buffers arena"
+    )

tests/unittest/_torch/executor/test_resource_manager.py

@@ -994,5 +994,25 @@ class MockStream:
             kv_cache_manager.shutdown()
 
 
+class TestResolveWindowSize(unittest.TestCase):
+
+    @staticmethod
+    def _make_manager(max_attention_window_vec, is_vswa):
+        mgr = KVCacheManager.__new__(KVCacheManager)
+        mgr.max_attention_window_vec = max_attention_window_vec
+        mgr.is_vswa = is_vswa
+        return mgr
+
+    def test_uniform_multi_layer_vector_does_not_raise(self):
+        mgr = self._make_manager([4096, 4096, 4096], is_vswa=False)
+        self.assertEqual(mgr._resolve_window_size(None), 4096)
+
+    def test_genuine_vswa_requires_window_size(self):
+        mgr = self._make_manager([4096, 8192], is_vswa=True)
+        with self.assertRaises(ValueError):
+            mgr._resolve_window_size(None)
+        self.assertEqual(mgr._resolve_window_size(8192), 8192)
+
+
 if __name__ == "__main__":
     unittest.main()

tests/unittest/disaggregated/test_cache_reuse_adapter.py

@@ -342,17 +342,6 @@ def test_defaults_to_prompt_len_when_no_extra(self):
         assert kv_slice.token_range is not None
         assert (kv_slice.token_range.start, kv_slice.token_range.end) == (0, prompt_len)
 
-    def test_respects_explicit_token_range(self):
-        prompt_len = 17
-        transceiver, req = _build_transceiver_for_kv_slice(
-            num_extra_kv_tokens=7, prompt_len=prompt_len
-        )
-        explicit = TokenRange(start=0, end=8)
-
-        kv_slice = transceiver._create_kv_slice(req, token_range=explicit)
-
-        assert kv_slice.token_range is explicit
-
 
 # ---------------------------------------------------------------------------
 # CacheReuseAdapter.get_cached_token_count_per_layer_group: SWA clamp.