rocketkvkv-cachellm-inferencesparse-attentionlong-contexttensorrt-llmtutorial

RocketKV Guide: Two-Stage KV Cache Compression for LLMs

RocketKV Architecture: Two-Stage KV Cache Compression with Coarse-Grain SnapKV Eviction and Fine-Grain Hybrid Sparse Attention

Processing extended context windows—from 32k up to 1M tokens—has become a foundational capability for enterprise software, automated code analysis, multi-document research copilots, and real-time knowledge ingestion engines powered by AI Fetchly. Whether summarizing entire legal repositories, navigating intricate corporate knowledge bases, or orchestrating web-scale intelligence through automated data fetching, large language models (LLMs) must sustain deep contextual awareness across tens of thousands of tokens without stalling inference servers.

However, scaling sequence length introduces a crippling bottleneck during inference: the rapid, unbounded growth of the Key-Value (KV) cache. In autoregressive Transformers like LLaMA-3, Mistral, and Qwen, every newly generated token requires attending to the cached key and value states of all preceding tokens. For an 8B parameter model serving a 128k context in FP16 precision, the KV cache alone consumes over 16 GB of GPU High-Bandwidth Memory (HBM) for a single concurrent stream. In multi-tenant enterprise deployments, this memory consumption forces batch sizes down to one, exhausts GPU capacity, and degrades throughput as memory bandwidth saturates.

To solve this memory crisis without requiring expensive model retraining or architectural re-engineering, researchers from NVIDIA introduced RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (arXiv:2502.14051, NVlabs/RocketKV, presented at ICML 2025 and officially merged into NVIDIA TensorRT-LLM).

RocketKV introduces a unified, two-stage compression pipeline that decouples prompt-level eviction from token-level dynamic selection. By pairing coarse-grain permanent pruning with a lightweight auxiliary KT Cache that estimates fine-grain attention scores across reduced head and sequence dimensions, RocketKV achieves up to 400x compression ratios, cuts peak memory usage by 32.6%, and delivers up to 3.7x end-to-end decoding speedup on NVIDIA A100 GPUs while maintaining near-oracle accuracy across standard long-context benchmarks.

In this comprehensive technical guide, we break down the mathematical foundations of RocketKV, explore the mechanics of Stage 1 coarse eviction and Stage 2 Hybrid Sparse Attention (HSA), contrast RocketKV with competing compression frameworks, walk through an end-to-end PyTorch implementation, and demonstrate how to configure RocketKV in NVIDIA TensorRT-LLM for production inference.

The Problem: The KV Cache Memory Bandwidth Wall in Long-Context LLMs

To understand why traditional KV cache management collapses under extended contexts, we must examine the physical and architectural constraints governing Transformer autoregressive decoding: memory capacity limits, memory bandwidth saturation (IO-bound decoding), and the false dilemma of existing compression techniques.

The Long-Context Autoregressive Bottleneck: Input Prompt: 128k Tokens | Concurrency: 8 Requests | Model: LLaMA-3-8B (FP16) │ ├── 1. Memory Capacity Explosion: │ └── KV Cache Size = 2 * n_layers (32) * n_kv_heads (8) * d_head (128) * seq_len (131,072) * 2 bytes │ └── Per Request: ~17.18 GB VRAM! │ └── 8 Concurrent Streams: 137.4 GB VRAM -> Out of Memory (OOM) on an 80GB A100! │ ├── 2. Memory Bandwidth Wall (Decoding Arithmetic Intensity): │ └── Arithmetic Intensity = FLOPs / Memory Access ≈ 2 * Params / KV_Cache_Size │ └── Standard Generation: Model loads 17GB of KV weights to calculate 1 token! │ └── GPU Compute Units sit idle waiting for HBM bus transfers. │ └── 3. The Compression Tradeoff Dilemma: ├── Static Eviction (H2O, Scissorhands): Drops tokens permanently; fails dynamic multi-hop reasoning. ├── Pure Top-K Sparsity (Quest, SparseKV): Full KV cache remains in VRAM; saves bandwidth but NOT memory capacity. └── RocketKV Two-Stage Resolution: Coarse Eviction (Capacity) + Fine Top-K (Bandwidth)!

The Linear Capacity Explosion

In standard Multi-Head Attention (MHA) and Grouped-Query Attention (GQA), the memory footprint of the KV cache grows strictly linearly with sequence length $S$:

$$\text{Memory}{\text{KV}} = 2 \times L \times H{\text{kv}} \times D_{\text{head}} \times S \times B_{\text{dtype}}$$

Where $L$ is the number of transformer layers, $H_{\text{kv}}$ is the number of key-value attention heads, $D_{\text{head}}$ is the head dimension, $S$ is the total sequence length, and $B_{\text{dtype}}$ is the byte width per element (2 bytes for FP16/BF16, 1 byte for FP8).

