diff --git a/tensorrt_llm/_torch/attention_backend/vanilla.py b/tensorrt_llm/_torch/attention_backend/vanilla.py index eef2fc520906..462921892933 100644 --- a/tensorrt_llm/_torch/attention_backend/vanilla.py +++ b/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,14 +133,66 @@ 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) @@ -147,31 +200,49 @@ def _single_request_update_kv_cache(self, 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) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 6149ba5ec0c3..9f6c0e5cbef2 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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. diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index 91d11be455df..19d9d7230dee 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/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, diff --git a/tests/unittest/_torch/attention/test_vanilla_attention.py b/tests/unittest/_torch/attention/test_vanilla_attention.py index 89d99d679ff2..7de13815c359 100644 --- a/tests/unittest/_torch/attention/test_vanilla_attention.py +++ b/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,