[None][feat] Add TriAttention KV-cache compression method by Hudayday · Pull Request #16957 · NVIDIA/TensorRT-LLM

Hudayday · 2026-07-28T15:51:22Z

Summary

Add TriAttention NVIDIA work @ ICML 2026 KV-cache compression support to the TensorRT-LLM.
TriAttention is a training-free decode-time eviction algorithm: every beta confirmed generation tokens, it scores the evictable decode region using offline per-head query calibration, keeps budget important tokens, compacts their K/V data in place, and releases unused tail pages.
Integrate TriAttention as a KV-cache compression resource manager on top of KVCacheManagerV2; decode continues to use the model's standard attention backend.
Add the score/select pipeline: V2 page-table snapshot → mean-phase gather → CuTe paged-K scoring → normalize/reduce → TopK/tie settlement → batched physical compaction.
Support union, per_head, and per_layer_perhead eviction modes.
Support GPT-OSS-style kernel-masked SWA and one-model MTP/EAGLE3 draft-cache co-compaction in union mode.
Preserve logical RoPE positions after physical eviction by disabling RoPE fusion for eviction-based compression.
Add configuration/API wiring, usage documentation, B200 CI registration, and focused tests.

Performance

On Qwen3-8B, TriAttention (B=2048, β=2048) delivers up to 1.61× the throughput of the Dense baseline at the same batch size (BS=32: 4,930.6 vs. 3,057.5 output tok/s per GPU; 154.1 vs. 95.5 tok/s per user), while its verified capacity frontier reaches 2.63× the Dense baseline’s maximum aggregate throughput (8,053.0 vs. 3,057.5 output tok/s per GPU).

Main Files

File	Purpose
`_torch/kv_cache_compression/triattention/triattention.py`	Request lifecycle, calibration loading, eviction cadence, score/select/compact orchestration, and cache resize
`triattention_cute_score_fused.py`	SM100 CuTe DSL paged-K score kernel and score pipeline
`triattention_cute_selection.py`	CuTe DSL score normalization and union reduction
`triattention_kernels.py`	Triton mean-phase, per-head reduction, TP-fold, and tie-settlement kernels
`batch_manager/kvCacheManagerV2Utils.{h,cpp}`	Batched V2 page-table snapshot used by the scorer and compactor
`llmapi/llm_args.py`	`TriAttentionKvCacheCompressionConfig` and KV-compression config dispatch
`examples/triattention/README.md`	Calibration, configuration, support matrix, and usage

Key Design Decisions

Physical compression, not sparse attention: TriAttention changes which K/V entries are stored. It does not introduce a sparse-attention backend or modify the normal decode attention kernel.
In-place V2 compaction: The implementation does not create a new KV-cache instance or memory pool. Selected K/V entries are moved within the existing paged pools through the generic compaction API from [None][feat] Batched physical KV-cache compaction for KV cache compression #16836, then the same cache is resized.
Direct paged-pool scoring: The scorer consumes a batched snapshot of the current V2 page table and reads paged K directly; no Torch gather/reformat path is used.
Position correctness: Prompt tokens, selected-token RoPE information, and protected overlap/speculative tails are preserved. py_num_compressed_tokens maintains the logical sequence length after the physical cache becomes shorter.
Reusable runtime: Score buffers and compiled CuTe launches are retained and only rebuilt when a newly admitted request raises the scorer capacity high-water mark.

Known Limitations

PyTorch backend, KVCacheManagerV2, and B200/SM100 only.
KV-cache block reuse must be disabled.
Native sliding-eviction layouts such as Gemma 4 are not supported; the current SWA path targets GPT-OSS-style full-length pools with kernel-side masking.
Attention DP and disaggregated serving are not supported.
TP beyond TP1 has not yet been validated end to end.

Test Plan

Focused TriAttention/KV suite: 55 passed
- CuTe score vs. Torch oracle
- all three eviction modes
- ragged lengths and unaligned final tokens
- deterministic heavy-tie settlement
- target/draft physical co-compaction
- GPT-OSS-style SWA
- overlap/page-table staging and cache resize
- RoPE-fusion gating
- score scratch larger than 2^31 elements
Fused-union boundary file: 9/9 passed
Focused executor integration: 25 passed