Consider deploying a production copilot using LLaMA-3.1-8B ($L=32, H_{\text{kv}}=8, D_{\text{head}}=128$):

  • At a modest context of 4k tokens: the KV cache consumes approximately $536 \text{ MB}$ per stream.
  • At an extended context of 128k tokens: the KV cache consumes $17.18 \text{ GB}$ per stream.
  • Serving a modest batch of 4 concurrent requests at 128k context requires $68.7 \text{ GB}$ of VRAM dedicated purely to KV storage, leaving virtually zero headroom for model weights ($16 \text{ GB}$) or activation buffers on an 80GB GPU.

The Memory Bandwidth Bottleneck During Autoregressive Decoding

While the prefill phase (processing the input prompt) is compute-bound and saturates GPU tensor cores with dense matrix multiplications, the autoregressive generation phase (producing tokens one by one) is strictly memory bandwidth-bound.

For every single token generated:

  1. The query vector $\mathbf{q}_t \in \mathbb{R}^{1 \times D}$ is projected.
  2. The entire historical KV tensor $\mathbf{K}, \mathbf{V} \in \mathbb{R}^{S \times D}$ must be streamed out of GPU High-Bandwidth Memory (HBM) into on-chip SRAM registers.
  3. The arithmetic intensity of this operation is extremely low ($\approx 1 \text{ FLOP/byte}$).

Because the GPU memory bus cannot transfer tens of gigabytes instantaneously, the generation speed drops to a fraction of its potential throughput.

The Failure of Single-Paradigm Compression Approaches

Prior research tackled this bottleneck through two divergent, isolated strategies—both of which introduce critical shortcomings:

  1. Permanent Token Eviction (e.g., StreamingLLM, H2O, SnapKV): These methods monitor attention weights and discard tokens deemed unimportant. While they successfully reclaim memory capacity by shrinking the physical KV tensor, permanent pruning is irreversible. If a query later in the conversation requires a piece of information discarded during prefill, the model hallucinates or fails completely.
  2. Dynamic Top-K Sparse Attention (e.g., Quest, SparseKV, Flash-Decoding++): These methods keep the entire KV cache in VRAM and dynamically gather only the top-$k$ relevant tokens for each decoding step. While dynamic selection preserves accuracy and slashes memory bandwidth consumption during computation, it does not solve the memory capacity crisis—the full multi-gigabyte KV cache must still reside in VRAM, preventing large batch sizes.

RocketKV resolves this fundamental contradiction by asking: Can we combine the physical capacity savings of permanent eviction with the adaptive retrieval accuracy of dynamic sparse attention?

RocketKV Architecture: Two-Stage Compression Mechanics

RocketKV formulates KV cache compression as a coordinated two-stage hierarchy: Stage 1 coarse-grain permanent eviction during context prefill, followed by Stage 2 fine-grain dynamic selection via Hybrid Sparse Attention (HSA) during autoregressive generation.

┌────────────────────────────────────────────────────────────────────────────────────────┐ │ ROCKETKV TWO-STAGE ARCHITECTURE │ ├────────────────────────────────────────────────────────────────────────────────────────┤ │ │ │ STAGE 1: CONTEXT PREFILL PHASE (Permanent Capacity Reduction) │ │ Input Prompt Tokens: [t_1, t_2, ..., t_N] (e.g., N = 128,000) │ │ │ │ │ ├── 1. Observation Window (W_obs = 32 tokens at prompt tail) │ │ ├── 2. Compute Attention Distribution: A_obs = Softmax(Q_obs · K^T / sqrt(d)) │ │ ├── 3. Feature Importance Pooling: Score_i = Sum_{w in W_obs} A_{w, i} │ │ └── 4. Coarse Eviction: Retain initial prefix + Top Prompt Budget (P = 1,024) │ │ ──▶ Physical KV Cache drops from 128,000 to 1,024 tokens! (125x reduction) │ │ │ │ TRANSITION: AUXILIARY KT CACHE GENERATION │ │ Build lightweight compressed Key tracking cache (KT Cache): │ │ ├── Head Dimension Pooling: Reduce H attention heads to representative heads │ │ └── Sequence Pooling: Downsample sequence length with 1D average pooling (k_down) │ │ ──▶ KT Cache footprint is < 1% of original KV cache size! │ │ │ │ STAGE 2: GENERATION PHASE (Dynamic Bandwidth Reduction via HSA) │ │ Generating New Token at step t with Query q_t: │ │ │ │ │ ├── 1. Query KT Cache: Fast Approximate Scores = q_t · (K_T)^T │ │ ├── 2. Dynamic Top-K Selection: Identify Top-K active token blocks (e.g., K = 128) │ │ ├── 3. Gather Physical KV Blocks from Stage 1 Cache │ │ └── 4. Exact Kernel Execution: Run FlashAttention-2 over gathered Top-K subset │ │ ──▶ Memory Bandwidth slashed by 8x - 16x during token decode! │ │ │ └────────────────────────────────────────────────────────────────────────────────────────┘

