[None][feat] Add paged KV cache support to Vanilla attention

tensorrt_llm/_torch/attention_backend/vanilla.py

@@ -8,6 +8,7 @@
 import torch.nn.functional as F
 
 from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
 
 try:
     from transformers.modeling_attn_mask_utils import AttentionMaskConverter
@@ -132,46 +133,116 @@ def _single_request_sparse_kv_predict(
             **kwargs) -> tuple[Optional[torch.Tensor], int]:
         raise NotImplementedError
 
+    @staticmethod
+    def _gather_paged_kv(kv_cache_tensor, block_ids, kv_idx, num_tokens,
+                         tokens_per_block):
+        """Materialize the first ``num_tokens`` logical K (``kv_idx=0``) or V
+        (``kv_idx=1``) tokens from the paged pool as a contiguous
+        ``[num_tokens, num_kv_heads, head_dim]`` tensor (NHD view).
+
+        A single block returns the contiguous view directly (no copy), which
+        keeps a within-one-page request allocation-free.
+
+        Invalid block IDs produce zeros without changing logical positions.
+        """
+        if num_tokens <= 0:
+            # Empty slice: content is irrelevant, but block_ids[0] may be
+            # BAD_PAGE_INDEX, so index the always-valid block 0.
+            return kv_cache_tensor[0, kv_idx, :0]
+        chunks = []
+        read = 0
+        while read < num_tokens:
+            blk = block_ids[read // tokens_per_block]
+            off = read % tokens_per_block
+            n = min(tokens_per_block - off, num_tokens - read)
+            if blk == BAD_PAGE_INDEX:
+                chunks.append(
+                    kv_cache_tensor.new_zeros((n, *kv_cache_tensor.shape[3:])))
+            else:
+                chunks.append(kv_cache_tensor[blk, kv_idx, off:off + n])
+            read += n
+        return chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=0)
+
+    @staticmethod
+    def _gather_paged_mla_latent(kv_cache, block_ids, kv_len):
+        """Materialize a request's MLA latent cache as a contiguous
+        ``[kv_len, kv_lora_rank + qk_rope_head_dim]`` tensor from the paged pool
+        (NHD ``[num_pages, 1, tokens_per_block, 1, d_latent]``). A single block
+        returns the view directly (no copy). Invalid block IDs produce zeros
+        without changing logical positions."""
+        tokens_per_block = kv_cache.shape[2]
+        chunks = []
+        read = 0
+        while read < kv_len:
+            blk = block_ids[read // tokens_per_block]
+            off = read % tokens_per_block
+            n = min(tokens_per_block - off, kv_len - read)
+            if blk == BAD_PAGE_INDEX:
+                chunks.append(kv_cache.new_zeros((n, kv_cache.shape[-1])))
+            else:
+                chunks.append(kv_cache[blk, 0, off:off + n, 0, :])
+            read += n
+        return chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=0)
+
     def _single_request_update_kv_cache(self,
                                         k,
                                         v,
                                         kv_cache_tensor,
                                         past_seen_token,
                                         kv_len,
-                                        cache_idx,
+                                        block_ids,
                                         sparse_kv_indices=None):
+        """Append new K/V tokens and gather the logical paged-cache sequence."""
         # select tokens using the sparse kv indices
         if sparse_kv_indices is not None:
             k_selected = triton_index_gather(k, sparse_kv_indices)
             v_selected = triton_index_gather(v, sparse_kv_indices)
         else:
             k_selected, v_selected = k, v
 
-        # get cache position
         seq_len = past_seen_token + kv_len
-        cache_position = torch.arange(past_seen_token,
-                                      seq_len,
-                                      device=kv_cache_tensor.device)
+        tokens_per_block = kv_cache_tensor.shape[2]
 
-        # get kv cache tensor
-        k_out = kv_cache_tensor[cache_idx, 0, :, :, :].unsqueeze(0)
-        v_out = kv_cache_tensor[cache_idx, 1, :, :, :].unsqueeze(0)
-
-        # update kv cache
         if k is not None and v is not None:
             access_type = self._access_type[k_selected.dtype.itemsize]
