diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp index 079de1a18893..e7e0946a2bf3 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp @@ -20,6 +20,7 @@ #include "tensorrt_llm/common/memoryUtils.h" #include #include +#include #include #include #include @@ -217,6 +218,40 @@ at::Tensor IndexMapper::getCopyIndex( return copyIndex_.slice(0, 0, numSeqs); } +void IndexMapper::gatherKBlockOffsets(at::Tensor const& source, at::Tensor destination, + std::vector const& requestIds, SizeType32 numBlocks) +{ + std::vector sourceRowsByRequest; + sourceRowsByRequest.reserve(requestIds.size()); + for (auto const requestId : requestIds) + { + sourceRowsByRequest.push_back(static_cast(getIndex(requestId)) * maxBeamWidth_); + } + + auto const* sourceData = source.data_ptr(); + auto* destinationData = destination.data_ptr(); + auto const sourceRows = source.size(1); + auto const sourcePlanes = source.size(2); + auto const sourceBlocks = source.size(3); + auto const destinationRows = destination.size(1); + auto const destinationPlanes = destination.size(2); + auto const destinationBlocks = destination.size(3); + auto const copyBytes = static_cast(numBlocks) * sizeof(int32_t); + + for (int64_t pool = 0; pool < source.size(0); ++pool) + { + for (size_t destinationRow = 0; destinationRow < sourceRowsByRequest.size(); ++destinationRow) + { + auto const sourceRow = sourceRowsByRequest[destinationRow]; + auto const sourceOffset = ((pool * sourceRows + sourceRow) * sourcePlanes) * sourceBlocks; + auto const destinationOffset + = ((pool * destinationRows + static_cast(destinationRow)) * destinationPlanes) + * destinationBlocks; + std::memcpy(destinationData + destinationOffset, sourceData + sourceOffset, copyBytes); + } + } +} + IndexMapper::IndexMapper(SizeType32 maxBatchSize, SizeType32 maxBeamWidth) : maxBeamWidth_(maxBeamWidth) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h index 5fa88a78468e..efa009ee4f8b 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -72,6 +72,10 @@ class IndexMapper at::Tensor getCopyIndex( std::vector const& requestIds, SizeType32 numContext, SizeType32 beamWidth); + //! Gathers each request's beam-0 K block offsets into a host snapshot. + void gatherKBlockOffsets(at::Tensor const& source, at::Tensor destination, + std::vector const& requestIds, SizeType32 numBlocks); + /// Number of sequences currently tracked (i.e. active IndexMapper slots). [[nodiscard]] SizeType32 size() const noexcept { diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp index 3e549ec6b7bb..83e9fd8053e8 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2Utils.cpp @@ -78,6 +78,8 @@ void KVCacheManagerV2UtilsBindings::initBindings(nb::module_& module) .def("get_index", &IndexMapper::getIndex) .def("remove_sequence", &IndexMapper::removeSequence) .def("get_copy_index", &IndexMapper::getCopyIndex) + .def("gather_k_block_offsets", &IndexMapper::gatherKBlockOffsets, nb::arg("source"), nb::arg("destination"), + nb::arg("request_ids"), nb::arg("num_blocks")) .def("size", &IndexMapper::size) .def("num_free_slots", &IndexMapper::numFreeSlots); diff --git a/examples/kv_cache_compression/triattention.md b/examples/kv_cache_compression/triattention.md new file mode 100644 index 000000000000..1a805d723c93 --- /dev/null +++ b/examples/kv_cache_compression/triattention.md @@ -0,0 +1,125 @@ +# TriAttention KV-Cache Compression + +This document describes enabling TriAttention KV-cache compression in TensorRT-LLM. + +TriAttention is a training-free, decode-time KV-cache eviction method for long-context LLM inference. During generation it periodically scores the cached tokens by a trigonometric importance measure derived from offline per-head query statistics (calibration), keeps the most important `budget` tokens, and physically compacts the cache — reducing KV-cache memory so more sequences fit on a GPU at once. + +For technical details see the paper [TriAttention](https://arxiv.org/abs/2604.04921) and the official implementation [github.com/WeianMao/triattention](https://github.com/WeianMao/triattention). + +## Overview + +TriAttention runs entirely in the generation phase and reuses the standard dense attention kernel over the compacted cache: + +1. **Calibration (offline, one-time per model).** The importance score needs each attention head's mean and magnitude of the pre-RoPE query, gathered over a small calibration corpus. **TensorRT-LLM does not compute calibration** — you produce it once with the official tool and pass the resulting `.pt` file. TensorRT-LLM loads and converts it when the compression manager is created. +2. **Periodic eviction (during generation).** Every `beta` confirmed generation tokens, once a sequence is over budget, TriAttention scores the evictable decode region, selects `budget` decode tokens to keep, preserves the prompt, and physically compacts the KV cache down to that set. A speculative iteration may confirm multiple tokens; crossing multiple periods in one update is coalesced into one eviction. + +TriAttention is integrated into TensorRT-LLM as a KV-cache compression manager on top of the `KVCacheManagerV2`. Scoring runs on CuTe DSL (SM100) and Triton kernels; compaction is a native CUDA kernel. + +## Support Matrix + +* NVIDIA B200 (SM100; the current validated target) +* Paged KV Cache (`KVCacheManagerV2`) +* PyTorch backend + +**Notes:** +1. TriAttention supports KV-cache block reuse. V2 reuses the committed prompt prefix, while TriAttention preserves that prefix and compacts only the generation suffix. +2. TriAttention requires the V2 KV-cache manager (`use_kv_cache_manager_v2=True`). +3. TriAttention does not compute calibration. Bring the official tool's calibration `.pt`; see [Calibration](#calibration). +4. The current SWA path covers models such as GPT-OSS whose V2 pools remain full length and whose attention kernel applies the window. Native sliding-eviction layouts such as Gemma 4, SSM/hybrid pools, and MLA caches are not supported. +5. Speculative decoding is supported for one-model MTP and EAGLE3 with `eviction_mode="union"`. Tensor parallelism beyond TP1, attention DP, and disaggregated serving have not yet been validated end to end. + +## Calibration + +The calibration file is produced once per model with the official tool, then reused for every inference run with that model. + +Generate the calibration file for your model with the official repository (for +example `qwen3-8b-calibration.pt` for Qwen3-8B), keep it anywhere on disk, and +point `calibration_path` at it: + +```bash +# Clone + install the official tool +git clone https://github.com/WeianMao/triattention.git +cd triattention && pip install -e . + +# Calibrate (writes the official {metadata, stats} .pt) +python3 scripts/calibrate.py \ + --model \ + --input data/calibration_text.txt \ + --output _calibration.pt \ + --max-length 32768 \ + --device cuda +``` + +TensorRT-LLM accepts that file directly: it reads the official `{metadata, stats}` layout and derives the model's RoPE tables from the model config, then converts everything to its runtime schema at load. (An already-converted flat `.pt` is also accepted.) + +## Usage + +To enable TriAttention, pass a `TriAttentionKvCacheCompressionConfig` (the eviction knobs + the calibration file) to the `LLM` constructor. TriAttention is a pure compression method — there is **no** sparse-attention config and no custom attention backend; decode runs the model's standard attention over the compacted cache. + +### Python API + +```python +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import (KvCacheConfig, + TriAttentionKvCacheCompressionConfig) + +# 1. Configure the eviction manager + point it at the calibration file. +compression_config = TriAttentionKvCacheCompressionConfig( + budget=2048, # tokens kept at each eviction (prompt is kept on top) + beta=64, # eviction period, in confirmed generation tokens + eviction_mode="union", + calibration_path="/path/to/qwen3-8b-calibration.pt", # official tool's output + model_path="", # used to derive the RoPE tables +) + +# 2. TriAttention needs the V2 KV-cache manager and supports block reuse. +kv_config = KvCacheConfig(enable_block_reuse=True, use_kv_cache_manager_v2=True) + +llm = LLM( + model="", + backend="pytorch", + kv_cache_compression_config=compression_config, + kv_cache_config=kv_config, +) + +# 3. Generate +prompts = ["To be or not to be, that is the question."] +sampling_params = SamplingParams(max_tokens=128) +outputs = llm.generate(prompts, sampling_params) +``` + +### Usage with `trtllm-bench` and `trtllm-serve` + +Pass the configs via `--config config.yaml`. The field names match the Python configs: + +```yaml +backend: pytorch +kv_cache_compression_config: + algorithm: triattention + budget: 2048 + beta: 64 + eviction_mode: union + calibration_path: /path/to/qwen3-8b-calibration.pt + model_path: +kv_cache_config: + enable_block_reuse: true + use_kv_cache_manager_v2: true +``` + +```bash +trtllm-eval --model --config config.yaml longbench_v2 --max_output_length 1024 ... +``` + +## Configuration Arguments + +`TriAttentionKvCacheCompressionConfig` controls the compression ratio and the eviction algorithm: + +* **`budget`** (int, default=2048): Tokens kept at each eviction. Prompt tokens are always preserved on top of this. Smaller `budget` → more compression. +* **`beta`** (int, default=128): Eviction period, in confirmed generation tokens (the upstream `divide_length`). Speculative acceptance advances the counter by `1 + accepted_draft_tokens`; at most one eviction is coalesced per final update. +* **`eviction_mode`** (str, default=`union`): Which token set each eviction keeps. + * `union`: union of each KV head's top-B, re-ranked by the per-token max score. Matches the official base setting. + * `per_head`: each KV head keeps its own set, shared across layers (mean of per-layer maxima). + * `per_layer_perhead`: each head keeps its own set, fully independent per layer. +* **`normalize_scores`** (bool, default=True): Z-normalize each head's scores over the decode region before selection (upstream default). `union` eviction always z-normalizes: `False` is overridden to `True` with a warning. +* **`calibration_path`** (str): Path to the calibration `.pt` from the official tool. Required — TensorRT-LLM does not compute calibration. +* **`model_path`** (str): Checkpoint path, used to derive the model's RoPE tables when converting the official calibration file and to classify kernel-masked sliding-window (SWA) layers from the model config. diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py deleted file mode 100644 index cd4e7bf32068..000000000000 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ /dev/null @@ -1,29 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from enum import IntEnum, auto -from typing import Optional - - -class KvCacheCompressionMode(IntEnum): - """Algorithm-level traits of a KV-cache compression method. - - Configs map their ``algorithm`` string to a member here; callers read the - ``is_*`` predicates instead of comparing strings. - """ - - NONE = auto() - - def is_eviction_method(self): - """Whether this method physically evicts cached tokens. Evicting - algorithms add their member and extend this predicate.""" - return False - - @staticmethod - def from_string(name: Optional[str]) -> "KvCacheCompressionMode": - if name is None: - return KvCacheCompressionMode.NONE - try: - return KvCacheCompressionMode[name.upper()] - except KeyError: - return KvCacheCompressionMode.NONE diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py new file mode 100644 index 000000000000..cc43f71441e9 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -0,0 +1,960 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""TriAttention KV-cache compression: periodic physical KV eviction during generation. + +Every ``beta`` confirmed tokens, cached tokens are scored with a trigonometric +importance score from offline calibration and tokens outside the top-``budget`` +keep set are physically deleted; decode runs the model's standard attention over +the compacted cache. Kept keys keep their original RoPE rotation (no re-RoPE). +KV pools must be read with ``kv_layout="HND"``. Calibration comes from the +official tool (github.com/WeianMao/triattention) and is converted at load. +""" + +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple + +import torch +import triton +from transformers import AutoConfig +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets_to_device, +) +from tensorrt_llm.logger import logger + +from ...distributed import allgather +from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from ...pyexecutor.llm_request import LlmRequestState +from ...pyexecutor.resource_manager import KVCacheCompressionManager +from ...utils import next_positive_power_of_2 +from ..compaction import build_compaction_params, compact +from .triattention_cute_score_fused import PADDED_HEAD_COLUMNS, build_score_pipeline +from .triattention_kernels import ( + fold_union_ranks, + gather_mean_phase, + reduce_per_head_scores, + settle_ties, +) + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig + + from ...pyexecutor.llm_request import LlmRequest + from ...pyexecutor.scheduler import ScheduledRequests + + +# Required keys for the calibration ``.pt`` consumed by TriAttention. +_REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) + +_MEAN_PHASE_OFFSETS = tuple(float(1 << exponent) for exponent in range(17)) + +# Physical TopK rows follow the 256-token reduce/tie kernel tiles. +_SELECTION_WIDTH_ALIGNMENT = 256 + + +class _EvictionRequest(NamedTuple): + """One due request and the cache state needed by its eviction round.""" + + request: "LlmRequest" + target_cache: object + draft_cache: Optional[object] + source_length: int + target_tail_length: int + + +_BLOCK_OFFSET_ALIGNMENT = 4 + + +def _allocate_block_offset_snapshot( + manager: KVCacheManagerV2, + anchor_pool: torch.Tensor, + *, + request_capacity: int, + token_capacity: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Allocate the bounded V2 page-table snapshot used by an eviction round.""" + required_blocks = triton.cdiv(token_capacity, int(manager.tokens_per_block)) + staged_blocks = min( + triton.cdiv(required_blocks, _BLOCK_OFFSET_ALIGNMENT) * _BLOCK_OFFSET_ALIGNMENT, + int(manager.max_blocks_per_seq), + ) + snapshot_shape = (int(manager.num_pools), request_capacity, 2, staged_blocks) + block_offsets_host = torch.empty( + snapshot_shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + ) + block_offsets_device = torch.empty(snapshot_shape, dtype=torch.int32, device=anchor_pool.device) + return block_offsets_host, block_offsets_device + + +_MEAN_PHASE_MAX_ROWS = 1 << 24 + + +class _MeanPhaseTable: + """Admission-sized mean-phase lookup used by every eviction round.""" + + def __init__(self, omega: torch.Tensor, device: torch.device) -> None: + self._omega = omega.to(device=device, dtype=torch.float32).contiguous() + self.cos: Optional[torch.Tensor] = None + self.sin: Optional[torch.Tensor] = None + self.rows = 0 + self.num_freqs = int(self._omega.numel()) + + def reserve(self, rows: int) -> None: + """Cover positions ``[0, rows)`` with a power-of-two table.""" + rows = int(rows) + if rows <= self.rows: + return + if rows > _MEAN_PHASE_MAX_ROWS: + raise ValueError(f"a {rows}-row mean-phase table exceeds the exact-FP32 position range") + target = next_positive_power_of_2(rows) + target = min(max(target, 2 * self.rows), _MEAN_PHASE_MAX_ROWS) + positions = torch.arange(target, device=self._omega.device, dtype=torch.float32) + cos_table = torch.zeros( + (target, self.num_freqs), + dtype=torch.float32, + device=self._omega.device, + ) + sin_table = torch.zeros_like(cos_table) + # Fixed summation order keeps table rebuilds bit-stable. + for offset in _MEAN_PHASE_OFFSETS: + angle = torch.outer(positions + offset, self._omega) + cos_table += torch.cos(angle) + sin_table += torch.sin(angle) + scale = 1.0 / len(_MEAN_PHASE_OFFSETS) + self.cos = cos_table.mul_(scale) + self.sin = sin_table.mul_(scale) + self.rows = target + + +class TriAttentionCompressionManager(KVCacheCompressionManager): + """KV-cache compression manager for periodic TriAttention eviction.""" + + # ---- construction ---- + + def __init__( + self, + config: "TriAttentionKvCacheCompressionConfig", + kv_cache_manager: KVCacheManagerV2, + draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, + ) -> None: + super().__init__(config, kv_cache_manager, draft_kv_cache_manager) + self.budget = config.budget + self.beta = config.beta + self.eviction_mode = config.eviction_mode + if self.eviction_mode == "union" and not config.normalize_scores: + logger.warning("TriAttention union mode enables score normalization") + self.normalize_scores = self.eviction_mode == "union" or config.normalize_scores + # Prompt always pinned; budget counts decode tokens only. + self.model_path = config.model_path + self.calibration_path = config.calibration_path + self._load_calibration() + + self._prepared_generation_batch: Optional["ScheduledRequests"] = None + # Manager-lifetime constants. + self._num_extra_kv_tokens = int(kv_cache_manager.num_extra_kv_tokens) + self._protected_tail_capacity = ( + self._num_extra_kv_tokens + int(kv_cache_manager._kv_reserve_draft_tokens) + 1 + ) + self._draft_protected_tail_capacity = 0 + if draft_kv_cache_manager is not None: + self._draft_protected_tail_capacity = ( + int(draft_kv_cache_manager.num_extra_kv_tokens) + + int(draft_kv_cache_manager._kv_reserve_draft_tokens) + + 1 + ) + # The next-step reservation size is fixed; overlap only changes which + # requests have it. + self._overlap_tail_length = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) + # Fixed buffer geometry. These are TriAttention scratch dimensions, + # not KV capacities owned by KVCacheManagerV2. + self._request_capacity = int(kv_cache_manager.max_batch_size) + max_draft_tokens = int(kv_cache_manager.max_total_draft_tokens) + # Crossing a cadence can overshoot by D accepted draft tokens; one + # suspended due round may resume with another 1 + D confirmed tokens. + required_selection_width = self.budget + self.beta + 2 * max_draft_tokens + 1 + self._selection_width_capacity = ( + triton.cdiv(required_selection_width, _SELECTION_WIDTH_ALIGNMENT) + * _SELECTION_WIDTH_ALIGNMENT + ) + max_tail_capacity = max( + self._protected_tail_capacity, + self._draft_protected_tail_capacity, + ) + if self._request_capacity * (self.budget + max_tail_capacity) >= 2**31: + raise ValueError("TriAttention compaction offsets exceed the int32 range") + # Manager-lifetime layer facts, resolved once: V2 fixes pp_layers at + # construction and the model config is immutable on disk. + self._global_layers = [int(layer) for layer in kv_cache_manager.pp_layers] + ( + self._dense_layers, + self._swa_layers, + self._swa_window, + ) = self._resolve_attention_layers() + self._initialize_eviction_state() + + def _resolve_attention_layers(self) -> Tuple[List[int], List[int], Optional[int]]: + """SWA layers here are stored at full length; the window applies only in the kernel.""" + model_path = self.model_path + global_layers = self._global_layers + num_layers = len(global_layers) + + config = AutoConfig.from_pretrained( + model_path, trust_remote_code=True, local_files_only=True + ) + config_values = config.get_text_config().to_dict() + layer_types = config_values.get("layer_types") + if not layer_types: + use_sliding_window = config_values.get("use_sliding_window") + has_swa_signal = ( + use_sliding_window + if isinstance(use_sliding_window, bool) + else any( + config_values.get(field) + for field in ( + "sliding_window", + "sliding_window_size", + "sliding_window_pattern", + "max_window_layers", + ) + ) + ) + if has_swa_signal: + raise ValueError( + "Model config exposes sliding-window metadata but no layer_types; " + "TriAttention cannot classify kernel-masked SWA layers safely" + ) + return (list(range(num_layers)), [], None) + if global_layers and max(global_layers) >= len(layer_types): + raise ValueError( + f"Model config has {len(layer_types)} layer_types entries, " + f"but this PP rank references global layer {max(global_layers)}" + ) + + swa_layers = [ + local_layer + for local_layer, global_layer in enumerate(global_layers) + if "sliding" in str(layer_types[global_layer]).lower() + ] + swa_set = set(swa_layers) + dense_layers = [layer for layer in range(num_layers) if layer not in swa_set] + # GPT-OSS SWA keeps full-length V2 pools and masks in the kernel; native + # sliding-eviction layouts such as Gemma 4 remain unsupported. + if not dense_layers: + raise ValueError("TriAttention requires at least one full-attention layer") + window_size = None + if swa_layers: + raw_window = config_values.get("sliding_window") + if not isinstance(raw_window, int) or raw_window <= 0: + raise ValueError( + "TriAttention requires a positive integer model sliding_window " + "when layer_types contains sliding attention" + ) + if self.budget < raw_window: + raise ValueError( + f"TriAttention budget={self.budget} must be at least " + f"the kernel-masked SWA window size {raw_window}" + ) + window_size = raw_window + return (dense_layers, swa_layers, window_size) + + def _load_calibration(self) -> None: + raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) + if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): + e_q = raw["E_q"] + e_q_norm = raw["E_q_norm"] + omega = raw["omega"] + freq_scale_sq = raw["freq_scale_sq"] + elif isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): + stats = raw["stats"] + metadata = raw["metadata"] + if "sampled_heads" in metadata: + heads = [(int(layer), int(head)) for layer, head in metadata["sampled_heads"]] + else: + heads = [ + ( + int(key[len("layer") : key.index("_head")]), + int(key[key.index("_head") + len("_head") :]), + ) + for key in stats + ] + num_layers = max(layer for layer, _ in heads) + 1 + num_heads = max(head for _, head in heads) + 1 + freq_count = int(next(iter(stats.values()))["q_mean_real"].numel()) + e_q = torch.zeros(num_layers, num_heads, freq_count, dtype=torch.complex64) + e_q_norm = torch.zeros(num_layers, num_heads, freq_count, dtype=torch.float32) + for layer, head in heads: + head_stats = stats[f"layer{layer:02d}_head{head:02d}"] + e_q[layer, head] = torch.complex( + head_stats["q_mean_real"].float(), + head_stats["q_mean_imag"].float(), + ) + e_q_norm[layer, head] = head_stats["q_abs_mean"].float() + + config = AutoConfig.from_pretrained( + self.model_path, trust_remote_code=True + ).get_text_config() + # transformers >= 5.5 folds rope_theta/rope_type into rope_parameters. + rope_params = config.to_dict()["rope_parameters"] + if all(isinstance(value, dict) for value in rope_params.values()): + raise ValueError( + "TriAttention does not support per-layer-type rope parameters " + f"({self.model_path})" + ) + rope_type = rope_params["rope_type"] + if rope_type == "default": + # "default" has no ROPE_INIT_FUNCTIONS entry. + head_dim = freq_count * 2 + base = float(rope_params["rope_theta"]) + positions = torch.arange(0, head_dim, 2, dtype=torch.float32) + omega = (1.0 / (base ** (positions / head_dim)))[:freq_count].clone() + attention_scale_sq = 1.0 + else: + inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](config, device="cpu") + omega = inv_freq.to(torch.float32)[:freq_count].clone() + attention_scale_sq = float(attention_factor) ** 2 + freq_scale_sq = torch.full((freq_count,), attention_scale_sq, dtype=torch.float32) + logger.info( + f"TriAttention: converted official calibration {self.calibration_path}" + f" -> E_q[L={num_layers}, H={num_heads}, F={freq_count}]" + ) + else: + got = sorted(raw) if isinstance(raw, dict) else type(raw).__name__ + raise ValueError( + f"Unrecognized calibration at {self.calibration_path}: expected the " + f"official {{metadata, stats}} layout or " + f"{sorted(_REQUIRED_CALIBRATION_KEYS)}; got {got}." + ) + + self._freq_scale_sq = freq_scale_sq.to(dtype=torch.float32) + self._omega = omega + # Pre-split query stats + MLR coefficient, shapes [L, H, F]. + self._calibration_q_real = e_q.real.to(torch.float32).contiguous() + self._calibration_q_imag = e_q.imag.to(torch.float32).contiguous() + self._calibration_mlr_coef = ( + e_q_norm.to(torch.float32) - e_q.abs().to(torch.float32) + ).contiguous() + + # ---- framework hooks (call order) ---- + + def on_request_init(self, request: "LlmRequest", **kwargs) -> None: + """Grow scorer state only when this request raises its capacity high-water mark.""" + manager = self.kv_cache_manager + prompt_length = int(request.py_prompt_len) + max_decode_tokens = min( + int(request.py_max_new_tokens), + max(int(manager.max_seq_len) - prompt_length, 0), + ) + first_evict_step = (self.budget // self.beta + 1) * self.beta + if max_decode_tokens < first_evict_step: + return + self._phase.reserve(prompt_length + max_decode_tokens + 1) + max_source_tokens = prompt_length + min(max_decode_tokens, self._selection_width_capacity) + if max_source_tokens <= self._score_token_capacity: + return + + # CuTe launches capture score buffers; retire the old capacity before replacing it. + if self._launch_score is not None: + self._compaction_done_event.synchronize() + + new_score_capacity = next_positive_power_of_2(max(max_source_tokens, 1024)) + new_score_capacity = min(new_score_capacity, int(manager.max_seq_len)) + score_tile_size = max(64, int(manager.tokens_per_block)) + new_score_capacity = triton.cdiv(new_score_capacity, score_tile_size) * score_tile_size + self._build_score_runtime(score_token_capacity=new_score_capacity) + + def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: + """Remember the next batch: overlap prepares it before updating the previous batch.""" + self._prepared_generation_batch = scheduled_batch + + def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: + """Compact after native KV-cache updates finalize the iteration. + + KVCacheManagerV2 must run first so capacity includes the written token and any rewind. + """ + self._evict_due_requests(scheduled_batch) + + # ---- eviction round ---- + + def _evict_due_requests( + self, + scheduled_batch: "ScheduledRequests", + ) -> None: + """Collect due requests, execute one eviction round, publish, and resize.""" + manager = self.kv_cache_manager + eviction_requests: List[_EvictionRequest] = [] + # With overlap, the next batch reserves KV before the previous batch is + # compacted. Those reserved slots are a byte-preserved tail, not score input. + prepared_batch = self._prepared_generation_batch + overlap_request_ids = ( + {request.py_request_id for request in prepared_batch.generation_requests} + if prepared_batch is not None and prepared_batch is not scheduled_batch + else set() + ) + for request in scheduled_batch.generation_requests: + if request.is_dummy or request.state in ( + LlmRequestState.GENERATION_COMPLETE, + LlmRequestState.CONTEXT_INIT, + ): + continue + request_id = request.py_request_id + target_cache = manager.kv_cache_map.get(request_id) + if target_cache is None or not target_cache.is_active: + # Overlap scheduling may suspend a cache mid-flight; defer + # this request (pre-launch) instead of failing the batch. + continue + draft_cache = None + if self.draft_kv_cache_manager is not None: + # A missing draft cache is a wiring bug: keep the precise KeyError. + draft_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] + if not draft_cache.is_active: + continue + target_tail_length = self._num_extra_kv_tokens + ( + self._overlap_tail_length if request_id in overlap_request_ids else 0 + ) + source_length = int(target_cache.capacity) - target_tail_length + if source_length < target_cache.history_length: + raise RuntimeError( + f"Request {request_id} KV length {source_length} is below " + f"finalized history {target_cache.history_length}" + ) + prompt_length = int(request.py_prompt_len) + # Restore the logical length from the physical cache and the + # eviction count already published to the model runtime. + compressed_tokens = int(request.py_num_compressed_tokens) + logical_source_length = source_length + compressed_tokens + confirmed_tokens = logical_source_length - prompt_length + # The last compact ended at budget + compressed_tokens. This + # watermark catches a beta boundary deferred by cache suspension. + if (self.budget + compressed_tokens) // self.beta >= (confirmed_tokens // self.beta): + continue + if source_length <= prompt_length + self.budget: + # Selection would be an identity: nothing to evict yet. + continue + decode_width = source_length - prompt_length + if decode_width > self._selection_width_capacity: + raise RuntimeError( + f"Request {request_id} TriAttention selection width " + f"{decode_width} exceeds compiled capacity " + f"{self._selection_width_capacity}" + ) + eviction_requests.append( + _EvictionRequest( + request=request, + target_cache=target_cache, + draft_cache=draft_cache, + source_length=source_length, + target_tail_length=target_tail_length, + ) + ) + if not eviction_requests: + return + + self._execute_eviction_round(eviction_requests) + for item in eviction_requests: + evicted = item.source_length - int(item.request.py_prompt_len) - self.budget + # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). + item.request.py_num_compressed_tokens += evicted + self._resize_compacted_caches(eviction_requests) + + def _execute_eviction_round( + self, + eviction_requests: Sequence[_EvictionRequest], + ) -> None: + """Score, select, and compact one due request group.""" + manager = self.kv_cache_manager + draft_manager = self.draft_kv_cache_manager + stream = torch.cuda.current_stream(self._block_offsets_device.device) + # PyExecutor already joins its execution stream before the final + # compression resource update, so the round can use caller current. + try: + request_ids = [item.request.py_request_id for item in eviction_requests] + logical_source_lengths = [ + item.source_length + int(item.request.py_num_compressed_tokens) + for item in eviction_requests + ] + prompt_lengths = [int(item.request.py_prompt_len) for item in eviction_requests] + source_lengths = [item.source_length for item in eviction_requests] + dense_move_offsets, swa_move_offsets, draft_move_offsets = ( + self._compute_compaction_move_offsets(eviction_requests) + ) + metadata_rows = ( + logical_source_lengths, + source_lengths, + prompt_lengths, + dense_move_offsets, + swa_move_offsets, + draft_move_offsets, + ) + # CPU may rewrite pinned staging only after its prior H2D completes. + self._staging_reuse_event.synchronize() + host_table = self._request_metadata_host_np + for row, values in enumerate(metadata_rows): + if values is not None: + host_table[row, : len(values)] = values + # Native compaction keeps fixed-capacity metadata views; make + # their unused request rows explicit no-ops. + host_table[:3, len(eviction_requests) :] = 0 + try: + self._stage_block_offset_snapshot( + manager, + request_ids, + self._block_offsets_host, + self._block_offsets_device, + ) + if draft_manager is not None: + self._stage_block_offset_snapshot( + draft_manager, + request_ids, + self._draft_block_offsets_host, + self._draft_block_offsets_device, + ) + self._request_metadata_device.copy_(self._request_metadata_host, non_blocking=True) + finally: + self._staging_reuse_event.record(stream) + + request_count = len(eviction_requests) + union = self.eviction_mode == "union" + # In-place refresh: the compiled score launches captured these pointers. + gather_mean_phase( + self._logical_source_lengths_device, + self._phase.cos, + self._phase.sin, + self._source_lengths_device, + self._prompt_lengths_device, + self._mean_cos, + self._mean_sin, + self._decode_lengths_device, + self._swa_destination_bases, + request_count=request_count, + swa_rebase_delta=self._swa_rebase_delta, + ) + self._launch_score(request_count) + if union and self._union_tp_mapping is not None: + # Max is order-free, so every TP rank keeps the same ordinals. + gathered = allgather( + self._selection_scores_rows[:request_count], + self._union_tp_mapping, + dim=0, + ) + fold_union_ranks( + gathered, + self._selection_scores_rows, + request_count=request_count, + ) + if not union: + reduce_per_head_scores( + self._score_scratch, + self._decode_lengths_device, + self._prompt_lengths_device, + self._row_mean, + self._row_inv_std, + self._selection_scores_rows, + self._selection_row_lengths, + request_count=request_count, + padded_head_columns=PADDED_HEAD_COLUMNS, + score_token_capacity=self._score_token_capacity, + per_layer=self.eviction_mode == "per_layer_perhead", + normalize_scores=self.normalize_scores, + ) + self._select_kept_ordinals(request_count) + compact(self._compaction_params, request_count) + finally: + # Target and draft V2 managers share this execution stream. + self._compaction_done_event.record(stream) + if manager._stream != stream: + manager._stream.wait_event(self._compaction_done_event) + + def _compute_compaction_move_offsets( + self, + eviction_requests: Sequence[_EvictionRequest], + ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: + """Build padded cumulative dense, SWA, and draft move offsets.""" + + def cumulative_offsets(move_counts: List[int]) -> List[int]: + offsets = [0] + for count in move_counts: + offsets.append(offsets[-1] + count) + offsets.extend(offsets[-1:] * (self._request_capacity - len(move_counts))) + return offsets + + tails = [int(item.target_tail_length) for item in eviction_requests] + dense_offsets = cumulative_offsets([self.budget + tail for tail in tails]) + swa_offsets = None + if self._swa_window is not None: + swa_offsets = cumulative_offsets([self._swa_window + tail for tail in tails]) + draft_offsets = None + if self.draft_kv_cache_manager is not None: + draft_offsets = cumulative_offsets( + [self.budget + self._draft_protected_tail_capacity] * len(eviction_requests) + ) + return dense_offsets, swa_offsets, draft_offsets + + def _stage_block_offset_snapshot( + self, + manager: KVCacheManagerV2, + request_ids: List[int], + host_block_offsets: torch.Tensor, + device_block_offsets: torch.Tensor, + ) -> None: + """Snapshot host block offsets before their asynchronous device copy.""" + manager.index_mapper.gather_k_block_offsets( + manager.host_kv_cache_block_offsets, + host_block_offsets, + request_ids, + host_block_offsets.shape[-1], + ) + copy_batch_block_offsets_to_device( + host_block_offsets, + device_block_offsets, + self._identity_copy_indices_host[: len(request_ids)], + manager.index_scales, + manager.kv_offset, + torch.cuda.current_stream(device_block_offsets.device).cuda_stream, + ) + + def _select_kept_ordinals(self, request_count: int) -> None: + """Select top-k tokens and settle score ties into kept-ordinal rows.""" + rows = request_count * self._selection_rows_per_request + # The trailing 1 is next_n: decode scores one query token per request. + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + self._selection_scores_rows[:rows], + self._selection_row_lengths[:rows], + self._provisional_rows[:rows], + self.budget, + 1, + ) + settle_ties( + self._selection_scores_rows, + self._selection_row_lengths, + self._prompt_lengths_device, + self._provisional_rows, + self._kept_ordinal_rows, + request_count=request_count, + selection_rows_per_request=self._selection_rows_per_request, + ) + + def _resize_compacted_caches(self, eviction_requests: Sequence[_EvictionRequest]) -> None: + for item in eviction_requests: + target_capacity = ( + int(item.request.py_prompt_len) + self.budget + item.target_tail_length + ) + if not item.target_cache.resize(target_capacity, None): + raise RuntimeError( + "Failed to resize compacted target KV cache for " + f"request {item.request.py_request_id} to " + f"{target_capacity} tokens" + ) + if self.draft_kv_cache_manager is None: + return + for item in eviction_requests: + draft_capacity = ( + int(item.request.py_prompt_len) + self.budget + self._draft_protected_tail_capacity + ) + if not item.draft_cache.resize(draft_capacity, None): + raise RuntimeError( + "Failed to resize compacted draft KV cache for " + f"request {item.request.py_request_id} to " + f"{draft_capacity} tokens" + ) + + # ---- persistent state + score runtime ---- + + def _initialize_eviction_state(self) -> None: + """Create manager-lifetime state once.""" + target_layout = self._create_kv_layout() + draft_layout = ( + self._create_kv_layout(draft=True) if self.draft_kv_cache_manager is not None else None + ) + self._target_layout = target_layout + self._draft_layout = draft_layout + + layer_pools = target_layout["layer_pools"] + dense_layers = target_layout["dense_layers"] + anchor_pool = layer_pools[dense_layers[0]] + device = anchor_pool.device + _, _, num_kv_heads, tokens_per_block, _ = anchor_pool.shape + + self._block_offsets_host = None + self._block_offsets_device = None + self._draft_block_offsets_host = None + self._draft_block_offsets_device = None + + global_layers = self._global_layers + if global_layers and max(global_layers) >= self._calibration_q_real.shape[0]: + raise ValueError( + f"TriAttention calibration has {self._calibration_q_real.shape[0]} layers, " + f"but this PP rank references global layer {max(global_layers)}" + ) + layer_ids = torch.as_tensor( + global_layers, + device=self._calibration_q_real.device, + dtype=torch.long, + ) + q_real = self._calibration_q_real.index_select(0, layer_ids) + q_imag = self._calibration_q_imag.index_select(0, layer_ids) + mlr_coef = self._calibration_mlr_coef.index_select(0, layer_ids) + mapping = self.kv_cache_manager.mapping + tp_size = 1 if mapping.enable_attention_dp else int(mapping.tp_size) + if tp_size > 1: + local_q_heads = int(q_real.shape[1]) // tp_size + heads = slice(mapping.tp_rank * local_q_heads, (mapping.tp_rank + 1) * local_q_heads) + q_real, q_imag, mlr_coef = q_real[:, heads], q_imag[:, heads], mlr_coef[:, heads] + q_real, q_imag, mlr_coef, self._freq_scale_sq = ( + tensor.to(device=device, dtype=torch.float32).contiguous() + for tensor in (q_real, q_imag, mlr_coef, self._freq_scale_sq) + ) + self._union_tp_mapping = ( + mapping if (self.eviction_mode == "union" and tp_size > 1) else None + ) + num_q_heads = int(q_real.shape[1]) + num_freqs = int(q_real.shape[2]) + self._score_q_real = q_real + self._score_q_imag = q_imag + self._score_mlr_coef = mlr_coef + + self._phase = _MeanPhaseTable(self._omega, device) + self._num_layers = len(dense_layers) + self._num_q_heads = num_q_heads + self._num_kv_heads = int(num_kv_heads) + self._allocate_metadata_buffers( + device, + num_freqs=num_freqs, + ) + self._allocate_selection_buffers(device, tp_size=tp_size) + + self._compaction_params = () + self._score_scratch = None + self._score_token_capacity = 0 + self._launch_score = None + + self._staging_reuse_event = torch.cuda.Event() + self._staging_reuse_event.record(torch.cuda.current_stream(device)) + self._compaction_done_event = torch.cuda.Event() + self._compaction_done_event.record(torch.cuda.current_stream(device)) + + logger.info( + f"TriAttention CuTe score configured: {self._num_q_heads}q/" + f"{self._num_kv_heads}kv heads, {num_freqs} freqs, " + f"{int(tokens_per_block)}-token pages" + ) + + def _allocate_metadata_buffers( + self, + device: torch.device, + *, + num_freqs: int, + ) -> None: + """Allocate fixed manager-lifetime host staging and device metadata.""" + row_count = 6 + request_capacity = self._request_capacity + self._request_metadata_host = torch.empty( + (row_count, request_capacity + 1), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._request_metadata_host_np = self._request_metadata_host.numpy() + self._identity_copy_indices_host = torch.arange( + request_capacity, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._request_metadata_device = torch.zeros( + (row_count, request_capacity + 1), dtype=torch.int32, device=device + ) + self._logical_source_lengths_device = self._request_metadata_device[0, :request_capacity] + self._source_lengths_device = self._request_metadata_device[1, :request_capacity] + self._prompt_lengths_device = self._request_metadata_device[2, :request_capacity] + self._dense_move_offsets_device = self._request_metadata_device[3] + self._swa_move_offsets_device = self._request_metadata_device[4] + self._draft_move_offsets_device = self._request_metadata_device[5] + + self._swa_destination_bases = ( + torch.empty_like(self._prompt_lengths_device) if self._swa_window is not None else None + ) + self._swa_rebase_delta = ( + self.budget - self._swa_window if self._swa_window is not None else 0 + ) + self._mean_cos = torch.empty( + (request_capacity, num_freqs), dtype=torch.float32, device=device + ) + self._mean_sin = torch.empty_like(self._mean_cos) + + def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None: + """Allocate fixed manager-lifetime TopK inputs and outputs.""" + request_capacity = self._request_capacity + selection_width = self._selection_width_capacity + union = self.eviction_mode == "union" + self._selection_rows_per_request = ( + 1 + if union + else self._num_kv_heads + * (self._num_layers if self.eviction_mode == "per_layer_perhead" else 1) + ) + selection_rows = request_capacity * self._selection_rows_per_request + selection_rect = selection_rows * selection_width + if union: + selection_rect = max( + selection_rect, + tp_size * request_capacity * selection_width, + ) + if selection_rect >= 2**31: + raise ValueError(f"selection rectangle overflows 32-bit indexing: {selection_rect}") + + self._decode_lengths_device = torch.full( + (request_capacity,), selection_width, dtype=torch.int32, device=device + ) + if union: + self._selection_scores_rows = torch.empty( + (request_capacity, selection_width), + dtype=torch.float32, + device=device, + ) + self._selection_row_lengths = self._decode_lengths_device + else: + score_shape = ( + request_capacity, + self._num_layers, + self._num_q_heads, + 1, + ) + self._row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) + self._row_inv_std = torch.empty_like(self._row_mean) + self._selection_scores_rows = torch.empty( + (selection_rows, selection_width), + dtype=torch.float32, + device=device, + ) + self._selection_row_lengths = torch.full( + (selection_rows,), + selection_width, + dtype=torch.int32, + device=device, + ) + + self._provisional_rows = torch.zeros( + (selection_rows, self.budget), dtype=torch.int32, device=device + ) + self._kept_ordinal_rows = torch.empty_like(self._provisional_rows) + + def _build_score_runtime( + self, + *, + score_token_capacity: int, + ) -> None: + """Build one scorer span and refresh its page-table bindings.""" + dense_layer = self._target_layout["dense_layers"][0] + anchor_pool = self._target_layout["layer_pools"][dense_layer] + request_capacity = self._request_capacity + block_offsets_host, block_offsets_device = _allocate_block_offset_snapshot( + self.kv_cache_manager, + anchor_pool, + request_capacity=request_capacity, + token_capacity=score_token_capacity + self._protected_tail_capacity, + ) + draft_block_offsets_host = None + draft_block_offsets_device = None + if self._draft_layout is not None: + draft_anchor_pool = self._draft_layout["layer_pools"][0] + draft_block_offsets_host, draft_block_offsets_device = _allocate_block_offset_snapshot( + self.draft_kv_cache_manager, + draft_anchor_pool, + request_capacity=request_capacity, + token_capacity=(score_token_capacity + self._draft_protected_tail_capacity), + ) + + score_scratch, launch_score = build_score_pipeline( + self._target_layout, + block_offsets=block_offsets_device, + source_lengths=self._source_lengths_device, + prompt_lengths=self._prompt_lengths_device, + mean_cos=self._mean_cos, + mean_sin=self._mean_sin, + q_real=self._score_q_real, + q_imag=self._score_q_imag, + mlr_coef=self._score_mlr_coef, + freq_scale_sq=self._freq_scale_sq, + score_token_capacity=score_token_capacity, + union_scores=(self._selection_scores_rows if self.eviction_mode == "union" else None), + ) + + compaction_params = [ + build_compaction_params( + self._target_layout, + block_offsets=block_offsets_device, + kept_ordinals=self._kept_ordinal_rows, + source_lengths=self._source_lengths_device, + dense_destination_bases=self._prompt_lengths_device, + dense_move_offsets=self._dense_move_offsets_device, + protected_tail_capacity=self._protected_tail_capacity, + swa_move_offsets=self._swa_move_offsets_device, + swa_destination_bases=self._swa_destination_bases, + ) + ] + if self._draft_layout is not None: + compaction_params.append( + build_compaction_params( + self._draft_layout, + block_offsets=draft_block_offsets_device, + kept_ordinals=self._kept_ordinal_rows, + source_lengths=self._source_lengths_device, + dense_destination_bases=self._prompt_lengths_device, + dense_move_offsets=self._draft_move_offsets_device, + protected_tail_capacity=self._draft_protected_tail_capacity, + ) + ) + + # Publish new score state only after every allocation and compile succeeds. + self._block_offsets_host = block_offsets_host + self._block_offsets_device = block_offsets_device + self._draft_block_offsets_host = draft_block_offsets_host + self._draft_block_offsets_device = draft_block_offsets_device + self._score_scratch = score_scratch + self._score_token_capacity = score_token_capacity + self._launch_score = launch_score + self._compaction_params = tuple(compaction_params) + + def _create_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: + """Resolve one manager-lifetime V2 pool layout.""" + manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager + + if draft: + global_layers = [int(layer) for layer in manager.pp_layers] + # The draft is never scored: all draft layers compact as dense. + dense_layers: List[int] = list(range(len(global_layers))) + swa_layers: List[int] = [] + swa_window: Optional[int] = None + else: + global_layers = self._global_layers + dense_layers = self._dense_layers + swa_layers = self._swa_layers + swa_window = self._swa_window + layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] + # Canonical pool IDs come from V2; its lookup errors are the precise ones. + layer_offsets = manager.layer_offsets + layer_to_pool = manager.layer_to_pool_mapping_dict + layer_pool_ids = tuple( + int(layer_to_pool[layer_offsets[global_layer]]) for global_layer in global_layers + ) + return dict( + layer_pools=layer_pools, + dense_layers=dense_layers, + swa_layers=swa_layers, + swa_window=swa_window, + layer_pool_ids=layer_pool_ids, + ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score_fused.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score_fused.py new file mode 100644 index 000000000000..90905663022f --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score_fused.py @@ -0,0 +1,1571 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SM100 CuTe-DSL score pipeline for TriAttention.""" + +from __future__ import annotations + +import threading +from typing import Callable, Dict, Optional, Tuple + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + + +@dsl_user_op +def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: + """Inline-PTX sqrt.approx.ftz.f32 (cute.math.sqrt's fast-sqrt kwarg varies across releases).""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Float32(value).ir_value(loc=loc, ip=ip)], + "sqrt.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +CTA_M = 128 +# PADDED_HEAD_COLUMNS is the minimum tcgen05 MMA tile N; GQA groups below 8 +# ride zero-padded head columns. +PADDED_HEAD_COLUMNS = 8 +THREADS = 256 +EPILOGUE_THREADS = 128 +RAW_PAGE_BUFFERS = 2 +# partial_stats: flat [stats_row, page_shard, {count, mean, m2}]; stats_row=segment*num_q_heads+q_head. +STATS_FIELDS = 3 +STATS_MEAN = 1 +STATS_M2 = 2 +# Stats smem scratch: PADDED_HEAD_COLUMNS score origins + one (sum, square-sum) pair per (warp, head column). +STATS_ORIGIN_SLOTS = PADDED_HEAD_COLUMNS +STATS_SCRATCH_ELEMENTS = STATS_ORIGIN_SLOTS + (EPILOGUE_THREADS // 32) * PADDED_HEAD_COLUMNS * 2 +# Staged block-offset entries encode physical_page * K_PLANES_PER_POOL_PAGE + plane. +K_PLANES_PER_POOL_PAGE = 2 + +RAW_K_VECTOR_ELEMENTS = 8 +TMA_DESCRIPTOR_QWORDS = 16 +_SUPPORTED_PAGE_SHARDS = (2, 3) +# Extra page shard for small workloads (few CTAs relative to the SM count). +SMALL_WORKLOAD_PAGE_SHARDS = 3 + + +class _TriAttentionScoreKernel: + """Assign one CTA to each segment/KV-head task and retain W across pages.""" + + # Accumulator dtype of the TMEM-to-global epilogue below. + acc_dtype = cutlass.Float32 + + def __init__( + self, + *, + num_layers: int, + score_token_capacity: int, + num_q_heads: int, + num_freqs: int, + pool_shape: tuple[int, int, int, int, int], + pool_strides: tuple[int, int, int, int, int], + page_shards: int, + write_partial_stats: bool = False, + ) -> None: + """Build the single validated production specialization.""" + self.num_physical_pages, _, num_kv_heads, tokens_per_block, pool_dim = pool_shape + if num_freqs not in (32, 64): + raise ValueError( + "TriAttention CuTe score requires 32 or 64 frequencies (head size 64/128)" + ) + if tokens_per_block not in (32, 128): + raise ValueError("TriAttention CuTe score requires 32- or 128-token pages") + if tokens_per_block > CTA_M: + # One page never spans multiple compute tiles. + raise ValueError("TriAttention CuTe score requires pages within one compute tile") + if num_q_heads % num_kv_heads or num_q_heads // num_kv_heads not in (4, 8): + raise ValueError("TriAttention CuTe score requires GQA group 4 or 8") + if page_shards not in _SUPPORTED_PAGE_SHARDS: + raise ValueError("TriAttention CuTe score has unsupported page shards") + + self.score_token_capacity = score_token_capacity + self.num_layers = num_layers + self.num_q_heads = num_q_heads + self.num_kv_heads = num_kv_heads + self.group_size = num_q_heads // num_kv_heads + self.page_shards = page_shards + self.write_partial_stats = write_partial_stats + self.num_freqs = num_freqs + # cos/sin/mlr coefficient planes per frequency. + self.k_coeff = 3 * num_freqs + self.tokens_per_block = tokens_per_block + # One tile = one page (128-token) or four page fragments (32-token), one TMA box each. + self.box_tokens = min(CTA_M, tokens_per_block) + self.fragments_per_phase = CTA_M // self.box_tokens + self.max_tiles = (score_token_capacity + CTA_M - 1) // CTA_M + + # Producer staging constants baked into the generated code. + self.prefetch_depth = 4 + self.raw_tma_feature_extent = num_freqs + # Barrier tx bytes per phase: the full 128-token tile of one coefficient plane. + self.raw_tma_copy_bytes = CTA_M * num_freqs * (cutlass.BFloat16.width // 8) + self.raw_tma_pipeline_stages = 2 * RAW_PAGE_BUFFERS if write_partial_stats else 1 + # Raw-K page buffers each specialization addresses: the fused union pipeline + # double-buffers across tiles; score-only reuses one buffer (stages 0/1) per tile. + self.raw_page_buffers = RAW_PAGE_BUFFERS if write_partial_stats else 1 + self.accumulator_pipeline_stages = 1 + self.producer_warp_id = 0 + + if pool_dim != 2 * num_freqs: + raise ValueError("K pool shape does not match the CuTe score specialization") + self.s_page, _, self.s_kv_head, self.s_token, self.s_dim = pool_strides + if self.s_token != 2 * num_freqs or self.s_dim != 1: + raise ValueError(f"K pages must be contiguous [{tokens_per_block}, {2 * num_freqs}]") + if self.s_page % RAW_K_VECTOR_ELEMENTS or self.s_kv_head % RAW_K_VECTOR_ELEMENTS: + raise ValueError("K page and KV-head strides must preserve 16-byte alignment") + + @cute.jit + def _stage_raw_band_copies( + self, + raw_tma_pipeline, + raw_tma_producer_state, + band, + first_page, + page_fragments, + shared_partition, + stage_args, + ): + """Stage one raw-K band's fragment copies (caller owns the pipeline acquire/advance).""" + raw_tma_atom, raw_tma_global_partition, kv_head, raw_tma_descriptor_ptr = stage_args + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = first_page + if cutlass.const_expr(fragment > 0): + fragment_page = page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + band, + 0, + (kv_head, fragment_page), + ) + ], + shared_partition[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + + @cute.jit + def __call__( + self, + block_offset_entries: cute.Tensor, + seg_page_off: cute.Tensor, + seg_req_id: cute.Tensor, + seg_layer_id: cute.Tensor, + source_lengths: cute.Tensor, + seg_out_offset: cute.Tensor, + prompt_lengths: cute.Tensor, + q_real: cute.Tensor, + q_imag: cute.Tensor, + mlr_coef: cute.Tensor, + mean_cos: cute.Tensor, + mean_sin: cute.Tensor, + freq_scale_sq: cute.Tensor, + output: cute.Tensor, + partial_stats: cute.Tensor, + anchor_pool: cute.Tensor, + raw_tma_descriptors: cute.Tensor, + request_count: cutlass.Int32, + stream: cuda.CUstream, + ): + self.c_dtype = output.element_type + self.c_layout = utils.LayoutEnum.COL_MAJOR + self.mma_tiler = (CTA_M, PADDED_HEAD_COLUMNS, self.k_coeff) + self.cta_tile_shape_mnk = self.mma_tiler + self.epi_tile = (CTA_M, PADDED_HEAD_COLUMNS) + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + cutlass.Float32, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + self.mma_tiler[:2], + ) + raw_bf16_tiled_mma = sm100_utils.make_trivial_tiled_mma( + cutlass.BFloat16, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + self.mma_tiler[:2], + ) + # Real + imag bf16 stages per raw-page buffer; swizzle follows the num_freqs row width. + raw_bf16_direct_a_smem_layout = sm100_utils.make_smem_layout_a( + raw_bf16_tiled_mma, + (CTA_M, PADDED_HEAD_COLUMNS, self.num_freqs), + cutlass.BFloat16, + 2 * self.raw_page_buffers, + ) + raw_tma_smem_layout = cute.make_composed_layout( + raw_bf16_direct_a_smem_layout.inner, + 0, + cute.make_layout( + (self.raw_tma_feature_extent, self.box_tokens), + stride=(1, self.raw_tma_feature_extent), + ), + ) + raw_tma_source_layout = cute.make_layout( + ( + 2 * self.num_freqs, + self.tokens_per_block, + (self.num_kv_heads, self.num_physical_pages), + ), + stride=( + self.s_dim, + self.s_token, + (self.s_kv_head, self.s_page), + ), + ) + raw_tma_source = cute.make_tensor( + anchor_pool.iterator, + raw_tma_source_layout, + ) + raw_tma_atom, raw_tma_tensor = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + raw_tma_source, + raw_tma_smem_layout, + (self.raw_tma_feature_extent, self.box_tokens), + ) + raw_bf16_b_smem_layout = sm100_utils.make_smem_layout_b( + raw_bf16_tiled_mma, + (CTA_M, PADDED_HEAD_COLUMNS, 2 * self.num_freqs), + cutlass.BFloat16, + 1, + ) + magnitude_lo_tiled_mma = sm100_utils.make_trivial_tiled_mma( + cutlass.Float16, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + cutlass.Float32, + tcgen05.CtaGroup.ONE, + self.mma_tiler[:2], + ) + magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( + magnitude_lo_tiled_mma, + (CTA_M, PADDED_HEAD_COLUMNS, self.num_freqs), + cutlass.Float16, + 1, + ) + magnitude_fp16_b_smem_layout = sm100_utils.make_smem_layout_b( + magnitude_lo_tiled_mma, + (CTA_M, PADDED_HEAD_COLUMNS, self.num_freqs), + cutlass.Float16, + 1, + ) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # One accumulator slot; explicit slot mode keeps the producer/consumer slicing protocol. + self.num_accumulator_slots = self.accumulator_pipeline_stages + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_accumulator_slots)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + # Real+imag bf16 stage pair per raw-page buffer (union double-buffers, score-only single). + raw_k_elements = CTA_M * 2 * self.num_freqs * self.raw_page_buffers + raw_bf16_b_elements = cute.cosize(raw_bf16_b_smem_layout.outer) + magnitude_fp16_a_elements = cute.cosize(magnitude_fp16_a_smem_layout.outer) + magnitude_fp16_b_elements = cute.cosize(magnitude_fp16_b_smem_layout.outer) + stats_scratch_elements = STATS_SCRATCH_ELEMENTS * int(self.write_partial_stats) + + @cute.struct + class SharedStorage: + # PipelineUmmaAsync uses one full and one empty barrier per stage. + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.accumulator_pipeline_stages * 2] + raw_tma_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, + 2 * self.raw_tma_pipeline_stages, + ] + tmem_holding_buf: cutlass.Int32 + sRawK: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_k_elements], + 1024, + ] + sRawBf16B0: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b_elements], + 1024, + ] + sRawBf16B1: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b_elements], + 1024, + ] + sMagnitudeFp16A0: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_fp16_a_elements], + 1024, + ] + sMagnitudeFp16A1: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_fp16_a_elements], + 1024, + ] + sMagnitudeFp16B0: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_fp16_b_elements], + 1024, + ] + sMagnitudeFp16B1: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_fp16_b_elements], + 1024, + ] + sStats: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, stats_scratch_elements], + 16, + ] + + self.shared_storage = SharedStorage + # The score-plane stride product can exceed 2^31, so it must reach the kernel as Int64. + segment_tokens = cutlass.Int64(request_count * self.num_layers * self.score_token_capacity) + num_ctas = request_count * self.num_layers * self.num_kv_heads * self.page_shards + self.kernel( + tiled_mma, + raw_bf16_tiled_mma, + magnitude_lo_tiled_mma, + raw_tma_atom, + raw_tma_tensor, + raw_tma_descriptors, + block_offset_entries, + seg_page_off, + seg_req_id, + seg_layer_id, + source_lengths, + seg_out_offset, + prompt_lengths, + q_real, + q_imag, + mlr_coef, + mean_cos, + mean_sin, + freq_scale_sq, + output, + partial_stats, + segment_tokens, + raw_bf16_direct_a_smem_layout, + raw_tma_smem_layout, + raw_bf16_b_smem_layout, + magnitude_fp16_a_smem_layout, + magnitude_fp16_b_smem_layout, + ).launch( + grid=(num_ctas, 1, 1), + block=(THREADS, 1, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + raw_bf16_tiled_mma: cute.TiledMma, + magnitude_lo_tiled_mma: cute.TiledMma, + raw_tma_atom: cute.CopyAtom, + raw_tma_source: cute.Tensor, + raw_tma_descriptors: cute.Tensor, + block_offset_entries: cute.Tensor, + seg_page_off: cute.Tensor, + seg_req_id: cute.Tensor, + seg_layer_id: cute.Tensor, + source_lengths: cute.Tensor, + seg_out_offset: cute.Tensor, + prompt_lengths: cute.Tensor, + q_real: cute.Tensor, + q_imag: cute.Tensor, + mlr_coef: cute.Tensor, + mean_cos: cute.Tensor, + mean_sin: cute.Tensor, + freq_scale_sq: cute.Tensor, + output: cute.Tensor, + partial_stats: cute.Tensor, + segment_tokens: cutlass.Int64, + raw_bf16_direct_a_smem_layout: cute.ComposedLayout, + raw_tma_smem_layout: cute.ComposedLayout, + raw_bf16_b_smem_layout: cute.ComposedLayout, + magnitude_fp16_a_smem_layout: cute.ComposedLayout, + magnitude_fp16_b_smem_layout: cute.ComposedLayout, + ): + tidx, _, _ = cute.arch.thread_idx() + cta_index, _, _ = cute.arch.block_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_idx = tidx % 32 + task = cta_index // self.page_shards + page_shard = cta_index % self.page_shards + segment = task // self.num_kv_heads + kv_head = task % self.num_kv_heads + req_id = seg_req_id[segment] + layer_id = seg_layer_id[segment] + source_length = source_lengths[req_id] + page_off = seg_page_off[segment] + out_base = seg_out_offset[segment] + # Per-request score window start; scratch writes stay absolute. + score_start = cutlass.Int32(prompt_lengths[req_id]) + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sMagnitudeFp16A0 = storage.sMagnitudeFp16A0.get_tensor( + magnitude_fp16_a_smem_layout.outer, + swizzle=magnitude_fp16_a_smem_layout.inner, + ) + sMagnitudeFp16A1 = storage.sMagnitudeFp16A1.get_tensor( + magnitude_fp16_a_smem_layout.outer, + swizzle=magnitude_fp16_a_smem_layout.inner, + ) + sMagnitudeFp16B0 = storage.sMagnitudeFp16B0.get_tensor( + magnitude_fp16_b_smem_layout.outer, + swizzle=magnitude_fp16_b_smem_layout.inner, + ) + sMagnitudeFp16B1 = storage.sMagnitudeFp16B1.get_tensor( + magnitude_fp16_b_smem_layout.outer, + swizzle=magnitude_fp16_b_smem_layout.inner, + ) + if cutlass.const_expr(self.write_partial_stats): + sStats = storage.sStats.get_tensor(cute.make_layout(STATS_SCRATCH_ELEMENTS)) + raw_k_storage = storage.sRawK + cpasync_raw_k_0 = raw_k_storage.get_tensor( + raw_bf16_direct_a_smem_layout.outer, + swizzle=raw_bf16_direct_a_smem_layout.inner, + ) + cpasync_raw_k_real = cpasync_raw_k_0[(None, None, None, 0)] + cpasync_raw_k_imag = cpasync_raw_k_0[(None, None, None, 1)] + if cutlass.const_expr(self.write_partial_stats): + # The second raw-page buffer exists only in the fused union specialization. + cpasync_raw_k_real_next = cpasync_raw_k_0[(None, None, None, 2)] + cpasync_raw_k_imag_next = cpasync_raw_k_0[(None, None, None, 3)] + # Use only the outer mapping for the TMA destination so the swizzle is not applied twice. + raw_tma_source_tiles = cute.local_tile( + raw_tma_source, + (self.raw_tma_feature_extent, self.box_tokens), + coord=(None, None, None), + ) + # One smem view/TMA partition per fragment; fragment offsets are whole swizzle periods. + raw_tma_shared_partition_real = [] + raw_tma_shared_partition_imag = [] + raw_tma_shared_partition_real_next = [] + raw_tma_shared_partition_imag_next = [] + raw_tma_global_partition = None + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_offset = fragment * self.box_tokens * self.raw_tma_feature_extent + fragment_real = cute.make_tensor( + cpasync_raw_k_real.iterator + fragment_offset, + raw_tma_smem_layout.outer, + ) + fragment_imag = cute.make_tensor( + cpasync_raw_k_imag.iterator + fragment_offset, + raw_tma_smem_layout.outer, + ) + if cutlass.const_expr(self.write_partial_stats): + fragment_real_next = cute.make_tensor( + cpasync_raw_k_real_next.iterator + fragment_offset, + raw_tma_smem_layout.outer, + ) + fragment_imag_next = cute.make_tensor( + cpasync_raw_k_imag_next.iterator + fragment_offset, + raw_tma_smem_layout.outer, + ) + partition_real, global_partition = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(fragment_real, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + partition_imag, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(fragment_imag, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + if cutlass.const_expr(self.write_partial_stats): + partition_real_next, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(fragment_real_next, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + partition_imag_next, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(fragment_imag_next, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + raw_tma_shared_partition_real.append(partition_real) + raw_tma_shared_partition_imag.append(partition_imag) + if cutlass.const_expr(self.write_partial_stats): + raw_tma_shared_partition_real_next.append(partition_real_next) + raw_tma_shared_partition_imag_next.append(partition_imag_next) + raw_tma_global_partition = global_partition + raw_tensormap_manager = utils.TensorMapManager( + utils.TensorMapUpdateMode.GMEM, + 128, + ) + raw_tma_descriptor_ptr = raw_tensormap_manager.get_tensormap_ptr( + (raw_tma_descriptors.iterator + layer_id * TMA_DESCRIPTOR_QWORDS).align(128), + cute.AddressSpace.generic, + ) + # Trace-time invariants of every raw-band stage copy; bound once, unpacked in the helper. + raw_stage_args = (raw_tma_atom, raw_tma_global_partition, kv_head, raw_tma_descriptor_ptr) + sRawBf16B0 = storage.sRawBf16B0.get_tensor( + raw_bf16_b_smem_layout.outer, + swizzle=raw_bf16_b_smem_layout.inner, + ) + sRawBf16B1 = storage.sRawBf16B1.get_tensor( + raw_bf16_b_smem_layout.outer, + swizzle=raw_bf16_b_smem_layout.inner, + ) + + tile_index = score_start // CTA_M + page_shard + tile_start_token = tile_index * CTA_M + shard_first_tile_start_token = tile_start_token + tiles_processed = cutlass.Int32(0) + if cutlass.const_expr(self.write_partial_stats): + stats_page_scores_m128 = cute.make_rmem_tensor((PADDED_HEAD_COLUMNS,), cutlass.Float32) + stats_origins_m128 = cute.make_rmem_tensor((PADDED_HEAD_COLUMNS,), cutlass.Float32) + stats_sums_m128 = cute.make_rmem_tensor((PADDED_HEAD_COLUMNS,), cutlass.Float32) + stats_square_sums_m128 = cute.make_rmem_tensor((PADDED_HEAD_COLUMNS,), cutlass.Float32) + for stats_head in cutlass.range_constexpr(PADDED_HEAD_COLUMNS): + stats_sums_m128[stats_head] = cutlass.Float32(0.0) + stats_square_sums_m128[stats_head] = cutlass.Float32(0.0) + producer_prefetched_page_id_lane0 = cutlass.Int32(0) + physical_fragments_arg = None + prefetched_fragments_arg = None + if cutlass.const_expr(self.fragments_per_phase > 1): + # Per-fragment page-id registers; slot 0 unused (fragment 0 uses the scalar registers). + producer_prefetched_page_ids_lane0 = cute.make_rmem_tensor( + (self.fragments_per_phase,), cutlass.Int32 + ) + physical_page_fragments = cute.make_rmem_tensor( + (self.fragments_per_phase,), cutlass.Int32 + ) + prefetched_page_fragments = cute.make_rmem_tensor( + (self.fragments_per_phase,), cutlass.Int32 + ) + physical_fragments_arg = physical_page_fragments + prefetched_fragments_arg = prefetched_page_fragments + shard_has_page = source_length > score_start and tile_start_token < source_length + empty_shard = source_length <= score_start or tile_start_token >= source_length + if cutlass.dynamic_expr(shard_has_page): + if warp_idx == self.producer_warp_id: + if lane_idx == 0: + # Staged entries encode physical_page * kv_factor; decode to the pool page. + producer_prefetched_page_id_lane0 = ( + cutlass.Int32( + block_offset_entries[page_off + tile_index * self.fragments_per_phase] + ) + // K_PLANES_PER_POOL_PAGE + ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.fragments_per_phase): + # Clamp tail-fragment pages so the TMA never reads an unstaged entry. + fragment_page_id = producer_prefetched_page_id_lane0 + if tile_start_token + fragment * self.box_tokens < source_length: + fragment_page_id = ( + cutlass.Int32( + block_offset_entries[ + page_off + + tile_index * self.fragments_per_phase + + fragment + ] + ) + // K_PLANES_PER_POOL_PAGE + ) + producer_prefetched_page_ids_lane0[fragment] = fragment_page_id + tCrRawBf16ASplit = raw_bf16_tiled_mma.make_fragment_A(cpasync_raw_k_0) + tCrRawBf16B0 = raw_bf16_tiled_mma.make_fragment_B(sRawBf16B0) + tCrRawBf16B1 = raw_bf16_tiled_mma.make_fragment_B(sRawBf16B1) + tCrMagnitudeFp16A0 = magnitude_lo_tiled_mma.make_fragment_A(sMagnitudeFp16A0) + tCrMagnitudeFp16A1 = magnitude_lo_tiled_mma.make_fragment_A(sMagnitudeFp16A1) + tCrMagnitudeFp16B0 = magnitude_lo_tiled_mma.make_fragment_B(sMagnitudeFp16B0) + tCrMagnitudeFp16B1 = magnitude_lo_tiled_mma.make_fragment_B(sMagnitudeFp16B1) + raw_tma_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.raw_tma_mbar_ptr.data_ptr(), + num_stages=self.raw_tma_pipeline_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, THREADS // 32), + tx_count=self.raw_tma_copy_bytes, + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + tidx=tidx, + defer_sync=True, + ) + raw_tma_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.raw_tma_pipeline_stages, + ) + raw_tma_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, + self.raw_tma_pipeline_stages, + ) + + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.accumulator_pipeline_stages, + producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, THREADS // 32), + cta_layout_vmnk=cute.make_layout((1, 1, 1, 1)), + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.accumulator_pipeline_stages + ) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.accumulator_pipeline_stages + ) + stats_epilogue_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=EPILOGUE_THREADS, + ) + cute.arch.mbarrier_init_fence() + # Per-(head, frequency) score coefficients, split into bf16/fp16 value+residual pairs. + for weight_round in cutlass.range_constexpr(PADDED_HEAD_COLUMNS * self.k_coeff // THREADS): + linear_index = tidx + weight_round * THREADS + qg = linear_index // self.k_coeff + feature = linear_index % self.k_coeff + coefficient_kind = feature // self.num_freqs + frequency = feature % self.num_freqs + mean_offset = req_id * self.num_freqs + frequency + # Padded GQA columns read the group's first head and force zero coefficients. + qg_read = qg + if cutlass.const_expr(self.group_size < PADDED_HEAD_COLUMNS): + if qg_read >= self.group_size: + qg_read = cutlass.Int32(0) + q_head = kv_head * self.group_size + qg_read + calib_offset = (layer_id * self.num_q_heads + q_head) * self.num_freqs + frequency + qr = cutlass.Float32(q_real[calib_offset]) + qi = cutlass.Float32(q_imag[calib_offset]) + mcos = cutlass.Float32(mean_cos[mean_offset]) + msin = cutlass.Float32(mean_sin[mean_offset]) + scale = cutlass.Float32(freq_scale_sq[frequency]) + value = cutlass.Float32(0.0) + if coefficient_kind == 0: + value = scale * (qr * mcos - qi * msin) + elif coefficient_kind == 1: + value = scale * (qr * msin + qi * mcos) + else: + value = scale * cutlass.Float32(mlr_coef[calib_offset]) + if cutlass.const_expr(self.group_size < PADDED_HEAD_COLUMNS): + if qg >= self.group_size: + value = cutlass.Float32(0.0) + raw_k_block = feature // 16 + magnitude_k_block = frequency // 16 + if coefficient_kind < 2: + value_bf16_0 = cutlass.BFloat16(value) + residual_1 = value - cutlass.Float32(value_bf16_0) + value_bf16_1 = cutlass.BFloat16(residual_1) + raw_coord = ( + (qg, feature % 16), + 0, + raw_k_block, + 0, + ) + sRawBf16B0[raw_coord] = value_bf16_0 + sRawBf16B1[raw_coord] = value_bf16_1 + else: + value_fp16_0 = cutlass.Float16(value) + value_fp16_1 = cutlass.Float16(value - cutlass.Float32(value_fp16_0)) + magnitude_coord_fp16 = ( + (qg, frequency % 16), + 0, + magnitude_k_block, + 0, + ) + sMagnitudeFp16B0[magnitude_coord_fp16] = value_fp16_0 + sMagnitudeFp16B1[magnitude_coord_fp16] = value_fp16_1 + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.barrier() + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_accumulator_slots)) + if warp_idx == 0: + cute.arch.alloc_tmem( + self.num_tmem_alloc_cols, + storage.tmem_holding_buf, + is_two_cta=False, + ) + cute.arch.barrier() + tmem_ptr = cute.arch.retrieve_tmem_ptr( + cutlass.Float32, + alignment=16, + ptr_to_buffer_holding_addr=storage.tmem_holding_buf, + ) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + if cutlass.dynamic_expr(empty_shard): + if warp_idx == self.producer_warp_id: + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + + thr_mma = tiled_mma.get_slice(0) + if cutlass.const_expr(self.write_partial_stats): + if cutlass.dynamic_expr(shard_has_page): + if warp_idx == self.producer_warp_id: + prefetched_physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, + ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.fragments_per_phase): + prefetched_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], + 0, + ) + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 0, + prefetched_physical_page, + prefetched_fragments_arg, + raw_tma_shared_partition_real, + raw_stage_args, + ) + raw_tma_producer_state.advance() + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 1, + prefetched_physical_page, + prefetched_fragments_arg, + raw_tma_shared_partition_imag, + raw_stage_args, + ) + raw_tma_producer_state.advance() + while ( + source_length > score_start + and tile_start_token < source_length + and tiles_processed < self.max_tiles + ): + physical_page = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, + ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.fragments_per_phase): + physical_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], + 0, + ) + raw_page_buffer = cutlass.Int32(0) + if cutlass.const_expr(self.write_partial_stats): + raw_page_buffer = tiles_processed % RAW_PAGE_BUFFERS + raw_real_stage = raw_page_buffer * 2 + raw_imag_stage = raw_real_stage + 1 + if cutlass.const_expr(not self.write_partial_stats): + if warp_idx == self.producer_warp_id: + # Phase 0 fills the packed real-band stage view. + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 0, + physical_page, + physical_fragments_arg, + raw_tma_shared_partition_real, + raw_stage_args, + ) + raw_tma_producer_state.advance() + raw_tma_pipeline.consumer_wait(raw_tma_consumer_state) + raw_tma_pipeline.consumer_release(raw_tma_consumer_state) + raw_tma_consumer_state.advance() + if cutlass.const_expr(not self.write_partial_stats): + if warp_idx == self.producer_warp_id: + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 1, + physical_page, + physical_fragments_arg, + raw_tma_shared_partition_imag, + raw_stage_args, + ) + raw_tma_producer_state.advance() + + next_page_id_lane0 = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + if lane_idx == 0: + next_tile_start_token = tile_start_token + CTA_M * self.page_shards + next_pages_processed = tiles_processed + 1 + if ( + next_tile_start_token < source_length + and next_pages_processed < self.max_tiles + ): + next_page_id_lane0 = ( + cutlass.Int32( + block_offset_entries[ + page_off + + (tile_index + self.page_shards) * self.fragments_per_phase + ] + ) + // K_PLANES_PER_POOL_PAGE + ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.fragments_per_phase): + # Tail-tile clamp: fall back to the first fragment's page. + next_fragment_page_id = next_page_id_lane0 + if ( + next_tile_start_token + fragment * self.box_tokens + < source_length + ): + next_fragment_page_id = ( + cutlass.Int32( + block_offset_entries[ + page_off + + (tile_index + self.page_shards) + * self.fragments_per_phase + + fragment + ] + ) + // K_PLANES_PER_POOL_PAGE + ) + producer_prefetched_page_ids_lane0[fragment] = next_fragment_page_id + producer_prefetched_page_id_lane0 = next_page_id_lane0 + + # Submit B0-real while the imaginary TMA is in flight. + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + if warp_idx == self.producer_warp_id: + acc_pipeline.producer_acquire(acc_producer_state) + raw_bf16_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + False, + ) + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, raw_real_stage)], + tCrRawBf16B0[(None, None, raw_k_block, 0)], + tCtAcc, + ) + raw_bf16_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + True, + ) + raw_tma_pipeline.consumer_wait(raw_tma_consumer_state) + if cutlass.const_expr(self.write_partial_stats): + if warp_idx == self.producer_warp_id: + next_tile_start_token = tile_start_token + CTA_M * self.page_shards + next_pages_processed = tiles_processed + 1 + prefetch_next_raw = ( + next_tile_start_token < source_length + and next_pages_processed < self.max_tiles + ) + prefetched_physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, + ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.fragments_per_phase): + prefetched_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], + 0, + ) + if cutlass.dynamic_expr(prefetch_next_raw): + # ONE acquire/advance per band, shared across the dynamic destination arms. + next_raw_page_buffer = (raw_page_buffer + 1) % RAW_PAGE_BUFFERS + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + if cutlass.dynamic_expr(next_raw_page_buffer == 0): + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 0, + prefetched_physical_page, + prefetched_fragments_arg, + raw_tma_shared_partition_real, + raw_stage_args, + ) + else: + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 0, + prefetched_physical_page, + prefetched_fragments_arg, + raw_tma_shared_partition_real_next, + raw_stage_args, + ) + raw_tma_producer_state.advance() + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + if cutlass.dynamic_expr(next_raw_page_buffer == 0): + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 1, + prefetched_physical_page, + prefetched_fragments_arg, + raw_tma_shared_partition_imag, + raw_stage_args, + ) + else: + self._stage_raw_band_copies( + raw_tma_pipeline, + raw_tma_producer_state, + 1, + prefetched_physical_page, + prefetched_fragments_arg, + raw_tma_shared_partition_imag_next, + raw_stage_args, + ) + raw_tma_producer_state.advance() + # Each lane stages one frequency per pass; 64-frequency heads take two passes. + for freq_rep in cutlass.range_constexpr(self.num_freqs // 32): + frequency = lane_idx + 32 * freq_rep + # Stage prefetch_depth independent token loads before consuming any of them. + for token_base in cutlass.range( + 0, + CTA_M // (THREADS // 32), + self.prefetch_depth, + unroll_full=False, + ): + staged_real = cute.make_rmem_tensor((self.prefetch_depth,), cutlass.Float32) + staged_imag = cute.make_rmem_tensor((self.prefetch_depth,), cutlass.Float32) + for prefetch_index in cutlass.range_constexpr(self.prefetch_depth): + token_round = token_base + prefetch_index + token = warp_idx + token_round * (THREADS // 32) + staged_real[prefetch_index] = cutlass.Float32( + cpasync_raw_k_0[ + ( + (token, frequency % 16), + 0, + frequency // 16, + raw_real_stage, + ) + ] + ) + staged_imag[prefetch_index] = cutlass.Float32( + cpasync_raw_k_0[ + ( + (token, frequency % 16), + 0, + frequency // 16, + raw_imag_stage, + ) + ] + ) + + for prefetch_index in cutlass.range_constexpr(self.prefetch_depth): + token_round = token_base + prefetch_index + token = warp_idx + token_round * (THREADS // 32) + real = staged_real[prefetch_index] + imag = staged_imag[prefetch_index] + norm2 = real * real + imag * imag + magnitude = _sqrt_approx_ftz(norm2) + magnitude_fp16_0 = cutlass.Float16(magnitude) + magnitude_fp16_1 = cutlass.Float16( + magnitude - cutlass.Float32(magnitude_fp16_0) + ) + magnitude_k_block_fp16 = frequency // 16 + magnitude_coord_fp16 = ( + (token, frequency % 16), + 0, + magnitude_k_block_fp16, + 0, + ) + sMagnitudeFp16A0[magnitude_coord_fp16] = magnitude_fp16_0 + sMagnitudeFp16A1[magnitude_coord_fp16] = magnitude_fp16_1 + + cute.arch.fence_proxy("async.shared", space="cta") + cute.arch.barrier() + + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + + if warp_idx == self.producer_warp_id: + # Finish B0-imag, then issue B1-real and B1-imag. + raw_bf16_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + True, + ) + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): + imag_b_block = self.num_freqs // 16 + raw_k_block + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, raw_imag_stage)], + tCrRawBf16B0[(None, None, imag_b_block, 0)], + tCtAcc, + ) + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, raw_real_stage)], + tCrRawBf16B1[(None, None, raw_k_block, 0)], + tCtAcc, + ) + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): + imag_b_block = self.num_freqs // 16 + raw_k_block + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, raw_imag_stage)], + tCrRawBf16B1[(None, None, imag_b_block, 0)], + tCtAcc, + ) + # Compensated FP16 magnitude: |K|*coeff = A0*B0 + A0*B1 + A1*B0 (A1*B1 dropped). + magnitude_lo_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + True, + ) + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): + cute.gemm( + magnitude_lo_tiled_mma, + tCtAcc, + tCrMagnitudeFp16A0[(None, None, magnitude_k_block, 0)], + tCrMagnitudeFp16B0[(None, None, magnitude_k_block, 0)], + tCtAcc, + ) + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): + cute.gemm( + magnitude_lo_tiled_mma, + tCtAcc, + tCrMagnitudeFp16A0[(None, None, magnitude_k_block, 0)], + tCrMagnitudeFp16B1[(None, None, magnitude_k_block, 0)], + tCtAcc, + ) + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): + cute.gemm( + magnitude_lo_tiled_mma, + tCtAcc, + tCrMagnitudeFp16A1[(None, None, magnitude_k_block, 0)], + tCrMagnitudeFp16B0[(None, None, magnitude_k_block, 0)], + tCtAcc, + ) + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + if tiles_processed == 0: + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + acc_pipeline.consumer_wait(acc_consumer_state) + if cutlass.const_expr(self.write_partial_stats): + # Release only the current imag phase after all of its async consumers finish. + raw_tma_pipeline.consumer_release(raw_tma_consumer_state) + raw_tma_consumer_state.advance() + # Every term multiplying segment_tokens must stay 64-bit; the plane stride can exceed 2^31. + output_offset = ( + cutlass.Int64(kv_head * PADDED_HEAD_COLUMNS) * segment_tokens + + out_base + + tile_start_token + ) + page_output = cute.make_tensor( + output.iterator + output_offset, + cute.make_layout( + (CTA_M, PADDED_HEAD_COLUMNS, 1), + stride=( + 1, + segment_tokens, + PADDED_HEAD_COLUMNS * segment_tokens, + ), + ), + ) + gC_mnl = cute.local_tile(page_output, self.epi_tile, (None, None, None)) + tCgC = thr_mma.partition_C(gC_mnl) + epilogue_tidx = tidx % EPILOGUE_THREADS + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + self.epi_tile, + False, + ) + accumulator_epilogue = cute.flat_divide( + tCtAcc[((None, None), 0, 0)], + self.epi_tile, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy( + copy_atom_t2r, + accumulator_epilogue[(None, None, 0, 0)], + ) + thread_copy = tiled_copy_t2r.get_slice(epilogue_tidx) + tTR_tAcc = thread_copy.partition_S(accumulator_epilogue) + output_epilogue = cute.flat_divide( + tCgC[((None, None), 0, 0, None, None, None)], + self.epi_tile, + ) + tTR_gC = thread_copy.partition_D(output_epilogue) + register_shape = tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape + tTR_rAcc = cute.make_rmem_tensor(register_shape, self.acc_dtype) + tTR_rC = cute.make_rmem_tensor(register_shape, self.c_dtype) + simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + tTR_gC = tTR_gC[(None, None, None, None, None, 0, 0, 0)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_gC = cute.group_modes(tTR_gC, 3, cute.rank(tTR_gC)) + if tidx < EPILOGUE_THREADS: + for subtile_idx in cutlass.range_constexpr(cute.size(tTR_tAcc.shape, mode=[3])): + cute.copy( + tiled_copy_t2r, + tTR_tAcc[(None, None, None, subtile_idx)], + tTR_rAcc, + ) + tTR_rC.store(tTR_rAcc.load().to(self.c_dtype)) + # Only the straddling first tile takes the per-token branch. + if cutlass.dynamic_expr( + tile_start_token >= score_start + and tile_start_token + CTA_M <= self.score_token_capacity + ): + cute.copy( + simt_atom, + tTR_rC, + tTR_gC[(None, None, None, subtile_idx)], + ) + else: + output_token = tile_start_token + epilogue_tidx + if cutlass.dynamic_expr( + output_token >= score_start and output_token < self.score_token_capacity + ): + cute.copy( + simt_atom, + tTR_rC, + tTR_gC[(None, None, None, subtile_idx)], + ) + if cutlass.const_expr(self.write_partial_stats): + stats_output = cute.coalesce(tTR_rC) + stats_head_base = subtile_idx * cute.size(stats_output) + for stats_value in cutlass.range_constexpr(cute.size(stats_output)): + stats_head = stats_head_base + stats_value + stats_page_scores_m128[stats_head] = cutlass.Float32( + stats_output[stats_value] + ) + if tiles_processed == 0: + if tidx == 0: + sStats[stats_head] = cutlass.Float32(stats_output[stats_value]) + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + if cutlass.const_expr(self.write_partial_stats): + if tidx < EPILOGUE_THREADS: + stats_epilogue_barrier.wait_unaligned() + else: + cute.arch.barrier() + if cutlass.const_expr(not self.write_partial_stats): + raw_tma_pipeline.consumer_release(raw_tma_consumer_state) + raw_tma_consumer_state.advance() + if cutlass.const_expr(self.write_partial_stats): + if tiles_processed == 0: + for stats_head in cutlass.range_constexpr(PADDED_HEAD_COLUMNS): + stats_origins_m128[stats_head] = sStats[stats_head] + stats_token = tile_start_token + tidx + if tidx < EPILOGUE_THREADS: + if cutlass.dynamic_expr( + stats_token >= score_start and stats_token < source_length + ): + for stats_head in cutlass.range_constexpr(PADDED_HEAD_COLUMNS): + stats_delta = ( + stats_page_scores_m128[stats_head] - stats_origins_m128[stats_head] + ) + stats_sums_m128[stats_head] = stats_sums_m128[stats_head] + stats_delta + stats_square_sums_m128[stats_head] = ( + stats_square_sums_m128[stats_head] + stats_delta * stats_delta + ) + tile_index += self.page_shards + tile_start_token += CTA_M * self.page_shards + tiles_processed += 1 + if warp_idx == self.producer_warp_id: + raw_tma_pipeline.producer_tail(raw_tma_producer_state) + acc_pipeline.producer_tail(acc_producer_state) + if cutlass.const_expr(self.write_partial_stats): + for stats_head in cutlass.range_constexpr(PADDED_HEAD_COLUMNS): + stats_sum = stats_sums_m128[stats_head] + stats_square_sum = stats_square_sums_m128[stats_head] + for stats_offset in (16, 8, 4, 2, 1): + stats_sum = stats_sum + cute.arch.shuffle_sync_bfly(stats_sum, stats_offset) + stats_square_sum = stats_square_sum + cute.arch.shuffle_sync_bfly( + stats_square_sum, stats_offset + ) + if lane_idx == 0 and warp_idx < EPILOGUE_THREADS // 32: + stats_scratch_base = ( + STATS_ORIGIN_SLOTS + (warp_idx * PADDED_HEAD_COLUMNS + stats_head) * 2 + ) + sStats[stats_scratch_base] = stats_sum + sStats[stats_scratch_base + 1] = stats_square_sum + cute.arch.barrier() + if warp_idx == 0: + # Only real heads merge into the compact rows (row = segment*num_q_heads + q_head). + if lane_idx < self.group_size: + stats_sum = cutlass.Float32(0.0) + stats_square_sum = cutlass.Float32(0.0) + for stats_warp in cutlass.range_constexpr(EPILOGUE_THREADS // 32): + stats_scratch_base = ( + STATS_ORIGIN_SLOTS + (stats_warp * PADDED_HEAD_COLUMNS + lane_idx) * 2 + ) + stats_sum = stats_sum + sStats[stats_scratch_base] + stats_square_sum = stats_square_sum + sStats[stats_scratch_base + 1] + stats_count_i32 = tiles_processed * CTA_M + if tiles_processed > 0: + stats_invalid_prefix = score_start - shard_first_tile_start_token + if cutlass.dynamic_expr(stats_invalid_prefix > 0): + stats_count_i32 = stats_count_i32 - stats_invalid_prefix + stats_last_tile_start_token = tile_start_token - CTA_M * self.page_shards + stats_invalid_tail = stats_last_tile_start_token + CTA_M - source_length + if cutlass.dynamic_expr(stats_invalid_tail > 0): + stats_count_i32 = stats_count_i32 - stats_invalid_tail + stats_count = cutlass.Float32(stats_count_i32) + stats_mean = cutlass.Float32(0.0) + stats_m2 = cutlass.Float32(0.0) + if cutlass.dynamic_expr(stats_count_i32 > 0): + inverse_count = cutlass.Float32(1.0) / stats_count + stats_origin = sStats[lane_idx] + stats_mean = stats_origin + stats_sum * inverse_count + stats_m2 = cute.arch.fmax( + stats_square_sum - stats_sum * stats_sum * inverse_count, + cutlass.Float32(0.0), + ) + stats_row = task * self.group_size + lane_idx + stats_base = (stats_row * self.page_shards + page_shard) * STATS_FIELDS + partial_stats[stats_base] = stats_count + partial_stats[stats_base + STATS_MEAN] = stats_mean + partial_stats[stats_base + STATS_M2] = stats_m2 + cute.arch.barrier() + if warp_idx == 0: + cute.arch.dealloc_tmem(tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=False) + + +_COMPILED_KERNELS: dict[tuple, object] = {} +_COMPILE_LOCK = threading.Lock() + + +def _encode_tma_descriptors( + layer_pools: list[torch.Tensor], + layer_indices: list[int], + num_freqs: int, + tokens_per_block: int, +) -> torch.Tensor: + anchor = layer_pools[layer_indices[0]] + active_layers = set(layer_indices) + uint32 = cuda.cuuint32_t + uint64 = cuda.cuuint64_t + descriptor_rows = [] + for layer, maybe_pool in enumerate(layer_pools): + pool = maybe_pool if layer in active_layers else anchor + if pool.dtype != torch.bfloat16: + raise TypeError("TriAttention CuTe score requires BF16 layer pools") + if tuple(pool.shape[1:]) != tuple(anchor.shape[1:]): + raise ValueError("TriAttention CuTe score requires uniform scored-layer pool geometry") + _, kv_factor, num_kv_heads, pool_tokens, head_dim = pool.shape + if (kv_factor, pool_tokens, head_dim) != ( + K_PLANES_PER_POOL_PAGE, + tokens_per_block, + 2 * num_freqs, + ): + raise ValueError( + f"TriAttention CuTe score requires [page, 2, Hkv, {tokens_per_block}, " + f"{2 * num_freqs}] pools" + ) + s_page, _, s_kv_head, s_token, s_dim = map(int, pool.stride()) + if s_dim != 1: + raise ValueError("TriAttention CuTe score requires contiguous K features") + + global_dims = [2 * num_freqs, tokens_per_block] + global_strides_bytes = [s_token * pool.element_size()] + if num_kv_heads > 1: + global_dims.append(int(num_kv_heads)) + global_strides_bytes.append(s_kv_head * pool.element_size()) + if pool.shape[0] > 1: + global_dims.append(int(pool.shape[0])) + global_strides_bytes.append(s_page * pool.element_size()) + tensor_rank = len(global_dims) + # One TMA box covers one coefficient plane of one page fragment. + box_dims = [num_freqs, min(CTA_M, tokens_per_block)] + [1] * (tensor_rank - 2) + status, tensor_map = cuda.cuTensorMapEncodeTiled( + cuda.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, + uint32(tensor_rank), + pool.data_ptr(), + [uint64(value) for value in global_dims], + [uint64(value) for value in global_strides_bytes], + [uint32(value) for value in box_dims], + [uint32(1) for _ in range(tensor_rank)], + cuda.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, + # The swizzle must match the destination smem layout (inner row = num_freqs bf16). + ( + cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B + if num_freqs * 2 == 64 + else cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B + ), + cuda.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_NONE, + cuda.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE, + ) + if status != cuda.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuTensorMapEncodeTiled failed for layer {layer}: {status}") + descriptor_rows.append( + [ + value if value < 1 << 63 else value - (1 << 64) + for value in map(int, tensor_map.opaque) + ] + ) + + descriptors = torch.tensor( + descriptor_rows, + dtype=torch.int64, + device=anchor.device, + ) + if descriptors.shape != (len(layer_pools), TMA_DESCRIPTOR_QWORDS): + raise AssertionError("each TriAttention TMA descriptor must occupy 128 bytes") + if descriptors.data_ptr() % 128 or descriptors.stride(0) != TMA_DESCRIPTOR_QWORDS: + raise AssertionError("TriAttention TMA descriptor rows must be 128-byte aligned") + return descriptors + + +def _tensor_spec(tensor: torch.Tensor) -> tuple: + return ( + tuple(int(value) for value in tensor.shape), + tuple(int(value) for value in tensor.stride()), + tensor.dtype, + tensor.device.type, + tensor.device.index, + ) + + +def _to_cute(tensor: torch.Tensor, *, assumed_align: int = 16) -> cute.Tensor: + return from_dlpack(tensor, assumed_align=assumed_align) + + +def _get_or_compile(cache_key: tuple, build: Callable[[], object]) -> object: + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + compiled = build() + _COMPILED_KERNELS[cache_key] = compiled + return compiled + + +def build_score_pipeline( + layout: Dict[str, object], + *, + block_offsets: torch.Tensor, + source_lengths: torch.Tensor, + prompt_lengths: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + freq_scale_sq: torch.Tensor, + score_token_capacity: int, + union_scores: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, Callable[[int], None]]: + """Compile a capacity-specific score pipeline and return its scratch and launcher.""" + layer_pools = tuple(layout["layer_pools"]) + scored_layers = tuple(int(layer) for layer in layout["dense_layers"]) + layer_pool_ids = tuple(int(slot) for slot in layout["layer_pool_ids"]) + anchor_pool = layer_pools[scored_layers[0]] + device = anchor_pool.device + request_capacity = int(source_lengths.numel()) + num_layers = len(scored_layers) + score_token_capacity = int(score_token_capacity) + + num_q_heads = int(q_real.shape[1]) + num_freqs = int(q_real.shape[2]) + _, _, num_kv_heads, tokens_per_block, _ = anchor_pool.shape + num_kv_heads = int(num_kv_heads) + tokens_per_block = int(tokens_per_block) + max_segments = request_capacity * num_layers + max_segment_offset = (max_segments - 1) * score_token_capacity + if max_segment_offset >= 2**31: + raise ValueError(f"score bucket overflows the int32 segment offsets: {max_segment_offset}") + + segment_request_ids = torch.arange( + request_capacity, dtype=torch.int32, device=device + ).repeat_interleave(num_layers) + segment_layer_ids = torch.tensor(scored_layers, dtype=torch.int32, device=device).repeat( + request_capacity + ) + segment_pool_slots = torch.tensor( + tuple(layer_pool_ids[layer] for layer in scored_layers), + dtype=torch.int64, + device=device, + ).repeat(request_capacity) + segment_page_offsets = segment_pool_slots * block_offsets.stride(0) + segment_request_ids.to( + torch.int64 + ) * block_offsets.stride(1) + segment_output_offsets = ( + torch.arange(max_segments, dtype=torch.int64, device=device) * score_token_capacity + ).to(torch.int32) + + score_scratch = torch.empty( + num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * score_token_capacity, + dtype=torch.float32, + device=device, + ) + partial_stats = torch.empty( + ( + request_capacity * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * STATS_FIELDS + if union_scores is not None + else 1 + ), + dtype=torch.float32, + device=device, + ) + tma_descriptors = _encode_tma_descriptors( + list(layer_pools), + list(scored_layers), + num_freqs, + tokens_per_block, + ) + + score_operands = ( + (block_offsets.view(-1), 16), + (segment_page_offsets, 16), + (segment_request_ids, 16), + (segment_layer_ids, 16), + (source_lengths, 4), + (segment_output_offsets, 16), + (prompt_lengths, 4), + (q_real.view(-1), 16), + (q_imag.view(-1), 16), + (mlr_coef.view(-1), 16), + (mean_cos.view(-1), 16), + (mean_sin.view(-1), 16), + (freq_scale_sq, 16), + (score_scratch, 16), + (partial_stats, 16), + (anchor_pool, 16), + (tma_descriptors, 128), + ) + score_args = tuple( + _to_cute(tensor, assumed_align=alignment) for tensor, alignment in score_operands + ) + tensor_specs = tuple(_tensor_spec(tensor) for tensor, _ in score_operands) + static_geometry = ( + request_capacity, + num_layers, + score_token_capacity, + num_q_heads, + num_kv_heads, + num_freqs, + tokens_per_block, + tuple(int(value) for value in anchor_pool.shape), + tuple(int(value) for value in anchor_pool.stride()), + ) + stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + sm_count = int(torch.cuda.get_device_properties(device).multi_processor_count) + variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] + if request_capacity > 1: + variants.append((request_capacity, 2)) + + compiled_scores: Dict[int, object] = {} + page_shards_by_request_count: Dict[int, int] = {} + variant_key = ( + "triattention_cute_score_stats" if union_scores is not None else "triattention_cute_score" + ) + for request_count, page_shards in variants: + cache_key = ( + variant_key, + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + compiled_scores[request_count] = _get_or_compile( + cache_key, + lambda page_shards=page_shards: cute.compile( + _TriAttentionScoreKernel( + num_layers=num_layers, + score_token_capacity=score_token_capacity, + num_q_heads=num_q_heads, + num_freqs=num_freqs, + pool_shape=tuple(int(value) for value in anchor_pool.shape), + pool_strides=tuple(int(value) for value in anchor_pool.stride()), + page_shards=page_shards, + write_partial_stats=union_scores is not None, + ), + *score_args, + cutlass.Int32(1), + stream, + ), + ) + page_shards_by_request_count[request_count] = page_shards + + if request_capacity > 1: + small = compiled_scores[1] + large = compiled_scores[request_capacity] + for request_count in range(1, request_capacity + 1): + use_extra_shard = request_count * num_layers * num_kv_heads * 2 < 2 * sm_count + compiled_scores[request_count] = small if use_extra_shard else large + page_shards_by_request_count[request_count] = ( + SMALL_WORKLOAD_PAGE_SHARDS if use_extra_shard else 2 + ) + + normalize_args: Tuple[object, ...] = () + compiled_normalizers: Dict[int, object] = {} + if union_scores is not None: + # Local import avoids a module cycle: selection imports the score + # module's shared layout constants. + from .triattention_cute_selection import ( + _select_normalize_union_config, + _TriAttentionNormalizeUnionKernel, + ) + + normalize_operands = ( + (partial_stats, 16), + (score_scratch, 16), + (source_lengths, 4), + (segment_output_offsets, 16), + (prompt_lengths, 4), + (union_scores, 16), + ) + normalize_args = tuple( + _to_cute(tensor, assumed_align=alignment) for tensor, alignment in normalize_operands + ) + for request_count in range(1, request_capacity + 1): + page_shards = page_shards_by_request_count[request_count] + config = _select_normalize_union_config(request_count, score_token_capacity, sm_count) + config_key = (page_shards, *config) + cache_key = ( + "triattention_cute_normalize_union", + static_geometry, + tensor_specs, + config_key, + _tensor_spec(union_scores), + ) + tokens_per_lane, token_subtiles, row_cluster_ctas = config + + def build_normalizer( + page_shards=page_shards, + tokens_per_lane=tokens_per_lane, + token_subtiles=token_subtiles, + row_cluster_ctas=row_cluster_ctas, + ): + return cute.compile( + _TriAttentionNormalizeUnionKernel( + num_layers=num_layers, + score_token_capacity=score_token_capacity, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + page_shards=page_shards, + tokens_per_lane=tokens_per_lane, + token_subtiles=token_subtiles, + row_cluster_ctas=row_cluster_ctas, + output_row_stride=int(union_scores.stride(0)), + ), + *normalize_args, + cutlass.Int32(1), + stream, + ) + + compiled_normalizers[request_count] = _get_or_compile(cache_key, build_normalizer) + + def launch_score(request_count: int) -> None: + current_stream = cuda.CUstream(torch.cuda.current_stream(device).cuda_stream) + compiled_scores[request_count](*score_args, request_count, current_stream) + if compiled_normalizers: + compiled_normalizers[request_count]( + *normalize_args, + request_count, + current_stream, + ) + + # The closure retains every DLPack-backed argument for the launcher's lifetime. + return score_scratch, launch_score diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py new file mode 100644 index 000000000000..0e89740e20a5 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -0,0 +1,537 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SM100 CuTe-DSL selection preparation for TriAttention scores.""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace +from cutlass.cute.typing import Int32 as CuteInt32 +from cutlass.cute.typing import Pointer as CutePointer +from cutlass.cutlass_dsl import T, dsl_user_op + +# Single-sourced constants: the score file owns N and the stats layout; Triton owns the epsilon. +from .triattention_cute_score_fused import PADDED_HEAD_COLUMNS, STATS_M2, STATS_MEAN +from .triattention_cute_score_fused import STATS_FIELDS as _STATS_FIELDS +from .triattention_kernels import STD_EPSILON as _STD_EPSILON + +_REDUCE_THREADS = 256 +_WARP_SIZE = 32 +_REDUCE_WARPS = _REDUCE_THREADS // _WARP_SIZE +_LARGE_TOKENS_PER_LANE = 4 +_LARGE_TOKEN_SUBTILES = 2 +_SMALL_TOKENS_PER_LANE = 2 +_SMALL_TOKEN_SUBTILES = 1 +_MAX_ROW_CLUSTER_CTAS = 4 +_SMALL_TILE_RESIDENT_CTAS_PER_SM = 6 + + +def _select_normalize_union_config( + request_count: int, + width: int, + sm_count: int, +) -> tuple[int, int, int]: + """Return (tokens_per_lane, token_subtiles, row_cluster_ctas).""" + row_cluster_ctas = max(1, _MAX_ROW_CLUSTER_CTAS // request_count) + small_token_tile = _WARP_SIZE * _SMALL_TOKENS_PER_LANE * _SMALL_TOKEN_SUBTILES + token_tiles = (width + small_token_tile - 1) // small_token_tile + grid_ctas = request_count * token_tiles * row_cluster_ctas + if grid_ctas <= sm_count * _SMALL_TILE_RESIDENT_CTAS_PER_SM: + return _SMALL_TOKENS_PER_LANE, _SMALL_TOKEN_SUBTILES, row_cluster_ctas + return _LARGE_TOKENS_PER_LANE, _LARGE_TOKEN_SUBTILES, 1 + + +@dsl_user_op +def _mapa_shared_cluster( + smem_ptr: CutePointer, + peer_rank: CuteInt32, + *, + loc=None, + ip=None, +) -> CuteInt32: + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [smem_ptr_i32, peer_rank.ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _ld_shared_cluster_f32( + mapped_addr: CuteInt32, + *, + loc=None, + ip=None, +) -> cutlass.Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [mapped_addr.ir_value(loc=loc, ip=ip)], + "ld.shared::cluster.f32 $0, [$1];", + "=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _gmem_lane_tile(iterator, flat_index, tokens_per_lane, assumed_align): + """One lane's fp32 gmem tile; folds the 64-bit index into the pointer before access.""" + return cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (iterator + flat_index).toint(), + AddressSpace.gmem, + assumed_align=assumed_align, + ), + cute.make_layout(tokens_per_lane), + ) + + +class _TriAttentionNormalizeUnionKernel: + """Merge row moments, normalize scores, and reduce their elementwise maximum.""" + + def __init__( + self, + *, + num_layers: int, + score_token_capacity: int, + num_q_heads: int, + num_kv_heads: int, + page_shards: int, + tokens_per_lane: int, + token_subtiles: int, + row_cluster_ctas: int, + output_row_stride: int, + ) -> None: + # Real head row q_head lives in score plane kv*8 + qg; partial-stats rows stay compact. + self.score_group_size = num_q_heads // num_kv_heads + self.score_head_pad = PADDED_HEAD_COLUMNS - self.score_group_size + self.num_layers = num_layers + self.score_token_capacity = score_token_capacity + self.output_row_stride = output_row_stride + self.num_q_heads = num_q_heads + self.num_rows = num_layers * num_q_heads + self.page_shards = page_shards + self.tokens_per_lane = tokens_per_lane + self.token_subtiles = token_subtiles + self.subtile_token_tile = _WARP_SIZE * self.tokens_per_lane + self.token_tile = self.subtile_token_tile * self.token_subtiles + self.reduce_threads = _REDUCE_THREADS + self.reduce_warps = _REDUCE_WARPS + self.row_cluster_ctas = row_cluster_ctas + # The widest score window (the whole bucket) sizes the token-tile grid; + # output rows are the TopK selection rows with their own stride. + self.num_token_tiles = (self.score_token_capacity + self.token_tile - 1) // self.token_tile + + @cute.jit + def __call__( + self, + partial_stats: cute.Tensor, + scores: cute.Tensor, + source_lengths: cute.Tensor, + seg_out_offset: cute.Tensor, + prompt_lengths: cute.Tensor, + union_scores: cute.Tensor, + request_count: cutlass.Int32, + stream: cuda.CUstream, + ): + kernel = self.kernel( + partial_stats, + scores, + source_lengths, + seg_out_offset, + prompt_lengths, + union_scores, + request_count, + ) + if cutlass.const_expr(self.row_cluster_ctas == 1): + kernel.launch( + grid=(request_count, self.num_token_tiles, 1), + block=(self.reduce_threads, 1, 1), + stream=stream, + ) + else: + # Cluster peers are consecutive CTAs; the kernel decode must keep this factor order. + kernel.launch( + grid=( + request_count * self.num_token_tiles * self.row_cluster_ctas, + 1, + 1, + ), + block=(self.reduce_threads, 1, 1), + cluster=(self.row_cluster_ctas, 1, 1), + stream=stream, + ) + + @cute.jit + def _reduce_and_store_union_rows( + self, + union_scores: cute.Tensor, + union_values: cute.Tensor, + warp_max: cute.Tensor, + warp_max_ptr, + score_copy_atom, + request_idx: cutlass.Int32, + decode_length: cutlass.Int32, + first_token: cutlass.Int32, + lane_idx: cutlass.Int32, + from_cluster_peers: cutlass.Constexpr, + ): + """Final peer reduce and union-row store (the peer source is picked at trace time).""" + for token_subtile in cutlass.range_constexpr(self.token_subtiles): + reduced_values = cute.make_rmem_tensor((self.tokens_per_lane,), cutlass.Float32) + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + if cutlass.const_expr(from_cluster_peers): + union_value = warp_max[(0, token_subtile, token_slot, lane_idx)] + # Offset derived from the same layout the smem tensor was built with. + shared_offset = cute.crd2idx( + (0, token_subtile, token_slot, lane_idx), warp_max.layout + ) + for peer_rank in cutlass.range_constexpr(1, self.row_cluster_ctas): + remote_addr = _mapa_shared_cluster(warp_max_ptr, cutlass.Int32(peer_rank)) + union_value = cute.arch.fmax( + union_value, + _ld_shared_cluster_f32( + remote_addr + shared_offset * (cutlass.Float32.width // 8) + ), + ) + else: + union_value = union_values[(token_subtile, token_slot)] + for other_warp in cutlass.range_constexpr(1, self.reduce_warps): + union_value = cute.arch.fmax( + union_value, + warp_max[(other_warp, token_subtile, token_slot, lane_idx)], + ) + reduced_values[token_slot] = union_value + subtile_first_token = first_token + token_subtile * self.subtile_token_tile + # Straddling subtiles store per token; the selection rows stay < 2^31 so i32 cannot wrap. + if cutlass.const_expr( + self.output_row_stride % self.tokens_per_lane == 0 + ) and cutlass.dynamic_expr(subtile_first_token + self.tokens_per_lane <= decode_length): + union_index = request_idx * self.output_row_stride + subtile_first_token + union_tile = _gmem_lane_tile( + union_scores.iterator, + union_index, + self.tokens_per_lane, + self.tokens_per_lane * 4, + ) + cute.copy( + score_copy_atom, + cute.coalesce(reduced_values), + cute.coalesce(union_tile), + ) + else: + union_index = request_idx * self.output_row_stride + subtile_first_token + union_tile = _gmem_lane_tile( + union_scores.iterator, + union_index, + self.tokens_per_lane, + 4, + ) + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + token = subtile_first_token + token_slot + if cutlass.dynamic_expr(token < decode_length): + union_tile[token_slot] = reduced_values[token_slot] + + @cute.kernel + def kernel( + self, + partial_stats: cute.Tensor, + scores: cute.Tensor, + source_lengths: cute.Tensor, + seg_out_offset: cute.Tensor, + prompt_lengths: cute.Tensor, + union_scores: cute.Tensor, + request_count: cutlass.Int32, + ): + tidx, _, _ = cute.arch.thread_idx() + block_idx_x, block_idx_y, _ = cute.arch.block_idx() + cta_rank = cutlass.Int32(0) + if cutlass.const_expr(self.row_cluster_ctas == 1): + request_idx = block_idx_x + token_tile_idx = block_idx_y + else: + cta_rank = cute.arch.block_idx_in_cluster() + cluster_idx = block_idx_x // self.row_cluster_ctas + request_idx = cluster_idx // self.num_token_tiles + token_tile_idx = cluster_idx - request_idx * self.num_token_tiles + warp_idx = tidx // _WARP_SIZE + lane_idx = tidx % _WARP_SIZE + first_token = token_tile_idx * self.token_tile + lane_idx * self.tokens_per_lane + first_segment = request_idx * self.num_layers + # The normalization domain and the union output row both cover [0, valid - start). + score_start = cutlass.Int32(prompt_lengths[request_idx]) + decode_length = source_lengths[request_idx] - score_start + warp_max_ptr = cute.arch.alloc_smem( + cutlass.Float32, + self.reduce_threads * self.tokens_per_lane * self.token_subtiles, + ) + warp_max = cute.make_tensor( + warp_max_ptr, + cute.make_layout( + ( + self.reduce_warps, + self.token_subtiles, + self.tokens_per_lane, + _WARP_SIZE, + ), + stride=( + self.token_subtiles * self.tokens_per_lane * _WARP_SIZE, + self.tokens_per_lane * _WARP_SIZE, + _WARP_SIZE, + 1, + ), + ), + ) + union_values = cute.make_rmem_tensor( + (self.token_subtiles, self.tokens_per_lane), + cutlass.Float32, + ) + score_value_tiles = tuple( + cute.make_rmem_tensor((self.tokens_per_lane,), cutlass.Float32) + for _ in range(self.token_subtiles) + ) + score_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=self.tokens_per_lane * cutlass.Float32.width, + ) + for token_subtile in cutlass.range_constexpr(self.token_subtiles): + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + union_values[(token_subtile, token_slot)] = cutlass.Float32(float("-inf")) + + common_count = cutlass.Float32(0.0) + mean_weight_1 = cutlass.Float32(0.0) + m2_cross_weight = cutlass.Float32(0.0) + shard_mean_weights = cute.make_rmem_tensor((self.page_shards,), cutlass.Float32) + shard_m2_cross_weights = cute.make_rmem_tensor((self.page_shards,), cutlass.Float32) + if cutlass.const_expr(self.page_shards == 2): + first_stats_row = first_segment * self.num_q_heads + first_stats_base = first_stats_row * 2 * _STATS_FIELDS + count_0 = partial_stats[first_stats_base] + count_1 = partial_stats[first_stats_base + _STATS_FIELDS] + common_count = count_0 + count_1 + if cutlass.dynamic_expr(common_count > 0.0): + mean_weight_1 = count_1 / common_count + m2_cross_weight = count_0 * count_1 / common_count + else: + first_stats_row = first_segment * self.num_q_heads + first_stats_base = first_stats_row * self.page_shards * _STATS_FIELDS + for page_shard in cutlass.range_constexpr(self.page_shards): + shard_count = partial_stats[first_stats_base + page_shard * _STATS_FIELDS] + merged_count = common_count + shard_count + mean_weight = cutlass.Float32(0.0) + m2_cross = cutlass.Float32(0.0) + if cutlass.dynamic_expr(merged_count > 0.0): + mean_weight = shard_count / merged_count + m2_cross = common_count * shard_count / merged_count + shard_mean_weights[page_shard] = mean_weight + shard_m2_cross_weights[page_shard] = m2_cross + common_count = merged_count + + first_logical_row = warp_idx + cta_rank * self.reduce_warps + logical_row_stride = self.reduce_warps * self.row_cluster_ctas + for logical_row in cutlass.range( + first_logical_row, + self.num_rows, + logical_row_stride, + unroll=1, + ): + layer_slot = logical_row // self.num_q_heads + q_head = logical_row - layer_slot * self.num_q_heads + segment = first_segment + layer_slot + stats_row = segment * self.num_q_heads + q_head + # Map the real head row onto its padded score plane; padded planes are never visited. + score_plane = q_head + if cutlass.const_expr(self.score_head_pad > 0): + score_plane = q_head + (q_head // self.score_group_size) * self.score_head_pad + + count = cutlass.Float32(0.0) + mean = cutlass.Float32(0.0) + m2 = cutlass.Float32(0.0) + delta = cutlass.Float32(0.0) + if cutlass.const_expr(self.page_shards == 2): + stats_base = stats_row * 2 * _STATS_FIELDS + mean_0 = partial_stats[stats_base + STATS_MEAN] + m2_0 = partial_stats[stats_base + STATS_M2] + mean_1 = partial_stats[stats_base + _STATS_FIELDS + STATS_MEAN] + m2_1 = partial_stats[stats_base + _STATS_FIELDS + STATS_M2] + delta = mean_1 - mean_0 + count = common_count + mean = mean_0 + delta * mean_weight_1 + m2 = m2_0 + m2_1 + delta * delta * m2_cross_weight + else: + count = common_count + for page_shard in cutlass.range_constexpr(self.page_shards): + stats_base = (stats_row * self.page_shards + page_shard) * _STATS_FIELDS + shard_mean = partial_stats[stats_base + STATS_MEAN] + shard_m2 = partial_stats[stats_base + STATS_M2] + delta = shard_mean - mean + mean = mean + delta * shard_mean_weights[page_shard] + m2 = m2 + shard_m2 + delta * delta * shard_m2_cross_weights[page_shard] + inv_std = cutlass.Float32(0.0) + if cutlass.dynamic_expr(count > 0.0): + variance = m2 / count + if cutlass.dynamic_expr(variance < _STD_EPSILON * _STD_EPSILON): + inv_std = cutlass.Float32(1.0 / _STD_EPSILON) + else: + inv_std = cute.math.rsqrt(variance) + + for token_subtile in cutlass.range_constexpr(self.token_subtiles): + subtile_first_token = first_token + token_subtile * self.subtile_token_tile + score_index = ( + cutlass.Int64(score_plane) + * request_count + * self.num_layers + * self.score_token_capacity + + seg_out_offset[segment] + + score_start + + subtile_first_token + ) + # The vectorized load needs the runtime start aligned to the lane width. + if cutlass.const_expr( + self.score_token_capacity % self.tokens_per_lane == 0 + ) and cutlass.dynamic_expr( + score_start % self.tokens_per_lane == 0 + and subtile_first_token + self.tokens_per_lane <= decode_length + ): + score_tile = _gmem_lane_tile( + scores.iterator, + score_index, + self.tokens_per_lane, + self.tokens_per_lane * 4, + ) + cute.copy( + score_copy_atom, + cute.coalesce(score_tile), + cute.coalesce(score_value_tiles[token_subtile]), + ) + else: + # Fold the 64-bit index into the pointer; the scratch exceeds 2^31 elements. + score_tail = _gmem_lane_tile( + scores.iterator, score_index, self.tokens_per_lane, 4 + ) + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + token = subtile_first_token + token_slot + if cutlass.dynamic_expr(token < decode_length): + score_value_tiles[token_subtile][token_slot] = score_tail[token_slot] + else: + score_value_tiles[token_subtile][token_slot] = cutlass.Float32( + float("-inf") + ) + + for token_subtile in cutlass.range_constexpr(self.token_subtiles): + if cutlass.const_expr(self.tokens_per_lane >= 2): + normalized_01 = cute.arch.sub_packed_f32x2( + ( + score_value_tiles[token_subtile][0], + score_value_tiles[token_subtile][1], + ), + (mean, mean), + ) + normalized_01 = cute.arch.mul_packed_f32x2( + normalized_01, + (inv_std, inv_std), + ) + union_values[(token_subtile, 0)] = cute.arch.fmax( + union_values[(token_subtile, 0)], normalized_01[0] + ) + union_values[(token_subtile, 1)] = cute.arch.fmax( + union_values[(token_subtile, 1)], normalized_01[1] + ) + if cutlass.const_expr(self.tokens_per_lane == 4): + normalized_23 = cute.arch.sub_packed_f32x2( + ( + score_value_tiles[token_subtile][2], + score_value_tiles[token_subtile][3], + ), + (mean, mean), + ) + normalized_23 = cute.arch.mul_packed_f32x2( + normalized_23, + (inv_std, inv_std), + ) + union_values[(token_subtile, 2)] = cute.arch.fmax( + union_values[(token_subtile, 2)], normalized_23[0] + ) + union_values[(token_subtile, 3)] = cute.arch.fmax( + union_values[(token_subtile, 3)], normalized_23[1] + ) + for token_subtile in cutlass.range_constexpr(self.token_subtiles): + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + warp_max[(warp_idx, token_subtile, token_slot, lane_idx)] = union_values[ + (token_subtile, token_slot) + ] + cute.arch.sync_threads() + if cutlass.const_expr(self.row_cluster_ctas == 1): + if warp_idx == 0: + self._reduce_and_store_union_rows( + union_scores, + union_values, + warp_max, + warp_max_ptr, + score_copy_atom, + request_idx, + decode_length, + first_token, + lane_idx, + False, + ) + else: + # Warp 0 reduces its CTA's rows; CTA 0 combines cluster maxima via distributed smem. + if warp_idx == 0: + for token_subtile in cutlass.range_constexpr(self.token_subtiles): + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + union_value = union_values[(token_subtile, token_slot)] + for other_warp in cutlass.range_constexpr(1, self.reduce_warps): + union_value = cute.arch.fmax( + union_value, + warp_max[(other_warp, token_subtile, token_slot, lane_idx)], + ) + warp_max[(0, token_subtile, token_slot, lane_idx)] = union_value + cute.arch.sync_threads() + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + if cta_rank == 0 and warp_idx == 0: + self._reduce_and_store_union_rows( + union_scores, + union_values, + warp_max, + warp_max_ptr, + score_copy_atom, + request_idx, + decode_length, + first_token, + lane_idx, + True, + ) + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py new file mode 100644 index 000000000000..8a923bc24d6e --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -0,0 +1,437 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Triton reduction, TP-fold, and selection kernels for TriAttention.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +# ---- Mean-phase gather: per-request phase-row fetch + width derivation ---- + +# Score z-normalization epsilon; must stay a plain float (the CuTe DSL traces it). +STD_EPSILON = 1e-6 + + +@triton.jit +def _gather_mean_phase_kernel( + logical_source_lengths, + phase_cos, + phase_sin, + source_lengths, + prompt_lengths, + mean_cos, + mean_sin, + decode_lengths, + swa_destination_bases, + swa_rebase_delta, + NUM_FREQS: tl.constexpr, + F_BLOCK: tl.constexpr, + HAS_SWA: tl.constexpr, +): + """Copy each request's phase-table row; derive decode lengths and SWA landing bases.""" + request = tl.program_id(0) + frequency = tl.arange(0, F_BLOCK) + frequency_mask = frequency < NUM_FREQS + table_row = tl.load(logical_source_lengths + request).to(tl.int64) + source_offset = table_row * NUM_FREQS + frequency + output_offset = request * NUM_FREQS + frequency + row_cos = tl.load(phase_cos + source_offset, mask=frequency_mask, other=0.0) + row_sin = tl.load(phase_sin + source_offset, mask=frequency_mask, other=0.0) + tl.store(mean_cos + output_offset, row_cos, mask=frequency_mask) + tl.store(mean_sin + output_offset, row_sin, mask=frequency_mask) + prompt_length = tl.load(prompt_lengths + request) + tl.store(decode_lengths + request, tl.load(source_lengths + request) - prompt_length) + if HAS_SWA: + tl.store(swa_destination_bases + request, prompt_length + swa_rebase_delta) + + +def gather_mean_phase( + logical_source_lengths: torch.Tensor, + phase_cos: torch.Tensor, + phase_sin: torch.Tensor, + source_lengths: torch.Tensor, + prompt_lengths: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + decode_lengths: torch.Tensor, + swa_destination_bases: torch.Tensor | None, + *, + request_count: int, + swa_rebase_delta: int, +) -> None: + """Gather mean-phase rows and derive per-request decode metadata.""" + num_freqs = int(phase_cos.shape[1]) + _gather_mean_phase_kernel[(request_count,)]( + logical_source_lengths, + phase_cos, + phase_sin, + source_lengths, + prompt_lengths, + mean_cos, + mean_sin, + decode_lengths, + swa_destination_bases, + swa_rebase_delta, + NUM_FREQS=num_freqs, + F_BLOCK=triton.next_power_of_2(num_freqs), + HAS_SWA=swa_destination_bases is not None, + num_warps=1, + ) + + +# ---- Selection: combine scores per mode, then finalize the top-k set ---- + + +@triton.jit +def _score_row_stats_kernel( + score_scratch, + decode_lengths, + prompt_lengths, + row_mean, + row_inv_std, + segment_tokens, + ROWS: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + PADDED_COLUMNS: tl.constexpr, + BUCKET: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr = 256, + # Triton rejects plain-global capture; the default binds the module float + # at def time (STD_EPSILON itself must stay a plain float for the CuTe import). + EPSILON: tl.constexpr = STD_EPSILON, +): + """Compute decode-window mean and inverse standard deviation for each score row.""" + QUERY_GROUP_SIZE: tl.constexpr = NUM_Q_HEADS // NUM_KV_HEADS + flat_row = tl.program_id(0) + request = flat_row // ROWS + row_in_request = flat_row % ROWS + layer = row_in_request // NUM_Q_HEADS + query_head = row_in_request % NUM_Q_HEADS + kv_head = query_head // QUERY_GROUP_SIZE + plane = kv_head * PADDED_COLUMNS + query_head % QUERY_GROUP_SIZE + decode_length = tl.load(decode_lengths + request) + prompt_start = tl.load(prompt_lengths + request) + score_row = ( + score_scratch + + plane.to(tl.int64) * segment_tokens + + ((request * NUM_LAYERS + layer) * BUCKET + prompt_start).to(tl.int64) + ) + lane = tl.arange(0, BLOCK) + score_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < decode_length + value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) + score_sum += tl.sum(value, axis=0) + mean = score_sum / decode_length + square_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < decode_length + value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) + centered = tl.where(valid, value - mean, 0.0) + square_sum += tl.sum(centered * centered, axis=0) + std = tl.sqrt(square_sum / decode_length) + tl.store(row_mean + flat_row, mean) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, EPSILON)) + + +@triton.jit +def _score_per_head_reduce_kernel( + score_scratch, + decode_lengths, + prompt_lengths, + row_mean, + row_inv_std, + selection_scores, + selection_row_lengths, + segment_tokens, + NUM_LAYERS: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + PADDED_COLUMNS: tl.constexpr, + BUCKET: tl.constexpr, + WIDTH: tl.constexpr, + PER_LAYER: tl.constexpr, + NORMALIZE: tl.constexpr, + BLOCK: tl.constexpr = 256, +): + """Reduce each KV-head decode window from score scratch into a selector row.""" + QUERY_GROUP_SIZE: tl.constexpr = NUM_Q_HEADS // NUM_KV_HEADS + SELECTION_ROWS: tl.constexpr = NUM_LAYERS * NUM_KV_HEADS if PER_LAYER else NUM_KV_HEADS + request = tl.program_id(0) + selection_row = tl.program_id(1) + token_block = tl.program_id(2) + token = token_block * BLOCK + tl.arange(0, BLOCK) + decode_length = tl.load(decode_lengths + request) + prompt_start = tl.load(prompt_lengths + request) + valid_token = token < decode_length + + if token_block == 0: + tl.store( + selection_row_lengths + request * SELECTION_ROWS + selection_row, + decode_length, + ) + + kv_head = selection_row % NUM_KV_HEADS + if PER_LAYER: + layer = selection_row // NUM_KV_HEADS + reduced = tl.full((BLOCK,), -float("inf"), tl.float32) + for query_in_group in tl.static_range(0, QUERY_GROUP_SIZE): + query_head = kv_head * QUERY_GROUP_SIZE + query_in_group + flat_row = (request * NUM_LAYERS + layer) * NUM_Q_HEADS + query_head + plane = kv_head * PADDED_COLUMNS + query_in_group + value = tl.load( + score_scratch + + plane.to(tl.int64) * segment_tokens + + ((request * NUM_LAYERS + layer) * BUCKET + prompt_start).to(tl.int64) + + token, + mask=valid_token, + other=-float("inf"), + ).to(tl.float32) + if NORMALIZE: + mean = tl.load(row_mean + flat_row) + inv_std = tl.load(row_inv_std + flat_row) + value = tl.where(valid_token, (value - mean) * inv_std, -float("inf")) + reduced = tl.maximum(reduced, value) + else: + reduced = tl.zeros((BLOCK,), tl.float32) + for layer in tl.static_range(0, NUM_LAYERS): + layer_max = tl.full((BLOCK,), -float("inf"), tl.float32) + for query_in_group in tl.static_range(0, QUERY_GROUP_SIZE): + query_head = kv_head * QUERY_GROUP_SIZE + query_in_group + flat_row = (request * NUM_LAYERS + layer) * NUM_Q_HEADS + query_head + plane = kv_head * PADDED_COLUMNS + query_in_group + value = tl.load( + score_scratch + + plane.to(tl.int64) * segment_tokens + + ((request * NUM_LAYERS + layer) * BUCKET + prompt_start).to(tl.int64) + + token, + mask=valid_token, + other=-float("inf"), + ).to(tl.float32) + if NORMALIZE: + mean = tl.load(row_mean + flat_row) + inv_std = tl.load(row_inv_std + flat_row) + value = tl.where(valid_token, (value - mean) * inv_std, -float("inf")) + layer_max = tl.maximum(layer_max, value) + reduced += layer_max + reduced /= NUM_LAYERS + + output = (request * SELECTION_ROWS + selection_row) * WIDTH + token + tl.store(selection_scores + output, reduced, mask=token < WIDTH) + + +def reduce_per_head_scores( + score_scratch: torch.Tensor, + decode_lengths: torch.Tensor, + prompt_lengths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + selection_scores_rows: torch.Tensor, + selection_row_lengths: torch.Tensor, + *, + request_count: int, + padded_head_columns: int, + score_token_capacity: int, + per_layer: bool, + normalize_scores: bool, +) -> None: + """Reduce score-scratch decode windows into per-head selection rows.""" + request_capacity = int(decode_lengths.numel()) + num_layers = int(row_mean.shape[1]) + num_q_heads = int(row_mean.shape[2]) + selection_rows = int(selection_scores_rows.shape[0]) // request_capacity + num_kv_heads = selection_rows // num_layers if per_layer else selection_rows + selection_width = int(selection_scores_rows.shape[1]) + rows = num_layers * num_q_heads + segment_tokens = request_count * num_layers * score_token_capacity + if normalize_scores: + _score_row_stats_kernel[(request_count * rows,)]( + score_scratch, + decode_lengths, + prompt_lengths, + row_mean, + row_inv_std, + segment_tokens, + ROWS=rows, + NUM_LAYERS=num_layers, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + PADDED_COLUMNS=padded_head_columns, + BUCKET=score_token_capacity, + WIDTH=selection_width, + ) + # 256-token tiles match the reduce kernel's BLOCK default. + _score_per_head_reduce_kernel[ + (request_count, selection_rows, triton.cdiv(selection_width, 256)) + ]( + score_scratch, + decode_lengths, + prompt_lengths, + row_mean, + row_inv_std, + selection_scores_rows, + selection_row_lengths, + segment_tokens, + NUM_LAYERS=num_layers, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + PADDED_COLUMNS=padded_head_columns, + BUCKET=score_token_capacity, + WIDTH=selection_width, + PER_LAYER=per_layer, + NORMALIZE=normalize_scores, + ) + + +@triton.jit +def _fold_union_ranks_kernel( + gathered_rows, + selection_scores_rows, + request_count, + TP_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr = 1024, +): + """Max-fold TP-gathered rank-local rows into each global union row.""" + request = tl.program_id(0) + token_block = tl.program_id(1) + token = token_block * BLOCK + tl.arange(0, BLOCK) + mask = token < WIDTH + folded = tl.full((BLOCK,), -float("inf"), tl.float32) + for rank in tl.static_range(0, TP_SIZE): + value = tl.load( + gathered_rows + (rank * request_count + request) * WIDTH + token, + mask=mask, + other=-float("inf"), + ) + folded = tl.maximum(folded, value) + tl.store(selection_scores_rows + request * WIDTH + token, folded, mask=mask) + + +def fold_union_ranks( + gathered_rows: torch.Tensor, + selection_scores_rows: torch.Tensor, + *, + request_count: int, +) -> None: + """Max-fold TP rank-local score rows into the global union rows.""" + width = int(selection_scores_rows.shape[1]) + tp_size = int(gathered_rows.shape[0]) // request_count + block = 1024 + _fold_union_ranks_kernel[(request_count, triton.cdiv(width, block))]( + gathered_rows, + selection_scores_rows, + request_count, + TP_SIZE=tp_size, + WIDTH=width, + BLOCK=block, + ) + + +@triton.jit +def _settle_ties_kernel( + selection_scores_rows, + selection_row_lengths, + prompt_lengths, + provisional_rows, + kept_ordinal_rows, + WIDTH: tl.constexpr, + KEEP_COUNT: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + BLOCK: tl.constexpr = 256, +): + """Settle score ties by lowest index and sort the kept-token indices.""" + request = tl.program_id(0) + selection_domain = tl.program_id(1) + row = request * SELECTION_ROWS + selection_domain + row_output = kept_ordinal_rows + row * KEEP_COUNT + row_scores = selection_scores_rows + row * WIDTH + row_selected = provisional_rows + row * KEEP_COUNT + # Rebases the decode-relative ordinals to absolute positions (per request: + # every selection row of a request shares its pinned prompt length). + prompt_length = tl.load(prompt_lengths + request) + + threshold = float("inf") + for start in tl.static_range(0, KEEP_COUNT, BLOCK): + selected_offset = start + tl.arange(0, BLOCK) + selected_mask = selected_offset < KEEP_COUNT + token_index = tl.load( + row_selected + selected_offset, + mask=selected_mask, + other=0, + ) + # Mask the top-k's -1 pad sentinels so no lane dereferences ``row_scores - 1``. + selected_valid = selected_mask & (token_index >= 0) + selected_score = tl.load( + row_scores + token_index, + mask=selected_valid, + other=float("inf"), + ).to(tl.float32) + threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) + + row_length = tl.load(selection_row_lengths + row) + greater_count = 0 + for start in tl.static_range(0, WIDTH, BLOCK): + token_index = start + tl.arange(0, BLOCK) + valid = (token_index < WIDTH) & (token_index < row_length) + score = tl.load( + row_scores + token_index, + mask=valid, + other=float("-inf"), + ).to(tl.float32) + greater_count += tl.sum((valid & (score > threshold)).to(tl.int32)) + + tie_quota = KEEP_COUNT - greater_count + output_count = 0 + ties_seen = 0 + for start in tl.static_range(0, WIDTH, BLOCK): + token_index = start + tl.arange(0, BLOCK) + valid = (token_index < WIDTH) & (token_index < row_length) + score = tl.load( + row_scores + token_index, + mask=valid, + other=float("-inf"), + ).to(tl.float32) + greater = valid & (score > threshold) + tied = valid & (score == threshold) + tied_i32 = tied.to(tl.int32) + tie_rank = ties_seen + tl.cumsum(tied_i32, axis=0) - tied_i32 + selected = greater | (tied & (tie_rank < tie_quota)) + selected_i32 = selected.to(tl.int32) + write_offset = output_count + tl.cumsum(selected_i32, axis=0) - selected_i32 + tl.store( + row_output + write_offset, + token_index + prompt_length, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) + + +def settle_ties( + selection_scores_rows: torch.Tensor, + selection_row_lengths: torch.Tensor, + prompt_lengths: torch.Tensor, + provisional_rows: torch.Tensor, + kept_ordinal_rows: torch.Tensor, + *, + request_count: int, + selection_rows_per_request: int, +) -> None: + """Settle TopK score ties into ascending absolute token ordinals.""" + _settle_ties_kernel[(request_count, selection_rows_per_request)]( + selection_scores_rows, + selection_row_lengths, + prompt_lengths, + provisional_rows, + kept_ordinal_rows, + WIDTH=int(selection_scores_rows.shape[1]), + KEEP_COUNT=int(kept_ordinal_rows.shape[1]), + SELECTION_ROWS=selection_rows_per_request, + ) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index e902a892af3e..53ab4b86bd44 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -636,6 +636,14 @@ def __init__( key="disable_rope_fusion_for_rocketkv") self.rope_fusion = False + if (config.kv_cache_compression_config is not None and + config.kv_cache_compression_config.changes_physical_kv_length): + logger.warning_once( + "KV-cache eviction changes the physical cache length; " + "setting rope_fusion=False.", + key="disable_rope_fusion_for_kv_cache_compression") + self.rope_fusion = False + if self.rope_fusion and not attn_cls.support_fused_rope(): logger.warning_once( "rope_fusion is true but the attention backend does not support it. Will disable rope_fusion.", diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9c0023cac331..e08bb27938f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -22,8 +22,8 @@ import tensorrt_llm import tensorrt_llm.bindings.executor as trtllm from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, - prefer_pinned, str_dtype_to_binding, - torch_dtype_to_str) + is_sm_100f, prefer_pinned, + str_dtype_to_binding, torch_dtype_to_str) from tensorrt_llm.bindings.executor import DecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams @@ -2292,24 +2292,29 @@ def _create_kv_cache_manager( return kv_cache_manager -def validate_kv_cache_compression_with_spec( +def validate_kv_cache_compression_compatibility( config: KvCacheCompressionConfig, + kv_cache_config: KvCacheConfig, spec_config: Optional[SpeculativeConfig], - draft_kv_cache_manager: Optional[KVCacheManagerV2], ) -> None: - """Reject speculative setups the compression method cannot run with.""" - if (spec_config is None - or not config.kv_cache_compression_mode.is_eviction_method()): + """Reject unsupported KV-cache compression feature combinations.""" + if kv_cache_config.enable_block_reuse and not config.supports_block_reuse(): + raise ValueError( + f"KV-cache compression algorithm {config.algorithm!r} does not " + "support KV-cache block reuse. Set " + "KvCacheConfig.enable_block_reuse=False.") + if spec_config is None: return - # Evicting methods co-compact the draft KV, so the draft must be a - # standard paged cache in the same forward (one-model speculation). + if not config.supports_speculative_decoding(): + raise ValueError( + f"KV-cache compression algorithm {config.algorithm!r} does not " + "support speculative decoding with its current configuration; " + "TriAttention requires eviction_mode='union'") mode = spec_config.spec_dec_mode if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): raise ValueError( - f"KV-cache compression algorithm {config.algorithm!r} does not " - f"support speculative decoding mode {mode.name}: the draft KV " - "must be a standard paged cache compacted together with the " - "target (one-model MTP/EAGLE3).") + f"KV-cache compression does not support speculative decoding " + f"mode {mode.name}; use one-model MTP or EAGLE3") def create_kv_cache_compression_manager( @@ -2322,9 +2327,23 @@ def create_kv_cache_compression_manager( Called from ``create_py_executor`` and registered as a resource manager, like the KV cache manager itself. Concrete algorithms add a dispatch branch - here; the framework ships none. Speculative-decoding compatibility is - checked by the caller via ``validate_kv_cache_compression_with_spec``. + here. Feature compatibility is checked before resource-manager construction. """ + if config.algorithm == "triattention": + if not is_sm_100f(): + raise RuntimeError( + "TriAttention requires an SM100-family device (SM100 or SM103)." + ) + # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. + from ..kv_cache_compression.triattention.triattention import \ + TriAttentionCompressionManager + + return TriAttentionCompressionManager( + config, + kv_cache_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " "a compression manager.", @@ -2579,9 +2598,6 @@ def create_py_executor_instance( if kv_cache_compression_config is not None: draft_kv_cache_manager = resources.get( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) - validate_kv_cache_compression_with_spec(kv_cache_compression_config, - spec_config, - draft_kv_cache_manager) compression_manager = create_kv_cache_compression_manager( kv_cache_compression_config, kv_cache_manager, @@ -3062,6 +3078,14 @@ def _adjust_torch_mem_fraction(): def validate_feature_combination(llm_args, model_engine, sampler_type): # Validate the flags for features' combination + compression_config = llm_args.kv_cache_compression_config + if compression_config is not None: + validate_kv_cache_compression_compatibility( + compression_config, + llm_args.kv_cache_config, + model_engine.spec_config, + ) + def init_feature_status(llm_args) -> Dict[str, bool]: assert isinstance( llm_args, TorchLlmArgs diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index b4faf969d0da..7500b6301036 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -19,8 +19,8 @@ from abc import ABC, abstractmethod from collections import OrderedDict, defaultdict, deque from dataclasses import dataclass -from typing import (TYPE_CHECKING, ClassVar, Dict, Iterable, List, Optional, - Sequence, Set, Tuple, Union) +from typing import (TYPE_CHECKING, Dict, Iterable, List, Optional, Sequence, + Set, Tuple, Union) import torch from mpi4py import MPI @@ -66,7 +66,8 @@ if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.interface import \ AttentionMetadata - from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + KvCacheCompressionConfig) from .kv_cache_manager_v2 import KVCacheManagerV2 @@ -2448,11 +2449,9 @@ class KVCacheCompressionManager(BaseResourceManager): engine subtracts that count when building ``num_cached_tokens_per_seq``. """ - adjusts_generation_kv_length: ClassVar[bool] = False - """Whether this manager can make target and logical KV lengths diverge.""" - def __init__( self, + config: "KvCacheCompressionConfig", kv_cache_manager: "KVCacheManagerV2", draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None, ): @@ -2466,18 +2465,12 @@ def __init__( "draft KV-cache compression requires KVCacheManagerV2") self.kv_cache_manager = kv_cache_manager self.draft_kv_cache_manager = draft_kv_cache_manager - # Compression evicts/rewrites stored keys and values, so a shared prefix - # block is no longer safe to reuse (same constraint as RocketKVCacheManager). - if kv_cache_manager.enable_block_reuse: - raise ValueError( - f"{type(self).__name__} changes stored keys and values and cannot " - f"run with KV-cache block reuse. Set " - f"KvCacheConfig.enable_block_reuse to False.") - kv_cache_manager.kv_compression_manages_history = self.adjusts_generation_kv_length + kv_cache_manager.kv_compression_manages_history = ( + config.changes_physical_kv_length) if draft_kv_cache_manager is not None: # The draft cache is compacted together with the target. draft_kv_cache_manager.kv_compression_manages_history = ( - self.adjusts_generation_kv_length) + config.changes_physical_kv_length) @property def has_independent_draft_kv_cache(self) -> bool: diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index e0967a5cc243..7a3583907041 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -25,6 +25,7 @@ SAEnhancerConfig, SaveHiddenStatesDecodingConfig, SchedulerConfig, SkipSoftmaxAttentionConfig, TorchCompileConfig, TorchLlmArgs, + TriAttentionKvCacheCompressionConfig, UserProvidedDecodingConfig) from .llm_utils import KvCacheRetentionConfig, QuantAlgo, QuantConfig from .mm_encoder import MultimodalEncoder @@ -89,6 +90,7 @@ 'MiniMaxM3SparseAttentionConfig', 'SchedulingParams', 'SkipSoftmaxAttentionConfig', + 'TriAttentionKvCacheCompressionConfig', 'PrometheusMetricsConfig', 'ThinkingBudgetLogitsProcessor', 'add_thinking_budget_logits_processor', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index ab5f561070c6..cb22ec639dab 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3561,19 +3561,83 @@ class KvCacheCompressionConfig(StrictBaseModel): as a resource manager in create_py_executor (_util.py), like the KV cache manager itself. Concrete algorithms subclass this and add their parameters. """ + + changes_physical_kv_length: ClassVar[bool] = False + """Whether physical and logical KV lengths can diverge.""" + algorithm: str = Field( description= "Name of the KV-cache compression algorithm to run; selects which " "compression manager is built. Concrete algorithm configs subclass this " "and set the value.") - @property - def kv_cache_compression_mode(self): - # The mode carries algorithm-level traits (``is_*`` predicates) the - # raw algorithm string does not. - from tensorrt_llm._torch.kv_cache_compression.interface import \ - KvCacheCompressionMode - return KvCacheCompressionMode.from_string(self.algorithm) + def supports_block_reuse(self) -> bool: + return False + + def supports_speculative_decoding(self) -> bool: + return False + + +class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): + """TriAttention KV-cache compression: periodic decode-time eviction. + + Scored by offline calibration (github.com/WeianMao/triattention; supply + the official .pt via ``calibration_path``). Pure compression — decode + runs the model's standard attention over the compacted cache. + """ + + changes_physical_kv_length: ClassVar[bool] = True + + algorithm: Literal["triattention"] = "triattention" + eviction_mode: Literal["union", "per_head", "per_layer_perhead"] = Field( + default="union", + description= + "Which token set each eviction round keeps. `union` (default) takes " + "the union of each KV head's top-B and re-ranks it by the per-token max " + "score; it matches the official base setting (per-head and " + "per-layer-per-head pruning both off). `per_head` keeps a per-KV-head " + "set shared across layers (mean of per-layer max); `per_layer_perhead` " + "keeps a fully independent set per (layer, KV head).") + normalize_scores: bool = Field( + default=True, + description="Z-normalize each head's scores over the decode region " + "before selection (upstream default). `union` eviction requires True: " + "its fused score+stats+union pipeline always normalizes.") + budget: int = Field( + default=2048, + gt=0, + description="Tokens kept at each periodic eviction; prompt tokens are " + "always preserved on top.") + beta: int = Field( + default=128, + gt=0, + description="Eviction period in confirmed generation tokens (upstream " + "`divide_length`): one speculative iteration may advance the counter " + "by multiple accepted tokens; at most one eviction is coalesced per update." + ) + model_path: str = Field( + min_length=1, + description="Checkpoint path used to derive RoPE tables when converting " + "the official calibration and to classify kernel-masked sliding-attention " + "layers.") + calibration_path: str = Field( + min_length=1, + description="Path to the official TriAttention calibration `.pt` " + "(produced by github.com/WeianMao/triattention). TRT-LLM does not " + "compute calibration; it converts this file to the runtime schema at " + "load.") + + def supports_block_reuse(self) -> bool: + return True + + def supports_speculative_decoding(self) -> bool: + return self.eviction_mode == "union" + + +KvCacheCompressionConfigType: TypeAlias = Annotated[ + Union[TriAttentionKvCacheCompressionConfig], + Field(discriminator="algorithm"), +] @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) @@ -4420,8 +4484,8 @@ class BaseLlmArgs(StrictBaseModel): status="prototype") # KV cache compression config (separate from sparse attention: changes which - # KV is stored, not the attention computation) - kv_cache_compression_config: Optional[KvCacheCompressionConfig] = Field( + # KV is stored, not the attention computation). + kv_cache_compression_config: Optional[KvCacheCompressionConfigType] = Field( default=None, description="KV-cache compression config; None disables compression.", status="prototype") diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 842660d2e743..44ec05f413c7 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -554,6 +554,47 @@ "kind": "value", "path": "iter_stats_max_iterations" }, + { + "allowed_values": [ + "triattention" + ], + "annotation": "Literal['triattention']", + "converter": "", + "kind": "categorical", + "path": "kv_cache_compression_config.algorithm" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.beta" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.budget" + }, + { + "allowed_values": [ + "union", + "per_head", + "per_layer_perhead" + ], + "annotation": "Literal['union', 'per_head', 'per_layer_perhead']", + "converter": "", + "kind": "categorical", + "path": "kv_cache_compression_config.eviction_mode" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.normalize_scores" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 51f7f6d821e3..ecc1bd12747a 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -75,6 +75,7 @@ l0_b200: - unittest/_torch/attention - unittest/_torch/compilation - unittest/_torch/debugger + - unittest/_torch/kv_cache_compression - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_2_model_mtp - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) diff --git a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py index 9af481b2f686..c6d1a236ff88 100644 --- a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py +++ b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -22,6 +22,7 @@ from unittest.mock import MagicMock import pytest +import torch from tensorrt_llm._torch.pyexecutor.py_executor import AsyncTransferManager, PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType @@ -196,6 +197,63 @@ def test_exhaustion_fixed_with_early_release(self, index_mapper): index_mapper.add_new_sequence(3) assert _has_sequence(index_mapper, 3) + def test_gather_k_block_offsets_uses_request_order_and_beam_zero(self): + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + IndexMapper, + ) + + index_mapper = IndexMapper(max_batch_size=3, max_beam_width=2) + index_mapper.add_new_sequence(11) + index_mapper.add_new_sequence(22) + index_mapper.remove_sequence(22) + index_mapper.add_new_sequence(33) + + source = torch.arange(2 * 6 * 2 * 5, dtype=torch.int32).reshape(2, 6, 2, 5) + destination = torch.full((2, 3, 2, 3), -1, dtype=torch.int32) + index_mapper.gather_k_block_offsets(source, destination, [33, 11], 3) + + torch.testing.assert_close(destination[:, 0, 0], source[:, 2, 0, :3]) + torch.testing.assert_close(destination[:, 1, 0], source[:, 0, 0, :3]) + assert torch.count_nonzero(destination[:, :, 1] != -1) == 0 + assert torch.count_nonzero(destination[:, 2, 0] != -1) == 0 + + def test_gather_k_block_offsets_rejects_unknown_request(self): + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + IndexMapper, + ) + + index_mapper = IndexMapper(max_batch_size=1, max_beam_width=1) + index_mapper.add_new_sequence(11) + source = torch.zeros((1, 1, 2, 4), dtype=torch.int32) + destination = torch.full((1, 2, 2, 3), -1, dtype=torch.int32) + + with pytest.raises(Exception, match="Request ID not found"): + index_mapper.gather_k_block_offsets(source, destination, [11, 12], 3) + assert torch.count_nonzero(destination != -1) == 0 + + def test_gather_k_block_offsets_matches_beam_zero_index_select(self): + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + IndexMapper, + ) + + index_mapper = IndexMapper(max_batch_size=4, max_beam_width=3) + for request_id in (101, 202, 303): + index_mapper.add_new_sequence(request_id) + index_mapper.remove_sequence(202) + index_mapper.add_new_sequence(404) + + request_ids = [404, 101, 404, 303] + source = torch.arange(3 * 12 * 2 * 7, dtype=torch.int32).reshape(3, 12, 2, 7) + destination = torch.full((3, 5, 2, 5), -1, dtype=torch.int32) + copy_index = index_mapper.get_copy_index(request_ids, 0, 1).to(torch.long) + expected = torch.index_select(source, 1, copy_index)[:, :, 0, :5] + + index_mapper.gather_k_block_offsets(source, destination, request_ids, 5) + + torch.testing.assert_close(destination[:, :4, 0], expected) + assert torch.count_nonzero(destination[:, :, 1] != -1) == 0 + assert torch.count_nonzero(destination[:, 4, 0] != -1) == 0 + class TestFreeResourcesDoubleReleaseSafety: """Test that free_resources handles already-released IndexMapper slots.""" diff --git a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py index 7c11357c3210..4f628d6fbb94 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -36,6 +36,7 @@ ResourceManager, ResourceManagerType, ) +from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig # ---------------------------------------------------------------------- # # Mock infra: in-memory managers / requests (avoid touching V2 / model). # @@ -47,7 +48,7 @@ class _RecordingMixin: translation without real algorithm side-effects.""" def __init__(self, kv_cache_manager, record_list, name="m"): - super().__init__(kv_cache_manager) + super().__init__(_compression_config(), kv_cache_manager) self._record_list = record_list self._name = name @@ -71,8 +72,17 @@ def on_request_finish(self, request): self._record("on_request_finish") -class _LengthAdjustingCompressionManager(KVCacheCompressionManager): - adjusts_generation_kv_length: ClassVar[bool] = True +class _PhysicalLengthChangingConfig(KvCacheCompressionConfig): + changes_physical_kv_length: ClassVar[bool] = True + + +class _BlockReuseCompatibleConfig(KvCacheCompressionConfig): + def supports_block_reuse(self) -> bool: + return True + + +def _compression_config() -> KvCacheCompressionConfig: + return KvCacheCompressionConfig(algorithm="test") def _v2_manager(*, is_draft: bool): @@ -87,8 +97,7 @@ def _v2_manager(*, is_draft: bool): @pytest.fixture def fake_kv_cache_manager(): - """A stand-in KVCacheManagerV2. The framework reads enable_block_reuse off - it in __init__; default it to False, like a normal run with reuse off.""" + """A stand-in KVCacheManagerV2 for compression-manager unit tests.""" return _v2_manager(is_draft=False) @@ -118,7 +127,7 @@ def test_inherits_base_resource_manager(self): assert issubclass(KVCacheCompressionManager, BaseResourceManager) def test_four_hooks_default_noop(self, fake_kv_cache_manager): - m = KVCacheCompressionManager(fake_kv_cache_manager) + m = KVCacheCompressionManager(_compression_config(), fake_kv_cache_manager) assert m.on_request_init(MagicMock()) is None assert m.on_context_step_end([MagicMock()]) is None assert m.on_generation_step_begin(MagicMock()) is None @@ -128,24 +137,25 @@ def test_four_hooks_default_noop(self, fake_kv_cache_manager): def test_hooks_accept_extra_kwargs(self, fake_kv_cache_manager): # **kwargs lets the framework pass new args later without breaking # existing overrides. - m = KVCacheCompressionManager(fake_kv_cache_manager) + m = KVCacheCompressionManager(_compression_config(), fake_kv_cache_manager) assert m.on_request_init(MagicMock(), future_arg=1) is None assert m.on_generation_step_end(MagicMock(), future_arg=1) is None def test_resource_counts_are_zero(self, fake_kv_cache_manager): - m = KVCacheCompressionManager(fake_kv_cache_manager) + m = KVCacheCompressionManager(_compression_config(), fake_kv_cache_manager) # The manager owns no physical resources (the V2 cache manager does), # so it must not gate the scheduler. assert m.get_max_resource_count() == 0 assert m.get_needed_resource_to_completion(MagicMock()) == 0 - def test_length_adjustment_marks_target_and_draft_v2(self): + def test_physical_length_change_marks_target_and_draft_v2(self): # The draft cache is compacted together with the target, so both # managers diverge from the logical length in the same way. target = _v2_manager(is_draft=False) draft = _v2_manager(is_draft=True) - manager = _LengthAdjustingCompressionManager(target, draft) + config = _PhysicalLengthChangingConfig(algorithm="test") + manager = KVCacheCompressionManager(config, target, draft) assert manager.kv_cache_manager is target assert manager.draft_kv_cache_manager is draft @@ -154,10 +164,11 @@ def test_length_adjustment_marks_target_and_draft_v2(self): assert draft.kv_compression_manages_history is True def test_rejects_non_v2_ownership(self): + config = _compression_config() with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - KVCacheCompressionManager(MagicMock()) + KVCacheCompressionManager(config, MagicMock()) with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - KVCacheCompressionManager(_v2_manager(is_draft=False), MagicMock()) + KVCacheCompressionManager(config, _v2_manager(is_draft=False), MagicMock()) def test_request_field_defaults_to_zero(self): """LlmRequest carries the compression count (the manager's only @@ -293,26 +304,43 @@ def test_factory_accepts_independent_draft_manager(self): is None ) - def test_eviction_method_predicate_defaults_false(self): - # Non-evicting methods (e.g. offloading) are never restricted by the - # speculative mode: the call-site gate reads this config predicate. - from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig + def test_triattention_requires_sm100_family(self, fake_kv_cache_manager): + cfg = MagicMock() + cfg.algorithm = "triattention" + with ( + patch.object(util_mod, "is_sm_100f", return_value=False), + pytest.raises(RuntimeError, match="SM100-family"), + ): + create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) + def test_capabilities_default_false(self): config = KvCacheCompressionConfig(algorithm="offload") - assert config.kv_cache_compression_mode.is_eviction_method() is False - m = KVCacheCompressionManager(_v2_manager(is_draft=False)) + target = _v2_manager(is_draft=False) + assert config.changes_physical_kv_length is False + assert config.supports_block_reuse() is False + assert config.supports_speculative_decoding() is False + m = KVCacheCompressionManager(config, target) + assert target.kv_compression_manages_history is False assert not hasattr(m, "spec_config") - def test_spec_gate_only_restricts_eviction_methods(self): - from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec + def test_spec_gate_uses_config_capability(self): + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_compatibility from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig - # Non-evicting methods pass with any speculative mode; no exception. config = KvCacheCompressionConfig(algorithm="offload") + kv_cache_config = SimpleNamespace(enable_block_reuse=False) spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.DFLASH) - validate_kv_cache_compression_with_spec(config, spec_config, None) - validate_kv_cache_compression_with_spec(config, None, None) + with pytest.raises(ValueError, match="speculative decoding"): + validate_kv_cache_compression_compatibility( + config, + kv_cache_config, + spec_config, + ) + validate_kv_cache_compression_compatibility( + config, + kv_cache_config, + None, + ) # ---------------------------------------------------------------------- # @@ -339,22 +367,31 @@ def test_names_not_in_sparse_module(self): # ---------------------------------------------------------------------- # -# 5. Block-reuse guard # +# 5. Compression compatibility gate # # ---------------------------------------------------------------------- # -class TestBlockReuseGuard: - """__init__ refuses block reuse for a method that changes the stored keys - and values, the same check RocketKVCacheManager makes.""" - - def _mgr(self, enable_block_reuse): - m = _v2_manager(is_draft=False) - m.enable_block_reuse = enable_block_reuse - return m - +class TestCompressionCompatibility: def test_raises_when_reuse_on(self): + config = _compression_config() with pytest.raises(ValueError, match="block reuse"): - KVCacheCompressionManager(self._mgr(enable_block_reuse=True)) + util_mod.validate_kv_cache_compression_compatibility( + config, + SimpleNamespace(enable_block_reuse=True), + None, + ) def test_ok_when_reuse_off(self): - KVCacheCompressionManager(self._mgr(enable_block_reuse=False)) # no raise + util_mod.validate_kv_cache_compression_compatibility( + _compression_config(), + SimpleNamespace(enable_block_reuse=False), + None, + ) + + def test_block_reuse_capability_allows_reuse(self): + config = _BlockReuseCompatibleConfig(algorithm="test") + util_mod.validate_kv_cache_compression_compatibility( + config, + SimpleNamespace(enable_block_reuse=True), + None, + ) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 4e475f35a674..b815b2270729 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -1,7 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared harness for the KV-cache compaction tests.""" +"""Shared harness for KV-cache compression tests.""" + +import json +import os +import tempfile +from contextlib import contextmanager +from types import SimpleNamespace +from typing import Optional +from unittest import mock import torch @@ -206,3 +214,432 @@ def run_compaction(compaction): out=compaction["swa_destination_bases"], ) compact(compaction["params"], compaction["request_count"]) + + +def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): + """A bare manager carrying only the page-table staging attributes.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, + ) + + staging = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) + staging.kv_cache_manager = None + staging.draft_kv_cache_manager = None + staging._request_capacity = max_requests + staging.budget = 4 + staging._swa_window = None + staging._draft_protected_tail_capacity = 0 + staging._compaction_done_event = torch.cuda.Event() + staging._staging_reuse_event = torch.cuda.Event() + staging._block_offsets_host = torch.empty( + 1, max_requests, 2, staged_blocks_per_seq, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging._identity_copy_indices_host = torch.arange( + max_requests, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging._block_offsets_device = torch.empty( + 1, max_requests, 2, staged_blocks_per_seq, dtype=torch.int32, device=device + ) + return staging + + +def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): + """The manager surface ``_stage_block_offset_snapshot`` consumes.""" + return SimpleNamespace( + host_kv_cache_block_offsets=host_table, + kv_factor=2, + index_mapper=SimpleNamespace(gather_k_block_offsets=gather), + index_scales=torch.full((num_slots,), 2, dtype=torch.int32, pin_memory=True), + kv_offset=torch.ones(num_slots, dtype=torch.int32, pin_memory=True), + _stream=manager_stream, + ) + + +def make_fake_v2(enable_block_reuse=False, *, is_draft=False): + """Build an unallocated V2 double with TriAttention's production contract.""" + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + + fake_v2 = KVCacheManagerV2.__new__(KVCacheManagerV2) + fake_v2.enable_block_reuse = enable_block_reuse + fake_v2.is_draft = is_draft + fake_v2.kv_compression_manages_history = False + fake_v2.kv_factor = 2 + fake_v2.max_beam_width = 1 + fake_v2.max_batch_size = 8 + fake_v2.num_extra_kv_tokens = 0 + fake_v2.max_draft_len = 0 + fake_v2.max_total_draft_tokens = 0 + fake_v2._kv_reserve_draft_tokens = 0 + fake_v2.max_seq_len = 65536 + fake_v2.tokens_per_block = 64 + fake_v2.max_blocks_per_seq = 1028 + fake_v2.get_num_available_tokens = lambda *, token_num_upper_bound, **_: token_num_upper_bound + fake_v2.max_attention_window_vec = [] + fake_v2.kv_cache_manager_py_config = SimpleNamespace(layers=[]) + fake_v2.impl = object() + fake_v2.kv_cache_map = {} + fake_v2.host_kv_cache_block_offsets = torch.zeros(1, 8, 2, 8, dtype=torch.int32) + fake_v2.pp_layers = [] + fake_v2.layer_offsets = {} + fake_v2.layer_to_pool_mapping_dict = {} + return fake_v2 + + +_TEST_MODEL_DIR: Optional[str] = None + + +def make_test_model_dir() -> str: + """Create a real dense-model config for production layer partitioning.""" + global _TEST_MODEL_DIR + if _TEST_MODEL_DIR is None: + _TEST_MODEL_DIR = tempfile.mkdtemp(prefix="triattention_test_model_") + config = { + "architectures": ["LlamaForCausalLM"], + "model_type": "llama", + "num_hidden_layers": 2, + "hidden_size": 64, + "num_attention_heads": 4, + } + with open(os.path.join(_TEST_MODEL_DIR, "config.json"), "w") as handle: + json.dump(config, handle) + return _TEST_MODEL_DIR + + +def make_test_calibration_pt() -> str: + """A real on-disk flat calibration file: construction loads it for real.""" + path = os.path.join(make_test_model_dir(), "calibration.pt") + if not os.path.exists(path): + num_layers, num_heads, freq_count = 2, 2, 4 + torch.save( + { + "E_q": torch.zeros(num_layers, num_heads, freq_count, dtype=torch.complex64), + "E_q_norm": torch.ones(num_layers, num_heads, freq_count), + "omega": torch.ones(freq_count), + "freq_scale_sq": torch.ones(freq_count), + }, + path, + ) + return path + + +def make_tri_config(**overrides): + """Build a real TriAttention config with test calibration inputs.""" + from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig + + options = { + "budget": 8, + "model_path": make_test_model_dir(), + "calibration_path": make_test_calibration_pt(), + } + options.update(overrides) + return TriAttentionKvCacheCompressionConfig(**options) + + +def make_triattention(**overrides): + """Construct a manager while isolating GPU-owned persistent state.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, + ) + + with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): + return TriAttentionCompressionManager(make_tri_config(**overrides), make_fake_v2()) + + +def make_eviction_request( + request=None, + *, + request_id=0, + source_length, + target_tail_length=0, + target_cache=None, + draft_cache=None, +): + """One due request shaped exactly like ``_evict_due_requests`` builds.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import _EvictionRequest + + if request is None: + request = SimpleNamespace( + py_request_id=request_id, + py_prompt_len=0, + py_num_compressed_tokens=0, + ) + return _EvictionRequest( + request=request, + target_cache=target_cache, + draft_cache=draft_cache, + source_length=int(source_length), + target_tail_length=int(target_tail_length), + ) + + +def make_request(request_id, **overrides): + """Build the explicit request fields consumed by TriAttention.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + + fields = { + "py_request_id": request_id, + "py_prompt_len": 0, + "py_max_new_tokens": 65536, + "py_draft_tokens": [], + "py_num_accepted_draft_tokens": 0, + "py_num_compressed_tokens": 0, + "is_dummy": False, + "state": LlmRequestState.GENERATION_IN_PROGRESS, + } + fields.update(overrides) + return SimpleNamespace(**fields) + + +@contextmanager +def mocked_eviction_internals(manager): + """Run the real ``_evict_due_requests`` transaction around a mocked round executor.""" + with mock.patch.object(manager, "_execute_eviction_round") as execute: + yield SimpleNamespace(execute=execute) + + +def torch_tri_score_oracle( + layer_pools, + page_ids, + seq_lens, + logical_source_lengths, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + layer_indices, +): + """Compute paged mean scores independently with Torch.""" + scores = [] + num_q_heads = int(q_real.shape[1]) + for request, seq_len in enumerate(seq_lens): + phase = (logical_source_lengths[request] + offsets[:, None]) * omega[None, :] + mean_cos = torch.cos(phase).mean(dim=0) + mean_sin = torch.sin(phase).mean(dim=0) + for layer in layer_indices: + pool = layer_pools[layer] + request_page_ids = ( + page_ids[layer][request] if isinstance(page_ids, dict) else page_ids[request] + ) + keys = ( + pool.index_select(0, request_page_ids)[:, 0] + .permute(1, 0, 2, 3) + .reshape(pool.shape[2], -1, pool.shape[4])[:, :seq_len] + .float() + ) + num_kv_heads = int(keys.shape[0]) + group_size = num_q_heads // num_kv_heads + head_scores = [] + for head in range(num_q_heads): + key = keys[head // group_size] + num_freqs = int(key.shape[-1]) // 2 + key_real = key[:, :num_freqs] + key_imag = key[:, num_freqs:] + product_real = q_real[layer, head] * key_real + q_imag[layer, head] * key_imag + product_imag = q_imag[layer, head] * key_real - q_real[layer, head] * key_imag + position = ( + freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) + ).sum(dim=-1) + mlr = ( + torch.sqrt(key_real.square() + key_imag.square()) + * mlr_coef[layer, head] + * freq_scale_sq + ).sum(dim=-1) + head_scores.append(position + mlr) + scores.append(torch.stack(head_scores)) + return scores + + +def make_phase_table(offsets, omega, initial_rows): + """Build the semantic phase-table surface consumed by an eviction round.""" + omega = omega.to(dtype=torch.float32).contiguous() + positions = torch.arange(max(int(initial_rows), 1), dtype=torch.float32, device=omega.device) + angles = (positions[:, None, None] + offsets[None, :, None]) * omega[None, None, :] + num_freqs = int(omega.numel()) + return SimpleNamespace( + cos=torch.cos(angles).mean(dim=1).contiguous(), + sin=torch.sin(angles).mean(dim=1).contiguous(), + num_freqs=num_freqs, + ) + + +def make_cute_buffers( + *, + eviction_mode, + layer_pools, + max_requests, + seq_len, + num_q_heads, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + decode_width=None, + keep_count=1, + protected_tail_capacity=0, + layer_pool_ids=None, + normalize_scores=True, +): + """Build a bare manager with a real score pipeline over test pools.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, + ) + + num_layers = len(layer_pools) + assert int(q_real.shape[1]) == num_q_heads + if decode_width is None: + decode_width = seq_len + if layer_pool_ids is None: + layer_pool_ids = [0] * num_layers + requested_tokens = seq_len + protected_tail_capacity + tokens_per_block = int(layer_pools[0].shape[3]) + source_blocks = -(-int(requested_tokens) // tokens_per_block) + source_blocks = (source_blocks + 3) // 4 * 4 + layout = dict( + layer_pools=layer_pools, + dense_layers=list(range(num_layers)), + swa_layers=[], + swa_window=None, + layer_pool_ids=layer_pool_ids, + ) + manager = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) + manager.kv_cache_manager = SimpleNamespace( + num_pools=max(layer_pool_ids) + 1, + tokens_per_block=tokens_per_block, + max_blocks_per_seq=source_blocks, + host_kv_cache_block_offsets=torch.empty(1, 1, 2, source_blocks, dtype=torch.int32), + mapping=SimpleNamespace(tp_size=1, tp_rank=0, enable_attention_dp=False), + ) + manager.draft_kv_cache_manager = None + manager._draft_protected_tail_capacity = 0 + manager.eviction_mode = eviction_mode + manager.normalize_scores = normalize_scores + manager._request_capacity = max_requests + manager._selection_width_capacity = decode_width + manager._phase = make_phase_table(offsets, omega, seq_len) + manager.budget = keep_count + manager._protected_tail_capacity = protected_tail_capacity + manager._freq_scale_sq = freq_scale_sq + manager._score_q_real = q_real + manager._score_q_imag = q_imag + manager._score_mlr_coef = mlr_coef + manager._target_layout = layout + manager._draft_layout = None + manager._num_layers = num_layers + manager._num_q_heads = num_q_heads + manager._num_kv_heads = int(layer_pools[0].shape[2]) + manager._union_tp_mapping = None + manager._swa_window = None + manager._allocate_metadata_buffers( + layer_pools[0].device, + num_freqs=int(q_real.shape[2]), + ) + manager._allocate_selection_buffers(layer_pools[0].device, tp_size=1) + manager._score_scratch = None + manager._score_token_capacity = 0 + manager._launch_score = None + manager._compaction_params = () + manager._staging_reuse_event = torch.cuda.Event() + manager._staging_reuse_event.record(torch.cuda.current_stream(layer_pools[0].device)) + manager._compaction_done_event = torch.cuda.Event() + manager._compaction_done_event.record(torch.cuda.current_stream(layer_pools[0].device)) + manager._build_score_runtime(score_token_capacity=seq_len) + return manager + + +def write_block_offsets(manager, encoded): + """Load a test page table into the staged block-offset plane.""" + manager._block_offsets_device.zero_() + manager._block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) + + +def rect_to_score_scratch(scores, num_kv_heads, padded_head_columns=8): + """Scatter rectangular scores into the fused scorer's scratch layout.""" + request_count, num_layers, num_q_heads, width = scores.shape + group = num_q_heads // num_kv_heads + scratch = torch.zeros( + num_kv_heads * padded_head_columns * request_count * num_layers * width, + dtype=torch.float32, + device=scores.device, + ) + view = scratch.view(num_kv_heads, padded_head_columns, request_count, num_layers, width) + view[:, :group] = scores.view(request_count, num_layers, num_kv_heads, group, width).permute( + 2, 3, 0, 1, 4 + ) + prompt_lengths = torch.zeros(request_count, dtype=torch.int32, device=scores.device) + return scratch, prompt_lengths + + +def stage_score_metadata(manager, request_count, source_lengths, decode_lengths, prompt_lengths): + """Stage per-round score metadata exactly as production does.""" + torch.sub( + source_lengths[:request_count], + prompt_lengths[:request_count], + out=decode_lengths[:request_count], + ) + manager._source_lengths_device[:request_count].copy_(source_lengths[:request_count]) + manager._prompt_lengths_device[:request_count].copy_(prompt_lengths[:request_count]) + + +def launch_split_scores( + manager, request_count, source_lengths, decode_lengths, prompt_lengths, mean_cos, mean_sin +): + """Run the score pipeline and gather its per-head decode-window rectangle.""" + stage_score_metadata(manager, request_count, source_lengths, decode_lengths, prompt_lengths) + manager._mean_cos[:request_count].copy_(mean_cos[:request_count]) + manager._mean_sin[:request_count].copy_(mean_sin[:request_count]) + manager._launch_score(request_count) + score_scratch = manager._score_scratch + score_token_capacity = manager._score_token_capacity + num_segments = request_count * manager._num_layers + group_size = manager._num_q_heads // manager._num_kv_heads + source = ( + score_scratch[: manager._num_kv_heads * 8 * num_segments * score_token_capacity] + .view( + manager._num_kv_heads, + 8, + request_count, + manager._num_layers, + score_token_capacity, + )[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + columns = prompt_lengths[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + torch.arange( + manager._selection_width_capacity, + dtype=torch.int64, + device=score_scratch.device, + ).view(1, 1, 1, 1, -1) + columns = columns.clamp_(max=score_token_capacity - 1).expand( + request_count, + manager._num_layers, + manager._num_kv_heads, + group_size, + manager._selection_width_capacity, + ) + output = torch.full( + ( + request_count, + manager._num_layers, + manager._num_q_heads, + manager._selection_width_capacity, + ), + float("nan"), + dtype=torch.float32, + device=score_scratch.device, + ) + torch.gather( + source, + 4, + columns, + out=output.view( + request_count, + manager._num_layers, + manager._num_kv_heads, + group_size, + manager._selection_width_capacity, + ), + ) + return output diff --git a/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py new file mode 100644 index 000000000000..144358fce6e7 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Physical KV-length changes force the unfused-RoPE path. + +When physical and logical KV lengths diverge, the fused path can no longer +derive rotary positions from physical KV length. The unfused path consumes the +engine's logical ``position_ids`` instead. +""" + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.attention import Attention +from tensorrt_llm.llmapi.llm_args import ( + KvCacheCompressionConfig, + TriAttentionKvCacheCompressionConfig, +) + + +def _make_attention(model_config: ModelConfig) -> Attention: + return Attention( + hidden_size=256, + num_attention_heads=8, + num_key_value_heads=8, + max_position_embeddings=1024, + bias=False, + pos_embd_params=None, + layer_idx=0, + dtype=torch.bfloat16, + config=model_config, + ) + + +def test_plain_attention_defaults_to_fused_rope() -> None: + attn = _make_attention(ModelConfig()) + + assert attn.rope_fusion is True + + +def test_physical_length_preserving_compression_keeps_fused_rope() -> None: + model_config = ModelConfig( + kv_cache_compression_config=KvCacheCompressionConfig(algorithm="test") + ) + attn = _make_attention(model_config) + + assert attn.rope_fusion is True + + +def test_physical_kv_length_change_forces_unfused_rope() -> None: + model_config = ModelConfig( + kv_cache_compression_config=TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt" + ) + ) + attn = _make_attention(model_config) + + assert attn.rope_fusion is False diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py new file mode 100644 index 000000000000..58c8a7afb4fb --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The SM100 TriAttention CuTe scorer (the only score path) vs oracles. + +The launch matrix drives the named production geometries against the +pure-PyTorch oracle; the contract test pins the no-fallback loud raise. +""" + +import pytest +import torch +from conftest import encode_block_offsets as _encode_block_offsets +from conftest import launch_split_scores as _launch_split_scores +from conftest import make_cute_buffers as _make_cute_buffers +from conftest import torch_tri_score_oracle as _torch_tri_score_oracle +from conftest import write_block_offsets as _write_block_offsets + +requires_sm100 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="TriAttention score requires SM100", +) + + +def _build_case( + *, + max_requests: int, + num_layers: int, + page_count: int, + tokens_per_block: int, + head_dim: int, + num_q_heads: int, + num_kv_heads: int, + prompt_len: int, + seed: int, + offsets: tuple = (1.0, 2.0, 4.0), +): + device = torch.device("cuda", torch.cuda.current_device()) + torch.manual_seed(seed) + num_freqs = head_dim // 2 + # 0.125 scaling keeps the kernel-vs-oracle tolerance tight. + pools = [ + ( + 0.125 + * torch.randn( + max_requests * page_count, + 2, + num_kv_heads, + tokens_per_block, + head_dim, + device=device, + ) + ).to(torch.bfloat16) + for _ in range(num_layers) + ] + page_ids = torch.randperm(max_requests * page_count).view(max_requests, page_count).to(device) + q_real = 0.125 * torch.randn(num_layers, num_q_heads, num_freqs, device=device) + q_imag = 0.125 * torch.randn(num_layers, num_q_heads, num_freqs, device=device) + mlr_coef = 0.125 * torch.randn(num_layers, num_q_heads, num_freqs, device=device) + freq_scale_sq = torch.rand(num_freqs, device=device) + 0.5 + omega = torch.rand(num_freqs, device=device) * 0.05 + offsets_t = torch.tensor(offsets, dtype=torch.float32, device=device) + capacity = page_count * tokens_per_block + tri = _make_cute_buffers( + eviction_mode="per_head", + layer_pools=pools, + max_requests=max_requests, + seq_len=capacity, + num_q_heads=num_q_heads, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + omega=omega, + offsets=offsets_t, + decode_width=capacity - prompt_len, + ) + _write_block_offsets(tri, _encode_block_offsets(page_ids)) + logical_source_lengths = ( + torch.arange(max_requests, dtype=torch.int32, device=device) + 9 + ).contiguous() + prompt_lengths = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) + # Mid-page/mid-tile tails; 58 leaves a fully-invalid trailing fragment. + tail_cuts = (0, 58, 3, 33) + seq_lens = [capacity - tail_cuts[request % len(tail_cuts)] for request in range(max_requests)] + source_lengths = torch.tensor(seq_lens, dtype=torch.int32, device=device) + phase = (logical_source_lengths.float()[:, None, None] + offsets_t[None, :, None]) * omega[ + None, None, : + ] + mean_cos = torch.cos(phase).mean(dim=1).contiguous() + mean_sin = torch.sin(phase).mean(dim=1).contiguous() + oracle_inputs = dict( + page_ids=page_ids, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + omega=omega, + offsets=offsets_t, + ) + return ( + tri, + pools, + prompt_lengths, + source_lengths, + seq_lens, + mean_cos, + mean_sin, + oracle_inputs, + ) + + +def _geometry(max_requests, num_layers, page_count, tokens_per_block, head_dim, num_q, num_kv): + return dict( + max_requests=max_requests, + num_layers=num_layers, + page_count=page_count, + tokens_per_block=tokens_per_block, + head_dim=head_dim, + num_q_heads=num_q, + num_kv_heads=num_kv, + ) + + +# One entry per supported production geometry. +_CASES = [ + pytest.param(_geometry(4, 2, 4, 32, 128, 8, 2), id="qwen3_f64_group4_tpb32"), + pytest.param(_geometry(2, 3, 4, 32, 64, 8, 1), id="gptoss_f32_group8_tpb32"), + pytest.param(_geometry(2, 2, 2, 128, 64, 8, 1), id="original_f32_group8_tpb128"), +] + + +@requires_sm100 +@pytest.mark.parametrize("case", _CASES) +def test_cute_kernel_matches_torch_oracle(case): + pytest.importorskip("cutlass") + case = dict(case) # parametrize reuses the dict across reruns + prompt_len = 5 + max_requests = case["max_requests"] + num_layers = case["num_layers"] + ( + tri, + pools, + prompt_lengths, + source_lengths, + seq_lens, + mean_cos, + mean_sin, + oracle_inputs, + ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) + device = tri._score_scratch.device + + oracle = _torch_tri_score_oracle( + pools, + oracle_inputs["page_ids"], + seq_lens, + [int(start) for start in range(9, 9 + max_requests)], + oracle_inputs["q_real"], + oracle_inputs["q_imag"], + oracle_inputs["mlr_coef"], + oracle_inputs["freq_scale_sq"], + oracle_inputs["omega"], + oracle_inputs["offsets"], + list(range(num_layers)), + ) + + # Every request count used by the runtime dispatches through the launcher. + for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): + decode_lengths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) + scores = _launch_split_scores( + tri, + request_count, + source_lengths, + decode_lengths, + prompt_lengths, + mean_cos, + mean_sin, + ) + assert scores.shape == ( + request_count, + num_layers, + case["num_q_heads"], + tri._selection_width_capacity, + ) + # The score leg owns the per-request decode widths the selection + # reduce kernels consume. + assert decode_lengths[:request_count].tolist() == [ + seq_lens[request] - prompt_len for request in range(request_count) + ] + for request in range(request_count): + width = seq_lens[request] - prompt_len + for layer in range(num_layers): + torch.testing.assert_close( + scores[request, layer, :, :width], + oracle[request * num_layers + layer][:, prompt_len : prompt_len + width], + rtol=5e-3, + atol=5e-3, + ) + + +def test_unsupported_geometry_raises_at_buffer_construction(): + pytest.importorskip("cutlass") + device = torch.device("cuda", torch.cuda.current_device()) + torch.manual_seed(20260722) + num_layers, max_requests, page_count, tokens_per_block, head_dim = 2, 2, 2, 4, 8 + num_freqs = head_dim // 2 + # fp32, 4-token pages, 4 freqs: outside the contract on every device. + pools = [ + torch.randn(max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device) + for _ in range(num_layers) + ] + calib = torch.randn(num_layers, 2, num_freqs, device=device) + # No rewrap: the buffer build's own contract error surfaces directly + # (the fp32 pools trip the BF16 gate first, at TMA descriptor encoding). + with pytest.raises(TypeError, match="BF16"): + _make_cute_buffers( + eviction_mode="per_head", + layer_pools=pools, + max_requests=max_requests, + seq_len=page_count * tokens_per_block, + num_q_heads=2, + q_real=calib, + q_imag=calib.clone(), + mlr_coef=calib.clone(), + freq_scale_sq=torch.rand(num_freqs, device=device) + 0.5, + omega=torch.rand(num_freqs, device=device) * 0.05, + offsets=torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), + decode_width=page_count * tokens_per_block - 1, + ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_union_fusion.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_union_fusion.py new file mode 100644 index 000000000000..b49b731085de --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_union_fusion.py @@ -0,0 +1,349 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Equivalence coverage for the fused score+stats+union pipeline. + +The reference leg gathers the production score rows and normalizes + +union-reduces them with a pure-torch float32 oracle; tolerances are +unchanged from the retired Triton reference copies. +""" + +import pytest +import torch +from conftest import launch_split_scores as _launch_split_scores +from conftest import make_cute_buffers as _make_cute_buffers +from conftest import stage_score_metadata as _stage_score_metadata +from conftest import write_block_offsets as _write_block_offsets + +_SM100_ONLY = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="TriAttention CuTe kernels require SM100", +) + + +def _run_fused_union( + tri, + request_count, + source_lengths, + decode_lengths, + prompt_lengths, + mean_cos, + mean_sin, + union_out, +): + """Run the fused score+stats+normalized-union pipeline.""" + _stage_score_metadata(tri, request_count, source_lengths, decode_lengths, prompt_lengths) + tri._mean_cos[:request_count].copy_(mean_cos[:request_count]) + tri._mean_sin[:request_count].copy_(mean_sin[:request_count]) + tri._launch_score(request_count) + columns = min(union_out.shape[1], tri._selection_scores_rows.shape[1]) + union_out[:request_count, :columns].copy_(tri._selection_scores_rows[:request_count, :columns]) + + +def _reference_union_scores( + scores_rows: torch.Tensor, decode_lengths: torch.Tensor +) -> torch.Tensor: + """Compute the normalized max-fold union score oracle.""" + request_count, _, width = scores_rows.shape + combined = torch.full( + (request_count, width), float("-inf"), dtype=torch.float32, device=scores_rows.device + ) + for request in range(request_count): + decode_length = int(decode_lengths[request]) + if decode_length <= 0: + continue + valid = scores_rows[request, :, :decode_length].to(torch.float32) + mean = valid.mean(dim=1, keepdim=True) + std = ((valid - mean).square().sum(dim=1, keepdim=True) / decode_length).sqrt() + combined[request, :decode_length] = ((valid - mean) / std.clamp_min(1e-6)).amax(dim=0) + return combined + + +@_SM100_ONLY +@pytest.mark.parametrize( + "tokens_per_block,num_freqs,num_q_heads,score_starts,valid_lens", + [ + # Representative rows per axis: the originally validated geometry + # (32 freqs, GQA group 8) at both page sizes with full-range and + # ragged page-aligned starts. + (32, 32, 8, 0, None), + (128, 32, 8, 128, [250, 230]), + # Qwen3 geometry: 128-element K rows (64 frequencies) and GQA group + # 4, which rides the MMA tile N=8 with zeroed padding columns. + (32, 64, 4, 37, [250, 198]), + (128, 64, 4, 128, [250, 230]), + # GQA group 4 with 32 frequencies: head columns pad up to the MMA + # tile N=8 with zeroed weights, the partial-stats epilogue writes + # only the real heads' rows, and the union finalizer maps head rows + # onto the padded score planes. + (128, 32, 4, 0, None), + # Mixed-prompt request group (one start mid-tile, one page-aligned) — the + # case the fused pipeline previously declined. Starts are per-request + # runtime reads, so one representative row covers the family. + (128, 64, 4, [37, 128], [250, 230]), + ], +) +def test_union_fusion_matches_split_pipeline( + tokens_per_block: int, + num_freqs: int, + num_q_heads: int, + score_starts: "int | list", + valid_lens: "list | None", +) -> None: + """Check fused union rows against the split score-normalize-union path.""" + pytest.importorskip("cutlass") + + torch.manual_seed(20260721) + device = torch.device("cuda") + seq_len = 256 + num_pages = seq_len // tokens_per_block + # Shuffled pages catch fragment/page mix-ups; ragged lengths land + # mid-tile. + page_permutation = {128: [0, 1], 32: [3, 1, 4, 7, 5, 0, 2, 6]}[tokens_per_block] + assert sorted(page_permutation) == list(range(num_pages)) + pool = ( + 0.125 * torch.randn(num_pages, 2, 1, tokens_per_block, 2 * num_freqs, device=device) + ).to(torch.bfloat16) + q_real = 0.125 * torch.randn(1, num_q_heads, num_freqs, device=device) + q_imag = 0.125 * torch.randn_like(q_real) + mlr_coef = 0.125 * torch.randn_like(q_real) + freq_scale_sq = torch.linspace(0.5, 1.5, num_freqs, device=device) + omega = torch.linspace(0.01, 0.03, num_freqs, device=device) + offsets = torch.tensor([1.0, 2.0, 4.0], device=device) + logical_source_lengths = torch.tensor([float(seq_len), float(seq_len + 1)], device=device) + phase = (logical_source_lengths[:, None, None] + offsets[None, :, None]) * omega[None, None] + mean_cos = torch.cos(phase).mean(dim=1).contiguous() + mean_sin = torch.sin(phase).mean(dim=1).contiguous() + + common = dict( + layer_pools=[pool], + max_requests=2, + seq_len=seq_len, + num_q_heads=num_q_heads, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + omega=omega, + offsets=offsets, + ) + tri = _make_cute_buffers(eviction_mode="union", **common) + # The split reference leg runs on its own score-only buffers. + ref_tri = _make_cute_buffers(eviction_mode="per_head", **common) + k_plane = [2 * page for page in page_permutation] + v_plane = [2 * page + 1 for page in page_permutation] + encoded = torch.tensor( + [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device + ) + _write_block_offsets(tri, encoded) + _write_block_offsets(ref_tri, encoded) + if valid_lens is None: + valid_lens = [seq_len, seq_len] + source_lengths = torch.tensor(valid_lens, dtype=torch.int32, device=device) + request_count = 2 + if isinstance(score_starts, int): + score_starts = [score_starts] * request_count + assert len(score_starts) == request_count + + # Reference: the production score gather over the same decode windows, + # then the pure-torch union oracle. + split_widths = torch.empty(request_count, dtype=torch.int32, device=device) + prompt_lengths = torch.tensor(score_starts, dtype=torch.int32, device=device) + per_head = _launch_split_scores( + ref_tri, + request_count, + source_lengths, + split_widths, + prompt_lengths, + mean_cos, + mean_sin, + ) + rows = per_head.shape[1] * per_head.shape[2] + scores_rows = per_head.reshape(request_count, rows, seq_len).contiguous() + expected = _reference_union_scores(scores_rows, split_widths) + + fused_widths = torch.empty(request_count, dtype=torch.int32, device=device) + fused_out = torch.full( + (request_count, seq_len), float("nan"), dtype=torch.float32, device=device + ) + _run_fused_union( + tri, + request_count, + source_lengths, + fused_widths, + prompt_lengths, + mean_cos, + mean_sin, + fused_out, + ) + assert torch.equal(fused_widths, split_widths) + for request in range(request_count): + width = int(valid_lens[request]) - int(score_starts[request]) + torch.testing.assert_close( + fused_out[request, :width], + expected[request, :width], + rtol=5.0e-3, + atol=5.0e-3, + ) + + +@_SM100_ONLY +def test_union_fusion_frequency_count_guard_raises() -> None: + """Reject unsupported 16-frequency fused-kernel geometry.""" + pytest.importorskip("cutlass") + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused import ( # noqa: E501 + _TriAttentionScoreKernel, + ) + + with pytest.raises(ValueError, match="frequencies"): + _TriAttentionScoreKernel( + num_layers=1, + score_token_capacity=256, + num_q_heads=8, + num_freqs=16, + pool_shape=(2, 2, 1, 128, 32), + pool_strides=(8192, 4096, 4096, 32, 1), + page_shards=3, + ) + + +@_SM100_ONLY +@pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_properties(0).total_memory < 32 * 1024**3, + reason="the giant-scratch geometry needs ~15 GiB of device memory", +) +@pytest.mark.parametrize( + "max_requests", + [ + # Qwen3-8B serve geometry at max_batch_size 64 with the 16384-token + # bucket: the score scratch spans 2,415,919,104 elements, past 2^31. + # Before the finalizer's fallback loads were folded into a 64-bit + # pointer, this leg died with an illegal memory access whenever a + # window start was not lane-aligned (the serve-mode eviction crash). + 64, + # Same shape at max_batch_size 32 stays below 2^31 and covers the + # boundary from the always-correct side. + 32, + ], +) +def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: + """Unaligned window starts must survive a past-2^31-element score scratch. + + The union finalizer reads the score scratch at flat offsets up to + ``plane * capacity * layers * bucket``; the production Qwen3-8B serve + shape (capacity 64, 36 layers, bucket 16384) pushes those offsets past + 2^31. Requests whose pinned prompt length is not a multiple of the lane + width take the per-token load branch, which must fold the Int64 offset + into the pointer instead of the DSL's 32-bit dynamic coordinate. The leg + launches at full capacity with zero-length tail rows, exactly like a + production eviction round, and checks the fused rows against the split + score-gather plus the pure-torch union oracle. + """ + pytest.importorskip("cutlass") + + torch.manual_seed(20260722) + device = torch.device("cuda") + num_layers = 36 + num_q_heads = 32 + num_kv_heads = 8 + num_freqs = 64 + tokens_per_block = 32 + seq_len = 16384 + decode_window = 8192 + # Window starts deliberately off the 4-token lane grid (real prompt + # lengths are arbitrary), one of them also off the page grid. + score_starts = [897, 641] + valid_lens = [start + decode_window for start in score_starts] + request_count = len(score_starts) + + # Every layer shares one physical pool: the scratch magnitude only needs + # the segment count, not distinct K content per layer. + num_pages = (max(valid_lens) + tokens_per_block - 1) // tokens_per_block + pool = ( + 0.125 + * torch.randn(num_pages, 2, num_kv_heads, tokens_per_block, 2 * num_freqs, device=device) + ).to(torch.bfloat16) + layer_pools = [pool] * num_layers + calib_shape = (num_layers, num_q_heads, num_freqs) + q_real = 0.125 * torch.randn(calib_shape, device=device) + q_imag = 0.125 * torch.randn_like(q_real) + mlr_coef = 0.125 * torch.randn_like(q_real) + freq_scale_sq = torch.linspace(0.5, 1.5, num_freqs, device=device) + omega = torch.linspace(0.01, 0.03, num_freqs, device=device) + offsets = torch.tensor([1.0, 2.0, 4.0], device=device) + logical_source_lengths = ( + torch.arange(max_requests, dtype=torch.float32, device=device) + seq_len + ) + phase = (logical_source_lengths[:, None, None] + offsets[None, :, None]) * omega[None, None] + mean_cos = torch.cos(phase).mean(dim=1).contiguous() + mean_sin = torch.sin(phase).mean(dim=1).contiguous() + + common = dict( + layer_pools=layer_pools, + seq_len=seq_len, + num_q_heads=num_q_heads, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + omega=omega, + offsets=offsets, + decode_width=decode_window, + ) + tri = _make_cute_buffers(eviction_mode="union", max_requests=max_requests, **common) + assert (tri._score_scratch.numel() > 2**31) == (max_requests == 64) + # The split reference leg only scores the two live requests; its own + # small per_head buffers keep the giant scratch on the union side. + ref_tri = _make_cute_buffers(eviction_mode="per_head", max_requests=request_count, **common) + page_ids = torch.arange(num_pages, dtype=torch.int32, device=device) + for staged in (tri, ref_tri): + staged._block_offsets_device.zero_() + staged._block_offsets_device[0, :request_count, 0, :num_pages] = 2 * page_ids + staged._block_offsets_device[0, :request_count, 1, :num_pages] = 2 * page_ids + 1 + + source_lengths = torch.zeros(max_requests, dtype=torch.int32, device=device) + prompt_lengths = torch.zeros(max_requests, dtype=torch.int32, device=device) + source_lengths[:request_count] = torch.tensor(valid_lens, dtype=torch.int32, device=device) + prompt_lengths[:request_count] = torch.tensor(score_starts, dtype=torch.int32, device=device) + + # Reference: the split score gather over the same decode windows, then + # the pure-torch union oracle. + split_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) + per_head = _launch_split_scores( + ref_tri, + request_count, + source_lengths, + split_widths, + prompt_lengths, + mean_cos, + mean_sin, + ) + rows = per_head.shape[1] * per_head.shape[2] + scores_rows = per_head.reshape(request_count, rows, decode_window).contiguous() + expected = _reference_union_scores(scores_rows, split_widths) + + # Fused pipeline at FULL capacity (zero-length tails), like production. + fused_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) + fused_out = torch.full( + (max_requests, seq_len), float("nan"), dtype=torch.float32, device=device + ) + _run_fused_union( + tri, + max_requests, + source_lengths, + fused_widths, + prompt_lengths, + mean_cos, + mean_sin, + fused_out, + ) + torch.cuda.synchronize() + assert torch.equal(fused_widths[:request_count], split_widths[:request_count]) + for request in range(request_count): + width = valid_lens[request] - score_starts[request] + torch.testing.assert_close( + fused_out[request, :width], + expected[request, :width], + rtol=5.0e-3, + atol=5.0e-3, + ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py new file mode 100644 index 000000000000..6824e5d1ce83 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TriAttention draft lifecycle, ordering, admission, and publication.""" + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from conftest import make_eviction_request as _make_eviction_request +from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_request as _make_request +from conftest import make_tri_config as _make_tri_config +from conftest import make_triattention as _make_triattention +from conftest import mocked_eviction_internals as _mocked_eviction_internals + +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, +) + + +def test_execute_eviction_round_uses_current_stream_and_hands_back_to_target(): + """Keep target and draft work on the caller stream before manager handoff.""" + event = mock.Mock() + host = torch.zeros(6, 9, dtype=torch.int32) + tri = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) + tri._request_capacity = 8 + tri.budget = 4 + tri._swa_window = None + tri._draft_protected_tail_capacity = 1 + tri._staging_reuse_event = mock.Mock() + tri._compaction_done_event = event + tri._request_metadata_host = host + tri._request_metadata_host_np = host.numpy() + metadata_device = mock.Mock() + tri._request_metadata_device = metadata_device + tri._block_offsets_host = None + tri._block_offsets_device = torch.zeros(1, dtype=torch.int32) + tri._draft_block_offsets_host = None + tri._draft_block_offsets_device = None + execution_stream = mock.Mock() + manager = SimpleNamespace(_stream=execution_stream) + draft_manager = SimpleNamespace( + _stream=execution_stream, num_extra_kv_tokens=0, _kv_reserve_draft_tokens=0 + ) + tri.kv_cache_manager = manager + tri.draft_kv_cache_manager = draft_manager + compute_stream = mock.Mock() + eviction_requests = [_make_eviction_request(request_id=7, source_length=8)] + + class Boom(RuntimeError): + pass + + metadata_device.copy_.side_effect = Boom + with ( + mock.patch.object( + torch.cuda, "current_stream", return_value=compute_stream + ) as current_stream, + mock.patch.object(tri, "_stage_block_offset_snapshot") as stage, + ): + with pytest.raises(Boom): + tri._execute_eviction_round(eviction_requests) + + # Both page-table planes were snapshotted before the round body fired. + assert stage.call_count == 2 + # One event records the current execution stream. Only the target manager + # owns the post-round resize/release handoff. + current_stream.assert_called_once_with(tri._block_offsets_device.device) + event.record.assert_called_once_with(compute_stream) + execution_stream.wait_event.assert_called_once_with(event) + + +@pytest.mark.parametrize( + "gate,match", + [ + ("callsite_dflash", "one-model MTP or EAGLE3"), + ("union_only_per_head", "eviction_mode='union'"), + ], +) +def test_speculative_admission_gates_raise(gate, match): + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_compatibility + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig, MTPDecodingConfig + + if gate == "callsite_dflash": + spec_config = DFlashDecodingConfig(max_draft_len=3) + else: + spec_config = MTPDecodingConfig(max_draft_len=1) + config = _make_tri_config( + budget=8, + eviction_mode="per_head" if gate == "union_only_per_head" else "union", + ) + + with pytest.raises(ValueError, match=match): + validate_kv_cache_compression_compatibility( + config, + SimpleNamespace(enable_block_reuse=False), + spec_config, + ) + + +def test_compressed_count_is_monotone_and_tracks_confirmed_length(): + manager = _make_triattention(budget=4, beta=4) + target = manager.kv_cache_manager + target._stream = mock.Mock() + target.pp_layers = [0, 1] + cache = SimpleNamespace( + capacity=0, + history_length=2, + is_active=True, + resize=mock.Mock(return_value=True), + ) + target.kv_cache_map = {7: cache} + draft_manager = _make_fake_v2(is_draft=True) + draft_cache = SimpleNamespace(is_active=True, resize=mock.Mock(return_value=True)) + draft_manager.kv_cache_map = {7: draft_cache} + draft_manager._stream = mock.Mock() + manager.draft_kv_cache_manager = draft_manager + # Injected post-construction: mirror the ctor-cached manager-lifetime tail. + manager._draft_protected_tail_capacity = 1 + + request = _make_request(7, py_prompt_len=2, py_num_accepted_draft_tokens=1) + batch = SimpleNamespace(generation_requests=[request]) + + # Every step confirms one sampled token plus one accepted draft token. + uncompressed = 6 + confirmed = uncompressed + cache.capacity = confirmed + previous_published = 0 + eviction_rounds = 0 + with _mocked_eviction_internals(manager) as internals: + for _ in range(6): + uncompressed += 2 + confirmed += 2 + cache.capacity = confirmed + + manager._evict_due_requests(batch) + + published = request.py_num_compressed_tokens + if published > previous_published: + # An eviction round compacted the cache to prompt + budget. + eviction_rounds += 1 + confirmed -= published - previous_published + cache.capacity = confirmed + assert confirmed == 2 + 4 + # The published count equals the uncompressed confirmed logical + # length minus the physical confirmed length, and never decreases. + assert published == uncompressed - confirmed + assert published >= previous_published + previous_published = published + + assert eviction_rounds == 3 + assert previous_published == 12 + # Each round the draft cache shrinks with the target, and the one + # executor call runs on the compression manager, which carries both cache + # managers while handing completion back through the target manager. + assert draft_cache.resize.call_args_list == [mock.call(7, None)] * eviction_rounds + assert len(internals.execute.call_args_list) == eviction_rounds + assert manager.kv_cache_manager is target + assert manager.draft_kv_cache_manager is draft_manager + for call in internals.execute.call_args_list: + assert len(call.args) == 1 + assert call.kwargs == {} + + +def test_request_admission_reserves_score_high_watermark(): + manager = _make_triattention(budget=128, beta=64) + assert manager._selection_width_capacity == 256 + manager._phase = mock.Mock() + manager._selection_width_capacity = 260 + manager._score_token_capacity = 0 + manager._launch_score = None + manager._compaction_done_event = mock.Mock() + manager.kv_cache_manager = SimpleNamespace(max_seq_len=65536, tokens_per_block=64) + + def publish_score_state(*, score_token_capacity): + manager._score_token_capacity = score_token_capacity + manager._launch_score = object() + + manager._build_score_runtime = mock.Mock(side_effect=publish_score_state) + requests = [ + _make_request(1, py_prompt_len=100, py_max_new_tokens=10000), + _make_request(2, py_prompt_len=700, py_max_new_tokens=10), + _make_request(3, py_prompt_len=900, py_max_new_tokens=200), + ] + + for request in requests: + manager.on_request_init(request) + + assert manager._build_score_runtime.call_args_list == [ + mock.call(score_token_capacity=1024), + mock.call(score_token_capacity=2048), + ] + manager._compaction_done_event.synchronize.assert_called_once_with() + assert manager._phase.reserve.call_args_list == [ + mock.call(10101), + mock.call(1101), + ] + + +def test_request_admission_aligns_clamped_score_bucket_to_tile(): + manager = _make_triattention(budget=128, beta=64) + manager._phase = mock.Mock() + manager._selection_width_capacity = 256 + manager._score_token_capacity = 0 + manager._launch_score = None + manager._compaction_done_event = mock.Mock() + manager.kv_cache_manager = SimpleNamespace(max_seq_len=1050, tokens_per_block=128) + manager._build_score_runtime = mock.Mock() + + manager.on_request_init(_make_request(1, py_prompt_len=850, py_max_new_tokens=200)) + + manager._build_score_runtime.assert_called_once_with(score_token_capacity=1152) + manager._compaction_done_event.synchronize.assert_not_called() + + +def test_block_offset_snapshot_width_is_aligned_and_capped(): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _allocate_block_offset_snapshot, + ) + + anchor_pool = torch.empty(1, 2, 1, 32, 4) + manager = SimpleNamespace( + num_pools=1, + tokens_per_block=32, + max_blocks_per_seq=4, + ) + host, device_table = _allocate_block_offset_snapshot( + manager, + anchor_pool, + request_capacity=2, + token_capacity=129, + ) + assert host.shape[-1] == 4 and device_table.shape[-1] == 4 + manager.max_blocks_per_seq = 64 + host, device_table = _allocate_block_offset_snapshot( + manager, + anchor_pool, + request_capacity=2, + token_capacity=129, + ) + assert host.shape[-1] == 8 and device_table.shape[-1] == 8 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_fused_settle_pack.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_fused_settle_pack.py new file mode 100644 index 000000000000..4dc6f9921b79 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_fused_settle_pack.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The TriAttention settle kernel vs a pure-torch integer oracle.""" + +import pytest +import torch + +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import settle_ties + +# Settle geometry rows: the width/keep axes flip the WIDTH and KEEP_COUNT +# static_range trip counts across the 256-lane BLOCK. +_WIDTH_KEEP_CASES = [ + # Small and ragged: rows shorter than the keep count, empty rows. + (21, 5), + # More than one 256-lane block along both the settle and move axes. + (350, 300), +] + + +def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, keep_count): + """Settle tied provisional scores deterministically into ascending ordinals.""" + rows_total, width = scores.shape + for row in range(rows_total): + lanes = [int(i) for i in provisional[row, :keep_count] if int(i) >= 0] + threshold = min((float(scores[row, i]) for i in lanes), default=float("inf")) + length = min(width, int(row_lengths[row])) + row_scores = scores[row, :length].tolist() + greater = [i for i, s in enumerate(row_scores) if s > threshold] + ties = [i for i, s in enumerate(row_scores) if s == threshold] + quota = max(0, keep_count - len(greater)) + selected = sorted(greater + ties[:quota]) + if selected: + prompt = int(row_prompt_offsets[row]) + output[row, : len(selected)] = torch.tensor( + [i + prompt for i in selected], dtype=output.dtype, device=output.device + ) + + +def _selection_rows_for(eviction_mode: str, num_layers: int, num_kv_heads: int) -> int: + if eviction_mode == "union": + return 1 + if eviction_mode == "per_head": + return num_kv_heads + return num_layers * num_kv_heads + + +def _make_settle_inputs(request_count, selection_rows, width, keep_count, seed, device): + """Build a seeded, tied, ragged settle problem with prompt rebasing.""" + rows_total = request_count * selection_rows + generator = torch.Generator(device=device).manual_seed(seed) + # Heavily tied integer scores force the tie-quota emission path. + scores = torch.randint( + -2, 3, (rows_total, width), generator=generator, dtype=torch.int32, device=device + ).to(torch.float32) + # Ragged rows: empty, shorter than the keep count (stale output + # entries survive), and full width. + row_lengths = torch.tensor( + [[0, keep_count - 2, width - 4, width][row % 4] for row in range(rows_total)], + dtype=torch.int32, + device=device, + ) + # Per-request prompt lengths (the kernel indexes them by request). + prompt_offsets = torch.tensor( + [3 * (request % 3) for request in range(request_count)], dtype=torch.int32, device=device + ) + # Stand-in for the CuTE top-k: in-range indices covering the top + # scores of each row with arbitrary tie breaking. + masked = scores.clone() + for row in range(rows_total): + masked[row, int(row_lengths[row]) :] = float("-inf") + provisional = torch.topk(masked, keep_count, dim=1).indices.to(torch.int32).contiguous() + return scores, row_lengths, prompt_offsets, provisional + + +# SELECTION_ROWS is stride/grid arithmetic only (no static branch): one +# single-row and one multi-row mode pin every settle path. +@pytest.mark.parametrize("eviction_mode", ["union", "per_layer_perhead"]) +@pytest.mark.parametrize("width,keep_count", _WIDTH_KEEP_CASES) +def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): + device = torch.device("cuda", torch.cuda.current_device()) + request_count, num_layers, num_kv_heads = 3, 2, 2 + selection_rows = _selection_rows_for(eviction_mode, num_layers, num_kv_heads) + rows_total = request_count * selection_rows + + for seed in range(5): + scores, row_lengths, prompt_offsets, provisional = _make_settle_inputs( + request_count, selection_rows, width, keep_count, seed, device + ) + row_prompt_offsets = prompt_offsets.repeat_interleave(selection_rows) + # Identical stale garbage on both sides so untouched regions must + # match too. + output_stale = torch.randint( + -(2**30), 2**30, (rows_total, keep_count), dtype=torch.int32, device=device + ) + output_reference = output_stale.clone() + _settle_oracle( + scores, row_lengths, row_prompt_offsets, provisional, output_reference, keep_count + ) + + output_actual = output_stale.clone() + settle_ties( + scores, + row_lengths, + prompt_offsets, + provisional, + output_actual, + request_count=request_count, + selection_rows_per_request=selection_rows, + ) + torch.cuda.synchronize(device) + + assert torch.equal(output_actual, output_reference), f"kept ordinals differ (seed {seed})" + + +def test_settle_handles_topk_sentinel_padding(): + """Rows shorter than KEEP_COUNT arrive -1-padded and must settle inertly. + + The production top-k pads a row shorter than KEEP_COUNT with -1 + sentinels (a zero-width padded row is all sentinels). The settle's + threshold gather must skip those lanes -- never touching the score byte + before the row -- while emitting exactly the real ordinals; the output + slots past a short row's length stay untouched (rows that short move + nothing downstream). Full rows must keep byte-identical behavior. + """ + device = torch.device("cuda", torch.cuda.current_device()) + rows_total, width, keep_count = 4, 33, 7 + generator = torch.Generator(device=device).manual_seed(23) + scores = torch.randint( + -2, 3, (rows_total, width), generator=generator, dtype=torch.int32, device=device + ).to(torch.float32) + row_lengths = torch.tensor([0, 3, 7, 33], dtype=torch.int32, device=device) + row_prompt_offsets = torch.tensor([5, 1, 2, 0], dtype=torch.int32, device=device) + + # Provisional rows exactly as the production top-k emits them: rows with + # length <= KEEP_COUNT carry [0..length) then -1 sentinels; longer rows + # carry a dense top-k. + provisional = torch.full((rows_total, keep_count), -1, dtype=torch.int32, device=device) + for row, length in enumerate(row_lengths.tolist()): + if length <= keep_count: + provisional[row, :length] = torch.arange(length, dtype=torch.int32, device=device) + else: + masked = scores[row].clone() + masked[length:] = float("-inf") + provisional[row] = torch.topk(masked, keep_count).indices.to(torch.int32) + + stale = 0x5EED + output = torch.full((rows_total, keep_count), stale, dtype=torch.int32, device=device) + settle_ties( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output, + request_count=rows_total, + selection_rows_per_request=1, + ) + torch.cuda.synchronize(device) + + for row, length in enumerate(row_lengths.tolist()): + prompt = int(row_prompt_offsets[row]) + emitted = min(length, keep_count) + if length > keep_count: + # Reference keep set: score-descending with lowest-index ties, + # emitted as ascending absolute ordinals. + order = sorted(range(length), key=lambda i: (-float(scores[row, i]), i)) + expected = sorted(order[:keep_count]) + else: + expected = list(range(length)) + expected_row = torch.tensor( + [ordinal + prompt for ordinal in expected], dtype=torch.int32, device=device + ) + assert torch.equal(output[row, :emitted], expected_row), f"row {row}" + assert (output[row, emitted:] == stale).all(), f"row {row} tail" diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py new file mode 100644 index 000000000000..2b45d3beb52d --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -0,0 +1,561 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the TriAttention compression-manager pipeline. + +Config, construction, eviction lifecycle, page-table staging, and admission-sized +score state; the manager publishes evicted counts via +``LlmRequest.py_num_compressed_tokens``. Draft contracts live in +``test_triattention_draft_cocompaction.py``. +""" + +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from conftest import make_bare_staging as _make_bare_staging +from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_request as _make_request +from conftest import make_staging_manager as _make_staging_manager +from conftest import make_tri_config as _make_tri_config +from conftest import make_triattention as _make_triattention +from conftest import mocked_eviction_internals as _mocked_eviction_internals + +# TriAttention lives in the kv_cache_compression package. It exposes only the +# compression manager -- no attention classes or KV-cache-manager subclass. +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, +) + +# Framework base class lives in pyexecutor.resource_manager; the factory lives +# in pyexecutor._util (next to _create_kv_cache_manager), matching #15106. +from tensorrt_llm._torch.pyexecutor._util import ( + create_kv_cache_compression_manager, + validate_kv_cache_compression_compatibility, +) + + +@pytest.fixture +def flat_calibration_pt(tmp_path): + """Build a minimal valid calibration ``.pt`` in our flat runtime schema.""" + path = tmp_path / "tri_calib.pt" + calibration = { + "E_q": torch.zeros(2, 2, 4, dtype=torch.complex64), + "E_q_norm": torch.ones(2, 2, 4, dtype=torch.float32), + "omega": torch.arange(4, dtype=torch.float32), + "freq_scale_sq": torch.ones(4, dtype=torch.float32), + } + torch.save(calibration, path) + return str(path) + + +def _make_hf_config(**values): + """Expose the normalized Hugging Face text-config contract.""" + text_config = SimpleNamespace(to_dict=lambda: dict(values)) + return SimpleNamespace(get_text_config=lambda: text_config) + + +class TestConfigAndFactory: + def test_factory_allows_block_reuse_and_propagates_config_fields(self): + # The factory contract is independent of GPU-owned persistent buffers. + fake_v2 = _make_fake_v2(enable_block_reuse=True) + cfg = _make_tri_config(budget=32, beta=16, eviction_mode="per_head") + validate_kv_cache_compression_compatibility( + cfg, + SimpleNamespace(enable_block_reuse=True), + None, + ) + with ( + mock.patch( + "tensorrt_llm._torch.pyexecutor._util.is_sm_100f", + return_value=True, + ), + mock.patch.object( + TriAttentionCompressionManager, "_initialize_eviction_state" + ) as initialize, + ): + mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) + assert isinstance(mgr, TriAttentionCompressionManager) + assert mgr.budget == 32 + assert mgr.beta == 16 + assert mgr.eviction_mode == "per_head" + assert mgr.kv_cache_manager is fake_v2 + assert fake_v2.kv_compression_manages_history + assert cfg.changes_physical_kv_length + assert cfg.supports_block_reuse() + assert not cfg.supports_speculative_decoding() + initialize.assert_called_once_with() + + +class TestTriAttentionCompressionManager: + def test_loads_flat_pt(self, flat_calibration_pt): + mgr = _make_triattention() + mgr.calibration_path = flat_calibration_pt + mgr.model_path = None + mgr._load_calibration() + + assert torch.equal(mgr._omega, torch.arange(4, dtype=torch.float32)) + assert torch.equal(mgr._freq_scale_sq, torch.ones(4)) + assert torch.equal(mgr._calibration_q_real, torch.zeros(2, 2, 4)) + assert torch.equal(mgr._calibration_q_imag, torch.zeros(2, 2, 4)) + assert torch.equal(mgr._calibration_mlr_coef, torch.ones(2, 2, 4)) + + def test_loads_official_layout(self, tmp_path): + # PRODUCT CONTRACT: the official R-KV {metadata, stats} layout is + # converted to the flat runtime schema at load; rope tables derive + # from the model config. + pytest.importorskip("transformers") + num_layers, num_heads, freq_count = 2, 2, 4 + stats, sampled = {}, [] + for layer in range(num_layers): + for head in range(num_heads): + stats[f"layer{layer:02d}_head{head:02d}"] = { + "q_mean_real": torch.full((freq_count,), float(10 * layer + head)), + "q_mean_imag": torch.full((freq_count,), float(layer - head)), + "q_abs_mean": torch.full((freq_count,), float(1 + layer + head)), + } + sampled.append((layer, head)) + path = tmp_path / "official.pt" + torch.save({"metadata": {"sampled_heads": sampled}, "stats": stats}, path) + mgr = _make_triattention() + mgr.calibration_path = str(path) + config = _make_hf_config(rope_parameters={"rope_type": "default", "rope_theta": 10000.0}) + + with mock.patch("transformers.AutoConfig.from_pretrained", return_value=config): + mgr._load_calibration() + + assert mgr._calibration_q_real.shape == (num_layers, num_heads, freq_count) + torch.testing.assert_close( + mgr._calibration_q_real[1, 0].cpu(), torch.full((freq_count,), 10.0) + ) + torch.testing.assert_close( + mgr._calibration_q_imag[1, 0].cpu(), torch.full((freq_count,), 1.0) + ) + torch.testing.assert_close( + mgr._calibration_mlr_coef[1, 1].cpu(), torch.full((freq_count,), -8.0) + ) + assert mgr._omega.numel() == freq_count + idx = torch.arange(0, 2 * freq_count, 2, dtype=torch.float32) + torch.testing.assert_close(mgr._omega.cpu(), 1.0 / (10000.0 ** (idx / (2 * freq_count)))) + assert torch.equal(mgr._freq_scale_sq.cpu(), torch.ones(freq_count)) + + def test_rope_tables_resolve_theta_and_attention_factor(self, tmp_path): + # transformers>=5.5 folds rope_theta into ``rope_parameters`` and drops + # "default" from ROPE_INIT_FUNCTIONS; resolution must find the true + # theta and scaled-rope attention factor on both config generations + # (the silent base-10000 analytic fallback was the B1 bug). + pytest.importorskip("transformers") + import json + + def config_dir(name, body): + d = tmp_path / name + d.mkdir() + (d / "config.json").write_text(json.dumps(body)) + return str(d) + + common = { + "model_type": "qwen3", + "architectures": ["Qwen3ForCausalLM"], + "hidden_size": 256, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_hidden_layers": 2, + "head_dim": 64, + "max_position_embeddings": 8192, + } + plain = config_dir("plain", {**common, "rope_theta": 1000000.0}) + yarn = config_dir( + "yarn", + { + **common, + "rope_theta": 150000.0, + "rope_scaling": { + "rope_type": "yarn", + "factor": 4.0, + "original_max_position_embeddings": 2048, + "attention_factor": 1.25, + }, + }, + ) + mgr = _make_triattention() + freq_count = 32 + calibration_path = tmp_path / "official.pt" + torch.save( + { + "metadata": {"sampled_heads": [(0, 0)]}, + "stats": { + "layer00_head00": { + "q_mean_real": torch.zeros(freq_count), + "q_mean_imag": torch.zeros(freq_count), + "q_abs_mean": torch.ones(freq_count), + } + }, + }, + calibration_path, + ) + mgr.calibration_path = str(calibration_path) + + mgr.model_path = plain + mgr._load_calibration() + omega = mgr._omega + freq_scale_sq = mgr._freq_scale_sq + idx = torch.arange(0, 64, 2, dtype=torch.float32) + torch.testing.assert_close(omega, (1.0 / (1000000.0 ** (idx / 64)))[:freq_count]) + assert torch.equal(freq_scale_sq, torch.ones(freq_count)) + + mgr.model_path = yarn + mgr._load_calibration() + omega_yarn = mgr._omega + freq_scale_sq_yarn = mgr._freq_scale_sq + # Routed through transformers' yarn init: the explicit attention + # factor lands squared, and the ladder leaves the plain-theta curve. + torch.testing.assert_close(freq_scale_sq_yarn, torch.full((freq_count,), 1.25**2)) + assert not torch.allclose(omega_yarn, omega) + + +class TestCompressedTokenPublication: + # The monotone publication contract is covered end to end in + # test_triattention_draft_cocompaction.py. + + def test_identity_selection_is_filtered_before_launch(self): + # Identity requests (seq_len == prompt + budget) are the pre-launch + # owner no-op: nothing launches and nothing is published. + manager = _make_triattention(budget=4, beta=4) + manager.kv_cache_manager._stream = mock.Mock() + request = _make_request(7, py_prompt_len=2) + # seq_len == prompt + budget: the due filter must drop the request. + cache = SimpleNamespace( + capacity=6, history_length=0, is_active=True, resize=mock.Mock(return_value=True) + ) + manager.kv_cache_manager.kv_cache_map = {7: cache} + + with _mocked_eviction_internals(manager) as internals: + manager._evict_due_requests(SimpleNamespace(generation_requests=[request])) + + internals.execute.assert_not_called() + assert request.py_num_compressed_tokens == 0 + cache.resize.assert_not_called() + + +class TestEvictionLifecycle: + def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): + manager = _make_triattention() + batch = SimpleNamespace( + context_requests=[], + context_requests_last_chunk=[], + generation_requests=[_make_request(7)], + ) + with mock.patch.object(manager, "_evict_due_requests") as evict_due: + manager.prepare_resources(batch) + evict_due.assert_not_called() + + manager.update_resources(batch) + + evict_due.assert_called_once_with(batch) + + @staticmethod + def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft_tokens=0): + # The growth and protected-tail capacity constants snapshot the + # manager at construction, so the reserve widths are set up front. + request = _make_request( + 7, + py_prompt_len=1024, + max_beam_num_tokens=seq_len + 1, + ) + batch = SimpleNamespace(generation_requests=[request]) + fake_v2 = _make_fake_v2() + fake_v2.num_extra_kv_tokens = num_extra_kv_tokens + fake_v2._kv_reserve_draft_tokens = kv_reserve_draft_tokens + with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): + mgr = TriAttentionCompressionManager(_make_tri_config(budget=8), fake_v2) + cache = SimpleNamespace( + capacity=seq_len, + history_length=1024, + is_active=True, + resize=mock.Mock(return_value=True), + ) + mgr.kv_cache_manager = SimpleNamespace( + get_buffers=lambda *args, **kwargs: None, + kv_cache_map={7: cache}, + pp_layers=[0, 1], + _stream=mock.Mock(), + num_extra_kv_tokens=num_extra_kv_tokens, + _kv_reserve_draft_tokens=kv_reserve_draft_tokens, + ) + mgr.beta = 128 + mgr.budget = 4096 + mgr._selection_width_capacity = mgr.budget + mgr.beta + 1 + return mgr, request, batch + + def test_suspended_cache_defers_that_request_pre_launch(self): + # A suspended cache is a legal overlap-scheduler transient: that + # request defers while the rest of the request group proceeds, then + # catches up to the missed cadence boundary when it resumes. + manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 128) + second_request = _make_request(8, py_prompt_len=1024) + second_cache = SimpleNamespace( + capacity=1024 + 4096 + 128, + history_length=1024, + is_active=False, + resize=mock.Mock(return_value=True), + ) + manager.kv_cache_manager.kv_cache_map[8] = second_cache + batch = SimpleNamespace(generation_requests=[first_request, second_request]) + + with _mocked_eviction_internals(manager) as internals: + manager._evict_due_requests(batch) + # Only the active request launched in the first round. + eviction_requests = internals.execute.call_args.args[0] + assert [item.request.py_request_id for item in eviction_requests] == [7] + assert second_request.py_num_compressed_tokens == 0 + + second_cache.is_active = True + # Resumption executes one more token before the next final update. + second_cache.capacity += 1 + manager._evict_due_requests(SimpleNamespace(generation_requests=[second_request])) + + resumed = internals.execute.call_args.args[0] + assert [item.request.py_request_id for item in resumed] == [8] + assert second_request.py_num_compressed_tokens == 129 + second_cache.resize.assert_called_once_with(1024 + 4096, None) + + def test_deferred_eviction_checks_the_compiled_selection_width(self): + manager, request, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 128 + 2) + + with _mocked_eviction_internals(manager) as internals: + with pytest.raises(RuntimeError, match="selection width"): + manager._evict_due_requests(batch) + + internals.execute.assert_not_called() + + # Accepted draft tokens may cross the same cadence boundary; they do not + # change the fixed overlap reservation. + @pytest.mark.parametrize("accepted", [0, 3]) + def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): + confirmed = 1024 + 4096 + 128 + reserve = 2 + current_growth = 4 + tail = reserve + current_growth + retained = 1024 + 4096 + # Growth constant = 1 + _kv_reserve_draft_tokens for batch members. + mgr, request, batch = self._make_due_decode_request( + seq_len=confirmed, + num_extra_kv_tokens=reserve, + kv_reserve_draft_tokens=current_growth - 1, + ) + request.py_num_accepted_draft_tokens = accepted + cache = mgr.kv_cache_manager.kv_cache_map[7] + cache.capacity = confirmed + tail + mgr.on_generation_step_begin(SimpleNamespace(generation_requests=[request])) + draft_manager = _make_fake_v2(is_draft=True) + draft_cache = SimpleNamespace(is_active=True, resize=mock.Mock(return_value=True)) + draft_manager.kv_cache_map = {7: draft_cache} + draft_manager._stream = mock.Mock() + mgr.draft_kv_cache_manager = draft_manager + # Injected post-construction: mirror the ctor-cached manager-lifetime tail. + mgr._draft_protected_tail_capacity = 1 + + with _mocked_eviction_internals(mgr) as internals: + mgr._evict_due_requests(batch) + + # Tail excluded from the source length; keep target = prompt + budget. + internals.execute.assert_called_once() + (launched,) = internals.execute.call_args.args[0] + assert launched.request is request + assert launched.target_cache is cache + assert launched.draft_cache is draft_cache + assert launched.source_length == confirmed + assert launched.target_tail_length == tail + assert request.py_num_compressed_tokens == confirmed - retained + cache.resize.assert_called_once_with(retained + tail, None) + draft_cache.resize.assert_called_once_with(retained + 1, None) + + def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): + # The due-branch source length must come from the physical capacity + # ledger (capacity minus the protected tail), never the logical length. + manager = _make_triattention(beta=128) + compressed_tokens = manager.beta - manager.budget + # Reachable second cadence boundary, within the compiled selection span. + physical_confirmed = 1024 + manager.budget + manager.beta + cache = SimpleNamespace( + capacity=physical_confirmed, + history_length=1024, + is_active=True, + resize=mock.Mock(return_value=True), + ) + manager.kv_cache_manager.kv_cache_map = {7: cache} + manager.kv_cache_manager.pp_layers = [0, 1] + request = _make_request( + 7, + py_prompt_len=1024, + py_num_compressed_tokens=compressed_tokens, + max_beam_num_tokens=physical_confirmed + compressed_tokens + 1, + py_draft_tokens=[1, 2, 3, 4], + ) + + with _mocked_eviction_internals(manager) as internals: + manager._evict_due_requests(SimpleNamespace(generation_requests=[request])) + + eviction_requests = internals.execute.call_args.args[0] + assert eviction_requests[0].source_length == physical_confirmed + assert request.py_num_compressed_tokens == ( + compressed_tokens + physical_confirmed - 1024 - manager.budget + ) + cache.resize.assert_called_once_with(1024 + manager.budget, None) + + @pytest.mark.parametrize("spec_mode", ["mtp", "eagle3"]) + def test_one_model_draft_co_compression_is_accepted(self, spec_mode): + draft_manager = _make_fake_v2(is_draft=True) + with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): + TriAttentionCompressionManager( + _make_tri_config(budget=8), + _make_fake_v2(), + draft_kv_cache_manager=draft_manager, + ) + + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_compatibility + from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig, MTPDecodingConfig + + spec_config = ( + MTPDecodingConfig(max_draft_len=1) + if spec_mode == "mtp" + else Eagle3DecodingConfig( + max_draft_len=1, + speculative_model="draft", + eagle3_one_model=True, + ) + ) + + validate_kv_cache_compression_compatibility( + _make_tri_config(budget=8), + SimpleNamespace(enable_block_reuse=False), + spec_config, + ) + + +class TestFixedScoreMetadata: + def test_union_forces_normalized_scores(self): + # Union eviction always z-normalizes: False is coerced to True at construction. + triattention = _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) + assert triattention.normalize_scores is True + + def test_bulk_page_table_copy_snapshots_on_current_stream(self): + """Snapshot and order bulk page-table copies on the caller's current stream.""" + device = torch.device("cuda", torch.cuda.current_device()) + current_stream = torch.cuda.current_stream(device) + host_table = torch.zeros( + 1, + 2, + 2, + 12, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + host_table[0, 0, 0, :5] = torch.tensor([3, 4, 5, 6, 7], dtype=torch.int32) + host_table[0, 1, 0, :5] = torch.tensor([8, 9, 10, 11, 12], dtype=torch.int32) + selected_slot = [0] + + def gather_k_block_offsets(source, destination, request_ids, num_blocks): + assert request_ids == [7] + destination[:, :1, 0, :num_blocks].copy_( + source[:, selected_slot[0], 0, :num_blocks].unsqueeze(1) + ) + + gather = mock.Mock(side_effect=gather_k_block_offsets) + staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) + manager = _make_staging_manager(host_table, gather, current_stream) + + def stage_once(): + # Raises on any staging failure; success returns None. + with torch.cuda.stream(current_stream): + staging._stage_block_offset_snapshot( + manager, + [7], + staging._block_offsets_host, + staging._block_offsets_device, + ) + + # Round 1: mutate the host table and the slot assignment right after + # staging. The staged result must still reflect the gathered snapshot. + with mock.patch.object( + torch, + "index_select", + side_effect=AssertionError("page-table staging used torch.index_select"), + ): + stage_once() + assert staging._block_offsets_host.shape == (1, 1, 2, 8) + host_table[0, 0, 0, :5] = torch.tensor([13, 14, 15, 16, 17], dtype=torch.int32) + selected_slot[0] = 1 + current_stream.synchronize() + + assert staging._block_offsets_device[0, 0, 0, :5].tolist() == [6, 8, 10, 12, 14] + assert staging._block_offsets_device[0, 0, 1, :5].tolist() == [7, 9, 11, 13, 15] + + # Round 2: same contract on a re-staged request group. + host_table[0, 0, 0, :5] = torch.tensor([18, 19, 20, 21, 22], dtype=torch.int32) + selected_slot[0] = 0 + stage_once() + host_table[0, 0, 0, :5] = torch.tensor([23, 24, 25, 26, 27], dtype=torch.int32) + selected_slot[0] = 1 + current_stream.synchronize() + + assert staging._block_offsets_device[0, 0, 0, :5].tolist() == [36, 38, 40, 42, 44] + assert staging._block_offsets_device[0, 0, 1, :5].tolist() == [37, 39, 41, 43, 45] + + # Round 3: a consumer queued before restaging sees the prior table. + selected_slot[0] = 0 + snapshot = torch.empty_like(staging._block_offsets_device) + snapshot.copy_(staging._block_offsets_device) + + stage_once() + current_stream.synchronize() + + assert snapshot[0, 0, 0, :5].tolist() == [36, 38, 40, 42, 44] + assert staging._block_offsets_device[0, 0, 0, :5].tolist() == [46, 48, 50, 52, 54] + + +class TestKernelMaskedSwa: + @pytest.mark.parametrize("budget,fits_window", [(128, True), (127, False)]) + def test_attention_layers_use_local_config_and_validate_window(self, budget, fits_window): + mgr = _make_triattention() + mgr.model_path = "/models/gpt-oss" + mgr.budget = budget + mgr._global_layers = [0, 1, 2, 3] + config = _make_hf_config( + layer_types=[ + "sliding_attention", + "full_attention", + "sliding_attention", + "full_attention", + ], + sliding_window=128, + ) + + with mock.patch("transformers.AutoConfig.from_pretrained", return_value=config) as load: + if not fits_window: + # The decode budget must cover the kernel-masked SWA window. + with pytest.raises(ValueError, match="budget=127"): + mgr._resolve_attention_layers() + return + dense, sliding, window = mgr._resolve_attention_layers() + + load.assert_called_once_with( + "/models/gpt-oss", trust_remote_code=True, local_files_only=True + ) + assert dense == [1, 3] + assert sliding == [0, 2] + assert window == 128 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py new file mode 100644 index 000000000000..a17a1b4c6731 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -0,0 +1,620 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +import pytest +import torch +from conftest import make_cute_buffers as _make_cute_buffers +from conftest import make_eviction_request as _make_eviction_request +from conftest import make_ramp_pools as _make_ramp_pools +from conftest import make_request as _make_request +from conftest import make_staging_manager as _make_staging_manager +from conftest import rect_to_score_scratch as _rect_to_score_scratch + +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, +) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + reduce_per_head_scores, +) + + +def _require_cute_topk_op() -> None: + """The CuTE TopK operation is a hard prerequisite for these tests.""" + assert hasattr(torch.ops.trtllm, "cute_dsl_indexer_topk_decode"), ( + "CuTE TopK operation is not loaded" + ) + + +# Tests that launch real scores run the SM100 CuTe score kernel -- the only +# score path -- so they are SM100-only, like the production feature itself. +requires_sm100 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="TriAttention score requires SM100", +) + + +def _make_selection_buffers( + *, + eviction_mode, + width, + keep_count, + device, + max_requests, + num_layers=1, + num_query_heads=1, + num_kv_heads=1, +): + """Allocate the production selection buffers without score or compaction.""" + tri = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) + tri.eviction_mode = eviction_mode + tri._request_capacity = max_requests + tri._selection_width_capacity = width + tri.budget = keep_count + tri._num_layers = num_layers + tri._num_q_heads = num_query_heads + tri._num_kv_heads = num_kv_heads + tri._prompt_lengths_device = torch.zeros(max_requests, dtype=torch.int32, device=device) + tri._allocate_selection_buffers(device, tp_size=1) + return tri + + +def _select_per_head(tri, scores, *, normalize_scores): + """The per-head selection flow: reduce kernels, then top-k settle.""" + request_count, _, _, width = scores.shape + score_scratch, prompt_lengths = _rect_to_score_scratch(scores, tri._num_kv_heads) + reduce_per_head_scores( + score_scratch, + tri._decode_lengths_device, + prompt_lengths, + tri._row_mean, + tri._row_inv_std, + tri._selection_scores_rows, + tri._selection_row_lengths, + request_count=request_count, + padded_head_columns=8, + score_token_capacity=width, + per_layer=tri.eviction_mode == "per_layer_perhead", + normalize_scores=normalize_scores, + ) + tri._select_kept_ordinals(tri._request_capacity) + + +def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: + values = row[:width].tolist() + selected = sorted(range(width), key=lambda index: (-values[index], index)) + return torch.tensor(selected[:keep_count], dtype=torch.int32, device=row.device) + + +def _per_head_keep_oracle( + scores: torch.Tensor, + decode_lengths: torch.Tensor, + keep_count: int, + eviction_mode: str, + normalize_scores: bool, +) -> torch.Tensor: + """Independent torch implementation of per-head selection.""" + request_count, num_layers, num_query_heads, width = scores.shape + num_kv_heads = 2 + rows = [] + for request in range(request_count): + decode_length = int(decode_lengths[request]) + valid = scores[request, ..., :decode_length].clone() + if normalize_scores: + mean = valid.mean(dim=-1, keepdim=True) + valid = valid - mean + std = valid.norm(dim=-1, keepdim=True) / (decode_length**0.5) + valid = valid / std.clamp_min(1e-6) + grouped = valid.view( + num_layers, num_kv_heads, num_query_heads // num_kv_heads, decode_length + ).amax(dim=2) + if eviction_mode == "per_head": + selection = grouped.mean(dim=0) + else: + selection = grouped.reshape(num_layers * num_kv_heads, decode_length) + rows.append( + torch.stack( + [ + torch.sort(_stable_topk(row, decode_length, keep_count)).values + for row in selection + ] + ) + ) + return torch.stack(rows) + + +@pytest.mark.parametrize( + "eviction_mode,normalize_scores", + [("per_head", True), ("per_layer_perhead", False)], +) +def test_per_head_selection_matches_torch_oracle_on_selector_stream( + eviction_mode, normalize_scores +): + _require_cute_topk_op() + request_count, layers, query_heads, kv_heads = 2, 3, 4, 2 + width, keep_count = 96, 64 + generator = torch.Generator().manual_seed(41) + scores_cpu = torch.randint( + -4, + 5, + (request_count, layers, query_heads, width), + generator=generator, + dtype=torch.int32, + ).to(torch.float32) + decode_lengths = torch.tensor([83, 91], dtype=torch.int32) + + expected = _per_head_keep_oracle( + scores_cpu, decode_lengths, keep_count, eviction_mode, normalize_scores + ) + + device = torch.device("cuda", torch.cuda.current_device()) + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + tri = _make_selection_buffers( + eviction_mode=eviction_mode, + width=width, + keep_count=keep_count, + device=device, + max_requests=request_count, + num_layers=layers, + num_query_heads=query_heads, + num_kv_heads=kv_heads, + ) + tri._decode_lengths_device.copy_(decode_lengths.to(device)) + scores = scores_cpu.to(device) + keep_shape = (request_count, tri._selection_rows_per_request, keep_count) + _select_per_head(tri, scores, normalize_scores=normalize_scores) + first = tri._kept_ordinal_rows.view(keep_shape).cpu() + _select_per_head(tri, scores, normalize_scores=normalize_scores) + second = tri._kept_ordinal_rows.view(keep_shape).cpu() + stream.synchronize() + + assert torch.equal(first, expected) + assert torch.equal(second, expected) + + +@pytest.mark.parametrize("keep_count,width", [(4, 64), (8192, 9216)]) +def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): + # Tied integer scores, ragged widths, per-request prompt rebase; the + # 8192-keep row is the large-k coverage. + _require_cute_topk_op() + device = torch.device("cuda", torch.cuda.current_device()) + prompt_len = 17 + request_count, rows = 2, 4 + generator = torch.Generator(device=device).manual_seed(keep_count) + scores = torch.randint( + -4, + 5, + (request_count, rows, width), + generator=generator, + dtype=torch.int32, + device=device, + ).to(torch.float32) + decode_lengths = (width, width - 32) + tri = _make_selection_buffers( + eviction_mode="union", + width=width, + keep_count=keep_count, + device=device, + max_requests=request_count, + ) + tri._decode_lengths_device.copy_(torch.tensor(decode_lengths, dtype=torch.int32, device=device)) + tri._prompt_lengths_device[:request_count].copy_( + torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) + ) + tri._selection_scores_rows.copy_(scores.amax(dim=1)) + tri._select_kept_ordinals(tri._request_capacity) + actual = tri._kept_ordinal_rows.cpu() + + combined = scores.amax(dim=1).cpu() + for request, decode_length in enumerate(decode_lengths): + expected_decode = torch.sort( + _stable_topk(combined[request], decode_length, keep_count).to(torch.int32) + prompt_len + ).values + assert torch.equal(actual[request], expected_decode) + + +@pytest.mark.parametrize("per_layer", [False, True]) +@pytest.mark.parametrize("normalize_scores", [False, True]) +def test_per_head_reduction_matches_ragged_torch_reference(per_layer, normalize_scores): + device = torch.device("cuda", torch.cuda.current_device()) + request_count, layers, query_heads, kv_heads, width = 2, 3, 4, 2, 97 + generator = torch.Generator(device=device).manual_seed(29) + scores = torch.randn( + request_count, + layers, + query_heads, + width, + generator=generator, + dtype=torch.float32, + device=device, + ) + decode_lengths = torch.tensor([83, 91], dtype=torch.int32, device=device) + row_mean = torch.empty( + request_count, layers, query_heads, 1, dtype=torch.float32, device=device + ) + row_inv_std = torch.empty_like(row_mean) + selection_rows = layers * kv_heads if per_layer else kv_heads + # Canonical row-major buffers, exactly like the product allocation. + selection_scores_rows = torch.empty( + request_count * selection_rows, width, dtype=torch.float32, device=device + ) + selection_row_lengths = torch.empty( + request_count * selection_rows, dtype=torch.int32, device=device + ) + + score_scratch, prompt_lengths = _rect_to_score_scratch(scores, kv_heads) + reduce_per_head_scores( + score_scratch, + decode_lengths, + prompt_lengths, + row_mean, + row_inv_std, + selection_scores_rows, + selection_row_lengths, + request_count=request_count, + padded_head_columns=8, + score_token_capacity=width, + per_layer=per_layer, + normalize_scores=normalize_scores, + ) + torch.cuda.synchronize(device) + + selection_scores = selection_scores_rows.view(request_count, selection_rows, width) + assert torch.equal( + selection_row_lengths.view(request_count, selection_rows).cpu(), + decode_lengths.cpu().view(request_count, 1).expand(-1, selection_rows), + ) + query_group_size = query_heads // kv_heads + for request, decode_length in enumerate(decode_lengths.tolist()): + valid = scores[request, :, :, :decode_length] + if normalize_scores: + mean = valid.mean(dim=-1, keepdim=True) + std = torch.linalg.vector_norm(valid - mean, dim=-1, keepdim=True) + std = (std / decode_length**0.5).clamp_min(1e-6) + valid = (valid - mean) / std + grouped = valid.view(layers, kv_heads, query_group_size, decode_length).amax(dim=2) + expected = grouped if per_layer else grouped.mean(dim=0) + expected = expected.reshape(selection_rows, decode_length) + assert torch.allclose( + selection_scores[request, :, :decode_length], + expected, + rtol=2e-5, + atol=2e-5, + ) + assert torch.isneginf(selection_scores[request, :, decode_length:]).all() + + +@requires_sm100 +def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): + """Keep score and compaction layer axes aligned across interleaved V2 pools.""" + pytest.importorskip("cutlass") + + device = torch.device("cuda", torch.cuda.current_device()) + num_layers = 3 + # Bucket capacity aligned to the 64-token compute tile; request stays 8. + bucket_capacity = 64 + seq_len = 8 + keep_count = 2 + # GQA group 8, zero calibration query, shared MLR: every head row + # carries the same |K|-driven score. + num_q_heads = 8 + # The two tables map the storage groups onto different physical pages; + # the layer-order alignment below depends on it. + tokens_per_block = 32 + head_dim = 64 + num_freqs = head_dim // 2 + page_tables = ( + torch.tensor([[1, 0]], dtype=torch.int32, device=device), + torch.tensor([[0, 1]], dtype=torch.int32, device=device), + ) + layer_tables = (page_tables[0], page_tables[1], page_tables[0]) + score_values = ( + (1, 8, 2, 3, 4, 5, 9, 6), + (2, 3, 8, 4, 5, 9, 6, 7), + (3, 4, 5, 9, 6, 7, 8, 10), + ) + expected_keep = torch.tensor([[[1, 6], [2, 5], [3, 7]]], dtype=torch.int32, device=device) + + pools = _make_ramp_pools(num_layers, num_kv_heads=1, pages=2, device=device) + for pool, (table, values) in zip(pools, zip(layer_tables, score_values)): + for token, value in enumerate(values): + page = int(table[0, token // tokens_per_block]) + slot = token % tokens_per_block + pool[page, 0, 0, slot, 0] = value + pool[page, 0, 0, slot, num_freqs] = 0 + initial_pools = [pool.clone() for pool in pools] + + q_real = torch.zeros(num_layers, num_q_heads, num_freqs, dtype=torch.float32, device=device) + q_imag = torch.zeros_like(q_real) + mlr_coef = torch.zeros_like(q_real) + mlr_coef[:, :, 0] = 1 + freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) + freq_scale_sq[0] = 1 + tri = _make_cute_buffers( + eviction_mode="per_layer_perhead", + layer_pools=pools, + max_requests=1, + seq_len=bucket_capacity, + num_q_heads=num_q_heads, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), + offsets=torch.zeros(1, dtype=torch.float32, device=device), + decode_width=bucket_capacity, + keep_count=keep_count, + layer_pool_ids=[0, 1, 0], + normalize_scores=False, + ) + # No SWA in this layout: no window, no rebase row for the phase gather. + assert tri._swa_window is None + assert tri._swa_destination_bases is None + # Native V2 staging contract: [pool, request, K/V, block] int32 pair with + # a 4-aligned block width (PackedInt copy ABI) and a pinned host snapshot. + assert tri._block_offsets_host.shape == tri._block_offsets_device.shape + assert tri._block_offsets_host.shape[:3] == (2, 1, 2) + assert tri._block_offsets_host.shape[-1] % 4 == 0 + assert tri._block_offsets_host.dtype == tri._block_offsets_device.dtype == torch.int32 + assert tri._block_offsets_host.is_contiguous() and tri._block_offsets_device.is_contiguous() + assert tri._block_offsets_host.is_pinned() + + # Stage through the round executor: the gather double writes both + # page-table slots' K page ids and the bulk copy encodes the K/V rows; + # the derived move offsets stage the buffers' own contract (keep_count + # moves per request, no protected tail). + def gather_k_block_offsets(host_table, source, request_ids, num_blocks): + assert request_ids == [7] + source[..., 0, :].zero_() + source[0, 0, 0, :2].copy_(page_tables[0][0].cpu()) + source[1, 0, 0, :2].copy_(page_tables[1][0].cpu()) + + manager = _make_staging_manager( + torch.zeros(2, 1, 2, 8, dtype=torch.int32), + gather_k_block_offsets, + torch.cuda.Stream(device=device), + num_slots=2, + ) + eviction_requests = [_make_eviction_request(request_id=7, source_length=seq_len)] + tri.kv_cache_manager = manager + tri._execute_eviction_round(eviction_requests) + assert torch.equal(tri._kept_ordinal_rows.view_as(expected_keep), expected_keep) + torch.cuda.synchronize(device) + + for before_pool, after_pool, table, layer in zip( + initial_pools, pools, layer_tables, range(num_layers) + ): + pages = table[0].to(torch.long) + # The logical view spans both pages (2 * tokens_per_block slots); the + # scored sequence occupies its first seq_len positions. + before = before_pool[pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, head_dim) + after = after_pool[pages].permute(1, 2, 0, 3, 4).reshape_as(before) + selected = expected_keep[0, layer].to(torch.long) + assert torch.equal(after[:, :, :keep_count], before.index_select(2, selected)) + + +@requires_sm100 +def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): + """Preserve bytes, tails, and V2 page reuse across two real eviction rounds.""" + pytest.importorskip("cutlass") + import tensorrt_llm + import tensorrt_llm.bindings + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import KvCacheConfig + from tensorrt_llm.mapping import Mapping + + device = torch.device("cuda", torch.cuda.current_device()) + request_id = 7 + prompt_len = 2 + # Bucket == confirmed length, 64-token-tile aligned; tail rides beyond. + seq_len = 64 + protected_tail = 2 + compacted_capacity = 36 + tokens_per_block = 32 + head_dim = 64 + num_freqs = head_dim // 2 + keep_count = compacted_capacity - prompt_len - protected_tail + manager = KVCacheManagerV2( + KvCacheConfig( + max_tokens=seq_len + protected_tail, + enable_block_reuse=False, + host_cache_size=0, + max_util_for_resume=1.0, + ), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=seq_len + protected_tail, + max_batch_size=2, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=tensorrt_llm.bindings.DataType.BF16, + vocab_size=128, + ) + + requests = [] + temporary_requests = [] + try: + created = manager.add_dummy_requests( + [request_id], + [seq_len + protected_tail], + ) + assert created is not None + requests = created + cache = manager.kv_cache_map[request_id] + assert cache.resize(seq_len + protected_tail, prompt_len) + manager.kv_compression_manages_history = True + pool = manager.get_buffers(0, kv_layout="HND") + + def page_ids(owner: int) -> torch.Tensor: + return torch.tensor( + manager.get_batch_cache_indices([owner])[0], + dtype=torch.long, + device=device, + ) + + def snapshot(length: int) -> torch.Tensor: + pages = page_ids(request_id) + return ( + pool.index_select(0, pages) + .permute(1, 2, 0, 3, 4) + .reshape(2, 1, -1, head_dim)[:, :, :length] + .clone() + ) + + def write_token(token: int, score: float) -> None: + pages = page_ids(request_id) + page = pages[token // tokens_per_block] + offset = token % tokens_per_block + # Shifted mod-251 ramp: bf16-exact, distinct per token. + payload = ( + ((torch.arange(2 * head_dim, dtype=torch.int32, device=device) + token * 37) % 251) + .reshape(2, head_dim) + .to(torch.bfloat16) + ) + payload[0, 0] = score + payload[0, num_freqs] = 0 + pool[page, :, 0, offset].copy_(payload) + + # Score mirror; 7 is invertible mod 64 so selection is tie-free. + token_scores = [0] * (seq_len + protected_tail) + for token in range(seq_len + protected_tail): + token_scores[token] = (token * 7) % 64 + 1 + write_token(token, token_scores[token]) + + def expected_keep() -> torch.Tensor: + decode = token_scores[prompt_len:seq_len] + order = sorted(range(len(decode)), key=lambda index: (-decode[index], index)) + return torch.tensor( + sorted(prompt_len + index for index in order[:keep_count]), + dtype=torch.long, + device=device, + ) + + # GQA group 8; zero calibration query and shared MLR coefficient. + num_q_heads = 8 + q_real = torch.zeros(1, num_q_heads, num_freqs, dtype=torch.float32, device=device) + q_imag = torch.zeros_like(q_real) + mlr_coef = torch.zeros_like(q_real) + mlr_coef[..., 0] = 1 + freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) + freq_scale_sq[0] = 1 + tri = _make_cute_buffers( + eviction_mode="union", + layer_pools=[pool], + max_requests=1, + seq_len=seq_len, + num_q_heads=num_q_heads, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), + offsets=torch.zeros(1, dtype=torch.float32, device=device), + decode_width=seq_len - prompt_len, + keep_count=keep_count, + protected_tail_capacity=protected_tail, + ) + tri.kv_cache_manager = manager + + def evict_once() -> tuple[torch.Tensor, torch.Tensor]: + before = snapshot(seq_len + protected_tail) + eviction_requests = [ + _make_eviction_request( + request=_make_request(request_id, py_prompt_len=prompt_len), + source_length=seq_len, + target_tail_length=protected_tail, + ) + ] + # THE union path (fused pipeline) through the one round executor; + # the derived move offsets stage keep_count + protected_tail + # moves. Z-normalization is monotonic per row, so the raw-score + # keep set is unchanged. + tri._execute_eviction_round(eviction_requests) + selected = tri._kept_ordinal_rows[0].clone().to(torch.long) + torch.cuda.synchronize(device) + assert cache.resize(compacted_capacity, None) + after = snapshot(compacted_capacity) + source = torch.cat( + ( + selected, + torch.arange(seq_len, seq_len + protected_tail, device=device), + ) + ) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + assert torch.equal( + after[:, :, prompt_len:], + before.index_select(2, source), + ) + assert cache.capacity == compacted_capacity + assert cache.history_length == prompt_len + return selected, after + + initial_pages = page_ids(request_id) + expected_first_keep = expected_keep() + first_keep, first_compacted = evict_once() + assert torch.equal(first_keep, expected_first_keep) + # The compacted cache spans two of the original three pages. + retained_pages = page_ids(request_id) + assert torch.equal(retained_pages, initial_pages[:2]) + released_page = initial_pages[2:] + + created = manager.add_dummy_requests([8], [tokens_per_block]) + assert created is not None + temporary_requests = created + assert torch.equal(page_ids(8), released_page) + manager.free_resources(temporary_requests[0]) + temporary_requests = [] + + assert cache.resize(seq_len + protected_tail, None) + assert cache.history_length == prompt_len + assert torch.equal(page_ids(request_id)[:2], retained_pages) + assert torch.equal(page_ids(request_id)[2:], released_page) + # Mirror the relayout; fresh tokens get a disjoint higher score band + # so round two must select differently. + survivors = list(range(prompt_len)) + first_keep.tolist() + [seq_len, seq_len + 1] + token_scores[:compacted_capacity] = [token_scores[source] for source in survivors] + for token in range(compacted_capacity, seq_len + protected_tail): + token_scores[token] = (token * 11) % 64 + 100 + assert torch.equal(snapshot(compacted_capacity), first_compacted) + for token in range(compacted_capacity, seq_len + protected_tail): + write_token(token, token_scores[token]) + + expected_second_keep = expected_keep() + second_keep, _ = evict_once() + assert torch.equal(second_keep, expected_second_keep) + assert not torch.equal(second_keep, first_keep) + + created = manager.add_dummy_requests([9], [tokens_per_block]) + assert created is not None + temporary_requests = created + assert torch.equal(page_ids(9), released_page) + finally: + for request in temporary_requests: + manager.free_resources(request) + for request in requests: + manager.free_resources(request) + manager.shutdown() + + +def test_fold_union_ranks_matches_max_oracle(): + """The TP union fold is an exact elementwise max over the gathered rank blocks.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + fold_union_ranks, + ) + + device = torch.device("cuda", torch.cuda.current_device()) + tp_size, request_count, width = 4, 3, 300 + generator = torch.Generator(device="cpu").manual_seed(46) + gathered = torch.randn(tp_size * request_count, width, generator=generator).to(device) + folded = torch.full((request_count, width), float("nan"), device=device) + fold_union_ranks( + gathered, + folded, + request_count=request_count, + ) + expected = gathered.view(tp_size, request_count, width).amax(dim=0) + torch.cuda.synchronize(device) + assert torch.equal(folded, expected) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index e30864532cb3..5cac19e246df 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -260,7 +260,7 @@ methods: default: null status: prototype kv_cache_compression_config: - annotation: Optional[tensorrt_llm.llmapi.llm_args.KvCacheCompressionConfig] + annotation: Union[tensorrt_llm.llmapi.llm_args.TriAttentionKvCacheCompressionConfig, NoneType] default: null status: prototype otlp_traces_endpoint: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 0edebc20c181..b29c65616790 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2947,6 +2947,35 @@ def test_no_custom_init_methods(self): ) +def test_kv_cache_compression_config_dispatches_by_algorithm(): + from tensorrt_llm.llmapi.llm_args import \ + TriAttentionKvCacheCompressionConfig + + config_dict = yaml.safe_load(""" +kv_cache_compression_config: + algorithm: triattention + budget: 32 + beta: 17 + eviction_mode: per_head + normalize_scores: false + model_path: /tmp/model + calibration_path: /tmp/calibration.pt +""") + + config = TorchLlmArgs(model="/tmp/dummy_model", + **config_dict).kv_cache_compression_config + + assert isinstance(config, TriAttentionKvCacheCompressionConfig) + assert config.budget == 32 + assert config.beta == 17 + assert config.eviction_mode == "per_head" + assert config.normalize_scores is False + assert config.changes_physical_kv_length + assert config.supports_block_reuse() + assert not config.supports_speculative_decoding() + assert "changes_physical_kv_length" not in config.model_dump() + + class TestSkipSoftmaxAttentionConfig: """Test LLM Skip Softmax Attention config behavior."""