Stage 1: Coarse-Grain Permanent Eviction via SnapKV

In the context prefill phase, the full prompt is processed in parallel. RocketKV leverages an enhanced version of SnapKV to perform irreversible, coarse-grain pruning:

  1. Observation Window ($W_{\text{obs}}$): Research demonstrates that tokens at the end of the input prompt (immediately before generation begins) exhibit a stable, highly predictive attention distribution over earlier context tokens. RocketKV defines an observation window of size $W_{\text{obs}} = 32$ tokens at the prompt boundary.

  2. Attention Aggregation: For each attention head $h$, the queries in the observation window $\mathbf{Q}{\text{obs}} \in \mathbb{R}^{W{\text{obs}} \times D}$ compute attention weights across all prompt keys $\mathbf{K} \in \mathbb{R}^{S \times D}$:

    $$\mathbf{A}{\text{obs}}^{(h)} = \text{Softmax}\left(\frac{\mathbf{Q}{\text{obs}}^{(h)} (\mathbf{K}^{(h)})^T}{\sqrt{D}}\right) \in \mathbb{R}^{W_{\text{obs}} \times S}$$

  3. Importance Scoring and Pruning: RocketKV sums the attention weights across the observation window and applies 1D pooling to capture contiguous semantic clusters. It retains a fixed prompt budget $P$ (e.g., $P = 1,024$ or $P = 2,048$ tokens out of $128,000$) consisting of:

    • Sink Tokens: The initial 4 to 32 prompt tokens, which act as attention sinks to preserve softmax stability.
    • Important Clusters: The highest-scoring token indices identified by the observation window.
    • Recent Window: The trailing tokens immediately preceding generation.

All remaining low-importance tokens are permanently deallocated from GPU HBM. This single step delivers a 30x to 120x reduction in physical KV cache memory capacity.

Stage 2: Fine-Grain Hybrid Sparse Attention (HSA)

While Stage 1 compresses the physical cache down to a compact budget $P$, loading $1,024$ or $2,048$ tokens on every decoding step still strains memory bandwidth for long generations or high batch concurrency. Furthermore, at any given generation step $t$, the current query $\mathbf{q}_t$ may only need a tiny subset (e.g., 64 or 128 tokens) of the remaining cache.

Stage 2 applies Hybrid Sparse Attention (HSA) to dynamically select only the top-$k$ relevant tokens for the active decoding step. To perform this selection with minimal compute overhead, RocketKV constructs a lightweight auxiliary cache called the KT Cache (Key Tracking Cache).

The KT Cache and Dual Dimensionality Reduction

Calculating exact query-key dot products across all remaining keys to find the top-$k$ tokens would defeat the purpose of sparse loading. RocketKV approximates attention scores by compressing keys along two distinct dimensions:

  1. Head Dimensionality Reduction: In modern models with Grouped-Query Attention (GQA), multiple query heads share a key-value head. Furthermore, attention patterns across heads exhibit significant correlation. RocketKV pools or selects representative heads to construct a low-dimensional key representation $\mathbf{K}_{\text{rep}}$.
  2. Sequence Dimensionality Reduction: RocketKV downsamples the sequence dimension using 1D average pooling with stride $k_{\text{down}}$ (typically 4 or 8 tokens per block). This groups adjacent tokens into semantic chunks and cuts the sequence evaluation length by $k_{\text{down}}\times$.

The resulting KT Cache $\mathbf{K}_T$ is minuscule—occupying less than 1% of the original KV cache storage.

Fast Top-K Scoring and Block Gathering

During each generation step $t$:

  1. The model computes fast approximate similarity scores:

    $$\mathbf{S}t = \mathbf{q}{t, \text{proj}} \mathbf{K}_T^T$$

  2. RocketKV selects the top-$k$ candidate blocks with the highest approximate scores.

  3. The inference engine gathers the full-precision keys and values only for these top-$k$ blocks from the Stage 1 physical cache.

  4. An optimized FlashAttention-2 or TensorRT-LLM sparse kernel executes exact self-attention over the selected tokens.

RocketKV-MT: Multi-Turn Conversation Adaptation

In multi-turn chat dialogues, a user might prompt the model, receive a response, and then ask a follow-up question referencing context that was initially irrelevant in Turn 1.

If coarse eviction from Stage 1 permanently dropped those tokens during Turn 1, Turn 2 would fail. To handle multi-turn workloads, NVIDIA designed RocketKV-MT:

  • In RocketKV-MT, unselected tokens are not immediately deleted from memory; instead, they are held in a secondary low-priority host or compressed memory tier.
  • As new turns arrive, the observation window shifts to the latest user prompt, allowing dynamically restored tokens to re-enter the primary Stage 1 cache while evicting completed conversational branches.

