[#14828][feat] AutoDeploy: support multi KV cache memory pool in trtllm attention

examples/auto_deploy/llmc/create_standalone_package.py

@@ -153,6 +153,8 @@
     "test_torch_gated_delta_rule_cache.py",
     "test_gated_delta_rule_cache.py",
     "test_kv_cache_transformers.py",
+    # trtllm attention backend (insert_cached_attention backend=trtllm) not available standalone
+    "test_kv_cache_trtllm_multipool.py",
     # Require TRT-LLM CUDA causal conv / mamba kernels (ops not registered standalone)
     "test_cuda_causal_conv_cached_op.py",
     "test_triton_causal_conv_cached_op.py",

tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.py

@@ -94,7 +94,19 @@ def __init__(self):
         self.context_lengths_gpu: Optional[torch.Tensor] = None  # [max_batch] int32 device
         # Persistent block_offsets buffer for CUDA graph compatibility.
         # Pre-allocated to max size so the tensor address is stable across replays.
+        # ``self.block_offsets`` is the group-0 buffer (kept for the spec-dec
+        # scratch path and backward compatibility); additional KV window groups
+        # (VSWA / non-uniform sliding window, e.g. gpt-oss) get their own
+        # persistent buffer keyed by the group's ``cache_loc`` input pointer in
+        # ``_block_offsets_by_cache_loc``. The transform invokes
+        # ``prepare_trtllm_metadata`` once per group with that group's
+        # ``cache_loc_g{i}`` / ``cu_num_pages_g{i}`` inputs, so without per-group
+        # buffers the groups would clobber a single shared buffer.
         self.block_offsets: Optional[torch.Tensor] = None
+        self._block_offsets_by_cache_loc: dict[int, torch.Tensor] = {}
+        # Shapes for lazy per-group buffer allocation (set in ``reset``).
+        self._max_batch: int = 0
+        self._max_blocks_per_seq: int = 0
         # Per-layer cache for tensors that must survive CUDA graph replay.
         # Keyed by kv_cache.data_ptr() (stable and unique per layer).
         self._layer_cache: dict[
@@ -148,9 +160,13 @@ def reset(self, device: torch.device, max_batch: int, max_blocks_per_seq: int) -
         self.host_request_types = torch.zeros(
             max_batch, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()
         )
+        self._max_batch = max_batch
+        self._max_blocks_per_seq = max_blocks_per_seq
         self.block_offsets = torch.zeros(
             1, max_batch, 2, max_blocks_per_seq, dtype=torch.int32, device=device
         )
+        # Group 0 reuses ``self.block_offsets``; it is registered under its
+        # ``cache_loc`` pointer on first use in ``_get_block_offsets_buffer``.
         self.host_past_kv_lengths = torch.zeros(
             max_batch, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()
         )
@@ -290,23 +306,71 @@ def refresh_batch_state(self, batch_info: BatchInfo) -> None:
         self.num_contexts = num_prefill
         self.num_ctx_tokens = batch_info.get_num_tokens()[0]
 
+    def _get_block_offsets_buffer(self, cache_loc: torch.Tensor) -> torch.Tensor:
+        """Return the persistent block_offsets buffer for this KV window group.
+
+        Each KV window group is driven by its own ``cache_loc`` input tensor
+        (group 0 uses ``cache_loc``; groups 1..N-1 use ``cache_loc_g{i}``), which
+        are persistent buffers with stable ``data_ptr()`` across CUDA-graph
+        replays. Keying by that pointer (same pattern as ``_layer_cache`` keyed
+        by ``kv_cache.data_ptr()``) gives each group an independent, address-stable
+        block_offsets buffer so per-group ``prepare_trtllm_metadata`` invocations
+        do not clobber each other.
+
+        Lazily allocates a buffer on first sight of a group's ``cache_loc``. This
+        must happen during warm-up (never mid-capture) so the tensor address is
+        stable for graph replay; group 0's buffer reuses the one already
+        allocated in ``reset``.
+        """
+        key = cache_loc.data_ptr()
+        buf = self._block_offsets_by_cache_loc.get(key)
+        if buf is None:
+            assert self.block_offsets is not None, (
+                "planner.reset() must run before _get_block_offsets_buffer()"
+            )
+            if not self._block_offsets_by_cache_loc:
+                # First group seen this run is group 0: reuse the reset() buffer.
+                buf = self.block_offsets
+            else:
+                assert (
+                    not torch.cuda.is_current_stream_capturing()
+                ) or cuda_graph_state.in_warm_up(), (
+                    "block_offsets buffer for a new KV window group must be "
+                    "allocated during warm-up, not during CUDA graph capture. "
+                    "Ensure warm-up exercises every KV pool."
+                )
+                buf = torch.zeros(
+                    1,
+                    self._max_batch,
+                    2,
+                    self._max_blocks_per_seq,
+                    dtype=torch.int32,
+                    device=self.block_offsets.device,
+                )
+            self._block_offsets_by_cache_loc[key] = buf
+        return buf
+
     def plan_device(
         self,
         num_seq: int,
         block_offset_multiplier: int,
         cu_num_pages: torch.Tensor,
         cache_loc: torch.Tensor,
-    ) -> None:
+    ) -> torch.Tensor:
         """Per-forward DEVICE metadata: block_offsets via Triton kernel (pure GPU).
 
         Called from the ``prepare_trtllm_metadata`` custom op (in the graph).
+        Returns the per-group block_offsets buffer that was populated, so the op
+        can flow it through the graph to that group's attention layers.
         """
-        k_slice = self.block_offsets[0, :, 0, :]  # [max_batch, M], stride [2*M, 1]
+        block_offsets = self._get_block_offsets_buffer(cache_loc)
+        k_slice = block_offsets[0, :, 0, :]  # [max_batch, M], stride [2*M, 1]
         torch.ops.auto_deploy.ragged_to_block_table_triton(
             cache_loc, cu_num_pages, k_slice, num_seq
         )
-        self.block_offsets[0, :num_seq, 0, :].mul_(block_offset_multiplier)
-        self.block_offsets[0, :num_seq, 1, :] = self.block_offsets[0, :num_seq, 0, :] + 1
+        block_offsets[0, :num_seq, 0, :].mul_(block_offset_multiplier)
+        block_offsets[0, :num_seq, 1, :] = block_offsets[0, :num_seq, 0, :] + 1
+        return block_offsets
 
 
 _GlobalTrtllmPlanner = _TrtllmPlanner()
@@ -479,14 +543,16 @@ def prepare_trtllm_metadata(
             _GlobalTrtllmPlanner.use_spec_decoding = batch_info.get_num_sequences()[2] == 0
     block_offset_multiplier = batch_info.get_block_offset_multiplier()
 
-    _GlobalTrtllmPlanner.plan_device(
+    block_offsets = _GlobalTrtllmPlanner.plan_device(
         num_seq=batch_info.get_total_num_sequences(),
         block_offset_multiplier=block_offset_multiplier,
         cu_num_pages=cu_num_pages,
         cache_loc=cache_loc,
     )
 
-    return [_GlobalTrtllmPlanner.block_offsets]
+    # Return this group's buffer (keyed by ``cache_loc``) so multi-pool
+    # (VSWA) deployments flow the correct block_offsets to each group's layers.
+    return [block_offsets]
 
 
 @prepare_trtllm_metadata.register_fake
@@ -571,7 +637,10 @@ def trtllm_mha_with_cache(
     num_tokens = batch_info.get_total_num_tokens()
     max_context_length = batch_info.get_max_context_length()
     max_num_requests = batch_info.get_max_batch_size()
-    # Use sliding_window for attention_window_size if provided, else full context length
+    # Use sliding_window for attention_window_size if provided, else full context length.
+    # The mask stays ``causal`` (matching the PyTorch backend, which never uses
+    # sliding_window_causal): the kernel honors the window via the cyclic
+    # attention-window handling driven by ``attention_window_size``.
     attention_window_size = (
         sliding_window
         if isinstance(sliding_window, int) and sliding_window > 0
@@ -800,6 +869,13 @@ class TrtllmAttention(AttentionDescriptor):
     Follows the same stateless descriptor pattern as ``FlashInferAttention``.
     """
 
+    @classmethod
+    def kernel_handles_cyclic_swa(cls) -> bool:
+        """thop.attention applies the sliding-window mask internally via cyclic
+        KV indexing, so the executor passes the full per-window block table and
+        global KV lengths (no host-side window slicing). See base class."""
+        return True
+
     @classmethod
     def get_attention_layout(cls) -> AttentionLayout:
         """Get the attention layout expected by the backend."""

tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py

@@ -2389,6 +2389,22 @@ def supports_shared_kv(cls) -> bool:
         """Whether this backend supports shared-KV cache aliasing."""
         return False
 
+    @classmethod
+    def kernel_handles_cyclic_swa(cls) -> bool:
+        """Whether the backend's kernel applies the sliding-window mask itself.
+
+        When ``True`` (e.g. the trtllm ``thop.attention`` kernel), the kernel
+        cyclically indexes the KV cache internally using the per-layer attention
+        window, so the executor must hand it the *full* per-window block table
+        and a *global* (un-window-capped) KV length -- the same contract as the
+        PyTorch backend.
+
+        When ``False`` (default; e.g. triton / flashinfer), the kernel does not
+        cyclic-index, so the executor must host-slice the block table down to the
+        live sliding-window view (see ``ad_executor._compute_window_local_view``).
+        """
+        return False
+
     @classmethod
     @abstractmethod
     def get_standard_metadata_args(cls) -> List[str]:

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -429,8 +429,18 @@ def disable_cudagraph_for_speculative_flashinfer(self):
     ### UTILITY METHODS ############################################################################
     @property
     def requires_uniform_kv_caches(self) -> bool:
-        """Whether CachedSequenceInterface must enforce a uniform KV cache mapping."""
-        return self.attn_backend.lower() == "trtllm"
+        """Whether CachedSequenceInterface must enforce a uniform KV cache mapping.
+
+        No attention backend currently requires this. The trtllm backend used to
+        return ``True`` here to force a single KV pool, but it now supports
+        multiple KV cache memory pools for non-uniform sliding-window models
+        (e.g. gpt-oss) -- the kernel applies the sliding-window mask internally
+        via cyclic indexing, so per-window pools route correctly. The flag is
+        kept (defaulting to ``False``) so the uniformity enforcement in
+        ``CachedSequenceInterface`` remains available should a future backend
+        need it.
+        """
+        return False
 
     def create_factory(self) -> ModelFactory:
         """Create a model factory from the arguments.

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -302,6 +302,40 @@ def _compute_window_local_view(
     return active_indices, extra_page, active_token_count, last_page_len
 
 
+def _compute_cyclic_full_view(
+    all_indices: Sequence[int],
+    end_compute_i: int,
+    tokens_per_block: int,
+) -> Tuple[List[int], int, int, int]:
+    """Compute the metadata view for a cyclic-SWA kernel (trtllm).
+
+    Unlike ``_compute_window_local_view`` (which slices the block table down to
+    the live sliding window for kernels that cannot cyclic-index), the trtllm
+    ``thop.attention`` kernel applies the sliding-window mask itself by wrapping
+    KV reads modulo the attention window. It therefore needs:
+
+      * the FULL per-window block table (``all_indices`` verbatim, including any
+        stale front-evicted entries -- the kernel's modulo indexing skips them),
+        and
+      * the GLOBAL (un-window-capped) KV length ``end_compute_i``.
+
+    This mirrors the PyTorch backend, which copies the manager's full block list
+    from index 0 and passes ``host_past_key_value_lengths == total KV length``.
+
+    Returns the same 4-tuple shape as ``_compute_window_local_view``:
+    ``(active_indices, extra_page, seq_len_with_cache, last_page_len)``.
+    ``extra_page`` is always -1: the full table already contains the next page,
+    so the overlap scheduler needs no deferred-page insertion.
+    """
+    active_indices = list(all_indices)
+    seq_len_with_cache = end_compute_i
+    if seq_len_with_cache > 0:
+        last_page_len = (seq_len_with_cache - 1) % tokens_per_block + 1
+    else:
+        last_page_len = 0
+    return active_indices, -1, seq_len_with_cache, last_page_len
+
+
 class ADEngine(ModelEngine):
     """The AutoDeploy Engine (ADEngine) is the main engine interface to execute AutoDeploy models.
 
@@ -770,6 +804,12 @@ def _prepare_inputs(
         # on SequenceInfo).  Per-window queries on the manager route to the
         # correct C++ pool via mLayerToWindowSize.
         kv_group_windows = self.cache_seq_interface.kv_group_windows
+        # When the attention kernel applies the sliding-window mask itself via
+        # cyclic KV indexing (trtllm), the executor must hand it the full
+        # per-window block table and a global (un-window-capped) KV length --
+        # the same contract as the PyTorch backend. Otherwise (triton /
+        # flashinfer) host-slice the block table to the live window below.
+        cyclic_swa = self.cache_seq_interface.kernel_handles_cyclic_swa
         # Cache hot lookups so the per-request loop avoids repeated C++
         # dispatch / hasattr calls.
         _tokens_per_block = kv_cache_manager.tokens_per_block
@@ -809,40 +849,56 @@ def _prepare_inputs(
 
             for pool_idx, group_window in enumerate(kv_group_windows):
                 all_indices = batch_cache_indices_per_pool[pool_idx][i]
-                # SWA front-eviction: get_batch_cache_indices returns the FULL
-                # historical page list including front-evicted entries (the
-                # C++ side bumps a counter rather than popping mCacheBlockIds).
-                # _compute_window_local_view slices it down to the live window
-                # in window-local coords.
-                front_removed = kv_cache_manager.get_num_front_blocks_removed(
-                    request.py_request_id, window_size=group_window
-                )
-                (
-                    active_indices,
-                    extra_page,
-                    active_token_count,
-                    lpl_i,
-                ) = _compute_window_local_view(
-                    all_indices,
-                    front_removed=front_removed,
-                    end_compute_i=end_compute_i,
-                    group_window=group_window,
-                    tokens_per_block=_tokens_per_block,
-                )
-                num_active = len(active_indices)
+                if cyclic_swa:
+                    # Cyclic-SWA kernels (trtllm) want the FULL per-window block
+                    # table and the GLOBAL KV length; the kernel masks the window
+                    # internally. No front-eviction slicing, so the
+                    # get_num_front_blocks_removed C++ dispatch is skipped here.
+                    (
+                        active_indices,
+                        extra_page,
+                        active_token_count,
+                        lpl_i,
+                    ) = _compute_cyclic_full_view(
+                        all_indices,
+                        end_compute_i=end_compute_i,
+                        tokens_per_block=_tokens_per_block,
+                    )
+                    num_active = len(active_indices)
+                else:
+                    # SWA front-eviction: get_batch_cache_indices returns the FULL
+                    # historical page list including front-evicted entries (the
+                    # C++ side bumps a counter rather than popping mCacheBlockIds).
+                    # _compute_window_local_view slices it down to the live window
+                    # in window-local coords.
+                    front_removed = kv_cache_manager.get_num_front_blocks_removed(
+                        request.py_request_id, window_size=group_window
+                    )
+                    (
+                        active_indices,
+                        extra_page,
+                        active_token_count,
+                        lpl_i,
+                    ) = _compute_window_local_view(
+                        all_indices,
+                        front_removed=front_removed,
+                        end_compute_i=end_compute_i,
+                        group_window=group_window,
+                        tokens_per_block=_tokens_per_block,
+                    )
+                    num_active = len(active_indices)
 
                 cache_loc_per_pool[pool_idx].extend(active_indices)
                 cu_num_pages_per_pool[pool_idx].append(
                     cu_num_pages_per_pool[pool_idx][i] + num_active
                 )
                 extra_page_per_seq_per_pool[pool_idx].append(extra_page)
-                # Window-local seq_len_with_cache / last_page_len for every
-                # pool (including 0).  For full-attention pools the helper
-                # returns the unclamped global value (group_window equals
-                # max_seq_len, no clamping kicks in), so this is identical to
-                # the legacy single-pool path for non-SWA models.  For SWA
-                # pools (whether pool 0 or pool 1+), it carries the
-                # window-local coords the kernel needs under front-eviction.
+                # seq_len_with_cache / last_page_len per pool (including 0).
+                # Cyclic-SWA (trtllm): the global KV length for every pool.
+                # Host-sliced (triton/flashinfer): the unclamped global value for
+                # full-attention pools (window == max_seq_len, no clamping), and
+                # the window-local coords for SWA pools under front-eviction --
+                # identical to the legacy single-pool path for non-SWA models.
                 seq_len_with_cache_per_pool[pool_idx].append(active_token_count)
                 last_page_len_per_pool[pool_idx].append(lpl_i)
 

tensorrt_llm/_torch/auto_deploy/shim/interface.py

@@ -147,6 +147,12 @@ def __init__(
         # same order as the C++ manager's internal pool ordering (i.e. the
         # insertion order of the per-window shape map keys).
         self._kv_group_windows: List[int] = []
+        # Whether the attention backend's kernel applies the sliding-window mask
+        # itself via cyclic KV indexing (trtllm). When True the executor passes
+        # the full per-window block table and global KV lengths instead of
+        # host-slicing to the live window. Set by the kvcache transform from the
+        # attention descriptor's ``kernel_handles_cyclic_swa()``.
+        self._kernel_handles_cyclic_swa: bool = False
         # lookup of unmanaged resources
         self._unmanaged_resources: List[str] = []
         self._spec_config = spec_config
@@ -1307,6 +1313,21 @@ def set_kv_groups(self, group_windows: List[int]) -> None:
         """
         self._kv_group_windows = list(group_windows)
 
+    @property
+    def kernel_handles_cyclic_swa(self) -> bool:
+        """Whether the attention kernel applies the sliding-window mask itself.
+
+        When True (trtllm), the executor passes the full per-window block table
+        and global KV lengths; when False (triton/flashinfer), it host-slices to
+        the live sliding window.
+        """
+        return self._kernel_handles_cyclic_swa
+
+    def set_kernel_handles_cyclic_swa(self, value: bool) -> None:
+        """Record the attention backend's cyclic-SWA capability (called by the
+        kvcache transform from ``AttentionDescriptor.kernel_handles_cyclic_swa``)."""
+        self._kernel_handles_cyclic_swa = bool(value)
+
     @property
     def kv_cache_manager(self) -> Optional[KVCacheManager]:
         """Return the unified KVCacheManager, or None if not initialized."""

tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py

@@ -349,6 +349,12 @@ def _apply(
                 skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True
             )
 
+        # Record whether this backend's kernel applies the sliding-window mask
+        # itself (cyclic KV indexing, e.g. trtllm). The executor uses this to
+        # decide between passing the full per-window block table + global KV
+        # lengths (cyclic) and host-slicing to the live window (triton/flashinfer).
+        cm.set_kernel_handles_cyclic_swa(attn_descriptor.kernel_handles_cyclic_swa())
+
         # get standard metadata nodes for all source attention nodes
         meta_nodes_std = self._process_metadata_std(gm, cm)