-            k_out.view(dtype=access_type).index_copy_(
-                1, cache_position, k_selected.view(dtype=access_type))
-            v_out.view(dtype=access_type).index_copy_(
-                1, cache_position, v_selected.view(dtype=access_type))
+            written = 0
+            while written < kv_len:
+                pos = past_seen_token + written
+                blk = block_ids[pos // tokens_per_block]
+                off = pos % tokens_per_block
+                n = min(tokens_per_block - off, kv_len - written)
+                # New tokens must land in a live page; fail loudly rather than
+                # writing into the last page via negative indexing.
+                assert blk != BAD_PAGE_INDEX, (
+                    f"Writing new KV into an evicted/invalid page (pos {pos}); "
+                    "block_ids/metadata are inconsistent.")
+                dst = torch.arange(off, off + n, device=kv_cache_tensor.device)
+                kv_cache_tensor[blk, 0].view(dtype=access_type).index_copy_(
+                    0, dst,
+                    k_selected[0, written:written + n].view(dtype=access_type))
+                kv_cache_tensor[blk, 1].view(dtype=access_type).index_copy_(
+                    0, dst,
+                    v_selected[0, written:written + n].view(dtype=access_type))
+                written += n
 
-        # return past kv and the dense kv tensors for sparse attention
         if sparse_kv_indices is not None:
-            k_states = torch.cat([k_out[:, :past_seen_token, :, :], k], dim=1)
-            v_states = torch.cat([v_out[:, :past_seen_token, :, :], v], dim=1)
+            k_states = torch.cat([
+                self._gather_paged_kv(kv_cache_tensor, block_ids, 0,
+                                      past_seen_token, tokens_per_block)[None],
+                k
+            ],
+                                 dim=1)
+            v_states = torch.cat([
+                self._gather_paged_kv(kv_cache_tensor, block_ids, 1,
+                                      past_seen_token, tokens_per_block)[None],
+                v
+            ],
+                                 dim=1)
         else:
-            k_states, v_states = k_out[:, :seq_len, :, :], v_out[:, :
-                                                                 seq_len, :, :]
+            k_states = self._gather_paged_kv(kv_cache_tensor, block_ids, 0,
+                                             seq_len, tokens_per_block)[None]
+            v_states = self._gather_paged_kv(kv_cache_tensor, block_ids, 1,
+                                             seq_len, tokens_per_block)[None]
         return k_states, v_states
 
     def _single_request_preprocess_inputs(self, q, k, v, kv_dtype):
@@ -281,7 +352,7 @@ def _single_request_forward(self,
                                 attention_mask: AttentionMask,
                                 kv_cache_tensor,
                                 past_seen_token,
-                                cache_idx,
+                                block_ids,
                                 sample_idx,
                                 metadata: AttentionMetadata,
                                 attention_window_size: Optional[int] = None):
@@ -297,7 +368,7 @@ def _single_request_forward(self,
 
         # update kv cache
         key_states, value_states = self._single_request_update_kv_cache(
-            k, v, kv_cache_tensor, past_seen_token, kv_len, cache_idx,
+            k, v, kv_cache_tensor, past_seen_token, kv_len, block_ids,
             sparse_kv_indices)
 
         # predict sparse attn indices
@@ -523,7 +594,7 @@ def _mla_forward_generation(self, fused_q: torch.Tensor,
         )
         past = metadata.kv_cache_params.num_cached_tokens_per_seq
         cache_indices = [
-            block_ids[0] for block_ids in metadata.block_ids_per_seq
+            list(block_ids) for block_ids in metadata.block_ids_per_seq
         ]
 
         # MLA scales by the q/k head_dim (qk_nope + qk_rope), not the latent dim.
@@ -534,12 +605,13 @@ def _mla_forward_generation(self, fused_q: torch.Tensor,
         offset = 0
         for i, q_len in enumerate(metadata.seq_lens.tolist()):
             past_i = int(past[i])
-            ci = cache_indices[i]
+            blocks_i = cache_indices[i]
             kv_len = past_i + q_len
 
             # K is the full latent ([compressed_kv | k_pe]); V is the kv_lora
             # slice. One latent head broadcast (MQA) to all query heads.
-            latent = kv_cache[ci, 0, :kv_len, 0, :].to(q.dtype)
+            latent = self._gather_paged_mla_latent(kv_cache, blocks_i,
+                                                   kv_len).to(q.dtype)
             k = latent[None, None]  # [1, 1, kv_len, d_latent]
             v = latent[None,
                        None, :, :self.kv_lora_rank]  # [1, 1, kv_len, kv_lora]
@@ -664,7 +736,7 @@ def forward(self,
 
         past_seen_tokens = metadata.kv_cache_params.num_cached_tokens_per_seq
         cache_indices = [
-            block_ids[0] for block_ids in metadata.block_ids_per_seq
+            list(block_ids) for block_ids in metadata.block_ids_per_seq
         ]
         kv_cache_tensor = metadata.kv_cache_manager.get_buffers(
             self.layer_idx, kv_layout=metadata.kv_layout)
@@ -688,11 +760,11 @@ def forward(self,
                 seq_len_kv] if v is not None and seq_len_kv != 0 else None
 
             past_seen_token = past_seen_tokens[sample_idx]
-            cache_idx = cache_indices[sample_idx]
+            block_ids = cache_indices[sample_idx]
 
             attn_output = self._single_request_forward(
                 single_q, single_k, single_v, forward_args.attention_mask,
-                kv_cache_tensor, past_seen_token, cache_idx, sample_idx,
+                kv_cache_tensor, past_seen_token, block_ids, sample_idx,
                 metadata, forward_args.attention_window_size)
 
             attn_outputs.append(attn_output)

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -417,8 +417,18 @@ def create_py_executor(
     ) = llm_args.get_runtime_sizes()
 
     tokens_per_block = kv_cache_config.tokens_per_block
-    if llm_args.attn_backend == "VANILLA":
+
+    # RocketKV's Vanilla path keeps its landmark (KT) cache in a single block per
+    # sequence: RocketVanillaAttention writes the whole sequence into
+    # kt_cache_block_offsets[0], and kt_tokens_per_block is derived from
+    # tokens_per_block. It does not support a paged KT cache, so force one block
+    # per sequence for it. Plain Vanilla attention supports paged KV cache and is
+    # left untouched.
+    sparse_config = llm_args.sparse_attention_config
+    if (llm_args.attn_backend == "VANILLA" and sparse_config is not None
+            and getattr(sparse_config, "algorithm", None) == "rocket"):
         tokens_per_block = max_num_tokens
+        kv_cache_config.tokens_per_block = tokens_per_block
 
     # The MSA kernels require a page size of 128; the Triton reference uses TRT-LLM's default
     # of 32.

tests/unittest/_torch/attention/backend_capability.py

@@ -55,7 +55,7 @@
         kv_layouts=("NHD", "HND"),  # selectable via metadata.kv_layout
     ),
     "VANILLA": dict(
-        paged=False,
+        paged=True,
         fp8_kv=True,
         fp4_kv=False,
         sliding_window=True,

tests/unittest/_torch/attention/test_vanilla_attention.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import unittest
 from unittest.mock import patch
 
@@ -10,13 +13,88 @@
 from tensorrt_llm._torch.attention_backend.interface import \
     PredefinedAttentionMask
 from tensorrt_llm._torch.metadata import KVCacheParams
+from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm.bindings.executor import KvCacheConfig
+from tensorrt_llm.llmapi.llm_args import KvCacheConfig as LlmKvCacheConfig
 from tensorrt_llm.mapping import Mapping
 
 
 class TestVanillaAttention(unittest.TestCase):
 
+    def test_kv_cache_manager_v2_sliding_window_eviction(self):
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
+        tokens_per_block = 2
+        mapping = Mapping(world_size=1, tp_size=1, rank=0)
+        manager = KVCacheManagerV2(
+            LlmKvCacheConfig(max_tokens=16,
+                             max_attention_window=[4],
+                             enable_block_reuse=False),
+            tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
+            num_layers=1,
+            num_kv_heads=1,
+            head_dim=2,
+            tokens_per_block=tokens_per_block,
+            max_seq_len=8,
+            max_batch_size=1,
+            mapping=mapping,
+            dtype=tensorrt_llm.bindings.DataType.BF16,
+        )
+
+        try:
+            manager.add_dummy_requests([0], [7], is_gen=True)
+            block_ids = manager.get_batch_cache_indices([0], 0)[0]
+            self.assertEqual(block_ids[0], -1)
+
+            cached_k = torch.tensor(
+                [[0.2, -0.1], [0.4, 0.3], [-0.2, 0.5], [0.1, 0.7]],
+                device=device,
+                dtype=dtype)
+            cached_v = torch.tensor(
+                [[0.6, -0.4], [0.8, 0.2], [-0.3, 0.9], [0.5, 0.1]],
+                device=device,
+                dtype=dtype)
+            kv_cache = manager.get_buffers(0, kv_layout="NHD")
+            for offset, position in enumerate(range(2, 6)):
+                block = block_ids[position // tokens_per_block]
+                block_offset = position % tokens_per_block
+                kv_cache[block, 0, block_offset, 0].copy_(cached_k[offset])
+                kv_cache[block, 1, block_offset, 0].copy_(cached_v[offset])
+
+            q = torch.tensor([[0.3, -0.6]], device=device, dtype=dtype)
+            new_k = torch.tensor([[0.7, 0.2]], device=device, dtype=dtype)
+            new_v = torch.tensor([[-0.5, 0.4]], device=device, dtype=dtype)
+            metadata = VanillaAttentionMetadata(
+                seq_lens=torch.tensor([1], dtype=torch.int),
+                num_contexts=0,
+                kv_cache_params=KVCacheParams(use_cache=True,
+                                              num_cached_tokens_per_seq=[6]),
+                max_num_requests=1,
+                max_num_tokens=1,
+                kv_cache_manager=manager,
+                request_ids=[0],
+            )
+            metadata.prepare()
+
+            vanilla_attn = VanillaAttention(layer_idx=0,
+                                            num_heads=1,
+                                            head_dim=2,
+                                            num_kv_heads=1)
+            actual = vanilla_attn.forward(q,
+                                          new_k,
+                                          new_v,
+                                          metadata,
+                                          attention_window_size=4)
+
+            reference_k = torch.cat((cached_k[1:], new_k)).view(1, 1, 4, 2)
+            reference_v = torch.cat((cached_v[1:], new_v)).view(1, 1, 4, 2)
+            expected = F.scaled_dot_product_attention(q.view(
+                1, 1, 1, 2), reference_k, reference_v).view_as(actual)
+            torch.testing.assert_close(actual, expected)
+        finally:
+            manager.shutdown()
+
     def test_sdpa_fallback_uses_metadata_cross_flag_for_causal_mask(self):
         vanilla_attn = VanillaAttention(layer_idx=0,
                                         num_heads=1,