Architectural Comparison: RocketKV vs. Competing KV Cache Optimizations

To contextualize RocketKV’s engineering innovations, we contrast it against the leading KV cache compression and sparse attention frameworks: Full Cache Baseline, StreamingLLM, H2O, Quest, and SnapKV.

Feature / MetricFull KV BaselineStreamingLLMH2O (Heavy Hitter)Quest (Sparse Attention)SnapKVRocketKV (ICML 2025)
Compression ParadigmNone (Dense)Static Window EvictionDynamic Eviction (LRA)Top-K Page SparsityObservation PruningTwo-Stage (Evict + HSA)
VRAM Capacity Savings0% (Unbounded)High (Fixed Window)High (Bounded)0% (Full Cache in VRAM)High (Prompt Budget)Up to 400x (Extreme)
Bandwidth Reduction0% (Reads all KV)HighHighHigh (Top-K loading)ModerateHigh (Dual Reduction)
Needle Retrieval Accuracy100%Fails on long contextModerate DegradationHigh (>98%)High (>97%)Near-Oracle (>99%)
Multi-Turn HandlingPerfect (Heavy)Amnesia on old turnsAmnesia on old turnsPreservedPruned PermanentlySupported (RocketKV-MT)
Training Required?NoneNoneNoneNoneNoneNone (Training-Free)
Decoding Speedup1.0x (Baseline)Up to 2.5xUp to 2.8xUp to 2.2xUp to 2.4xUp to 3.7x (A100)
Production IntegrationStandard vLLM / SGLangCustom KernelsCustom ScriptsAcademic CUDAHuggingFace ForksNVIDIA TensorRT-LLM

Why RocketKV Outperforms Pure Sparse Attention (Quest)

Frameworks like Quest divide the KV cache into pages and use min-max bounding boxes to load only relevant pages. While Quest reduces memory bandwidth consumption, the entire 128k-token cache must remain resident in GPU VRAM. As a result, Quest cannot increase serving batch size on memory-constrained GPUs. RocketKV’s Stage 1 solves this by physically freeing over 90% of VRAM before Stage 2 begins.

Why RocketKV Outperforms Pure Eviction (SnapKV)

SnapKV compresses the prompt down to a fixed budget (e.g., 2,048 tokens). However, during decoding, SnapKV reads all 2,048 tokens on every single generation step. RocketKV’s Stage 2 reduces this active set to 128 or 256 tokens using the KT Cache, delivering an additional 2x–4x reduction in decoding memory bandwidth.

End-to-End PyTorch Implementation: Building RocketKV

To understand how RocketKV operates under the hood, let us build a complete, modular PyTorch implementation featuring Stage 1 SnapKV coarse eviction, auxiliary KT Cache construction, and Stage 2 Hybrid Sparse Attention (HSA).

""" RocketKV: Two-Stage KV Cache Compression Implementation Includes: 1. Stage 1: Coarse-grain permanent eviction via Observation Window (SnapKV) 2. Auxiliary Key Tracking Cache (KT Cache) with Dual Dimensionality Reduction 3. Stage 2: Dynamic Top-K Hybrid Sparse Attention (HSA) """ import torch import torch.nn as nn import torch.nn.functional as F from typing import Tuple, Optional, Dict, List class RocketKVCacheManager: """ Manages the two-stage lifecycle of Key-Value states in RocketKV. """ def __init__( self, num_layers: int = 32, num_heads: int = 32, num_kv_heads: int = 8, head_dim: int = 128, obs_window_size: int = 32, prompt_budget: int = 1024, top_k_blocks: int = 16, block_size: int = 16, kernel_downsample: int = 4, device: str = "cuda" if torch.cuda.is_available() else "cpu", dtype: torch.dtype = torch.float16 ): self.num_layers = num_layers self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = head_dim self.obs_window_size = obs_window_size self.prompt_budget = prompt_budget self.top_k_blocks = top_k_blocks self.block_size = block_size self.kernel_downsample = kernel_downsample self.device = device self.dtype = dtype # Physical Stage 1 Storage: Layer -> {"k": Tensor, "v": Tensor} self.physical_cache: Dict[int, Dict[str, torch.Tensor]] = {} # Auxiliary Stage 2 KT Storage: Layer -> Tensor [Batch, KV_Heads, Seq_Down, Head_Dim] self.kt_cache: Dict[int, torch.Tensor] = {} def stage1_coarse_eviction( self, layer_idx: int, key_states: torch.Tensor, value_states: torch.Tensor, query_states: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """ Stage 1: Prunes full prefill keys/values using observation window attention. Args: key_states: [batch_size, num_kv_heads, seq_len, head_dim] value_states: [batch_size, num_kv_heads, seq_len, head_dim] query_states: [batch_size, num_heads, seq_len, head_dim] Returns: Compressed (key_states, value_states) within prompt_budget """ batch_size, num_kv_heads, seq_len, head_dim = key_states.shape if seq_len <= self.prompt_budget: self.physical_cache[layer_idx] = {"k": key_states, "v": value_states} self._build_kt_cache(layer_idx, key_states) return key_states, value_states # Extract observation window queries from tail obs_q = query_states[:, :, -self.obs_window_size:, :] # [B, H, W_obs, D] # Repeat KV heads for GQA alignment if needed gqa_ratio = self.num_heads // num_kv_heads k_expanded = key_states.repeat_interleave(gqa_ratio, dim=1) # [B, H, S, D] # Compute raw attention weights across observation window attn_scores = torch.matmul(obs_q, k_expanded.transpose(-1, -2)) / (head_dim ** 0.5) attn_probs = F.softmax(attn_scores, dim=-1) # [B, H, W_obs, S] # Aggregate importance scores across observation window queries importance_scores = attn_probs.sum(dim=2) # [B, H, S] # Pool across grouped heads for GQA consistency importance_scores = importance_scores.view(batch_size, num_kv_heads, gqa_ratio, seq_len).mean(dim=2) # importance_scores: [B, num_kv_heads, S] # Guarantee attention sinks (first 16 tokens) and recent tokens (32 tokens) sink_count = 16 recent_count = self.obs_window_size cluster_budget = self.prompt_budget - sink_count - recent_count # Mask sinks and recents from cluster selection score_mask = importance_scores.clone() score_mask[:, :, :sink_count] = -float("inf") score_mask[:, :, -recent_count:] = -float("inf") # Top-K cluster selection _, topk_cluster_indices = torch.topk(score_mask, k=cluster_budget, dim=-1, largest=True) sink_indices = torch.arange(sink_count, device=self.device).expand(batch_size, num_kv_heads, sink_count) recent_indices = torch.arange(seq_len - recent_count, seq_len, device=self.device).expand(batch_size, num_kv_heads, recent_count) # Merge and sort indices to preserve causal token ordering selected_indices = torch.cat([sink_indices, topk_cluster_indices, recent_indices], dim=-1) selected_indices, _ = torch.sort(selected_indices, dim=-1) # [B, num_kv_heads, prompt_budget] # Gather pruned tensors idx_expanded = selected_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim) pruned_k = torch.gather(key_states, dim=2, index=idx_expanded) pruned_v = torch.gather(value_states, dim=2, index=idx_expanded) self.physical_cache[layer_idx] = {"k": pruned_k, "v": pruned_v} self._build_kt_cache(layer_idx, pruned_k) return pruned_k, pruned_v def _build_kt_cache(self, layer_idx: int, key_states: torch.Tensor): """ Builds the auxiliary KT Cache using sequence downsampling. Args: key_states: [B, num_kv_heads, P_budget, head_dim] """ # 1D average pooling along the sequence dimension # Reshape to [B * num_kv_heads, head_dim, P_budget] for 1D pooling b, h, p, d = key_states.shape k_transposed = key_states.permute(0, 1, 3, 2).reshape(b * h, d, p) pooled_k = F.avg_pool1d( k_transposed, kernel_size=self.kernel_downsample, stride=self.kernel_downsample, ceil_mode=True ) # [B * h, d, P_down] pooled_k = pooled_k.view(b, h, d, -1).permute(0, 1, 3, 2) # pooled_k: [B, num_kv_heads, P_down, head_dim] self.kt_cache[layer_idx] = pooled_k def stage2_hybrid_sparse_decode( self, layer_idx: int, query_token: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """ Stage 2: Dynamically selects active KV blocks for current token decoding. Args: query_token: [batch_size, num_heads, 1, head_dim] Returns: gathered_k, gathered_v: Active subset of KV states for FlashAttention """ phys_k = self.physical_cache[layer_idx]["k"] phys_v = self.physical_cache[layer_idx]["v"] kt_k = self.kt_cache[layer_idx] b, num_kv_heads, p_down, d = kt_k.shape gqa_ratio = self.num_heads // num_kv_heads # Group query heads to match KV heads for lightweight scoring q_grouped = query_token.view(b, num_kv_heads, gqa_ratio, 1, d).mean(dim=2) # [B, num_kv_heads, 1, d] # Fast approximate similarity using KT Cache approx_scores = torch.matmul(q_grouped, kt_k.transpose(-1, -2)).squeeze(2) # [B, num_kv_heads, P_down] # Top-K downsampled block selection k_blocks = min(self.top_k_blocks, p_down) _, topk_block_idx = torch.topk(approx_scores, k=k_blocks, dim=-1, largest=True) # [B, num_kv_heads, k_blocks] # Expand downsampled blocks to physical indices # Each block corresponds to kernel_downsample tokens expanded_indices = [] for i in range(self.kernel_downsample): expanded_indices.append(topk_block_idx * self.kernel_downsample + i) active_indices = torch.cat(expanded_indices, dim=-1) # Clamp to physical sequence length p_len = phys_k.shape[2] active_indices = torch.clamp(active_indices, max=p_len - 1) active_indices, _ = torch.sort(active_indices, dim=-1) # Gather active KV states idx_gather = active_indices.unsqueeze(-1).expand(-1, -1, -1, d) active_k = torch.gather(phys_k, dim=2, index=idx_gather) active_v = torch.gather(phys_v, dim=2, index=idx_gather) return active_k, active_v class RocketKVAttentionLayer(nn.Module): """ Transformer Attention Layer wrapped with RocketKV Two-Stage Acceleration. """ def __init__( self, embed_dim: int = 4096, num_heads: int = 32, num_kv_heads: int = 8, layer_idx: int = 0, manager: Optional[RocketKVCacheManager] = None ): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = embed_dim // num_heads self.layer_idx = layer_idx self.manager = manager or RocketKVCacheManager() self.q_proj = nn.Linear(embed_dim, num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(embed_dim, num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(embed_dim, num_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(embed_dim, embed_dim, bias=False) def forward( self, hidden_states: torch.Tensor, is_prefill: bool = True ) -> torch.Tensor: b, seq_len, _ = hidden_states.shape q = self.q_proj(hidden_states).view(b, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(hidden_states).view(b, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(hidden_states).view(b, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) if is_prefill: # Stage 1: Coarse-grain permanent eviction comp_k, comp_v = self.manager.stage1_coarse_eviction(self.layer_idx, k, v, q) # Execute prefill attention over compressed prompt gqa_ratio = self.num_heads // self.num_kv_heads k_full = comp_k.repeat_interleave(gqa_ratio, dim=1) v_full = comp_v.repeat_interleave(gqa_ratio, dim=1) attn_out = F.scaled_dot_product_attention(q, k_full, v_full, is_causal=True) else: # Stage 2: Fine-grain dynamic top-K selection via HSA active_k, active_v = self.manager.stage2_hybrid_sparse_decode(self.layer_idx, q) gqa_ratio = self.num_heads // self.num_kv_heads active_k_full = active_k.repeat_interleave(gqa_ratio, dim=1) active_v_full = active_v.repeat_interleave(gqa_ratio, dim=1) attn_out = F.scaled_dot_product_attention(q, active_k_full, active_v_full) attn_out = attn_out.transpose(1, 2).reshape(b, seq_len, self.embed_dim) return self.o_proj(attn_out) if __name__ == "__main__": device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Executing RocketKV Validation on {device}...") manager = RocketKVCacheManager( num_layers=1, num_heads=32, num_kv_heads=8, head_dim=128, obs_window_size=32, prompt_budget=256, top_k_blocks=8, block_size=16, kernel_downsample=4, device=device ) layer = RocketKVAttentionLayer( embed_dim=4096, num_heads=32, num_kv_heads=8, layer_idx=0, manager=manager ).to(device) # 1. Simulate 4k-token prefill prompt prompt_seq_len = 4096 x_prefill = torch.randn(1, prompt_seq_len, 4096, device=device) out_prefill = layer(x_prefill, is_prefill=True) print(f"[Prefill] Input: {x_prefill.shape} -> Output: {out_prefill.shape}") print(f"[Stage 1] Physical KV Cache pruned to: {manager.physical_cache[0]['k'].shape}") print(f"[KT Cache] Downsampled Key Cache built: {manager.kt_cache[0].shape}") # 2. Simulate single-token autoregressive decoding step x_decode = torch.randn(1, 1, 4096, device=device) out_decode = layer(x_decode, is_prefill=False) print(f"[Decode] Input: {x_decode.shape} -> Output: {out_decode.shape}") print("RocketKV end-to-end execution verified successfully!")

Production Deployment: Configuring RocketKV in NVIDIA TensorRT-LLM

RocketKV has been officially integrated into NVIDIA’s enterprise inference compiler, TensorRT-LLM (TensorRT-LLM Sparse Attention Guide). This allows engineers to deploy RocketKV in high-throughput production clusters without writing custom CUDA kernels.

┌──────────────────────────────────────────────────────────────────────────────────┐ │ TENSORRT-LLM ROCKETKV SERVING ARCHITECTURE │ ├──────────────────────────────────────────────────────────────────────────────────┤ │ │ │ Inference Request (128k Tokens) │ │ │ │ │ ▼ │ │ ┌────────────────────────────────────────────────────────────────────────┐ │ │ │ TensorRT-LLM Context Engine (Prefill) │ │ │ │ - Evaluates Prompt Tokens via FlashAttention │ │ │ │ - Invokes `KVCacheManager` with `prompt_budget=1024` │ │ │ │ - Performs SnapKV Coarse Eviction into Paged KV Cache Pool │ │ │ └────────────────────────────────────────────────────────────────────────┘ │ │ │ │ │ ├── Physical Paged KV Cache (Truncated to 1,024 slots) │ │ └── Lightweight KT Cache (< 1% Memory Overhead) │ │ │ │ │ ▼ │ │ ┌────────────────────────────────────────────────────────────────────────┐ │ │ │ TensorRT-LLM Generation Engine (Decode) │ │ │ │ - Dynamic Top-K Selection Kernel queries KT Cache │ │ │ │ - Loads sparse Top-K blocks (e.g., 128 active tokens) │ │ │ │ - Fused Paged Sparse Attention Execution │ │ │ │ - Delivers 3.7x Token-to-Token Speedup! │ │ │ └────────────────────────────────────────────────────────────────────────┘ │ │ │ └──────────────────────────────────────────────────────────────────────────────────┘

Step 1: Building the TensorRT Engine with Sparse Attention

To enable RocketKV during engine compilation, configure the model build command with sparse attention flags:

# Clone and build TensorRT-LLM with sparse attention support git clone https://github.com/NVIDIA/TensorRT-LLM.git cd TensorRT-LLM/examples/sparse_attention # Convert HuggingFace LLaMA-3.1-8B-Instruct to TensorRT checkpoint python3 convert_checkpoint.py \ --model_dir ./Meta-Llama-3.1-8B-Instruct \ --output_dir ./tllm_checkpoint_llama_8b \ --dtype bfloat16 # Build TensorRT engine enabling RocketKV sparse attention trtllm-build \ --checkpoint_dir ./tllm_checkpoint_llama_8b \ --output_dir ./engine_rocketkv_llama_8b \ --gemm_plugin bfloat16 \ --max_batch_size 16 \ --max_input_len 131072 \ --max_seq_len 135168 \ --tokens_per_block 64 \ --use_custom_all_reduce enable

Step 2: Runtime Configuration Script

Deploying RocketKV in TensorRT-LLM requires configuring the sparse_attention_config parameter within the model runner:

""" TensorRT-LLM RocketKV Production Serving Script """ import tensorrt_llm from tensorrt_llm.runtime import ModelRunner import torch # Define RocketKV Sparse Attention Parameters rocketkv_config = { "sparse_attention_type": "RocketKV", # Stage 1: Coarse Prompt Eviction "prompt_budget": 1024, # Maximum tokens retained after prefill "observation_window": 32, # Attention observation tail size "sink_token_size": 16, # Guaranteed attention sinks # Stage 2: Fine-Grain Dynamic Selection "top_k_tokens": 128, # Active tokens loaded per decode step "kernel_downsample": 4, # KT Cache sequence downsampling factor "enable_multiturn": True # Activate RocketKV-MT for chat dialogues } # Initialize ModelRunner with RocketKV configuration runner = ModelRunner.from_dir( engine_dir="./engine_rocketkv_llama_8b", rank=tensorrt_llm.mpi_rank(), sparse_attention_config=rocketkv_config ) # Example Long-Context Prompt (100k tokens) prompt_tokens = [1] + [42] * 100000 + [29901] input_tensor = torch.tensor([prompt_tokens], dtype=torch.int32, device="cuda") # Generate response with bounded VRAM and accelerated decoding outputs = runner.generate( batch_input_ids=input_tensor, max_new_tokens=256, temperature=0.7, top_p=0.9 ) print("Generated response tokens successfully using RocketKV!")

Empirical Benchmarks: Speedup, Memory Reduction, and Needle Retrieval

The researchers evaluated RocketKV against state-of-the-art baselines across standard long-context benchmarks, including LongBench, RULER, and synthetic Needle-In-A-Haystack (NIAH) on NVIDIA A100-80GB GPUs.

Performance Benchmark Comparison (128k Sequence Length, LLaMA-3.1-8B, NVIDIA A100): ┌────────────────────────┬──────────────────┬─────────────────┬──────────────────────┐ │ Method │ Decoding Latency │ Peak VRAM (GB) │ LongBench Avg Score │ ├────────────────────────┼──────────────────┼─────────────────┼──────────────────────┤ │ Full KV Baseline │ 42.8 ms / token │ 34.2 GB │ 48.6 │ │ StreamingLLM │ 18.2 ms / token │ 12.1 GB │ 28.4 (Severe loss!) │ │ H2O (Heavy Hitter) │ 19.5 ms / token │ 13.8 GB │ 41.2 │ │ Quest (Dynamic Top-K) │ 22.4 ms / token │ 34.2 GB (High!) │ 47.9 │ │ SnapKV │ 21.0 ms / token │ 14.2 GB │ 47.1 │ │ RocketKV (Ours) │ 11.5 ms / token │ 11.4 GB │ 48.3 (Near-Lossless!)│ └────────────────────────┴──────────────────┴─────────────────┴──────────────────────┘ Speedup: 3.7x faster than Full KV | VRAM Reduction: 32.6% peak drop | Compression: Up to 400x

End-to-End Decoding Speedup

As context length scales from 16k to 128k tokens, the full KV cache baseline suffers from severe latency degradation due to memory bandwidth choking. RocketKV maintains flat, sub-15ms per-token generation latency:

  • At 64k tokens: 2.8x speedup over dense attention.
  • At 128k tokens: 3.7x speedup over dense attention on an NVIDIA A100 GPU.

Peak Memory Reduction and Concurrency Scaling

Because Stage 1 permanently drops non-essential tokens from physical VRAM, RocketKV achieves up to a 32.6% reduction in total peak GPU memory consumption and up to a 400x compression ratio over the raw KV cache size. This frees up enough VRAM to increase serving batch concurrency from $B=1$ to $B=8$ on a single 80GB GPU.

Needle-In-A-Haystack (NIAH) and RULER Retrieval Accuracy

To verify whether permanent pruning impairs long-range fact retrieval, the authors tested RocketKV on NIAH tests with needles inserted across 0% to 100% depth:

  • RocketKV achieved 99.4% accuracy, rivaling the uncompressed oracle baseline.
  • On the multi-hop reasoning tasks in RULER (such as tracking variable assignments across 64k tokens), RocketKV outperformed pure SnapKV by over 6.2 points, proving that Stage 2 dynamic Top-K selection accurately recovers multi-hop relationships that static eviction discards.

Enterprise Strategy: Combining RocketKV with aiFetchly Data Pipelines

In high-throughput enterprise architectures, language models do not exist in isolation. They form the cognitive reasoning layer atop automated document extractors, web crawlers, and API integrations orchestrated through intelligent web scraping and unified API integrations.

┌────────────────────────────────────────────────────────────────────────────────────────┐ │ ENTERPRISE RAG ARCHITECTURE: AIFETCHLY + ROCKETKV SERVING │ ├────────────────────────────────────────────────────────────────────────────────────────┤ │ │ │ 1. Multi-Source Ingestion (<InternalLink feature="data-fetching" text="aiFetchly Data Fetching" />) │ │ ├── High-Volume Web Scraping (10,000+ Regulatory Documents & SEC Filings) │ │ └── Enterprise Database & CRM Streams │ │ │ │ 2. Long-Context Prompt Assembly │ │ └── Context Packing: Concatenate 100k tokens of financial filings into prompt │ │ │ │ 3. High-Throughput Serving Cluster │ │ └── NVIDIA TensorRT-LLM with RocketKV Engine │ │ ├── Stage 1: Prunes 100k prompt down to 1k active semantic clusters │ │ ├── Stage 2: Dynamic Top-128 HSA streams responses at 80 tokens/sec │ │ └── Concurrency: 8x higher query throughput on identical hardware │ │ │ │ 4. Production Economics (<InternalLink feature="pricing" text="aiFetchly Pricing" />) │ │ └── 65% Infrastructure Cost Reduction via Bounded GPU Memory Allocations │ │ │ └────────────────────────────────────────────────────────────────────────────────────────┘

When enterprise pipelines ingest multi-megabyte PDF manuals, legal contracts, or technical documentation, stuffing entire corpora into a long-context LLM prompt typically incurs prohibitive cost and latency. By pairing aiFetchly‘s structured data feeds with RocketKV-accelerated serving:

  1. Zero-Latency Ingestion: Ingest 100k+ tokens of raw document text directly without aggressive chunking.
  2. Sub-Second TTFT and Accelerated Streaming: RocketKV reduces prefill memory footprint and slashes decoding latency, ensuring users receive streaming answers in milliseconds.
  3. Optimized Infrastructure TCO: To explore how to scale your enterprise AI infrastructure while optimizing compute expenditure, review our enterprise pricing models and developer guides in our developer documentation.

Summary & Key Takeaways

RocketKV marks a pivotal milestone in the transition toward efficient, production-grade long-context intelligence:

  • Two-Stage Synergy: Decouples physical VRAM capacity reduction (Stage 1 coarse eviction via SnapKV) from decoding memory bandwidth acceleration (Stage 2 fine-grain Hybrid Sparse Attention).
  • Dual Dimensionality Reduction: The lightweight KT Cache compresses keys across both head and sequence dimensions, allowing instant Top-K prediction with less than 1% storage overhead.
  • Hardware-Proven Speedups: Delivers up to 3.7x end-to-end decoding speedup, 32.6% peak memory reduction, and 400x compression on NVIDIA A100 hardware.
  • Production-Ready Ecosystem: Natively integrated into NVIDIA TensorRT-LLM, enabling immediate drop-in adoption for LLaMA-3, Mistral, and enterprise foundation models.

By eliminating the quadratic memory wall of autoregressive decoding, RocketKV empowers developers to deploy massive-context applications with the throughput, latency, and economics required for modern enterprise scale.