From 69372f353da3659436ce02c8b633b59984e78a9b Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 14 Jul 2026 21:40:05 -0700 Subject: [PATCH 001/178] [None][feat] Request-carried KV-cache compression lengths Signed-off-by: tianruih --- .../_torch/kv_cache_compression/__init__.py | 0 tensorrt_llm/_torch/model_config.py | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 24 ++- .../_torch/pyexecutor/kv_cache_manager_v2.py | 9 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 2 + .../_torch/pyexecutor/model_engine.py | 36 ++-- .../_torch/pyexecutor/model_loader.py | 2 + .../_torch/pyexecutor/resource_manager.py | 80 ++++++--- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../test_kv_cache_compression_manager.py | 161 ++++++++++++++++-- .../test_kv_cache_v2_capacity_only.py | 154 +++++++++++++++++ 11 files changed, 416 insertions(+), 59 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/__init__.py create mode 100644 tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index dd92337543c6..0bd6b64e1fcd 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -51,8 +51,9 @@ if TYPE_CHECKING: from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp - from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, LoraConfig, - SparseAttentionConfig, + from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, + KvCacheCompressionConfig, + LoraConfig, SparseAttentionConfig, SpeculativeConfig) TConfig = TypeVar("TConfig", bound=transformers.PretrainedConfig) @@ -162,6 +163,7 @@ class ModelConfig(Generic[TConfig]): spec_config: Optional["DecodingBaseConfig"] = None lora_config: Optional["LoraConfig"] = None sparse_attention_config: Optional["SparseAttentionConfig"] = None + kv_cache_compression_config: Optional["KvCacheCompressionConfig"] = None is_generation: bool = True is_encoder_decoder: bool = False diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 48d0c0dccdc4..feacedf85151 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2035,13 +2035,17 @@ def _create_kv_cache_manager( def create_kv_cache_compression_manager( config: KvCacheCompressionConfig, kv_cache_manager: KVCacheManagerV2, + draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, + spec_config: Optional[SpeculativeConfig] = None, ) -> Optional[BaseKVCacheCompressionManager]: """Build the KV-cache compression manager for ``config.algorithm``, or return None if no algorithm matches. 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. + here; the framework ships none. Speculative-decoding compatibility is also + decided here: a compression manager is only created when the speculative + mode supports it, otherwise the run stays uncompressed. """ logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " @@ -2256,7 +2260,12 @@ def create_py_executor_instance( "kv_cache_compression_config", None) if kv_cache_compression_config is not None: compression_manager = create_kv_cache_compression_manager( - kv_cache_compression_config, kv_cache_manager) + kv_cache_compression_config, + kv_cache_manager, + draft_kv_cache_manager=resources.get( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER), + spec_config=spec_config, + ) if compression_manager is not None: resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( compression_manager) @@ -2268,17 +2277,16 @@ def create_py_executor_instance( if kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.KV_CACHE_MANAGER, last=True) - # Compression manager runs after the cache manager: reconciles history once it's resized. - if (ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER - in resource_manager.resource_managers): - resource_manager.resource_managers.move_to_end( - ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, last=True) - cross_kv_cache_manager = resources.get( ResourceManagerType.CROSS_KV_CACHE_MANAGER) if cross_kv_cache_manager is not None: resource_manager.resource_managers.move_to_end( ResourceManagerType.CROSS_KV_CACHE_MANAGER, last=True) + # Compression is the final reconciler after every native KV manager. + if (ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER + in resource_manager.resource_managers): + resource_manager.resource_managers.move_to_end( + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER, last=True) # When scheduler_capacity == 1, attention dp dummy request will prevent the scheduling of DISAGG_GENERATION_INIT. # Enlarge scheduler capacity to avoid DISAGG_GENERATION_INIT stuck in the scheduler. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 5d22a2d395f2..260d2526d649 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -647,6 +647,8 @@ def __init__( layer_mask=layer_mask, ) self.is_draft = is_draft + # Set True by a compression manager; generation-step resize then leaves history untouched. + self.kv_compression_manages_history: bool = False self.enable_swa_scratch_reuse = ( kv_cache_config.enable_swa_scratch_reuse and not self.is_draft ) @@ -3061,12 +3063,15 @@ def update_resources( if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) else kv_cache.capacity - req.py_rewind_len ) - success = kv_cache.resize(new_capacity, req.max_beam_num_tokens - 1) + history_length = ( + None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 + ) + success = kv_cache.resize(new_capacity, history_length) if not success: raise ValueError( f"Failed to resize KV cache for request {req.py_request_id} " f"to capacity {new_capacity} and history length " - f"{req.max_beam_num_tokens - 1} tokens at generation update" + f"{history_length} tokens at generation update" ) def copy_batch_block_offsets( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 63bfb318e284..00d01f65df69 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -730,6 +730,8 @@ def __init__( self.py_batch_idx = None self.py_draft_pages_allocated = 0 self.py_rewind_len = 0 + # Tokens physically evicted by KV-cache compression; deducted in the engine. + self.py_num_compressed_tokens = 0 self.py_draft_tokens = [] if self.draft_tokens is None else self.draft_tokens self.py_last_context_chunk = (None, None) self.py_draft_logits = None diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b1a4f1b49cd6..a5f1d079bc46 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3087,6 +3087,10 @@ def _apply_incremental_update_target( else: prompt_lengths[idx] = request.py_prompt_len + # Physical KV length for the kernels: subtract the tokens a + # KV-cache compression manager evicted (tracked on the request, + # 0 without compression). Position ids and the cached_tokens stat + # keep the logical count. if request.is_dummy: num_cached_tokens_per_seq[idx] = base_past_seen request.cached_tokens = base_past_seen @@ -3097,9 +3101,11 @@ def _apply_incremental_update_target( num_previous_batch] = request.py_batch_idx num_previous_batch += 1 - num_cached_tokens_per_seq[ - idx] = base_past_seen + num_tokens_per_extend_request - request.cached_tokens = num_cached_tokens_per_seq[idx].item() + request.cached_tokens = (base_past_seen + + num_tokens_per_extend_request) + num_cached_tokens_per_seq[idx] = ( + base_past_seen + num_tokens_per_extend_request - + request.py_num_compressed_tokens) request.py_batch_idx = request.py_seq_slot @@ -3341,8 +3347,9 @@ def append_cross_attention_state(request: LlmRequest, py_request_id] = request.py_num_accepted_draft_tokens_indices prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute - num_cached_tokens_per_seq.append(past_seen_token_num) - request.cached_tokens = num_cached_tokens_per_seq[-1] + num_cached_tokens_per_seq.append(past_seen_token_num - + request.py_num_compressed_tokens) + request.cached_tokens = past_seen_token_num append_cross_attention_state( request, project_encoder_output=not request.py_skip_cross_kv_projection @@ -3527,8 +3534,9 @@ def append_cross_attention_state(request: LlmRequest, list( range(past_seen_token_num, past_seen_token_num + 1 + num_draft_tokens))) - num_cached_tokens_per_seq.append(past_seen_token_num) - request.cached_tokens = num_cached_tokens_per_seq[-1] + num_cached_tokens_per_seq.append( + past_seen_token_num - request.py_num_compressed_tokens) + request.cached_tokens = past_seen_token_num # update batch index request.py_batch_idx = request.py_seq_slot else: @@ -3555,9 +3563,11 @@ def append_cross_attention_state(request: LlmRequest, previous_pos_indices.extend([previous_batch_idx] * runtime_tokens_per_gen_step) - num_cached_tokens_per_seq.append(past_seen_token_num + - runtime_tokens_per_gen_step) - request.cached_tokens = num_cached_tokens_per_seq[-1] + num_cached_tokens_per_seq.append( + past_seen_token_num + runtime_tokens_per_gen_step - + request.py_num_compressed_tokens) + request.cached_tokens = (past_seen_token_num + + runtime_tokens_per_gen_step) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: prompt_lengths.append(runtime_tokens_per_gen_step) @@ -3614,7 +3624,8 @@ def append_cross_attention_state(request: LlmRequest, py_request_id] = request.py_num_accepted_draft_tokens_indices prompt_lengths.append(request.py_prompt_len) past_seen_token_num = begin_compute - num_cached_tokens_per_seq.append(past_seen_token_num) + num_cached_tokens_per_seq.append(past_seen_token_num - + request.py_num_compressed_tokens) append_cross_attention_state(request, project_encoder_output=False) # update batch index @@ -3691,7 +3702,8 @@ def append_cross_attention_state(request: LlmRequest, request.cached_tokens = past_seen_token_num for beam in range(beam_width): position_ids.append(position_id) - num_cached_tokens_per_seq.append(past_seen_token_num) + num_cached_tokens_per_seq.append( + past_seen_token_num - request.py_num_compressed_tokens) prompt_lengths.append(request.py_prompt_len) gather_ids.append(len(position_ids) - 1) diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 4aa8ce703d4c..b0c2ed44ef34 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -1166,6 +1166,8 @@ def _load_and_validate_config( force_dynamic_quantization=self.llm_args.force_dynamic_quantization, spec_config=self.spec_config, sparse_attention_config=self.sparse_attention_config, + kv_cache_compression_config=( + self.llm_args.kv_cache_compression_config), max_num_tokens=self.max_num_tokens, max_seq_len=self.max_seq_len, moe_max_num_tokens=self.llm_args.moe_config.max_num_tokens, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 486c43e8e269..c5f3b1947069 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, Dict, Iterable, List, Optional, Sequence, - Set, Tuple, Union) +from typing import (TYPE_CHECKING, ClassVar, Dict, Iterable, List, Optional, + Sequence, Set, Tuple, Union) import torch from mpi4py import MPI @@ -2341,14 +2341,37 @@ class BaseKVCacheCompressionManager(BaseResourceManager): base implementations below translate those callbacks into the lifecycle hooks. - Concrete compression methods subclass this directly. All 4 hooks default to + Concrete compression methods subclass this directly. The hooks default to no-op; subclasses override what they need. The manager never inherits from any cache manager because this layer decides *how* the physical KV is used, not *what* physical KV exists. Subclasses hold ``KVCacheManagerV2`` as a tool. + + A subclass compacts through the ``KVCacheManagerV2`` it holds and records + the evicted count on ``LlmRequest.py_num_compressed_tokens``; the model + engine subtracts that count when building ``num_cached_tokens_per_seq``. """ - def __init__(self, kv_cache_manager: "KVCacheManagerV2"): + adjusts_generation_kv_length: ClassVar[bool] = False + """Whether this manager can make target and logical KV lengths diverge.""" + + physically_evicts_cached_tokens: ClassVar[bool] = False + """True for evicting methods; attention modules then keep RoPE unfused.""" + + def __init__( + self, + kv_cache_manager: "KVCacheManagerV2", + draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None, + ): + from .kv_cache_manager_v2 import KVCacheManagerV2 + + if not isinstance(kv_cache_manager, KVCacheManagerV2): + raise TypeError("KV-cache compression requires KVCacheManagerV2") + if draft_kv_cache_manager is not None and not isinstance( + draft_kv_cache_manager, KVCacheManagerV2): + raise TypeError( + "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: @@ -2356,9 +2379,23 @@ def __init__(self, kv_cache_manager: "KVCacheManagerV2"): 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 + 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) + + @classmethod + def is_eviction_method(cls) -> bool: + """Whether this method physically evicts cached tokens.""" + return cls.physically_evicts_cached_tokens + + @property + def has_independent_draft_kv_cache(self) -> bool: + return self.draft_kv_cache_manager is not None # ================================================================== # - # KV-cache lifecycle hooks (4, in temporal order). # + # KV-cache lifecycle hooks (5, in temporal order). # # Subclasses override what they need; all default to no-op. # # ================================================================== # @@ -2369,20 +2406,23 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: scoring buffers). """ - def on_context_step_end( + def on_context_step_end(self, requests: List["LlmRequest"], + **kwargs) -> None: + """Fired once per iteration with the requests whose prefill finished + (their final chunk) this step. Batched like the generation hook so a + one-shot prefill-end eviction can process the cohort in one launch. + """ + + def on_generation_step_begin( self, - request: "LlmRequest", - metadata: "AttentionMetadata", + scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per request, when its prefill finishes (its final - chunk). Override for a one-shot prefill-end eviction. - """ + """Fired once per generation step before this step's forward.""" def on_generation_step_end( self, scheduled_batch: "ScheduledRequests", - attn_metadata: "AttentionMetadata", **kwargs, ) -> None: """Fired once per generation step, after every layer's forward @@ -2422,6 +2462,7 @@ def prepare_resources(self, scheduled_batch: "ScheduledRequests") -> None: for req in scheduled_batch.context_requests: if req.is_first_context_chunk: self.on_request_init(req) + self.on_generation_step_begin(scheduled_batch) def update_resources( self, @@ -2429,8 +2470,8 @@ def update_resources( attn_metadata: Optional["AttentionMetadata"] = None, kv_cache_dtype_byte_size: Optional[float] = None, ) -> None: - """Fire :meth:`on_context_step_end` once per request, on the iteration its - final prefill chunk runs, then :meth:`on_generation_step_end` once. + """Fire :meth:`on_context_step_end` with the requests whose final + prefill chunk ran this iteration, then :meth:`on_generation_step_end`. Uses the scheduler's ``context_requests_last_chunk`` split (computed at schedule time from ``is_last_context_chunk``) rather than tracking @@ -2441,9 +2482,10 @@ def update_resources( managers so PyExecutor passes ``attn_metadata`` / ``kv_cache_dtype_byte_size`` through transparently. """ - for req in scheduled_batch.context_requests_last_chunk: - self.on_context_step_end(req, attn_metadata) - self.on_generation_step_end(scheduled_batch, attn_metadata) + if scheduled_batch.context_requests_last_chunk: + self.on_context_step_end( + scheduled_batch.context_requests_last_chunk) + self.on_generation_step_end(scheduled_batch) def free_resources(self, request: "LlmRequest") -> None: """Fire :meth:`on_request_finish`.""" @@ -2480,9 +2522,9 @@ def update_resources( attn_metadata: Optional["AttentionMetadata"] = None, kv_cache_dtype_byte_size: Optional[float] = None, ): - for _, resource_manager in self.resource_managers.items(): + for resource_type, resource_manager in self.resource_managers.items(): if hasattr(resource_manager, "update_resources"): - if isinstance(resource_manager, KVCacheManager): + if resource_type == ResourceManagerType.KV_CACHE_MANAGER: resource_manager.update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1e6836fc7cd5..cdd02b720890 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -37,6 +37,7 @@ l0_a10: - unittest/_torch/executor/test_kv_pool_rebalance.py - unittest/_torch/executor/test_disagg_index_mapper_early_release.py - unittest/_torch/executor/test_kv_cache_compression_manager.py + - unittest/_torch/executor/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/test_error_classification.py - unittest/_torch/modules/dwdp/test_dwdp_fixup_moe_backends.py - unittest/_torch/modules/dwdp/test_dwdp_manager.py 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 13cc7fd9d6d2..e7448b79ff33 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -12,8 +12,8 @@ - The resource-manager API -> lifecycle-hook translation, gated on PyExecutor's own signals: ``prepare_resources`` fires ``on_request_init`` on each request's first prefill chunk (``is_first_context_chunk``); - ``update_resources`` fires ``on_context_step_end`` for each request in - ``context_requests_last_chunk`` + one ``on_generation_step_end`` per + ``update_resources`` fires ``on_context_step_end`` once with the + ``context_requests_last_chunk`` list + one ``on_generation_step_end`` per iteration; ``free_resources`` fires ``on_request_finish``. - :func:`create_kv_cache_compression_manager` factory. @@ -22,6 +22,8 @@ lives in ``_util.py`` next to ``_create_kv_cache_manager``. """ +from types import SimpleNamespace +from typing import ClassVar from unittest.mock import MagicMock, patch import pytest @@ -31,6 +33,8 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseKVCacheCompressionManager, BaseResourceManager, + ResourceManager, + ResourceManagerType, ) # ---------------------------------------------------------------------- # @@ -57,23 +61,35 @@ class _MockCompressionManager(_RecordingMixin, BaseKVCacheCompressionManager): def on_request_init(self, request): self._record("on_request_init") - def on_context_step_end(self, request, metadata): - self._record("on_context_step_end") + def on_context_step_end(self, requests): + self._record(f"on_context_step_end[{len(requests)}]") - def on_generation_step_end(self, scheduled_batch, attn_metadata): + def on_generation_step_end(self, scheduled_batch): self._record("on_generation_step_end") def on_request_finish(self, request): self._record("on_request_finish") +class _LengthAdjustingCompressionManager(BaseKVCacheCompressionManager): + adjusts_generation_kv_length: ClassVar[bool] = True + + +def _v2_manager(*, is_draft: bool): + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.enable_block_reuse = False + manager.kv_compression_manages_history = False + manager.is_draft = is_draft + return manager + + @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.""" - m = MagicMock(name="fake_KVCacheManagerV2") - m.enable_block_reuse = False - return m + return _v2_manager(is_draft=False) def _req(rid, first_chunk=True): @@ -103,10 +119,10 @@ def test_inherits_base_resource_manager(self): def test_four_hooks_default_noop(self, fake_kv_cache_manager): m = BaseKVCacheCompressionManager(fake_kv_cache_manager) - meta = MagicMock() assert m.on_request_init(MagicMock()) is None - assert m.on_context_step_end(MagicMock(), meta) is None - assert m.on_generation_step_end(MagicMock(), meta) is None + assert m.on_context_step_end([MagicMock()]) is None + assert m.on_generation_step_begin(MagicMock()) is None + assert m.on_generation_step_end(MagicMock()) is None assert m.on_request_finish(MagicMock()) is None def test_hooks_accept_extra_kwargs(self, fake_kv_cache_manager): @@ -114,7 +130,7 @@ def test_hooks_accept_extra_kwargs(self, fake_kv_cache_manager): # existing overrides. m = BaseKVCacheCompressionManager(fake_kv_cache_manager) assert m.on_request_init(MagicMock(), future_arg=1) is None - assert m.on_generation_step_end(MagicMock(), 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 = BaseKVCacheCompressionManager(fake_kv_cache_manager) @@ -123,6 +139,42 @@ def test_resource_counts_are_zero(self, fake_kv_cache_manager): 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): + # 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) + + assert manager.kv_cache_manager is target + assert manager.draft_kv_cache_manager is draft + assert manager.has_independent_draft_kv_cache + assert target.kv_compression_manages_history is True + assert draft.kv_compression_manages_history is True + + def test_rejects_non_v2_ownership(self): + with pytest.raises(TypeError, match="requires KVCacheManagerV2"): + BaseKVCacheCompressionManager(MagicMock()) + with pytest.raises(TypeError, match="requires KVCacheManagerV2"): + BaseKVCacheCompressionManager(_v2_manager(is_draft=False), MagicMock()) + + def test_request_field_defaults_to_zero(self): + """LlmRequest carries the compression count (the manager's only + channel to the runtime); a fresh request must default to 0 so runs + without a compression manager are unchanged.""" + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm.bindings import SamplingConfig + + request = LlmRequest( + request_id=1, + max_new_tokens=8, + input_tokens=[1, 2, 3], + sampling_config=SamplingConfig(), + is_streaming=False, + ) + assert request.py_num_compressed_tokens == 0 + # ---------------------------------------------------------------------- # # 2. Resource-manager API -> lifecycle-hook translation # @@ -131,6 +183,46 @@ def test_resource_counts_are_zero(self, fake_kv_cache_manager): class TestResourceManagerAPI: + def test_target_update_receives_metadata_before_final_compression(self): + calls = [] + metadata = MagicMock(name="attention_metadata") + draft = MagicMock(name="draft_kv_cache_manager") + target = MagicMock(name="target_kv_cache_manager") + compression = MagicMock(name="compression_manager") + draft.update_resources.side_effect = lambda *args: calls.append(("draft", args)) + target.update_resources.side_effect = lambda *args: calls.append(("target", args)) + compression.update_resources.side_effect = lambda *args: calls.append(("compression", args)) + manager = ResourceManager( + { + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: draft, + ResourceManagerType.KV_CACHE_MANAGER: target, + ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER: compression, + } + ) + batch = _batch(generation=[_req(1)]) + + manager.update_resources(batch, metadata, 2.0) + + assert calls == [ + ("draft", (batch,)), + ("target", (batch, metadata, 2.0)), + ("compression", (batch,)), + ] + + def test_real_v2_target_receives_relocation_metadata(self): + from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module + + target = _v2_manager(is_draft=False) + target.kv_cache_map = {} + batch = _batch(generation=[_req(1)]) + metadata = MagicMock(name="attention_metadata") + manager = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: target}) + + with patch.object(kv_cache_v2_module, "_update_kv_cache_draft_token_location") as relocate: + manager.update_resources(batch, metadata, 2.0) + + relocate.assert_called_once_with(target, batch, metadata, 2.0) + def test_prepare_fires_init_on_first_chunk_only(self, fake_kv_cache_manager): rec = [] m = _MockCompressionManager(fake_kv_cache_manager, rec, "s") @@ -144,9 +236,10 @@ def test_update_fires_context_end_on_last_chunk(self, fake_kv_cache_manager): rec = [] m = _MockCompressionManager(fake_kv_cache_manager, rec, "s") req = _req(1) - # Request's final prefill chunk this iteration -> context_step_end fires. - m.update_resources(_batch(generation=[req], last_chunk=[req]), attn_metadata=MagicMock()) - assert "s:on_context_step_end" in rec + # Final prefill chunks this iteration -> one batched context_step_end. + req2 = _req(2) + m.update_resources(_batch(generation=[req], last_chunk=[req, req2])) + assert "s:on_context_step_end[2]" in rec assert rec[-1] == "s:on_generation_step_end" # Subsequent decode iteration (not in last_chunk) -> no context_step_end. rec.clear() @@ -185,6 +278,42 @@ def test_warns_for_unregistered_algorithm(self, fake_kv_cache_manager): create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) mock_logger.warning.assert_called_once() + def test_factory_accepts_independent_draft_manager(self): + from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + + cfg = MagicMock() + cfg.algorithm = "made_up_method" + target = _v2_manager(is_draft=False) + draft = _v2_manager(is_draft=True) + spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.EAGLE3_ONE_MODEL) + + assert ( + create_kv_cache_compression_manager( + cfg, + target, + draft_kv_cache_manager=draft, + spec_config=spec_config, + ) + is None + ) + + def test_eviction_method_predicate_defaults_false(self): + # The base is not an evicting method, so the factory's speculative + # mode gate never restricts it: methods that do not touch the draft + # KV (e.g. offloading) work with any speculative mode. + m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) + assert m.is_eviction_method() is False + assert not hasattr(m, "spec_config") + + def test_eviction_method_predicate_follows_class_flag(self): + # The factory's speculative gate reads this predicate: an evicting + # method (physically_evicts_cached_tokens True) only accepts modes + # whose draft KV is a standard paged cache in the same forward. + class _EvictingManager(BaseKVCacheCompressionManager): + physically_evicts_cached_tokens = True + + assert _EvictingManager.is_eviction_method() is True + # ---------------------------------------------------------------------- # # 4. Canonical names live in resource_manager, not in the sparse module # @@ -219,7 +348,7 @@ class TestBlockReuseGuard: and values, the same check RocketKVCacheManager makes.""" def _mgr(self, enable_block_reuse): - m = MagicMock(name="KVCacheManagerV2") + m = _v2_manager(is_draft=False) m.enable_block_reuse = enable_block_reuse return m diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py new file mode 100644 index 000000000000..df38d7bba8fc --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +import tensorrt_llm +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor import kv_cache_manager_v2 as kv_cache_v2_module +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + + +def _manager(*, is_draft: bool, kv_compression_manages_history: bool = False) -> KVCacheManagerV2: + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + manager.is_draft = is_draft + manager.kv_compression_manages_history = kv_compression_manages_history + manager.kv_cache_map = {} + return manager + + +def _request(request_id: int, *, rewind: int = 0, complete: bool = False) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=request_id, + py_rewind_len=rewind, + max_beam_num_tokens=201, + state=LlmRequestState.GENERATION_COMPLETE + if complete + else LlmRequestState.GENERATION_IN_PROGRESS, + ) + + +def _cache(*, capacity: int = 256, active: bool = True) -> MagicMock: + cache = MagicMock() + cache.capacity = capacity + cache.is_active = active + cache.resize.return_value = True + return cache + + +@pytest.fixture(autouse=True) +def _disable_draft_token_relocation(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(kv_cache_v2_module, "_update_kv_cache_draft_token_location", MagicMock()) + + +def test_manager_initializes_capacity_only_policy_to_false() -> None: + class StopInitialization(RuntimeError): + pass + + class StopAfterPolicyConfig: + @property + def enable_swa_scratch_reuse(self): + raise StopInitialization + + manager = KVCacheManagerV2.__new__(KVCacheManagerV2) + mapping = SimpleNamespace(cp_config={}) + + with ( + patch.object(kv_cache_v2_module, "get_pp_layers", return_value=([0], 1)), + pytest.raises(StopInitialization), + ): + manager.__init__( + StopAfterPolicyConfig(), + kv_cache_v2_module.CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=128, + tokens_per_block=64, + max_seq_len=256, + max_batch_size=1, + mapping=mapping, + ) + + assert manager.kv_compression_manages_history is False + + +def test_default_generation_resize_updates_capacity_and_history() -> None: + manager = _manager(is_draft=False) + request = _request(1, rewind=3) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(253, 200) + + +def test_capacity_only_is_scoped_to_target_manager() -> None: + request = _request(1, rewind=3) + batch = SimpleNamespace(generation_requests=[request]) + target = _manager(is_draft=False, kv_compression_manages_history=True) + draft = _manager(is_draft=True) + target_cache = _cache() + draft_cache = _cache() + target.kv_cache_map[request.py_request_id] = target_cache + draft.kv_cache_map[request.py_request_id] = draft_cache + + draft.update_resources(batch) + target.update_resources(batch) + + draft_cache.resize.assert_called_once_with(253, 200) + target_cache.resize.assert_called_once_with(253, None) + + +def test_capacity_only_completion_preserves_history() -> None: + manager = _manager(is_draft=False, kv_compression_manages_history=True) + request = _request(1, complete=True) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(None, None) + + +def test_capacity_only_skips_suspended_cache() -> None: + manager = _manager(is_draft=False, kv_compression_manages_history=True) + request = _request(1, rewind=3) + cache = _cache(active=False) + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_not_called() + + +def test_generation_update_has_no_request_compaction_marker() -> None: + manager = _manager(is_draft=False, kv_compression_manages_history=True) + request = _request(1, rewind=3) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + assert "py_kv_cache_kv_compression_manages_history" not in vars(request) + assert "py_kv_cache_compaction" not in vars(request) + + +def test_llm_request_has_no_compression_consumer_marker() -> None: + request = LlmRequest( + request_id=1, + max_new_tokens=1, + input_tokens=[1], + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + + assert "py_kv_cache_kv_compression_manages_history" not in vars(request) + assert "py_kv_cache_compaction" not in vars(request) From 8fc7cc9cd5cec9d1cccd875bb419e6fff256a80d Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 14 Jul 2026 21:40:07 -0700 Subject: [PATCH 002/178] [None][feat] General KV eviction primitives Signed-off-by: tianruih --- .../batch_manager/kvCacheManagerV2Utils.cpp | 55 ++++ .../batch_manager/kvCacheManagerV2Utils.h | 6 +- .../kernels/unfusedAttentionKernels.h | 13 +- .../unfusedAttentionKernels_2_bf16_bf16.cu | 3 +- .../unfusedAttentionKernels_2_float_float.cu | 3 +- .../unfusedAttentionKernels_2_half_half.cu | 3 +- .../unfusedAttentionKernels_2_template.h | 158 +++++++++- .../batch_manager/kvCacheManagerV2Utils.cpp | 2 + cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + .../thop/sparseKvCacheCompactOp.cpp | 165 ++++++++++ .../_torch/custom_ops/cute_dsl_custom_ops.py | 23 +- .../top_k/filtered_top_k_decode_varlen.py | 6 +- .../top_k/filtered_top_k_varlen_util.py | 5 +- .../test_disagg_index_mapper_early_release.py | 65 +++- .../_torch/thop/parallel/test_indexer_topk.py | 127 ++++++++ .../serial/test_sparse_kv_cache_compact.py | 297 ++++++++++++++++++ 16 files changed, 897 insertions(+), 35 deletions(-) create mode 100644 cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp create mode 100644 tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp index 079de1a18893..1f735fa112e1 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include +#include #include #include #include @@ -217,6 +218,60 @@ 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) +{ + TLLM_CHECK_WITH_INFO(source.device().is_cpu() && destination.device().is_cpu(), + "Block-offset gather requires CPU tensors"); + TLLM_CHECK_WITH_INFO(source.scalar_type() == at::kInt && destination.scalar_type() == at::kInt, + "Block-offset gather requires int32 tensors"); + TLLM_CHECK_WITH_INFO(source.dim() == 4 && destination.dim() == 4 && source.size(2) == 2 + && destination.size(2) == 2, + "Block-offset gather requires [pool, sequence, K/V, block] tensors"); + TLLM_CHECK_WITH_INFO(source.is_contiguous() && destination.is_contiguous(), + "Block-offset gather requires contiguous tensors"); + TLLM_CHECK_WITH_INFO(source.storage().data_ptr().get() != destination.storage().data_ptr().get(), + "Block-offset gather requires distinct source and destination storage"); + TLLM_CHECK_WITH_INFO(source.size(0) == destination.size(0), "Block-offset gather pool counts must match"); + TLLM_CHECK_WITH_INFO(!requestIds.empty(), "Block-offset gather requires at least one request"); + TLLM_CHECK_WITH_INFO(numBlocks > 0 && numBlocks <= source.size(3) && numBlocks <= destination.size(3), + "Block-offset gather block count exceeds the source or destination capacity"); + TLLM_CHECK_WITH_INFO(static_cast(requestIds.size()) <= destination.size(1), + "Block-offset gather request count exceeds the destination capacity"); + + auto const sourceRows = source.size(1); + std::vector sourceRowsByRequest; + sourceRowsByRequest.reserve(requestIds.size()); + for (auto const requestId : requestIds) + { + auto const sourceRow = static_cast(getIndex(requestId)) * maxBeamWidth_; + TLLM_CHECK_WITH_INFO(sourceRow < sourceRows, "IndexMapper slot exceeds the source tensor capacity"); + sourceRowsByRequest.push_back(sourceRow); + } + + auto const* sourceData = source.data_ptr(); + auto* destinationData = destination.data_ptr(); + 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 e32c727d0d81..9117259552d0 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"); @@ -77,6 +77,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/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 302c278f7a2a..708f90c970ef 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-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. @@ -412,6 +412,17 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams template void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams params, cudaStream_t stream); +//! Adapt a uniform group of KVCacheManagerV2 layer pools to the existing +//! sparse-KV updater. Device pointer arrays allow one layered launch while the +//! updater remains the only implementation of the in-place copy loop. Within +//! each request and head, source ordinals must increase strictly and satisfy +//! destinationBase + move <= source[move]. +template +void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, + int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, + int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t destinationBase, int32_t batchSize, + int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream); + // Debug function to test basic parameter access template void invokeDebugSparseKvCacheParams( diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu index 5d006ef4a979..ec1b92aad4fa 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -26,6 +26,7 @@ namespace kernels #ifdef ENABLE_BF16 INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(__nv_bfloat16, __nv_bfloat16, KVBlockArray); INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(__nv_bfloat16, __nv_bfloat16, KVLinearBuffer); +INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(__nv_bfloat16); #endif } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu index 55e3e8756afe..4b413712d787 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -25,6 +25,7 @@ namespace kernels INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(float, float, KVBlockArray); INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(float, float, KVLinearBuffer); +INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(float); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu index 5abd544359d1..943207c2e589 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -25,6 +25,7 @@ namespace kernels INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(half, half, KVBlockArray); INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(half, half, KVLinearBuffer); +INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(half); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 560a153ffa70..f39196a3b985 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -1754,7 +1754,63 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams //////////////////////////////////////////////////////////////////////////////////////////////////// -template +//! Layered KVCacheManagerV2 layout policy for the existing sparse-KV updater. +//! blockIdx.x selects an independent per-layer pool view. Layers in this launch +//! share one V2 block-offset table and the production FMHA updater copy loop. +struct KvCacheV2LayersBuffer +{ + int64_t const* poolPointers; + int32_t const* pageTable; + int32_t const* sourceLayerIndices; + int64_t sourceLayerStride; + int64_t pageTableRequestStride; + int32_t tokensPerBlock; + int32_t destinationBase; + size_t bytesPerPage; + size_t bytesPerKvHalf; + + __device__ __forceinline__ uint8_t* getPool() const + { + return reinterpret_cast(static_cast(poolPointers[blockIdx.x])); + } + + __device__ __forceinline__ void* getKBlockPtr(int32_t batchIdx, int32_t tokenIdx) const + { + int32_t const blockOffset = pageTable[batchIdx * pageTableRequestStride + tokenIdx / tokensPerBlock]; + int32_t const page = blockOffset / 2; + return getPool() + static_cast(page) * bytesPerPage; + } + + __device__ __forceinline__ void* getVBlockPtr(int32_t batchIdx, int32_t tokenIdx) const + { + int32_t const blockOffset = pageTable[batchIdx * pageTableRequestStride + tokenIdx / tokensPerBlock]; + int32_t const page = blockOffset / 2; + return getPool() + static_cast(page) * bytesPerPage + bytesPerKvHalf; + } + + __device__ __forceinline__ int32_t getKVLocalIdx( + int32_t tokenIdx, int32_t headIdx, int32_t valuesPerHead, int32_t headValueIdx) const + { + return (headIdx * tokensPerBlock + (tokenIdx % tokensPerBlock)) * valuesPerHead + headValueIdx; + } + + __device__ __forceinline__ int32_t getSparseKvSourceToken( + int32_t const* sourceIndices, int32_t headIdx, int32_t totalMoves, int32_t globalMove) const + { + int32_t const layer + = sourceLayerIndices == nullptr ? static_cast(blockIdx.x) : sourceLayerIndices[blockIdx.x]; + int64_t const offset + = static_cast(layer) * sourceLayerStride + static_cast(headIdx) * totalMoves + globalMove; + return sourceIndices[offset]; + } + + __device__ __forceinline__ int32_t getSparseKvDestinationToken(int32_t requestMove) const + { + return destinationBase + requestMove; + } +}; + +template __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( QKVPreprocessingParams params) { @@ -1785,9 +1841,17 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( if (sparse_token_offset < num_sparse_tokens) { int const global_sparse_idx = sparse_start_idx + sparse_token_offset; - int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; - - int const src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; + int src_token_idx; + if constexpr (Layered) + { + src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( + params.sparse_kv_indices, kv_head_idx, total_num_sparse_kv_tokens, global_sparse_idx); + } + else + { + int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; + src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; + } void* src_k_ptr = params.kv_cache_buffer.getKBlockPtr(batch_idx, src_token_idx); void* src_v_ptr = params.kv_cache_buffer.getVBlockPtr(batch_idx, src_token_idx); @@ -1810,10 +1874,20 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( if (sparse_token_offset < num_sparse_tokens) { int const global_sparse_idx = sparse_start_idx + sparse_token_offset; - int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; - - int const src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; - int const dst_token_idx = sparse_token_offset; + int src_token_idx; + int dst_token_idx; + if constexpr (Layered) + { + src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( + params.sparse_kv_indices, kv_head_idx, total_num_sparse_kv_tokens, global_sparse_idx); + dst_token_idx = params.kv_cache_buffer.getSparseKvDestinationToken(sparse_token_offset); + } + else + { + int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; + src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; + dst_token_idx = sparse_token_offset; + } if (src_token_idx != dst_token_idx) { @@ -1851,7 +1925,8 @@ void kernelSparseDispatchHeadSize(QKVPreprocessingParams param // grid.x is always 1 to avoid data races dim3 grid(1, params.kv_head_num, params.batch_size); - updateSparseKvCacheAfterFmha<<>>(params); + updateSparseKvCacheAfterFmha + <<>>(params); } template @@ -1878,6 +1953,64 @@ void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams //////////////////////////////////////////////////////////////////////////////////////////////////// +template +void launchSparseKvCacheCompactV2Layers( + QKVPreprocessingParams params, int32_t numLayers, cudaStream_t stream) +{ + constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / 16; + constexpr int32_t kDefaultSharedMemoryBytes = 48 * 1024; + constexpr int32_t kSharedBytesPerToken = 2 * kVectorsPerHead * sizeof(uint4); + constexpr bool kUseD64VectorThreads = HeadDim == 64 && sizeof(T) == 2; + constexpr int32_t kVectorThreads = kUseD64VectorThreads ? kVectorsPerHead : 32; + constexpr int32_t kTokensPerTile + = kUseD64VectorThreads ? 32 : (kSharedBytesPerToken * 32 <= kDefaultSharedMemoryBytes ? 32 : 16); + constexpr int32_t kBlockSize = kVectorThreads * kTokensPerTile; + static_assert(kSharedBytesPerToken * kTokensPerTile <= kDefaultSharedMemoryBytes); + dim3 const block(kVectorThreads, kTokensPerTile); + dim3 const grid(numLayers, params.kv_head_num, params.batch_size); + size_t const sharedBytes = 2 * block.y * kVectorsPerHead * sizeof(uint4); + updateSparseKvCacheAfterFmha + <<>>(params); +} + +template +void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, + int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, + int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t destinationBase, int32_t batchSize, + int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream) +{ + KvCacheV2LayersBuffer buffer{}; + buffer.poolPointers = poolPointers; + buffer.pageTable = pageTable; + buffer.sourceLayerIndices = sourceLayerIndices; + buffer.sourceLayerStride = sourceLayerStride; + buffer.pageTableRequestStride = pageTableRequestStride; + buffer.tokensPerBlock = tokensPerBlock; + buffer.destinationBase = destinationBase; + buffer.bytesPerKvHalf = static_cast(numKvHeads) * tokensPerBlock * headDim * sizeof(T); + buffer.bytesPerPage = 2 * buffer.bytesPerKvHalf; + + QKVPreprocessingParams params{}; + params.kv_cache_buffer = buffer; + params.sparse_kv_indices = sparseKvIndices; + params.sparse_kv_offsets = sparseKvOffsets; + params.batch_size = batchSize; + params.kv_head_num = numKvHeads; + params.size_per_head = headDim; + + switch (headDim) + { + case 16: launchSparseKvCacheCompactV2Layers<16>(params, numLayers, stream); break; + case 32: launchSparseKvCacheCompactV2Layers<32>(params, numLayers, stream); break; + case 64: launchSparseKvCacheCompactV2Layers<64>(params, numLayers, stream); break; + case 128: launchSparseKvCacheCompactV2Layers<128>(params, numLayers, stream); break; + case 256: launchSparseKvCacheCompactV2Layers<256>(params, numLayers, stream); break; + default: TLLM_CHECK_WITH_INFO(false, "Sparse KV compaction does not support head size %d", headDim); + } +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + #define INSTANTIATE_ATTENTION_INPUT_PROCESSING(T, TCache, KVCacheBuffer) \ template void invokeApplyBiasRopeUpdateKVCacheDispatch( \ QKVPreprocessingParams params, cudaStream_t stream); @@ -1891,6 +2024,11 @@ void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams QKVPreprocessingParams params, cudaStream_t stream); \ //////////////////////////////////////////////////////////////////////////////////////////////////// +#define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(T) \ + template void invokeSparseKvCacheCompactV2Layers(int64_t const*, int32_t const*, int32_t, int64_t, \ + int32_t const*, int32_t const*, int64_t, int32_t const*, int32_t, int32_t, int32_t, int32_t, int32_t, \ + cudaStream_t); + } // namespace kernels TRTLLM_NAMESPACE_END 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/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 369bf721b099..e1b8e21dd8b7 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -109,6 +109,7 @@ add_library( IndexerKCacheGatherOp.cpp IndexerKCacheScatterOp.cpp IndexerTopKOp.cpp + sparseKvCacheCompactOp.cpp mlaRopeInplaceOp.cpp ncclCommunicatorOp.cpp allocateOutput.cpp diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp new file mode 100644 index 000000000000..be7dc7d75bc3 --- /dev/null +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -0,0 +1,165 @@ +/* + * 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. + */ + +#include "tensorrt_llm/common/opUtils.h" +#include "tensorrt_llm/kernels/unfusedAttentionKernels.h" +#include "tensorrt_llm/runtime/torchUtils.h" + +#include +#include +#include + +namespace th = torch; +namespace tk = tensorrt_llm::kernels; + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Adapt one uniform group of KVCacheManagerV2 HND layer pools to the +//! existing sparse-KV post-FMHA updater. Layers share one V2 block-offset table; +//! destinationBase replaces the former arbitrary +//! destination tensor because every TriAttention move targets one interval. +//! Within each request and KV head, TriAttention supplies increasing source +//! ordinals with destinationBase + move <= source[move], which makes the +//! updater's forward tiled in-place copy safe. +void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor const& poolPointers, + th::Tensor const& pageTable, th::Tensor const& sourceIndices, th::Tensor const& sourceOffsets, + std::optional const& sourceLayerIndices, int64_t destinationBase) +{ + TORCH_CHECK(!pools.empty(), "sparse_kv_cache_compact_layers: pools must be non-empty"); + + auto const& firstPool = pools.front(); + TORCH_CHECK(firstPool.is_cuda() && firstPool.dim() == 5 && firstPool.size(1) == 2 && firstPool.is_contiguous(), + "sparse_kv_cache_compact_layers: pools must be contiguous CUDA " + "[pages, 2, kv_heads, tokens_per_block, head_dim] tensors"); + TORCH_CHECK(pageTable.is_cuda() && pageTable.dim() == 2 && pageTable.scalar_type() == th::kInt32 + && pageTable.stride(1) == 1, + "sparse_kv_cache_compact_layers: K block offsets must be CUDA int32 " + "[batch, max_pages] tensors with a contiguous block dimension"); + + auto const device = firstPool.get_device(); + auto const dtype = firstPool.scalar_type(); + auto const numLayers = static_cast(pools.size()); + auto const numKvHeads = static_cast(firstPool.size(2)); + auto const tokensPerBlock = static_cast(firstPool.size(3)); + auto const headDim = static_cast(firstPool.size(4)); + auto const batchSize = static_cast(pageTable.size(0)); + auto const pageTableRequestStride = pageTable.stride(0); + + for (int32_t layer = 0; layer < numLayers; ++layer) + { + auto const& pool = pools[layer]; + TORCH_CHECK(pool.is_cuda() && pool.get_device() == device && pool.scalar_type() == dtype && pool.dim() == 5 + && pool.size(1) == 2 && pool.is_contiguous(), + "sparse_kv_cache_compact_layers: all pools must have one device, dtype, layout, and contiguous storage"); + TORCH_CHECK(pool.size(2) == numKvHeads && pool.size(3) == tokensPerBlock && pool.size(4) == headDim, + "sparse_kv_cache_compact_layers: all pools must share KV-head, block, and head-dimension geometry"); + } + TORCH_CHECK( + pageTable.get_device() == device, "sparse_kv_cache_compact_layers: block offsets must be on the pool device"); + + auto const checkPointerArray = [device, numLayers](th::Tensor const& pointers, char const* name) + { + TORCH_CHECK(pointers.is_cuda() && pointers.get_device() == device && pointers.scalar_type() == th::kInt64 + && pointers.dim() == 1 && pointers.size(0) == numLayers && pointers.is_contiguous(), + "sparse_kv_cache_compact_layers: ", name, " must be contiguous CUDA int64 [num_layers]"); + }; + checkPointerArray(poolPointers, "pool_pointers"); + + TORCH_CHECK(sourceIndices.is_cuda() && sourceIndices.get_device() == device + && sourceIndices.scalar_type() == th::kInt32 && sourceIndices.is_contiguous() + && (sourceIndices.dim() == 2 || sourceIndices.dim() == 3), + "sparse_kv_cache_compact_layers: source_indices must be contiguous CUDA int32 " + "[kv_heads, total] or [source_layers, kv_heads, total]"); + int64_t sourceLayerStride = 0; + int32_t const* sourceLayerPtr = nullptr; + if (sourceIndices.dim() == 2) + { + TORCH_CHECK(sourceIndices.size(0) == numKvHeads, + "sparse_kv_cache_compact_layers: source_indices KV-head dimension mismatch"); + } + else + { + TORCH_CHECK(sourceIndices.size(0) > 0 && sourceIndices.size(1) == numKvHeads, + "sparse_kv_cache_compact_layers: per-layer source_indices geometry mismatch"); + TORCH_CHECK(sourceLayerIndices.has_value(), + "sparse_kv_cache_compact_layers: per-layer source_indices require source_layer_indices"); + sourceLayerStride = sourceIndices.stride(0); + } + if (sourceLayerIndices.has_value()) + { + auto const& layerIndices = *sourceLayerIndices; + TORCH_CHECK(layerIndices.is_cuda() && layerIndices.get_device() == device + && layerIndices.scalar_type() == th::kInt32 && layerIndices.is_contiguous() && layerIndices.dim() == 1 + && layerIndices.size(0) == numLayers, + "sparse_kv_cache_compact_layers: source_layer_indices must be contiguous CUDA int32 [num_layers]"); + sourceLayerPtr = layerIndices.data_ptr(); + } + + TORCH_CHECK(sourceOffsets.is_cuda() && sourceOffsets.get_device() == device + && sourceOffsets.scalar_type() == th::kInt32 && sourceOffsets.is_contiguous() && sourceOffsets.dim() == 1 + && sourceOffsets.size(0) == batchSize + 1, + "sparse_kv_cache_compact_layers: source_offsets must be contiguous CUDA int32 [batch + 1]"); + TORCH_CHECK(destinationBase >= 0 && destinationBase <= std::numeric_limits::max(), + "sparse_kv_cache_compact_layers: destination_base must fit a non-negative int32"); + + auto const stream = at::cuda::getCurrentCUDAStream(device); + auto const base = static_cast(destinationBase); + if (dtype == th::kBFloat16) + { + tk::invokeSparseKvCacheCompactV2Layers<__nv_bfloat16>(poolPointers.data_ptr(), + pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), + sourceLayerPtr, sourceLayerStride, sourceOffsets.data_ptr(), base, batchSize, numKvHeads, + tokensPerBlock, headDim, stream); + } + else if (dtype == th::kHalf) + { + tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), + numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, + sourceOffsets.data_ptr(), base, batchSize, numKvHeads, tokensPerBlock, headDim, stream); + } + else if (dtype == th::kFloat) + { + tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), + numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, + sourceOffsets.data_ptr(), base, batchSize, numKvHeads, tokensPerBlock, headDim, stream); + } + else + { + TORCH_CHECK(false, "sparse_kv_cache_compact_layers: unsupported pool dtype ", dtype); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "sparse_kv_cache_compact_layers(Tensor(a!)[] pools, Tensor pool_pointers, Tensor page_table, Tensor " + "source_indices, Tensor source_offsets, Tensor? source_layer_indices=None, " + "int destination_base=0) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("sparse_kv_cache_compact_layers", &tensorrt_llm::torch_ext::sparseKvCacheCompactLayers); +} diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 84d55a1f4c79..a0fc3a534212 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -4823,7 +4823,6 @@ class CuteDSLTopKDecodeSingleCTARunner: Note: - Requires Blackwell architecture (SM100+) - - Maximum tested top_k is 2048 (see kernel documentation for larger values) - Supports fp16, bf16, and fp32 dtypes - Automatically selects occupancy optimization based on batch size """ @@ -5019,7 +5018,7 @@ def cute_dsl_topk_decode_blackwell( Args: input_values: Input logits tensor [batch_size * next_n, vocab_size] seq_lens: Sequence lengths for each batch [batch_size] - top_k: Number of top elements to select (max 2048) + top_k: Number of top elements to select (max 4096) next_n: Number of candidates per sequence (for speculative decoding) num_copy_bits: Number of bits for vectorized memory copy (128 or 256) load_balance: Enable persistent dynamic scheduling for load balancing @@ -5029,7 +5028,7 @@ def cute_dsl_topk_decode_blackwell( Note: This function requires Blackwell architecture (SM100+) and CuTE DSL support. - Maximum supported top_k is 2048. + Maximum supported top_k is 4096. """ # Validate SM version sm_version = get_sm_version() @@ -5039,10 +5038,11 @@ def cute_dsl_topk_decode_blackwell( "Use standard top-k implementation for older architectures.") # Validate inputs - if top_k <= 0 or top_k > 2048: + if top_k <= 0 or top_k > 4096: raise ValueError( - f"top_k must be in range [1, 2048], got {top_k}. " - "Maximum supported top_k is 2048 for Blackwell architecture.") + f"top_k must be in range [1, 4096], got {top_k}. " + "Maximum supported top_k is 4096 (filtered_topk_max_k staging raised from 2048)." + ) if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") @@ -5754,7 +5754,7 @@ def cute_dsl_topk_decode_multi_cta_blackwell( Args: input_values: Input logits tensor [batch_size * next_n, vocab_size] seq_lens: Sequence lengths for each batch [batch_size] - top_k: Number of top elements to select (max 2048) + top_k: Number of top elements to select (max 4096) next_n: Number of candidates per sequence (for speculative decoding) num_copy_bits: Number of bits for vectorized memory copy (128 or 256) chunk_size_per_cta: Number of columns each CTA processes @@ -5774,10 +5774,11 @@ def cute_dsl_topk_decode_multi_cta_blackwell( "Use standard top-k implementation for older architectures.") # Validate inputs - if top_k <= 0 or top_k > 2048: + if top_k <= 0 or top_k > 4096: raise ValueError( - f"top_k must be in range [1, 2048], got {top_k}. " - "Maximum supported top_k is 2048 for Blackwell architecture.") + f"top_k must be in range [1, 4096], got {top_k}. " + "Maximum supported top_k is 4096 (filtered_topk_max_k staging raised from 2048)." + ) if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") @@ -5885,7 +5886,7 @@ def cute_dsl_indexer_topk_decode( input_values: Input logits tensor [batch_size * next_n, vocab_size] seq_lens: Sequence lengths for each batch [batch_size] output_indices: Pre-allocated output buffer [batch_size * next_n, top_k] - top_k: Number of top elements to select (max 2048) + top_k: Number of top elements to select next_n: Number of candidates per sequence (for speculative decoding) num_copy_bits: Number of bits for vectorized memory copy (128 or 256) dynamic: Use dynamic multi-CTA scheduling (for 2-pass multi-CTA) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py index 66b79c6986f8..ab16d033dba2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_decode_varlen.py @@ -56,7 +56,6 @@ --top_k 2048 --do_ref_check --return_val --do_benchmark Constraints for this example: -* The problem size of top_k <= 2048. * The input tensor has data contiguous on the n dimension (row-major). * The supported input data types are Float32, Float16, or BFloat16. """ @@ -346,7 +345,7 @@ def filtered_topk_kernel( g_num_input = None s_indices = smem.allocate_tensor( element_type=self.index_type, - layout=cute.make_ordered_layout((self.filtered_topk_max_k,), order=(0)), + layout=cute.make_ordered_layout((self.top_k,), order=(0)), byte_alignment=128, ) s_input_idx = smem.allocate_tensor( @@ -1393,9 +1392,6 @@ def run_topk_decode( parser.add_argument("--use_cold_l2", action="store_true", default=True, help="Use cold L2") args = parser.parse_args() - if args.top_k % 2 != 0: - parser.error("top_k must be a multiple of 2 (got top_k={})".format(args.top_k)) - run_topk_decode( dtype=args.dtype, batch_size=args.batch_size, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py index be796763e8e2..71d8bcbe79f2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/filtered_top_k_varlen_util.py @@ -58,8 +58,6 @@ def __init__( self.num_ctas_per_row = num_ctas_per_row self.merge_blocks = merge_blocks - # Note: now we only support top_k <= 2048, we could change the code here to support larger top_k. - self.filtered_topk_max_k = 2048 # 8 bits for radix-based filter. self.radix = 256 @@ -963,13 +961,14 @@ def filtered_topk_kernel_per_row( cute.arch.barrier() # Phase 3: Output phase + output_vector_width = 2 if self.top_k % 2 == 0 else 1 vecsize_out = cutlass.const_expr( min( self.top_k, cute.ceil_div(self.top_k, self.num_threads_per_cta), self.num_copy_bits // self.dtype.width, # TODO: only tested for float32. need to check for other dtypes. - 2, + output_vector_width, ) ) assert self.top_k % vecsize_out == 0 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..fdf4ecfb75b1 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,68 @@ 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_invalid_request_or_block_count(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.zeros((1, 2, 2, 3), dtype=torch.int32) + + destination.fill_(-1) + 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 + with pytest.raises(Exception, match="block count"): + index_mapper.gather_k_block_offsets(source, destination, [11], 4) + with pytest.raises(Exception, match="distinct source and destination storage"): + index_mapper.gather_k_block_offsets(source, source, [11], 4) + + 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/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index ee0e2bc2129c..09d23ff9be56 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -1387,6 +1387,133 @@ def run_fn(logits, seq_lens): ) +def _large_k_indexer_case(top_k: int, generation: int): + """Build a deterministic, tie-heavy input for the direct indexer op.""" + batch_size = 2 + next_n = 2 + num_rows = batch_size * next_n + num_tokens = ((top_k + 521 + 31) // 32) * 32 + columns = torch.arange(num_tokens, dtype=torch.int64, device="cuda") + rows = torch.arange(num_rows, dtype=torch.int64, device="cuda")[:, None] + logits = ((columns[None, :] * 17 + rows * 29 + generation * 43) % 257).to(torch.float32) + seq_lens = torch.tensor( + [num_tokens - 3 - generation, num_tokens - 17 + generation], + dtype=torch.int32, + device="cuda", + ) + return logits.contiguous(), seq_lens, next_n + + +def _assert_large_k_indexer_result( + logits: torch.Tensor, + seq_lens: torch.Tensor, + indices: torch.Tensor, + top_k: int, + next_n: int, +): + """Compare indices with a tie-safe, independent torch.topk oracle.""" + num_rows = logits.shape[0] + row_offsets = torch.arange(num_rows, device="cuda") % next_n + row_seq_lens = seq_lens.repeat_interleave(next_n) - next_n + row_offsets + 1 + + assert indices.shape == (num_rows, top_k) + assert indices.dtype == torch.int32 + assert indices.is_contiguous() + assert torch.all(indices >= 0) + assert torch.all(indices < row_seq_lens[:, None]) + + sorted_indices = indices.sort(dim=1).values + assert torch.all(sorted_indices[:, 1:] > sorted_indices[:, :-1]) + + selected_values = logits.gather(1, indices.to(torch.int64)).sort(dim=1, descending=True).values + oracle_values = [] + for row in range(num_rows): + row_seq_len = int(row_seq_lens[row].item()) + oracle_values.append( + torch.topk(logits[row, :row_seq_len], top_k).values.sort(descending=True).values + ) + oracle_values = torch.stack(oracle_values) + torch.testing.assert_close(selected_values, oracle_values, rtol=0.0, atol=0.0) + + +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") +@skip_pre_blackwell +@pytest.mark.parametrize("index_topk", [4096, 4097, 8192]) +def test_cute_dsl_indexer_topk_decode_large_k_cuda_graph(index_topk): + """Validate eager and captured direct-indexer execution above k=2048.""" + eager_logits, eager_seq_lens, next_n = _large_k_indexer_case(index_topk, 0) + num_rows = eager_logits.shape[0] + eager_indices = torch.empty(num_rows, index_topk, dtype=torch.int32, device="cuda") + eager_pointer = eager_indices.data_ptr() + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + input_values=eager_logits, + seq_lens=eager_seq_lens, + output_indices=eager_indices, + top_k=index_topk, + next_n=next_n, + num_copy_bits=256, + ) + torch.cuda.synchronize() + assert eager_indices.data_ptr() == eager_pointer + _assert_large_k_indexer_result(eager_logits, eager_seq_lens, eager_indices, index_topk, next_n) + + graph_logits, graph_seq_lens, _ = _large_k_indexer_case(index_topk, 1) + graph_indices = torch.empty_like(eager_indices) + graph_pointers = ( + graph_logits.data_ptr(), + graph_seq_lens.data_ptr(), + graph_indices.data_ptr(), + ) + + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + input_values=graph_logits, + seq_lens=graph_seq_lens, + output_indices=graph_indices, + top_k=index_topk, + next_n=next_n, + num_copy_bits=256, + ) + torch.cuda.current_stream().wait_stream(warmup_stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + input_values=graph_logits, + seq_lens=graph_seq_lens, + output_indices=graph_indices, + top_k=index_topk, + next_n=next_n, + num_copy_bits=256, + ) + torch.cuda.synchronize() + assert graph_pointers == ( + graph_logits.data_ptr(), + graph_seq_lens.data_ptr(), + graph_indices.data_ptr(), + ) + _assert_large_k_indexer_result(graph_logits, graph_seq_lens, graph_indices, index_topk, next_n) + + for generation in (2, 3): + replay_logits, replay_seq_lens, _ = _large_k_indexer_case(index_topk, generation) + graph_logits.copy_(replay_logits) + graph_seq_lens.copy_(replay_seq_lens) + graph_indices.fill_(-1) + graph.replay() + torch.cuda.synchronize() + assert graph_pointers == ( + graph_logits.data_ptr(), + graph_seq_lens.data_ptr(), + graph_indices.data_ptr(), + ) + _assert_large_k_indexer_result( + graph_logits, graph_seq_lens, graph_indices, index_topk, next_n + ) + + @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") @skip_pre_blackwell @pytest.mark.parametrize("batch_size", [1, 16, 256]) diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py new file mode 100644 index 000000000000..cf46fe186c15 --- /dev/null +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the layered V2 adapter over the existing sparse-KV updater.""" + +from typing import NamedTuple, Optional + +import pytest +import torch + +import tensorrt_llm # noqa: F401 # Register torch.ops.trtllm operators. + +_TOKENS_PER_BLOCK = 4 +_NUM_KV_HEADS = 2 +_BATCH_SIZE = 2 +_MAX_PAGES_PER_SEQUENCE = 3 +_NUM_PAGES = _BATCH_SIZE * _MAX_PAGES_PER_SEQUENCE +_PAGE_INDEX_DIVISOR = 2 + + +def _encode_k_block_offsets( + page_table: torch.Tensor, page_index_scale: int = _PAGE_INDEX_DIVISOR +) -> torch.Tensor: + encoded = torch.empty( + page_table.shape[0], + 2, + page_table.shape[1], + dtype=torch.int32, + device=page_table.device, + ) + encoded[:, 0] = page_table * page_index_scale + encoded[:, 1] = encoded[:, 0] + 1 + return encoded[:, 0] + + +class _DeviceArguments(NamedTuple): + pool_pointers: torch.Tensor + source_indices: torch.Tensor + source_offsets: torch.Tensor + source_layer_indices: Optional[torch.Tensor] + + +def _make_pools( + num_layers: int, + dtype: torch.dtype, + head_dim: int, + page_index_scale: int = _PAGE_INDEX_DIVISOR, +) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: + num_pages = _NUM_PAGES * page_index_scale // _PAGE_INDEX_DIVISOR + shape = ( + num_pages, + 2, + _NUM_KV_HEADS, + _TOKENS_PER_BLOCK, + head_dim, + ) + numel = torch.Size(shape).numel() + pools_cpu = [ + ((torch.arange(numel, dtype=torch.int32) + layer * 37) % 251).reshape(shape).to(dtype) + for layer in range(num_layers) + ] + pools = [pool.cuda() for pool in pools_cpu] + raw_pages = [[4, 1, 5], [2, 0, 3]] + assert set(raw_pages[0]).isdisjoint(raw_pages[1]) + raw_page_table = torch.tensor(raw_pages, dtype=torch.int32, device="cuda") + page_table = _encode_k_block_offsets(raw_page_table, page_index_scale) + page_tables = [page_table] * num_layers + assert page_tables[0].stride(0) == 2 * page_tables[0].shape[1] + return pools_cpu, pools, page_tables + + +def _device_arguments( + pools: list[torch.Tensor], + source_indices: torch.Tensor, + source_offsets: torch.Tensor, + source_layer_indices: Optional[torch.Tensor] = None, +) -> _DeviceArguments: + device = pools[0].device + return _DeviceArguments( + pool_pointers=torch.tensor( + [pool.data_ptr() for pool in pools], dtype=torch.int64, device=device + ), + source_indices=source_indices.to(device), + source_offsets=source_offsets.to(device), + source_layer_indices=( + None if source_layer_indices is None else source_layer_indices.to(device) + ), + ) + + +def _reference_compact( + pools: list[torch.Tensor], + page_tables: list[torch.Tensor], + source_indices: torch.Tensor, + source_offsets: torch.Tensor, + destination_base: int, + source_layer_indices: Optional[torch.Tensor] = None, +) -> list[torch.Tensor]: + original = [pool.clone() for pool in pools] + expected = [pool.clone() for pool in pools] + for group_layer, (source_pool, destination_pool, page_table) in enumerate( + zip(original, expected, page_tables) + ): + raw_page_table = page_table // _PAGE_INDEX_DIVISOR + if source_indices.ndim == 2: + layer_sources = source_indices + else: + assert source_layer_indices is not None + layer_sources = source_indices[int(source_layer_indices[group_layer])] + for request in range(_BATCH_SIZE): + begin = int(source_offsets[request]) + end = int(source_offsets[request + 1]) + for head in range(_NUM_KV_HEADS): + for request_move, global_move in enumerate(range(begin, end)): + source_token = int(layer_sources[head, global_move]) + destination_token = destination_base + request_move + source_page = int(raw_page_table[request, source_token // _TOKENS_PER_BLOCK]) + destination_page = int( + raw_page_table[request, destination_token // _TOKENS_PER_BLOCK] + ) + destination_pool[ + destination_page, + :, + head, + destination_token % _TOKENS_PER_BLOCK, + :, + ] = source_pool[ + source_page, + :, + head, + source_token % _TOKENS_PER_BLOCK, + :, + ] + return expected + + +def _compact( + pools: list[torch.Tensor], + page_tables: list[torch.Tensor], + arguments: _DeviceArguments, + destination_base: int, +) -> None: + torch.ops.trtllm.sparse_kv_cache_compact_layers( + pools, + arguments.pool_pointers, + page_tables[0], + arguments.source_indices, + arguments.source_offsets, + arguments.source_layer_indices, + destination_base, + ) + + +@pytest.mark.parametrize( + "dtype,head_dim", + [ + (torch.float16, 16), + (torch.bfloat16, 32), + (torch.float32, 64), + (torch.float16, 128), + (torch.bfloat16, 256), + (torch.float32, 256), + ], +) +@pytest.mark.parametrize( + "destination_base,page_index_scale", + [(0, 2), (2, 2), (2, 4)], +) +def test_sparse_kv_cache_compact_layers(dtype, head_dim, destination_base, page_index_scale): + pools_cpu, pools, page_tables = _make_pools(3, dtype, head_dim, page_index_scale) + page_tables_cpu = [page_table.cpu() for page_table in page_tables] + source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) + source_row = torch.tensor([2, 5, 8, 3, 7, 10], dtype=torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + expected = _reference_compact( + pools_cpu, + page_tables_cpu, + source_indices, + source_offsets, + destination_base, + ) + arguments = _device_arguments(pools, source_indices, source_offsets) + + _compact(pools, page_tables, arguments, destination_base) + torch.cuda.synchronize() + + for actual, reference in zip(pools, expected): + assert torch.equal(actual.cpu(), reference) + + +def test_sparse_kv_cache_compact_layers_per_layer_source(): + pools_cpu, pools, page_tables = _make_pools(2, torch.bfloat16, 64) + page_tables_cpu = [page_table.cpu() for page_table in page_tables] + source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) + source_indices = torch.tensor( + [ + [[2, 5, 8, 3, 7, 10], [3, 6, 9, 2, 5, 8]], + [[3, 7, 10, 2, 6, 9], [2, 5, 9, 3, 6, 10]], + [[4, 7, 9, 3, 6, 8], [3, 5, 8, 4, 7, 10]], + ], + dtype=torch.int32, + ) + source_layer_indices = torch.tensor([2, 0], dtype=torch.int32) + destination_base = 2 + expected = _reference_compact( + pools_cpu, + page_tables_cpu, + source_indices, + source_offsets, + destination_base, + source_layer_indices, + ) + arguments = _device_arguments( + pools, + source_indices, + source_offsets, + source_layer_indices, + ) + + _compact(pools, page_tables, arguments, destination_base) + torch.cuda.synchronize() + + for actual, reference in zip(pools, expected): + assert torch.equal(actual.cpu(), reference) + + +def test_sparse_kv_cache_compact_layers_multiple_tiles(): + num_layers = 2 + max_pages_per_sequence = 24 + num_pages = _BATCH_SIZE * max_pages_per_sequence + shape = (num_pages, 2, _NUM_KV_HEADS, _TOKENS_PER_BLOCK, 64) + numel = torch.Size(shape).numel() + pools_cpu = [ + ((torch.arange(numel, dtype=torch.int32) + layer * 37) % 251) + .reshape(shape) + .to(torch.bfloat16) + for layer in range(num_layers) + ] + pools = [pool.cuda() for pool in pools_cpu] + raw_page_table = torch.arange(num_pages, dtype=torch.int32, device="cuda").reshape( + _BATCH_SIZE, max_pages_per_sequence + ) + page_table = _encode_k_block_offsets(raw_page_table) + page_tables = [page_table] * num_layers + assert page_tables[0].stride(0) == 2 * max_pages_per_sequence + page_tables_cpu = [table.cpu() for table in page_tables] + source_offsets = torch.tensor([0, 40, 75], dtype=torch.int32) + source_row = torch.cat((torch.arange(40, 80), torch.arange(36, 71))).to(torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + destination_base = 2 + expected = _reference_compact( + pools_cpu, + page_tables_cpu, + source_indices, + source_offsets, + destination_base, + ) + arguments = _device_arguments(pools, source_indices, source_offsets) + + _compact(pools, page_tables, arguments, destination_base) + torch.cuda.synchronize() + + for actual, reference in zip(pools, expected): + assert torch.equal(actual.cpu(), reference) + + +def test_sparse_kv_cache_compact_layers_cuda_graph_replay(): + """Check operation-level capture safety, not a standalone TriAttention graph.""" + pools_cpu, pools, page_tables = _make_pools(3, torch.bfloat16, 64) + page_tables_cpu = [page_table.cpu() for page_table in page_tables] + source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) + source_row = torch.tensor([2, 5, 8, 3, 7, 10], dtype=torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + replay_row = torch.tensor([3, 6, 9, 2, 5, 8], dtype=torch.int32) + replay_indices = replay_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + destination_base = 2 + expected = _reference_compact( + pools_cpu, + page_tables_cpu, + replay_indices, + source_offsets, + destination_base, + ) + arguments = _device_arguments(pools, source_indices, source_offsets) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _compact(pools, page_tables, arguments, destination_base) + + for pool, initial in zip(pools, pools_cpu): + pool.copy_(initial) + arguments.source_indices.copy_(replay_indices) + graph.replay() + torch.cuda.synchronize() + + for actual, reference in zip(pools, expected): + assert torch.equal(actual.cpu(), reference) From 25db68879c4b77613c1cdef2859afd687f5f85b5 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 14 Jul 2026 21:40:08 -0700 Subject: [PATCH 003/178] [None][feat] Port TriAttention onto the request-carried compression channel Eager eviction pipeline (the standalone compression CUDA graph was removed), optimized compaction kernels, and draft KV co-compression: one-model speculation applies the target union keep set to the draft cache and protects the unconfirmed draft tail; modes whose draft KV cannot join the unified eviction fail fast at admission. Signed-off-by: tianruih --- examples/triattention/README.md | 128 + .../triattention/__init__.py | 21 + .../triattention/compaction.py | 626 ++++ .../triattention/triattention.py | 2721 +++++++++++++++++ .../triattention/triattention_kernels.py | 1057 +++++++ tensorrt_llm/_torch/pyexecutor/_util.py | 85 +- tensorrt_llm/llmapi/__init__.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 74 +- .../usage/llm_args_golden_manifest.json | 124 + .../test_lists/test-db/l0_b200.yml | 1 + .../test_rope_fusion_gate.py | 91 + .../test_triattention_draft_cocompaction.py | 534 ++++ .../test_triattention_eager.py | 1015 ++++++ .../test_triattention_pipeline.py | 1839 +++++++++++ .../api_stability/references/llm.yaml | 2 +- 15 files changed, 8310 insertions(+), 12 deletions(-) create mode 100644 examples/triattention/README.md create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py create mode 100644 tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py diff --git a/examples/triattention/README.md b/examples/triattention/README.md new file mode 100644 index 000000000000..c152607c5006 --- /dev/null +++ b/examples/triattention/README.md @@ -0,0 +1,128 @@ +# 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 `top_B` 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 it and converts it to its runtime schema at the first request. +2. **Periodic eviction (Stage during generation).** Every `beta` confirmed generation tokens, once a sequence is over budget, TriAttention scores the whole cache, selects `top_B` tokens to keep (the prompt tokens are preserved on top of the budget), and physically compacts the KV cache down to the kept 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`. The scoring and compaction kernels are implemented in **Triton**. + +## Support Matrix + +* GPU Compute Capability >= 9.0 (Hopper or newer) +* FP16 / BF16 +* Paged KV Cache (`KVCacheManagerV2`) +* Tensor Parallel +* PyTorch backend + +**Notes:** +1. TriAttention requires `enable_block_reuse=False` in the KV-cache configuration — the eviction physically rewrites stored keys, which is incompatible with block reuse. The construction step rejects a cache manager that has block reuse enabled. +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). + +## 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( + top_B=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 block reuse disabled. +kv_config = KvCacheConfig(enable_block_reuse=False, use_kv_cache_manager_v2=True) + +llm = LLM( + model="", + backend="pytorch", + kv_cache_compression_config=compression_config, + kv_cache_config=kv_config, +) + +# 4. 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 + top_B: 2048 + beta: 64 + eviction_mode: union + calibration_path: /path/to/qwen3-8b-calibration.pt + model_path: +kv_cache_config: + enable_block_reuse: false + 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: + +* **`top_B`** (int, default=1024): Tokens kept at each eviction (the upstream `budget`). Prompt tokens are always preserved on top of this. Smaller `top_B` → 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=`per_layer`): Which token set each eviction keeps. + * `per_layer`: score a layer, average over heads, keep one set per layer (the simplest variant). + * `per_head`: each KV head keeps its own set, shared across layers (mean of per-layer maxima). The upstream AIME default. + * `per_layer_perhead`: each head keeps its own set, fully independent per layer. + * `union`: union of the per-head top-k, then re-ranked. +* **`window_size`** (int, default=128): Most-recent tokens always preserved from eviction. Prevents the scorer from evicting freshly generated tokens. +* **`normalize_scores`** (bool, default=True): Z-normalize each head's scores over the decode region before selection (upstream default). Used by `per_head` / `per_layer_perhead` / `union`; ignored by `per_layer`. +* **`pin_prefill`** (bool, default=True): Always preserve the prompt (prefill) tokens; only decode tokens compete for the budget (upstream behaviour). Used by `per_head` / `per_layer_perhead` / `union`; `per_layer` uses the recency window instead. +* **`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 only to derive the model's RoPE tables when converting the official calibration file. 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..d40381e114cc --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py @@ -0,0 +1,21 @@ +"""TriAttention KV-cache compression: periodic physical KV eviction driven by +trigonometric importance scoring. + +TriAttention is a pure KV-cache compression method. Decode still runs the model's +standard attention over the compacted cache. The manager publishes each request's +cumulative evicted count on ``LlmRequest.py_num_compressed_tokens``; the model +engine subtracts it when building the cached-token metadata. With one-model +speculative decoding, the separate draft KV cache is compacted in the same round +with the target's kept token set (union mode only), so target and draft always +share one physical KV length. + +Public surface: + - ``TriAttention`` -- the ``BaseKVCacheCompressionManager`` (the eviction + manager; snapshots allocation metadata before forward and compacts the + finalized prefix in ``on_generation_step_end``). It uses V2 capacity-only + decode, so there is no KV-cache-manager subclass. +""" + +from .triattention import TriAttention + +__all__ = ["TriAttention"] diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py new file mode 100644 index 000000000000..0ba1292636a0 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -0,0 +1,626 @@ +# 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. + +"""Batched physical KV-cache compaction for eviction-based compression. + +Given each request's kept-token ordinals, its valid sequence length, and the +staged V2 block offsets, this module packs per-request move indices with one +prepared Triton launch and then moves the surviving KV in place with batched +C++ compact launches. Inputs are plain tensors, so any eviction method that +produces a kept-token set per request can drive it. +""" + +from collections import OrderedDict +from typing import Callable, Dict, List, NamedTuple, Optional, Tuple + +import torch + +_SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + + +class _PreparedTritonKernelLaunch: + """Replay one Triton kernel launch frozen at build time. + + ``warmup`` JIT-compiles the kernel once for a fixed grid, bound tensor + set, and constexpr set; ``__call__`` re-launches the compiled binary + directly, skipping Triton's per-call dispatch. Constexpr values are + passed positionally on replay, so their dict order must match the + kernel's constexpr parameter declaration order. + """ + + def __init__( + self, + triton_kernel, + bound_tensors: Tuple[torch.Tensor, ...], + constexpr_values: Dict[str, object], + *, + grid: Tuple[int, ...], + num_warps: int, + ) -> None: + self.device = bound_tensors[0].device + self.bound_tensors = tuple(bound_tensors) + self.constexpr_values = dict(constexpr_values) + with torch.cuda.device(self.device): + self.build_stream = torch.cuda.current_stream(self.device) + compiled = triton_kernel.warmup( + *self.bound_tensors, + **self.constexpr_values, + num_warps=num_warps, + grid=grid, + ) + self.compiled_kernel_runner = compiled[grid] + + def __call__(self, *replay_tensors: torch.Tensor) -> None: + """Replay the launch; ``replay_tensors``, if given, substitute the bound tensors.""" + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.build_stream.device, + self.build_stream.cuda_stream, + ): + raise RuntimeError( + "a prepared Triton kernel launch must run on the stream it was built on" + ) + self.compiled_kernel_runner( + *(replay_tensors if replay_tensors else self.bound_tensors), + *self.constexpr_values.values(), + stream=self.build_stream.cuda_stream, + ) + + +class _CppCompactGroup(NamedTuple): + """One layered sparse-KV updater launch over pools sharing a block table.""" + + pools: Tuple[torch.Tensor, ...] + page_table: torch.Tensor + pool_pointers: torch.Tensor + source_layer_indices: Optional[torch.Tensor] + + def launch(self, source: torch.Tensor, offsets: torch.Tensor, destination_base: int) -> None: + torch.ops.trtllm.sparse_kv_cache_compact_layers( + list(self.pools), + self.pool_pointers, + self.page_table, + source, + offsets, + self.source_layer_indices, + destination_base, + ) + + +class _SingleCacheCompaction(NamedTuple): + """One compacted cache family (target dense, target SWA, or draft). + + Holds the prepared launch that packs this family's move indices (None + when an earlier family's pack launch fills them in the same call), the + C++ launch groups that consume them, and the destination base the moved + tokens land at. + """ + + prepared_move_index_pack: Optional[_PreparedTritonKernelLaunch] + cpp_launch_groups: Tuple[_CppCompactGroup, ...] + move_source_indices: torch.Tensor + move_source_offsets: torch.Tensor + destination_base: int + + def launch(self) -> None: + if self.prepared_move_index_pack is not None: + self.prepared_move_index_pack() + for group in self.cpp_launch_groups: + group.launch(self.move_source_indices, self.move_source_offsets, self.destination_base) + + +def _cuda_int32_contiguous(tensors: Tuple[torch.Tensor, ...], device: torch.device) -> bool: + return all( + tensor.is_cuda + and tensor.dtype == torch.int32 + and tensor.device == device + and tensor.is_contiguous() + for tensor in tensors + ) + + +def _validated_kv_head_count( + pools: List[torch.Tensor], + layers: Tuple[int, ...], + device: torch.device, + what: str, +) -> int: + """Return the common KV-head count of one launch side's pools. + + The C++ compact op reads the interleaved V2 layout + ``[page, K/V, head, token, dim]`` and takes the KV-head count from each + launch's pool shape, so every layer on one side must agree on it. + """ + first = pools[layers[0]] + num_kv_heads = int(first.shape[2]) if first.ndim == 5 and first.shape[2] > 0 else -1 + if not all( + pools[layer].ndim == 5 + and pools[layer].shape[1] == 2 + and pools[layer].device == device + and int(pools[layer].shape[2]) == num_kv_heads + and pools[layer].is_contiguous() + and pools[layer].dtype in _SUPPORTED_POOL_DTYPES + for layer in layers + ): + raise ValueError( + f"{what} requires contiguous interleaved BF16/FP16/FP32 pools " + "with one common KV-head count" + ) + return num_kv_heads + + +def _make_move_buffers( + index_prefix: Tuple[int, ...], + moves_per_request: List[int], + device: torch.device, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Allocate the packed source-index buffer and its per-request offsets.""" + offsets = [0] + for count in moves_per_request: + offsets.append(offsets[-1] + count) + indices = torch.empty((*index_prefix, offsets[-1]), dtype=torch.int32, device=device) + return indices, torch.tensor(offsets, dtype=torch.int32, device=device) + + +def _validated_tail_lengths( + tail_lengths: Optional[List[int]], + request_count: int, + what: str, +) -> Tuple[int, ...]: + if tail_lengths is None: + tail_lengths = [0] * request_count + if len(tail_lengths) != request_count or any(length < 0 for length in tail_lengths): + raise ValueError(f"{what} lengths must match the request count") + return tuple(int(length) for length in tail_lengths) + + +def _page_table_provider( + page_table_slots: Dict[int, int], + kv_block_offsets: torch.Tensor, + device: torch.device, + request_count: int, + what: str, +) -> Callable[[int], torch.Tensor]: + """Return validated per-slot K block-offset views, cached per slot.""" + tables: Dict[int, torch.Tensor] = {} + + def page_table_for(representative: int) -> torch.Tensor: + slot = page_table_slots[representative] + if slot not in tables: + block_offsets = kv_block_offsets[slot, :request_count, 0] + if block_offsets.device != device or block_offsets.dtype != torch.int32: + raise ValueError(f"{what} block offsets must be int32 tensors on the pool device") + if block_offsets.ndim != 2 or block_offsets.stride(1) != 1: + raise ValueError(f"{what} K block offsets must have a contiguous block dimension") + tables[slot] = block_offsets + return tables[slot] + + return page_table_for + + +def _compact_groups( + entries: List[Tuple[int, torch.Tensor, torch.Tensor]], + mode: str, + pool_keys: Tuple[object, ...], + device: torch.device, + per_layer_slots: Optional[Dict[int, int]] = None, +) -> Tuple[_CppCompactGroup, ...]: + """Batch layers into one C++ launch per uniform V2 pool. + + ``per_layer_slots`` maps each layer to its selection row; it is only set + when every dense layer keeps its own token set (per-layer eviction). + """ + grouped = OrderedDict() + for layer, pool, page_table in entries: + key = ( + pool_keys[layer], + mode, + str(pool.dtype), + str(pool.device), + tuple(int(value) for value in pool.shape[1:]), + tuple(int(value) for value in page_table.shape), + ) + grouped.setdefault(key, []).append((layer, pool, page_table)) + + result = [] + for group_entries in grouped.values(): + layers = tuple(entry[0] for entry in group_entries) + pools = tuple(entry[1] for entry in group_entries) + page_tables = tuple(entry[2] for entry in group_entries) + if len({int(pool.data_ptr()) for pool in pools}) != len(pools): + raise ValueError("layered compaction requires a distinct pool view for every layer") + if len({int(page_table.data_ptr()) for page_table in page_tables}) != 1: + raise ValueError("layers in one V2 pool must share one block-offset table") + source_layer_indices = None + if per_layer_slots is not None: + source_layer_indices = torch.tensor( + [per_layer_slots[layer] for layer in layers], + dtype=torch.int32, + device=device, + ) + result.append( + _CppCompactGroup( + pools=pools, + page_table=page_tables[0], + pool_pointers=torch.tensor( + [pool.data_ptr() for pool in pools], + dtype=torch.int64, + device=device, + ), + source_layer_indices=source_layer_indices, + ) + ) + return tuple(result) + + +def _prepared_move_index_pack_launch( + kept_token_ordinals: torch.Tensor, + valid_sequence_lengths: torch.Tensor, + move_source_offsets: torch.Tensor, + move_source_indices: torch.Tensor, + *, + eviction_mode: str, + prompt_len: int, + keep_count: int, + num_dense_layers: int, + num_kv_heads: int, + max_protected_tail: int, + swa_window: int, + swa_move_source_offsets: Optional[torch.Tensor], + swa_move_source_indices: Optional[torch.Tensor], +) -> _PreparedTritonKernelLaunch: + """Build one prepared launch of the move-index packing kernel. + + The kernel reads the kept-token ordinals and each request's valid length + and writes the packed per-(layer, head) move source indices consumed by + the C++ compact launches. Only the caller-provided selection tensors are + validated here; the move buffers are allocated by this module. + """ + per_layer = eviction_mode == "per_layer_perhead" + union = eviction_mode == "union" + request_count = int(kept_token_ordinals.shape[0]) if kept_token_ordinals.ndim else 0 + if union: + selection_rows = 1 + elif per_layer: + selection_rows = num_dense_layers * num_kv_heads + else: + selection_rows = num_kv_heads + selection_prefix = (request_count,) if union else (request_count, selection_rows) + if ( + request_count <= 0 + or tuple(kept_token_ordinals.shape) != (*selection_prefix, prompt_len + keep_count) + or valid_sequence_lengths.shape != (request_count,) + ): + raise ValueError("prepared compaction packing requires one valid fixed geometry") + + device = kept_token_ordinals.device + if not _cuda_int32_contiguous((kept_token_ordinals, valid_sequence_lengths), device): + raise ValueError("prepared compaction packing requires contiguous CUDA int32 tensors") + + if swa_move_source_indices is not None: + swa_offsets_arg = swa_move_source_offsets + swa_indices_arg = swa_move_source_indices + swa_total = int(swa_move_source_indices.shape[-1]) + else: + # HAS_SWA specializes all corresponding loads and stores away. + swa_offsets_arg = move_source_offsets + swa_indices_arg = move_source_indices + swa_total = 0 + + from .triattention_kernels import _pack_compaction_sources_kernel + + block = 256 + max_move = keep_count + max_protected_tail + if swa_total: + max_move = max(max_move, swa_window + max_protected_tail) + packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads + grid = (request_count, packed_row_count, (max_move + block - 1) // block) + bound_tensors = ( + kept_token_ordinals, + valid_sequence_lengths, + move_source_offsets, + move_source_indices, + swa_offsets_arg, + swa_indices_arg, + ) + # Ordered to match the kernel's constexpr parameter declaration: the + # prepared launch replays these by position. + constexpr_values = dict( + DENSE_TOTAL=int(move_source_indices.shape[-1]), + SWA_TOTAL=swa_total, + SELECTION_ROWS=selection_rows, + SELECTION_STRIDE=prompt_len + keep_count, + KEEP_COUNT=keep_count, + PROMPT_LEN=prompt_len, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=swa_total > 0, + BLOCK=block, + ) + return _PreparedTritonKernelLaunch( + _pack_compaction_sources_kernel, + bound_tensors, + constexpr_values, + grid=grid, + num_warps=4, + ) + + +class BatchedKVCacheCompaction: + """Batched physical compaction of the KV caches for one fixed geometry. + + Dense layers keep the prompt in place and compact the selected decode + tokens plus any target KV reserved for the next overlapped forward; + kernel-masked SWA layers keep the latest window plus the same protected + tail. A co-compressed draft cache reuses the target's kept token ordinals + (broadcast over the draft's own KV-head count) plus the draft's own + protected tail, landing at the same destination base. + """ + + def __init__( + self, + *, + eviction_mode: str, + layer_pools: List[torch.Tensor], + dense_layers: List[int], + swa_layers: List[int], + layer_group_representative: Dict[int, int], + kept_token_ordinals: torch.Tensor, + valid_sequence_lengths: torch.Tensor, + kv_block_offsets: torch.Tensor, + page_table_slots: Dict[int, int], + request_count: int, + prompt_len: int, + decode_keep_count: int, + swa_window: Optional[int], + protected_tail_lengths: Optional[List[int]] = None, + layer_pool_keys: Optional[List[object]] = None, + draft_layer_pools: Optional[List[torch.Tensor]] = None, + draft_layers: Optional[List[int]] = None, + draft_layer_group_representative: Optional[Dict[int, int]] = None, + draft_layer_pool_keys: Optional[List[object]] = None, + draft_protected_tail_lengths: Optional[List[int]] = None, + draft_kv_block_offsets: Optional[torch.Tensor] = None, + draft_page_table_slots: Optional[Dict[int, int]] = None, + ) -> None: + if eviction_mode not in ("union", "per_head", "per_layer_perhead"): + raise ValueError(f"unsupported compaction mode: {eviction_mode}") + if request_count <= 0 or decode_keep_count <= 0 or prompt_len < 0: + raise ValueError("batched compaction requires requests and retained tokens") + if not dense_layers: + raise ValueError("batched compaction requires at least one dense layer") + + self.eviction_mode = eviction_mode + self.device = layer_pools[dense_layers[0]].device + self.request_count = int(request_count) + self.prompt_len = int(prompt_len) + self.decode_keep_count = int(decode_keep_count) + self.keep_count = self.prompt_len + self.decode_keep_count + self.protected_tail_lengths = _validated_tail_lengths( + protected_tail_lengths, self.request_count, "protected-tail" + ) + self.max_protected_tail = max(self.protected_tail_lengths, default=0) + self.dense_layers = tuple(int(layer) for layer in dense_layers) + self.swa_layers = tuple(int(layer) for layer in swa_layers) + if layer_pool_keys is None: + layer_pool_keys = [("layer", layer) for layer in range(len(layer_pools))] + if len(layer_pool_keys) != len(layer_pools): + raise ValueError("pool keys must match the layer-pool count") + self.layer_pool_keys = tuple(layer_pool_keys) + + per_layer = self.eviction_mode == "per_layer_perhead" + self.num_kv_heads = _validated_kv_head_count( + layer_pools, + (*self.dense_layers, *self.swa_layers), + self.device, + "batched compaction", + ) + dense_index_prefix = ( + (len(self.dense_layers), self.num_kv_heads) if per_layer else (self.num_kv_heads,) + ) + dense_move_indices, dense_move_offsets = _make_move_buffers( + dense_index_prefix, + [self.decode_keep_count + length for length in self.protected_tail_lengths], + self.device, + ) + page_table_for = _page_table_provider( + page_table_slots, + kv_block_offsets, + self.device, + self.request_count, + "compaction", + ) + dense_entries = [ + (layer, layer_pools[layer], page_table_for(layer_group_representative[layer])) + for layer in self.dense_layers + ] + + self.swa_window = 0 + swa_move_indices = None + swa_move_offsets = None + swa_entries = [] + if self.swa_layers: + if swa_window is None or swa_window <= 0 or self.keep_count < swa_window: + raise ValueError("SWA compaction requires a valid retained window") + self.swa_window = int(swa_window) + swa_move_indices, swa_move_offsets = _make_move_buffers( + (self.num_kv_heads,), + [self.swa_window + length for length in self.protected_tail_lengths], + self.device, + ) + # SWA layers are staged as their own page-table representatives. + swa_entries = [ + (layer, layer_pools[layer], page_table_for(layer)) for layer in self.swa_layers + ] + + dense_slots = ( + {layer: slot for slot, layer in enumerate(self.dense_layers)} if per_layer else None + ) + dense_pack = _prepared_move_index_pack_launch( + kept_token_ordinals, + valid_sequence_lengths, + dense_move_offsets, + dense_move_indices, + eviction_mode=self.eviction_mode, + prompt_len=self.prompt_len, + keep_count=self.decode_keep_count, + num_dense_layers=len(self.dense_layers), + num_kv_heads=self.num_kv_heads, + max_protected_tail=self.max_protected_tail, + swa_window=self.swa_window, + swa_move_source_offsets=swa_move_offsets, + swa_move_source_indices=swa_move_indices, + ) + self.target_dense_compaction = _SingleCacheCompaction( + prepared_move_index_pack=dense_pack, + cpp_launch_groups=_compact_groups( + dense_entries, "dense", self.layer_pool_keys, self.device, dense_slots + ), + move_source_indices=dense_move_indices, + move_source_offsets=dense_move_offsets, + destination_base=self.prompt_len, + ) + # The dense pack launch fills the SWA move buffers in the same call. + self.target_swa_compaction = None + if self.swa_layers: + self.target_swa_compaction = _SingleCacheCompaction( + prepared_move_index_pack=None, + cpp_launch_groups=_compact_groups( + swa_entries, "swa", self.layer_pool_keys, self.device + ), + move_source_indices=swa_move_indices, + move_source_offsets=swa_move_offsets, + destination_base=self.keep_count - self.swa_window, + ) + + self.draft_compaction = None + if draft_layers: + self.draft_compaction = self._build_draft_compaction( + kept_token_ordinals, + valid_sequence_lengths, + draft_layer_pools=draft_layer_pools, + draft_layers=draft_layers, + draft_layer_group_representative=draft_layer_group_representative, + draft_layer_pool_keys=draft_layer_pool_keys, + draft_protected_tail_lengths=draft_protected_tail_lengths, + draft_kv_block_offsets=draft_kv_block_offsets, + draft_page_table_slots=draft_page_table_slots, + ) + + self.cache_compactions = tuple( + compaction + for compaction in ( + self.target_dense_compaction, + self.target_swa_compaction, + self.draft_compaction, + ) + if compaction is not None + ) + + def _build_draft_compaction( + self, + kept_token_ordinals: torch.Tensor, + valid_sequence_lengths: torch.Tensor, + *, + draft_layer_pools: Optional[List[torch.Tensor]], + draft_layers: List[int], + draft_layer_group_representative: Optional[Dict[int, int]], + draft_layer_pool_keys: Optional[List[object]], + draft_protected_tail_lengths: Optional[List[int]], + draft_kv_block_offsets: Optional[torch.Tensor], + draft_page_table_slots: Optional[Dict[int, int]], + ) -> _SingleCacheCompaction: + """Build the co-compressed draft cache's own pack and launch groups. + + The draft forms its own launch groups so it may use a different + KV-head count than the target. + """ + if self.eviction_mode != "union": + raise ValueError("draft co-compaction supports only union eviction") + if ( + draft_layer_pools is None + or draft_layer_group_representative is None + or draft_layer_pool_keys is None + ): + raise ValueError("draft co-compaction requires the full draft layout") + if draft_kv_block_offsets is None or draft_page_table_slots is None: + raise ValueError("draft co-compaction requires staged draft page tables") + draft_tail_lengths = _validated_tail_lengths( + draft_protected_tail_lengths, self.request_count, "draft protected-tail" + ) + if len(draft_layer_pool_keys) != len(draft_layer_pools): + raise ValueError("draft pool keys must match the draft layer-pool count") + draft_layers = tuple(int(layer) for layer in draft_layers) + draft_num_kv_heads = _validated_kv_head_count( + draft_layer_pools, + draft_layers, + self.device, + "draft co-compaction", + ) + draft_move_indices, draft_move_offsets = _make_move_buffers( + (draft_num_kv_heads,), + [self.decode_keep_count + length for length in draft_tail_lengths], + self.device, + ) + draft_page_table_for = _page_table_provider( + draft_page_table_slots, + draft_kv_block_offsets, + self.device, + self.request_count, + "draft", + ) + draft_entries = [ + ( + layer, + draft_layer_pools[layer], + draft_page_table_for(draft_layer_group_representative[layer]), + ) + for layer in draft_layers + ] + # In union mode the pack kernel reads selection row 0 for every + # packed row, so one more prepared launch broadcasts the target keep + # set over the draft KV heads and appends the draft's own tail + # ordinals (valid_seq_len + 0..tail-1). + draft_pack = _prepared_move_index_pack_launch( + kept_token_ordinals, + valid_sequence_lengths, + draft_move_offsets, + draft_move_indices, + eviction_mode="union", + prompt_len=self.prompt_len, + keep_count=self.decode_keep_count, + num_dense_layers=1, + num_kv_heads=draft_num_kv_heads, + max_protected_tail=max(draft_tail_lengths, default=0), + swa_window=0, + swa_move_source_offsets=None, + swa_move_source_indices=None, + ) + return _SingleCacheCompaction( + prepared_move_index_pack=draft_pack, + cpp_launch_groups=_compact_groups( + draft_entries, "draft", tuple(draft_layer_pool_keys), self.device + ), + move_source_indices=draft_move_indices, + move_source_offsets=draft_move_offsets, + destination_base=self.prompt_len, + ) + + def launch(self) -> None: + """Pack the move indices, then run every cache family's C++ compacts.""" + for compaction in self.cache_compactions: + compaction.launch() 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..1c9f2c867dab --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -0,0 +1,2721 @@ +# 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. + +Every ``beta`` confirmed generation tokens TriAttention scores each cached token with a +trigonometric importance score (computed from offline-calibrated statistics of +the model's pre-RoPE query vectors) and physically deletes the tokens below the +top-B keep set. There is no context-phase work and no per-step attention mask: +the eviction runs in the compression manager's final +``on_generation_step_end`` hook. + +TriAttention is a :class:`BaseKVCacheCompressionManager` and nothing more -- it +has no attention backend of its own; decode runs the model's standard dense +kernel over the compacted cache. TriAttention derives each request's effective +confirmed physical length after V2's native update/rewind and publishes the +cumulative evicted count on ``LlmRequest.py_num_compressed_tokens``; attention +metadata reads it back through the KV cache manager on the next step. +Physical reclaim uses V2's existing resize path directly after compaction. An +already-enqueued speculative suffix is excluded from scoring and appended +unchanged to the retained prefix by the same per-layer compact operation. +With one-model speculative decoding, the separate draft KV cache is compacted +in the same round with the target's kept token set (union mode only), so +target and draft always share one physical KV length. + +KV layout: the decode kernel stores keys in HND layout +``[num_pages, kv_factor, num_kv_heads, tokens_per_block, head_dim]``. The Python +gather / score / compact code MUST read ``get_buffers`` with ``kv_layout="HND"``; +reading the default NHD silently swaps the token and head axes and scrambles the +cache. + +Position handling: kept keys retain their original RoPE rotation (no re-RoPE on +compaction). The model engine keeps the decode query at its true absolute +position while the attention metadata uses the compacted physical length, so a +query against a kept key at its original rotation still yields the correct +relative distance. + +Calibration is NOT computed here: the user calibrates with the official tool +(github.com/WeianMao/triattention) and passes that .pt via ``calibration_path``; +the manager converts it to our runtime schema at load (see _resolve_calibration). +The scoring math follows the same upstream reference (``methods/pruning_utils.py``). +""" + +from collections import OrderedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union + +import torch + +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( + _PreparedTritonKernelLaunch, +) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.resource_manager import BaseKVCacheCompressionManager +from tensorrt_llm._utils import nvtx_range, nvtx_range_debug, 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 tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig + +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from tensorrt_llm._torch.pyexecutor.scheduler.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"}) + +# Bound eager staging memory. A large due cohort is processed as consecutive +# request chunks with identical results. +_EAGER_REQUEST_CHUNK_SIZE = 32 +_EAGER_RESOURCE_CACHE_LIMIT = 3 +_EAGER_COMPACTION_CACHE_LIMIT = 6 +_INT32_MAX = torch.iinfo(torch.int32).max + + +def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Tensor: + """Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <=max].""" + if max_length < 1: + raise ValueError("offset_max_length must be >= 1") + offsets: List[float] = [] + value = 1 + while value <= max_length: + offsets.append(float(value)) + value *= 2 + return torch.tensor(offsets, device=device, dtype=torch.float32) + + +def _canonical_device(device: torch.device) -> torch.device: + device = torch.device(device) + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _topk_indices_into( + scores: torch.Tensor, + seq_lens: torch.Tensor, + indices_i32: torch.Tensor, + keep_count: int, +) -> None: + """Write per-row top-k indices with the CuTE-DSL selector.""" + torch.ops.trtllm.cute_dsl_indexer_topk_decode(scores, seq_lens, indices_i32, keep_count, 1) + + +class _PreparedCuteTopK: + """Run the existing fixed-shape CuTe TopK with owned scratch storage.""" + + def __init__(self, max_rows: int, width: int, keep_count: int, device: torch.device) -> None: + if max_rows <= 0 or width <= 0 or not 1 <= keep_count <= width: + raise ValueError("prepared CuTe TopK requires valid fixed dimensions") + if device.type != "cuda": + raise ValueError("prepared CuTe TopK requires a CUDA device") + + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + self.max_rows = max_rows + self.width = width + self.keep_count = keep_count + self.device = device + with torch.cuda.device(device): + self.stream = torch.cuda.current_stream(device) + self.scratch = torch.empty((max_rows, 2, width), dtype=torch.int32, device=device) + runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner + key = ( + cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], + 1 << (width - 1).bit_length(), + keep_count, + 1, + False, + 256, + False, + max_rows > cute_dsl_custom_ops._get_num_sms(), + ) + runner._compile(*key) + self.compiled = runner.kernel_cache[key] + + def __call__( + self, + scores: torch.Tensor, + seq_lens: torch.Tensor, + output_indices: torch.Tensor, + ) -> None: + rows = int(scores.shape[0]) + if ( + rows != self.max_rows + or scores.shape != (self.max_rows, self.width) + or scores.dtype != torch.float32 + or scores.device != self.device + or seq_lens.shape != (self.max_rows,) + or seq_lens.dtype != torch.int32 + or seq_lens.device != self.device + or output_indices.shape != (self.max_rows, self.keep_count) + or output_indices.dtype != torch.int32 + or output_indices.device != self.device + ): + raise ValueError("prepared CuTe TopK inputs do not match their fixed buffers") + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.stream.device, + self.stream.cuda_stream, + ): + raise RuntimeError("prepared CuTe TopK must run on the stream it was built on") + self.compiled( + scores, + None, + self.scratch, + None, + seq_lens, + output_indices, + None, + ) + + +class _PreparedTopKFinalizer: + """Launch the fixed deterministic finalizer without Triton JIT dispatch.""" + + def __init__( + self, + scores: torch.Tensor, + seq_lens: torch.Tensor, + provisional_indices: torch.Tensor, + output_indices: torch.Tensor, + keep_count: int, + prompt_len: int, + ) -> None: + rows, width = scores.shape + output_width = prompt_len + keep_count + if ( + not scores.is_cuda + or scores.dtype != torch.float32 + or not scores.is_contiguous() + or seq_lens.shape != (rows,) + or seq_lens.dtype != torch.int32 + or seq_lens.device != scores.device + or provisional_indices.shape != (rows, keep_count) + or provisional_indices.dtype != torch.int32 + or provisional_indices.device != scores.device + or output_indices.shape != (rows, output_width) + or output_indices.dtype != torch.int32 + or output_indices.device != scores.device + or not seq_lens.is_contiguous() + or not provisional_indices.is_contiguous() + or not output_indices.is_contiguous() + ): + raise ValueError("prepared TopK finalizer tensors do not share one fixed geometry") + + from .triattention_kernels import _finalize_topk_indices_kernel + + self._prepared_launch = _PreparedTritonKernelLaunch( + _finalize_topk_indices_kernel, + (scores, seq_lens, provisional_indices, output_indices), + dict( + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=output_width, + PROMPT_LEN=prompt_len, + BLOCK=256, + ), + grid=(rows, 1, 1), + num_warps=4, + ) + + def __call__( + self, + scores: torch.Tensor, + seq_lens: torch.Tensor, + provisional_indices: torch.Tensor, + output_indices: torch.Tensor, + ) -> None: + self._prepared_launch(scores, seq_lens, provisional_indices, output_indices) + + +class _PreparedUnionScores: + """Launch fixed union score preparation without Triton JIT dispatch.""" + + def __init__( + self, + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + combined: torch.Tensor, + *, + normalize_scores: bool, + ) -> None: + if scores.ndim != 3: + raise ValueError("prepared union scores require request-major rows") + request_count, rows, width = scores.shape + if ( + not scores.is_cuda + or scores.dtype != torch.float32 + or not scores.is_contiguous() + or valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.shape != (request_count, rows, 1) + or row_mean.dtype != torch.float32 + or row_mean.device != scores.device + or row_inv_std.shape != row_mean.shape + or row_inv_std.dtype != torch.float32 + or row_inv_std.device != scores.device + or combined.shape != (request_count, width) + or combined.dtype != torch.float32 + or combined.device != scores.device + or not valid_widths.is_contiguous() + or not row_mean.is_contiguous() + or not row_inv_std.is_contiguous() + or not combined.is_contiguous() + ): + raise ValueError("prepared union score tensors do not share one fixed geometry") + + from .triattention_kernels import _score_row_stats_kernel, _score_union_kernel + + normalize_scores = bool(normalize_scores) + # Callers verify score-tensor identity and the normalize flag against + # this launcher before dispatching to it. + self.scores = scores + self.normalize_scores = normalize_scores + self._prepared_stats_launch = None + if normalize_scores: + stats_grid = (request_count * rows, 1, 1) + self._prepared_stats_launch = _PreparedTritonKernelLaunch( + _score_row_stats_kernel, + (scores, valid_widths, row_mean, row_inv_std), + dict(ROWS=rows, WIDTH=width, BLOCK=256), + grid=stats_grid, + num_warps=4, + ) + union_grid = (request_count, (width + 31) // 32, 1) + self._prepared_union_launch = _PreparedTritonKernelLaunch( + _score_union_kernel, + (scores, valid_widths, row_mean, row_inv_std, combined), + dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), + grid=union_grid, + num_warps=1, + ) + + def __call__(self) -> None: + if self._prepared_stats_launch is not None: + self._prepared_stats_launch() + self._prepared_union_launch() + + +def _deterministic_topk_indices_into( + scores: torch.Tensor, + seq_lens: torch.Tensor, + provisional_indices_i32: torch.Tensor, + output_indices_i32: torch.Tensor, + keep_count: int, + prompt_len: int, + prepared_topk: Optional[_PreparedCuteTopK] = None, + prepared_finalizer: Optional[_PreparedTopKFinalizer] = None, +) -> None: + """Write stable, increasing physical indices around one CuTE TopK call.""" + if scores.is_cuda: + if prepared_topk is None or prepared_finalizer is None: + raise ValueError("CUDA selection requires prepared TopK launchers") + prepared_topk(scores, seq_lens, provisional_indices_i32) + prepared_finalizer( + scores, + seq_lens, + provisional_indices_i32, + output_indices_i32, + ) + return + + # CPU selectors only exercise the selector contract in unit tests. + _topk_indices_into(scores, seq_lens, provisional_indices_i32, keep_count) + for row_index, row_scores in enumerate(scores): + valid_width = int(seq_lens[row_index]) + selected = provisional_indices_i32[row_index].to(torch.long) + threshold = torch.amin(row_scores[selected]) + valid_scores = row_scores[:valid_width] + higher = torch.nonzero(valid_scores > threshold, as_tuple=False).flatten() + tied = torch.nonzero(valid_scores == threshold, as_tuple=False).flatten() + tie_count = keep_count - int(higher.numel()) + ordered = torch.sort(torch.cat((higher, tied[:tie_count]))).values + output_indices_i32[row_index, prompt_len : prompt_len + keep_count].copy_( + ordered.to(torch.int32).add(prompt_len) + ) + + +class _CrossRequestSelectionPlan(NamedTuple): + """Selection dimensions used to allocate reusable eager buffers.""" + + eviction_mode: str + dense_layers: Tuple[int, ...] + num_query_heads: int + num_kv_heads: int + rows: int + width: int + keep_count: int + prompt_len: int + dtype: torch.dtype + device: torch.device + max_requests: int + + +class _RuntimeKVLayout(NamedTuple): + """Manager-lifetime layer and pool views used by every eviction.""" + + manager: object + num_layers: int + global_layers: List[int] + layer_pools: List[torch.Tensor] + dense_layers: List[int] + swa_layers: List[int] + swa_window: Optional[int] + storage_groups: Dict[object, List[int]] + layer_group_representative: Dict[int, int] + layer_pool_keys: Tuple[object, ...] + pool_representatives: Tuple[int, ...] + pool_page_counts: Tuple[int, ...] + pool_view_fingerprint: Tuple[tuple, ...] + + +class _BatchedUnionKeepSetSelector: + """Persistent ``[request, ...]`` buffers for union selection.""" + + def __init__( + self, + rows: int, + width: int, + keep_count: int, + prompt_len: int, + *, + dtype: torch.dtype, + device: torch.device, + max_requests: int, + dense_layers: Tuple[int, ...] = (), + num_query_heads: int = 0, + num_kv_heads: int = 0, + input_scores: Optional[torch.Tensor] = None, + normalize_scores: bool = True, + ) -> None: + if rows <= 0 or width <= keep_count or keep_count <= 0: + raise ValueError("cross-request selection requires rows > 0 and width > keep_count > 0") + if max_requests <= 0: + raise ValueError("cross-request selection requires a positive request capacity") + self.max_requests = max_requests + self.eviction_mode = "union" + self.dense_layers = tuple(dense_layers) + self.num_query_heads = int(num_query_heads) + self.num_kv_heads = int(num_kv_heads) + self.rows = rows + self.width = width + self.keep_count = keep_count + self.prompt_len = prompt_len + self.total_keep = prompt_len + keep_count + self.dtype = dtype + self.device = _canonical_device(device) + if self.device.type == "cuda" and input_scores is None: + raise ValueError("CUDA union selection requires its fixed score input") + + self.row_mean = torch.empty((max_requests, rows, 1), dtype=dtype, device=self.device) + self.row_std = torch.empty_like(self.row_mean) + self.combined = torch.empty((max_requests, width), dtype=dtype, device=self.device) + self.final_indices = torch.empty( + (max_requests, keep_count), dtype=torch.int32, device=self.device + ) + self.keep = torch.empty( + (max_requests, self.total_keep), dtype=torch.int32, device=self.device + ) + self.valid_widths = torch.full( + (max_requests,), width, dtype=torch.int32, device=self.device + ) + if self.device.type == "cpu": + self.valid_scale = torch.empty((max_requests, 1, 1), dtype=dtype, device=self.device) + self.token_indices = torch.arange(width, dtype=torch.int32, device=self.device) + self.invalid_mask = torch.empty( + (max_requests, 1, width), dtype=torch.bool, device=self.device + ) + else: + self.valid_scale = None + self.token_indices = None + self.invalid_mask = None + self.prepared_topk = ( + _PreparedCuteTopK(max_requests, width, keep_count, self.device) + if self.device.type == "cuda" + else None + ) + self.prepared_finalizer = ( + _PreparedTopKFinalizer( + self.combined, + self.valid_widths, + self.final_indices, + self.keep, + keep_count, + prompt_len, + ) + if self.device.type == "cuda" + else None + ) + self.prepared_scores = ( + _PreparedUnionScores( + input_scores, + self.valid_widths, + self.row_mean, + self.row_std, + self.combined, + normalize_scores=normalize_scores, + ) + if self.device.type == "cuda" and input_scores is not None + else None + ) + if prompt_len: + prompt = torch.arange(prompt_len, dtype=torch.int32, device=self.device) + self.keep[:, :prompt_len].copy_(prompt.expand(max_requests, -1)) + + def _select_input_scores( + self, + input_scores: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, + ) -> None: + valid_widths = self.valid_widths[:request_count] + combined = self.combined[:request_count] + if input_scores.is_cuda: + raise RuntimeError("CUDA union scores must use their prepared fixed launcher") + else: + self._select_input_scores_reference( + input_scores, + request_count, + normalize_scores=normalize_scores, + ) + + final_indices = self.final_indices[:request_count] + _deterministic_topk_indices_into( + combined, + valid_widths, + final_indices, + self.keep[:request_count], + self.keep_count, + self.prompt_len, + self.prepared_topk, + self.prepared_finalizer, + ) + + def select_prepared_requests(self) -> None: + """Select from the CUDA score tensor bound to this fixed selector.""" + if self.prepared_scores is None: + raise RuntimeError("prepared union scores are unavailable") + self.prepared_scores() + final_indices = self.final_indices + _deterministic_topk_indices_into( + self.combined, + self.valid_widths, + final_indices, + self.keep, + self.keep_count, + self.prompt_len, + self.prepared_topk, + self.prepared_finalizer, + ) + + def _select_input_scores_reference( + self, + input_scores: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, + ) -> None: + """Keep a CPU-only reference for selector contract tests.""" + valid_widths = self.valid_widths[:request_count] + assert self.invalid_mask is not None + assert self.token_indices is not None + assert self.valid_scale is not None + invalid_mask = self.invalid_mask[:request_count] + torch.ge( + self.token_indices.view(1, 1, self.width), + valid_widths.view(request_count, 1, 1), + out=invalid_mask, + ) + if normalize_scores: + row_mean = self.row_mean[:request_count] + row_std = self.row_std[:request_count] + input_scores.masked_fill_(invalid_mask, 0.0) + torch.sum(input_scores, dim=2, keepdim=True, out=row_mean) + self.valid_scale[:request_count].view(request_count).copy_(valid_widths) + row_mean.div_(self.valid_scale[:request_count]) + torch.sub(input_scores, row_mean, out=input_scores) + input_scores.masked_fill_(invalid_mask, 0.0) + torch.linalg.vector_norm( + input_scores, + dim=2, + keepdim=True, + out=row_std, + ) + self.valid_scale[:request_count].sqrt_() + row_std.div_(self.valid_scale[:request_count]) + row_std.clamp_min_(1e-6) + torch.div(input_scores, row_std, out=input_scores) + input_scores.masked_fill_(invalid_mask, float("-inf")) + + torch.amax(input_scores, dim=1, out=self.combined[:request_count]) + + def select_requests( + self, + scores: torch.Tensor, + *, + normalize_scores: bool, + ) -> None: + """Select from request-major score output without repacking it.""" + request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 + if request_count <= 0 or request_count > self.max_requests: + raise ValueError("request count exceeds the cross-request selection capacity") + if scores.is_cuda and request_count != self.max_requests: + raise ValueError("CUDA selection requires the selector's fixed request count") + if ( + scores.numel() != request_count * self.rows * self.width + or int(scores.shape[-1]) != self.width + or scores.dtype != self.dtype + or scores.device != self.device + or not scores.is_contiguous() + ): + raise ValueError("cross-request scores do not match the selector geometry") + if scores.is_cuda: + if ( + self.prepared_scores is None + or scores is not self.prepared_scores.scores + or bool(normalize_scores) != self.prepared_scores.normalize_scores + ): + raise ValueError("CUDA union scores do not match their prepared fixed launcher") + self.select_prepared_requests() + return + self._select_input_scores( + scores.view(request_count, self.rows, self.width), + request_count, + normalize_scores=normalize_scores, + ) + + +class _BatchedPerHeadKeepSetSelector: + """Fixed ``[request, ...]`` selector for both per-head modes.""" + + def __init__( + self, + *, + eviction_mode: str, + dense_layers: Tuple[int, ...], + num_query_heads: int, + num_kv_heads: int, + width: int, + keep_count: int, + prompt_len: int, + dtype: torch.dtype, + device: torch.device, + max_requests: int, + ) -> None: + if eviction_mode not in ("per_head", "per_layer_perhead"): + raise ValueError(f"unsupported per-head eviction mode: {eviction_mode}") + if not dense_layers or min(num_query_heads, num_kv_heads, max_requests) <= 0: + raise ValueError("per-head selection requires positive layer, head, and request counts") + if num_query_heads % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + if width <= keep_count or keep_count <= 0: + raise ValueError("per-head selection requires width > keep_count > 0") + self.eviction_mode = eviction_mode + self.dense_layers = tuple(int(layer) for layer in dense_layers) + self.num_layers = len(self.dense_layers) + self.num_query_heads = int(num_query_heads) + self.num_kv_heads = int(num_kv_heads) + self.query_group_size = self.num_query_heads // self.num_kv_heads + self.rows = self.num_layers * self.num_query_heads + self.selection_rows = ( + self.num_kv_heads + if eviction_mode == "per_head" + else self.num_layers * self.num_kv_heads + ) + self.width = int(width) + self.keep_count = int(keep_count) + self.prompt_len = int(prompt_len) + self.total_keep = self.prompt_len + self.keep_count + self.dtype = dtype + self.device = _canonical_device(device) + self.max_requests = int(max_requests) + + score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) + grouped_shape = (self.max_requests, self.num_layers, self.num_kv_heads, self.width) + self.row_mean = torch.empty(score_shape[:-1] + (1,), dtype=dtype, device=self.device) + self.row_std = torch.empty_like(self.row_mean) + self.valid_widths = torch.full( + (self.max_requests,), self.width, dtype=torch.int32, device=self.device + ) + self.selection_scores = torch.empty( + (self.max_requests, self.selection_rows, self.width), + dtype=dtype, + device=self.device, + ) + if self.device.type == "cpu": + self.valid_scale = torch.empty( + (self.max_requests, 1, 1, 1), dtype=dtype, device=self.device + ) + self.token_indices = torch.arange(self.width, dtype=torch.long, device=self.device) + self.invalid_mask = torch.empty( + (self.max_requests, 1, 1, self.width), dtype=torch.bool, device=self.device + ) + self.grouped_scores = torch.empty(grouped_shape, dtype=dtype, device=self.device) + else: + self.valid_scale = None + self.token_indices = None + self.invalid_mask = None + self.grouped_scores = None + self.prepared_topk = ( + _PreparedCuteTopK( + self.max_requests * self.selection_rows, + self.width, + self.keep_count, + self.device, + ) + if self.device.type == "cuda" + else None + ) + self.row_seq_lens = torch.full( + (self.max_requests, self.selection_rows), + self.width, + dtype=torch.int32, + device=self.device, + ) + selection_shape = (self.max_requests, self.selection_rows, self.keep_count) + self.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=self.device) + self.keep = torch.empty( + (self.max_requests, self.selection_rows, self.total_keep), + dtype=torch.int32, + device=self.device, + ) + self.selection_scores_flat = self.selection_scores.view( + self.max_requests * self.selection_rows, self.width + ) + self.row_seq_lens_flat = self.row_seq_lens.view(-1) + self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) + self.keep_flat = self.keep.view(-1, self.total_keep) + self.prepared_finalizer = ( + _PreparedTopKFinalizer( + self.selection_scores_flat, + self.row_seq_lens_flat, + self.top_indices_i32_flat, + self.keep_flat, + self.keep_count, + self.prompt_len, + ) + if self.device.type == "cuda" + else None + ) + if self.prompt_len: + prompt = torch.arange(self.prompt_len, dtype=torch.int32, device=self.device) + self.keep[:, :, : self.prompt_len].copy_( + prompt.view(1, 1, -1).expand(self.max_requests, self.selection_rows, -1) + ) + + def _select_input_scores( + self, + input_scores: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, + ) -> None: + valid_widths = self.valid_widths[:request_count] + selection_scores = self.selection_scores[:request_count] + row_seq_lens = self.row_seq_lens[:request_count] + if input_scores.is_cuda: + from .triattention_kernels import prepare_per_head_scores + + prepare_per_head_scores( + input_scores, + valid_widths, + self.row_mean[:request_count], + self.row_std[:request_count], + selection_scores, + row_seq_lens, + request_count, + num_kv_heads=self.num_kv_heads, + per_layer=self.eviction_mode == "per_layer_perhead", + normalize_scores=normalize_scores, + ) + else: + self._select_input_scores_reference( + input_scores, + request_count, + normalize_scores=normalize_scores, + ) + + _deterministic_topk_indices_into( + self.selection_scores_flat, + self.row_seq_lens_flat, + self.top_indices_i32_flat, + self.keep_flat, + self.keep_count, + self.prompt_len, + self.prepared_topk, + self.prepared_finalizer, + ) + + def _select_input_scores_reference( + self, + input_scores: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, + ) -> None: + """Keep the explicit PyTorch implementation as the CPU oracle.""" + valid_widths = self.valid_widths[:request_count] + assert self.invalid_mask is not None + assert self.token_indices is not None + assert self.valid_scale is not None + assert self.grouped_scores is not None + invalid_mask = self.invalid_mask[:request_count] + torch.ge( + self.token_indices.view(1, 1, 1, self.width), + valid_widths.view(request_count, 1, 1, 1), + out=invalid_mask, + ) + if normalize_scores: + row_mean = self.row_mean[:request_count] + row_std = self.row_std[:request_count] + input_scores.masked_fill_(invalid_mask, 0.0) + torch.sum(input_scores, dim=3, keepdim=True, out=row_mean) + self.valid_scale[:request_count].view(request_count).copy_(valid_widths) + row_mean.div_(self.valid_scale[:request_count]) + torch.sub(input_scores, row_mean, out=input_scores) + input_scores.masked_fill_(invalid_mask, 0.0) + torch.linalg.vector_norm(input_scores, dim=3, keepdim=True, out=row_std) + self.valid_scale[:request_count].sqrt_() + row_std.div_(self.valid_scale[:request_count]) + row_std.clamp_min_(1e-6) + torch.div(input_scores, row_std, out=input_scores) + input_scores.masked_fill_(invalid_mask, float("-inf")) + + grouped_scores = self.grouped_scores[:request_count] + torch.amax( + input_scores.view( + request_count, + self.num_layers, + self.num_kv_heads, + self.query_group_size, + self.width, + ), + dim=3, + out=grouped_scores, + ) + if self.eviction_mode == "per_head": + torch.mean(grouped_scores, dim=1, out=self.selection_scores[:request_count]) + else: + self.selection_scores[:request_count].copy_( + grouped_scores.view(request_count, self.selection_rows, self.width) + ) + row_seq_lens = self.row_seq_lens[:request_count] + row_seq_lens.copy_(valid_widths.view(request_count, 1).expand(-1, self.selection_rows)) + + def select_requests( + self, + scores: torch.Tensor, + *, + normalize_scores: bool, + ) -> None: + request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 + if request_count <= 0 or request_count > self.max_requests: + raise ValueError("request count exceeds the per-head selection capacity") + if scores.is_cuda and request_count != self.max_requests: + raise ValueError("CUDA selection requires the selector's fixed request count") + expected_shape = ( + request_count, + self.num_layers, + self.num_query_heads, + self.width, + ) + if ( + tuple(scores.shape) != expected_shape + or scores.dtype != self.dtype + or scores.device != self.device + or not scores.is_contiguous() + ): + raise ValueError("per-head scores do not match the selector geometry") + self._select_input_scores(scores, request_count, normalize_scores=normalize_scores) + + +class _FixedScoreStreamMismatch(RuntimeError): + """Raised when a fixed score staging buffers are used from another CUDA stream.""" + + +class _FixedScoreStagingBuffers: + """Pool-bound fixed score metadata with one nonblocking page-table upload.""" + + @staticmethod + def _page_table_slot_layout( + page_representatives: List[int], + page_table_keys: List[object], + ) -> Tuple[Dict[int, int], int]: + if len(page_table_keys) != len(page_representatives): + raise ValueError("page-table keys must match the representative count") + use_pool_ids = all( + isinstance(key, tuple) + and len(key) == 2 + and key[0] == "pool" + and isinstance(key[1], int) + and key[1] >= 0 + for key in page_table_keys + ) + unique_slots = [] + key_to_slot = {} + representative_slots = {} + for representative, key in zip(page_representatives, page_table_keys): + slot = key_to_slot.get(key) + if slot is None: + slot = int(key[1]) if use_pool_ids else len(key_to_slot) + key_to_slot[key] = slot + unique_slots.append(slot) + representative_slots[representative] = slot + slot_count = max(unique_slots, default=-1) + 1 + return representative_slots, slot_count + + def __init__( + self, + layer_pools: List[torch.Tensor], + dense_groups: List[List[int]], + dense_layers: List[int], + page_representatives: List[int], + max_requests: int, + seq_len: int, + num_q_heads: int, + num_freqs: int, + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + freq_scale_sq: torch.Tensor, + offsets: torch.Tensor, + omega: torch.Tensor, + page_table_keys: Optional[List[object]] = None, + num_page_table_slots: Optional[int] = None, + prompt_len: int = 0, + page_table_token_capacity: Optional[int] = None, + draft_layer_pools: Optional[List[torch.Tensor]] = None, + draft_page_representatives: Optional[List[int]] = None, + draft_page_table_keys: Optional[List[object]] = None, + draft_num_page_table_slots: Optional[int] = None, + draft_page_table_token_capacity: Optional[int] = None, + ) -> None: + from .triattention_kernels import _FixedScoreGroup + + if not dense_groups or not dense_layers or not page_representatives or max_requests <= 0: + raise ValueError("fixed score metadata requires non-empty positive geometry") + grouped_layers = [layer for layers in dense_groups for layer in layers] + if ( + len(grouped_layers) != len(dense_layers) + or len(set(grouped_layers)) != len(grouped_layers) + or len(set(dense_layers)) != len(dense_layers) + or set(dense_layers) != set(grouped_layers) + ): + raise ValueError("dense layer order must cover every grouped layer exactly once") + self.device = _canonical_device(layer_pools[page_representatives[0]].device) + if self.device.type != "cuda": + raise ValueError("fixed score metadata is CUDA-only") + self.max_requests = max_requests + self.bucket_seq_len = seq_len + if page_table_token_capacity is None: + page_table_token_capacity = seq_len + if page_table_token_capacity < seq_len: + raise ValueError("page-table capacity cannot be smaller than the score bucket") + self.page_table_token_capacity = int(page_table_token_capacity) + if prompt_len < 0 or prompt_len > seq_len: + raise ValueError("fixed score metadata prompt length is outside its bucket") + self.prompt_len = prompt_len + q_real = q_real.to(device=self.device, dtype=torch.float32).contiguous() + q_imag = q_imag.to(device=self.device, dtype=torch.float32).contiguous() + mlr_coef = mlr_coef.to(device=self.device, dtype=torch.float32).contiguous() + freq_scale_sq = freq_scale_sq.to(device=self.device, dtype=torch.float32).contiguous() + offsets = offsets.to(device=self.device, dtype=torch.float32).contiguous() + omega = omega.to(device=self.device, dtype=torch.float32).contiguous() + if page_table_keys is None: + page_table_keys = list(range(len(page_representatives))) + self.representative_slots, minimum_page_table_slots = self._page_table_slot_layout( + page_representatives, page_table_keys + ) + if num_page_table_slots is None: + num_page_table_slots = minimum_page_table_slots + if num_page_table_slots < minimum_page_table_slots: + raise ValueError("page-table slot capacity does not cover every V2 pool") + tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) + if int(layer_pools[page_representatives[0]].shape[1]) != 2: + raise ValueError("fixed score metadata requires an interleaved K/V pool") + self.page_count = ( + self.page_table_token_capacity + tokens_per_block - 1 + ) // tokens_per_block + self.copy_block_count = (self.page_count + 3) // 4 * 4 + if any( + (self.page_table_token_capacity + int(layer_pools[layer].shape[3]) - 1) + // int(layer_pools[layer].shape[3]) + != self.page_count + for layer in page_representatives + ): + raise ValueError("fixed score metadata requires a uniform page count") + device_page_shape = ( + num_page_table_slots, + max_requests, + 2, + self.copy_block_count, + ) + self.request_metadata_host = torch.empty( + (2, max_requests), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._bulk_copy_idx_src = torch.arange( + max_requests, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._bulk_offsets_src = torch.empty( + device_page_shape, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.block_offsets_device = torch.empty( + device_page_shape, + dtype=torch.int32, + device=self.device, + ) + # Optional second block-offset staging plane for a co-compressed draft + # KV cache. The draft is never scored: these offsets feed only the + # draft compact launches. + self.draft_block_offsets_device: Optional[torch.Tensor] = None + self._draft_bulk_offsets_src: Optional[torch.Tensor] = None + self.draft_representative_slots: Dict[int, int] = {} + self.draft_copy_block_count = 0 + self.draft_page_table_token_capacity = 0 + if draft_layer_pools is not None: + if ( + not draft_page_representatives + or draft_page_table_keys is None + or draft_page_table_token_capacity is None + or draft_page_table_token_capacity <= 0 + ): + raise ValueError( + "draft page-table staging requires representatives, keys, and capacity" + ) + ( + self.draft_representative_slots, + minimum_draft_slots, + ) = self._page_table_slot_layout(draft_page_representatives, draft_page_table_keys) + if draft_num_page_table_slots is None: + draft_num_page_table_slots = minimum_draft_slots + if draft_num_page_table_slots < minimum_draft_slots: + raise ValueError("draft page-table slot capacity does not cover every V2 pool") + draft_tokens_per_block = int(draft_layer_pools[draft_page_representatives[0]].shape[3]) + if int(draft_layer_pools[draft_page_representatives[0]].shape[1]) != 2: + raise ValueError("draft page-table staging requires an interleaved K/V pool") + self.draft_page_table_token_capacity = int(draft_page_table_token_capacity) + draft_capacity = self.draft_page_table_token_capacity + draft_page_count = ( + draft_capacity + draft_tokens_per_block - 1 + ) // draft_tokens_per_block + if any( + (draft_capacity + int(draft_layer_pools[layer].shape[3]) - 1) + // int(draft_layer_pools[layer].shape[3]) + != draft_page_count + for layer in draft_page_representatives + ): + raise ValueError("draft page-table staging requires a uniform page count") + self.draft_copy_block_count = (draft_page_count + 3) // 4 * 4 + draft_page_shape = ( + draft_num_page_table_slots, + max_requests, + 2, + self.draft_copy_block_count, + ) + self._draft_bulk_offsets_src = torch.empty( + draft_page_shape, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.draft_block_offsets_device = torch.empty( + draft_page_shape, + dtype=torch.int32, + device=self.device, + ) + self.request_metadata_device = torch.empty( + (2, max_requests), dtype=torch.int32, device=self.device + ) + self.round_starts_device = self.request_metadata_device[0] + self.valid_seq_lens_device = self.request_metadata_device[1] + self.mean_cos = torch.empty( + (max_requests, num_freqs), dtype=torch.float32, device=self.device + ) + self.mean_sin = torch.empty_like(self.mean_cos) + self.offsets = offsets + self.omega = omega + # ONE fused group across ALL dense layers: segments carry their own + # layer base address and page-table slot, so distinct per-layer + # storages/block tables no longer force one launch per storage group. + _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} + _page_table_slots = [self.representative_slots[_rep_of[layer]] for layer in dense_layers] + self.fused_group = _FixedScoreGroup( + layer_pools, + dense_layers, + max_requests, + self.page_count, + seq_len, + num_q_heads, + self.block_offsets_device, + _page_table_slots, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + prompt_len=prompt_len, + ) + self.copy_done = torch.cuda.Event() + # First record publishes constructor allocations to the V2 copy stream; + # later records protect pinned metadata before the next cohort reuses it. + self.copy_done.record(torch.cuda.current_stream(self.device)) + self.bulk_copy_done = torch.cuda.Event() + self.bulk_consume_done = torch.cuda.Event() + self.copy_pending = False + self.page_tables_active = False + self.stream = None + self._phase_runner = None + self._phase_args: tuple = () + self._score_runner = None + self._score_args: tuple = () + + def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str) -> None: + """Compile and bind phase/score launches for this exact resource bucket.""" + from .triattention_kernels import _prepare_mean_phase_kernel, _tri_score_perhead_kernel + + if self._score_runner is not None: + raise RuntimeError("TriAttention score launcher is already bound") + if score_aggregation not in ("mean", "max"): + raise ValueError(f"unsupported score aggregation: {score_aggregation}") + group = self.fused_group + if ( + valid_widths.shape != (self.max_requests,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != self.device + or not valid_widths.is_contiguous() + ): + raise ValueError("prepared score lengths do not match their exact bucket") + + frequency_block = 1 << (group.num_freqs - 1).bit_length() + phase_pointer_args = ( + self.round_starts_device, + self.offsets, + self.omega, + self.mean_cos, + self.mean_sin, + ) + phase_constants = ( + group.num_freqs, + int(self.offsets.numel()), + frequency_block, + ) + score_pointer_args = ( + *group.pointer_prefix, + self.valid_seq_lens_device, + valid_widths, + self.round_starts_device, + *group.pointer_middle, + self.mean_cos.view(-1), + self.mean_sin.view(-1), + *group.pointer_tail, + group.output, + ) + score_geometry = ( + group.output_width, + group.num_layers, + *group.geometry_args, + ) + score_constants = ( + score_aggregation == "max", + group.prompt_len, + group.token_block, + frequency_block, + ) + self._phase_args = (*phase_pointer_args, *phase_constants) + self._score_args = (*score_pointer_args, *score_geometry, *score_constants) + phase_grid = (self.max_requests, 1, 1) + score_grid = ( + self.max_requests * group.num_layers, + group.max_ntblk, + group.num_kv_heads, + ) + with torch.cuda.device(self.device): + self.stream = torch.cuda.current_stream(self.device) + if score_aggregation == "mean": + compiled = _prepare_mean_phase_kernel.warmup( + *phase_pointer_args, + NUM_FREQS=phase_constants[0], + NUM_OFFSETS=phase_constants[1], + F_BLOCK=phase_constants[2], + num_warps=1, + grid=phase_grid, + ) + self._phase_runner = compiled[phase_grid] + compiled = _tri_score_perhead_kernel.warmup( + *score_pointer_args, + *score_geometry, + USE_MAX=score_constants[0], + TOKEN_START=score_constants[1], + T_BLOCK=score_constants[2], + F_BLOCK=score_constants[3], + grid=score_grid, + ) + self._score_runner = compiled[score_grid] + + def launch_prepared_score(self) -> torch.Tensor: + """Launch the phase and score runners bound to this exact bucket.""" + if self._score_runner is None or self.stream is None: + raise RuntimeError("TriAttention score launcher is not bound") + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.stream.device, + self.stream.cuda_stream, + ): + raise _FixedScoreStreamMismatch( + "TriAttention prepared score is bound to its staging CUDA stream" + ) + if self._phase_runner is not None: + self._phase_runner(*self._phase_args, stream=self.stream.cuda_stream) + self._score_runner(*self._score_args, stream=self.stream.cuda_stream) + return self.fused_group.output + + def stage( + self, + manager: KVCacheManagerV2, + request_ids: List[int], + round_starts: List[int], + seq_lens: Optional[List[int]] = None, + page_table_seq_lens: Optional[List[int]] = None, + draft_manager: Optional[KVCacheManagerV2] = None, + ) -> bool: + """Copy one eager eviction cohort into reusable device buffers.""" + request_count = len(request_ids) + if ( + request_count == 0 + or request_count > self.max_requests + or len(round_starts) != request_count + or any( + round_start != round_start + or round_start < 0 + or round_start > _INT32_MAX + or round_start != int(round_start) + for round_start in round_starts + ) + ): + return False + if (draft_manager is None) != (self.draft_block_offsets_device is None): + return False + stream = torch.cuda.current_stream(self.device) + if self.stream is None: + self.stream = stream + elif (stream.device, stream.cuda_stream) != ( + self.stream.device, + self.stream.cuda_stream, + ): + raise _FixedScoreStreamMismatch( + "TriAttention fixed score metadata is bound to its first CUDA stream" + ) + if self.page_tables_active: + raise RuntimeError("previous page-table cohort is still active") + if seq_lens is None: + seq_lens = [self.bucket_seq_len] * request_count + if len(seq_lens) != request_count or any( + seq_len <= 0 or seq_len > self.bucket_seq_len for seq_len in seq_lens + ): + return False + if page_table_seq_lens is None: + page_table_seq_lens = seq_lens + if len(page_table_seq_lens) != request_count or any( + page_seq_len < seq_len or page_seq_len > self.page_table_token_capacity + for seq_len, page_seq_len in zip(seq_lens, page_table_seq_lens) + ): + return False + if manager.enable_swa_scratch_reuse: + raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") + try: + request_metadata = torch.as_tensor((round_starts, seq_lens), dtype=torch.int32) + except (OverflowError, RuntimeError, TypeError, ValueError): + return False + if not self._stage_page_tables_bulk( + manager, + request_ids, + stream, + self._bulk_offsets_src, + self.block_offsets_device, + self.copy_block_count, + ): + return False + if draft_manager is not None: + if draft_manager.enable_swa_scratch_reuse: + raise RuntimeError( + "TriAttention does not support V2 SWA scratch page-table remapping" + ) + assert self._draft_bulk_offsets_src is not None + assert self.draft_block_offsets_device is not None + if not self._stage_page_tables_bulk( + draft_manager, + request_ids, + stream, + self._draft_bulk_offsets_src, + self.draft_block_offsets_device, + self.draft_copy_block_count, + ): + return False + self.request_metadata_host[:, :request_count].copy_(request_metadata) + try: + # Copy the fixed backing once. Only the first ``request_count`` + # columns are consumed by this cohort. + self.request_metadata_device.copy_(self.request_metadata_host, non_blocking=True) + finally: + # Guard the pinned metadata until its asynchronous copies complete. + # Page-table device-buffer reuse is guarded separately after compact. + self.copy_done.record(stream) + self.copy_pending = True + self.page_tables_active = True + return True + + def _stage_page_tables_bulk( + self, + manager: KVCacheManagerV2, + request_ids: List[int], + current_stream: torch.cuda.Stream, + source: torch.Tensor, + destination: torch.Tensor, + copy_block_count: int, + ) -> bool: + """Copy one request group's V2 block offsets before live compaction. + + Uses the V2 block-offset kernel with an immutable pinned snapshot of the + selected host-table rows. The snapshot is required because + this method enqueues asynchronous host-memory reads; TriAttention later + resizes the same cache, which mutates the manager's table in place. + The IndexMapper synchronously resolves request slots and gathers only + their beam-0 K block offsets, decoupling both live inputs before the + native asynchronous copy consumes the snapshot with identity indices. + ``dst[pool, r, 0(K), :]`` holds ``base_page * index_scales``. Score and + compact decode that K plane inline, avoiding any conversion kernel. + """ + if not request_ids or len(request_ids) > self.max_requests: + return False + + host_table = manager.host_kv_cache_block_offsets + num_pools, _, kv_planes, max_blocks = host_table.shape + if ( + host_table.dtype != torch.int32 + or kv_planes != 2 + or copy_block_count > max_blocks + or int(manager.kv_factor) != 2 + or num_pools != destination.shape[0] + ): + return False + request_count = len(request_ids) + submitted = False + try: + if self.copy_pending and not self.copy_done.query(): + self.copy_done.synchronize() + # The native device copy reads only K and derives V with kv_offset. + manager.index_mapper.gather_k_block_offsets( + host_table, + source, + request_ids, + copy_block_count, + ) + manager._stream.wait_event(self.copy_done) + copy_batch_block_offsets_to_device( + source, + destination, + self._bulk_copy_idx_src[:request_count], + manager.index_scales, + manager.kv_offset, + manager._stream.cuda_stream, + ) + submitted = True + self.bulk_copy_done.record(manager._stream) + current_stream.wait_event(self.bulk_copy_done) + except (AttributeError, IndexError, KeyError, RuntimeError, TypeError, ValueError) as exc: + if submitted: + raise RuntimeError( + "TriAttention bulk page-table copy failed after GPU submission" + ) from exc + logger.warning(f"TriAttention bulk page-table staging failed: {exc}") + return False + return True + + def mark_page_tables_consumed(self, *manager_streams: torch.cuda.Stream) -> None: + """Order V2 page-table reuse and resize after this cohort's compact. + + Every passed manager stream (target, and the draft when co-compressed) + waits on one event recorded after the compact launches, so neither + cache can free or reallocate pages this cohort is still reading. + """ + if not self.page_tables_active: + raise RuntimeError("TriAttention page tables were not staged") + self.bulk_consume_done.record(torch.cuda.current_stream(self.device)) + for manager_stream in manager_streams: + manager_stream.wait_event(self.bulk_consume_done) + self.page_tables_active = False + + +@dataclass(frozen=True, kw_only=True, slots=True) +class _PreparedEviction: + """Request metadata validated before eager score, select, and compact.""" + + request: "LlmRequest" + request_id: int + seq_len: int + round_start: int + expected_keep_count: int + protected_tail: int + + +@dataclass(kw_only=True, slots=True) +class _RequestCompressionState: + """Mutable compression state owned by one live request.""" + + generation_steps: int = 0 + evicted_tokens: int = 0 + confirmed_kv_length: Optional[int] = None + + +@dataclass(kw_only=True, slots=True) +class _PreparedGenerationBatch: + """Target growth reserved by the most recently prepared generation batch.""" + + batch: "ScheduledRequests" + growth_by_request: Dict[int, int] + + +@dataclass(kw_only=True, slots=True) +class _EvictionBucketResources: + """Reusable eager score and selection buffers for one runtime shape.""" + + score_staging: _FixedScoreStagingBuffers + keep_set_selector: Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector] + + +class TriAttention(BaseKVCacheCompressionManager): + """Periodic physical KV eviction driven by trigonometric importance scoring. + + Overrides ``on_generation_step_end``: every ``beta`` confirmed generation tokens it + reads the cached keys through the ``KVCacheManagerV2``, scores each token + with offline-calibrated stats, and physically evicts the tokens below the + keep set. Full-attention layers are scored; kernel-masked SWA layers preserve + their latest window in the same compacted prefix. Every layer ends with the + same request-wide cached length. + """ + + adjusts_generation_kv_length = True + physically_evicts_cached_tokens = True + + def __init__( + self, + kv_cache_manager: KVCacheManagerV2, + top_B: int, + draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, + beta: int = 128, + model_path: Optional[str] = None, + calibration_path: Optional[str] = None, + offset_max_length: int = 65536, + score_aggregation: str = "mean", + eviction_mode: str = "union", + normalize_scores: bool = True, + pin_prefill: bool = True, + count_prompt_tokens: bool = False, + ): + super().__init__(kv_cache_manager, draft_kv_cache_manager) + self.top_B = top_B + self.beta = beta + if self.top_B <= 0 or self.beta <= 0: + raise ValueError("TriAttention top_B and beta must both be positive") + # Which token set each eviction round keeps (all reproduce the upstream + # selection: z-normalize scores, pin the prompt tokens, no recency window): + # union -- union of every KV head's top-B, re-ranked by the + # per-token max score. Default; matches the official + # base setting (per-head and per-layer-per-head + # pruning both off). + # per_head -- each KV head keeps its own set, shared across + # layers (mean of per-layer max). + # per_layer_perhead -- each (layer, KV head) keeps its own set, fully + # independent per layer. + self.eviction_mode = eviction_mode + if self.eviction_mode not in ("union", "per_head", "per_layer_perhead"): + raise ValueError( + f"Unknown eviction_mode {self.eviction_mode!r}; expected one of " + "'union', 'per_head', 'per_layer_perhead'" + ) + self.normalize_scores = bool(normalize_scores) + self.pin_prefill = bool(pin_prefill) + # cpt=False (default): budget counts DECODE tokens only (pinned prompt is + # extra). cpt=True: budget INCLUDES the pinned prompt. + self.count_prompt_tokens = bool(count_prompt_tokens) + if not self.pin_prefill or self.count_prompt_tokens: + raise ValueError( + "TriAttention physical KV reclaim requires pin_prefill=True and " + "count_prompt_tokens=False so finalized prompt KV is preserved" + ) + # All physical moves use the C++ V2 compaction operation. + # No other compaction path exists. + self.score_aggregation = score_aggregation + # Calibration is the OFFICIAL TriAttention .pt (passed via + # calibration_path), resolved + converted on the first request + # (on_request_init). TRT-LLM does NOT compute calibration; model_path is + # used for RoPE tables and local layer_types/sliding_window metadata. + self.model_path = model_path + if self.model_path is None: + raise ValueError( + "TriAttention requires model_path so kernel-masked " + "sliding-attention layers can be classified safely" + ) + self.calibration_path = calibration_path + self.calibration: Optional[Dict[str, torch.Tensor]] = None + self._calibrated = False + # Calibration-derived dims + stats, filled in on_request_init. + self._H: Optional[int] = None + self._F: Optional[int] = None + self._freq_scale_sq: Optional[torch.Tensor] = None + + # Geometric integration offsets (built lazily on first eviction so the + # device matches the cache pool). + self._offset_max_length = offset_max_length + self._offsets: Optional[torch.Tensor] = None + + # Request presence records successful initialization. The record also + # owns the counters and physical length cleared at request finish. + self._request_states: Dict[int, _RequestCompressionState] = {} + # The overlap executor prepares B(n) before finalizing B(n-1). Keep the + # exact fixed-linear generation width for that currently in-flight batch; + # the final hook treats those slots as an opaque suffix. + self._prepared_generation_batch: Optional[_PreparedGenerationBatch] = None + # Eager score/selection buffers are built from the first live cohort and + # reused for subsequent evictions with the same runtime geometry. + self._eviction_buckets = OrderedDict() + self._batched_compactions = OrderedDict() + self._local_to_global_layers_cache: Optional[List[int]] = None + self._attention_layer_partition_cache: Optional[ + Tuple[List[int], List[int], Optional[int]] + ] = None + self._runtime_kv_layout_cache: Optional[_RuntimeKVLayout] = None + self._draft_runtime_kv_layout_cache: Optional[_RuntimeKVLayout] = None + + def on_request_init(self, request: "LlmRequest", **kwargs) -> None: + """Mark capacity-only decode and resolve calibration once. + + Loads the user-supplied OFFICIAL calibration .pt and converts it to our + runtime schema (see _resolve_calibration). TRT-LLM does not calibrate. + """ + request_id = request.py_request_id + if request_id not in self._request_states: + self._validate_v2_compatibility() + self._validate_request_capacity(request) + num_layers = self._num_layers_from_manager() + self._attention_layer_partition(num_layers) + self._request_states[request_id] = _RequestCompressionState() + self._ensure_calibrated() + + def _validate_request_capacity(self, request: "LlmRequest") -> None: + """Require enough target page-table capacity to reach first eviction.""" + manager = self.kv_cache_manager + # V2 mirrors the resolved speculative draft length (0 without spec). + speculative_overshoot = int(manager.max_draft_len) + first_eviction_decode_length = ( + self.top_B // self.beta + 1 + ) * self.beta + speculative_overshoot + decode_capacity = min(int(request.py_max_new_tokens), first_eviction_decode_length) + confirmed_capacity = int(request.py_prompt_len) + decode_capacity + protected_tail_capacity = self._configured_protected_tail_capacity() + required_capacity = confirmed_capacity + protected_tail_capacity + pool_confirmed_capacity = manager.get_num_available_tokens( + token_num_upper_bound=confirmed_capacity, + max_num_draft_tokens=int(manager._kv_reserve_draft_tokens) + 1, + ) + table_capacity = manager.max_blocks_per_seq * manager.tokens_per_block + if confirmed_capacity > pool_confirmed_capacity or required_capacity > table_capacity: + raise ValueError( + "TriAttention target KV capacity is too small to reach its first " + f"eviction: request requires {required_capacity} tokens " + f"(prompt={request.py_prompt_len}, budget={self.top_B}, " + f"beta={self.beta}, decode before eviction or completion=" + f"{decode_capacity}, speculative overshoot=" + f"{speculative_overshoot}, protected tail=" + f"{protected_tail_capacity}), " + f"but the V2 pool covers {pool_confirmed_capacity + protected_tail_capacity} " + f"tokens and its page table covers {table_capacity} tokens" + ) + draft_manager = self.draft_kv_cache_manager + if draft_manager is None: + return + draft_protected_tail = self._draft_protected_tail_capacity() + draft_required_capacity = confirmed_capacity + draft_protected_tail + draft_pool_capacity = draft_manager.get_num_available_tokens( + token_num_upper_bound=confirmed_capacity, + max_num_draft_tokens=int(draft_manager._kv_reserve_draft_tokens) + 1, + ) + draft_table_capacity = draft_manager.max_blocks_per_seq * draft_manager.tokens_per_block + if ( + confirmed_capacity > draft_pool_capacity + or draft_required_capacity > draft_table_capacity + ): + raise ValueError( + "TriAttention draft KV capacity is too small to reach the first " + f"co-compression: request requires {draft_required_capacity} " + f"tokens (prompt={request.py_prompt_len}, budget={self.top_B}, " + f"beta={self.beta}, decode before eviction or completion=" + f"{decode_capacity}, draft protected tail={draft_protected_tail}), " + f"but the draft V2 pool covers " + f"{draft_pool_capacity + draft_protected_tail} tokens and its " + f"page table covers {draft_table_capacity} tokens" + ) + + def _draft_protected_tail_capacity(self) -> int: + """Return the draft tail moved and re-reserved by every co-compression.""" + draft_manager = self.draft_kv_cache_manager + capacity = ( + int(draft_manager.num_extra_kv_tokens) + int(draft_manager._kv_reserve_draft_tokens) + 1 + ) + if capacity <= 0: + raise RuntimeError("draft KVCacheManagerV2 exposes an invalid protected-tail capacity") + return capacity + + def _ensure_calibrated(self) -> None: + """Resolve calibration once for the first request.""" + if self._calibrated: + return + self.calibration = self._resolve_calibration() + self._H = int(self.calibration["E_q"].shape[1]) + self._F = int(self.calibration["E_q"].shape[2]) + # Squared per-frequency RoPE scaling factor (required calibration key). + self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) + # Pre-split query stats + MLR coefficient for the Triton score kernel so + # it doesn't recompute (E_q_norm - |E_q|) per call. Shapes [L, H, F]. + _Eq = self.calibration["E_q"] + self._triattn_q_real = _Eq.real.to(torch.float32).contiguous() + self._triattn_q_imag = _Eq.imag.to(torch.float32).contiguous() + self._triattn_mlr_coef = ( + self.calibration["E_q_norm"].to(torch.float32) - _Eq.abs().to(torch.float32) + ).contiguous() + self._calibrated = True + + def _validate_v2_compatibility(self) -> None: + """Reject runtime modes outside the V2 physical-compaction contract.""" + manager = self.kv_cache_manager + if not isinstance(manager, KVCacheManagerV2): + raise ValueError("TriAttention physical eviction requires KVCacheManagerV2") + if manager.kv_factor != 2: + raise ValueError( + "TriAttention requires a standard key/value KV cache; " + "MLA/SELFKONLY caches are not supported" + ) + if manager.mapping.enable_attention_dp: + raise ValueError("TriAttention does not support attention DP") + if manager.is_disagg: + raise ValueError("TriAttention does not support disaggregated serving") + if manager.max_beam_width != 1: + raise ValueError("TriAttention requires beam-width-one decoding") + # Speculative feature gates (resolved draft length, linear drafting, + # mode whitelist) run in the factory, where spec_config lives. The + # draft cache itself is validated here whenever one is attached. + draft_manager = self.draft_kv_cache_manager + if draft_manager is not None: + if not draft_manager.is_draft: + raise ValueError( + "TriAttention speculative compatibility requires the actual " + "separate draft KV cache manager" + ) + if draft_manager.kv_factor != 2: + raise ValueError( + "TriAttention compresses the draft KV cache together with " + "the target, so the draft cache must be a standard " + "key/value cache" + ) + if self.eviction_mode != "union": + raise ValueError( + "TriAttention draft KV co-compression supports only " + "eviction_mode='union'; per-head keep sets are not defined " + "for draft layers, which are never scored" + ) + if any(window is not None for window in draft_manager.max_attention_window_vec) or any( + not isinstance(layer, AttentionLayerConfig) or layer.sliding_window_size is not None + for layer in draft_manager.kv_cache_manager_py_config.layers + ): + raise ValueError( + "TriAttention draft KV co-compression requires full-attention " + "draft V2 lifecycles" + ) + if any(window is not None for window in manager.max_attention_window_vec) or any( + not isinstance(layer, AttentionLayerConfig) or layer.sliding_window_size is not None + for layer in manager.kv_cache_manager_py_config.layers + ): + raise ValueError( + "TriAttention requires full-attention V2 lifecycles; native SWA, " + "VSWA, and SSM pools are not supported" + ) + + # The framework drives all request-lifecycle hooks. TriAttention resolves + # calibration on request init, evicts periodically at generation-step end, + # and removes per-request state at finish. It scores from offline + # calibration, not from live queries or attention scores, so it needs no + # per-layer attention hook: the whole eviction runs once per period in + # on_generation_step_end, which loops the layers and reads each layer's keys + # straight from the KV pool. + + def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: + """Compact after native KV-cache updates have finalized this iteration. + + The compression manager is ordered after KVCacheManagerV2, so capacity + already reflects the written token and any rewind. The overlap scheduler + may already have enqueued the next forward; CUDA stream ordering keeps + compaction after that reader. The resize happens only after compaction; + it detaches the compacted tail without blocking the host, while V2's + per-slot finish events prevent early page reuse. + """ + with nvtx_range_debug("triattention.generation_step_end", color="blue"): + self._periodic_evict(scheduled_batch) + + def prepare_resources(self, scheduled_batch: "ScheduledRequests") -> None: + """Snapshot fixed-linear target growth; mutation remains in final update.""" + super().prepare_resources(scheduled_batch) + generation_growth = {} + for request in scheduled_batch.generation_requests: + request_id = request.py_request_id + growth = 1 + max( + get_draft_token_length(request), + self.kv_cache_manager._kv_reserve_draft_tokens, + ) + generation_growth[request_id] = growth + self._prepared_generation_batch = _PreparedGenerationBatch( + batch=scheduled_batch, + growth_by_request=generation_growth, + ) + + def _inflight_generation_growth( + self, scheduled_batch: "ScheduledRequests", request_id: int + ) -> int: + """Return exact newer target allocation width under overlap scheduling.""" + prepared = self._prepared_generation_batch + if prepared is None or scheduled_batch is prepared.batch: + return 0 + return prepared.growth_by_request.get(request_id, 0) + + def _periodic_evict( + self, + scheduled_batch: "ScheduledRequests", + ) -> None: + """Count confirmed tokens; every ``beta`` tokens score the cache + and physically evict to the pinned prompt plus top-B decode tokens.""" + gen_requests = scheduled_batch.generation_requests + if not gen_requests: + return + active_requests = [] + for request in gen_requests: + if request.is_dummy or request.state in ( + LlmRequestState.GENERATION_COMPLETE, + LlmRequestState.CONTEXT_INIT, + ): + continue + kv_cache = self.kv_cache_manager.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + continue + if not kv_cache.is_active: + raise RuntimeError( + "TriAttention cannot finalize a suspended target KV cache; " + f"request {request.py_request_id} must be resumed before " + "the final update hook" + ) + if request.py_request_id not in self._request_states: + self.on_request_init(request) + active_requests.append(request) + if not active_requests or not self._calibrated: + return + mgr = self.kv_cache_manager + num_layers = self._num_layers_from_manager() + protected_tails: Dict[int, int] = {} + + # (1) bump per-request step counters; collect who evicts THIS step. + evict_now = [] + for request in active_requests: + rid = request.py_request_id + kv_cache = mgr.kv_cache_map.get(rid) + if kv_cache is None or not kv_cache.is_active: + continue + raw_capacity = int(kv_cache.capacity) + # One-engine speculative decoding keeps a fixed reserve E. Under + # overlap, B(n) is allocated/enqueued before finalizing B(n-1), so + # its exact scheduler growth Q is also opaque. Both spans are + # contiguous after the stable target prefix and move byte-for-byte. + protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( + scheduled_batch, rid + ) + seq_len = raw_capacity - protected_tail + if seq_len < 0 or protected_tail < 0: + raise RuntimeError( + f"Request {rid} has an inconsistent protected target tail: " + f"confirmed={seq_len}, capacity={raw_capacity}, " + f"protected_tail={protected_tail}" + ) + if seq_len < kv_cache.history_length: + raise RuntimeError( + f"Request {rid} KV length {seq_len} is below finalized " + f"history {kv_cache.history_length}" + ) + request_state = self._request_states[rid] + request_state.confirmed_kv_length = seq_len + protected_tails[rid] = protected_tail + previous_step = request_state.generation_steps + confirmed_delta = 1 + int(request.py_num_accepted_draft_tokens) + step = previous_step + confirmed_delta + request_state.generation_steps = step + if previous_step // self.beta < step // self.beta: + if seq_len > self._minimum_evictable_length(request, seq_len): + if self.draft_kv_cache_manager is not None: + draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(rid) + if draft_kv_cache is None or not draft_kv_cache.is_active: + raise RuntimeError( + "TriAttention cannot co-compress a missing or " + f"suspended draft KV cache; request {rid} must " + "be resumed before the final update hook" + ) + evict_now.append((request, rid)) + + # (2) Compact all affected dense and kernel-masked SWA layers, then release + # the unreachable tail directly through V2's public resize primitive. + if not evict_now: + return + protected_tail_lengths = {rid: protected_tails[rid] for _, rid in evict_now} + # Prompt and retained geometry define selection and destination layout. + # Requests with the same geometry execute eagerly in bounded chunks. + # Chunking limits staging memory. + eviction_groups = {} + for request, rid in evict_now: + seq_len = self._request_states[rid].confirmed_kv_length + if seq_len is None: + raise RuntimeError(f"Missing confirmed KV length for request {rid}") + prompt_len = min(int(request.py_prompt_len), seq_len) + keep_count = self._minimum_evictable_length(request, seq_len) + key = (prompt_len, keep_count) + eviction_groups.setdefault(key, []).append((request, rid)) + + for group in eviction_groups.values(): + for begin in range(0, len(group), _EAGER_REQUEST_CHUNK_SIZE): + chunk = group[begin : begin + _EAGER_REQUEST_CHUNK_SIZE] + chunk_tails = {rid: protected_tail_lengths[rid] for _, rid in chunk} + with nvtx_range_debug("triattention.evict_request_group", color="purple"): + capacity_targets = self._evict_requests( + chunk, + num_layers, + protected_tail_lengths=chunk_tails, + ) + self._resize_compacted_requests(capacity_targets, protected_tails) + + def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: + if not capacity_targets: + return + mgr = self.kv_cache_manager + draft_manager = self.draft_kv_cache_manager + with nvtx_range("triattention.resize", color="red"): + with nvtx_range_debug("triattention.v2_resize", color="red"): + for rid, target_capacity in capacity_targets: + kv_cache = mgr.kv_cache_map.get(rid) + if kv_cache is None or not kv_cache.is_active: + continue + if target_capacity > kv_cache.capacity: + raise RuntimeError( + f"Request {rid} compacted capacity {target_capacity} exceeds " + f"current capacity {kv_cache.capacity}" + ) + protected_tail = protected_tails[rid] + resized_capacity = target_capacity + protected_tail + if not kv_cache.resize(resized_capacity, None): + raise RuntimeError( + f"Failed to resize compacted KV cache for request {rid} " + f"to {resized_capacity} tokens" + ) + if draft_manager is not None: + # The draft cache was compacted with the same kept token + # set, so it shrinks to the same retained length plus its + # own protected tail. + draft_protected_tail = self._draft_protected_tail_capacity() + for rid, target_capacity in capacity_targets: + draft_kv_cache = draft_manager.kv_cache_map.get(rid) + if draft_kv_cache is None or not draft_kv_cache.is_active: + continue + draft_capacity = target_capacity + draft_protected_tail + if not draft_kv_cache.resize(draft_capacity, None): + raise RuntimeError( + "Failed to resize co-compressed draft KV cache " + f"for request {rid} to {draft_capacity} tokens" + ) + + def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: + """Return the largest cache length for which selection is an identity. + + With a decode-only budget, pinned prompt tokens do not consume ``top_B``. + Selection therefore keeps every token until the cache exceeds + ``prompt_len + top_B``. + """ + if self.pin_prefill and not self.count_prompt_tokens: + prompt_len = min(int(request.py_prompt_len), seq_len) + return prompt_len + self.top_B + return self.top_B + + def _local_score_calibration( + self, + num_layers: int, + global_layers: List[int], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return calibration tensors indexed in this PP rank's local layer order.""" + if global_layers and max(global_layers) >= self._triattn_q_real.shape[0]: + raise ValueError( + f"TriAttention calibration has {self._triattn_q_real.shape[0]} layers, " + f"but this PP rank references global layer {max(global_layers)}" + ) + if global_layers == list(range(global_layers[0], global_layers[0] + num_layers)): + layer_slice = slice(global_layers[0], global_layers[0] + num_layers) + return ( + self._triattn_q_real[layer_slice], + self._triattn_q_imag[layer_slice], + self._triattn_mlr_coef[layer_slice], + ) + layer_ids = torch.as_tensor( + global_layers, + device=self._triattn_q_real.device, + dtype=torch.long, + ) + return ( + self._triattn_q_real.index_select(0, layer_ids), + self._triattn_q_imag.index_select(0, layer_ids), + self._triattn_mlr_coef.index_select(0, layer_ids), + ) + + def _configured_protected_tail_capacity(self) -> int: + """Return the largest target tail reserved by the native V2 lifecycle.""" + capacity = ( + int(self.kv_cache_manager.num_extra_kv_tokens) + + int(self.kv_cache_manager._kv_reserve_draft_tokens) + + 1 + ) + if capacity <= 0: + raise RuntimeError("KVCacheManagerV2 exposes an invalid protected-tail capacity") + return capacity + + @staticmethod + def _build_cross_request_keep_set_selector( + plan: _CrossRequestSelectionPlan, + *, + input_scores: Optional[torch.Tensor] = None, + normalize_scores: bool = True, + ) -> Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector]: + """Allocate one fixed ``[request, ...]`` keep-set selector.""" + if plan.eviction_mode == "union": + return _BatchedUnionKeepSetSelector( + plan.rows, + plan.width, + plan.keep_count, + plan.prompt_len, + dtype=plan.dtype, + device=plan.device, + max_requests=plan.max_requests, + dense_layers=plan.dense_layers, + num_query_heads=plan.num_query_heads, + num_kv_heads=plan.num_kv_heads, + input_scores=input_scores, + normalize_scores=normalize_scores, + ) + return _BatchedPerHeadKeepSetSelector( + eviction_mode=plan.eviction_mode, + dense_layers=plan.dense_layers, + num_query_heads=plan.num_query_heads, + num_kv_heads=plan.num_kv_heads, + width=plan.width, + keep_count=plan.keep_count, + prompt_len=plan.prompt_len, + dtype=plan.dtype, + device=plan.device, + max_requests=plan.max_requests, + ) + + def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: + """Drop this request's per-request length and eviction state.""" + request_id = request.py_request_id + self._request_states.pop(request_id, None) + prepared = self._prepared_generation_batch + if prepared is not None: + prepared.growth_by_request.pop(request_id, None) + if not self._request_states: + self._eviction_buckets.clear() + self._batched_compactions.clear() + + # ================================================================== # + # Helpers (eviction / scoring / V2 cache access / calibration) # + # ================================================================== # + + # --- Upstream-faithful eviction modes (per_head / per_layer_perhead / union) --- + # + # These reproduce github.com/WeianMao/triattention's selection: scores are NOT + # averaged over heads (each KV head keeps its own token set), they are + # z-normalized per head over the decode region, the prompt (prefill) tokens are + # pinned, and there is no recency window. The kept COUNT stays uniform (= top_B) + # so paged attention + # and the num_cached bookkeeping are unchanged; only the kept SET differs per + # head. Kept K keeps its original RoPE rotation (scored post-RoPE), so a head + # holding a different token set still scores the correct relative distance + # and no per-head position tracking is needed. + + def _local_to_global_layers(self, num_layers: int) -> List[int]: + """Return V2's global layer id for every local TriAttention layer slot.""" + cached = self._local_to_global_layers_cache + if cached is not None: + if len(cached) != num_layers: + raise ValueError( + f"TriAttention layer count changed from {len(cached)} to {num_layers}" + ) + return cached + + global_layers = [int(layer) for layer in self.kv_cache_manager.pp_layers] + if len(global_layers) != num_layers: + raise ValueError( + f"KVCacheManagerV2 exposes {len(global_layers)} PP layers, " + f"but TriAttention received {num_layers} local layers" + ) + self._local_to_global_layers_cache = global_layers + return global_layers + + @staticmethod + def _has_sliding_window_signal(config: Dict[str, object]) -> bool: + """Return whether config metadata hints at sliding attention.""" + use_sliding_window = config.get("use_sliding_window") + if isinstance(use_sliding_window, bool): + return use_sliding_window + for field in ( + "sliding_window", + "sliding_window_size", + "sliding_window_pattern", + "max_window_layers", + ): + value = config.get(field) + if isinstance(value, bool): + if value: + return True + elif isinstance(value, (int, float)): + if value > 0: + return True + elif value: + return True + return False + + def _attention_layer_partition( + self, num_layers: int + ) -> Tuple[List[int], List[int], Optional[int]]: + """Return dense layers, kernel-masked SWA layers, and the SWA window. + + TriAttention initialization has already rejected real V2 windowed + lifecycles. A sliding layer found here is therefore stored at full length + and applies its window only in the attention kernel. + """ + cached = self._attention_layer_partition_cache + if cached is not None: + return cached + + model_path = self.model_path + if model_path is None: + raise ValueError("TriAttention requires model_path") + + try: + from transformers import AutoConfig + + config = AutoConfig.from_pretrained( + model_path, trust_remote_code=True, local_files_only=True + ) + except Exception as exc: + raise ValueError( + f"TriAttention could not load the local model config from {model_path!r}" + ) from exc + config_values = config.get_text_config().to_dict() + layer_types = config_values.get("layer_types") + if not layer_types: + if self._has_sliding_window_signal(config_values): + raise ValueError( + "Model config exposes sliding-window metadata but no layer_types; " + "TriAttention cannot classify kernel-masked SWA layers safely" + ) + result = (list(range(num_layers)), [], None) + self._attention_layer_partition_cache = result + return result + global_layers = self._local_to_global_layers(num_layers) + 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] + 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.top_B < raw_window: + raise ValueError( + f"TriAttention decode budget top_B={self.top_B} must be at least " + f"the kernel-masked SWA window size {raw_window}" + ) + window_size = raw_window + result = (dense_layers, swa_layers, window_size) + self._attention_layer_partition_cache = result + return result + + def _runtime_kv_layout(self, num_layers: int) -> _RuntimeKVLayout: + """Return stable V2 pool views and layer groups for eager eviction. + + KVCacheManagerV2 keeps GPU virtual addresses and layer geometry stable, + while opt-in pool rebalance can change the page dimension. Cache all + layer views, then query the live page count for one representative per + physical pool before reuse. This avoids rebuilding TensorWrapper views + on every eviction while retaining the same fail-closed rebalance check. + """ + cached = self._runtime_kv_layout_cache + manager = self.kv_cache_manager + if cached is not None: + if cached.num_layers != num_layers: + raise ValueError( + f"TriAttention layer count changed from {cached.num_layers} to {num_layers}" + ) + if cached.manager is not manager: + raise RuntimeError("TriAttention target KV cache manager changed at runtime") + current_page_counts = self._pool_page_counts( + manager, + cached.global_layers, + cached.pool_representatives, + ) + if current_page_counts != cached.pool_page_counts: + raise RuntimeError( + "TriAttention V2 pool layout changed after the layout was built; " + "KV pool rebalance is not supported" + ) + return cached + + global_layers = self._local_to_global_layers(num_layers) + dense_layers, swa_layers, swa_window = self._attention_layer_partition(num_layers) + if not dense_layers: + raise ValueError("TriAttention requires at least one full-attention layer") + layout = self._build_runtime_kv_layout( + manager, + global_layers, + dense_layers=dense_layers, + swa_layers=swa_layers, + swa_window=swa_window, + dense_storage_groups=self._dense_layer_pool_groups(dense_layers, global_layers), + what="", + ) + self._runtime_kv_layout_cache = layout + return layout + + def _build_runtime_kv_layout( + self, + manager: KVCacheManagerV2, + global_layers: List[int], + *, + dense_layers: List[int], + swa_layers: List[int], + swa_window: Optional[int], + dense_storage_groups: Optional[Dict[object, List[int]]], + what: str, + ) -> _RuntimeKVLayout: + """Build the manager-lifetime layer and pool views one eviction reads. + + ``dense_storage_groups`` restricts the compaction groups to the dense + layers (target cache); None groups every layer (draft cache, which has + no SWA partition). ``what`` prefixes error messages ("" or "draft "). + """ + num_layers = len(global_layers) + maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] + if any(pool is None for pool in maybe_layer_pools): + missing = [ + layer for layer, pool in zip(global_layers, maybe_layer_pools) if pool is None + ] + raise RuntimeError(f"Missing {what}KV pools for attention layers {missing}") + layer_pools = [pool for pool in maybe_layer_pools if pool is not None] + all_layers = list(range(num_layers)) + layer_pool_keys = tuple( + self._page_table_pool_keys(all_layers, global_layers, manager=manager) + ) + all_storage_groups: Dict[object, List[int]] = {} + for layer, pool_key in zip(all_layers, layer_pool_keys): + all_storage_groups.setdefault(pool_key, []).append(layer) + storage_groups = ( + dense_storage_groups if dense_storage_groups is not None else all_storage_groups + ) + layer_group_representative = { + layer: layers[0] for layers in storage_groups.values() for layer in layers + } + pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) + return _RuntimeKVLayout( + manager=manager, + num_layers=num_layers, + global_layers=global_layers, + layer_pools=layer_pools, + dense_layers=dense_layers, + swa_layers=swa_layers, + swa_window=swa_window, + storage_groups=storage_groups, + layer_group_representative=layer_group_representative, + layer_pool_keys=layer_pool_keys, + pool_representatives=pool_representatives, + pool_page_counts=tuple( + int(layer_pools[layer].shape[0]) for layer in pool_representatives + ), + pool_view_fingerprint=self._pool_view_fingerprint( + [layer_pools[layer] for layer in pool_representatives] + ), + ) + + def _draft_runtime_kv_layout(self) -> _RuntimeKVLayout: + """Return stable draft V2 pool views, mirroring ``_runtime_kv_layout``. + + The draft cache is compacted with the target's kept token set, so its + layout has no scoring role: every draft layer is dense and there is no + SWA partition. + """ + manager = self.draft_kv_cache_manager + if manager is None: + raise RuntimeError("TriAttention has no draft KV cache manager to lay out") + cached = self._draft_runtime_kv_layout_cache + if cached is not None: + if cached.manager is not manager: + raise RuntimeError("TriAttention draft KV cache manager changed at runtime") + current_page_counts = self._pool_page_counts( + manager, + cached.global_layers, + cached.pool_representatives, + ) + if current_page_counts != cached.pool_page_counts: + raise RuntimeError( + "TriAttention draft V2 pool layout changed after the layout " + "was built; KV pool rebalance is not supported" + ) + return cached + + global_layers = [int(layer) for layer in manager.pp_layers] + if not global_layers: + raise RuntimeError("TriAttention draft KV cache manager exposes no layers") + layout = self._build_runtime_kv_layout( + manager, + global_layers, + dense_layers=list(range(len(global_layers))), + swa_layers=[], + swa_window=None, + dense_storage_groups=None, + what="draft ", + ) + self._draft_runtime_kv_layout_cache = layout + return layout + + @staticmethod + def _pool_page_counts( + manager: KVCacheManagerV2, + global_layers: Sequence[int], + pool_representatives: Sequence[int], + ) -> Tuple[int, ...]: + """Read the only pool-view dimension that V2 rebalance can change.""" + return tuple( + int( + manager.impl.get_page_index_upper_bound( + manager.layer_offsets[global_layers[layer]], + Role.KEY, + ) + ) + // int(manager.kv_factor) + for layer in pool_representatives + ) + + @staticmethod + def _pool_view_fingerprint(pools: List[torch.Tensor]) -> Tuple[tuple, ...]: + """Identify the V2 pool properties consumed by score and compact kernels.""" + return tuple( + ( + pool.data_ptr(), + tuple(int(value) for value in pool.shape), + tuple(int(value) for value in pool.stride()), + pool.dtype, + pool.device, + ) + for pool in pools + ) + + def _eager_resources_for( + self, + layout: _RuntimeKVLayout, + prepared: Sequence[_PreparedEviction], + ) -> _EvictionBucketResources: + """Build or reuse eager score and selection buffers for one cohort.""" + if not prepared: + raise ValueError("TriAttention eviction requires at least one request") + prompt_lens = {min(int(item.request.py_prompt_len), item.seq_len) for item in prepared} + if len(prompt_lens) != 1: + raise ValueError("TriAttention batches require one common prompt length") + prompt_len = next(iter(prompt_lens)) + seq_len = max(item.seq_len for item in prepared) + page_table_token_capacity = max(item.seq_len + item.protected_tail for item in prepared) + request_count = len(prepared) + dense_groups = list(layout.storage_groups.values()) + representatives = [group[0] for group in dense_groups] + representatives.extend(layer for layer in layout.swa_layers if layer not in representatives) + draft_key = None + draft_kwargs = {} + if self.draft_kv_cache_manager is not None: + draft_layout = self._draft_runtime_kv_layout() + draft_tail_capacity = self._draft_protected_tail_capacity() + draft_representatives = list(draft_layout.pool_representatives) + draft_kwargs = dict( + draft_layer_pools=draft_layout.layer_pools, + draft_page_representatives=draft_representatives, + draft_page_table_keys=[ + draft_layout.layer_pool_keys[layer] for layer in draft_representatives + ], + draft_num_page_table_slots=self.draft_kv_cache_manager.num_pools, + draft_page_table_token_capacity=seq_len + draft_tail_capacity, + ) + draft_key = ( + draft_tail_capacity, + draft_layout.pool_page_counts, + draft_layout.pool_view_fingerprint, + ) + key = ( + "triattention.eager.v1", + self.eviction_mode, + request_count, + seq_len, + prompt_len, + page_table_token_capacity, + self.top_B, + tuple(layout.dense_layers), + layout.pool_view_fingerprint, + draft_key, + ) + resources = self._eviction_buckets.get(key) + if resources is not None: + self._eviction_buckets.move_to_end(key) + return resources + + first_pool = layout.layer_pools[layout.dense_layers[0]] + if self._offsets is None: + self._offsets = _build_geometric_offsets(self._offset_max_length, first_pool.device) + q_real, q_imag, mlr_coef = self._local_score_calibration( + layout.num_layers, layout.global_layers + ) + score_staging = _FixedScoreStagingBuffers( + layout.layer_pools, + dense_groups=dense_groups, + dense_layers=layout.dense_layers, + page_representatives=representatives, + max_requests=request_count, + seq_len=seq_len, + num_q_heads=int(self._H), + num_freqs=int(self._F), + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=self._freq_scale_sq, + offsets=self._offsets, + omega=self.calibration["omega"], + page_table_keys=self._page_table_pool_keys(representatives, layout.global_layers), + num_page_table_slots=layout.manager.num_pools, + prompt_len=prompt_len, + page_table_token_capacity=page_table_token_capacity, + **draft_kwargs, + ) + keep_set_selector = self._build_cross_request_keep_set_selector( + _CrossRequestSelectionPlan( + eviction_mode=self.eviction_mode, + dense_layers=tuple(layout.dense_layers), + num_query_heads=int(self._H), + num_kv_heads=int(first_pool.shape[2]), + rows=len(layout.dense_layers) * int(self._H), + width=seq_len - prompt_len, + keep_count=self.top_B, + prompt_len=prompt_len, + dtype=torch.float32, + device=first_pool.device, + max_requests=request_count, + ), + input_scores=score_staging.fused_group.output.view( + request_count, + len(layout.dense_layers) * int(self._H), + seq_len - prompt_len, + ), + normalize_scores=self.normalize_scores, + ) + score_staging.bind_score_launcher( + keep_set_selector.valid_widths, + self.score_aggregation, + ) + resources = _EvictionBucketResources( + score_staging=score_staging, + keep_set_selector=keep_set_selector, + ) + self._eviction_buckets[key] = resources + while len(self._eviction_buckets) > _EAGER_RESOURCE_CACHE_LIMIT: + _, stale = self._eviction_buckets.popitem(last=False) + stale_ids = (id(stale.score_staging), id(stale.keep_set_selector)) + for compaction_key in tuple(self._batched_compactions): + if compaction_key[:2] == stale_ids: + del self._batched_compactions[compaction_key] + return resources + + def _batched_compaction_for( + self, + *, + layout: _RuntimeKVLayout, + prepared: Sequence[_PreparedEviction], + score_staging: _FixedScoreStagingBuffers, + keep_set_selector: Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector], + ): + """Build or reuse the eager C++ compaction launches for one cohort.""" + from .compaction import BatchedKVCacheCompaction + + protected_tail_lengths = tuple(item.protected_tail for item in prepared) + draft_tail_lengths: Optional[Tuple[int, ...]] = None + draft_kwargs = {} + if self.draft_kv_cache_manager is not None: + draft_layout = self._draft_runtime_kv_layout() + draft_tail_lengths = (self._draft_protected_tail_capacity(),) * len(prepared) + draft_kwargs = dict( + draft_layer_pools=draft_layout.layer_pools, + draft_layers=draft_layout.dense_layers, + draft_layer_group_representative=draft_layout.layer_group_representative, + draft_layer_pool_keys=list(draft_layout.layer_pool_keys), + draft_protected_tail_lengths=list(draft_tail_lengths), + draft_kv_block_offsets=score_staging.draft_block_offsets_device, + draft_page_table_slots=score_staging.draft_representative_slots, + ) + key = ( + id(score_staging), + id(keep_set_selector), + protected_tail_lengths, + draft_tail_lengths, + ) + batched_compaction = self._batched_compactions.get(key) + if batched_compaction is not None: + self._batched_compactions.move_to_end(key) + else: + batched_compaction = BatchedKVCacheCompaction( + eviction_mode=self.eviction_mode, + layer_pools=layout.layer_pools, + dense_layers=layout.dense_layers, + swa_layers=layout.swa_layers, + layer_group_representative=layout.layer_group_representative, + layer_pool_keys=list(layout.layer_pool_keys), + kept_token_ordinals=keep_set_selector.keep[: len(prepared)], + valid_sequence_lengths=score_staging.valid_seq_lens_device[: len(prepared)], + kv_block_offsets=score_staging.block_offsets_device, + page_table_slots=score_staging.representative_slots, + request_count=len(prepared), + prompt_len=keep_set_selector.prompt_len, + decode_keep_count=self.top_B, + swa_window=layout.swa_window, + protected_tail_lengths=list(protected_tail_lengths), + **draft_kwargs, + ) + self._batched_compactions[key] = batched_compaction + while len(self._batched_compactions) > _EAGER_COMPACTION_CACHE_LIMIT: + self._batched_compactions.popitem(last=False) + return batched_compaction + + def _page_table_pool_keys( + self, + representatives: List[int], + global_layers: List[int], + manager: Optional[KVCacheManagerV2] = None, + ) -> List[object]: + """Return stable V2-pool keys for the representative layers.""" + if manager is None: + manager = self.kv_cache_manager + layer_offsets = manager.layer_offsets + layer_to_pool = manager.layer_to_pool_mapping_dict + try: + return [ + ("pool", int(layer_to_pool[layer_offsets[global_layers[layer]]])) + for layer in representatives + ] + except (IndexError, KeyError, TypeError, ValueError) as exc: + raise RuntimeError("KVCacheManagerV2 exposes an invalid layer-to-pool mapping") from exc + + def _dense_layer_pool_groups( + self, + dense_layers: List[int], + global_layers: List[int], + ) -> Dict[object, List[int]]: + """Group layers that use the same V2 page table.""" + groups: Dict[object, List[int]] = {} + for layer, pool_key in zip( + dense_layers, + self._page_table_pool_keys(dense_layers, global_layers), + ): + groups.setdefault(pool_key, []).append(layer) + return groups + + def _attach_page_ids( + self, + prepared: Sequence[_PreparedEviction], + staging: _FixedScoreStagingBuffers, + ) -> None: + try: + staged = staging.stage( + self.kv_cache_manager, + [item.request_id for item in prepared], + [item.round_start for item in prepared], + [item.seq_len for item in prepared], + [item.seq_len + item.protected_tail for item in prepared], + draft_manager=self.draft_kv_cache_manager, + ) + except _FixedScoreStreamMismatch: + raise + except Exception as exc: + raise RuntimeError("TriAttention score staging failed") from exc + if not staged: + raise RuntimeError("TriAttention page-table staging rejected the cohort") + + def _evict_requests( + self, + evict_reqs, + num_layers: int, + protected_tail_lengths: Optional[Dict[int, int]] = None, + ) -> List[Tuple[int, int]]: + """Score and compact requests, returning ``(request_id, capacity)`` targets. + + Only full-attention layers participate in scoring. For kernel-masked SWA + layers, the latest model window is rebased to the tail of the common + compacted prefix before the request-wide capacity is reduced. + """ + if protected_tail_lengths is None: + protected_tail_lengths = {} + protected_tail_capacity = self._configured_protected_tail_capacity() + with nvtx_range_debug("triattention.resolve_layout", color="blue"): + layout = self._runtime_kv_layout(num_layers) + + # Resolve request length and page metadata before mutating any layer. + prepared: List[_PreparedEviction] = [] + with nvtx_range("triattention.metadata", color="cyan"): + for request, rid in evict_reqs: + request_state = self._request_states.get(rid) + seq_len = None if request_state is None else request_state.confirmed_kv_length + if seq_len is None: + raise RuntimeError(f"Missing confirmed KV length for request {rid}") + # Restore the uncompressed confirmed logical position from the + # physical prefix and cumulative eviction count. + round_start = seq_len + request_state.evicted_tokens + if seq_len <= self._minimum_evictable_length(request, seq_len): + continue + expected_keep_count = self._minimum_evictable_length(request, seq_len) + protected_tail = int(protected_tail_lengths.get(rid, 0)) + if protected_tail < 0 or protected_tail > protected_tail_capacity: + raise RuntimeError( + f"Request {rid} protected tail {protected_tail} exceeds " + f"configured capacity {protected_tail_capacity}" + ) + prepared.append( + _PreparedEviction( + request=request, + request_id=rid, + seq_len=int(seq_len), + round_start=int(round_start), + expected_keep_count=expected_keep_count, + protected_tail=protected_tail, + ) + ) + if not prepared: + return [] + with nvtx_range_debug("triattention.staging_lookup", color="blue"): + resources = self._eager_resources_for(layout, prepared) + score_staging = resources.score_staging + keep_set_selector = resources.keep_set_selector + batched_compaction = self._batched_compaction_for( + layout=layout, + prepared=prepared, + score_staging=score_staging, + keep_set_selector=keep_set_selector, + ) + with nvtx_range_debug("triattention.page_table_stage", color="orange"): + self._attach_page_ids(prepared, score_staging) + + try: + with nvtx_range("triattention.score", color="blue"): + per_head = score_staging.launch_prepared_score() + with nvtx_range("triattention.select", color="yellow"): + if isinstance(keep_set_selector, _BatchedUnionKeepSetSelector): + keep_set_selector.select_prepared_requests() + else: + keep_set_selector.select_requests( + per_head, + normalize_scores=self.normalize_scores, + ) + with nvtx_range("triattention.compact", color="purple"): + batched_compaction.launch() + finally: + consumer_streams = [self.kv_cache_manager._stream] + if self.draft_kv_cache_manager is not None: + consumer_streams.append(self.draft_kv_cache_manager._stream) + score_staging.mark_page_tables_consumed(*consumer_streams) + + capacity_targets = [] + for item in prepared: + keep_count = item.expected_keep_count + evicted = item.seq_len - keep_count + if evicted <= 0: + raise RuntimeError("TriAttention attempted an identity compaction") + request_state = self._request_states[item.request_id] + request_state.evicted_tokens += evicted + request_state.confirmed_kv_length = keep_count + # Publish the cumulative count on the request: this is the + # manager's only channel to the runtime. The model engine + # reads it back where it builds num_cached_tokens_per_seq, + # so the kernels see the compacted KV length next step. + item.request.py_num_compressed_tokens = request_state.evicted_tokens + capacity_targets.append((item.request_id, keep_count)) + return capacity_targets + + def _num_layers_from_manager(self) -> int: + return len(self.kv_cache_manager.pp_layers) + + # ------------------------------------------------------------------ # + # Helpers: calibration loading # + # ------------------------------------------------------------------ # + + def _resolve_calibration(self) -> Dict[str, torch.Tensor]: + """Load the user-supplied calibration .pt and return our runtime schema. + + TriAttention does NOT compute calibration -- the user calibrates with the + official tool (github.com/WeianMao/triattention) and passes that file via + ``calibration_path``; we only run inference. Both the official R-KV layout + (``{metadata, stats{"layerLL_headHH": {q_mean_real, q_mean_imag, + q_abs_mean}}}``) and our already-converted flat layout are accepted -- the + official one is converted here.""" + if self.calibration_path is None: + raise ValueError( + "TriAttention requires `calibration_path`: a calibration .pt from " + "the official tool (github.com/WeianMao/triattention). TRT-LLM does " + "not compute calibration -- see examples/ for the Qwen3-8B file and " + "the official calibration instructions." + ) + raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) + if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): + calib = {k: (v.to("cuda") if torch.is_tensor(v) else v) for k, v in raw.items()} + self._validate_calibration(calib) + return calib + if isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): + return self._convert_official_calibration(raw) + got = sorted(raw.keys()) 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}." + ) + + def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: + """Convert the official per-(layer, head) stats to our flat runtime schema. + + ``E_q[l,h] = q_mean_real + i*q_mean_imag`` and ``E_q_norm[l,h] = + q_abs_mean`` are the same statistic, just restacked into ``[L, H, F]``. + ``omega`` / ``freq_scale_sq`` are not in the official file (its runtime + recomputes them from the model rotary), so we derive them from the model + config -- model-intrinsic and corpus-independent.""" + stats = raw["stats"] + meta = raw.get("metadata", {}) + if "sampled_heads" in meta: + heads = [(int(a), int(b)) for a, b in meta["sampled_heads"]] + else: + heads = [ + (int(k[len("layer") : k.index("_head")]), int(k[k.index("_head") + len("_head") :])) + for k in stats + ] + num_layers = max(layer for layer, _ in heads) + 1 + num_heads = max(h for _, h 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, h in heads: + s = stats[f"layer{layer:02d}_head{h:02d}"] + E_q[layer, h] = torch.complex(s["q_mean_real"].float(), s["q_mean_imag"].float()) + E_q_norm[layer, h] = s["q_abs_mean"].float() + omega, freq_scale_sq = self._rope_tables(freq_count) + calib = { + "E_q": E_q.to("cuda"), + "E_q_norm": E_q_norm.to("cuda"), + "omega": omega.to("cuda"), + "freq_scale_sq": freq_scale_sq.to("cuda"), + } + self._validate_calibration(calib) + logger.info( + f"TriAttention: converted official calibration {self.calibration_path}" + f" -> E_q[L={num_layers}, H={num_heads}, F={freq_count}]" + ) + return calib + + def _rope_tables(self, freq_count: int): + """RoPE ``omega`` (inv_freq) + ``freq_scale_sq`` (squared position-0 + amplitude) from the model config -- model-intrinsic, corpus-independent + (the official file does not store them). transformers' rope-init handles + plain and scaled RoPE; plain RoPE has attention_factor 1 so freq_scale_sq + is all ones. Falls back to the analytic inv_freq if rope-init is absent.""" + if self.model_path is None: + raise ValueError( + "TriAttention needs `model_path` to derive the RoPE tables " + "(omega / freq_scale_sq) when converting the official calibration." + ) + from transformers import AutoConfig + + cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() + config_values = cfg.to_dict() + head_dim = freq_count * 2 + base = float(config_values.get("rope_theta", 10000.0)) + try: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + scaling = config_values.get("rope_scaling") or {} + rope_type = scaling.get("rope_type") or scaling.get("type") or "default" + inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") + omega = inv_freq.to(torch.float32)[:freq_count].clone() + scale_sq = float(attention_factor) ** 2 + except Exception: + idx = torch.arange(0, head_dim, 2, dtype=torch.float32) + omega = (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() + scale_sq = 1.0 + return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) + + def _validate_calibration(self, calibration: Dict[str, torch.Tensor]) -> None: + """Verify the calibration dict has the expected keys.""" + missing = _REQUIRED_CALIBRATION_KEYS - set(calibration.keys()) + if missing: + raise ValueError( + f"TriAttention calibration is missing keys: {sorted(missing)}; " + f"got {sorted(calibration.keys())}." + ) 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..cd830079a61e --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -0,0 +1,1057 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU kernels for the TriAttention KV-eviction pipeline. + +The production path uses one fixed-shape trig-score launch across all dense +layers, CuTE-DSL TopK selection, and grouped C++ compaction. This module owns +the score kernel and its persistent launcher; selection and compaction live in +their respective runtime modules. + +House rules honored throughout: + * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). + * int64 for every page/stride offset that can exceed 2^31 (paged-pool reads). + * mask seq tails not divisible by ``tokens_per_block`` (and freq/dim tails). + * the kernels are vendored in this module (no lazy-load hub). +""" + +from __future__ import annotations + +from typing import List, Optional + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _prepare_mean_phase_kernel( + round_starts, + offsets, + omega, + mean_cos, + mean_sin, + NUM_FREQS: tl.constexpr, + NUM_OFFSETS: tl.constexpr, + F_BLOCK: tl.constexpr, +): + """Collapse all offset phases for one request into reusable frequency means.""" + request = tl.program_id(0) + frequency = tl.arange(0, F_BLOCK) + frequency_mask = frequency < NUM_FREQS + round_start = tl.load(round_starts + request) + angular_frequency = tl.load(omega + frequency, mask=frequency_mask, other=0.0) + cos_sum = tl.zeros((F_BLOCK,), tl.float32) + sin_sum = tl.zeros((F_BLOCK,), tl.float32) + for offset_index in tl.static_range(0, NUM_OFFSETS): + offset = tl.load(offsets + offset_index) + phase = (round_start + offset) * angular_frequency + cos_sum += tl.cos(phase) + sin_sum += tl.sin(phase) + output_offset = request * NUM_FREQS + frequency + scale = 1.0 / NUM_OFFSETS + tl.store(mean_cos + output_offset, cos_sum * scale, mask=frequency_mask) + tl.store(mean_sin + output_offset, sin_sum * scale, mask=frequency_mask) + + +def prepare_mean_phase( + round_starts: torch.Tensor, + offsets: torch.Tensor, + omega: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + request_count: int, +) -> None: + """Prepare mean score phases in one launch without intermediate tensors.""" + request_count = int(request_count) + if request_count <= 0 or request_count > round_starts.numel(): + raise ValueError("phase preparation request count is outside its fixed buffers") + num_freqs = int(omega.numel()) + num_offsets = int(offsets.numel()) + if ( + num_freqs <= 0 + or num_offsets <= 0 + or mean_cos.ndim != 2 + or mean_cos.shape[0] < request_count + or mean_cos.shape[1] != num_freqs + or mean_sin.shape != mean_cos.shape + or any( + tensor.device != round_starts.device for tensor in (offsets, omega, mean_cos, mean_sin) + ) + or round_starts.dtype != torch.int32 + or any(tensor.dtype != torch.float32 for tensor in (offsets, omega, mean_cos, mean_sin)) + ): + raise ValueError("phase preparation tensors do not share one valid FP32 geometry") + _prepare_mean_phase_kernel[(request_count,)]( + round_starts, + offsets, + omega, + mean_cos, + mean_sin, + NUM_FREQS=num_freqs, + NUM_OFFSETS=num_offsets, + F_BLOCK=triton.next_power_of_2(num_freqs), + num_warps=1, + ) + + +@triton.jit +def _score_row_stats_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + """Compute one valid-prefix mean and inverse standard deviation per score row.""" + flat_row = tl.program_id(0) + request = flat_row // ROWS + valid_width = tl.load(valid_widths + request) + score_row = scores + flat_row * WIDTH + lane = tl.arange(0, BLOCK) + score_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < valid_width + value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) + score_sum += tl.sum(value, axis=0) + mean = score_sum / valid_width + square_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < valid_width + 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 / valid_width) + tl.store(row_mean + flat_row, mean) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) + + +@triton.jit +def _score_union_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, + NORMALIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Normalize score rows and reduce them directly to one request-level union.""" + request = tl.program_id(0) + token = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + valid_width = tl.load(valid_widths + request) + valid_token = token < valid_width + union_max = tl.full((BLOCK,), -float("inf"), tl.float32) + for row in tl.range(0, ROWS): + flat_row = request * ROWS + row + value = tl.load( + scores + flat_row * WIDTH + 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")) + union_max = tl.maximum(union_max, value) + tl.store(combined + request * WIDTH + token, union_max, mask=token < WIDTH) + + +def prepare_union_scores( + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + combined: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, +) -> None: + """Mask, normalize, and union-reduce score rows in two or three launches.""" + request_count = int(request_count) + if not scores.is_cuda or scores.ndim != 3 or scores.dtype != torch.float32: + raise ValueError("union score preparation requires contiguous CUDA FP32 rows") + if not scores.is_contiguous() or request_count != scores.shape[0]: + raise ValueError("union score preparation request geometry does not match") + _, rows, width = scores.shape + if ( + valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.numel() < request_count * rows + or row_inv_std.shape != row_mean.shape + or combined.shape != (request_count, width) + ): + raise ValueError("union score preparation buffers do not match") + stats_block = 256 + if normalize_scores: + _score_row_stats_kernel[(request_count * rows,)]( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS=rows, + WIDTH=width, + BLOCK=stats_block, + num_warps=4, + ) + union_block = 32 + _score_union_kernel[(request_count, triton.cdiv(width, union_block))]( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + ROWS=rows, + WIDTH=width, + NORMALIZE=normalize_scores, + BLOCK=union_block, + num_warps=1, + ) + + +@triton.jit +def _score_per_head_reduce_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + selection_scores, + selection_seq_lens, + NUM_LAYERS: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + QUERY_GROUP_SIZE: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + WIDTH: tl.constexpr, + PER_LAYER: tl.constexpr, + NORMALIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Reduce query-head score rows into one selector row per KV-head domain.""" + 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) + valid_width = tl.load(valid_widths + request) + valid_token = token < valid_width + + if token_block == 0: + tl.store( + selection_seq_lens + request * SELECTION_ROWS + selection_row, + valid_width, + ) + + 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_QUERY_HEADS + query_head + value = tl.load( + scores + flat_row * WIDTH + 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_QUERY_HEADS + query_head + value = tl.load( + scores + flat_row * WIDTH + 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 prepare_per_head_scores( + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + selection_scores: torch.Tensor, + selection_seq_lens: torch.Tensor, + request_count: int, + *, + num_kv_heads: int, + per_layer: bool, + normalize_scores: bool, +) -> None: + """Normalize and reduce score rows for either per-head eviction mode.""" + request_count = int(request_count) + num_kv_heads = int(num_kv_heads) + if not scores.is_cuda or scores.ndim != 4 or scores.dtype != torch.float32: + raise ValueError("per-head score preparation requires CUDA FP32 rows") + if not scores.is_contiguous() or request_count != scores.shape[0]: + raise ValueError("per-head score preparation request geometry does not match") + _, num_layers, num_query_heads, width = scores.shape + if num_kv_heads <= 0 or num_query_heads % num_kv_heads: + raise ValueError("per-head score preparation requires valid GQA geometry") + selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads + if ( + valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.numel() < request_count * num_layers * num_query_heads + or row_inv_std.shape != row_mean.shape + or selection_scores.shape != (request_count, selection_rows, width) + or selection_scores.dtype != torch.float32 + or selection_scores.device != scores.device + or selection_seq_lens.shape != (request_count, selection_rows) + or selection_seq_lens.dtype != torch.int32 + or selection_seq_lens.device != scores.device + ): + raise ValueError("per-head score preparation buffers do not match") + + stats_block = 256 + rows = num_layers * num_query_heads + if normalize_scores: + _score_row_stats_kernel[(request_count * rows,)]( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS=rows, + WIDTH=width, + BLOCK=stats_block, + num_warps=4, + ) + reduction_block = 256 + _score_per_head_reduce_kernel[ + (request_count, selection_rows, triton.cdiv(width, reduction_block)) + ]( + scores, + valid_widths, + row_mean, + row_inv_std, + selection_scores, + selection_seq_lens, + NUM_LAYERS=num_layers, + NUM_QUERY_HEADS=num_query_heads, + NUM_KV_HEADS=num_kv_heads, + QUERY_GROUP_SIZE=num_query_heads // num_kv_heads, + SELECTION_ROWS=selection_rows, + WIDTH=width, + PER_LAYER=per_layer, + NORMALIZE=normalize_scores, + BLOCK=reduction_block, + num_warps=4, + ) + + +@triton.jit +def _pack_compaction_sources_kernel( + selected_indices, + valid_seq_lens, + dense_offsets, + dense_indices, + swa_offsets, + swa_indices, + DENSE_TOTAL: tl.constexpr, + SWA_TOTAL: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + SELECTION_STRIDE: tl.constexpr, + KEEP_COUNT: tl.constexpr, + PROMPT_LEN: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + SWA_WINDOW: tl.constexpr, + UNION: tl.constexpr, + PER_LAYER: tl.constexpr, + HAS_SWA: tl.constexpr, + BLOCK: tl.constexpr, +): + """Pack selected decode ordinals and protected tails for the C++ updater.""" + request = tl.program_id(0) + domain = tl.program_id(1) + move = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + + dense_begin = tl.load(dense_offsets + request) + dense_end = tl.load(dense_offsets + request + 1) + dense_count = dense_end - dense_begin + seq_len = tl.load(valid_seq_lens + request) + + if UNION: + selection_domain = 0 + else: + selection_domain = domain + selection_row = request * SELECTION_ROWS + selection_domain + selected = tl.load( + selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + PROMPT_LEN + move, + mask=move < KEEP_COUNT, + other=0, + ) + dense_source = tl.where(move < KEEP_COUNT, selected, seq_len + move - KEEP_COUNT) + dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) + + if HAS_SWA: + # Per-layer selection has one dense domain per (layer, head). SWA uses + # one shared source row per head, so only the first layer writes it. + if PER_LAYER: + write_swa = domain < NUM_KV_HEADS + else: + write_swa = move >= 0 + swa_begin = tl.load(swa_offsets + request) + swa_end = tl.load(swa_offsets + request + 1) + swa_count = swa_end - swa_begin + head = domain % NUM_KV_HEADS + swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move + swa_source = seq_len - SWA_WINDOW + move + tl.store( + swa_indices + swa_output, + swa_source, + mask=write_swa & (move < swa_count), + ) + + +def pack_compaction_sources( + selected_indices: torch.Tensor, + valid_seq_lens: torch.Tensor, + dense_offsets: torch.Tensor, + dense_indices: torch.Tensor, + *, + eviction_mode: str, + prompt_len: int, + keep_count: int, + num_dense_layers: int, + num_kv_heads: int, + max_protected_tail: int, + swa_window: int = 0, + swa_offsets: Optional[torch.Tensor] = None, + swa_indices: Optional[torch.Tensor] = None, +) -> None: + """Pack all dynamic dense and optional SWA source ordinals in one launch.""" + if eviction_mode not in ("union", "per_head", "per_layer_perhead"): + raise ValueError(f"unsupported compaction mode: {eviction_mode}") + prompt_len = int(prompt_len) + keep_count = int(keep_count) + num_dense_layers = int(num_dense_layers) + num_kv_heads = int(num_kv_heads) + max_protected_tail = int(max_protected_tail) + swa_window = int(swa_window) + request_count = int(selected_indices.shape[0]) if selected_indices.ndim else 0 + if ( + request_count <= 0 + or min(keep_count, num_dense_layers, num_kv_heads) <= 0 + or min(prompt_len, max_protected_tail, swa_window) < 0 + or selected_indices.shape[-1] != prompt_len + keep_count + ): + raise ValueError("compaction packing requires valid positive geometry") + + tensors = (selected_indices, valid_seq_lens, dense_offsets, dense_indices) + if any( + not tensor.is_cuda + or tensor.dtype != torch.int32 + or not tensor.is_contiguous() + or tensor.device != selected_indices.device + for tensor in tensors + ): + raise ValueError("compaction packing requires contiguous CUDA int32 tensors") + if valid_seq_lens.shape != (request_count,) or dense_offsets.shape != (request_count + 1,): + raise ValueError("compaction lengths and offsets do not match the request count") + + per_layer = eviction_mode == "per_layer_perhead" + union = eviction_mode == "union" + if union: + selection_rows = 1 + elif per_layer: + selection_rows = num_dense_layers * num_kv_heads + else: + selection_rows = num_kv_heads + if union: + expected_selection_prefix = (request_count,) + else: + expected_selection_prefix = (request_count, selection_rows) + if tuple(selected_indices.shape[:-1]) != expected_selection_prefix: + raise ValueError("selected indices do not match the compaction mode") + + domain_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads + expected_dense_prefix = (num_dense_layers, num_kv_heads) if per_layer else (num_kv_heads,) + if ( + dense_indices.ndim != len(expected_dense_prefix) + 1 + or tuple(dense_indices.shape[:-1]) != expected_dense_prefix + ): + raise ValueError("dense source buffer does not match the compaction mode") + dense_total = int(dense_indices.shape[-1]) + + has_swa = swa_indices is not None or swa_offsets is not None + if has_swa: + if swa_indices is None or swa_offsets is None or swa_window <= 0: + raise ValueError("SWA packing requires indices, offsets, and a positive window") + if ( + not swa_indices.is_cuda + or not swa_offsets.is_cuda + or swa_indices.dtype != torch.int32 + or swa_offsets.dtype != torch.int32 + or not swa_indices.is_contiguous() + or not swa_offsets.is_contiguous() + or swa_indices.device != selected_indices.device + or swa_offsets.device != selected_indices.device + or swa_indices.ndim != 2 + or tuple(swa_indices.shape[:-1]) != (num_kv_heads,) + or swa_offsets.shape != (request_count + 1,) + ): + raise ValueError("SWA source buffers do not match the compaction geometry") + swa_total = int(swa_indices.shape[-1]) + swa_indices_arg = swa_indices + swa_offsets_arg = swa_offsets + else: + if swa_window != 0: + raise ValueError("SWA window requires SWA source buffers") + swa_total = 0 + # HAS_SWA specializes the corresponding loads and stores away. + swa_indices_arg = dense_indices + swa_offsets_arg = dense_offsets + + max_move = keep_count + max_protected_tail + if has_swa: + max_move = max(max_move, swa_window + max_protected_tail) + block = 256 + _pack_compaction_sources_kernel[(request_count, domain_count, triton.cdiv(max_move, block))]( + selected_indices, + valid_seq_lens, + dense_offsets, + dense_indices, + swa_offsets_arg, + swa_indices_arg, + DENSE_TOTAL=dense_total, + SWA_TOTAL=swa_total, + SELECTION_ROWS=selection_rows, + SELECTION_STRIDE=prompt_len + keep_count, + KEEP_COUNT=keep_count, + PROMPT_LEN=prompt_len, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=has_swa, + BLOCK=block, + num_warps=4, + ) + + +@triton.jit +def _finalize_topk_indices_kernel( + scores, + seq_lens, + provisional_indices, + output_indices, + WIDTH: tl.constexpr, + KEEP_COUNT: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + PROMPT_LEN: tl.constexpr, + BLOCK: tl.constexpr, +): + """Resolve boundary ties and emit increasing physical token indices.""" + row = tl.program_id(0) + row_scores = scores + row * WIDTH + row_selected = provisional_indices + row * KEEP_COUNT + row_output = output_indices + row * OUTPUT_WIDTH + PROMPT_LEN + + 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, + ) + selected_score = tl.load( + row_scores + token_index, + mask=selected_mask, + other=float("inf"), + ).to(tl.float32) + threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) + + seq_len = tl.load(seq_lens + 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 < seq_len) + 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 < seq_len) + 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_LEN, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) + + +def finalize_topk_indices( + scores: torch.Tensor, + seq_lens: torch.Tensor, + provisional_indices: torch.Tensor, + output_indices: torch.Tensor, + keep_count: int, + prompt_len: int, +) -> None: + """Finalize one provisional TopK set without changing the CuTE selector. + + The kernel derives the provisional set's threshold, keeps all strictly + better scores, resolves the remaining boundary ties by lower token index, + and writes increasing physical ordinals directly into ``output_indices``. + """ + keep_count = int(keep_count) + prompt_len = int(prompt_len) + if not scores.is_cuda: + raise ValueError("deterministic TopK finalization requires CUDA scores") + if scores.ndim != 2 or scores.dtype != torch.float32 or not scores.is_contiguous(): + raise ValueError("TopK finalization requires contiguous two-dimensional FP32 scores") + rows, width = scores.shape + if rows <= 0 or not 1 <= keep_count <= width or prompt_len < 0: + raise ValueError("TopK finalization requires valid rows, keep count, and prompt length") + if ( + seq_lens.shape != (rows,) + or seq_lens.dtype != torch.int32 + or seq_lens.device != scores.device + or not seq_lens.is_contiguous() + ): + raise ValueError("TopK finalization sequence lengths do not match the score rows") + if ( + provisional_indices.shape != (rows, keep_count) + or provisional_indices.dtype != torch.int32 + or provisional_indices.device != scores.device + or not provisional_indices.is_contiguous() + ): + raise ValueError("provisional TopK indices do not match the requested selection") + if ( + output_indices.ndim != 2 + or output_indices.shape[0] != rows + or output_indices.shape[1] < prompt_len + keep_count + or output_indices.dtype != torch.int32 + or output_indices.device != scores.device + or not output_indices.is_contiguous() + ): + raise ValueError("TopK output does not fit the requested physical indices") + _finalize_topk_indices_kernel[(rows,)]( + scores, + seq_lens, + provisional_indices, + output_indices, + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=output_indices.shape[1], + PROMPT_LEN=prompt_len, + BLOCK=256, + num_warps=4, + ) + + +@triton.jit +def _tri_score_perhead_kernel( + pool_anchor_ptr, # typed pool pointer; used ONLY to infer the element type + # for the int->pointer cast below (its data is never read through it). + layer_base_addrs, # [num_layers] int64: ABSOLUTE device address of each + # scored layer's HND base. Layers do NOT need to share one storage: + # each segment casts its own layer's address back to a typed pointer, so + # all address arithmetic stays inside that layer's own allocation. + block_offsets_ptr, # Native V2 [pool, request, K/V, block] int32 offsets. + seg_page_off, # [nseg] int64: offset of this segment's page table into + # block_offsets_ptr. + # per-SEGMENT metadata (seg = req_slot*L_scored + layer_slot), idx by pid(0): + seg_req_id, # [nseg] int32: request slot (round_start / mean phase lookup) + seg_layer_id, # [nseg] int32: ABSOLUTE layer id (indexes layer_base_addrs + calib) + req_seq_len, # [num_requests] int32 + req_valid_width_out, # [num_requests] int32: decode-only length for selection + req_round_start, # [num_requests] int32 logical token position + # per-LAYER calibration, [L,H,F] flattened layer-major: + q_real_ptr, # [L*H*F] fp32 + q_imag_ptr, # [L*H*F] fp32 + mlr_coef_ptr, # [L*H*F] fp32 + # per-REQUEST offset-collapsed phase ('mean' path), [num_requests,F] flattened: + mean_cos_ptr, # [num_requests*F] fp32 + mean_sin_ptr, # [num_requests*F] fp32 + # shared freq vectors: + freq_scale_sq_ptr, # [F] fp32 + omega_ptr, # [F] fp32 ('max' path only) + offsets_ptr, # [O] fp32 ('max' path only) + out_ptr, # [request, layer, query_head, decode_token] fp32 + output_width, + # scalars uniform across the batch: + num_layers, + num_q_heads, + num_kv_heads, + num_freqs, # F = head_dim // 2 + head_dim, + tokens_per_block, + kv_factor, + num_offsets, # O ('max' path only) + # per-layer HND element strides (uniform across scored layers): + s_page, + s_kv_head, + s_slot, + s_dim, + USE_MAX: tl.constexpr, + TOKEN_START: tl.constexpr, + T_BLOCK: tl.constexpr, + F_BLOCK: tl.constexpr, +): + seg = tl.program_id(0) + t_blk = tl.program_id(1) + # KV heads are grid-parallel (axis 2): iterations of the former kv_head + # loop shared NO data (each KV head reads its own K and writes its own + # output rows), so hoisting it onto the grid multiplies parallelism with + # zero extra HBM traffic. The q-in-group loop below stays inside the + # program because it REUSES this head's K from registers (GQA dedup). + kv_head = tl.program_id(2) + + req_id = tl.load(seg_req_id + seg) + seq_len = tl.load(req_seq_len + req_id) + if (seg % num_layers == 0) & (t_blk == 0) & (kv_head == 0): + tl.store(req_valid_width_out + req_id, seq_len - TOKEN_START) + # Derive the ragged launch bound in the score program instead of staging + # one replicated length and block count for every request/layer segment. + n_tblk = (seq_len - TOKEN_START + T_BLOCK - 1) // T_BLOCK + if t_blk >= n_tblk: + return + + layer_id = tl.load(seg_layer_id + seg) + rstart = tl.load(req_round_start + req_id) + # This segment's layer base: an absolute address cast back to a pool-typed + # pointer. TRT-LLM V2 exposes every layer as its own TensorWrapper storage, + # so "element offset relative to one shared storage" does not exist; the + # per-layer absolute address is the same device-pointer-array pattern the + # C++ backends use (KVBlockArray / grouped-GEMM pointer arrays). + layer_ptr = tl.load(layer_base_addrs + layer_id).to( + tl.pointer_type(pool_anchor_ptr.dtype.element_ty) + ) + page_off = tl.load(seg_page_off + seg) + + f = tl.arange(0, F_BLOCK) + f_mask = f < num_freqs + f64 = f.to(tl.int64) + + # ---- token tile of THIS segment ---- + t = t_blk * T_BLOCK + tl.arange(0, T_BLOCK) + absolute_t = t + TOKEN_START + t_mask = absolute_t < seq_len + blk_in_seq = absolute_t // tokens_per_block + slot = (absolute_t % tokens_per_block).to(tl.int64) + # The native attention page-table copy encodes K offsets in units of the + # underlying K/V role pages. Convert that value to the HND pool page inline + # instead of materializing a second page table before scoring. + encoded_page = tl.load( + block_offsets_ptr + page_off + blk_in_seq, + mask=t_mask, + other=0, + ) + phys_page = (encoded_page // kv_factor).to(tl.int64) + + # element offset into THIS layer's pool for (page, KEY=0, *, slot). + # KEY half is kv_factor index 0 -> its stride term is 0 (matches reference). + tok_base = phys_page * s_page + slot * s_slot # [T_BLOCK] int64 + + # per-request 'mean'-path phase + shared freq scale. + mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) + msin = tl.load(mean_sin_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) + fss = tl.load(freq_scale_sq_ptr + f, mask=f_mask, other=0.0) + + # ---- PER-HEAD (position + mlr), GQA-deduped, NO head reduction ---- + # This program scores ONE KV head's token tile for the group_size q-heads + # that share it. K (and |K|) is loaded ONCE and reused across the group; + # h = kv_head*group_size + qg keeps query-head order 0..num_q_heads-1, so + # every head's math is bit-for-bit identical to the looped variant. + group_size = num_q_heads // num_kv_heads + load_mask = t_mask[:, None] & f_mask[None, :] + off_re = f64[None, :] * s_dim + off_im = (num_freqs + f64[None, :]) * s_dim + + base = tok_base + kv_head.to(tl.int64) * s_kv_head # [T_BLOCK] + # paged K loaded ONCE for this KV head (shared by group_size q-heads). + k_re = tl.load(layer_ptr + base[:, None] + off_re, mask=load_mask, other=0.0).to(tl.float32) + k_im = tl.load(layer_ptr + base[:, None] + off_im, mask=load_mask, other=0.0).to(tl.float32) + kmag = tl.sqrt(k_re * k_re + k_im * k_im) # once per KV head + + qg = 0 + while qg < group_size: + h = kv_head * group_size + qg + calib_off = (layer_id.to(tl.int64) * num_q_heads + h) * num_freqs + qre = tl.load(q_real_ptr + calib_off + f, mask=f_mask, other=0.0) + qim = tl.load(q_imag_ptr + calib_off + f, mask=f_mask, other=0.0) + mlrc = tl.load(mlr_coef_ptr + calib_off + f, mask=f_mask, other=0.0) + + # complex product Q . conj(K) -- the trig importance score. + prod_real = qre[None, :] * k_re + qim[None, :] * k_im + prod_imag = qim[None, :] * k_re - qre[None, :] * k_im + + if USE_MAX: + # max over O offsets does NOT commute through the freq-sum; + # explicit O loop reducing max over the per-offset F-sum. + score = tl.full((T_BLOCK,), -float("inf"), tl.float32) + o = 0 + while o < num_offsets: + off = tl.load(offsets_ptr + o) + om = tl.load(omega_ptr + f, mask=f_mask, other=0.0) + phase = (rstart + off) * om + cphase = tl.cos(phase) + sphase = tl.sin(phase) + per_f = fss[None, :] * (prod_real * cphase[None, :] - prod_imag * sphase[None, :]) + offset_score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) + score = tl.maximum(score, offset_score) + o += 1 + else: + # 'mean': offset loop collapsed into mean_cos/mean_sin. + per_f = fss[None, :] * (prod_real * mcos[None, :] - prod_imag * msin[None, :]) + score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) + + # position-INDEPENDENT MLR term (reuses the per-KV-head |K|). + mlr_f = kmag * mlrc[None, :] * fss[None, :] + mlr = tl.sum(tl.where(f_mask[None, :], mlr_f, 0.0), axis=1) + + # Segments are request-major then layer-major. Write the decode-only + # score directly in the selector's [request, layer, head, token] layout. + out_offset = (seg.to(tl.int64) * num_q_heads + h) * output_width + t + tl.store(out_ptr + out_offset, score + mlr, mask=t_mask) + qg += 1 + + +def _launch_tri_score_perhead( + grid: tuple, + pointer_args: tuple, + geometry_args: tuple, + *, + score_aggregation: str, + token_start: int, + token_block: int, + num_freqs: int, +) -> None: + """Launch the shared score ABI for eager and fixed metadata owners.""" + if score_aggregation not in ("mean", "max"): + raise ValueError(f"unsupported score aggregation: {score_aggregation}") + _tri_score_perhead_kernel[grid]( + *pointer_args, + *geometry_args, + USE_MAX=(score_aggregation == "max"), + TOKEN_START=token_start, + T_BLOCK=token_block, + F_BLOCK=triton.next_power_of_2(num_freqs), + ) + + +class _FixedScoreGroup: + """Persistent score metadata/output for one sequence bucket. + + Since the per-layer absolute-address ABI, ONE group can span dense layers + living in DISTINCT storages with DISTINCT block tables. ``block_offsets`` + uses the native TRT-LLM attention layout and ``page_table_slots`` maps each + scored layer to its V2 pool slot. + """ + + def __init__( + self, + layer_pools: List[torch.Tensor], + layer_indices: List[int], + max_requests: int, + page_count: int, + seq_len: int, + num_q_heads: int, + block_offsets: torch.Tensor, # [num_pools, max_requests, 2, copied_blocks] int32 + page_table_slots: List[int], # per scored layer: pool slot into block_offsets + q_real_LHF: torch.Tensor, + q_imag_LHF: torch.Tensor, + mlr_coef_LHF: torch.Tensor, + freq_scale_sq: torch.Tensor, + omega: torch.Tensor, + offsets: torch.Tensor, + prompt_len: int = 0, + ) -> None: + if not layer_indices or min(max_requests, page_count, seq_len) <= 0: + raise ValueError("fixed score group requires non-empty positive geometry") + if prompt_len < 0 or prompt_len >= seq_len: + raise ValueError("fixed score prompt length must leave a non-empty decode region") + if len(page_table_slots) != len(layer_indices): + raise ValueError("page_table_slots must align with layer_indices") + self.max_requests = max_requests + self.prompt_len = prompt_len + self.output_width = seq_len - prompt_len + self.num_layers = len(layer_indices) + p0 = layer_pools[layer_indices[0]] + if p0.ndim != 5: + raise ValueError("fixed score group requires HND pools") + device = p0.device + q_real_LHF = q_real_LHF.to(device=device, dtype=torch.float32).contiguous() + q_imag_LHF = q_imag_LHF.to(device=device, dtype=torch.float32).contiguous() + mlr_coef_LHF = mlr_coef_LHF.to(device=device, dtype=torch.float32).contiguous() + freq_scale_sq = freq_scale_sq.to(device=device, dtype=torch.float32).contiguous() + omega = omega.to(device=device, dtype=torch.float32).contiguous() + offsets = offsets.to(device=device, dtype=torch.float32).contiguous() + _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape + if num_q_heads % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + self.num_kv_heads = int(num_kv_heads) + self.num_freqs = head_dim // 2 + strides = tuple(int(value) for value in p0.stride()) + self.geometry_args = ( + num_q_heads, + num_kv_heads, + self.num_freqs, + head_dim, + tokens_per_block, + kv_factor, + int(offsets.numel()), + strides[0], + strides[2], + strides[3], + strides[4], + ) + # Per-layer ABSOLUTE base addresses. Layers may live in distinct + # storages (V2 TensorWrapper-per-layer); only geometry must be uniform. + element_size = p0.element_size() + layer_base_addrs = torch.zeros(len(layer_pools), dtype=torch.int64, device=device) + for layer in layer_indices: + pool = layer_pools[layer] + if ( + tuple(pool.shape[1:]) != tuple(p0.shape[1:]) + or tuple(pool.stride()) != strides + or pool.dtype != p0.dtype + ): + raise ValueError("fixed score layers must share one uniform geometry") + address = int(pool.data_ptr()) + if address % element_size: + raise ValueError("fixed score layer base is not element-aligned") + layer_base_addrs[layer] = address + # The anchor pool is passed as a typed kernel argument ONLY so the + # kernel can recover the element type for the int->pointer cast. + seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( + self.num_layers + ) + seg_layer = torch.tensor(layer_indices, dtype=torch.int32, device=device).repeat( + max_requests + ) + # Each segment reads the K plane for one request from the same native + # block-offset buffer used by TRT-LLM attention metadata preparation. + if ( + block_offsets.ndim != 4 + or tuple(block_offsets.shape[1:3]) != (max_requests, 2) + or block_offsets.shape[3] < page_count + or block_offsets.dtype != torch.int32 + or block_offsets.device != device + ): + raise ValueError("block offsets do not match fixed score geometry") + if not block_offsets.is_contiguous(): + raise ValueError("fixed score block offsets must be contiguous") + slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) + if int(slots_t.max()) >= int(block_offsets.shape[0]): + raise ValueError("page table slot exceeds staged page-id planes") + req_idx = torch.arange(max_requests, dtype=torch.int64, device=device).repeat_interleave( + self.num_layers + ) + slot_idx = slots_t.repeat(max_requests) + seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) + self.token_block = 64 + self.max_ntblk = (self.output_width + self.token_block - 1) // self.token_block + self.output = torch.empty( + max_requests, + self.num_layers, + num_q_heads, + self.output_width, + dtype=torch.float32, + device=device, + ) + self.pointer_prefix = ( + p0, + layer_base_addrs, + block_offsets.view(-1), + seg_page_off, + seg_req, + seg_layer, + ) + self.pointer_middle = ( + q_real_LHF.view(-1), + q_imag_LHF.view(-1), + mlr_coef_LHF.view(-1), + ) + self.pointer_tail = (freq_scale_sq, omega, offsets) + + def launch( + self, + request_count: int, + valid_seq_lens: torch.Tensor, + valid_widths: torch.Tensor, + round_starts_device: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + score_aggregation: str, + ) -> torch.Tensor: + """Return decode-only scores as ``[request, layer, head, token]``.""" + if request_count <= 0 or request_count > self.max_requests: + raise ValueError("request count exceeds fixed score capacity") + if ( + valid_widths.ndim != 1 + or valid_widths.numel() < request_count + or valid_widths.dtype != torch.int32 + or valid_widths.device != self.output.device + ): + raise ValueError("score output lengths do not fit the keep-set selector") + num_segments = request_count * self.num_layers + output = self.output[:request_count] + _launch_tri_score_perhead( + (num_segments, self.max_ntblk, self.num_kv_heads), + ( + *self.pointer_prefix, + valid_seq_lens, + valid_widths, + round_starts_device, + *self.pointer_middle, + mean_cos.view(-1), + mean_sin.view(-1), + *self.pointer_tail, + output, + ), + (self.output_width, self.num_layers, *self.geometry_args), + score_aggregation=score_aggregation, + token_start=self.prompt_len, + token_block=self.token_block, + num_freqs=self.num_freqs, + ) + return output diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index feacedf85151..573919f46914 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -32,7 +32,7 @@ CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, PeftCacheConfig, SamplerType, SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, - TorchLlmArgs, WaitingQueuePolicy) + TorchLlmArgs, TriAttentionKvCacheCompressionConfig, WaitingQueuePolicy) # isort: on from tensorrt_llm.logger import logger from tensorrt_llm.lora_helper import (LoraConfig, @@ -2043,15 +2043,82 @@ 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 also - decided here: a compression manager is only created when the speculative - mode supports it, otherwise the run stays uncompressed. + here; the framework ships none. Speculative-decoding compatibility is + decided here: evicting methods only accept modes whose draft KV is a + standard paged cache compacted with the target. """ - logger.warning( - "KV-cache compression algorithm '%s' is not registered; running without " - "a compression manager.", - config.algorithm, - ) + manager_class = None + if config.algorithm == "triattention": + from tensorrt_llm._torch.kv_cache_compression.triattention import \ + TriAttention + manager_class = TriAttention + + if manager_class is None: + logger.warning( + "KV-cache compression algorithm '%s' is not registered; running without " + "a compression manager.", + config.algorithm, + ) + return None + + # Evicting methods co-compact the draft KV, so they only support spec + # modes whose draft KV is a standard paged cache in the same forward + # (one-model speculation). For any other mode, no compression manager + # is created and the run stays uncompressed. + if (manager_class.is_eviction_method() and spec_config is not None + and not (spec_config.spec_dec_mode.is_mtp_one_model() + or spec_config.spec_dec_mode.is_eagle3_one_model())): + logger.warning( + "KV-cache compression algorithm '%s' evicts cached tokens and does " + "not support speculative decoding mode %s (the draft KV must be a " + "standard paged cache compacted together with the target); running " + "without a compression manager.", + config.algorithm, + spec_config.spec_dec_mode.name, + ) + return None + + if config.algorithm == "triattention": + if spec_config is not None: + if spec_config.max_draft_len is None: + raise ValueError( + "TriAttention speculative compatibility requires a " + "resolved max_draft_len") + if not spec_config.is_linear_tree: + raise ValueError( + "TriAttention speculative compatibility requires linear " + "drafting") + if spec_config.draft_len_schedule is not None: + raise ValueError( + "TriAttention does not yet support dynamic speculative " + "draft lengths") + if (spec_config.acceptance_window is not None + or spec_config.acceptance_length_threshold is not None): + raise ValueError( + "TriAttention does not support runtime speculative " + "acceptance gating") + if draft_kv_cache_manager is None: + raise ValueError( + "TriAttention speculative compatibility requires a " + "separate draft KV cache; shared target/draft pools " + "cannot be compacted safely") + triattention_config = ( + config if isinstance(config, TriAttentionKvCacheCompressionConfig) + else TriAttentionKvCacheCompressionConfig.model_validate( + config.model_dump())) + return TriAttention( + kv_cache_manager, + draft_kv_cache_manager=draft_kv_cache_manager, + top_B=triattention_config.top_B, + beta=triattention_config.beta, + model_path=triattention_config.model_path, + calibration_path=triattention_config.calibration_path, + eviction_mode=triattention_config.eviction_mode, + normalize_scores=triattention_config.normalize_scores, + pin_prefill=triattention_config.pin_prefill, + count_prompt_tokens=triattention_config.count_prompt_tokens, + ) + return None diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 1682eb51e71c..39e06c2b156f 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -24,7 +24,8 @@ SADecodingConfig, SAEnhancerConfig, SaveHiddenStatesDecodingConfig, SchedulerConfig, SkipSoftmaxAttentionConfig, TorchCompileConfig, - TorchLlmArgs, TrtLlmArgs, UserProvidedDecodingConfig) + TorchLlmArgs, TriAttentionKvCacheCompressionConfig, + TrtLlmArgs, UserProvidedDecodingConfig) from .llm_utils import (BuildConfig, KvCacheRetentionConfig, QuantAlgo, QuantConfig) from .mm_encoder import MultimodalEncoder @@ -90,6 +91,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 fd3276163831..a38db3370ac8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3269,6 +3269,78 @@ class KvCacheCompressionConfig(StrictBaseModel): "and set the value.") +class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): + """KV-cache compression config for TriAttention. + + TriAttention periodically evicts cached tokens during generation, guided by + an offline-calibrated trigonometric importance score + (github.com/WeianMao/triattention). It runs on the KV-cache compression + framework with the standard ``KVCacheManagerV2``, whose ``update_resources`` + returns eviction-freed blocks to the pool for the capacity gain. TRT-LLM does + not compute calibration: supply the official tool's ``.pt`` via ``calibration_path`` and + it is converted to the runtime schema at load. TriAttention is a pure + compression method: it has no sparse-attention config and no attention + backend of its own -- decode runs the model's standard attention over the + compacted cache, and the manager reconciles the cached-token count via the + framework's ``adjust_attention_metadata`` hook. + """ + 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).") + pin_prefill: bool = Field( + default=True, + description="Always preserve the prompt (prefill) tokens; only decode " + "tokens compete for the budget (upstream behaviour).") + top_B: int = Field( + default=2048, + description="Tokens kept at each periodic eviction (upstream `budget`; " + "prompt tokens are always preserved on top).") + beta: int = Field( + default=128, + 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: Optional[str] = Field( + default=None, + description="Checkpoint path used to derive RoPE tables when converting " + "the official calibration and to classify kernel-masked sliding-attention " + "layers. Required by TriAttention.") + calibration_path: Optional[str] = Field( + default=None, + 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.") + window_size: int = Field( + default=128, + description="Compatibility field retained for existing configs. The " + "implemented calibration-based selection does not use a separate " + "recency window.") + count_prompt_tokens: bool = Field( + default=False, + description="If False (default), the KV budget counts only DECODE tokens " + "(the pinned prompt is kept on top). Physical capacity reclaim currently " + "requires False.") + + +KvCacheCompressionConfigType: TypeAlias = Union[ + TriAttentionKvCacheCompressionConfig, + KvCacheCompressionConfig, +] + + @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) class AgentTreeConfig(StrictBaseModel, PybindMirror): """Configuration for agent tree scheduling. @@ -3987,7 +4059,7 @@ class BaseLlmArgs(StrictBaseModel): # 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_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 db1c7772fdfb..c491b8d30315 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -543,6 +543,68 @@ "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.count_prompt_tokens" + }, + { + "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": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.pin_prefill" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.top_B" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.window_size" + }, { "allowed_values": [], "annotation": "", @@ -2961,6 +3023,68 @@ "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.count_prompt_tokens" + }, + { + "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": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.pin_prefill" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.top_B" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.window_size" + }, { "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 71e081414933..212f59770402 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -106,6 +106,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/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) 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..4989508f73c7 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""KV-cache compression forces the unfused-RoPE path. + +Compression physically evicts cached tokens, so the KV length stops matching +the logical sequence length. The fused path derives each new token's rotary +position from the KV length inside the attention kernel; the unfused path +consumes the engine's logical ``position_ids``. With compression enabled the +attention module must therefore keep RoPE unfused so rotary positions stay +logical (original absolute positions, matching the official TriAttention +implementations) while the shortened KV length only bounds attention extent. +""" + +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 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_kv_cache_compression_forces_unfused_rope() -> None: + model_config = ModelConfig(kv_cache_compression_config=TriAttentionKvCacheCompressionConfig()) + attn = _make_attention(model_config) + + assert attn.rope_fusion is False + + +def test_unfused_yarn_rope_is_applied_exactly_once() -> None: + """With rope_fusion=False the Python-side rotary module owns RoPE, so the + backend must receive no position-embedding params. yarn is not listed in + PositionEmbeddingType.is_rope(), which used to leak the params through and + made the C++ QKV preprocess rotate a second time (double RoPE).""" + from tensorrt_llm._torch.attention_backend.interface import ( + PositionalEmbeddingParams, + RopeParams, + ) + from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + yarn_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.yarn, + rope=RopeParams( + dim=32, + theta=150000, + scale_type=RotaryScalingType.yarn, + scale=32.0, + max_positions=1024, + original_max_positions=256, + beta_fast=32, + beta_slow=1, + duplicate_data=False, + ), + is_neox=True, + ) + model_config = ModelConfig(kv_cache_compression_config=TriAttentionKvCacheCompressionConfig()) + attn = Attention( + hidden_size=256, + num_attention_heads=8, + num_key_value_heads=8, + max_position_embeddings=1024, + bias=False, + pos_embd_params=yarn_params, + layer_idx=0, + dtype=torch.bfloat16, + config=model_config, + ) + + assert attn.rope_fusion is False + assert attn.rotary_emb is not None + # The TRTLLM backend keeps the type as an int; 0 means no position + # embedding was handed to the kernel side. + assert attn.attn.position_embedding_type == 0 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..647a7ae35f31 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -0,0 +1,534 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Draft KV co-compression tests for TriAttention eviction. + +With one-model speculative decoding, TriAttention compacts the separate draft +KV cache in the same round as the target: the target's union keep set is +broadcast over the draft's own KV heads, the draft's own protected tail is +appended as ordinals ``valid_seq_len + 0..tail-1``, and both caches land at +``destination_base = prompt_len``. These tests cover the physical draft moves, +the packed move indices, stream ordering across both cache managers, the +speculative admission gates, the published compressed-token invariant, and +prepared-compaction cache invalidation. +""" + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch.kv_cache_compression.triattention import TriAttention +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( + BatchedKVCacheCompaction, +) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + _PreparedEviction, + _RequestCompressionState, +) +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + + +def _encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: + """Build the native V2 [pool, request, K/V, block] layout.""" + if page_ids.ndim == 2: + page_ids = page_ids.unsqueeze(0) + encoded = torch.empty( + page_ids.shape[0], + page_ids.shape[1], + 2, + page_ids.shape[2], + dtype=torch.int32, + device=page_ids.device, + ) + encoded[:, :, 0] = page_ids.to(torch.int32) * 2 + encoded[:, :, 1] = encoded[:, :, 0] + 1 + return encoded + + +def _make_fake_v2(*, 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 = False + fake_v2.is_draft = is_draft + fake_v2.kv_compression_manages_history = False + fake_v2.kv_factor = 2 + fake_v2.mapping = SimpleNamespace(enable_attention_dp=False) + fake_v2.is_disagg = False + 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.empty(1, dtype=torch.int64) + fake_v2.pp_layers = [] + fake_v2.layer_offsets = {} + fake_v2.layer_to_pool_mapping_dict = {} + return fake_v2 + + +def _make_triattention(**overrides): + options = {"top_B": 8, "model_path": "/models/test"} + options.update(overrides) + return TriAttention(_make_fake_v2(), **options) + + +def _make_request(request_id, **overrides): + 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) + + +def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: + """Gather one request's pages into [K/V, head, token, dim] order.""" + num_kv_heads = int(pool.shape[2]) + head_dim = int(pool.shape[4]) + return pool.index_select(0, pages).permute(1, 2, 0, 3, 4).reshape(2, num_kv_heads, -1, head_dim) + + +def _launched_draft_compaction(draft_protected_tails): + """Build target and draft pools with distinct head counts, then compact.""" + device = torch.device("cuda") + request_count = 2 + target_kv_heads = 2 + draft_kv_heads = 4 + prompt_len = 2 + decode_keep_count = 4 + tokens_per_block = 4 + head_dim = 16 + target_protected_tails = [2, 1] + valid_seq_lens = [10, 9] + + target_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + draft_tables = torch.tensor([[1, 0, 2], [5, 4, 3]], dtype=torch.int32, device=device) + target_pools = [ + ( + torch.arange( + 6 * 2 * target_kv_heads * tokens_per_block * head_dim, + dtype=torch.float32, + device=device, + ).view(6, 2, target_kv_heads, tokens_per_block, head_dim) + + layer * 100_000.0 + ) + for layer in range(2) + ] + draft_pool = ( + torch.arange( + 6 * 2 * draft_kv_heads * tokens_per_block * head_dim, + dtype=torch.float32, + device=device, + ).view(6, 2, draft_kv_heads, tokens_per_block, head_dim) + + 900_000.0 + ) + assert target_pools[0].shape[2] != draft_pool.shape[2] + initial_target = [pool.clone() for pool in target_pools] + initial_draft = draft_pool.clone() + + prompt = torch.tensor([0, 1], dtype=torch.int64, device=device) + union_decode = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) + keep = torch.cat((prompt.view(1, -1).expand(request_count, -1), union_decode), dim=1) + + compaction = BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=target_pools, + dense_layers=[0, 1], + swa_layers=[], + layer_group_representative={0: 0, 1: 1}, + layer_pool_keys=[("pool", 0), ("pool", 0)], + kept_token_ordinals=keep.to(torch.int32), + valid_sequence_lengths=torch.tensor(valid_seq_lens, dtype=torch.int32, device=device), + kv_block_offsets=_encode_block_offsets(target_tables), + page_table_slots={0: 0, 1: 0}, + request_count=request_count, + prompt_len=prompt_len, + decode_keep_count=decode_keep_count, + swa_window=None, + protected_tail_lengths=target_protected_tails, + draft_layer_pools=[draft_pool], + draft_layers=[0], + draft_layer_group_representative={0: 0}, + draft_layer_pool_keys=[("draft_pool", 0)], + draft_protected_tail_lengths=draft_protected_tails, + draft_kv_block_offsets=_encode_block_offsets(draft_tables), + draft_page_table_slots={0: 0}, + ) + compaction.launch() + torch.cuda.synchronize(device) + + return SimpleNamespace( + device=device, + request_count=request_count, + prompt_len=prompt_len, + keep=keep, + valid_seq_lens=valid_seq_lens, + target_protected_tails=target_protected_tails, + draft_protected_tails=draft_protected_tails, + target_tables=target_tables, + draft_tables=draft_tables, + target_pools=target_pools, + draft_pool=draft_pool, + initial_target=initial_target, + initial_draft=initial_draft, + compaction=compaction, + ) + + +def test_draft_pools_receive_target_union_keep_set_and_own_tail(): + built = _launched_draft_compaction(draft_protected_tails=[1, 2]) + device = built.device + prompt_len = built.prompt_len + + for request in range(built.request_count): + valid = built.valid_seq_lens[request] + # Target dense layers compact the union keep set plus the target tail. + target_pages = built.target_tables[request].to(torch.long) + target_tail = torch.arange( + valid, + valid + built.target_protected_tails[request], + dtype=torch.int64, + device=device, + ) + target_source = torch.cat((built.keep[request, prompt_len:], target_tail)) + target_destination = torch.arange( + prompt_len, + prompt_len + target_source.numel(), + dtype=torch.int64, + device=device, + ) + for before_pool, after_pool in zip(built.initial_target, built.target_pools): + before = _logical_view(before_pool, target_pages) + after = _logical_view(after_pool, target_pages) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + assert torch.equal( + after.index_select(2, target_destination), + before.index_select(2, target_source), + ) + + # The draft compacts the SAME kept ordinals through its OWN page + # table, over its own head count, with its own protected tail, at + # destination_base = prompt_len. + draft_pages = built.draft_tables[request].to(torch.long) + draft_tail = torch.arange( + valid, + valid + built.draft_protected_tails[request], + dtype=torch.int64, + device=device, + ) + draft_source = torch.cat((built.keep[request, prompt_len:], draft_tail)) + draft_destination = torch.arange( + prompt_len, + prompt_len + draft_source.numel(), + dtype=torch.int64, + device=device, + ) + before = _logical_view(built.initial_draft, draft_pages) + after = _logical_view(built.draft_pool, draft_pages) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + for head in range(int(built.draft_pool.shape[2])): + assert torch.equal( + after[:, head].index_select(1, draft_destination), + before[:, head].index_select(1, draft_source), + ) + + +@pytest.mark.parametrize("draft_protected_tails", [[1, 1], [1, 2]]) +def test_draft_pack_matches_keep_broadcast_and_tail_ordinal_oracle(draft_protected_tails): + built = _launched_draft_compaction(draft_protected_tails=draft_protected_tails) + draft_compaction = built.compaction.draft_compaction + + expected_offsets = [0] + expected_moves = [] + for request in range(built.request_count): + decode = built.keep[request, built.prompt_len :].to(torch.int32) + tail = torch.arange( + built.valid_seq_lens[request], + built.valid_seq_lens[request] + draft_protected_tails[request], + dtype=torch.int32, + device=built.device, + ) + moves = torch.cat((decode, tail)) + expected_moves.append(moves) + expected_offsets.append(expected_offsets[-1] + int(moves.numel())) + expected_row = torch.cat(expected_moves) + + assert draft_compaction.move_source_offsets.cpu().tolist() == expected_offsets + draft_indices = draft_compaction.move_source_indices + assert draft_indices.shape == (int(built.draft_pool.shape[2]), expected_offsets[-1]) + for head in range(int(draft_indices.shape[0])): + # Union mode broadcasts one keep set over every draft KV head. + assert torch.equal(draft_indices[head], expected_row) + + +def test_mark_page_tables_consumed_orders_both_manager_streams(): + staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + staging.device = torch.device("cuda") + staging.page_tables_active = True + event = mock.Mock() + staging.bulk_consume_done = event + target_stream = mock.Mock() + draft_stream = mock.Mock() + compute_stream = SimpleNamespace() + + with mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream): + staging.mark_page_tables_consumed(target_stream, draft_stream) + + # One event records the compact launches; BOTH cache managers wait on it, + # so neither can free or reallocate pages this cohort is still reading. + event.record.assert_called_once_with(compute_stream) + target_stream.wait_event.assert_called_once_with(event) + draft_stream.wait_event.assert_called_once_with(event) + assert staging.page_tables_active is False + + with pytest.raises(RuntimeError, match="not staged"): + staging.mark_page_tables_consumed(target_stream, draft_stream) + + +@pytest.mark.parametrize( + "gate,match", + [ + ("union_only", "union"), + ("draft_kv_factor", "standard key/value cache"), + ("full_attention_draft", "full-attention draft"), + ("dflash", "standard paged cache compacted together"), + ], +) +def test_draft_admission_gates_raise(gate, match): + draft_manager = _make_fake_v2(is_draft=True) + if gate == "full_attention_draft": + draft_manager.max_attention_window_vec = [128] + if gate == "dflash": + # DFlash reads cross-attention context buffers, not a paged KV cache; + # the factory's eviction-method speculative-mode gate declines to + # create a manager and the run stays uncompressed. + from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + from tensorrt_llm.llmapi.llm_args import ( + DFlashDecodingConfig, + TriAttentionKvCacheCompressionConfig, + ) + + manager = create_kv_cache_compression_manager( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + _make_fake_v2(), + draft_kv_cache_manager=draft_manager, + spec_config=DFlashDecodingConfig(max_draft_len=3), + ) + assert manager is None + return + manager = TriAttention( + _make_fake_v2(), + top_B=8, + model_path="/models/test", + eviction_mode="per_head" if gate == "union_only" else "union", + draft_kv_cache_manager=draft_manager, + ) + if gate == "draft_kv_factor": + # Flipping kv_factor after construction exercises TriAttention's own + # runtime gate. + draft_manager.kv_factor = 1 + + with pytest.raises(ValueError, match=match): + manager._validate_v2_compatibility() + + +@contextmanager +def _mocked_eviction_internals(manager): + """Run the real ``_evict_requests`` body around mocked GPU launches.""" + score_staging = SimpleNamespace( + launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), + mark_page_tables_consumed=mock.Mock(), + ) + keep_set_selector = SimpleNamespace(select_requests=mock.Mock()) + resources = SimpleNamespace( + score_staging=score_staging, + keep_set_selector=keep_set_selector, + ) + batched_compaction = SimpleNamespace(launch=mock.Mock()) + with ( + mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), + mock.patch.object(manager, "_eager_resources_for", return_value=resources), + mock.patch.object( + manager, + "_batched_compaction_for", + return_value=batched_compaction, + ), + mock.patch.object(manager, "_attach_page_ids"), + ): + yield score_staging + + +def test_compressed_count_is_monotone_and_tracks_confirmed_length(): + manager = _make_triattention(top_B=4, beta=4) + manager._calibrated = True + manager._attention_layer_partition_cache = ([0, 1], [], None) + 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 + + request = _make_request(7, py_prompt_len=2, py_num_accepted_draft_tokens=1) + manager._request_states[7] = _RequestCompressionState() + 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 score_staging: + for _ in range(6): + uncompressed += 2 + confirmed += 2 + cache.capacity = confirmed + + manager._periodic_evict(batch) + + state = manager._request_states[7] + if state.confirmed_kv_length < confirmed: + # An eviction round compacted the cache to prompt + budget. + eviction_rounds += 1 + confirmed = state.confirmed_kv_length + 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 request.py_num_compressed_tokens == uncompressed - confirmed + assert request.py_num_compressed_tokens >= previous_published + previous_published = request.py_num_compressed_tokens + + assert eviction_rounds == 3 + assert previous_published == 12 + # Each round the draft cache shrinks with the target and both manager + # streams are ordered after the compact launches. + assert draft_cache.resize.call_args_list == [mock.call(7, None)] * eviction_rounds + assert score_staging.mark_page_tables_consumed.call_args_list == ( + [mock.call(target._stream, draft_manager._stream)] * eviction_rounds + ) + + +def test_lru_score_staging_eviction_drops_dependent_batched_compactions(): + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module + + manager = _make_triattention(top_B=4) + manager._H = 2 + manager._F = 2 + manager._freq_scale_sq = torch.ones(2) + manager._offsets = torch.ones(2) + manager.calibration = {"omega": torch.ones(2)} + manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) + manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) + draft_manager = _make_fake_v2(is_draft=True) + draft_manager.num_pools = 1 + manager.draft_kv_cache_manager = draft_manager + manager._draft_runtime_kv_layout = mock.Mock( + return_value=SimpleNamespace( + layer_pools=[], + pool_representatives=(), + layer_pool_keys=(), + pool_page_counts=(4,), + pool_view_fingerprint=(), + ) + ) + pool = torch.empty(8, 2, 1, 4, 4) + layout = SimpleNamespace( + manager=SimpleNamespace(num_pools=1), + num_layers=2, + global_layers=[0, 1], + layer_pools=[pool, pool], + dense_layers=[0, 1], + swa_layers=[], + storage_groups={0: [0, 1]}, + pool_view_fingerprint=(("fixed",),), + ) + + stale = [SimpleNamespace(score_staging=object(), keep_set_selector=object()) for _ in range(3)] + for index, resources in enumerate(stale): + manager._eviction_buckets[("stale", index)] = resources + oldest_ids = (id(stale[0].score_staging), id(stale[0].keep_set_selector)) + dependent_draft_key = (*oldest_ids, (0,), (1,)) + dependent_plain_key = (*oldest_ids, (2,), None) + surviving_key = ( + id(stale[1].score_staging), + id(stale[1].keep_set_selector), + (0,), + (1,), + ) + manager._batched_compactions[dependent_draft_key] = object() + manager._batched_compactions[dependent_plain_key] = object() + manager._batched_compactions[surviving_key] = object() + + score_staging = SimpleNamespace( + fused_group=SimpleNamespace(output=torch.empty(1, 4, 8)), + bind_score_launcher=mock.Mock(), + ) + keep_set_selector = SimpleNamespace(valid_widths=torch.empty(1, dtype=torch.int32)) + prepared = [ + _PreparedEviction( + request=_make_request(7), + request_id=7, + seq_len=8, + round_start=8, + expected_keep_count=4, + protected_tail=0, + ) + ] + + with ( + mock.patch.object( + module, + "_FixedScoreStagingBuffers", + return_value=score_staging, + ) as score_cls, + mock.patch.object( + manager, + "_build_cross_request_keep_set_selector", + return_value=keep_set_selector, + ), + ): + resources = manager._eager_resources_for(layout, prepared) + + assert resources.score_staging is score_staging + # The new bucket carries a draft page-table plane sized for the draft tail. + assert score_cls.call_args.kwargs["draft_page_table_token_capacity"] == 8 + 1 + # The oldest score staging fell out of the LRU, taking every dependent + # compaction staging (draft-carrying included) with it. + assert ("stale", 0) not in manager._eviction_buckets + assert dependent_draft_key not in manager._batched_compactions + assert dependent_plain_key not in manager._batched_compactions + assert surviving_key in manager._batched_compactions + assert ("stale", 1) in manager._eviction_buckets + assert ("stale", 2) in manager._eviction_buckets diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py new file mode 100644 index 000000000000..7a9508c982d2 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -0,0 +1,1015 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( + BatchedKVCacheCompaction, +) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _BatchedPerHeadKeepSetSelector, + _BatchedUnionKeepSetSelector, +) + + +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" + ) + + +def _encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: + """Build the native V2 [pool, request, K/V, block] layout.""" + encoded = torch.empty( + page_ids.shape[0], + page_ids.shape[1], + 2, + page_ids.shape[2], + dtype=torch.int32, + device=page_ids.device, + ) + encoded[:, :, 0] = page_ids.to(torch.int32) * 2 + encoded[:, :, 1] = encoded[:, :, 0] + 1 + return encoded + + +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 _fake_cute_topk(scores, seq_lens, output, top_k, next_n): + assert next_n == 1 + for row_index, row in enumerate(scores): + output[row_index].copy_(_stable_topk(row, int(seq_lens[row_index]), int(top_k))) + + +class _AdversarialTieTopK: + """Prefer high-index boundary ties so finalization must correct membership.""" + + def __init__(self): + self.calls = 0 + + def __call__(self, scores, seq_lens, output, top_k, next_n): + assert next_n == 1 + self.calls += 1 + for row_index, row in enumerate(scores): + width = int(seq_lens[row_index]) + values = row[:width] + threshold = torch.sort(values, descending=True).values[int(top_k) - 1] + higher = torch.nonzero(values > threshold, as_tuple=False).flatten() + tied = torch.nonzero(values == threshold, as_tuple=False).flatten() + tied = tied.flip(0) + remaining = int(top_k) - int(higher.numel()) + output[row_index].copy_(torch.cat((higher, tied[:remaining])).to(torch.int32)) + + +def _legacy_union(scores: torch.Tensor, keep_count: int) -> torch.Tensor: + row_top = { + int(index) + for row in scores + for index in _stable_topk(row, row.numel(), keep_count).tolist() + } + combined = scores.max(dim=0).values + ordered = sorted( + row_top, + key=lambda index: (-float(combined[index]), index), + ) + return torch.tensor(ordered[:keep_count], dtype=torch.long) + + +@pytest.mark.parametrize("rows,width,keep_count", [(2, 8, 4), (5, 17, 7)]) +def test_direct_union_topk_matches_legacy_union_with_heavy_ties(rows, width, keep_count): + for seed in range(25): + generator = torch.Generator().manual_seed(seed) + scores = torch.randint( + -2, + 3, + (rows, width), + generator=generator, + dtype=torch.int32, + ).to(torch.float32) + combined = scores.max(dim=0).values + direct = _stable_topk(combined, width, keep_count).to(torch.long) + assert torch.equal(direct, _legacy_union(scores, keep_count)) + + +@pytest.mark.parametrize("keep_count,width", [(4, 8), (4096, 4224), (8192, 9216)]) +def test_union_eager_uses_one_deterministic_cute_selection(keep_count, width): + prompt_len = 17 + generator = torch.Generator().manual_seed(keep_count) + scores = torch.randint( + -8, + 9, + (2, width), + generator=generator, + dtype=torch.int32, + ).to(torch.float32) + selector = _BatchedUnionKeepSetSelector( + rows=2, + width=width, + keep_count=keep_count, + prompt_len=prompt_len, + dtype=torch.float32, + device=torch.device("cpu"), + max_requests=1, + ) + raw_topk = _AdversarialTieTopK() + with ( + mock.patch.object( + torch.ops.trtllm, + "cute_dsl_indexer_topk_decode", + side_effect=raw_topk, + create=True, + ), + mock.patch.object( + torch.ops.trtllm, + "indexer_topk_decode", + side_effect=AssertionError("legacy selector was called"), + create=True, + ), + ): + selector.select_requests(scores.unsqueeze(0), normalize_scores=False) + + expected = _stable_topk(scores.max(dim=0).values, width, keep_count) + assert torch.equal( + selector.keep[0, prompt_len:], + torch.sort(expected + prompt_len).values, + ) + assert raw_topk.calls == 1 + + +@pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) +def test_per_head_eager_keeps_stable_indices(eviction_mode): + selector = _BatchedPerHeadKeepSetSelector( + eviction_mode=eviction_mode, + dense_layers=(0, 1), + num_query_heads=4, + num_kv_heads=2, + width=16, + keep_count=5, + prompt_len=3, + dtype=torch.float32, + device=torch.device("cpu"), + max_requests=1, + ) + scores = torch.arange(2 * 4 * 16, dtype=torch.float32).reshape(2, 4, 16) + with mock.patch.object( + torch.ops.trtllm, + "cute_dsl_indexer_topk_decode", + side_effect=_fake_cute_topk, + create=True, + ): + selector.select_requests(scores.unsqueeze(0), normalize_scores=False) + assert tuple(selector.keep.shape) == ( + 1, + selector.selection_rows, + selector.total_keep, + ) + assert torch.all(selector.keep[..., 1:] >= selector.keep[..., :-1]) + + +@pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) +@pytest.mark.parametrize("normalize_scores", [False, True]) +def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( + eviction_mode, normalize_scores +): + request_count, layers, query_heads, kv_heads = 2, 3, 4, 2 + width, keep_count, prompt_len = 96, 64, 0 + 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) + valid_widths = torch.tensor([83, 91], dtype=torch.int32) + + reference = _BatchedPerHeadKeepSetSelector( + eviction_mode=eviction_mode, + dense_layers=tuple(range(layers)), + num_query_heads=query_heads, + num_kv_heads=kv_heads, + width=width, + keep_count=keep_count, + prompt_len=prompt_len, + dtype=torch.float32, + device=torch.device("cpu"), + max_requests=request_count, + ) + reference.valid_widths.copy_(valid_widths) + with mock.patch.object( + torch.ops.trtllm, + "cute_dsl_indexer_topk_decode", + side_effect=_fake_cute_topk, + create=True, + ): + reference.select_requests(scores_cpu, normalize_scores=normalize_scores) + expected = reference.keep.clone() + + device = torch.device("cuda") + stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(stream): + selector = _BatchedPerHeadKeepSetSelector( + eviction_mode=eviction_mode, + dense_layers=tuple(range(layers)), + num_query_heads=query_heads, + num_kv_heads=kv_heads, + width=width, + keep_count=keep_count, + prompt_len=prompt_len, + dtype=torch.float32, + device=device, + max_requests=request_count, + ) + selector.valid_widths.copy_(valid_widths.to(device)) + scores = scores_cpu.to(device) + selector.select_requests(scores, normalize_scores=normalize_scores) + first = selector.keep.cpu() + selector.select_requests(scores, normalize_scores=normalize_scores) + second = selector.keep.cpu() + stream.synchronize() + + assert torch.equal(first, expected) + assert torch.equal(second, expected) + + +def test_union_eager_runs_the_registered_cute_op(): + _require_cute_topk_op() + device = torch.device("cuda") + scores = torch.randn(2, 4, 96, dtype=torch.float32, device=device) + selector = _BatchedUnionKeepSetSelector( + rows=4, + width=96, + keep_count=64, + prompt_len=0, + dtype=torch.float32, + device=device, + max_requests=2, + input_scores=scores, + normalize_scores=True, + ) + selector.select_prepared_requests() + torch.cuda.synchronize(device) + assert torch.all(selector.keep[:, 1:] >= selector.keep[:, :-1]) + + +@pytest.mark.parametrize("normalize_scores", [False, True]) +def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_scores): + _require_cute_topk_op() + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels + + device = torch.device("cuda") + request_count, rows, width, keep_count = 2, 7, 97, 64 + generator = torch.Generator(device=device).manual_seed(53) + scores = torch.randn( + request_count, + rows, + width, + generator=generator, + dtype=torch.float32, + device=device, + ) + valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) + reference_mean = torch.empty(request_count, rows, 1, dtype=torch.float32, device=device) + reference_inv_std = torch.empty_like(reference_mean) + reference_combined = torch.empty(request_count, width, dtype=torch.float32, device=device) + triattention_kernels.prepare_union_scores( + scores, + valid_widths, + reference_mean, + reference_inv_std, + reference_combined, + request_count, + normalize_scores=normalize_scores, + ) + + selector = _BatchedUnionKeepSetSelector( + rows=rows, + width=width, + keep_count=keep_count, + prompt_len=0, + dtype=torch.float32, + device=device, + max_requests=request_count, + input_scores=scores, + normalize_scores=normalize_scores, + ) + selector.valid_widths.copy_(valid_widths) + with mock.patch.object( + triattention_kernels, + "prepare_union_scores", + side_effect=AssertionError("checked Triton wrapper was called"), + ): + selector.select_prepared_requests() + actual_combined = selector.combined.cpu() + actual_keep = selector.keep.cpu() + expected_combined = reference_combined.cpu() + torch.cuda.synchronize(device) + + assert torch.equal(actual_combined, expected_combined) + for request, valid_width in enumerate(valid_widths.cpu().tolist()): + expected_keep = torch.sort( + _stable_topk(expected_combined[request], valid_width, keep_count) + ).values + assert torch.equal(actual_keep[request], expected_keep) + + +@pytest.mark.parametrize("keep_count,width", [(4096, 4224), (8192, 9216)]) +def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): + _require_cute_topk_op() + device = torch.device("cuda") + 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) + valid_widths = (width, width - 32) + selector = _BatchedUnionKeepSetSelector( + rows=rows, + width=width, + keep_count=keep_count, + prompt_len=prompt_len, + dtype=torch.float32, + device=device, + max_requests=request_count, + input_scores=scores, + normalize_scores=False, + ) + selector.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) + selector.select_prepared_requests() + actual = selector.keep.cpu() + + expected_prompt = torch.arange(prompt_len, dtype=torch.int32) + combined = scores.amax(dim=1).cpu() + for request, valid_width in enumerate(valid_widths): + expected_decode = torch.sort( + _stable_topk(combined[request], valid_width, keep_count).to(torch.int32) + prompt_len + ).values + assert torch.equal(actual[request, :prompt_len], expected_prompt) + assert torch.equal(actual[request, prompt_len:], expected_decode) + + +def test_fused_union_preparation_matches_ragged_torch_reference(): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + prepare_union_scores, + ) + + device = torch.device("cuda") + request_count, rows, width = 2, 7, 97 + generator = torch.Generator(device=device).manual_seed(17) + scores = torch.randn( + request_count, + rows, + width, + generator=generator, + dtype=torch.float32, + device=device, + ) + valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) + row_mean = torch.empty(request_count, rows, 1, dtype=torch.float32, device=device) + row_inv_std = torch.empty_like(row_mean) + combined = torch.empty(request_count, width, device=device) + + prepare_union_scores( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + request_count, + normalize_scores=True, + ) + torch.cuda.synchronize(device) + + expected = torch.full_like(combined, float("-inf")) + for request, valid_width in enumerate(valid_widths.tolist()): + valid_scores = scores[request, :, :valid_width] + mean = valid_scores.mean(dim=1, keepdim=True) + std = torch.linalg.vector_norm(valid_scores - mean, dim=1, keepdim=True) + std = (std / valid_width**0.5).clamp_min(1e-6) + expected[request, :valid_width] = ((valid_scores - mean) / std).amax(dim=0) + assert torch.allclose(combined, expected, rtol=2e-5, atol=2e-5) + + +@pytest.mark.parametrize("per_layer", [False, True]) +@pytest.mark.parametrize("normalize_scores", [False, True]) +def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, normalize_scores): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + prepare_per_head_scores, + ) + + device = torch.device("cuda") + 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, + ) + valid_widths = 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 + selection_scores = torch.empty( + request_count, selection_rows, width, dtype=torch.float32, device=device + ) + selection_seq_lens = torch.empty( + request_count, selection_rows, dtype=torch.int32, device=device + ) + + prepare_per_head_scores( + scores, + valid_widths, + row_mean, + row_inv_std, + selection_scores, + selection_seq_lens, + request_count, + num_kv_heads=kv_heads, + per_layer=per_layer, + normalize_scores=normalize_scores, + ) + torch.cuda.synchronize(device) + + assert torch.equal( + selection_seq_lens.cpu(), + valid_widths.cpu().view(request_count, 1).expand(-1, selection_rows), + ) + query_group_size = query_heads // kv_heads + for request, valid_width in enumerate(valid_widths.tolist()): + valid = scores[request, :, :, :valid_width] + if normalize_scores: + mean = valid.mean(dim=-1, keepdim=True) + std = torch.linalg.vector_norm(valid - mean, dim=-1, keepdim=True) + std = (std / valid_width**0.5).clamp_min(1e-6) + valid = (valid - mean) / std + grouped = valid.view(layers, kv_heads, query_group_size, valid_width).amax(dim=2) + expected = grouped if per_layer else grouped.mean(dim=0) + expected = expected.reshape(selection_rows, valid_width) + assert torch.allclose( + selection_scores[request, :, :valid_width], + expected, + rtol=2e-5, + atol=2e-5, + ) + assert torch.isneginf(selection_scores[request, :, valid_width:]).all() + + +@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) +def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode): + device = torch.device("cuda") + request_count = 2 + num_layers = 2 + num_kv_heads = 2 + prompt_len = 2 + decode_keep_count = 4 + seq_len = 10 + tokens_per_block = 4 + pages_per_request = 3 + head_dim = 16 + protected_tails = [2, 1] + page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + initial_pools = [ + ( + torch.arange( + 6 * 2 * num_kv_heads * tokens_per_block * head_dim, + dtype=torch.float32, + device=device, + ).view(6, 2, num_kv_heads, tokens_per_block, head_dim) + + layer * 100_000.0 + ) + for layer in range(num_layers) + ] + pools = [pool.clone() for pool in initial_pools] + + prompt = torch.tensor([0, 1], dtype=torch.int64, device=device) + union_decode = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) + if eviction_mode == "union": + keep = torch.cat((prompt.view(1, -1).expand(request_count, -1), union_decode), dim=1) + selection_rows = 1 + else: + selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads + decode = torch.empty( + request_count, + selection_rows, + decode_keep_count, + dtype=torch.int64, + device=device, + ) + for request in range(request_count): + for row in range(selection_rows): + decode[request, row] = torch.tensor( + sorted( + { + 2 + ((request + row + offset * 2) % 8) + for offset in range(decode_keep_count) + } + ), + dtype=torch.int64, + device=device, + ) + keep = torch.cat( + ( + prompt.view(1, 1, -1).expand(request_count, selection_rows, -1), + decode, + ), + dim=2, + ) + + compaction = BatchedKVCacheCompaction( + eviction_mode=eviction_mode, + layer_pools=pools, + dense_layers=[0, 1], + swa_layers=[], + layer_group_representative={0: 0, 1: 1}, + layer_pool_keys=[("dense", 0), ("dense", 0)], + kept_token_ordinals=keep.to(torch.int32), + valid_sequence_lengths=torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device), + kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), + page_table_slots={0: 0, 1: 0}, + request_count=request_count, + prompt_len=prompt_len, + decode_keep_count=decode_keep_count, + swa_window=None, + protected_tail_lengths=protected_tails, + ) + compaction.launch() + torch.cuda.synchronize(device) + + for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): + for request in range(request_count): + pages = page_tables[request].to(torch.long) + before = ( + before_pool[pages] + .permute(1, 2, 0, 3, 4) + .reshape(2, num_kv_heads, pages_per_request * tokens_per_block, head_dim) + ) + after = after_pool[pages].permute(1, 2, 0, 3, 4).reshape_as(before) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + for head in range(num_kv_heads): + if eviction_mode == "union": + selected = keep[request, prompt_len:] + elif eviction_mode == "per_head": + selected = keep[request, head, prompt_len:] + else: + selected = keep[request, layer * num_kv_heads + head, prompt_len:] + tail = torch.arange( + seq_len, + seq_len + protected_tails[request], + dtype=torch.int64, + device=device, + ) + source = torch.cat((selected, tail)) + destination = torch.arange( + prompt_len, + prompt_len + source.numel(), + dtype=torch.int64, + device=device, + ) + assert torch.equal( + after[:, head].index_select(1, destination), + before[:, head].index_select(1, source), + ) + + +def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): + """Keep score and compaction layer axes aligned across interleaved V2 pools.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + ) + + device = torch.device("cuda") + num_layers = 3 + seq_len = 8 + keep_count = 2 + tokens_per_block = 4 + head_dim = 16 + num_freqs = head_dim // 2 + dense_layers = [0, 1, 2] + dense_groups = [[0, 2], [1]] + layer_group_representative = {0: 0, 1: 1, 2: 0} + 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 = [] + for layer, (table, values) in enumerate(zip(layer_tables, score_values)): + pool = ( + torch.arange( + 2 * 2 * tokens_per_block * head_dim, + dtype=torch.float32, + device=device, + ).view(2, 2, 1, tokens_per_block, head_dim) + + layer * 10_000 + ) + 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 + pools.append(pool) + initial_pools = [pool.clone() for pool in pools] + + q_real = torch.zeros(num_layers, 1, 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 + score_staging = _FixedScoreStagingBuffers( + pools, + dense_groups, + dense_layers, + [0, 1], + 1, + seq_len, + 1, + num_freqs, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + torch.zeros(1, dtype=torch.float32, device=device), + torch.zeros(num_freqs, dtype=torch.float32, device=device), + page_table_keys=[("pool", 0), ("pool", 1)], + num_page_table_slots=2, + ) + score_staging.block_offsets_device.zero_() + score_staging.block_offsets_device[..., :2].copy_( + _encode_block_offsets(torch.stack(page_tables)) + ) + score_staging.round_starts_device.fill_(0) + score_staging.valid_seq_lens_device.fill_(seq_len) + keep_set_selector = _BatchedPerHeadKeepSetSelector( + eviction_mode="per_layer_perhead", + dense_layers=tuple(dense_layers), + num_query_heads=1, + num_kv_heads=1, + width=seq_len, + keep_count=keep_count, + prompt_len=0, + dtype=torch.float32, + device=device, + max_requests=1, + ) + score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") + + scores = score_staging.launch_prepared_score() + keep_set_selector.select_requests(scores, normalize_scores=False) + assert torch.equal(keep_set_selector.keep, expected_keep) + + batched_compaction = BatchedKVCacheCompaction( + eviction_mode="per_layer_perhead", + layer_pools=pools, + dense_layers=dense_layers, + swa_layers=[], + layer_group_representative=layer_group_representative, + layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], + kept_token_ordinals=keep_set_selector.keep[:1], + valid_sequence_lengths=score_staging.valid_seq_lens_device[:1], + kv_block_offsets=score_staging.block_offsets_device, + page_table_slots=score_staging.representative_slots, + request_count=1, + prompt_len=0, + decode_keep_count=keep_count, + swa_window=None, + protected_tail_lengths=[0], + ) + batched_compaction.launch() + torch.cuda.synchronize(device) + + for layer, (before_pool, after_pool, table) in enumerate( + zip(initial_pools, pools, layer_tables) + ): + pages = table[0].to(torch.long) + before = before_pool[pages].permute(1, 2, 0, 3, 4).reshape(2, 1, seq_len, 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)) + + +def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): + """Run two real eviction rounds through one live V2 cache.""" + import tensorrt_llm + import tensorrt_llm.bindings + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + ) + 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") + request_id = 7 + prompt_len = 2 + seq_len = 10 + protected_tail = 2 + compacted_capacity = 8 + tokens_per_block = 4 + head_dim = 16 + num_freqs = head_dim // 2 + 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.HALF, + 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 + payload = ( + torch.arange(2 * head_dim, dtype=torch.float16, device=device) + .reshape(2, head_dim) + .add_(token * 64) + ) + payload[0, 0] = score + payload[0, num_freqs] = 0 + pool[page, :, 0, offset].copy_(payload) + + first_scores = [0, 0, 8, 1, 7, 2, 6, 3, 5, 4, 9, 0] + for token, score in enumerate(first_scores): + write_token(token, score) + + q_real = torch.zeros(1, 1, 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 + score_staging = _FixedScoreStagingBuffers( + [pool], + [[0]], + [0], + [0], + 1, + seq_len, + 1, + num_freqs, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + torch.zeros(1, dtype=torch.float32, device=device), + torch.zeros(num_freqs, dtype=torch.float32, device=device), + page_table_keys=[("pool", 0)], + num_page_table_slots=1, + prompt_len=prompt_len, + page_table_token_capacity=seq_len + protected_tail, + ) + keep_set_selector = _BatchedUnionKeepSetSelector( + rows=1, + width=seq_len - prompt_len, + keep_count=compacted_capacity - prompt_len - protected_tail, + prompt_len=prompt_len, + dtype=torch.float32, + device=device, + max_requests=1, + dense_layers=(0,), + num_query_heads=1, + num_kv_heads=1, + input_scores=score_staging.fused_group.output.view(1, 1, seq_len - prompt_len), + normalize_scores=False, + ) + score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") + batched_compaction = BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=[pool], + dense_layers=[0], + swa_layers=[], + layer_group_representative={0: 0}, + layer_pool_keys=[("pool", 0)], + kept_token_ordinals=keep_set_selector.keep[:1], + valid_sequence_lengths=score_staging.valid_seq_lens_device[:1], + kv_block_offsets=score_staging.block_offsets_device, + page_table_slots=score_staging.representative_slots, + request_count=1, + prompt_len=prompt_len, + decode_keep_count=compacted_capacity - prompt_len - protected_tail, + swa_window=None, + protected_tail_lengths=[protected_tail], + ) + + def evict_once() -> tuple[torch.Tensor, torch.Tensor]: + before = snapshot(seq_len + protected_tail) + assert score_staging.stage( + manager, + [request_id], + [0], + [seq_len], + [seq_len + protected_tail], + ) + score_staging.launch_prepared_score() + keep_set_selector.select_prepared_requests() + selected = keep_set_selector.keep[0].clone().to(torch.long) + batched_compaction.launch() + score_staging.mark_page_tables_consumed(manager._stream) + torch.cuda.synchronize(device) + assert cache.resize(compacted_capacity, None) + after = snapshot(compacted_capacity) + source = torch.cat( + ( + selected[prompt_len:], + 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) + first_keep, first_compacted = evict_once() + assert torch.equal( + first_keep, + torch.tensor([0, 1, 2, 4, 6, 8], dtype=torch.long, device=device), + ) + 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) + # The first protected tail becomes confirmed input to round two. Only + # later generated tokens and the next protected tail are written here. + write_token(8, 10) + write_token(9, 4.5) + write_token(10, 11) + write_token(11, 0.5) + assert torch.equal(snapshot(8), first_compacted) + + second_keep, _ = evict_once() + assert torch.equal( + second_keep, + torch.tensor([0, 1, 2, 3, 6, 8], dtype=torch.long, device=device), + ) + 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_eager_compaction_rebases_masked_swa_window_and_tail(): + device = torch.device("cuda") + dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) + swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) + initial_pools = [ + torch.arange(6 * 2 * 1 * 4 * 16, dtype=torch.float32, device=device).view(6, 2, 1, 4, 16), + torch.arange(6 * 2 * 1 * 4 * 16, dtype=torch.float32, device=device).view(6, 2, 1, 4, 16) + + 1000.0, + ] + pools = [pool.clone() for pool in initial_pools] + keep = torch.tensor( + [[0, 1, 2, 4, 5, 7], [0, 1, 2, 3, 5, 6]], + dtype=torch.int64, + device=device, + ) + valid_seq_lens = torch.tensor([8, 7], dtype=torch.int32, device=device) + protected_tails = [2, 1] + compaction = BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=pools, + dense_layers=[0], + swa_layers=[1], + layer_group_representative={0: 0}, + layer_pool_keys=[("dense", 0), ("swa", 0)], + kept_token_ordinals=keep.to(torch.int32), + valid_sequence_lengths=valid_seq_lens, + kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), + page_table_slots={0: 0, 1: 1}, + request_count=2, + prompt_len=2, + decode_keep_count=4, + swa_window=2, + protected_tail_lengths=protected_tails, + ) + compaction.launch() + torch.cuda.synchronize(device) + + for request, (valid_seq_len, tail_length) in enumerate( + zip(valid_seq_lens.tolist(), protected_tails) + ): + dense_pages = dense_tables[request].to(torch.long) + swa_pages = swa_tables[request].to(torch.long) + dense_before = initial_pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 16) + dense_after = pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape_as(dense_before) + swa_before = initial_pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 16) + swa_after = pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape_as(swa_before) + tail = torch.arange( + valid_seq_len, + valid_seq_len + tail_length, + dtype=torch.int64, + device=device, + ) + dense_source = torch.cat((keep[request, 2:], tail)) + dense_destination = torch.arange( + 2, 2 + dense_source.numel(), dtype=torch.int64, device=device + ) + swa_source = torch.arange( + valid_seq_len - 2, + valid_seq_len + tail_length, + dtype=torch.int64, + device=device, + ) + swa_destination = torch.arange(4, 4 + swa_source.numel(), dtype=torch.int64, device=device) + assert torch.equal(dense_after[:, :, :2], dense_before[:, :, :2]) + assert torch.equal(swa_after[:, :, :2], swa_before[:, :, :2]) + assert torch.equal( + dense_after.index_select(2, dense_destination), + dense_before.index_select(2, dense_source), + ) + assert torch.equal( + swa_after.index_select(2, swa_destination), + swa_before.index_select(2, swa_source), + ) 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..2042cd5a2f91 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -0,0 +1,1839 @@ +# 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. + +TriAttention is a pure KV-cache compression method on the PR-15106 framework: it +has NO sparse-attention config and NO attention backend of its own. Decode runs +the model's standard attention over the compacted cache; the manager publishes +the cumulative evicted count on ``LlmRequest.py_num_compressed_tokens`` and the +model engine subtracts it where it builds ``num_cached_tokens_per_seq``. These +tests cover the config, construction, compressed-count publication, eager +selection, page-table staging, bounded request chunks, and request lifecycle. +Model-level correctness is covered by separate end-to-end tests. +""" + +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch +from pydantic import ValidationError + +# 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 import TriAttention +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _BatchedUnionKeepSetSelector, + _PreparedEviction, + _PreparedGenerationBatch, + _RequestCompressionState, + _RuntimeKVLayout, +) + +# 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 +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm.llmapi.llm_args import ( + KvCacheCompressionConfig, + TriAttentionKvCacheCompressionConfig, +) + +_TORCH_TOPK_ORACLE = torch.topk + + +def _encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: + """Build the native V2 [pool, request, K/V, block] layout.""" + if page_ids.ndim == 2: + page_ids = page_ids.unsqueeze(0) + encoded = torch.empty( + page_ids.shape[0], + page_ids.shape[1], + 2, + page_ids.shape[2], + dtype=torch.int32, + device=page_ids.device, + ) + encoded[:, :, 0] = page_ids.to(torch.int32) * 2 + encoded[:, :, 1] = encoded[:, :, 0] + 1 + return encoded + + +def _set_request_state( + manager, + request_id, + *, + generation_steps=0, + evicted_tokens=0, + confirmed_kv_length=None, +): + state = _RequestCompressionState( + generation_steps=generation_steps, + evicted_tokens=evicted_tokens, + confirmed_kv_length=confirmed_kv_length, + ) + manager._request_states[request_id] = state + return state + + +def _prepared_eviction( + request, + *, + seq_len, + expected_keep_count, + protected_tail=0, + request_id=0, + round_start=None, +): + return _PreparedEviction( + request=request, + request_id=request_id, + seq_len=seq_len, + round_start=int(seq_len if round_start is None else round_start), + expected_keep_count=expected_keep_count, + protected_tail=protected_tail, + ) + + +def _fake_cute_dsl_topk( + values: torch.Tensor, + seq_lens: torch.Tensor, + output: torch.Tensor, + top_k: int, + next_n: int, +) -> None: + """CPU oracle for the CUDA-only CuTE-DSL selector custom op.""" + assert next_n == 1 + for row in range(int(values.shape[0])): + width = int(seq_lens[row]) + selected = _TORCH_TOPK_ORACLE( + values[row, :width], + top_k, + sorted=False, + ).indices + output[row].copy_(selected.to(torch.int32)) + + +@contextmanager +def _mock_cute_topk_without_fallbacks(): + """Provide the CuTE op while making both retired fallbacks fatal.""" + with ( + mock.patch.object( + torch.ops.trtllm, + "cute_dsl_indexer_topk_decode", + side_effect=_fake_cute_dsl_topk, + create=True, + ) as cute_topk, + mock.patch.object( + torch.ops.trtllm, + "indexer_topk_decode", + side_effect=AssertionError("native IndexerTopK fallback is forbidden"), + create=True, + ), + mock.patch.object( + torch, + "topk", + side_effect=AssertionError("torch.topk production fallback is forbidden"), + ), + ): + yield cute_topk + + +def _union_oracle(scores: torch.Tensor, keep_count: int) -> torch.Tensor: + """Independent expected-result implementation of union selection.""" + combined = scores.max(dim=0).values + row_top = _TORCH_TOPK_ORACLE( + scores, + keep_count, + dim=1, + sorted=False, + ).indices + union_mask = torch.zeros(scores.shape[1], dtype=torch.bool, device=scores.device) + union_mask.scatter_(0, row_top.reshape(-1), True) + union_indices = torch.nonzero(union_mask, as_tuple=False).flatten() + if union_indices.numel() >= keep_count: + candidates = combined.index_select(0, union_indices) + relative = _TORCH_TOPK_ORACLE( + candidates, + keep_count, + sorted=False, + ).indices + return torch.sort(union_indices.index_select(0, relative)).values + + remaining = keep_count - int(union_indices.numel()) + residual = combined.clone() + residual[union_mask] = float("-inf") + extra = _TORCH_TOPK_ORACLE( + residual, + remaining, + sorted=False, + ).indices + return torch.sort(torch.cat((union_indices, extra))).values + + +def _distinct_topk_scores(width: int, rows: int = 2) -> torch.Tensor: + """Create deterministic finite rows without top-k boundary ties.""" + token = torch.arange(width, dtype=torch.float32) + return torch.stack( + [ + torch.sin(token * (0.0017 + row * 0.0003)) + + token * (0.00011 + row * 0.000013) + + row * 0.000001 + for row in range(rows) + ] + ) + + +@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_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.mapping = SimpleNamespace(enable_attention_dp=False) + fake_v2.is_disagg = False + 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.empty(1, dtype=torch.int64) + fake_v2.pp_layers = [] + fake_v2.layer_offsets = {} + fake_v2.layer_to_pool_mapping_dict = {} + return fake_v2 + + +def _make_triattention(**overrides): + """Construct a fully initialized manager for method-level unit tests.""" + options = {"top_B": 8, "model_path": "/models/test"} + options.update(overrides) + return TriAttention(_make_fake_v2(), **options) + + +def _make_request(request_id, **overrides): + """Build the explicit request fields consumed by TriAttention.""" + 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) + + +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) + + +def _torch_tri_score_oracle( + layer_pools, + page_ids, + seq_lens, + round_starts, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + layer_indices, + aggregation, +): + """Independent Torch implementation of the paged TriAttention score.""" + scores = [] + num_q_heads = int(q_real.shape[1]) + for request, seq_len in enumerate(seq_lens): + phase = (round_starts[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 + if aggregation == "mean": + position = ( + freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) + ).sum(dim=-1) + else: + position = ( + ( + freq_scale_sq[None, None, :] + * ( + product_real[None] * torch.cos(phase)[:, None, :] + - product_imag[None] * torch.sin(phase)[:, None, :] + ) + ) + .sum(dim=-1) + .max(dim=0) + .values + ) + 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 + + +class TestKvCacheCompressionConfig: + def test_llm_args_dispatches_concrete_and_unknown_algorithms(self): + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + tri_args = TorchLlmArgs( + model="dummy", + kv_cache_compression_config={"algorithm": "triattention"}, + ) + assert isinstance( + tri_args.kv_cache_compression_config, + TriAttentionKvCacheCompressionConfig, + ) + assert tri_args.kv_cache_compression_config.top_B == 2048 + assert tri_args.kv_cache_compression_config.beta == 128 + + unknown_args = TorchLlmArgs( + model="dummy", + kv_cache_compression_config={"algorithm": "future_method"}, + ) + assert type(unknown_args.kv_cache_compression_config) is KvCacheCompressionConfig + + def test_eviction_mode_validated(self): + with pytest.raises(ValidationError): + TriAttentionKvCacheCompressionConfig(eviction_mode="made_up_mode") + + +class TestTriAttentionClass: + def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): + page_count_query = mock.Mock(side_effect=[8, 16, 8, 18]) + manager = SimpleNamespace( + get_buffers=mock.Mock(side_effect=AssertionError("pool view was rebuilt")), + impl=SimpleNamespace(get_page_index_upper_bound=page_count_query), + kv_factor=2, + layer_offsets={10: 100, 11: 101, 12: 102}, + ) + triattention = _make_triattention() + triattention.kv_cache_manager = manager + cached = _RuntimeKVLayout( + manager=manager, + num_layers=3, + global_layers=[10, 11, 12], + layer_pools=[torch.empty(4), torch.empty(8), torch.empty(4)], + dense_layers=[0, 1, 2], + swa_layers=[], + swa_window=None, + storage_groups={0: [0, 2], 1: [1]}, + layer_group_representative={0: 0, 1: 1, 2: 0}, + layer_pool_keys=(0, 1, 0), + # These are local layer slots. Layer 2 shares layer 0's pool. + pool_representatives=(0, 1), + pool_page_counts=(4, 8), + pool_view_fingerprint=(), + ) + triattention._runtime_kv_layout_cache = cached + + assert triattention._runtime_kv_layout(3) is cached + manager.get_buffers.assert_not_called() + assert page_count_query.call_args_list == [ + mock.call(100, Role.KEY), + mock.call(101, Role.KEY), + ] + + with pytest.raises(RuntimeError, match="pool layout changed"): + triattention._runtime_kv_layout(3) + manager.get_buffers.assert_not_called() + assert page_count_query.call_args_list == [ + mock.call(100, Role.KEY), + mock.call(101, Role.KEY), + mock.call(100, Role.KEY), + mock.call(101, Role.KEY), + ] + + def test_triattention_enables_capacity_only_on_target_manager(self): + manager = _make_fake_v2() + triattention = TriAttention(manager, top_B=8, model_path="/models/test") + triattention._attention_layer_partition_cache = ([], [], None) + triattention._calibrated = True + first = _make_request(11) + second = _make_request(12) + + triattention.on_request_init(first) + triattention.on_request_init(second) + + assert triattention.adjusts_generation_kv_length is True + assert manager.kv_compression_manages_history + assert set(triattention._request_states) == {11, 12} + + def test_request_init_accepts_speculative_capacity(self): + manager = _make_fake_v2() + manager.num_extra_kv_tokens = 4 + manager._kv_reserve_draft_tokens = 4 + triattention = TriAttention(manager, top_B=8, model_path="/models/test") + triattention._attention_layer_partition_cache = ([], [], None) + triattention._calibrated = True + request = _make_request(11) + + triattention.on_request_init(request) + + assert manager.kv_compression_manages_history + assert set(triattention._request_states) == {11} + + def test_resolve_accepts_flat_pt(self, flat_calibration_pt): + mgr = _make_triattention() + mgr.calibration_path = flat_calibration_pt + mgr.model_path = None + loaded = mgr._resolve_calibration() + for key in ("E_q", "E_q_norm", "omega", "freq_scale_sq"): + assert key in loaded + + +# --------------------------------------------------------------------------- +# Eviction publishes the cumulative evicted count on the request; the model +# engine reads it back where it builds num_cached_tokens_per_seq. +# --------------------------------------------------------------------------- + + +class TestCompressedTokenPublication: + def test_manager_is_marked_capacity_only_and_requests_default_to_zero(self): + mgr = _make_triattention() + manager = mgr.kv_cache_manager + # The compression-manager base marks the target manager so V2 sizing + # keeps logical max_seq_len while capacity is reclaimed and reused. + assert manager.kv_compression_manages_history + request = _make_request(7) + # Default 0 keeps the engine's num_cached subtraction a no-op until + # the first eviction publishes a count. + assert request.py_num_compressed_tokens == 0 + + @contextmanager + def _mocked_eviction_internals(self, manager): + """Run the real ``_evict_requests`` body around mocked GPU launches.""" + score_staging = SimpleNamespace( + launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), + mark_page_tables_consumed=mock.Mock(), + ) + keep_set_selector = SimpleNamespace(select_requests=mock.Mock()) + resources = SimpleNamespace( + score_staging=score_staging, + keep_set_selector=keep_set_selector, + ) + batched_compaction = SimpleNamespace(launch=mock.Mock()) + with ( + mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), + mock.patch.object(manager, "_eager_resources_for", return_value=resources), + mock.patch.object( + manager, + "_batched_compaction_for", + return_value=batched_compaction, + ), + mock.patch.object(manager, "_attach_page_ids") as attach, + ): + yield SimpleNamespace( + score_staging=score_staging, + keep_set_selector=keep_set_selector, + batched_compaction=batched_compaction, + attach=attach, + ) + + def test_eviction_bookkeeping_publishes_cumulative_count(self): + # The eviction bookkeeping writes the cumulative evicted count on the + # request in the same step that compacts the cache; this is the + # channel's only producer. + manager = _make_triattention(top_B=4) + manager.kv_cache_manager._stream = mock.Mock() + request = _make_request(7, py_prompt_len=2) + _set_request_state(manager, 7, confirmed_kv_length=10) + + with self._mocked_eviction_internals(manager) as internals: + first = manager._evict_requests([(request, 7)], 2) + + assert first == [(7, 6)] + # 10 confirmed - (2 pinned prompt + 4 decode budget) = 4 evicted. + assert request.py_num_compressed_tokens == 4 + assert manager._request_states[7].confirmed_kv_length == 6 + internals.batched_compaction.launch.assert_called_once_with() + internals.score_staging.mark_page_tables_consumed.assert_called_once_with( + manager.kv_cache_manager._stream + ) + + # Round two: 6 retained + 8 newly confirmed decode tokens. + manager._request_states[7].confirmed_kv_length = 14 + with self._mocked_eviction_internals(manager) as internals: + second = manager._evict_requests([(request, 7)], 2) + + assert second == [(7, 6)] + # The count is cumulative and never decreases: 4 + (14 - 6) = 12. + assert request.py_num_compressed_tokens == 12 + # The staged logical position restores the uncompressed length. + prepared = internals.attach.call_args.args[0] + assert prepared[0].round_start == 14 + 4 + + def test_identity_compaction_is_rejected_instead_of_published(self): + manager = _make_triattention(top_B=4) + manager.kv_cache_manager._stream = mock.Mock() + request = _make_request(7, py_prompt_len=2) + # Selection keeps every token: seq_len == prompt + budget + 1 evicts + # one token; seq_len == prompt + budget must never publish. + _set_request_state(manager, 7, confirmed_kv_length=6) + + with self._mocked_eviction_internals(manager): + assert manager._evict_requests([(request, 7)], 2) == [] + assert request.py_num_compressed_tokens == 0 + + +class TestStepEndHookRefactor: + def test_triattention_prepare_only_snapshots_and_update_uses_final_hook(self): + assert "prepare_resources" in TriAttention.__dict__ + assert "update_resources" not in TriAttention.__dict__ + assert "on_generation_step_end" in TriAttention.__dict__ + + def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): + manager = _make_triattention() + request = _make_request(7) + batch = SimpleNamespace( + context_requests=[], + context_requests_last_chunk=[], + generation_requests=[request], + ) + + with mock.patch.object(manager, "_periodic_evict") as periodic_evict: + manager.prepare_resources(batch) + periodic_evict.assert_not_called() + + manager.update_resources(batch) + + periodic_evict.assert_called_once_with(batch) + + @pytest.mark.parametrize("top_B", [511, 512]) + def test_non_v2_manager_is_always_rejected(self, top_B): + with pytest.raises(TypeError, match="requires KVCacheManagerV2"): + TriAttention(SimpleNamespace(), top_B=top_B) + + @staticmethod + def _make_due_decode_request(seq_len): + request = _make_request( + 7, + py_prompt_len=1024, + max_beam_num_tokens=seq_len + 1, + ) + batch = SimpleNamespace(generation_requests=[request]) + mgr = _make_triattention() + mgr._calibrated = True + 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=0, + ) + mgr._L = 2 + mgr._request_states = {} + _set_request_state(mgr, 7, generation_steps=127) + mgr.beta = 128 + mgr.top_B = 4096 + mgr.pin_prefill = True + mgr.count_prompt_tokens = False + return mgr, request, batch + + def test_identity_gate_preserves_real_eviction_round(self): + import contextlib + + import tensorrt_llm._torch.kv_cache_compression.triattention.triattention as tri_module + + mgr, request, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) + timeline = [] + cache = mgr.kv_cache_manager.kv_cache_map[7] + + def compact(*args, protected_tail_lengths, **_kwargs): + assert protected_tail_lengths == {7: 0} + timeline.append("compact_dispatch") + return [(7, 1024 + 4096)] + + @contextlib.contextmanager + def track_range(name, **kwargs): + timeline.append(f"enter:{name}") + yield + timeline.append(f"exit:{name}") + + with ( + mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict, + mock.patch.object(tri_module, "nvtx_range", side_effect=track_range), + ): + mgr._periodic_evict(batch) + + evict.assert_called_once_with( + [(request, 7)], + 2, + protected_tail_lengths={7: 0}, + ) + mgr.kv_cache_manager._stream.wait_event.assert_not_called() + cache.resize.assert_called_once_with(1024 + 4096, None) + assert timeline == [ + "compact_dispatch", + "enter:triattention.resize", + "exit:triattention.resize", + ] + + def test_eager_eviction_chunks_large_due_cohort(self): + manager, _, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) + requests = [] + caches = {} + for request_id in range(65): + request = _make_request(request_id, py_prompt_len=1024) + requests.append(request) + caches[request_id] = SimpleNamespace( + capacity=1024 + 4096 + 1, + history_length=1024, + is_active=True, + ) + _set_request_state(manager, request_id, generation_steps=127) + manager.kv_cache_manager.kv_cache_map = caches + batch = SimpleNamespace(generation_requests=requests) + + with ( + mock.patch.object(manager, "_evict_requests", return_value=[]) as evict, + mock.patch.object(manager, "_resize_compacted_requests") as resize, + ): + manager._periodic_evict(batch) + + assert [len(call.args[0]) for call in evict.call_args_list] == [32, 32, 1] + assert resize.call_count == 3 + + def test_last_request_finish_releases_eager_buffers(self): + manager = _make_triattention() + request = _make_request(7) + _set_request_state(manager, 7) + manager._eviction_buckets[("score",)] = object() + manager._batched_compactions[("compact",)] = object() + + manager.on_request_finish(request) + + assert manager._eviction_buckets == {} + assert manager._batched_compactions == {} + + @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) + def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): + confirmed = 1024 + 4096 + 1 + accepted + reserve = 2 + current_growth = 4 + tail = reserve + current_growth + retained = 1024 + 4096 + mgr, request, batch = self._make_due_decode_request(seq_len=confirmed) + request.py_num_accepted_draft_tokens = accepted + cache = mgr.kv_cache_manager.kv_cache_map[7] + cache.capacity = confirmed + tail + mgr.kv_cache_manager.num_extra_kv_tokens = reserve + mgr._prepared_generation_batch = _PreparedGenerationBatch( + batch=SimpleNamespace(generation_requests=[request]), + growth_by_request={7: current_growth}, + ) + 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 + + def compact(*_args, **_kwargs): + mgr._request_states[7].confirmed_kv_length = retained + return [(7, retained)] + + with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: + mgr._periodic_evict(batch) + + evict.assert_called_once_with( + [(request, 7)], + 2, + protected_tail_lengths={7: tail}, + ) + assert mgr._request_states[7].confirmed_kv_length == retained + cache.resize.assert_called_once_with(retained + tail, None) + # The draft cache shrinks in the same round, to the same retained + # length plus the draft's own protected tail. + draft_cache.resize.assert_called_once_with(retained + 1, None) + + def test_missing_draft_cache_fails_the_due_eviction_round(self): + mgr, request, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) + mgr.draft_kv_cache_manager = _make_fake_v2(is_draft=True) + + with mock.patch.object(mgr, "_evict_requests") as evict: + with pytest.raises(RuntimeError, match="missing or.*suspended draft KV cache"): + mgr._periodic_evict(batch) + evict.assert_not_called() + + def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): + physical_confirmed = 6100 + manager = _make_triattention(beta=128) + manager._calibrated = True + _set_request_state(manager, 7, evicted_tokens=100) + 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] + manager.kv_cache_manager.num_extra_kv_tokens = 0 + request = _make_request( + 7, + py_prompt_len=1024, + max_beam_num_tokens=999999, + py_draft_tokens=[1, 2, 3, 4], + ) + + manager._periodic_evict(SimpleNamespace(generation_requests=[request])) + + assert manager._request_states[7].confirmed_kv_length == physical_confirmed + cache.resize.assert_not_called() + + def test_mla_selfkonly_cache_is_rejected(self): + manager = _make_triattention() + manager.kv_cache_manager.kv_factor = 1 + + with pytest.raises(ValueError, match="standard key/value KV cache"): + manager._validate_v2_compatibility() + + def test_one_model_mtp_co_compression_contract_is_accepted(self): + draft_manager = _make_fake_v2(is_draft=True) + manager = TriAttention( + _make_fake_v2(), + top_B=8, + model_path="/models/test", + draft_kv_cache_manager=draft_manager, + ) + + manager._validate_v2_compatibility() + assert manager.kv_cache_manager.kv_compression_manages_history is True + # The draft cache is compacted together with the target, so its + # physical length diverges from the logical length the same way. + assert draft_manager.kv_compression_manages_history is True + + def test_draft_co_compression_accepts_smaller_draft_max_seq_len(self): + # Co-compression keeps the draft's physical length equal to the + # target's, so the draft does not have to cover the target's logical + # maximum sequence length. + draft_manager = _make_fake_v2(is_draft=True) + draft_manager.max_seq_len = 8192 + manager = TriAttention( + _make_fake_v2(), + top_B=8, + model_path="/models/test", + draft_kv_cache_manager=draft_manager, + ) + + manager._validate_v2_compatibility() + + @pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) + def test_draft_co_compression_requires_union_mode(self, eviction_mode): + manager = TriAttention( + _make_fake_v2(), + top_B=8, + model_path="/models/test", + eviction_mode=eviction_mode, + draft_kv_cache_manager=_make_fake_v2(is_draft=True), + ) + + with pytest.raises(ValueError, match="union"): + manager._validate_v2_compatibility() + + def test_draft_co_compression_rejects_mla_draft_cache(self): + draft_manager = _make_fake_v2(is_draft=True) + manager = TriAttention( + _make_fake_v2(), + top_B=8, + model_path="/models/test", + draft_kv_cache_manager=draft_manager, + ) + # The base class already rejects a mismatched draft kv_factor at + # construction; this guards against later runtime divergence too. + draft_manager.kv_factor = 1 + + with pytest.raises(ValueError, match="standard key/value cache"): + manager._validate_v2_compatibility() + + def test_draft_co_compression_requires_full_attention_draft(self): + draft_manager = _make_fake_v2(is_draft=True) + draft_manager.max_attention_window_vec = [128] + manager = TriAttention( + _make_fake_v2(), + top_B=8, + model_path="/models/test", + draft_kv_cache_manager=draft_manager, + ) + + with pytest.raises(ValueError, match="full-attention draft"): + manager._validate_v2_compatibility() + + def test_resize_shrinks_draft_cache_with_its_own_protected_tail(self): + retained = 1024 + 4096 + manager = _make_triattention() + target_cache = SimpleNamespace( + capacity=retained + 10, + is_active=True, + resize=mock.Mock(return_value=True), + ) + manager.kv_cache_manager = SimpleNamespace( + kv_cache_map={7: target_cache}, + ) + draft_manager = _make_fake_v2(is_draft=True) + draft_manager.num_extra_kv_tokens = 2 + draft_manager._kv_reserve_draft_tokens = 3 + draft_cache = SimpleNamespace(is_active=True, resize=mock.Mock(return_value=True)) + draft_manager.kv_cache_map = {7: draft_cache} + manager.draft_kv_cache_manager = draft_manager + + manager._resize_compacted_requests([(7, retained)], {7: 4}) + + target_cache.resize.assert_called_once_with(retained + 4, None) + # Draft protected tail = num_extra_kv_tokens + reserved draft width + 1. + draft_cache.resize.assert_called_once_with(retained + 6, None) + + def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): + # The factory runs the speculative feature gates and then builds the + # manager; a one-model MTP contract must pass both. + from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + from tensorrt_llm.llmapi.llm_args import ( + MTPDecodingConfig, + TriAttentionKvCacheCompressionConfig, + ) + + manager = create_kv_cache_compression_manager( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + _make_fake_v2(), + draft_kv_cache_manager=_make_fake_v2(is_draft=True), + spec_config=MTPDecodingConfig(max_draft_len=1), + ) + + manager._validate_v2_compatibility() + + @pytest.mark.parametrize("mode", ["draft_target", "pard"]) + def test_unvalidated_paged_draft_tail_contracts_remain_fail_closed(self, mode): + from tensorrt_llm.llmapi.llm_args import DraftTargetDecodingConfig, PARDDecodingConfig + + if mode == "draft_target": + spec_config = DraftTargetDecodingConfig( + max_draft_len=3, + speculative_model="/tmp/draft-target-model", + ) + else: + spec_config = PARDDecodingConfig(max_draft_len=3) + # The factory's eviction-method speculative-mode gate declines to + # create a manager; the run stays uncompressed. + from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig + + manager = create_kv_cache_compression_manager( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + _make_fake_v2(), + draft_kv_cache_manager=_make_fake_v2(is_draft=True), + spec_config=spec_config, + ) + assert manager is None + + def test_dflash_spec_mode_is_rejected(self): + # Policy: the DFlash draft reads cross-attention context buffers, not + # a paged KV cache, so compression cannot cover it. The factory's + # eviction-method speculative-mode gate declines to create a manager. + from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + from tensorrt_llm.llmapi.llm_args import ( + DFlashDecodingConfig, + TriAttentionKvCacheCompressionConfig, + ) + + manager = create_kv_cache_compression_manager( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + _make_fake_v2(), + draft_kv_cache_manager=_make_fake_v2(is_draft=True), + spec_config=DFlashDecodingConfig(max_draft_len=3), + ) + assert manager is None + + def test_prepare_snapshots_fixed_linear_generation_growth(self): + manager = _make_fake_v2() + manager.num_extra_kv_tokens = 2 + manager.kv_cache_map = { + 7: SimpleNamespace(capacity=106, is_active=True), + } + triattention = TriAttention(manager, top_B=8, model_path="/models/test") + batch = SimpleNamespace( + context_requests=[], + generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], + ) + + triattention.prepare_resources(batch) + + assert triattention._prepared_generation_batch.batch is batch + assert triattention._prepared_generation_batch.growth_by_request == {7: 4} + + def test_prepare_protects_reserved_draft_width(self): + manager = _make_fake_v2() + manager._kv_reserve_draft_tokens = 6 + manager.kv_cache_map = { + 7: SimpleNamespace(capacity=106, is_active=True), + } + triattention = TriAttention(manager, top_B=8, model_path="/models/test") + batch = SimpleNamespace( + context_requests=[], + generation_requests=[_make_request(7, py_draft_tokens=[1, 2])], + ) + + triattention.prepare_resources(batch) + + assert triattention._prepared_generation_batch.growth_by_request == {7: 7} + + def test_request_finish_clears_compression_state(self): + request = SimpleNamespace(py_request_id=7) + mgr = _make_triattention() + _set_request_state( + mgr, + 7, + generation_steps=1, + evicted_tokens=127, + confirmed_kv_length=128, + ) + mgr._prepared_generation_batch = _PreparedGenerationBatch( + batch=SimpleNamespace(), + growth_by_request={7: 1}, + ) + + mgr.on_request_finish(request) + + assert mgr._request_states == {} + assert mgr._prepared_generation_batch.growth_by_request == {} + + +class TestTopKRouting: + @pytest.mark.parametrize("keep_count", [4096, 8192]) + def test_cross_request_union_uses_cute_without_fallback(self, keep_count): + width = keep_count + 64 + request_scores = [ + _distinct_topk_scores(width), + _distinct_topk_scores(width).roll(17, dims=1) + 0.000007, + ] + expected = [_union_oracle(scores, keep_count) for scores in request_scores] + selector = _BatchedUnionKeepSetSelector( + request_scores[0].shape[0], + width, + keep_count, + 0, + dtype=request_scores[0].dtype, + device=request_scores[0].device, + max_requests=len(request_scores), + ) + + with _mock_cute_topk_without_fallbacks() as cute_topk: + selector.select_requests( + torch.stack(request_scores), + normalize_scores=False, + ) + selected = selector.keep[: len(request_scores)].clone() + + assert cute_topk.call_count == 1 + for actual, expected_keep in zip(selected, expected): + assert torch.equal(actual, expected_keep) + + +class TestFixedScoreMetadata: + @pytest.mark.parametrize("normalize_scores", [False, True]) + @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) + def test_eager_bucket_binds_score_after_selection(self, eviction_mode, normalize_scores): + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module + + manager = _make_triattention( + top_B=4, + eviction_mode=eviction_mode, + normalize_scores=normalize_scores, + ) + manager._H = 2 + manager._F = 2 + manager._freq_scale_sq = torch.ones(2) + manager._offsets = torch.ones(2) + manager.calibration = {"omega": torch.ones(2)} + manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) + manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) + pool = torch.empty(8, 2, 1, 4, 4) + layout = SimpleNamespace( + manager=SimpleNamespace(num_pools=1), + num_layers=2, + global_layers=[0, 1], + layer_pools=[pool, pool], + dense_layers=[0, 1], + swa_layers=[], + storage_groups={0: [0, 1]}, + pool_view_fingerprint=(("fixed",),), + ) + score_staging = SimpleNamespace( + fused_group=SimpleNamespace(output=torch.empty(1, 4, 8)), + bind_score_launcher=mock.Mock(), + ) + keep_set_selector = SimpleNamespace(valid_widths=torch.empty(1, dtype=torch.int32)) + prepared = [ + _prepared_eviction( + _make_request(7), + request_id=7, + seq_len=8, + expected_keep_count=4, + ) + ] + + with ( + mock.patch.object( + module, + "_FixedScoreStagingBuffers", + return_value=score_staging, + ), + mock.patch.object( + manager, + "_build_cross_request_keep_set_selector", + return_value=keep_set_selector, + ) as build_selection, + ): + resources = manager._eager_resources_for(layout, prepared) + + score_staging.bind_score_launcher.assert_called_once_with( + keep_set_selector.valid_widths, + manager.score_aggregation, + ) + plan = build_selection.call_args.args[0] + assert plan.eviction_mode == eviction_mode + assert build_selection.call_args.kwargs["normalize_scores"] is normalize_scores + assert resources.score_staging is score_staging + assert resources.keep_set_selector is keep_set_selector + + def test_bulk_page_table_copy_uses_immutable_host_snapshots(self): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + ) + + device = torch.device("cuda") + current_stream = torch.cuda.current_stream(device) + manager_stream = torch.cuda.Stream(device=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 = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + staging.device = device + staging.max_requests = 1 + staging.page_count = 5 + staging.copy_block_count = 8 + staging.bulk_copy_done = torch.cuda.Event() + staging.bulk_consume_done = torch.cuda.Event() + staging.page_tables_active = False + staging.copy_done = torch.cuda.Event() + staging.copy_pending = False + staging._bulk_offsets_src = torch.empty( + 1, 1, 2, 8, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging._bulk_copy_idx_src = torch.arange( + 1, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging.block_offsets_device = torch.empty(1, 1, 2, 8, dtype=torch.int32, device=device) + staging.copy_done.record(current_stream) + + manager = SimpleNamespace( + host_kv_cache_block_offsets=host_table, + kv_factor=2, + layer_offsets={10: 0}, + layer_to_pool_mapping_dict={0: 0}, + index_mapper=SimpleNamespace(gather_k_block_offsets=gather), + index_scales=torch.tensor([2], dtype=torch.int32, device="cpu", pin_memory=True), + kv_offset=torch.tensor([1], dtype=torch.int32, device="cpu", pin_memory=True), + _stream=manager_stream, + ) + + with torch.cuda.stream(manager_stream): + torch.cuda._sleep(50_000_000) + with mock.patch.object( + torch, + "index_select", + side_effect=AssertionError("page-table staging used torch.index_select"), + ): + assert staging._stage_page_tables_bulk( + manager, + [7], + current_stream, + staging._bulk_offsets_src, + staging.block_offsets_device, + staging.copy_block_count, + ) + assert staging._bulk_offsets_src.shape[-1] == 8 + assert staging._bulk_offsets_src.shape[1] == 1 + + # Mutate both persistent V2 host inputs before the delayed kernel reads. + # The staged result must still reflect row 0 values [3, 4, 5, 6, 7]. + 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] + + host_table[0, 0, 0, :5] = torch.tensor([18, 19, 20, 21, 22], dtype=torch.int32) + selected_slot[0] = 0 + with torch.cuda.stream(manager_stream): + torch.cuda._sleep(50_000_000) + assert staging._stage_page_tables_bulk( + manager, + [7], + current_stream, + staging._bulk_offsets_src, + staging.block_offsets_device, + staging.copy_block_count, + ) + 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] + + def test_next_bulk_copy_waits_for_page_table_consumers(self): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + ) + + device = torch.device("cuda") + current_stream = torch.cuda.current_stream(device) + manager_stream = torch.cuda.Stream(device=device) + host_table = torch.zeros(1, 1, 2, 4, dtype=torch.int32, device="cpu", pin_memory=True) + host_table[0, 0, 0] = torch.tensor([1, 2, 3, 4], dtype=torch.int32) + + def gather_k_block_offsets(source, destination, request_ids, num_blocks): + assert request_ids == [7] + destination[:, :1, 0, :num_blocks].copy_(source[:, :1, 0, :num_blocks]) + + staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + staging.device = device + staging.max_requests = 1 + staging.copy_block_count = 4 + staging.bulk_copy_done = torch.cuda.Event() + staging.bulk_consume_done = torch.cuda.Event() + staging.page_tables_active = False + staging.copy_done = torch.cuda.Event() + staging.copy_pending = False + staging._bulk_offsets_src = torch.empty( + 1, 1, 2, 4, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging._bulk_copy_idx_src = torch.arange( + 1, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging.block_offsets_device = torch.empty(1, 1, 2, 4, dtype=torch.int32, device=device) + staging.copy_done.record(current_stream) + manager = SimpleNamespace( + host_kv_cache_block_offsets=host_table, + kv_factor=2, + index_mapper=SimpleNamespace( + gather_k_block_offsets=mock.Mock(side_effect=gather_k_block_offsets) + ), + index_scales=torch.tensor([2], dtype=torch.int32, device="cpu", pin_memory=True), + kv_offset=torch.tensor([1], dtype=torch.int32, device="cpu", pin_memory=True), + _stream=manager_stream, + ) + + assert staging._stage_page_tables_bulk( + manager, + [7], + current_stream, + staging._bulk_offsets_src, + staging.block_offsets_device, + staging.copy_block_count, + ) + manager_stream.synchronize() + snapshot = torch.empty_like(staging.block_offsets_device) + torch.cuda._sleep(20_000_000) + snapshot.copy_(staging.block_offsets_device) + staging.page_tables_active = True + staging.mark_page_tables_consumed(manager_stream) + + host_table[0, 0, 0] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) + assert staging._stage_page_tables_bulk( + manager, + [7], + current_stream, + staging._bulk_offsets_src, + staging.block_offsets_device, + staging.copy_block_count, + ) + current_stream.synchronize() + + assert snapshot[0, 0, 0].tolist() == [2, 4, 6, 8] + assert staging.block_offsets_device[0, 0, 0].tolist() == [10, 12, 14, 16] + + def test_cross_stream_staging_is_rejected_before_page_table_query(self): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + _FixedScoreStreamMismatch, + ) + + staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + staging.device = torch.device("cuda") + staging.max_requests = 8 + staging.stream = SimpleNamespace(device=torch.device("cuda:0"), cuda_stream=4) + staging.page_tables_active = False + staging.copy_pending = False + staging.copy_done = SimpleNamespace(query=mock.Mock(), synchronize=mock.Mock()) + staging.draft_block_offsets_device = None + manager = mock.Mock() + + other_stream = SimpleNamespace(device=torch.device("cuda:0"), cuda_stream=5) + with mock.patch.object(torch.cuda, "current_stream", return_value=other_stream): + with pytest.raises(_FixedScoreStreamMismatch, match="first CUDA stream"): + staging.stage(manager, [1], [8.0]) + staging.copy_done.query.assert_not_called() + staging.copy_done.synchronize.assert_not_called() + + def test_staged_page_tables_bypass_per_request_cuda_materialization(self): + manager = _make_triattention() + get_batch = mock.Mock() + manager.kv_cache_manager = SimpleNamespace(get_batch_cache_indices=get_batch) + staging = SimpleNamespace( + stage=mock.Mock(return_value=True), + ) + prepared = [ + _prepared_eviction( + SimpleNamespace(), + request_id=7, + round_start=8, + seq_len=8, + expected_keep_count=6, + protected_tail=2, + ), + _prepared_eviction( + SimpleNamespace(), + request_id=8, + round_start=9, + seq_len=9, + expected_keep_count=6, + protected_tail=3, + ), + ] + + manager._attach_page_ids(prepared, staging) + + staging.stage.assert_called_once_with( + manager.kv_cache_manager, + [7, 8], + [8, 9], + [8, 9], + [10, 12], + draft_manager=None, + ) + assert all(not hasattr(item, "page_ids") for item in prepared) + + @pytest.mark.parametrize("request_count", [1, 7, 8]) + def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, request_count): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + _FixedScoreStreamMismatch, + ) + + device = torch.device("cuda") + max_requests = 8 + page_count = 3 + seq_len = 7 + page_table_token_capacity = 11 + layer_elements = max_requests * page_count * 2 * 1 * 4 * 4 + shared = torch.randn(2 * layer_elements, device=device) + pools = [ + shared[:layer_elements].view(max_requests * page_count, 2, 1, 4, 4), + shared[layer_elements:].view(max_requests * page_count, 2, 1, 4, 4), + torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device), + torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device), + ] + dense_groups = [[0, 1], [2]] + representatives = [0, 2, 3] + q_real = torch.randn(4, 2, 4, dtype=torch.float64, device=device)[..., ::2] + q_imag = torch.randn(4, 2, 4, dtype=torch.float64, device=device)[..., ::2] + mlr = torch.randn(4, 2, 4, dtype=torch.float64, device=device)[..., ::2] + freq = torch.tensor([1.0, 0.0, 1.0, 0.0], dtype=torch.float64, device=device)[::2] + omega = torch.tensor([0.01, 0.0, 0.03, 0.0], dtype=torch.float64, device=device)[::2] + offsets = torch.tensor([1.0, 0.0, 2.0, 0.0], dtype=torch.float64, device=device)[::2] + assert not q_real.is_contiguous() + assert not freq.is_contiguous() + assert not omega.is_contiguous() + assert not offsets.is_contiguous() + staging = _FixedScoreStagingBuffers( + pools, + dense_groups, + [0, 1, 2], + representatives, + max_requests, + seq_len, + 2, + 2, + q_real, + q_imag, + mlr, + freq, + offsets, + omega, + page_table_token_capacity=page_table_token_capacity, + ) + assert staging.bucket_seq_len == seq_len + assert staging.page_table_token_capacity == page_table_token_capacity + assert staging.page_count == page_count + assert staging.offsets.dtype == torch.float32 + assert staging.offsets.is_contiguous() + assert staging.omega.dtype == torch.float32 + assert staging.omega.is_contiguous() + fused = staging.fused_group + for calibration in (*fused.pointer_middle[2:], *fused.pointer_tail): + assert calibration.dtype == torch.float32 + assert calibration.is_contiguous() + tables = { + 10: [ + [3 * request, 3 * request + 1, 3 * request + 2] for request in range(request_count) + ], + 12: [ + [3 * request + 2, 3 * request + 1, 3 * request] for request in range(request_count) + ], + 13: [ + [23 - 3 * request, 22 - 3 * request, 21 - 3 * request] + for request in range(request_count) + ], + } + + request_ids = list(range(request_count)) + round_starts = [131_071 + request for request in request_ids] + host_table = torch.zeros( + 3, + max_requests, + 2, + staging.copy_block_count, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + for slot, global_layer in enumerate((10, 12, 13)): + host_table[slot, :request_count, 0, :page_count].copy_( + torch.tensor(tables[global_layer], dtype=torch.int32) + ) + + def gather_k_block_offsets(source, destination, requested_ids, num_blocks): + for destination_row, request_id in enumerate(requested_ids): + destination[:, destination_row, 0, :num_blocks].copy_( + source[:, request_id, 0, :num_blocks] + ) + + gather = mock.Mock(side_effect=gather_k_block_offsets) + manager = SimpleNamespace( + enable_swa_scratch_reuse=False, + host_kv_cache_block_offsets=host_table, + kv_factor=2, + index_mapper=SimpleNamespace(gather_k_block_offsets=gather), + index_scales=torch.full((3,), 2, dtype=torch.int32, pin_memory=True), + kv_offset=torch.ones(3, dtype=torch.int32, pin_memory=True), + _stream=torch.cuda.Stream(device=device), + ) + + assert not staging.stage( + manager, + request_ids, + [2**31] * request_count, + [seq_len] * request_count, + [10] * request_count, + ) + assert gather.call_count == 0 + with mock.patch.object( + torch, + "index_select", + side_effect=AssertionError("page-table staging used torch.index_select"), + ): + assert staging.stage( + manager, + request_ids, + round_starts, + [seq_len] * request_count, + [10] * request_count, + ) + torch.cuda.current_stream(device).synchronize() + assert staging.round_starts_device.untyped_storage().data_ptr() == ( + staging.valid_seq_lens_device.untyped_storage().data_ptr() + ) + assert torch.equal( + staging.round_starts_device[:request_count], + torch.tensor(round_starts, dtype=torch.int32, device=device), + ) + assert torch.equal( + staging.valid_seq_lens_device[:request_count], + torch.full((request_count,), seq_len, dtype=torch.int32, device=device), + ) + for slot, global_layer in enumerate((10, 12, 13)): + expected = torch.tensor(tables[global_layer], dtype=torch.int32, device=device) * 2 + assert torch.equal( + staging.block_offsets_device[slot, :request_count, 0, :page_count], + expected, + ) + calls = gather.call_count + other_stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(other_stream): + with pytest.raises(_FixedScoreStreamMismatch, match="first CUDA stream"): + staging.stage(manager, request_ids, round_starts) + assert gather.call_count == calls + + @pytest.mark.parametrize("request_count", [1, 7, 8]) + @pytest.mark.parametrize("aggregation", ["mean", "max"]) + def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, aggregation): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + ) + + device = torch.device("cuda") + torch.manual_seed(20260703 + request_count) + max_requests = 8 + page_count = 2 + seq_len = 7 + prompt_len = 2 + page_ids = torch.arange(max_requests * page_count, dtype=torch.int64, device=device).view( + max_requests, page_count + ) + layer_elements = max_requests * page_count * 2 * 1 * 4 * 4 + shared = torch.randn(2 * layer_elements, device=device) + pools = [ + shared[:layer_elements].view(max_requests * page_count, 2, 1, 4, 4), + shared[layer_elements:].view(max_requests * page_count, 2, 1, 4, 4), + torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device), + ] + storage_groups = [[0, 1], [2]] + q_real = torch.randn(3, 2, 4, device=device)[..., ::2] + q_imag = torch.randn(3, 2, 4, device=device)[..., ::2] + mlr = torch.randn(3, 2, 4, device=device)[..., ::2] + freq = torch.tensor([0.7, 0.0, 1.3, 0.0], device=device)[::2] + omega = torch.tensor([0.013, 0.0, 0.071, 0.0], device=device)[::2] + offsets = torch.tensor([1.0, 0.0, 2.0, 0.0, 4.0, 0.0], device=device)[::2] + assert not q_real.is_contiguous() + assert not q_imag.is_contiguous() + assert not mlr.is_contiguous() + assert not freq.is_contiguous() + assert not omega.is_contiguous() + assert not offsets.is_contiguous() + round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 + round_starts = round_device[:request_count].tolist() + seq_lens = [seq_len - request % 2 for request in range(request_count)] + phase = (round_device[:, None, None] + offsets[None, :, None]) * omega[None, None] + oracle = _torch_tri_score_oracle( + pools, + page_ids[:request_count], + seq_lens, + round_starts, + q_real, + q_imag, + mlr, + freq, + omega, + offsets, + [0, 1, 2], + aggregation, + ) + for layers in storage_groups: + group = _FixedScoreGroup( + pools, + layers, + max_requests, + page_count, + seq_len, + 2, + _encode_block_offsets(page_ids), + [0] * len(layers), + q_real, + q_imag, + mlr, + freq, + omega, + offsets, + prompt_len=prompt_len, + ) + valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) + fixed = group.launch( + request_count, + torch.tensor(seq_lens, dtype=torch.int32, device=device), + valid_widths, + round_device, + torch.cos(phase).mean(dim=1), + torch.sin(phase).mean(dim=1), + aggregation, + ) + assert valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + assert fixed.shape == ( + request_count, + len(layers), + 2, + seq_len - prompt_len, + ) + for request in range(request_count): + for layer_slot, layer in enumerate(layers): + valid_width = seq_lens[request] - prompt_len + segment = fixed[request, layer_slot, :, :valid_width] + expected = oracle[request * len(pools) + layer][:, prompt_len:] + torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) + selected = torch.topk(segment.max(dim=0).values, 3).indices.sort().values + expected_selected = ( + torch.topk(expected.max(dim=0).values, 3).indices.sort().values + ) + assert torch.equal(selected, expected_selected) + + @pytest.mark.parametrize("request_count", [1, 7, 8]) + @pytest.mark.parametrize("aggregation", ["mean", "max"]) + def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count, aggregation): + """ONE launch over layers in DISTINCT storages with DISTINCT block tables. + + This is the production V2 shape: get_buffers wraps every layer as its + own TensorWrapper storage and every layer allocates its own pages, so + the fused path must not assume a shared storage anchor or a shared + per-request block table. + """ + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + _FixedScoreStreamMismatch, + ) + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + ) + + device = torch.device("cuda") + torch.manual_seed(20260707 + request_count) + max_requests = request_count + page_count = 2 + seq_len = 7 + prompt_len = 1 + num_layers = 3 + # Three SEPARATE allocations (distinct storages, like V2 TensorWrapper). + pools = [ + torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device) + for _ in range(num_layers) + ] + assert len({pool.untyped_storage().data_ptr() for pool in pools}) == num_layers + # A DIFFERENT block table per layer (per-layer page allocation). + generator = torch.Generator(device="cpu").manual_seed(7 + request_count) + page_ids_3d = torch.stack( + [ + torch.randperm(max_requests * page_count, generator=generator)[ + : max_requests * page_count + ] + .view(max_requests, page_count) + .to(device=device, dtype=torch.int64) + for _ in range(num_layers) + ] + ).contiguous() + q_real = torch.randn(num_layers, 2, 2, device=device) + q_imag = torch.randn(num_layers, 2, 2, device=device) + mlr = torch.randn(num_layers, 2, 2, device=device) + freq = torch.tensor([0.7, 1.3], device=device) + omega = torch.tensor([0.013, 0.071], device=device) + offsets = torch.tensor([1.0, 2.0, 4.0], device=device) + round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 + round_starts = round_device[:request_count].tolist() + seq_lens = [seq_len - request % 2 for request in range(request_count)] + layer_order = list(range(num_layers)) + block_offsets = _encode_block_offsets(page_ids_3d) + group = _FixedScoreGroup( + pools, + layer_order, + max_requests, + page_count, + seq_len, + 2, + block_offsets, + layer_order, # slot i holds layer i's tables + q_real, + q_imag, + mlr, + freq, + omega, + offsets, + prompt_len=prompt_len, + ) + valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) + mean_cos = torch.empty(request_count, 2, dtype=torch.float32, device=device) + mean_sin = torch.empty_like(mean_cos) + if aggregation == "mean": + triattention_kernels.prepare_mean_phase( + round_device, + offsets, + omega, + mean_cos, + mean_sin, + request_count, + ) + score_sentinel = -12345.0 + group.output.fill_(score_sentinel) + checked = group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_device, + mean_cos, + mean_sin, + aggregation, + ).clone() + staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + staging.device = group.output.device + staging.max_requests = request_count + staging.fused_group = group + staging.round_starts_device = round_device + staging.valid_seq_lens_device = valid_seq_lens + staging.mean_cos = mean_cos + staging.mean_sin = mean_sin + staging.offsets = offsets + staging.omega = omega + staging.stream = None + staging._phase_runner = None + staging._phase_args = () + staging._score_runner = None + staging._score_args = () + staging.bind_score_launcher(valid_widths, aggregation) + group.output.fill_(score_sentinel) + with ( + mock.patch.object( + triattention_kernels, + "prepare_mean_phase", + side_effect=AssertionError("checked phase wrapper was called"), + ), + mock.patch.object( + group, + "launch", + side_effect=AssertionError("checked score wrapper was called"), + ), + ): + fixed = staging.launch_prepared_score().clone() + torch.testing.assert_close(fixed, checked, rtol=0, atol=0) + assert valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + + # The deployed fused score must agree with the independent Torch oracle + # when every layer owns a distinct V2 block table. + oracle = _torch_tri_score_oracle( + pools, + {layer: page_ids_3d[layer, :request_count] for layer in layer_order}, + seq_lens, + round_starts, + q_real, + q_imag, + mlr, + freq, + omega, + offsets, + layer_order, + aggregation, + ) + for request in range(request_count): + for layer_slot, layer in enumerate(layer_order): + valid_width = seq_lens[request] - prompt_len + segment = fixed[request, layer_slot, :, :valid_width] + expected = oracle[request * num_layers + layer][:, prompt_len:] + torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) + + round_device.add_(17) + block_offsets.copy_(_encode_block_offsets(page_ids_3d.roll(1, dims=2))) + valid_seq_lens.copy_( + torch.tensor( + [seq_len - (request + 1) % 2 for request in range(request_count)], + dtype=torch.int32, + device=device, + ) + ) + if aggregation == "mean": + triattention_kernels.prepare_mean_phase( + round_device, + offsets, + omega, + mean_cos, + mean_sin, + request_count, + ) + expected_second_widths = valid_seq_lens - prompt_len + group.output.fill_(score_sentinel) + valid_widths.fill_(-1) + checked_second = group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_device, + mean_cos, + mean_sin, + aggregation, + ).clone() + group.output.fill_(score_sentinel) + valid_widths.fill_(-1) + with ( + mock.patch.object( + triattention_kernels, + "prepare_mean_phase", + side_effect=AssertionError("checked phase wrapper was called"), + ), + mock.patch.object( + group, + "launch", + side_effect=AssertionError("checked score wrapper was called"), + ), + ): + second_launch = staging.launch_prepared_score().clone() + torch.testing.assert_close(second_launch, checked_second, rtol=0, atol=0) + assert torch.equal(valid_widths, expected_second_widths) + assert not torch.equal(second_launch, fixed) + + other_stream = torch.cuda.Stream(device=device) + with torch.cuda.stream(other_stream): + with pytest.raises(_FixedScoreStreamMismatch, match="staging CUDA stream"): + staging.launch_prepared_score() + + +class TestKernelMaskedSwa: + def test_layer_partition_uses_local_model_config(self): + mgr = _make_triattention() + mgr.model_path = "/models/gpt-oss" + mgr.top_B = 128 + mgr.kv_cache_manager = SimpleNamespace(pp_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: + dense, sliding, window = mgr._attention_layer_partition(4) + + 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 + + def test_layer_partition_rejects_decode_budget_smaller_than_window(self): + mgr = _make_triattention() + mgr.model_path = "/models/gpt-oss" + mgr.top_B = 127 + mgr.kv_cache_manager = SimpleNamespace(pp_layers=[0, 1]) + config = _make_hf_config( + layer_types=["sliding_attention", "full_attention"], + sliding_window=128, + ) + + with ( + mock.patch("transformers.AutoConfig.from_pretrained", return_value=config), + pytest.raises(ValueError, match="decode budget top_B=127"), + ): + mgr._attention_layer_partition(2) + + +class TestFactory: + def test_returns_triattention_instance_with_v2(self): + # A plain V2 manager (block reuse off) yields a TriAttention instance. + # Calibration is deferred to the first request, so construction needs + # no calibration file or CUDA. + fake_v2 = _make_fake_v2(enable_block_reuse=False) + cfg = TriAttentionKvCacheCompressionConfig(top_B=32, beta=16, model_path="/models/test") + mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) + assert isinstance(mgr, TriAttention) + assert mgr.top_B == 32 + assert mgr.beta == 16 + assert mgr.kv_cache_manager is fake_v2 + + def test_factory_propagates_eviction_mode(self): + cfg = TriAttentionKvCacheCompressionConfig( + top_B=64, + beta=8, + eviction_mode="per_head", + model_path="/models/test", + ) + mgr = create_kv_cache_compression_manager( + cfg, kv_cache_manager=_make_fake_v2(enable_block_reuse=False) + ) + assert isinstance(mgr, TriAttention) + assert mgr.eviction_mode == "per_head" diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 1610f8700662..e2059936849b 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -252,7 +252,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, tensorrt_llm.llmapi.llm_args.KvCacheCompressionConfig, NoneType] default: null status: prototype otlp_traces_endpoint: From 43c2efcdd41856fd48f29cc736dfbeb40d6feb9d Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 05:25:55 -0700 Subject: [PATCH 004/178] [None][fix] Keep RoPE unfused under KV-cache compression Signed-off-by: tianruih --- .../_torch/kv_cache_compression/__init__.py | 14 +++++++++++++ .../_torch/models/modeling_gpt_oss.py | 6 +++++- tensorrt_llm/_torch/modules/attention.py | 21 +++++++++++++++++++ tensorrt_llm/functional.py | 3 ++- 4 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/__init__.py index e69de29bb2d1..64a6c31f2016 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/__init__.py +++ b/tensorrt_llm/_torch/kv_cache_compression/__init__.py @@ -0,0 +1,14 @@ +def compression_manager_class(algorithm: str): + """Return the compression-manager class registered for ``algorithm``. + + Mirrors the dispatch in ``_util.create_kv_cache_compression_manager`` + without building a manager, for callers that only need class-level + declarations (e.g. ``physically_evicts_cached_tokens`` in the attention + RoPE gate). Returns None for unknown algorithms; the factory is the one + that rejects them. + """ + if algorithm == "triattention": + from .triattention import TriAttention + + return TriAttention + return None diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index 00b5c77c8951..6398574ad11b 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -60,7 +60,11 @@ def __init__( beta_fast=pretrained_config.rope_scaling['beta_fast'], beta_slow=pretrained_config.rope_scaling['beta_slow'], duplicate_data=False), - is_neox=False, + # GPT-OSS pairs rotary dims NeoX-style (HF chunks halves). The + # fused kernel ignores this flag for YaRN and always applies the + # NeoX split, which masked the wrong value; the unfused Python + # path honors it, so False breaks any unfused-RoPE run. + is_neox=True, ) super().__init__( diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 80e9c67d74d3..4540893d67aa 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -633,6 +633,27 @@ def __init__( key="disable_rope_fusion_for_rocketkv") self.rope_fusion = False + if config.kv_cache_compression_config is not None: + from tensorrt_llm._torch.kv_cache_compression import \ + compression_manager_class + manager_class = compression_manager_class( + config.kv_cache_compression_config.algorithm) + if (manager_class is not None + and manager_class.physically_evicts_cached_tokens): + # The configured manager physically evicts cached tokens, so + # the KV length no longer equals the logical sequence length. + # The fused path derives each new token's rotary position from + # the KV length inside the kernel; the unfused path consumes + # the engine's logical position_ids, so surviving keys retain + # their original phases. + logger.warning_once( + "disable rope_fusion for KV-cache compression " + f"({config.kv_cache_compression_config.algorithm}): " + "rotary positions must come from logical position_ids, " + "not the compression-shortened KV length.", + 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/functional.py b/tensorrt_llm/functional.py index df80459359aa..a6e1ba3ec7c5 100755 --- a/tensorrt_llm/functional.py +++ b/tensorrt_llm/functional.py @@ -709,7 +709,8 @@ class PositionEmbeddingType(IntEnum): def is_rope(self) -> bool: return self in [ - self.rope_gptj, self.rope_gpt_neox, self.long_rope, self.mrope + self.rope_gptj, self.rope_gpt_neox, self.long_rope, self.mrope, + self.yarn ] def is_mrope(self) -> bool: From da8113ccf2f9d0abbd3424f442298734471e62f3 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 05:25:57 -0700 Subject: [PATCH 005/178] [None][feat] Raise the CuTE-DSL top-k decode limit to 16384 Signed-off-by: tianruih --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 22 +++---- .../_torch/thop/parallel/test_indexer_topk.py | 57 +++++++++++++++++++ 2 files changed, 68 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index a0fc3a534212..472d1a0689ae 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -5018,7 +5018,7 @@ def cute_dsl_topk_decode_blackwell( Args: input_values: Input logits tensor [batch_size * next_n, vocab_size] seq_lens: Sequence lengths for each batch [batch_size] - top_k: Number of top elements to select (max 4096) + top_k: Number of top elements to select (max 16384) next_n: Number of candidates per sequence (for speculative decoding) num_copy_bits: Number of bits for vectorized memory copy (128 or 256) load_balance: Enable persistent dynamic scheduling for load balancing @@ -5028,7 +5028,7 @@ def cute_dsl_topk_decode_blackwell( Note: This function requires Blackwell architecture (SM100+) and CuTE DSL support. - Maximum supported top_k is 4096. + Maximum supported top_k is 16384. """ # Validate SM version sm_version = get_sm_version() @@ -5038,11 +5038,11 @@ def cute_dsl_topk_decode_blackwell( "Use standard top-k implementation for older architectures.") # Validate inputs - if top_k <= 0 or top_k > 4096: + if top_k <= 0 or top_k > 16384: raise ValueError( - f"top_k must be in range [1, 4096], got {top_k}. " - "Maximum supported top_k is 4096 (filtered_topk_max_k staging raised from 2048)." - ) + f"top_k must be in range [1, 16384], got {top_k}. " + "16384 is the largest top_k verified bit-exact against torch.topk " + "on Blackwell (unit-tested at 8192 and 16384).") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") @@ -5754,7 +5754,7 @@ def cute_dsl_topk_decode_multi_cta_blackwell( Args: input_values: Input logits tensor [batch_size * next_n, vocab_size] seq_lens: Sequence lengths for each batch [batch_size] - top_k: Number of top elements to select (max 4096) + top_k: Number of top elements to select (max 16384) next_n: Number of candidates per sequence (for speculative decoding) num_copy_bits: Number of bits for vectorized memory copy (128 or 256) chunk_size_per_cta: Number of columns each CTA processes @@ -5774,11 +5774,11 @@ def cute_dsl_topk_decode_multi_cta_blackwell( "Use standard top-k implementation for older architectures.") # Validate inputs - if top_k <= 0 or top_k > 4096: + if top_k <= 0 or top_k > 16384: raise ValueError( - f"top_k must be in range [1, 4096], got {top_k}. " - "Maximum supported top_k is 4096 (filtered_topk_max_k staging raised from 2048)." - ) + f"top_k must be in range [1, 16384], got {top_k}. " + "16384 is the largest top_k verified bit-exact against torch.topk " + "on Blackwell (unit-tested at 8192 and 16384).") if next_n <= 0: raise ValueError(f"next_n must be positive, got {next_n}") diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index 09d23ff9be56..6285917bc39c 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -811,6 +811,63 @@ def run_fn(logits, seq_lens): ) +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") +@skip_pre_blackwell +@pytest.mark.parametrize("batch_size", [1, 8]) +@pytest.mark.parametrize("index_topk", [8192, 16384]) +@pytest.mark.parametrize("num_tokens", [32768, 131072]) +def test_cute_dsl_topk_decode_high_k(batch_size, index_topk, num_tokens): + """top_k above 4096 stays bit-exact (wrapper guard raised to 16384). + + Covers the three decode entry points: the single-CTA and multi-CTA + Blackwell wrappers (whose top_k guard this test backs) and the indexer + variant used by KV-cache eviction with large keep budgets. + """ + + def run_single_cta(logits, seq_lens): + return torch.ops.trtllm.cute_dsl_topk_decode_blackwell( + input_values=logits, + seq_lens=seq_lens, + top_k=index_topk, + next_n=1, + num_copy_bits=256, + load_balance=False, + ) + + def run_multi_cta(logits, seq_lens): + return torch.ops.trtllm.cute_dsl_topk_decode_multi_cta_blackwell( + input_values=logits, + seq_lens=seq_lens, + top_k=index_topk, + next_n=1, + num_copy_bits=256, + chunk_size_per_cta=16384, + dynamic=False, + ) + + def run_indexer(logits, seq_lens): + output_indices = torch.empty(batch_size, index_topk, dtype=torch.int32, device="cuda") + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + input_values=logits, + seq_lens=seq_lens, + output_indices=output_indices, + top_k=index_topk, + next_n=1, + num_copy_bits=256, + ) + return output_indices + + for run_fn in (run_single_cta, run_multi_cta, run_indexer): + _run_cute_dsl_topk_test( + batch_size, + 1, + index_topk, + num_tokens, + torch.float32, + run_fn, + ) + + @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") @skip_pre_blackwell @pytest.mark.parametrize("batch_size", [1, 4, 8, 16, 256]) From ec747810cd32194daddad0c3db2270dd576ff494 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:16:27 -0700 Subject: [PATCH 006/178] [None][refactor] Decide speculative compatibility at the executor call site Signed-off-by: tianruih --- tensorrt_llm/_torch/pyexecutor/_util.py | 54 ++++++++++++++----- .../_torch/pyexecutor/resource_manager.py | 8 --- tensorrt_llm/llmapi/llm_args.py | 4 ++ .../test_kv_cache_compression_manager.py | 29 +++++----- 4 files changed, 58 insertions(+), 37 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index feacedf85151..f9775b4ee6f6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2032,20 +2032,43 @@ def _create_kv_cache_manager( return kv_cache_manager +def kv_cache_compression_supported_with_spec( + config: KvCacheCompressionConfig, + spec_config: Optional[SpeculativeConfig], + draft_kv_cache_manager: Optional[KVCacheManagerV2], +) -> bool: + """Decide, before any manager is created, whether the configured + compression method supports the speculative setup. Logs the reason and + returns False when it does not; the run then stays uncompressed.""" + if spec_config is None or not config.is_eviction_method(): + return True + # Evicting methods co-compact the draft KV, so they only support spec + # modes whose draft KV is a standard paged cache in the same forward + # (one-model speculation). + mode = spec_config.spec_dec_mode + if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): + logger.warning( + "KV-cache compression algorithm '%s' evicts cached tokens and " + "does not support speculative decoding mode %s (the draft KV must " + "be a standard paged cache compacted together with the target); " + "running without a compression manager.", config.algorithm, + mode.name) + return False + return True + + def create_kv_cache_compression_manager( config: KvCacheCompressionConfig, kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, - spec_config: Optional[SpeculativeConfig] = None, ) -> Optional[BaseKVCacheCompressionManager]: """Build the KV-cache compression manager for ``config.algorithm``, or return None if no algorithm matches. 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 also - decided here: a compression manager is only created when the speculative - mode supports it, otherwise the run stays uncompressed. + here; the framework ships none. Speculative-decoding compatibility is + decided by the caller via ``kv_cache_compression_supported_with_spec``. """ logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " @@ -2259,16 +2282,19 @@ def create_py_executor_instance( kv_cache_compression_config = getattr(llm_args, "kv_cache_compression_config", None) if kv_cache_compression_config is not None: - compression_manager = create_kv_cache_compression_manager( - kv_cache_compression_config, - kv_cache_manager, - draft_kv_cache_manager=resources.get( - ResourceManagerType.DRAFT_KV_CACHE_MANAGER), - spec_config=spec_config, - ) - if compression_manager is not None: - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_manager) + draft_kv_cache_manager = resources.get( + ResourceManagerType.DRAFT_KV_CACHE_MANAGER) + if kv_cache_compression_supported_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, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if compression_manager is not None: + resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( + compression_manager) resource_manager = ResourceManager(resources) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index c5f3b1947069..58ec6ab41485 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2354,9 +2354,6 @@ class BaseKVCacheCompressionManager(BaseResourceManager): adjusts_generation_kv_length: ClassVar[bool] = False """Whether this manager can make target and logical KV lengths diverge.""" - physically_evicts_cached_tokens: ClassVar[bool] = False - """True for evicting methods; attention modules then keep RoPE unfused.""" - def __init__( self, kv_cache_manager: "KVCacheManagerV2", @@ -2385,11 +2382,6 @@ def __init__( draft_kv_cache_manager.kv_compression_manages_history = ( self.adjusts_generation_kv_length) - @classmethod - def is_eviction_method(cls) -> bool: - """Whether this method physically evicts cached tokens.""" - return cls.physically_evicts_cached_tokens - @property def has_independent_draft_kv_cache(self) -> bool: return self.draft_kv_cache_manager is not None diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fd3276163831..5fa1c4012264 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3268,6 +3268,10 @@ class KvCacheCompressionConfig(StrictBaseModel): "compression manager is built. Concrete algorithm configs subclass this " "and set the value.") + def is_eviction_method(self) -> bool: + """Whether this method physically evicts cached tokens.""" + return False + @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) class AgentTreeConfig(StrictBaseModel, PybindMirror): 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 e7448b79ff33..b95810fe503c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -279,40 +279,39 @@ def test_warns_for_unregistered_algorithm(self, fake_kv_cache_manager): mock_logger.warning.assert_called_once() def test_factory_accepts_independent_draft_manager(self): - from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - cfg = MagicMock() cfg.algorithm = "made_up_method" target = _v2_manager(is_draft=False) draft = _v2_manager(is_draft=True) - spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.EAGLE3_ONE_MODEL) assert ( create_kv_cache_compression_manager( cfg, target, draft_kv_cache_manager=draft, - spec_config=spec_config, ) is None ) def test_eviction_method_predicate_defaults_false(self): - # The base is not an evicting method, so the factory's speculative - # mode gate never restricts it: methods that do not touch the draft - # KV (e.g. offloading) work with any speculative mode. + # 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 + + config = KvCacheCompressionConfig(algorithm="offload") + assert config.is_eviction_method() is False m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) - assert m.is_eviction_method() is False assert not hasattr(m, "spec_config") - def test_eviction_method_predicate_follows_class_flag(self): - # The factory's speculative gate reads this predicate: an evicting - # method (physically_evicts_cached_tokens True) only accepts modes - # whose draft KV is a standard paged cache in the same forward. - class _EvictingManager(BaseKVCacheCompressionManager): - physically_evicts_cached_tokens = True + def test_spec_gate_only_restricts_eviction_methods(self): + from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig - assert _EvictingManager.is_eviction_method() is True + config = KvCacheCompressionConfig(algorithm="offload") + spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.DFLASH) + assert kv_cache_compression_supported_with_spec(config, spec_config, None) is True + assert kv_cache_compression_supported_with_spec(config, None, None) is True # ---------------------------------------------------------------------- # From 9eec866a828adee6fab2542ca21476ccf1fc1ab1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:20:01 -0700 Subject: [PATCH 007/178] [None][refactor] Carry algorithm traits on a KV-cache compression mode Signed-off-by: tianruih --- .../_torch/kv_cache_compression/interface.py | 30 +++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/_util.py | 3 +- tensorrt_llm/llmapi/llm_args.py | 10 +++++-- .../test_kv_cache_compression_manager.py | 2 +- 4 files changed, 40 insertions(+), 5 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/interface.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py new file mode 100644 index 000000000000..8ac47c98a640 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -0,0 +1,30 @@ +# 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. + + Mirrors ``SpeculativeDecodingMode``: configs map their ``algorithm`` + string to a member here, and callers read ``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/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f9775b4ee6f6..8ce2d2564f85 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2040,7 +2040,8 @@ def kv_cache_compression_supported_with_spec( """Decide, before any manager is created, whether the configured compression method supports the speculative setup. Logs the reason and returns False when it does not; the run then stays uncompressed.""" - if spec_config is None or not config.is_eviction_method(): + if (spec_config is None + or not config.kv_cache_compression_mode.is_eviction_method()): return True # Evicting methods co-compact the draft KV, so they only support spec # modes whose draft KV is a standard paged cache in the same forward diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5fa1c4012264..63002d47e7b3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3268,9 +3268,13 @@ class KvCacheCompressionConfig(StrictBaseModel): "compression manager is built. Concrete algorithm configs subclass this " "and set the value.") - def is_eviction_method(self) -> bool: - """Whether this method physically evicts cached tokens.""" - return False + @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) @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) 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 b95810fe503c..8a9d16974146 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -299,7 +299,7 @@ def test_eviction_method_predicate_defaults_false(self): from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig config = KvCacheCompressionConfig(algorithm="offload") - assert config.is_eviction_method() is False + assert config.kv_cache_compression_mode.is_eviction_method() is False m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) assert not hasattr(m, "spec_config") From 1ab4e5cf9a4c9364703212b0008a5c6477adaa96 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:22:33 -0700 Subject: [PATCH 008/178] [None][perf] Port register-staged compaction and streamline TriAttention - stage layered KV compaction vectors in registers (D64/D128 half-precision) - single-pass eviction cohort bookkeeping - drop the unprepared pack-compaction wrapper - shared keep-set selector base for union and per-head modes - adopt KvCacheCompressionMode and the call-site speculative gate Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 52 ++- .../_torch/kv_cache_compression/__init__.py | 14 - .../_torch/kv_cache_compression/interface.py | 9 +- .../triattention/triattention.py | 332 +++++++++--------- .../triattention/triattention_kernels.py | 128 +------ tensorrt_llm/_torch/modules/attention.py | 36 +- .../test_triattention_draft_cocompaction.py | 19 +- .../test_triattention_pipeline.py | 89 +++-- 8 files changed, 312 insertions(+), 367 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index f39196a3b985..4dd7d2fb8839 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1817,6 +1817,9 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( // The number of 16B vectors per head size in the kv cache. constexpr int VECS_PER_HEAD = Dh * sizeof(TCache) / 16; static_assert(BLOCK_SIZE % VECS_PER_HEAD == 0, "Kernel block should be able to handle entire heads."); + // D64 and D128 map one complete K/V vector to each x lane, so registers can + // preserve the read-before-write value across the in-place ordering barrier. + constexpr bool use_register_staging = Layered && (Dh == 64 || Dh == 128) && sizeof(TCache) == 2; int const batch_idx = blockIdx.z; int const kv_head_idx = blockIdx.y; @@ -1836,6 +1839,8 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( for (int token_block_offset = 0; token_block_offset < num_sparse_tokens; token_block_offset += tokens_per_block) { + uint4 key_vector; + uint4 value_vector; int const sparse_token_offset = token_block_offset + threadIdx.y; if (sparse_token_offset < num_sparse_tokens) @@ -1858,15 +1863,27 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( auto const src_k_block_ptr = reinterpret_cast(src_k_ptr); auto const src_v_block_ptr = reinterpret_cast(src_v_ptr); - for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) + if constexpr (use_register_staging) { + int const head_vec_idx = threadIdx.x; auto const src_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); auto const src_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - - k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_k_block_ptr[src_k_vec_idx]; - v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_v_block_ptr[src_v_vec_idx]; + key_vector = src_k_block_ptr[src_k_vec_idx]; + value_vector = src_v_block_ptr[src_v_vec_idx]; + } + else + { + for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) + { + auto const src_k_vec_idx + = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + auto const src_v_vec_idx + = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_k_block_ptr[src_k_vec_idx]; + v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_v_block_ptr[src_v_vec_idx]; + } } } __syncthreads(); @@ -1896,14 +1913,27 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( auto const dst_k_block_ptr = reinterpret_cast(dst_k_ptr); auto const dst_v_block_ptr = reinterpret_cast(dst_v_ptr); - for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) + if constexpr (use_register_staging) { + int const head_vec_idx = threadIdx.x; auto const dst_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx(dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); auto const dst_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx(dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - dst_k_block_ptr[dst_k_vec_idx] = k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; - dst_v_block_ptr[dst_v_vec_idx] = v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; + dst_k_block_ptr[dst_k_vec_idx] = key_vector; + dst_v_block_ptr[dst_v_vec_idx] = value_vector; + } + else + { + for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) + { + auto const dst_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx( + dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + auto const dst_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx( + dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + dst_k_block_ptr[dst_k_vec_idx] = k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; + dst_v_block_ptr[dst_v_vec_idx] = v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; + } } } } @@ -1960,15 +1990,15 @@ void launchSparseKvCacheCompactV2Layers( constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / 16; constexpr int32_t kDefaultSharedMemoryBytes = 48 * 1024; constexpr int32_t kSharedBytesPerToken = 2 * kVectorsPerHead * sizeof(uint4); - constexpr bool kUseD64VectorThreads = HeadDim == 64 && sizeof(T) == 2; - constexpr int32_t kVectorThreads = kUseD64VectorThreads ? kVectorsPerHead : 32; + constexpr bool kUseRegisterStaging = (HeadDim == 64 || HeadDim == 128) && sizeof(T) == 2; + constexpr int32_t kVectorThreads = kUseRegisterStaging ? kVectorsPerHead : 32; constexpr int32_t kTokensPerTile - = kUseD64VectorThreads ? 32 : (kSharedBytesPerToken * 32 <= kDefaultSharedMemoryBytes ? 32 : 16); + = kUseRegisterStaging ? 16 : (kSharedBytesPerToken * 32 <= kDefaultSharedMemoryBytes ? 32 : 16); constexpr int32_t kBlockSize = kVectorThreads * kTokensPerTile; static_assert(kSharedBytesPerToken * kTokensPerTile <= kDefaultSharedMemoryBytes); dim3 const block(kVectorThreads, kTokensPerTile); dim3 const grid(numLayers, params.kv_head_num, params.batch_size); - size_t const sharedBytes = 2 * block.y * kVectorsPerHead * sizeof(uint4); + size_t const sharedBytes = kUseRegisterStaging ? 0 : 2 * block.y * kVectorsPerHead * sizeof(uint4); updateSparseKvCacheAfterFmha <<>>(params); } diff --git a/tensorrt_llm/_torch/kv_cache_compression/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/__init__.py index 64a6c31f2016..e69de29bb2d1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/__init__.py +++ b/tensorrt_llm/_torch/kv_cache_compression/__init__.py @@ -1,14 +0,0 @@ -def compression_manager_class(algorithm: str): - """Return the compression-manager class registered for ``algorithm``. - - Mirrors the dispatch in ``_util.create_kv_cache_compression_manager`` - without building a manager, for callers that only need class-level - declarations (e.g. ``physically_evicts_cached_tokens`` in the attention - RoPE gate). Returns None for unknown algorithms; the factory is the one - that rejects them. - """ - if algorithm == "triattention": - from .triattention import TriAttention - - return TriAttention - return None diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index 8ac47c98a640..f10cb574fef4 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -13,12 +13,15 @@ class KvCacheCompressionMode(IntEnum): comparing strings. """ + TRIATTENTION = auto() NONE = auto() + def is_triattention(self): + return self == KvCacheCompressionMode.TRIATTENTION + def is_eviction_method(self): - """Whether this method physically evicts cached tokens. Evicting - algorithms add their member and extend this predicate.""" - return False + """Whether this method physically evicts cached tokens.""" + return self == KvCacheCompressionMode.TRIATTENTION @staticmethod def from_string(name: Optional[str]) -> "KvCacheCompressionMode": diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 1c9f2c867dab..748ff64ed88a 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -389,7 +389,97 @@ class _RuntimeKVLayout(NamedTuple): pool_view_fingerprint: Tuple[tuple, ...] -class _BatchedUnionKeepSetSelector: +class _BatchedKeepSetSelectorBase: + """Shared fixed buffers and prepared launchers for keep-set selectors.""" + + def __init__( + self, + *, + eviction_mode: str, + dense_layers: Tuple[int, ...], + num_query_heads: int, + num_kv_heads: int, + width: int, + keep_count: int, + prompt_len: int, + dtype: torch.dtype, + device: torch.device, + max_requests: int, + ) -> None: + if width <= keep_count or keep_count <= 0: + raise ValueError("keep-set selection requires width > keep_count > 0") + if max_requests <= 0: + raise ValueError("keep-set selection requires a positive request capacity") + self.eviction_mode = eviction_mode + self.dense_layers = tuple(int(layer) for layer in dense_layers) + self.num_query_heads = int(num_query_heads) + self.num_kv_heads = int(num_kv_heads) + self.width = int(width) + self.keep_count = int(keep_count) + self.prompt_len = int(prompt_len) + self.total_keep = self.prompt_len + self.keep_count + self.dtype = dtype + self.device = _canonical_device(device) + self.max_requests = int(max_requests) + self.valid_widths = torch.full( + (self.max_requests,), self.width, dtype=torch.int32, device=self.device + ) + + def _allocate_cpu_reference_buffers( + self, + scale_shape: Tuple[int, ...], + mask_shape: Tuple[int, ...], + index_dtype: torch.dtype, + ) -> None: + """CPU-only oracle scratch; None on CUDA, where prepared kernels run.""" + if self.device.type == "cpu": + self.valid_scale = torch.empty(scale_shape, dtype=self.dtype, device=self.device) + self.token_indices = torch.arange(self.width, dtype=index_dtype, device=self.device) + self.invalid_mask = torch.empty(mask_shape, dtype=torch.bool, device=self.device) + else: + self.valid_scale = None + self.token_indices = None + self.invalid_mask = None + + def _build_prepared_selection_launchers( + self, + scores_rows: torch.Tensor, + row_lengths: torch.Tensor, + provisional_indices: torch.Tensor, + keep_rows: torch.Tensor, + ) -> None: + """Bind the CuTE topk and the deterministic finalizer over row-major views.""" + if self.device.type != "cuda": + self.prepared_topk = None + self.prepared_finalizer = None + return + self.prepared_topk = _PreparedCuteTopK( + int(scores_rows.shape[0]), self.width, self.keep_count, self.device + ) + self.prepared_finalizer = _PreparedTopKFinalizer( + scores_rows, + row_lengths, + provisional_indices, + keep_rows, + self.keep_count, + self.prompt_len, + ) + + def _prefill_prompt_ordinals(self, keep: torch.Tensor) -> None: + if self.prompt_len: + prompt = torch.arange(self.prompt_len, dtype=torch.int32, device=self.device) + keep[..., : self.prompt_len].copy_(prompt.expand(*keep.shape[:-1], self.prompt_len)) + + def _validated_request_count(self, scores: torch.Tensor, what: str) -> int: + request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 + if request_count <= 0 or request_count > self.max_requests: + raise ValueError(f"request count exceeds the {what} selection capacity") + if scores.is_cuda and request_count != self.max_requests: + raise ValueError("CUDA selection requires the selector's fixed request count") + return request_count + + +class _BatchedUnionKeepSetSelector(_BatchedKeepSetSelectorBase): """Persistent ``[request, ...]`` buffers for union selection.""" def __init__( @@ -408,22 +498,21 @@ def __init__( input_scores: Optional[torch.Tensor] = None, normalize_scores: bool = True, ) -> None: - if rows <= 0 or width <= keep_count or keep_count <= 0: - raise ValueError("cross-request selection requires rows > 0 and width > keep_count > 0") - if max_requests <= 0: - raise ValueError("cross-request selection requires a positive request capacity") - self.max_requests = max_requests - self.eviction_mode = "union" - self.dense_layers = tuple(dense_layers) - self.num_query_heads = int(num_query_heads) - self.num_kv_heads = int(num_kv_heads) + if rows <= 0: + raise ValueError("cross-request selection requires rows > 0") + super().__init__( + eviction_mode="union", + dense_layers=dense_layers, + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + width=width, + keep_count=keep_count, + prompt_len=prompt_len, + dtype=dtype, + device=device, + max_requests=max_requests, + ) self.rows = rows - self.width = width - self.keep_count = keep_count - self.prompt_len = prompt_len - self.total_keep = prompt_len + keep_count - self.dtype = dtype - self.device = _canonical_device(device) if self.device.type == "cuda" and input_scores is None: raise ValueError("CUDA union selection requires its fixed score input") @@ -436,35 +525,11 @@ def __init__( self.keep = torch.empty( (max_requests, self.total_keep), dtype=torch.int32, device=self.device ) - self.valid_widths = torch.full( - (max_requests,), width, dtype=torch.int32, device=self.device + self._allocate_cpu_reference_buffers( + (max_requests, 1, 1), (max_requests, 1, width), torch.int32 ) - if self.device.type == "cpu": - self.valid_scale = torch.empty((max_requests, 1, 1), dtype=dtype, device=self.device) - self.token_indices = torch.arange(width, dtype=torch.int32, device=self.device) - self.invalid_mask = torch.empty( - (max_requests, 1, width), dtype=torch.bool, device=self.device - ) - else: - self.valid_scale = None - self.token_indices = None - self.invalid_mask = None - self.prepared_topk = ( - _PreparedCuteTopK(max_requests, width, keep_count, self.device) - if self.device.type == "cuda" - else None - ) - self.prepared_finalizer = ( - _PreparedTopKFinalizer( - self.combined, - self.valid_widths, - self.final_indices, - self.keep, - keep_count, - prompt_len, - ) - if self.device.type == "cuda" - else None + self._build_prepared_selection_launchers( + self.combined, self.valid_widths, self.final_indices, self.keep ) self.prepared_scores = ( _PreparedUnionScores( @@ -478,9 +543,7 @@ def __init__( if self.device.type == "cuda" and input_scores is not None else None ) - if prompt_len: - prompt = torch.arange(prompt_len, dtype=torch.int32, device=self.device) - self.keep[:, :prompt_len].copy_(prompt.expand(max_requests, -1)) + self._prefill_prompt_ordinals(self.keep) def _select_input_scores( self, @@ -577,11 +640,7 @@ def select_requests( normalize_scores: bool, ) -> None: """Select from request-major score output without repacking it.""" - request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 - if request_count <= 0 or request_count > self.max_requests: - raise ValueError("request count exceeds the cross-request selection capacity") - if scores.is_cuda and request_count != self.max_requests: - raise ValueError("CUDA selection requires the selector's fixed request count") + request_count = self._validated_request_count(scores, "cross-request") if ( scores.numel() != request_count * self.rows * self.width or int(scores.shape[-1]) != self.width @@ -606,7 +665,7 @@ def select_requests( ) -class _BatchedPerHeadKeepSetSelector: +class _BatchedPerHeadKeepSetSelector(_BatchedKeepSetSelectorBase): """Fixed ``[request, ...]`` selector for both per-head modes.""" def __init__( @@ -629,13 +688,19 @@ def __init__( raise ValueError("per-head selection requires positive layer, head, and request counts") if num_query_heads % num_kv_heads: raise ValueError("query heads must be divisible by KV heads") - if width <= keep_count or keep_count <= 0: - raise ValueError("per-head selection requires width > keep_count > 0") - self.eviction_mode = eviction_mode - self.dense_layers = tuple(int(layer) for layer in dense_layers) + super().__init__( + eviction_mode=eviction_mode, + dense_layers=dense_layers, + num_query_heads=num_query_heads, + num_kv_heads=num_kv_heads, + width=width, + keep_count=keep_count, + prompt_len=prompt_len, + dtype=dtype, + device=device, + max_requests=max_requests, + ) self.num_layers = len(self.dense_layers) - self.num_query_heads = int(num_query_heads) - self.num_kv_heads = int(num_kv_heads) self.query_group_size = self.num_query_heads // self.num_kv_heads self.rows = self.num_layers * self.num_query_heads self.selection_rows = ( @@ -643,48 +708,22 @@ def __init__( if eviction_mode == "per_head" else self.num_layers * self.num_kv_heads ) - self.width = int(width) - self.keep_count = int(keep_count) - self.prompt_len = int(prompt_len) - self.total_keep = self.prompt_len + self.keep_count - self.dtype = dtype - self.device = _canonical_device(device) - self.max_requests = int(max_requests) score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) grouped_shape = (self.max_requests, self.num_layers, self.num_kv_heads, self.width) self.row_mean = torch.empty(score_shape[:-1] + (1,), dtype=dtype, device=self.device) self.row_std = torch.empty_like(self.row_mean) - self.valid_widths = torch.full( - (self.max_requests,), self.width, dtype=torch.int32, device=self.device - ) self.selection_scores = torch.empty( (self.max_requests, self.selection_rows, self.width), dtype=dtype, device=self.device, ) - if self.device.type == "cpu": - self.valid_scale = torch.empty( - (self.max_requests, 1, 1, 1), dtype=dtype, device=self.device - ) - self.token_indices = torch.arange(self.width, dtype=torch.long, device=self.device) - self.invalid_mask = torch.empty( - (self.max_requests, 1, 1, self.width), dtype=torch.bool, device=self.device - ) - self.grouped_scores = torch.empty(grouped_shape, dtype=dtype, device=self.device) - else: - self.valid_scale = None - self.token_indices = None - self.invalid_mask = None - self.grouped_scores = None - self.prepared_topk = ( - _PreparedCuteTopK( - self.max_requests * self.selection_rows, - self.width, - self.keep_count, - self.device, - ) - if self.device.type == "cuda" + self._allocate_cpu_reference_buffers( + (self.max_requests, 1, 1, 1), (self.max_requests, 1, 1, self.width), torch.long + ) + self.grouped_scores = ( + torch.empty(grouped_shape, dtype=dtype, device=self.device) + if self.device.type == "cpu" else None ) self.row_seq_lens = torch.full( @@ -706,23 +745,13 @@ def __init__( self.row_seq_lens_flat = self.row_seq_lens.view(-1) self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) self.keep_flat = self.keep.view(-1, self.total_keep) - self.prepared_finalizer = ( - _PreparedTopKFinalizer( - self.selection_scores_flat, - self.row_seq_lens_flat, - self.top_indices_i32_flat, - self.keep_flat, - self.keep_count, - self.prompt_len, - ) - if self.device.type == "cuda" - else None + self._build_prepared_selection_launchers( + self.selection_scores_flat, + self.row_seq_lens_flat, + self.top_indices_i32_flat, + self.keep_flat, ) - if self.prompt_len: - prompt = torch.arange(self.prompt_len, dtype=torch.int32, device=self.device) - self.keep[:, :, : self.prompt_len].copy_( - prompt.view(1, 1, -1).expand(self.max_requests, self.selection_rows, -1) - ) + self._prefill_prompt_ordinals(self.keep) def _select_input_scores( self, @@ -829,11 +858,7 @@ def select_requests( *, normalize_scores: bool, ) -> None: - request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 - if request_count <= 0 or request_count > self.max_requests: - raise ValueError("request count exceeds the per-head selection capacity") - if scores.is_cuda and request_count != self.max_requests: - raise ValueError("CUDA selection requires the selector's fixed request count") + request_count = self._validated_request_count(scores, "per-head") expected_shape = ( request_count, self.num_layers, @@ -1426,7 +1451,6 @@ class TriAttention(BaseKVCacheCompressionManager): """ adjusts_generation_kv_length = True - physically_evicts_cached_tokens = True def __init__( self, @@ -1726,99 +1750,89 @@ def _periodic_evict( gen_requests = scheduled_batch.generation_requests if not gen_requests: return - active_requests = [] + mgr = self.kv_cache_manager + resolved_requests = [] for request in gen_requests: if request.is_dummy or request.state in ( LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT, ): continue - kv_cache = self.kv_cache_manager.kv_cache_map.get(request.py_request_id) + request_id = request.py_request_id + kv_cache = mgr.kv_cache_map.get(request_id) if kv_cache is None: continue if not kv_cache.is_active: raise RuntimeError( "TriAttention cannot finalize a suspended target KV cache; " - f"request {request.py_request_id} must be resumed before " + f"request {request_id} must be resumed before " "the final update hook" ) - if request.py_request_id not in self._request_states: + if request_id not in self._request_states: self.on_request_init(request) - active_requests.append(request) - if not active_requests or not self._calibrated: + resolved_requests.append((request, request_id, kv_cache)) + if not resolved_requests or not self._calibrated: return - mgr = self.kv_cache_manager - num_layers = self._num_layers_from_manager() protected_tails: Dict[int, int] = {} + eviction_groups = {} - # (1) bump per-request step counters; collect who evicts THIS step. - evict_now = [] - for request in active_requests: - rid = request.py_request_id - kv_cache = mgr.kv_cache_map.get(rid) - if kv_cache is None or not kv_cache.is_active: - continue + # Resolve every active target cache before changing cadence state. The + # captured cache objects also avoid repeating the V2 map lookup here. + for request, request_id, kv_cache in resolved_requests: raw_capacity = int(kv_cache.capacity) # One-engine speculative decoding keeps a fixed reserve E. Under # overlap, B(n) is allocated/enqueued before finalizing B(n-1), so # its exact scheduler growth Q is also opaque. Both spans are # contiguous after the stable target prefix and move byte-for-byte. protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( - scheduled_batch, rid + scheduled_batch, request_id ) seq_len = raw_capacity - protected_tail if seq_len < 0 or protected_tail < 0: raise RuntimeError( - f"Request {rid} has an inconsistent protected target tail: " + f"Request {request_id} has an inconsistent protected target tail: " f"confirmed={seq_len}, capacity={raw_capacity}, " f"protected_tail={protected_tail}" ) if seq_len < kv_cache.history_length: raise RuntimeError( - f"Request {rid} KV length {seq_len} is below finalized " + f"Request {request_id} KV length {seq_len} is below finalized " f"history {kv_cache.history_length}" ) - request_state = self._request_states[rid] + request_state = self._request_states[request_id] request_state.confirmed_kv_length = seq_len - protected_tails[rid] = protected_tail previous_step = request_state.generation_steps confirmed_delta = 1 + int(request.py_num_accepted_draft_tokens) step = previous_step + confirmed_delta request_state.generation_steps = step - if previous_step // self.beta < step // self.beta: - if seq_len > self._minimum_evictable_length(request, seq_len): - if self.draft_kv_cache_manager is not None: - draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(rid) - if draft_kv_cache is None or not draft_kv_cache.is_active: - raise RuntimeError( - "TriAttention cannot co-compress a missing or " - f"suspended draft KV cache; request {rid} must " - "be resumed before the final update hook" - ) - evict_now.append((request, rid)) + if previous_step // self.beta >= step // self.beta: + continue + keep_count = self._minimum_evictable_length(request, seq_len) + if seq_len <= keep_count: + continue + if self.draft_kv_cache_manager is not None: + draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) + if draft_kv_cache is None or not draft_kv_cache.is_active: + raise RuntimeError( + "TriAttention cannot co-compress a missing or " + f"suspended draft KV cache; request {request_id} must " + "be resumed before the final update hook" + ) + protected_tails[request_id] = protected_tail + prompt_len = min(int(request.py_prompt_len), seq_len) + key = (prompt_len, keep_count) + eviction_groups.setdefault(key, []).append((request, request_id)) # (2) Compact all affected dense and kernel-masked SWA layers, then release # the unreachable tail directly through V2's public resize primitive. - if not evict_now: + if not eviction_groups: return - protected_tail_lengths = {rid: protected_tails[rid] for _, rid in evict_now} - # Prompt and retained geometry define selection and destination layout. - # Requests with the same geometry execute eagerly in bounded chunks. - # Chunking limits staging memory. - eviction_groups = {} - for request, rid in evict_now: - seq_len = self._request_states[rid].confirmed_kv_length - if seq_len is None: - raise RuntimeError(f"Missing confirmed KV length for request {rid}") - prompt_len = min(int(request.py_prompt_len), seq_len) - keep_count = self._minimum_evictable_length(request, seq_len) - key = (prompt_len, keep_count) - eviction_groups.setdefault(key, []).append((request, rid)) + num_layers = self._num_layers_from_manager() for group in eviction_groups.values(): for begin in range(0, len(group), _EAGER_REQUEST_CHUNK_SIZE): chunk = group[begin : begin + _EAGER_REQUEST_CHUNK_SIZE] - chunk_tails = {rid: protected_tail_lengths[rid] for _, rid in chunk} + chunk_tails = {rid: protected_tails[rid] for _, rid in chunk} with nvtx_range_debug("triattention.evict_request_group", color="purple"): capacity_targets = self._evict_requests( chunk, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index cd830079a61e..df0d4dca0680 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import List, Optional +from typing import List import torch import triton @@ -428,132 +428,6 @@ def _pack_compaction_sources_kernel( ) -def pack_compaction_sources( - selected_indices: torch.Tensor, - valid_seq_lens: torch.Tensor, - dense_offsets: torch.Tensor, - dense_indices: torch.Tensor, - *, - eviction_mode: str, - prompt_len: int, - keep_count: int, - num_dense_layers: int, - num_kv_heads: int, - max_protected_tail: int, - swa_window: int = 0, - swa_offsets: Optional[torch.Tensor] = None, - swa_indices: Optional[torch.Tensor] = None, -) -> None: - """Pack all dynamic dense and optional SWA source ordinals in one launch.""" - if eviction_mode not in ("union", "per_head", "per_layer_perhead"): - raise ValueError(f"unsupported compaction mode: {eviction_mode}") - prompt_len = int(prompt_len) - keep_count = int(keep_count) - num_dense_layers = int(num_dense_layers) - num_kv_heads = int(num_kv_heads) - max_protected_tail = int(max_protected_tail) - swa_window = int(swa_window) - request_count = int(selected_indices.shape[0]) if selected_indices.ndim else 0 - if ( - request_count <= 0 - or min(keep_count, num_dense_layers, num_kv_heads) <= 0 - or min(prompt_len, max_protected_tail, swa_window) < 0 - or selected_indices.shape[-1] != prompt_len + keep_count - ): - raise ValueError("compaction packing requires valid positive geometry") - - tensors = (selected_indices, valid_seq_lens, dense_offsets, dense_indices) - if any( - not tensor.is_cuda - or tensor.dtype != torch.int32 - or not tensor.is_contiguous() - or tensor.device != selected_indices.device - for tensor in tensors - ): - raise ValueError("compaction packing requires contiguous CUDA int32 tensors") - if valid_seq_lens.shape != (request_count,) or dense_offsets.shape != (request_count + 1,): - raise ValueError("compaction lengths and offsets do not match the request count") - - per_layer = eviction_mode == "per_layer_perhead" - union = eviction_mode == "union" - if union: - selection_rows = 1 - elif per_layer: - selection_rows = num_dense_layers * num_kv_heads - else: - selection_rows = num_kv_heads - if union: - expected_selection_prefix = (request_count,) - else: - expected_selection_prefix = (request_count, selection_rows) - if tuple(selected_indices.shape[:-1]) != expected_selection_prefix: - raise ValueError("selected indices do not match the compaction mode") - - domain_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads - expected_dense_prefix = (num_dense_layers, num_kv_heads) if per_layer else (num_kv_heads,) - if ( - dense_indices.ndim != len(expected_dense_prefix) + 1 - or tuple(dense_indices.shape[:-1]) != expected_dense_prefix - ): - raise ValueError("dense source buffer does not match the compaction mode") - dense_total = int(dense_indices.shape[-1]) - - has_swa = swa_indices is not None or swa_offsets is not None - if has_swa: - if swa_indices is None or swa_offsets is None or swa_window <= 0: - raise ValueError("SWA packing requires indices, offsets, and a positive window") - if ( - not swa_indices.is_cuda - or not swa_offsets.is_cuda - or swa_indices.dtype != torch.int32 - or swa_offsets.dtype != torch.int32 - or not swa_indices.is_contiguous() - or not swa_offsets.is_contiguous() - or swa_indices.device != selected_indices.device - or swa_offsets.device != selected_indices.device - or swa_indices.ndim != 2 - or tuple(swa_indices.shape[:-1]) != (num_kv_heads,) - or swa_offsets.shape != (request_count + 1,) - ): - raise ValueError("SWA source buffers do not match the compaction geometry") - swa_total = int(swa_indices.shape[-1]) - swa_indices_arg = swa_indices - swa_offsets_arg = swa_offsets - else: - if swa_window != 0: - raise ValueError("SWA window requires SWA source buffers") - swa_total = 0 - # HAS_SWA specializes the corresponding loads and stores away. - swa_indices_arg = dense_indices - swa_offsets_arg = dense_offsets - - max_move = keep_count + max_protected_tail - if has_swa: - max_move = max(max_move, swa_window + max_protected_tail) - block = 256 - _pack_compaction_sources_kernel[(request_count, domain_count, triton.cdiv(max_move, block))]( - selected_indices, - valid_seq_lens, - dense_offsets, - dense_indices, - swa_offsets_arg, - swa_indices_arg, - DENSE_TOTAL=dense_total, - SWA_TOTAL=swa_total, - SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=prompt_len + keep_count, - KEEP_COUNT=keep_count, - PROMPT_LEN=prompt_len, - NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window, - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=has_swa, - BLOCK=block, - num_warps=4, - ) - - @triton.jit def _finalize_topk_indices_kernel( scores, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 4540893d67aa..d2a84d24b2cd 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -633,26 +633,22 @@ def __init__( key="disable_rope_fusion_for_rocketkv") self.rope_fusion = False - if config.kv_cache_compression_config is not None: - from tensorrt_llm._torch.kv_cache_compression import \ - compression_manager_class - manager_class = compression_manager_class( - config.kv_cache_compression_config.algorithm) - if (manager_class is not None - and manager_class.physically_evicts_cached_tokens): - # The configured manager physically evicts cached tokens, so - # the KV length no longer equals the logical sequence length. - # The fused path derives each new token's rotary position from - # the KV length inside the kernel; the unfused path consumes - # the engine's logical position_ids, so surviving keys retain - # their original phases. - logger.warning_once( - "disable rope_fusion for KV-cache compression " - f"({config.kv_cache_compression_config.algorithm}): " - "rotary positions must come from logical position_ids, " - "not the compression-shortened KV length.", - key="disable_rope_fusion_for_kv_cache_compression") - self.rope_fusion = False + if (config.kv_cache_compression_config is not None + and config.kv_cache_compression_config. + kv_cache_compression_mode.is_eviction_method()): + # The configured method physically evicts cached tokens, so the + # KV length no longer equals the logical sequence length. The + # fused path derives each new token's rotary position from the KV + # length inside the kernel; the unfused path consumes the engine's + # logical position_ids, so surviving keys retain their original + # phases. + logger.warning_once( + "disable rope_fusion for KV-cache compression " + f"({config.kv_cache_compression_config.algorithm}): " + "rotary positions must come from logical position_ids, " + "not the compression-shortened KV length.", + 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( 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 index 647a7ae35f31..c53bade407d5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -321,21 +321,22 @@ def test_draft_admission_gates_raise(gate, match): draft_manager.max_attention_window_vec = [128] if gate == "dflash": # DFlash reads cross-attention context buffers, not a paged KV cache; - # the factory's eviction-method speculative-mode gate declines to - # create a manager and the run stays uncompressed. - from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + # the call-site speculative gate declines to create a manager and the + # run stays uncompressed. + from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, TriAttentionKvCacheCompressionConfig, ) - manager = create_kv_cache_compression_manager( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), - _make_fake_v2(), - draft_kv_cache_manager=draft_manager, - spec_config=DFlashDecodingConfig(max_draft_len=3), + assert ( + kv_cache_compression_supported_with_spec( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + DFlashDecodingConfig(max_draft_len=3), + draft_manager, + ) + is False ) - assert manager is None return manager = TriAttention( _make_fake_v2(), diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 2042cd5a2f91..4b32b4164341 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -646,6 +646,34 @@ def track_range(name, **kwargs): "exit:triattention.resize", ] + def test_suspended_cache_rejects_batch_before_cadence_mutation(self): + manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) + second_request = _make_request(8, py_prompt_len=1024) + manager.kv_cache_manager.kv_cache_map[8] = SimpleNamespace(is_active=False) + first_state = manager._request_states[7] + second_state = _set_request_state(manager, 8, generation_steps=127) + batch = SimpleNamespace(generation_requests=[first_request, second_request]) + + with pytest.raises(RuntimeError, match="request 8 must be resumed"): + manager._periodic_evict(batch) + + assert first_state.generation_steps == 127 + assert first_state.confirmed_kv_length is None + assert second_state.generation_steps == 127 + assert second_state.confirmed_kv_length is None + + def test_non_boundary_step_skips_eviction_geometry(self): + manager, _, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) + state = manager._request_states[7] + state.generation_steps = 126 + + with mock.patch.object(manager, "_minimum_evictable_length") as keep_count: + manager._periodic_evict(batch) + + keep_count.assert_not_called() + assert state.generation_steps == 127 + assert state.confirmed_kv_length == 1024 + 4096 + 1 + def test_eager_eviction_chunks_large_due_cohort(self): manager, _, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) requests = [] @@ -861,19 +889,30 @@ def test_resize_shrinks_draft_cache_with_its_own_protected_tail(self): draft_cache.resize.assert_called_once_with(retained + 6, None) def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): - # The factory runs the speculative feature gates and then builds the - # manager; a one-model MTP contract must pass both. - from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + # A one-model MTP contract passes the call-site speculative gate, and + # the factory then builds a manager that validates cleanly. + from tensorrt_llm._torch.pyexecutor._util import ( + create_kv_cache_compression_manager, + kv_cache_compression_supported_with_spec, + ) from tensorrt_llm.llmapi.llm_args import ( MTPDecodingConfig, TriAttentionKvCacheCompressionConfig, ) + config = TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8) + assert config.kv_cache_compression_mode.is_eviction_method() is True + draft_manager = _make_fake_v2(is_draft=True) + assert ( + kv_cache_compression_supported_with_spec( + config, MTPDecodingConfig(max_draft_len=1), draft_manager + ) + is True + ) manager = create_kv_cache_compression_manager( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + config, _make_fake_v2(), - draft_kv_cache_manager=_make_fake_v2(is_draft=True), - spec_config=MTPDecodingConfig(max_draft_len=1), + draft_kv_cache_manager=draft_manager, ) manager._validate_v2_compatibility() @@ -889,36 +928,38 @@ def test_unvalidated_paged_draft_tail_contracts_remain_fail_closed(self, mode): ) else: spec_config = PARDDecodingConfig(max_draft_len=3) - # The factory's eviction-method speculative-mode gate declines to - # create a manager; the run stays uncompressed. - from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + # The call-site speculative gate declines before any manager is + # created; the run stays uncompressed. + from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig - manager = create_kv_cache_compression_manager( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), - _make_fake_v2(), - draft_kv_cache_manager=_make_fake_v2(is_draft=True), - spec_config=spec_config, + assert ( + kv_cache_compression_supported_with_spec( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + spec_config, + _make_fake_v2(is_draft=True), + ) + is False ) - assert manager is None def test_dflash_spec_mode_is_rejected(self): # Policy: the DFlash draft reads cross-attention context buffers, not - # a paged KV cache, so compression cannot cover it. The factory's - # eviction-method speculative-mode gate declines to create a manager. - from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager + # a paged KV cache, so compression cannot cover it. The call-site + # speculative gate declines to create a manager. + from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, TriAttentionKvCacheCompressionConfig, ) - manager = create_kv_cache_compression_manager( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), - _make_fake_v2(), - draft_kv_cache_manager=_make_fake_v2(is_draft=True), - spec_config=DFlashDecodingConfig(max_draft_len=3), + assert ( + kv_cache_compression_supported_with_spec( + TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + DFlashDecodingConfig(max_draft_len=3), + _make_fake_v2(is_draft=True), + ) + is False ) - assert manager is None def test_prepare_snapshots_fixed_linear_generation_growth(self): manager = _make_fake_v2() From 262cdbb3b9e44962321b0b45e7c03f4ec76bb1de Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:28:39 -0700 Subject: [PATCH 009/178] [None][chore] Trim the compression-mode docstring Signed-off-by: tianruih --- tensorrt_llm/_torch/kv_cache_compression/interface.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index f10cb574fef4..c50891eb9529 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -8,9 +8,8 @@ class KvCacheCompressionMode(IntEnum): """Algorithm-level traits of a KV-cache compression method. - Mirrors ``SpeculativeDecodingMode``: configs map their ``algorithm`` - string to a member here, and callers read ``is_*`` predicates instead of - comparing strings. + Configs map their ``algorithm`` string to a member here; callers read the + ``is_*`` predicates instead of comparing strings. """ TRIATTENTION = auto() From 6cedd04dc11357383c438cec887b0de2e90fe4f5 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:28:42 -0700 Subject: [PATCH 010/178] [None][chore] Trim the compression-mode docstring Signed-off-by: tianruih --- tensorrt_llm/_torch/kv_cache_compression/interface.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index 8ac47c98a640..cd4e7bf32068 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -8,9 +8,8 @@ class KvCacheCompressionMode(IntEnum): """Algorithm-level traits of a KV-cache compression method. - Mirrors ``SpeculativeDecodingMode``: configs map their ``algorithm`` - string to a member here, and callers read ``is_*`` predicates instead of - comparing strings. + Configs map their ``algorithm`` string to a member here; callers read the + ``is_*`` predicates instead of comparing strings. """ NONE = auto() From 51afc9965f069e97099d57c4097959dc76e97a7e Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 21:49:05 -0700 Subject: [PATCH 011/178] [None][chore] Describe the real compressed-token mechanism in the config docstring Signed-off-by: tianruih --- tensorrt_llm/llmapi/llm_args.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 10f98d6cb1a8..020f4986cc85 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3289,8 +3289,8 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): it is converted to the runtime schema at load. TriAttention is a pure compression method: it has no sparse-attention config and no attention backend of its own -- decode runs the model's standard attention over the - compacted cache, and the manager reconciles the cached-token count via the - framework's ``adjust_attention_metadata`` hook. + compacted cache, and the manager publishes each request's evicted count on + ``LlmRequest.py_num_compressed_tokens`` for the engine to subtract. """ algorithm: Literal["triattention"] = "triattention" eviction_mode: Literal["union", "per_head", "per_layer_perhead"] = Field( From c4d80d3523c029e7d396786ab31bfbb15506d4b2 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 22:04:37 -0700 Subject: [PATCH 012/178] [None][fix] Fail fast on speculative modes compression cannot support Signed-off-by: tianruih --- tensorrt_llm/_torch/pyexecutor/_util.py | 52 ++++++++----------- .../test_kv_cache_compression_manager.py | 7 +-- 2 files changed, 27 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 8ce2d2564f85..673ad835a67e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2032,30 +2032,24 @@ def _create_kv_cache_manager( return kv_cache_manager -def kv_cache_compression_supported_with_spec( +def validate_kv_cache_compression_with_spec( config: KvCacheCompressionConfig, spec_config: Optional[SpeculativeConfig], draft_kv_cache_manager: Optional[KVCacheManagerV2], -) -> bool: - """Decide, before any manager is created, whether the configured - compression method supports the speculative setup. Logs the reason and - returns False when it does not; the run then stays uncompressed.""" +) -> 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()): - return True - # Evicting methods co-compact the draft KV, so they only support spec - # modes whose draft KV is a standard paged cache in the same forward - # (one-model speculation). + return + # Evicting methods co-compact the draft KV, so the draft must be a + # standard paged cache in the same forward (one-model speculation). mode = spec_config.spec_dec_mode if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): - logger.warning( - "KV-cache compression algorithm '%s' evicts cached tokens and " - "does not support speculative decoding mode %s (the draft KV must " - "be a standard paged cache compacted together with the target); " - "running without a compression manager.", config.algorithm, - mode.name) - return False - return True + 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).") def create_kv_cache_compression_manager( @@ -2069,7 +2063,7 @@ 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 - decided by the caller via ``kv_cache_compression_supported_with_spec``. + checked by the caller via ``validate_kv_cache_compression_with_spec``. """ logger.warning( "KV-cache compression algorithm '%s' is not registered; running without " @@ -2285,17 +2279,17 @@ def create_py_executor_instance( if kv_cache_compression_config is not None: draft_kv_cache_manager = resources.get( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) - if kv_cache_compression_supported_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, - draft_kv_cache_manager=draft_kv_cache_manager, - ) - if compression_manager is not None: - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_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, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if compression_manager is not None: + resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( + compression_manager) resource_manager = ResourceManager(resources) 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 8a9d16974146..5c6bed5af1b6 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -304,14 +304,15 @@ def test_eviction_method_predicate_defaults_false(self): assert not hasattr(m, "spec_config") def test_spec_gate_only_restricts_eviction_methods(self): - from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec 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") spec_config = SimpleNamespace(spec_dec_mode=SpeculativeDecodingMode.DFLASH) - assert kv_cache_compression_supported_with_spec(config, spec_config, None) is True - assert kv_cache_compression_supported_with_spec(config, None, None) is True + validate_kv_cache_compression_with_spec(config, spec_config, None) + validate_kv_cache_compression_with_spec(config, None, None) # ---------------------------------------------------------------------- # From 8dab5cc82a1a63f54977b0609b8c4ced4863af59 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 15 Jul 2026 22:06:47 -0700 Subject: [PATCH 013/178] [None][fix] Fail fast on speculative modes compression cannot support Signed-off-by: tianruih --- .../test_triattention_draft_cocompaction.py | 12 ++++---- .../test_triattention_pipeline.py | 30 +++++++------------ 2 files changed, 16 insertions(+), 26 deletions(-) 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 index c53bade407d5..4ea7527ee08c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -321,22 +321,20 @@ def test_draft_admission_gates_raise(gate, match): draft_manager.max_attention_window_vec = [128] if gate == "dflash": # DFlash reads cross-attention context buffers, not a paged KV cache; - # the call-site speculative gate declines to create a manager and the - # run stays uncompressed. - from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + # the call-site speculative gate rejects before any manager is + # created. + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, TriAttentionKvCacheCompressionConfig, ) - assert ( - kv_cache_compression_supported_with_spec( + with pytest.raises(ValueError, match=match): + validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), DFlashDecodingConfig(max_draft_len=3), draft_manager, ) - is False - ) return manager = TriAttention( _make_fake_v2(), diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 4b32b4164341..72179ee2733d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -893,7 +893,7 @@ def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): # the factory then builds a manager that validates cleanly. from tensorrt_llm._torch.pyexecutor._util import ( create_kv_cache_compression_manager, - kv_cache_compression_supported_with_spec, + validate_kv_cache_compression_with_spec, ) from tensorrt_llm.llmapi.llm_args import ( MTPDecodingConfig, @@ -903,11 +903,8 @@ def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): config = TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8) assert config.kv_cache_compression_mode.is_eviction_method() is True draft_manager = _make_fake_v2(is_draft=True) - assert ( - kv_cache_compression_supported_with_spec( - config, MTPDecodingConfig(max_draft_len=1), draft_manager - ) - is True + validate_kv_cache_compression_with_spec( + config, MTPDecodingConfig(max_draft_len=1), draft_manager ) manager = create_kv_cache_compression_manager( config, @@ -928,38 +925,33 @@ def test_unvalidated_paged_draft_tail_contracts_remain_fail_closed(self, mode): ) else: spec_config = PARDDecodingConfig(max_draft_len=3) - # The call-site speculative gate declines before any manager is - # created; the run stays uncompressed. - from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + # The call-site speculative gate rejects before any manager is created. + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig - assert ( - kv_cache_compression_supported_with_spec( + with pytest.raises(ValueError, match="standard paged cache compacted together"): + validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), spec_config, _make_fake_v2(is_draft=True), ) - is False - ) def test_dflash_spec_mode_is_rejected(self): # Policy: the DFlash draft reads cross-attention context buffers, not # a paged KV cache, so compression cannot cover it. The call-site - # speculative gate declines to create a manager. - from tensorrt_llm._torch.pyexecutor._util import kv_cache_compression_supported_with_spec + # speculative gate rejects before any manager is created. + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, TriAttentionKvCacheCompressionConfig, ) - assert ( - kv_cache_compression_supported_with_spec( + with pytest.raises(ValueError, match="standard paged cache compacted together"): + validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), DFlashDecodingConfig(max_draft_len=3), _make_fake_v2(is_draft=True), ) - is False - ) def test_prepare_snapshots_fixed_linear_generation_growth(self): manager = _make_fake_v2() From cf1b3ad6074542ab892af9e14fddcd391dfd2074 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 01:12:16 -0700 Subject: [PATCH 014/178] [None][chore] Tighten the compression manager step-hook surface update_resources declared attn_metadata/kv_cache_dtype_byte_size but the executor dispatches those extras to the KV cache manager slot only, and the method body never read them; drop the dead parameters and the docstring sentence that claimed a transparent pass-through. Also document that on_generation_step_begin fires on context-only iterations too, and that it should be split into context/generation-step hooks once an algorithm needs the distinction. Signed-off-by: tianruih --- .../_torch/pyexecutor/resource_manager.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 58ec6ab41485..c299169cebb6 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2410,7 +2410,11 @@ def on_generation_step_begin( scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per generation step before this step's forward.""" + """Fired once per iteration before the forward pass, including + iterations that schedule only context requests. No algorithm overrides + it yet; split it into separate context-step and generation-step hooks + once a subclass needs the distinction. + """ def on_generation_step_end( self, @@ -2456,12 +2460,7 @@ def prepare_resources(self, scheduled_batch: "ScheduledRequests") -> None: self.on_request_init(req) self.on_generation_step_begin(scheduled_batch) - def update_resources( - self, - scheduled_batch: "ScheduledRequests", - attn_metadata: Optional["AttentionMetadata"] = None, - kv_cache_dtype_byte_size: Optional[float] = None, - ) -> None: + def update_resources(self, scheduled_batch: "ScheduledRequests") -> None: """Fire :meth:`on_context_step_end` with the requests whose final prefill chunk ran this iteration, then :meth:`on_generation_step_end`. @@ -2470,9 +2469,7 @@ def update_resources( request-state transitions: it is iteration-exact and immune to a short-output request going straight to ``GENERATION_TO_COMPLETE`` (which, under the overlap scheduler, never passes through - ``GENERATION_IN_PROGRESS``). Signature matches the other resource - managers so PyExecutor passes ``attn_metadata`` / - ``kv_cache_dtype_byte_size`` through transparently. + ``GENERATION_IN_PROGRESS``). """ if scheduled_batch.context_requests_last_chunk: self.on_context_step_end( From 1cda81063caef98ad89944a8592ff0c5e16a46e2 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 01:13:45 -0700 Subject: [PATCH 015/178] [None][refactor] Tighten the KV-cache compression config surface - Dispatch the config union on the algorithm tag (discriminated union, like sparse attention and speculative configs) and drop the abstract base from the union: a mistyped algorithm now fails config validation with the list of valid tags instead of silently validating as the base class. The factory hard-cast that guarded against base instances is gone with it. - Rename the base class to BaseKvCacheCompressionConfig and give the plain KvCacheCompressionConfig name to the user-facing union alias, matching the sparse/speculative naming convention. - Raise on a config whose algorithm has no registered manager instead of warning and silently running without compression. - Require model_path and calibration_path at config validation time; both are consumed at manager construction and previously failed deep in executor setup (or worse, fell back to an unrelated calibration). - Drop the dead window_size field (nothing reads it; it also occupied a telemetry manifest entry) and its README row, and fix the README rows that still described a removed eviction mode. - Add gt=0 to top_B/beta and collect the manager constructor kwargs in KvCacheCompressionConfig.to_manager_kwargs(). Signed-off-by: tianruih --- examples/triattention/README.md | 12 +++--- tensorrt_llm/_torch/pyexecutor/_util.py | 38 +++++----------- tensorrt_llm/llmapi/llm_args.py | 43 ++++++++++++++----- .../test_kv_cache_compression_manager.py | 29 +++++-------- .../test_rope_fusion_gate.py | 6 ++- .../test_triattention_draft_cocompaction.py | 3 +- .../test_triattention_pipeline.py | 36 ++++++++++------ .../api_stability/references/llm.yaml | 2 +- 8 files changed, 88 insertions(+), 81 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index c152607c5006..97deea79eb81 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -116,13 +116,11 @@ trtllm-eval --model --config config.yaml longbench_v2 --max_outp * **`top_B`** (int, default=1024): Tokens kept at each eviction (the upstream `budget`). Prompt tokens are always preserved on top of this. Smaller `top_B` → 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=`per_layer`): Which token set each eviction keeps. - * `per_layer`: score a layer, average over heads, keep one set per layer (the simplest variant). - * `per_head`: each KV head keeps its own set, shared across layers (mean of per-layer maxima). The upstream AIME default. +* **`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. - * `union`: union of the per-head top-k, then re-ranked. -* **`window_size`** (int, default=128): Most-recent tokens always preserved from eviction. Prevents the scorer from evicting freshly generated tokens. -* **`normalize_scores`** (bool, default=True): Z-normalize each head's scores over the decode region before selection (upstream default). Used by `per_head` / `per_layer_perhead` / `union`; ignored by `per_layer`. -* **`pin_prefill`** (bool, default=True): Always preserve the prompt (prefill) tokens; only decode tokens compete for the budget (upstream behaviour). Used by `per_head` / `per_layer_perhead` / `union`; `per_layer` uses the recency window instead. +* **`normalize_scores`** (bool, default=True): Z-normalize each head's scores over the decode region before selection (upstream default). +* **`pin_prefill`** (bool, default=True): Always preserve the prompt (prefill) tokens; only decode tokens compete for the budget (upstream behaviour). * **`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 only to derive the model's RoPE tables when converting the official calibration file. diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b9c57cd7d130..44e8985da674 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -32,7 +32,7 @@ CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, PeftCacheConfig, SamplerType, SchedulerConfig, SparseAttentionConfig, SpeculativeConfig, - TorchLlmArgs, TriAttentionKvCacheCompressionConfig, WaitingQueuePolicy) + TorchLlmArgs, WaitingQueuePolicy) # isort: on from tensorrt_llm.logger import logger from tensorrt_llm.lora_helper import (LoraConfig, @@ -2077,42 +2077,27 @@ def create_kv_cache_compression_manager( config: KvCacheCompressionConfig, kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, -) -> Optional[BaseKVCacheCompressionManager]: - """Build the KV-cache compression manager for ``config.algorithm``, or return - None if no algorithm matches. +) -> BaseKVCacheCompressionManager: + """Build the KV-cache compression manager for ``config.algorithm``. 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. Speculative-decoding compatibility is checked by the caller via + ``validate_kv_cache_compression_with_spec``. """ if config.algorithm == "triattention": from tensorrt_llm._torch.kv_cache_compression.triattention import \ TriAttention - triattention_config = ( - config if isinstance(config, TriAttentionKvCacheCompressionConfig) - else TriAttentionKvCacheCompressionConfig.model_validate( - config.model_dump())) return TriAttention( kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, - top_B=triattention_config.top_B, - beta=triattention_config.beta, - model_path=triattention_config.model_path, - calibration_path=triattention_config.calibration_path, - eviction_mode=triattention_config.eviction_mode, - normalize_scores=triattention_config.normalize_scores, - pin_prefill=triattention_config.pin_prefill, - count_prompt_tokens=triattention_config.count_prompt_tokens, + **config.to_manager_kwargs(), ) - logger.warning( - "KV-cache compression algorithm '%s' is not registered; running without " - "a compression manager.", - config.algorithm, - ) - return None + raise ValueError( + f"KV-cache compression algorithm {config.algorithm!r} has a config but " + "no registered compression manager.") def create_py_executor_instance( @@ -2329,9 +2314,8 @@ def create_py_executor_instance( kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, ) - if compression_manager is not None: - resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( - compression_manager) + resources[ResourceManagerType.KV_CACHE_COMPRESSION_MANAGER] = ( + compression_manager) resource_manager = ResourceManager(resources) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 020f4986cc85..c1cde0329baf 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3253,7 +3253,7 @@ def supports_backend(self, backend: str) -> bool: ] -class KvCacheCompressionConfig(StrictBaseModel): +class BaseKvCacheCompressionConfig(StrictBaseModel): """Config for KV-cache compression: a compression manager runs a KV-reduction algorithm (e.g. periodic token eviction) alongside KVCacheManagerV2. @@ -3277,7 +3277,7 @@ def kv_cache_compression_mode(self): return KvCacheCompressionMode.from_string(self.algorithm) -class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): +class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): """KV-cache compression config for TriAttention. TriAttention periodically evicts cached tokens during generation, guided by @@ -3312,10 +3312,12 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): "tokens compete for the budget (upstream behaviour).") top_B: int = Field( default=2048, + gt=0, description="Tokens kept at each periodic eviction (upstream `budget`; " "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." @@ -3331,21 +3333,40 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): "(produced by github.com/WeianMao/triattention). TRT-LLM does not " "compute calibration; it converts this file to the runtime schema at " "load.") - window_size: int = Field( - default=128, - description="Compatibility field retained for existing configs. The " - "implemented calibration-based selection does not use a separate " - "recency window.") count_prompt_tokens: bool = Field( default=False, description="If False (default), the KV budget counts only DECODE tokens " "(the pinned prompt is kept on top). Physical capacity reclaim currently " "requires False.") + @model_validator(mode="after") + def _require_calibration_inputs(self): + # Both paths are consumed at manager construction; failing here surfaces + # the error at config-validation time instead of deep in executor setup. + if not self.model_path or not self.calibration_path: + raise ValueError( + "TriAttention requires both model_path and calibration_path; " + "TRT-LLM consumes an official calibration file and does not " + "compute one.") + return self -KvCacheCompressionConfigType: TypeAlias = Union[ - TriAttentionKvCacheCompressionConfig, - KvCacheCompressionConfig, + def to_manager_kwargs(self) -> dict: + """Constructor kwargs for the TriAttention manager.""" + return { + "top_B": self.top_B, + "beta": self.beta, + "model_path": self.model_path, + "calibration_path": self.calibration_path, + "eviction_mode": self.eviction_mode, + "normalize_scores": self.normalize_scores, + "pin_prefill": self.pin_prefill, + "count_prompt_tokens": self.count_prompt_tokens, + } + + +KvCacheCompressionConfig: TypeAlias = Annotated[ + Union[TriAttentionKvCacheCompressionConfig], + Field(discriminator="algorithm"), ] @@ -4067,7 +4088,7 @@ class BaseLlmArgs(StrictBaseModel): # KV cache compression config (separate from sparse attention: changes which # KV is stored, not the attention computation) - kv_cache_compression_config: Optional[KvCacheCompressionConfigType] = Field( + kv_cache_compression_config: Optional[KvCacheCompressionConfig] = Field( default=None, description="KV-cache compression config; None disables compression.", status="prototype") 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 5c6bed5af1b6..53a9b9bd7eee 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -28,7 +28,6 @@ import pytest -from tensorrt_llm._torch.pyexecutor import _util as util_mod from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseKVCacheCompressionManager, @@ -265,40 +264,34 @@ def test_free_fires_finish(self, fake_kv_cache_manager): class TestFactory: - def test_returns_none_when_no_algorithm_registered(self, fake_kv_cache_manager): - # Framework-only: no concrete algorithm ships, so any config -> None. + def test_raises_when_no_algorithm_registered(self, fake_kv_cache_manager): + # A config whose algorithm has no registered manager is a developer + # error (config subclass added without a factory branch): fail loudly + # instead of silently running without compression. cfg = MagicMock() cfg.algorithm = "made_up_method" - assert create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) is None - - def test_warns_for_unregistered_algorithm(self, fake_kv_cache_manager): - cfg = MagicMock() - cfg.algorithm = "made_up_method" - with patch.object(util_mod, "logger") as mock_logger: + with pytest.raises(ValueError, match="no registered compression manager"): create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) - mock_logger.warning.assert_called_once() - def test_factory_accepts_independent_draft_manager(self): + def test_unregistered_algorithm_raises_with_draft_manager_too(self): cfg = MagicMock() cfg.algorithm = "made_up_method" target = _v2_manager(is_draft=False) draft = _v2_manager(is_draft=True) - assert ( + with pytest.raises(ValueError, match="no registered compression manager"): create_kv_cache_compression_manager( cfg, target, draft_kv_cache_manager=draft, ) - 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 + from tensorrt_llm.llmapi.llm_args import BaseKvCacheCompressionConfig - config = KvCacheCompressionConfig(algorithm="offload") + config = BaseKvCacheCompressionConfig(algorithm="offload") assert config.kv_cache_compression_mode.is_eviction_method() is False m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) assert not hasattr(m, "spec_config") @@ -306,10 +299,10 @@ def test_eviction_method_predicate_defaults_false(self): def test_spec_gate_only_restricts_eviction_methods(self): from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig + from tensorrt_llm.llmapi.llm_args import BaseKvCacheCompressionConfig # Non-evicting methods pass with any speculative mode; no exception. - config = KvCacheCompressionConfig(algorithm="offload") + config = BaseKvCacheCompressionConfig(algorithm="offload") 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) 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 index 4989508f73c7..b68e2be63b51 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py +++ b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py @@ -39,7 +39,8 @@ def test_plain_attention_defaults_to_fused_rope() -> None: def test_kv_cache_compression_forces_unfused_rope() -> None: - model_config = ModelConfig(kv_cache_compression_config=TriAttentionKvCacheCompressionConfig()) + 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 @@ -71,7 +72,8 @@ def test_unfused_yarn_rope_is_applied_exactly_once() -> None: ), is_neox=True, ) - model_config = ModelConfig(kv_cache_compression_config=TriAttentionKvCacheCompressionConfig()) + model_config = ModelConfig(kv_cache_compression_config=TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt")) attn = Attention( hidden_size=256, num_attention_heads=8, 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 index 4ea7527ee08c..a946ffb5fae5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -331,7 +331,8 @@ def test_draft_admission_gates_raise(gate, match): with pytest.raises(ValueError, match=match): validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8), DFlashDecodingConfig(max_draft_len=3), draft_manager, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 72179ee2733d..1f32fa261a3e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -49,10 +49,7 @@ from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm.llmapi.llm_args import ( - KvCacheCompressionConfig, - TriAttentionKvCacheCompressionConfig, -) +from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig _TORCH_TOPK_ORACLE = torch.topk @@ -349,7 +346,11 @@ def test_llm_args_dispatches_concrete_and_unknown_algorithms(self): tri_args = TorchLlmArgs( model="dummy", - kv_cache_compression_config={"algorithm": "triattention"}, + kv_cache_compression_config={ + "algorithm": "triattention", + "model_path": "/models/test", + "calibration_path": "/calib/test.pt", + }, ) assert isinstance( tri_args.kv_cache_compression_config, @@ -358,11 +359,13 @@ def test_llm_args_dispatches_concrete_and_unknown_algorithms(self): assert tri_args.kv_cache_compression_config.top_B == 2048 assert tri_args.kv_cache_compression_config.beta == 128 - unknown_args = TorchLlmArgs( - model="dummy", - kv_cache_compression_config={"algorithm": "future_method"}, - ) - assert type(unknown_args.kv_cache_compression_config) is KvCacheCompressionConfig + # The union dispatches on the algorithm tag, so an unknown algorithm + # fails config validation instead of falling back to a base config. + with pytest.raises(ValidationError): + TorchLlmArgs( + model="dummy", + kv_cache_compression_config={"algorithm": "future_method"}, + ) def test_eviction_mode_validated(self): with pytest.raises(ValidationError): @@ -900,7 +903,8 @@ def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): TriAttentionKvCacheCompressionConfig, ) - config = TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8) + config = TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8) assert config.kv_cache_compression_mode.is_eviction_method() is True draft_manager = _make_fake_v2(is_draft=True) validate_kv_cache_compression_with_spec( @@ -931,7 +935,8 @@ def test_unvalidated_paged_draft_tail_contracts_remain_fail_closed(self, mode): with pytest.raises(ValueError, match="standard paged cache compacted together"): validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8), spec_config, _make_fake_v2(is_draft=True), ) @@ -948,7 +953,8 @@ def test_dflash_spec_mode_is_rejected(self): with pytest.raises(ValueError, match="standard paged cache compacted together"): validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig(model_path="/models/test", top_B=8), + TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8), DFlashDecodingConfig(max_draft_len=3), _make_fake_v2(is_draft=True), ) @@ -1851,7 +1857,8 @@ def test_returns_triattention_instance_with_v2(self): # Calibration is deferred to the first request, so construction needs # no calibration file or CUDA. fake_v2 = _make_fake_v2(enable_block_reuse=False) - cfg = TriAttentionKvCacheCompressionConfig(top_B=32, beta=16, model_path="/models/test") + cfg = TriAttentionKvCacheCompressionConfig( + top_B=32, beta=16, model_path="/models/test", calibration_path="/calib/test.pt") mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) assert mgr.top_B == 32 @@ -1864,6 +1871,7 @@ def test_factory_propagates_eviction_mode(self): beta=8, eviction_mode="per_head", model_path="/models/test", + calibration_path="/calib/test.pt", ) mgr = create_kv_cache_compression_manager( cfg, kv_cache_manager=_make_fake_v2(enable_block_reuse=False) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index e2059936849b..7651a7a1cdb8 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -252,7 +252,7 @@ methods: default: null status: prototype kv_cache_compression_config: - annotation: Union[tensorrt_llm.llmapi.llm_args.TriAttentionKvCacheCompressionConfig, tensorrt_llm.llmapi.llm_args.KvCacheCompressionConfig, NoneType] + annotation: Union[tensorrt_llm.llmapi.llm_args.TriAttentionKvCacheCompressionConfig, NoneType] default: null status: prototype otlp_traces_endpoint: From 10ffc1aef3e706da75f7a6b8676e680d7c268862 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 01:22:44 -0700 Subject: [PATCH 016/178] [None][chore] Regenerate the LLM args telemetry manifest The dropped window_size field leaves the manifest; no other entries change. Signed-off-by: tianruih --- tensorrt_llm/usage/llm_args_golden_manifest.json | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index c491b8d30315..374ab751ae52 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -598,13 +598,6 @@ "kind": "value", "path": "kv_cache_compression_config.top_B" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_cache_compression_config.window_size" - }, { "allowed_values": [], "annotation": "", @@ -3078,13 +3071,6 @@ "kind": "value", "path": "kv_cache_compression_config.top_B" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_cache_compression_config.window_size" - }, { "allowed_values": [], "annotation": "", From cd94f7646712b7be4ad36611ee550945505ade6c Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 05:14:09 -0700 Subject: [PATCH 017/178] [None][perf] Put score-kernel token tiles on the fastest grid axis Adjacent programs now walk consecutive K pages of one (request, layer, head) and reuse its calibration/phase rows in L2. Median score-kernel time drops 1.8-2.0% across two model geometries and both sequential and shuffled page layouts, bit-exact against the segment-major order. Segments move to the y grid axis, which CUDA caps at 65535; both launch sites validate the request*layer count against that limit so the unbounded x axis stays with the token tiles of long sequences. Signed-off-by: tianruih --- .../triattention/triattention.py | 7 ++++++- .../triattention/triattention_kernels.py | 13 ++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 748ff64ed88a..d7c2b7650c83 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1180,9 +1180,14 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str self._phase_args = (*phase_pointer_args, *phase_constants) self._score_args = (*score_pointer_args, *score_geometry, *score_constants) phase_grid = (self.max_requests, 1, 1) + score_segments = self.max_requests * group.num_layers + if score_segments > 65535: + # Segments sit on the y grid axis (CUDA caps y/z at 65535) so the + # unbounded x axis can hold the token tiles of long sequences. + raise ValueError("request*layer segment count exceeds the CUDA grid limit") score_grid = ( - self.max_requests * group.num_layers, group.max_ntblk, + score_segments, group.num_kv_heads, ) with torch.cuda.device(self.device): diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index df0d4dca0680..e988cf28d548 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -610,8 +610,11 @@ def _tri_score_perhead_kernel( T_BLOCK: tl.constexpr, F_BLOCK: tl.constexpr, ): - seg = tl.program_id(0) - t_blk = tl.program_id(1) + # Token tiles ride the fastest grid axis: adjacent programs then walk + # consecutive K pages of one (request, layer, head) and reuse its + # calibration/phase rows in L2 (~2% faster than segment-major order). + seg = tl.program_id(1) + t_blk = tl.program_id(0) # KV heads are grid-parallel (axis 2): iterations of the former kv_head # loop shared NO data (each KV head reads its own K and writes its own # output rows), so hoisting it onto the grid multiplies parallelism with @@ -908,9 +911,13 @@ def launch( ): raise ValueError("score output lengths do not fit the keep-set selector") num_segments = request_count * self.num_layers + if num_segments > 65535: + # Segments sit on the y grid axis (CUDA caps y/z at 65535) so the + # unbounded x axis can hold the token tiles of long sequences. + raise ValueError("request*layer segment count exceeds the CUDA grid limit") output = self.output[:request_count] _launch_tri_score_perhead( - (num_segments, self.max_ntblk, self.num_kv_heads), + (self.max_ntblk, num_segments, self.num_kv_heads), ( *self.pointer_prefix, valid_seq_lens, From 0f128a0cfde26766af05a9f763dc6a7cff87c8c2 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 05:56:33 -0700 Subject: [PATCH 018/178] [None][chore] Polish the batched compaction module - Promote the prepared Triton launch helper to its own module with a public name (it also backs the score, stats, and union selection launches) and validate the constexpr ordering against the kernel declaration at build time instead of by comment. - Fail fast on combinations that previously degraded silently: selection tensors on the wrong device, swa_window without SWA layers, draft pools without draft layers, and (in the C++ op) flat source indices combined with per-layer source indices, with a unit test for the rejected op combination. - Name the pack kernel launch-shape constants, split combined error messages, rename the pack launch keep_count to decode_keep_count, drop the grouping parameter that never affected grouping, and document the constructor staging-snapshot inputs. - Deduplicate the block-offset encoder into the test conftest, reword generic C++ comments that named one consumer, and document the V2 2*page+plane offset decode in the updater kernel. Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 3 + .../thop/sparseKvCacheCompactOp.cpp | 24 +-- .../triattention/compaction.py | 146 +++++++++--------- .../triattention/prepared_launch.py | 82 ++++++++++ .../triattention/triattention.py | 10 +- .../_torch/kv_cache_compression/conftest.py | 37 +++++ .../test_triattention_draft_cocompaction.py | 21 +-- .../test_triattention_eager.py | 16 +- .../serial/test_sparse_kv_cache_compact.py | 17 ++ 9 files changed, 232 insertions(+), 124 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py create mode 100644 tests/unittest/_torch/kv_cache_compression/conftest.py diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 4dd7d2fb8839..df59045fee31 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1774,6 +1774,9 @@ struct KvCacheV2LayersBuffer return reinterpret_cast(static_cast(poolPointers[blockIdx.x])); } + // KVCacheManagerV2 block offsets encode page and K/V plane as + // 2*page + plane (K = 2p, V = 2p + 1); this table carries the K plane, + // so dividing by 2 recovers the page for both halves. __device__ __forceinline__ void* getKBlockPtr(int32_t batchIdx, int32_t tokenIdx) const { int32_t const blockOffset = pageTable[batchIdx * pageTableRequestStride + tokenIdx / tokensPerBlock]; diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index be7dc7d75bc3..733ea45516ee 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -33,10 +33,10 @@ namespace torch_ext { //! Adapt one uniform group of KVCacheManagerV2 HND layer pools to the -//! existing sparse-KV post-FMHA updater. Layers share one V2 block-offset table; -//! destinationBase replaces the former arbitrary -//! destination tensor because every TriAttention move targets one interval. -//! Within each request and KV head, TriAttention supplies increasing source +//! existing sparse-KV post-FMHA updater. Layers share one V2 block-offset +//! table; destinationBase replaces the former arbitrary destination tensor +//! because every compaction move targets one contiguous interval. Within +//! each request and KV head, the caller must supply increasing source //! ordinals with destinationBase + move <= source[move], which makes the //! updater's forward tiled in-place copy safe. void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor const& poolPointers, @@ -75,13 +75,10 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor TORCH_CHECK( pageTable.get_device() == device, "sparse_kv_cache_compact_layers: block offsets must be on the pool device"); - auto const checkPointerArray = [device, numLayers](th::Tensor const& pointers, char const* name) - { - TORCH_CHECK(pointers.is_cuda() && pointers.get_device() == device && pointers.scalar_type() == th::kInt64 - && pointers.dim() == 1 && pointers.size(0) == numLayers && pointers.is_contiguous(), - "sparse_kv_cache_compact_layers: ", name, " must be contiguous CUDA int64 [num_layers]"); - }; - checkPointerArray(poolPointers, "pool_pointers"); + TORCH_CHECK(poolPointers.is_cuda() && poolPointers.get_device() == device + && poolPointers.scalar_type() == th::kInt64 && poolPointers.dim() == 1 + && poolPointers.size(0) == numLayers && poolPointers.is_contiguous(), + "sparse_kv_cache_compact_layers: pool_pointers must be contiguous CUDA int64 [num_layers]"); TORCH_CHECK(sourceIndices.is_cuda() && sourceIndices.get_device() == device && sourceIndices.scalar_type() == th::kInt32 && sourceIndices.is_contiguous() @@ -94,6 +91,8 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor { TORCH_CHECK(sourceIndices.size(0) == numKvHeads, "sparse_kv_cache_compact_layers: source_indices KV-head dimension mismatch"); + TORCH_CHECK(!sourceLayerIndices.has_value(), + "sparse_kv_cache_compact_layers: source_layer_indices require 3-D per-layer source_indices"); } else { @@ -113,6 +112,9 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor sourceLayerPtr = layerIndices.data_ptr(); } + // The last element of source_offsets is the total move count and must equal + // source_indices.size(-1); it lives on device, so checking it here would + // force a sync -- the kernel trusts the caller. TORCH_CHECK(sourceOffsets.is_cuda() && sourceOffsets.get_device() == device && sourceOffsets.scalar_type() == th::kInt32 && sourceOffsets.is_contiguous() && sourceOffsets.dim() == 1 && sourceOffsets.size(0) == batchSize + 1, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 0ba1292636a0..fa278b661785 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -27,56 +27,14 @@ import torch -_SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) - - -class _PreparedTritonKernelLaunch: - """Replay one Triton kernel launch frozen at build time. +from .prepared_launch import PreparedTritonKernelLaunch - ``warmup`` JIT-compiles the kernel once for a fixed grid, bound tensor - set, and constexpr set; ``__call__`` re-launches the compiled binary - directly, skipping Triton's per-call dispatch. Constexpr values are - passed positionally on replay, so their dict order must match the - kernel's constexpr parameter declaration order. - """ +_SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) - def __init__( - self, - triton_kernel, - bound_tensors: Tuple[torch.Tensor, ...], - constexpr_values: Dict[str, object], - *, - grid: Tuple[int, ...], - num_warps: int, - ) -> None: - self.device = bound_tensors[0].device - self.bound_tensors = tuple(bound_tensors) - self.constexpr_values = dict(constexpr_values) - with torch.cuda.device(self.device): - self.build_stream = torch.cuda.current_stream(self.device) - compiled = triton_kernel.warmup( - *self.bound_tensors, - **self.constexpr_values, - num_warps=num_warps, - grid=grid, - ) - self.compiled_kernel_runner = compiled[grid] - - def __call__(self, *replay_tensors: torch.Tensor) -> None: - """Replay the launch; ``replay_tensors``, if given, substitute the bound tensors.""" - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.build_stream.device, - self.build_stream.cuda_stream, - ): - raise RuntimeError( - "a prepared Triton kernel launch must run on the stream it was built on" - ) - self.compiled_kernel_runner( - *(replay_tensors if replay_tensors else self.bound_tensors), - *self.constexpr_values.values(), - stream=self.build_stream.cuda_stream, - ) +# Launch shape of the move-index packing kernel: tokens per program along the +# move axis, and its warp count. +_PACK_BLOCK_TOKENS = 256 +_PACK_NUM_WARPS = 4 class _CppCompactGroup(NamedTuple): @@ -108,7 +66,7 @@ class _SingleCacheCompaction(NamedTuple): tokens land at. """ - prepared_move_index_pack: Optional[_PreparedTritonKernelLaunch] + prepared_move_index_pack: Optional[PreparedTritonKernelLaunch] cpp_launch_groups: Tuple[_CppCompactGroup, ...] move_source_indices: torch.Tensor move_source_offsets: torch.Tensor @@ -144,7 +102,13 @@ def _validated_kv_head_count( launch's pool shape, so every layer on one side must agree on it. """ first = pools[layers[0]] - num_kv_heads = int(first.shape[2]) if first.ndim == 5 and first.shape[2] > 0 else -1 + if first.ndim != 5 or first.shape[2] <= 0: + raise ValueError( + f"{what} pools must be 5-D interleaved V2 pools " + f"[pages, K/V, heads, tokens, dim]; layer {layers[0]} has shape " + f"{tuple(first.shape)}" + ) + num_kv_heads = int(first.shape[2]) if not all( pools[layer].ndim == 5 and pools[layer].shape[1] == 2 @@ -181,8 +145,10 @@ def _validated_tail_lengths( ) -> Tuple[int, ...]: if tail_lengths is None: tail_lengths = [0] * request_count - if len(tail_lengths) != request_count or any(length < 0 for length in tail_lengths): + if len(tail_lengths) != request_count: raise ValueError(f"{what} lengths must match the request count") + if any(length < 0 for length in tail_lengths): + raise ValueError(f"{what} lengths must be non-negative") return tuple(int(length) for length in tail_lengths) @@ -212,7 +178,6 @@ def page_table_for(representative: int) -> torch.Tensor: def _compact_groups( entries: List[Tuple[int, torch.Tensor, torch.Tensor]], - mode: str, pool_keys: Tuple[object, ...], device: torch.device, per_layer_slots: Optional[Dict[int, int]] = None, @@ -226,7 +191,6 @@ def _compact_groups( for layer, pool, page_table in entries: key = ( pool_keys[layer], - mode, str(pool.dtype), str(pool.device), tuple(int(value) for value in pool.shape[1:]), @@ -273,14 +237,14 @@ def _prepared_move_index_pack_launch( *, eviction_mode: str, prompt_len: int, - keep_count: int, + decode_keep_count: int, num_dense_layers: int, num_kv_heads: int, max_protected_tail: int, swa_window: int, swa_move_source_offsets: Optional[torch.Tensor], swa_move_source_indices: Optional[torch.Tensor], -) -> _PreparedTritonKernelLaunch: +) -> PreparedTritonKernelLaunch: """Build one prepared launch of the move-index packing kernel. The kernel reads the kept-token ordinals and each request's valid length @@ -298,12 +262,18 @@ def _prepared_move_index_pack_launch( else: selection_rows = num_kv_heads selection_prefix = (request_count,) if union else (request_count, selection_rows) + expected_selection = (*selection_prefix, prompt_len + decode_keep_count) if ( request_count <= 0 - or tuple(kept_token_ordinals.shape) != (*selection_prefix, prompt_len + keep_count) + or tuple(kept_token_ordinals.shape) != expected_selection or valid_sequence_lengths.shape != (request_count,) ): - raise ValueError("prepared compaction packing requires one valid fixed geometry") + raise ValueError( + f"prepared compaction packing expects kept ordinals of shape " + f"{expected_selection} and one valid length per request; got " + f"{tuple(kept_token_ordinals.shape)} and " + f"{tuple(valid_sequence_lengths.shape)}" + ) device = kept_token_ordinals.device if not _cuda_int32_contiguous((kept_token_ordinals, valid_sequence_lengths), device): @@ -321,12 +291,15 @@ def _prepared_move_index_pack_launch( from .triattention_kernels import _pack_compaction_sources_kernel - block = 256 - max_move = keep_count + max_protected_tail + max_move = decode_keep_count + max_protected_tail if swa_total: max_move = max(max_move, swa_window + max_protected_tail) packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads - grid = (request_count, packed_row_count, (max_move + block - 1) // block) + grid = ( + request_count, + packed_row_count, + (max_move + _PACK_BLOCK_TOKENS - 1) // _PACK_BLOCK_TOKENS, + ) bound_tensors = ( kept_token_ordinals, valid_sequence_lengths, @@ -341,22 +314,22 @@ def _prepared_move_index_pack_launch( DENSE_TOTAL=int(move_source_indices.shape[-1]), SWA_TOTAL=swa_total, SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=prompt_len + keep_count, - KEEP_COUNT=keep_count, + SELECTION_STRIDE=prompt_len + decode_keep_count, + KEEP_COUNT=decode_keep_count, PROMPT_LEN=prompt_len, NUM_KV_HEADS=num_kv_heads, SWA_WINDOW=swa_window, UNION=union, PER_LAYER=per_layer, HAS_SWA=swa_total > 0, - BLOCK=block, + BLOCK=_PACK_BLOCK_TOKENS, ) - return _PreparedTritonKernelLaunch( + return PreparedTritonKernelLaunch( _pack_compaction_sources_kernel, bound_tensors, constexpr_values, grid=grid, - num_warps=4, + num_warps=_PACK_NUM_WARPS, ) @@ -369,6 +342,22 @@ class BatchedKVCacheCompaction: tail. A co-compressed draft cache reuses the target's kept token ordinals (broadcast over the draft's own KV-head count) plus the draft's own protected tail, landing at the same destination base. + + Key constructor inputs: + `kept_token_ordinals`: increasing kept ordinals per request; shape + `[requests, prompt+keep]` for `union`, with a selection-row + dimension in between for the per-head modes. + `kv_block_offsets`: the staged V2 block-offset snapshot laid out as + `[slot, request, K/V, block]`, where a block offset encodes + page and K/V plane as `2*page + plane`. + `page_table_slots` / `layer_group_representative`: map each layer's + group representative to its snapshot slot; layers that share a + slot must share one block-offset table. + `protected_tail_lengths`: per-request KV positions past the valid + length reserved for a forward already in flight; they move with + the kept tokens. + `draft_*`: co-compressed draft-cache layout (union mode only); the + draft reuses the target keep set and pins the same prompt. """ def __init__( @@ -403,9 +392,19 @@ def __init__( raise ValueError("batched compaction requires requests and retained tokens") if not dense_layers: raise ValueError("batched compaction requires at least one dense layer") + if draft_layers and eviction_mode != "union": + raise ValueError("draft co-compaction supports only union eviction") + if draft_layer_pools is not None and not draft_layers: + raise ValueError("draft pools were given without any draft layers") + if not swa_layers and swa_window: + raise ValueError("swa_window was given without any SWA layers") self.eviction_mode = eviction_mode self.device = layer_pools[dense_layers[0]].device + # The move buffers are allocated on the pool device, so the selection + # tensors feeding the pack kernel must already live there. + if kept_token_ordinals.device != self.device: + raise ValueError("kept-token ordinals must live on the pool device") self.request_count = int(request_count) self.prompt_len = int(prompt_len) self.decode_keep_count = int(decode_keep_count) @@ -477,7 +476,7 @@ def __init__( dense_move_indices, eviction_mode=self.eviction_mode, prompt_len=self.prompt_len, - keep_count=self.decode_keep_count, + decode_keep_count=self.decode_keep_count, num_dense_layers=len(self.dense_layers), num_kv_heads=self.num_kv_heads, max_protected_tail=self.max_protected_tail, @@ -488,7 +487,7 @@ def __init__( self.target_dense_compaction = _SingleCacheCompaction( prepared_move_index_pack=dense_pack, cpp_launch_groups=_compact_groups( - dense_entries, "dense", self.layer_pool_keys, self.device, dense_slots + dense_entries, self.layer_pool_keys, self.device, dense_slots ), move_source_indices=dense_move_indices, move_source_offsets=dense_move_offsets, @@ -499,9 +498,7 @@ def __init__( if self.swa_layers: self.target_swa_compaction = _SingleCacheCompaction( prepared_move_index_pack=None, - cpp_launch_groups=_compact_groups( - swa_entries, "swa", self.layer_pool_keys, self.device - ), + cpp_launch_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), move_source_indices=swa_move_indices, move_source_offsets=swa_move_offsets, destination_base=self.keep_count - self.swa_window, @@ -547,10 +544,9 @@ def _build_draft_compaction( """Build the co-compressed draft cache's own pack and launch groups. The draft forms its own launch groups so it may use a different - KV-head count than the target. + KV-head count than the target. Union-only eviction is enforced by + the constructor before dense groups are built. """ - if self.eviction_mode != "union": - raise ValueError("draft co-compaction supports only union eviction") if ( draft_layer_pools is None or draft_layer_group_representative is None @@ -602,7 +598,7 @@ def _build_draft_compaction( draft_move_indices, eviction_mode="union", prompt_len=self.prompt_len, - keep_count=self.decode_keep_count, + decode_keep_count=self.decode_keep_count, num_dense_layers=1, num_kv_heads=draft_num_kv_heads, max_protected_tail=max(draft_tail_lengths, default=0), @@ -613,7 +609,7 @@ def _build_draft_compaction( return _SingleCacheCompaction( prepared_move_index_pack=draft_pack, cpp_launch_groups=_compact_groups( - draft_entries, "draft", tuple(draft_layer_pool_keys), self.device + draft_entries, tuple(draft_layer_pool_keys), self.device ), move_source_indices=draft_move_indices, move_source_offsets=draft_move_offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py new file mode 100644 index 000000000000..3af30eb75151 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py @@ -0,0 +1,82 @@ +# 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. + +"""Replayable Triton kernel launches frozen at build time.""" + +from typing import Dict, Tuple + +import torch + + +class PreparedTritonKernelLaunch: + """Replay one Triton kernel launch frozen at build time. + + ``warmup`` JIT-compiles the kernel once for a fixed grid, bound tensor + set, and constexpr set; ``__call__`` re-launches the compiled binary + directly, skipping Triton's per-call dispatch. Constexpr values are + passed positionally on replay, so their order is validated against the + kernel's constexpr parameter declaration order at build time. + """ + + def __init__( + self, + triton_kernel, + bound_tensors: Tuple[torch.Tensor, ...], + constexpr_values: Dict[str, object], + *, + grid: Tuple[int, ...], + num_warps: int, + ) -> None: + params = getattr(triton_kernel, "params", None) + if params is not None: + declared = [param.name for param in params if param.is_constexpr] + if list(constexpr_values.keys()) != declared: + raise ValueError( + f"constexpr order {list(constexpr_values.keys())} must match the " + f"kernel's declaration order {declared}: replay passes them " + "positionally" + ) + self.device = bound_tensors[0].device + self.bound_tensors = tuple(bound_tensors) + self.constexpr_values = dict(constexpr_values) + with torch.cuda.device(self.device): + self.build_stream = torch.cuda.current_stream(self.device) + # warmup() then indexing the compiled cache by grid is the + # documented-by-use Triton pattern for dispatch-free replay; if a + # Triton upgrade changes it, this raises here at build time rather + # than corrupting a replay. + compiled = triton_kernel.warmup( + *self.bound_tensors, + **self.constexpr_values, + num_warps=num_warps, + grid=grid, + ) + self.compiled_kernel_runner = compiled[grid] + + def __call__(self, *replay_tensors: torch.Tensor) -> None: + """Replay the launch; ``replay_tensors``, if given, substitute the bound tensors.""" + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.build_stream.device, + self.build_stream.cuda_stream, + ): + raise RuntimeError( + "a prepared Triton kernel launch must run on the stream it was built on" + ) + self.compiled_kernel_runner( + *(replay_tensors if replay_tensors else self.bound_tensors), + *self.constexpr_values.values(), + stream=self.build_stream.cuda_stream, + ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d7c2b7650c83..944b543d21c0 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -59,8 +59,8 @@ import torch -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - _PreparedTritonKernelLaunch, +from tensorrt_llm._torch.kv_cache_compression.triattention.prepared_launch import ( + PreparedTritonKernelLaunch, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, get_draft_token_length @@ -221,7 +221,7 @@ def __init__( from .triattention_kernels import _finalize_topk_indices_kernel - self._prepared_launch = _PreparedTritonKernelLaunch( + self._prepared_launch = PreparedTritonKernelLaunch( _finalize_topk_indices_kernel, (scores, seq_lens, provisional_indices, output_indices), dict( @@ -294,7 +294,7 @@ def __init__( self._prepared_stats_launch = None if normalize_scores: stats_grid = (request_count * rows, 1, 1) - self._prepared_stats_launch = _PreparedTritonKernelLaunch( + self._prepared_stats_launch = PreparedTritonKernelLaunch( _score_row_stats_kernel, (scores, valid_widths, row_mean, row_inv_std), dict(ROWS=rows, WIDTH=width, BLOCK=256), @@ -302,7 +302,7 @@ def __init__( num_warps=4, ) union_grid = (request_count, (width + 31) // 32, 1) - self._prepared_union_launch = _PreparedTritonKernelLaunch( + self._prepared_union_launch = PreparedTritonKernelLaunch( _score_union_kernel, (scores, valid_widths, row_mean, row_inv_std, combined), dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py new file mode 100644 index 000000000000..e2eefb7f034d --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -0,0 +1,37 @@ +# 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. + +import torch + + +def encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: + """Build the native V2 [pool, request, K/V, block] layout. + + Accepts ``[request, block]`` page ids (one pool) or ``[pool, request, + block]``; K offsets encode as ``2*page`` and V as ``2*page + 1``. + """ + if page_ids.ndim == 2: + page_ids = page_ids.unsqueeze(0) + encoded = torch.empty( + page_ids.shape[0], + page_ids.shape[1], + 2, + page_ids.shape[2], + dtype=torch.int32, + device=page_ids.device, + ) + encoded[:, :, 0] = page_ids.to(torch.int32) * 2 + encoded[:, :, 1] = encoded[:, :, 0] + 1 + return encoded 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 index a946ffb5fae5..d0225105e84e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -19,6 +19,7 @@ import pytest import torch +from conftest import encode_block_offsets as _encode_block_offsets from tensorrt_llm._torch.kv_cache_compression.triattention import TriAttention from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( @@ -32,23 +33,6 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -def _encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: - """Build the native V2 [pool, request, K/V, block] layout.""" - if page_ids.ndim == 2: - page_ids = page_ids.unsqueeze(0) - encoded = torch.empty( - page_ids.shape[0], - page_ids.shape[1], - 2, - page_ids.shape[2], - dtype=torch.int32, - device=page_ids.device, - ) - encoded[:, :, 0] = page_ids.to(torch.int32) * 2 - encoded[:, :, 1] = encoded[:, :, 0] + 1 - return encoded - - def _make_fake_v2(*, is_draft=False): """Build an unallocated V2 double with TriAttention's production contract.""" from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 @@ -332,7 +316,8 @@ def test_draft_admission_gates_raise(gate, match): with pytest.raises(ValueError, match=match): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8), + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + ), DFlashDecodingConfig(max_draft_len=3), draft_manager, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index 7a9508c982d2..1e20b5ab2f53 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -5,6 +5,7 @@ import pytest import torch +from conftest import encode_block_offsets as _encode_block_offsets from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( BatchedKVCacheCompaction, @@ -22,21 +23,6 @@ def _require_cute_topk_op() -> None: ) -def _encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: - """Build the native V2 [pool, request, K/V, block] layout.""" - encoded = torch.empty( - page_ids.shape[0], - page_ids.shape[1], - 2, - page_ids.shape[2], - dtype=torch.int32, - device=page_ids.device, - ) - encoded[:, :, 0] = page_ids.to(torch.int32) * 2 - encoded[:, :, 1] = encoded[:, :, 0] + 1 - return encoded - - 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)) diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index cf46fe186c15..c38563f175d3 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -101,6 +101,9 @@ def _reference_compact( for group_layer, (source_pool, destination_pool, page_table) in enumerate( zip(original, expected, page_tables) ): + # The kernel decodes K offsets as offset // 2 (the V2 2*page+plane + # encoding) regardless of the scale the table was built with; the + # reference mirrors the kernel, not the encoder. raw_page_table = page_table // _PAGE_INDEX_DIVISOR if source_indices.ndim == 2: layer_sources = source_indices @@ -224,6 +227,20 @@ def test_sparse_kv_cache_compact_layers_per_layer_source(): assert torch.equal(actual.cpu(), reference) +def test_sparse_kv_cache_compact_layers_rejects_flat_source_with_layer_indices(): + # A flat [kv_heads, total] source with per-layer indices would silently + # read layer 0 for every launch; the op rejects the combination instead. + _, pools, page_tables = _make_pools(2, torch.bfloat16, 64) + source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) + source_row = torch.tensor([2, 5, 8, 3, 7, 10], dtype=torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + source_layer_indices = torch.tensor([0, 0], dtype=torch.int32) + arguments = _device_arguments(pools, source_indices, source_offsets, source_layer_indices) + + with pytest.raises(RuntimeError, match="require 3-D per-layer source_indices"): + _compact(pools, page_tables, arguments, 0) + + def test_sparse_kv_cache_compact_layers_multiple_tiles(): num_layers = 2 max_pages_per_sequence = 24 From 16cac9904141c7e338c2f8b4a99b670a15482d9c Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 09:45:26 -0700 Subject: [PATCH 019/178] [None][chore] Regenerate the LLM args telemetry manifest after the main merge The merge resolved the manifest to the main side; regenerating restores the KV-cache compression entries alongside the fields main added. Signed-off-by: tianruih --- .../usage/llm_args_golden_manifest.json | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index ca6e762490cf..d82e9627cf68 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -544,6 +544,61 @@ "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.count_prompt_tokens" + }, + { + "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": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.pin_prefill" + }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.top_B" + }, { "allowed_values": [], "annotation": "", From fd0d761c78fc981f8192b09f69403c0242b57ae3 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 10:01:04 -0700 Subject: [PATCH 020/178] [None][fix] Track the reworked speculative acceptance fields in the spec gate Upstream replaced acceptance_window/acceptance_length_threshold with use_rejection_sampling (base config) and use_relaxed_acceptance_for_thinking (MTP only); the compression spec gate now rejects those, with the same fail-fast intent: eviction is only validated with greedy acceptance. Signed-off-by: tianruih --- tensorrt_llm/_torch/pyexecutor/_util.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 60f152359fde..ce185c847bc2 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2097,10 +2097,13 @@ def validate_kv_cache_compression_with_spec( if spec_config.draft_len_schedule is not None: raise ValueError("TriAttention does not yet support dynamic " "speculative draft lengths") - if (spec_config.acceptance_window is not None - or spec_config.acceptance_length_threshold is not None): - raise ValueError("TriAttention does not support runtime " - "speculative acceptance gating") + # Rejection sampling (and MTP's relaxed thinking acceptance, an + # MTP-only field) makes the per-step accepted-token count stochastic; + # compression eviction is only validated with greedy acceptance. + if spec_config.use_rejection_sampling or getattr( + spec_config, "use_relaxed_acceptance_for_thinking", False): + raise ValueError("TriAttention does not support speculative " + "rejection sampling or relaxed acceptance") if draft_kv_cache_manager is None: raise ValueError( "TriAttention speculative compatibility requires a separate " From 213532de7d734f4378010f69b27b6a70c3868bf2 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 21:25:39 -0700 Subject: [PATCH 021/178] [None][perf] Batch one eviction round across mixed prompt lengths Eviction cohorts were grouped by (prompt_len, keep_count) because the batched pipeline treated the pinned prompt length as scalar geometry: the selection rectangle carried a materialized prompt prefix, the score kernel baked TOKEN_START as a compile-time constant, and the C++ compact launch took one destination base. Real batches mix prompt lengths, so a 32-request round split into up to 15 sub-rounds; a profiled step with 13 sub-rounds spent 39.8 ms in update_resources against 2.3 ms for a single round. Prompt lengths are now per-request runtime metadata end to end: - the score kernel reads each request's decode-window origin from a staged token-starts row instead of a constexpr; - selection rectangles are decode-only ([request(, rows), top_B]) with the per-row prompt offset rebasing emitted ordinals to absolute positions; - the pack kernel drops its prompt constants (moves never touch the prompt); - the compact op takes per-request destination bases (contiguous CUDA int32 [batch]), and SWA landing positions rebase per request each round; - the cohort key is gone: one due round launches once, chunked only by the staging memory bound. A cohort mixing prompt lengths is asserted byte-identical to per-request compactions, and the op gains a mixed-bases unit test. Signed-off-by: tianruih --- .../kernels/unfusedAttentionKernels.h | 5 +- .../unfusedAttentionKernels_2_template.h | 13 +- .../thop/sparseKvCacheCompactOp.cpp | 31 ++- .../triattention/compaction.py | 70 +++-- .../triattention/triattention.py | 246 ++++++++++++------ .../triattention/triattention_kernels.py | 69 ++--- .../test_triattention_draft_cocompaction.py | 21 +- .../test_triattention_eager.py | 192 +++++++++++--- .../test_triattention_pipeline.py | 48 +++- .../serial/test_sparse_kv_cache_compact.py | 46 +++- 10 files changed, 528 insertions(+), 213 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 708f90c970ef..b51d8e8dc036 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -416,11 +416,12 @@ void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams //! sparse-KV updater. Device pointer arrays allow one layered launch while the //! updater remains the only implementation of the in-place copy loop. Within //! each request and head, source ordinals must increase strictly and satisfy -//! destinationBase + move <= source[move]. +//! destinationBases[request] + move <= source[move]; the per-request bases +//! let one launch cover a cohort with mixed pinned-prompt lengths. template void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, - int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t destinationBase, int32_t batchSize, + int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream); // Debug function to test basic parameter access diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index df59045fee31..4eceb5467338 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1765,7 +1765,7 @@ struct KvCacheV2LayersBuffer int64_t sourceLayerStride; int64_t pageTableRequestStride; int32_t tokensPerBlock; - int32_t destinationBase; + int32_t const* destinationBases; size_t bytesPerPage; size_t bytesPerKvHalf; @@ -1807,9 +1807,10 @@ struct KvCacheV2LayersBuffer return sourceIndices[offset]; } - __device__ __forceinline__ int32_t getSparseKvDestinationToken(int32_t requestMove) const + __device__ __forceinline__ int32_t getSparseKvDestinationToken(int32_t batchIdx, int32_t requestMove) const { - return destinationBase + requestMove; + // Per-request landing positions: cohorts may mix prompt lengths. + return destinationBases[batchIdx] + requestMove; } }; @@ -1900,7 +1901,7 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( { src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( params.sparse_kv_indices, kv_head_idx, total_num_sparse_kv_tokens, global_sparse_idx); - dst_token_idx = params.kv_cache_buffer.getSparseKvDestinationToken(sparse_token_offset); + dst_token_idx = params.kv_cache_buffer.getSparseKvDestinationToken(batch_idx, sparse_token_offset); } else { @@ -2009,7 +2010,7 @@ void launchSparseKvCacheCompactV2Layers( template void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, - int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t destinationBase, int32_t batchSize, + int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream) { KvCacheV2LayersBuffer buffer{}; @@ -2019,7 +2020,7 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con buffer.sourceLayerStride = sourceLayerStride; buffer.pageTableRequestStride = pageTableRequestStride; buffer.tokensPerBlock = tokensPerBlock; - buffer.destinationBase = destinationBase; + buffer.destinationBases = destinationBases; buffer.bytesPerKvHalf = static_cast(numKvHeads) * tokensPerBlock * headDim * sizeof(T); buffer.bytesPerPage = 2 * buffer.bytesPerKvHalf; diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index 733ea45516ee..126cf32b4b85 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -34,14 +34,15 @@ namespace torch_ext //! Adapt one uniform group of KVCacheManagerV2 HND layer pools to the //! existing sparse-KV post-FMHA updater. Layers share one V2 block-offset -//! table; destinationBase replaces the former arbitrary destination tensor -//! because every compaction move targets one contiguous interval. Within -//! each request and KV head, the caller must supply increasing source -//! ordinals with destinationBase + move <= source[move], which makes the -//! updater's forward tiled in-place copy safe. +//! table; per-request destinationBases replace the former arbitrary +//! destination tensor because every compaction move targets one contiguous +//! interval per request. Within each request and KV head, the caller must +//! supply increasing source ordinals with destinationBases[request] + move +//! <= source[move], which makes the updater's forward tiled in-place copy +//! safe. void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor const& poolPointers, th::Tensor const& pageTable, th::Tensor const& sourceIndices, th::Tensor const& sourceOffsets, - std::optional const& sourceLayerIndices, int64_t destinationBase) + std::optional const& sourceLayerIndices, th::Tensor const& destinationBases) { TORCH_CHECK(!pools.empty(), "sparse_kv_cache_compact_layers: pools must be non-empty"); @@ -119,29 +120,33 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor && sourceOffsets.scalar_type() == th::kInt32 && sourceOffsets.is_contiguous() && sourceOffsets.dim() == 1 && sourceOffsets.size(0) == batchSize + 1, "sparse_kv_cache_compact_layers: source_offsets must be contiguous CUDA int32 [batch + 1]"); - TORCH_CHECK(destinationBase >= 0 && destinationBase <= std::numeric_limits::max(), - "sparse_kv_cache_compact_layers: destination_base must fit a non-negative int32"); + // Per-request landing positions: one launch covers a cohort with mixed + // pinned-prompt lengths. Values live on device; the kernel trusts them. + TORCH_CHECK(destinationBases.is_cuda() && destinationBases.get_device() == device + && destinationBases.scalar_type() == th::kInt32 && destinationBases.dim() == 1 + && destinationBases.size(0) == batchSize && destinationBases.is_contiguous(), + "sparse_kv_cache_compact_layers: destination_bases must be contiguous CUDA int32 [batch]"); auto const stream = at::cuda::getCurrentCUDAStream(device); - auto const base = static_cast(destinationBase); + auto const* bases = destinationBases.data_ptr(); if (dtype == th::kBFloat16) { tk::invokeSparseKvCacheCompactV2Layers<__nv_bfloat16>(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), - sourceLayerPtr, sourceLayerStride, sourceOffsets.data_ptr(), base, batchSize, numKvHeads, + sourceLayerPtr, sourceLayerStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); } else if (dtype == th::kHalf) { tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, - sourceOffsets.data_ptr(), base, batchSize, numKvHeads, tokensPerBlock, headDim, stream); + sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); } else if (dtype == th::kFloat) { tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, - sourceOffsets.data_ptr(), base, batchSize, numKvHeads, tokensPerBlock, headDim, stream); + sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); } else { @@ -158,7 +163,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "sparse_kv_cache_compact_layers(Tensor(a!)[] pools, Tensor pool_pointers, Tensor page_table, Tensor " "source_indices, Tensor source_offsets, Tensor? source_layer_indices=None, " - "int destination_base=0) -> ()"); + "Tensor destination_bases) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index fa278b661785..3cb472678b3b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -45,7 +45,9 @@ class _CppCompactGroup(NamedTuple): pool_pointers: torch.Tensor source_layer_indices: Optional[torch.Tensor] - def launch(self, source: torch.Tensor, offsets: torch.Tensor, destination_base: int) -> None: + def launch( + self, source: torch.Tensor, offsets: torch.Tensor, destination_bases: torch.Tensor + ) -> None: torch.ops.trtllm.sparse_kv_cache_compact_layers( list(self.pools), self.pool_pointers, @@ -53,7 +55,7 @@ def launch(self, source: torch.Tensor, offsets: torch.Tensor, destination_base: source, offsets, self.source_layer_indices, - destination_base, + destination_bases, ) @@ -70,13 +72,17 @@ class _SingleCacheCompaction(NamedTuple): cpp_launch_groups: Tuple[_CppCompactGroup, ...] move_source_indices: torch.Tensor move_source_offsets: torch.Tensor - destination_base: int + # Per-request landing positions; may alias staged prompt lengths so the + # values track the current round without a refresh. + destination_bases: torch.Tensor def launch(self) -> None: if self.prepared_move_index_pack is not None: self.prepared_move_index_pack() for group in self.cpp_launch_groups: - group.launch(self.move_source_indices, self.move_source_offsets, self.destination_base) + group.launch( + self.move_source_indices, self.move_source_offsets, self.destination_bases + ) def _cuda_int32_contiguous(tensors: Tuple[torch.Tensor, ...], device: torch.device) -> bool: @@ -236,7 +242,6 @@ def _prepared_move_index_pack_launch( move_source_indices: torch.Tensor, *, eviction_mode: str, - prompt_len: int, decode_keep_count: int, num_dense_layers: int, num_kv_heads: int, @@ -262,7 +267,9 @@ def _prepared_move_index_pack_launch( else: selection_rows = num_kv_heads selection_prefix = (request_count,) if union else (request_count, selection_rows) - expected_selection = (*selection_prefix, prompt_len + decode_keep_count) + # Selection rows carry decode-only kept ordinals (already absolute), so + # the rectangle is prompt-length independent. + expected_selection = (*selection_prefix, decode_keep_count) if ( request_count <= 0 or tuple(kept_token_ordinals.shape) != expected_selection @@ -314,9 +321,8 @@ def _prepared_move_index_pack_launch( DENSE_TOTAL=int(move_source_indices.shape[-1]), SWA_TOTAL=swa_total, SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=prompt_len + decode_keep_count, + SELECTION_STRIDE=decode_keep_count, KEEP_COUNT=decode_keep_count, - PROMPT_LEN=prompt_len, NUM_KV_HEADS=num_kv_heads, SWA_WINDOW=swa_window, UNION=union, @@ -344,9 +350,12 @@ class BatchedKVCacheCompaction: protected tail, landing at the same destination base. Key constructor inputs: - `kept_token_ordinals`: increasing kept ordinals per request; shape - `[requests, prompt+keep]` for `union`, with a selection-row - dimension in between for the per-head modes. + `kept_token_ordinals`: increasing kept decode ordinals (absolute + positions) per request; shape `[requests, keep]` for `union`, + with a selection-row dimension in between for the per-head + modes. Prompt tokens never move, so the rectangle is + prompt-length independent and one cohort may mix prompt sizes; + `prompt_offsets` carries each request's pinned prompt length. `kv_block_offsets`: the staged V2 block-offset snapshot laid out as `[slot, request, K/V, block]`, where a block offset encodes page and K/V plane as `2*page + plane`. @@ -373,7 +382,7 @@ def __init__( kv_block_offsets: torch.Tensor, page_table_slots: Dict[int, int], request_count: int, - prompt_len: int, + prompt_offsets: torch.Tensor, decode_keep_count: int, swa_window: Optional[int], protected_tail_lengths: Optional[List[int]] = None, @@ -388,7 +397,7 @@ def __init__( ) -> None: if eviction_mode not in ("union", "per_head", "per_layer_perhead"): raise ValueError(f"unsupported compaction mode: {eviction_mode}") - if request_count <= 0 or decode_keep_count <= 0 or prompt_len < 0: + if request_count <= 0 or decode_keep_count <= 0: raise ValueError("batched compaction requires requests and retained tokens") if not dense_layers: raise ValueError("batched compaction requires at least one dense layer") @@ -406,9 +415,18 @@ def __init__( if kept_token_ordinals.device != self.device: raise ValueError("kept-token ordinals must live on the pool device") self.request_count = int(request_count) - self.prompt_len = int(prompt_len) + # Per-request pinned prompt lengths; this usually aliases the staged + # prompt buffer, so the values track the current round. Only the + # geometry is validated here. + if ( + prompt_offsets.shape != (self.request_count,) + or prompt_offsets.dtype != torch.int32 + or prompt_offsets.device != self.device + or not prompt_offsets.is_contiguous() + ): + raise ValueError("per-request prompt offsets do not match the cohort") + self.prompt_offsets = prompt_offsets self.decode_keep_count = int(decode_keep_count) - self.keep_count = self.prompt_len + self.decode_keep_count self.protected_tail_lengths = _validated_tail_lengths( protected_tail_lengths, self.request_count, "protected-tail" ) @@ -449,13 +467,17 @@ def __init__( ] self.swa_window = 0 + self.swa_destination_bases = None swa_move_indices = None swa_move_offsets = None swa_entries = [] if self.swa_layers: - if swa_window is None or swa_window <= 0 or self.keep_count < swa_window: + if swa_window is None or swa_window <= 0: raise ValueError("SWA compaction requires a valid retained window") + # Per-request window validity (prompt + decode keep >= window) is + # prompt-dependent and checked by the caller each round. self.swa_window = int(swa_window) + self.swa_destination_bases = torch.empty_like(self.prompt_offsets) swa_move_indices, swa_move_offsets = _make_move_buffers( (self.num_kv_heads,), [self.swa_window + length for length in self.protected_tail_lengths], @@ -475,7 +497,6 @@ def __init__( dense_move_offsets, dense_move_indices, eviction_mode=self.eviction_mode, - prompt_len=self.prompt_len, decode_keep_count=self.decode_keep_count, num_dense_layers=len(self.dense_layers), num_kv_heads=self.num_kv_heads, @@ -491,7 +512,7 @@ def __init__( ), move_source_indices=dense_move_indices, move_source_offsets=dense_move_offsets, - destination_base=self.prompt_len, + destination_bases=self.prompt_offsets, ) # The dense pack launch fills the SWA move buffers in the same call. self.target_swa_compaction = None @@ -501,7 +522,7 @@ def __init__( cpp_launch_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), move_source_indices=swa_move_indices, move_source_offsets=swa_move_offsets, - destination_base=self.keep_count - self.swa_window, + destination_bases=self.swa_destination_bases, ) self.draft_compaction = None @@ -597,7 +618,6 @@ def _build_draft_compaction( draft_move_offsets, draft_move_indices, eviction_mode="union", - prompt_len=self.prompt_len, decode_keep_count=self.decode_keep_count, num_dense_layers=1, num_kv_heads=draft_num_kv_heads, @@ -613,10 +633,18 @@ def _build_draft_compaction( ), move_source_indices=draft_move_indices, move_source_offsets=draft_move_offsets, - destination_base=self.prompt_len, + destination_bases=self.prompt_offsets, ) def launch(self) -> None: """Pack the move indices, then run every cache family's C++ compacts.""" + if self.swa_destination_bases is not None: + # The prompt offsets may have been re-staged since construction; + # rebase the SWA landing positions for this round. + torch.add( + self.prompt_offsets, + self.decode_keep_count - self.swa_window, + out=self.swa_destination_bases, + ) for compaction in self.cache_compactions: compaction.launch() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 944b543d21c0..d5caf2958dee 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -193,13 +193,12 @@ def __init__( self, scores: torch.Tensor, seq_lens: torch.Tensor, + prompt_offsets: torch.Tensor, provisional_indices: torch.Tensor, output_indices: torch.Tensor, keep_count: int, - prompt_len: int, ) -> None: rows, width = scores.shape - output_width = prompt_len + keep_count if ( not scores.is_cuda or scores.dtype != torch.float32 @@ -207,13 +206,17 @@ def __init__( or seq_lens.shape != (rows,) or seq_lens.dtype != torch.int32 or seq_lens.device != scores.device + or prompt_offsets.shape != (rows,) + or prompt_offsets.dtype != torch.int32 + or prompt_offsets.device != scores.device or provisional_indices.shape != (rows, keep_count) or provisional_indices.dtype != torch.int32 or provisional_indices.device != scores.device - or output_indices.shape != (rows, output_width) + or output_indices.shape != (rows, keep_count) or output_indices.dtype != torch.int32 or output_indices.device != scores.device or not seq_lens.is_contiguous() + or not prompt_offsets.is_contiguous() or not provisional_indices.is_contiguous() or not output_indices.is_contiguous() ): @@ -223,12 +226,11 @@ def __init__( self._prepared_launch = PreparedTritonKernelLaunch( _finalize_topk_indices_kernel, - (scores, seq_lens, provisional_indices, output_indices), + (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), dict( WIDTH=width, KEEP_COUNT=keep_count, - OUTPUT_WIDTH=output_width, - PROMPT_LEN=prompt_len, + OUTPUT_WIDTH=keep_count, BLOCK=256, ), grid=(rows, 1, 1), @@ -239,10 +241,13 @@ def __call__( self, scores: torch.Tensor, seq_lens: torch.Tensor, + prompt_offsets: torch.Tensor, provisional_indices: torch.Tensor, output_indices: torch.Tensor, ) -> None: - self._prepared_launch(scores, seq_lens, provisional_indices, output_indices) + self._prepared_launch( + scores, seq_lens, prompt_offsets, provisional_indices, output_indices + ) class _PreparedUnionScores: @@ -319,14 +324,18 @@ def __call__(self) -> None: def _deterministic_topk_indices_into( scores: torch.Tensor, seq_lens: torch.Tensor, + prompt_offsets: torch.Tensor, provisional_indices_i32: torch.Tensor, output_indices_i32: torch.Tensor, keep_count: int, - prompt_len: int, prepared_topk: Optional[_PreparedCuteTopK] = None, prepared_finalizer: Optional[_PreparedTopKFinalizer] = None, ) -> None: - """Write stable, increasing physical indices around one CuTE TopK call.""" + """Write stable, increasing physical indices around one CuTE TopK call. + + Scores are decode-relative; ``prompt_offsets`` rebases each row's emitted + ordinals to absolute positions, so one call may mix prompt lengths. + """ if scores.is_cuda: if prepared_topk is None or prepared_finalizer is None: raise ValueError("CUDA selection requires prepared TopK launchers") @@ -334,6 +343,7 @@ def _deterministic_topk_indices_into( prepared_finalizer( scores, seq_lens, + prompt_offsets, provisional_indices_i32, output_indices_i32, ) @@ -343,6 +353,7 @@ def _deterministic_topk_indices_into( _topk_indices_into(scores, seq_lens, provisional_indices_i32, keep_count) for row_index, row_scores in enumerate(scores): valid_width = int(seq_lens[row_index]) + prompt_len = int(prompt_offsets[row_index]) selected = provisional_indices_i32[row_index].to(torch.long) threshold = torch.amin(row_scores[selected]) valid_scores = row_scores[:valid_width] @@ -350,7 +361,7 @@ def _deterministic_topk_indices_into( tied = torch.nonzero(valid_scores == threshold, as_tuple=False).flatten() tie_count = keep_count - int(higher.numel()) ordered = torch.sort(torch.cat((higher, tied[:tie_count]))).values - output_indices_i32[row_index, prompt_len : prompt_len + keep_count].copy_( + output_indices_i32[row_index, :keep_count].copy_( ordered.to(torch.int32).add(prompt_len) ) @@ -365,7 +376,6 @@ class _CrossRequestSelectionPlan(NamedTuple): rows: int width: int keep_count: int - prompt_len: int dtype: torch.dtype device: torch.device max_requests: int @@ -401,7 +411,8 @@ def __init__( num_kv_heads: int, width: int, keep_count: int, - prompt_len: int, + selection_rows_per_request: int = 1, + prompt_offsets_buffer: Optional[torch.Tensor] = None, dtype: torch.dtype, device: torch.device, max_requests: int, @@ -416,14 +427,61 @@ def __init__( self.num_kv_heads = int(num_kv_heads) self.width = int(width) self.keep_count = int(keep_count) - self.prompt_len = int(prompt_len) - self.total_keep = self.prompt_len + self.keep_count self.dtype = dtype self.device = _canonical_device(device) self.max_requests = int(max_requests) self.valid_widths = torch.full( (self.max_requests,), self.width, dtype=torch.int32, device=self.device ) + # Per-request pinned prompt lengths, refreshed each round: scores are + # decode-relative and these offsets rebase emitted ordinals, so one + # cohort may mix prompt lengths. ``row_prompt_offsets`` is the + # row-major expansion consumed by the finalizer. + self.selection_rows_per_request = int(selection_rows_per_request) + if prompt_offsets_buffer is not None: + # Share the staging buffers' per-request prompt lengths so the + # values are written once per round. + if ( + prompt_offsets_buffer.shape != (self.max_requests,) + or prompt_offsets_buffer.dtype != torch.int32 + or prompt_offsets_buffer.device != self.device + ): + raise ValueError("prompt offsets buffer does not match the selector geometry") + self.prompt_offsets = prompt_offsets_buffer + else: + self.prompt_offsets = torch.zeros( + (self.max_requests,), dtype=torch.int32, device=self.device + ) + if self.selection_rows_per_request == 1: + self.row_prompt_offsets = self.prompt_offsets + else: + self.row_prompt_offsets = torch.zeros( + (self.max_requests * self.selection_rows_per_request,), + dtype=torch.int32, + device=self.device, + ) + + def set_prompt_offsets(self, prompt_lens: torch.Tensor) -> None: + """Refresh the per-request prompt offsets for the coming round.""" + count = int(prompt_lens.numel()) + if count > self.max_requests: + raise ValueError("prompt offsets exceed the selector's request capacity") + self.prompt_offsets[:count].copy_(prompt_lens, non_blocking=True) + self.refresh_row_prompt_offsets() + + def refresh_row_prompt_offsets(self) -> None: + """Re-expand the per-request prompt offsets into their row-major view. + + Called after the shared per-request buffer was staged externally. + """ + if self.row_prompt_offsets is not self.prompt_offsets: + self.row_prompt_offsets.view( + self.max_requests, self.selection_rows_per_request + ).copy_( + self.prompt_offsets.unsqueeze(1).expand( + -1, self.selection_rows_per_request + ) + ) def _allocate_cpu_reference_buffers( self, @@ -459,17 +517,12 @@ def _build_prepared_selection_launchers( self.prepared_finalizer = _PreparedTopKFinalizer( scores_rows, row_lengths, + self.row_prompt_offsets, provisional_indices, keep_rows, self.keep_count, - self.prompt_len, ) - def _prefill_prompt_ordinals(self, keep: torch.Tensor) -> None: - if self.prompt_len: - prompt = torch.arange(self.prompt_len, dtype=torch.int32, device=self.device) - keep[..., : self.prompt_len].copy_(prompt.expand(*keep.shape[:-1], self.prompt_len)) - def _validated_request_count(self, scores: torch.Tensor, what: str) -> int: request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 if request_count <= 0 or request_count > self.max_requests: @@ -487,7 +540,6 @@ def __init__( rows: int, width: int, keep_count: int, - prompt_len: int, *, dtype: torch.dtype, device: torch.device, @@ -497,6 +549,7 @@ def __init__( num_kv_heads: int = 0, input_scores: Optional[torch.Tensor] = None, normalize_scores: bool = True, + prompt_offsets_buffer: Optional[torch.Tensor] = None, ) -> None: if rows <= 0: raise ValueError("cross-request selection requires rows > 0") @@ -507,7 +560,7 @@ def __init__( num_kv_heads=num_kv_heads, width=width, keep_count=keep_count, - prompt_len=prompt_len, + prompt_offsets_buffer=prompt_offsets_buffer, dtype=dtype, device=device, max_requests=max_requests, @@ -522,8 +575,10 @@ def __init__( self.final_indices = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=self.device ) + # Kept decode ordinals only: rows are prompt-length independent, so + # one selector serves cohorts with mixed prompt lengths. self.keep = torch.empty( - (max_requests, self.total_keep), dtype=torch.int32, device=self.device + (max_requests, self.keep_count), dtype=torch.int32, device=self.device ) self._allocate_cpu_reference_buffers( (max_requests, 1, 1), (max_requests, 1, width), torch.int32 @@ -543,7 +598,6 @@ def __init__( if self.device.type == "cuda" and input_scores is not None else None ) - self._prefill_prompt_ordinals(self.keep) def _select_input_scores( self, @@ -567,10 +621,10 @@ def _select_input_scores( _deterministic_topk_indices_into( combined, valid_widths, + self.prompt_offsets[:request_count], final_indices, self.keep[:request_count], self.keep_count, - self.prompt_len, self.prepared_topk, self.prepared_finalizer, ) @@ -584,10 +638,10 @@ def select_prepared_requests(self) -> None: _deterministic_topk_indices_into( self.combined, self.valid_widths, + self.prompt_offsets, final_indices, self.keep, self.keep_count, - self.prompt_len, self.prepared_topk, self.prepared_finalizer, ) @@ -677,10 +731,10 @@ def __init__( num_kv_heads: int, width: int, keep_count: int, - prompt_len: int, dtype: torch.dtype, device: torch.device, max_requests: int, + prompt_offsets_buffer: Optional[torch.Tensor] = None, ) -> None: if eviction_mode not in ("per_head", "per_layer_perhead"): raise ValueError(f"unsupported per-head eviction mode: {eviction_mode}") @@ -688,6 +742,11 @@ def __init__( raise ValueError("per-head selection requires positive layer, head, and request counts") if num_query_heads % num_kv_heads: raise ValueError("query heads must be divisible by KV heads") + selection_rows = ( + num_kv_heads + if eviction_mode == "per_head" + else len(dense_layers) * num_kv_heads + ) super().__init__( eviction_mode=eviction_mode, dense_layers=dense_layers, @@ -695,7 +754,8 @@ def __init__( num_kv_heads=num_kv_heads, width=width, keep_count=keep_count, - prompt_len=prompt_len, + selection_rows_per_request=selection_rows, + prompt_offsets_buffer=prompt_offsets_buffer, dtype=dtype, device=device, max_requests=max_requests, @@ -703,11 +763,7 @@ def __init__( self.num_layers = len(self.dense_layers) self.query_group_size = self.num_query_heads // self.num_kv_heads self.rows = self.num_layers * self.num_query_heads - self.selection_rows = ( - self.num_kv_heads - if eviction_mode == "per_head" - else self.num_layers * self.num_kv_heads - ) + self.selection_rows = selection_rows score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) grouped_shape = (self.max_requests, self.num_layers, self.num_kv_heads, self.width) @@ -734,8 +790,10 @@ def __init__( ) selection_shape = (self.max_requests, self.selection_rows, self.keep_count) self.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=self.device) + # Kept decode ordinals only: rows are prompt-length independent, so + # one selector serves cohorts with mixed prompt lengths. self.keep = torch.empty( - (self.max_requests, self.selection_rows, self.total_keep), + (self.max_requests, self.selection_rows, self.keep_count), dtype=torch.int32, device=self.device, ) @@ -744,14 +802,13 @@ def __init__( ) self.row_seq_lens_flat = self.row_seq_lens.view(-1) self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) - self.keep_flat = self.keep.view(-1, self.total_keep) + self.keep_flat = self.keep.view(-1, self.keep_count) self._build_prepared_selection_launchers( self.selection_scores_flat, self.row_seq_lens_flat, self.top_indices_i32_flat, self.keep_flat, ) - self._prefill_prompt_ordinals(self.keep) def _select_input_scores( self, @@ -788,10 +845,10 @@ def _select_input_scores( _deterministic_topk_indices_into( self.selection_scores_flat, self.row_seq_lens_flat, + self.row_prompt_offsets, self.top_indices_i32_flat, self.keep_flat, self.keep_count, - self.prompt_len, self.prepared_topk, self.prepared_finalizer, ) @@ -928,7 +985,7 @@ def __init__( omega: torch.Tensor, page_table_keys: Optional[List[object]] = None, num_page_table_slots: Optional[int] = None, - prompt_len: int = 0, + min_prompt_len: int = 0, page_table_token_capacity: Optional[int] = None, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_page_representatives: Optional[List[int]] = None, @@ -958,9 +1015,11 @@ def __init__( if page_table_token_capacity < seq_len: raise ValueError("page-table capacity cannot be smaller than the score bucket") self.page_table_token_capacity = int(page_table_token_capacity) - if prompt_len < 0 or prompt_len > seq_len: + # The smallest cohort prompt only sizes the widest decode window; + # per-request prompt lengths are staged runtime metadata. + if min_prompt_len < 0 or min_prompt_len > seq_len: raise ValueError("fixed score metadata prompt length is outside its bucket") - self.prompt_len = prompt_len + self.min_prompt_len = min_prompt_len q_real = q_real.to(device=self.device, dtype=torch.float32).contiguous() q_imag = q_imag.to(device=self.device, dtype=torch.float32).contiguous() mlr_coef = mlr_coef.to(device=self.device, dtype=torch.float32).contiguous() @@ -997,7 +1056,7 @@ def __init__( self.copy_block_count, ) self.request_metadata_host = torch.empty( - (2, max_requests), + (3, max_requests), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned(), @@ -1079,10 +1138,13 @@ def __init__( device=self.device, ) self.request_metadata_device = torch.empty( - (2, max_requests), dtype=torch.int32, device=self.device + (3, max_requests), dtype=torch.int32, device=self.device ) self.round_starts_device = self.request_metadata_device[0] self.valid_seq_lens_device = self.request_metadata_device[1] + # Per-request pinned prompt lengths: the score kernel starts each + # request's decode window here, so one bucket may mix prompt lengths. + self.token_starts_device = self.request_metadata_device[2] self.mean_cos = torch.empty( (max_requests, num_freqs), dtype=torch.float32, device=self.device ) @@ -1109,7 +1171,7 @@ def __init__( freq_scale_sq, omega, offsets, - prompt_len=prompt_len, + min_prompt_len=min_prompt_len, ) self.copy_done = torch.cuda.Event() # First record publishes constructor allocations to the V2 copy stream; @@ -1160,6 +1222,7 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str self.valid_seq_lens_device, valid_widths, self.round_starts_device, + self.token_starts_device, *group.pointer_middle, self.mean_cos.view(-1), self.mean_sin.view(-1), @@ -1173,7 +1236,6 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str ) score_constants = ( score_aggregation == "max", - group.prompt_len, group.token_block, frequency_block, ) @@ -1206,9 +1268,8 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str *score_pointer_args, *score_geometry, USE_MAX=score_constants[0], - TOKEN_START=score_constants[1], - T_BLOCK=score_constants[2], - F_BLOCK=score_constants[3], + T_BLOCK=score_constants[1], + F_BLOCK=score_constants[2], grid=score_grid, ) self._score_runner = compiled[score_grid] @@ -1235,16 +1296,23 @@ def stage( manager: KVCacheManagerV2, request_ids: List[int], round_starts: List[int], + token_starts: List[int], seq_lens: Optional[List[int]] = None, page_table_seq_lens: Optional[List[int]] = None, draft_manager: Optional[KVCacheManagerV2] = None, ) -> bool: - """Copy one eager eviction cohort into reusable device buffers.""" + """Copy one eager eviction cohort into reusable device buffers. + + ``token_starts`` carries each request's pinned prompt length; the + score kernel starts that request's decode window there, so the cohort + may mix prompt lengths. + """ request_count = len(request_ids) if ( request_count == 0 or request_count > self.max_requests or len(round_starts) != request_count + or len(token_starts) != request_count or any( round_start != round_start or round_start < 0 @@ -1252,6 +1320,10 @@ def stage( or round_start != int(round_start) for round_start in round_starts ) + or any( + token_start < 0 or token_start > _INT32_MAX or token_start != int(token_start) + for token_start in token_starts + ) ): return False if (draft_manager is None) != (self.draft_block_offsets_device is None): @@ -1284,7 +1356,9 @@ def stage( if manager.enable_swa_scratch_reuse: raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") try: - request_metadata = torch.as_tensor((round_starts, seq_lens), dtype=torch.int32) + request_metadata = torch.as_tensor( + (round_starts, seq_lens, token_starts), dtype=torch.int32 + ) except (OverflowError, RuntimeError, TypeError, ValueError): return False if not self._stage_page_tables_bulk( @@ -1415,6 +1489,7 @@ class _PreparedEviction: request_id: int seq_len: int round_start: int + prompt_len: int expected_keep_count: int protected_tail: int @@ -1779,7 +1854,7 @@ def _periodic_evict( if not resolved_requests or not self._calibrated: return protected_tails: Dict[int, int] = {} - eviction_groups = {} + due_requests = [] # Resolve every active target cache before changing cadence state. The # captured cache objects also avoid repeating the V2 map lookup here. @@ -1812,8 +1887,7 @@ def _periodic_evict( request_state.generation_steps = step if previous_step // self.beta >= step // self.beta: continue - keep_count = self._minimum_evictable_length(request, seq_len) - if seq_len <= keep_count: + if seq_len <= self._minimum_evictable_length(request, seq_len): continue if self.draft_kv_cache_manager is not None: draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) @@ -1824,27 +1898,26 @@ def _periodic_evict( "be resumed before the final update hook" ) protected_tails[request_id] = protected_tail - prompt_len = min(int(request.py_prompt_len), seq_len) - key = (prompt_len, keep_count) - eviction_groups.setdefault(key, []).append((request, request_id)) + due_requests.append((request, request_id)) # (2) Compact all affected dense and kernel-masked SWA layers, then release # the unreachable tail directly through V2's public resize primitive. - if not eviction_groups: + # Prompt lengths are per-request metadata, so the whole due cohort runs + # as one batched round (chunked only by the staging memory bound). + if not due_requests: return num_layers = self._num_layers_from_manager() - for group in eviction_groups.values(): - for begin in range(0, len(group), _EAGER_REQUEST_CHUNK_SIZE): - chunk = group[begin : begin + _EAGER_REQUEST_CHUNK_SIZE] - chunk_tails = {rid: protected_tails[rid] for _, rid in chunk} - with nvtx_range_debug("triattention.evict_request_group", color="purple"): - capacity_targets = self._evict_requests( - chunk, - num_layers, - protected_tail_lengths=chunk_tails, - ) - self._resize_compacted_requests(capacity_targets, protected_tails) + for begin in range(0, len(due_requests), _EAGER_REQUEST_CHUNK_SIZE): + chunk = due_requests[begin : begin + _EAGER_REQUEST_CHUNK_SIZE] + chunk_tails = {rid: protected_tails[rid] for _, rid in chunk} + with nvtx_range_debug("triattention.evict_request_group", color="purple"): + capacity_targets = self._evict_requests( + chunk, + num_layers, + protected_tail_lengths=chunk_tails, + ) + self._resize_compacted_requests(capacity_targets, protected_tails) def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: if not capacity_targets: @@ -1943,6 +2016,7 @@ def _build_cross_request_keep_set_selector( *, input_scores: Optional[torch.Tensor] = None, normalize_scores: bool = True, + prompt_offsets_buffer: Optional[torch.Tensor] = None, ) -> Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector]: """Allocate one fixed ``[request, ...]`` keep-set selector.""" if plan.eviction_mode == "union": @@ -1950,7 +2024,6 @@ def _build_cross_request_keep_set_selector( plan.rows, plan.width, plan.keep_count, - plan.prompt_len, dtype=plan.dtype, device=plan.device, max_requests=plan.max_requests, @@ -1959,6 +2032,7 @@ def _build_cross_request_keep_set_selector( num_kv_heads=plan.num_kv_heads, input_scores=input_scores, normalize_scores=normalize_scores, + prompt_offsets_buffer=prompt_offsets_buffer, ) return _BatchedPerHeadKeepSetSelector( eviction_mode=plan.eviction_mode, @@ -1967,10 +2041,10 @@ def _build_cross_request_keep_set_selector( num_kv_heads=plan.num_kv_heads, width=plan.width, keep_count=plan.keep_count, - prompt_len=plan.prompt_len, dtype=plan.dtype, device=plan.device, max_requests=plan.max_requests, + prompt_offsets_buffer=prompt_offsets_buffer, ) def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: @@ -2298,10 +2372,9 @@ def _eager_resources_for( """Build or reuse eager score and selection buffers for one cohort.""" if not prepared: raise ValueError("TriAttention eviction requires at least one request") - prompt_lens = {min(int(item.request.py_prompt_len), item.seq_len) for item in prepared} - if len(prompt_lens) != 1: - raise ValueError("TriAttention batches require one common prompt length") - prompt_len = next(iter(prompt_lens)) + # Prompt lengths are per-request runtime metadata; the smallest one + # only sizes the widest decode window of the shared buffers. + min_prompt_len = min(item.prompt_len for item in prepared) seq_len = max(item.seq_len for item in prepared) page_table_token_capacity = max(item.seq_len + item.protected_tail for item in prepared) request_count = len(prepared) @@ -2333,7 +2406,7 @@ def _eager_resources_for( self.eviction_mode, request_count, seq_len, - prompt_len, + min_prompt_len, page_table_token_capacity, self.top_B, tuple(layout.dense_layers), @@ -2368,7 +2441,7 @@ def _eager_resources_for( omega=self.calibration["omega"], page_table_keys=self._page_table_pool_keys(representatives, layout.global_layers), num_page_table_slots=layout.manager.num_pools, - prompt_len=prompt_len, + min_prompt_len=min_prompt_len, page_table_token_capacity=page_table_token_capacity, **draft_kwargs, ) @@ -2379,9 +2452,8 @@ def _eager_resources_for( num_query_heads=int(self._H), num_kv_heads=int(first_pool.shape[2]), rows=len(layout.dense_layers) * int(self._H), - width=seq_len - prompt_len, + width=seq_len - min_prompt_len, keep_count=self.top_B, - prompt_len=prompt_len, dtype=torch.float32, device=first_pool.device, max_requests=request_count, @@ -2389,9 +2461,10 @@ def _eager_resources_for( input_scores=score_staging.fused_group.output.view( request_count, len(layout.dense_layers) * int(self._H), - seq_len - prompt_len, + seq_len - min_prompt_len, ), normalize_scores=self.normalize_scores, + prompt_offsets_buffer=score_staging.token_starts_device, ) score_staging.bind_score_launcher( keep_set_selector.valid_widths, @@ -2436,6 +2509,16 @@ def _batched_compaction_for( draft_kv_block_offsets=score_staging.draft_block_offsets_device, draft_page_table_slots=score_staging.draft_representative_slots, ) + if layout.swa_layers and layout.swa_window: + # SWA landing positions are prompt-dependent; reject a request + # whose retained span cannot cover the model window this round. + for item in prepared: + if item.prompt_len + self.top_B < int(layout.swa_window): + raise ValueError( + f"Request {item.request_id} retains " + f"{item.prompt_len + self.top_B} tokens, below the " + f"sliding window {layout.swa_window}" + ) key = ( id(score_staging), id(keep_set_selector), @@ -2458,7 +2541,7 @@ def _batched_compaction_for( kv_block_offsets=score_staging.block_offsets_device, page_table_slots=score_staging.representative_slots, request_count=len(prepared), - prompt_len=keep_set_selector.prompt_len, + prompt_offsets=score_staging.token_starts_device[: len(prepared)], decode_keep_count=self.top_B, swa_window=layout.swa_window, protected_tail_lengths=list(protected_tail_lengths), @@ -2512,6 +2595,7 @@ def _attach_page_ids( self.kv_cache_manager, [item.request_id for item in prepared], [item.round_start for item in prepared], + [item.prompt_len for item in prepared], [item.seq_len for item in prepared], [item.seq_len + item.protected_tail for item in prepared], draft_manager=self.draft_kv_cache_manager, @@ -2567,6 +2651,7 @@ def _evict_requests( request_id=rid, seq_len=int(seq_len), round_start=int(round_start), + prompt_len=min(int(request.py_prompt_len), int(seq_len)), expected_keep_count=expected_keep_count, protected_tail=protected_tail, ) @@ -2585,6 +2670,9 @@ def _evict_requests( ) with nvtx_range_debug("triattention.page_table_stage", color="orange"): self._attach_page_ids(prepared, score_staging) + # The staged per-request prompt lengths are shared with the + # selector; per-head modes re-expand them to selection rows here. + keep_set_selector.refresh_row_prompt_offsets() try: with nvtx_range("triattention.score", color="blue"): diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index e988cf28d548..729479d720d2 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -376,7 +376,6 @@ def _pack_compaction_sources_kernel( SELECTION_ROWS: tl.constexpr, SELECTION_STRIDE: tl.constexpr, KEEP_COUNT: tl.constexpr, - PROMPT_LEN: tl.constexpr, NUM_KV_HEADS: tl.constexpr, SWA_WINDOW: tl.constexpr, UNION: tl.constexpr, @@ -398,9 +397,11 @@ def _pack_compaction_sources_kernel( selection_domain = 0 else: selection_domain = domain + # Selection rows carry decode-only kept ordinals (already absolute), so + # rows are prompt-length independent and one cohort may mix prompt sizes. selection_row = request * SELECTION_ROWS + selection_domain selected = tl.load( - selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + PROMPT_LEN + move, + selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + move, mask=move < KEEP_COUNT, other=0, ) @@ -432,19 +433,23 @@ def _pack_compaction_sources_kernel( def _finalize_topk_indices_kernel( scores, seq_lens, + prompt_offsets, provisional_indices, output_indices, WIDTH: tl.constexpr, KEEP_COUNT: tl.constexpr, OUTPUT_WIDTH: tl.constexpr, - PROMPT_LEN: tl.constexpr, BLOCK: tl.constexpr, ): """Resolve boundary ties and emit increasing physical token indices.""" row = tl.program_id(0) row_scores = scores + row * WIDTH row_selected = provisional_indices + row * KEEP_COUNT - row_output = output_indices + row * OUTPUT_WIDTH + PROMPT_LEN + row_output = output_indices + row * OUTPUT_WIDTH + # Scores are decode-relative; this row's pinned prompt length rebases the + # emitted ordinals to absolute positions (per row, so one launch may mix + # prompt lengths). + prompt_len = tl.load(prompt_offsets + row) threshold = float("inf") for start in tl.static_range(0, KEEP_COUNT, BLOCK): @@ -494,7 +499,7 @@ def _finalize_topk_indices_kernel( write_offset = output_count + tl.cumsum(selected_i32, axis=0) - selected_i32 tl.store( row_output + write_offset, - token_index + PROMPT_LEN, + token_index + prompt_len, mask=selected, ) output_count += tl.sum(selected_i32) @@ -504,33 +509,34 @@ def _finalize_topk_indices_kernel( def finalize_topk_indices( scores: torch.Tensor, seq_lens: torch.Tensor, + prompt_offsets: torch.Tensor, provisional_indices: torch.Tensor, output_indices: torch.Tensor, keep_count: int, - prompt_len: int, ) -> None: """Finalize one provisional TopK set without changing the CuTE selector. The kernel derives the provisional set's threshold, keeps all strictly better scores, resolves the remaining boundary ties by lower token index, - and writes increasing physical ordinals directly into ``output_indices``. + and writes increasing physical ordinals (decode index + this row's prompt + offset) directly into ``output_indices``. """ keep_count = int(keep_count) - prompt_len = int(prompt_len) if not scores.is_cuda: raise ValueError("deterministic TopK finalization requires CUDA scores") if scores.ndim != 2 or scores.dtype != torch.float32 or not scores.is_contiguous(): raise ValueError("TopK finalization requires contiguous two-dimensional FP32 scores") rows, width = scores.shape - if rows <= 0 or not 1 <= keep_count <= width or prompt_len < 0: - raise ValueError("TopK finalization requires valid rows, keep count, and prompt length") - if ( - seq_lens.shape != (rows,) - or seq_lens.dtype != torch.int32 - or seq_lens.device != scores.device - or not seq_lens.is_contiguous() - ): - raise ValueError("TopK finalization sequence lengths do not match the score rows") + if rows <= 0 or not 1 <= keep_count <= width: + raise ValueError("TopK finalization requires valid rows and keep count") + for name, tensor in (("sequence lengths", seq_lens), ("prompt offsets", prompt_offsets)): + if ( + tensor.shape != (rows,) + or tensor.dtype != torch.int32 + or tensor.device != scores.device + or not tensor.is_contiguous() + ): + raise ValueError(f"TopK finalization {name} do not match the score rows") if ( provisional_indices.shape != (rows, keep_count) or provisional_indices.dtype != torch.int32 @@ -541,7 +547,7 @@ def finalize_topk_indices( if ( output_indices.ndim != 2 or output_indices.shape[0] != rows - or output_indices.shape[1] < prompt_len + keep_count + or output_indices.shape[1] < keep_count or output_indices.dtype != torch.int32 or output_indices.device != scores.device or not output_indices.is_contiguous() @@ -550,12 +556,12 @@ def finalize_topk_indices( _finalize_topk_indices_kernel[(rows,)]( scores, seq_lens, + prompt_offsets, provisional_indices, output_indices, WIDTH=width, KEEP_COUNT=keep_count, OUTPUT_WIDTH=output_indices.shape[1], - PROMPT_LEN=prompt_len, BLOCK=256, num_warps=4, ) @@ -578,6 +584,9 @@ def _tri_score_perhead_kernel( req_seq_len, # [num_requests] int32 req_valid_width_out, # [num_requests] int32: decode-only length for selection req_round_start, # [num_requests] int32 logical token position + req_token_start, # [num_requests] int32: pinned prompt length; scoring + # starts at this decode-region origin, so prompt lengths may differ + # across the cohort. # per-LAYER calibration, [L,H,F] flattened layer-major: q_real_ptr, # [L*H*F] fp32 q_imag_ptr, # [L*H*F] fp32 @@ -606,7 +615,6 @@ def _tri_score_perhead_kernel( s_slot, s_dim, USE_MAX: tl.constexpr, - TOKEN_START: tl.constexpr, T_BLOCK: tl.constexpr, F_BLOCK: tl.constexpr, ): @@ -624,11 +632,12 @@ def _tri_score_perhead_kernel( req_id = tl.load(seg_req_id + seg) seq_len = tl.load(req_seq_len + req_id) + token_start = tl.load(req_token_start + req_id) if (seg % num_layers == 0) & (t_blk == 0) & (kv_head == 0): - tl.store(req_valid_width_out + req_id, seq_len - TOKEN_START) + tl.store(req_valid_width_out + req_id, seq_len - token_start) # Derive the ragged launch bound in the score program instead of staging # one replicated length and block count for every request/layer segment. - n_tblk = (seq_len - TOKEN_START + T_BLOCK - 1) // T_BLOCK + n_tblk = (seq_len - token_start + T_BLOCK - 1) // T_BLOCK if t_blk >= n_tblk: return @@ -650,7 +659,7 @@ def _tri_score_perhead_kernel( # ---- token tile of THIS segment ---- t = t_blk * T_BLOCK + tl.arange(0, T_BLOCK) - absolute_t = t + TOKEN_START + absolute_t = t + token_start t_mask = absolute_t < seq_len blk_in_seq = absolute_t // tokens_per_block slot = (absolute_t % tokens_per_block).to(tl.int64) @@ -738,7 +747,6 @@ def _launch_tri_score_perhead( geometry_args: tuple, *, score_aggregation: str, - token_start: int, token_block: int, num_freqs: int, ) -> None: @@ -749,7 +757,6 @@ def _launch_tri_score_perhead( *pointer_args, *geometry_args, USE_MAX=(score_aggregation == "max"), - TOKEN_START=token_start, T_BLOCK=token_block, F_BLOCK=triton.next_power_of_2(num_freqs), ) @@ -780,17 +787,18 @@ def __init__( freq_scale_sq: torch.Tensor, omega: torch.Tensor, offsets: torch.Tensor, - prompt_len: int = 0, + min_prompt_len: int = 0, ) -> None: if not layer_indices or min(max_requests, page_count, seq_len) <= 0: raise ValueError("fixed score group requires non-empty positive geometry") - if prompt_len < 0 or prompt_len >= seq_len: + if min_prompt_len < 0 or min_prompt_len >= seq_len: raise ValueError("fixed score prompt length must leave a non-empty decode region") if len(page_table_slots) != len(layer_indices): raise ValueError("page_table_slots must align with layer_indices") self.max_requests = max_requests - self.prompt_len = prompt_len - self.output_width = seq_len - prompt_len + # Prompt lengths are per-request kernel inputs; the smallest one only + # sizes the widest possible decode window of the output buffer. + self.output_width = seq_len - min_prompt_len self.num_layers = len(layer_indices) p0 = layer_pools[layer_indices[0]] if p0.ndim != 5: @@ -896,6 +904,7 @@ def launch( valid_seq_lens: torch.Tensor, valid_widths: torch.Tensor, round_starts_device: torch.Tensor, + token_starts_device: torch.Tensor, mean_cos: torch.Tensor, mean_sin: torch.Tensor, score_aggregation: str, @@ -923,6 +932,7 @@ def launch( valid_seq_lens, valid_widths, round_starts_device, + token_starts_device, *self.pointer_middle, mean_cos.view(-1), mean_sin.view(-1), @@ -931,7 +941,6 @@ def launch( ), (self.output_width, self.num_layers, *self.geometry_args), score_aggregation=score_aggregation, - token_start=self.prompt_len, token_block=self.token_block, num_freqs=self.num_freqs, ) 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 index d0225105e84e..7f8afdca0034 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -131,9 +131,9 @@ def _launched_draft_compaction(draft_protected_tails): initial_target = [pool.clone() for pool in target_pools] initial_draft = draft_pool.clone() - prompt = torch.tensor([0, 1], dtype=torch.int64, device=device) - union_decode = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) - keep = torch.cat((prompt.view(1, -1).expand(request_count, -1), union_decode), dim=1) + # Kept ordinals are decode-only but absolute; the pinned prompt tokens + # never appear in the selection rectangle. + keep = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) compaction = BatchedKVCacheCompaction( eviction_mode="union", @@ -147,7 +147,7 @@ def _launched_draft_compaction(draft_protected_tails): kv_block_offsets=_encode_block_offsets(target_tables), page_table_slots={0: 0, 1: 0}, request_count=request_count, - prompt_len=prompt_len, + prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, swa_window=None, protected_tail_lengths=target_protected_tails, @@ -195,7 +195,7 @@ def test_draft_pools_receive_target_union_keep_set_and_own_tail(): dtype=torch.int64, device=device, ) - target_source = torch.cat((built.keep[request, prompt_len:], target_tail)) + target_source = torch.cat((built.keep[request], target_tail)) target_destination = torch.arange( prompt_len, prompt_len + target_source.numel(), @@ -221,7 +221,7 @@ def test_draft_pools_receive_target_union_keep_set_and_own_tail(): dtype=torch.int64, device=device, ) - draft_source = torch.cat((built.keep[request, prompt_len:], draft_tail)) + draft_source = torch.cat((built.keep[request], draft_tail)) draft_destination = torch.arange( prompt_len, prompt_len + draft_source.numel(), @@ -246,7 +246,7 @@ def test_draft_pack_matches_keep_broadcast_and_tail_ordinal_oracle(draft_protect expected_offsets = [0] expected_moves = [] for request in range(built.request_count): - decode = built.keep[request, built.prompt_len :].to(torch.int32) + decode = built.keep[request].to(torch.int32) tail = torch.arange( built.valid_seq_lens[request], built.valid_seq_lens[request] + draft_protected_tails[request], @@ -345,7 +345,10 @@ def _mocked_eviction_internals(manager): launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), mark_page_tables_consumed=mock.Mock(), ) - keep_set_selector = SimpleNamespace(select_requests=mock.Mock()) + keep_set_selector = SimpleNamespace( + select_requests=mock.Mock(), + refresh_row_prompt_offsets=mock.Mock(), + ) resources = SimpleNamespace( score_staging=score_staging, keep_set_selector=keep_set_selector, @@ -479,6 +482,7 @@ def test_lru_score_staging_eviction_drops_dependent_batched_compactions(): score_staging = SimpleNamespace( fused_group=SimpleNamespace(output=torch.empty(1, 4, 8)), bind_score_launcher=mock.Mock(), + token_starts_device=torch.zeros(1, dtype=torch.int32), ) keep_set_selector = SimpleNamespace(valid_widths=torch.empty(1, dtype=torch.int32)) prepared = [ @@ -487,6 +491,7 @@ def test_lru_score_staging_eviction_drops_dependent_batched_compactions(): request_id=7, seq_len=8, round_start=8, + prompt_len=0, expected_keep_count=4, protected_tail=0, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index 1e20b5ab2f53..de5626468567 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -100,11 +100,11 @@ def test_union_eager_uses_one_deterministic_cute_selection(keep_count, width): rows=2, width=width, keep_count=keep_count, - prompt_len=prompt_len, dtype=torch.float32, device=torch.device("cpu"), max_requests=1, ) + selector.set_prompt_offsets(torch.tensor([prompt_len], dtype=torch.int32)) raw_topk = _AdversarialTieTopK() with ( mock.patch.object( @@ -124,7 +124,7 @@ def test_union_eager_uses_one_deterministic_cute_selection(keep_count, width): expected = _stable_topk(scores.max(dim=0).values, width, keep_count) assert torch.equal( - selector.keep[0, prompt_len:], + selector.keep[0], torch.sort(expected + prompt_len).values, ) assert raw_topk.calls == 1 @@ -139,11 +139,11 @@ def test_per_head_eager_keeps_stable_indices(eviction_mode): num_kv_heads=2, width=16, keep_count=5, - prompt_len=3, dtype=torch.float32, device=torch.device("cpu"), max_requests=1, ) + selector.set_prompt_offsets(torch.tensor([3], dtype=torch.int32)) scores = torch.arange(2 * 4 * 16, dtype=torch.float32).reshape(2, 4, 16) with mock.patch.object( torch.ops.trtllm, @@ -155,9 +155,15 @@ def test_per_head_eager_keeps_stable_indices(eviction_mode): assert tuple(selector.keep.shape) == ( 1, selector.selection_rows, - selector.total_keep, + selector.keep_count, + ) + # Scores increase with the token index, so every row keeps the last five + # decode ordinals, rebased by the pinned prompt length. + expected_row = torch.arange(16 - 5, 16, dtype=torch.int32) + 3 + assert torch.equal( + selector.keep, + expected_row.expand(1, selector.selection_rows, -1), ) - assert torch.all(selector.keep[..., 1:] >= selector.keep[..., :-1]) @pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) @@ -166,7 +172,7 @@ def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( eviction_mode, normalize_scores ): request_count, layers, query_heads, kv_heads = 2, 3, 4, 2 - width, keep_count, prompt_len = 96, 64, 0 + width, keep_count = 96, 64 generator = torch.Generator().manual_seed(41) scores_cpu = torch.randint( -4, @@ -184,7 +190,6 @@ def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( num_kv_heads=kv_heads, width=width, keep_count=keep_count, - prompt_len=prompt_len, dtype=torch.float32, device=torch.device("cpu"), max_requests=request_count, @@ -209,7 +214,6 @@ def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( num_kv_heads=kv_heads, width=width, keep_count=keep_count, - prompt_len=prompt_len, dtype=torch.float32, device=device, max_requests=request_count, @@ -234,7 +238,6 @@ def test_union_eager_runs_the_registered_cute_op(): rows=4, width=96, keep_count=64, - prompt_len=0, dtype=torch.float32, device=device, max_requests=2, @@ -280,7 +283,6 @@ def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_ rows=rows, width=width, keep_count=keep_count, - prompt_len=0, dtype=torch.float32, device=device, max_requests=request_count, @@ -327,7 +329,6 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid rows=rows, width=width, keep_count=keep_count, - prompt_len=prompt_len, dtype=torch.float32, device=device, max_requests=request_count, @@ -335,17 +336,18 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid normalize_scores=False, ) selector.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) + selector.set_prompt_offsets( + torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) + ) selector.select_prepared_requests() actual = selector.keep.cpu() - expected_prompt = torch.arange(prompt_len, dtype=torch.int32) combined = scores.amax(dim=1).cpu() for request, valid_width in enumerate(valid_widths): expected_decode = torch.sort( _stable_topk(combined[request], valid_width, keep_count).to(torch.int32) + prompt_len ).values - assert torch.equal(actual[request, :prompt_len], expected_prompt) - assert torch.equal(actual[request, prompt_len:], expected_decode) + assert torch.equal(actual[request], expected_decode) def test_fused_union_preparation_matches_ragged_torch_reference(): @@ -487,14 +489,15 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) ] pools = [pool.clone() for pool in initial_pools] - prompt = torch.tensor([0, 1], dtype=torch.int64, device=device) + # Kept ordinals are decode-only but hold absolute positions; the pinned + # prompt tokens never appear in the selection rectangle. union_decode = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) if eviction_mode == "union": - keep = torch.cat((prompt.view(1, -1).expand(request_count, -1), union_decode), dim=1) + keep = union_decode selection_rows = 1 else: selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads - decode = torch.empty( + keep = torch.empty( request_count, selection_rows, decode_keep_count, @@ -503,7 +506,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) ) for request in range(request_count): for row in range(selection_rows): - decode[request, row] = torch.tensor( + keep[request, row] = torch.tensor( sorted( { 2 + ((request + row + offset * 2) % 8) @@ -513,13 +516,6 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) dtype=torch.int64, device=device, ) - keep = torch.cat( - ( - prompt.view(1, 1, -1).expand(request_count, selection_rows, -1), - decode, - ), - dim=2, - ) compaction = BatchedKVCacheCompaction( eviction_mode=eviction_mode, @@ -533,7 +529,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), page_table_slots={0: 0, 1: 0}, request_count=request_count, - prompt_len=prompt_len, + prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, swa_window=None, protected_tail_lengths=protected_tails, @@ -553,11 +549,11 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) for head in range(num_kv_heads): if eviction_mode == "union": - selected = keep[request, prompt_len:] + selected = keep[request] elif eviction_mode == "per_head": - selected = keep[request, head, prompt_len:] + selected = keep[request, head] else: - selected = keep[request, layer * num_kv_heads + head, prompt_len:] + selected = keep[request, layer * num_kv_heads + head] tail = torch.arange( seq_len, seq_len + protected_tails[request], @@ -577,6 +573,119 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) ) +def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): + """One union cohort mixing prompt lengths compacts byte-identically to + running the same two requests as two single-request compactions.""" + device = torch.device("cuda") + request_count = 2 + num_layers = 2 + num_kv_heads = 2 + seq_len = 10 + decode_keep_count = 3 + prompt_lens = [2, 5] + protected_tails = [2, 1] + tokens_per_block = 4 + head_dim = 16 + decode_widths = [seq_len - prompt_len for prompt_len in prompt_lens] + width = max(decode_widths) + + # CPU-oracle selection: decode-relative scores per request, rebased to + # absolute ordinals by each request's own prompt offset. + selector = _BatchedUnionKeepSetSelector( + rows=1, + width=width, + keep_count=decode_keep_count, + dtype=torch.float32, + device=torch.device("cpu"), + max_requests=request_count, + ) + selector.valid_widths.copy_(torch.tensor(decode_widths, dtype=torch.int32)) + selector.set_prompt_offsets(torch.tensor(prompt_lens, dtype=torch.int32)) + generator = torch.Generator().manual_seed(11) + scores = torch.randint( + -8, + 9, + (request_count, 1, width), + generator=generator, + dtype=torch.int32, + ).to(torch.float32) + reference_scores = scores.clone() + with mock.patch.object( + torch.ops.trtllm, + "cute_dsl_indexer_topk_decode", + side_effect=_fake_cute_topk, + create=True, + ): + selector.select_requests(scores, normalize_scores=False) + keep = selector.keep.clone() + for request, (prompt_len, decode_width) in enumerate(zip(prompt_lens, decode_widths)): + expected = torch.sort( + _stable_topk(reference_scores[request, 0], decode_width, decode_keep_count) + prompt_len + ).values + assert torch.equal(keep[request], expected) + + page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + initial_pools = [ + ( + torch.arange( + 6 * 2 * num_kv_heads * tokens_per_block * head_dim, + dtype=torch.float32, + device=device, + ).view(6, 2, num_kv_heads, tokens_per_block, head_dim) + + layer * 100_000.0 + ) + for layer in range(num_layers) + ] + cohort_pools = [pool.clone() for pool in initial_pools] + keep_cuda = keep.to(device) + valid_seq_lens = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) + cohort_compaction = BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=cohort_pools, + dense_layers=[0, 1], + swa_layers=[], + layer_group_representative={0: 0, 1: 1}, + layer_pool_keys=[("dense", 0), ("dense", 0)], + kept_token_ordinals=keep_cuda, + valid_sequence_lengths=valid_seq_lens, + kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), + page_table_slots={0: 0, 1: 0}, + request_count=request_count, + prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), + decode_keep_count=decode_keep_count, + swa_window=None, + protected_tail_lengths=protected_tails, + ) + cohort_compaction.launch() + + expected_pools = [pool.clone() for pool in initial_pools] + for request in range(request_count): + single_compaction = BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=expected_pools, + dense_layers=[0, 1], + swa_layers=[], + layer_group_representative={0: 0, 1: 1}, + layer_pool_keys=[("dense", 0), ("dense", 0)], + kept_token_ordinals=keep_cuda[request : request + 1], + valid_sequence_lengths=valid_seq_lens[request : request + 1], + kv_block_offsets=_encode_block_offsets(page_tables[request : request + 1].unsqueeze(0)), + page_table_slots={0: 0, 1: 0}, + request_count=1, + prompt_offsets=torch.tensor([prompt_lens[request]], dtype=torch.int32, device=device), + decode_keep_count=decode_keep_count, + swa_window=None, + protected_tail_lengths=[protected_tails[request]], + ) + single_compaction.launch() + torch.cuda.synchronize(device) + + # The two requests own disjoint pages, so whole-pool equality proves the + # cohort produced exactly the two single-request results. + for cohort_pool, expected_pool in zip(cohort_pools, expected_pools): + assert torch.equal(cohort_pool, expected_pool) + + def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): """Keep score and compaction layer axes aligned across interleaved V2 pools.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( @@ -653,6 +762,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): ) score_staging.round_starts_device.fill_(0) score_staging.valid_seq_lens_device.fill_(seq_len) + score_staging.token_starts_device.fill_(0) keep_set_selector = _BatchedPerHeadKeepSetSelector( eviction_mode="per_layer_perhead", dense_layers=tuple(dense_layers), @@ -660,7 +770,6 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): num_kv_heads=1, width=seq_len, keep_count=keep_count, - prompt_len=0, dtype=torch.float32, device=device, max_requests=1, @@ -683,7 +792,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): kv_block_offsets=score_staging.block_offsets_device, page_table_slots=score_staging.representative_slots, request_count=1, - prompt_len=0, + prompt_offsets=torch.zeros(1, dtype=torch.int32, device=device), decode_keep_count=keep_count, swa_window=None, protected_tail_lengths=[0], @@ -810,14 +919,13 @@ def write_token(token: int, score: float) -> None: torch.zeros(num_freqs, dtype=torch.float32, device=device), page_table_keys=[("pool", 0)], num_page_table_slots=1, - prompt_len=prompt_len, + min_prompt_len=prompt_len, page_table_token_capacity=seq_len + protected_tail, ) keep_set_selector = _BatchedUnionKeepSetSelector( rows=1, width=seq_len - prompt_len, keep_count=compacted_capacity - prompt_len - protected_tail, - prompt_len=prompt_len, dtype=torch.float32, device=device, max_requests=1, @@ -826,6 +934,7 @@ def write_token(token: int, score: float) -> None: num_kv_heads=1, input_scores=score_staging.fused_group.output.view(1, 1, seq_len - prompt_len), normalize_scores=False, + prompt_offsets_buffer=score_staging.token_starts_device, ) score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") batched_compaction = BatchedKVCacheCompaction( @@ -840,7 +949,7 @@ def write_token(token: int, score: float) -> None: kv_block_offsets=score_staging.block_offsets_device, page_table_slots=score_staging.representative_slots, request_count=1, - prompt_len=prompt_len, + prompt_offsets=score_staging.token_starts_device[:1], decode_keep_count=compacted_capacity - prompt_len - protected_tail, swa_window=None, protected_tail_lengths=[protected_tail], @@ -852,9 +961,11 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: manager, [request_id], [0], + [prompt_len], [seq_len], [seq_len + protected_tail], ) + keep_set_selector.refresh_row_prompt_offsets() score_staging.launch_prepared_score() keep_set_selector.select_prepared_requests() selected = keep_set_selector.keep[0].clone().to(torch.long) @@ -865,7 +976,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: after = snapshot(compacted_capacity) source = torch.cat( ( - selected[prompt_len:], + selected, torch.arange(seq_len, seq_len + protected_tail, device=device), ) ) @@ -882,7 +993,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: first_keep, first_compacted = evict_once() assert torch.equal( first_keep, - torch.tensor([0, 1, 2, 4, 6, 8], dtype=torch.long, device=device), + torch.tensor([2, 4, 6, 8], dtype=torch.long, device=device), ) retained_pages = page_ids(request_id) assert torch.equal(retained_pages, initial_pages[:2]) @@ -910,7 +1021,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: second_keep, _ = evict_once() assert torch.equal( second_keep, - torch.tensor([0, 1, 2, 3, 6, 8], dtype=torch.long, device=device), + torch.tensor([2, 3, 6, 8], dtype=torch.long, device=device), ) assert not torch.equal(second_keep, first_keep) @@ -936,8 +1047,9 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): + 1000.0, ] pools = [pool.clone() for pool in initial_pools] + # Decode-only kept ordinals holding absolute positions past the prompt. keep = torch.tensor( - [[0, 1, 2, 4, 5, 7], [0, 1, 2, 3, 5, 6]], + [[2, 4, 5, 7], [2, 3, 5, 6]], dtype=torch.int64, device=device, ) @@ -955,7 +1067,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), page_table_slots={0: 0, 1: 1}, request_count=2, - prompt_len=2, + prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), decode_keep_count=4, swa_window=2, protected_tail_lengths=protected_tails, @@ -978,7 +1090,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): dtype=torch.int64, device=device, ) - dense_source = torch.cat((keep[request, 2:], tail)) + dense_source = torch.cat((keep[request], tail)) dense_destination = torch.arange( 2, 2 + dense_source.numel(), dtype=torch.int64, device=device ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 1f32fa261a3e..9f6500562b6f 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -96,12 +96,14 @@ def _prepared_eviction( protected_tail=0, request_id=0, round_start=None, + prompt_len=0, ): return _PreparedEviction( request=request, request_id=request_id, seq_len=seq_len, round_start=int(seq_len if round_start is None else round_start), + prompt_len=prompt_len, expected_keep_count=expected_keep_count, protected_tail=protected_tail, ) @@ -481,7 +483,10 @@ def _mocked_eviction_internals(self, manager): launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), mark_page_tables_consumed=mock.Mock(), ) - keep_set_selector = SimpleNamespace(select_requests=mock.Mock()) + keep_set_selector = SimpleNamespace( + select_requests=mock.Mock(), + refresh_row_prompt_offsets=mock.Mock(), + ) resources = SimpleNamespace( score_staging=score_staging, keep_set_selector=keep_set_selector, @@ -904,7 +909,8 @@ def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): ) config = TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8) + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + ) assert config.kv_cache_compression_mode.is_eviction_method() is True draft_manager = _make_fake_v2(is_draft=True) validate_kv_cache_compression_with_spec( @@ -936,7 +942,8 @@ def test_unvalidated_paged_draft_tail_contracts_remain_fail_closed(self, mode): with pytest.raises(ValueError, match="standard paged cache compacted together"): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8), + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + ), spec_config, _make_fake_v2(is_draft=True), ) @@ -954,7 +961,8 @@ def test_dflash_spec_mode_is_rejected(self): with pytest.raises(ValueError, match="standard paged cache compacted together"): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8), + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + ), DFlashDecodingConfig(max_draft_len=3), _make_fake_v2(is_draft=True), ) @@ -1026,7 +1034,6 @@ def test_cross_request_union_uses_cute_without_fallback(self, keep_count): request_scores[0].shape[0], width, keep_count, - 0, dtype=request_scores[0].dtype, device=request_scores[0].device, max_requests=len(request_scores), @@ -1076,6 +1083,7 @@ def test_eager_bucket_binds_score_after_selection(self, eviction_mode, normalize score_staging = SimpleNamespace( fused_group=SimpleNamespace(output=torch.empty(1, 4, 8)), bind_score_launcher=mock.Mock(), + token_starts_device=torch.zeros(1, dtype=torch.int32), ) keep_set_selector = SimpleNamespace(valid_widths=torch.empty(1, dtype=torch.int32)) prepared = [ @@ -1307,7 +1315,7 @@ def test_cross_stream_staging_is_rejected_before_page_table_query(self): other_stream = SimpleNamespace(device=torch.device("cuda:0"), cuda_stream=5) with mock.patch.object(torch.cuda, "current_stream", return_value=other_stream): with pytest.raises(_FixedScoreStreamMismatch, match="first CUDA stream"): - staging.stage(manager, [1], [8.0]) + staging.stage(manager, [1], [8.0], [0]) staging.copy_done.query.assert_not_called() staging.copy_done.synchronize.assert_not_called() @@ -1324,6 +1332,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): request_id=7, round_start=8, seq_len=8, + prompt_len=3, expected_keep_count=6, protected_tail=2, ), @@ -1332,6 +1341,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): request_id=8, round_start=9, seq_len=9, + prompt_len=5, expected_keep_count=6, protected_tail=3, ), @@ -1343,6 +1353,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): manager.kv_cache_manager, [7, 8], [8, 9], + [3, 5], [8, 9], [10, 12], draft_manager=None, @@ -1424,6 +1435,8 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re request_ids = list(range(request_count)) round_starts = [131_071 + request for request in request_ids] + # Per-request pinned prompt lengths: one cohort may mix them. + token_starts = list(request_ids) host_table = torch.zeros( 3, max_requests, @@ -1459,6 +1472,7 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): manager, request_ids, [2**31] * request_count, + token_starts, [seq_len] * request_count, [10] * request_count, ) @@ -1472,6 +1486,7 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): manager, request_ids, round_starts, + token_starts, [seq_len] * request_count, [10] * request_count, ) @@ -1487,6 +1502,10 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): staging.valid_seq_lens_device[:request_count], torch.full((request_count,), seq_len, dtype=torch.int32, device=device), ) + assert torch.equal( + staging.token_starts_device[:request_count], + torch.tensor(token_starts, dtype=torch.int32, device=device), + ) for slot, global_layer in enumerate((10, 12, 13)): expected = torch.tensor(tables[global_layer], dtype=torch.int32, device=device) * 2 assert torch.equal( @@ -1497,7 +1516,7 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): other_stream = torch.cuda.Stream(device=device) with torch.cuda.stream(other_stream): with pytest.raises(_FixedScoreStreamMismatch, match="first CUDA stream"): - staging.stage(manager, request_ids, round_starts) + staging.stage(manager, request_ids, round_starts, token_starts) assert gather.call_count == calls @pytest.mark.parametrize("request_count", [1, 7, 8]) @@ -1538,6 +1557,9 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, assert not offsets.is_contiguous() round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 round_starts = round_device[:request_count].tolist() + token_starts_device = torch.full( + (max_requests,), prompt_len, dtype=torch.int32, device=device + ) seq_lens = [seq_len - request % 2 for request in range(request_count)] phase = (round_device[:, None, None] + offsets[None, :, None]) * omega[None, None] oracle = _torch_tri_score_oracle( @@ -1570,7 +1592,7 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, freq, omega, offsets, - prompt_len=prompt_len, + min_prompt_len=prompt_len, ) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) fixed = group.launch( @@ -1578,6 +1600,7 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, torch.tensor(seq_lens, dtype=torch.int32, device=device), valid_widths, round_device, + token_starts_device, torch.cos(phase).mean(dim=1), torch.sin(phase).mean(dim=1), aggregation, @@ -1653,6 +1676,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun offsets = torch.tensor([1.0, 2.0, 4.0], device=device) round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 round_starts = round_device[:request_count].tolist() + token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) block_offsets = _encode_block_offsets(page_ids_3d) @@ -1671,7 +1695,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun freq, omega, offsets, - prompt_len=prompt_len, + min_prompt_len=prompt_len, ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) @@ -1693,6 +1717,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun valid_seq_lens, valid_widths, round_device, + token_starts, mean_cos, mean_sin, aggregation, @@ -1703,6 +1728,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun staging.fused_group = group staging.round_starts_device = round_device staging.valid_seq_lens_device = valid_seq_lens + staging.token_starts_device = token_starts staging.mean_cos = mean_cos staging.mean_sin = mean_sin staging.offsets = offsets @@ -1779,6 +1805,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun valid_seq_lens, valid_widths, round_device, + token_starts, mean_cos, mean_sin, aggregation, @@ -1858,7 +1885,8 @@ def test_returns_triattention_instance_with_v2(self): # no calibration file or CUDA. fake_v2 = _make_fake_v2(enable_block_reuse=False) cfg = TriAttentionKvCacheCompressionConfig( - top_B=32, beta=16, model_path="/models/test", calibration_path="/calib/test.pt") + top_B=32, beta=16, model_path="/models/test", calibration_path="/calib/test.pt" + ) mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) assert mgr.top_B == 32 diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index c38563f175d3..59f700956549 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -93,7 +93,7 @@ def _reference_compact( page_tables: list[torch.Tensor], source_indices: torch.Tensor, source_offsets: torch.Tensor, - destination_base: int, + destination_base: "int | list[int]", source_layer_indices: Optional[torch.Tensor] = None, ) -> list[torch.Tensor]: original = [pool.clone() for pool in pools] @@ -113,10 +113,15 @@ def _reference_compact( for request in range(_BATCH_SIZE): begin = int(source_offsets[request]) end = int(source_offsets[request + 1]) + request_base = ( + destination_base[request] + if isinstance(destination_base, (list, tuple)) + else destination_base + ) for head in range(_NUM_KV_HEADS): for request_move, global_move in enumerate(range(begin, end)): source_token = int(layer_sources[head, global_move]) - destination_token = destination_base + request_move + destination_token = request_base + request_move source_page = int(raw_page_table[request, source_token // _TOKENS_PER_BLOCK]) destination_page = int( raw_page_table[request, destination_token // _TOKENS_PER_BLOCK] @@ -141,8 +146,16 @@ def _compact( pools: list[torch.Tensor], page_tables: list[torch.Tensor], arguments: _DeviceArguments, - destination_base: int, + destination_base: "int | list[int]", ) -> None: + # The op takes per-request destination bases; scalar test parameters are + # broadcast to the batch here. torch.full stays CUDA-graph-capturable. + if isinstance(destination_base, int): + destination_bases = torch.full( + (_BATCH_SIZE,), destination_base, dtype=torch.int32, device="cuda" + ) + else: + destination_bases = torch.tensor(destination_base, dtype=torch.int32, device="cuda") torch.ops.trtllm.sparse_kv_cache_compact_layers( pools, arguments.pool_pointers, @@ -150,7 +163,7 @@ def _compact( arguments.source_indices, arguments.source_offsets, arguments.source_layer_indices, - destination_base, + destination_bases, ) @@ -191,6 +204,31 @@ def test_sparse_kv_cache_compact_layers(dtype, head_dim, destination_base, page_ assert torch.equal(actual.cpu(), reference) +def test_sparse_kv_cache_compact_layers_per_request_destination_bases(): + # One launch may mix pinned-prompt lengths: each request lands at its own + # destination base. + pools_cpu, pools, page_tables = _make_pools(2, torch.bfloat16, 64) + page_tables_cpu = [page_table.cpu() for page_table in page_tables] + source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) + source_row = torch.tensor([3, 5, 8, 6, 7, 10], dtype=torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + destination_bases = [2, 5] + expected = _reference_compact( + pools_cpu, + page_tables_cpu, + source_indices, + source_offsets, + destination_bases, + ) + arguments = _device_arguments(pools, source_indices, source_offsets) + + _compact(pools, page_tables, arguments, destination_bases) + torch.cuda.synchronize() + + for actual, reference in zip(pools, expected): + assert torch.equal(actual.cpu(), reference) + + def test_sparse_kv_cache_compact_layers_per_layer_source(): pools_cpu, pools, page_tables = _make_pools(2, torch.bfloat16, 64) page_tables_cpu = [page_table.cpu() for page_table in page_tables] From 3c7f3a29c5956993051f1335b4693a4974378149 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 21:36:31 -0700 Subject: [PATCH 022/178] [None][fix] Update the compact-layers instantiation macro signature The explicit instantiations still declared the scalar destination base; match the per-request destinationBases pointer. Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 4eceb5467338..fbf04669d702 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -2060,7 +2060,7 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con #define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(T) \ template void invokeSparseKvCacheCompactV2Layers(int64_t const*, int32_t const*, int32_t, int64_t, \ - int32_t const*, int32_t const*, int64_t, int32_t const*, int32_t, int32_t, int32_t, int32_t, int32_t, \ + int32_t const*, int32_t const*, int64_t, int32_t const*, int32_t const*, int32_t, int32_t, int32_t, int32_t, \ cudaStream_t); } // namespace kernels From f88c95d4fba0fce73994101cc589d70ddc724595 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 21:55:29 -0700 Subject: [PATCH 023/178] [None][fix] Order destination_bases before the optional layer indices Torch schemas reject a non-default positional parameter after a defaulted one; stub generation aborted on import. Signed-off-by: tianruih --- cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp | 6 +++--- .../_torch/kv_cache_compression/triattention/compaction.py | 2 +- .../_torch/thop/serial/test_sparse_kv_cache_compact.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index 126cf32b4b85..c01a396b21bc 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -42,7 +42,7 @@ namespace torch_ext //! safe. void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor const& poolPointers, th::Tensor const& pageTable, th::Tensor const& sourceIndices, th::Tensor const& sourceOffsets, - std::optional const& sourceLayerIndices, th::Tensor const& destinationBases) + th::Tensor const& destinationBases, std::optional const& sourceLayerIndices) { TORCH_CHECK(!pools.empty(), "sparse_kv_cache_compact_layers: pools must be non-empty"); @@ -162,8 +162,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "sparse_kv_cache_compact_layers(Tensor(a!)[] pools, Tensor pool_pointers, Tensor page_table, Tensor " - "source_indices, Tensor source_offsets, Tensor? source_layer_indices=None, " - "Tensor destination_bases) -> ()"); + "source_indices, Tensor source_offsets, Tensor destination_bases, " + "Tensor? source_layer_indices=None) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 3cb472678b3b..bcf99b978af7 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -54,8 +54,8 @@ def launch( self.page_table, source, offsets, - self.source_layer_indices, destination_bases, + self.source_layer_indices, ) diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 59f700956549..3075f2bef739 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -162,8 +162,8 @@ def _compact( page_tables[0], arguments.source_indices, arguments.source_offsets, - arguments.source_layer_indices, destination_bases, + arguments.source_layer_indices, ) From 31e55fcb5da8fd769f7e105817617683ed6378ae Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 22:43:46 -0700 Subject: [PATCH 024/178] [None][refactor] Replace per-shape eviction caches with one resident buffer set The eviction path kept LRU caches of score staging, keep-set selectors, and batched compaction launches keyed by exact cohort shapes, and split large due cohorts into fixed 32-request chunks. Every new shape paid an allocation storm next to a nearly full KV pool, and the caches were only correct while their keys stayed exact. Now one buffer set is built at first eviction and stays resident: the request capacity follows the executor max batch size, the decode width follows the eviction bound (top_B plus two eviction periods plus draft tokens), and the page-table plane follows max_seq_len plus the protected tail capacity. Buffers rebuild only when the V2 pool views change or a round outgrows them. One generation batch is one eviction round; padding rows carry zero widths and move nothing. Protected tails become per-round state: BatchedKVCacheCompaction sizes its move buffers by a tail capacity and set_protected_tails() loads each cohort per-request tails into the move offsets consumed by the pack kernel and the C++ compact launches. Signed-off-by: tianruih --- .../triattention/compaction.py | 104 ++++--- .../triattention/triattention.py | 259 +++++++++--------- .../triattention/triattention_kernels.py | 12 +- .../test_triattention_draft_cocompaction.py | 69 ++--- .../test_triattention_eager.py | 18 +- .../test_triattention_pipeline.py | 33 ++- 6 files changed, 274 insertions(+), 221 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index bcf99b978af7..aa28b4f3c87c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -80,9 +80,7 @@ def launch(self) -> None: if self.prepared_move_index_pack is not None: self.prepared_move_index_pack() for group in self.cpp_launch_groups: - group.launch( - self.move_source_indices, self.move_source_offsets, self.destination_bases - ) + group.launch(self.move_source_indices, self.move_source_offsets, self.destination_bases) def _cuda_int32_contiguous(tensors: Tuple[torch.Tensor, ...], device: torch.device) -> bool: @@ -144,20 +142,6 @@ def _make_move_buffers( return indices, torch.tensor(offsets, dtype=torch.int32, device=device) -def _validated_tail_lengths( - tail_lengths: Optional[List[int]], - request_count: int, - what: str, -) -> Tuple[int, ...]: - if tail_lengths is None: - tail_lengths = [0] * request_count - if len(tail_lengths) != request_count: - raise ValueError(f"{what} lengths must match the request count") - if any(length < 0 for length in tail_lengths): - raise ValueError(f"{what} lengths must be non-negative") - return tuple(int(length) for length in tail_lengths) - - def _page_table_provider( page_table_slots: Dict[int, int], kv_block_offsets: torch.Tensor, @@ -362,9 +346,11 @@ class BatchedKVCacheCompaction: `page_table_slots` / `layer_group_representative`: map each layer's group representative to its snapshot slot; layers that share a slot must share one block-offset table. - `protected_tail_lengths`: per-request KV positions past the valid - length reserved for a forward already in flight; they move with - the kept tokens. + `protected_tail_capacity`: widest per-request protected tail this + object must support. A protected tail covers KV positions past + the valid length reserved for a forward already in flight; the + actual per-round lengths are loaded via `set_protected_tails` + and move with the kept tokens. `draft_*`: co-compressed draft-cache layout (union mode only); the draft reuses the target keep set and pins the same prompt. """ @@ -385,13 +371,13 @@ def __init__( prompt_offsets: torch.Tensor, decode_keep_count: int, swa_window: Optional[int], - protected_tail_lengths: Optional[List[int]] = None, + protected_tail_capacity: int = 0, layer_pool_keys: Optional[List[object]] = None, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_layers: Optional[List[int]] = None, draft_layer_group_representative: Optional[Dict[int, int]] = None, draft_layer_pool_keys: Optional[List[object]] = None, - draft_protected_tail_lengths: Optional[List[int]] = None, + draft_protected_tail_capacity: Optional[int] = None, draft_kv_block_offsets: Optional[torch.Tensor] = None, draft_page_table_slots: Optional[Dict[int, int]] = None, ) -> None: @@ -427,10 +413,9 @@ def __init__( raise ValueError("per-request prompt offsets do not match the cohort") self.prompt_offsets = prompt_offsets self.decode_keep_count = int(decode_keep_count) - self.protected_tail_lengths = _validated_tail_lengths( - protected_tail_lengths, self.request_count, "protected-tail" - ) - self.max_protected_tail = max(self.protected_tail_lengths, default=0) + if protected_tail_capacity < 0: + raise ValueError("the protected-tail capacity must be non-negative") + self.protected_tail_capacity = int(protected_tail_capacity) self.dense_layers = tuple(int(layer) for layer in dense_layers) self.swa_layers = tuple(int(layer) for layer in swa_layers) if layer_pool_keys is None: @@ -451,7 +436,7 @@ def __init__( ) dense_move_indices, dense_move_offsets = _make_move_buffers( dense_index_prefix, - [self.decode_keep_count + length for length in self.protected_tail_lengths], + [self.decode_keep_count + self.protected_tail_capacity] * self.request_count, self.device, ) page_table_for = _page_table_provider( @@ -480,7 +465,7 @@ def __init__( self.swa_destination_bases = torch.empty_like(self.prompt_offsets) swa_move_indices, swa_move_offsets = _make_move_buffers( (self.num_kv_heads,), - [self.swa_window + length for length in self.protected_tail_lengths], + [self.swa_window + self.protected_tail_capacity] * self.request_count, self.device, ) # SWA layers are staged as their own page-table representatives. @@ -500,7 +485,7 @@ def __init__( decode_keep_count=self.decode_keep_count, num_dense_layers=len(self.dense_layers), num_kv_heads=self.num_kv_heads, - max_protected_tail=self.max_protected_tail, + max_protected_tail=self.protected_tail_capacity, swa_window=self.swa_window, swa_move_source_offsets=swa_move_offsets, swa_move_source_indices=swa_move_indices, @@ -526,6 +511,7 @@ def __init__( ) self.draft_compaction = None + self.draft_protected_tail_capacity = 0 if draft_layers: self.draft_compaction = self._build_draft_compaction( kept_token_ordinals, @@ -534,7 +520,7 @@ def __init__( draft_layers=draft_layers, draft_layer_group_representative=draft_layer_group_representative, draft_layer_pool_keys=draft_layer_pool_keys, - draft_protected_tail_lengths=draft_protected_tail_lengths, + draft_protected_tail_capacity=draft_protected_tail_capacity, draft_kv_block_offsets=draft_kv_block_offsets, draft_page_table_slots=draft_page_table_slots, ) @@ -558,7 +544,7 @@ def _build_draft_compaction( draft_layers: List[int], draft_layer_group_representative: Optional[Dict[int, int]], draft_layer_pool_keys: Optional[List[object]], - draft_protected_tail_lengths: Optional[List[int]], + draft_protected_tail_capacity: Optional[int], draft_kv_block_offsets: Optional[torch.Tensor], draft_page_table_slots: Optional[Dict[int, int]], ) -> _SingleCacheCompaction: @@ -576,9 +562,9 @@ def _build_draft_compaction( raise ValueError("draft co-compaction requires the full draft layout") if draft_kv_block_offsets is None or draft_page_table_slots is None: raise ValueError("draft co-compaction requires staged draft page tables") - draft_tail_lengths = _validated_tail_lengths( - draft_protected_tail_lengths, self.request_count, "draft protected-tail" - ) + if draft_protected_tail_capacity is not None and draft_protected_tail_capacity < 0: + raise ValueError("the draft protected-tail capacity must be non-negative") + self.draft_protected_tail_capacity = int(draft_protected_tail_capacity or 0) if len(draft_layer_pool_keys) != len(draft_layer_pools): raise ValueError("draft pool keys must match the draft layer-pool count") draft_layers = tuple(int(layer) for layer in draft_layers) @@ -590,7 +576,7 @@ def _build_draft_compaction( ) draft_move_indices, draft_move_offsets = _make_move_buffers( (draft_num_kv_heads,), - [self.decode_keep_count + length for length in draft_tail_lengths], + [self.decode_keep_count + self.draft_protected_tail_capacity] * self.request_count, self.device, ) draft_page_table_for = _page_table_provider( @@ -621,7 +607,7 @@ def _build_draft_compaction( decode_keep_count=self.decode_keep_count, num_dense_layers=1, num_kv_heads=draft_num_kv_heads, - max_protected_tail=max(draft_tail_lengths, default=0), + max_protected_tail=self.draft_protected_tail_capacity, swa_window=0, swa_move_source_offsets=None, swa_move_source_indices=None, @@ -636,6 +622,52 @@ def _build_draft_compaction( destination_bases=self.prompt_offsets, ) + def _write_move_offsets(self, offsets: torch.Tensor, moves_per_request: List[int]) -> None: + cumulative = [0] + for count in moves_per_request: + cumulative.append(cumulative[-1] + count) + # Rows past the cohort are padding and contribute no moves. + cumulative.extend(cumulative[-1:] * (self.request_count - len(moves_per_request))) + offsets.copy_(torch.tensor(cumulative, dtype=torch.int32), non_blocking=True) + + def set_protected_tails( + self, + tail_lengths: List[int], + draft_tail_lengths: Optional[List[int]] = None, + ) -> None: + """Load this cohort's per-request protected tails into the move offsets. + + The pack kernel and the C++ compacts read every request's move range + from these offsets, so refreshing them retargets the fixed buffers to + the cohort at hand without any reallocation. + """ + if len(tail_lengths) > self.request_count: + raise ValueError("the cohort exceeds the compaction request capacity") + if any(tail < 0 or tail > self.protected_tail_capacity for tail in tail_lengths): + raise ValueError("a protected tail exceeds the configured capacity") + self._write_move_offsets( + self.target_dense_compaction.move_source_offsets, + [self.decode_keep_count + int(tail) for tail in tail_lengths], + ) + if self.target_swa_compaction is not None: + self._write_move_offsets( + self.target_swa_compaction.move_source_offsets, + [self.swa_window + int(tail) for tail in tail_lengths], + ) + if self.draft_compaction is not None: + if draft_tail_lengths is None: + draft_tail_lengths = [0] * len(tail_lengths) + if len(draft_tail_lengths) != len(tail_lengths): + raise ValueError("draft protected tails must match the cohort") + if any( + tail < 0 or tail > self.draft_protected_tail_capacity for tail in draft_tail_lengths + ): + raise ValueError("a draft protected tail exceeds the configured capacity") + self._write_move_offsets( + self.draft_compaction.move_source_offsets, + [self.decode_keep_count + int(tail) for tail in draft_tail_lengths], + ) + def launch(self) -> None: """Pack the move indices, then run every cache family's C++ compacts.""" if self.swa_destination_bases is not None: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d5caf2958dee..a8d31db19799 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -53,7 +53,6 @@ The scoring math follows the same upstream reference (``methods/pruning_utils.py``). """ -from collections import OrderedDict from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union @@ -80,11 +79,6 @@ # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) -# Bound eager staging memory. A large due cohort is processed as consecutive -# request chunks with identical results. -_EAGER_REQUEST_CHUNK_SIZE = 32 -_EAGER_RESOURCE_CACHE_LIMIT = 3 -_EAGER_COMPACTION_CACHE_LIMIT = 6 _INT32_MAX = torch.iinfo(torch.int32).max @@ -245,9 +239,7 @@ def __call__( provisional_indices: torch.Tensor, output_indices: torch.Tensor, ) -> None: - self._prepared_launch( - scores, seq_lens, prompt_offsets, provisional_indices, output_indices - ) + self._prepared_launch(scores, seq_lens, prompt_offsets, provisional_indices, output_indices) class _PreparedUnionScores: @@ -361,9 +353,7 @@ def _deterministic_topk_indices_into( tied = torch.nonzero(valid_scores == threshold, as_tuple=False).flatten() tie_count = keep_count - int(higher.numel()) ordered = torch.sort(torch.cat((higher, tied[:tie_count]))).values - output_indices_i32[row_index, :keep_count].copy_( - ordered.to(torch.int32).add(prompt_len) - ) + output_indices_i32[row_index, :keep_count].copy_(ordered.to(torch.int32).add(prompt_len)) class _CrossRequestSelectionPlan(NamedTuple): @@ -475,12 +465,8 @@ def refresh_row_prompt_offsets(self) -> None: Called after the shared per-request buffer was staged externally. """ if self.row_prompt_offsets is not self.prompt_offsets: - self.row_prompt_offsets.view( - self.max_requests, self.selection_rows_per_request - ).copy_( - self.prompt_offsets.unsqueeze(1).expand( - -1, self.selection_rows_per_request - ) + self.row_prompt_offsets.view(self.max_requests, self.selection_rows_per_request).copy_( + self.prompt_offsets.unsqueeze(1).expand(-1, self.selection_rows_per_request) ) def _allocate_cpu_reference_buffers( @@ -743,9 +729,7 @@ def __init__( if num_query_heads % num_kv_heads: raise ValueError("query heads must be divisible by KV heads") selection_rows = ( - num_kv_heads - if eviction_mode == "per_head" - else len(dense_layers) * num_kv_heads + num_kv_heads if eviction_mode == "per_head" else len(dense_layers) * num_kv_heads ) super().__init__( eviction_mode=eviction_mode, @@ -985,7 +969,7 @@ def __init__( omega: torch.Tensor, page_table_keys: Optional[List[object]] = None, num_page_table_slots: Optional[int] = None, - min_prompt_len: int = 0, + decode_width: int = 0, page_table_token_capacity: Optional[int] = None, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_page_representatives: Optional[List[int]] = None, @@ -1015,11 +999,11 @@ def __init__( if page_table_token_capacity < seq_len: raise ValueError("page-table capacity cannot be smaller than the score bucket") self.page_table_token_capacity = int(page_table_token_capacity) - # The smallest cohort prompt only sizes the widest decode window; - # per-request prompt lengths are staged runtime metadata. - if min_prompt_len < 0 or min_prompt_len > seq_len: - raise ValueError("fixed score metadata prompt length is outside its bucket") - self.min_prompt_len = min_prompt_len + # Decode-width capacity of the score buffers; per-request prompt + # lengths are staged runtime metadata. + if decode_width <= 0 or decode_width > seq_len: + raise ValueError("fixed score decode width is outside its bucket") + self.decode_width = int(decode_width) q_real = q_real.to(device=self.device, dtype=torch.float32).contiguous() q_imag = q_imag.to(device=self.device, dtype=torch.float32).contiguous() mlr_coef = mlr_coef.to(device=self.device, dtype=torch.float32).contiguous() @@ -1171,7 +1155,7 @@ def __init__( freq_scale_sq, omega, offsets, - min_prompt_len=min_prompt_len, + output_width=decode_width, ) self.copy_done = torch.cuda.Event() # First record publishes constructor allocations to the V2 copy stream; @@ -1387,6 +1371,9 @@ def stage( ): return False self.request_metadata_host[:, :request_count].copy_(request_metadata) + # Rows past this cohort are padding: zero lengths keep the score + # kernel and selection inert for them. + self.request_metadata_host[:, request_count:].zero_() try: # Copy the fixed backing once. Only the first ``request_count`` # columns are consumed by this cohort. @@ -1512,7 +1499,7 @@ class _PreparedGenerationBatch: @dataclass(kw_only=True, slots=True) -class _EvictionBucketResources: +class _EvictionBuffers: """Reusable eager score and selection buffers for one runtime shape.""" score_staging: _FixedScoreStagingBuffers @@ -1611,10 +1598,11 @@ def __init__( # exact fixed-linear generation width for that currently in-flight batch; # the final hook treats those slots as an opaque suffix. self._prepared_generation_batch: Optional[_PreparedGenerationBatch] = None - # Eager score/selection buffers are built from the first live cohort and - # reused for subsequent evictions with the same runtime geometry. - self._eviction_buckets = OrderedDict() - self._batched_compactions = OrderedDict() + # Eviction buffers are built once at the first eviction, sized to + # capacity bounds, and reused for the manager's lifetime. + self._eviction_resources: Optional[_EvictionBuffers] = None + self._eviction_pool_fingerprint: Optional[tuple] = None + self._batched_compaction = None self._local_to_global_layers_cache: Optional[List[int]] = None self._attention_layer_partition_cache: Optional[ Tuple[List[int], List[int], Optional[int]] @@ -1902,22 +1890,19 @@ def _periodic_evict( # (2) Compact all affected dense and kernel-masked SWA layers, then release # the unreachable tail directly through V2's public resize primitive. - # Prompt lengths are per-request metadata, so the whole due cohort runs - # as one batched round (chunked only by the staging memory bound). + # Prompt lengths and tails are per-request metadata, so the whole due + # cohort runs as one batched round (the workspace holds max_batch_size + # requests, which bounds any generation batch). if not due_requests: return num_layers = self._num_layers_from_manager() - - for begin in range(0, len(due_requests), _EAGER_REQUEST_CHUNK_SIZE): - chunk = due_requests[begin : begin + _EAGER_REQUEST_CHUNK_SIZE] - chunk_tails = {rid: protected_tails[rid] for _, rid in chunk} - with nvtx_range_debug("triattention.evict_request_group", color="purple"): - capacity_targets = self._evict_requests( - chunk, - num_layers, - protected_tail_lengths=chunk_tails, - ) - self._resize_compacted_requests(capacity_targets, protected_tails) + with nvtx_range_debug("triattention.evict_request_group", color="purple"): + capacity_targets = self._evict_requests( + due_requests, + num_layers, + protected_tail_lengths=protected_tails, + ) + self._resize_compacted_requests(capacity_targets, protected_tails) def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: if not capacity_targets: @@ -2054,9 +2039,9 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: prepared = self._prepared_generation_batch if prepared is not None: prepared.growth_by_request.pop(request_id, None) - if not self._request_states: - self._eviction_buckets.clear() - self._batched_compactions.clear() + # The workspace stays resident across idle periods: its memory is a + # deliberate one-time cost and rebuilding it per burst would reintroduce + # allocation on the decode hot path. # ================================================================== # # Helpers (eviction / scoring / V2 cache access / calibration) # @@ -2368,20 +2353,63 @@ def _eager_resources_for( self, layout: _RuntimeKVLayout, prepared: Sequence[_PreparedEviction], - ) -> _EvictionBucketResources: - """Build or reuse eager score and selection buffers for one cohort.""" + ) -> _EvictionBuffers: + """Return the eviction buffers, building them once at first use. + + The request capacity follows the executor's max batch size (memory + scales linearly with it) and the decode-width capacity follows the + eviction bound (compaction keeps the scored decode region near + ``top_B`` plus one period of growth), so one set of buffers serves + every round. They are rebuilt only when the pool views change or a + round outgrows them (prompt-counting budgets score the prompt, whose + length is workload-defined). + """ if not prepared: raise ValueError("TriAttention eviction requires at least one request") - # Prompt lengths are per-request runtime metadata; the smallest one - # only sizes the widest decode window of the shared buffers. - min_prompt_len = min(item.prompt_len for item in prepared) - seq_len = max(item.seq_len for item in prepared) - page_table_token_capacity = max(item.seq_len + item.protected_tail for item in prepared) - request_count = len(prepared) + needed_width = max(item.seq_len - item.prompt_len for item in prepared) + needed_page_tokens = max(item.seq_len + item.protected_tail for item in prepared) + needed_requests = len(prepared) + draft_fingerprint = None + if self.draft_kv_cache_manager is not None: + draft_layout = self._draft_runtime_kv_layout() + draft_fingerprint = ( + draft_layout.pool_page_counts, + draft_layout.pool_view_fingerprint, + ) + fingerprint = ( + self.eviction_mode, + self.top_B, + tuple(layout.dense_layers), + layout.pool_view_fingerprint, + draft_fingerprint, + ) + resources = self._eviction_resources + if resources is not None: + staging = resources.score_staging + if ( + self._eviction_pool_fingerprint == fingerprint + and needed_width <= staging.decode_width + and needed_page_tokens <= staging.page_table_token_capacity + and needed_requests <= staging.max_requests + ): + return resources + # Pools changed or this round outgrew the buffers: rebuild. + self._eviction_resources = None + self._batched_compaction = None + + mgr = self.kv_cache_manager + tail_capacity = self._configured_protected_tail_capacity() + request_capacity = max(needed_requests, int(mgr.max_batch_size)) + decode_width = max( + needed_width, + self.top_B + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), + ) + seq_capacity = max(needed_page_tokens, int(mgr.max_seq_len)) + page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) + dense_groups = list(layout.storage_groups.values()) representatives = [group[0] for group in dense_groups] representatives.extend(layer for layer in layout.swa_layers if layer not in representatives) - draft_key = None draft_kwargs = {} if self.draft_kv_cache_manager is not None: draft_layout = self._draft_runtime_kv_layout() @@ -2394,29 +2422,8 @@ def _eager_resources_for( draft_layout.layer_pool_keys[layer] for layer in draft_representatives ], draft_num_page_table_slots=self.draft_kv_cache_manager.num_pools, - draft_page_table_token_capacity=seq_len + draft_tail_capacity, + draft_page_table_token_capacity=seq_capacity + draft_tail_capacity, ) - draft_key = ( - draft_tail_capacity, - draft_layout.pool_page_counts, - draft_layout.pool_view_fingerprint, - ) - key = ( - "triattention.eager.v1", - self.eviction_mode, - request_count, - seq_len, - min_prompt_len, - page_table_token_capacity, - self.top_B, - tuple(layout.dense_layers), - layout.pool_view_fingerprint, - draft_key, - ) - resources = self._eviction_buckets.get(key) - if resources is not None: - self._eviction_buckets.move_to_end(key) - return resources first_pool = layout.layer_pools[layout.dense_layers[0]] if self._offsets is None: @@ -2429,8 +2436,8 @@ def _eager_resources_for( dense_groups=dense_groups, dense_layers=layout.dense_layers, page_representatives=representatives, - max_requests=request_count, - seq_len=seq_len, + max_requests=request_capacity, + seq_len=seq_capacity, num_q_heads=int(self._H), num_freqs=int(self._F), q_real=q_real, @@ -2441,7 +2448,7 @@ def _eager_resources_for( omega=self.calibration["omega"], page_table_keys=self._page_table_pool_keys(representatives, layout.global_layers), num_page_table_slots=layout.manager.num_pools, - min_prompt_len=min_prompt_len, + decode_width=decode_width, page_table_token_capacity=page_table_token_capacity, **draft_kwargs, ) @@ -2452,35 +2459,36 @@ def _eager_resources_for( num_query_heads=int(self._H), num_kv_heads=int(first_pool.shape[2]), rows=len(layout.dense_layers) * int(self._H), - width=seq_len - min_prompt_len, + width=decode_width, keep_count=self.top_B, dtype=torch.float32, device=first_pool.device, - max_requests=request_count, + max_requests=request_capacity, ), input_scores=score_staging.fused_group.output.view( - request_count, + request_capacity, len(layout.dense_layers) * int(self._H), - seq_len - min_prompt_len, + decode_width, ), normalize_scores=self.normalize_scores, prompt_offsets_buffer=score_staging.token_starts_device, ) + # Padded rows carry zero valid width; their provisional TopK entries + # must still be in-range ordinals for the finalizer's score gather. + provisional = getattr(keep_set_selector, "final_indices", None) + if provisional is None: + provisional = keep_set_selector.top_indices_i32 + provisional.zero_() score_staging.bind_score_launcher( keep_set_selector.valid_widths, self.score_aggregation, ) - resources = _EvictionBucketResources( + resources = _EvictionBuffers( score_staging=score_staging, keep_set_selector=keep_set_selector, ) - self._eviction_buckets[key] = resources - while len(self._eviction_buckets) > _EAGER_RESOURCE_CACHE_LIMIT: - _, stale = self._eviction_buckets.popitem(last=False) - stale_ids = (id(stale.score_staging), id(stale.keep_set_selector)) - for compaction_key in tuple(self._batched_compactions): - if compaction_key[:2] == stale_ids: - del self._batched_compactions[compaction_key] + self._eviction_resources = resources + self._eviction_pool_fingerprint = fingerprint return resources def _batched_compaction_for( @@ -2494,21 +2502,6 @@ def _batched_compaction_for( """Build or reuse the eager C++ compaction launches for one cohort.""" from .compaction import BatchedKVCacheCompaction - protected_tail_lengths = tuple(item.protected_tail for item in prepared) - draft_tail_lengths: Optional[Tuple[int, ...]] = None - draft_kwargs = {} - if self.draft_kv_cache_manager is not None: - draft_layout = self._draft_runtime_kv_layout() - draft_tail_lengths = (self._draft_protected_tail_capacity(),) * len(prepared) - draft_kwargs = dict( - draft_layer_pools=draft_layout.layer_pools, - draft_layers=draft_layout.dense_layers, - draft_layer_group_representative=draft_layout.layer_group_representative, - draft_layer_pool_keys=list(draft_layout.layer_pool_keys), - draft_protected_tail_lengths=list(draft_tail_lengths), - draft_kv_block_offsets=score_staging.draft_block_offsets_device, - draft_page_table_slots=score_staging.draft_representative_slots, - ) if layout.swa_layers and layout.swa_window: # SWA landing positions are prompt-dependent; reject a request # whose retained span cannot cover the model window this round. @@ -2519,16 +2512,20 @@ def _batched_compaction_for( f"{item.prompt_len + self.top_B} tokens, below the " f"sliding window {layout.swa_window}" ) - key = ( - id(score_staging), - id(keep_set_selector), - protected_tail_lengths, - draft_tail_lengths, - ) - batched_compaction = self._batched_compactions.get(key) - if batched_compaction is not None: - self._batched_compactions.move_to_end(key) - else: + batched_compaction = self._batched_compaction + if batched_compaction is None: + draft_kwargs = {} + if self.draft_kv_cache_manager is not None: + draft_layout = self._draft_runtime_kv_layout() + draft_kwargs = dict( + draft_layer_pools=draft_layout.layer_pools, + draft_layers=draft_layout.dense_layers, + draft_layer_group_representative=draft_layout.layer_group_representative, + draft_layer_pool_keys=list(draft_layout.layer_pool_keys), + draft_protected_tail_capacity=self._draft_protected_tail_capacity(), + draft_kv_block_offsets=score_staging.draft_block_offsets_device, + draft_page_table_slots=score_staging.draft_representative_slots, + ) batched_compaction = BatchedKVCacheCompaction( eviction_mode=self.eviction_mode, layer_pools=layout.layer_pools, @@ -2536,20 +2533,26 @@ def _batched_compaction_for( swa_layers=layout.swa_layers, layer_group_representative=layout.layer_group_representative, layer_pool_keys=list(layout.layer_pool_keys), - kept_token_ordinals=keep_set_selector.keep[: len(prepared)], - valid_sequence_lengths=score_staging.valid_seq_lens_device[: len(prepared)], + kept_token_ordinals=keep_set_selector.keep, + valid_sequence_lengths=score_staging.valid_seq_lens_device, kv_block_offsets=score_staging.block_offsets_device, page_table_slots=score_staging.representative_slots, - request_count=len(prepared), - prompt_offsets=score_staging.token_starts_device[: len(prepared)], + request_count=score_staging.max_requests, + prompt_offsets=score_staging.token_starts_device, decode_keep_count=self.top_B, swa_window=layout.swa_window, - protected_tail_lengths=list(protected_tail_lengths), + protected_tail_capacity=self._configured_protected_tail_capacity(), **draft_kwargs, ) - self._batched_compactions[key] = batched_compaction - while len(self._batched_compactions) > _EAGER_COMPACTION_CACHE_LIMIT: - self._batched_compactions.popitem(last=False) + self._batched_compaction = batched_compaction + # Tails vary per round (in-flight growth); padded rows move nothing. + draft_tails = None + if self.draft_kv_cache_manager is not None: + draft_tails = [self._draft_protected_tail_capacity()] * len(prepared) + batched_compaction.set_protected_tails( + [item.protected_tail for item in prepared], + draft_tails, + ) return batched_compaction def _page_table_pool_keys( diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 729479d720d2..4d84254ea53e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -763,7 +763,7 @@ def _launch_tri_score_perhead( class _FixedScoreGroup: - """Persistent score metadata/output for one sequence bucket. + """Persistent score metadata/output for one fixed geometry. Since the per-layer absolute-address ABI, ONE group can span dense layers living in DISTINCT storages with DISTINCT block tables. ``block_offsets`` @@ -787,18 +787,18 @@ def __init__( freq_scale_sq: torch.Tensor, omega: torch.Tensor, offsets: torch.Tensor, - min_prompt_len: int = 0, + output_width: int = 0, ) -> None: if not layer_indices or min(max_requests, page_count, seq_len) <= 0: raise ValueError("fixed score group requires non-empty positive geometry") - if min_prompt_len < 0 or min_prompt_len >= seq_len: - raise ValueError("fixed score prompt length must leave a non-empty decode region") + if output_width <= 0 or output_width > seq_len: + raise ValueError("fixed score group requires a decode width within its capacity") if len(page_table_slots) != len(layer_indices): raise ValueError("page_table_slots must align with layer_indices") self.max_requests = max_requests - # Prompt lengths are per-request kernel inputs; the smallest one only + # Prompt lengths are per-request kernel inputs; this capacity only # sizes the widest possible decode window of the output buffer. - self.output_width = seq_len - min_prompt_len + self.output_width = int(output_width) self.num_layers = len(layer_indices) p0 = layer_pools[layer_indices[0]] if p0.ndim != 5: 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 index 7f8afdca0034..0b393653f9b5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -150,15 +150,16 @@ def _launched_draft_compaction(draft_protected_tails): prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, swa_window=None, - protected_tail_lengths=target_protected_tails, + protected_tail_capacity=max(target_protected_tails), draft_layer_pools=[draft_pool], draft_layers=[0], draft_layer_group_representative={0: 0}, draft_layer_pool_keys=[("draft_pool", 0)], - draft_protected_tail_lengths=draft_protected_tails, + draft_protected_tail_capacity=max(draft_protected_tails), draft_kv_block_offsets=_encode_block_offsets(draft_tables), draft_page_table_slots={0: 0}, ) + compaction.set_protected_tails(target_protected_tails, draft_protected_tails) compaction.launch() torch.cuda.synchronize(device) @@ -428,7 +429,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): ) -def test_lru_score_staging_eviction_drops_dependent_batched_compactions(): +def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention(top_B=4) @@ -463,28 +464,18 @@ def test_lru_score_staging_eviction_drops_dependent_batched_compactions(): pool_view_fingerprint=(("fixed",),), ) - stale = [SimpleNamespace(score_staging=object(), keep_set_selector=object()) for _ in range(3)] - for index, resources in enumerate(stale): - manager._eviction_buckets[("stale", index)] = resources - oldest_ids = (id(stale[0].score_staging), id(stale[0].keep_set_selector)) - dependent_draft_key = (*oldest_ids, (0,), (1,)) - dependent_plain_key = (*oldest_ids, (2,), None) - surviving_key = ( - id(stale[1].score_staging), - id(stale[1].keep_set_selector), - (0,), - (1,), - ) - manager._batched_compactions[dependent_draft_key] = object() - manager._batched_compactions[dependent_plain_key] = object() - manager._batched_compactions[surviving_key] = object() - score_staging = SimpleNamespace( - fused_group=SimpleNamespace(output=torch.empty(1, 4, 8)), + fused_group=SimpleNamespace(output=torch.empty(8, 4, 260)), bind_score_launcher=mock.Mock(), - token_starts_device=torch.zeros(1, dtype=torch.int32), + token_starts_device=torch.zeros(8, dtype=torch.int32), + decode_width=260, + page_table_token_capacity=65537, + max_requests=8, + ) + keep_set_selector = SimpleNamespace( + valid_widths=torch.empty(8, dtype=torch.int32), + top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), ) - keep_set_selector = SimpleNamespace(valid_widths=torch.empty(1, dtype=torch.int32)) prepared = [ _PreparedEviction( request=_make_request(7), @@ -511,14 +502,26 @@ def test_lru_score_staging_eviction_drops_dependent_batched_compactions(): ): resources = manager._eager_resources_for(layout, prepared) - assert resources.score_staging is score_staging - # The new bucket carries a draft page-table plane sized for the draft tail. - assert score_cls.call_args.kwargs["draft_page_table_token_capacity"] == 8 + 1 - # The oldest score staging fell out of the LRU, taking every dependent - # compaction staging (draft-carrying included) with it. - assert ("stale", 0) not in manager._eviction_buckets - assert dependent_draft_key not in manager._batched_compactions - assert dependent_plain_key not in manager._batched_compactions - assert surviving_key in manager._batched_compactions - assert ("stale", 1) in manager._eviction_buckets - assert ("stale", 2) in manager._eviction_buckets + # The buffers follow the executor limits, not this one-request cohort. + assert resources.score_staging is score_staging + assert score_cls.call_args.kwargs["max_requests"] == 8 + assert score_cls.call_args.kwargs["decode_width"] == 4 + 2 * 128 + assert score_cls.call_args.kwargs["seq_len"] == 65536 + assert score_cls.call_args.kwargs["page_table_token_capacity"] == 65536 + 1 + assert score_cls.call_args.kwargs["draft_page_table_token_capacity"] == 65536 + 1 + + # A second round with unchanged pools reuses the resident buffers and + # keeps the cached compaction launches. + cached_compaction = object() + manager._batched_compaction = cached_compaction + assert manager._eager_resources_for(layout, prepared) is resources + assert score_cls.call_count == 1 + assert manager._batched_compaction is cached_compaction + + # A pool change invalidates both the buffers and the compaction + # launches that alias them. + layout.pool_view_fingerprint = (("moved",),) + rebuilt = manager._eager_resources_for(layout, prepared) + assert rebuilt is not resources + assert score_cls.call_count == 2 + assert manager._batched_compaction is None diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index de5626468567..438b1bf68cd3 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -532,8 +532,9 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, swa_window=None, - protected_tail_lengths=protected_tails, + protected_tail_capacity=max(protected_tails), ) + compaction.set_protected_tails(protected_tails) compaction.launch() torch.cuda.synchronize(device) @@ -654,8 +655,9 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, swa_window=None, - protected_tail_lengths=protected_tails, + protected_tail_capacity=max(protected_tails), ) + cohort_compaction.set_protected_tails(protected_tails) cohort_compaction.launch() expected_pools = [pool.clone() for pool in initial_pools] @@ -675,8 +677,9 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): prompt_offsets=torch.tensor([prompt_lens[request]], dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, swa_window=None, - protected_tail_lengths=[protected_tails[request]], + protected_tail_capacity=protected_tails[request], ) + single_compaction.set_protected_tails([protected_tails[request]]) single_compaction.launch() torch.cuda.synchronize(device) @@ -795,8 +798,9 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): prompt_offsets=torch.zeros(1, dtype=torch.int32, device=device), decode_keep_count=keep_count, swa_window=None, - protected_tail_lengths=[0], + protected_tail_capacity=0, ) + batched_compaction.set_protected_tails([0]) batched_compaction.launch() torch.cuda.synchronize(device) @@ -952,8 +956,9 @@ def write_token(token: int, score: float) -> None: prompt_offsets=score_staging.token_starts_device[:1], decode_keep_count=compacted_capacity - prompt_len - protected_tail, swa_window=None, - protected_tail_lengths=[protected_tail], + protected_tail_capacity=protected_tail, ) + batched_compaction.set_protected_tails([protected_tail]) def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) @@ -1070,8 +1075,9 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), decode_keep_count=4, swa_window=2, - protected_tail_lengths=protected_tails, + protected_tail_capacity=max(protected_tails), ) + compaction.set_protected_tails(protected_tails) compaction.launch() torch.cuda.synchronize(device) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 9f6500562b6f..4c8eac9d39b1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -682,7 +682,7 @@ def test_non_boundary_step_skips_eviction_geometry(self): assert state.generation_steps == 127 assert state.confirmed_kv_length == 1024 + 4096 + 1 - def test_eager_eviction_chunks_large_due_cohort(self): + def test_eager_eviction_runs_large_due_cohort_in_one_round(self): manager, _, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) requests = [] caches = {} @@ -704,20 +704,24 @@ def test_eager_eviction_chunks_large_due_cohort(self): ): manager._periodic_evict(batch) - assert [len(call.args[0]) for call in evict.call_args_list] == [32, 32, 1] - assert resize.call_count == 3 + assert [len(call.args[0]) for call in evict.call_args_list] == [65] + assert resize.call_count == 1 - def test_last_request_finish_releases_eager_buffers(self): + def test_request_finish_keeps_eviction_buffers_resident(self): manager = _make_triattention() request = _make_request(7) _set_request_state(manager, 7) - manager._eviction_buckets[("score",)] = object() - manager._batched_compactions[("compact",)] = object() + buffers = object() + compaction = object() + manager._eviction_resources = buffers + manager._batched_compaction = compaction manager.on_request_finish(request) - assert manager._eviction_buckets == {} - assert manager._batched_compactions == {} + # The buffers are sized for the executor limits, not one cohort, so + # they stay resident for the next generation batch. + assert manager._eviction_resources is buffers + assert manager._batched_compaction is compaction @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): @@ -1054,7 +1058,7 @@ def test_cross_request_union_uses_cute_without_fallback(self, keep_count): class TestFixedScoreMetadata: @pytest.mark.parametrize("normalize_scores", [False, True]) @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) - def test_eager_bucket_binds_score_after_selection(self, eviction_mode, normalize_scores): + def test_eager_buffers_bind_score_after_selection(self, eviction_mode, normalize_scores): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention( @@ -1080,12 +1084,17 @@ def test_eager_bucket_binds_score_after_selection(self, eviction_mode, normalize storage_groups={0: [0, 1]}, pool_view_fingerprint=(("fixed",),), ) + # The buffers follow the executor limits: eight requests (max batch + # size) by 260 decode tokens (top_B plus two eviction periods). score_staging = SimpleNamespace( - fused_group=SimpleNamespace(output=torch.empty(1, 4, 8)), + fused_group=SimpleNamespace(output=torch.empty(8, 4, 260)), bind_score_launcher=mock.Mock(), - token_starts_device=torch.zeros(1, dtype=torch.int32), + token_starts_device=torch.zeros(8, dtype=torch.int32), + ) + keep_set_selector = SimpleNamespace( + valid_widths=torch.empty(8, dtype=torch.int32), + top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), ) - keep_set_selector = SimpleNamespace(valid_widths=torch.empty(1, dtype=torch.int32)) prepared = [ _prepared_eviction( _make_request(7), From dc1107e92af7f3ce9e4d3be63b450b4836381ca7 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 23:04:08 -0700 Subject: [PATCH 025/178] [None][chore] Empty the TriAttention package __init__ like the framework package The framework kv_cache_compression package __init__ was emptied on review (no re-exports); apply the same rule to the triattention subpackage. The factory and the tests import the manager module directly. Signed-off-by: tianruih --- .../triattention/__init__.py | 21 ------------------- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- .../test_triattention_draft_cocompaction.py | 2 +- .../test_triattention_pipeline.py | 2 +- 4 files changed, 3 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py index d40381e114cc..e69de29bb2d1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/__init__.py @@ -1,21 +0,0 @@ -"""TriAttention KV-cache compression: periodic physical KV eviction driven by -trigonometric importance scoring. - -TriAttention is a pure KV-cache compression method. Decode still runs the model's -standard attention over the compacted cache. The manager publishes each request's -cumulative evicted count on ``LlmRequest.py_num_compressed_tokens``; the model -engine subtracts it when building the cached-token metadata. With one-model -speculative decoding, the separate draft KV cache is compacted in the same round -with the target's kept token set (union mode only), so target and draft always -share one physical KV length. - -Public surface: - - ``TriAttention`` -- the ``BaseKVCacheCompressionManager`` (the eviction - manager; snapshots allocation metadata before forward and compacts the - finalized prefix in ``on_generation_step_end``). It uses V2 capacity-only - decode, so there is no KV-cache-manager subclass. -""" - -from .triattention import TriAttention - -__all__ = ["TriAttention"] diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ce185c847bc2..fa6dfefa2853 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2124,7 +2124,7 @@ def create_kv_cache_compression_manager( ``validate_kv_cache_compression_with_spec``. """ if config.algorithm == "triattention": - from tensorrt_llm._torch.kv_cache_compression.triattention import \ + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import \ TriAttention return TriAttention( 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 index 0b393653f9b5..872229a970da 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -21,11 +21,11 @@ import torch from conftest import encode_block_offsets as _encode_block_offsets -from tensorrt_llm._torch.kv_cache_compression.triattention import TriAttention from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( BatchedKVCacheCompaction, ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttention, _FixedScoreStagingBuffers, _PreparedEviction, _RequestCompressionState, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 4c8eac9d39b1..8e60176015b5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -35,8 +35,8 @@ # 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 import TriAttention from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttention, _BatchedUnionKeepSetSelector, _PreparedEviction, _PreparedGenerationBatch, From 0097566bd62be830bcbb29fb4782de9564020cda Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 23:22:03 -0700 Subject: [PATCH 026/178] [None][fix] Pass the move-index head stride to the compact kernel explicitly The sparse compact kernel took the head-plane stride of the packed move indices from the last source-offsets entry (the round total move count). The move buffers are now allocated once at their widest size and only the offsets are refreshed per round, so the buffer width and the round total diverge: the packing kernel writes head planes at the allocation width while the compact kernel read them at the round total, corrupting every head plane past the first (illegal memory access in the co-compaction tests). Pass the stride from the host, where the tensor shape is known. Signed-off-by: tianruih --- .../kernels/unfusedAttentionKernels.h | 7 +++-- .../unfusedAttentionKernels_2_template.h | 29 ++++++++++++------- .../thop/sparseKvCacheCompactOp.cpp | 18 +++++++----- 3 files changed, 35 insertions(+), 19 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index b51d8e8dc036..b2068538b1eb 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -418,11 +418,14 @@ void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams //! each request and head, source ordinals must increase strictly and satisfy //! destinationBases[request] + move <= source[move]; the per-request bases //! let one launch cover a cohort with mixed pinned-prompt lengths. +//! sourceHeadStride is the head-plane stride of sparseKvIndices: the index +//! buffers may be wider than one round's total move count. template void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, - int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t const* destinationBases, int32_t batchSize, - int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream); + int64_t sourceLayerStride, int64_t sourceHeadStride, int32_t const* sparseKvOffsets, + int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, + cudaStream_t stream); // Debug function to test basic parameter access template diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index fbf04669d702..eccc4e3cf984 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1763,6 +1763,10 @@ struct KvCacheV2LayersBuffer int32_t const* pageTable; int32_t const* sourceLayerIndices; int64_t sourceLayerStride; + // Head-plane stride of the packed move-source indices. The move buffers + // may be wider than one round's move count, so the stride comes from the + // host-side allocation width, not from the last source-offsets entry. + int64_t sourceHeadStride; int64_t pageTableRequestStride; int32_t tokensPerBlock; int32_t const* destinationBases; @@ -1798,12 +1802,12 @@ struct KvCacheV2LayersBuffer } __device__ __forceinline__ int32_t getSparseKvSourceToken( - int32_t const* sourceIndices, int32_t headIdx, int32_t totalMoves, int32_t globalMove) const + int32_t const* sourceIndices, int32_t headIdx, int32_t globalMove) const { int32_t const layer = sourceLayerIndices == nullptr ? static_cast(blockIdx.x) : sourceLayerIndices[blockIdx.x]; - int64_t const offset - = static_cast(layer) * sourceLayerStride + static_cast(headIdx) * totalMoves + globalMove; + int64_t const offset = static_cast(layer) * sourceLayerStride + + static_cast(headIdx) * sourceHeadStride + globalMove; return sourceIndices[offset]; } @@ -1828,7 +1832,10 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( int const batch_idx = blockIdx.z; int const kv_head_idx = blockIdx.y; - int const total_num_sparse_kv_tokens = params.sparse_kv_offsets[params.batch_size]; + // Head-plane stride of the packed indices for the non-layered layout; the + // layered layout carries its stride on the buffer (its index buffers may + // be wider than one round's move count). + [[maybe_unused]] int const total_num_sparse_kv_tokens = params.sparse_kv_offsets[params.batch_size]; int const sparse_start_idx = params.sparse_kv_offsets[batch_idx]; int const sparse_end_idx = params.sparse_kv_offsets[batch_idx + 1]; @@ -1854,7 +1861,7 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( if constexpr (Layered) { src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( - params.sparse_kv_indices, kv_head_idx, total_num_sparse_kv_tokens, global_sparse_idx); + params.sparse_kv_indices, kv_head_idx, global_sparse_idx); } else { @@ -1900,7 +1907,7 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( if constexpr (Layered) { src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( - params.sparse_kv_indices, kv_head_idx, total_num_sparse_kv_tokens, global_sparse_idx); + params.sparse_kv_indices, kv_head_idx, global_sparse_idx); dst_token_idx = params.kv_cache_buffer.getSparseKvDestinationToken(batch_idx, sparse_token_offset); } else @@ -2010,14 +2017,16 @@ void launchSparseKvCacheCompactV2Layers( template void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, - int64_t sourceLayerStride, int32_t const* sparseKvOffsets, int32_t const* destinationBases, int32_t batchSize, - int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream) + int64_t sourceLayerStride, int64_t sourceHeadStride, int32_t const* sparseKvOffsets, + int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, + cudaStream_t stream) { KvCacheV2LayersBuffer buffer{}; buffer.poolPointers = poolPointers; buffer.pageTable = pageTable; buffer.sourceLayerIndices = sourceLayerIndices; buffer.sourceLayerStride = sourceLayerStride; + buffer.sourceHeadStride = sourceHeadStride; buffer.pageTableRequestStride = pageTableRequestStride; buffer.tokensPerBlock = tokensPerBlock; buffer.destinationBases = destinationBases; @@ -2060,8 +2069,8 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con #define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(T) \ template void invokeSparseKvCacheCompactV2Layers(int64_t const*, int32_t const*, int32_t, int64_t, \ - int32_t const*, int32_t const*, int64_t, int32_t const*, int32_t const*, int32_t, int32_t, int32_t, int32_t, \ - cudaStream_t); + int32_t const*, int32_t const*, int64_t, int64_t, int32_t const*, int32_t const*, int32_t, int32_t, int32_t, \ + int32_t, cudaStream_t); } // namespace kernels diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index c01a396b21bc..c8fac21f431d 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -113,9 +113,10 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor sourceLayerPtr = layerIndices.data_ptr(); } - // The last element of source_offsets is the total move count and must equal - // source_indices.size(-1); it lives on device, so checking it here would - // force a sync -- the kernel trusts the caller. + // source_offsets carve each request's move range out of source_indices; + // the values live on device and the kernel trusts them. The index buffer + // may be wider than one round's total move count, so the head-plane + // stride passed below comes from the tensor shape, not from the offsets. TORCH_CHECK(sourceOffsets.is_cuda() && sourceOffsets.get_device() == device && sourceOffsets.scalar_type() == th::kInt32 && sourceOffsets.is_contiguous() && sourceOffsets.dim() == 1 && sourceOffsets.size(0) == batchSize + 1, @@ -129,24 +130,27 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor auto const stream = at::cuda::getCurrentCUDAStream(device); auto const* bases = destinationBases.data_ptr(); + auto const sourceHeadStride = sourceIndices.size(-1); if (dtype == th::kBFloat16) { tk::invokeSparseKvCacheCompactV2Layers<__nv_bfloat16>(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), - sourceLayerPtr, sourceLayerStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, - tokensPerBlock, headDim, stream); + sourceLayerPtr, sourceLayerStride, sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, + numKvHeads, tokensPerBlock, headDim, stream); } else if (dtype == th::kHalf) { tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, - sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); + sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, + stream); } else if (dtype == th::kFloat) { tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, - sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); + sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, + stream); } else { From ef5dd55148fc7ad4cc062ea19752138f2cb23a0c Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 23:28:22 -0700 Subject: [PATCH 027/178] [None][chore] Name the compaction entry points for what they physically do launch() read like CUDA-graph vocabulary; the methods run the physical KV compaction, so call them compact() at every level (C++ launch group, cache family, batched compaction). Signed-off-by: tianruih --- .../triattention/compaction.py | 24 ++++++++++--------- .../triattention/triattention.py | 2 +- .../test_triattention_draft_cocompaction.py | 4 ++-- .../test_triattention_eager.py | 12 +++++----- .../test_triattention_pipeline.py | 4 ++-- 5 files changed, 24 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index aa28b4f3c87c..d0114c45824a 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -45,7 +45,7 @@ class _CppCompactGroup(NamedTuple): pool_pointers: torch.Tensor source_layer_indices: Optional[torch.Tensor] - def launch( + def compact( self, source: torch.Tensor, offsets: torch.Tensor, destination_bases: torch.Tensor ) -> None: torch.ops.trtllm.sparse_kv_cache_compact_layers( @@ -64,23 +64,25 @@ class _SingleCacheCompaction(NamedTuple): Holds the prepared launch that packs this family's move indices (None when an earlier family's pack launch fills them in the same call), the - C++ launch groups that consume them, and the destination base the moved + C++ compact groups that consume them, and the destination base the moved tokens land at. """ prepared_move_index_pack: Optional[PreparedTritonKernelLaunch] - cpp_launch_groups: Tuple[_CppCompactGroup, ...] + cpp_compact_groups: Tuple[_CppCompactGroup, ...] move_source_indices: torch.Tensor move_source_offsets: torch.Tensor # Per-request landing positions; may alias staged prompt lengths so the # values track the current round without a refresh. destination_bases: torch.Tensor - def launch(self) -> None: + def compact(self) -> None: if self.prepared_move_index_pack is not None: self.prepared_move_index_pack() - for group in self.cpp_launch_groups: - group.launch(self.move_source_indices, self.move_source_offsets, self.destination_bases) + for group in self.cpp_compact_groups: + group.compact( + self.move_source_indices, self.move_source_offsets, self.destination_bases + ) def _cuda_int32_contiguous(tensors: Tuple[torch.Tensor, ...], device: torch.device) -> bool: @@ -492,7 +494,7 @@ def __init__( ) self.target_dense_compaction = _SingleCacheCompaction( prepared_move_index_pack=dense_pack, - cpp_launch_groups=_compact_groups( + cpp_compact_groups=_compact_groups( dense_entries, self.layer_pool_keys, self.device, dense_slots ), move_source_indices=dense_move_indices, @@ -504,7 +506,7 @@ def __init__( if self.swa_layers: self.target_swa_compaction = _SingleCacheCompaction( prepared_move_index_pack=None, - cpp_launch_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), + cpp_compact_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), move_source_indices=swa_move_indices, move_source_offsets=swa_move_offsets, destination_bases=self.swa_destination_bases, @@ -614,7 +616,7 @@ def _build_draft_compaction( ) return _SingleCacheCompaction( prepared_move_index_pack=draft_pack, - cpp_launch_groups=_compact_groups( + cpp_compact_groups=_compact_groups( draft_entries, tuple(draft_layer_pool_keys), self.device ), move_source_indices=draft_move_indices, @@ -668,7 +670,7 @@ def set_protected_tails( [self.decode_keep_count + int(tail) for tail in draft_tail_lengths], ) - def launch(self) -> None: + def compact(self) -> None: """Pack the move indices, then run every cache family's C++ compacts.""" if self.swa_destination_bases is not None: # The prompt offsets may have been re-staged since construction; @@ -679,4 +681,4 @@ def launch(self) -> None: out=self.swa_destination_bases, ) for compaction in self.cache_compactions: - compaction.launch() + compaction.compact() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index a8d31db19799..83054ab97bed 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -2689,7 +2689,7 @@ def _evict_requests( normalize_scores=self.normalize_scores, ) with nvtx_range("triattention.compact", color="purple"): - batched_compaction.launch() + batched_compaction.compact() finally: consumer_streams = [self.kv_cache_manager._stream] if self.draft_kv_cache_manager is not None: 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 index 872229a970da..f33dba37947e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -160,7 +160,7 @@ def _launched_draft_compaction(draft_protected_tails): draft_page_table_slots={0: 0}, ) compaction.set_protected_tails(target_protected_tails, draft_protected_tails) - compaction.launch() + compaction.compact() torch.cuda.synchronize(device) return SimpleNamespace( @@ -354,7 +354,7 @@ def _mocked_eviction_internals(manager): score_staging=score_staging, keep_set_selector=keep_set_selector, ) - batched_compaction = SimpleNamespace(launch=mock.Mock()) + batched_compaction = SimpleNamespace(compact=mock.Mock()) with ( mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), mock.patch.object(manager, "_eager_resources_for", return_value=resources), diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index 438b1bf68cd3..c4b55a5b2078 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -535,7 +535,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) protected_tail_capacity=max(protected_tails), ) compaction.set_protected_tails(protected_tails) - compaction.launch() + compaction.compact() torch.cuda.synchronize(device) for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): @@ -658,7 +658,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): protected_tail_capacity=max(protected_tails), ) cohort_compaction.set_protected_tails(protected_tails) - cohort_compaction.launch() + cohort_compaction.compact() expected_pools = [pool.clone() for pool in initial_pools] for request in range(request_count): @@ -680,7 +680,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): protected_tail_capacity=protected_tails[request], ) single_compaction.set_protected_tails([protected_tails[request]]) - single_compaction.launch() + single_compaction.compact() torch.cuda.synchronize(device) # The two requests own disjoint pages, so whole-pool equality proves the @@ -801,7 +801,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): protected_tail_capacity=0, ) batched_compaction.set_protected_tails([0]) - batched_compaction.launch() + batched_compaction.compact() torch.cuda.synchronize(device) for layer, (before_pool, after_pool, table) in enumerate( @@ -974,7 +974,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: score_staging.launch_prepared_score() keep_set_selector.select_prepared_requests() selected = keep_set_selector.keep[0].clone().to(torch.long) - batched_compaction.launch() + batched_compaction.compact() score_staging.mark_page_tables_consumed(manager._stream) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) @@ -1078,7 +1078,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): protected_tail_capacity=max(protected_tails), ) compaction.set_protected_tails(protected_tails) - compaction.launch() + compaction.compact() torch.cuda.synchronize(device) for request, (valid_seq_len, tail_length) in enumerate( diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 8e60176015b5..82ee6503a1aa 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -491,7 +491,7 @@ def _mocked_eviction_internals(self, manager): score_staging=score_staging, keep_set_selector=keep_set_selector, ) - batched_compaction = SimpleNamespace(launch=mock.Mock()) + batched_compaction = SimpleNamespace(compact=mock.Mock()) with ( mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), mock.patch.object(manager, "_eager_resources_for", return_value=resources), @@ -525,7 +525,7 @@ def test_eviction_bookkeeping_publishes_cumulative_count(self): # 10 confirmed - (2 pinned prompt + 4 decode budget) = 4 evicted. assert request.py_num_compressed_tokens == 4 assert manager._request_states[7].confirmed_kv_length == 6 - internals.batched_compaction.launch.assert_called_once_with() + internals.batched_compaction.compact.assert_called_once_with() internals.score_staging.mark_page_tables_consumed.assert_called_once_with( manager.kv_cache_manager._stream ) From 9cd3fdfce0ed0cca1ae0923c8f1d42c75eca1d97 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 23:40:37 -0700 Subject: [PATCH 028/178] [None][chore] Fold the frozen Triton kernel call into the kernels module PreparedTritonKernelLaunch read like CUDA-graph vocabulary and lived in its own file only so both the score and compaction modules could import it without pulling triton at module scope. Name it FrozenTritonKernelCall for what it does (freeze grid, tensors, and constexprs once; skip Triton dispatch on every later call) and house it with the kernels it launches. The manager module is only imported by the factory when compression is enabled, so the module-scope triton import is paid there. Signed-off-by: tianruih --- .../triattention/compaction.py | 38 ++++----- .../triattention/prepared_launch.py | 82 ------------------- .../triattention/triattention.py | 22 ++--- .../triattention/triattention_kernels.py | 64 ++++++++++++++- 4 files changed, 93 insertions(+), 113 deletions(-) delete mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index d0114c45824a..a3454404326e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -17,7 +17,7 @@ Given each request's kept-token ordinals, its valid sequence length, and the staged V2 block offsets, this module packs per-request move indices with one -prepared Triton launch and then moves the surviving KV in place with batched +frozen Triton kernel call and then moves the surviving KV in place with batched C++ compact launches. Inputs are plain tensors, so any eviction method that produces a kept-token set per request can drive it. """ @@ -27,7 +27,7 @@ import torch -from .prepared_launch import PreparedTritonKernelLaunch +from .triattention_kernels import FrozenTritonKernelCall _SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) @@ -62,13 +62,13 @@ def compact( class _SingleCacheCompaction(NamedTuple): """One compacted cache family (target dense, target SWA, or draft). - Holds the prepared launch that packs this family's move indices (None - when an earlier family's pack launch fills them in the same call), the + Holds the frozen kernel call that packs this family's move indices (None + when an earlier family's pack call fills them in the same run), the C++ compact groups that consume them, and the destination base the moved tokens land at. """ - prepared_move_index_pack: Optional[PreparedTritonKernelLaunch] + frozen_move_index_pack: Optional[FrozenTritonKernelCall] cpp_compact_groups: Tuple[_CppCompactGroup, ...] move_source_indices: torch.Tensor move_source_offsets: torch.Tensor @@ -77,8 +77,8 @@ class _SingleCacheCompaction(NamedTuple): destination_bases: torch.Tensor def compact(self) -> None: - if self.prepared_move_index_pack is not None: - self.prepared_move_index_pack() + if self.frozen_move_index_pack is not None: + self.frozen_move_index_pack() for group in self.cpp_compact_groups: group.compact( self.move_source_indices, self.move_source_offsets, self.destination_bases @@ -221,7 +221,7 @@ def _compact_groups( return tuple(result) -def _prepared_move_index_pack_launch( +def _frozen_move_index_pack( kept_token_ordinals: torch.Tensor, valid_sequence_lengths: torch.Tensor, move_source_offsets: torch.Tensor, @@ -235,8 +235,8 @@ def _prepared_move_index_pack_launch( swa_window: int, swa_move_source_offsets: Optional[torch.Tensor], swa_move_source_indices: Optional[torch.Tensor], -) -> PreparedTritonKernelLaunch: - """Build one prepared launch of the move-index packing kernel. +) -> FrozenTritonKernelCall: + """Build one frozen call of the move-index packing kernel. The kernel reads the kept-token ordinals and each request's valid length and writes the packed per-(layer, head) move source indices consumed by @@ -302,7 +302,7 @@ def _prepared_move_index_pack_launch( swa_indices_arg, ) # Ordered to match the kernel's constexpr parameter declaration: the - # prepared launch replays these by position. + # frozen call passes these by position. constexpr_values = dict( DENSE_TOTAL=int(move_source_indices.shape[-1]), SWA_TOTAL=swa_total, @@ -316,7 +316,7 @@ def _prepared_move_index_pack_launch( HAS_SWA=swa_total > 0, BLOCK=_PACK_BLOCK_TOKENS, ) - return PreparedTritonKernelLaunch( + return FrozenTritonKernelCall( _pack_compaction_sources_kernel, bound_tensors, constexpr_values, @@ -478,7 +478,7 @@ def __init__( dense_slots = ( {layer: slot for slot, layer in enumerate(self.dense_layers)} if per_layer else None ) - dense_pack = _prepared_move_index_pack_launch( + dense_pack = _frozen_move_index_pack( kept_token_ordinals, valid_sequence_lengths, dense_move_offsets, @@ -493,7 +493,7 @@ def __init__( swa_move_source_indices=swa_move_indices, ) self.target_dense_compaction = _SingleCacheCompaction( - prepared_move_index_pack=dense_pack, + frozen_move_index_pack=dense_pack, cpp_compact_groups=_compact_groups( dense_entries, self.layer_pool_keys, self.device, dense_slots ), @@ -501,11 +501,11 @@ def __init__( move_source_offsets=dense_move_offsets, destination_bases=self.prompt_offsets, ) - # The dense pack launch fills the SWA move buffers in the same call. + # The dense pack call fills the SWA move buffers in the same run. self.target_swa_compaction = None if self.swa_layers: self.target_swa_compaction = _SingleCacheCompaction( - prepared_move_index_pack=None, + frozen_move_index_pack=None, cpp_compact_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), move_source_indices=swa_move_indices, move_source_offsets=swa_move_offsets, @@ -597,10 +597,10 @@ def _build_draft_compaction( for layer in draft_layers ] # In union mode the pack kernel reads selection row 0 for every - # packed row, so one more prepared launch broadcasts the target keep + # packed row, so one more frozen kernel call broadcasts the target keep # set over the draft KV heads and appends the draft's own tail # ordinals (valid_seq_len + 0..tail-1). - draft_pack = _prepared_move_index_pack_launch( + draft_pack = _frozen_move_index_pack( kept_token_ordinals, valid_sequence_lengths, draft_move_offsets, @@ -615,7 +615,7 @@ def _build_draft_compaction( swa_move_source_indices=None, ) return _SingleCacheCompaction( - prepared_move_index_pack=draft_pack, + frozen_move_index_pack=draft_pack, cpp_compact_groups=_compact_groups( draft_entries, tuple(draft_layer_pool_keys), self.device ), diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py deleted file mode 100644 index 3af30eb75151..000000000000 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/prepared_launch.py +++ /dev/null @@ -1,82 +0,0 @@ -# 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. - -"""Replayable Triton kernel launches frozen at build time.""" - -from typing import Dict, Tuple - -import torch - - -class PreparedTritonKernelLaunch: - """Replay one Triton kernel launch frozen at build time. - - ``warmup`` JIT-compiles the kernel once for a fixed grid, bound tensor - set, and constexpr set; ``__call__`` re-launches the compiled binary - directly, skipping Triton's per-call dispatch. Constexpr values are - passed positionally on replay, so their order is validated against the - kernel's constexpr parameter declaration order at build time. - """ - - def __init__( - self, - triton_kernel, - bound_tensors: Tuple[torch.Tensor, ...], - constexpr_values: Dict[str, object], - *, - grid: Tuple[int, ...], - num_warps: int, - ) -> None: - params = getattr(triton_kernel, "params", None) - if params is not None: - declared = [param.name for param in params if param.is_constexpr] - if list(constexpr_values.keys()) != declared: - raise ValueError( - f"constexpr order {list(constexpr_values.keys())} must match the " - f"kernel's declaration order {declared}: replay passes them " - "positionally" - ) - self.device = bound_tensors[0].device - self.bound_tensors = tuple(bound_tensors) - self.constexpr_values = dict(constexpr_values) - with torch.cuda.device(self.device): - self.build_stream = torch.cuda.current_stream(self.device) - # warmup() then indexing the compiled cache by grid is the - # documented-by-use Triton pattern for dispatch-free replay; if a - # Triton upgrade changes it, this raises here at build time rather - # than corrupting a replay. - compiled = triton_kernel.warmup( - *self.bound_tensors, - **self.constexpr_values, - num_warps=num_warps, - grid=grid, - ) - self.compiled_kernel_runner = compiled[grid] - - def __call__(self, *replay_tensors: torch.Tensor) -> None: - """Replay the launch; ``replay_tensors``, if given, substitute the bound tensors.""" - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.build_stream.device, - self.build_stream.cuda_stream, - ): - raise RuntimeError( - "a prepared Triton kernel launch must run on the stream it was built on" - ) - self.compiled_kernel_runner( - *(replay_tensors if replay_tensors else self.bound_tensors), - *self.constexpr_values.values(), - stream=self.build_stream.cuda_stream, - ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 83054ab97bed..8e13a5f4586b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -58,8 +58,8 @@ import torch -from tensorrt_llm._torch.kv_cache_compression.triattention.prepared_launch import ( - PreparedTritonKernelLaunch, +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + FrozenTritonKernelCall, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, get_draft_token_length @@ -218,7 +218,7 @@ def __init__( from .triattention_kernels import _finalize_topk_indices_kernel - self._prepared_launch = PreparedTritonKernelLaunch( + self._frozen_call = FrozenTritonKernelCall( _finalize_topk_indices_kernel, (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), dict( @@ -239,7 +239,7 @@ def __call__( provisional_indices: torch.Tensor, output_indices: torch.Tensor, ) -> None: - self._prepared_launch(scores, seq_lens, prompt_offsets, provisional_indices, output_indices) + self._frozen_call(scores, seq_lens, prompt_offsets, provisional_indices, output_indices) class _PreparedUnionScores: @@ -288,10 +288,10 @@ def __init__( # this launcher before dispatching to it. self.scores = scores self.normalize_scores = normalize_scores - self._prepared_stats_launch = None + self._frozen_stats_call = None if normalize_scores: stats_grid = (request_count * rows, 1, 1) - self._prepared_stats_launch = PreparedTritonKernelLaunch( + self._frozen_stats_call = FrozenTritonKernelCall( _score_row_stats_kernel, (scores, valid_widths, row_mean, row_inv_std), dict(ROWS=rows, WIDTH=width, BLOCK=256), @@ -299,7 +299,7 @@ def __init__( num_warps=4, ) union_grid = (request_count, (width + 31) // 32, 1) - self._prepared_union_launch = PreparedTritonKernelLaunch( + self._frozen_union_call = FrozenTritonKernelCall( _score_union_kernel, (scores, valid_widths, row_mean, row_inv_std, combined), dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), @@ -308,9 +308,9 @@ def __init__( ) def __call__(self) -> None: - if self._prepared_stats_launch is not None: - self._prepared_stats_launch() - self._prepared_union_launch() + if self._frozen_stats_call is not None: + self._frozen_stats_call() + self._frozen_union_call() def _deterministic_topk_indices_into( @@ -390,7 +390,7 @@ class _RuntimeKVLayout(NamedTuple): class _BatchedKeepSetSelectorBase: - """Shared fixed buffers and prepared launchers for keep-set selectors.""" + """Shared fixed buffers and frozen kernel calls for keep-set selectors.""" def __init__( self, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 4d84254ea53e..40c69829f604 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -16,13 +16,75 @@ from __future__ import annotations -from typing import List +from typing import Dict, List, Tuple import torch import triton import triton.language as tl +class FrozenTritonKernelCall: + """Call one Triton kernel frozen at build time. + + Triton's standard dispatch costs tens of microseconds of host work per + call; eviction fires its kernels every round, so the grid, bound tensor + set, and constexpr set are frozen once here. ``warmup`` (Triton's own + API) JIT-compiles the kernel at build time and ``__call__`` runs the + compiled binary directly. Constexpr values are passed positionally on + each call, so their order is validated against the kernel's constexpr + parameter declaration order at build time. + """ + + def __init__( + self, + triton_kernel, + bound_tensors: Tuple[torch.Tensor, ...], + constexpr_values: Dict[str, object], + *, + grid: Tuple[int, ...], + num_warps: int, + ) -> None: + params = getattr(triton_kernel, "params", None) + if params is not None: + declared = [param.name for param in params if param.is_constexpr] + if list(constexpr_values.keys()) != declared: + raise ValueError( + f"constexpr order {list(constexpr_values.keys())} must match the " + f"kernel's declaration order {declared}: the frozen call passes " + "them positionally" + ) + self.device = bound_tensors[0].device + self.bound_tensors = tuple(bound_tensors) + self.constexpr_values = dict(constexpr_values) + with torch.cuda.device(self.device): + self.build_stream = torch.cuda.current_stream(self.device) + # warmup() then indexing the compiled cache by grid is the + # documented-by-use Triton pattern for dispatch-free calls; if a + # Triton upgrade changes it, this raises here at build time rather + # than corrupting a later call. + compiled = triton_kernel.warmup( + *self.bound_tensors, + **self.constexpr_values, + num_warps=num_warps, + grid=grid, + ) + self.compiled_kernel_runner = compiled[grid] + + def __call__(self, *call_tensors: torch.Tensor) -> None: + """Run the kernel; ``call_tensors``, if given, substitute the bound tensors.""" + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.build_stream.device, + self.build_stream.cuda_stream, + ): + raise RuntimeError("a frozen Triton kernel call must run on the stream it was built on") + self.compiled_kernel_runner( + *(call_tensors if call_tensors else self.bound_tensors), + *self.constexpr_values.values(), + stream=self.build_stream.cuda_stream, + ) + + @triton.jit def _prepare_mean_phase_kernel( round_starts, From 9de7766664c0e4f58f32511f51737bb7c281240a Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 23:51:46 -0700 Subject: [PATCH 029/178] [None][fix] Catch the score-buffer tests up with the resident-buffer surface The width parameters grown by the resident-buffer rework left the direct score-staging and score-group tests behind: give both widths a whole-sequence default, move the three call sites off the removed min_prompt_len parameter, and teach the draft pack test the capacity-sized index buffer (moves are packed where the offsets point). Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 9 ++++++--- .../triattention/triattention_kernels.py | 5 ++++- .../test_triattention_draft_cocompaction.py | 9 +++++++-- .../kv_cache_compression/test_triattention_eager.py | 2 +- .../kv_cache_compression/test_triattention_pipeline.py | 4 ++-- 5 files changed, 20 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 8e13a5f4586b..3f2c1ebc41cd 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -969,7 +969,7 @@ def __init__( omega: torch.Tensor, page_table_keys: Optional[List[object]] = None, num_page_table_slots: Optional[int] = None, - decode_width: int = 0, + decode_width: Optional[int] = None, page_table_token_capacity: Optional[int] = None, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_page_representatives: Optional[List[int]] = None, @@ -1000,9 +1000,12 @@ def __init__( raise ValueError("page-table capacity cannot be smaller than the score bucket") self.page_table_token_capacity = int(page_table_token_capacity) # Decode-width capacity of the score buffers; per-request prompt - # lengths are staged runtime metadata. + # lengths are staged runtime metadata. Default: the whole sequence + # capacity is scorable. + if decode_width is None: + decode_width = int(seq_len) if decode_width <= 0 or decode_width > seq_len: - raise ValueError("fixed score decode width is outside its bucket") + raise ValueError("fixed score decode width exceeds the sequence capacity") self.decode_width = int(decode_width) q_real = q_real.to(device=self.device, dtype=torch.float32).contiguous() q_imag = q_imag.to(device=self.device, dtype=torch.float32).contiguous() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 40c69829f604..2966f680a84b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -849,10 +849,13 @@ def __init__( freq_scale_sq: torch.Tensor, omega: torch.Tensor, offsets: torch.Tensor, - output_width: int = 0, + output_width: int | None = None, ) -> None: if not layer_indices or min(max_requests, page_count, seq_len) <= 0: raise ValueError("fixed score group requires non-empty positive geometry") + # Default: the whole sequence capacity is scorable. + if output_width is None: + output_width = int(seq_len) if output_width <= 0 or output_width > seq_len: raise ValueError("fixed score group requires a decode width within its capacity") if len(page_table_slots) != len(layer_indices): 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 index f33dba37947e..5e2057604a7a 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -261,10 +261,15 @@ def test_draft_pack_matches_keep_broadcast_and_tail_ordinal_oracle(draft_protect assert draft_compaction.move_source_offsets.cpu().tolist() == expected_offsets draft_indices = draft_compaction.move_source_indices - assert draft_indices.shape == (int(built.draft_pool.shape[2]), expected_offsets[-1]) + # The index buffer is sized for the widest tail (the capacity); this + # round's moves are packed at the front, where the offsets point. + capacity_total = built.request_count * ( + int(built.keep.shape[1]) + max(built.draft_protected_tails) + ) + assert draft_indices.shape == (int(built.draft_pool.shape[2]), capacity_total) for head in range(int(draft_indices.shape[0])): # Union mode broadcasts one keep set over every draft KV head. - assert torch.equal(draft_indices[head], expected_row) + assert torch.equal(draft_indices[head, : expected_offsets[-1]], expected_row) def test_mark_page_tables_consumed_orders_both_manager_streams(): diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index c4b55a5b2078..8a2477f7adc7 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -923,7 +923,7 @@ def write_token(token: int, score: float) -> None: torch.zeros(num_freqs, dtype=torch.float32, device=device), page_table_keys=[("pool", 0)], num_page_table_slots=1, - min_prompt_len=prompt_len, + decode_width=seq_len - prompt_len, page_table_token_capacity=seq_len + protected_tail, ) keep_set_selector = _BatchedUnionKeepSetSelector( diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 82ee6503a1aa..fe2ab0a47a33 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1601,7 +1601,7 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, freq, omega, offsets, - min_prompt_len=prompt_len, + output_width=seq_len - prompt_len, ) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) fixed = group.launch( @@ -1704,7 +1704,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun freq, omega, offsets, - min_prompt_len=prompt_len, + output_width=seq_len - prompt_len, ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) From 165e03b76cc4ec83ad764d7960b462450b940135 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 16 Jul 2026 23:57:20 -0700 Subject: [PATCH 030/178] [None][chore] Drop the device-normalization helper CUDA tensors always report an indexed device, and production only ever derives devices from tensors, so normalizing an index-less cuda device protected nothing but hand-built test inputs. Follow the upstream convention instead: take devices from tensors, and have the tests build indexed devices. Signed-off-by: tianruih --- .../triattention/triattention.py | 11 ++-------- .../test_triattention_draft_cocompaction.py | 4 ++-- .../test_triattention_eager.py | 22 +++++++++---------- .../test_triattention_pipeline.py | 12 +++++----- 4 files changed, 21 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 3f2c1ebc41cd..b856b1b40c49 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -94,13 +94,6 @@ def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Ten return torch.tensor(offsets, device=device, dtype=torch.float32) -def _canonical_device(device: torch.device) -> torch.device: - device = torch.device(device) - if device.type == "cuda" and device.index is None: - return torch.device("cuda", torch.cuda.current_device()) - return device - - def _topk_indices_into( scores: torch.Tensor, seq_lens: torch.Tensor, @@ -418,7 +411,7 @@ def __init__( self.width = int(width) self.keep_count = int(keep_count) self.dtype = dtype - self.device = _canonical_device(device) + self.device = device self.max_requests = int(max_requests) self.valid_widths = torch.full( (self.max_requests,), self.width, dtype=torch.int32, device=self.device @@ -989,7 +982,7 @@ def __init__( or set(dense_layers) != set(grouped_layers) ): raise ValueError("dense layer order must cover every grouped layer exactly once") - self.device = _canonical_device(layer_pools[page_representatives[0]].device) + self.device = layer_pools[page_representatives[0]].device if self.device.type != "cuda": raise ValueError("fixed score metadata is CUDA-only") self.max_requests = max_requests 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 index 5e2057604a7a..1a8ad27752c8 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -95,7 +95,7 @@ def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: def _launched_draft_compaction(draft_protected_tails): """Build target and draft pools with distinct head counts, then compact.""" - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 target_kv_heads = 2 draft_kv_heads = 4 @@ -274,7 +274,7 @@ def test_draft_pack_matches_keep_broadcast_and_tail_ordinal_oracle(draft_protect def test_mark_page_tables_consumed_orders_both_manager_streams(): staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = torch.device("cuda") + staging.device = torch.device("cuda", torch.cuda.current_device()) staging.page_tables_active = True event = mock.Mock() staging.bulk_consume_done = event diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index 8a2477f7adc7..c3c1ca84a542 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -204,7 +204,7 @@ def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( reference.select_requests(scores_cpu, normalize_scores=normalize_scores) expected = reference.keep.clone() - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) stream = torch.cuda.Stream(device=device) with torch.cuda.stream(stream): selector = _BatchedPerHeadKeepSetSelector( @@ -232,7 +232,7 @@ def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( def test_union_eager_runs_the_registered_cute_op(): _require_cute_topk_op() - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) scores = torch.randn(2, 4, 96, dtype=torch.float32, device=device) selector = _BatchedUnionKeepSetSelector( rows=4, @@ -254,7 +254,7 @@ def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_ _require_cute_topk_op() from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) request_count, rows, width, keep_count = 2, 7, 97, 64 generator = torch.Generator(device=device).manual_seed(53) scores = torch.randn( @@ -312,7 +312,7 @@ def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_ @pytest.mark.parametrize("keep_count,width", [(4096, 4224), (8192, 9216)]) def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): _require_cute_topk_op() - device = torch.device("cuda") + 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) @@ -355,7 +355,7 @@ def test_fused_union_preparation_matches_ragged_torch_reference(): prepare_union_scores, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) request_count, rows, width = 2, 7, 97 generator = torch.Generator(device=device).manual_seed(17) scores = torch.randn( @@ -399,7 +399,7 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no prepare_per_head_scores, ) - device = torch.device("cuda") + 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( @@ -464,7 +464,7 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode): - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 num_layers = 2 num_kv_heads = 2 @@ -577,7 +577,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): """One union cohort mixing prompt lengths compacts byte-identically to running the same two requests as two single-request compactions.""" - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 num_layers = 2 num_kv_heads = 2 @@ -695,7 +695,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): _FixedScoreStagingBuffers, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) num_layers = 3 seq_len = 8 keep_count = 2 @@ -825,7 +825,7 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) request_id = 7 prompt_len = 2 seq_len = 10 @@ -1043,7 +1043,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: def test_eager_compaction_rebases_masked_swa_window_and_tail(): - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) initial_pools = [ diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index fe2ab0a47a33..685135796fdb 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1133,7 +1133,7 @@ def test_bulk_page_table_copy_uses_immutable_host_snapshots(self): _FixedScoreStagingBuffers, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) current_stream = torch.cuda.current_stream(device) manager_stream = torch.cuda.Stream(device=device) host_table = torch.zeros( @@ -1238,7 +1238,7 @@ def test_next_bulk_copy_waits_for_page_table_consumers(self): _FixedScoreStagingBuffers, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) current_stream = torch.cuda.current_stream(device) manager_stream = torch.cuda.Stream(device=device) host_table = torch.zeros(1, 1, 2, 4, dtype=torch.int32, device="cpu", pin_memory=True) @@ -1312,7 +1312,7 @@ def test_cross_stream_staging_is_rejected_before_page_table_query(self): ) staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = torch.device("cuda") + staging.device = torch.device("cuda", torch.cuda.current_device()) staging.max_requests = 8 staging.stream = SimpleNamespace(device=torch.device("cuda:0"), cuda_stream=4) staging.page_tables_active = False @@ -1376,7 +1376,7 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re _FixedScoreStreamMismatch, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) max_requests = 8 page_count = 3 seq_len = 7 @@ -1535,7 +1535,7 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, _FixedScoreGroup, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(20260703 + request_count) max_requests = 8 page_count = 2 @@ -1652,7 +1652,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun _FixedScoreGroup, ) - device = torch.device("cuda") + device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(20260707 + request_count) max_requests = request_count page_count = 2 From b4070f81a08622ba84751baef188bbfa59763377 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 00:05:07 -0700 Subject: [PATCH 031/178] [None][chore] Call the CuTE top-k operation directly A one-line wrapper with a single caller; the sparse-attention backends call their operations inline, so follow that convention. The trailing next_n=1 argument gets a comment instead. Signed-off-by: tianruih --- .../triattention/triattention.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b856b1b40c49..654b27b01f40 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -94,16 +94,6 @@ def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Ten return torch.tensor(offsets, device=device, dtype=torch.float32) -def _topk_indices_into( - scores: torch.Tensor, - seq_lens: torch.Tensor, - indices_i32: torch.Tensor, - keep_count: int, -) -> None: - """Write per-row top-k indices with the CuTE-DSL selector.""" - torch.ops.trtllm.cute_dsl_indexer_topk_decode(scores, seq_lens, indices_i32, keep_count, 1) - - class _PreparedCuteTopK: """Run the existing fixed-shape CuTe TopK with owned scratch storage.""" @@ -335,7 +325,10 @@ def _deterministic_topk_indices_into( return # CPU selectors only exercise the selector contract in unit tests. - _topk_indices_into(scores, seq_lens, provisional_indices_i32, keep_count) + # The trailing 1 is next_n: decode scores one query token per request. + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + scores, seq_lens, provisional_indices_i32, keep_count, 1 + ) for row_index, row_scores in enumerate(scores): valid_width = int(seq_lens[row_index]) prompt_len = int(prompt_offsets[row_index]) From 201978e70a6ea462bbf7ce12c7ce3621c36d9871 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 00:46:22 -0700 Subject: [PATCH 032/178] [None][refactor] Make the keep-set selectors CUDA-only The selectors carried a parallel CPU implementation, CPU-only scratch buffers, per-element input sanitizing, and device-type branches -- all of it test scaffolding living in the production module (the sparse-attention backends carry none of this). Selection now always runs the compiled CuTE top-k plus the frozen deterministic finalizer as one prepared object, and the tests compare that real path against plain torch oracles of their own, which also covers tie-breaking without mocked kernels. Cohort staging trusts its internal callers: the int32 conversion already rejects overflow, so the per-element scans are gone. Signed-off-by: tianruih --- .../triattention/triattention.py | 522 +++--------------- .../test_triattention_eager.py | 197 ++----- .../test_triattention_pipeline.py | 66 +-- 3 files changed, 136 insertions(+), 649 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 654b27b01f40..372f23fc186f 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -79,8 +79,6 @@ # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) -_INT32_MAX = torch.iinfo(torch.int32).max - def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Tensor: """Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <=max].""" @@ -94,77 +92,16 @@ def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Ten return torch.tensor(offsets, device=device, dtype=torch.float32) -class _PreparedCuteTopK: - """Run the existing fixed-shape CuTe TopK with owned scratch storage.""" - - def __init__(self, max_rows: int, width: int, keep_count: int, device: torch.device) -> None: - if max_rows <= 0 or width <= 0 or not 1 <= keep_count <= width: - raise ValueError("prepared CuTe TopK requires valid fixed dimensions") - if device.type != "cuda": - raise ValueError("prepared CuTe TopK requires a CUDA device") - - from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops - - self.max_rows = max_rows - self.width = width - self.keep_count = keep_count - self.device = device - with torch.cuda.device(device): - self.stream = torch.cuda.current_stream(device) - self.scratch = torch.empty((max_rows, 2, width), dtype=torch.int32, device=device) - runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner - key = ( - cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], - 1 << (width - 1).bit_length(), - keep_count, - 1, - False, - 256, - False, - max_rows > cute_dsl_custom_ops._get_num_sms(), - ) - runner._compile(*key) - self.compiled = runner.kernel_cache[key] - - def __call__( - self, - scores: torch.Tensor, - seq_lens: torch.Tensor, - output_indices: torch.Tensor, - ) -> None: - rows = int(scores.shape[0]) - if ( - rows != self.max_rows - or scores.shape != (self.max_rows, self.width) - or scores.dtype != torch.float32 - or scores.device != self.device - or seq_lens.shape != (self.max_rows,) - or seq_lens.dtype != torch.int32 - or seq_lens.device != self.device - or output_indices.shape != (self.max_rows, self.keep_count) - or output_indices.dtype != torch.int32 - or output_indices.device != self.device - ): - raise ValueError("prepared CuTe TopK inputs do not match their fixed buffers") - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.stream.device, - self.stream.cuda_stream, - ): - raise RuntimeError("prepared CuTe TopK must run on the stream it was built on") - self.compiled( - scores, - None, - self.scratch, - None, - seq_lens, - output_indices, - None, - ) - +class _PreparedDeterministicTopK: + """Deterministic top-k over fixed row-major buffers, built once. -class _PreparedTopKFinalizer: - """Launch the fixed deterministic finalizer without Triton JIT dispatch.""" + The CuTE top-k kernel is fast but breaks score ties arbitrarily and emits + indices in arbitrary order, so a frozen Triton finalizer recomputes the + threshold membership with lowest-index-wins ties, rebases each row by its + prompt offset, and writes sorted ordinals. The CuTE kernel is compiled + here once with owned scratch storage; every later call runs both kernels + over the bound buffers with no dispatch work and no allocations. + """ def __init__( self, @@ -176,32 +113,34 @@ def __init__( keep_count: int, ) -> None: rows, width = scores.shape - if ( - not scores.is_cuda - or scores.dtype != torch.float32 - or not scores.is_contiguous() - or seq_lens.shape != (rows,) - or seq_lens.dtype != torch.int32 - or seq_lens.device != scores.device - or prompt_offsets.shape != (rows,) - or prompt_offsets.dtype != torch.int32 - or prompt_offsets.device != scores.device - or provisional_indices.shape != (rows, keep_count) - or provisional_indices.dtype != torch.int32 - or provisional_indices.device != scores.device - or output_indices.shape != (rows, keep_count) - or output_indices.dtype != torch.int32 - or output_indices.device != scores.device - or not seq_lens.is_contiguous() - or not prompt_offsets.is_contiguous() - or not provisional_indices.is_contiguous() - or not output_indices.is_contiguous() - ): - raise ValueError("prepared TopK finalizer tensors do not share one fixed geometry") + if not 1 <= keep_count <= width: + raise ValueError("deterministic top-k requires 1 <= keep_count <= width") + + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops from .triattention_kernels import _finalize_topk_indices_kernel - self._frozen_call = FrozenTritonKernelCall( + self.device = scores.device + self.scores = scores + self.seq_lens = seq_lens + self.provisional_indices = provisional_indices + with torch.cuda.device(self.device): + self.stream = torch.cuda.current_stream(self.device) + self.scratch = torch.empty((rows, 2, width), dtype=torch.int32, device=self.device) + runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner + key = ( + cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], + 1 << (width - 1).bit_length(), + keep_count, + 1, + False, + 256, + False, + rows > cute_dsl_custom_ops._get_num_sms(), + ) + runner._compile(*key) + self.compiled_topk = runner.kernel_cache[key] + self.frozen_finalize = FrozenTritonKernelCall( _finalize_topk_indices_kernel, (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), dict( @@ -214,15 +153,17 @@ def __init__( num_warps=4, ) - def __call__( - self, - scores: torch.Tensor, - seq_lens: torch.Tensor, - prompt_offsets: torch.Tensor, - provisional_indices: torch.Tensor, - output_indices: torch.Tensor, - ) -> None: - self._frozen_call(scores, seq_lens, prompt_offsets, provisional_indices, output_indices) + def __call__(self) -> None: + self.compiled_topk( + self.scores, + None, + self.scratch, + None, + self.seq_lens, + self.provisional_indices, + None, + ) + self.frozen_finalize() class _PreparedUnionScores: @@ -296,52 +237,6 @@ def __call__(self) -> None: self._frozen_union_call() -def _deterministic_topk_indices_into( - scores: torch.Tensor, - seq_lens: torch.Tensor, - prompt_offsets: torch.Tensor, - provisional_indices_i32: torch.Tensor, - output_indices_i32: torch.Tensor, - keep_count: int, - prepared_topk: Optional[_PreparedCuteTopK] = None, - prepared_finalizer: Optional[_PreparedTopKFinalizer] = None, -) -> None: - """Write stable, increasing physical indices around one CuTE TopK call. - - Scores are decode-relative; ``prompt_offsets`` rebases each row's emitted - ordinals to absolute positions, so one call may mix prompt lengths. - """ - if scores.is_cuda: - if prepared_topk is None or prepared_finalizer is None: - raise ValueError("CUDA selection requires prepared TopK launchers") - prepared_topk(scores, seq_lens, provisional_indices_i32) - prepared_finalizer( - scores, - seq_lens, - prompt_offsets, - provisional_indices_i32, - output_indices_i32, - ) - return - - # CPU selectors only exercise the selector contract in unit tests. - # The trailing 1 is next_n: decode scores one query token per request. - torch.ops.trtllm.cute_dsl_indexer_topk_decode( - scores, seq_lens, provisional_indices_i32, keep_count, 1 - ) - for row_index, row_scores in enumerate(scores): - valid_width = int(seq_lens[row_index]) - prompt_len = int(prompt_offsets[row_index]) - selected = provisional_indices_i32[row_index].to(torch.long) - threshold = torch.amin(row_scores[selected]) - valid_scores = row_scores[:valid_width] - higher = torch.nonzero(valid_scores > threshold, as_tuple=False).flatten() - tied = torch.nonzero(valid_scores == threshold, as_tuple=False).flatten() - tie_count = keep_count - int(higher.numel()) - ordered = torch.sort(torch.cat((higher, tied[:tie_count]))).values - output_indices_i32[row_index, :keep_count].copy_(ordered.to(torch.int32).add(prompt_len)) - - class _CrossRequestSelectionPlan(NamedTuple): """Selection dimensions used to allocate reusable eager buffers.""" @@ -455,22 +350,6 @@ def refresh_row_prompt_offsets(self) -> None: self.prompt_offsets.unsqueeze(1).expand(-1, self.selection_rows_per_request) ) - def _allocate_cpu_reference_buffers( - self, - scale_shape: Tuple[int, ...], - mask_shape: Tuple[int, ...], - index_dtype: torch.dtype, - ) -> None: - """CPU-only oracle scratch; None on CUDA, where prepared kernels run.""" - if self.device.type == "cpu": - self.valid_scale = torch.empty(scale_shape, dtype=self.dtype, device=self.device) - self.token_indices = torch.arange(self.width, dtype=index_dtype, device=self.device) - self.invalid_mask = torch.empty(mask_shape, dtype=torch.bool, device=self.device) - else: - self.valid_scale = None - self.token_indices = None - self.invalid_mask = None - def _build_prepared_selection_launchers( self, scores_rows: torch.Tensor, @@ -478,15 +357,8 @@ def _build_prepared_selection_launchers( provisional_indices: torch.Tensor, keep_rows: torch.Tensor, ) -> None: - """Bind the CuTE topk and the deterministic finalizer over row-major views.""" - if self.device.type != "cuda": - self.prepared_topk = None - self.prepared_finalizer = None - return - self.prepared_topk = _PreparedCuteTopK( - int(scores_rows.shape[0]), self.width, self.keep_count, self.device - ) - self.prepared_finalizer = _PreparedTopKFinalizer( + """Bind the deterministic top-k over row-major views of the buffers.""" + self.prepared_topk = _PreparedDeterministicTopK( scores_rows, row_lengths, self.row_prompt_offsets, @@ -495,14 +367,6 @@ def _build_prepared_selection_launchers( self.keep_count, ) - def _validated_request_count(self, scores: torch.Tensor, what: str) -> int: - request_count = int(scores.shape[0]) if scores.ndim >= 1 else 0 - if request_count <= 0 or request_count > self.max_requests: - raise ValueError(f"request count exceeds the {what} selection capacity") - if scores.is_cuda and request_count != self.max_requests: - raise ValueError("CUDA selection requires the selector's fixed request count") - return request_count - class _BatchedUnionKeepSetSelector(_BatchedKeepSetSelectorBase): """Persistent ``[request, ...]`` buffers for union selection.""" @@ -538,8 +402,8 @@ def __init__( max_requests=max_requests, ) self.rows = rows - if self.device.type == "cuda" and input_scores is None: - raise ValueError("CUDA union selection requires its fixed score input") + if input_scores is None: + raise ValueError("union selection requires its fixed score input") self.row_mean = torch.empty((max_requests, rows, 1), dtype=dtype, device=self.device) self.row_std = torch.empty_like(self.row_mean) @@ -552,112 +416,22 @@ def __init__( self.keep = torch.empty( (max_requests, self.keep_count), dtype=torch.int32, device=self.device ) - self._allocate_cpu_reference_buffers( - (max_requests, 1, 1), (max_requests, 1, width), torch.int32 - ) self._build_prepared_selection_launchers( self.combined, self.valid_widths, self.final_indices, self.keep ) - self.prepared_scores = ( - _PreparedUnionScores( - input_scores, - self.valid_widths, - self.row_mean, - self.row_std, - self.combined, - normalize_scores=normalize_scores, - ) - if self.device.type == "cuda" and input_scores is not None - else None - ) - - def _select_input_scores( - self, - input_scores: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, - ) -> None: - valid_widths = self.valid_widths[:request_count] - combined = self.combined[:request_count] - if input_scores.is_cuda: - raise RuntimeError("CUDA union scores must use their prepared fixed launcher") - else: - self._select_input_scores_reference( - input_scores, - request_count, - normalize_scores=normalize_scores, - ) - - final_indices = self.final_indices[:request_count] - _deterministic_topk_indices_into( - combined, - valid_widths, - self.prompt_offsets[:request_count], - final_indices, - self.keep[:request_count], - self.keep_count, - self.prepared_topk, - self.prepared_finalizer, + self.prepared_scores = _PreparedUnionScores( + input_scores, + self.valid_widths, + self.row_mean, + self.row_std, + self.combined, + normalize_scores=normalize_scores, ) def select_prepared_requests(self) -> None: """Select from the CUDA score tensor bound to this fixed selector.""" - if self.prepared_scores is None: - raise RuntimeError("prepared union scores are unavailable") self.prepared_scores() - final_indices = self.final_indices - _deterministic_topk_indices_into( - self.combined, - self.valid_widths, - self.prompt_offsets, - final_indices, - self.keep, - self.keep_count, - self.prepared_topk, - self.prepared_finalizer, - ) - - def _select_input_scores_reference( - self, - input_scores: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, - ) -> None: - """Keep a CPU-only reference for selector contract tests.""" - valid_widths = self.valid_widths[:request_count] - assert self.invalid_mask is not None - assert self.token_indices is not None - assert self.valid_scale is not None - invalid_mask = self.invalid_mask[:request_count] - torch.ge( - self.token_indices.view(1, 1, self.width), - valid_widths.view(request_count, 1, 1), - out=invalid_mask, - ) - if normalize_scores: - row_mean = self.row_mean[:request_count] - row_std = self.row_std[:request_count] - input_scores.masked_fill_(invalid_mask, 0.0) - torch.sum(input_scores, dim=2, keepdim=True, out=row_mean) - self.valid_scale[:request_count].view(request_count).copy_(valid_widths) - row_mean.div_(self.valid_scale[:request_count]) - torch.sub(input_scores, row_mean, out=input_scores) - input_scores.masked_fill_(invalid_mask, 0.0) - torch.linalg.vector_norm( - input_scores, - dim=2, - keepdim=True, - out=row_std, - ) - self.valid_scale[:request_count].sqrt_() - row_std.div_(self.valid_scale[:request_count]) - row_std.clamp_min_(1e-6) - torch.div(input_scores, row_std, out=input_scores) - input_scores.masked_fill_(invalid_mask, float("-inf")) - - torch.amax(input_scores, dim=1, out=self.combined[:request_count]) + self.prepared_topk() def select_requests( self, @@ -665,30 +439,13 @@ def select_requests( *, normalize_scores: bool, ) -> None: - """Select from request-major score output without repacking it.""" - request_count = self._validated_request_count(scores, "cross-request") + """Select from the score tensor bound to this fixed selector.""" if ( - scores.numel() != request_count * self.rows * self.width - or int(scores.shape[-1]) != self.width - or scores.dtype != self.dtype - or scores.device != self.device - or not scores.is_contiguous() + scores is not self.prepared_scores.scores + or bool(normalize_scores) != self.prepared_scores.normalize_scores ): - raise ValueError("cross-request scores do not match the selector geometry") - if scores.is_cuda: - if ( - self.prepared_scores is None - or scores is not self.prepared_scores.scores - or bool(normalize_scores) != self.prepared_scores.normalize_scores - ): - raise ValueError("CUDA union scores do not match their prepared fixed launcher") - self.select_prepared_requests() - return - self._select_input_scores( - scores.view(request_count, self.rows, self.width), - request_count, - normalize_scores=normalize_scores, - ) + raise ValueError("union scores do not match their prepared fixed launcher") + self.select_prepared_requests() class _BatchedPerHeadKeepSetSelector(_BatchedKeepSetSelectorBase): @@ -731,12 +488,10 @@ def __init__( max_requests=max_requests, ) self.num_layers = len(self.dense_layers) - self.query_group_size = self.num_query_heads // self.num_kv_heads self.rows = self.num_layers * self.num_query_heads self.selection_rows = selection_rows score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) - grouped_shape = (self.max_requests, self.num_layers, self.num_kv_heads, self.width) self.row_mean = torch.empty(score_shape[:-1] + (1,), dtype=dtype, device=self.device) self.row_std = torch.empty_like(self.row_mean) self.selection_scores = torch.empty( @@ -744,14 +499,6 @@ def __init__( dtype=dtype, device=self.device, ) - self._allocate_cpu_reference_buffers( - (self.max_requests, 1, 1, 1), (self.max_requests, 1, 1, self.width), torch.long - ) - self.grouped_scores = ( - torch.empty(grouped_shape, dtype=dtype, device=self.device) - if self.device.type == "cpu" - else None - ) self.row_seq_lens = torch.full( (self.max_requests, self.selection_rows), self.width, @@ -780,126 +527,35 @@ def __init__( self.keep_flat, ) - def _select_input_scores( - self, - input_scores: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, - ) -> None: - valid_widths = self.valid_widths[:request_count] - selection_scores = self.selection_scores[:request_count] - row_seq_lens = self.row_seq_lens[:request_count] - if input_scores.is_cuda: - from .triattention_kernels import prepare_per_head_scores - - prepare_per_head_scores( - input_scores, - valid_widths, - self.row_mean[:request_count], - self.row_std[:request_count], - selection_scores, - row_seq_lens, - request_count, - num_kv_heads=self.num_kv_heads, - per_layer=self.eviction_mode == "per_layer_perhead", - normalize_scores=normalize_scores, - ) - else: - self._select_input_scores_reference( - input_scores, - request_count, - normalize_scores=normalize_scores, - ) - - _deterministic_topk_indices_into( - self.selection_scores_flat, - self.row_seq_lens_flat, - self.row_prompt_offsets, - self.top_indices_i32_flat, - self.keep_flat, - self.keep_count, - self.prepared_topk, - self.prepared_finalizer, - ) - - def _select_input_scores_reference( - self, - input_scores: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, - ) -> None: - """Keep the explicit PyTorch implementation as the CPU oracle.""" - valid_widths = self.valid_widths[:request_count] - assert self.invalid_mask is not None - assert self.token_indices is not None - assert self.valid_scale is not None - assert self.grouped_scores is not None - invalid_mask = self.invalid_mask[:request_count] - torch.ge( - self.token_indices.view(1, 1, 1, self.width), - valid_widths.view(request_count, 1, 1, 1), - out=invalid_mask, - ) - if normalize_scores: - row_mean = self.row_mean[:request_count] - row_std = self.row_std[:request_count] - input_scores.masked_fill_(invalid_mask, 0.0) - torch.sum(input_scores, dim=3, keepdim=True, out=row_mean) - self.valid_scale[:request_count].view(request_count).copy_(valid_widths) - row_mean.div_(self.valid_scale[:request_count]) - torch.sub(input_scores, row_mean, out=input_scores) - input_scores.masked_fill_(invalid_mask, 0.0) - torch.linalg.vector_norm(input_scores, dim=3, keepdim=True, out=row_std) - self.valid_scale[:request_count].sqrt_() - row_std.div_(self.valid_scale[:request_count]) - row_std.clamp_min_(1e-6) - torch.div(input_scores, row_std, out=input_scores) - input_scores.masked_fill_(invalid_mask, float("-inf")) - - grouped_scores = self.grouped_scores[:request_count] - torch.amax( - input_scores.view( - request_count, - self.num_layers, - self.num_kv_heads, - self.query_group_size, - self.width, - ), - dim=3, - out=grouped_scores, - ) - if self.eviction_mode == "per_head": - torch.mean(grouped_scores, dim=1, out=self.selection_scores[:request_count]) - else: - self.selection_scores[:request_count].copy_( - grouped_scores.view(request_count, self.selection_rows, self.width) - ) - row_seq_lens = self.row_seq_lens[:request_count] - row_seq_lens.copy_(valid_widths.view(request_count, 1).expand(-1, self.selection_rows)) - def select_requests( self, scores: torch.Tensor, *, normalize_scores: bool, ) -> None: - request_count = self._validated_request_count(scores, "per-head") + from .triattention_kernels import prepare_per_head_scores + expected_shape = ( - request_count, + self.max_requests, self.num_layers, self.num_query_heads, self.width, ) - if ( - tuple(scores.shape) != expected_shape - or scores.dtype != self.dtype - or scores.device != self.device - or not scores.is_contiguous() - ): + if tuple(scores.shape) != expected_shape or not scores.is_contiguous(): raise ValueError("per-head scores do not match the selector geometry") - self._select_input_scores(scores, request_count, normalize_scores=normalize_scores) + prepare_per_head_scores( + scores, + self.valid_widths, + self.row_mean, + self.row_std, + self.selection_scores, + self.row_seq_lens, + self.max_requests, + num_kv_heads=self.num_kv_heads, + per_layer=self.eviction_mode == "per_layer_perhead", + normalize_scores=normalize_scores, + ) + self.prepared_topk() class _FixedScoreStreamMismatch(RuntimeError): @@ -1286,17 +942,6 @@ def stage( or request_count > self.max_requests or len(round_starts) != request_count or len(token_starts) != request_count - or any( - round_start != round_start - or round_start < 0 - or round_start > _INT32_MAX - or round_start != int(round_start) - for round_start in round_starts - ) - or any( - token_start < 0 or token_start > _INT32_MAX or token_start != int(token_start) - for token_start in token_starts - ) ): return False if (draft_manager is None) != (self.draft_block_offsets_device is None): @@ -1315,16 +960,9 @@ def stage( raise RuntimeError("previous page-table cohort is still active") if seq_lens is None: seq_lens = [self.bucket_seq_len] * request_count - if len(seq_lens) != request_count or any( - seq_len <= 0 or seq_len > self.bucket_seq_len for seq_len in seq_lens - ): - return False if page_table_seq_lens is None: page_table_seq_lens = seq_lens - if len(page_table_seq_lens) != request_count or any( - page_seq_len < seq_len or page_seq_len > self.page_table_token_capacity - for seq_len, page_seq_len in zip(seq_lens, page_table_seq_lens) - ): + if len(seq_lens) != request_count or len(page_table_seq_lens) != request_count: return False if manager.enable_swa_scratch_reuse: raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index c3c1ca84a542..99ba128e30de 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -29,30 +29,38 @@ def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor return torch.tensor(selected[:keep_count], dtype=torch.int32, device=row.device) -def _fake_cute_topk(scores, seq_lens, output, top_k, next_n): - assert next_n == 1 - for row_index, row in enumerate(scores): - output[row_index].copy_(_stable_topk(row, int(seq_lens[row_index]), int(top_k))) - - -class _AdversarialTieTopK: - """Prefer high-index boundary ties so finalization must correct membership.""" - - def __init__(self): - self.calls = 0 - - def __call__(self, scores, seq_lens, output, top_k, next_n): - assert next_n == 1 - self.calls += 1 - for row_index, row in enumerate(scores): - width = int(seq_lens[row_index]) - values = row[:width] - threshold = torch.sort(values, descending=True).values[int(top_k) - 1] - higher = torch.nonzero(values > threshold, as_tuple=False).flatten() - tied = torch.nonzero(values == threshold, as_tuple=False).flatten() - tied = tied.flip(0) - remaining = int(top_k) - int(higher.numel()) - output[row_index].copy_(torch.cat((higher, tied[:remaining])).to(torch.int32)) +def _per_head_keep_oracle( + scores: torch.Tensor, + valid_widths: 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): + valid_width = int(valid_widths[request]) + valid = scores[request, ..., :valid_width].clone() + if normalize_scores: + mean = valid.mean(dim=-1, keepdim=True) + valid = valid - mean + std = valid.norm(dim=-1, keepdim=True) / (valid_width**0.5) + valid = valid / std.clamp_min(1e-6) + grouped = valid.view( + num_layers, num_kv_heads, num_query_heads // num_kv_heads, valid_width + ).amax(dim=2) + if eviction_mode == "per_head": + selection = grouped.mean(dim=0) + else: + selection = grouped.reshape(num_layers * num_kv_heads, valid_width) + rows.append( + torch.stack( + [torch.sort(_stable_topk(row, valid_width, keep_count)).values for row in selection] + ) + ) + return torch.stack(rows) def _legacy_union(scores: torch.Tensor, keep_count: int) -> torch.Tensor: @@ -85,92 +93,12 @@ def test_direct_union_topk_matches_legacy_union_with_heavy_ties(rows, width, kee assert torch.equal(direct, _legacy_union(scores, keep_count)) -@pytest.mark.parametrize("keep_count,width", [(4, 8), (4096, 4224), (8192, 9216)]) -def test_union_eager_uses_one_deterministic_cute_selection(keep_count, width): - prompt_len = 17 - generator = torch.Generator().manual_seed(keep_count) - scores = torch.randint( - -8, - 9, - (2, width), - generator=generator, - dtype=torch.int32, - ).to(torch.float32) - selector = _BatchedUnionKeepSetSelector( - rows=2, - width=width, - keep_count=keep_count, - dtype=torch.float32, - device=torch.device("cpu"), - max_requests=1, - ) - selector.set_prompt_offsets(torch.tensor([prompt_len], dtype=torch.int32)) - raw_topk = _AdversarialTieTopK() - with ( - mock.patch.object( - torch.ops.trtllm, - "cute_dsl_indexer_topk_decode", - side_effect=raw_topk, - create=True, - ), - mock.patch.object( - torch.ops.trtllm, - "indexer_topk_decode", - side_effect=AssertionError("legacy selector was called"), - create=True, - ), - ): - selector.select_requests(scores.unsqueeze(0), normalize_scores=False) - - expected = _stable_topk(scores.max(dim=0).values, width, keep_count) - assert torch.equal( - selector.keep[0], - torch.sort(expected + prompt_len).values, - ) - assert raw_topk.calls == 1 - - -@pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) -def test_per_head_eager_keeps_stable_indices(eviction_mode): - selector = _BatchedPerHeadKeepSetSelector( - eviction_mode=eviction_mode, - dense_layers=(0, 1), - num_query_heads=4, - num_kv_heads=2, - width=16, - keep_count=5, - dtype=torch.float32, - device=torch.device("cpu"), - max_requests=1, - ) - selector.set_prompt_offsets(torch.tensor([3], dtype=torch.int32)) - scores = torch.arange(2 * 4 * 16, dtype=torch.float32).reshape(2, 4, 16) - with mock.patch.object( - torch.ops.trtllm, - "cute_dsl_indexer_topk_decode", - side_effect=_fake_cute_topk, - create=True, - ): - selector.select_requests(scores.unsqueeze(0), normalize_scores=False) - assert tuple(selector.keep.shape) == ( - 1, - selector.selection_rows, - selector.keep_count, - ) - # Scores increase with the token index, so every row keeps the last five - # decode ordinals, rebased by the pinned prompt length. - expected_row = torch.arange(16 - 5, 16, dtype=torch.int32) + 3 - assert torch.equal( - selector.keep, - expected_row.expand(1, selector.selection_rows, -1), - ) - - @pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) @pytest.mark.parametrize("normalize_scores", [False, True]) -def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( +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) @@ -183,26 +111,9 @@ def test_per_head_eager_cuda_matches_cpu_reference_on_selector_stream( ).to(torch.float32) valid_widths = torch.tensor([83, 91], dtype=torch.int32) - reference = _BatchedPerHeadKeepSetSelector( - eviction_mode=eviction_mode, - dense_layers=tuple(range(layers)), - num_query_heads=query_heads, - num_kv_heads=kv_heads, - width=width, - keep_count=keep_count, - dtype=torch.float32, - device=torch.device("cpu"), - max_requests=request_count, + expected = _per_head_keep_oracle( + scores_cpu, valid_widths, keep_count, eviction_mode, normalize_scores ) - reference.valid_widths.copy_(valid_widths) - with mock.patch.object( - torch.ops.trtllm, - "cute_dsl_indexer_topk_decode", - side_effect=_fake_cute_topk, - create=True, - ): - reference.select_requests(scores_cpu, normalize_scores=normalize_scores) - expected = reference.keep.clone() device = torch.device("cuda", torch.cuda.current_device()) stream = torch.cuda.Stream(device=device) @@ -309,7 +220,7 @@ def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_ assert torch.equal(actual_keep[request], expected_keep) -@pytest.mark.parametrize("keep_count,width", [(4096, 4224), (8192, 9216)]) +@pytest.mark.parametrize("keep_count,width", [(4, 64), (4096, 4224), (8192, 9216)]) def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): _require_cute_topk_op() device = torch.device("cuda", torch.cuda.current_device()) @@ -590,18 +501,8 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): decode_widths = [seq_len - prompt_len for prompt_len in prompt_lens] width = max(decode_widths) - # CPU-oracle selection: decode-relative scores per request, rebased to + # Oracle selection: decode-relative scores per request, rebased to # absolute ordinals by each request's own prompt offset. - selector = _BatchedUnionKeepSetSelector( - rows=1, - width=width, - keep_count=decode_keep_count, - dtype=torch.float32, - device=torch.device("cpu"), - max_requests=request_count, - ) - selector.valid_widths.copy_(torch.tensor(decode_widths, dtype=torch.int32)) - selector.set_prompt_offsets(torch.tensor(prompt_lens, dtype=torch.int32)) generator = torch.Generator().manual_seed(11) scores = torch.randint( -8, @@ -610,20 +511,14 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): generator=generator, dtype=torch.int32, ).to(torch.float32) - reference_scores = scores.clone() - with mock.patch.object( - torch.ops.trtllm, - "cute_dsl_indexer_topk_decode", - side_effect=_fake_cute_topk, - create=True, - ): - selector.select_requests(scores, normalize_scores=False) - keep = selector.keep.clone() - for request, (prompt_len, decode_width) in enumerate(zip(prompt_lens, decode_widths)): - expected = torch.sort( - _stable_topk(reference_scores[request, 0], decode_width, decode_keep_count) + prompt_len - ).values - assert torch.equal(keep[request], expected) + keep = torch.stack( + [ + torch.sort( + _stable_topk(scores[request, 0], decode_width, decode_keep_count) + prompt_len + ).values + for request, (prompt_len, decode_width) in enumerate(zip(prompt_lens, decode_widths)) + ] + ) page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) initial_pools = [ diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 685135796fdb..b1b33d682401 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -109,50 +109,6 @@ def _prepared_eviction( ) -def _fake_cute_dsl_topk( - values: torch.Tensor, - seq_lens: torch.Tensor, - output: torch.Tensor, - top_k: int, - next_n: int, -) -> None: - """CPU oracle for the CUDA-only CuTE-DSL selector custom op.""" - assert next_n == 1 - for row in range(int(values.shape[0])): - width = int(seq_lens[row]) - selected = _TORCH_TOPK_ORACLE( - values[row, :width], - top_k, - sorted=False, - ).indices - output[row].copy_(selected.to(torch.int32)) - - -@contextmanager -def _mock_cute_topk_without_fallbacks(): - """Provide the CuTE op while making both retired fallbacks fatal.""" - with ( - mock.patch.object( - torch.ops.trtllm, - "cute_dsl_indexer_topk_decode", - side_effect=_fake_cute_dsl_topk, - create=True, - ) as cute_topk, - mock.patch.object( - torch.ops.trtllm, - "indexer_topk_decode", - side_effect=AssertionError("native IndexerTopK fallback is forbidden"), - create=True, - ), - mock.patch.object( - torch, - "topk", - side_effect=AssertionError("torch.topk production fallback is forbidden"), - ), - ): - yield cute_topk - - def _union_oracle(scores: torch.Tensor, keep_count: int) -> torch.Tensor: """Independent expected-result implementation of union selection.""" combined = scores.max(dim=0).values @@ -1027,32 +983,30 @@ def test_request_finish_clears_compression_state(self): class TestTopKRouting: @pytest.mark.parametrize("keep_count", [4096, 8192]) - def test_cross_request_union_uses_cute_without_fallback(self, keep_count): + def test_cross_request_union_matches_oracle_at_high_keep_counts(self, keep_count): width = keep_count + 64 request_scores = [ _distinct_topk_scores(width), _distinct_topk_scores(width).roll(17, dims=1) + 0.000007, ] expected = [_union_oracle(scores, keep_count) for scores in request_scores] + device = torch.device("cuda", torch.cuda.current_device()) + scores = torch.stack(request_scores).to(device) selector = _BatchedUnionKeepSetSelector( request_scores[0].shape[0], width, keep_count, - dtype=request_scores[0].dtype, - device=request_scores[0].device, + dtype=scores.dtype, + device=device, max_requests=len(request_scores), + input_scores=scores, + normalize_scores=False, ) + selector.select_requests(scores, normalize_scores=False) + selected = selector.keep.cpu() - with _mock_cute_topk_without_fallbacks() as cute_topk: - selector.select_requests( - torch.stack(request_scores), - normalize_scores=False, - ) - selected = selector.keep[: len(request_scores)].clone() - - assert cute_topk.call_count == 1 for actual, expected_keep in zip(selected, expected): - assert torch.equal(actual, expected_keep) + assert torch.equal(actual, expected_keep.to(torch.int32)) class TestFixedScoreMetadata: From b5a0365c47105397def7c53dcc4fa355c70959d1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 00:57:33 -0700 Subject: [PATCH 033/178] [None][chore] Drop the unused checked entry of the top-k finalizer Production reaches the finalizer kernel through the frozen call bound at build time; the checked wrapper and its validation wall had no callers left. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 61 ------------------- 1 file changed, 61 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 2966f680a84b..5f547a6e7cda 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -568,67 +568,6 @@ def _finalize_topk_indices_kernel( ties_seen += tl.sum(tied_i32) -def finalize_topk_indices( - scores: torch.Tensor, - seq_lens: torch.Tensor, - prompt_offsets: torch.Tensor, - provisional_indices: torch.Tensor, - output_indices: torch.Tensor, - keep_count: int, -) -> None: - """Finalize one provisional TopK set without changing the CuTE selector. - - The kernel derives the provisional set's threshold, keeps all strictly - better scores, resolves the remaining boundary ties by lower token index, - and writes increasing physical ordinals (decode index + this row's prompt - offset) directly into ``output_indices``. - """ - keep_count = int(keep_count) - if not scores.is_cuda: - raise ValueError("deterministic TopK finalization requires CUDA scores") - if scores.ndim != 2 or scores.dtype != torch.float32 or not scores.is_contiguous(): - raise ValueError("TopK finalization requires contiguous two-dimensional FP32 scores") - rows, width = scores.shape - if rows <= 0 or not 1 <= keep_count <= width: - raise ValueError("TopK finalization requires valid rows and keep count") - for name, tensor in (("sequence lengths", seq_lens), ("prompt offsets", prompt_offsets)): - if ( - tensor.shape != (rows,) - or tensor.dtype != torch.int32 - or tensor.device != scores.device - or not tensor.is_contiguous() - ): - raise ValueError(f"TopK finalization {name} do not match the score rows") - if ( - provisional_indices.shape != (rows, keep_count) - or provisional_indices.dtype != torch.int32 - or provisional_indices.device != scores.device - or not provisional_indices.is_contiguous() - ): - raise ValueError("provisional TopK indices do not match the requested selection") - if ( - output_indices.ndim != 2 - or output_indices.shape[0] != rows - or output_indices.shape[1] < keep_count - or output_indices.dtype != torch.int32 - or output_indices.device != scores.device - or not output_indices.is_contiguous() - ): - raise ValueError("TopK output does not fit the requested physical indices") - _finalize_topk_indices_kernel[(rows,)]( - scores, - seq_lens, - prompt_offsets, - provisional_indices, - output_indices, - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=output_indices.shape[1], - BLOCK=256, - num_warps=4, - ) - - @triton.jit def _tri_score_perhead_kernel( pool_anchor_ptr, # typed pool pointer; used ONLY to infer the element type From 32c8bb9d918d5ca133aa82ef8c53b59982512e82 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 01:00:56 -0700 Subject: [PATCH 034/178] [None][chore] Order the kernels module along the eviction pipeline Pure move: launch infrastructure, then scoring, selection, and compaction sections in execution order, with section banners. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 1191 +++++++++-------- 1 file changed, 603 insertions(+), 588 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 5f547a6e7cda..9c30ff397cc0 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -85,6 +85,11 @@ def __call__(self, *call_tensors: torch.Tensor) -> None: ) +# --------------------------------------------------------------------------- # +# Scoring: trig-score every cached token across all dense layers. # +# --------------------------------------------------------------------------- # + + @triton.jit def _prepare_mean_phase_kernel( round_starts, @@ -157,614 +162,202 @@ def prepare_mean_phase( @triton.jit -def _score_row_stats_kernel( - scores, - valid_widths, - row_mean, - row_inv_std, - ROWS: tl.constexpr, - WIDTH: tl.constexpr, - BLOCK: tl.constexpr, +def _tri_score_perhead_kernel( + pool_anchor_ptr, # typed pool pointer; used ONLY to infer the element type + # for the int->pointer cast below (its data is never read through it). + layer_base_addrs, # [num_layers] int64: ABSOLUTE device address of each + # scored layer's HND base. Layers do NOT need to share one storage: + # each segment casts its own layer's address back to a typed pointer, so + # all address arithmetic stays inside that layer's own allocation. + block_offsets_ptr, # Native V2 [pool, request, K/V, block] int32 offsets. + seg_page_off, # [nseg] int64: offset of this segment's page table into + # block_offsets_ptr. + # per-SEGMENT metadata (seg = req_slot*L_scored + layer_slot), idx by pid(0): + seg_req_id, # [nseg] int32: request slot (round_start / mean phase lookup) + seg_layer_id, # [nseg] int32: ABSOLUTE layer id (indexes layer_base_addrs + calib) + req_seq_len, # [num_requests] int32 + req_valid_width_out, # [num_requests] int32: decode-only length for selection + req_round_start, # [num_requests] int32 logical token position + req_token_start, # [num_requests] int32: pinned prompt length; scoring + # starts at this decode-region origin, so prompt lengths may differ + # across the cohort. + # per-LAYER calibration, [L,H,F] flattened layer-major: + q_real_ptr, # [L*H*F] fp32 + q_imag_ptr, # [L*H*F] fp32 + mlr_coef_ptr, # [L*H*F] fp32 + # per-REQUEST offset-collapsed phase ('mean' path), [num_requests,F] flattened: + mean_cos_ptr, # [num_requests*F] fp32 + mean_sin_ptr, # [num_requests*F] fp32 + # shared freq vectors: + freq_scale_sq_ptr, # [F] fp32 + omega_ptr, # [F] fp32 ('max' path only) + offsets_ptr, # [O] fp32 ('max' path only) + out_ptr, # [request, layer, query_head, decode_token] fp32 + output_width, + # scalars uniform across the batch: + num_layers, + num_q_heads, + num_kv_heads, + num_freqs, # F = head_dim // 2 + head_dim, + tokens_per_block, + kv_factor, + num_offsets, # O ('max' path only) + # per-layer HND element strides (uniform across scored layers): + s_page, + s_kv_head, + s_slot, + s_dim, + USE_MAX: tl.constexpr, + T_BLOCK: tl.constexpr, + F_BLOCK: tl.constexpr, ): - """Compute one valid-prefix mean and inverse standard deviation per score row.""" - flat_row = tl.program_id(0) - request = flat_row // ROWS - valid_width = tl.load(valid_widths + request) - score_row = scores + flat_row * WIDTH - lane = tl.arange(0, BLOCK) - score_sum = 0.0 - for start in tl.static_range(0, WIDTH, BLOCK): - token = start + lane - valid = token < valid_width - value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) - score_sum += tl.sum(value, axis=0) - mean = score_sum / valid_width - square_sum = 0.0 - for start in tl.static_range(0, WIDTH, BLOCK): - token = start + lane - valid = token < valid_width - 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 / valid_width) - tl.store(row_mean + flat_row, mean) - tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) + # Token tiles ride the fastest grid axis: adjacent programs then walk + # consecutive K pages of one (request, layer, head) and reuse its + # calibration/phase rows in L2 (~2% faster than segment-major order). + seg = tl.program_id(1) + t_blk = tl.program_id(0) + # KV heads are grid-parallel (axis 2): iterations of the former kv_head + # loop shared NO data (each KV head reads its own K and writes its own + # output rows), so hoisting it onto the grid multiplies parallelism with + # zero extra HBM traffic. The q-in-group loop below stays inside the + # program because it REUSES this head's K from registers (GQA dedup). + kv_head = tl.program_id(2) + req_id = tl.load(seg_req_id + seg) + seq_len = tl.load(req_seq_len + req_id) + token_start = tl.load(req_token_start + req_id) + if (seg % num_layers == 0) & (t_blk == 0) & (kv_head == 0): + tl.store(req_valid_width_out + req_id, seq_len - token_start) + # Derive the ragged launch bound in the score program instead of staging + # one replicated length and block count for every request/layer segment. + n_tblk = (seq_len - token_start + T_BLOCK - 1) // T_BLOCK + if t_blk >= n_tblk: + return -@triton.jit -def _score_union_kernel( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - ROWS: tl.constexpr, - WIDTH: tl.constexpr, - NORMALIZE: tl.constexpr, - BLOCK: tl.constexpr, -): - """Normalize score rows and reduce them directly to one request-level union.""" - request = tl.program_id(0) - token = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - valid_width = tl.load(valid_widths + request) - valid_token = token < valid_width - union_max = tl.full((BLOCK,), -float("inf"), tl.float32) - for row in tl.range(0, ROWS): - flat_row = request * ROWS + row - value = tl.load( - scores + flat_row * WIDTH + 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")) - union_max = tl.maximum(union_max, value) - tl.store(combined + request * WIDTH + token, union_max, mask=token < WIDTH) + layer_id = tl.load(seg_layer_id + seg) + rstart = tl.load(req_round_start + req_id) + # This segment's layer base: an absolute address cast back to a pool-typed + # pointer. TRT-LLM V2 exposes every layer as its own TensorWrapper storage, + # so "element offset relative to one shared storage" does not exist; the + # per-layer absolute address is the same device-pointer-array pattern the + # C++ backends use (KVBlockArray / grouped-GEMM pointer arrays). + layer_ptr = tl.load(layer_base_addrs + layer_id).to( + tl.pointer_type(pool_anchor_ptr.dtype.element_ty) + ) + page_off = tl.load(seg_page_off + seg) + f = tl.arange(0, F_BLOCK) + f_mask = f < num_freqs + f64 = f.to(tl.int64) -def prepare_union_scores( - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - combined: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, -) -> None: - """Mask, normalize, and union-reduce score rows in two or three launches.""" - request_count = int(request_count) - if not scores.is_cuda or scores.ndim != 3 or scores.dtype != torch.float32: - raise ValueError("union score preparation requires contiguous CUDA FP32 rows") - if not scores.is_contiguous() or request_count != scores.shape[0]: - raise ValueError("union score preparation request geometry does not match") - _, rows, width = scores.shape - if ( - valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.numel() < request_count * rows - or row_inv_std.shape != row_mean.shape - or combined.shape != (request_count, width) - ): - raise ValueError("union score preparation buffers do not match") - stats_block = 256 - if normalize_scores: - _score_row_stats_kernel[(request_count * rows,)]( - scores, - valid_widths, - row_mean, - row_inv_std, - ROWS=rows, - WIDTH=width, - BLOCK=stats_block, - num_warps=4, - ) - union_block = 32 - _score_union_kernel[(request_count, triton.cdiv(width, union_block))]( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - ROWS=rows, - WIDTH=width, - NORMALIZE=normalize_scores, - BLOCK=union_block, - num_warps=1, + # ---- token tile of THIS segment ---- + t = t_blk * T_BLOCK + tl.arange(0, T_BLOCK) + absolute_t = t + token_start + t_mask = absolute_t < seq_len + blk_in_seq = absolute_t // tokens_per_block + slot = (absolute_t % tokens_per_block).to(tl.int64) + # The native attention page-table copy encodes K offsets in units of the + # underlying K/V role pages. Convert that value to the HND pool page inline + # instead of materializing a second page table before scoring. + encoded_page = tl.load( + block_offsets_ptr + page_off + blk_in_seq, + mask=t_mask, + other=0, ) + phys_page = (encoded_page // kv_factor).to(tl.int64) + # element offset into THIS layer's pool for (page, KEY=0, *, slot). + # KEY half is kv_factor index 0 -> its stride term is 0 (matches reference). + tok_base = phys_page * s_page + slot * s_slot # [T_BLOCK] int64 -@triton.jit -def _score_per_head_reduce_kernel( - scores, - valid_widths, - row_mean, - row_inv_std, - selection_scores, - selection_seq_lens, - NUM_LAYERS: tl.constexpr, - NUM_QUERY_HEADS: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - QUERY_GROUP_SIZE: tl.constexpr, - SELECTION_ROWS: tl.constexpr, - WIDTH: tl.constexpr, - PER_LAYER: tl.constexpr, - NORMALIZE: tl.constexpr, - BLOCK: tl.constexpr, -): - """Reduce query-head score rows into one selector row per KV-head domain.""" - 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) - valid_width = tl.load(valid_widths + request) - valid_token = token < valid_width - - if token_block == 0: - tl.store( - selection_seq_lens + request * SELECTION_ROWS + selection_row, - valid_width, - ) + # per-request 'mean'-path phase + shared freq scale. + mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) + msin = tl.load(mean_sin_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) + fss = tl.load(freq_scale_sq_ptr + f, mask=f_mask, other=0.0) - 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_QUERY_HEADS + query_head - value = tl.load( - scores + flat_row * WIDTH + 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_QUERY_HEADS + query_head - value = tl.load( - scores + flat_row * WIDTH + 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 + # ---- PER-HEAD (position + mlr), GQA-deduped, NO head reduction ---- + # This program scores ONE KV head's token tile for the group_size q-heads + # that share it. K (and |K|) is loaded ONCE and reused across the group; + # h = kv_head*group_size + qg keeps query-head order 0..num_q_heads-1, so + # every head's math is bit-for-bit identical to the looped variant. + group_size = num_q_heads // num_kv_heads + load_mask = t_mask[:, None] & f_mask[None, :] + off_re = f64[None, :] * s_dim + off_im = (num_freqs + f64[None, :]) * s_dim - output = (request * SELECTION_ROWS + selection_row) * WIDTH + token - tl.store(selection_scores + output, reduced, mask=token < WIDTH) + base = tok_base + kv_head.to(tl.int64) * s_kv_head # [T_BLOCK] + # paged K loaded ONCE for this KV head (shared by group_size q-heads). + k_re = tl.load(layer_ptr + base[:, None] + off_re, mask=load_mask, other=0.0).to(tl.float32) + k_im = tl.load(layer_ptr + base[:, None] + off_im, mask=load_mask, other=0.0).to(tl.float32) + kmag = tl.sqrt(k_re * k_re + k_im * k_im) # once per KV head + qg = 0 + while qg < group_size: + h = kv_head * group_size + qg + calib_off = (layer_id.to(tl.int64) * num_q_heads + h) * num_freqs + qre = tl.load(q_real_ptr + calib_off + f, mask=f_mask, other=0.0) + qim = tl.load(q_imag_ptr + calib_off + f, mask=f_mask, other=0.0) + mlrc = tl.load(mlr_coef_ptr + calib_off + f, mask=f_mask, other=0.0) -def prepare_per_head_scores( - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - selection_scores: torch.Tensor, - selection_seq_lens: torch.Tensor, - request_count: int, - *, - num_kv_heads: int, - per_layer: bool, - normalize_scores: bool, -) -> None: - """Normalize and reduce score rows for either per-head eviction mode.""" - request_count = int(request_count) - num_kv_heads = int(num_kv_heads) - if not scores.is_cuda or scores.ndim != 4 or scores.dtype != torch.float32: - raise ValueError("per-head score preparation requires CUDA FP32 rows") - if not scores.is_contiguous() or request_count != scores.shape[0]: - raise ValueError("per-head score preparation request geometry does not match") - _, num_layers, num_query_heads, width = scores.shape - if num_kv_heads <= 0 or num_query_heads % num_kv_heads: - raise ValueError("per-head score preparation requires valid GQA geometry") - selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - if ( - valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.numel() < request_count * num_layers * num_query_heads - or row_inv_std.shape != row_mean.shape - or selection_scores.shape != (request_count, selection_rows, width) - or selection_scores.dtype != torch.float32 - or selection_scores.device != scores.device - or selection_seq_lens.shape != (request_count, selection_rows) - or selection_seq_lens.dtype != torch.int32 - or selection_seq_lens.device != scores.device - ): - raise ValueError("per-head score preparation buffers do not match") + # complex product Q . conj(K) -- the trig importance score. + prod_real = qre[None, :] * k_re + qim[None, :] * k_im + prod_imag = qim[None, :] * k_re - qre[None, :] * k_im - stats_block = 256 - rows = num_layers * num_query_heads - if normalize_scores: - _score_row_stats_kernel[(request_count * rows,)]( - scores, - valid_widths, - row_mean, - row_inv_std, - ROWS=rows, - WIDTH=width, - BLOCK=stats_block, - num_warps=4, - ) - reduction_block = 256 - _score_per_head_reduce_kernel[ - (request_count, selection_rows, triton.cdiv(width, reduction_block)) - ]( - scores, - valid_widths, - row_mean, - row_inv_std, - selection_scores, - selection_seq_lens, - NUM_LAYERS=num_layers, - NUM_QUERY_HEADS=num_query_heads, - NUM_KV_HEADS=num_kv_heads, - QUERY_GROUP_SIZE=num_query_heads // num_kv_heads, - SELECTION_ROWS=selection_rows, - WIDTH=width, - PER_LAYER=per_layer, - NORMALIZE=normalize_scores, - BLOCK=reduction_block, - num_warps=4, - ) + if USE_MAX: + # max over O offsets does NOT commute through the freq-sum; + # explicit O loop reducing max over the per-offset F-sum. + score = tl.full((T_BLOCK,), -float("inf"), tl.float32) + o = 0 + while o < num_offsets: + off = tl.load(offsets_ptr + o) + om = tl.load(omega_ptr + f, mask=f_mask, other=0.0) + phase = (rstart + off) * om + cphase = tl.cos(phase) + sphase = tl.sin(phase) + per_f = fss[None, :] * (prod_real * cphase[None, :] - prod_imag * sphase[None, :]) + offset_score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) + score = tl.maximum(score, offset_score) + o += 1 + else: + # 'mean': offset loop collapsed into mean_cos/mean_sin. + per_f = fss[None, :] * (prod_real * mcos[None, :] - prod_imag * msin[None, :]) + score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) + # position-INDEPENDENT MLR term (reuses the per-KV-head |K|). + mlr_f = kmag * mlrc[None, :] * fss[None, :] + mlr = tl.sum(tl.where(f_mask[None, :], mlr_f, 0.0), axis=1) -@triton.jit -def _pack_compaction_sources_kernel( - selected_indices, - valid_seq_lens, - dense_offsets, - dense_indices, - swa_offsets, - swa_indices, - DENSE_TOTAL: tl.constexpr, - SWA_TOTAL: tl.constexpr, - SELECTION_ROWS: tl.constexpr, - SELECTION_STRIDE: tl.constexpr, - KEEP_COUNT: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - SWA_WINDOW: tl.constexpr, - UNION: tl.constexpr, - PER_LAYER: tl.constexpr, - HAS_SWA: tl.constexpr, - BLOCK: tl.constexpr, -): - """Pack selected decode ordinals and protected tails for the C++ updater.""" - request = tl.program_id(0) - domain = tl.program_id(1) - move = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + # Segments are request-major then layer-major. Write the decode-only + # score directly in the selector's [request, layer, head, token] layout. + out_offset = (seg.to(tl.int64) * num_q_heads + h) * output_width + t + tl.store(out_ptr + out_offset, score + mlr, mask=t_mask) + qg += 1 - dense_begin = tl.load(dense_offsets + request) - dense_end = tl.load(dense_offsets + request + 1) - dense_count = dense_end - dense_begin - seq_len = tl.load(valid_seq_lens + request) - if UNION: - selection_domain = 0 - else: - selection_domain = domain - # Selection rows carry decode-only kept ordinals (already absolute), so - # rows are prompt-length independent and one cohort may mix prompt sizes. - selection_row = request * SELECTION_ROWS + selection_domain - selected = tl.load( - selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + move, - mask=move < KEEP_COUNT, - other=0, +def _launch_tri_score_perhead( + grid: tuple, + pointer_args: tuple, + geometry_args: tuple, + *, + score_aggregation: str, + token_block: int, + num_freqs: int, +) -> None: + """Launch the shared score ABI for eager and fixed metadata owners.""" + if score_aggregation not in ("mean", "max"): + raise ValueError(f"unsupported score aggregation: {score_aggregation}") + _tri_score_perhead_kernel[grid]( + *pointer_args, + *geometry_args, + USE_MAX=(score_aggregation == "max"), + T_BLOCK=token_block, + F_BLOCK=triton.next_power_of_2(num_freqs), ) - dense_source = tl.where(move < KEEP_COUNT, selected, seq_len + move - KEEP_COUNT) - dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move - tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) - if HAS_SWA: - # Per-layer selection has one dense domain per (layer, head). SWA uses - # one shared source row per head, so only the first layer writes it. - if PER_LAYER: - write_swa = domain < NUM_KV_HEADS - else: - write_swa = move >= 0 - swa_begin = tl.load(swa_offsets + request) - swa_end = tl.load(swa_offsets + request + 1) - swa_count = swa_end - swa_begin - head = domain % NUM_KV_HEADS - swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move - swa_source = seq_len - SWA_WINDOW + move - tl.store( - swa_indices + swa_output, - swa_source, - mask=write_swa & (move < swa_count), - ) - -@triton.jit -def _finalize_topk_indices_kernel( - scores, - seq_lens, - prompt_offsets, - provisional_indices, - output_indices, - WIDTH: tl.constexpr, - KEEP_COUNT: tl.constexpr, - OUTPUT_WIDTH: tl.constexpr, - BLOCK: tl.constexpr, -): - """Resolve boundary ties and emit increasing physical token indices.""" - row = tl.program_id(0) - row_scores = scores + row * WIDTH - row_selected = provisional_indices + row * KEEP_COUNT - row_output = output_indices + row * OUTPUT_WIDTH - # Scores are decode-relative; this row's pinned prompt length rebases the - # emitted ordinals to absolute positions (per row, so one launch may mix - # prompt lengths). - prompt_len = tl.load(prompt_offsets + row) - - 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, - ) - selected_score = tl.load( - row_scores + token_index, - mask=selected_mask, - other=float("inf"), - ).to(tl.float32) - threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) - - seq_len = tl.load(seq_lens + 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 < seq_len) - 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 < seq_len) - 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_len, - mask=selected, - ) - output_count += tl.sum(selected_i32) - ties_seen += tl.sum(tied_i32) - - -@triton.jit -def _tri_score_perhead_kernel( - pool_anchor_ptr, # typed pool pointer; used ONLY to infer the element type - # for the int->pointer cast below (its data is never read through it). - layer_base_addrs, # [num_layers] int64: ABSOLUTE device address of each - # scored layer's HND base. Layers do NOT need to share one storage: - # each segment casts its own layer's address back to a typed pointer, so - # all address arithmetic stays inside that layer's own allocation. - block_offsets_ptr, # Native V2 [pool, request, K/V, block] int32 offsets. - seg_page_off, # [nseg] int64: offset of this segment's page table into - # block_offsets_ptr. - # per-SEGMENT metadata (seg = req_slot*L_scored + layer_slot), idx by pid(0): - seg_req_id, # [nseg] int32: request slot (round_start / mean phase lookup) - seg_layer_id, # [nseg] int32: ABSOLUTE layer id (indexes layer_base_addrs + calib) - req_seq_len, # [num_requests] int32 - req_valid_width_out, # [num_requests] int32: decode-only length for selection - req_round_start, # [num_requests] int32 logical token position - req_token_start, # [num_requests] int32: pinned prompt length; scoring - # starts at this decode-region origin, so prompt lengths may differ - # across the cohort. - # per-LAYER calibration, [L,H,F] flattened layer-major: - q_real_ptr, # [L*H*F] fp32 - q_imag_ptr, # [L*H*F] fp32 - mlr_coef_ptr, # [L*H*F] fp32 - # per-REQUEST offset-collapsed phase ('mean' path), [num_requests,F] flattened: - mean_cos_ptr, # [num_requests*F] fp32 - mean_sin_ptr, # [num_requests*F] fp32 - # shared freq vectors: - freq_scale_sq_ptr, # [F] fp32 - omega_ptr, # [F] fp32 ('max' path only) - offsets_ptr, # [O] fp32 ('max' path only) - out_ptr, # [request, layer, query_head, decode_token] fp32 - output_width, - # scalars uniform across the batch: - num_layers, - num_q_heads, - num_kv_heads, - num_freqs, # F = head_dim // 2 - head_dim, - tokens_per_block, - kv_factor, - num_offsets, # O ('max' path only) - # per-layer HND element strides (uniform across scored layers): - s_page, - s_kv_head, - s_slot, - s_dim, - USE_MAX: tl.constexpr, - T_BLOCK: tl.constexpr, - F_BLOCK: tl.constexpr, -): - # Token tiles ride the fastest grid axis: adjacent programs then walk - # consecutive K pages of one (request, layer, head) and reuse its - # calibration/phase rows in L2 (~2% faster than segment-major order). - seg = tl.program_id(1) - t_blk = tl.program_id(0) - # KV heads are grid-parallel (axis 2): iterations of the former kv_head - # loop shared NO data (each KV head reads its own K and writes its own - # output rows), so hoisting it onto the grid multiplies parallelism with - # zero extra HBM traffic. The q-in-group loop below stays inside the - # program because it REUSES this head's K from registers (GQA dedup). - kv_head = tl.program_id(2) - - req_id = tl.load(seg_req_id + seg) - seq_len = tl.load(req_seq_len + req_id) - token_start = tl.load(req_token_start + req_id) - if (seg % num_layers == 0) & (t_blk == 0) & (kv_head == 0): - tl.store(req_valid_width_out + req_id, seq_len - token_start) - # Derive the ragged launch bound in the score program instead of staging - # one replicated length and block count for every request/layer segment. - n_tblk = (seq_len - token_start + T_BLOCK - 1) // T_BLOCK - if t_blk >= n_tblk: - return - - layer_id = tl.load(seg_layer_id + seg) - rstart = tl.load(req_round_start + req_id) - # This segment's layer base: an absolute address cast back to a pool-typed - # pointer. TRT-LLM V2 exposes every layer as its own TensorWrapper storage, - # so "element offset relative to one shared storage" does not exist; the - # per-layer absolute address is the same device-pointer-array pattern the - # C++ backends use (KVBlockArray / grouped-GEMM pointer arrays). - layer_ptr = tl.load(layer_base_addrs + layer_id).to( - tl.pointer_type(pool_anchor_ptr.dtype.element_ty) - ) - page_off = tl.load(seg_page_off + seg) - - f = tl.arange(0, F_BLOCK) - f_mask = f < num_freqs - f64 = f.to(tl.int64) - - # ---- token tile of THIS segment ---- - t = t_blk * T_BLOCK + tl.arange(0, T_BLOCK) - absolute_t = t + token_start - t_mask = absolute_t < seq_len - blk_in_seq = absolute_t // tokens_per_block - slot = (absolute_t % tokens_per_block).to(tl.int64) - # The native attention page-table copy encodes K offsets in units of the - # underlying K/V role pages. Convert that value to the HND pool page inline - # instead of materializing a second page table before scoring. - encoded_page = tl.load( - block_offsets_ptr + page_off + blk_in_seq, - mask=t_mask, - other=0, - ) - phys_page = (encoded_page // kv_factor).to(tl.int64) - - # element offset into THIS layer's pool for (page, KEY=0, *, slot). - # KEY half is kv_factor index 0 -> its stride term is 0 (matches reference). - tok_base = phys_page * s_page + slot * s_slot # [T_BLOCK] int64 - - # per-request 'mean'-path phase + shared freq scale. - mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) - msin = tl.load(mean_sin_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) - fss = tl.load(freq_scale_sq_ptr + f, mask=f_mask, other=0.0) - - # ---- PER-HEAD (position + mlr), GQA-deduped, NO head reduction ---- - # This program scores ONE KV head's token tile for the group_size q-heads - # that share it. K (and |K|) is loaded ONCE and reused across the group; - # h = kv_head*group_size + qg keeps query-head order 0..num_q_heads-1, so - # every head's math is bit-for-bit identical to the looped variant. - group_size = num_q_heads // num_kv_heads - load_mask = t_mask[:, None] & f_mask[None, :] - off_re = f64[None, :] * s_dim - off_im = (num_freqs + f64[None, :]) * s_dim - - base = tok_base + kv_head.to(tl.int64) * s_kv_head # [T_BLOCK] - # paged K loaded ONCE for this KV head (shared by group_size q-heads). - k_re = tl.load(layer_ptr + base[:, None] + off_re, mask=load_mask, other=0.0).to(tl.float32) - k_im = tl.load(layer_ptr + base[:, None] + off_im, mask=load_mask, other=0.0).to(tl.float32) - kmag = tl.sqrt(k_re * k_re + k_im * k_im) # once per KV head - - qg = 0 - while qg < group_size: - h = kv_head * group_size + qg - calib_off = (layer_id.to(tl.int64) * num_q_heads + h) * num_freqs - qre = tl.load(q_real_ptr + calib_off + f, mask=f_mask, other=0.0) - qim = tl.load(q_imag_ptr + calib_off + f, mask=f_mask, other=0.0) - mlrc = tl.load(mlr_coef_ptr + calib_off + f, mask=f_mask, other=0.0) - - # complex product Q . conj(K) -- the trig importance score. - prod_real = qre[None, :] * k_re + qim[None, :] * k_im - prod_imag = qim[None, :] * k_re - qre[None, :] * k_im - - if USE_MAX: - # max over O offsets does NOT commute through the freq-sum; - # explicit O loop reducing max over the per-offset F-sum. - score = tl.full((T_BLOCK,), -float("inf"), tl.float32) - o = 0 - while o < num_offsets: - off = tl.load(offsets_ptr + o) - om = tl.load(omega_ptr + f, mask=f_mask, other=0.0) - phase = (rstart + off) * om - cphase = tl.cos(phase) - sphase = tl.sin(phase) - per_f = fss[None, :] * (prod_real * cphase[None, :] - prod_imag * sphase[None, :]) - offset_score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) - score = tl.maximum(score, offset_score) - o += 1 - else: - # 'mean': offset loop collapsed into mean_cos/mean_sin. - per_f = fss[None, :] * (prod_real * mcos[None, :] - prod_imag * msin[None, :]) - score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) - - # position-INDEPENDENT MLR term (reuses the per-KV-head |K|). - mlr_f = kmag * mlrc[None, :] * fss[None, :] - mlr = tl.sum(tl.where(f_mask[None, :], mlr_f, 0.0), axis=1) - - # Segments are request-major then layer-major. Write the decode-only - # score directly in the selector's [request, layer, head, token] layout. - out_offset = (seg.to(tl.int64) * num_q_heads + h) * output_width + t - tl.store(out_ptr + out_offset, score + mlr, mask=t_mask) - qg += 1 - - -def _launch_tri_score_perhead( - grid: tuple, - pointer_args: tuple, - geometry_args: tuple, - *, - score_aggregation: str, - token_block: int, - num_freqs: int, -) -> None: - """Launch the shared score ABI for eager and fixed metadata owners.""" - if score_aggregation not in ("mean", "max"): - raise ValueError(f"unsupported score aggregation: {score_aggregation}") - _tri_score_perhead_kernel[grid]( - *pointer_args, - *geometry_args, - USE_MAX=(score_aggregation == "max"), - T_BLOCK=token_block, - F_BLOCK=triton.next_power_of_2(num_freqs), - ) - - -class _FixedScoreGroup: - """Persistent score metadata/output for one fixed geometry. +class _FixedScoreGroup: + """Persistent score metadata/output for one fixed geometry. Since the per-layer absolute-address ABI, ONE group can span dense layers living in DISTINCT storages with DISTINCT block tables. ``block_offsets`` @@ -949,3 +542,425 @@ def launch( num_freqs=self.num_freqs, ) return output + + +# --------------------------------------------------------------------------- # +# Selection: combine scores per mode, then finalize the top-k set. # +# --------------------------------------------------------------------------- # + + +@triton.jit +def _score_row_stats_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + """Compute one valid-prefix mean and inverse standard deviation per score row.""" + flat_row = tl.program_id(0) + request = flat_row // ROWS + valid_width = tl.load(valid_widths + request) + score_row = scores + flat_row * WIDTH + lane = tl.arange(0, BLOCK) + score_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < valid_width + value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) + score_sum += tl.sum(value, axis=0) + mean = score_sum / valid_width + square_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < valid_width + 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 / valid_width) + tl.store(row_mean + flat_row, mean) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) + + +@triton.jit +def _score_union_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, + NORMALIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Normalize score rows and reduce them directly to one request-level union.""" + request = tl.program_id(0) + token = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + valid_width = tl.load(valid_widths + request) + valid_token = token < valid_width + union_max = tl.full((BLOCK,), -float("inf"), tl.float32) + for row in tl.range(0, ROWS): + flat_row = request * ROWS + row + value = tl.load( + scores + flat_row * WIDTH + 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")) + union_max = tl.maximum(union_max, value) + tl.store(combined + request * WIDTH + token, union_max, mask=token < WIDTH) + + +def prepare_union_scores( + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + combined: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, +) -> None: + """Mask, normalize, and union-reduce score rows in two or three launches.""" + request_count = int(request_count) + if not scores.is_cuda or scores.ndim != 3 or scores.dtype != torch.float32: + raise ValueError("union score preparation requires contiguous CUDA FP32 rows") + if not scores.is_contiguous() or request_count != scores.shape[0]: + raise ValueError("union score preparation request geometry does not match") + _, rows, width = scores.shape + if ( + valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.numel() < request_count * rows + or row_inv_std.shape != row_mean.shape + or combined.shape != (request_count, width) + ): + raise ValueError("union score preparation buffers do not match") + stats_block = 256 + if normalize_scores: + _score_row_stats_kernel[(request_count * rows,)]( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS=rows, + WIDTH=width, + BLOCK=stats_block, + num_warps=4, + ) + union_block = 32 + _score_union_kernel[(request_count, triton.cdiv(width, union_block))]( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + ROWS=rows, + WIDTH=width, + NORMALIZE=normalize_scores, + BLOCK=union_block, + num_warps=1, + ) + + +@triton.jit +def _score_per_head_reduce_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + selection_scores, + selection_seq_lens, + NUM_LAYERS: tl.constexpr, + NUM_QUERY_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + QUERY_GROUP_SIZE: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + WIDTH: tl.constexpr, + PER_LAYER: tl.constexpr, + NORMALIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Reduce query-head score rows into one selector row per KV-head domain.""" + 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) + valid_width = tl.load(valid_widths + request) + valid_token = token < valid_width + + if token_block == 0: + tl.store( + selection_seq_lens + request * SELECTION_ROWS + selection_row, + valid_width, + ) + + 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_QUERY_HEADS + query_head + value = tl.load( + scores + flat_row * WIDTH + 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_QUERY_HEADS + query_head + value = tl.load( + scores + flat_row * WIDTH + 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 prepare_per_head_scores( + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + selection_scores: torch.Tensor, + selection_seq_lens: torch.Tensor, + request_count: int, + *, + num_kv_heads: int, + per_layer: bool, + normalize_scores: bool, +) -> None: + """Normalize and reduce score rows for either per-head eviction mode.""" + request_count = int(request_count) + num_kv_heads = int(num_kv_heads) + if not scores.is_cuda or scores.ndim != 4 or scores.dtype != torch.float32: + raise ValueError("per-head score preparation requires CUDA FP32 rows") + if not scores.is_contiguous() or request_count != scores.shape[0]: + raise ValueError("per-head score preparation request geometry does not match") + _, num_layers, num_query_heads, width = scores.shape + if num_kv_heads <= 0 or num_query_heads % num_kv_heads: + raise ValueError("per-head score preparation requires valid GQA geometry") + selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads + if ( + valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.numel() < request_count * num_layers * num_query_heads + or row_inv_std.shape != row_mean.shape + or selection_scores.shape != (request_count, selection_rows, width) + or selection_scores.dtype != torch.float32 + or selection_scores.device != scores.device + or selection_seq_lens.shape != (request_count, selection_rows) + or selection_seq_lens.dtype != torch.int32 + or selection_seq_lens.device != scores.device + ): + raise ValueError("per-head score preparation buffers do not match") + + stats_block = 256 + rows = num_layers * num_query_heads + if normalize_scores: + _score_row_stats_kernel[(request_count * rows,)]( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS=rows, + WIDTH=width, + BLOCK=stats_block, + num_warps=4, + ) + reduction_block = 256 + _score_per_head_reduce_kernel[ + (request_count, selection_rows, triton.cdiv(width, reduction_block)) + ]( + scores, + valid_widths, + row_mean, + row_inv_std, + selection_scores, + selection_seq_lens, + NUM_LAYERS=num_layers, + NUM_QUERY_HEADS=num_query_heads, + NUM_KV_HEADS=num_kv_heads, + QUERY_GROUP_SIZE=num_query_heads // num_kv_heads, + SELECTION_ROWS=selection_rows, + WIDTH=width, + PER_LAYER=per_layer, + NORMALIZE=normalize_scores, + BLOCK=reduction_block, + num_warps=4, + ) + + +@triton.jit +def _finalize_topk_indices_kernel( + scores, + seq_lens, + prompt_offsets, + provisional_indices, + output_indices, + WIDTH: tl.constexpr, + KEEP_COUNT: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + """Resolve boundary ties and emit increasing physical token indices.""" + row = tl.program_id(0) + row_scores = scores + row * WIDTH + row_selected = provisional_indices + row * KEEP_COUNT + row_output = output_indices + row * OUTPUT_WIDTH + # Scores are decode-relative; this row's pinned prompt length rebases the + # emitted ordinals to absolute positions (per row, so one launch may mix + # prompt lengths). + prompt_len = tl.load(prompt_offsets + row) + + 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, + ) + selected_score = tl.load( + row_scores + token_index, + mask=selected_mask, + other=float("inf"), + ).to(tl.float32) + threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) + + seq_len = tl.load(seq_lens + 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 < seq_len) + 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 < seq_len) + 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_len, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) + + +# --------------------------------------------------------------------------- # +# Compaction: pack the kept ordinals into per-request move indices. # +# --------------------------------------------------------------------------- # + + +@triton.jit +def _pack_compaction_sources_kernel( + selected_indices, + valid_seq_lens, + dense_offsets, + dense_indices, + swa_offsets, + swa_indices, + DENSE_TOTAL: tl.constexpr, + SWA_TOTAL: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + SELECTION_STRIDE: tl.constexpr, + KEEP_COUNT: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + SWA_WINDOW: tl.constexpr, + UNION: tl.constexpr, + PER_LAYER: tl.constexpr, + HAS_SWA: tl.constexpr, + BLOCK: tl.constexpr, +): + """Pack selected decode ordinals and protected tails for the C++ updater.""" + request = tl.program_id(0) + domain = tl.program_id(1) + move = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + + dense_begin = tl.load(dense_offsets + request) + dense_end = tl.load(dense_offsets + request + 1) + dense_count = dense_end - dense_begin + seq_len = tl.load(valid_seq_lens + request) + + if UNION: + selection_domain = 0 + else: + selection_domain = domain + # Selection rows carry decode-only kept ordinals (already absolute), so + # rows are prompt-length independent and one cohort may mix prompt sizes. + selection_row = request * SELECTION_ROWS + selection_domain + selected = tl.load( + selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + move, + mask=move < KEEP_COUNT, + other=0, + ) + dense_source = tl.where(move < KEEP_COUNT, selected, seq_len + move - KEEP_COUNT) + dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) + + if HAS_SWA: + # Per-layer selection has one dense domain per (layer, head). SWA uses + # one shared source row per head, so only the first layer writes it. + if PER_LAYER: + write_swa = domain < NUM_KV_HEADS + else: + write_swa = move >= 0 + swa_begin = tl.load(swa_offsets + request) + swa_end = tl.load(swa_offsets + request + 1) + swa_count = swa_end - swa_begin + head = domain % NUM_KV_HEADS + swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move + swa_source = seq_len - SWA_WINDOW + move + tl.store( + swa_indices + swa_output, + swa_source, + mask=write_swa & (move < swa_count), + ) From 1ed3f594ecd22ff16a840bf0a2fd9a99a2ea6846 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 01:07:38 -0700 Subject: [PATCH 035/178] [None][chore] Name the tie-settling kernel by its position after the top-k finalize_topk_indices read as if this module ranked tokens itself; the kernel runs after the CuTE top-k and only settles its boundary ties deterministically. Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 8 ++++---- .../triattention/triattention_kernels.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 372f23fc186f..963e51e5c3df 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -118,7 +118,7 @@ def __init__( from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops - from .triattention_kernels import _finalize_topk_indices_kernel + from .triattention_kernels import _settle_ties_after_topk_kernel self.device = scores.device self.scores = scores @@ -140,8 +140,8 @@ def __init__( ) runner._compile(*key) self.compiled_topk = runner.kernel_cache[key] - self.frozen_finalize = FrozenTritonKernelCall( - _finalize_topk_indices_kernel, + self.frozen_settle_ties = FrozenTritonKernelCall( + _settle_ties_after_topk_kernel, (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), dict( WIDTH=width, @@ -163,7 +163,7 @@ def __call__(self) -> None: self.provisional_indices, None, ) - self.frozen_finalize() + self.frozen_settle_ties() class _PreparedUnionScores: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 9c30ff397cc0..85c345c87c70 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -819,7 +819,7 @@ def prepare_per_head_scores( @triton.jit -def _finalize_topk_indices_kernel( +def _settle_ties_after_topk_kernel( scores, seq_lens, prompt_offsets, From 6d2d3de2a763ac49d7ea9398ca78007c2c20a218 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 01:36:05 -0700 Subject: [PATCH 036/178] [None][chore] House the prepared kernel launchers with their kernels Pure move: the union-score and deterministic top-k launcher objects are kernel-launch machinery, so they live in the kernels module next to FrozenTritonKernelCall and the kernels they bind. The manager module keeps only runtime state: buffers, staging, and the eviction lifecycle. Signed-off-by: tianruih --- .../triattention/triattention.py | 148 +----------------- .../triattention/triattention_kernels.py | 141 +++++++++++++++++ 2 files changed, 143 insertions(+), 146 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 963e51e5c3df..47dffb2e4ab9 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -59,7 +59,8 @@ import torch from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - FrozenTritonKernelCall, + _PreparedDeterministicTopK, + _PreparedUnionScores, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, get_draft_token_length @@ -92,151 +93,6 @@ def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Ten return torch.tensor(offsets, device=device, dtype=torch.float32) -class _PreparedDeterministicTopK: - """Deterministic top-k over fixed row-major buffers, built once. - - The CuTE top-k kernel is fast but breaks score ties arbitrarily and emits - indices in arbitrary order, so a frozen Triton finalizer recomputes the - threshold membership with lowest-index-wins ties, rebases each row by its - prompt offset, and writes sorted ordinals. The CuTE kernel is compiled - here once with owned scratch storage; every later call runs both kernels - over the bound buffers with no dispatch work and no allocations. - """ - - def __init__( - self, - scores: torch.Tensor, - seq_lens: torch.Tensor, - prompt_offsets: torch.Tensor, - provisional_indices: torch.Tensor, - output_indices: torch.Tensor, - keep_count: int, - ) -> None: - rows, width = scores.shape - if not 1 <= keep_count <= width: - raise ValueError("deterministic top-k requires 1 <= keep_count <= width") - - from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops - - from .triattention_kernels import _settle_ties_after_topk_kernel - - self.device = scores.device - self.scores = scores - self.seq_lens = seq_lens - self.provisional_indices = provisional_indices - with torch.cuda.device(self.device): - self.stream = torch.cuda.current_stream(self.device) - self.scratch = torch.empty((rows, 2, width), dtype=torch.int32, device=self.device) - runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner - key = ( - cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], - 1 << (width - 1).bit_length(), - keep_count, - 1, - False, - 256, - False, - rows > cute_dsl_custom_ops._get_num_sms(), - ) - runner._compile(*key) - self.compiled_topk = runner.kernel_cache[key] - self.frozen_settle_ties = FrozenTritonKernelCall( - _settle_ties_after_topk_kernel, - (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), - dict( - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, - BLOCK=256, - ), - grid=(rows, 1, 1), - num_warps=4, - ) - - def __call__(self) -> None: - self.compiled_topk( - self.scores, - None, - self.scratch, - None, - self.seq_lens, - self.provisional_indices, - None, - ) - self.frozen_settle_ties() - - -class _PreparedUnionScores: - """Launch fixed union score preparation without Triton JIT dispatch.""" - - def __init__( - self, - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - combined: torch.Tensor, - *, - normalize_scores: bool, - ) -> None: - if scores.ndim != 3: - raise ValueError("prepared union scores require request-major rows") - request_count, rows, width = scores.shape - if ( - not scores.is_cuda - or scores.dtype != torch.float32 - or not scores.is_contiguous() - or valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.shape != (request_count, rows, 1) - or row_mean.dtype != torch.float32 - or row_mean.device != scores.device - or row_inv_std.shape != row_mean.shape - or row_inv_std.dtype != torch.float32 - or row_inv_std.device != scores.device - or combined.shape != (request_count, width) - or combined.dtype != torch.float32 - or combined.device != scores.device - or not valid_widths.is_contiguous() - or not row_mean.is_contiguous() - or not row_inv_std.is_contiguous() - or not combined.is_contiguous() - ): - raise ValueError("prepared union score tensors do not share one fixed geometry") - - from .triattention_kernels import _score_row_stats_kernel, _score_union_kernel - - normalize_scores = bool(normalize_scores) - # Callers verify score-tensor identity and the normalize flag against - # this launcher before dispatching to it. - self.scores = scores - self.normalize_scores = normalize_scores - self._frozen_stats_call = None - if normalize_scores: - stats_grid = (request_count * rows, 1, 1) - self._frozen_stats_call = FrozenTritonKernelCall( - _score_row_stats_kernel, - (scores, valid_widths, row_mean, row_inv_std), - dict(ROWS=rows, WIDTH=width, BLOCK=256), - grid=stats_grid, - num_warps=4, - ) - union_grid = (request_count, (width + 31) // 32, 1) - self._frozen_union_call = FrozenTritonKernelCall( - _score_union_kernel, - (scores, valid_widths, row_mean, row_inv_std, combined), - dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), - grid=union_grid, - num_warps=1, - ) - - def __call__(self) -> None: - if self._frozen_stats_call is not None: - self._frozen_stats_call() - self._frozen_union_call() - - class _CrossRequestSelectionPlan(NamedTuple): """Selection dimensions used to allocate reusable eager buffers.""" diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 85c345c87c70..06812a1fad92 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -895,6 +895,147 @@ def _settle_ties_after_topk_kernel( ties_seen += tl.sum(tied_i32) +class _PreparedDeterministicTopK: + """Deterministic top-k over fixed row-major buffers, built once. + + The CuTE top-k kernel is fast but breaks score ties arbitrarily and emits + indices in arbitrary order, so a frozen Triton finalizer recomputes the + threshold membership with lowest-index-wins ties, rebases each row by its + prompt offset, and writes sorted ordinals. The CuTE kernel is compiled + here once with owned scratch storage; every later call runs both kernels + over the bound buffers with no dispatch work and no allocations. + """ + + def __init__( + self, + scores: torch.Tensor, + seq_lens: torch.Tensor, + prompt_offsets: torch.Tensor, + provisional_indices: torch.Tensor, + output_indices: torch.Tensor, + keep_count: int, + ) -> None: + rows, width = scores.shape + if not 1 <= keep_count <= width: + raise ValueError("deterministic top-k requires 1 <= keep_count <= width") + + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + self.device = scores.device + self.scores = scores + self.seq_lens = seq_lens + self.provisional_indices = provisional_indices + with torch.cuda.device(self.device): + self.stream = torch.cuda.current_stream(self.device) + self.scratch = torch.empty((rows, 2, width), dtype=torch.int32, device=self.device) + runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner + key = ( + cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], + 1 << (width - 1).bit_length(), + keep_count, + 1, + False, + 256, + False, + rows > cute_dsl_custom_ops._get_num_sms(), + ) + runner._compile(*key) + self.compiled_topk = runner.kernel_cache[key] + self.frozen_settle_ties = FrozenTritonKernelCall( + _settle_ties_after_topk_kernel, + (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), + dict( + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + BLOCK=256, + ), + grid=(rows, 1, 1), + num_warps=4, + ) + + def __call__(self) -> None: + self.compiled_topk( + self.scores, + None, + self.scratch, + None, + self.seq_lens, + self.provisional_indices, + None, + ) + self.frozen_settle_ties() + + +class _PreparedUnionScores: + """Launch fixed union score preparation without Triton JIT dispatch.""" + + def __init__( + self, + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + combined: torch.Tensor, + *, + normalize_scores: bool, + ) -> None: + if scores.ndim != 3: + raise ValueError("prepared union scores require request-major rows") + request_count, rows, width = scores.shape + if ( + not scores.is_cuda + or scores.dtype != torch.float32 + or not scores.is_contiguous() + or valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.shape != (request_count, rows, 1) + or row_mean.dtype != torch.float32 + or row_mean.device != scores.device + or row_inv_std.shape != row_mean.shape + or row_inv_std.dtype != torch.float32 + or row_inv_std.device != scores.device + or combined.shape != (request_count, width) + or combined.dtype != torch.float32 + or combined.device != scores.device + or not valid_widths.is_contiguous() + or not row_mean.is_contiguous() + or not row_inv_std.is_contiguous() + or not combined.is_contiguous() + ): + raise ValueError("prepared union score tensors do not share one fixed geometry") + + normalize_scores = bool(normalize_scores) + # Callers verify score-tensor identity and the normalize flag against + # this launcher before dispatching to it. + self.scores = scores + self.normalize_scores = normalize_scores + self._frozen_stats_call = None + if normalize_scores: + stats_grid = (request_count * rows, 1, 1) + self._frozen_stats_call = FrozenTritonKernelCall( + _score_row_stats_kernel, + (scores, valid_widths, row_mean, row_inv_std), + dict(ROWS=rows, WIDTH=width, BLOCK=256), + grid=stats_grid, + num_warps=4, + ) + union_grid = (request_count, (width + 31) // 32, 1) + self._frozen_union_call = FrozenTritonKernelCall( + _score_union_kernel, + (scores, valid_widths, row_mean, row_inv_std, combined), + dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), + grid=union_grid, + num_warps=1, + ) + + def __call__(self) -> None: + if self._frozen_stats_call is not None: + self._frozen_stats_call() + self._frozen_union_call() + + # --------------------------------------------------------------------------- # # Compaction: pack the kept ordinals into per-request move indices. # # --------------------------------------------------------------------------- # From bbb907a7ae49db83b65690e3d4fc7fa3eb8e18ef Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 02:00:54 -0700 Subject: [PATCH 037/178] [None][refactor] Gather every kernel warmup into one module with one trigger All build-time kernel compilation now lives in triattention_kernel_warmup.py: the frozen Triton call, the score launch, the deterministic top-k, the union score preparation, and the compaction move-index pack. The kernels module keeps only kernels, the compaction and manager modules keep only runtime state and thin calls, and the new TriAttention.warmup_kernels() compiles the whole eviction pipeline in one call so it can later be driven from engine warmup or manager creation. Signed-off-by: tianruih --- .../triattention/compaction.py | 115 +---- .../triattention/triattention.py | 152 ++---- .../triattention_kernel_warmup.py | 451 ++++++++++++++++++ .../triattention/triattention_kernels.py | 206 +------- .../test_triattention_pipeline.py | 5 +- 5 files changed, 501 insertions(+), 428 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index a3454404326e..03ccb4aef6e6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -27,15 +27,10 @@ import torch -from .triattention_kernels import FrozenTritonKernelCall +from .triattention_kernel_warmup import FrozenTritonKernelCall, frozen_move_index_pack _SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) -# Launch shape of the move-index packing kernel: tokens per program along the -# move axis, and its warp count. -_PACK_BLOCK_TOKENS = 256 -_PACK_NUM_WARPS = 4 - class _CppCompactGroup(NamedTuple): """One layered sparse-KV updater launch over pools sharing a block table.""" @@ -221,110 +216,6 @@ def _compact_groups( return tuple(result) -def _frozen_move_index_pack( - kept_token_ordinals: torch.Tensor, - valid_sequence_lengths: torch.Tensor, - move_source_offsets: torch.Tensor, - move_source_indices: torch.Tensor, - *, - eviction_mode: str, - decode_keep_count: int, - num_dense_layers: int, - num_kv_heads: int, - max_protected_tail: int, - swa_window: int, - swa_move_source_offsets: Optional[torch.Tensor], - swa_move_source_indices: Optional[torch.Tensor], -) -> FrozenTritonKernelCall: - """Build one frozen call of the move-index packing kernel. - - The kernel reads the kept-token ordinals and each request's valid length - and writes the packed per-(layer, head) move source indices consumed by - the C++ compact launches. Only the caller-provided selection tensors are - validated here; the move buffers are allocated by this module. - """ - per_layer = eviction_mode == "per_layer_perhead" - union = eviction_mode == "union" - request_count = int(kept_token_ordinals.shape[0]) if kept_token_ordinals.ndim else 0 - if union: - selection_rows = 1 - elif per_layer: - selection_rows = num_dense_layers * num_kv_heads - else: - selection_rows = num_kv_heads - selection_prefix = (request_count,) if union else (request_count, selection_rows) - # Selection rows carry decode-only kept ordinals (already absolute), so - # the rectangle is prompt-length independent. - expected_selection = (*selection_prefix, decode_keep_count) - if ( - request_count <= 0 - or tuple(kept_token_ordinals.shape) != expected_selection - or valid_sequence_lengths.shape != (request_count,) - ): - raise ValueError( - f"prepared compaction packing expects kept ordinals of shape " - f"{expected_selection} and one valid length per request; got " - f"{tuple(kept_token_ordinals.shape)} and " - f"{tuple(valid_sequence_lengths.shape)}" - ) - - device = kept_token_ordinals.device - if not _cuda_int32_contiguous((kept_token_ordinals, valid_sequence_lengths), device): - raise ValueError("prepared compaction packing requires contiguous CUDA int32 tensors") - - if swa_move_source_indices is not None: - swa_offsets_arg = swa_move_source_offsets - swa_indices_arg = swa_move_source_indices - swa_total = int(swa_move_source_indices.shape[-1]) - else: - # HAS_SWA specializes all corresponding loads and stores away. - swa_offsets_arg = move_source_offsets - swa_indices_arg = move_source_indices - swa_total = 0 - - from .triattention_kernels import _pack_compaction_sources_kernel - - max_move = decode_keep_count + max_protected_tail - if swa_total: - max_move = max(max_move, swa_window + max_protected_tail) - packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads - grid = ( - request_count, - packed_row_count, - (max_move + _PACK_BLOCK_TOKENS - 1) // _PACK_BLOCK_TOKENS, - ) - bound_tensors = ( - kept_token_ordinals, - valid_sequence_lengths, - move_source_offsets, - move_source_indices, - swa_offsets_arg, - swa_indices_arg, - ) - # Ordered to match the kernel's constexpr parameter declaration: the - # frozen call passes these by position. - constexpr_values = dict( - DENSE_TOTAL=int(move_source_indices.shape[-1]), - SWA_TOTAL=swa_total, - SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=decode_keep_count, - KEEP_COUNT=decode_keep_count, - NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window, - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=swa_total > 0, - BLOCK=_PACK_BLOCK_TOKENS, - ) - return FrozenTritonKernelCall( - _pack_compaction_sources_kernel, - bound_tensors, - constexpr_values, - grid=grid, - num_warps=_PACK_NUM_WARPS, - ) - - class BatchedKVCacheCompaction: """Batched physical compaction of the KV caches for one fixed geometry. @@ -478,7 +369,7 @@ def __init__( dense_slots = ( {layer: slot for slot, layer in enumerate(self.dense_layers)} if per_layer else None ) - dense_pack = _frozen_move_index_pack( + dense_pack = frozen_move_index_pack( kept_token_ordinals, valid_sequence_lengths, dense_move_offsets, @@ -600,7 +491,7 @@ def _build_draft_compaction( # packed row, so one more frozen kernel call broadcasts the target keep # set over the draft KV heads and appends the draft's own tail # ordinals (valid_seq_len + 0..tail-1). - draft_pack = _frozen_move_index_pack( + draft_pack = frozen_move_index_pack( kept_token_ordinals, valid_sequence_lengths, draft_move_offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 47dffb2e4ab9..61584631f5cd 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -58,7 +58,9 @@ import torch -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernel_warmup import ( + FrozenScoreCall, + _FixedScoreStreamMismatch, _PreparedDeterministicTopK, _PreparedUnionScores, ) @@ -414,10 +416,6 @@ def select_requests( self.prepared_topk() -class _FixedScoreStreamMismatch(RuntimeError): - """Raised when a fixed score staging buffers are used from another CUDA stream.""" - - class _FixedScoreStagingBuffers: """Pool-bound fixed score metadata with one nonblocking page-table upload.""" @@ -667,114 +665,22 @@ def __init__( self.copy_pending = False self.page_tables_active = False self.stream = None - self._phase_runner = None - self._phase_args: tuple = () - self._score_runner = None - self._score_args: tuple = () + self._frozen_score: Optional[FrozenScoreCall] = None def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str) -> None: - """Compile and bind phase/score launches for this exact resource bucket.""" - from .triattention_kernels import _prepare_mean_phase_kernel, _tri_score_perhead_kernel - - if self._score_runner is not None: + """Compile and bind the phase and score launches for these buffers.""" + if self._frozen_score is not None: raise RuntimeError("TriAttention score launcher is already bound") - if score_aggregation not in ("mean", "max"): - raise ValueError(f"unsupported score aggregation: {score_aggregation}") - group = self.fused_group - if ( - valid_widths.shape != (self.max_requests,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != self.device - or not valid_widths.is_contiguous() - ): - raise ValueError("prepared score lengths do not match their exact bucket") - - frequency_block = 1 << (group.num_freqs - 1).bit_length() - phase_pointer_args = ( - self.round_starts_device, - self.offsets, - self.omega, - self.mean_cos, - self.mean_sin, - ) - phase_constants = ( - group.num_freqs, - int(self.offsets.numel()), - frequency_block, - ) - score_pointer_args = ( - *group.pointer_prefix, - self.valid_seq_lens_device, - valid_widths, - self.round_starts_device, - self.token_starts_device, - *group.pointer_middle, - self.mean_cos.view(-1), - self.mean_sin.view(-1), - *group.pointer_tail, - group.output, - ) - score_geometry = ( - group.output_width, - group.num_layers, - *group.geometry_args, - ) - score_constants = ( - score_aggregation == "max", - group.token_block, - frequency_block, - ) - self._phase_args = (*phase_pointer_args, *phase_constants) - self._score_args = (*score_pointer_args, *score_geometry, *score_constants) - phase_grid = (self.max_requests, 1, 1) - score_segments = self.max_requests * group.num_layers - if score_segments > 65535: - # Segments sit on the y grid axis (CUDA caps y/z at 65535) so the - # unbounded x axis can hold the token tiles of long sequences. - raise ValueError("request*layer segment count exceeds the CUDA grid limit") - score_grid = ( - group.max_ntblk, - score_segments, - group.num_kv_heads, - ) - with torch.cuda.device(self.device): - self.stream = torch.cuda.current_stream(self.device) - if score_aggregation == "mean": - compiled = _prepare_mean_phase_kernel.warmup( - *phase_pointer_args, - NUM_FREQS=phase_constants[0], - NUM_OFFSETS=phase_constants[1], - F_BLOCK=phase_constants[2], - num_warps=1, - grid=phase_grid, - ) - self._phase_runner = compiled[phase_grid] - compiled = _tri_score_perhead_kernel.warmup( - *score_pointer_args, - *score_geometry, - USE_MAX=score_constants[0], - T_BLOCK=score_constants[1], - F_BLOCK=score_constants[2], - grid=score_grid, - ) - self._score_runner = compiled[score_grid] + self._frozen_score = FrozenScoreCall(self, valid_widths, score_aggregation) + # stage() binds this staging to its first CUDA stream; the frozen + # score call was just built on that same stream. + self.stream = self._frozen_score.stream def launch_prepared_score(self) -> torch.Tensor: - """Launch the phase and score runners bound to this exact bucket.""" - if self._score_runner is None or self.stream is None: + """Launch the phase and score calls bound to these buffers.""" + if self._frozen_score is None: raise RuntimeError("TriAttention score launcher is not bound") - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.stream.device, - self.stream.cuda_stream, - ): - raise _FixedScoreStreamMismatch( - "TriAttention prepared score is bound to its staging CUDA stream" - ) - if self._phase_runner is not None: - self._phase_runner(*self._phase_args, stream=self.stream.cuda_stream) - self._score_runner(*self._score_args, stream=self.stream.cuda_stream) - return self.fused_group.output + return self._frozen_score() def stage( self, @@ -1172,6 +1078,38 @@ def _draft_protected_tail_capacity(self) -> int: raise RuntimeError("draft KVCacheManagerV2 exposes an invalid protected-tail capacity") return capacity + def warmup_kernels(self) -> None: + """Build every eviction buffer set and compile all its kernels now. + + One call compiles the whole pipeline: score, selection (CuTE top-k + plus the tie-settling call), and the compaction pack. The eviction + kernels never run inside the engine warmup forwards, so without this + the first due eviction round pays the one-time JIT compilation while + serving. Callable as soon as the KV pools exist; idempotent. + """ + self._ensure_calibrated() + layout = self._runtime_kv_layout(self._num_layers_from_manager()) + # A synthetic one-request cohort: the buffers are sized from the + # executor limits, so its geometry only has to be admissible (the + # prompt must cover the SWA landing window when SWA layers exist). + prompt_len = int(layout.swa_window or 0) + synthetic = _PreparedEviction( + request=None, + request_id=-1, + seq_len=prompt_len + 1, + round_start=prompt_len + 1, + prompt_len=prompt_len, + expected_keep_count=prompt_len + self.top_B, + protected_tail=0, + ) + resources = self._eager_resources_for(layout, [synthetic]) + self._batched_compaction_for( + layout=layout, + prepared=[synthetic], + score_staging=resources.score_staging, + keep_set_selector=resources.keep_set_selector, + ) + def _ensure_calibrated(self) -> None: """Resolve calibration once for the first request.""" if self._calibrated: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py new file mode 100644 index 000000000000..5f7fdd2f32e5 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py @@ -0,0 +1,451 @@ +# 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. + +"""One home for every kernel warmup in the TriAttention eviction pipeline. + +Triton and CuTE-DSL kernels JIT-compile on first call, and the eviction +kernels never run inside the engine warmup forwards, so each launcher here +compiles its kernels once at build time and then calls the compiled binaries +directly: no per-call dispatch work and no allocations on the eviction path. +Build one launcher per fixed buffer set; call it once per eviction round. +""" + +from typing import Dict, Optional, Tuple + +import torch + +from .triattention_kernels import ( + _score_row_stats_kernel, + _score_union_kernel, + _settle_ties_after_topk_kernel, +) + + +class _FixedScoreStreamMismatch(RuntimeError): + """Raised when fixed score staging buffers are used from another CUDA stream.""" + + +class FrozenTritonKernelCall: + """Call one Triton kernel frozen at build time. + + Triton's standard dispatch costs tens of microseconds of host work per + call; eviction fires its kernels every round, so the grid, bound tensor + set, and constexpr set are frozen once here. ``warmup`` (Triton's own + API) JIT-compiles the kernel at build time and ``__call__`` runs the + compiled binary directly. Constexpr values are passed positionally on + each call, so their order is validated against the kernel's constexpr + parameter declaration order at build time. + """ + + def __init__( + self, + triton_kernel, + bound_tensors: Tuple[torch.Tensor, ...], + constexpr_values: Dict[str, object], + *, + grid: Tuple[int, ...], + num_warps: int, + ) -> None: + params = getattr(triton_kernel, "params", None) + if params is not None: + declared = [param.name for param in params if param.is_constexpr] + if list(constexpr_values.keys()) != declared: + raise ValueError( + f"constexpr order {list(constexpr_values.keys())} must match the " + f"kernel's declaration order {declared}: the frozen call passes " + "them positionally" + ) + self.device = bound_tensors[0].device + self.bound_tensors = tuple(bound_tensors) + self.constexpr_values = dict(constexpr_values) + with torch.cuda.device(self.device): + self.build_stream = torch.cuda.current_stream(self.device) + # warmup() then indexing the compiled cache by grid is the + # documented-by-use Triton pattern for dispatch-free calls; if a + # Triton upgrade changes it, this raises here at build time rather + # than corrupting a later call. + compiled = triton_kernel.warmup( + *self.bound_tensors, + **self.constexpr_values, + num_warps=num_warps, + grid=grid, + ) + self.compiled_kernel_runner = compiled[grid] + + def __call__(self, *call_tensors: torch.Tensor) -> None: + """Run the kernel; ``call_tensors``, if given, substitute the bound tensors.""" + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.build_stream.device, + self.build_stream.cuda_stream, + ): + raise RuntimeError("a frozen Triton kernel call must run on the stream it was built on") + self.compiled_kernel_runner( + *(call_tensors if call_tensors else self.bound_tensors), + *self.constexpr_values.values(), + stream=self.build_stream.cuda_stream, + ) + + +class FrozenScoreCall: + """Phase and trig-score launches frozen over one staging buffer set.""" + + def __init__(self, staging, valid_widths: torch.Tensor, score_aggregation: str) -> None: + from .triattention_kernels import _prepare_mean_phase_kernel, _tri_score_perhead_kernel + + if score_aggregation not in ("mean", "max"): + raise ValueError(f"unsupported score aggregation: {score_aggregation}") + group = staging.fused_group + if ( + valid_widths.shape != (staging.max_requests,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != staging.device + or not valid_widths.is_contiguous() + ): + raise ValueError("prepared score lengths do not match their fixed buffers") + + frequency_block = 1 << (group.num_freqs - 1).bit_length() + phase_pointer_args = ( + staging.round_starts_device, + staging.offsets, + staging.omega, + staging.mean_cos, + staging.mean_sin, + ) + phase_constants = ( + group.num_freqs, + int(staging.offsets.numel()), + frequency_block, + ) + score_pointer_args = ( + *group.pointer_prefix, + staging.valid_seq_lens_device, + valid_widths, + staging.round_starts_device, + staging.token_starts_device, + *group.pointer_middle, + staging.mean_cos.view(-1), + staging.mean_sin.view(-1), + *group.pointer_tail, + group.output, + ) + score_geometry = ( + group.output_width, + group.num_layers, + *group.geometry_args, + ) + score_constants = ( + score_aggregation == "max", + group.token_block, + frequency_block, + ) + self._phase_args = (*phase_pointer_args, *phase_constants) + self._score_args = (*score_pointer_args, *score_geometry, *score_constants) + phase_grid = (staging.max_requests, 1, 1) + score_segments = staging.max_requests * group.num_layers + if score_segments > 65535: + # Segments sit on the y grid axis (CUDA caps y/z at 65535) so the + # unbounded x axis can hold the token tiles of long sequences. + raise ValueError("request*layer segment count exceeds the CUDA grid limit") + score_grid = ( + group.max_ntblk, + score_segments, + group.num_kv_heads, + ) + self.device = staging.device + self.output = group.output + self._phase_runner = None + with torch.cuda.device(self.device): + self.stream = torch.cuda.current_stream(self.device) + if score_aggregation == "mean": + compiled = _prepare_mean_phase_kernel.warmup( + *phase_pointer_args, + NUM_FREQS=phase_constants[0], + NUM_OFFSETS=phase_constants[1], + F_BLOCK=phase_constants[2], + num_warps=1, + grid=phase_grid, + ) + self._phase_runner = compiled[phase_grid] + compiled = _tri_score_perhead_kernel.warmup( + *score_pointer_args, + *score_geometry, + USE_MAX=score_constants[0], + T_BLOCK=score_constants[1], + F_BLOCK=score_constants[2], + grid=score_grid, + ) + self._score_runner = compiled[score_grid] + + def __call__(self) -> torch.Tensor: + current_stream = torch.cuda.current_stream(self.device) + if (current_stream.device, current_stream.cuda_stream) != ( + self.stream.device, + self.stream.cuda_stream, + ): + raise _FixedScoreStreamMismatch( + "TriAttention prepared score is bound to its staging CUDA stream" + ) + if self._phase_runner is not None: + self._phase_runner(*self._phase_args, stream=self.stream.cuda_stream) + self._score_runner(*self._score_args, stream=self.stream.cuda_stream) + return self.output + + +class _PreparedDeterministicTopK: + """Deterministic top-k over fixed row-major buffers, built once. + + The CuTE top-k kernel is fast but breaks score ties arbitrarily and emits + indices in arbitrary order, so a frozen Triton finalizer recomputes the + threshold membership with lowest-index-wins ties, rebases each row by its + prompt offset, and writes sorted ordinals. The CuTE kernel is compiled + here once with owned scratch storage; every later call runs both kernels + over the bound buffers with no dispatch work and no allocations. + """ + + def __init__( + self, + scores: torch.Tensor, + seq_lens: torch.Tensor, + prompt_offsets: torch.Tensor, + provisional_indices: torch.Tensor, + output_indices: torch.Tensor, + keep_count: int, + ) -> None: + rows, width = scores.shape + if not 1 <= keep_count <= width: + raise ValueError("deterministic top-k requires 1 <= keep_count <= width") + + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + self.device = scores.device + self.scores = scores + self.seq_lens = seq_lens + self.provisional_indices = provisional_indices + with torch.cuda.device(self.device): + self.stream = torch.cuda.current_stream(self.device) + self.scratch = torch.empty((rows, 2, width), dtype=torch.int32, device=self.device) + runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner + key = ( + cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], + 1 << (width - 1).bit_length(), + keep_count, + 1, + False, + 256, + False, + rows > cute_dsl_custom_ops._get_num_sms(), + ) + runner._compile(*key) + self.compiled_topk = runner.kernel_cache[key] + self.frozen_settle_ties = FrozenTritonKernelCall( + _settle_ties_after_topk_kernel, + (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), + dict( + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + BLOCK=256, + ), + grid=(rows, 1, 1), + num_warps=4, + ) + + def __call__(self) -> None: + self.compiled_topk( + self.scores, + None, + self.scratch, + None, + self.seq_lens, + self.provisional_indices, + None, + ) + self.frozen_settle_ties() + + +class _PreparedUnionScores: + """Launch fixed union score preparation without Triton JIT dispatch.""" + + def __init__( + self, + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + combined: torch.Tensor, + *, + normalize_scores: bool, + ) -> None: + if scores.ndim != 3: + raise ValueError("prepared union scores require request-major rows") + request_count, rows, width = scores.shape + if ( + not scores.is_cuda + or scores.dtype != torch.float32 + or not scores.is_contiguous() + or valid_widths.shape != (request_count,) + or valid_widths.dtype != torch.int32 + or valid_widths.device != scores.device + or row_mean.shape != (request_count, rows, 1) + or row_mean.dtype != torch.float32 + or row_mean.device != scores.device + or row_inv_std.shape != row_mean.shape + or row_inv_std.dtype != torch.float32 + or row_inv_std.device != scores.device + or combined.shape != (request_count, width) + or combined.dtype != torch.float32 + or combined.device != scores.device + or not valid_widths.is_contiguous() + or not row_mean.is_contiguous() + or not row_inv_std.is_contiguous() + or not combined.is_contiguous() + ): + raise ValueError("prepared union score tensors do not share one fixed geometry") + + normalize_scores = bool(normalize_scores) + # Callers verify score-tensor identity and the normalize flag against + # this launcher before dispatching to it. + self.scores = scores + self.normalize_scores = normalize_scores + self._frozen_stats_call = None + if normalize_scores: + stats_grid = (request_count * rows, 1, 1) + self._frozen_stats_call = FrozenTritonKernelCall( + _score_row_stats_kernel, + (scores, valid_widths, row_mean, row_inv_std), + dict(ROWS=rows, WIDTH=width, BLOCK=256), + grid=stats_grid, + num_warps=4, + ) + union_grid = (request_count, (width + 31) // 32, 1) + self._frozen_union_call = FrozenTritonKernelCall( + _score_union_kernel, + (scores, valid_widths, row_mean, row_inv_std, combined), + dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), + grid=union_grid, + num_warps=1, + ) + + def __call__(self) -> None: + if self._frozen_stats_call is not None: + self._frozen_stats_call() + self._frozen_union_call() + + +# Launch shape of the move-index packing kernel: tokens per program along the +# move axis, and its warp count. +_PACK_BLOCK_TOKENS = 256 +_PACK_NUM_WARPS = 4 + + +def frozen_move_index_pack( + kept_token_ordinals: torch.Tensor, + valid_sequence_lengths: torch.Tensor, + move_source_offsets: torch.Tensor, + move_source_indices: torch.Tensor, + *, + eviction_mode: str, + decode_keep_count: int, + num_dense_layers: int, + num_kv_heads: int, + max_protected_tail: int, + swa_window: int, + swa_move_source_offsets: Optional[torch.Tensor], + swa_move_source_indices: Optional[torch.Tensor], +) -> FrozenTritonKernelCall: + """Build one frozen call of the move-index packing kernel. + + The kernel reads the kept-token ordinals and each request's valid length + and writes the packed per-(layer, head) move source indices consumed by + the C++ compact launches. Only the caller-provided selection tensors are + validated here; the move buffers are allocated by this module. + """ + per_layer = eviction_mode == "per_layer_perhead" + union = eviction_mode == "union" + request_count = int(kept_token_ordinals.shape[0]) if kept_token_ordinals.ndim else 0 + if union: + selection_rows = 1 + elif per_layer: + selection_rows = num_dense_layers * num_kv_heads + else: + selection_rows = num_kv_heads + selection_prefix = (request_count,) if union else (request_count, selection_rows) + # Selection rows carry decode-only kept ordinals (already absolute), so + # the rectangle is prompt-length independent. + expected_selection = (*selection_prefix, decode_keep_count) + if ( + request_count <= 0 + or tuple(kept_token_ordinals.shape) != expected_selection + or valid_sequence_lengths.shape != (request_count,) + ): + raise ValueError( + f"prepared compaction packing expects kept ordinals of shape " + f"{expected_selection} and one valid length per request; got " + f"{tuple(kept_token_ordinals.shape)} and " + f"{tuple(valid_sequence_lengths.shape)}" + ) + + if swa_move_source_indices is not None: + swa_offsets_arg = swa_move_source_offsets + swa_indices_arg = swa_move_source_indices + swa_total = int(swa_move_source_indices.shape[-1]) + else: + # HAS_SWA specializes all corresponding loads and stores away. + swa_offsets_arg = move_source_offsets + swa_indices_arg = move_source_indices + swa_total = 0 + + from .triattention_kernels import _pack_compaction_sources_kernel + + max_move = decode_keep_count + max_protected_tail + if swa_total: + max_move = max(max_move, swa_window + max_protected_tail) + packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads + grid = ( + request_count, + packed_row_count, + (max_move + _PACK_BLOCK_TOKENS - 1) // _PACK_BLOCK_TOKENS, + ) + bound_tensors = ( + kept_token_ordinals, + valid_sequence_lengths, + move_source_offsets, + move_source_indices, + swa_offsets_arg, + swa_indices_arg, + ) + # Ordered to match the kernel's constexpr parameter declaration: the + # frozen call passes these by position. + constexpr_values = dict( + DENSE_TOTAL=int(move_source_indices.shape[-1]), + SWA_TOTAL=swa_total, + SELECTION_ROWS=selection_rows, + SELECTION_STRIDE=decode_keep_count, + KEEP_COUNT=decode_keep_count, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=swa_total > 0, + BLOCK=_PACK_BLOCK_TOKENS, + ) + return FrozenTritonKernelCall( + _pack_compaction_sources_kernel, + bound_tensors, + constexpr_values, + grid=grid, + num_warps=_PACK_NUM_WARPS, + ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 06812a1fad92..4f4e93513883 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -16,75 +16,12 @@ from __future__ import annotations -from typing import Dict, List, Tuple +from typing import List import torch import triton import triton.language as tl - -class FrozenTritonKernelCall: - """Call one Triton kernel frozen at build time. - - Triton's standard dispatch costs tens of microseconds of host work per - call; eviction fires its kernels every round, so the grid, bound tensor - set, and constexpr set are frozen once here. ``warmup`` (Triton's own - API) JIT-compiles the kernel at build time and ``__call__`` runs the - compiled binary directly. Constexpr values are passed positionally on - each call, so their order is validated against the kernel's constexpr - parameter declaration order at build time. - """ - - def __init__( - self, - triton_kernel, - bound_tensors: Tuple[torch.Tensor, ...], - constexpr_values: Dict[str, object], - *, - grid: Tuple[int, ...], - num_warps: int, - ) -> None: - params = getattr(triton_kernel, "params", None) - if params is not None: - declared = [param.name for param in params if param.is_constexpr] - if list(constexpr_values.keys()) != declared: - raise ValueError( - f"constexpr order {list(constexpr_values.keys())} must match the " - f"kernel's declaration order {declared}: the frozen call passes " - "them positionally" - ) - self.device = bound_tensors[0].device - self.bound_tensors = tuple(bound_tensors) - self.constexpr_values = dict(constexpr_values) - with torch.cuda.device(self.device): - self.build_stream = torch.cuda.current_stream(self.device) - # warmup() then indexing the compiled cache by grid is the - # documented-by-use Triton pattern for dispatch-free calls; if a - # Triton upgrade changes it, this raises here at build time rather - # than corrupting a later call. - compiled = triton_kernel.warmup( - *self.bound_tensors, - **self.constexpr_values, - num_warps=num_warps, - grid=grid, - ) - self.compiled_kernel_runner = compiled[grid] - - def __call__(self, *call_tensors: torch.Tensor) -> None: - """Run the kernel; ``call_tensors``, if given, substitute the bound tensors.""" - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.build_stream.device, - self.build_stream.cuda_stream, - ): - raise RuntimeError("a frozen Triton kernel call must run on the stream it was built on") - self.compiled_kernel_runner( - *(call_tensors if call_tensors else self.bound_tensors), - *self.constexpr_values.values(), - stream=self.build_stream.cuda_stream, - ) - - # --------------------------------------------------------------------------- # # Scoring: trig-score every cached token across all dense layers. # # --------------------------------------------------------------------------- # @@ -895,147 +832,6 @@ def _settle_ties_after_topk_kernel( ties_seen += tl.sum(tied_i32) -class _PreparedDeterministicTopK: - """Deterministic top-k over fixed row-major buffers, built once. - - The CuTE top-k kernel is fast but breaks score ties arbitrarily and emits - indices in arbitrary order, so a frozen Triton finalizer recomputes the - threshold membership with lowest-index-wins ties, rebases each row by its - prompt offset, and writes sorted ordinals. The CuTE kernel is compiled - here once with owned scratch storage; every later call runs both kernels - over the bound buffers with no dispatch work and no allocations. - """ - - def __init__( - self, - scores: torch.Tensor, - seq_lens: torch.Tensor, - prompt_offsets: torch.Tensor, - provisional_indices: torch.Tensor, - output_indices: torch.Tensor, - keep_count: int, - ) -> None: - rows, width = scores.shape - if not 1 <= keep_count <= width: - raise ValueError("deterministic top-k requires 1 <= keep_count <= width") - - from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops - - self.device = scores.device - self.scores = scores - self.seq_lens = seq_lens - self.provisional_indices = provisional_indices - with torch.cuda.device(self.device): - self.stream = torch.cuda.current_stream(self.device) - self.scratch = torch.empty((rows, 2, width), dtype=torch.int32, device=self.device) - runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner - key = ( - cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], - 1 << (width - 1).bit_length(), - keep_count, - 1, - False, - 256, - False, - rows > cute_dsl_custom_ops._get_num_sms(), - ) - runner._compile(*key) - self.compiled_topk = runner.kernel_cache[key] - self.frozen_settle_ties = FrozenTritonKernelCall( - _settle_ties_after_topk_kernel, - (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), - dict( - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, - BLOCK=256, - ), - grid=(rows, 1, 1), - num_warps=4, - ) - - def __call__(self) -> None: - self.compiled_topk( - self.scores, - None, - self.scratch, - None, - self.seq_lens, - self.provisional_indices, - None, - ) - self.frozen_settle_ties() - - -class _PreparedUnionScores: - """Launch fixed union score preparation without Triton JIT dispatch.""" - - def __init__( - self, - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - combined: torch.Tensor, - *, - normalize_scores: bool, - ) -> None: - if scores.ndim != 3: - raise ValueError("prepared union scores require request-major rows") - request_count, rows, width = scores.shape - if ( - not scores.is_cuda - or scores.dtype != torch.float32 - or not scores.is_contiguous() - or valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.shape != (request_count, rows, 1) - or row_mean.dtype != torch.float32 - or row_mean.device != scores.device - or row_inv_std.shape != row_mean.shape - or row_inv_std.dtype != torch.float32 - or row_inv_std.device != scores.device - or combined.shape != (request_count, width) - or combined.dtype != torch.float32 - or combined.device != scores.device - or not valid_widths.is_contiguous() - or not row_mean.is_contiguous() - or not row_inv_std.is_contiguous() - or not combined.is_contiguous() - ): - raise ValueError("prepared union score tensors do not share one fixed geometry") - - normalize_scores = bool(normalize_scores) - # Callers verify score-tensor identity and the normalize flag against - # this launcher before dispatching to it. - self.scores = scores - self.normalize_scores = normalize_scores - self._frozen_stats_call = None - if normalize_scores: - stats_grid = (request_count * rows, 1, 1) - self._frozen_stats_call = FrozenTritonKernelCall( - _score_row_stats_kernel, - (scores, valid_widths, row_mean, row_inv_std), - dict(ROWS=rows, WIDTH=width, BLOCK=256), - grid=stats_grid, - num_warps=4, - ) - union_grid = (request_count, (width + 31) // 32, 1) - self._frozen_union_call = FrozenTritonKernelCall( - _score_union_kernel, - (scores, valid_widths, row_mean, row_inv_std, combined), - dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), - grid=union_grid, - num_warps=1, - ) - - def __call__(self) -> None: - if self._frozen_stats_call is not None: - self._frozen_stats_call() - self._frozen_union_call() - - # --------------------------------------------------------------------------- # # Compaction: pack the kept ordinals into per-request move indices. # # --------------------------------------------------------------------------- # diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index b1b33d682401..eac4e7e7ece4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1697,10 +1697,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun staging.offsets = offsets staging.omega = omega staging.stream = None - staging._phase_runner = None - staging._phase_args = () - staging._score_runner = None - staging._score_args = () + staging._frozen_score = None staging.bind_score_launcher(valid_widths, aggregation) group.output.fill_(score_sentinel) with ( From 07827322a8c9febb29378b97a2d2b5bc749380f9 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 03:11:34 -0700 Subject: [PATCH 038/178] [None][perf] Carry the move offsets on the staged metadata table The compaction move offsets were refreshed with one to three separate small host-to-device copies per eviction round. The staging table grows one row per compacted cache family, the manager computes the padded offsets on the host, and the single metadata copy now carries every per-round host value: one table, one copy, one C++ launch per pool group. Standalone compactions still allocate their own offsets and refresh them through set_protected_tails. Signed-off-by: tianruih --- .../triattention/compaction.py | 18 +++- .../triattention/triattention.py | 90 +++++++++++++++---- 2 files changed, 91 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 03ccb4aef6e6..8251d7a85114 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -130,12 +130,20 @@ def _make_move_buffers( index_prefix: Tuple[int, ...], moves_per_request: List[int], device: torch.device, + external_offsets: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Allocate the packed source-index buffer and its per-request offsets.""" + """Allocate the packed source-index buffer and its per-request offsets. + + ``external_offsets`` shares a caller-owned device row (refreshed together + with the round metadata in one copy) instead of allocating one here; the + index buffer is always sized for the widest per-request move counts. + """ offsets = [0] for count in moves_per_request: offsets.append(offsets[-1] + count) indices = torch.empty((*index_prefix, offsets[-1]), dtype=torch.int32, device=device) + if external_offsets is not None: + return indices, external_offsets return indices, torch.tensor(offsets, dtype=torch.int32, device=device) @@ -273,6 +281,9 @@ def __init__( draft_protected_tail_capacity: Optional[int] = None, draft_kv_block_offsets: Optional[torch.Tensor] = None, draft_page_table_slots: Optional[Dict[int, int]] = None, + dense_move_offsets: Optional[torch.Tensor] = None, + swa_move_offsets: Optional[torch.Tensor] = None, + draft_move_offsets: Optional[torch.Tensor] = None, ) -> None: if eviction_mode not in ("union", "per_head", "per_layer_perhead"): raise ValueError(f"unsupported compaction mode: {eviction_mode}") @@ -331,6 +342,7 @@ def __init__( dense_index_prefix, [self.decode_keep_count + self.protected_tail_capacity] * self.request_count, self.device, + external_offsets=dense_move_offsets, ) page_table_for = _page_table_provider( page_table_slots, @@ -360,6 +372,7 @@ def __init__( (self.num_kv_heads,), [self.swa_window + self.protected_tail_capacity] * self.request_count, self.device, + external_offsets=swa_move_offsets, ) # SWA layers are staged as their own page-table representatives. swa_entries = [ @@ -416,6 +429,7 @@ def __init__( draft_protected_tail_capacity=draft_protected_tail_capacity, draft_kv_block_offsets=draft_kv_block_offsets, draft_page_table_slots=draft_page_table_slots, + draft_move_offsets=draft_move_offsets, ) self.cache_compactions = tuple( @@ -440,6 +454,7 @@ def _build_draft_compaction( draft_protected_tail_capacity: Optional[int], draft_kv_block_offsets: Optional[torch.Tensor], draft_page_table_slots: Optional[Dict[int, int]], + draft_move_offsets: Optional[torch.Tensor] = None, ) -> _SingleCacheCompaction: """Build the co-compressed draft cache's own pack and launch groups. @@ -471,6 +486,7 @@ def _build_draft_compaction( (draft_num_kv_heads,), [self.decode_keep_count + self.draft_protected_tail_capacity] * self.request_count, self.device, + external_offsets=draft_move_offsets, ) draft_page_table_for = _page_table_provider( draft_page_table_slots, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 61584631f5cd..1658fad5459a 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -538,8 +538,13 @@ def __init__( 2, self.copy_block_count, ) + # One table carries every per-round host value: three metadata rows + # (logical position, valid length, prompt length) plus one move-offsets + # row per compacted cache family, so each round pays exactly one + # host-to-device copy. Offsets rows have request_capacity + 1 entries, + # hence the extra column. self.request_metadata_host = torch.empty( - (3, max_requests), + (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned(), @@ -621,13 +626,18 @@ def __init__( device=self.device, ) self.request_metadata_device = torch.empty( - (3, max_requests), dtype=torch.int32, device=self.device + (6, max_requests + 1), dtype=torch.int32, device=self.device ) - self.round_starts_device = self.request_metadata_device[0] - self.valid_seq_lens_device = self.request_metadata_device[1] + self.round_starts_device = self.request_metadata_device[0, :max_requests] + self.valid_seq_lens_device = self.request_metadata_device[1, :max_requests] # Per-request pinned prompt lengths: the score kernel starts each # request's decode window here, so one bucket may mix prompt lengths. - self.token_starts_device = self.request_metadata_device[2] + self.token_starts_device = self.request_metadata_device[2, :max_requests] + # Per-family move offsets consumed by the compaction pack kernel and + # the C++ compact launches; refreshed with the metadata each round. + self.dense_move_offsets = self.request_metadata_device[3] + self.swa_move_offsets = self.request_metadata_device[4] + self.draft_move_offsets = self.request_metadata_device[5] self.mean_cos = torch.empty( (max_requests, num_freqs), dtype=torch.float32, device=self.device ) @@ -691,6 +701,9 @@ def stage( seq_lens: Optional[List[int]] = None, page_table_seq_lens: Optional[List[int]] = None, draft_manager: Optional[KVCacheManagerV2] = None, + dense_move_offsets: Optional[List[int]] = None, + swa_move_offsets: Optional[List[int]] = None, + draft_move_offsets: Optional[List[int]] = None, ) -> bool: """Copy one eager eviction cohort into reusable device buffers. @@ -759,10 +772,21 @@ def stage( self.draft_copy_block_count, ): return False - self.request_metadata_host[:, :request_count].copy_(request_metadata) + self.request_metadata_host[:3, :request_count].copy_(request_metadata) # Rows past this cohort are padding: zero lengths keep the score # kernel and selection inert for them. - self.request_metadata_host[:, request_count:].zero_() + self.request_metadata_host[:3, request_count:].zero_() + # This round's per-family move offsets ride the same table, so the + # single device copy below carries them too. + for row, family_offsets in ( + (3, dense_move_offsets), + (4, swa_move_offsets), + (5, draft_move_offsets), + ): + if family_offsets is not None: + self.request_metadata_host[row, : len(family_offsets)].copy_( + torch.as_tensor(family_offsets, dtype=torch.int32) + ) try: # Copy the fixed backing once. Only the first ``request_count`` # columns are consumed by this cohort. @@ -1963,19 +1987,46 @@ def _batched_compaction_for( decode_keep_count=self.top_B, swa_window=layout.swa_window, protected_tail_capacity=self._configured_protected_tail_capacity(), + dense_move_offsets=score_staging.dense_move_offsets, + swa_move_offsets=score_staging.swa_move_offsets, + draft_move_offsets=score_staging.draft_move_offsets, **draft_kwargs, ) self._batched_compaction = batched_compaction - # Tails vary per round (in-flight growth); padded rows move nothing. - draft_tails = None - if self.draft_kv_cache_manager is not None: - draft_tails = [self._draft_protected_tail_capacity()] * len(prepared) - batched_compaction.set_protected_tails( - [item.protected_tail for item in prepared], - draft_tails, - ) + # Tails vary per round (in-flight growth), so the per-family move + # offsets ride the staged metadata table each round. return batched_compaction + def _move_offsets_for( + self, + layout: _RuntimeKVLayout, + prepared: Sequence[_PreparedEviction], + capacity: int, + ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: + """Build this round's per-family move offsets, padded to the capacity. + + Rows past the cohort repeat the final offset, so padded requests move + nothing in the pack kernel and the C++ compact launches. + """ + + def padded_offsets(moves_per_request: List[int]) -> List[int]: + offsets = [0] + for moves in moves_per_request: + offsets.append(offsets[-1] + moves) + offsets.extend(offsets[-1:] * (capacity - len(moves_per_request))) + return offsets + + tails = [item.protected_tail for item in prepared] + dense = padded_offsets([self.top_B + tail for tail in tails]) + swa = None + if layout.swa_layers and layout.swa_window: + swa = padded_offsets([int(layout.swa_window) + tail for tail in tails]) + draft = None + if self.draft_kv_cache_manager is not None: + draft_tail = self._draft_protected_tail_capacity() + draft = padded_offsets([self.top_B + draft_tail] * len(prepared)) + return dense, swa, draft + def _page_table_pool_keys( self, representatives: List[int], @@ -2013,7 +2064,11 @@ def _attach_page_ids( self, prepared: Sequence[_PreparedEviction], staging: _FixedScoreStagingBuffers, + layout: _RuntimeKVLayout, ) -> None: + dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( + layout, prepared, staging.max_requests + ) try: staged = staging.stage( self.kv_cache_manager, @@ -2023,6 +2078,9 @@ def _attach_page_ids( [item.seq_len for item in prepared], [item.seq_len + item.protected_tail for item in prepared], draft_manager=self.draft_kv_cache_manager, + dense_move_offsets=dense_offsets, + swa_move_offsets=swa_offsets, + draft_move_offsets=draft_offsets, ) except _FixedScoreStreamMismatch: raise @@ -2093,7 +2151,7 @@ def _evict_requests( keep_set_selector=keep_set_selector, ) with nvtx_range_debug("triattention.page_table_stage", color="orange"): - self._attach_page_ids(prepared, score_staging) + self._attach_page_ids(prepared, score_staging, layout) # The staged per-request prompt lengths are shared with the # selector; per-head modes re-expand them to selection rows here. keep_set_selector.refresh_row_prompt_offsets() From c7ca77b30900a3174183e4451872862277586168 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 03:19:59 -0700 Subject: [PATCH 039/178] [None][chore] Catch the page-table staging test up with the offsets-carrying stage Signed-off-by: tianruih --- .../kv_cache_compression/test_triattention_pipeline.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index eac4e7e7ece4..dbb6dae26e25 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1288,6 +1288,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): manager.kv_cache_manager = SimpleNamespace(get_batch_cache_indices=get_batch) staging = SimpleNamespace( stage=mock.Mock(return_value=True), + max_requests=8, ) prepared = [ _prepared_eviction( @@ -1310,8 +1311,11 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ), ] - manager._attach_page_ids(prepared, staging) + layout = SimpleNamespace(swa_layers=[], swa_window=None) + manager._attach_page_ids(prepared, staging, layout) + # top_B=8: per-request moves are keep + tail = [10, 11]; padded rows + # repeat the final offset out to the request capacity. staging.stage.assert_called_once_with( manager.kv_cache_manager, [7, 8], @@ -1320,6 +1324,9 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): [8, 9], [10, 12], draft_manager=None, + dense_move_offsets=[0, 10, 21, 21, 21, 21, 21, 21, 21], + swa_move_offsets=None, + draft_move_offsets=None, ) assert all(not hasattr(item, "page_ids") for item in prepared) From 59330896fbee8f99cb2dd1828b3658cdebb5f082 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 06:56:33 -0700 Subject: [PATCH 040/178] [None][refactor] Drop the kernel warmup machinery and dispatch directly Per review: the eviction kernels now launch through standard Triton dispatch and the CuTE operator entry each round, like every other kernel consumer in the backend. This removes the warmup module, the frozen launch objects, and the manager warmup entry; the one-time JIT compilation lands on the first eviction round, and if the per-round dispatch cost ever matters the work moves into the CUDA kernels instead. The warmup implementation is preserved on the kvcache-v2-triattention-warmup-backup branch. Signed-off-by: tianruih --- .../triattention/compaction.py | 132 ++++- .../triattention/triattention.py | 165 ++++--- .../triattention_kernel_warmup.py | 451 ------------------ .../test_triattention_eager.py | 8 +- .../test_triattention_pipeline.py | 31 +- 5 files changed, 214 insertions(+), 573 deletions(-) delete mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 8251d7a85114..b14ebfe4a03d 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -17,7 +17,7 @@ Given each request's kept-token ordinals, its valid sequence length, and the staged V2 block offsets, this module packs per-request move indices with one -frozen Triton kernel call and then moves the surviving KV in place with batched +Triton launch and then moves the surviving KV in place with batched C++ compact launches. Inputs are plain tensors, so any eviction method that produces a kept-token set per request can drive it. """ @@ -27,8 +27,6 @@ import torch -from .triattention_kernel_warmup import FrozenTritonKernelCall, frozen_move_index_pack - _SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) @@ -57,13 +55,13 @@ def compact( class _SingleCacheCompaction(NamedTuple): """One compacted cache family (target dense, target SWA, or draft). - Holds the frozen kernel call that packs this family's move indices (None + Holds the launch that packs this family's move indices (None when an earlier family's pack call fills them in the same run), the C++ compact groups that consume them, and the destination base the moved tokens land at. """ - frozen_move_index_pack: Optional[FrozenTritonKernelCall] + move_index_pack: Optional[Callable[[], None]] cpp_compact_groups: Tuple[_CppCompactGroup, ...] move_source_indices: torch.Tensor move_source_offsets: torch.Tensor @@ -72,8 +70,8 @@ class _SingleCacheCompaction(NamedTuple): destination_bases: torch.Tensor def compact(self) -> None: - if self.frozen_move_index_pack is not None: - self.frozen_move_index_pack() + if self.move_index_pack is not None: + self.move_index_pack() for group in self.cpp_compact_groups: group.compact( self.move_source_indices, self.move_source_offsets, self.destination_bases @@ -224,6 +222,114 @@ def _compact_groups( return tuple(result) +# Launch shape of the move-index packing kernel: tokens per program along the +# move axis, and its warp count. +_PACK_BLOCK_TOKENS = 256 +_PACK_NUM_WARPS = 4 + + +def _move_index_pack_launcher( + kept_token_ordinals: torch.Tensor, + valid_sequence_lengths: torch.Tensor, + move_source_offsets: torch.Tensor, + move_source_indices: torch.Tensor, + *, + eviction_mode: str, + decode_keep_count: int, + num_dense_layers: int, + num_kv_heads: int, + max_protected_tail: int, + swa_window: int, + swa_move_source_offsets: Optional[torch.Tensor], + swa_move_source_indices: Optional[torch.Tensor], +) -> Callable[[], None]: + """Build one launch of the move-index packing kernel. + + The kernel reads the kept-token ordinals and each request's valid length + and writes the packed per-(layer, head) move source indices consumed by + the C++ compact launches. Only the caller-provided selection tensors are + validated here; the move buffers are allocated by this module. + """ + per_layer = eviction_mode == "per_layer_perhead" + union = eviction_mode == "union" + request_count = int(kept_token_ordinals.shape[0]) if kept_token_ordinals.ndim else 0 + if union: + selection_rows = 1 + elif per_layer: + selection_rows = num_dense_layers * num_kv_heads + else: + selection_rows = num_kv_heads + selection_prefix = (request_count,) if union else (request_count, selection_rows) + # Selection rows carry decode-only kept ordinals (already absolute), so + # the rectangle is prompt-length independent. + expected_selection = (*selection_prefix, decode_keep_count) + if ( + request_count <= 0 + or tuple(kept_token_ordinals.shape) != expected_selection + or valid_sequence_lengths.shape != (request_count,) + ): + raise ValueError( + f"prepared compaction packing expects kept ordinals of shape " + f"{expected_selection} and one valid length per request; got " + f"{tuple(kept_token_ordinals.shape)} and " + f"{tuple(valid_sequence_lengths.shape)}" + ) + + if swa_move_source_indices is not None: + swa_offsets_arg = swa_move_source_offsets + swa_indices_arg = swa_move_source_indices + swa_total = int(swa_move_source_indices.shape[-1]) + else: + # HAS_SWA specializes all corresponding loads and stores away. + swa_offsets_arg = move_source_offsets + swa_indices_arg = move_source_indices + swa_total = 0 + + from .triattention_kernels import _pack_compaction_sources_kernel + + max_move = decode_keep_count + max_protected_tail + if swa_total: + max_move = max(max_move, swa_window + max_protected_tail) + packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads + grid = ( + request_count, + packed_row_count, + (max_move + _PACK_BLOCK_TOKENS - 1) // _PACK_BLOCK_TOKENS, + ) + bound_tensors = ( + kept_token_ordinals, + valid_sequence_lengths, + move_source_offsets, + move_source_indices, + swa_offsets_arg, + swa_indices_arg, + ) + # Ordered to match the kernel's constexpr parameter declaration: the + # ordered to match the kernel constexpr declaration. + constexpr_values = dict( + DENSE_TOTAL=int(move_source_indices.shape[-1]), + SWA_TOTAL=swa_total, + SELECTION_ROWS=selection_rows, + SELECTION_STRIDE=decode_keep_count, + KEEP_COUNT=decode_keep_count, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=swa_total > 0, + BLOCK=_PACK_BLOCK_TOKENS, + ) + + def launch_pack() -> None: + _pack_compaction_sources_kernel[grid]( + *bound_tensors, + **constexpr_values, + num_warps=_PACK_NUM_WARPS, + ) + + return launch_pack + + class BatchedKVCacheCompaction: """Batched physical compaction of the KV caches for one fixed geometry. @@ -382,7 +488,7 @@ def __init__( dense_slots = ( {layer: slot for slot, layer in enumerate(self.dense_layers)} if per_layer else None ) - dense_pack = frozen_move_index_pack( + dense_pack = _move_index_pack_launcher( kept_token_ordinals, valid_sequence_lengths, dense_move_offsets, @@ -397,7 +503,7 @@ def __init__( swa_move_source_indices=swa_move_indices, ) self.target_dense_compaction = _SingleCacheCompaction( - frozen_move_index_pack=dense_pack, + move_index_pack=dense_pack, cpp_compact_groups=_compact_groups( dense_entries, self.layer_pool_keys, self.device, dense_slots ), @@ -409,7 +515,7 @@ def __init__( self.target_swa_compaction = None if self.swa_layers: self.target_swa_compaction = _SingleCacheCompaction( - frozen_move_index_pack=None, + move_index_pack=None, cpp_compact_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), move_source_indices=swa_move_indices, move_source_offsets=swa_move_offsets, @@ -504,10 +610,10 @@ def _build_draft_compaction( for layer in draft_layers ] # In union mode the pack kernel reads selection row 0 for every - # packed row, so one more frozen kernel call broadcasts the target keep + # packed row, so one more pack launch broadcasts the target keep # set over the draft KV heads and appends the draft's own tail # ordinals (valid_seq_len + 0..tail-1). - draft_pack = frozen_move_index_pack( + draft_pack = _move_index_pack_launcher( kept_token_ordinals, valid_sequence_lengths, draft_move_offsets, @@ -522,7 +628,7 @@ def _build_draft_compaction( swa_move_source_indices=None, ) return _SingleCacheCompaction( - frozen_move_index_pack=draft_pack, + move_index_pack=draft_pack, cpp_compact_groups=_compact_groups( draft_entries, tuple(draft_layer_pool_keys), self.device ), diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 1658fad5459a..e4d36491ce9a 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -58,12 +58,6 @@ import torch -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernel_warmup import ( - FrozenScoreCall, - _FixedScoreStreamMismatch, - _PreparedDeterministicTopK, - _PreparedUnionScores, -) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, get_draft_token_length from tensorrt_llm._torch.pyexecutor.resource_manager import BaseKVCacheCompressionManager @@ -95,6 +89,10 @@ def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Ten return torch.tensor(offsets, device=device, dtype=torch.float32) +class _FixedScoreStreamMismatch(RuntimeError): + """Raised when fixed score staging buffers are used from another CUDA stream.""" + + class _CrossRequestSelectionPlan(NamedTuple): """Selection dimensions used to allocate reusable eager buffers.""" @@ -129,7 +127,7 @@ class _RuntimeKVLayout(NamedTuple): class _BatchedKeepSetSelectorBase: - """Shared fixed buffers and frozen kernel calls for keep-set selectors.""" + """Shared fixed buffers and row views for keep-set selectors.""" def __init__( self, @@ -208,21 +206,49 @@ def refresh_row_prompt_offsets(self) -> None: self.prompt_offsets.unsqueeze(1).expand(-1, self.selection_rows_per_request) ) - def _build_prepared_selection_launchers( + def _bind_selection_rows( self, scores_rows: torch.Tensor, row_lengths: torch.Tensor, provisional_indices: torch.Tensor, keep_rows: torch.Tensor, ) -> None: - """Bind the deterministic top-k over row-major views of the buffers.""" - self.prepared_topk = _PreparedDeterministicTopK( - scores_rows, - row_lengths, - self.row_prompt_offsets, - provisional_indices, - keep_rows, + """Keep row-major views of the buffers the top-k selection reads.""" + self._selection_scores_rows = scores_rows + self._selection_row_lengths = row_lengths + self._provisional_rows = provisional_indices + self._keep_rows = keep_rows + + def _select_top_tokens(self) -> None: + """Pick the top-k with the CuTE selector, then settle its output. + + The CuTE top-k is fast but breaks score ties arbitrarily and emits + indices in arbitrary order; the settle kernel recomputes the threshold + membership with lowest-index-wins ties, rebases each row by its prompt + offset, and writes sorted ordinals. + """ + from .triattention_kernels import _settle_ties_after_topk_kernel + + rows = int(self._selection_scores_rows.shape[0]) + # 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, + self._selection_row_lengths, + self._provisional_rows, self.keep_count, + 1, + ) + _settle_ties_after_topk_kernel[(rows,)]( + self._selection_scores_rows, + self._selection_row_lengths, + self.row_prompt_offsets, + self._provisional_rows, + self._keep_rows, + WIDTH=self.width, + KEEP_COUNT=self.keep_count, + OUTPUT_WIDTH=self.keep_count, + BLOCK=256, + num_warps=4, ) @@ -274,22 +300,25 @@ def __init__( self.keep = torch.empty( (max_requests, self.keep_count), dtype=torch.int32, device=self.device ) - self._build_prepared_selection_launchers( - self.combined, self.valid_widths, self.final_indices, self.keep - ) - self.prepared_scores = _PreparedUnionScores( - input_scores, + self._bind_selection_rows(self.combined, self.valid_widths, self.final_indices, self.keep) + # Callers select from exactly this tensor with exactly this flag. + self.input_scores = input_scores + self.normalize_scores = bool(normalize_scores) + + def select_prepared_requests(self) -> None: + """Select from the CUDA score tensor bound to this fixed selector.""" + from .triattention_kernels import prepare_union_scores + + prepare_union_scores( + self.input_scores, self.valid_widths, self.row_mean, self.row_std, self.combined, - normalize_scores=normalize_scores, + self.max_requests, + normalize_scores=self.normalize_scores, ) - - def select_prepared_requests(self) -> None: - """Select from the CUDA score tensor bound to this fixed selector.""" - self.prepared_scores() - self.prepared_topk() + self._select_top_tokens() def select_requests( self, @@ -298,11 +327,8 @@ def select_requests( normalize_scores: bool, ) -> None: """Select from the score tensor bound to this fixed selector.""" - if ( - scores is not self.prepared_scores.scores - or bool(normalize_scores) != self.prepared_scores.normalize_scores - ): - raise ValueError("union scores do not match their prepared fixed launcher") + if scores is not self.input_scores or bool(normalize_scores) != self.normalize_scores: + raise ValueError("union scores do not match this selector's bound input") self.select_prepared_requests() @@ -378,7 +404,7 @@ def __init__( self.row_seq_lens_flat = self.row_seq_lens.view(-1) self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) self.keep_flat = self.keep.view(-1, self.keep_count) - self._build_prepared_selection_launchers( + self._bind_selection_rows( self.selection_scores_flat, self.row_seq_lens_flat, self.top_indices_i32_flat, @@ -413,7 +439,7 @@ def select_requests( per_layer=self.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - self.prepared_topk() + self._select_top_tokens() class _FixedScoreStagingBuffers: @@ -675,22 +701,43 @@ def __init__( self.copy_pending = False self.page_tables_active = False self.stream = None - self._frozen_score: Optional[FrozenScoreCall] = None + self._score_valid_widths: Optional[torch.Tensor] = None + self._score_aggregation: Optional[str] = None def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str) -> None: - """Compile and bind the phase and score launches for these buffers.""" - if self._frozen_score is not None: + """Bind the per-row score widths and aggregation for these buffers.""" + if self._score_aggregation is not None: raise RuntimeError("TriAttention score launcher is already bound") - self._frozen_score = FrozenScoreCall(self, valid_widths, score_aggregation) - # stage() binds this staging to its first CUDA stream; the frozen - # score call was just built on that same stream. - self.stream = self._frozen_score.stream + if score_aggregation not in ("mean", "max"): + raise ValueError(f"unsupported score aggregation: {score_aggregation}") + self._score_valid_widths = valid_widths + self._score_aggregation = score_aggregation def launch_prepared_score(self) -> torch.Tensor: - """Launch the phase and score calls bound to these buffers.""" - if self._frozen_score is None: + """Launch the phase and score kernels over these buffers.""" + from .triattention_kernels import prepare_mean_phase + + if self._score_aggregation is None: raise RuntimeError("TriAttention score launcher is not bound") - return self._frozen_score() + if self._score_aggregation == "mean": + prepare_mean_phase( + self.round_starts_device, + self.offsets, + self.omega, + self.mean_cos, + self.mean_sin, + self.max_requests, + ) + return self.fused_group.launch( + self.max_requests, + self.valid_seq_lens_device, + self._score_valid_widths, + self.round_starts_device, + self.token_starts_device, + self.mean_cos, + self.mean_sin, + self._score_aggregation, + ) def stage( self, @@ -1102,38 +1149,6 @@ def _draft_protected_tail_capacity(self) -> int: raise RuntimeError("draft KVCacheManagerV2 exposes an invalid protected-tail capacity") return capacity - def warmup_kernels(self) -> None: - """Build every eviction buffer set and compile all its kernels now. - - One call compiles the whole pipeline: score, selection (CuTE top-k - plus the tie-settling call), and the compaction pack. The eviction - kernels never run inside the engine warmup forwards, so without this - the first due eviction round pays the one-time JIT compilation while - serving. Callable as soon as the KV pools exist; idempotent. - """ - self._ensure_calibrated() - layout = self._runtime_kv_layout(self._num_layers_from_manager()) - # A synthetic one-request cohort: the buffers are sized from the - # executor limits, so its geometry only has to be admissible (the - # prompt must cover the SWA landing window when SWA layers exist). - prompt_len = int(layout.swa_window or 0) - synthetic = _PreparedEviction( - request=None, - request_id=-1, - seq_len=prompt_len + 1, - round_start=prompt_len + 1, - prompt_len=prompt_len, - expected_keep_count=prompt_len + self.top_B, - protected_tail=0, - ) - resources = self._eager_resources_for(layout, [synthetic]) - self._batched_compaction_for( - layout=layout, - prepared=[synthetic], - score_staging=resources.score_staging, - keep_set_selector=resources.keep_set_selector, - ) - def _ensure_calibrated(self) -> None: """Resolve calibration once for the first request.""" if self._calibrated: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py deleted file mode 100644 index 5f7fdd2f32e5..000000000000 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernel_warmup.py +++ /dev/null @@ -1,451 +0,0 @@ -# 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. - -"""One home for every kernel warmup in the TriAttention eviction pipeline. - -Triton and CuTE-DSL kernels JIT-compile on first call, and the eviction -kernels never run inside the engine warmup forwards, so each launcher here -compiles its kernels once at build time and then calls the compiled binaries -directly: no per-call dispatch work and no allocations on the eviction path. -Build one launcher per fixed buffer set; call it once per eviction round. -""" - -from typing import Dict, Optional, Tuple - -import torch - -from .triattention_kernels import ( - _score_row_stats_kernel, - _score_union_kernel, - _settle_ties_after_topk_kernel, -) - - -class _FixedScoreStreamMismatch(RuntimeError): - """Raised when fixed score staging buffers are used from another CUDA stream.""" - - -class FrozenTritonKernelCall: - """Call one Triton kernel frozen at build time. - - Triton's standard dispatch costs tens of microseconds of host work per - call; eviction fires its kernels every round, so the grid, bound tensor - set, and constexpr set are frozen once here. ``warmup`` (Triton's own - API) JIT-compiles the kernel at build time and ``__call__`` runs the - compiled binary directly. Constexpr values are passed positionally on - each call, so their order is validated against the kernel's constexpr - parameter declaration order at build time. - """ - - def __init__( - self, - triton_kernel, - bound_tensors: Tuple[torch.Tensor, ...], - constexpr_values: Dict[str, object], - *, - grid: Tuple[int, ...], - num_warps: int, - ) -> None: - params = getattr(triton_kernel, "params", None) - if params is not None: - declared = [param.name for param in params if param.is_constexpr] - if list(constexpr_values.keys()) != declared: - raise ValueError( - f"constexpr order {list(constexpr_values.keys())} must match the " - f"kernel's declaration order {declared}: the frozen call passes " - "them positionally" - ) - self.device = bound_tensors[0].device - self.bound_tensors = tuple(bound_tensors) - self.constexpr_values = dict(constexpr_values) - with torch.cuda.device(self.device): - self.build_stream = torch.cuda.current_stream(self.device) - # warmup() then indexing the compiled cache by grid is the - # documented-by-use Triton pattern for dispatch-free calls; if a - # Triton upgrade changes it, this raises here at build time rather - # than corrupting a later call. - compiled = triton_kernel.warmup( - *self.bound_tensors, - **self.constexpr_values, - num_warps=num_warps, - grid=grid, - ) - self.compiled_kernel_runner = compiled[grid] - - def __call__(self, *call_tensors: torch.Tensor) -> None: - """Run the kernel; ``call_tensors``, if given, substitute the bound tensors.""" - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.build_stream.device, - self.build_stream.cuda_stream, - ): - raise RuntimeError("a frozen Triton kernel call must run on the stream it was built on") - self.compiled_kernel_runner( - *(call_tensors if call_tensors else self.bound_tensors), - *self.constexpr_values.values(), - stream=self.build_stream.cuda_stream, - ) - - -class FrozenScoreCall: - """Phase and trig-score launches frozen over one staging buffer set.""" - - def __init__(self, staging, valid_widths: torch.Tensor, score_aggregation: str) -> None: - from .triattention_kernels import _prepare_mean_phase_kernel, _tri_score_perhead_kernel - - if score_aggregation not in ("mean", "max"): - raise ValueError(f"unsupported score aggregation: {score_aggregation}") - group = staging.fused_group - if ( - valid_widths.shape != (staging.max_requests,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != staging.device - or not valid_widths.is_contiguous() - ): - raise ValueError("prepared score lengths do not match their fixed buffers") - - frequency_block = 1 << (group.num_freqs - 1).bit_length() - phase_pointer_args = ( - staging.round_starts_device, - staging.offsets, - staging.omega, - staging.mean_cos, - staging.mean_sin, - ) - phase_constants = ( - group.num_freqs, - int(staging.offsets.numel()), - frequency_block, - ) - score_pointer_args = ( - *group.pointer_prefix, - staging.valid_seq_lens_device, - valid_widths, - staging.round_starts_device, - staging.token_starts_device, - *group.pointer_middle, - staging.mean_cos.view(-1), - staging.mean_sin.view(-1), - *group.pointer_tail, - group.output, - ) - score_geometry = ( - group.output_width, - group.num_layers, - *group.geometry_args, - ) - score_constants = ( - score_aggregation == "max", - group.token_block, - frequency_block, - ) - self._phase_args = (*phase_pointer_args, *phase_constants) - self._score_args = (*score_pointer_args, *score_geometry, *score_constants) - phase_grid = (staging.max_requests, 1, 1) - score_segments = staging.max_requests * group.num_layers - if score_segments > 65535: - # Segments sit on the y grid axis (CUDA caps y/z at 65535) so the - # unbounded x axis can hold the token tiles of long sequences. - raise ValueError("request*layer segment count exceeds the CUDA grid limit") - score_grid = ( - group.max_ntblk, - score_segments, - group.num_kv_heads, - ) - self.device = staging.device - self.output = group.output - self._phase_runner = None - with torch.cuda.device(self.device): - self.stream = torch.cuda.current_stream(self.device) - if score_aggregation == "mean": - compiled = _prepare_mean_phase_kernel.warmup( - *phase_pointer_args, - NUM_FREQS=phase_constants[0], - NUM_OFFSETS=phase_constants[1], - F_BLOCK=phase_constants[2], - num_warps=1, - grid=phase_grid, - ) - self._phase_runner = compiled[phase_grid] - compiled = _tri_score_perhead_kernel.warmup( - *score_pointer_args, - *score_geometry, - USE_MAX=score_constants[0], - T_BLOCK=score_constants[1], - F_BLOCK=score_constants[2], - grid=score_grid, - ) - self._score_runner = compiled[score_grid] - - def __call__(self) -> torch.Tensor: - current_stream = torch.cuda.current_stream(self.device) - if (current_stream.device, current_stream.cuda_stream) != ( - self.stream.device, - self.stream.cuda_stream, - ): - raise _FixedScoreStreamMismatch( - "TriAttention prepared score is bound to its staging CUDA stream" - ) - if self._phase_runner is not None: - self._phase_runner(*self._phase_args, stream=self.stream.cuda_stream) - self._score_runner(*self._score_args, stream=self.stream.cuda_stream) - return self.output - - -class _PreparedDeterministicTopK: - """Deterministic top-k over fixed row-major buffers, built once. - - The CuTE top-k kernel is fast but breaks score ties arbitrarily and emits - indices in arbitrary order, so a frozen Triton finalizer recomputes the - threshold membership with lowest-index-wins ties, rebases each row by its - prompt offset, and writes sorted ordinals. The CuTE kernel is compiled - here once with owned scratch storage; every later call runs both kernels - over the bound buffers with no dispatch work and no allocations. - """ - - def __init__( - self, - scores: torch.Tensor, - seq_lens: torch.Tensor, - prompt_offsets: torch.Tensor, - provisional_indices: torch.Tensor, - output_indices: torch.Tensor, - keep_count: int, - ) -> None: - rows, width = scores.shape - if not 1 <= keep_count <= width: - raise ValueError("deterministic top-k requires 1 <= keep_count <= width") - - from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops - - self.device = scores.device - self.scores = scores - self.seq_lens = seq_lens - self.provisional_indices = provisional_indices - with torch.cuda.device(self.device): - self.stream = torch.cuda.current_stream(self.device) - self.scratch = torch.empty((rows, 2, width), dtype=torch.int32, device=self.device) - runner = cute_dsl_custom_ops.CuteDSLTopKDecodeSingleCTARunner - key = ( - cute_dsl_custom_ops._TORCH_TO_CUTLASS_DTYPE[torch.float32], - 1 << (width - 1).bit_length(), - keep_count, - 1, - False, - 256, - False, - rows > cute_dsl_custom_ops._get_num_sms(), - ) - runner._compile(*key) - self.compiled_topk = runner.kernel_cache[key] - self.frozen_settle_ties = FrozenTritonKernelCall( - _settle_ties_after_topk_kernel, - (scores, seq_lens, prompt_offsets, provisional_indices, output_indices), - dict( - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, - BLOCK=256, - ), - grid=(rows, 1, 1), - num_warps=4, - ) - - def __call__(self) -> None: - self.compiled_topk( - self.scores, - None, - self.scratch, - None, - self.seq_lens, - self.provisional_indices, - None, - ) - self.frozen_settle_ties() - - -class _PreparedUnionScores: - """Launch fixed union score preparation without Triton JIT dispatch.""" - - def __init__( - self, - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - combined: torch.Tensor, - *, - normalize_scores: bool, - ) -> None: - if scores.ndim != 3: - raise ValueError("prepared union scores require request-major rows") - request_count, rows, width = scores.shape - if ( - not scores.is_cuda - or scores.dtype != torch.float32 - or not scores.is_contiguous() - or valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.shape != (request_count, rows, 1) - or row_mean.dtype != torch.float32 - or row_mean.device != scores.device - or row_inv_std.shape != row_mean.shape - or row_inv_std.dtype != torch.float32 - or row_inv_std.device != scores.device - or combined.shape != (request_count, width) - or combined.dtype != torch.float32 - or combined.device != scores.device - or not valid_widths.is_contiguous() - or not row_mean.is_contiguous() - or not row_inv_std.is_contiguous() - or not combined.is_contiguous() - ): - raise ValueError("prepared union score tensors do not share one fixed geometry") - - normalize_scores = bool(normalize_scores) - # Callers verify score-tensor identity and the normalize flag against - # this launcher before dispatching to it. - self.scores = scores - self.normalize_scores = normalize_scores - self._frozen_stats_call = None - if normalize_scores: - stats_grid = (request_count * rows, 1, 1) - self._frozen_stats_call = FrozenTritonKernelCall( - _score_row_stats_kernel, - (scores, valid_widths, row_mean, row_inv_std), - dict(ROWS=rows, WIDTH=width, BLOCK=256), - grid=stats_grid, - num_warps=4, - ) - union_grid = (request_count, (width + 31) // 32, 1) - self._frozen_union_call = FrozenTritonKernelCall( - _score_union_kernel, - (scores, valid_widths, row_mean, row_inv_std, combined), - dict(ROWS=rows, WIDTH=width, NORMALIZE=normalize_scores, BLOCK=32), - grid=union_grid, - num_warps=1, - ) - - def __call__(self) -> None: - if self._frozen_stats_call is not None: - self._frozen_stats_call() - self._frozen_union_call() - - -# Launch shape of the move-index packing kernel: tokens per program along the -# move axis, and its warp count. -_PACK_BLOCK_TOKENS = 256 -_PACK_NUM_WARPS = 4 - - -def frozen_move_index_pack( - kept_token_ordinals: torch.Tensor, - valid_sequence_lengths: torch.Tensor, - move_source_offsets: torch.Tensor, - move_source_indices: torch.Tensor, - *, - eviction_mode: str, - decode_keep_count: int, - num_dense_layers: int, - num_kv_heads: int, - max_protected_tail: int, - swa_window: int, - swa_move_source_offsets: Optional[torch.Tensor], - swa_move_source_indices: Optional[torch.Tensor], -) -> FrozenTritonKernelCall: - """Build one frozen call of the move-index packing kernel. - - The kernel reads the kept-token ordinals and each request's valid length - and writes the packed per-(layer, head) move source indices consumed by - the C++ compact launches. Only the caller-provided selection tensors are - validated here; the move buffers are allocated by this module. - """ - per_layer = eviction_mode == "per_layer_perhead" - union = eviction_mode == "union" - request_count = int(kept_token_ordinals.shape[0]) if kept_token_ordinals.ndim else 0 - if union: - selection_rows = 1 - elif per_layer: - selection_rows = num_dense_layers * num_kv_heads - else: - selection_rows = num_kv_heads - selection_prefix = (request_count,) if union else (request_count, selection_rows) - # Selection rows carry decode-only kept ordinals (already absolute), so - # the rectangle is prompt-length independent. - expected_selection = (*selection_prefix, decode_keep_count) - if ( - request_count <= 0 - or tuple(kept_token_ordinals.shape) != expected_selection - or valid_sequence_lengths.shape != (request_count,) - ): - raise ValueError( - f"prepared compaction packing expects kept ordinals of shape " - f"{expected_selection} and one valid length per request; got " - f"{tuple(kept_token_ordinals.shape)} and " - f"{tuple(valid_sequence_lengths.shape)}" - ) - - if swa_move_source_indices is not None: - swa_offsets_arg = swa_move_source_offsets - swa_indices_arg = swa_move_source_indices - swa_total = int(swa_move_source_indices.shape[-1]) - else: - # HAS_SWA specializes all corresponding loads and stores away. - swa_offsets_arg = move_source_offsets - swa_indices_arg = move_source_indices - swa_total = 0 - - from .triattention_kernels import _pack_compaction_sources_kernel - - max_move = decode_keep_count + max_protected_tail - if swa_total: - max_move = max(max_move, swa_window + max_protected_tail) - packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads - grid = ( - request_count, - packed_row_count, - (max_move + _PACK_BLOCK_TOKENS - 1) // _PACK_BLOCK_TOKENS, - ) - bound_tensors = ( - kept_token_ordinals, - valid_sequence_lengths, - move_source_offsets, - move_source_indices, - swa_offsets_arg, - swa_indices_arg, - ) - # Ordered to match the kernel's constexpr parameter declaration: the - # frozen call passes these by position. - constexpr_values = dict( - DENSE_TOTAL=int(move_source_indices.shape[-1]), - SWA_TOTAL=swa_total, - SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=decode_keep_count, - KEEP_COUNT=decode_keep_count, - NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window, - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=swa_total > 0, - BLOCK=_PACK_BLOCK_TOKENS, - ) - return FrozenTritonKernelCall( - _pack_compaction_sources_kernel, - bound_tensors, - constexpr_values, - grid=grid, - num_warps=_PACK_NUM_WARPS, - ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index 99ba128e30de..1a1181b2b9ad 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from unittest import mock import pytest import torch @@ -201,12 +200,7 @@ def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_ normalize_scores=normalize_scores, ) selector.valid_widths.copy_(valid_widths) - with mock.patch.object( - triattention_kernels, - "prepare_union_scores", - side_effect=AssertionError("checked Triton wrapper was called"), - ): - selector.select_prepared_requests() + selector.select_prepared_requests() actual_combined = selector.combined.cpu() actual_keep = selector.keep.cpu() expected_combined = reference_combined.cpu() diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index dbb6dae26e25..d81ec96061e6 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1704,22 +1704,11 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun staging.offsets = offsets staging.omega = omega staging.stream = None - staging._frozen_score = None + staging._score_valid_widths = None + staging._score_aggregation = None staging.bind_score_launcher(valid_widths, aggregation) group.output.fill_(score_sentinel) - with ( - mock.patch.object( - triattention_kernels, - "prepare_mean_phase", - side_effect=AssertionError("checked phase wrapper was called"), - ), - mock.patch.object( - group, - "launch", - side_effect=AssertionError("checked score wrapper was called"), - ), - ): - fixed = staging.launch_prepared_score().clone() + fixed = staging.launch_prepared_score().clone() torch.testing.assert_close(fixed, checked, rtol=0, atol=0) assert valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] @@ -1779,19 +1768,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun ).clone() group.output.fill_(score_sentinel) valid_widths.fill_(-1) - with ( - mock.patch.object( - triattention_kernels, - "prepare_mean_phase", - side_effect=AssertionError("checked phase wrapper was called"), - ), - mock.patch.object( - group, - "launch", - side_effect=AssertionError("checked score wrapper was called"), - ), - ): - second_launch = staging.launch_prepared_score().clone() + second_launch = staging.launch_prepared_score().clone() torch.testing.assert_close(second_launch, checked_second, rtol=0, atol=0) assert torch.equal(valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) From 04a7ad556180c6b5b16978d0dc71b7b6cc32eba0 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 08:02:58 -0700 Subject: [PATCH 041/178] [None][fix] Keep score launches on the staging CUDA stream Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index e4d36491ce9a..f2df658e7352 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -719,6 +719,16 @@ def launch_prepared_score(self) -> torch.Tensor: if self._score_aggregation is None: raise RuntimeError("TriAttention score launcher is not bound") + stream = torch.cuda.current_stream(self.device) + if self.stream is None: + self.stream = stream + elif (stream.device, stream.cuda_stream) != ( + self.stream.device, + self.stream.cuda_stream, + ): + raise _FixedScoreStreamMismatch( + "TriAttention score launches must stay on the staging CUDA stream" + ) if self._score_aggregation == "mean": prepare_mean_phase( self.round_starts_device, From f4a895e54fc8678448e4e2431fe4d82db83f864a Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 08:45:27 -0700 Subject: [PATCH 042/178] [None][fix] Bind the SWA compact family to the staged move offsets A local initializer shadowed the constructor-supplied offsets row, so the sliding-window family kept its construction-time offsets and compacted every workbench slot each round. Padded slots then produced negative source ordinals and an illegal memory access on sliding-window models. Signed-off-by: tianruih --- .../triattention/compaction.py | 6 ++- .../test_triattention_eager.py | 50 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index b14ebfe4a03d..ab94a4a5f69f 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -465,8 +465,12 @@ def __init__( self.swa_window = 0 self.swa_destination_bases = None swa_move_indices = None - swa_move_offsets = None swa_entries = [] + if not self.swa_layers: + # No SWA family: drop the unused offsets row. With SWA layers the + # constructor argument must stay live so the family reads the + # per-round staged offsets instead of its construction-time sizes. + swa_move_offsets = None if self.swa_layers: if swa_window is None or swa_window <= 0: raise ValueError("SWA compaction requires a valid retained window") diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py index 1a1181b2b9ad..0a21d9b176e4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py @@ -1006,3 +1006,53 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): swa_after.index_select(2, swa_destination), swa_before.index_select(2, swa_source), ) + + +def test_workbench_families_read_the_staged_move_offsets_rows(): + """Every cache family must consume the caller-staged offsets row. + + A family that silently falls back to its construction-time offsets + compacts workbench slots that are not in the staged cohort; on + sliding-window models the padded slots then produce negative source + ordinals and an illegal memory access. Binding the staged row by + reference is part of the constructor contract. + """ + device = torch.device("cuda", torch.cuda.current_device()) + dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) + swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) + pools = [ + torch.zeros(6, 2, 1, 4, 16, dtype=torch.float32, device=device), + torch.zeros(6, 2, 1, 4, 16, dtype=torch.float32, device=device), + ] + keep = torch.tensor([[2, 4, 5, 7], [2, 3, 5, 6]], dtype=torch.int32, device=device) + staged_rows = torch.zeros(2, 3, dtype=torch.int32, device=device) + dense_offsets_row = staged_rows[0] + swa_offsets_row = staged_rows[1] + compaction = BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=pools, + dense_layers=[0], + swa_layers=[1], + layer_group_representative={0: 0}, + layer_pool_keys=[("dense", 0), ("swa", 0)], + kept_token_ordinals=keep, + valid_sequence_lengths=torch.tensor([8, 7], dtype=torch.int32, device=device), + kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), + page_table_slots={0: 0, 1: 1}, + request_count=2, + prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), + decode_keep_count=4, + swa_window=2, + protected_tail_capacity=2, + dense_move_offsets=dense_offsets_row, + swa_move_offsets=swa_offsets_row, + ) + assert ( + compaction.target_dense_compaction.move_source_offsets.data_ptr() + == dense_offsets_row.data_ptr() + ) + assert compaction.target_swa_compaction is not None + assert ( + compaction.target_swa_compaction.move_source_offsets.data_ptr() + == swa_offsets_row.data_ptr() + ) From 1e035c686d4e4044dc1e6fdc9398083b6e818fad Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 17 Jul 2026 19:40:10 -0700 Subject: [PATCH 043/178] [None][fix] Skip the steady-generation fast prepare under KV compression The steady-state decode fast path assumes positions and cached-token counts advance in lockstep and only refreshes them by one per step. A KV-cache compression manager shrinks the physical cache mid-generation, so both values it serves go stale after the first eviction round: the kernels then read past the compacted region every step. Refuse to record the steady cache when the cache manager is compression-managed, which keeps every step on the full prepare path that rebuilds both values. Signed-off-by: tianruih --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6b09c03f87c2..904ad2f22b5f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4570,7 +4570,13 @@ def previous_seq_slots_device(): and not _has_any_multimodal_request and not multimodal_params_list and not lora_params and attn_metadata.padded_num_tokens is None - and self._get_position_id_offset() == 0): + and self._get_position_id_offset() == 0 + # A KV-cache compression manager shrinks the physical + # cache mid-generation, so positions and cached-token + # counts stop advancing in lockstep; keep the full + # prepare path, which rebuilds both every step. + and not getattr(kv_cache_manager, + "kv_compression_manages_history", False)): self._steady_gen_positions_pinned[:_n_gen].copy_( torch.as_tensor(num_cached_tokens_snapshot, dtype=torch.int)) From a6bf4eae45951addfb6cf7763153071639aa1aa1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Sat, 18 Jul 2026 11:32:41 -0700 Subject: [PATCH 044/178] [None][perf] Cut TriAttention score kernel time 1.5x via launch shape and mask elision The scoring kernel used 64-token tiles, which pushed register usage to the 255-per-thread ceiling (with spills at 2 warps) and collapsed SM occupancy; its launch also relied on Triton default warp count. Shrink the token tile to 16 and pin num_warps=4 (measured fastest across the sweep). Additionally compile out the frequency-axis masking when the padded frequency block already equals the real frequency count (F_EXACT), removing no-op predicates and selects from the hot loop. Outputs are bit-identical (torch.equal on scores and valid widths across the launch sweep). Measured on B200: score kernel 7.86ms -> 5.26ms per eviction round at batch 32; round decode stall 10.72ms -> 8.16ms. AIME bands reproduce exactly (Qwen3-8B 55.00 n=8, GPT-OSS-20B 48.75 n=8); eagle3 end-to-end leg passes. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 62 +++++++++++++++---- 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 4f4e93513883..4b88170e7a38 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -148,6 +148,9 @@ def _tri_score_perhead_kernel( USE_MAX: tl.constexpr, T_BLOCK: tl.constexpr, F_BLOCK: tl.constexpr, + F_EXACT: tl.constexpr, # F_BLOCK == num_freqs: every f-lane is real, so the + # f-axis masks/wheres below are no-ops and are compiled out (fewer + # predicates and selects on the hot path; bit-identical results). ): # Token tiles ride the fastest grid axis: adjacent programs then walk # consecutive K pages of one (request, layer, head) and reuse its @@ -209,9 +212,14 @@ def _tri_score_perhead_kernel( tok_base = phys_page * s_page + slot * s_slot # [T_BLOCK] int64 # per-request 'mean'-path phase + shared freq scale. - mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) - msin = tl.load(mean_sin_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) - fss = tl.load(freq_scale_sq_ptr + f, mask=f_mask, other=0.0) + if F_EXACT: + mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f) + msin = tl.load(mean_sin_ptr + req_id * num_freqs + f) + fss = tl.load(freq_scale_sq_ptr + f) + else: + mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) + msin = tl.load(mean_sin_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) + fss = tl.load(freq_scale_sq_ptr + f, mask=f_mask, other=0.0) # ---- PER-HEAD (position + mlr), GQA-deduped, NO head reduction ---- # This program scores ONE KV head's token tile for the group_size q-heads @@ -219,7 +227,10 @@ def _tri_score_perhead_kernel( # h = kv_head*group_size + qg keeps query-head order 0..num_q_heads-1, so # every head's math is bit-for-bit identical to the looped variant. group_size = num_q_heads // num_kv_heads - load_mask = t_mask[:, None] & f_mask[None, :] + if F_EXACT: + load_mask = tl.broadcast_to(t_mask[:, None], (T_BLOCK, F_BLOCK)) + else: + load_mask = t_mask[:, None] & f_mask[None, :] off_re = f64[None, :] * s_dim off_im = (num_freqs + f64[None, :]) * s_dim @@ -233,9 +244,14 @@ def _tri_score_perhead_kernel( while qg < group_size: h = kv_head * group_size + qg calib_off = (layer_id.to(tl.int64) * num_q_heads + h) * num_freqs - qre = tl.load(q_real_ptr + calib_off + f, mask=f_mask, other=0.0) - qim = tl.load(q_imag_ptr + calib_off + f, mask=f_mask, other=0.0) - mlrc = tl.load(mlr_coef_ptr + calib_off + f, mask=f_mask, other=0.0) + if F_EXACT: + qre = tl.load(q_real_ptr + calib_off + f) + qim = tl.load(q_imag_ptr + calib_off + f) + mlrc = tl.load(mlr_coef_ptr + calib_off + f) + else: + qre = tl.load(q_real_ptr + calib_off + f, mask=f_mask, other=0.0) + qim = tl.load(q_imag_ptr + calib_off + f, mask=f_mask, other=0.0) + mlrc = tl.load(mlr_coef_ptr + calib_off + f, mask=f_mask, other=0.0) # complex product Q . conj(K) -- the trig importance score. prod_real = qre[None, :] * k_re + qim[None, :] * k_im @@ -248,22 +264,34 @@ def _tri_score_perhead_kernel( o = 0 while o < num_offsets: off = tl.load(offsets_ptr + o) - om = tl.load(omega_ptr + f, mask=f_mask, other=0.0) + if F_EXACT: + om = tl.load(omega_ptr + f) + else: + om = tl.load(omega_ptr + f, mask=f_mask, other=0.0) phase = (rstart + off) * om cphase = tl.cos(phase) sphase = tl.sin(phase) per_f = fss[None, :] * (prod_real * cphase[None, :] - prod_imag * sphase[None, :]) - offset_score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) + if F_EXACT: + offset_score = tl.sum(per_f, axis=1) + else: + offset_score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) score = tl.maximum(score, offset_score) o += 1 else: # 'mean': offset loop collapsed into mean_cos/mean_sin. per_f = fss[None, :] * (prod_real * mcos[None, :] - prod_imag * msin[None, :]) - score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) + if F_EXACT: + score = tl.sum(per_f, axis=1) + else: + score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) # position-INDEPENDENT MLR term (reuses the per-KV-head |K|). mlr_f = kmag * mlrc[None, :] * fss[None, :] - mlr = tl.sum(tl.where(f_mask[None, :], mlr_f, 0.0), axis=1) + if F_EXACT: + mlr = tl.sum(mlr_f, axis=1) + else: + mlr = tl.sum(tl.where(f_mask[None, :], mlr_f, 0.0), axis=1) # Segments are request-major then layer-major. Write the decode-only # score directly in the selector's [request, layer, head, token] layout. @@ -284,12 +312,17 @@ def _launch_tri_score_perhead( """Launch the shared score ABI for eager and fixed metadata owners.""" if score_aggregation not in ("mean", "max"): raise ValueError(f"unsupported score aggregation: {score_aggregation}") + f_block = triton.next_power_of_2(num_freqs) _tri_score_perhead_kernel[grid]( *pointer_args, *geometry_args, USE_MAX=(score_aggregation == "max"), T_BLOCK=token_block, - F_BLOCK=triton.next_power_of_2(num_freqs), + F_BLOCK=f_block, + F_EXACT=(f_block == num_freqs), + # 4 warps measured fastest across token blocks; the default was never + # tuned for this kernel (sibling kernels pin their warp counts too). + num_warps=4, ) @@ -407,7 +440,10 @@ def __init__( ) slot_idx = slots_t.repeat(max_requests) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - self.token_block = 64 + # 16-token tiles keep the per-program fp32 working set small enough to + # avoid the register-pressure occupancy collapse measured at 64-token + # tiles (255 regs/thread + spills -> 63 regs, ~1.5x faster end to end). + self.token_block = 16 self.max_ntblk = (self.output_width + self.token_block - 1) // self.token_block self.output = torch.empty( max_requests, From 87484dac076ee281ebbfcd153d4cf8b37982c0b0 Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 19 Jul 2026 03:35:56 -0700 Subject: [PATCH 045/178] [None][perf] Move TriAttention scoring to a folded CUDA op (4.4x score kernel) The eviction scoring kernel was instruction-issue bound: register-pressure occupancy limits and per-token re-evaluation of token-independent trig math kept it ~10x off its memory roofline. Fold the token-independent factors into per-(request, layer, head, frequency) coefficients computed once per round by a small kernel, then score with a one-thread-per-token CUDA kernel (vectorized 16-byte loads on the standard bf16/fp16 layout, generic strided scalar path otherwise). The scoring stage is shared by all three selection modes, which speed up identically. Measured on B200, Qwen3-8B, batch 32: score kernel 7.86ms -> 1.77ms per eviction round (4.4x), round decode stall 10.72ms -> 5.06ms; batch 256 single-wave: 62.8ms -> 13.8ms, stall 78.5ms -> 29.6ms. GPT-OSS-20B scoring drops 2.7x (its round cost is dominated by host-side resize, unchanged). Outputs are not bit-identical to the previous kernel (floating-point reassociation, ~1e-5): AIME bands reproduce within run noise (Qwen3-8B 57.08 vs 55.00, GPT-OSS-20B 47.08 vs 48.75, n=8 each), and the eagle3 and speculative-decode batteries pass. The op covers every geometry the Triton kernel supported (arbitrary frequency counts, page sizes, strides, mean and max aggregation, bf16/fp16/ fp32 pools, generic GQA group sizes), so the Triton scoring kernel is removed; unit tests now compare the op against an independent PyTorch oracle. Quantized fp8/int8 pools are supported functionally: per-layer dequantization scales fold into the coefficients at zero hot-loop cost (kernel-level tests only; no end-to-end quantized producer exists yet). Signed-off-by: tianruih --- .../triAttentionScoreKernels.cu | 698 ++++++++++++++++++ .../triAttentionScoreKernels.h | 144 ++++ cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp | 323 ++++++++ .../triattention/triattention.py | 2 +- .../triattention/triattention_kernels.py | 444 +++++------ .../test_triattention_score_ops.py | 690 +++++++++++++++++ 7 files changed, 2056 insertions(+), 248 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu create mode 100644 cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h create mode 100644 cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu new file mode 100644 index 000000000000..48f96760d8ba --- /dev/null +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu @@ -0,0 +1,698 @@ +/* + * 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. + */ + +// ============================================================================ +// TriAttention folded score kernels +// ============================================================================ +// +// Scores every cached decode token of every scored layer for KV eviction. +// The per-round trigonometry is folded into coefficient tables first +// (foldScoreCoefficientsLaunch), so the hot kernel is a pure fused +// multiply-add stream over paged KV: +// +// score(t, h) = sum_f K_re(t,f)*c_re + K_im(t,f)*c_im + |K(t,f)|*c_mlr +// +// One thread scores one token across all frequencies; a 128-thread CTA covers +// 128 consecutive tokens of one (request, layer, KV-head) segment. There are +// no shuffles, no shared memory, and no barriers: each thread keeps one fused +// accumulator per query head of its GQA group ("mean" aggregation) or one +// partial sum per offset plane ("max" aggregation, where max over offsets +// does not commute through the frequency sum). Coefficient loads are +// lane-uniform 16-byte reads served by L1 broadcast; K loads are 16-byte +// chunks of 8 frequencies when the pool layout allows it, otherwise a fully +// strided scalar path runs the same math. +// +// The kernel reassociates the frequency reduction relative to the Triton +// reference (sequential chunks instead of a block-wide tree), so results are +// tolerance-equal, not bit-equal. The valid-width side store IS replicated +// bit-exactly: first tile, first head, first segment of each request, thread +// 0, before any early-out. +// +// This file must NOT be compiled with --use_fast_math: the fold kernel's +// cosf/sinf and the scalar-path precision are part of the accuracy contract. +// The approximate square root below is an explicit, scoped opt-in instead. +// ============================================================================ + +#include +#include +#include +#include +#include +#include + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::tri_attention_score +{ + +namespace +{ + +// |K| uses the hardware approximate square root (one MUFU op). This matches +// the Triton reference's tl.sqrt lowering and is gated by the unit suite's +// tolerance comparison. Define TRTLLM_TRI_ATTENTION_IEEE_SQRT to restore the +// IEEE sqrtf sequence if a future geometry needs the extra bits. +__device__ __forceinline__ float triSqrtApprox(float x) +{ +#ifdef TRTLLM_TRI_ATTENTION_IEEE_SQRT + return sqrtf(x); +#else + float y; + asm("sqrt.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +#endif +} + +template +__device__ __forceinline__ float toFloat(T value); + +template <> +__device__ __forceinline__ float toFloat<__nv_bfloat16>(__nv_bfloat16 value) +{ + return __bfloat162float(value); +} + +template <> +__device__ __forceinline__ float toFloat(half value) +{ + return __half2float(value); +} + +template <> +__device__ __forceinline__ float toFloat(float value) +{ + return value; +} + +// Quantized pool elements are converted RAW (no scale applied): the per-layer +// dequantization scale is folded into the coefficient tables at fold time, so +// the hot loop stays a pure convert-and-FMA stream. +template <> +__device__ __forceinline__ float toFloat<__nv_fp8_e4m3>(__nv_fp8_e4m3 value) +{ + return static_cast(value); +} + +template <> +__device__ __forceinline__ float toFloat(int8_t value) +{ + return static_cast(value); +} + +// Unpack one 16-byte K chunk (8 consecutive frequencies) to fp32. +template +__device__ __forceinline__ void unpackChunk8(uint4 v, float* dst); + +template <> +__device__ __forceinline__ void unpackChunk8<__nv_bfloat16>(uint4 v, float* dst) +{ + auto const* p = reinterpret_cast<__nv_bfloat162 const*>(&v); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + float2 f2 = __bfloat1622float2(p[i]); + dst[2 * i] = f2.x; + dst[2 * i + 1] = f2.y; + } +} + +template <> +__device__ __forceinline__ void unpackChunk8(uint4 v, float* dst) +{ + auto const* p = reinterpret_cast<__half2 const*>(&v); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + float2 f2 = __half22float2(p[i]); + dst[2 * i] = f2.x; + dst[2 * i + 1] = f2.y; + } +} + +// Predicated 16-byte loads of one 8-frequency chunk of one token row. Row +// layout with unit frequency stride is [re f0..F-1 | im f0..F-1], so the +// imaginary half of chunk c sits sizeof(T) * numFreqs bytes further. +template +__device__ __forceinline__ void scoreLoadChunk(char const* row, bool valid, int numFreqs, int c, uint4& re4, uint4& im4) +{ + re4 = make_uint4(0u, 0u, 0u, 0u); + im4 = make_uint4(0u, 0u, 0u, 0u); + if (valid) + { + int const fByte = c * 8 * static_cast(sizeof(T)); + re4 = __ldg(reinterpret_cast(row + fByte)); + im4 = __ldg(reinterpret_cast(row + static_cast(sizeof(T)) * numFreqs + fByte)); + } +} + +// Accumulate one 8-frequency chunk into this thread's per-head accumulators. +// coff0 = flat index of (request, layer, first head of this block, chunk +// frequency 0) in the coefficient tables; all coefficient reads are +// lane-uniform 16-byte loads. |K| is computed once per (token, frequency) +// BEFORE the head loop so the GROUP heads share it from registers. +template +__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, int64_t coff0, int64_t planeStride, + uint4 re4, uint4 im4, float* accMean, float* accMlr, float* accPos) +{ + float kRe[8], kIm[8], kMag[8]; + unpackChunk8(re4, kRe); + unpackChunk8(im4, kIm); +#pragma unroll + for (int i = 0; i < 8; ++i) + { + kMag[i] = triSqrtApprox(kRe[i] * kRe[i] + kIm[i] * kIm[i]); + } +#pragma unroll + for (int hg = 0; hg < GROUP; ++hg) + { + int64_t const coff = coff0 + static_cast(hg) * a.numFreqs; + float const* cmp = a.cMlr + coff; + float4 const cm0 = __ldg(reinterpret_cast(cmp)); + float4 const cm1 = __ldg(reinterpret_cast(cmp + 4)); + float const cml[8] = {cm0.x, cm0.y, cm0.z, cm0.w, cm1.x, cm1.y, cm1.z, cm1.w}; + if constexpr (USE_MAX) + { + // The |K| term is offset independent: keep it in its own + // accumulator and add it after the offset max at store time. + float m = accMlr[hg]; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + m = fmaf(kMag[i], cml[i], m); + } + accMlr[hg] = m; + // Unrolled with a live-plane guard so accPos indexing stays + // static (register-resident) despite the runtime offset count. +#pragma unroll + for (int o = 0; o < kMaxScoreOffsets; ++o) + { + if (o < a.numOffsets) + { + float const* crp = a.cRe + o * planeStride + coff; + float const* cip = a.cIm + o * planeStride + coff; + float4 const cr0 = __ldg(reinterpret_cast(crp)); + float4 const cr1 = __ldg(reinterpret_cast(crp + 4)); + float4 const ci0 = __ldg(reinterpret_cast(cip)); + float4 const ci1 = __ldg(reinterpret_cast(cip + 4)); + float const cre[8] = {cr0.x, cr0.y, cr0.z, cr0.w, cr1.x, cr1.y, cr1.z, cr1.w}; + float const cim[8] = {ci0.x, ci0.y, ci0.z, ci0.w, ci1.x, ci1.y, ci1.z, ci1.w}; + float p = accPos[hg * kMaxScoreOffsets + o]; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + p = fmaf(kRe[i], cre[i], fmaf(kIm[i], cim[i], p)); + } + accPos[hg * kMaxScoreOffsets + o] = p; + } + } + } + else + { + float const* crp = a.cRe + coff; + float const* cip = a.cIm + coff; + float4 const cr0 = __ldg(reinterpret_cast(crp)); + float4 const cr1 = __ldg(reinterpret_cast(crp + 4)); + float4 const ci0 = __ldg(reinterpret_cast(cip)); + float4 const ci1 = __ldg(reinterpret_cast(cip + 4)); + float const cre[8] = {cr0.x, cr0.y, cr0.z, cr0.w, cr1.x, cr1.y, cr1.z, cr1.w}; + float const cim[8] = {ci0.x, ci0.y, ci0.z, ci0.w, ci1.x, ci1.y, ci1.z, ci1.w}; + float t = accMean[hg]; +#pragma unroll + for (int i = 0; i < 8; ++i) + { + // 3 fused multiply-adds per (token, frequency, head): the + // position and |K| terms share one accumulator chain, valid + // because the mean path has a single coefficient plane. + t = fmaf(kRe[i], cre[i], fmaf(kIm[i], cim[i], fmaf(kMag[i], cml[i], t))); + } + accMean[hg] = t; + } + } +} + +// Vectorized token-per-thread score kernel. Grid: x = 128-token tiles of the +// page-aligned decode span, y = (request, layer) segment, z = KV head (or one +// query head when a.zIsQueryHead covers GQA group sizes with no dedicated +// GROUP instantiation). STATIC_CHUNKS == 8 pins the production 64-frequency +// shape at compile time (fully unrolled chunk loop, the tuned register +// budget); STATIC_CHUNKS == 0 loops numFreqs / 8 chunks at runtime. +// +// minBlocksPerMultiprocessor = 7: tighter caps force ptxas into ~48-56 +// registers with stack spills in the fully unrolled inner loops; 7 CTAs/SM +// admits the ~72-register spill-free allocation this kernel was tuned at. +template +__global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKernel(FoldedScoreParams a) +{ + int const seg = blockIdx.y; + int const reqId = a.segRequestIds[seg]; + int const seqLen = a.requestSeqLens[reqId]; + int const tokenStart = a.requestTokenStarts[reqId]; + // Valid-width side store: same predicate as the Triton reference, before + // any early-out, so it fires exactly once per request even when the + // decode region is empty. + if (blockIdx.x == 0 && blockIdx.z == 0 && (seg % a.numLayers) == 0 && threadIdx.x == 0) + { + a.validWidthOut[reqId] = seqLen - tokenStart; + } + // Tiles start on the first page of the decode region: with 32-token pages + // every warp then covers exactly one page and its lane-identical page + // lookup collapses to one L1 broadcast. + int const alignedStart = (tokenStart / a.tokensPerBlock) * a.tokensPerBlock; + if (alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads >= seqLen) + { + return; // CTA-uniform: the whole tile is past this sequence + } + int const absT = alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads + static_cast(threadIdx.x); + bool const valid = absT >= tokenStart && absT < seqLen; + + int kvHead; + int headBase; + if (a.zIsQueryHead) + { + headBase = static_cast(blockIdx.z); + kvHead = headBase / (a.numQueryHeads / a.numKvHeads); + } + else + { + kvHead = static_cast(blockIdx.z); + headBase = kvHead * GROUP; + } + + int const layerId = a.segLayerIds[seg]; + int const page = absT / a.tokensPerBlock; + int const slot = absT - page * a.tokensPerBlock; + // Threads past the sequence tail must not touch the page table (their + // page ordinal may exceed the staged row); their K loads are predicated + // off below, so page 0 is a safe placeholder. + int encoded = 0; + if (absT < seqLen) + { + encoded = a.blockOffsets[a.segPageOffsets[seg] + page]; + } + // Page-table entries count K/V role pages; kvFactor converts to the + // layer pool page holding the K plane. + auto const physPage = static_cast(encoded / a.kvFactor); + auto const* layerBase = reinterpret_cast(a.layerBaseAddrs[layerId]); + char const* row = layerBase + + static_cast(sizeof(T)) + * (physPage * a.stridePage + static_cast(kvHead) * a.strideKvHead + + static_cast(slot) * a.strideSlot); + + int64_t const coff0 = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + + static_cast(headBase) * a.numFreqs; + int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads + * static_cast(a.numFreqs); + + float accMean[GROUP]; + float accMlr[USE_MAX ? GROUP : 1]; + float accPos[USE_MAX ? GROUP * kMaxScoreOffsets : 1]; +#pragma unroll + for (int hg = 0; hg < GROUP; ++hg) + { + accMean[hg] = 0.0f; + } + if constexpr (USE_MAX) + { +#pragma unroll + for (int hg = 0; hg < GROUP; ++hg) + { + accMlr[hg] = 0.0f; + } +#pragma unroll + for (int i = 0; i < GROUP * kMaxScoreOffsets; ++i) + { + accPos[i] = 0.0f; + } + } + + if constexpr (STATIC_CHUNKS > 0) + { +#pragma unroll + for (int c = 0; c < STATIC_CHUNKS; ++c) + { + uint4 re4, im4; + scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); + scoreComputeChunk(a, coff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + } + } + else + { + int const chunkCount = a.numFreqs / 8; + for (int c = 0; c < chunkCount; ++c) + { + uint4 re4, im4; + scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); + scoreComputeChunk(a, coff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + } + } + + // Store: per (tile, head) the CTA writes one contiguous fp32 row; the + // predicate replicates the reference token mask and decode-region clip. + int const tDec = absT - tokenStart; + if (tDec >= 0 && absT < seqLen) + { +#pragma unroll + for (int hg = 0; hg < GROUP; ++hg) + { + float score; + if constexpr (USE_MAX) + { + float best = -INFINITY; +#pragma unroll + for (int o = 0; o < kMaxScoreOffsets; ++o) + { + if (o < a.numOffsets) + { + best = fmaxf(best, accPos[hg * kMaxScoreOffsets + o]); + } + } + score = best + accMlr[hg]; + } + else + { + score = accMean[hg]; + } + a.out[(static_cast(seg) * a.numQueryHeads + headBase + hg) * a.outputWidth + tDec] = score; + } + } +} + +// Strided scalar score kernel: same token-per-thread mapping and math as the +// vectorized kernel, but every K element is loaded through the full runtime +// stride set (any frequency count, any element stride, fp32 pools included). +// The GQA head loop runs at runtime, so any group size is covered. |K| is +// recomputed per head from the same loads — bit-identical to hoisting it. +template +__global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(FoldedScoreParams a) +{ + int const seg = blockIdx.y; + int const reqId = a.segRequestIds[seg]; + int const seqLen = a.requestSeqLens[reqId]; + int const tokenStart = a.requestTokenStarts[reqId]; + // Valid-width side store: identical contract to the vectorized kernel. + if (blockIdx.x == 0 && blockIdx.z == 0 && (seg % a.numLayers) == 0 && threadIdx.x == 0) + { + a.validWidthOut[reqId] = seqLen - tokenStart; + } + int const alignedStart = (tokenStart / a.tokensPerBlock) * a.tokensPerBlock; + if (alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads >= seqLen) + { + return; + } + int const absT = alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads + static_cast(threadIdx.x); + bool const valid = absT >= tokenStart && absT < seqLen; + + int const kvHead = blockIdx.z; + int const groupSize = a.numQueryHeads / a.numKvHeads; + int const layerId = a.segLayerIds[seg]; + int const page = absT / a.tokensPerBlock; + int const slot = absT - page * a.tokensPerBlock; + int encoded = 0; + if (absT < seqLen) + { + encoded = a.blockOffsets[a.segPageOffsets[seg] + page]; + } + auto const physPage = static_cast(encoded / a.kvFactor); + auto const* row = reinterpret_cast(a.layerBaseAddrs[layerId]) + physPage * a.stridePage + + static_cast(kvHead) * a.strideKvHead + static_cast(slot) * a.strideSlot; + + int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads + * static_cast(a.numFreqs); + int const tDec = absT - tokenStart; + bool const store = tDec >= 0 && absT < seqLen; + + for (int hg = 0; hg < groupSize; ++hg) + { + int const h = kvHead * groupSize + hg; + int64_t const coff + = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + + static_cast(h) * a.numFreqs; + float acc = 0.0f; + float accMlr = 0.0f; + float accPos[kMaxScoreOffsets]; +#pragma unroll + for (int o = 0; o < kMaxScoreOffsets; ++o) + { + accPos[o] = 0.0f; + } + for (int f = 0; f < a.numFreqs; ++f) + { + float kRe = 0.0f; + float kIm = 0.0f; + if (valid) + { + kRe = toFloat(row[static_cast(f) * a.strideDim]); + kIm = toFloat(row[static_cast(a.numFreqs + f) * a.strideDim]); + } + float const kMag = triSqrtApprox(kRe * kRe + kIm * kIm); + if constexpr (USE_MAX) + { + accMlr = fmaf(kMag, a.cMlr[coff + f], accMlr); +#pragma unroll + for (int o = 0; o < kMaxScoreOffsets; ++o) + { + if (o < a.numOffsets) + { + accPos[o] = fmaf(kRe, a.cRe[o * planeStride + coff + f], + fmaf(kIm, a.cIm[o * planeStride + coff + f], accPos[o])); + } + } + } + else + { + acc = fmaf(kRe, a.cRe[coff + f], fmaf(kIm, a.cIm[coff + f], fmaf(kMag, a.cMlr[coff + f], acc))); + } + } + if (store) + { + float score; + if constexpr (USE_MAX) + { + float best = -INFINITY; +#pragma unroll + for (int o = 0; o < kMaxScoreOffsets; ++o) + { + if (o < a.numOffsets) + { + best = fmaxf(best, accPos[o]); + } + } + score = best + accMlr; + } + else + { + score = acc; + } + a.out[(static_cast(seg) * a.numQueryHeads + h) * a.outputWidth + tDec] = score; + } + } +} + +// Per-round coefficient fold: one thread per (request, layer, head, freq) +// element. On the max path each thread additionally writes one c_re/c_im +// value per offset plane (planes are `total` elements apart). +__global__ void triFoldScoreCoefficientsKernel(float const* __restrict__ qReal, float const* __restrict__ qImag, + float const* __restrict__ mlrCoef, float const* __restrict__ freqScaleSq, float const* __restrict__ meanCos, + float const* __restrict__ meanSin, float const* __restrict__ omega, float const* __restrict__ offsets, + int32_t const* __restrict__ roundStarts, float const* __restrict__ kvScales, float* __restrict__ cRe, + float* __restrict__ cIm, float* __restrict__ cMlr, int32_t numCalibratedLayers, int32_t numQueryHeads, + int32_t numFreqs, int32_t numOffsets, bool useMax, int64_t total) +{ + int64_t const idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) + { + return; + } + auto const f = static_cast(idx % numFreqs); + int64_t rest = idx / numFreqs; + auto const h = static_cast(rest % numQueryHeads); + rest /= numQueryHeads; + auto const l = static_cast(rest % numCalibratedLayers); + auto const req = static_cast(rest / numCalibratedLayers); + // Calibration tables are per (layer, head, freq); the request axis only + // enters through the phase terms below. + int64_t const cIdx = (static_cast(l) * numQueryHeads + h) * numFreqs + f; + float const qre = qReal[cIdx]; + float const qim = qImag[cIdx]; + float s = freqScaleSq[f]; + // Quantized-pool dequantization fold: K_real = scale_l * K_quant, so + // multiplying scale_l into ALL coefficient tables (c_mlr and every + // c_re/c_im plane below, since they all carry s) lets the score kernel + // read raw quantized elements. The |K| term relies on scale > 0 + // (validated host-side). Guarded (not "* 1.0f") so the float-pool path is + // instruction-identical to before this parameter existed. + if (kvScales != nullptr) + { + s *= kvScales[l]; + } + cMlr[idx] = mlrCoef[cIdx] * s; + if (!useMax) + { + float const mc = meanCos[static_cast(req) * numFreqs + f]; + float const ms = meanSin[static_cast(req) * numFreqs + f]; + cRe[idx] = s * (qre * mc - qim * ms); + cIm[idx] = s * (qim * mc + qre * ms); + } + else + { + float const om = omega[f]; + auto const rs = static_cast(roundStarts[req]); + for (int32_t o = 0; o < numOffsets; ++o) + { + float const phase = (rs + offsets[o]) * om; + float const cp = cosf(phase); + float const sp = sinf(phase); + cRe[o * total + idx] = s * (qre * cp - qim * sp); + cIm[o * total + idx] = s * (qim * cp + qre * sp); + } + } +} + +template +void launchVectorized(FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useMax, cudaStream_t stream) +{ + bool const staticChunks = params.numFreqs == 64; + int32_t const effectiveGroup = params.zIsQueryHead ? 1 : groupSize; +#define TRTLLM_TRI_SCORE_LAUNCH_GROUP(GROUP_V) \ + do \ + { \ + if (staticChunks) \ + { \ + if (useMax) \ + triScoreVectorizedKernel<<>>(params); \ + else \ + triScoreVectorizedKernel<<>>(params); \ + } \ + else \ + { \ + if (useMax) \ + triScoreVectorizedKernel<<>>(params); \ + else \ + triScoreVectorizedKernel<<>>(params); \ + } \ + } while (0) + switch (effectiveGroup) + { + case 1: TRTLLM_TRI_SCORE_LAUNCH_GROUP(1); break; + case 2: TRTLLM_TRI_SCORE_LAUNCH_GROUP(2); break; + case 4: TRTLLM_TRI_SCORE_LAUNCH_GROUP(4); break; + case 8: TRTLLM_TRI_SCORE_LAUNCH_GROUP(8); break; + default: + TLLM_CHECK_WITH_INFO(false, + "tri_attention_score: vectorized GQA group size must be 1/2/4/8 or use the per-query-head mapping (got " + "%d)", + effectiveGroup); + } +#undef TRTLLM_TRI_SCORE_LAUNCH_GROUP +} + +template +void launchScalar(FoldedScoreParams const& params, dim3 grid, bool useMax, cudaStream_t stream) +{ + if (useMax) + { + triScoreScalarKernel<<>>(params); + } + else + { + triScoreScalarKernel<<>>(params); + } +} + +} // namespace + +void foldScoreCoefficientsLaunch(float const* qReal, float const* qImag, float const* mlrCoef, float const* freqScaleSq, + float const* meanCos, float const* meanSin, float const* omega, float const* offsets, int32_t const* roundStarts, + float const* kvScales, float* cRe, float* cIm, float* cMlr, int32_t numRequests, int32_t numCalibratedLayers, + int32_t numQueryHeads, int32_t numFreqs, int32_t numOffsets, bool useMax, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(useMax ? (omega != nullptr && offsets != nullptr && roundStarts != nullptr) + : (meanCos != nullptr && meanSin != nullptr), + "tri_attention_score_fold: aggregation-path inputs are missing"); + int64_t const total = static_cast(numRequests) * numCalibratedLayers * numQueryHeads * numFreqs; + int32_t const threads = 256; + auto const blocks = static_cast((total + threads - 1) / threads); + triFoldScoreCoefficientsKernel<<>>(qReal, qImag, mlrCoef, freqScaleSq, meanCos, meanSin, + omega, offsets, roundStarts, kvScales, cRe, cIm, cMlr, numCalibratedLayers, numQueryHeads, numFreqs, numOffsets, + useMax, total); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, + int32_t numSegments, bool useVectorized, bool useMax, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(numSegments > 0 && numSegments <= 65535, + "tri_attention_score: request*layer segment count exceeds the CUDA grid limit"); + TLLM_CHECK_WITH_INFO(params.numOffsets >= 1 && params.numOffsets <= kMaxScoreOffsets, + "tri_attention_score: offset planes exceed the per-thread accumulator budget"); + TLLM_CHECK_WITH_INFO(!useVectorized || (params.numFreqs % 8 == 0 && params.strideDim == 1), + "tri_attention_score: vectorized path requires 8-frequency chunks with unit stride"); + // Tile count covers the decode span plus the worst-case page-alignment + // slack (tokenStart may sit up to tokensPerBlock - 1 tokens into a page). + auto const tiles = static_cast( + (params.outputWidth + params.tokensPerBlock - 1 + kScoreBlockThreads - 1) / kScoreBlockThreads); + uint32_t const headBlocks = params.zIsQueryHead && useVectorized ? params.numQueryHeads : params.numKvHeads; + dim3 const grid(tiles, static_cast(numSegments), headBlocks); + switch (poolType) + { + case PoolElementType::kBFloat16: + if (useVectorized) + { + launchVectorized<__nv_bfloat16>(params, groupSize, grid, useMax, stream); + } + else + { + launchScalar<__nv_bfloat16>(params, grid, useMax, stream); + } + break; + case PoolElementType::kHalf: + if (useVectorized) + { + launchVectorized(params, groupSize, grid, useMax, stream); + } + else + { + launchScalar(params, grid, useMax, stream); + } + break; + case PoolElementType::kFloat32: + // fp32 pools have 32-byte 8-frequency rows; the 16-byte chunk path + // does not apply, so they always take the strided scalar kernel. + TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: fp32 pools must use the scalar path"); + launchScalar(params, grid, useMax, stream); + break; + case PoolElementType::kFloat8E4M3: + // Quantized pools are functional-only: no vectorized instantiation + // exists for them by design (their dequant scale is folded into the + // coefficients, so only the scalar load path knows how to read them). + TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); + launchScalar<__nv_fp8_e4m3>(params, grid, useMax, stream); + break; + case PoolElementType::kInt8: + TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); + launchScalar(params, grid, useMax, stream); + break; + } + TLLM_CUDA_CHECK(cudaGetLastError()); +} + +} // namespace kernels::tri_attention_score + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h new file mode 100644 index 000000000000..226d54275e81 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h @@ -0,0 +1,144 @@ +/* + * 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. + */ +#pragma once + +#include +#include + +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::tri_attention_score +{ + +// The TriAttention trig score of one cached token t for query head h is +// score(t, h) = sum_f K_re(t, f) * c_re(h, f) +// + K_im(t, f) * c_im(h, f) +// + |K(t, f)| * c_mlr(h, f) +// where the c tables fold the per-round query calibration and phase terms so +// the per-token kernel touches no trigonometry. The fold runs once per +// eviction round; the score kernel then reads paged KV directly (one thread +// per token) and writes fp32 rows in the selector's layout. + +// Element type of the paged KV pools. The score kernel reads layers through +// raw per-layer base addresses (V2 exposes each layer as its own storage), so +// the caller passes the shared element type explicitly instead of a tensor. +// The quantized types (fp8/int8) are functional-only: they run the strided +// scalar kernel exclusively, and their per-layer dequantization scale is +// folded into the score coefficients (see foldScoreCoefficientsLaunch), so +// the score kernel reads raw quantized elements with zero hot-loop cost. +enum class PoolElementType : int32_t +{ + kBFloat16 = 0, + kHalf = 1, + kFloat32 = 2, + kFloat8E4M3 = 3, + kInt8 = 4, +}; + +// Upper bound on per-offset accumulators held by one score thread on the +// "max" aggregation path. Production configurations use 4 offsets; 8 leaves +// headroom without bloating the per-thread register budget. +inline constexpr int32_t kMaxScoreOffsets = 8; + +// Threads per score CTA; one thread scores one cached token. +inline constexpr int32_t kScoreBlockThreads = 128; + +// Fold the per-round score coefficients: +// c_re = fss * (q_re * cos - q_im * sin) +// c_im = fss * (q_im * cos + q_re * sin) +// c_mlr = mlr * fss +// Mean aggregation collapses all offsets into meanCos/meanSin beforehand and +// writes one plane (numOffsets == 1). Max aggregation cannot collapse (max +// does not commute through the frequency sum), so it writes one c_re/c_im +// plane per offset with cos/sin((round_start + offset) * omega); c_mlr is +// offset independent either way. Output layout per plane is +// [numRequests, numCalibratedLayers, numQueryHeads, numFreqs] fp32, indexed +// by ABSOLUTE layer id (matching the calibration tables). +// +// kvScales (nullable) carries one fp32 dequantization scale per ABSOLUTE +// layer id for quantized (fp8/int8) KV pools: K_real = scale_l * K_quant, so +// multiplying scale_l into all three coefficient tables (every per-offset +// plane included) lets the score kernel consume raw quantized elements. The +// |K| term relies on |scale * K_q| == scale * |K_q|, which only holds for +// scale > 0 — the host wrapper validates positivity before launch. +void foldScoreCoefficientsLaunch(float const* qReal, // [L_cal * HQ * F] + float const* qImag, // [L_cal * HQ * F] + float const* mlrCoef, // [L_cal * HQ * F] + float const* freqScaleSq, // [F] + float const* meanCos, // [numRequests * F] (mean path, else nullptr) + float const* meanSin, // [numRequests * F] (mean path, else nullptr) + float const* omega, // [F] (max path, else nullptr) + float const* offsets, // [numOffsets] (max path, else nullptr) + int32_t const* roundStarts, // [numRequests] (max path, else nullptr) + float const* kvScales, // [L_cal] per-layer dequant scale (quantized pools, else nullptr) + float* cRe, // [numOffsets, numRequests, L_cal, HQ, F] + float* cIm, // [numOffsets, numRequests, L_cal, HQ, F] + float* cMlr, // [numRequests, L_cal, HQ, F] + int32_t numRequests, int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, int32_t numOffsets, + bool useMax, cudaStream_t stream); + +// Everything one folded-score launch needs. One "segment" is one +// (request, scored layer) pair; segments are request-major so +// seg % numLayers == 0 identifies each request's first segment. +struct FoldedScoreParams +{ + int64_t const* layerBaseAddrs; // [num pools] absolute device addresses, ABSOLUTE layer id indexed + int32_t const* blockOffsets; // flattened native V2 page table + int64_t const* segPageOffsets; // [numSegments] offset of each segment's page row into blockOffsets + int32_t const* segRequestIds; // [numSegments] + int32_t const* segLayerIds; // [numSegments] ABSOLUTE layer ids + int32_t const* requestSeqLens; // [numRequests] + int32_t* validWidthOut; // [numRequests] side-store: seqLen - tokenStart, once per request + int32_t const* requestTokenStarts; // [numRequests] pinned prompt length = decode-region origin + float const* cRe; // fold output (see foldScoreCoefficientsLaunch) + float const* cIm; + float const* cMlr; + float* out; // [segment, numQueryHeads, outputWidth] fp32 decode-only scores + int32_t outputWidth; + int32_t numLayers; // scored layers per request (the segment period) + int32_t numRequests; + int32_t numCalibratedLayers; + int32_t numQueryHeads; + int32_t numKvHeads; + int32_t numFreqs; + int32_t tokensPerBlock; + int32_t kvFactor; // page-table entries encode role pages; entry / kvFactor = pool page + int32_t numOffsets; // effective c_re/c_im planes (1 on the mean path) + // Grid mapping for GQA group sizes without a dedicated template + // instantiation: grid.z indexes single query heads instead of KV heads + // (the KV plane is derived per block; K traffic is repeated per head). + bool zIsQueryHead; + // HND pool element strides (elements, not bytes) shared by all layers. + int64_t stridePage; + int64_t strideKvHead; + int64_t strideSlot; + int64_t strideDim; +}; + +// Launch the folded score over paged KV. useVectorized selects 16-byte +// 8-frequency chunk loads (requires numFreqs % 8 == 0, strideDim == 1, +// bf16/fp16 pools, and 16-byte aligned bases/strides — the caller audits +// alignment); otherwise a fully strided scalar path runs the same math. +// groupSize = numQueryHeads / numKvHeads must be 1, 2, 4, or 8 unless +// params.zIsQueryHead maps grid.z to single query heads. +void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, + int32_t numSegments, bool useVectorized, bool useMax, cudaStream_t stream); + +} // namespace kernels::tri_attention_score + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 08985d48660e..1784c94400a4 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -138,7 +138,8 @@ add_library( trtllmGenQKVProcessOp.cpp inplaceSliceCopyOp.cpp mhcOp.cpp - compressorOp.cpp) + compressorOp.cpp + triAttentionScoreOp.cpp) set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_libraries( th_common PRIVATE ${TORCH_LIBRARIES} th_utils ${Python3_LIBRARIES} diff --git a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp new file mode 100644 index 000000000000..e87c819f83d9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp @@ -0,0 +1,323 @@ +/* + * 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. + */ + +#include "tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +namespace tk = tensorrt_llm::kernels::tri_attention_score; + +// Two ops rather than one fused op, matching the file-level granularity of +// sibling kernel wrappers (one op per kernel launch): the coefficient fold +// runs once per eviction round into persistent buffers whose plane count +// depends on the aggregation mode, while the score launch consumes those +// buffers; callers time and re-plan them independently. + +namespace +{ + +void checkContiguousCudaFloat(torch::Tensor const& tensor, char const* name) +{ + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == at::kFloat, name, + " must be a contiguous fp32 CUDA tensor"); +} + +void checkContiguousCudaInt(torch::Tensor const& tensor, char const* name) +{ + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == at::kInt, name, + " must be a contiguous int32 CUDA tensor"); +} + +void checkContiguousCudaLong(torch::Tensor const& tensor, char const* name) +{ + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == at::kLong, name, + " must be a contiguous int64 CUDA tensor"); +} + +// Validate the optional per-layer dequantization scales for quantized +// (fp8/int8) KV pools and return the device pointer (nullptr when absent). +// Scales are indexed by ABSOLUTE layer id, matching layer_base_addrs and the +// calibration tables, so the extent must cover every calibrated layer. The +// positivity check runs HOST-side (one small sync per eviction round on a +// functional-only path): the |K| coefficient fold assumes +// |scale * K_q| == scale * |K_q|, which silently corrupts scores for +// scale <= 0, so a loud host error here is required rather than a +// device-side assert. +float const* checkKvScales( + std::optional const& kv_scales, int64_t num_calibrated_layers, char const* op_name) +{ + if (!kv_scales.has_value()) + { + return nullptr; + } + checkContiguousCudaFloat(*kv_scales, "kv_scales"); + TORCH_CHECK(kv_scales->numel() >= num_calibrated_layers, op_name, + ": kv_scales must carry one scale per calibrated layer (absolute layer id indexed)"); + TORCH_CHECK(kv_scales->min().item() > 0.0f, op_name, + ": kv_scales must be strictly positive (the |K| dequantization fold is only valid for positive scales)"); + return kv_scales->data_ptr(); +} + +// Fold the per-round TriAttention score coefficients into c_re/c_im/c_mlr +// (fp32, [num_offsets?, num_requests, num_calibrated_layers, heads, freqs]). +// The mean aggregation consumes offset-collapsed mean_cos/mean_sin and writes +// one plane; the max aggregation consumes omega/offsets/round_starts and +// writes one c_re/c_im plane per offset. kv_scales (quantized pools only) +// folds the per-layer dequantization scale into every coefficient table; the +// paired score op then reads raw quantized elements. This op cannot see the +// pool dtype, so presence-iff-quantized is enforced by the score op. +void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, torch::Tensor c_mlr, + torch::Tensor q_real, torch::Tensor q_imag, torch::Tensor mlr_coef, torch::Tensor freq_scale_sq, + std::optional mean_cos, std::optional mean_sin, std::optional omega, + std::optional offsets, std::optional round_starts, int64_t num_requests, + int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_freqs, int64_t num_offsets, bool use_max, + std::optional kv_scales) +{ + checkContiguousCudaFloat(c_re, "c_re"); + checkContiguousCudaFloat(c_im, "c_im"); + checkContiguousCudaFloat(c_mlr, "c_mlr"); + checkContiguousCudaFloat(q_real, "q_real"); + checkContiguousCudaFloat(q_imag, "q_imag"); + checkContiguousCudaFloat(mlr_coef, "mlr_coef"); + checkContiguousCudaFloat(freq_scale_sq, "freq_scale_sq"); + TORCH_CHECK(num_requests > 0 && num_calibrated_layers > 0 && num_query_heads > 0 && num_freqs > 0, + "tri_attention_fold_score_coefficients: fold extents must be positive"); + TORCH_CHECK(num_offsets >= 1 && num_offsets <= tk::kMaxScoreOffsets, + "tri_attention_fold_score_coefficients: num_offsets must be in [1, ", tk::kMaxScoreOffsets, "], got ", + num_offsets); + int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; + int64_t const planes = use_max ? num_offsets : 1; + TORCH_CHECK(c_re.numel() >= planes * total && c_im.numel() >= planes * total && c_mlr.numel() >= total, + "tri_attention_fold_score_coefficients: coefficient buffers are undersized"); + int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; + TORCH_CHECK(q_real.numel() >= calibration && q_imag.numel() >= calibration && mlr_coef.numel() >= calibration + && freq_scale_sq.numel() >= num_freqs, + "tri_attention_fold_score_coefficients: calibration tensors are undersized for the fold extent"); + + float const* meanCosPtr = nullptr; + float const* meanSinPtr = nullptr; + float const* omegaPtr = nullptr; + float const* offsetsPtr = nullptr; + int32_t const* roundStartsPtr = nullptr; + if (use_max) + { + TORCH_CHECK(omega.has_value() && offsets.has_value() && round_starts.has_value(), + "tri_attention_fold_score_coefficients: max aggregation requires omega, offsets, and round_starts"); + checkContiguousCudaFloat(*omega, "omega"); + checkContiguousCudaFloat(*offsets, "offsets"); + checkContiguousCudaInt(*round_starts, "round_starts"); + TORCH_CHECK( + omega->numel() >= num_freqs && offsets->numel() >= num_offsets && round_starts->numel() >= num_requests, + "tri_attention_fold_score_coefficients: max-path inputs are undersized for the folded request count"); + omegaPtr = omega->data_ptr(); + offsetsPtr = offsets->data_ptr(); + roundStartsPtr = round_starts->data_ptr(); + } + else + { + TORCH_CHECK(mean_cos.has_value() && mean_sin.has_value(), + "tri_attention_fold_score_coefficients: mean aggregation requires mean_cos and mean_sin"); + checkContiguousCudaFloat(*mean_cos, "mean_cos"); + checkContiguousCudaFloat(*mean_sin, "mean_sin"); + TORCH_CHECK(mean_cos->numel() >= num_requests * num_freqs && mean_sin->numel() >= num_requests * num_freqs, + "tri_attention_fold_score_coefficients: mean_cos/mean_sin are undersized (the fold iterates one row per " + "folded request)"); + meanCosPtr = mean_cos->data_ptr(); + meanSinPtr = mean_sin->data_ptr(); + } + float const* kvScalesPtr = checkKvScales(kv_scales, num_calibrated_layers, "tri_attention_fold_score_coefficients"); + + auto stream = at::cuda::getCurrentCUDAStream(); + tk::foldScoreCoefficientsLaunch(q_real.data_ptr(), q_imag.data_ptr(), mlr_coef.data_ptr(), + freq_scale_sq.data_ptr(), meanCosPtr, meanSinPtr, omegaPtr, offsetsPtr, roundStartsPtr, kvScalesPtr, + c_re.data_ptr(), c_im.data_ptr(), c_mlr.data_ptr(), static_cast(num_requests), + static_cast(num_calibrated_layers), static_cast(num_query_heads), + static_cast(num_freqs), static_cast(num_offsets), use_max, stream); +} + +// Score every cached decode token of every (request, layer) segment against +// the folded coefficient tables, writing fp32 [segment, head, token] rows and +// each request's decode width. pool_anchor is one of the scored layer pools: +// the kernel reads all layers through layer_base_addrs (V2 exposes each layer +// as its own storage), and the anchor only supplies their common element type +// and the device; its data is never read through this argument. +void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_base_addrs, torch::Tensor block_offsets, + torch::Tensor seg_page_offsets, torch::Tensor seg_request_ids, torch::Tensor seg_layer_ids, + torch::Tensor request_seq_lens, torch::Tensor valid_widths, torch::Tensor request_token_starts, torch::Tensor c_re, + torch::Tensor c_im, torch::Tensor c_mlr, torch::Tensor out, int64_t output_width, int64_t num_layers, + int64_t num_requests, int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_kv_heads, + int64_t num_freqs, int64_t tokens_per_block, int64_t kv_factor, int64_t num_offsets, int64_t stride_page, + int64_t stride_kv_head, int64_t stride_slot, int64_t stride_dim, int64_t num_segments, bool use_max, + bool use_vectorized, std::optional kv_scales) +{ + TORCH_CHECK(use_max || num_offsets == 1, + "tri_attention_paged_score: mean aggregation consumes exactly one folded coefficient plane"); + checkContiguousCudaLong(layer_base_addrs, "layer_base_addrs"); + checkContiguousCudaInt(block_offsets, "block_offsets"); + checkContiguousCudaLong(seg_page_offsets, "seg_page_offsets"); + checkContiguousCudaInt(seg_request_ids, "seg_request_ids"); + checkContiguousCudaInt(seg_layer_ids, "seg_layer_ids"); + checkContiguousCudaInt(request_seq_lens, "request_seq_lens"); + checkContiguousCudaInt(valid_widths, "valid_widths"); + checkContiguousCudaInt(request_token_starts, "request_token_starts"); + checkContiguousCudaFloat(c_re, "c_re"); + checkContiguousCudaFloat(c_im, "c_im"); + checkContiguousCudaFloat(c_mlr, "c_mlr"); + checkContiguousCudaFloat(out, "out"); + TORCH_CHECK(pool_anchor.is_cuda(), "tri_attention_paged_score: pool anchor must be a CUDA tensor"); + + TORCH_CHECK(num_segments > 0 && num_segments <= 65535, + "tri_attention_paged_score: request*layer segment count exceeds the CUDA grid limit"); + TORCH_CHECK(output_width > 0 && num_layers > 0 && num_requests > 0 && num_calibrated_layers > 0 + && tokens_per_block > 0 && kv_factor > 0 && num_freqs > 0, + "tri_attention_paged_score: geometry extents must be positive"); + TORCH_CHECK(num_kv_heads > 0 && num_query_heads % num_kv_heads == 0, + "tri_attention_paged_score: query heads must be divisible by KV heads"); + // Production configurations use 4 score offsets; 8 is the per-thread + // accumulator budget baked into the kernels. + TORCH_CHECK(num_offsets >= 1 && num_offsets <= tk::kMaxScoreOffsets, + "tri_attention_paged_score: num_offsets must be in [1, ", tk::kMaxScoreOffsets, "], got ", num_offsets); + TORCH_CHECK(seg_page_offsets.numel() >= num_segments && seg_request_ids.numel() >= num_segments + && seg_layer_ids.numel() >= num_segments, + "tri_attention_paged_score: segment metadata is undersized"); + TORCH_CHECK(request_seq_lens.numel() >= num_requests && valid_widths.numel() >= num_requests + && request_token_starts.numel() >= num_requests, + "tri_attention_paged_score: per-request metadata is undersized"); + int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; + TORCH_CHECK(c_re.numel() >= num_offsets * total && c_im.numel() >= num_offsets * total && c_mlr.numel() >= total, + "tri_attention_paged_score: folded coefficient buffers are undersized"); + TORCH_CHECK(out.numel() >= num_segments * num_query_heads * output_width, + "tri_attention_paged_score: score output buffer is undersized"); + + auto const dtype = pool_anchor.scalar_type(); + auto poolType = tk::PoolElementType::kBFloat16; + if (dtype == at::kBFloat16) + { + poolType = tk::PoolElementType::kBFloat16; + } + else if (dtype == at::kHalf) + { + poolType = tk::PoolElementType::kHalf; + } + else if (dtype == at::kFloat) + { + poolType = tk::PoolElementType::kFloat32; + } + else if (dtype == at::kFloat8_e4m3fn) + { + poolType = tk::PoolElementType::kFloat8E4M3; + } + else if (dtype == at::kChar) + { + poolType = tk::PoolElementType::kInt8; + } + else + { + TORCH_CHECK(false, "tri_attention_paged_score: unsupported KV pool dtype ", dtype, + " (supported: bf16, fp16, fp32, fp8_e4m3fn, int8)"); + } + // The score kernel never applies kv_scales itself (the fold op already + // multiplied them into the coefficient tables), but this op is the only + // one that sees the pool dtype, so it owns the presence contract: + // quantized elements without scales would be scored as raw integers, and + // scales alongside float pools would silently double-scale. + bool const quantizedPool = poolType == tk::PoolElementType::kFloat8E4M3 || poolType == tk::PoolElementType::kInt8; + TORCH_CHECK(!quantizedPool || kv_scales.has_value(), + "tri_attention_paged_score: quantized (fp8/int8) KV pools require per-layer kv_scales"); + TORCH_CHECK(quantizedPool || !kv_scales.has_value(), + "tri_attention_paged_score: kv_scales are only valid for quantized (fp8/int8) KV pools"); + checkKvScales(kv_scales, num_calibrated_layers, "tri_attention_paged_score"); + TORCH_CHECK(!use_vectorized || dtype == at::kBFloat16 || dtype == at::kHalf, + "tri_attention_paged_score: the vectorized path requires bf16 or fp16 pools"); + TORCH_CHECK(!use_vectorized || (num_freqs % 8 == 0 && stride_dim == 1), + "tri_attention_paged_score: the vectorized path requires num_freqs % 8 == 0 and a unit frequency stride"); + + auto const groupSize = static_cast(num_query_heads / num_kv_heads); + bool const vectorizedGroup = groupSize == 1 || groupSize == 2 || groupSize == 4 || groupSize == 8; + // Other GQA group sizes run the vectorized math one query head per grid.z + // block instead of one KV head (a runtime mapping, no extra template). + bool const zIsQueryHead = use_vectorized && !vectorizedGroup; + TORCH_CHECK((zIsQueryHead ? num_query_heads : num_kv_heads) <= 65535, + "tri_attention_paged_score: head count exceeds the CUDA grid limit"); + + tk::FoldedScoreParams params; + params.layerBaseAddrs = layer_base_addrs.data_ptr(); + params.blockOffsets = block_offsets.data_ptr(); + params.segPageOffsets = seg_page_offsets.data_ptr(); + params.segRequestIds = seg_request_ids.data_ptr(); + params.segLayerIds = seg_layer_ids.data_ptr(); + params.requestSeqLens = request_seq_lens.data_ptr(); + params.validWidthOut = valid_widths.data_ptr(); + params.requestTokenStarts = request_token_starts.data_ptr(); + params.cRe = c_re.data_ptr(); + params.cIm = c_im.data_ptr(); + params.cMlr = c_mlr.data_ptr(); + params.out = out.data_ptr(); + params.outputWidth = static_cast(output_width); + params.numLayers = static_cast(num_layers); + params.numRequests = static_cast(num_requests); + params.numCalibratedLayers = static_cast(num_calibrated_layers); + params.numQueryHeads = static_cast(num_query_heads); + params.numKvHeads = static_cast(num_kv_heads); + params.numFreqs = static_cast(num_freqs); + params.tokensPerBlock = static_cast(tokens_per_block); + params.kvFactor = static_cast(kv_factor); + params.numOffsets = static_cast(num_offsets); + params.zIsQueryHead = zIsQueryHead; + params.stridePage = stride_page; + params.strideKvHead = stride_kv_head; + params.strideSlot = stride_slot; + params.strideDim = stride_dim; + + auto stream = at::cuda::getCurrentCUDAStream(); + tk::foldedScoreLaunch( + params, poolType, groupSize, static_cast(num_segments), use_vectorized, use_max, stream); +} + +} // anonymous namespace + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "tri_attention_fold_score_coefficients(" + "Tensor(a!) c_re, Tensor(b!) c_im, Tensor(c!) c_mlr, " + "Tensor q_real, Tensor q_imag, Tensor mlr_coef, Tensor freq_scale_sq, " + "Tensor? mean_cos, Tensor? mean_sin, " + "Tensor? omega, Tensor? offsets, Tensor? round_starts, " + "int num_requests, int num_calibrated_layers, " + "int num_query_heads, int num_freqs, " + "int num_offsets, bool use_max, Tensor? kv_scales=None) -> ()"); + + m.def( + "tri_attention_paged_score(" + "Tensor pool_anchor, Tensor layer_base_addrs, Tensor block_offsets, " + "Tensor seg_page_offsets, Tensor seg_request_ids, Tensor seg_layer_ids, " + "Tensor request_seq_lens, Tensor(a!) valid_widths, Tensor request_token_starts, " + "Tensor c_re, Tensor c_im, Tensor c_mlr, Tensor(b!) out, " + "int output_width, int num_layers, int num_requests, int num_calibrated_layers, " + "int num_query_heads, int num_kv_heads, int num_freqs, int tokens_per_block, " + "int kv_factor, int num_offsets, int stride_page, int stride_kv_head, " + "int stride_slot, int stride_dim, int num_segments, bool use_max, bool use_vectorized, " + "Tensor? kv_scales=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("tri_attention_fold_score_coefficients", &triAttentionFoldScoreCoefficientsOp); + m.impl("tri_attention_paged_score", &triAttentionPagedScoreOp); +} diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index f2df658e7352..d9654deb5adf 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1168,7 +1168,7 @@ def _ensure_calibrated(self) -> None: self._F = int(self.calibration["E_q"].shape[2]) # Squared per-frequency RoPE scaling factor (required calibration key). self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) - # Pre-split query stats + MLR coefficient for the Triton score kernel so + # Pre-split query stats + MLR coefficient for the score kernel so # it doesn't recompute (E_q_norm - |E_q|) per call. Shapes [L, H, F]. _Eq = self.calibration["E_q"] self._triattn_q_real = _Eq.real.to(torch.float32).contiguous() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 4b88170e7a38..4a23c1f17432 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -4,8 +4,11 @@ The production path uses one fixed-shape trig-score launch across all dense layers, CuTE-DSL TopK selection, and grouped C++ compaction. This module owns -the score kernel and its persistent launcher; selection and compaction live in -their respective runtime modules. +the score launcher and its persistent metadata; scoring itself runs through +the compiled ``trtllm`` CUDA ops (coefficient fold + folded paged score) for +every supported geometry. The original Triton score kernel has been deleted; +the unit tests validate the CUDA ops against an independent PyTorch oracle. +Selection and compaction live in their respective runtime modules. House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). @@ -98,231 +101,118 @@ def prepare_mean_phase( ) -@triton.jit -def _tri_score_perhead_kernel( - pool_anchor_ptr, # typed pool pointer; used ONLY to infer the element type - # for the int->pointer cast below (its data is never read through it). - layer_base_addrs, # [num_layers] int64: ABSOLUTE device address of each - # scored layer's HND base. Layers do NOT need to share one storage: - # each segment casts its own layer's address back to a typed pointer, so - # all address arithmetic stays inside that layer's own allocation. - block_offsets_ptr, # Native V2 [pool, request, K/V, block] int32 offsets. - seg_page_off, # [nseg] int64: offset of this segment's page table into - # block_offsets_ptr. - # per-SEGMENT metadata (seg = req_slot*L_scored + layer_slot), idx by pid(0): - seg_req_id, # [nseg] int32: request slot (round_start / mean phase lookup) - seg_layer_id, # [nseg] int32: ABSOLUTE layer id (indexes layer_base_addrs + calib) - req_seq_len, # [num_requests] int32 - req_valid_width_out, # [num_requests] int32: decode-only length for selection - req_round_start, # [num_requests] int32 logical token position - req_token_start, # [num_requests] int32: pinned prompt length; scoring - # starts at this decode-region origin, so prompt lengths may differ - # across the cohort. - # per-LAYER calibration, [L,H,F] flattened layer-major: - q_real_ptr, # [L*H*F] fp32 - q_imag_ptr, # [L*H*F] fp32 - mlr_coef_ptr, # [L*H*F] fp32 - # per-REQUEST offset-collapsed phase ('mean' path), [num_requests,F] flattened: - mean_cos_ptr, # [num_requests*F] fp32 - mean_sin_ptr, # [num_requests*F] fp32 - # shared freq vectors: - freq_scale_sq_ptr, # [F] fp32 - omega_ptr, # [F] fp32 ('max' path only) - offsets_ptr, # [O] fp32 ('max' path only) - out_ptr, # [request, layer, query_head, decode_token] fp32 - output_width, - # scalars uniform across the batch: - num_layers, - num_q_heads, - num_kv_heads, - num_freqs, # F = head_dim // 2 - head_dim, - tokens_per_block, - kv_factor, - num_offsets, # O ('max' path only) - # per-layer HND element strides (uniform across scored layers): - s_page, - s_kv_head, - s_slot, - s_dim, - USE_MAX: tl.constexpr, - T_BLOCK: tl.constexpr, - F_BLOCK: tl.constexpr, - F_EXACT: tl.constexpr, # F_BLOCK == num_freqs: every f-lane is real, so the - # f-axis masks/wheres below are no-ops and are compiled out (fewer - # predicates and selects on the hot path; bit-identical results). -): - # Token tiles ride the fastest grid axis: adjacent programs then walk - # consecutive K pages of one (request, layer, head) and reuse its - # calibration/phase rows in L2 (~2% faster than segment-major order). - seg = tl.program_id(1) - t_blk = tl.program_id(0) - # KV heads are grid-parallel (axis 2): iterations of the former kv_head - # loop shared NO data (each KV head reads its own K and writes its own - # output rows), so hoisting it onto the grid multiplies parallelism with - # zero extra HBM traffic. The q-in-group loop below stays inside the - # program because it REUSES this head's K from registers (GQA dedup). - kv_head = tl.program_id(2) - - req_id = tl.load(seg_req_id + seg) - seq_len = tl.load(req_seq_len + req_id) - token_start = tl.load(req_token_start + req_id) - if (seg % num_layers == 0) & (t_blk == 0) & (kv_head == 0): - tl.store(req_valid_width_out + req_id, seq_len - token_start) - # Derive the ragged launch bound in the score program instead of staging - # one replicated length and block count for every request/layer segment. - n_tblk = (seq_len - token_start + T_BLOCK - 1) // T_BLOCK - if t_blk >= n_tblk: - return - - layer_id = tl.load(seg_layer_id + seg) - rstart = tl.load(req_round_start + req_id) - # This segment's layer base: an absolute address cast back to a pool-typed - # pointer. TRT-LLM V2 exposes every layer as its own TensorWrapper storage, - # so "element offset relative to one shared storage" does not exist; the - # per-layer absolute address is the same device-pointer-array pattern the - # C++ backends use (KVBlockArray / grouped-GEMM pointer arrays). - layer_ptr = tl.load(layer_base_addrs + layer_id).to( - tl.pointer_type(pool_anchor_ptr.dtype.element_ty) - ) - page_off = tl.load(seg_page_off + seg) - - f = tl.arange(0, F_BLOCK) - f_mask = f < num_freqs - f64 = f.to(tl.int64) - - # ---- token tile of THIS segment ---- - t = t_blk * T_BLOCK + tl.arange(0, T_BLOCK) - absolute_t = t + token_start - t_mask = absolute_t < seq_len - blk_in_seq = absolute_t // tokens_per_block - slot = (absolute_t % tokens_per_block).to(tl.int64) - # The native attention page-table copy encodes K offsets in units of the - # underlying K/V role pages. Convert that value to the HND pool page inline - # instead of materializing a second page table before scoring. - encoded_page = tl.load( - block_offsets_ptr + page_off + blk_in_seq, - mask=t_mask, - other=0, - ) - phys_page = (encoded_page // kv_factor).to(tl.int64) - - # element offset into THIS layer's pool for (page, KEY=0, *, slot). - # KEY half is kv_factor index 0 -> its stride term is 0 (matches reference). - tok_base = phys_page * s_page + slot * s_slot # [T_BLOCK] int64 - - # per-request 'mean'-path phase + shared freq scale. - if F_EXACT: - mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f) - msin = tl.load(mean_sin_ptr + req_id * num_freqs + f) - fss = tl.load(freq_scale_sq_ptr + f) - else: - mcos = tl.load(mean_cos_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) - msin = tl.load(mean_sin_ptr + req_id * num_freqs + f, mask=f_mask, other=0.0) - fss = tl.load(freq_scale_sq_ptr + f, mask=f_mask, other=0.0) - - # ---- PER-HEAD (position + mlr), GQA-deduped, NO head reduction ---- - # This program scores ONE KV head's token tile for the group_size q-heads - # that share it. K (and |K|) is loaded ONCE and reused across the group; - # h = kv_head*group_size + qg keeps query-head order 0..num_q_heads-1, so - # every head's math is bit-for-bit identical to the looped variant. - group_size = num_q_heads // num_kv_heads - if F_EXACT: - load_mask = tl.broadcast_to(t_mask[:, None], (T_BLOCK, F_BLOCK)) - else: - load_mask = t_mask[:, None] & f_mask[None, :] - off_re = f64[None, :] * s_dim - off_im = (num_freqs + f64[None, :]) * s_dim - - base = tok_base + kv_head.to(tl.int64) * s_kv_head # [T_BLOCK] - # paged K loaded ONCE for this KV head (shared by group_size q-heads). - k_re = tl.load(layer_ptr + base[:, None] + off_re, mask=load_mask, other=0.0).to(tl.float32) - k_im = tl.load(layer_ptr + base[:, None] + off_im, mask=load_mask, other=0.0).to(tl.float32) - kmag = tl.sqrt(k_re * k_re + k_im * k_im) # once per KV head - - qg = 0 - while qg < group_size: - h = kv_head * group_size + qg - calib_off = (layer_id.to(tl.int64) * num_q_heads + h) * num_freqs - if F_EXACT: - qre = tl.load(q_real_ptr + calib_off + f) - qim = tl.load(q_imag_ptr + calib_off + f) - mlrc = tl.load(mlr_coef_ptr + calib_off + f) - else: - qre = tl.load(q_real_ptr + calib_off + f, mask=f_mask, other=0.0) - qim = tl.load(q_imag_ptr + calib_off + f, mask=f_mask, other=0.0) - mlrc = tl.load(mlr_coef_ptr + calib_off + f, mask=f_mask, other=0.0) - - # complex product Q . conj(K) -- the trig importance score. - prod_real = qre[None, :] * k_re + qim[None, :] * k_im - prod_imag = qim[None, :] * k_re - qre[None, :] * k_im - - if USE_MAX: - # max over O offsets does NOT commute through the freq-sum; - # explicit O loop reducing max over the per-offset F-sum. - score = tl.full((T_BLOCK,), -float("inf"), tl.float32) - o = 0 - while o < num_offsets: - off = tl.load(offsets_ptr + o) - if F_EXACT: - om = tl.load(omega_ptr + f) - else: - om = tl.load(omega_ptr + f, mask=f_mask, other=0.0) - phase = (rstart + off) * om - cphase = tl.cos(phase) - sphase = tl.sin(phase) - per_f = fss[None, :] * (prod_real * cphase[None, :] - prod_imag * sphase[None, :]) - if F_EXACT: - offset_score = tl.sum(per_f, axis=1) - else: - offset_score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) - score = tl.maximum(score, offset_score) - o += 1 - else: - # 'mean': offset loop collapsed into mean_cos/mean_sin. - per_f = fss[None, :] * (prod_real * mcos[None, :] - prod_imag * msin[None, :]) - if F_EXACT: - score = tl.sum(per_f, axis=1) - else: - score = tl.sum(tl.where(f_mask[None, :], per_f, 0.0), axis=1) - - # position-INDEPENDENT MLR term (reuses the per-KV-head |K|). - mlr_f = kmag * mlrc[None, :] * fss[None, :] - if F_EXACT: - mlr = tl.sum(mlr_f, axis=1) - else: - mlr = tl.sum(tl.where(f_mask[None, :], mlr_f, 0.0), axis=1) - - # Segments are request-major then layer-major. Write the decode-only - # score directly in the selector's [request, layer, head, token] layout. - out_offset = (seg.to(tl.int64) * num_q_heads + h) * output_width + t - tl.store(out_ptr + out_offset, score + mlr, mask=t_mask) - qg += 1 - - def _launch_tri_score_perhead( - grid: tuple, - pointer_args: tuple, - geometry_args: tuple, + group: "_FixedScoreGroup", + request_count: int, + num_segments: int, + valid_seq_lens: torch.Tensor, + valid_widths: torch.Tensor, + round_starts_device: torch.Tensor, + token_starts_device: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, *, score_aggregation: str, - token_block: int, - num_freqs: int, ) -> None: - """Launch the shared score ABI for eager and fixed metadata owners.""" + """Fold the per-round coefficients, then score paged KV via the C++ ops. + + The compiled ``trtllm`` score ops are THE implementation for every + geometry this launcher accepts; unsupported inputs fail loudly inside the + ops (TORCH_CHECK) instead of routing to another kernel. The unit tests + validate them against an independent PyTorch oracle. + """ if score_aggregation not in ("mean", "max"): raise ValueError(f"unsupported score aggregation: {score_aggregation}") - f_block = triton.next_power_of_2(num_freqs) - _tri_score_perhead_kernel[grid]( - *pointer_args, - *geometry_args, - USE_MAX=(score_aggregation == "max"), - T_BLOCK=token_block, - F_BLOCK=f_block, - F_EXACT=(f_block == num_freqs), - # 4 warps measured fastest across token blocks; the default was never - # tuned for this kernel (sibling kernels pin their warp counts too). - num_warps=4, + if not ( + hasattr(torch.ops.trtllm, "tri_attention_fold_score_coefficients") + and hasattr(torch.ops.trtllm, "tri_attention_paged_score") + ): + raise RuntimeError( + "this TensorRT-LLM build is missing the TriAttention score ops; rebuild the C++ " + "th_common extension (there is deliberately no Triton fallback: a loud failure " + "here beats silently scoring through a slower path)" + ) + use_max = score_aggregation == "max" + ( + num_q_heads, + num_kv_heads, + num_freqs, + tokens_per_block, + kv_factor, + num_offsets, + s_page, + s_kv_head, + s_slot, + s_dim, + ) = group.geometry_args + # Mean aggregation collapses all offsets into mean_cos/mean_sin (one + # coefficient plane); max keeps one c_re/c_im plane per offset because + # max does not commute through the frequency sum. + offset_planes = num_offsets if use_max else 1 + c_re, c_im, c_mlr = group._fold_coefficient_buffers(offset_planes) + q_real, q_imag, mlr_coef = group.pointer_middle + freq_scale_sq, omega, offsets = group.pointer_tail + torch.ops.trtllm.tri_attention_fold_score_coefficients( + c_re, + c_im, + c_mlr, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + None if use_max else mean_cos.view(-1), + None if use_max else mean_sin.view(-1), + omega if use_max else None, + offsets if use_max else None, + round_starts_device if use_max else None, + request_count, + group._num_calibrated_layers, + num_q_heads, + num_freqs, + offset_planes, + use_max, + # Per-layer dequant scales (quantized pools only, else None): the fold + # multiplies them into the coefficient tables so the score op below + # reads raw quantized elements at zero hot-loop cost. + group._kv_scales, + ) + pool_anchor, layer_base_addrs, block_offsets, seg_page_off, seg_req, seg_layer = ( + group.pointer_prefix + ) + torch.ops.trtllm.tri_attention_paged_score( + pool_anchor, + layer_base_addrs, + block_offsets, + seg_page_off, + seg_req, + seg_layer, + valid_seq_lens, + valid_widths, + token_starts_device, + c_re, + c_im, + c_mlr, + group.output, + group.output_width, + group.num_layers, + request_count, + group._num_calibrated_layers, + num_q_heads, + num_kv_heads, + num_freqs, + tokens_per_block, + kv_factor, + offset_planes, + s_page, + s_kv_head, + s_slot, + s_dim, + num_segments, + use_max, + group._use_vectorized, + # Validation-only here (presence must match the pool dtype; the fold + # op above already consumed the values). + group._kv_scales, ) @@ -333,6 +223,13 @@ class _FixedScoreGroup: living in DISTINCT storages with DISTINCT block tables. ``block_offsets`` uses the native TRT-LLM attention layout and ``page_table_slots`` maps each scored layer to its V2 pool slot. + + LIFETIME CONTRACT: the group captures the scored layer pools as raw device + addresses (``layer_base_addrs``) and keeps a reference only to the anchor + pool (the score op's dtype witness). The caller owns ``layer_pools`` and must + keep every scored pool alive for as long as it launches through this group + (in production the V2 KV-cache manager does); a dropped pool leaves its + address dangling and scores read allocator-recycled memory. """ def __init__( @@ -352,6 +249,7 @@ def __init__( omega: torch.Tensor, offsets: torch.Tensor, output_width: int | None = None, + kv_scales: torch.Tensor | None = None, ) -> None: if not layer_indices or min(max_requests, page_count, seq_len) <= 0: raise ValueError("fixed score group requires non-empty positive geometry") @@ -380,14 +278,12 @@ def __init__( _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape if num_q_heads % num_kv_heads: raise ValueError("query heads must be divisible by KV heads") - self.num_kv_heads = int(num_kv_heads) self.num_freqs = head_dim // 2 strides = tuple(int(value) for value in p0.stride()) self.geometry_args = ( num_q_heads, num_kv_heads, self.num_freqs, - head_dim, tokens_per_block, kv_factor, int(offsets.numel()), @@ -399,6 +295,7 @@ def __init__( # Per-layer ABSOLUTE base addresses. Layers may live in distinct # storages (V2 TensorWrapper-per-layer); only geometry must be uniform. element_size = p0.element_size() + bases_16b_aligned = True layer_base_addrs = torch.zeros(len(layer_pools), dtype=torch.int64, device=device) for layer in layer_indices: pool = layer_pools[layer] @@ -411,9 +308,51 @@ def __init__( address = int(pool.data_ptr()) if address % element_size: raise ValueError("fixed score layer base is not element-aligned") + bases_16b_aligned &= address % 16 == 0 layer_base_addrs[layer] = address - # The anchor pool is passed as a typed kernel argument ONLY so the - # kernel can recover the element type for the int->pointer cast. + # The score op runs 16-byte 8-frequency K loads when the fixed layout + # guarantees aligned rows, and its strided scalar path otherwise. + # Audited ONCE here: bases and strides never change for this group. + strides_16b_aligned = all( + (element_size * stride) % 16 == 0 for stride in (strides[0], strides[2], strides[3]) + ) + self._use_vectorized = ( + p0.dtype in (torch.bfloat16, torch.float16) + and self.num_freqs % 8 == 0 + and strides[4] == 1 + and bases_16b_aligned + and strides_16b_aligned + ) + # Quantized (fp8/int8) pools are FUNCTIONAL-ONLY and scalar-path-only + # (the dtype gate above already excludes them from the vectorized + # path). Their per-layer dequantization scale is folded into the + # score coefficients at launch time, so it must be present up front; + # conversely, scales alongside a float pool would double-scale the + # coefficients, so that pairing is rejected just as loudly. + quantized_pool = p0.dtype in (torch.float8_e4m3fn, torch.int8) + if quantized_pool and kv_scales is None: + raise ValueError("quantized (fp8/int8) KV pools require per-layer kv_scales") + if not quantized_pool and kv_scales is not None: + raise ValueError("kv_scales are only valid for quantized (fp8/int8) KV pools") + self._kv_scales = ( + None + if kv_scales is None + else kv_scales.to(device=device, dtype=torch.float32).contiguous().view(-1) + ) + # Calibration tables span every model layer; segments index them by + # ABSOLUTE layer id, so the fold covers the full calibrated extent. + self._num_calibrated_layers = q_real_LHF.numel() // (int(num_q_heads) * self.num_freqs) + # Segment layer ids index the fold tables ON DEVICE where they cannot + # be range-checked; validate the extent once here, loudly. + if min(layer_indices) < 0 or max(layer_indices) >= self._num_calibrated_layers: + raise ValueError("scored layer index exceeds the calibrated layer extent") + # Folded per-round coefficient tables, allocated on first launch and + # keyed by plane count so switching aggregation (mean: one plane; + # max: one plane per offset) re-shapes without churn. + self._fold_buffers: dict = {} + # The anchor pool is passed to the CUDA score op ONLY as its dtype + # witness: the op recovers the pool element type from it and never + # reads data through it. seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( self.num_layers ) @@ -440,11 +379,6 @@ def __init__( ) slot_idx = slots_t.repeat(max_requests) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - # 16-token tiles keep the per-program fp32 working set small enough to - # avoid the register-pressure occupancy collapse measured at 64-token - # tiles (255 regs/thread + spills -> 63 regs, ~1.5x faster end to end). - self.token_block = 16 - self.max_ntblk = (self.output_width + self.token_block - 1) // self.token_block self.output = torch.empty( max_requests, self.num_layers, @@ -468,6 +402,31 @@ def __init__( ) self.pointer_tail = (freq_scale_sq, omega, offsets) + def _fold_coefficient_buffers( + self, offset_planes: int + ) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]": + """Return (c_re, c_im, c_mlr) fold tables for one plane count. + + Sized on ``max_requests`` so any launch's active ``request_count`` + fits without reallocation (each launch folds only its active rows); + c_mlr is offset independent so it never grows planes. + """ + buffers = self._fold_buffers.get(offset_planes) + if buffers is None: + elements = ( + self.max_requests + * self._num_calibrated_layers + * int(self.geometry_args[0]) + * self.num_freqs + ) + device = self.output.device + c_re = torch.empty(offset_planes * elements, dtype=torch.float32, device=device) + c_im = torch.empty_like(c_re) + c_mlr = torch.empty(elements, dtype=torch.float32, device=device) + buffers = (c_re, c_im, c_mlr) + self._fold_buffers[offset_planes] = buffers + return buffers + def launch( self, request_count: int, @@ -496,23 +455,16 @@ def launch( raise ValueError("request*layer segment count exceeds the CUDA grid limit") output = self.output[:request_count] _launch_tri_score_perhead( - (self.max_ntblk, num_segments, self.num_kv_heads), - ( - *self.pointer_prefix, - valid_seq_lens, - valid_widths, - round_starts_device, - token_starts_device, - *self.pointer_middle, - mean_cos.view(-1), - mean_sin.view(-1), - *self.pointer_tail, - output, - ), - (self.output_width, self.num_layers, *self.geometry_args), + self, + request_count, + num_segments, + valid_seq_lens, + valid_widths, + round_starts_device, + token_starts_device, + mean_cos, + mean_sin, score_aggregation=score_aggregation, - token_block=self.token_block, - num_freqs=self.num_freqs, ) return output diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py new file mode 100644 index 000000000000..11d9361b0b15 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py @@ -0,0 +1,690 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CUDA score ops vs an independent PyTorch oracle. + +`_launch_tri_score_perhead` calls the compiled `trtllm` fold + paged-score +ops unconditionally (no Triton fallback; the original Triton score kernel has +been deleted). These tests score the same paged pools through a pure-PyTorch +oracle and compare the two implementations across the geometry matrix the +launcher must cover: both CUDA load paths (vectorized 8-frequency chunks and +strided scalar), both aggregations, every supported pool dtype (bf16/fp16/ +fp32, plus functional-only fp8_e4m3fn/int8 with per-layer dequantization +scales), and GQA group sizes with and without a dedicated template +instantiation. +""" + +import pytest +import torch +from conftest import encode_block_offsets as _encode_block_offsets + +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, +) + + +def _require_score_ops() -> None: + """The compiled score ops are a hard prerequisite for these tests.""" + assert hasattr(torch.ops.trtllm, "tri_attention_fold_score_coefficients"), ( + "TriAttention fold op is not loaded" + ) + assert hasattr(torch.ops.trtllm, "tri_attention_paged_score"), ( + "TriAttention paged score op is not loaded" + ) + + +# Copied verbatim from test_triattention_pipeline.py (_torch_tri_score_oracle) +# so this module stays runnable standalone: an independent pure-PyTorch +# implementation of the paged TriAttention score, covering both aggregations +# (mean and max), GQA head mapping via head // group_size, and the +# position-independent MLR term. +def _torch_tri_score_oracle( + layer_pools, + page_ids, + seq_lens, + round_starts, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + layer_indices, + aggregation, +): + """Independent Torch implementation of the paged TriAttention score.""" + scores = [] + num_q_heads = int(q_real.shape[1]) + for request, seq_len in enumerate(seq_lens): + phase = (round_starts[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 + if aggregation == "mean": + position = ( + freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) + ).sum(dim=-1) + else: + position = ( + ( + freq_scale_sq[None, None, :] + * ( + product_real[None] * torch.cos(phase)[:, None, :] + - product_imag[None] * torch.sin(phase)[:, None, :] + ) + ) + .sum(dim=-1) + .max(dim=0) + .values + ) + 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 _oracle_reference( + group: _FixedScoreGroup, + pools: list, + oracle_inputs: dict, + request_count: int, + seq_lens: list, + round_starts: torch.Tensor, + prompt_len: int, + aggregation: str, + sentinel: float, +) -> torch.Tensor: + """Sentinel-filled oracle scores in the ops' [request, layer, head, token] layout. + + The oracle scores every cached token in [0, seq_len); the score ops write + only the decode region [prompt_len, seq_len) at column origin 0. Slice the + prompt columns off each oracle row and leave every column the ops must not + touch at the sentinel, so the comparison covers the write MASK as well as + the values. + """ + num_layers = group.num_layers + oracle = _torch_tri_score_oracle( + pools, + oracle_inputs["page_ids"][:request_count], + seq_lens, + round_starts[:request_count].tolist(), + 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)), + aggregation, + ) + reference = torch.full_like(group.output, sentinel) + for request in range(request_count): + width = seq_lens[request] - prompt_len + for layer in range(num_layers): + reference[request, layer, :, :width] = oracle[request * num_layers + layer][ + :, prompt_len: + ] + return reference + + +def _build_case( + *, + request_count: int, + max_requests: int, + num_layers: int, + page_count: int, + tokens_per_block: int, + head_dim: int, + num_q_heads: int, + num_kv_heads: int, + dtype: torch.dtype, + offsets: list, + prompt_len: int, + seed: int, + pools: list | None = None, + kv_scales: torch.Tensor | None = None, +): + device = torch.device("cuda", torch.cuda.current_device()) + torch.manual_seed(seed) + num_freqs = head_dim // 2 + # Callers may inject prebuilt pools: the quantized tests build a quantized + # pool (plus its dequantized fp32 twin for the oracle reference leg) from + # ONE set of randoms and pass the quantized pool in here. + if pools is None: + pools = [ + torch.randn( + max_requests * page_count, + 2, + num_kv_heads, + tokens_per_block, + head_dim, + device=device, + ).to(dtype) + for _ in range(num_layers) + ] + page_ids = torch.randperm(max_requests * page_count).view(max_requests, page_count).to(device) + q_real = torch.randn(num_layers, num_q_heads, num_freqs, device=device) + q_imag = torch.randn(num_layers, num_q_heads, num_freqs, device=device) + mlr_coef = 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 + group = _FixedScoreGroup( + pools, + list(range(num_layers)), + max_requests, + page_count, + capacity, + num_q_heads, + _encode_block_offsets(page_ids), + [0] * num_layers, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets_t, + output_width=capacity - prompt_len, + kv_scales=kv_scales, + ) + # LIFETIME: the group records only raw device ADDRESSES of the scored + # layer pools (layer_base_addrs) and references just the anchor pool + # (pools[0], the kernel dtype witness). Production pools are owned by the + # KV-cache manager, so the group deliberately does not hold them; the + # test must keep the whole pool list alive itself. Dropping it here frees + # every non-anchor layer pool, whose blocks the caching allocator then + # recycles for the fold/output/reference tensors allocated later in the + # test — a use-after-free that reads back sentinel/coefficient bytes as + # K data (observed as layer>=1 inf/NaN score garbage). + group.test_pools_keepalive = pools + round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() + token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) + seq_lens = [capacity - ((request * 3) % 5) for request in range(request_count)] + valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + phase = (round_starts.float()[:, None, None] + offsets_t[None, :, None]) * omega[None, None, :] + mean_cos = torch.cos(phase).mean(dim=1) + mean_sin = torch.sin(phase).mean(dim=1) + # Everything the PyTorch oracle needs to rebuild the reference leg + # independently (it recomputes its own mean phases from these). + 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 ( + group, + round_starts, + token_starts, + valid_seq_lens, + seq_lens, + mean_cos, + mean_sin, + oracle_inputs, + ) + + +# --------------------------------------------------------------------------- +# Quantized (fp8_e4m3fn / int8) KV pools — FUNCTIONAL kernel-level coverage +# ONLY. End-to-end quantized-KV eviction is NOT validated here or anywhere +# else yet: nothing in the pipeline produces quantized pools, so these tests +# exercise the scalar score path + coefficient scale fold in isolation. The +# reference leg dequantizes the SAME quantized elements to fp32 pools and +# runs the PyTorch oracle on them, so quantization error cancels and only the +# scale-fold/loading math is under test. +# --------------------------------------------------------------------------- + +_QUANTIZED_GEOMETRY = dict( + request_count=3, + max_requests=4, + num_layers=2, + page_count=4, + tokens_per_block=4, + head_dim=8, + num_q_heads=4, + num_kv_heads=2, + prompt_len=5, + offsets=[1.0, 2.0, 4.0], +) + + +def _quantize_pools(raw_pools: list, dtype: torch.dtype): + """Per-layer amax quantization of fp32 pools. + + Returns (quantized pools, dequantized fp32 twin pools, per-layer scales). + The twin holds values dequantized FROM the quantized elements — NOT the + pre-quantization randoms — so quantization error is present in both legs + identically and the CUDA-vs-reference tolerance can stay tight. + """ + quant_max = 448.0 if dtype == torch.float8_e4m3fn else 127.0 + quantized, dequantized, scales = [], [], [] + for raw in raw_pools: + scale = raw.abs().amax().float() / quant_max + if dtype == torch.int8: + quant = torch.round(raw / scale).clamp(-127, 127).to(torch.int8) + else: + quant = (raw / scale).to(torch.float8_e4m3fn) + quantized.append(quant) + dequantized.append(quant.to(torch.float32) * scale) + scales.append(scale) + return quantized, dequantized, torch.stack(scales) + + +def _build_quantized_raw_pools(seed: int, device: torch.device) -> list: + """The shared fp32 randoms both quantized-test pools derive from.""" + torch.manual_seed(seed) + g = _QUANTIZED_GEOMETRY + return [ + torch.randn( + g["max_requests"] * g["page_count"], + 2, + g["num_kv_heads"], + g["tokens_per_block"], + g["head_dim"], + device=device, + ) + for _ in range(g["num_layers"]) + ] + + +# One entry per geometry class the launcher must cover. expected_vectorized +# white-boxes the launch-time path selection so the matrix provably exercises +# both CUDA load paths. +_CASES = [ + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=8, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + True, + id="production_bf16_f64_group4_mean", + ), + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=8, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0, 8.0], + aggregation="max", + ), + True, + id="max_aggregation_four_offsets", + ), + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=4, + num_kv_heads=2, + dtype=torch.float16, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + True, + id="fp16_pool", + ), + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=6, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + True, + id="group3_generic_head_mapping", + ), + pytest.param( + dict( + head_dim=32, + tokens_per_block=32, + num_q_heads=4, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + True, + id="f16_runtime_chunk_count", + ), + pytest.param( + dict( + head_dim=8, + tokens_per_block=4, + num_q_heads=4, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + False, + id="tiny_f4_scalar", + ), + pytest.param( + dict( + head_dim=8, + tokens_per_block=4, + num_q_heads=4, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0], + aggregation="max", + ), + False, + id="tiny_f4_scalar_max", + ), + pytest.param( + dict( + head_dim=12, + tokens_per_block=4, + num_q_heads=4, + num_kv_heads=2, + dtype=torch.bfloat16, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + False, + id="f6_nonpow2_masked_tail", + ), + # fp32 pools always take the scalar path; group=3 additionally exercises + # its runtime GQA loop against the reference (the other generic-group + # case is vectorized, the other scalar cases use a templated group size). + pytest.param( + dict( + head_dim=8, + tokens_per_block=4, + num_q_heads=6, + num_kv_heads=2, + dtype=torch.float32, + offsets=[1.0, 2.0, 4.0], + aggregation="mean", + ), + False, + id="fp32_scalar_generic_group3", + ), +] + + +class TestTriAttentionScoreOps: + @pytest.mark.parametrize("case,expected_vectorized", _CASES) + def test_cuda_ops_match_torch_oracle(self, case, expected_vectorized): + _require_score_ops() + case = dict(case) # parametrize reuses the dict across reruns + aggregation = case.pop("aggregation") + request_count = 3 + prompt_len = 5 + ( + group, + round_starts, + token_starts, + valid_seq_lens, + seq_lens, + mean_cos, + mean_sin, + oracle_inputs, + ) = _build_case( + request_count=request_count, + max_requests=4, + num_layers=2, + page_count=4, + prompt_len=prompt_len, + seed=20260719, + **case, + ) + assert group._use_vectorized == expected_vectorized + device = group.output.device + sentinel = -54321.0 + + group.output.fill_(sentinel) + valid_widths_cuda = torch.empty(request_count, dtype=torch.int32, device=device) + cuda_scores = group.launch( + request_count, + valid_seq_lens, + valid_widths_cuda, + round_starts, + token_starts, + mean_cos, + mean_sin, + aggregation, + ).clone() + + # The oracle reads the same stored pool elements (up-cast to fp32, + # like the ops' loads), so the legs differ only by coefficient-fold + # association and reduction order. + reference = _oracle_reference( + group, + group.test_pools_keepalive, + oracle_inputs, + request_count, + seq_lens, + round_starts, + prompt_len, + aggregation, + sentinel, + ) + + assert valid_widths_cuda.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + # Sentinel-filled outputs make the comparison cover the write MASK as + # well as the values: any stray or missing store breaks equality. + # The ops' fp32 math tracks this oracle to ~2e-6 on these geometries, + # so 1e-4 is a tight gate with ample margin. + torch.testing.assert_close(cuda_scores, reference[:request_count], rtol=1e-4, atol=1e-4) + + @pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.int8], ids=["fp8_e4m3fn", "int8"]) + @pytest.mark.parametrize("aggregation", ["mean", "max"]) + def test_quantized_pool_matches_dequantized_reference(self, dtype, aggregation): + """Scale-folded scoring of quantized pools == dense scoring of their fp32 twin. + + The max aggregation additionally covers the per-offset coefficient + planes, which must all carry the folded per-layer scale. + """ + _require_score_ops() + device = torch.device("cuda", torch.cuda.current_device()) + seed = 20260721 + raw_pools = _build_quantized_raw_pools(seed, device) + quant_pools, dequant_pools, kv_scales = _quantize_pools(raw_pools, dtype) + ( + quant_group, + round_starts, + token_starts, + valid_seq_lens, + seq_lens, + mean_cos, + mean_sin, + oracle_inputs, + ) = _build_case( + dtype=dtype, + seed=seed, + pools=quant_pools, + kv_scales=kv_scales, + **_QUANTIZED_GEOMETRY, + ) + # Quantized pools must never select the vectorized load path. + assert quant_group._use_vectorized is False + request_count = _QUANTIZED_GEOMETRY["request_count"] + sentinel = -54321.0 + + quant_group.output.fill_(sentinel) + valid_widths_cuda = torch.empty(request_count, dtype=torch.int32, device=device) + cuda_scores = quant_group.launch( + request_count, + valid_seq_lens, + valid_widths_cuda, + round_starts, + token_starts, + mean_cos, + mean_sin, + aggregation, + ).clone() + + # Reference leg: the ORACLE over the dequantized-from-quantized fp32 + # twin pools (same page tables and calibration as the quantized group). + reference = _oracle_reference( + quant_group, + dequant_pools, + oracle_inputs, + request_count, + seq_lens, + round_starts, + _QUANTIZED_GEOMETRY["prompt_len"], + aggregation, + sentinel, + ) + + assert valid_widths_cuda.tolist() == [ + seq_len - _QUANTIZED_GEOMETRY["prompt_len"] for seq_len in seq_lens + ] + # Quantization error is identical in both legs (the reference pool is + # dequantized from the quantized values), so the only differences are + # scale-fold association (q*(s*c) vs (q*s)*c), the approximate sqrt, + # and reduction order — hence a tolerance close to the float cases'. + torch.testing.assert_close(cuda_scores, reference[:request_count], rtol=3e-3, atol=3e-3) + + def test_quantized_pool_missing_scales_raises(self): + with pytest.raises(ValueError, match="require per-layer kv_scales"): + _build_case(dtype=torch.int8, seed=20260722, **_QUANTIZED_GEOMETRY) + + def test_scales_with_float_pool_raises(self): + device = torch.device("cuda", torch.cuda.current_device()) + scales = torch.ones(_QUANTIZED_GEOMETRY["num_layers"], device=device) + with pytest.raises(ValueError, match="only valid for quantized"): + _build_case( + dtype=torch.bfloat16, seed=20260722, kv_scales=scales, **_QUANTIZED_GEOMETRY + ) + + def test_negative_scale_raises(self): + """Positivity is enforced host-side in the C++ op, at launch time. + + The |K| coefficient fold assumes |scale * K_q| == scale * |K_q|, which + breaks silently for non-positive scales, so the op must refuse them. + """ + _require_score_ops() + device = torch.device("cuda", torch.cuda.current_device()) + seed = 20260722 + raw_pools = _build_quantized_raw_pools(seed, device) + quant_pools, _, kv_scales = _quantize_pools(raw_pools, torch.int8) + bad_scales = kv_scales.clone() + bad_scales[0] = -bad_scales[0] + group, round_starts, token_starts, valid_seq_lens, _, mean_cos, mean_sin, _ = _build_case( + dtype=torch.int8, + seed=seed, + pools=quant_pools, + kv_scales=bad_scales, + **_QUANTIZED_GEOMETRY, + ) + request_count = _QUANTIZED_GEOMETRY["request_count"] + valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) + with pytest.raises(RuntimeError, match="strictly positive"): + group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_starts, + token_starts, + mean_cos, + mean_sin, + "mean", + ) + + def test_short_scales_raise(self): + """kv_scales must cover every calibrated layer, enforced at launch. + + The Python group only gates presence; the extent contract lives in the + C++ ops (segments index the fold tables by absolute layer id on + device, where a short scale tensor could not be range-checked). + """ + _require_score_ops() + device = torch.device("cuda", torch.cuda.current_device()) + seed = 20260722 + raw_pools = _build_quantized_raw_pools(seed, device) + quant_pools, _, kv_scales = _quantize_pools(raw_pools, torch.int8) + short_scales = kv_scales[:1] # the geometry calibrates two layers + group, round_starts, token_starts, valid_seq_lens, _, mean_cos, mean_sin, _ = _build_case( + dtype=torch.int8, + seed=seed, + pools=quant_pools, + kv_scales=short_scales, + **_QUANTIZED_GEOMETRY, + ) + request_count = _QUANTIZED_GEOMETRY["request_count"] + valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) + with pytest.raises(RuntimeError, match="one scale per calibrated layer"): + group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_starts, + token_starts, + mean_cos, + mean_sin, + "mean", + ) + + def test_unsupported_pool_dtype_raises(self): + _require_score_ops() + # fp32 pools stay supported (the existing tiny-geometry unit suite + # drives them through this launcher); fp64 is genuinely outside the + # op's coverage and must fail loudly instead of routing elsewhere. + request_count = 2 + group, round_starts, token_starts, valid_seq_lens, _, mean_cos, mean_sin, _ = _build_case( + request_count=request_count, + max_requests=2, + num_layers=1, + page_count=2, + prompt_len=1, + seed=20260720, + head_dim=8, + tokens_per_block=4, + num_q_heads=2, + num_kv_heads=1, + dtype=torch.float64, + offsets=[1.0, 2.0], + ) + valid_widths = torch.empty(request_count, dtype=torch.int32, device=group.output.device) + with pytest.raises(RuntimeError, match="unsupported KV pool dtype"): + group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_starts, + token_starts, + mean_cos, + mean_sin, + "mean", + ) From aba76091613bd5db153e498c0aef6c205790629a Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 19 Jul 2026 22:01:16 -0700 Subject: [PATCH 046/178] [None][chore] Trim redundant TriAttention code and dedupe test fixtures Remove dead selector setters and unreachable branches, require the constructor arguments every caller already passes, collapse three identical tensor-check helpers in the score op into one, and fix comments that no longer matched the code. Shared test fixtures (fake V2 manager factory, PyTorch score oracle, eviction-internals mock) move into the package conftest; the eager-workbench test file is renamed to test_triattention_selection_compaction to describe what it actually covers. Net -198 lines, no behavior change. Signed-off-by: tianruih --- .../triAttentionScoreKernels.cu | 81 ++++--- .../triAttentionScoreKernels.h | 7 +- cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp | 77 +++--- .../_torch/kv_cache_compression/interface.py | 8 +- .../triattention/compaction.py | 32 +-- .../triattention/triattention.py | 73 ++---- .../triattention/triattention_kernels.py | 9 +- .../test_kv_cache_compression_manager.py | 14 +- .../_torch/kv_cache_compression/conftest.py | 170 +++++++++++++ .../test_triattention_draft_cocompaction.py | 97 +------- .../test_triattention_pipeline.py | 226 ++---------------- .../test_triattention_score_ops.py | 75 +----- ...test_triattention_selection_compaction.py} | 9 +- 13 files changed, 331 insertions(+), 547 deletions(-) rename tests/unittest/_torch/kv_cache_compression/{test_triattention_eager.py => test_triattention_selection_compaction.py} (99%) diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu index 48f96760d8ba..a2a4120bc23c 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu @@ -35,11 +35,12 @@ // chunks of 8 frequencies when the pool layout allows it, otherwise a fully // strided scalar path runs the same math. // -// The kernel reassociates the frequency reduction relative to the Triton -// reference (sequential chunks instead of a block-wide tree), so results are -// tolerance-equal, not bit-equal. The valid-width side store IS replicated -// bit-exactly: first tile, first head, first segment of each request, thread -// 0, before any early-out. +// The kernel accumulates the frequency reduction in sequential chunks (not a +// block-wide tree), so results are tolerance-equal, not bit-equal, against +// the unit tests' PyTorch oracle (the in-tree reference; the original Triton +// score kernel has been deleted). The valid-width side store IS exact: first +// tile, first head, first segment of each request, thread 0, before any +// early-out. // // This file must NOT be compiled with --use_fast_math: the fold kernel's // cosf/sinf and the scalar-path precision are part of the accuracy contract. @@ -65,10 +66,11 @@ namespace kernels::tri_attention_score namespace { -// |K| uses the hardware approximate square root (one MUFU op). This matches -// the Triton reference's tl.sqrt lowering and is gated by the unit suite's -// tolerance comparison. Define TRTLLM_TRI_ATTENTION_IEEE_SQRT to restore the -// IEEE sqrtf sequence if a future geometry needs the extra bits. +// |K| uses the hardware approximate square root (one MUFU op), gated by the +// unit suite's tolerance comparison against the PyTorch oracle (this matched +// the deleted Triton score kernel's tl.sqrt lowering). Define +// TRTLLM_TRI_ATTENTION_IEEE_SQRT to restore the IEEE sqrtf sequence if a +// future geometry needs the extra bits. __device__ __forceinline__ float triSqrtApprox(float x) { #ifdef TRTLLM_TRI_ATTENTION_IEEE_SQRT @@ -264,9 +266,8 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne int const reqId = a.segRequestIds[seg]; int const seqLen = a.requestSeqLens[reqId]; int const tokenStart = a.requestTokenStarts[reqId]; - // Valid-width side store: same predicate as the Triton reference, before - // any early-out, so it fires exactly once per request even when the - // decode region is empty. + // Valid-width side store: evaluated before any early-out, so it fires + // exactly once per request even when the decode region is empty. if (blockIdx.x == 0 && blockIdx.z == 0 && (seg % a.numLayers) == 0 && threadIdx.x == 0) { a.validWidthOut[reqId] = seqLen - tokenStart; @@ -616,6 +617,33 @@ void launchScalar(FoldedScoreParams const& params, dim3 grid, bool useMax, cudaS } } +// Launch flavor for bf16/fp16 pools, the only element types owning both load +// paths (the vectorized 16-byte chunk kernel and the strided scalar kernel). +template +void launchVectorizedOrScalar( + FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useVectorized, bool useMax, cudaStream_t stream) +{ + if (useVectorized) + { + launchVectorized(params, groupSize, grid, useMax, stream); + } + else + { + launchScalar(params, grid, useMax, stream); + } +} + +// Quantized pools are functional-only: no vectorized instantiation exists for +// them by design (their dequant scale is folded into the coefficients, so +// only the scalar load path knows how to read them). +template +void launchQuantizedScalar( + FoldedScoreParams const& params, dim3 grid, bool useVectorized, bool useMax, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); + launchScalar(params, grid, useMax, stream); +} + } // namespace void foldScoreCoefficientsLaunch(float const* qReal, float const* qImag, float const* mlrCoef, float const* freqScaleSq, @@ -653,24 +681,10 @@ void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType switch (poolType) { case PoolElementType::kBFloat16: - if (useVectorized) - { - launchVectorized<__nv_bfloat16>(params, groupSize, grid, useMax, stream); - } - else - { - launchScalar<__nv_bfloat16>(params, grid, useMax, stream); - } + launchVectorizedOrScalar<__nv_bfloat16>(params, groupSize, grid, useVectorized, useMax, stream); break; case PoolElementType::kHalf: - if (useVectorized) - { - launchVectorized(params, groupSize, grid, useMax, stream); - } - else - { - launchScalar(params, grid, useMax, stream); - } + launchVectorizedOrScalar(params, groupSize, grid, useVectorized, useMax, stream); break; case PoolElementType::kFloat32: // fp32 pools have 32-byte 8-frequency rows; the 16-byte chunk path @@ -679,16 +693,9 @@ void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType launchScalar(params, grid, useMax, stream); break; case PoolElementType::kFloat8E4M3: - // Quantized pools are functional-only: no vectorized instantiation - // exists for them by design (their dequant scale is folded into the - // coefficients, so only the scalar load path knows how to read them). - TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); - launchScalar<__nv_fp8_e4m3>(params, grid, useMax, stream); - break; - case PoolElementType::kInt8: - TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); - launchScalar(params, grid, useMax, stream); + launchQuantizedScalar<__nv_fp8_e4m3>(params, grid, useVectorized, useMax, stream); break; + case PoolElementType::kInt8: launchQuantizedScalar(params, grid, useVectorized, useMax, stream); break; } TLLM_CUDA_CHECK(cudaGetLastError()); } diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h index 226d54275e81..2219b8388e02 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h @@ -51,8 +51,11 @@ enum class PoolElementType : int32_t }; // Upper bound on per-offset accumulators held by one score thread on the -// "max" aggregation path. Production configurations use 4 offsets; 8 leaves -// headroom without bloating the per-thread register budget. +// "max" aggregation path. The production "mean" path folds every offset into +// ONE coefficient plane, so this bound never constrains it; "max" needs one +// plane per offset, and the default geometric offset table exceeds 8, so a +// "max" run trips the fold op's TORCH_CHECK unless the offset budget is +// reduced. 8 keeps headroom without bloating the per-thread register budget. inline constexpr int32_t kMaxScoreOffsets = 8; // Threads per score CTA; one thread scores one cached token. diff --git a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp index e87c819f83d9..401c815f289c 100644 --- a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp +++ b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp @@ -15,7 +15,6 @@ */ #include "tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h" -#include "tensorrt_llm/thop/thUtils.h" #include #include @@ -31,22 +30,12 @@ namespace tk = tensorrt_llm::kernels::tri_attention_score; namespace { -void checkContiguousCudaFloat(torch::Tensor const& tensor, char const* name) +// One dtype-parameterized operand validator; dtypeName keeps the message +// spelling (fp32 / int32 / int64) for each expected scalar type. +void checkContiguousCuda(torch::Tensor const& tensor, at::ScalarType dtype, char const* dtypeName, char const* name) { - TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == at::kFloat, name, - " must be a contiguous fp32 CUDA tensor"); -} - -void checkContiguousCudaInt(torch::Tensor const& tensor, char const* name) -{ - TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == at::kInt, name, - " must be a contiguous int32 CUDA tensor"); -} - -void checkContiguousCudaLong(torch::Tensor const& tensor, char const* name) -{ - TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == at::kLong, name, - " must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == dtype, name, + " must be a contiguous ", dtypeName, " CUDA tensor"); } // Validate the optional per-layer dequantization scales for quantized @@ -65,7 +54,7 @@ float const* checkKvScales( { return nullptr; } - checkContiguousCudaFloat(*kv_scales, "kv_scales"); + checkContiguousCuda(*kv_scales, at::kFloat, "fp32", "kv_scales"); TORCH_CHECK(kv_scales->numel() >= num_calibrated_layers, op_name, ": kv_scales must carry one scale per calibrated layer (absolute layer id indexed)"); TORCH_CHECK(kv_scales->min().item() > 0.0f, op_name, @@ -88,13 +77,13 @@ void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_freqs, int64_t num_offsets, bool use_max, std::optional kv_scales) { - checkContiguousCudaFloat(c_re, "c_re"); - checkContiguousCudaFloat(c_im, "c_im"); - checkContiguousCudaFloat(c_mlr, "c_mlr"); - checkContiguousCudaFloat(q_real, "q_real"); - checkContiguousCudaFloat(q_imag, "q_imag"); - checkContiguousCudaFloat(mlr_coef, "mlr_coef"); - checkContiguousCudaFloat(freq_scale_sq, "freq_scale_sq"); + checkContiguousCuda(c_re, at::kFloat, "fp32", "c_re"); + checkContiguousCuda(c_im, at::kFloat, "fp32", "c_im"); + checkContiguousCuda(c_mlr, at::kFloat, "fp32", "c_mlr"); + checkContiguousCuda(q_real, at::kFloat, "fp32", "q_real"); + checkContiguousCuda(q_imag, at::kFloat, "fp32", "q_imag"); + checkContiguousCuda(mlr_coef, at::kFloat, "fp32", "mlr_coef"); + checkContiguousCuda(freq_scale_sq, at::kFloat, "fp32", "freq_scale_sq"); TORCH_CHECK(num_requests > 0 && num_calibrated_layers > 0 && num_query_heads > 0 && num_freqs > 0, "tri_attention_fold_score_coefficients: fold extents must be positive"); TORCH_CHECK(num_offsets >= 1 && num_offsets <= tk::kMaxScoreOffsets, @@ -118,9 +107,9 @@ void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, { TORCH_CHECK(omega.has_value() && offsets.has_value() && round_starts.has_value(), "tri_attention_fold_score_coefficients: max aggregation requires omega, offsets, and round_starts"); - checkContiguousCudaFloat(*omega, "omega"); - checkContiguousCudaFloat(*offsets, "offsets"); - checkContiguousCudaInt(*round_starts, "round_starts"); + checkContiguousCuda(*omega, at::kFloat, "fp32", "omega"); + checkContiguousCuda(*offsets, at::kFloat, "fp32", "offsets"); + checkContiguousCuda(*round_starts, at::kInt, "int32", "round_starts"); TORCH_CHECK( omega->numel() >= num_freqs && offsets->numel() >= num_offsets && round_starts->numel() >= num_requests, "tri_attention_fold_score_coefficients: max-path inputs are undersized for the folded request count"); @@ -132,8 +121,8 @@ void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, { TORCH_CHECK(mean_cos.has_value() && mean_sin.has_value(), "tri_attention_fold_score_coefficients: mean aggregation requires mean_cos and mean_sin"); - checkContiguousCudaFloat(*mean_cos, "mean_cos"); - checkContiguousCudaFloat(*mean_sin, "mean_sin"); + checkContiguousCuda(*mean_cos, at::kFloat, "fp32", "mean_cos"); + checkContiguousCuda(*mean_sin, at::kFloat, "fp32", "mean_sin"); TORCH_CHECK(mean_cos->numel() >= num_requests * num_freqs && mean_sin->numel() >= num_requests * num_freqs, "tri_attention_fold_score_coefficients: mean_cos/mean_sin are undersized (the fold iterates one row per " "folded request)"); @@ -167,18 +156,18 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas { TORCH_CHECK(use_max || num_offsets == 1, "tri_attention_paged_score: mean aggregation consumes exactly one folded coefficient plane"); - checkContiguousCudaLong(layer_base_addrs, "layer_base_addrs"); - checkContiguousCudaInt(block_offsets, "block_offsets"); - checkContiguousCudaLong(seg_page_offsets, "seg_page_offsets"); - checkContiguousCudaInt(seg_request_ids, "seg_request_ids"); - checkContiguousCudaInt(seg_layer_ids, "seg_layer_ids"); - checkContiguousCudaInt(request_seq_lens, "request_seq_lens"); - checkContiguousCudaInt(valid_widths, "valid_widths"); - checkContiguousCudaInt(request_token_starts, "request_token_starts"); - checkContiguousCudaFloat(c_re, "c_re"); - checkContiguousCudaFloat(c_im, "c_im"); - checkContiguousCudaFloat(c_mlr, "c_mlr"); - checkContiguousCudaFloat(out, "out"); + checkContiguousCuda(layer_base_addrs, at::kLong, "int64", "layer_base_addrs"); + checkContiguousCuda(block_offsets, at::kInt, "int32", "block_offsets"); + checkContiguousCuda(seg_page_offsets, at::kLong, "int64", "seg_page_offsets"); + checkContiguousCuda(seg_request_ids, at::kInt, "int32", "seg_request_ids"); + checkContiguousCuda(seg_layer_ids, at::kInt, "int32", "seg_layer_ids"); + checkContiguousCuda(request_seq_lens, at::kInt, "int32", "request_seq_lens"); + checkContiguousCuda(valid_widths, at::kInt, "int32", "valid_widths"); + checkContiguousCuda(request_token_starts, at::kInt, "int32", "request_token_starts"); + checkContiguousCuda(c_re, at::kFloat, "fp32", "c_re"); + checkContiguousCuda(c_im, at::kFloat, "fp32", "c_im"); + checkContiguousCuda(c_mlr, at::kFloat, "fp32", "c_mlr"); + checkContiguousCuda(out, at::kFloat, "fp32", "out"); TORCH_CHECK(pool_anchor.is_cuda(), "tri_attention_paged_score: pool anchor must be a CUDA tensor"); TORCH_CHECK(num_segments > 0 && num_segments <= 65535, @@ -188,8 +177,10 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas "tri_attention_paged_score: geometry extents must be positive"); TORCH_CHECK(num_kv_heads > 0 && num_query_heads % num_kv_heads == 0, "tri_attention_paged_score: query heads must be divisible by KV heads"); - // Production configurations use 4 score offsets; 8 is the per-thread - // accumulator budget baked into the kernels. + // kMaxScoreOffsets is the per-thread accumulator budget baked into the + // kernels. It only constrains the "max" path (one coefficient plane per + // offset); the "mean" path always folds every offset into one plane, so + // the default geometric offset table (which is larger) still passes here. TORCH_CHECK(num_offsets >= 1 && num_offsets <= tk::kMaxScoreOffsets, "tri_attention_paged_score: num_offsets must be in [1, ", tk::kMaxScoreOffsets, "], got ", num_offsets); TORCH_CHECK(seg_page_offsets.numel() >= num_segments && seg_request_ids.numel() >= num_segments diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index c50891eb9529..06890572f0fe 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -8,16 +8,14 @@ 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. + Configs map their ``algorithm`` string to a member here; feature gates + read the ``is_*`` trait predicates (algorithm dispatch itself matches the + config's ``algorithm`` string). """ TRIATTENTION = auto() NONE = auto() - def is_triattention(self): - return self == KvCacheCompressionMode.TRIATTENTION - def is_eviction_method(self): """Whether this method physically evicts cached tokens.""" return self == KvCacheCompressionMode.TRIATTENTION diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index ab94a4a5f69f..fedca37eda59 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -17,9 +17,11 @@ Given each request's kept-token ordinals, its valid sequence length, and the staged V2 block offsets, this module packs per-request move indices with one -Triton launch and then moves the surviving KV in place with batched -C++ compact launches. Inputs are plain tensors, so any eviction method that -produces a kept-token set per request can drive it. +Triton launch per compacted cache (one launch covers the target's dense and +SWA families; a co-compressed draft adds a second) and then moves the +surviving KV in place with batched C++ compact launches. Inputs are plain +tensors, so any eviction method that produces a kept-token set per request +can drive it. """ from collections import OrderedDict @@ -78,16 +80,6 @@ def compact(self) -> None: ) -def _cuda_int32_contiguous(tensors: Tuple[torch.Tensor, ...], device: torch.device) -> bool: - return all( - tensor.is_cuda - and tensor.dtype == torch.int32 - and tensor.device == device - and tensor.is_contiguous() - for tensor in tensors - ) - - def _validated_kv_head_count( pools: List[torch.Tensor], layers: Tuple[int, ...], @@ -304,8 +296,7 @@ def _move_index_pack_launcher( swa_offsets_arg, swa_indices_arg, ) - # Ordered to match the kernel's constexpr parameter declaration: the - # ordered to match the kernel constexpr declaration. + # Ordered to match the kernel's constexpr parameter declaration. constexpr_values = dict( DENSE_TOTAL=int(move_source_indices.shape[-1]), SWA_TOTAL=swa_total, @@ -355,9 +346,10 @@ class BatchedKVCacheCompaction: slot must share one block-offset table. `protected_tail_capacity`: widest per-request protected tail this object must support. A protected tail covers KV positions past - the valid length reserved for a forward already in flight; the - actual per-round lengths are loaded via `set_protected_tails` - and move with the kept tokens. + the valid length reserved for a forward already in flight; each + round's actual lengths arrive through the per-family move-offset + rows staged with the round metadata (`set_protected_tails` fills + the same rows for standalone use) and move with the kept tokens. `draft_*`: co-compressed draft-cache layout (union mode only); the draft reuses the target keep set and pins the same prompt. """ @@ -378,8 +370,8 @@ def __init__( prompt_offsets: torch.Tensor, decode_keep_count: int, swa_window: Optional[int], + layer_pool_keys: List[object], protected_tail_capacity: int = 0, - layer_pool_keys: Optional[List[object]] = None, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_layers: Optional[List[int]] = None, draft_layer_group_representative: Optional[Dict[int, int]] = None, @@ -428,8 +420,6 @@ def __init__( self.protected_tail_capacity = int(protected_tail_capacity) self.dense_layers = tuple(int(layer) for layer in dense_layers) self.swa_layers = tuple(int(layer) for layer in swa_layers) - if layer_pool_keys is None: - layer_pool_keys = [("layer", layer) for layer in range(len(layer_pools))] if len(layer_pool_keys) != len(layer_pools): raise ValueError("pool keys must match the layer-pool count") self.layer_pool_keys = tuple(layer_pool_keys) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d9654deb5adf..ee9a4a4bc6ee 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -94,7 +94,7 @@ class _FixedScoreStreamMismatch(RuntimeError): class _CrossRequestSelectionPlan(NamedTuple): - """Selection dimensions used to allocate reusable eager buffers.""" + """Selection dimensions used to allocate reusable fixed buffers.""" eviction_mode: str dense_layers: Tuple[int, ...] @@ -154,7 +154,6 @@ def __init__( self.num_kv_heads = int(num_kv_heads) self.width = int(width) self.keep_count = int(keep_count) - self.dtype = dtype self.device = device self.max_requests = int(max_requests) self.valid_widths = torch.full( @@ -188,14 +187,6 @@ def __init__( device=self.device, ) - def set_prompt_offsets(self, prompt_lens: torch.Tensor) -> None: - """Refresh the per-request prompt offsets for the coming round.""" - count = int(prompt_lens.numel()) - if count > self.max_requests: - raise ValueError("prompt offsets exceed the selector's request capacity") - self.prompt_offsets[:count].copy_(prompt_lens, non_blocking=True) - self.refresh_row_prompt_offsets() - def refresh_row_prompt_offsets(self) -> None: """Re-expand the per-request prompt offsets into their row-major view. @@ -285,7 +276,6 @@ def __init__( device=device, max_requests=max_requests, ) - self.rows = rows if input_scores is None: raise ValueError("union selection requires its fixed score input") @@ -320,17 +310,6 @@ def select_prepared_requests(self) -> None: ) self._select_top_tokens() - def select_requests( - self, - scores: torch.Tensor, - *, - normalize_scores: bool, - ) -> None: - """Select from the score tensor bound to this fixed selector.""" - if scores is not self.input_scores or bool(normalize_scores) != self.normalize_scores: - raise ValueError("union scores do not match this selector's bound input") - self.select_prepared_requests() - class _BatchedPerHeadKeepSetSelector(_BatchedKeepSetSelectorBase): """Fixed ``[request, ...]`` selector for both per-head modes.""" @@ -372,7 +351,6 @@ def __init__( max_requests=max_requests, ) self.num_layers = len(self.dense_layers) - self.rows = self.num_layers * self.num_query_heads self.selection_rows = selection_rows score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) @@ -762,7 +740,7 @@ def stage( swa_move_offsets: Optional[List[int]] = None, draft_move_offsets: Optional[List[int]] = None, ) -> bool: - """Copy one eager eviction cohort into reusable device buffers. + """Copy one eviction cohort into reusable device buffers. ``token_starts`` carries each request's pinned prompt length; the score kernel starts that request's decode window there, so the cohort @@ -940,7 +918,7 @@ def mark_page_tables_consumed(self, *manager_streams: torch.cuda.Stream) -> None @dataclass(frozen=True, kw_only=True, slots=True) class _PreparedEviction: - """Request metadata validated before eager score, select, and compact.""" + """Request metadata validated before score, select, and compact.""" request: "LlmRequest" request_id: int @@ -970,7 +948,7 @@ class _PreparedGenerationBatch: @dataclass(kw_only=True, slots=True) class _EvictionBuffers: - """Reusable eager score and selection buffers for one runtime shape.""" + """Reusable fixed score and selection buffers for one runtime shape.""" score_staging: _FixedScoreStagingBuffers keep_set_selector: Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector] @@ -1009,16 +987,9 @@ def __init__( self.beta = beta if self.top_B <= 0 or self.beta <= 0: raise ValueError("TriAttention top_B and beta must both be positive") - # Which token set each eviction round keeps (all reproduce the upstream - # selection: z-normalize scores, pin the prompt tokens, no recency window): - # union -- union of every KV head's top-B, re-ranked by the - # per-token max score. Default; matches the official - # base setting (per-head and per-layer-per-head - # pruning both off). - # per_head -- each KV head keeps its own set, shared across - # layers (mean of per-layer max). - # per_layer_perhead -- each (layer, KV head) keeps its own set, fully - # independent per layer. + # Which token set each eviction round keeps. The user-facing meaning of + # each mode is documented on TriAttentionKvCacheCompressionConfig + # (llm_args); implementation notes live above the selection helpers. self.eviction_mode = eviction_mode if self.eviction_mode not in ("union", "per_head", "per_layer_perhead"): raise ValueError( @@ -1418,12 +1389,11 @@ def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: With a decode-only budget, pinned prompt tokens do not consume ``top_B``. Selection therefore keeps every token until the cache exceeds - ``prompt_len + top_B``. + ``prompt_len + top_B``. The constructor guarantees the decode-only + budget (``pin_prefill=True``, ``count_prompt_tokens=False``). """ - if self.pin_prefill and not self.count_prompt_tokens: - prompt_len = min(int(request.py_prompt_len), seq_len) - return prompt_len + self.top_B - return self.top_B + prompt_len = min(int(request.py_prompt_len), seq_len) + return prompt_len + self.top_B def _local_score_calibration( self, @@ -1642,7 +1612,7 @@ def _attention_layer_partition( return result def _runtime_kv_layout(self, num_layers: int) -> _RuntimeKVLayout: - """Return stable V2 pool views and layer groups for eager eviction. + """Return stable V2 pool views and layer groups for eviction. KVCacheManagerV2 keeps GPU virtual addresses and layer geometry stable, while opt-in pool rebalance can change the page dimension. Cache all @@ -1819,7 +1789,7 @@ def _pool_view_fingerprint(pools: List[torch.Tensor]) -> Tuple[tuple, ...]: for pool in pools ) - def _eager_resources_for( + def _fixed_resources_for( self, layout: _RuntimeKVLayout, prepared: Sequence[_PreparedEviction], @@ -1831,8 +1801,7 @@ def _eager_resources_for( eviction bound (compaction keeps the scored decode region near ``top_B`` plus one period of growth), so one set of buffers serves every round. They are rebuilt only when the pool views change or a - round outgrows them (prompt-counting budgets score the prompt, whose - length is workload-defined). + round outgrows them. """ if not prepared: raise ValueError("TriAttention eviction requires at least one request") @@ -1969,7 +1938,7 @@ def _batched_compaction_for( score_staging: _FixedScoreStagingBuffers, keep_set_selector: Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector], ): - """Build or reuse the eager C++ compaction launches for one cohort.""" + """Build or reuse the C++ compaction launches for one cohort.""" from .compaction import BatchedKVCacheCompaction if layout.swa_layers and layout.swa_window: @@ -2143,9 +2112,10 @@ def _evict_requests( # Restore the uncompressed confirmed logical position from the # physical prefix and cumulative eviction count. round_start = seq_len + request_state.evicted_tokens - if seq_len <= self._minimum_evictable_length(request, seq_len): + minimum_evictable_length = self._minimum_evictable_length(request, seq_len) + if seq_len <= minimum_evictable_length: continue - expected_keep_count = self._minimum_evictable_length(request, seq_len) + expected_keep_count = minimum_evictable_length protected_tail = int(protected_tail_lengths.get(rid, 0)) if protected_tail < 0 or protected_tail > protected_tail_capacity: raise RuntimeError( @@ -2166,7 +2136,7 @@ def _evict_requests( if not prepared: return [] with nvtx_range_debug("triattention.staging_lookup", color="blue"): - resources = self._eager_resources_for(layout, prepared) + resources = self._fixed_resources_for(layout, prepared) score_staging = resources.score_staging keep_set_selector = resources.keep_set_selector batched_compaction = self._batched_compaction_for( @@ -2300,11 +2270,6 @@ def _rope_tables(self, freq_count: int): (the official file does not store them). transformers' rope-init handles plain and scaled RoPE; plain RoPE has attention_factor 1 so freq_scale_sq is all ones. Falls back to the analytic inv_freq if rope-init is absent.""" - if self.model_path is None: - raise ValueError( - "TriAttention needs `model_path` to derive the RoPE tables " - "(omega / freq_scale_sq) when converting the official calibration." - ) from transformers import AutoConfig cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 4a23c1f17432..f1d87ae795ff 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -12,8 +12,8 @@ House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). - * int64 for every page/stride offset that can exceed 2^31 (paged-pool reads). - * mask seq tails not divisible by ``tokens_per_block`` (and freq/dim tails). + * int64 for every flat buffer offset that can exceed 2^31. + * mask ragged valid-width tails (and frequency tails) in every load and store. * the kernels are vendored in this module (no lazy-load hub). """ @@ -248,14 +248,11 @@ def __init__( freq_scale_sq: torch.Tensor, omega: torch.Tensor, offsets: torch.Tensor, - output_width: int | None = None, + output_width: int, kv_scales: torch.Tensor | None = None, ) -> None: if not layer_indices or min(max_requests, page_count, seq_len) <= 0: raise ValueError("fixed score group requires non-empty positive geometry") - # Default: the whole sequence capacity is scorable. - if output_width is None: - output_width = int(seq_len) if output_width <= 0 or output_width > seq_len: raise ValueError("fixed score group requires a decode width within its capacity") if len(page_table_slots) != len(layer_indices): 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 53a9b9bd7eee..165cfabf1fd8 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -267,23 +267,17 @@ class TestFactory: def test_raises_when_no_algorithm_registered(self, fake_kv_cache_manager): # A config whose algorithm has no registered manager is a developer # error (config subclass added without a factory branch): fail loudly - # instead of silently running without compression. + # instead of silently running without compression. The draft-manager + # kwarg does not change the dispatch, so both forms raise identically. cfg = MagicMock() cfg.algorithm = "made_up_method" with pytest.raises(ValueError, match="no registered compression manager"): create_kv_cache_compression_manager(cfg, fake_kv_cache_manager) - - def test_unregistered_algorithm_raises_with_draft_manager_too(self): - cfg = MagicMock() - cfg.algorithm = "made_up_method" - target = _v2_manager(is_draft=False) - draft = _v2_manager(is_draft=True) - with pytest.raises(ValueError, match="no registered compression manager"): create_kv_cache_compression_manager( cfg, - target, - draft_kv_cache_manager=draft, + fake_kv_cache_manager, + draft_kv_cache_manager=_v2_manager(is_draft=True), ) def test_eviction_method_predicate_defaults_false(self): diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index e2eefb7f034d..69ffaf9b6e84 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -13,6 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + import torch @@ -35,3 +39,169 @@ def encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: encoded[:, :, 0] = page_ids.to(torch.int32) * 2 encoded[:, :, 1] = encoded[:, :, 0] + 1 return encoded + + +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.mapping = SimpleNamespace(enable_attention_dp=False) + fake_v2.is_disagg = False + 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.empty(1, dtype=torch.int64) + fake_v2.pp_layers = [] + fake_v2.layer_offsets = {} + fake_v2.layer_to_pool_mapping_dict = {} + return fake_v2 + + +def make_triattention(**overrides): + """Construct a fully initialized manager for method-level unit tests.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention + + options = {"top_B": 8, "model_path": "/models/test"} + options.update(overrides) + return TriAttention(make_fake_v2(), **options) + + +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_requests`` body around mocked GPU launches.""" + score_staging = SimpleNamespace( + launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), + mark_page_tables_consumed=mock.Mock(), + ) + keep_set_selector = SimpleNamespace( + select_requests=mock.Mock(), + refresh_row_prompt_offsets=mock.Mock(), + ) + resources = SimpleNamespace( + score_staging=score_staging, + keep_set_selector=keep_set_selector, + ) + batched_compaction = SimpleNamespace(compact=mock.Mock()) + with ( + mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), + mock.patch.object(manager, "_fixed_resources_for", return_value=resources), + mock.patch.object( + manager, + "_batched_compaction_for", + return_value=batched_compaction, + ), + mock.patch.object(manager, "_attach_page_ids") as attach, + ): + yield SimpleNamespace( + score_staging=score_staging, + keep_set_selector=keep_set_selector, + batched_compaction=batched_compaction, + attach=attach, + ) + + +def torch_tri_score_oracle( + layer_pools, + page_ids, + seq_lens, + round_starts, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + layer_indices, + aggregation, +): + """Independent Torch implementation of the paged TriAttention score. + + Covers both aggregations (mean and max), GQA head mapping via + ``head // group_size``, and the position-independent MLR term. + """ + scores = [] + num_q_heads = int(q_real.shape[1]) + for request, seq_len in enumerate(seq_lens): + phase = (round_starts[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 + if aggregation == "mean": + position = ( + freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) + ).sum(dim=-1) + else: + position = ( + ( + freq_scale_sq[None, None, :] + * ( + product_real[None] * torch.cos(phase)[:, None, :] + - product_imag[None] * torch.sin(phase)[:, None, :] + ) + ) + .sum(dim=-1) + .max(dim=0) + .values + ) + 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 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 index 1a8ad27752c8..4d7af4ea8ae4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -13,13 +13,16 @@ prepared-compaction cache invalidation. """ -from contextlib import contextmanager from types import SimpleNamespace from unittest import mock import pytest import torch from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_request as _make_request +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.compaction import ( BatchedKVCacheCompaction, @@ -30,60 +33,6 @@ _PreparedEviction, _RequestCompressionState, ) -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState - - -def _make_fake_v2(*, 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 = False - fake_v2.is_draft = is_draft - fake_v2.kv_compression_manages_history = False - fake_v2.kv_factor = 2 - fake_v2.mapping = SimpleNamespace(enable_attention_dp=False) - fake_v2.is_disagg = False - 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.empty(1, dtype=torch.int64) - fake_v2.pp_layers = [] - fake_v2.layer_offsets = {} - fake_v2.layer_to_pool_mapping_dict = {} - return fake_v2 - - -def _make_triattention(**overrides): - options = {"top_B": 8, "model_path": "/models/test"} - options.update(overrides) - return TriAttention(_make_fake_v2(), **options) - - -def _make_request(request_id, **overrides): - 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) def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: @@ -344,35 +293,6 @@ def test_draft_admission_gates_raise(gate, match): manager._validate_v2_compatibility() -@contextmanager -def _mocked_eviction_internals(manager): - """Run the real ``_evict_requests`` body around mocked GPU launches.""" - score_staging = SimpleNamespace( - launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), - mark_page_tables_consumed=mock.Mock(), - ) - keep_set_selector = SimpleNamespace( - select_requests=mock.Mock(), - refresh_row_prompt_offsets=mock.Mock(), - ) - resources = SimpleNamespace( - score_staging=score_staging, - keep_set_selector=keep_set_selector, - ) - batched_compaction = SimpleNamespace(compact=mock.Mock()) - with ( - mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), - mock.patch.object(manager, "_eager_resources_for", return_value=resources), - mock.patch.object( - manager, - "_batched_compaction_for", - return_value=batched_compaction, - ), - mock.patch.object(manager, "_attach_page_ids"), - ): - yield score_staging - - def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager = _make_triattention(top_B=4, beta=4) manager._calibrated = True @@ -403,7 +323,8 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): cache.capacity = confirmed previous_published = 0 eviction_rounds = 0 - with _mocked_eviction_internals(manager) as score_staging: + with _mocked_eviction_internals(manager) as internals: + score_staging = internals.score_staging for _ in range(6): uncompressed += 2 confirmed += 2 @@ -505,7 +426,7 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): return_value=keep_set_selector, ), ): - resources = manager._eager_resources_for(layout, prepared) + resources = manager._fixed_resources_for(layout, prepared) # The buffers follow the executor limits, not this one-request cohort. assert resources.score_staging is score_staging @@ -519,14 +440,14 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): # keeps the cached compaction launches. cached_compaction = object() manager._batched_compaction = cached_compaction - assert manager._eager_resources_for(layout, prepared) is resources + assert manager._fixed_resources_for(layout, prepared) is resources assert score_cls.call_count == 1 assert manager._batched_compaction is cached_compaction # A pool change invalidates both the buffers and the compaction # launches that alias them. layout.pool_view_fingerprint = (("moved",),) - rebuilt = manager._eager_resources_for(layout, prepared) + rebuilt = manager._fixed_resources_for(layout, prepared) assert rebuilt is not resources assert score_cls.call_count == 2 assert manager._batched_compaction is None diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index d81ec96061e6..39871d396015 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -25,12 +25,17 @@ Model-level correctness is covered by separate end-to-end tests. """ -from contextlib import contextmanager from types import SimpleNamespace from unittest import mock import pytest import torch +from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_request as _make_request +from conftest import make_triattention as _make_triattention +from conftest import mocked_eviction_internals as _mocked_eviction_internals +from conftest import torch_tri_score_oracle as _torch_tri_score_oracle from pydantic import ValidationError # TriAttention lives in the kv_cache_compression package. It exposes only the @@ -48,29 +53,11 @@ # in pyexecutor._util (next to _create_kv_cache_manager), matching #15106. from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig _TORCH_TOPK_ORACLE = torch.topk -def _encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: - """Build the native V2 [pool, request, K/V, block] layout.""" - if page_ids.ndim == 2: - page_ids = page_ids.unsqueeze(0) - encoded = torch.empty( - page_ids.shape[0], - page_ids.shape[1], - 2, - page_ids.shape[2], - dtype=torch.int32, - device=page_ids.device, - ) - encoded[:, :, 0] = page_ids.to(torch.int32) * 2 - encoded[:, :, 1] = encoded[:, :, 0] + 1 - return encoded - - def _set_request_state( manager, request_id, @@ -168,136 +155,12 @@ def flat_calibration_pt(tmp_path): return str(path) -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.mapping = SimpleNamespace(enable_attention_dp=False) - fake_v2.is_disagg = False - 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.empty(1, dtype=torch.int64) - fake_v2.pp_layers = [] - fake_v2.layer_offsets = {} - fake_v2.layer_to_pool_mapping_dict = {} - return fake_v2 - - -def _make_triattention(**overrides): - """Construct a fully initialized manager for method-level unit tests.""" - options = {"top_B": 8, "model_path": "/models/test"} - options.update(overrides) - return TriAttention(_make_fake_v2(), **options) - - -def _make_request(request_id, **overrides): - """Build the explicit request fields consumed by TriAttention.""" - 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) - - 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) -def _torch_tri_score_oracle( - layer_pools, - page_ids, - seq_lens, - round_starts, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - layer_indices, - aggregation, -): - """Independent Torch implementation of the paged TriAttention score.""" - scores = [] - num_q_heads = int(q_real.shape[1]) - for request, seq_len in enumerate(seq_lens): - phase = (round_starts[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 - if aggregation == "mean": - position = ( - freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) - ).sum(dim=-1) - else: - position = ( - ( - freq_scale_sq[None, None, :] - * ( - product_real[None] * torch.cos(phase)[:, None, :] - - product_imag[None] * torch.sin(phase)[:, None, :] - ) - ) - .sum(dim=-1) - .max(dim=0) - .values - ) - 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 - - class TestKvCacheCompressionConfig: def test_llm_args_dispatches_concrete_and_unknown_algorithms(self): from tensorrt_llm.llmapi.llm_args import TorchLlmArgs @@ -432,39 +295,6 @@ def test_manager_is_marked_capacity_only_and_requests_default_to_zero(self): # the first eviction publishes a count. assert request.py_num_compressed_tokens == 0 - @contextmanager - def _mocked_eviction_internals(self, manager): - """Run the real ``_evict_requests`` body around mocked GPU launches.""" - score_staging = SimpleNamespace( - launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), - mark_page_tables_consumed=mock.Mock(), - ) - keep_set_selector = SimpleNamespace( - select_requests=mock.Mock(), - refresh_row_prompt_offsets=mock.Mock(), - ) - resources = SimpleNamespace( - score_staging=score_staging, - keep_set_selector=keep_set_selector, - ) - batched_compaction = SimpleNamespace(compact=mock.Mock()) - with ( - mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), - mock.patch.object(manager, "_eager_resources_for", return_value=resources), - mock.patch.object( - manager, - "_batched_compaction_for", - return_value=batched_compaction, - ), - mock.patch.object(manager, "_attach_page_ids") as attach, - ): - yield SimpleNamespace( - score_staging=score_staging, - keep_set_selector=keep_set_selector, - batched_compaction=batched_compaction, - attach=attach, - ) - def test_eviction_bookkeeping_publishes_cumulative_count(self): # The eviction bookkeeping writes the cumulative evicted count on the # request in the same step that compacts the cache; this is the @@ -474,7 +304,7 @@ def test_eviction_bookkeeping_publishes_cumulative_count(self): request = _make_request(7, py_prompt_len=2) _set_request_state(manager, 7, confirmed_kv_length=10) - with self._mocked_eviction_internals(manager) as internals: + with _mocked_eviction_internals(manager) as internals: first = manager._evict_requests([(request, 7)], 2) assert first == [(7, 6)] @@ -488,7 +318,7 @@ def test_eviction_bookkeeping_publishes_cumulative_count(self): # Round two: 6 retained + 8 newly confirmed decode tokens. manager._request_states[7].confirmed_kv_length = 14 - with self._mocked_eviction_internals(manager) as internals: + with _mocked_eviction_internals(manager) as internals: second = manager._evict_requests([(request, 7)], 2) assert second == [(7, 6)] @@ -506,12 +336,12 @@ def test_identity_compaction_is_rejected_instead_of_published(self): # one token; seq_len == prompt + budget must never publish. _set_request_state(manager, 7, confirmed_kv_length=6) - with self._mocked_eviction_internals(manager): + with _mocked_eviction_internals(manager): assert manager._evict_requests([(request, 7)], 2) == [] assert request.py_num_compressed_tokens == 0 -class TestStepEndHookRefactor: +class TestEvictionLifecycle: def test_triattention_prepare_only_snapshots_and_update_uses_final_hook(self): assert "prepare_resources" in TriAttention.__dict__ assert "update_resources" not in TriAttention.__dict__ @@ -534,10 +364,9 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): periodic_evict.assert_called_once_with(batch) - @pytest.mark.parametrize("top_B", [511, 512]) - def test_non_v2_manager_is_always_rejected(self, top_B): + def test_non_v2_manager_is_always_rejected(self): with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - TriAttention(SimpleNamespace(), top_B=top_B) + TriAttention(SimpleNamespace(), top_B=8) @staticmethod def _make_due_decode_request(seq_len): @@ -567,8 +396,6 @@ def _make_due_decode_request(seq_len): _set_request_state(mgr, 7, generation_steps=127) mgr.beta = 128 mgr.top_B = 4096 - mgr.pin_prefill = True - mgr.count_prompt_tokens = False return mgr, request, batch def test_identity_gate_preserves_real_eviction_round(self): @@ -1002,7 +829,7 @@ def test_cross_request_union_matches_oracle_at_high_keep_counts(self, keep_count input_scores=scores, normalize_scores=False, ) - selector.select_requests(scores, normalize_scores=False) + selector.select_prepared_requests() selected = selector.keep.cpu() for actual, expected_keep in zip(selected, expected): @@ -1012,7 +839,7 @@ def test_cross_request_union_matches_oracle_at_high_keep_counts(self, keep_count class TestFixedScoreMetadata: @pytest.mark.parametrize("normalize_scores", [False, True]) @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) - def test_eager_buffers_bind_score_after_selection(self, eviction_mode, normalize_scores): + def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize_scores): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention( @@ -1070,7 +897,7 @@ def test_eager_buffers_bind_score_after_selection(self, eviction_mode, normalize return_value=keep_set_selector, ) as build_selection, ): - resources = manager._eager_resources_for(layout, prepared) + resources = manager._fixed_resources_for(layout, prepared) score_staging.bind_score_launcher.assert_called_once_with( keep_set_selector.valid_widths, @@ -1823,30 +1650,21 @@ def test_layer_partition_rejects_decode_budget_smaller_than_window(self): class TestFactory: - def test_returns_triattention_instance_with_v2(self): + def test_returns_triattention_instance_and_propagates_config_fields(self): # A plain V2 manager (block reuse off) yields a TriAttention instance. # Calibration is deferred to the first request, so construction needs # no calibration file or CUDA. fake_v2 = _make_fake_v2(enable_block_reuse=False) cfg = TriAttentionKvCacheCompressionConfig( - top_B=32, beta=16, model_path="/models/test", calibration_path="/calib/test.pt" - ) - mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) - assert isinstance(mgr, TriAttention) - assert mgr.top_B == 32 - assert mgr.beta == 16 - assert mgr.kv_cache_manager is fake_v2 - - def test_factory_propagates_eviction_mode(self): - cfg = TriAttentionKvCacheCompressionConfig( - top_B=64, - beta=8, + top_B=32, + beta=16, eviction_mode="per_head", model_path="/models/test", calibration_path="/calib/test.pt", ) - mgr = create_kv_cache_compression_manager( - cfg, kv_cache_manager=_make_fake_v2(enable_block_reuse=False) - ) + mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) + assert mgr.top_B == 32 + assert mgr.beta == 16 assert mgr.eviction_mode == "per_head" + assert mgr.kv_cache_manager is fake_v2 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py index 11d9361b0b15..2f2c39ce57c8 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py @@ -16,6 +16,7 @@ import pytest import torch from conftest import encode_block_offsets as _encode_block_offsets +from conftest import torch_tri_score_oracle as _torch_tri_score_oracle from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _FixedScoreGroup, @@ -32,80 +33,6 @@ def _require_score_ops() -> None: ) -# Copied verbatim from test_triattention_pipeline.py (_torch_tri_score_oracle) -# so this module stays runnable standalone: an independent pure-PyTorch -# implementation of the paged TriAttention score, covering both aggregations -# (mean and max), GQA head mapping via head // group_size, and the -# position-independent MLR term. -def _torch_tri_score_oracle( - layer_pools, - page_ids, - seq_lens, - round_starts, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - layer_indices, - aggregation, -): - """Independent Torch implementation of the paged TriAttention score.""" - scores = [] - num_q_heads = int(q_real.shape[1]) - for request, seq_len in enumerate(seq_lens): - phase = (round_starts[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 - if aggregation == "mean": - position = ( - freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) - ).sum(dim=-1) - else: - position = ( - ( - freq_scale_sq[None, None, :] - * ( - product_real[None] * torch.cos(phase)[:, None, :] - - product_imag[None] * torch.sin(phase)[:, None, :] - ) - ) - .sum(dim=-1) - .max(dim=0) - .values - ) - 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 _oracle_reference( group: _FixedScoreGroup, pools: list, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py similarity index 99% rename from tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py rename to tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py index 0a21d9b176e4..8aa7be2995fa 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_eager.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -241,9 +241,12 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid normalize_scores=False, ) selector.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) - selector.set_prompt_offsets( + # Write the shared per-request prompt lengths the way production staging + # does: fill the bound buffer, then re-expand the row-major view. + selector.prompt_offsets[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) + selector.refresh_row_prompt_offsets() selector.select_prepared_requests() actual = selector.keep.cpu() @@ -1008,11 +1011,11 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): ) -def test_workbench_families_read_the_staged_move_offsets_rows(): +def test_cache_families_read_the_staged_move_offsets_rows(): """Every cache family must consume the caller-staged offsets row. A family that silently falls back to its construction-time offsets - compacts workbench slots that are not in the staged cohort; on + compacts request slots that are not in the staged cohort; on sliding-window models the padded slots then produce negative source ordinals and an illegal memory access. Binding the staged row by reference is part of the constructor contract. From 6e0417fba2f9631d4c63f66f606a957e49152b63 Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 19 Jul 2026 22:01:17 -0700 Subject: [PATCH 047/178] [None][perf] Add pipelined bf16 KV-compact fast path behind a temporary A/B knob Port Fanrong Li's cp.async double-buffered compact kernels (store the current tile while the next one loads) alongside the existing register-staging kernel. Adapted to this branch's ABI: shared flat page table with per-request stride and 2*page entry decode, per-request destination bases read on device, and an explicit allocation-wide head stride for the move buffers. Geometry coverage extends the original 128-token/64-dim tiles to tokens_per_block 32 and head_dim 128 so the production Qwen3 and GPT-OSS configs take the fast path. TLLM_SPARSE_KV_COMPACT_FAST (default on, "0" restores the old kernel) is read on every launch so the same build can A/B both kernels; the knob is marked REMOVE-BEFORE-PR. Tests run every case under both kernel selections, assert byte-identical pools both ways, and use a profiler probe to prove the selected kernel actually launched. Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 491 ++++++++++++++++++ .../serial/test_sparse_kv_cache_compact.py | 257 ++++++++- 2 files changed, 741 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index eccc4e3cf984..21a90c3e3280 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -29,6 +29,10 @@ #include "tensorrt_llm/kernels/quantization.cuh" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" +#include +#include +#include + using namespace tensorrt_llm::common; TRTLLM_NAMESPACE_BEGIN @@ -1818,6 +1822,444 @@ struct KvCacheV2LayersBuffer } }; +#ifdef ENABLE_BF16 + +// Optimized bf16 compaction fast path, ported from Fanrong Li's optimized +// compact kernels (snapshot 2026-07-19). The port keeps the double-buffered +// cp.async pipeline intact and adapts only the addressing to this tree's ABI: +// (a) the per-layer page-table pointer array became one flat int32 V2 +// K-plane block-offset table shared by all layers (entries encode +// 2 * page + plane with plane == 0, so >> 1 recovers the page), strided +// per request; +// (b) the host-scalar destination base became per-request device bases, read +// once per CTA (one launch covers a cohort with mixed prompt lengths); +// (c) the head-plane stride of the move-source indices is an explicit +// parameter instead of being derived from sourceOffsets[batchSize] on +// device: the move buffers are allocation-wide, so a device-derived +// stride would silently read the wrong plane for every KV head above +// head 0. +// The original kernels were written for 128-token pages and Dh = 64; the port +// additionally parameterizes the page and head-vector math so 32-token pages +// and Dh = 128 (this tree's production geometry) take the same pipeline. + +namespace compact_detail +{ +// Vendored cp.async wrappers, equivalent to the ones in +// cpp/kernels/xqa/ldgsts.cuh. That header cannot be included from this +// widely-included template header because it drags in xqa's cuda_hint.cuh / +// barriers.cuh, whose macros and helpers would leak into every translation +// unit that includes this file. +template +__device__ __forceinline__ void copyAsync(void* dst, void const* src, uint32_t srcSize = size) +{ + static_assert(size == 16, "only the 16B cp.async variant is vendored"); + // srcSize == 0 turns the copy into a shared-memory zero fill; predicated + // lanes use it so ragged tiles never touch global memory. Nulling src is + // the same workaround as the xqa original, which observed speculative + // global reads without it. + if (srcSize == 0) + { + src = nullptr; + } + asm volatile( + "cp.async.cg.shared.global [%0], [%1], 16, %2;\n" ::"l"(__cvta_generic_to_shared(dst)), "l"(src), "r"(srcSize)); +} + +__device__ __forceinline__ void commitGroup() +{ + asm volatile("cp.async.commit_group;\n"); +} + +// Wait until at most InFlightGroups cp.async groups remain in flight. +template +__device__ __forceinline__ void waitGroup() +{ + asm volatile("cp.async.wait_group %0;\n" ::"n"(InFlightGroups)); +} +} // namespace compact_detail + +// One pipeline stage moves a 32-token tile regardless of the page geometry. +constexpr int32_t kSparseKvCompactFastTokensPerTile = 32; + +//! Launch parameters for the pipelined bf16 fast-path compaction kernels. +struct SparseKvCacheCompactV2Bf16Params +{ + int64_t const* poolPointers; + int32_t const* pageTable; + int32_t const* sourceIndices; + int32_t const* sourceOffsets; + int32_t const* sourceLayerIndices; + int32_t const* destinationBases; + int64_t sourceLayerStride; + int64_t sourceHeadStride; + int64_t pageTableRequestStride; + int32_t numLayers; + int32_t batchSize; + int32_t numKvHeads; + size_t bytesPerKvHalf; + size_t bytesPerPage; +}; + +//! Double-buffered cp.async pipeline: while the current 32-token tile drains +//! from shared memory into its destination pages, the next tile's K/V vectors +//! are already streaming global -> shared into the other buffer. One CTA per +//! (layer, KV head, request); threadIdx.x walks the 16B vectors of one head, +//! threadIdx.y walks the tokens of a tile. +template +__global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) + * kSparseKvCompactFastTokensPerTile) void sparseKvCacheCompactV2Bf16PipelineKernel(SparseKvCacheCompactV2Bf16Params + params) +{ + static_assert(std::is_same_v); + static_assert(HeadDim == 64 || HeadDim == 128); + // 128-token pages are the geometry the kernel was written for; 32-token + // pages cover this tree's production configuration (one tile == one page). + static_assert(TokensPerBlock == 32 || TokensPerBlock == 128); + // 16B vectors per head: Dh64 -> 8 lanes (block 8x32 = 256 threads), + // Dh128 -> 16 lanes (block 16x32 = 512 threads). + constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); + constexpr int32_t kTokensPerTile = kSparseKvCompactFastTokensPerTile; + constexpr int32_t kVectorsPerTile = kTokensPerTile * kVectorsPerHead; + // A buffer holds one K tile plus one V tile; two buffers ping-pong. + constexpr int32_t kVectorsPerBuffer = 2 * kVectorsPerTile; + + int32_t const layerIdx = static_cast(blockIdx.x); + int32_t const kvHeadIdx = static_cast(blockIdx.y); + int32_t const batchIdx = static_cast(blockIdx.z); + int32_t const moveBegin = params.sourceOffsets[batchIdx]; + int32_t const moveEnd = params.sourceOffsets[batchIdx + 1]; + int32_t const moveCount = moveEnd - moveBegin; + if (moveCount <= 0) + { + return; + } + + // Layer resolution matches KvCacheV2LayersBuffer::getSparseKvSourceToken: + // without an explicit map, launch layer i reads source plane i (the flat + // layout passes sourceLayerStride == 0, which collapses the term). + int32_t const sourceLayer = params.sourceLayerIndices == nullptr ? layerIdx : params.sourceLayerIndices[layerIdx]; + // Tip-ABI adaptation (c): head planes are strided by the allocation width + // of the move buffers; this request's range within a plane starts at + // moveBegin. + int64_t const sourceMoveBase = static_cast(sourceLayer) * params.sourceLayerStride + + static_cast(kvHeadIdx) * params.sourceHeadStride + moveBegin; + // Tip-ABI adaptation (b): per-request landing position. + int32_t const destinationBase = params.destinationBases[batchIdx]; + auto* const pool = reinterpret_cast(static_cast(params.poolPointers[layerIdx])); + // Tip-ABI adaptation (a): flat V2 K-plane block-offset table; each lookup + // below decodes an entry to a page with >> 1. TokensPerBlock is a + // compile-time power of two, so / and % lower to shifts and masks. + int32_t const* const pageTable = params.pageTable + static_cast(batchIdx) * params.pageTableRequestStride; + + extern __shared__ uint4 sharedVectors[]; + int32_t const sharedVector + = static_cast(threadIdx.y) * kVectorsPerHead + static_cast(threadIdx.x); + int32_t currentRequestMove = static_cast(threadIdx.y); + bool currentValid = currentRequestMove < moveCount; + int32_t currentSourceToken = currentValid ? params.sourceIndices[sourceMoveBase + currentRequestMove] : -1; + uint4* currentSharedK = sharedVectors; + uint4* currentSharedV = currentSharedK + kVectorsPerTile; + + // Prologue: explicitly wait for tile 0 and synchronize the CTA before any thread stores it. + uint4 const* currentSourceKVector = nullptr; + uint4 const* currentSourceVVector = nullptr; + if (currentValid) + { + int32_t const sourcePage = pageTable[currentSourceToken / TokensPerBlock] >> 1; + auto* const sourcePageBase = pool + static_cast(sourcePage) * params.bytesPerPage; + auto const* const sourceK = reinterpret_cast(sourcePageBase); + auto const* const sourceV = reinterpret_cast(sourcePageBase + params.bytesPerKvHalf); + int32_t const localVector = (kvHeadIdx * TokensPerBlock + currentSourceToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + currentSourceKVector = &sourceK[localVector]; + currentSourceVVector = &sourceV[localVector]; + } + uint32_t const currentSourceBytes = currentValid ? sizeof(uint4) : 0U; + compact_detail::copyAsync(¤tSharedK[sharedVector], currentSourceKVector, currentSourceBytes); + compact_detail::copyAsync(¤tSharedV[sharedVector], currentSourceVVector, currentSourceBytes); + compact_detail::commitGroup(); + compact_detail::waitGroup<0>(); + __syncthreads(); + + for (int32_t nextTileBegin = kTokensPerTile; nextTileBegin < moveCount; nextTileBegin += kTokensPerTile) + { + int32_t const nextRequestMove = nextTileBegin + static_cast(threadIdx.y); + bool const nextValid = nextRequestMove < moveCount; + int32_t const nextSourceToken = nextValid ? params.sourceIndices[sourceMoveBase + nextRequestMove] : -1; + int32_t const nextBuffer = (nextTileBegin / kTokensPerTile) & 1; + uint4* const nextSharedK = sharedVectors + nextBuffer * kVectorsPerBuffer; + uint4* const nextSharedV = nextSharedK + kVectorsPerTile; + + uint4 const* nextSourceKVector = nullptr; + uint4 const* nextSourceVVector = nullptr; + if (nextValid) + { + int32_t const sourcePage = pageTable[nextSourceToken / TokensPerBlock] >> 1; + auto* const sourcePageBase = pool + static_cast(sourcePage) * params.bytesPerPage; + auto const* const sourceK = reinterpret_cast(sourcePageBase); + auto const* const sourceV = reinterpret_cast(sourcePageBase + params.bytesPerKvHalf); + int32_t const localVector + = (kvHeadIdx * TokensPerBlock + nextSourceToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + nextSourceKVector = &sourceK[localVector]; + nextSourceVVector = &sourceV[localVector]; + } + uint32_t const nextSourceBytes = nextValid ? sizeof(uint4) : 0U; + compact_detail::copyAsync(&nextSharedK[sharedVector], nextSourceKVector, nextSourceBytes); + compact_detail::copyAsync(&nextSharedV[sharedVector], nextSourceVVector, nextSourceBytes); + compact_detail::commitGroup(); + + // The compaction contract provides strictly increasing sources per request/head and + // dst(i) = destinationBase + i <= src(i). For current i and future j, i < j implies + // dst(i) < destinationBase + j <= src(j), so current stores cannot alias future prefetch sources. + // The current tile itself completed its wait and CTA barrier before reaching this store phase. + int32_t const destinationToken = destinationBase + currentRequestMove; + if (currentValid && currentSourceToken != destinationToken) + { + int32_t const destinationPage = pageTable[destinationToken / TokensPerBlock] >> 1; + auto* const destinationPageBase = pool + static_cast(destinationPage) * params.bytesPerPage; + auto* const destinationK = reinterpret_cast(destinationPageBase); + auto* const destinationV = reinterpret_cast(destinationPageBase + params.bytesPerKvHalf); + int32_t const localVector + = (kvHeadIdx * TokensPerBlock + destinationToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + destinationK[localVector] = currentSharedK[sharedVector]; + destinationV[localVector] = currentSharedV[sharedVector]; + } + + // commitGroup only closes the group; waitGroup<0> completes this thread's next-tile copies. The CTA + // barrier then makes the ping-pong buffer visible to all threads before it becomes current. + compact_detail::waitGroup<0>(); + __syncthreads(); + currentRequestMove = nextRequestMove; + currentValid = nextValid; + currentSourceToken = nextSourceToken; + currentSharedK = nextSharedK; + currentSharedV = nextSharedV; + } + + // Epilogue: the final tile already completed its async wait and CTA barrier. + int32_t const destinationToken = destinationBase + currentRequestMove; + if (currentValid && currentSourceToken != destinationToken) + { + int32_t const destinationPage = pageTable[destinationToken / TokensPerBlock] >> 1; + auto* const destinationPageBase = pool + static_cast(destinationPage) * params.bytesPerPage; + auto* const destinationK = reinterpret_cast(destinationPageBase); + auto* const destinationV = reinterpret_cast(destinationPageBase + params.bytesPerKvHalf); + int32_t const localVector = (kvHeadIdx * TokensPerBlock + destinationToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + destinationK[localVector] = currentSharedK[sharedVector]; + destinationV[localVector] = currentSharedV[sharedVector]; + } +} + +//! Variant of the pipeline kernel that stages the destination page index in +//! shared memory once per tile instead of having every thread look it up. +//! Only valid when each request's destination base is 32-token tile aligned, +//! so a whole tile lands in one destination page (both supported page sizes +//! are multiples of the tile). Compiled but not yet dispatched: the bases +//! live on device, so the host cannot prove alignment (see the dispatch +//! helper below). +template +__global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) + * kSparseKvCompactFastTokensPerTile) void sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel(SparseKvCacheCompactV2Bf16Params + params) +{ + static_assert(std::is_same_v); + static_assert(HeadDim == 64 || HeadDim == 128); + static_assert(TokensPerBlock == 32 || TokensPerBlock == 128); + constexpr int32_t kBufferCount = 2; + constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); + constexpr int32_t kTokensPerTile = kSparseKvCompactFastTokensPerTile; + constexpr int32_t kVectorsPerTile = kTokensPerTile * kVectorsPerHead; + constexpr int32_t kVectorsPerBuffer = 2 * kVectorsPerTile; + + int32_t const layerIdx = static_cast(blockIdx.x); + int32_t const kvHeadIdx = static_cast(blockIdx.y); + int32_t const batchIdx = static_cast(blockIdx.z); + int32_t const moveBegin = params.sourceOffsets[batchIdx]; + int32_t const moveEnd = params.sourceOffsets[batchIdx + 1]; + int32_t const moveCount = moveEnd - moveBegin; + if (moveCount <= 0) + { + return; + } + + int32_t const sourceLayer = params.sourceLayerIndices == nullptr ? layerIdx : params.sourceLayerIndices[layerIdx]; + int64_t const sourceMoveBase = static_cast(sourceLayer) * params.sourceLayerStride + + static_cast(kvHeadIdx) * params.sourceHeadStride + moveBegin; + int32_t const destinationBase = params.destinationBases[batchIdx]; + auto* const pool = reinterpret_cast(static_cast(params.poolPointers[layerIdx])); + int32_t const* const pageTable = params.pageTable + static_cast(batchIdx) * params.pageTableRequestStride; + + extern __shared__ uint4 sharedVectors[]; + auto* const sharedDestinationPages = reinterpret_cast(sharedVectors + kBufferCount * kVectorsPerBuffer); + int32_t const sharedVector + = static_cast(threadIdx.y) * kVectorsPerHead + static_cast(threadIdx.x); + int32_t currentBuffer = 0; + int32_t currentRequestMove = static_cast(threadIdx.y); + bool currentValid = currentRequestMove < moveCount; + int32_t currentSourceToken = currentValid ? params.sourceIndices[sourceMoveBase + currentRequestMove] : -1; + uint4* currentSharedK = sharedVectors; + uint4* currentSharedV = currentSharedK + kVectorsPerTile; + + // The dispatch guard tile-aligns the destination base, so every 32-token + // tile lies within one destination page. One producer stages the decoded + // page beside buffer 0; the existing prologue barrier publishes both + // products. + if (threadIdx.x == 0 && threadIdx.y == 0) + { + sharedDestinationPages[currentBuffer] = pageTable[destinationBase / TokensPerBlock] >> 1; + } + uint4 const* currentSourceKVector = nullptr; + uint4 const* currentSourceVVector = nullptr; + if (currentValid) + { + int32_t const sourcePage = pageTable[currentSourceToken / TokensPerBlock] >> 1; + auto* const sourcePageBase = pool + static_cast(sourcePage) * params.bytesPerPage; + auto const* const sourceK = reinterpret_cast(sourcePageBase); + auto const* const sourceV = reinterpret_cast(sourcePageBase + params.bytesPerKvHalf); + int32_t const localVector = (kvHeadIdx * TokensPerBlock + currentSourceToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + currentSourceKVector = &sourceK[localVector]; + currentSourceVVector = &sourceV[localVector]; + } + uint32_t const currentSourceBytes = currentValid ? sizeof(uint4) : 0U; + compact_detail::copyAsync(¤tSharedK[sharedVector], currentSourceKVector, currentSourceBytes); + compact_detail::copyAsync(¤tSharedV[sharedVector], currentSourceVVector, currentSourceBytes); + compact_detail::commitGroup(); + compact_detail::waitGroup<0>(); + __syncthreads(); + + for (int32_t nextTileBegin = kTokensPerTile; nextTileBegin < moveCount; nextTileBegin += kTokensPerTile) + { + int32_t const nextRequestMove = nextTileBegin + static_cast(threadIdx.y); + bool const nextValid = nextRequestMove < moveCount; + int32_t const nextSourceToken = nextValid ? params.sourceIndices[sourceMoveBase + nextRequestMove] : -1; + int32_t const nextBuffer = (nextTileBegin / kTokensPerTile) & 1; + uint4* const nextSharedK = sharedVectors + nextBuffer * kVectorsPerBuffer; + uint4* const nextSharedV = nextSharedK + kVectorsPerTile; + + // Produce the page paired with the next ping-pong buffer. The loop's existing final barrier publishes it + // before the buffer becomes current, so destination staging does not add a CTA synchronization. + if (threadIdx.x == 0 && threadIdx.y == 0) + { + int32_t const nextDestinationToken = destinationBase + nextTileBegin; + sharedDestinationPages[nextBuffer] = pageTable[nextDestinationToken / TokensPerBlock] >> 1; + } + uint4 const* nextSourceKVector = nullptr; + uint4 const* nextSourceVVector = nullptr; + if (nextValid) + { + int32_t const sourcePage = pageTable[nextSourceToken / TokensPerBlock] >> 1; + auto* const sourcePageBase = pool + static_cast(sourcePage) * params.bytesPerPage; + auto const* const sourceK = reinterpret_cast(sourcePageBase); + auto const* const sourceV = reinterpret_cast(sourcePageBase + params.bytesPerKvHalf); + int32_t const localVector + = (kvHeadIdx * TokensPerBlock + nextSourceToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + nextSourceKVector = &sourceK[localVector]; + nextSourceVVector = &sourceV[localVector]; + } + uint32_t const nextSourceBytes = nextValid ? sizeof(uint4) : 0U; + compact_detail::copyAsync(&nextSharedK[sharedVector], nextSourceKVector, nextSourceBytes); + compact_detail::copyAsync(&nextSharedV[sharedVector], nextSourceVVector, nextSourceBytes); + compact_detail::commitGroup(); + + // The compaction contract provides strictly increasing sources per request/head and + // dst(i) = destinationBase + i <= src(i). For current i and future j, i < j implies + // dst(i) < destinationBase + j <= src(j), so current stores cannot alias future prefetch sources. + int32_t const destinationToken = destinationBase + currentRequestMove; + if (currentValid && currentSourceToken != destinationToken) + { + int32_t const destinationPage = sharedDestinationPages[currentBuffer]; + auto* const destinationPageBase = pool + static_cast(destinationPage) * params.bytesPerPage; + auto* const destinationK = reinterpret_cast(destinationPageBase); + auto* const destinationV = reinterpret_cast(destinationPageBase + params.bytesPerKvHalf); + int32_t const localVector + = (kvHeadIdx * TokensPerBlock + destinationToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + destinationK[localVector] = currentSharedK[sharedVector]; + destinationV[localVector] = currentSharedV[sharedVector]; + } + + compact_detail::waitGroup<0>(); + __syncthreads(); + currentBuffer = nextBuffer; + currentRequestMove = nextRequestMove; + currentValid = nextValid; + currentSourceToken = nextSourceToken; + currentSharedK = nextSharedK; + currentSharedV = nextSharedV; + } + + // The final tile and its destination page scalar already completed the existing wait and CTA barrier. + int32_t const destinationToken = destinationBase + currentRequestMove; + if (currentValid && currentSourceToken != destinationToken) + { + int32_t const destinationPage = sharedDestinationPages[currentBuffer]; + auto* const destinationPageBase = pool + static_cast(destinationPage) * params.bytesPerPage; + auto* const destinationK = reinterpret_cast(destinationPageBase); + auto* const destinationV = reinterpret_cast(destinationPageBase + params.bytesPerKvHalf); + int32_t const localVector = (kvHeadIdx * TokensPerBlock + destinationToken % TokensPerBlock) * kVectorsPerHead + + static_cast(threadIdx.x); + destinationK[localVector] = currentSharedK[sharedVector]; + destinationV[localVector] = currentSharedV[sharedVector]; + } +} + +template +void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactV2Bf16Params const& params, cudaStream_t stream) +{ + constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); + dim3 const block(kVectorsPerHead, kSparseKvCompactFastTokensPerTile); + dim3 const grid(params.numLayers, params.numKvHeads, params.batchSize); + // Two ping-pong buffers x (K tile + V tile) of 32 tokens x kVectorsPerHead + // 16B vectors: + // Dh64: 4 * 32 * 8 * 16 B = 16 KiB + // Dh128: 4 * 32 * 16 * 16 B = 32 KiB + // Both fit the 48 KiB per-CTA dynamic shared memory default, so no + // cudaFuncSetAttribute opt-in is required. + size_t const sharedBytes = 4 * kSparseKvCompactFastTokensPerTile * kVectorsPerHead * sizeof(uint4); + sparseKvCacheCompactV2Bf16PipelineKernel<<>>(params); +} + +template +void launchSparseKvCacheCompactV2Bf16DestinationPagePipeline( + SparseKvCacheCompactV2Bf16Params const& params, cudaStream_t stream) +{ + constexpr int32_t kDestinationPageBuffers = 2; + constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); + dim3 const block(kVectorsPerHead, kSparseKvCompactFastTokensPerTile); + dim3 const grid(params.numLayers, params.numKvHeads, params.batchSize); + // Same 16/32 KiB tile buffers as the plain pipeline plus two staged + // destination page indices; still far below the 48 KiB default. + size_t const sharedBytes = 4 * kSparseKvCompactFastTokensPerTile * kVectorsPerHead * sizeof(uint4) + + kDestinationPageBuffers * sizeof(int32_t); + sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel + <<>>(params); +} + +template +void dispatchSparseKvCacheCompactV2FastGeometry(SparseKvCacheCompactV2Bf16Params const& params, cudaStream_t stream) +{ + // The destination-page variant requires every request's destination base + // to be 32-token tile aligned so each tile lands in one page, but the + // bases live on device where the host cannot check them. Until a host-side + // alignment flag is plumbed through compaction.py, always launch the plain + // pipeline. A plain `if` (not `if constexpr`) keeps the destination-page + // kernel instantiated and compiling. + constexpr bool kDestinationPageVariantEnabled = false; + if (kDestinationPageVariantEnabled) + { + launchSparseKvCacheCompactV2Bf16DestinationPagePipeline(params, stream); + return; + } + launchSparseKvCacheCompactV2Bf16Pipeline(params, stream); +} + +#endif // ENABLE_BF16 + template __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( QKVPreprocessingParams params) @@ -2021,6 +2463,55 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream) { +#ifdef ENABLE_BF16 + if constexpr (std::is_same_v) + { + // REMOVE-BEFORE-PR: A/B knob. TLLM_SPARSE_KV_COMPACT_FAST=0 forces the + // register-staging path below on the same inputs; unset or any other + // value selects the pipelined fast path. Read on every launch rather + // than cached in a static: compaction runs once per eviction round + // (~10/s), so the getenv is free, and tests can flip the knob within + // one process through os.environ. + char const* const fastKnob = std::getenv("TLLM_SPARSE_KV_COMPACT_FAST"); + bool const fastEnabled = fastKnob == nullptr || std::strcmp(fastKnob, "0") != 0; + if (fastEnabled && (headDim == 64 || headDim == 128) && (tokensPerBlock == 32 || tokensPerBlock == 128)) + { + SparseKvCacheCompactV2Bf16Params fastParams{}; + fastParams.poolPointers = poolPointers; + fastParams.pageTable = pageTable; + fastParams.sourceIndices = sparseKvIndices; + fastParams.sourceOffsets = sparseKvOffsets; + fastParams.sourceLayerIndices = sourceLayerIndices; + fastParams.destinationBases = destinationBases; + fastParams.sourceLayerStride = sourceLayerStride; + fastParams.sourceHeadStride = sourceHeadStride; + fastParams.pageTableRequestStride = pageTableRequestStride; + fastParams.numLayers = numLayers; + fastParams.batchSize = batchSize; + fastParams.numKvHeads = numKvHeads; + fastParams.bytesPerKvHalf = static_cast(numKvHeads) * tokensPerBlock * headDim * sizeof(T); + fastParams.bytesPerPage = 2 * fastParams.bytesPerKvHalf; + if (headDim == 64 && tokensPerBlock == 32) + { + dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + } + else if (headDim == 64 && tokensPerBlock == 128) + { + dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + } + else if (headDim == 128 && tokensPerBlock == 32) + { + dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + } + else + { + dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + } + return; + } + } +#endif // ENABLE_BF16 + KvCacheV2LayersBuffer buffer{}; buffer.poolPointers = poolPointers; buffer.pageTable = pageTable; diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 3075f2bef739..2dbd7a6fad77 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -17,6 +17,22 @@ _NUM_PAGES = _BATCH_SIZE * _MAX_PAGES_PER_SEQUENCE _PAGE_INDEX_DIVISOR = 2 +_FAST_KNOB_ENV = "TLLM_SPARSE_KV_COMPACT_FAST" + + +@pytest.fixture(autouse=True, params=["0", "1"], ids=["existing_kernel", "fast_kernel"]) +def sparse_compact_kernel_knob(request, monkeypatch): + """Run every test in this module under both kernel selections. + + The C++ dispatcher re-reads TLLM_SPARSE_KV_COMPACT_FAST on every launch + (temporary A/B knob), so flipping the process environment per test is + enough; no re-import or subprocess is needed. Cases outside the fast-path + gate (dtype, head_dim, or page size not covered) take the existing kernel + under both values and simply run twice. + """ + monkeypatch.setenv(_FAST_KNOB_ENV, request.param) + return request.param + def _encode_k_block_offsets( page_table: torch.Tensor, page_index_scale: int = _PAGE_INDEX_DIVISOR @@ -95,6 +111,8 @@ def _reference_compact( source_offsets: torch.Tensor, destination_base: "int | list[int]", source_layer_indices: Optional[torch.Tensor] = None, + tokens_per_block: int = _TOKENS_PER_BLOCK, + batch_size: int = _BATCH_SIZE, ) -> list[torch.Tensor]: original = [pool.clone() for pool in pools] expected = [pool.clone() for pool in pools] @@ -110,7 +128,7 @@ def _reference_compact( else: assert source_layer_indices is not None layer_sources = source_indices[int(source_layer_indices[group_layer])] - for request in range(_BATCH_SIZE): + for request in range(batch_size): begin = int(source_offsets[request]) end = int(source_offsets[request + 1]) request_base = ( @@ -118,25 +136,25 @@ def _reference_compact( if isinstance(destination_base, (list, tuple)) else destination_base ) - for head in range(_NUM_KV_HEADS): + for head in range(layer_sources.shape[0]): for request_move, global_move in enumerate(range(begin, end)): source_token = int(layer_sources[head, global_move]) destination_token = request_base + request_move - source_page = int(raw_page_table[request, source_token // _TOKENS_PER_BLOCK]) + source_page = int(raw_page_table[request, source_token // tokens_per_block]) destination_page = int( - raw_page_table[request, destination_token // _TOKENS_PER_BLOCK] + raw_page_table[request, destination_token // tokens_per_block] ) destination_pool[ destination_page, :, head, - destination_token % _TOKENS_PER_BLOCK, + destination_token % tokens_per_block, :, ] = source_pool[ source_page, :, head, - source_token % _TOKENS_PER_BLOCK, + source_token % tokens_per_block, :, ] return expected @@ -147,12 +165,13 @@ def _compact( page_tables: list[torch.Tensor], arguments: _DeviceArguments, destination_base: "int | list[int]", + batch_size: int = _BATCH_SIZE, ) -> None: # The op takes per-request destination bases; scalar test parameters are # broadcast to the batch here. torch.full stays CUDA-graph-capturable. if isinstance(destination_base, int): destination_bases = torch.full( - (_BATCH_SIZE,), destination_base, dtype=torch.int32, device="cuda" + (batch_size,), destination_base, dtype=torch.int32, device="cuda" ) else: destination_bases = torch.tensor(destination_base, dtype=torch.int32, device="cuda") @@ -350,3 +369,227 @@ def test_sparse_kv_cache_compact_layers_cuda_graph_replay(): for actual, reference in zip(pools, expected): assert torch.equal(actual.cpu(), reference) + + +# --- Production-shaped geometry for the pipelined bf16 fast path ---------- +# +# The fast path only dispatches for bf16 pools with head_dim 64/128 and +# 32/128-token pages, so the cases below use their own builder instead of the +# 4-token-page fixtures above. + +_FAST_BATCH_SIZE = 3 +# Per-request move counts: two ragged tiles plus a full one (pipeline steady +# state), an empty request (kernel early-return), and a single ragged tile +# (prologue/epilogue only). +_FAST_MOVE_COUNTS = (71, 0, 29) +# Mixed prompt lengths: none tile- or page-aligned. +_FAST_DESTINATION_BASES = [3, 9, 17] +# The production move-index buffers are allocation-wide: pad the head-plane +# width beyond this round's total move count so a kernel that derives the +# stride on device (instead of honoring the explicit sourceHeadStride) reads +# padding for heads > 0 and fails the byte-compare. +_FAST_SOURCE_PAD = 37 +_FAST_IDENTITY_MOVES = 5 + + +class _FastGeometryCase(NamedTuple): + pools_cpu: list[torch.Tensor] + pools: list[torch.Tensor] + page_tables_cpu: list[torch.Tensor] + page_tables: list[torch.Tensor] + source_indices: torch.Tensor + source_offsets: torch.Tensor + source_layer_indices: Optional[torch.Tensor] + destination_bases: list[int] + tokens_per_block: int + batch_size: int + + +def _fast_sources_row( + base: int, count: int, limit: int, generator: torch.Generator +) -> torch.Tensor: + """Distinct sorted source tokens >= base, so src(i) >= base + i (the op's + in-place-safety contract). The first few moves are identities (src == dst), + which the kernel skips storing.""" + if count == 0: + return torch.empty(0, dtype=torch.int32) + identity = min(_FAST_IDENTITY_MOVES, count) + candidates = torch.arange(base + identity, limit, dtype=torch.int32) + picks = torch.randperm(candidates.numel(), generator=generator)[: count - identity] + tail = candidates[picks].sort().values + return torch.cat((torch.arange(base, base + identity, dtype=torch.int32), tail)) + + +def _make_fast_geometry_case( + head_dim: int, + tokens_per_block: int, + dtype: torch.dtype = torch.bfloat16, + num_layers: int = 2, + per_layer_sources: bool = False, +) -> _FastGeometryCase: + batch_size = _FAST_BATCH_SIZE + pages_per_seq = max(2, 128 // tokens_per_block) + tokens_per_seq = pages_per_seq * tokens_per_block + num_pages = batch_size * pages_per_seq + shape = (num_pages, 2, _NUM_KV_HEADS, tokens_per_block, head_dim) + numel = torch.Size(shape).numel() + pools_cpu = [ + ((torch.arange(numel, dtype=torch.int32) + layer * 37) % 251).reshape(shape).to(dtype) + for layer in range(num_layers) + ] + pools = [pool.cuda() for pool in pools_cpu] + + # Deterministic per-geometry inputs: the cross-kernel test rebuilds the + # identical case for each knob value. + generator = torch.Generator().manual_seed(20260720 + head_dim * 1000 + tokens_per_block) + raw_page_table = ( + torch.randperm(num_pages, generator=generator) + .to(torch.int32) + .reshape(batch_size, pages_per_seq) + .cuda() + ) + page_table = _encode_k_block_offsets(raw_page_table) + page_tables = [page_table] * num_layers + assert page_tables[0].stride(0) == 2 * pages_per_seq + page_tables_cpu = [table.cpu() for table in page_tables] + + offsets = [0] + for count in _FAST_MOVE_COUNTS: + offsets.append(offsets[-1] + count) + source_offsets = torch.tensor(offsets, dtype=torch.int32) + width = offsets[-1] + _FAST_SOURCE_PAD + source_layers = 3 if per_layer_sources else 1 + # Padding is a valid token id so a stride bug corrupts output (caught by + # the byte-compare) instead of faulting. + rows = torch.full((source_layers, _NUM_KV_HEADS, width), tokens_per_seq - 1, dtype=torch.int32) + for layer in range(source_layers): + for head in range(_NUM_KV_HEADS): + cursor = 0 + for request, count in enumerate(_FAST_MOVE_COUNTS): + rows[layer, head, cursor : cursor + count] = _fast_sources_row( + _FAST_DESTINATION_BASES[request], count, tokens_per_seq, generator + ) + cursor += count + if per_layer_sources: + source_indices = rows.contiguous() + source_layer_indices = torch.tensor([2, 0], dtype=torch.int32) + else: + source_indices = rows[0].contiguous() + source_layer_indices = None + + return _FastGeometryCase( + pools_cpu=pools_cpu, + pools=pools, + page_tables_cpu=page_tables_cpu, + page_tables=page_tables, + source_indices=source_indices, + source_offsets=source_offsets, + source_layer_indices=source_layer_indices, + destination_bases=list(_FAST_DESTINATION_BASES), + tokens_per_block=tokens_per_block, + batch_size=batch_size, + ) + + +def _run_fast_geometry_case(case: _FastGeometryCase) -> list[torch.Tensor]: + expected = _reference_compact( + case.pools_cpu, + case.page_tables_cpu, + case.source_indices, + case.source_offsets, + case.destination_bases, + case.source_layer_indices, + tokens_per_block=case.tokens_per_block, + batch_size=case.batch_size, + ) + arguments = _device_arguments( + case.pools, case.source_indices, case.source_offsets, case.source_layer_indices + ) + _compact( + case.pools, + case.page_tables, + arguments, + case.destination_bases, + batch_size=case.batch_size, + ) + torch.cuda.synchronize() + return expected + + +# The full fast-path gate matrix. 128-token pages are the geometry the ported +# kernel was written for; 32-token pages and head_dim 128 are this tree's +# production configuration. +_FAST_GEOMETRY_MATRIX = [(64, 32), (128, 32), (64, 128), (128, 128)] + + +@pytest.mark.parametrize("head_dim,tokens_per_block", _FAST_GEOMETRY_MATRIX) +def test_sparse_kv_cache_compact_layers_fast_geometry(head_dim, tokens_per_block): + # The autouse knob fixture runs this both through the pipelined fast path + # and through the existing register-staging kernel. + case = _make_fast_geometry_case(head_dim, tokens_per_block) + expected = _run_fast_geometry_case(case) + for actual, reference in zip(case.pools, expected): + assert torch.equal(actual.cpu(), reference) + + +def test_sparse_kv_cache_compact_layers_fast_geometry_per_layer_source(): + case = _make_fast_geometry_case(64, 32, per_layer_sources=True) + expected = _run_fast_geometry_case(case) + for actual, reference in zip(case.pools, expected): + assert torch.equal(actual.cpu(), reference) + + +@pytest.mark.parametrize( + "dtype,head_dim,tokens_per_block", + [ + (torch.float16, 64, 32), # dtype outside the bf16-only gate + (torch.bfloat16, 256, 32), # head_dim outside the gate + (torch.bfloat16, 64, 16), # page size outside the gate + ], +) +def test_sparse_kv_cache_compact_layers_fast_gate_fallback(dtype, head_dim, tokens_per_block): + # Near-miss geometries must fall back to the existing kernel and stay + # byte-correct under both knob values. + case = _make_fast_geometry_case(head_dim, tokens_per_block, dtype=dtype) + expected = _run_fast_geometry_case(case) + for actual, reference in zip(case.pools, expected): + assert torch.equal(actual.cpu(), reference) + + +@pytest.mark.parametrize("head_dim,tokens_per_block", _FAST_GEOMETRY_MATRIX) +def test_sparse_kv_cache_compact_layers_fast_path_actually_runs( + head_dim, tokens_per_block, sparse_compact_kernel_knob +): + # Guard against the fast-path gate silently never firing: every + # byte-equality test in this module would still pass if the dispatcher + # fell through to the existing kernel under both knob values. Assert via + # the profiler that the selected kernel is the one that actually ran, + # for each of the four static geometry dispatch branches. + case = _make_fast_geometry_case(head_dim, tokens_per_block) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: + _run_fast_geometry_case(case) + names = [event.name for event in profiler.events()] + fast_fired = any("sparseKvCacheCompactV2Bf16PipelineKernel" in name for name in names) + existing_fired = any("updateSparseKvCacheAfterFmha" in name for name in names) + if sparse_compact_kernel_knob == "1": + assert fast_fired and not existing_fired + else: + assert existing_fired and not fast_fired + + +@pytest.mark.parametrize("head_dim,tokens_per_block", [(64, 32), (128, 128)]) +def test_sparse_kv_cache_compact_layers_fast_matches_existing_kernel( + head_dim, tokens_per_block, monkeypatch +): + # The same inputs through both kernel selections must produce + # byte-identical pools, and both must match the reference. + outputs = {} + expected = None + for knob in ("0", "1"): + monkeypatch.setenv(_FAST_KNOB_ENV, knob) + case = _make_fast_geometry_case(head_dim, tokens_per_block) + expected = _run_fast_geometry_case(case) + outputs[knob] = [pool.cpu() for pool in case.pools] + for existing_pool, fast_pool, reference in zip(outputs["0"], outputs["1"], expected): + assert torch.equal(fast_pool, existing_pool) + assert torch.equal(fast_pool, reference) From af786d3e9aabeb82203c7151714973e4c316d885 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 00:19:47 -0700 Subject: [PATCH 048/178] [None][perf] Make the pipelined compact kernel the default bf16 path Drop the temporary TLLM_SPARSE_KV_COMPACT_FAST knob: the pipelined kernels won the A/B comparison everywhere (nsys, same build, production config: 1.47x at batch 1, 1.09-1.30x at batch 32, 1.09-1.13x at batch 256, byte-identical pools; clean-throughput legs show no regression and up to +3.7% on GPT-OSS), so eligible geometry (bf16, head size 64/128, 32/128-token pages) now dispatches them unconditionally. The register-staging kernel remains only as the fallback for other dtypes and geometries; the pre-replacement implementation is preserved on the backup/compact-pre-fanrong-20260720 branch. Tests assert the dispatch routing on both sides of the gate via profiler probes. Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 19 ++--- .../serial/test_sparse_kv_cache_compact.py | 77 ++++++------------- 2 files changed, 33 insertions(+), 63 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 21a90c3e3280..7482b0fc1882 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -29,8 +29,6 @@ #include "tensorrt_llm/kernels/quantization.cuh" #include "tensorrt_llm/kernels/unfusedAttentionKernels.h" -#include -#include #include using namespace tensorrt_llm::common; @@ -2466,15 +2464,14 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con #ifdef ENABLE_BF16 if constexpr (std::is_same_v) { - // REMOVE-BEFORE-PR: A/B knob. TLLM_SPARSE_KV_COMPACT_FAST=0 forces the - // register-staging path below on the same inputs; unset or any other - // value selects the pipelined fast path. Read on every launch rather - // than cached in a static: compaction runs once per eviction round - // (~10/s), so the getenv is free, and tests can flip the knob within - // one process through os.environ. - char const* const fastKnob = std::getenv("TLLM_SPARSE_KV_COMPACT_FAST"); - bool const fastEnabled = fastKnob == nullptr || std::strcmp(fastKnob, "0") != 0; - if (fastEnabled && (headDim == 64 || headDim == 128) && (tokensPerBlock == 32 || tokensPerBlock == 128)) + // Production path for bf16 pools with head size 64/128 and 32/128-token + // pages: the pipelined kernels won the A/B comparison against the + // register-staging kernel everywhere (verified 2026-07-20: 1.47x at + // batch 1, 1.09-1.30x at batch 32, 1.09-1.13x at batch 256, with + // byte-identical outputs), so they dispatch unconditionally here. The + // register-staging path below remains only as the fallback for other + // dtypes and geometries. + if ((headDim == 64 || headDim == 128) && (tokensPerBlock == 32 || tokensPerBlock == 128)) { SparseKvCacheCompactV2Bf16Params fastParams{}; fastParams.poolPointers = poolPointers; diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 2dbd7a6fad77..e601f1aa034a 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -17,21 +17,13 @@ _NUM_PAGES = _BATCH_SIZE * _MAX_PAGES_PER_SEQUENCE _PAGE_INDEX_DIVISOR = 2 -_FAST_KNOB_ENV = "TLLM_SPARSE_KV_COMPACT_FAST" - - -@pytest.fixture(autouse=True, params=["0", "1"], ids=["existing_kernel", "fast_kernel"]) -def sparse_compact_kernel_knob(request, monkeypatch): - """Run every test in this module under both kernel selections. - - The C++ dispatcher re-reads TLLM_SPARSE_KV_COMPACT_FAST on every launch - (temporary A/B knob), so flipping the process environment per test is - enough; no re-import or subprocess is needed. Cases outside the fast-path - gate (dtype, head_dim, or page size not covered) take the existing kernel - under both values and simply run twice. - """ - monkeypatch.setenv(_FAST_KNOB_ENV, request.param) - return request.param +# Kernel-name substrings for the profiler probes below: the pipelined bf16 +# kernel dispatches unconditionally on eligible geometry, the register-staging +# kernel covers everything else. The byte-compare tests would pass no matter +# which kernel ran, so the probes pin down that the dispatch gate routes each +# case to the intended kernel. +_FAST_KERNEL_NAME = "sparseKvCacheCompactV2Bf16PipelineKernel" +_EXISTING_KERNEL_NAME = "updateSparseKvCacheAfterFmha" def _encode_k_block_offsets( @@ -439,8 +431,7 @@ def _make_fast_geometry_case( ] pools = [pool.cuda() for pool in pools_cpu] - # Deterministic per-geometry inputs: the cross-kernel test rebuilds the - # identical case for each knob value. + # Deterministic per-geometry inputs so any failure reproduces exactly. generator = torch.Generator().manual_seed(20260720 + head_dim * 1000 + tokens_per_block) raw_page_table = ( torch.randperm(num_pages, generator=generator) @@ -524,8 +515,8 @@ def _run_fast_geometry_case(case: _FastGeometryCase) -> list[torch.Tensor]: @pytest.mark.parametrize("head_dim,tokens_per_block", _FAST_GEOMETRY_MATRIX) def test_sparse_kv_cache_compact_layers_fast_geometry(head_dim, tokens_per_block): - # The autouse knob fixture runs this both through the pipelined fast path - # and through the existing register-staging kernel. + # Eligible geometry always dispatches the pipelined kernel; the + # byte-compare against the CPU reference is the correctness net. case = _make_fast_geometry_case(head_dim, tokens_per_block) expected = _run_fast_geometry_case(case) for actual, reference in zip(case.pools, expected): @@ -548,48 +539,30 @@ def test_sparse_kv_cache_compact_layers_fast_geometry_per_layer_source(): ], ) def test_sparse_kv_cache_compact_layers_fast_gate_fallback(dtype, head_dim, tokens_per_block): - # Near-miss geometries must fall back to the existing kernel and stay - # byte-correct under both knob values. + # Near-miss geometries must fall back to the register-staging kernel and + # stay byte-correct. The profiler probe proves the gate actually rejected + # the case: a gate widened by accident would hand the pipelined kernel a + # geometry it was never built for. case = _make_fast_geometry_case(head_dim, tokens_per_block, dtype=dtype) - expected = _run_fast_geometry_case(case) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: + expected = _run_fast_geometry_case(case) + names = [event.name for event in profiler.events()] + assert any(_EXISTING_KERNEL_NAME in name for name in names) + assert not any(_FAST_KERNEL_NAME in name for name in names) for actual, reference in zip(case.pools, expected): assert torch.equal(actual.cpu(), reference) @pytest.mark.parametrize("head_dim,tokens_per_block", _FAST_GEOMETRY_MATRIX) -def test_sparse_kv_cache_compact_layers_fast_path_actually_runs( - head_dim, tokens_per_block, sparse_compact_kernel_knob -): +def test_sparse_kv_cache_compact_layers_fast_path_actually_runs(head_dim, tokens_per_block): # Guard against the fast-path gate silently never firing: every # byte-equality test in this module would still pass if the dispatcher - # fell through to the existing kernel under both knob values. Assert via - # the profiler that the selected kernel is the one that actually ran, - # for each of the four static geometry dispatch branches. + # fell through to the register-staging kernel. Assert via the profiler + # that the pipelined kernel ran (and the fallback did not) for each of + # the four static geometry dispatch branches. case = _make_fast_geometry_case(head_dim, tokens_per_block) with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: _run_fast_geometry_case(case) names = [event.name for event in profiler.events()] - fast_fired = any("sparseKvCacheCompactV2Bf16PipelineKernel" in name for name in names) - existing_fired = any("updateSparseKvCacheAfterFmha" in name for name in names) - if sparse_compact_kernel_knob == "1": - assert fast_fired and not existing_fired - else: - assert existing_fired and not fast_fired - - -@pytest.mark.parametrize("head_dim,tokens_per_block", [(64, 32), (128, 128)]) -def test_sparse_kv_cache_compact_layers_fast_matches_existing_kernel( - head_dim, tokens_per_block, monkeypatch -): - # The same inputs through both kernel selections must produce - # byte-identical pools, and both must match the reference. - outputs = {} - expected = None - for knob in ("0", "1"): - monkeypatch.setenv(_FAST_KNOB_ENV, knob) - case = _make_fast_geometry_case(head_dim, tokens_per_block) - expected = _run_fast_geometry_case(case) - outputs[knob] = [pool.cpu() for pool in case.pools] - for existing_pool, fast_pool, reference in zip(outputs["0"], outputs["1"], expected): - assert torch.equal(fast_pool, existing_pool) - assert torch.equal(fast_pool, reference) + assert any(_FAST_KERNEL_NAME in name for name in names) + assert not any(_EXISTING_KERNEL_NAME in name for name in names) From ae3be2298fff6b4f2f2f76e50fa414b4fecbd48d Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 03:32:48 -0700 Subject: [PATCH 049/178] [None][perf] port TriAttention CuTe score kernel onto the rebuilt score path Port of Fanrong Li's SM100 CuTe-DSL mean-score specialization (original patch 0001-None-perf-integrate-optimized-TriAttention-CuTe-score-kernel, md5 e45e0417f1ec2f7d8514f96cf4c13466, authored against the pre-rebuild lineage at 661cbbb0d4) onto the rebuilt scoring stack, where the Triton scorer no longer exists and scoring runs through the compiled C++ ops. What is ported: - triattention_cute_score.py: the 1248-line kernel and runner, verbatim except for the page-table decode below. - Page-table contract: the kernel now consumes the flattened NATIVE block-offset staging buffer ([pool_slot, request, K/V plane, block] int32) through the group's existing seg_page_off, decoding the K-plane entries as physical_page = entry // 2 at the two producer read sites, exactly mirroring the C++ score op's encoded / kvFactor decode (triAttentionScoreKernels.cu). No staging conversion pass is needed. - _FixedScoreGroup wiring, rewritten against the current class: env-gated prepare_cute_score() (TRTLLM_TRIATTENTION_CUTE_SCORE=1, default off) plus a mean-aggregation dispatch in launch(). The kernel scores the full sequence from token zero into a private head-major scratch; the dispatch then gathers each request's decode window (per-request pinned prompt starts) into the group output, preserving the C++ op's output contract. Global page-aligned score_start plumbing from the original patch is dropped (constant 0 at the runner boundary); per-request token starts are strictly more general. - One prepare_cute_score() call in _FixedScoreStagingBuffers so compilation happens outside CUDA graph capture. - Unit test: the original Torch mean oracle, with the construction and launch call sites adapted to the current _FixedScoreGroup API (native int32 block offsets, output_width, 8-argument launch). Dropped from the original patch: Triton scorer edits (kernel deleted on this lineage), test_triattention_cuda_graph.py hunks (file never existed here), the _dummy_pool_like reflow, and the pipeline-test score_start assert (no score_start plumbing survives). Geometry gap, stated plainly: the kernel's supported contract (SM100 exactly, BF16, 128-token pages, 64-element K rows / 32 frequencies, 8 query heads per KV head) matches none of the current production models (Qwen3: 128-element K rows, 32-token pages, group 4; GPT-OSS: 32-token pages). Today it fires only on the synthetic unit-test geometry; with the env knob unset the integration is a no-op and nothing is imported. Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Signed-off-by: tianruih --- .../triattention/triattention.py | 5 + .../triattention/triattention_cute_score.py | 1270 +++++++++++++++++ .../triattention/triattention_kernels.py | 154 +- .../test_triattention_cute_score.py | 100 ++ 4 files changed, 1528 insertions(+), 1 deletion(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index ee9a4a4bc6ee..d98e3be94056 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -670,6 +670,11 @@ def __init__( offsets, output_width=decode_width, ) + # Compile the optional SM100 CuTe score specialization here, outside + # any CUDA graph capture (compilation allocates and synchronizes). + # Default off: without TRTLLM_TRIATTENTION_CUTE_SCORE=1 this is a + # no-op and scoring stays on the compiled C++ score ops. + self.fused_group.prepare_cute_score(self.mean_cos, self.mean_sin) self.copy_done = torch.cuda.Event() # First record publishes constructor allocations to the V2 copy stream; # later records protect pinned metadata before the next cohort reuses it. diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py new file mode 100644 index 000000000000..9580bf8ca6d4 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -0,0 +1,1270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SM100 CuTe-DSL scorer for the TriAttention mean-score path. + +This is the production specialization of the final workbench kernel. It +uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, and +producer-only page-ID lookahead. The public integration keeps the compiled +C++ score ops as the implementation for every geometry outside the exact +contract validated here. + +Page-table contract: ``page_ids`` is the flattened native block-offset +staging buffer ([pool_slot, request, K/V plane, block] int32) shared with +the C++ score op; K-plane entries encode ``physical_page * kv_factor`` and +are decoded inline (kv_factor == 2), so no per-round conversion pass is +needed. +""" + +from __future__ import annotations + +import threading + +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.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack + +CTA_M = 64 +K = 96 +N = 8 +NUM_FREQS = 32 +THREADS = 128 + +PAGE_TOKENS = 128 +RAW_K_HALF_ELEMENTS = CTA_M * 2 * NUM_FREQS +RAW_K_VECTOR_ELEMENTS = 8 +RAW_K_SPLIT_PHASE_ELEMENTS = CTA_M * NUM_FREQS +RAW_K_SPLIT_TMA_COPY_BYTES = RAW_K_SPLIT_PHASE_ELEMENTS * (cutlass.BFloat16.width // 8) +TMA_DESCRIPTOR_QWORDS = 16 + + +class _TriScoreEpilogue: + """Minimal TMEM-to-global epilogue for the score specialization.""" + + def __init__(self) -> None: + self.acc_dtype = cutlass.Float32 + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + accumulator: cute.Tensor, + output: cute.Tensor, + epilogue_tile: cute.Tile, + use_2cta_instrs: bool, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + copy_atom = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epilogue_tile, + use_2cta_instrs, + ) + accumulator_epilogue = cute.flat_divide( + accumulator[((None, None), 0, 0)], + epilogue_tile, + ) + tiled_copy = tcgen05.make_tmem_copy( + copy_atom, + accumulator_epilogue[(None, None, 0, 0)], + ) + thread_copy = tiled_copy.get_slice(tidx) + thread_accumulator = thread_copy.partition_S(accumulator_epilogue) + output_epilogue = cute.flat_divide( + output[((None, None), 0, 0, None, None, None)], + epilogue_tile, + ) + thread_output = thread_copy.partition_D(output_epilogue) + register_accumulator = cute.make_rmem_tensor( + thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.acc_dtype, + ) + return tiled_copy, thread_accumulator, register_accumulator + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tiled_copy: cute.TiledCopy, + output: cute.Tensor, + epilogue_tile: cute.Tile, + _unused_smem: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + output_epilogue = cute.flat_divide( + output[((None, None), 0, 0, None, None, None)], + epilogue_tile, + ) + thread_copy = tiled_copy.get_slice(tidx) + thread_output = thread_copy.partition_D(output_epilogue) + register_output = cute.make_rmem_tensor( + thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.c_dtype, + ) + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return copy_atom, register_output, thread_output + + +class _TriAttentionScoreKernel(_TriScoreEpilogue): + """Assign one CTA to each segment/KV-head task and retain W across pages.""" + + def __init__( + self, + *, + num_segments: int, + seq_len: int, + score_start: int, + num_q_heads: int, + num_kv_heads: int, + num_freqs: int, + tokens_per_block: int, + pool_shape: tuple[int, int, int, int, int], + pool_strides: tuple[int, int, int, int, int], + pool_dtype: type[cutlass.Numeric], + page_shards: int, + ) -> None: + """Build the single validated production specialization.""" + super().__init__() + if pool_dtype is not cutlass.BFloat16: + raise ValueError("TriAttention CuTe score requires BF16 K pages") + if num_freqs != NUM_FREQS: + raise ValueError("TriAttention CuTe score requires 32 frequencies") + if tokens_per_block != PAGE_TOKENS: + raise ValueError("TriAttention CuTe score requires 128-token pages") + if num_q_heads % num_kv_heads or num_q_heads // num_kv_heads != N: + raise ValueError("TriAttention CuTe score requires GQA group 8") + if score_start % PAGE_TOKENS: + raise ValueError("TriAttention CuTe score requires page-aligned score_start") + if page_shards not in (2, 3): + raise ValueError("TriAttention CuTe score requires two or three page shards") + + self.score_start = score_start + self.num_q_heads = num_q_heads + self.num_kv_heads = num_kv_heads + self.group_size = num_q_heads // num_kv_heads + self.sum_seq = num_segments * seq_len + self.num_tasks = num_segments * num_kv_heads + self.page_shards = page_shards + self.num_ctas = self.num_tasks * page_shards + self.max_pages = (seq_len + PAGE_TOKENS - 1) // PAGE_TOKENS + self.halves_per_page = PAGE_TOKENS // CTA_M + + # Measured final choices that still shape layouts or generated code. + self.prefetch_depth = 4 + self.sqrt_mode = "approx" + self.k_staging_mode = "half_page_tma" + self.use_tma = True + self.cpasync_schedule = "sync_each_half" + self.split_raw_tma = True + self.raw_tma_feature_extent = NUM_FREQS + self.raw_tma_copy_bytes = RAW_K_SPLIT_TMA_COPY_BYTES + self.raw_tma_pipeline_stages = 1 + self.accumulator_pipeline_stages = 1 + self.umma_accumulator_partitions = 1 + self.raw_cpasync_direct_a = True + self.weight_builder_mode = "coefficient_scalar_bf16_two_term" + self.main_operand_mode = "bf16_raw_three_term_weight" + self.numerical_policy = "three_term" + self.magnitude_residual_mode = "fp16_mma_two_term_single_commit" + self.fp16_magnitude_two_term = True + self.magnitude_sqrt_ftz = True + self.producer_page_id_prefetch = True + self.producer_warp_id = 0 + self.physical_threads = THREADS + self.shared_a_raw_alias = False + self.compact_token_loop = True + + self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape + if pool_kv_heads != num_kv_heads or pool_tokens != PAGE_TOKENS or 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_slot, self.s_dim = pool_strides + if self.s_slot != 2 * NUM_FREQS or self.s_dim != 1: + raise ValueError("K pages must be contiguous [128, 64]") + 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 __call__( + self, + page_ids: cute.Tensor, + seg_page_off: cute.Tensor, + seg_req_id: cute.Tensor, + seg_layer_id: cute.Tensor, + seg_seq_len: cute.Tensor, + seg_out_offset: 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, + pool_template: cute.Tensor, + raw_tma_descriptors: cute.Tensor, + stream: cuda.CUstream, + ): + self.c_dtype = output.element_type + self.c_layout = utils.LayoutEnum.COL_MAJOR + self.mma_tiler = (CTA_M, N, K) + self.cta_tile_shape_mnk = self.mma_tiler + self.epi_tile = (CTA_M, N) + + 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], + ) + main_a_shape = ( + (CTA_M, N, NUM_FREQS) + if self.main_operand_mode == "bf16_raw_three_term_weight" + else self.mma_tiler + ) + a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, main_a_shape, cutlass.Float32, 1) + raw_bf16_a_smem_layout = sm100_utils.make_smem_layout_a( + raw_bf16_tiled_mma, + (CTA_M, N, 2 * NUM_FREQS), + cutlass.BFloat16, + 1, + ) + raw_bf16_split_a_smem_layout = sm100_utils.make_smem_layout_a( + raw_bf16_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.BFloat16, + 2, + ) + # The full transport uses one K_SW128 8-KiB tile. The split transport + # packs two K_SW64 4-KiB stages into that same allocation, one each for + # real and imaginary data; the compile-time schedule selects the view. + raw_bf16_direct_a_smem_layout = ( + raw_bf16_split_a_smem_layout if self.split_raw_tma else raw_bf16_a_smem_layout + ) + raw_tma_smem_layout = cute.make_composed_layout( + raw_bf16_direct_a_smem_layout.inner, + 0, + cute.make_layout( + (self.raw_tma_feature_extent, CTA_M), + stride=(1, self.raw_tma_feature_extent), + ), + ) + raw_tma_source_layout = cute.make_layout( + ( + 2 * NUM_FREQS, + PAGE_TOKENS, + (self.num_kv_heads, self.num_physical_pages), + ), + stride=( + self.s_dim, + self.s_slot, + (self.s_kv_head, self.s_page), + ), + ) + raw_tma_source = cute.make_tensor( + pool_template.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, CTA_M), + ) + raw_bf16_b_smem_layout = sm100_utils.make_smem_layout_b( + raw_bf16_tiled_mma, + (CTA_M, N, 2 * NUM_FREQS), + cutlass.BFloat16, + 1, + ) + # The magnitude residual has only K=32. A separate compact descriptor + # lets the producer issue its four UMMA steps before the first commit, + # rather than waiting for and overwriting the K=96 main A tile. + magnitude_lo_smem_layout = sm100_utils.make_smem_layout_a( + tiled_mma, (CTA_M, N, NUM_FREQS), cutlass.Float32, 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_lo_fp16_smem_layout = sm100_utils.make_smem_layout_a( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + magnitude_hi_fp16_smem_layout = sm100_utils.make_smem_layout_b( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + magnitude_fp16_b_smem_layout = sm100_utils.make_smem_layout_b( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + main_b_shape = ( + (CTA_M, N, NUM_FREQS) + if self.main_operand_mode == "bf16_raw_three_term_weight" + else self.mma_tiler + ) + b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, main_b_shape, cutlass.Float32, 1) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # Keep an explicit stage mode even for the one-stage control. The + # singleton mode folds away in codegen and lets both specializations + # share the same producer/consumer slicing protocol. + self.num_accumulator_slots = ( + self.accumulator_pipeline_stages * self.umma_accumulator_partitions + ) + 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) + + b_hi_elements = cute.cosize(b_smem_layout.outer) * int(not self.fp16_magnitude_two_term) + b_lo_elements = cute.cosize(b_smem_layout.outer) * int( + self.numerical_policy == "three_term" and not self.fp16_magnitude_two_term + ) + magnitude_lo_elements = cute.cosize(magnitude_lo_fp16_smem_layout.outer) * int( + self.numerical_policy == "three_term" + and self.magnitude_residual_mode in ("fp16_smem", "fp16_mma_single_commit") + ) + magnitude_lo_fp32_elements = cute.cosize(magnitude_lo_smem_layout.outer) * int( + self.numerical_policy == "three_term" + and self.magnitude_residual_mode == "fp32_smem_single_commit" + ) + magnitude_hi_fp16_elements = cute.cosize(magnitude_hi_fp16_smem_layout.outer) * int( + self.numerical_policy == "three_term" + and self.magnitude_residual_mode == "fp16_mma_single_commit" + ) + a_elements = cute.cosize(a_smem_layout.outer) * int( + not self.shared_a_raw_alias and not self.fp16_magnitude_two_term + ) + raw_k_elements = RAW_K_HALF_ELEMENTS * int( + self.k_staging_mode in ("half_page_cpasync", "half_page_tma") + and not self.shared_a_raw_alias + ) + alias_a_elements = cute.cosize(a_smem_layout.outer) * int(self.shared_a_raw_alias) + alias_raw_k_elements = RAW_K_HALF_ELEMENTS * int(self.shared_a_raw_alias) + raw_bf16_a_elements = cute.cosize(raw_bf16_a_smem_layout.outer) * int( + self.main_operand_mode == "bf16_raw_three_term_weight" + and ( + not self.raw_cpasync_direct_a + or self.cpasync_schedule not in ("sync_each_half", "intra_half_overlap") + ) + ) + raw_bf16_b_elements = cute.cosize(raw_bf16_b_smem_layout.outer) * int( + self.main_operand_mode == "bf16_raw_three_term_weight" + ) + raw_bf16_b2_elements = raw_bf16_b_elements * int( + self.weight_builder_mode != "coefficient_scalar_bf16_two_term" + ) + magnitude_fp16_a_elements = cute.cosize(magnitude_fp16_a_smem_layout.outer) * int( + self.fp16_magnitude_two_term + ) + magnitude_fp16_b_elements = cute.cosize(magnitude_fp16_b_smem_layout.outer) * int( + self.fp16_magnitude_two_term + ) + + @cute.union + class SharedARawAlias: + # The two descriptors are byte-identical in size (8 KiB) and have + # disjoint lifetimes in the alias specialization. + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, alias_a_elements], + 1024, + ] + sRawK: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, alias_raw_k_elements], + 16, + ] + + @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 * int(self.use_tma), + ] + tmem_holding_buf: cutlass.Int32 + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, a_elements], + 1024, + ] + sB_hi: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, b_hi_elements], + 1024, + ] + sB_lo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, b_lo_elements], + 1024, + ] + sMagnitudeLo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_lo_elements], + 1024, + ] + sMagnitudeLoFp32: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, magnitude_lo_fp32_elements], + 1024, + ] + sMagnitudeHiFp16: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_hi_fp16_elements], + 1024, + ] + sRawK: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_k_elements], + 1024, + ] + sRawBf16A: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_bf16_a_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, + ] + sRawBf16B2: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b2_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, + ] + sARawAlias: SharedARawAlias + + self.shared_storage = SharedStorage + self.kernel( + tiled_mma, + raw_bf16_tiled_mma, + magnitude_lo_tiled_mma, + raw_tma_atom, + raw_tma_tensor, + raw_tma_descriptors, + page_ids, + seg_page_off, + seg_req_id, + seg_layer_id, + seg_seq_len, + seg_out_offset, + q_real, + q_imag, + mlr_coef, + mean_cos, + mean_sin, + freq_scale_sq, + output, + a_smem_layout, + raw_bf16_a_smem_layout, + raw_bf16_direct_a_smem_layout, + raw_bf16_split_a_smem_layout, + raw_tma_smem_layout, + raw_bf16_b_smem_layout, + magnitude_lo_smem_layout, + magnitude_lo_fp16_smem_layout, + magnitude_hi_fp16_smem_layout, + magnitude_fp16_a_smem_layout, + magnitude_fp16_b_smem_layout, + b_smem_layout, + ).launch( + grid=(self.num_ctas, 1, 1), + block=(self.physical_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, + page_ids: cute.Tensor, + seg_page_off: cute.Tensor, + seg_req_id: cute.Tensor, + seg_layer_id: cute.Tensor, + seg_seq_len: cute.Tensor, + seg_out_offset: 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, + a_smem_layout: cute.ComposedLayout, + raw_bf16_a_smem_layout: cute.ComposedLayout, + raw_bf16_direct_a_smem_layout: cute.ComposedLayout, + raw_bf16_split_a_smem_layout: cute.ComposedLayout, + raw_tma_smem_layout: cute.ComposedLayout, + raw_bf16_b_smem_layout: cute.ComposedLayout, + magnitude_lo_smem_layout: cute.ComposedLayout, + magnitude_lo_fp16_smem_layout: cute.ComposedLayout, + magnitude_hi_fp16_smem_layout: cute.ComposedLayout, + magnitude_fp16_a_smem_layout: cute.ComposedLayout, + magnitude_fp16_b_smem_layout: cute.ComposedLayout, + 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] + valid_seq_len = seg_seq_len[segment] + page_off = seg_page_off[segment] + out_base = seg_out_offset[segment] + + 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, + ) + 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)] + # Each stage slice retains the K_SW64 pointer flags. Reuse only + # the feature-first outer mapping for the corresponding TMA + # destination so the swizzle is not applied twice. + raw_tma_shared_real = cute.make_tensor( + cpasync_raw_k_real.iterator, + raw_tma_smem_layout.outer, + ) + raw_tma_shared_imag = cute.make_tensor( + cpasync_raw_k_imag.iterator, + raw_tma_smem_layout.outer, + ) + raw_tma_source_tiles = cute.local_tile( + raw_tma_source, + (self.raw_tma_feature_extent, CTA_M), + coord=(None, None, None), + ) + raw_tma_shared_partition_real, raw_tma_global_partition = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_tma_shared_real, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + raw_tma_shared_partition_imag, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_tma_shared_imag, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + 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, + ) + 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, + ) + + page_index = self.score_start // PAGE_TOKENS + page_shard + page_start = page_index * PAGE_TOKENS + pages_processed = cutlass.Int32(0) + producer_prefetched_page_id_lane0 = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + if lane_idx == 0: + # ``page_ids`` is the flattened native block-offset staging + # buffer ([pool_slot, request, K/V plane, block] int32) and + # ``page_off`` points at one request's K plane. K-plane + # entries encode ``physical_page * kv_factor`` (kv_factor is + # 2 for the interleaved K/V pools this kernel requires); the + # C++ score op decodes the same buffer with + # ``encoded / kvFactor`` (triAttentionScoreKernels.cu), so + # divide by two here as well. V-plane entries are never read. + producer_prefetched_page_id_lane0 = ( + cutlass.Int32(page_ids[page_off + page_index]) // 2 + ) + 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 + ) + cute.arch.mbarrier_init_fence() + for weight_round in cutlass.range_constexpr(N * K // THREADS): + linear_index = tidx + weight_round * THREADS + qg = linear_index // K + feature = linear_index % K + coefficient_kind = feature // NUM_FREQS + frequency = feature % NUM_FREQS + mean_offset = req_id * NUM_FREQS + frequency + q_head = kv_head * self.group_size + qg + calib_offset = (layer_id * self.num_q_heads + q_head) * 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]) + 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) + + thr_mma = tiled_mma.get_slice(0) + while page_start < valid_seq_len and pages_processed < self.max_pages: + physical_page = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, + ) + for page_half in cutlass.range_constexpr(self.halves_per_page): + if warp_idx == self.producer_warp_id: + # Phase 0 fills the packed 4-KiB K_SW64 real + # view. Every producer-warp lane participates + # in the PipelineTmaAsync barrier election. + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + page_half, + (kv_head, physical_page), + ) + ], + raw_tma_shared_partition_real, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + 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 warp_idx == self.producer_warp_id: + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + page_half, + (kv_head, physical_page), + ) + ], + raw_tma_shared_partition_imag, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + + if cutlass.const_expr(self.producer_page_id_prefetch and page_half == 1): + next_page_id_lane0 = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + if lane_idx == 0: + next_page_start = page_start + PAGE_TOKENS * self.page_shards + next_pages_processed = pages_processed + 1 + if ( + next_page_start < valid_seq_len + and next_pages_processed < self.max_pages + ): + # Same K-plane decode as the initial + # prefetch: entries are physical_page * 2. + next_page_id_lane0 = ( + cutlass.Int32( + page_ids[page_off + page_index + self.page_shards] + ) + // 2 + ) + 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(NUM_FREQS // 16): + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, 0)], + 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) + frequency = lane_idx + # Issue several independent token loads before consuming + # any of them. This bounded RMEM window is unchanged; the + # optional half-page staging only switches its K source + # from global to the single raw shared buffer. + for token_base in cutlass.range( + 0, + CTA_M // 4, + self.prefetch_depth, + unroll_full=not self.compact_token_loop, + ): + 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 * 4 + staged_real[prefetch_index] = cutlass.Float32( + cpasync_raw_k_0[ + ( + (token, frequency % 16), + 0, + frequency // 16, + 0, + ) + ] + ) + staged_imag[prefetch_index] = cutlass.Float32( + cpasync_raw_k_0[ + ( + (token, frequency % 16), + 0, + frequency // 16, + 1, + ) + ] + ) + + for prefetch_index in cutlass.range_constexpr(self.prefetch_depth): + token_round = token_base + prefetch_index + token = warp_idx + token_round * 4 + real = staged_real[prefetch_index] + imag = staged_imag[prefetch_index] + norm2 = real * real + imag * imag + magnitude = cute.math.sqrt( + norm2, + approx=self.sqrt_mode == "approx", + ftz=self.magnitude_sqrt_ftz, + ) + 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(NUM_FREQS // 16): + imag_b_block = NUM_FREQS // 16 + raw_k_block + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, 1)], + tCrRawBf16B0[(None, None, imag_b_block, 0)], + tCtAcc, + ) + for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, 0)], + tCrRawBf16B1[(None, None, raw_k_block, 0)], + tCtAcc, + ) + for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + imag_b_block = NUM_FREQS // 16 + raw_k_block + cute.gemm( + raw_bf16_tiled_mma, + tCtAcc, + tCrRawBf16ASplit[(None, None, raw_k_block, 1)], + tCrRawBf16B1[(None, None, imag_b_block, 0)], + tCtAcc, + ) + # Keep the full four-product control unchanged. + # The independent omit_a1b1 mode drops only the + # second-order residual product; all other FP16 + # K16 products retain their original order. + magnitude_lo_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + True, + ) + for magnitude_k_block in cutlass.range_constexpr(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(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(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 pages_processed == 0: + if cutlass.const_expr(page_half == 0): + cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) + acc_pipeline.consumer_wait(acc_consumer_state) + output_offset = ( + kv_head * self.group_size * self.sum_seq + + out_base + + page_start + + page_half * CTA_M + ) + page_output = cute.make_tensor( + output.iterator + output_offset, + cute.make_layout( + (CTA_M, N, 1), + stride=( + 1, + self.sum_seq, + self.group_size * self.sum_seq, + ), + ), + ) + gC_mnl = cute.local_tile(page_output, self.epi_tile, (None, None, None)) + tCgC = thr_mma.partition_C(gC_mnl) + tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( + tidx, tCtAcc, tCgC, self.epi_tile, False + ) + simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( + tidx, tiled_copy_t2r, tCgC, self.epi_tile, None + ) + 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)) + for subtile_idx in range(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)) + cute.copy( + simt_atom, + tTR_rC, + tTR_gC[(None, None, None, subtile_idx)], + ) + + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + cute.arch.barrier() + raw_tma_pipeline.consumer_release(raw_tma_consumer_state) + raw_tma_consumer_state.advance() + + page_index += self.page_shards + page_start += PAGE_TOKENS * self.page_shards + pages_processed += 1 + if warp_idx == self.producer_warp_id: + raw_tma_pipeline.producer_tail(raw_tma_producer_state) + if warp_idx == self.producer_warp_id: + acc_pipeline.producer_tail(acc_producer_state) + 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], +) -> torch.Tensor: + """Encode one immutable feature-first TensorMap per layer index.""" + 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, tokens_per_block, head_dim = pool.shape + if (kv_factor, tokens_per_block, head_dim) != (2, PAGE_TOKENS, 2 * NUM_FREQS): + raise ValueError("TriAttention CuTe score requires [page, 2, Hkv, 128, 64] 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, PAGE_TOKENS] + 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) + box_dims = [NUM_FREQS, CTA_M] + [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, + cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B, + 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) + + +class TriAttentionCuteScoreRunner: + """Compile and launch the exact SM100 mean-score specialization.""" + + def __init__( + self, + *, + layer_pools: list[torch.Tensor], + layer_indices: list[int], + max_requests: int, + num_layers: int, + seq_len: int, + score_start: int, + num_q_heads: int, + num_kv_heads: int, + num_freqs: int, + tokens_per_block: int, + page_ids: torch.Tensor, + seg_page_off: torch.Tensor, + seg_req_id: torch.Tensor, + seg_layer_id: torch.Tensor, + seg_seq_len: torch.Tensor, + seg_out_offset: torch.Tensor, + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + freq_scale_sq: torch.Tensor, + output: torch.Tensor, + ) -> None: + self.max_requests = int(max_requests) + self.num_layers = int(num_layers) + self.num_kv_heads = int(num_kv_heads) + self.descriptors = _encode_tma_descriptors(layer_pools, layer_indices) + self._torch_prefix = ( + page_ids, + seg_page_off, + seg_req_id, + seg_layer_id, + seg_seq_len, + seg_out_offset, + q_real, + q_imag, + mlr_coef, + ) + self._torch_tail = ( + freq_scale_sq, + output, + layer_pools[layer_indices[0]], + self.descriptors, + ) + self._cute_prefix = tuple(_to_cute(tensor) for tensor in self._torch_prefix) + self._cute_tail = ( + _to_cute(freq_scale_sq), + _to_cute(output), + _to_cute(layer_pools[layer_indices[0]]), + _to_cute(self.descriptors, assumed_align=128), + ) + self._compiled: dict[int, object] = {} + static_geometry = ( + max_requests * num_layers, + seq_len, + score_start, + num_q_heads, + num_kv_heads, + num_freqs, + tokens_per_block, + tuple(int(value) for value in layer_pools[layer_indices[0]].shape), + tuple(int(value) for value in layer_pools[layer_indices[0]].stride()), + ) + tensor_specs = tuple( + _tensor_spec(tensor) + for tensor in ( + *self._torch_prefix, + mean_cos.view(-1), + mean_sin.view(-1), + *self._torch_tail, + ) + ) + variants = [(1, 3)] + if max_requests > 1: + variants.append((max_requests, 2)) + for request_count, page_shards in variants: + cache_key = ( + "triattention_cute_score", + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + kernel = _TriAttentionScoreKernel( + num_segments=request_count * num_layers, + seq_len=seq_len, + score_start=score_start, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + pool_shape=tuple( + int(value) for value in layer_pools[layer_indices[0]].shape + ), + pool_strides=tuple( + int(value) for value in layer_pools[layer_indices[0]].stride() + ), + pool_dtype=cutlass.BFloat16, + page_shards=page_shards, + ) + stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) + compiled = cute.compile( + kernel, + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled + self._compiled[request_count] = compiled + + def supports(self, request_count: int) -> bool: + """Return whether an exact static specialization was precompiled.""" + return request_count in self._compiled + + def launch( + self, + request_count: int, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + ) -> None: + """Launch the CuTe score kernel on the current PyTorch stream.""" + stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) + self._compiled[request_count]( + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + stream, + ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index f1d87ae795ff..cc6b5e56ab44 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -8,7 +8,10 @@ the compiled ``trtllm`` CUDA ops (coefficient fold + folded paged score) for every supported geometry. The original Triton score kernel has been deleted; the unit tests validate the CUDA ops against an independent PyTorch oracle. -Selection and compaction live in their respective runtime modules. +Selection and compaction live in their respective runtime modules. An +optional SM100 CuTe-DSL specialization (``triattention_cute_score.py``, +default off behind ``TRTLLM_TRIATTENTION_CUTE_SCORE=1``) can take over the +mean-aggregation score launch for one exactly-validated geometry. House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). @@ -19,6 +22,8 @@ from __future__ import annotations +import os +import warnings from typing import List import torch @@ -398,6 +403,115 @@ def __init__( mlr_coef_LHF.view(-1), ) self.pointer_tail = (freq_scale_sq, omega, offsets) + # Optional SM100 CuTe mean-score specialization (see + # triattention_cute_score.py), compiled by the first + # ``prepare_cute_score`` call and default OFF behind the + # TRTLLM_TRIATTENTION_CUTE_SCORE environment knob (read once here). + # The CuTe runner encodes TMA descriptors from the actual pool + # tensors, so pool references are retained ONLY when the knob is on; + # the default path keeps the raw-address-only lifetime contract + # documented above. + self.seq_len = int(seq_len) + self._cute_score_runner = None + self._cute_score_attempted = False + cute_score_enabled = os.environ.get("TRTLLM_TRIATTENTION_CUTE_SCORE", "0") == "1" + self._cute_layer_pools = list(layer_pools) if cute_score_enabled else None + self._cute_layer_indices = [int(layer) for layer in layer_indices] + + def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> None: + """Compile the optional SM100 CuTe mean-score specialization once. + + Call this outside CUDA graph capture: compilation allocates memory + and synchronizes. With the environment knob unset (the default) or on + any unsupported geometry this returns without importing the CuTe + module, and every launch keeps using the compiled C++ score ops. + + Geometry reality check: the supported contract below (SM100 exactly, + BF16 pools, 128-token pages, 64-element K rows / 32 frequencies, + 8 query heads per KV head) matches none of the current production + models (Qwen3 uses 128-element K rows, 32-token pages, and 4 query + heads per KV head; GPT-OSS uses 32-token pages), so today the kernel + fires only on the synthetic unit-test geometry. This wiring exists to + validate the kernel end to end while wider geometry support lands. + """ + if self._cute_score_attempted: + return + self._cute_score_attempted = True + if self._cute_layer_pools is None: + return + anchor = self.pointer_prefix[0] + num_q_heads, num_kv_heads, num_freqs, tokens_per_block, kv_factor = self.geometry_args[:5] + max_segments = self.max_requests * self.num_layers + supported = ( + torch.cuda.get_device_capability(anchor.device) == (10, 0) + and anchor.dtype == torch.bfloat16 + and kv_factor == 2 + and tokens_per_block == 128 + and num_freqs == 32 + and num_q_heads == num_kv_heads * 8 + and int(anchor.stride(-1)) == 1 + # The kernel computes flat score offsets in 32-bit arithmetic. + and num_q_heads * max_segments * self.seq_len < 2**31 + ) + if not supported: + return + device = anchor.device + try: + from .triattention_cute_score import TriAttentionCuteScoreRunner + + # The kernel scores the FULL sequence from physical token zero + # into its own head-major scratch (row = query head, column = + # segment * seq_len + token); ``launch`` gathers each request's + # decode window from that scratch into ``self.output``. All + # buffers below are persistent because the compiled kernel + # captures their device pointers. + scratch = torch.empty( + num_q_heads * max_segments * self.seq_len, + dtype=torch.float32, + device=device, + ) + seg_seq_len = torch.zeros(max_segments, dtype=torch.int32, device=device) + seg_out_offset = ( + torch.arange(max_segments, dtype=torch.int64, device=device) * self.seq_len + ).to(torch.int32) + gather_columns = torch.arange(self.output_width, dtype=torch.int64, device=device) + self._cute_score_runner = TriAttentionCuteScoreRunner( + layer_pools=self._cute_layer_pools, + layer_indices=self._cute_layer_indices, + max_requests=self.max_requests, + num_layers=self.num_layers, + seq_len=self.seq_len, + # Always score from physical token zero: per-request prompt + # windows are applied by the gather in ``launch`` instead of + # one global page-aligned start. + score_start=0, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + page_ids=self.pointer_prefix[2], + seg_page_off=self.pointer_prefix[3], + seg_req_id=self.pointer_prefix[4], + seg_layer_id=self.pointer_prefix[5], + seg_seq_len=seg_seq_len, + seg_out_offset=seg_out_offset, + q_real=self.pointer_middle[0], + q_imag=self.pointer_middle[1], + mlr_coef=self.pointer_middle[2], + mean_cos=mean_cos, + mean_sin=mean_sin, + freq_scale_sq=self.pointer_tail[0], + output=scratch, + ) + self._cute_scratch = scratch + self._cute_seg_seq_len = seg_seq_len + self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) + except (ImportError, RuntimeError, ValueError, AssertionError) as error: + warnings.warn( + f"TriAttention CuTe score setup failed; using the C++ score ops: {error}", + RuntimeWarning, + stacklevel=2, + ) def _fold_coefficient_buffers( self, offset_planes: int @@ -451,6 +565,44 @@ def launch( # unbounded x axis can hold the token tiles of long sequences. raise ValueError("request*layer segment count exceeds the CUDA grid limit") output = self.output[:request_count] + if score_aggregation == "mean": + # Lazy compile covers groups used without their owning workspace + # (unit tests); production compiles in the workspace constructor, + # outside CUDA graph capture. Default off: one attribute check. + self.prepare_cute_score(mean_cos, mean_sin) + runner = self._cute_score_runner + if runner is not None and runner.supports(request_count): + # Stage per-segment valid lengths (segment = request x layer). + torch.index_select( + valid_seq_lens, + 0, + self.pointer_prefix[4][:num_segments], + out=self._cute_seg_seq_len[:num_segments], + ) + runner.launch(request_count, mean_cos, mean_sin) + # The kernel wrote full-sequence scores from physical token + # zero into its head-major scratch. Gather each request's + # decode window (starting at its pinned prompt length) into + # the group output so callers see exactly the layout the C++ + # score ops produce. This costs one extra read+write of the + # score volume per round, only on this opt-in path; columns + # past a request's valid width carry unscored scratch data, + # matching the C++ op, whose consumers mask by valid width. + num_q_heads = int(self.geometry_args[0]) + source = ( + self._cute_scratch[: num_q_heads * num_segments * self.seq_len] + .view(num_q_heads, request_count, self.num_layers, self.seq_len) + .permute(1, 2, 0, 3) + ) + columns = ( + token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1) + + self._cute_gather_columns + ) + columns = columns.clamp_(max=self.seq_len - 1).expand( + request_count, self.num_layers, num_q_heads, self.output_width + ) + torch.gather(source, 3, columns, out=output) + return output _launch_tri_score_perhead( self, request_count, 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..b61a0cbe4c8b --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Correctness coverage for the optional SM100 TriAttention CuTe scorer.""" + +import pytest +import torch + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="TriAttention CuTe score kernel requires SM100", +) +def test_cute_score_matches_torch_mean_oracle(monkeypatch: pytest.MonkeyPatch) -> None: + pytest.importorskip("cutlass") + monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_SCORE", "1") + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + ) + + torch.manual_seed(20260720) + device = torch.device("cuda") + seq_len = 256 + num_q_heads = 8 + num_freqs = 32 + pool = (0.125 * torch.randn(2, 2, 1, 128, 64, 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) + round_starts = torch.tensor([float(seq_len), float(seq_len + 1)], device=device) + phase = (round_starts[:, 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() + + # Native block-offset staging layout ([pool_slot, request, K/V plane, + # block] int32): K-plane entries encode physical_page * kv_factor with + # kv_factor == 2. Both requests read pool pages [0, 1]. + block_offsets = torch.tensor( + [[[[0, 2], [1, 3]], [[0, 2], [1, 3]]]], dtype=torch.int32, device=device + ) + group = _FixedScoreGroup( + [pool], + [0], + 2, + 2, + seq_len, + num_q_heads, + block_offsets, + [0], + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + output_width=seq_len, + ) + keys = pool[:, 0, 0].reshape(seq_len, 2 * num_freqs).float() + k_real = keys[:, :num_freqs] + k_imag = keys[:, num_freqs:] + magnitude = torch.sqrt(k_real.square() + k_imag.square()) + valid_seq_lens = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) + valid_widths = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) + round_starts_device = torch.tensor([seq_len, seq_len + 1], dtype=torch.int32, device=device) + token_starts_device = torch.zeros(2, dtype=torch.int32, device=device) + for request_count in (1, 2): + group.output.fill_(float("nan")) + actual = group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_starts_device, + token_starts_device, + mean_cos, + mean_sin, + "mean", + ) + assert actual.shape == (request_count, 1, num_q_heads, seq_len) + for request in range(request_count): + rotated_real = freq_scale_sq * (k_real * mean_cos[request] + k_imag * mean_sin[request]) + rotated_imag = freq_scale_sq * (k_imag * mean_cos[request] - k_real * mean_sin[request]) + expected = ( + q_real[0, :, None] * rotated_real[None] + + q_imag[0, :, None] * rotated_imag[None] + + mlr_coef[0, :, None] * freq_scale_sq[None, None] * magnitude[None] + ).sum(dim=-1) + torch.testing.assert_close( + actual[request, 0], + expected, + rtol=5.0e-3, + atol=5.0e-3, + ) + + torch.cuda.synchronize() + # Fails loudly if setup silently fell back to the C++ score ops (whose + # scores would also match the oracle here). + assert group._cute_score_runner is not None From ae8ed9526da40b88d3db435852e0fd2f9172d6d5 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 05:44:06 -0700 Subject: [PATCH 050/178] [None][perf] tabulate TriAttention mean-score phases at init, rotate per round Replace the mean-aggregation half of the per-round TriAttention score coefficient preparation (the Triton offset-mean phase kernel plus the mean branch of the CUDA fold kernel) with one gather-and-rotate CUDA kernel over RoPE-style phase tables built once at workspace construction. Identity: the mean fold computes c_re = fss_f * kv_scale_l * (q_re * mean_cos(rs) - q_im * mean_sin(rs)) c_im = fss_f * kv_scale_l * (q_im * mean_cos(rs) + q_re * mean_sin(rs)) with mean_cos/sin(rs) = (1/O) * sum_o cos/sin((rs + offset_o) * omega_f). Both scale factors distribute over the complex rotation, so tabulating C_cos[pos, f] = fss_f * (1/O) * sum_o cos((pos + offset_o) * omega_f) (and sin likewise) for every position in [0, max_position) at initialization (float64 accumulation, stored fp32) and pre-scaling the calibration query by kv_scale (identity for float pools) reduces every eviction round to one table-row gather per request plus four multiply-adds per element: zero trigonometry at runtime. max_position derives from the manager's known max sequence length at init, and the op TORCH_CHECKs every round start against it before gathering. The MLR coefficient has no position term, so its fold (kv_scale * freq_scale_sq * mlr) is fully static: the request axis is removed, nothing about it is recomputed per round, and the score kernels now read c_mlr through a request-independent [layer, head, freq] index (their only change; c_re/c_im reads are untouched). The "max" aggregation keeps the existing trigonometric fold kernel and launch path unchanged. The per-round Triton phase kernel now runs only when the CuTe score environment knob compiled a runner (it refreshes the mean_cos/mean_sin buffers whose device pointers the compiled CuTe kernel captured); the default production path launches zero phase kernels per round. Adds one focused unit test comparing the table+rotation coefficients against a direct float64 trigonometric fold at round starts {0, mid, max_position - 1} (allclose 1e-5). Design: Fanrong Li (torch-graph review, 2026-07-20). Signed-off-by: tianruih --- .../triAttentionScoreKernels.cu | 90 +++++++-- .../triAttentionScoreKernels.h | 38 +++- cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp | 81 +++++++- .../triattention/triattention.py | 7 +- .../triattention/triattention_kernels.py | 179 +++++++++++++----- .../test_triattention_phase_rotation.py | 106 +++++++++++ 6 files changed, 433 insertions(+), 68 deletions(-) create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu index a2a4120bc23c..16c0f4869db0 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu @@ -166,12 +166,14 @@ __device__ __forceinline__ void scoreLoadChunk(char const* row, bool valid, int // Accumulate one 8-frequency chunk into this thread's per-head accumulators. // coff0 = flat index of (request, layer, first head of this block, chunk -// frequency 0) in the coefficient tables; all coefficient reads are -// lane-uniform 16-byte loads. |K| is computed once per (token, frequency) -// BEFORE the head loop so the GROUP heads share it from registers. +// frequency 0) in the c_re/c_im tables; mlrOff0 = the matching index in the +// static request-independent [layer, head, freq] c_mlr table. All coefficient +// reads are lane-uniform 16-byte loads. |K| is computed once per +// (token, frequency) BEFORE the head loop so the GROUP heads share it from +// registers. template -__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, int64_t coff0, int64_t planeStride, - uint4 re4, uint4 im4, float* accMean, float* accMlr, float* accPos) +__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, int64_t coff0, int64_t mlrOff0, + int64_t planeStride, uint4 re4, uint4 im4, float* accMean, float* accMlr, float* accPos) { float kRe[8], kIm[8], kMag[8]; unpackChunk8(re4, kRe); @@ -185,7 +187,7 @@ __device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, in for (int hg = 0; hg < GROUP; ++hg) { int64_t const coff = coff0 + static_cast(hg) * a.numFreqs; - float const* cmp = a.cMlr + coff; + float const* cmp = a.cMlr + mlrOff0 + static_cast(hg) * a.numFreqs; float4 const cm0 = __ldg(reinterpret_cast(cmp)); float4 const cm1 = __ldg(reinterpret_cast(cmp + 4)); float const cml[8] = {cm0.x, cm0.y, cm0.z, cm0.w, cm1.x, cm1.y, cm1.z, cm1.w}; @@ -318,6 +320,9 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne int64_t const coff0 = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + static_cast(headBase) * a.numFreqs; + // The MLR coefficient is position independent, so its table is folded once + // at initialization without a request axis: [layer, head, freq]. + int64_t const mlrOff0 = (static_cast(layerId) * a.numQueryHeads + headBase) * a.numFreqs; int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads * static_cast(a.numFreqs); @@ -350,7 +355,8 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne { uint4 re4, im4; scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); - scoreComputeChunk(a, coff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + scoreComputeChunk( + a, coff0 + c * 8, mlrOff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); } } else @@ -360,7 +366,8 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne { uint4 re4, im4; scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); - scoreComputeChunk(a, coff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + scoreComputeChunk( + a, coff0 + c * 8, mlrOff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); } } @@ -445,6 +452,8 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde int64_t const coff = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + static_cast(h) * a.numFreqs; + // Static request-independent [layer, head, freq] MLR table index. + int64_t const mlrOff = (static_cast(layerId) * a.numQueryHeads + h) * a.numFreqs; float acc = 0.0f; float accMlr = 0.0f; float accPos[kMaxScoreOffsets]; @@ -465,7 +474,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde float const kMag = triSqrtApprox(kRe * kRe + kIm * kIm); if constexpr (USE_MAX) { - accMlr = fmaf(kMag, a.cMlr[coff + f], accMlr); + accMlr = fmaf(kMag, a.cMlr[mlrOff + f], accMlr); #pragma unroll for (int o = 0; o < kMaxScoreOffsets; ++o) { @@ -478,7 +487,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde } else { - acc = fmaf(kRe, a.cRe[coff + f], fmaf(kIm, a.cIm[coff + f], fmaf(kMag, a.cMlr[coff + f], acc))); + acc = fmaf(kRe, a.cRe[coff + f], fmaf(kIm, a.cIm[coff + f], fmaf(kMag, a.cMlr[mlrOff + f], acc))); } } if (store) @@ -508,7 +517,10 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde // Per-round coefficient fold: one thread per (request, layer, head, freq) // element. On the max path each thread additionally writes one c_re/c_im -// value per offset plane (planes are `total` elements apart). +// value per offset plane (planes are `total` elements apart). The production +// mean path no longer runs this kernel (see the rotation kernel below); the +// c_mlr rows written here are request-identical, and the score kernels read +// only the leading [layer, head, freq] block of that buffer. __global__ void triFoldScoreCoefficientsKernel(float const* __restrict__ qReal, float const* __restrict__ qImag, float const* __restrict__ mlrCoef, float const* __restrict__ freqScaleSq, float const* __restrict__ meanCos, float const* __restrict__ meanSin, float const* __restrict__ omega, float const* __restrict__ offsets, @@ -566,6 +578,50 @@ __global__ void triFoldScoreCoefficientsKernel(float const* __restrict__ qReal, } } +// Per-round mean-path replacement for the fold above: all trigonometry is +// tabulated once at initialization (RoPE-style position tables), so the round +// reduces to one table-row gather per request plus four multiplies and two +// adds per element. With +// phaseCos[pos, f] = freq_scale_sq[f] * (1/O) * sum_o cos((pos + offset_o) * omega_f) +// phaseSin[pos, f] = freq_scale_sq[f] * (1/O) * sum_o sin((pos + offset_o) * omega_f) +// qRealScaled / qImagScaled = kv_scale_l * q (identity for float pools) +// the rotation +// c_re = qRealScaled * phaseCos[rs] - qImagScaled * phaseSin[rs] +// c_im = qImagScaled * phaseCos[rs] + qRealScaled * phaseSin[rs] +// equals the mean fold's freq_scale_sq * kv_scale_l * (q rotated by the +// offset-mean phase at round start rs), because both scale factors distribute +// over the complex product. The position-independent c_mlr fold is fully +// static (folded once at initialization, request axis removed), so this +// kernel never writes it. Grid mirrors the fold kernel: one thread per +// (request, layer, head, freq) element. The host wrapper guarantees every +// round start indexes inside the tables. +// Design: Fanrong Li (torch-graph review, 2026-07-20). +__global__ void triRotateMeanScoreCoefficientsKernel(float const* __restrict__ qRealScaled, + float const* __restrict__ qImagScaled, float const* __restrict__ phaseCos, float const* __restrict__ phaseSin, + int32_t const* __restrict__ roundStarts, float* __restrict__ cRe, float* __restrict__ cIm, + int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, int64_t total) +{ + int64_t const idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) + { + return; + } + auto const f = static_cast(idx % numFreqs); + int64_t const rest = idx / numFreqs; + int64_t const calibrationRows = static_cast(numCalibratedLayers) * numQueryHeads; + auto const req = static_cast(rest / calibrationRows); + // Calibration tables are per (layer, head, freq); the request-major flat + // index reduces onto them modulo the calibration extent. + int64_t const cIdx = (rest % calibrationRows) * numFreqs + f; + int64_t const phaseIdx = static_cast(roundStarts[req]) * numFreqs + f; + float const pc = phaseCos[phaseIdx]; + float const ps = phaseSin[phaseIdx]; + float const qre = qRealScaled[cIdx]; + float const qim = qImagScaled[cIdx]; + cRe[idx] = qre * pc - qim * ps; + cIm[idx] = qim * pc + qre * ps; +} + template void launchVectorized(FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useMax, cudaStream_t stream) { @@ -663,6 +719,18 @@ void foldScoreCoefficientsLaunch(float const* qReal, float const* qImag, float c TLLM_CUDA_CHECK(cudaGetLastError()); } +void rotateMeanScoreCoefficientsLaunch(float const* qRealScaled, float const* qImagScaled, float const* phaseCos, + float const* phaseSin, int32_t const* roundStarts, float* cRe, float* cIm, int32_t numRequests, + int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, cudaStream_t stream) +{ + int64_t const total = static_cast(numRequests) * numCalibratedLayers * numQueryHeads * numFreqs; + int32_t const threads = 256; + auto const blocks = static_cast((total + threads - 1) / threads); + triRotateMeanScoreCoefficientsKernel<<>>(qRealScaled, qImagScaled, phaseCos, phaseSin, + roundStarts, cRe, cIm, numCalibratedLayers, numQueryHeads, numFreqs, total); + TLLM_CUDA_CHECK(cudaGetLastError()); +} + void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, int32_t numSegments, bool useVectorized, bool useMax, cudaStream_t stream) { diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h index 2219b8388e02..bcf96b0d0ede 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h @@ -79,6 +79,13 @@ inline constexpr int32_t kScoreBlockThreads = 128; // plane included) lets the score kernel consume raw quantized elements. The // |K| term relies on |scale * K_q| == scale * |K_q|, which only holds for // scale > 0 — the host wrapper validates positivity before launch. +// +// The c_mlr rows this kernel writes are request-identical (the request axis +// only enters through the phase terms), and the score kernels consume c_mlr +// through a request-independent [L_cal, HQ, F] index, i.e. only the leading +// calibration block of the buffer. The production mean path skips this fold +// entirely (see rotateMeanScoreCoefficientsLaunch below); this launch remains +// the "max" aggregation path. void foldScoreCoefficientsLaunch(float const* qReal, // [L_cal * HQ * F] float const* qImag, // [L_cal * HQ * F] float const* mlrCoef, // [L_cal * HQ * F] @@ -95,6 +102,33 @@ void foldScoreCoefficientsLaunch(float const* qReal, // [L_cal * HQ * F] int32_t numRequests, int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, int32_t numOffsets, bool useMax, cudaStream_t stream); +// Mean-aggregation replacement for the per-round fold above: rotate the +// pre-scaled calibration query by tabulated phases instead of computing any +// trigonometry per round. phaseCos/phaseSin hold, for every possible round +// start position (built once at initialization, float64 accumulation), +// phaseCos[pos, f] = freq_scale_sq[f] * (1/O) * sum_o cos((pos + offset_o) * omega_f) +// phaseSin[pos, f] = freq_scale_sq[f] * (1/O) * sum_o sin((pos + offset_o) * omega_f) +// and qRealScaled/qImagScaled carry the per-layer KV dequantization scale +// (identity for float pools). Each thread gathers its request's table row and +// writes +// c_re = q_re_s * phaseCos[rs] - q_im_s * phaseSin[rs] +// c_im = q_im_s * phaseCos[rs] + q_re_s * phaseSin[rs] +// into the same [numRequests, L_cal, HQ, F] planes the fold produces, because +// freq_scale_sq and kv_scale distribute over the complex rotation. c_mlr has +// no position term, so on this path it is folded once at initialization into +// a static [L_cal, HQ, F] table and never rewritten per round. The "max" +// aggregation keeps foldScoreCoefficientsLaunch unchanged. Every round start +// must lie in [0, maxPosition); the host wrapper enforces that loudly before +// launch. Design: Fanrong Li (torch-graph review, 2026-07-20). +void rotateMeanScoreCoefficientsLaunch(float const* qRealScaled, // [L_cal * HQ * F] + float const* qImagScaled, // [L_cal * HQ * F] + float const* phaseCos, // [maxPosition * F] + float const* phaseSin, // [maxPosition * F] + int32_t const* roundStarts, // [numRequests], each in [0, maxPosition) + float* cRe, // [numRequests, L_cal, HQ, F] + float* cIm, // [numRequests, L_cal, HQ, F] + int32_t numRequests, int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, cudaStream_t stream); + // Everything one folded-score launch needs. One "segment" is one // (request, scored layer) pair; segments are request-major so // seg % numLayers == 0 identifies each request's first segment. @@ -108,9 +142,9 @@ struct FoldedScoreParams int32_t const* requestSeqLens; // [numRequests] int32_t* validWidthOut; // [numRequests] side-store: seqLen - tokenStart, once per request int32_t const* requestTokenStarts; // [numRequests] pinned prompt length = decode-region origin - float const* cRe; // fold output (see foldScoreCoefficientsLaunch) + float const* cRe; // per-round output (see foldScoreCoefficientsLaunch / rotateMeanScoreCoefficientsLaunch) float const* cIm; - float const* cMlr; + float const* cMlr; // static, request-independent [L_cal, HQ, F] MLR table float* out; // [segment, numQueryHeads, outputWidth] fp32 decode-only scores int32_t outputWidth; int32_t numLayers; // scored layers per request (the segment period) diff --git a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp index 401c815f289c..c34f3600af1d 100644 --- a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp +++ b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp @@ -21,11 +21,11 @@ namespace tk = tensorrt_llm::kernels::tri_attention_score; -// Two ops rather than one fused op, matching the file-level granularity of -// sibling kernel wrappers (one op per kernel launch): the coefficient fold -// runs once per eviction round into persistent buffers whose plane count -// depends on the aggregation mode, while the score launch consumes those -// buffers; callers time and re-plan them independently. +// One op per kernel launch, matching the file-level granularity of sibling +// kernel wrappers: the per-round coefficient preparation (the rotation op on +// the mean path, the trigonometric fold on the max path) writes persistent +// buffers whose plane count depends on the aggregation mode, while the score +// launch consumes those buffers; callers time and re-plan them independently. namespace { @@ -69,7 +69,10 @@ float const* checkKvScales( // writes one c_re/c_im plane per offset. kv_scales (quantized pools only) // folds the per-layer dequantization scale into every coefficient table; the // paired score op then reads raw quantized elements. This op cannot see the -// pool dtype, so presence-iff-quantized is enforced by the score op. +// pool dtype, so presence-iff-quantized is enforced by the score op. The +// production mean path prepares its coefficients through the rotation op +// below instead; this fold remains the max-aggregation path (and keeps its +// mean branch for that kernel's documented contract). void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, torch::Tensor c_mlr, torch::Tensor q_real, torch::Tensor q_imag, torch::Tensor mlr_coef, torch::Tensor freq_scale_sq, std::optional mean_cos, std::optional mean_sin, std::optional omega, @@ -139,6 +142,57 @@ void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, static_cast(num_freqs), static_cast(num_offsets), use_max, stream); } +// Per-round mean-path coefficient rotation: gather each request's row of the +// initialization-time phase tables (freq_scale_sq, the offset mean, and — via +// the pre-scaled calibration query — the per-layer KV dequantization scale +// are already baked in) and rotate the query by it, writing the same +// c_re/c_im planes the mean fold produced, with zero trigonometry per round. +// c_mlr has no position term, so on this path it is folded once at +// initialization and this op never touches it. Round starts index the phase +// tables directly, so a start at or past max_position (the tabulated position +// count) must fail loudly here rather than gather out of range; like the +// kv_scales positivity check above, the bounds reduction costs one small +// host sync per eviction round. Design: Fanrong Li (torch-graph review, +// 2026-07-20). +void triAttentionRotateMeanScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, torch::Tensor q_real_scaled, + torch::Tensor q_imag_scaled, torch::Tensor phase_cos, torch::Tensor phase_sin, torch::Tensor round_starts, + int64_t num_requests, int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_freqs, + int64_t max_position) +{ + checkContiguousCuda(c_re, at::kFloat, "fp32", "c_re"); + checkContiguousCuda(c_im, at::kFloat, "fp32", "c_im"); + checkContiguousCuda(q_real_scaled, at::kFloat, "fp32", "q_real_scaled"); + checkContiguousCuda(q_imag_scaled, at::kFloat, "fp32", "q_imag_scaled"); + checkContiguousCuda(phase_cos, at::kFloat, "fp32", "phase_cos"); + checkContiguousCuda(phase_sin, at::kFloat, "fp32", "phase_sin"); + checkContiguousCuda(round_starts, at::kInt, "int32", "round_starts"); + TORCH_CHECK( + num_requests > 0 && num_calibrated_layers > 0 && num_query_heads > 0 && num_freqs > 0 && max_position > 0, + "tri_attention_rotate_mean_score_coefficients: rotation extents must be positive"); + int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; + int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; + TORCH_CHECK(c_re.numel() >= total && c_im.numel() >= total, + "tri_attention_rotate_mean_score_coefficients: coefficient buffers are undersized"); + TORCH_CHECK(q_real_scaled.numel() >= calibration && q_imag_scaled.numel() >= calibration, + "tri_attention_rotate_mean_score_coefficients: scaled calibration tensors are undersized"); + TORCH_CHECK(phase_cos.numel() >= max_position * num_freqs && phase_sin.numel() >= max_position * num_freqs, + "tri_attention_rotate_mean_score_coefficients: phase tables do not cover max_position rows"); + TORCH_CHECK(round_starts.numel() >= num_requests, + "tri_attention_rotate_mean_score_coefficients: round_starts are undersized"); + auto const [minStart, maxStart] = round_starts.narrow(0, 0, num_requests).aminmax(); + TORCH_CHECK(minStart.item() >= 0 && maxStart.item() < max_position, + "tri_attention_rotate_mean_score_coefficients: a round start lies outside the tabulated position range " + "[0, ", + max_position, ")"); + + auto stream = at::cuda::getCurrentCUDAStream(); + tk::rotateMeanScoreCoefficientsLaunch(q_real_scaled.data_ptr(), q_imag_scaled.data_ptr(), + phase_cos.data_ptr(), phase_sin.data_ptr(), round_starts.data_ptr(), + c_re.data_ptr(), c_im.data_ptr(), static_cast(num_requests), + static_cast(num_calibrated_layers), static_cast(num_query_heads), + static_cast(num_freqs), stream); +} + // Score every cached decode token of every (request, layer) segment against // the folded coefficient tables, writing fp32 [segment, head, token] rows and // each request's decode width. pool_anchor is one of the scored layer pools: @@ -190,7 +244,11 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas && request_token_starts.numel() >= num_requests, "tri_attention_paged_score: per-request metadata is undersized"); int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; - TORCH_CHECK(c_re.numel() >= num_offsets * total && c_im.numel() >= num_offsets * total && c_mlr.numel() >= total, + // c_mlr is request independent: the score kernels index it as one static + // [layer, head, freq] table, so only the calibration extent is required. + int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; + TORCH_CHECK( + c_re.numel() >= num_offsets * total && c_im.numel() >= num_offsets * total && c_mlr.numel() >= calibration, "tri_attention_paged_score: folded coefficient buffers are undersized"); TORCH_CHECK(out.numel() >= num_segments * num_query_heads * output_width, "tri_attention_paged_score: score output buffer is undersized"); @@ -294,6 +352,14 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "int num_query_heads, int num_freqs, " "int num_offsets, bool use_max, Tensor? kv_scales=None) -> ()"); + m.def( + "tri_attention_rotate_mean_score_coefficients(" + "Tensor(a!) c_re, Tensor(b!) c_im, " + "Tensor q_real_scaled, Tensor q_imag_scaled, " + "Tensor phase_cos, Tensor phase_sin, Tensor round_starts, " + "int num_requests, int num_calibrated_layers, " + "int num_query_heads, int num_freqs, int max_position) -> ()"); + m.def( "tri_attention_paged_score(" "Tensor pool_anchor, Tensor layer_base_addrs, Tensor block_offsets, " @@ -310,5 +376,6 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { m.impl("tri_attention_fold_score_coefficients", &triAttentionFoldScoreCoefficientsOp); + m.impl("tri_attention_rotate_mean_score_coefficients", &triAttentionRotateMeanScoreCoefficientsOp); m.impl("tri_attention_paged_score", &triAttentionPagedScoreOp); } diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d98e3be94056..3fa8a2e8932e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -712,7 +712,12 @@ def launch_prepared_score(self) -> torch.Tensor: raise _FixedScoreStreamMismatch( "TriAttention score launches must stay on the staging CUDA stream" ) - if self._score_aggregation == "mean": + if self._score_aggregation == "mean" and self.fused_group._cute_score_runner is not None: + # mean_cos/mean_sin feed ONLY the opt-in CuTe score runner, whose + # compiled kernel captured their device pointers, so they must be + # refreshed before it launches. The default C++ mean path rotates + # init-time phase tables inside its coefficient op instead, so + # production rounds launch zero phase kernels. prepare_mean_phase( self.round_starts_device, self.offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index cc6b5e56ab44..219d45867eb8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -5,8 +5,10 @@ The production path uses one fixed-shape trig-score launch across all dense layers, CuTE-DSL TopK selection, and grouped C++ compaction. This module owns the score launcher and its persistent metadata; scoring itself runs through -the compiled ``trtllm`` CUDA ops (coefficient fold + folded paged score) for -every supported geometry. The original Triton score kernel has been deleted; +the compiled ``trtllm`` CUDA ops (per-round coefficient rotation over +init-time phase tables on the mean path, trigonometric coefficient fold on +the max path, then the folded paged score) for every supported geometry. The +original Triton score kernel has been deleted; the unit tests validate the CUDA ops against an independent PyTorch oracle. Selection and compaction live in their respective runtime modules. An optional SM100 CuTe-DSL specialization (``triattention_cute_score.py``, @@ -114,12 +116,10 @@ def _launch_tri_score_perhead( valid_widths: torch.Tensor, round_starts_device: torch.Tensor, token_starts_device: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, *, score_aggregation: str, ) -> None: - """Fold the per-round coefficients, then score paged KV via the C++ ops. + """Prepare the per-round coefficients, then score paged KV via the C++ ops. The compiled ``trtllm`` score ops are THE implementation for every geometry this launcher accepts; unsupported inputs fail loudly inside the @@ -130,6 +130,7 @@ def _launch_tri_score_perhead( raise ValueError(f"unsupported score aggregation: {score_aggregation}") if not ( hasattr(torch.ops.trtllm, "tri_attention_fold_score_coefficients") + and hasattr(torch.ops.trtllm, "tri_attention_rotate_mean_score_coefficients") and hasattr(torch.ops.trtllm, "tri_attention_paged_score") ): raise RuntimeError( @@ -150,37 +151,59 @@ def _launch_tri_score_perhead( s_slot, s_dim, ) = group.geometry_args - # Mean aggregation collapses all offsets into mean_cos/mean_sin (one - # coefficient plane); max keeps one c_re/c_im plane per offset because - # max does not commute through the frequency sum. + # Mean aggregation collapses all offsets into one coefficient plane; max + # keeps one c_re/c_im plane per offset because max does not commute + # through the frequency sum. offset_planes = num_offsets if use_max else 1 - c_re, c_im, c_mlr = group._fold_coefficient_buffers(offset_planes) - q_real, q_imag, mlr_coef = group.pointer_middle - freq_scale_sq, omega, offsets = group.pointer_tail - torch.ops.trtllm.tri_attention_fold_score_coefficients( - c_re, - c_im, - c_mlr, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - None if use_max else mean_cos.view(-1), - None if use_max else mean_sin.view(-1), - omega if use_max else None, - offsets if use_max else None, - round_starts_device if use_max else None, - request_count, - group._num_calibrated_layers, - num_q_heads, - num_freqs, - offset_planes, - use_max, - # Per-layer dequant scales (quantized pools only, else None): the fold - # multiplies them into the coefficient tables so the score op below - # reads raw quantized elements at zero hot-loop cost. - group._kv_scales, - ) + c_re, c_im, c_mlr = group._fold_coefficient_buffers(offset_planes, with_mlr=use_max) + if use_max: + q_real, q_imag, mlr_coef = group.pointer_middle + freq_scale_sq, omega, offsets = group.pointer_tail + torch.ops.trtllm.tri_attention_fold_score_coefficients( + c_re, + c_im, + c_mlr, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + None, + None, + omega, + offsets, + round_starts_device, + request_count, + group._num_calibrated_layers, + num_q_heads, + num_freqs, + offset_planes, + True, + # Per-layer dequant scales (quantized pools only, else None): the + # fold multiplies them into the coefficient tables so the score op + # below reads raw quantized elements at zero hot-loop cost. + group._kv_scales, + ) + else: + # Mean aggregation rotates the pre-scaled calibration query by the + # tabulated offset-mean phase of each request's round start (tables + # built once at group construction; design: Fanrong Li, torch-graph + # review 2026-07-20). c_mlr is the static init-time table, so the + # round writes only c_re/c_im and runs zero trigonometry. + c_mlr = group._mlr_fold + torch.ops.trtllm.tri_attention_rotate_mean_score_coefficients( + c_re, + c_im, + group._q_real_scaled, + group._q_imag_scaled, + group._phase_cos, + group._phase_sin, + round_starts_device, + request_count, + group._num_calibrated_layers, + num_q_heads, + num_freqs, + group._max_position, + ) pool_anchor, layer_base_addrs, block_offsets, seg_page_off, seg_req, seg_layer = ( group.pointer_prefix ) @@ -215,8 +238,9 @@ def _launch_tri_score_perhead( num_segments, use_max, group._use_vectorized, - # Validation-only here (presence must match the pool dtype; the fold - # op above already consumed the values). + # Validation-only here (presence must match the pool dtype; the max + # fold above or the init-time pre-scaled mean tables already consumed + # the values). group._kv_scales, ) @@ -348,6 +372,66 @@ def __init__( # be range-checked; validate the extent once here, loudly. if min(layer_indices) < 0 or max(layer_indices) >= self._num_calibrated_layers: raise ValueError("scored layer index exceeds the calibrated layer extent") + # Init-once tables for the mean-aggregation coefficient rotation + # (design: Fanrong Li, torch-graph review 2026-07-20). The + # offset-averaged phase of every possible round-start position is + # tabulated once, RoPE-style (float64 accumulation, stored fp32): + # phase_cos[pos, f] = freq_scale_sq[f] * mean_o cos((pos + offset_o) * omega_f) + # phase_sin[pos, f] = freq_scale_sq[f] * mean_o sin((pos + offset_o) * omega_f) + # and the per-layer KV dequantization scale is pre-multiplied into + # the calibration query (identity for float pools). The MLR + # coefficient has no position term, so its fold + # (kv_scale * freq_scale_sq * mlr) is fully static: the request axis + # disappears and nothing about it is recomputed per round. Every + # eviction round then gathers one table row per request and rotates + # the static query by it -- zero trigonometry at runtime. + # + # Positions can legitimately reach the full sequence capacity, so the + # table covers [0, seq_len] inclusive. The 64-row floor keeps tiny + # synthetic unit-test geometries (whose logical round starts exceed + # their physical bucket) inside the table at negligible cost; every + # row is exact for its position, and production buckets (the + # manager's max sequence length) always exceed the floor. + self._max_position = max(int(seq_len), 63) + 1 + omega64 = omega.view(-1)[: self.num_freqs].to(torch.float64) + freq_scale_sq64 = freq_scale_sq.view(-1)[: self.num_freqs].to(torch.float64) + positions = torch.arange(self._max_position, dtype=torch.float64, device=device) + cos_acc = torch.zeros( + self._max_position, self.num_freqs, dtype=torch.float64, device=device + ) + sin_acc = torch.zeros_like(cos_acc) + for offset in offsets.to(torch.float64).tolist(): + angle = (positions + offset).unsqueeze(1) * omega64.unsqueeze(0) + cos_acc += torch.cos(angle) + sin_acc += torch.sin(angle) + phase_scale = freq_scale_sq64.unsqueeze(0) / float(offsets.numel()) + self._phase_cos = (cos_acc * phase_scale).to(torch.float32).contiguous() + self._phase_sin = (sin_acc * phase_scale).to(torch.float32).contiguous() + calibration_shape = (self._num_calibrated_layers, int(num_q_heads), self.num_freqs) + mlr_fold64 = mlr_coef_LHF.view(calibration_shape).to(torch.float64) * freq_scale_sq64 + if self._kv_scales is None: + # kv_scale is 1.0 for float pools: the pre-scaled query IS the + # calibration query (aliased, not copied). + self._q_real_scaled = q_real_LHF.view(-1) + self._q_imag_scaled = q_imag_LHF.view(-1) + else: + layer_scales64 = ( + self._kv_scales[: self._num_calibrated_layers].to(torch.float64).view(-1, 1, 1) + ) + self._q_real_scaled = ( + (q_real_LHF.view(calibration_shape).to(torch.float64) * layer_scales64) + .to(torch.float32) + .contiguous() + .view(-1) + ) + self._q_imag_scaled = ( + (q_imag_LHF.view(calibration_shape).to(torch.float64) * layer_scales64) + .to(torch.float32) + .contiguous() + .view(-1) + ) + mlr_fold64 = mlr_fold64 * layer_scales64 + self._mlr_fold = mlr_fold64.to(torch.float32).contiguous().view(-1) # Folded per-round coefficient tables, allocated on first launch and # keyed by plane count so switching aggregation (mean: one plane; # max: one plane per offset) re-shapes without churn. @@ -514,15 +598,18 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> ) def _fold_coefficient_buffers( - self, offset_planes: int - ) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor]": - """Return (c_re, c_im, c_mlr) fold tables for one plane count. + self, offset_planes: int, with_mlr: bool + ) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]": + """Return (c_re, c_im, c_mlr) per-round scratch for one plane count. Sized on ``max_requests`` so any launch's active ``request_count`` - fits without reallocation (each launch folds only its active rows); - c_mlr is offset independent so it never grows planes. + fits without reallocation (each launch prepares only its active + rows). Only the max aggregation still writes a per-round c_mlr; the + mean path reads the static init-time MLR table instead (request axis + removed), so its scratch skips c_mlr entirely. """ - buffers = self._fold_buffers.get(offset_planes) + key = (offset_planes, with_mlr) + buffers = self._fold_buffers.get(key) if buffers is None: elements = ( self.max_requests @@ -533,9 +620,9 @@ def _fold_coefficient_buffers( device = self.output.device c_re = torch.empty(offset_planes * elements, dtype=torch.float32, device=device) c_im = torch.empty_like(c_re) - c_mlr = torch.empty(elements, dtype=torch.float32, device=device) + c_mlr = torch.empty(elements, dtype=torch.float32, device=device) if with_mlr else None buffers = (c_re, c_im, c_mlr) - self._fold_buffers[offset_planes] = buffers + self._fold_buffers[key] = buffers return buffers def launch( @@ -611,8 +698,6 @@ def launch( valid_widths, round_starts_device, token_starts_device, - mean_cos, - mean_sin, score_aggregation=score_aggregation, ) return output diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py new file mode 100644 index 000000000000..87ef6eb816f1 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mean-path phase-table rotation vs a direct trigonometric fold. + +The mean-aggregation coefficient preparation tabulates the offset-averaged +phase of every possible round-start position once at initialization (float64 +accumulation, stored fp32) and rotates the pre-scaled calibration query by +one gathered table row per request. This test rebuilds the same coefficients +with direct float64 trigonometry at the exact round starts and compares. +""" + +import torch + +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, +) + + +def test_phase_table_rotation_matches_direct_trig_fold(): + """Table + rotation == direct trig fold at positions {0, mid, last}. + + The last position exercises phases of tens of thousands of radians, + where fp32 runtime trigonometry would already have lost several digits + to argument reduction; the float64-built table must not. + """ + assert hasattr(torch.ops.trtllm, "tri_attention_rotate_mean_score_coefficients"), ( + "TriAttention rotation op is not loaded" + ) + device = torch.device("cuda", torch.cuda.current_device()) + torch.manual_seed(20260720) + num_layers = 2 + num_q_heads = 4 + head_dim = 16 + num_freqs = head_dim // 2 + # Sizes the phase table: the group tabulates [0, seq_len] inclusive. + seq_len = 32768 + pools = [ + torch.randn(2, 2, 2, 4, head_dim, device=device).to(torch.bfloat16) + for _ in range(num_layers) + ] + block_offsets = torch.zeros(1, 3, 2, 1, dtype=torch.int32, device=device) + q_real = torch.randn(num_layers, num_q_heads, num_freqs, device=device) + q_imag = torch.randn(num_layers, num_q_heads, num_freqs, device=device) + mlr_coef = torch.randn(num_layers, num_q_heads, num_freqs, device=device) + freq_scale_sq = torch.rand(num_freqs, device=device) + 0.5 + # RoPE-style inverse frequencies: omega[0] == 1.0 makes the last tabulated + # position a genuinely large trigonometric argument. + omega = 10000.0 ** (-torch.arange(num_freqs, device=device, dtype=torch.float32) / num_freqs) + offsets = torch.tensor([1.0, 2.0, 4.0, 8.0], dtype=torch.float32, device=device) + group = _FixedScoreGroup( + pools, + list(range(num_layers)), + 3, + 1, + seq_len, + num_q_heads, + block_offsets, + [0] * num_layers, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + output_width=4, + ) + max_position = group._max_position + assert max_position == seq_len + 1 + + positions = [0, max_position // 2, max_position - 1] + round_starts = torch.tensor(positions, dtype=torch.int32, device=device) + total = len(positions) * num_layers * num_q_heads * num_freqs + c_re = torch.full((total,), float("nan"), dtype=torch.float32, device=device) + c_im = torch.full_like(c_re, float("nan")) + torch.ops.trtllm.tri_attention_rotate_mean_score_coefficients( + c_re, + c_im, + group._q_real_scaled, + group._q_imag_scaled, + group._phase_cos, + group._phase_sin, + round_starts, + len(positions), + num_layers, + num_q_heads, + num_freqs, + max_position, + ) + + # Direct trigonometric fold: average the offset phases at each round + # start, then rotate and scale the calibration query (float64 throughout, + # cast to fp32 only at the end, exactly like the tables were built). + phase = ( + round_starts.to(torch.float64)[:, None, None] + offsets.to(torch.float64)[None, :, None] + ) * omega.to(torch.float64)[None, None, :] + mean_cos = torch.cos(phase).mean(dim=1)[:, None, None, :] + mean_sin = torch.sin(phase).mean(dim=1)[:, None, None, :] + fss = freq_scale_sq.to(torch.float64) + q_re = q_real.to(torch.float64)[None] + q_im = q_imag.to(torch.float64)[None] + reference_re = (fss * (q_re * mean_cos - q_im * mean_sin)).to(torch.float32) + reference_im = (fss * (q_im * mean_cos + q_re * mean_sin)).to(torch.float32) + + shape = (len(positions), num_layers, num_q_heads, num_freqs) + torch.testing.assert_close(c_re.view(shape), reference_re, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(c_im.view(shape), reference_im, rtol=1e-5, atol=1e-5) From 9e67032168747d58dbda1601fb20545d6e83d8a3 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 06:23:20 -0700 Subject: [PATCH 051/178] [None][perf] fuse TriAttention keep-set settle and move-source packing Merge the per-round tie-settlement kernel and the compaction move-source packing kernel into one Triton launch: the program that finalizes each selection row's kept ordinals immediately packs that row's dense/SWA move source indices, removing one kernel launch per eviction round. BatchedKVCacheCompaction now exports its dense/SWA packing description (hand_move_source_pack_to_selection) and drops its own dense pack launch, so each round packs exactly once; the keep-set selector validates the packing against its own keep buffer and selection geometry before fusing it in. The co-compressed draft keeps its own pack launch because it broadcasts the finalized keep set over the draft's KV-head layout. The standalone settle and pack kernels stay in the module as the references the unit tests compare the fused kernel against. Bit-exactness: the fused kernel reproduces the original two-kernel sequence torch.equal on kept ordinals, dense move sources, and SWA move sources -- including buffer regions neither path overwrites -- across union/per_head/per_layer_perhead, with and without SWA, ragged rows, and multi-block widths (GPU validation: 181 checks, 0 failures; the new unit test replays the same cases), and the settle-only specialization matches the standalone settle kernel exactly. Design: Fanrong Li (torch-graph review, 2026-07-20). Signed-off-by: tianruih --- .../triattention/compaction.py | 109 +++++- .../triattention/triattention.py | 83 ++++- .../triattention/triattention_kernels.py | 182 +++++++++- .../test_triattention_fused_settle_pack.py | 326 ++++++++++++++++++ 4 files changed, 680 insertions(+), 20 deletions(-) create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_fused_settle_pack.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index fedca37eda59..635842bef7d7 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -21,7 +21,9 @@ SWA families; a co-compressed draft adds a second) and then moves the surviving KV in place with batched C++ compact launches. Inputs are plain tensors, so any eviction method that produces a kept-token set per request -can drive it. +can drive it. A driver that finalizes the keep set in its own GPU launch can +take the target's dense/SWA packing over into that launch instead (see +``hand_move_source_pack_to_selection``); the draft always packs here. """ from collections import OrderedDict @@ -220,6 +222,38 @@ def _compact_groups( _PACK_NUM_WARPS = 4 +class _MoveSourcePackArguments(NamedTuple): + """One cache family's move-index packing, described as plain launch data. + + The selection-side fused settle-and-pack launch (design suggested by + Fanrong Li, torch-graph review 2026-07-20) consumes this to pack the + dense/SWA move sources in the same kernel that finalizes the kept + ordinals, instead of a second launch at compaction time. + """ + + kept_token_ordinals: torch.Tensor + valid_sequence_lengths: torch.Tensor + dense_offsets: torch.Tensor + dense_indices: torch.Tensor + # With no SWA family these alias the dense tensors and ``has_swa`` + # specializes every SWA load and store away. + swa_offsets: torch.Tensor + swa_indices: torch.Tensor + dense_total: int + swa_total: int + selection_rows: int + keep_count: int + request_count: int + num_kv_heads: int + swa_window: int + # Widest per-request move count any staged offsets may express; the + # packing loop covers exactly this many move slots per packed row. + move_capacity: int + union: bool + per_layer: bool + has_swa: bool + + def _move_index_pack_launcher( kept_token_ordinals: torch.Tensor, valid_sequence_lengths: torch.Tensor, @@ -234,13 +268,15 @@ def _move_index_pack_launcher( swa_window: int, swa_move_source_offsets: Optional[torch.Tensor], swa_move_source_indices: Optional[torch.Tensor], -) -> Callable[[], None]: +) -> Tuple[Callable[[], None], _MoveSourcePackArguments]: """Build one launch of the move-index packing kernel. The kernel reads the kept-token ordinals and each request's valid length and writes the packed per-(layer, head) move source indices consumed by the C++ compact launches. Only the caller-provided selection tensors are - validated here; the move buffers are allocated by this module. + validated here; the move buffers are allocated by this module. The + returned arguments bundle describes the same packing so a fused + selection-side launch can take it over. """ per_layer = eviction_mode == "per_layer_perhead" union = eviction_mode == "union" @@ -282,6 +318,25 @@ def _move_index_pack_launcher( max_move = decode_keep_count + max_protected_tail if swa_total: max_move = max(max_move, swa_window + max_protected_tail) + pack_arguments = _MoveSourcePackArguments( + kept_token_ordinals=kept_token_ordinals, + valid_sequence_lengths=valid_sequence_lengths, + dense_offsets=move_source_offsets, + dense_indices=move_source_indices, + swa_offsets=swa_offsets_arg, + swa_indices=swa_indices_arg, + dense_total=int(move_source_indices.shape[-1]), + swa_total=swa_total, + selection_rows=selection_rows, + keep_count=decode_keep_count, + request_count=request_count, + num_kv_heads=num_kv_heads, + swa_window=swa_window, + move_capacity=max_move, + union=union, + per_layer=per_layer, + has_swa=swa_total > 0, + ) packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads grid = ( request_count, @@ -298,16 +353,16 @@ def _move_index_pack_launcher( ) # Ordered to match the kernel's constexpr parameter declaration. constexpr_values = dict( - DENSE_TOTAL=int(move_source_indices.shape[-1]), - SWA_TOTAL=swa_total, - SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=decode_keep_count, - KEEP_COUNT=decode_keep_count, - NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window, - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=swa_total > 0, + DENSE_TOTAL=pack_arguments.dense_total, + SWA_TOTAL=pack_arguments.swa_total, + SELECTION_ROWS=pack_arguments.selection_rows, + SELECTION_STRIDE=pack_arguments.keep_count, + KEEP_COUNT=pack_arguments.keep_count, + NUM_KV_HEADS=pack_arguments.num_kv_heads, + SWA_WINDOW=pack_arguments.swa_window, + UNION=pack_arguments.union, + PER_LAYER=pack_arguments.per_layer, + HAS_SWA=pack_arguments.has_swa, BLOCK=_PACK_BLOCK_TOKENS, ) @@ -318,7 +373,7 @@ def launch_pack() -> None: num_warps=_PACK_NUM_WARPS, ) - return launch_pack + return launch_pack, pack_arguments class BatchedKVCacheCompaction: @@ -482,7 +537,7 @@ def __init__( dense_slots = ( {layer: slot for slot, layer in enumerate(self.dense_layers)} if per_layer else None ) - dense_pack = _move_index_pack_launcher( + dense_pack, self._dense_pack_arguments = _move_index_pack_launcher( kept_token_ordinals, valid_sequence_lengths, dense_move_offsets, @@ -542,6 +597,28 @@ def __init__( if compaction is not None ) + def hand_move_source_pack_to_selection(self) -> _MoveSourcePackArguments: + """Hand the dense/SWA move packing over to the selection launch. + + Returns the packing description and drops this object's own dense + pack launch, so each round packs exactly once: the caller's fused + settle-and-pack kernel fills the move buffers when it finalizes the + kept ordinals, and ``compact`` then only runs the C++ moves. The + co-compressed draft keeps its own pack launch because it broadcasts + the finalized keep set over the draft's own KV-head layout. + """ + self.target_dense_compaction = self.target_dense_compaction._replace(move_index_pack=None) + self.cache_compactions = tuple( + compaction + for compaction in ( + self.target_dense_compaction, + self.target_swa_compaction, + self.draft_compaction, + ) + if compaction is not None + ) + return self._dense_pack_arguments + def _build_draft_compaction( self, kept_token_ordinals: torch.Tensor, @@ -607,7 +684,7 @@ def _build_draft_compaction( # packed row, so one more pack launch broadcasts the target keep # set over the draft KV heads and appends the draft's own tail # ordinals (valid_seq_len + 0..tail-1). - draft_pack = _move_index_pack_launcher( + draft_pack, _ = _move_index_pack_launcher( kept_token_ordinals, valid_sequence_lengths, draft_move_offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 3fa8a2e8932e..6cdf5ea985bf 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -164,6 +164,9 @@ def __init__( # cohort may mix prompt lengths. ``row_prompt_offsets`` is the # row-major expansion consumed by the finalizer. self.selection_rows_per_request = int(selection_rows_per_request) + # Optional compaction move packing fused into the settle launch; set + # once the compaction buffers exist (see ``fuse_move_source_pack``). + self._move_source_pack = None if prompt_offsets_buffer is not None: # Share the staging buffers' per-request prompt lengths so the # values are written once per round. @@ -210,15 +213,41 @@ def _bind_selection_rows( self._provisional_rows = provisional_indices self._keep_rows = keep_rows + def fuse_move_source_pack(self, pack_arguments) -> None: + """Pack compaction move sources inside this selector's settle launch. + + ``pack_arguments`` is the dense/SWA packing description exported by + ``BatchedKVCacheCompaction.hand_move_source_pack_to_selection`` + (fusion suggested by Fanrong Li, torch-graph review 2026-07-20). The + fused kernel reads back the kept ordinals it just wrote, so the + packing must read this selector's own keep buffer, and the packing + geometry must match the selection rows this selector settles. + """ + if ( + pack_arguments.kept_token_ordinals.data_ptr() != self._keep_rows.data_ptr() + or pack_arguments.kept_token_ordinals.numel() != self._keep_rows.numel() + ): + raise ValueError("fused move packing must read this selector's keep buffer") + if ( + pack_arguments.selection_rows != self.selection_rows_per_request + or pack_arguments.keep_count != self.keep_count + or pack_arguments.request_count * self.selection_rows_per_request + != int(self._keep_rows.shape[0]) + ): + raise ValueError("fused move packing does not match the selector geometry") + self._move_source_pack = pack_arguments + def _select_top_tokens(self) -> None: """Pick the top-k with the CuTE selector, then settle its output. The CuTE top-k is fast but breaks score ties arbitrarily and emits indices in arbitrary order; the settle kernel recomputes the threshold membership with lowest-index-wins ties, rebases each row by its prompt - offset, and writes sorted ordinals. + offset, and writes sorted ordinals. When a compaction move packing is + fused in, the same launch also packs each request's dense/SWA move + source indices from the ordinals it just settled. """ - from .triattention_kernels import _settle_ties_after_topk_kernel + from .triattention_kernels import _settle_ties_and_pack_compaction_sources_kernel rows = int(self._selection_scores_rows.shape[0]) # The trailing 1 is next_n: decode scores one query token per request. @@ -229,15 +258,56 @@ def _select_top_tokens(self) -> None: self.keep_count, 1, ) - _settle_ties_after_topk_kernel[(rows,)]( + pack = self._move_source_pack + if pack is None: + # Settle only: the pack half is compiled away, so its tensor + # parameters are never read; any resident tensor stands in. + placeholder = self._selection_row_lengths + pack_tensors = (placeholder,) * 5 + pack_shape = dict( + DENSE_TOTAL=0, + SWA_TOTAL=0, + MOVE_CAPACITY=0, + NUM_KV_HEADS=1, + SWA_WINDOW=0, + UNION=False, + PER_LAYER=False, + HAS_SWA=False, + HAS_PACK=False, + ) + else: + pack_tensors = ( + pack.valid_sequence_lengths, + pack.dense_offsets, + pack.dense_indices, + pack.swa_offsets, + pack.swa_indices, + ) + pack_shape = dict( + DENSE_TOTAL=pack.dense_total, + SWA_TOTAL=pack.swa_total, + MOVE_CAPACITY=pack.move_capacity, + NUM_KV_HEADS=pack.num_kv_heads, + SWA_WINDOW=pack.swa_window, + UNION=pack.union, + PER_LAYER=pack.per_layer, + HAS_SWA=pack.has_swa, + HAS_PACK=True, + ) + _settle_ties_and_pack_compaction_sources_kernel[ + (rows // self.selection_rows_per_request, self.selection_rows_per_request) + ]( self._selection_scores_rows, self._selection_row_lengths, self.row_prompt_offsets, self._provisional_rows, self._keep_rows, + *pack_tensors, WIDTH=self.width, KEEP_COUNT=self.keep_count, OUTPUT_WIDTH=self.keep_count, + SELECTION_ROWS=self.selection_rows_per_request, + **pack_shape, BLOCK=256, num_warps=4, ) @@ -1996,6 +2066,13 @@ def _batched_compaction_for( draft_move_offsets=score_staging.draft_move_offsets, **draft_kwargs, ) + # One launch settles the kept ordinals and packs the dense/SWA + # move sources; the compaction keeps only its C++ moves (plus the + # draft's own pack). Both caches are invalidated together, so the + # fused packing always points at the live compaction buffers. + keep_set_selector.fuse_move_source_pack( + batched_compaction.hand_move_source_pack_to_selection() + ) self._batched_compaction = batched_compaction # Tails vary per round (in-flight growth), so the per-family move # offsets ride the staged metadata table each round. diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 219d45867eb8..7f3a220707fc 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -989,7 +989,12 @@ def _settle_ties_after_topk_kernel( OUTPUT_WIDTH: tl.constexpr, BLOCK: tl.constexpr, ): - """Resolve boundary ties and emit increasing physical token indices.""" + """Resolve boundary ties and emit increasing physical token indices. + + The production selectors launch the fused settle-and-pack kernel below; + this standalone version is kept as the reference the unit tests compare + the fused kernel against. + """ row = tl.program_id(0) row_scores = scores + row * WIDTH row_selected = provisional_indices + row * KEEP_COUNT @@ -1123,3 +1128,178 @@ def _pack_compaction_sources_kernel( swa_source, mask=write_swa & (move < swa_count), ) + + +# --------------------------------------------------------------------------- # +# Fused finalize: settle the top-k ties and pack the move indices in one # +# launch (fusion suggested by Fanrong Li, torch-graph review 2026-07-20). # +# --------------------------------------------------------------------------- # + + +@triton.jit +def _settle_ties_and_pack_compaction_sources_kernel( + scores, + seq_lens, + prompt_offsets, + provisional_indices, + output_indices, + valid_seq_lens, + dense_offsets, + dense_indices, + swa_offsets, + swa_indices, + WIDTH: tl.constexpr, + KEEP_COUNT: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + DENSE_TOTAL: tl.constexpr, + SWA_TOTAL: tl.constexpr, + MOVE_CAPACITY: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + SWA_WINDOW: tl.constexpr, + UNION: tl.constexpr, + PER_LAYER: tl.constexpr, + HAS_SWA: tl.constexpr, + HAS_PACK: tl.constexpr, + BLOCK: tl.constexpr, +): + """Settle one selection row's ties, then pack its compaction move sources. + + One program per (request, selection row). The first half repeats + ``_settle_ties_after_topk_kernel`` verbatim: recover the top-k threshold + from the provisional selection, count the strictly greater scores, then + emit the kept ordinals in increasing order, rebased by the row's pinned + prompt length. With ``HAS_PACK`` the same program then continues with the + ``_pack_compaction_sources_kernel`` work for the packed rows this + selection row feeds: the kept ordinals it just wrote, followed by the + request's protected tail, plus the SWA rows (latest window) under the + same conditions as the standalone kernel. Union selection has one row per + request feeding every KV head's packed row, so that single program writes + all of them. ``HAS_PACK=False`` compiles the second half away, leaving + exactly the standalone settle. Fusing the two launches was suggested by + Fanrong Li (torch-graph review 2026-07-20). + """ + request = tl.program_id(0) + selection_domain = tl.program_id(1) + row = request * SELECTION_ROWS + selection_domain + row_scores = scores + row * WIDTH + row_selected = provisional_indices + row * KEEP_COUNT + row_output = output_indices + row * OUTPUT_WIDTH + # Scores are decode-relative; this row's pinned prompt length rebases the + # emitted ordinals to absolute positions (per row, so one launch may mix + # prompt lengths). + prompt_len = tl.load(prompt_offsets + row) + + 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, + ) + selected_score = tl.load( + row_scores + token_index, + mask=selected_mask, + other=float("inf"), + ).to(tl.float32) + threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) + + seq_len = tl.load(seq_lens + 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 < seq_len) + 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 < seq_len) + 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_len, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) + + if HAS_PACK: + # The emission above scatters through other lanes of this program; + # make those global stores visible to every lane before the pack + # half reads the row back. + tl.debug_barrier() + dense_begin = tl.load(dense_offsets + request) + dense_end = tl.load(dense_offsets + request + 1) + dense_count = dense_end - dense_begin + valid_len = tl.load(valid_seq_lens + request) + if HAS_SWA: + swa_begin = tl.load(swa_offsets + request) + swa_end = tl.load(swa_offsets + request + 1) + swa_count = swa_end - swa_begin + for move_start in tl.static_range(0, MOVE_CAPACITY, BLOCK): + move = move_start + tl.arange(0, BLOCK) + selected = tl.load( + row_output + move, + mask=move < KEEP_COUNT, + other=0, + ) + dense_source = tl.where(move < KEEP_COUNT, selected, valid_len + move - KEEP_COUNT) + if UNION: + # The one union row per request feeds every KV head's packed + # row with the same move sources. + for head in tl.static_range(0, NUM_KV_HEADS): + tl.store( + dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, + dense_source, + mask=move < dense_count, + ) + else: + domain = tl.program_id(1) + dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) + if HAS_SWA: + swa_source = valid_len - SWA_WINDOW + move + if UNION: + for head in tl.static_range(0, NUM_KV_HEADS): + tl.store( + swa_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, + swa_source, + mask=move < swa_count, + ) + else: + domain = tl.program_id(1) + # Per-layer selection has one dense domain per (layer, + # head). SWA uses one shared source row per head, so only + # the first layer writes it. + if PER_LAYER: + write_swa = domain < NUM_KV_HEADS + else: + write_swa = move >= 0 + head = domain % NUM_KV_HEADS + swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move + tl.store( + swa_indices + swa_output, + swa_source, + mask=write_swa & (move < swa_count), + ) 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..816074f669ac --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_fused_settle_pack.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The fused settle-and-pack kernel must reproduce the original two-kernel +sequence byte for byte. + +The original tie settlement and move-source packing kernels stay in the +module as the reference here: every case launches them on one set of buffers +and the fused kernel on an identically initialized set, then requires +``torch.equal`` on the kept ordinals, the dense move sources, and the SWA +move sources -- including the buffer regions neither path overwrites (rows +shorter than the keep count leave stale entries behind, and the packing +forwards those stale entries the same way in both paths). +""" + +import pytest +import torch +from conftest import encode_block_offsets as _encode_block_offsets + +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( + BatchedKVCacheCompaction, +) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _BatchedUnionKeepSetSelector, +) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _pack_compaction_sources_kernel, + _settle_ties_after_topk_kernel, + _settle_ties_and_pack_compaction_sources_kernel, +) + +_BLOCK = 256 +_NUM_WARPS = 4 + + +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 _staged_offsets(counts, device): + offsets = [0] + for count in counts: + offsets.append(offsets[-1] + count) + return torch.tensor(offsets, dtype=torch.int32, device=device) + + +@pytest.mark.parametrize("has_swa", [False, True]) +@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) +@pytest.mark.parametrize( + "width,keep_count", + [ + # 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 test_fused_settle_pack_matches_two_kernel_sequence(eviction_mode, has_swa, width, keep_count): + device = torch.device("cuda", torch.cuda.current_device()) + request_count, num_layers, num_kv_heads = 3, 2, 2 + union = eviction_mode == "union" + per_layer = eviction_mode == "per_layer_perhead" + selection_rows = _selection_rows_for(eviction_mode, num_layers, num_kv_heads) + rows_total = request_count * selection_rows + packed_rows = num_layers * num_kv_heads if per_layer else num_kv_heads + + # Move geometry: the last request is a padded row that moves nothing. + protected_tails = [2, 0, 0] + tail_capacity = max(protected_tails) + dense_counts = [keep_count + protected_tails[0], keep_count + protected_tails[1], 0] + swa_window = 6 + swa_counts = [swa_window + protected_tails[0], swa_window + protected_tails[1], 0] + move_capacity = keep_count + tail_capacity + if has_swa: + move_capacity = max(move_capacity, swa_window + tail_capacity) + dense_total = request_count * (keep_count + tail_capacity) + swa_total = request_count * (swa_window + tail_capacity) + valid_seq_lens = torch.tensor([10, 8, 0], dtype=torch.int32, device=device) + dense_offsets = _staged_offsets(dense_counts, device) + swa_offsets = _staged_offsets(swa_counts, device) + + for seed in range(5): + 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, + ) + row_prompt_offsets = torch.tensor( + [3 * (row % 3) for row in range(rows_total)], 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() + + # 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 + ) + dense_stale = torch.randint( + -(2**30), 2**30, (packed_rows, dense_total), dtype=torch.int32, device=device + ) + swa_stale = torch.randint( + -(2**30), 2**30, (num_kv_heads, swa_total), dtype=torch.int32, device=device + ) + + output_reference = output_stale.clone() + dense_reference = dense_stale.clone() + swa_reference = swa_stale.clone() + _settle_ties_after_topk_kernel[(rows_total,)]( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output_reference, + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + BLOCK=_BLOCK, + num_warps=_NUM_WARPS, + ) + swa_offsets_arg = swa_offsets if has_swa else dense_offsets + swa_reference_arg = swa_reference if has_swa else dense_reference + _pack_compaction_sources_kernel[ + (request_count, packed_rows, (move_capacity + _BLOCK - 1) // _BLOCK) + ]( + output_reference, + valid_seq_lens, + dense_offsets, + dense_reference, + swa_offsets_arg, + swa_reference_arg, + DENSE_TOTAL=dense_total, + SWA_TOTAL=swa_total if has_swa else 0, + SELECTION_ROWS=selection_rows, + SELECTION_STRIDE=keep_count, + KEEP_COUNT=keep_count, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window if has_swa else 0, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=has_swa, + BLOCK=_BLOCK, + num_warps=_NUM_WARPS, + ) + + output_fused = output_stale.clone() + dense_fused = dense_stale.clone() + swa_fused = swa_stale.clone() + swa_fused_arg = swa_fused if has_swa else dense_fused + _settle_ties_and_pack_compaction_sources_kernel[(request_count, selection_rows)]( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output_fused, + valid_seq_lens, + dense_offsets, + dense_fused, + swa_offsets_arg, + swa_fused_arg, + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + SELECTION_ROWS=selection_rows, + DENSE_TOTAL=dense_total, + SWA_TOTAL=swa_total if has_swa else 0, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window if has_swa else 0, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=has_swa, + HAS_PACK=True, + BLOCK=_BLOCK, + num_warps=_NUM_WARPS, + ) + torch.cuda.synchronize(device) + + assert torch.equal(output_fused, output_reference), f"kept ordinals differ (seed {seed})" + assert torch.equal(dense_fused, dense_reference), f"dense moves differ (seed {seed})" + assert torch.equal(swa_fused, swa_reference), f"SWA moves differ (seed {seed})" + + +def test_fused_kernel_without_pack_matches_standalone_settle(): + """``HAS_PACK=False`` must leave exactly the standalone settle kernel.""" + device = torch.device("cuda", torch.cuda.current_device()) + rows_total, width, keep_count = 6, 33, 7 + generator = torch.Generator(device=device).manual_seed(11) + scores = torch.randint( + -2, 3, (rows_total, width), generator=generator, dtype=torch.int32, device=device + ).to(torch.float32) + row_lengths = torch.tensor([0, 3, 9, 17, 33, 33], dtype=torch.int32, device=device) + row_prompt_offsets = torch.tensor([5, 0, 2, 0, 1, 4], dtype=torch.int32, device=device) + 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() + output_stale = torch.randint( + -(2**30), 2**30, (rows_total, keep_count), dtype=torch.int32, device=device + ) + + output_reference = output_stale.clone() + _settle_ties_after_topk_kernel[(rows_total,)]( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output_reference, + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + BLOCK=_BLOCK, + num_warps=_NUM_WARPS, + ) + + output_fused = output_stale.clone() + placeholder = row_lengths + _settle_ties_and_pack_compaction_sources_kernel[(rows_total, 1)]( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output_fused, + placeholder, + placeholder, + placeholder, + placeholder, + placeholder, + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + SELECTION_ROWS=1, + DENSE_TOTAL=0, + SWA_TOTAL=0, + MOVE_CAPACITY=0, + NUM_KV_HEADS=1, + SWA_WINDOW=0, + UNION=False, + PER_LAYER=False, + HAS_SWA=False, + HAS_PACK=False, + BLOCK=_BLOCK, + num_warps=_NUM_WARPS, + ) + torch.cuda.synchronize(device) + + assert torch.equal(output_fused, output_reference) + + +def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buffers(): + """The handoff exports the live move buffers, drops the compaction-time + dense pack launch, and the selector only accepts a packing that reads its + own keep buffer.""" + device = torch.device("cuda", torch.cuda.current_device()) + request_count, num_kv_heads, keep_count, width = 2, 2, 4, 16 + tokens_per_block, head_dim = 4, 8 + pools = [ + torch.zeros(6, 2, num_kv_heads, tokens_per_block, head_dim, device=device) for _ in range(2) + ] + page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + + selector = _BatchedUnionKeepSetSelector( + rows=3, + width=width, + keep_count=keep_count, + dtype=torch.float32, + device=device, + max_requests=request_count, + input_scores=torch.zeros(request_count, 3, width, device=device), + normalize_scores=False, + ) + + def build_compaction(kept_token_ordinals): + return BatchedKVCacheCompaction( + eviction_mode="union", + layer_pools=pools, + dense_layers=[0, 1], + swa_layers=[], + layer_group_representative={0: 0, 1: 1}, + layer_pool_keys=[("dense", 0), ("dense", 0)], + kept_token_ordinals=kept_token_ordinals, + valid_sequence_lengths=torch.full( + (request_count,), 10, dtype=torch.int32, device=device + ), + kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), + page_table_slots={0: 0, 1: 0}, + request_count=request_count, + prompt_offsets=torch.zeros(request_count, dtype=torch.int32, device=device), + decode_keep_count=keep_count, + swa_window=None, + protected_tail_capacity=1, + ) + + compaction = build_compaction(selector.keep) + assert compaction.target_dense_compaction.move_index_pack is not None + + pack_arguments = compaction.hand_move_source_pack_to_selection() + assert compaction.target_dense_compaction.move_index_pack is None + assert len(compaction.cache_compactions) == 1 + assert compaction.cache_compactions[0] is compaction.target_dense_compaction + assert pack_arguments.dense_indices is compaction.target_dense_compaction.move_source_indices + assert pack_arguments.dense_offsets is compaction.target_dense_compaction.move_source_offsets + + selector.fuse_move_source_pack(pack_arguments) + assert selector._move_source_pack is pack_arguments + + # A packing built over any other keep buffer must be rejected: the fused + # kernel reads back the ordinals it just wrote. + foreign = build_compaction(torch.zeros_like(selector.keep)) + with pytest.raises(ValueError, match="keep buffer"): + selector.fuse_move_source_pack(foreign.hand_move_source_pack_to_selection()) From 538dcb9089a12d210e02ce6d6038468141d63cf6 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 08:14:05 -0700 Subject: [PATCH 052/178] [None][fix] Version-adapt CuTe sqrt and restore score kernel register profile Two independent fixes on the round-2 kernel branch: 1) CuTe score sqrt shim: cute.math.sqrt only accepts the approx/ftz keywords on newer CuTe DSL releases; the container's DSL raises TypeError at cute.compile time. Probe the installed signature once at import (inspect.signature; a trace-time TypeError cannot be caught inside the jit) and branch on the resulting trace-time constant: supported DSLs keep the measured approx+ftz configuration unchanged, older DSLs fall back to the plain (IEEE) sqrt, which is strictly more accurate; the unit test's 5e-3 oracle tolerance absorbs the difference. No other change to the CuTe kernel. 2) Score kernel +11% regression (score_main 1002us vs 902us baseline while compact stayed flat): the round-2 switch of c_mlr to a static [layer, head, freq] table threaded a second flat offset (mlrOff0) through scoreComputeChunk, adding a live int64 pair on top of the c_re/c_im offset in a kernel tuned to exactly 72 registers at __launch_bounds__(128, 7). Hoist the static row pointer once in the kernel prologue (no request stride) and pass that single pre-offset pointer through the chunk instead: one pointer register replaces the old base + request-strided offset pair, restoring the baseline hot-loop instruction pattern while keeping the table static. Same treatment on the scalar fallback kernel. No functional change; the c_mlr addressing still resolves to the identical static table row. Signed-off-by: tianruih --- .../triAttentionScoreKernels.cu | 36 +++++++++++------- .../triattention/triattention_cute_score.py | 37 ++++++++++++++++--- 2 files changed, 55 insertions(+), 18 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu index 16c0f4869db0..2c9166c9e793 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu @@ -166,13 +166,17 @@ __device__ __forceinline__ void scoreLoadChunk(char const* row, bool valid, int // Accumulate one 8-frequency chunk into this thread's per-head accumulators. // coff0 = flat index of (request, layer, first head of this block, chunk -// frequency 0) in the c_re/c_im tables; mlrOff0 = the matching index in the -// static request-independent [layer, head, freq] c_mlr table. All coefficient -// reads are lane-uniform 16-byte loads. |K| is computed once per +// frequency 0) in the c_re/c_im tables; cMlr = the matching chunk pointer +// into the static request-independent [layer, head, freq] c_mlr table, +// pre-offset by the caller. Passing the resolved pointer (instead of a +// second flat offset) keeps this loop's live set at the tuned ~72-register +// baseline: one pointer register replaces the base + request-strided offset +// pair the c_mlr reads consumed before the table went static. All +// coefficient reads are lane-uniform 16-byte loads. |K| is computed once per // (token, frequency) BEFORE the head loop so the GROUP heads share it from // registers. template -__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, int64_t coff0, int64_t mlrOff0, +__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, int64_t coff0, float const* cMlr, int64_t planeStride, uint4 re4, uint4 im4, float* accMean, float* accMlr, float* accPos) { float kRe[8], kIm[8], kMag[8]; @@ -187,7 +191,7 @@ __device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, in for (int hg = 0; hg < GROUP; ++hg) { int64_t const coff = coff0 + static_cast(hg) * a.numFreqs; - float const* cmp = a.cMlr + mlrOff0 + static_cast(hg) * a.numFreqs; + float const* cmp = cMlr + static_cast(hg) * a.numFreqs; float4 const cm0 = __ldg(reinterpret_cast(cmp)); float4 const cm1 = __ldg(reinterpret_cast(cmp + 4)); float const cml[8] = {cm0.x, cm0.y, cm0.z, cm0.w, cm1.x, cm1.y, cm1.z, cm1.w}; @@ -321,8 +325,10 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne int64_t const coff0 = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + static_cast(headBase) * a.numFreqs; // The MLR coefficient is position independent, so its table is folded once - // at initialization without a request axis: [layer, head, freq]. - int64_t const mlrOff0 = (static_cast(layerId) * a.numQueryHeads + headBase) * a.numFreqs; + // at initialization without a request axis: [layer, head, freq]. Hoist the + // row pointer here (no request stride) so the fully unrolled chunk loop + // sees a single pre-offset pointer, not an extra live flat offset. + float const* const cMlrRow = a.cMlr + (static_cast(layerId) * a.numQueryHeads + headBase) * a.numFreqs; int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads * static_cast(a.numFreqs); @@ -356,7 +362,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne uint4 re4, im4; scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); scoreComputeChunk( - a, coff0 + c * 8, mlrOff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + a, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); } } else @@ -367,7 +373,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne uint4 re4, im4; scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); scoreComputeChunk( - a, coff0 + c * 8, mlrOff0 + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + a, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); } } @@ -443,6 +449,11 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads * static_cast(a.numFreqs); + // The static request-independent [layer, head, freq] c_mlr table: hoist + // this block's first-head row pointer once (no request stride), mirroring + // the vectorized kernel; the head loop advances it by numFreqs per head. + float const* const cMlrRow + = a.cMlr + (static_cast(layerId) * a.numQueryHeads + kvHead * groupSize) * a.numFreqs; int const tDec = absT - tokenStart; bool const store = tDec >= 0 && absT < seqLen; @@ -452,8 +463,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde int64_t const coff = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + static_cast(h) * a.numFreqs; - // Static request-independent [layer, head, freq] MLR table index. - int64_t const mlrOff = (static_cast(layerId) * a.numQueryHeads + h) * a.numFreqs; + float const* const cml = cMlrRow + static_cast(hg) * a.numFreqs; float acc = 0.0f; float accMlr = 0.0f; float accPos[kMaxScoreOffsets]; @@ -474,7 +484,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde float const kMag = triSqrtApprox(kRe * kRe + kIm * kIm); if constexpr (USE_MAX) { - accMlr = fmaf(kMag, a.cMlr[mlrOff + f], accMlr); + accMlr = fmaf(kMag, cml[f], accMlr); #pragma unroll for (int o = 0; o < kMaxScoreOffsets; ++o) { @@ -487,7 +497,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde } else { - acc = fmaf(kRe, a.cRe[coff + f], fmaf(kIm, a.cIm[coff + f], fmaf(kMag, a.cMlr[mlrOff + f], acc))); + acc = fmaf(kRe, a.cRe[coff + f], fmaf(kIm, a.cIm[coff + f], fmaf(kMag, cml[f], acc))); } } if (store) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py index 9580bf8ca6d4..0ecd19bd6007 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -17,6 +17,7 @@ from __future__ import annotations +import inspect import threading import cuda.bindings.driver as cuda @@ -29,6 +30,24 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.runtime import from_dlpack + +def _cute_sqrt_supports_approx_ftz() -> bool: + """Probe whether this CuTe DSL's ``cute.math.sqrt`` takes ``approx``/``ftz``. + + Older DSL releases expose a plain one-argument ``sqrt``; passing the + keywords there raises TypeError at trace time (inside ``cute.compile``, + where it cannot be caught), so the capability is probed once at import + time via signature inspection and folded into a trace-time constant. + """ + try: + parameters = inspect.signature(cute.math.sqrt).parameters + except (TypeError, ValueError): + return False + return "approx" in parameters and "ftz" in parameters + + +_CUTE_SQRT_HAS_APPROX_FTZ = _cute_sqrt_supports_approx_ftz() + CTA_M = 64 K = 96 N = 8 @@ -887,11 +906,19 @@ def kernel( real = staged_real[prefetch_index] imag = staged_imag[prefetch_index] norm2 = real * real + imag * imag - magnitude = cute.math.sqrt( - norm2, - approx=self.sqrt_mode == "approx", - ftz=self.magnitude_sqrt_ftz, - ) + if cutlass.const_expr(_CUTE_SQRT_HAS_APPROX_FTZ): + magnitude = cute.math.sqrt( + norm2, + approx=self.sqrt_mode == "approx", + ftz=self.magnitude_sqrt_ftz, + ) + else: + # DSLs without the keywords get the plain (IEEE) + # sqrt, which is strictly MORE accurate than the + # measured approx+ftz choice above; the unit + # test's 5e-3 oracle tolerance absorbs the + # difference. + magnitude = cute.math.sqrt(norm2) magnitude_fp16_0 = cutlass.Float16(magnitude) magnitude_fp16_1 = cutlass.Float16( magnitude - cutlass.Float32(magnitude_fp16_0) From d72bb15bc5735ad8951d4794324a978311483a01 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 08:31:38 -0700 Subject: [PATCH 053/178] [None][fix] Map the CuTe sqrt shim onto cutlass 4.5 fastmath cutlass 4.5 renamed the approximate-sqrt keywords (approx/ftz) to a single fastmath flag; probe all three spellings so the authored approximate-sqrt behavior is preserved on 4.5 instead of falling back to the plain IEEE sqrt. Signed-off-by: tianruih --- .../triattention/triattention_cute_score.py | 34 ++++++++++++------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py index 0ecd19bd6007..19ec1f885885 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -31,22 +31,28 @@ from cutlass.cute.runtime import from_dlpack -def _cute_sqrt_supports_approx_ftz() -> bool: - """Probe whether this CuTe DSL's ``cute.math.sqrt`` takes ``approx``/``ftz``. +def _cute_sqrt_keyword_mode() -> str: + """Probe which fast-sqrt spelling this CuTe DSL's ``cute.math.sqrt`` takes. - Older DSL releases expose a plain one-argument ``sqrt``; passing the - keywords there raises TypeError at trace time (inside ``cute.compile``, + The approximate-sqrt control was renamed across DSL releases: some expose + ``approx``/``ftz`` keywords, cutlass 4.5 exposes a single ``fastmath`` + flag, and older releases expose a plain one-argument ``sqrt``. Passing an + unknown keyword raises TypeError at trace time (inside ``cute.compile``, where it cannot be caught), so the capability is probed once at import time via signature inspection and folded into a trace-time constant. """ try: parameters = inspect.signature(cute.math.sqrt).parameters except (TypeError, ValueError): - return False - return "approx" in parameters and "ftz" in parameters + return "plain" + if "approx" in parameters and "ftz" in parameters: + return "approx_ftz" + if "fastmath" in parameters: + return "fastmath" + return "plain" -_CUTE_SQRT_HAS_APPROX_FTZ = _cute_sqrt_supports_approx_ftz() +_CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() CTA_M = 64 K = 96 @@ -906,18 +912,22 @@ def kernel( real = staged_real[prefetch_index] imag = staged_imag[prefetch_index] norm2 = real * real + imag * imag - if cutlass.const_expr(_CUTE_SQRT_HAS_APPROX_FTZ): + if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): magnitude = cute.math.sqrt( norm2, approx=self.sqrt_mode == "approx", ftz=self.magnitude_sqrt_ftz, ) + elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): + # cutlass 4.5 renamed the approximate-sqrt control + # to ``fastmath``; map the measured approx choice + # onto it to preserve the authored behavior. + magnitude = cute.math.sqrt(norm2, fastmath=self.sqrt_mode == "approx") else: - # DSLs without the keywords get the plain (IEEE) + # DSLs with neither spelling get the plain (IEEE) # sqrt, which is strictly MORE accurate than the - # measured approx+ftz choice above; the unit - # test's 5e-3 oracle tolerance absorbs the - # difference. + # measured approx choice above; the unit test's + # 5e-3 oracle tolerance absorbs the difference. magnitude = cute.math.sqrt(norm2) magnitude_fp16_0 = cutlass.Float16(magnitude) magnitude_fp16_1 = cutlass.Float16( From 72332122ff2c4ccdde10d5debdc62aab8099aaa1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 08:07:03 -0700 Subject: [PATCH 054/178] [None][perf] rotate TriAttention mean score coefficients in the score kernel prologue Builds on Fanrong Li's phase-table design (torch-graph review 2026-07-20): the mean-aggregation eviction round no longer launches the standalone coefficient rotation kernel. Instead each score CTA gathers its request's phase-table row and its (layer, head-block) slice of the pre-scaled calibration query, rotates them into shared memory in a short prologue, and reads coefficients from there. One launch per round, no per-round c_re/c_im scratch at all. The prologue and the standalone kernel compile one shared rotation expression, so scores are bit-identical between the two preparation flavors; the standalone kernel and its op stay as the unit tests' equality reference leg and as the fallback for coefficient blocks past the score launch's 48KB dynamic shared memory bound. The max aggregation path is untouched. Rebased over the score kernel register fix (538dcb9089): scoreComputeChunk keeps the fix's pre-offset c_mlr row pointer (cMlr/cMlrRow) through the new ROTATE_IN_CTA signature. Signed-off-by: tianruih --- .../triAttentionScoreKernels.cu | 284 ++++++++++++++---- .../triAttentionScoreKernels.h | 40 ++- cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp | 109 +++++-- .../triattention/triattention.py | 5 +- .../triattention/triattention_kernels.py | 77 ++++- .../test_triattention_phase_rotation.py | 136 ++++++++- 6 files changed, 550 insertions(+), 101 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu index 2c9166c9e793..381f9bcb4a4f 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu @@ -26,14 +26,16 @@ // score(t, h) = sum_f K_re(t,f)*c_re + K_im(t,f)*c_im + |K(t,f)|*c_mlr // // One thread scores one token across all frequencies; a 128-thread CTA covers -// 128 consecutive tokens of one (request, layer, KV-head) segment. There are -// no shuffles, no shared memory, and no barriers: each thread keeps one fused +// 128 consecutive tokens of one (request, layer, KV-head) segment. The hot +// loop has no shuffles and no barriers: each thread keeps one fused // accumulator per query head of its GQA group ("mean" aggregation) or one // partial sum per offset plane ("max" aggregation, where max over offsets // does not commute through the frequency sum). Coefficient loads are -// lane-uniform 16-byte reads served by L1 broadcast; K loads are 16-byte -// chunks of 8 frequencies when the pool layout allows it, otherwise a fully -// strided scalar path runs the same math. +// lane-uniform 16-byte reads served by L1 broadcast (or by shared memory on +// the mean path's default in-CTA rotation, whose prologue is the one place a +// score CTA uses shared memory and a barrier); K loads are 16-byte chunks of +// 8 frequencies when the pool layout allows it, otherwise a fully strided +// scalar path runs the same math. // // The kernel accumulates the frequency reduction in sequential chunks (not a // block-wide tree), so results are tolerance-equal, not bit-equal, against @@ -82,6 +84,22 @@ __device__ __forceinline__ float triSqrtApprox(float x) #endif } +// The ONE mean-path coefficient rotation, compiled by BOTH the standalone +// triRotateMeanScoreCoefficientsKernel and the score kernels' in-CTA +// prologue. The operand order is load-bearing: a single source expression +// makes the compiler emit the same multiply/FMA contraction in every +// inlining site, so the in-CTA coefficients are bit-identical to the +// standalone kernel's output and end-to-end mean scores are bit-identical +// between the two preparation paths (the unit suite proves this with +// torch.equal). Do not reorder or "simplify" these expressions. +__device__ __forceinline__ float2 rotateCoefficientPair(float qre, float qim, float pc, float ps) +{ + float2 c; + c.x = qre * pc - qim * ps; + c.y = qim * pc + qre * ps; + return c; +} + template __device__ __forceinline__ float toFloat(T value); @@ -164,21 +182,53 @@ __device__ __forceinline__ void scoreLoadChunk(char const* row, bool valid, int } } +// Mean-path in-CTA rotation prologue shared by the vectorized and scalar +// score kernels: rebuild the [coefCount = group * numFreqs] coefficient block +// this CTA will read — gather the request's phase-table row (L2-resident, +// shared by all of the request's CTAs) and the CTA's slice of the static +// pre-scaled calibration query, rotate with rotateCoefficientPair (the SAME +// arithmetic the standalone rotation kernel compiles, the bit-equality +// contract), and stage the results in shared memory. Thread-strided plain +// loads + FMAs only, one barrier; the prologue's temporaries die here, so +// ptxas can fold them into the main loop's register budget. +__device__ __forceinline__ void rotateCoefficientsIntoShared( + FoldedScoreParams const& a, int reqId, int layerId, int headBase, int coefCount, float* sCRe, float* sCIm) +{ + int64_t const phaseRow = static_cast(a.roundStarts[reqId]) * a.numFreqs; + int64_t const qBase = (static_cast(layerId) * a.numQueryHeads + headBase) * a.numFreqs; + for (int i = threadIdx.x; i < coefCount; i += kScoreBlockThreads) + { + int const f = i % a.numFreqs; + float const pc = a.phaseCos[phaseRow + f]; + float const ps = a.phaseSin[phaseRow + f]; + float2 const c = rotateCoefficientPair(a.qReS[qBase + i], a.qImS[qBase + i], pc, ps); + sCRe[i] = c.x; + sCIm[i] = c.y; + } + __syncthreads(); +} + // Accumulate one 8-frequency chunk into this thread's per-head accumulators. +// cRe/cIm point at the coefficient source: the global per-round tables (with // coff0 = flat index of (request, layer, first head of this block, chunk -// frequency 0) in the c_re/c_im tables; cMlr = the matching chunk pointer -// into the static request-independent [layer, head, freq] c_mlr table, -// pre-offset by the caller. Passing the resolved pointer (instead of a -// second flat offset) keeps this loop's live set at the tuned ~72-register -// baseline: one pointer register replaces the base + request-strided offset -// pair the c_mlr reads consumed before the table went static. All -// coefficient reads are lane-uniform 16-byte loads. |K| is computed once per -// (token, frequency) BEFORE the head loop so the GROUP heads share it from -// registers. -template -__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, int64_t coff0, float const* cMlr, - int64_t planeStride, uint4 re4, uint4 im4, float* accMean, float* accMlr, float* accPos) +// frequency 0)) or, under ROTATE_IN_CTA, this CTA's shared-memory block +// (with coff0 = the chunk's frequency offset — the head-block origin is the +// shared block itself). cMlr = the matching chunk pointer into the static +// request-independent [layer, head, freq] c_mlr table, pre-offset by the +// caller. Passing the resolved pointer (instead of a second flat offset) +// keeps this loop's live set at the tuned ~72-register baseline: one pointer +// register replaces the base + request-strided offset pair the c_mlr reads +// consumed before the table went static. All coefficient reads are +// lane-uniform 16-byte loads (__ldg is global-memory-only, so the shared +// block reads through plain loads — same values either way). |K| is computed +// once per (token, frequency) BEFORE the head loop so the GROUP heads share +// it from registers. +template +__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, float const* cRe, float const* cIm, + int64_t coff0, float const* cMlr, int64_t planeStride, uint4 re4, uint4 im4, float* accMean, float* accMlr, + float* accPos) { + static_assert(!(ROTATE_IN_CTA && USE_MAX), "in-CTA coefficient rotation exists only on the mean path"); float kRe[8], kIm[8], kMag[8]; unpackChunk8(re4, kRe); unpackChunk8(im4, kIm); @@ -213,8 +263,8 @@ __device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, in { if (o < a.numOffsets) { - float const* crp = a.cRe + o * planeStride + coff; - float const* cip = a.cIm + o * planeStride + coff; + float const* crp = cRe + o * planeStride + coff; + float const* cip = cIm + o * planeStride + coff; float4 const cr0 = __ldg(reinterpret_cast(crp)); float4 const cr1 = __ldg(reinterpret_cast(crp + 4)); float4 const ci0 = __ldg(reinterpret_cast(cip)); @@ -233,12 +283,26 @@ __device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, in } else { - float const* crp = a.cRe + coff; - float const* cip = a.cIm + coff; - float4 const cr0 = __ldg(reinterpret_cast(crp)); - float4 const cr1 = __ldg(reinterpret_cast(crp + 4)); - float4 const ci0 = __ldg(reinterpret_cast(cip)); - float4 const ci1 = __ldg(reinterpret_cast(cip + 4)); + float const* crp = cRe + coff; + float const* cip = cIm + coff; + float4 cr0, cr1, ci0, ci1; + if constexpr (ROTATE_IN_CTA) + { + // Shared-memory coefficient block written by this CTA's + // rotation prologue (16-byte aligned: numFreqs % 8 == 0 on + // the vectorized path). + cr0 = *reinterpret_cast(crp); + cr1 = *reinterpret_cast(crp + 4); + ci0 = *reinterpret_cast(cip); + ci1 = *reinterpret_cast(cip + 4); + } + else + { + cr0 = __ldg(reinterpret_cast(crp)); + cr1 = __ldg(reinterpret_cast(crp + 4)); + ci0 = __ldg(reinterpret_cast(cip)); + ci1 = __ldg(reinterpret_cast(cip + 4)); + } float const cre[8] = {cr0.x, cr0.y, cr0.z, cr0.w, cr1.x, cr1.y, cr1.z, cr1.w}; float const cim[8] = {ci0.x, ci0.y, ci0.z, ci0.w, ci1.x, ci1.y, ci1.z, ci1.w}; float t = accMean[hg]; @@ -265,9 +329,15 @@ __device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, in // minBlocksPerMultiprocessor = 7: tighter caps force ptxas into ~48-56 // registers with stack spills in the fully unrolled inner loops; 7 CTAs/SM // admits the ~72-register spill-free allocation this kernel was tuned at. -template +// +// ROTATE_IN_CTA (mean path only) prepends the shared-memory coefficient +// rotation prologue and points the mean coefficient reads at it; the +// standalone-rotation read path (ROTATE_IN_CTA == false) stays compiled for +// the max aggregation and for the unit tests' bit-equality reference leg. +template __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKernel(FoldedScoreParams a) { + static_assert(!(ROTATE_IN_CTA && USE_MAX), "in-CTA coefficient rotation exists only on the mean path"); int const seg = blockIdx.y; int const reqId = a.segRequestIds[seg]; int const seqLen = a.requestSeqLens[reqId]; @@ -303,6 +373,20 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne } int const layerId = a.segLayerIds[seg]; + // Mean-path in-CTA rotation: stage this CTA's GROUP x numFreqs + // coefficient block in shared memory (the early-out above is + // CTA-uniform, so the barrier inside is safe here and skipped CTAs do no + // rotation work). + float* sCRe = nullptr; + float* sCIm = nullptr; + if constexpr (ROTATE_IN_CTA) + { + extern __shared__ float coefficientSmem[]; + int const coefCount = GROUP * a.numFreqs; + sCRe = coefficientSmem; + sCIm = coefficientSmem + coefCount; + rotateCoefficientsIntoShared(a, reqId, layerId, headBase, coefCount, sCRe, sCIm); + } int const page = absT / a.tokensPerBlock; int const slot = absT - page * a.tokensPerBlock; // Threads past the sequence tail must not touch the page table (their @@ -322,8 +406,18 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne * (physPage * a.stridePage + static_cast(kvHead) * a.strideKvHead + static_cast(slot) * a.strideSlot); - int64_t const coff0 = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + // Coefficient source: the global per-round tables, or this CTA's shared + // block (whose head-block origin is index 0). + float const* coefRe = a.cRe; + float const* coefIm = a.cIm; + int64_t coff0 = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + static_cast(headBase) * a.numFreqs; + if constexpr (ROTATE_IN_CTA) + { + coefRe = sCRe; + coefIm = sCIm; + coff0 = 0; + } // The MLR coefficient is position independent, so its table is folded once // at initialization without a request axis: [layer, head, freq]. Hoist the // row pointer here (no request stride) so the fully unrolled chunk loop @@ -361,8 +455,8 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne { uint4 re4, im4; scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); - scoreComputeChunk( - a, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + scoreComputeChunk( + a, coefRe, coefIm, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); } } else @@ -372,8 +466,8 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne { uint4 re4, im4; scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); - scoreComputeChunk( - a, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); + scoreComputeChunk( + a, coefRe, coefIm, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); } } @@ -413,9 +507,12 @@ __global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKerne // stride set (any frequency count, any element stride, fp32 pools included). // The GQA head loop runs at runtime, so any group size is covered. |K| is // recomputed per head from the same loads — bit-identical to hoisting it. -template +// ROTATE_IN_CTA carries the same mean-path shared-memory rotation prologue +// as the vectorized kernel (here sized by the runtime GQA group). +template __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(FoldedScoreParams a) { + static_assert(!(ROTATE_IN_CTA && USE_MAX), "in-CTA coefficient rotation exists only on the mean path"); int const seg = blockIdx.y; int const reqId = a.segRequestIds[seg]; int const seqLen = a.requestSeqLens[reqId]; @@ -436,6 +533,21 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde int const kvHead = blockIdx.z; int const groupSize = a.numQueryHeads / a.numKvHeads; int const layerId = a.segLayerIds[seg]; + // Mean-path in-CTA rotation: stage this CTA's groupSize x numFreqs + // coefficient block in shared memory (the early-out above is + // CTA-uniform, so the barrier inside is safe here). + float const* coefRe = a.cRe; + float const* coefIm = a.cIm; + if constexpr (ROTATE_IN_CTA) + { + extern __shared__ float coefficientSmem[]; + int const coefCount = groupSize * a.numFreqs; + float* sCRe = coefficientSmem; + float* sCIm = coefficientSmem + coefCount; + rotateCoefficientsIntoShared(a, reqId, layerId, kvHead * groupSize, coefCount, sCRe, sCIm); + coefRe = sCRe; + coefIm = sCIm; + } int const page = absT / a.tokensPerBlock; int const slot = absT - page * a.tokensPerBlock; int encoded = 0; @@ -460,9 +572,13 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde for (int hg = 0; hg < groupSize; ++hg) { int const h = kvHead * groupSize + hg; - int64_t const coff - = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + int64_t coff = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs + static_cast(h) * a.numFreqs; + if constexpr (ROTATE_IN_CTA) + { + // The shared coefficient block's origin is this CTA's head block. + coff = static_cast(hg) * a.numFreqs; + } float const* const cml = cMlrRow + static_cast(hg) * a.numFreqs; float acc = 0.0f; float accMlr = 0.0f; @@ -497,7 +613,7 @@ __global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(Folde } else { - acc = fmaf(kRe, a.cRe[coff + f], fmaf(kIm, a.cIm[coff + f], fmaf(kMag, cml[f], acc))); + acc = fmaf(kRe, coefRe[coff + f], fmaf(kIm, coefIm[coff + f], fmaf(kMag, cml[f], acc))); } } if (store) @@ -606,6 +722,13 @@ __global__ void triFoldScoreCoefficientsKernel(float const* __restrict__ qReal, // (request, layer, head, freq) element. The host wrapper guarantees every // round start indexes inside the tables. // Design: Fanrong Li (torch-graph review, 2026-07-20). +// +// The production mean path now runs this rotation inside the score kernels' +// CTA prologue instead (rotateCoefficientsIntoShared, sharing +// rotateCoefficientPair with this kernel so the two paths stay bit-identical +// end to end). This standalone kernel remains as the unit tests' equality +// reference leg and as the fallback for coefficient blocks past the score +// launch's dynamic shared memory bound. __global__ void triRotateMeanScoreCoefficientsKernel(float const* __restrict__ qRealScaled, float const* __restrict__ qImagScaled, float const* __restrict__ phaseCos, float const* __restrict__ phaseSin, int32_t const* __restrict__ roundStarts, float* __restrict__ cRe, float* __restrict__ cIm, @@ -626,33 +749,52 @@ __global__ void triRotateMeanScoreCoefficientsKernel(float const* __restrict__ q int64_t const phaseIdx = static_cast(roundStarts[req]) * numFreqs + f; float const pc = phaseCos[phaseIdx]; float const ps = phaseSin[phaseIdx]; - float const qre = qRealScaled[cIdx]; - float const qim = qImagScaled[cIdx]; - cRe[idx] = qre * pc - qim * ps; - cIm[idx] = qim * pc + qre * ps; + // rotateCoefficientPair is the ONE rotation expression shared with the + // score kernels' in-CTA prologue — the bit-equality contract between the + // two mean-path preparation flavors. + float2 const c = rotateCoefficientPair(qRealScaled[cIdx], qImagScaled[cIdx], pc, ps); + cRe[idx] = c.x; + cIm[idx] = c.y; } +// The in-CTA rotation flavor adds one instantiation per (T, GROUP, chunk +// mode) — mean-only, so the instantiation count per element type grows from +// 4 to 6 per GROUP (binary-size cost of keeping the standalone-rotation read +// path compiled for the max aggregation and the unit tests' equality leg). template -void launchVectorized(FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useMax, cudaStream_t stream) +void launchVectorized( + FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useMax, bool rotateInCta, cudaStream_t stream) { bool const staticChunks = params.numFreqs == 64; int32_t const effectiveGroup = params.zIsQueryHead ? 1 : groupSize; + // Dynamic shared memory for the rotation prologue's coefficient block + // (re + im); zero on the other paths. + size_t const smemBytes + = rotateInCta ? sizeof(float) * 2 * static_cast(effectiveGroup) * params.numFreqs : 0; #define TRTLLM_TRI_SCORE_LAUNCH_GROUP(GROUP_V) \ do \ { \ if (staticChunks) \ { \ if (useMax) \ - triScoreVectorizedKernel<<>>(params); \ + triScoreVectorizedKernel<<>>(params); \ + else if (rotateInCta) \ + triScoreVectorizedKernel \ + <<>>(params); \ else \ - triScoreVectorizedKernel<<>>(params); \ + triScoreVectorizedKernel \ + <<>>(params); \ } \ else \ { \ if (useMax) \ - triScoreVectorizedKernel<<>>(params); \ + triScoreVectorizedKernel<<>>(params); \ + else if (rotateInCta) \ + triScoreVectorizedKernel \ + <<>>(params); \ else \ - triScoreVectorizedKernel<<>>(params); \ + triScoreVectorizedKernel \ + <<>>(params); \ } \ } while (0) switch (effectiveGroup) @@ -671,31 +813,40 @@ void launchVectorized(FoldedScoreParams const& params, int32_t groupSize, dim3 g } template -void launchScalar(FoldedScoreParams const& params, dim3 grid, bool useMax, cudaStream_t stream) +void launchScalar(FoldedScoreParams const& params, dim3 grid, bool useMax, bool rotateInCta, cudaStream_t stream) { + // The scalar rotation prologue sizes its shared coefficient block by the + // runtime GQA group (any group size is covered here). + size_t const smemBytes = rotateInCta + ? sizeof(float) * 2 * static_cast(params.numQueryHeads / params.numKvHeads) * params.numFreqs + : 0; if (useMax) { - triScoreScalarKernel<<>>(params); + triScoreScalarKernel<<>>(params); + } + else if (rotateInCta) + { + triScoreScalarKernel<<>>(params); } else { - triScoreScalarKernel<<>>(params); + triScoreScalarKernel<<>>(params); } } // Launch flavor for bf16/fp16 pools, the only element types owning both load // paths (the vectorized 16-byte chunk kernel and the strided scalar kernel). template -void launchVectorizedOrScalar( - FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useVectorized, bool useMax, cudaStream_t stream) +void launchVectorizedOrScalar(FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useVectorized, + bool useMax, bool rotateInCta, cudaStream_t stream) { if (useVectorized) { - launchVectorized(params, groupSize, grid, useMax, stream); + launchVectorized(params, groupSize, grid, useMax, rotateInCta, stream); } else { - launchScalar(params, grid, useMax, stream); + launchScalar(params, grid, useMax, rotateInCta, stream); } } @@ -704,10 +855,10 @@ void launchVectorizedOrScalar( // only the scalar load path knows how to read them). template void launchQuantizedScalar( - FoldedScoreParams const& params, dim3 grid, bool useVectorized, bool useMax, cudaStream_t stream) + FoldedScoreParams const& params, dim3 grid, bool useVectorized, bool useMax, bool rotateInCta, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); - launchScalar(params, grid, useMax, stream); + launchScalar(params, grid, useMax, rotateInCta, stream); } } // namespace @@ -742,7 +893,7 @@ void rotateMeanScoreCoefficientsLaunch(float const* qRealScaled, float const* qI } void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, - int32_t numSegments, bool useVectorized, bool useMax, cudaStream_t stream) + int32_t numSegments, bool useVectorized, bool useMax, bool rotateInCta, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(numSegments > 0 && numSegments <= 65535, "tri_attention_score: request*layer segment count exceeds the CUDA grid limit"); @@ -750,6 +901,19 @@ void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType "tri_attention_score: offset planes exceed the per-thread accumulator budget"); TLLM_CHECK_WITH_INFO(!useVectorized || (params.numFreqs % 8 == 0 && params.strideDim == 1), "tri_attention_score: vectorized path requires 8-frequency chunks with unit stride"); + TLLM_CHECK_WITH_INFO(!(rotateInCta && useMax), "tri_attention_score: in-CTA coefficient rotation is mean-only"); + TLLM_CHECK_WITH_INFO(!rotateInCta + || (params.phaseCos != nullptr && params.phaseSin != nullptr && params.qReS != nullptr + && params.qImS != nullptr && params.roundStarts != nullptr), + "tri_attention_score: in-CTA rotation inputs are missing"); + // Launches request dynamic shared memory without opting into the + // above-48KB attribute, so oversized coefficient blocks must fail loudly + // here; such geometries can still score through the standalone rotation + // launch (rotateInCta == false). The scalar prologue's runtime GQA group + // is the worst case (the vectorized per-query-head mapping uses 1). + TLLM_CHECK_WITH_INFO( + !rotateInCta || sizeof(float) * 2 * static_cast(groupSize) * params.numFreqs <= 48u * 1024u, + "tri_attention_score: in-CTA rotation coefficient block exceeds the 48KB dynamic shared memory bound"); // Tile count covers the decode span plus the worst-case page-alignment // slack (tokenStart may sit up to tokensPerBlock - 1 tokens into a page). auto const tiles = static_cast( @@ -759,21 +923,23 @@ void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType switch (poolType) { case PoolElementType::kBFloat16: - launchVectorizedOrScalar<__nv_bfloat16>(params, groupSize, grid, useVectorized, useMax, stream); + launchVectorizedOrScalar<__nv_bfloat16>(params, groupSize, grid, useVectorized, useMax, rotateInCta, stream); break; case PoolElementType::kHalf: - launchVectorizedOrScalar(params, groupSize, grid, useVectorized, useMax, stream); + launchVectorizedOrScalar(params, groupSize, grid, useVectorized, useMax, rotateInCta, stream); break; case PoolElementType::kFloat32: // fp32 pools have 32-byte 8-frequency rows; the 16-byte chunk path // does not apply, so they always take the strided scalar kernel. TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: fp32 pools must use the scalar path"); - launchScalar(params, grid, useMax, stream); + launchScalar(params, grid, useMax, rotateInCta, stream); break; case PoolElementType::kFloat8E4M3: - launchQuantizedScalar<__nv_fp8_e4m3>(params, grid, useVectorized, useMax, stream); + launchQuantizedScalar<__nv_fp8_e4m3>(params, grid, useVectorized, useMax, rotateInCta, stream); + break; + case PoolElementType::kInt8: + launchQuantizedScalar(params, grid, useVectorized, useMax, rotateInCta, stream); break; - case PoolElementType::kInt8: launchQuantizedScalar(params, grid, useVectorized, useMax, stream); break; } TLLM_CUDA_CHECK(cudaGetLastError()); } diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h index bcf96b0d0ede..43192f2f770d 100644 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h +++ b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h @@ -120,6 +120,15 @@ void foldScoreCoefficientsLaunch(float const* qReal, // [L_cal * HQ * F] // aggregation keeps foldScoreCoefficientsLaunch unchanged. Every round start // must lie in [0, maxPosition); the host wrapper enforces that loudly before // launch. Design: Fanrong Li (torch-graph review, 2026-07-20). +// +// The production mean path no longer runs this standalone launch either: the +// score kernels rebuild the same coefficients in an in-CTA shared-memory +// prologue (see FoldedScoreParams and foldedScoreLaunch's rotateInCta). The +// prologue compiles the SAME rotation arithmetic as this kernel, so the two +// preparation paths produce bit-identical scores; this launch stays as that +// equality proof's reference leg (the unit suite compares them with +// torch.equal) and as the fallback for coefficient blocks past the dynamic +// shared memory bound. void rotateMeanScoreCoefficientsLaunch(float const* qRealScaled, // [L_cal * HQ * F] float const* qImagScaled, // [L_cal * HQ * F] float const* phaseCos, // [maxPosition * F] @@ -142,12 +151,26 @@ struct FoldedScoreParams int32_t const* requestSeqLens; // [numRequests] int32_t* validWidthOut; // [numRequests] side-store: seqLen - tokenStart, once per request int32_t const* requestTokenStarts; // [numRequests] pinned prompt length = decode-region origin - float const* cRe; // per-round output (see foldScoreCoefficientsLaunch / rotateMeanScoreCoefficientsLaunch) + // Per-round coefficient planes (see foldScoreCoefficientsLaunch / + // rotateMeanScoreCoefficientsLaunch); nullptr when the mean-path kernels + // rotate their coefficients in-CTA instead (rotateInCta below). + float const* cRe; float const* cIm; float const* cMlr; // static, request-independent [L_cal, HQ, F] MLR table - float* out; // [segment, numQueryHeads, outputWidth] fp32 decode-only scores + // Mean-path in-CTA rotation inputs, consumed only by rotateInCta launches + // (nullptr otherwise). Table contract and arithmetic are EXACTLY those of + // rotateMeanScoreCoefficientsLaunch; round starts must be host-validated + // against the tabulated position extent BEFORE launch (the kernels gather + // phase rows unguarded — a device-side bound would need its own check + // kernel). + float const* phaseCos; // [maxPosition, F] + float const* phaseSin; // [maxPosition, F] + float const* qReS; // [L_cal, HQ, F] pre-scaled calibration query, real part + float const* qImS; // [L_cal, HQ, F] pre-scaled calibration query, imaginary part + int32_t const* roundStarts; // [numRequests], each in [0, maxPosition) + float* out; // [segment, numQueryHeads, outputWidth] fp32 decode-only scores int32_t outputWidth; - int32_t numLayers; // scored layers per request (the segment period) + int32_t numLayers; // scored layers per request (the segment period) int32_t numRequests; int32_t numCalibratedLayers; int32_t numQueryHeads; @@ -173,8 +196,17 @@ struct FoldedScoreParams // alignment); otherwise a fully strided scalar path runs the same math. // groupSize = numQueryHeads / numKvHeads must be 1, 2, 4, or 8 unless // params.zIsQueryHead maps grid.z to single query heads. +// +// rotateInCta (mean aggregation only) makes each score CTA rotate its own +// coefficient block from params.phaseCos/phaseSin/qReS/qImS/roundStarts into +// dynamic shared memory instead of reading pre-rotated global c_re/c_im +// planes, eliminating the standalone rotation launch and its per-round +// coefficient scratch. The prologue shares the standalone kernel's rotation +// arithmetic, so scores are bit-identical between the two preparation paths. +// The coefficient block (2 * group * numFreqs fp32) must fit the default +// 48KB dynamic shared memory bound — enforced loudly here. void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, - int32_t numSegments, bool useVectorized, bool useMax, cudaStream_t stream); + int32_t numSegments, bool useVectorized, bool useMax, bool rotateInCta, cudaStream_t stream); } // namespace kernels::tri_attention_score diff --git a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp index c34f3600af1d..61a13dedd5e1 100644 --- a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp +++ b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp @@ -23,9 +23,11 @@ namespace tk = tensorrt_llm::kernels::tri_attention_score; // One op per kernel launch, matching the file-level granularity of sibling // kernel wrappers: the per-round coefficient preparation (the rotation op on -// the mean path, the trigonometric fold on the max path) writes persistent -// buffers whose plane count depends on the aggregation mode, while the score -// launch consumes those buffers; callers time and re-plan them independently. +// the legacy mean path, the trigonometric fold on the max path) writes +// persistent buffers whose plane count depends on the aggregation mode, while +// the score launch consumes those buffers; callers time and re-plan them +// independently. The default mean path collapses to the score op alone (its +// kernels rotate coefficients in-CTA), so no preparation op runs there. namespace { @@ -154,6 +156,12 @@ void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, // kv_scales positivity check above, the bounds reduction costs one small // host sync per eviction round. Design: Fanrong Li (torch-graph review, // 2026-07-20). +// +// The production mean path no longer calls this op: the paged score op below +// performs the same rotation in its kernels' CTA prologue (bit-identically — +// the two paths share one device rotation expression). This op stays as the +// unit tests' equality reference leg and as the fallback for coefficient +// blocks past the score launch's dynamic shared memory bound. void triAttentionRotateMeanScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, torch::Tensor q_real_scaled, torch::Tensor q_imag_scaled, torch::Tensor phase_cos, torch::Tensor phase_sin, torch::Tensor round_starts, int64_t num_requests, int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_freqs, @@ -199,14 +207,30 @@ void triAttentionRotateMeanScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor // the kernel reads all layers through layer_base_addrs (V2 exposes each layer // as its own storage), and the anchor only supplies their common element type // and the device; its data is never read through this argument. +// +// Mean-path coefficient preparation comes in two flavors. Passing phase_cos/ +// phase_sin/q_real_scaled/q_imag_scaled/round_starts (+ max_position) makes +// the score kernels rotate their own coefficient block in an in-CTA prologue +// — c_re/c_im must then be omitted (no per-round coefficient scratch exists +// at all). Passing c_re/c_im instead scores against pre-rotated global planes +// (the rotation op above, or the fold op on the max path). The two mean +// flavors share one rotation expression device-side, so their scores are +// bit-identical. The round-start bounds contract of the rotation op moves +// here on the in-CTA flavor: a start at or past max_position would gather +// past the phase tables, and a device-side guard would need its own check +// kernel, so the same host-side aminmax reduction (one small sync per +// eviction round) runs before the launch. void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_base_addrs, torch::Tensor block_offsets, torch::Tensor seg_page_offsets, torch::Tensor seg_request_ids, torch::Tensor seg_layer_ids, - torch::Tensor request_seq_lens, torch::Tensor valid_widths, torch::Tensor request_token_starts, torch::Tensor c_re, - torch::Tensor c_im, torch::Tensor c_mlr, torch::Tensor out, int64_t output_width, int64_t num_layers, - int64_t num_requests, int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_kv_heads, - int64_t num_freqs, int64_t tokens_per_block, int64_t kv_factor, int64_t num_offsets, int64_t stride_page, - int64_t stride_kv_head, int64_t stride_slot, int64_t stride_dim, int64_t num_segments, bool use_max, - bool use_vectorized, std::optional kv_scales) + torch::Tensor request_seq_lens, torch::Tensor valid_widths, torch::Tensor request_token_starts, + std::optional c_re, std::optional c_im, torch::Tensor c_mlr, torch::Tensor out, + int64_t output_width, int64_t num_layers, int64_t num_requests, int64_t num_calibrated_layers, + int64_t num_query_heads, int64_t num_kv_heads, int64_t num_freqs, int64_t tokens_per_block, int64_t kv_factor, + int64_t num_offsets, int64_t stride_page, int64_t stride_kv_head, int64_t stride_slot, int64_t stride_dim, + int64_t num_segments, bool use_max, bool use_vectorized, std::optional kv_scales, + std::optional phase_cos, std::optional phase_sin, + std::optional q_real_scaled, std::optional q_imag_scaled, + std::optional round_starts, int64_t max_position) { TORCH_CHECK(use_max || num_offsets == 1, "tri_attention_paged_score: mean aggregation consumes exactly one folded coefficient plane"); @@ -218,8 +242,6 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas checkContiguousCuda(request_seq_lens, at::kInt, "int32", "request_seq_lens"); checkContiguousCuda(valid_widths, at::kInt, "int32", "valid_widths"); checkContiguousCuda(request_token_starts, at::kInt, "int32", "request_token_starts"); - checkContiguousCuda(c_re, at::kFloat, "fp32", "c_re"); - checkContiguousCuda(c_im, at::kFloat, "fp32", "c_im"); checkContiguousCuda(c_mlr, at::kFloat, "fp32", "c_mlr"); checkContiguousCuda(out, at::kFloat, "fp32", "out"); TORCH_CHECK(pool_anchor.is_cuda(), "tri_attention_paged_score: pool anchor must be a CUDA tensor"); @@ -247,9 +269,46 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas // c_mlr is request independent: the score kernels index it as one static // [layer, head, freq] table, so only the calibration extent is required. int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; - TORCH_CHECK( - c_re.numel() >= num_offsets * total && c_im.numel() >= num_offsets * total && c_mlr.numel() >= calibration, - "tri_attention_paged_score: folded coefficient buffers are undersized"); + bool const rotateInCta = phase_cos.has_value() || phase_sin.has_value() || q_real_scaled.has_value() + || q_imag_scaled.has_value() || round_starts.has_value(); + if (rotateInCta) + { + TORCH_CHECK(!use_max, "tri_attention_paged_score: in-CTA coefficient rotation is mean-only"); + TORCH_CHECK(phase_cos.has_value() && phase_sin.has_value() && q_real_scaled.has_value() + && q_imag_scaled.has_value() && round_starts.has_value(), + "tri_attention_paged_score: in-CTA rotation requires phase_cos, phase_sin, q_real_scaled, " + "q_imag_scaled, and round_starts together"); + TORCH_CHECK(!c_re.has_value() && !c_im.has_value(), + "tri_attention_paged_score: in-CTA rotation reads no c_re/c_im planes; omit them"); + checkContiguousCuda(*phase_cos, at::kFloat, "fp32", "phase_cos"); + checkContiguousCuda(*phase_sin, at::kFloat, "fp32", "phase_sin"); + checkContiguousCuda(*q_real_scaled, at::kFloat, "fp32", "q_real_scaled"); + checkContiguousCuda(*q_imag_scaled, at::kFloat, "fp32", "q_imag_scaled"); + checkContiguousCuda(*round_starts, at::kInt, "int32", "round_starts"); + TORCH_CHECK(max_position > 0, "tri_attention_paged_score: in-CTA rotation requires a positive max_position"); + TORCH_CHECK(phase_cos->numel() >= max_position * num_freqs && phase_sin->numel() >= max_position * num_freqs, + "tri_attention_paged_score: phase tables do not cover max_position rows"); + TORCH_CHECK(q_real_scaled->numel() >= calibration && q_imag_scaled->numel() >= calibration, + "tri_attention_paged_score: scaled calibration tensors are undersized"); + TORCH_CHECK(round_starts->numel() >= num_requests, "tri_attention_paged_score: round_starts are undersized"); + // Same host-side bounds reduction as the standalone rotation op (one + // small sync per eviction round): the kernels gather phase rows + // unguarded, so an out-of-range start must fail loudly here. + auto const [minStart, maxStart] = round_starts->narrow(0, 0, num_requests).aminmax(); + TORCH_CHECK(minStart.item() >= 0 && maxStart.item() < max_position, + "tri_attention_paged_score: a round start lies outside the tabulated position range [0, ", max_position, + ")"); + } + else + { + TORCH_CHECK(c_re.has_value() && c_im.has_value(), + "tri_attention_paged_score: pre-rotated scoring requires c_re and c_im"); + checkContiguousCuda(*c_re, at::kFloat, "fp32", "c_re"); + checkContiguousCuda(*c_im, at::kFloat, "fp32", "c_im"); + TORCH_CHECK(c_re->numel() >= num_offsets * total && c_im->numel() >= num_offsets * total, + "tri_attention_paged_score: folded coefficient buffers are undersized"); + } + TORCH_CHECK(c_mlr.numel() >= calibration, "tri_attention_paged_score: folded coefficient buffers are undersized"); TORCH_CHECK(out.numel() >= num_segments * num_query_heads * output_width, "tri_attention_paged_score: score output buffer is undersized"); @@ -313,9 +372,14 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas params.requestSeqLens = request_seq_lens.data_ptr(); params.validWidthOut = valid_widths.data_ptr(); params.requestTokenStarts = request_token_starts.data_ptr(); - params.cRe = c_re.data_ptr(); - params.cIm = c_im.data_ptr(); + params.cRe = rotateInCta ? nullptr : c_re->data_ptr(); + params.cIm = rotateInCta ? nullptr : c_im->data_ptr(); params.cMlr = c_mlr.data_ptr(); + params.phaseCos = rotateInCta ? phase_cos->data_ptr() : nullptr; + params.phaseSin = rotateInCta ? phase_sin->data_ptr() : nullptr; + params.qReS = rotateInCta ? q_real_scaled->data_ptr() : nullptr; + params.qImS = rotateInCta ? q_imag_scaled->data_ptr() : nullptr; + params.roundStarts = rotateInCta ? round_starts->data_ptr() : nullptr; params.out = out.data_ptr(); params.outputWidth = static_cast(output_width); params.numLayers = static_cast(num_layers); @@ -335,7 +399,7 @@ void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_bas auto stream = at::cuda::getCurrentCUDAStream(); tk::foldedScoreLaunch( - params, poolType, groupSize, static_cast(num_segments), use_vectorized, use_max, stream); + params, poolType, groupSize, static_cast(num_segments), use_vectorized, use_max, rotateInCta, stream); } } // anonymous namespace @@ -360,17 +424,24 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "int num_requests, int num_calibrated_layers, " "int num_query_heads, int num_freqs, int max_position) -> ()"); + // c_re/c_im are optional: the mean path's default in-CTA coefficient + // rotation (phase_cos .. round_starts + max_position, all-or-nothing) + // replaces them entirely; the max path and the legacy mean path keep + // passing pre-rotated planes. m.def( "tri_attention_paged_score(" "Tensor pool_anchor, Tensor layer_base_addrs, Tensor block_offsets, " "Tensor seg_page_offsets, Tensor seg_request_ids, Tensor seg_layer_ids, " "Tensor request_seq_lens, Tensor(a!) valid_widths, Tensor request_token_starts, " - "Tensor c_re, Tensor c_im, Tensor c_mlr, Tensor(b!) out, " + "Tensor? c_re, Tensor? c_im, Tensor c_mlr, Tensor(b!) out, " "int output_width, int num_layers, int num_requests, int num_calibrated_layers, " "int num_query_heads, int num_kv_heads, int num_freqs, int tokens_per_block, " "int kv_factor, int num_offsets, int stride_page, int stride_kv_head, " "int stride_slot, int stride_dim, int num_segments, bool use_max, bool use_vectorized, " - "Tensor? kv_scales=None) -> ()"); + "Tensor? kv_scales=None, " + "Tensor? phase_cos=None, Tensor? phase_sin=None, " + "Tensor? q_real_scaled=None, Tensor? q_imag_scaled=None, " + "Tensor? round_starts=None, int max_position=0) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 6cdf5ea985bf..13dbab12da47 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -786,8 +786,9 @@ def launch_prepared_score(self) -> torch.Tensor: # mean_cos/mean_sin feed ONLY the opt-in CuTe score runner, whose # compiled kernel captured their device pointers, so they must be # refreshed before it launches. The default C++ mean path rotates - # init-time phase tables inside its coefficient op instead, so - # production rounds launch zero phase kernels. + # init-time phase tables inside the score kernels' own CTA + # prologue instead, so production rounds launch zero phase or + # coefficient kernels. prepare_mean_phase( self.round_starts_device, self.offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 7f3a220707fc..776b8d5bfc2c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -5,9 +5,10 @@ The production path uses one fixed-shape trig-score launch across all dense layers, CuTE-DSL TopK selection, and grouped C++ compaction. This module owns the score launcher and its persistent metadata; scoring itself runs through -the compiled ``trtllm`` CUDA ops (per-round coefficient rotation over -init-time phase tables on the mean path, trigonometric coefficient fold on -the max path, then the folded paged score) for every supported geometry. The +the compiled ``trtllm`` CUDA ops (on the mean path ONE folded paged-score +launch whose CTAs rotate their own coefficients from init-time phase tables; +on the max path a trigonometric coefficient fold, then the folded paged +score) for every supported geometry. The original Triton score kernel has been deleted; the unit tests validate the CUDA ops against an independent PyTorch oracle. Selection and compaction live in their respective runtime modules. An @@ -118,6 +119,7 @@ def _launch_tri_score_perhead( token_starts_device: torch.Tensor, *, score_aggregation: str, + mean_rotate_in_cta: bool = True, ) -> None: """Prepare the per-round coefficients, then score paged KV via the C++ ops. @@ -125,6 +127,15 @@ def _launch_tri_score_perhead( geometry this launcher accepts; unsupported inputs fail loudly inside the ops (TORCH_CHECK) instead of routing to another kernel. The unit tests validate them against an independent PyTorch oracle. + + ``mean_rotate_in_cta`` (mean aggregation only, default on) hands the + per-round coefficient rotation to the score kernels' own CTA prologue: + one launch per round, zero coefficient scratch. The two-launch + preparation (standalone rotation kernel writing global tables the score + kernels then read) stays selectable because the unit tests prove the two + paths produce bit-identical scores — the kernels share one rotation + expression — and because geometries whose coefficient block exceeds the + score launch's shared-memory bound must fall back to it. """ if score_aggregation not in ("mean", "max"): raise ValueError(f"unsupported score aggregation: {score_aggregation}") @@ -155,8 +166,11 @@ def _launch_tri_score_perhead( # keeps one c_re/c_im plane per offset because max does not commute # through the frequency sum. offset_planes = num_offsets if use_max else 1 - c_re, c_im, c_mlr = group._fold_coefficient_buffers(offset_planes, with_mlr=use_max) + # Keyword operands of the score op that trigger its in-CTA coefficient + # rotation; empty when pre-rotated c_re/c_im planes are passed instead. + rotate_in_cta_kwargs: dict = {} if use_max: + c_re, c_im, c_mlr = group._fold_coefficient_buffers(offset_planes, with_mlr=True) q_real, q_imag, mlr_coef = group.pointer_middle freq_scale_sq, omega, offsets = group.pointer_tail torch.ops.trtllm.tri_attention_fold_score_coefficients( @@ -183,12 +197,32 @@ def _launch_tri_score_perhead( # below reads raw quantized elements at zero hot-loop cost. group._kv_scales, ) + elif mean_rotate_in_cta: + # Production mean path: the score kernels rotate the pre-scaled + # calibration query by the tabulated offset-mean phase of each + # request's round start in their own CTA prologue (tables built once + # at group construction; phase-table design: Fanrong Li, torch-graph + # review 2026-07-20). No rotation launch, no per-round coefficient + # scratch; c_mlr is the static init-time table. The prologue compiles + # the standalone rotation kernel's exact arithmetic, so scores are + # bit-identical to the two-launch path below. + c_re = None + c_im = None + c_mlr = group._mlr_fold + rotate_in_cta_kwargs = dict( + phase_cos=group._phase_cos, + phase_sin=group._phase_sin, + q_real_scaled=group._q_real_scaled, + q_imag_scaled=group._q_imag_scaled, + round_starts=round_starts_device, + max_position=group._max_position, + ) else: - # Mean aggregation rotates the pre-scaled calibration query by the - # tabulated offset-mean phase of each request's round start (tables - # built once at group construction; design: Fanrong Li, torch-graph - # review 2026-07-20). c_mlr is the static init-time table, so the - # round writes only c_re/c_im and runs zero trigonometry. + # Two-launch mean preparation: the standalone rotation kernel writes + # global c_re/c_im planes the score kernels then read. Kept callable + # as the unit tests' bit-equality reference leg and as the fallback + # for coefficient blocks past the score launch's shared-memory bound. + c_re, c_im, _ = group._fold_coefficient_buffers(offset_planes, with_mlr=False) c_mlr = group._mlr_fold torch.ops.trtllm.tri_attention_rotate_mean_score_coefficients( c_re, @@ -242,6 +276,7 @@ def _launch_tri_score_perhead( # fold above or the init-time pre-scaled mean tables already consumed # the values). group._kv_scales, + **rotate_in_cta_kwargs, ) @@ -384,7 +419,10 @@ def __init__( # (kv_scale * freq_scale_sq * mlr) is fully static: the request axis # disappears and nothing about it is recomputed per round. Every # eviction round then gathers one table row per request and rotates - # the static query by it -- zero trigonometry at runtime. + # the static query by it -- zero trigonometry at runtime. The score + # kernels perform that rotation themselves in an in-CTA prologue by + # default; the standalone rotation kernel remains the two-launch + # reference flavor (bit-identical scores, proven by the unit tests). # # Positions can legitimately reach the full sequence capacity, so the # table covers [0, seq_len] inclusive. The 64-row floor keeps tiny @@ -604,9 +642,11 @@ def _fold_coefficient_buffers( Sized on ``max_requests`` so any launch's active ``request_count`` fits without reallocation (each launch prepares only its active - rows). Only the max aggregation still writes a per-round c_mlr; the - mean path reads the static init-time MLR table instead (request axis - removed), so its scratch skips c_mlr entirely. + rows). The production mean path allocates NOTHING here: its score + kernels rotate coefficients in-CTA, so only the max aggregation + (which also writes a per-round c_mlr) and the two-launch mean + reference path (static MLR table, c_re/c_im planes only) ever call + this. """ key = (offset_planes, with_mlr) buffers = self._fold_buffers.get(key) @@ -635,8 +675,16 @@ def launch( mean_cos: torch.Tensor, mean_sin: torch.Tensor, score_aggregation: str, + *, + mean_rotate_in_cta: bool = True, ) -> torch.Tensor: - """Return decode-only scores as ``[request, layer, head, token]``.""" + """Return decode-only scores as ``[request, layer, head, token]``. + + ``mean_rotate_in_cta=False`` selects the two-launch mean coefficient + preparation (standalone rotation kernel + global-table score reads); + see ``_launch_tri_score_perhead``. Scores are bit-identical either + way — the unit tests compare the two with ``torch.equal``. + """ if request_count <= 0 or request_count > self.max_requests: raise ValueError("request count exceeds fixed score capacity") if ( @@ -699,6 +747,7 @@ def launch( round_starts_device, token_starts_device, score_aggregation=score_aggregation, + mean_rotate_in_cta=mean_rotate_in_cta, ) return output diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py index 87ef6eb816f1..3274a3336646 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py @@ -1,15 +1,22 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Mean-path phase-table rotation vs a direct trigonometric fold. +"""Mean-path phase-table rotation: numerical fidelity and path equality. The mean-aggregation coefficient preparation tabulates the offset-averaged phase of every possible round-start position once at initialization (float64 accumulation, stored fp32) and rotates the pre-scaled calibration query by -one gathered table row per request. This test rebuilds the same coefficients -with direct float64 trigonometry at the exact round starts and compares. +one gathered table row per request. The first test rebuilds the same +coefficients with direct float64 trigonometry at the exact round starts and +compares. The second proves the two places that rotation can run — the +standalone rotation kernel writing global coefficient planes, and the score +kernels' own in-CTA shared-memory prologue (the production default) — yield +BIT-IDENTICAL paged-score outputs, because both compile one shared rotation +expression. """ +import pytest import torch +from test_triattention_score_ops import _build_case as _build_score_case from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _FixedScoreGroup, @@ -104,3 +111,126 @@ def test_phase_table_rotation_matches_direct_trig_fold(): shape = (len(positions), num_layers, num_q_heads, num_freqs) torch.testing.assert_close(c_re.view(shape), reference_re, rtol=1e-5, atol=1e-5) torch.testing.assert_close(c_im.view(shape), reference_im, rtol=1e-5, atol=1e-5) + + +# One geometry per code path the in-CTA rotation prologue compiles into: +# both CUDA load paths (16-byte vectorized chunks and strided scalar), GQA +# groups 2/4/8 with dedicated template instantiations, group 3 through the +# per-query-head grid mapping, and every float pool dtype. The shared +# builder's per-request sequence lengths are ragged, so the write mask is +# compared too (sentinel-filled outputs). +_EQUALITY_CASES = [ + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=8, + num_kv_heads=2, + dtype=torch.bfloat16, + ), + id="vectorized_bf16_group4", + ), + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=8, + num_kv_heads=1, + dtype=torch.bfloat16, + ), + id="vectorized_bf16_group8", + ), + pytest.param( + dict( + head_dim=128, + tokens_per_block=32, + num_q_heads=6, + num_kv_heads=2, + dtype=torch.float16, + ), + id="vectorized_fp16_group3_per_query_head", + ), + pytest.param( + dict( + head_dim=8, + tokens_per_block=4, + num_q_heads=4, + num_kv_heads=2, + dtype=torch.bfloat16, + ), + id="scalar_bf16_group2", + ), + pytest.param( + dict( + head_dim=8, + tokens_per_block=4, + num_q_heads=6, + num_kv_heads=2, + dtype=torch.float32, + ), + id="scalar_fp32_runtime_group3", + ), +] + + +@pytest.mark.parametrize("case", _EQUALITY_CASES) +def test_in_cta_rotation_scores_bit_equal_to_standalone_rotation(case): + """Full mean-path score outputs are torch.equal between rotation flavors. + + Bit-equality (not a tolerance) is the contract: the score kernels' in-CTA + prologue and the standalone rotation kernel share one rotation + expression, and the score accumulation downstream is identical, so any + single-ulp drift means the arithmetic diverged and must fail here. + """ + assert hasattr(torch.ops.trtllm, "tri_attention_paged_score"), ( + "TriAttention paged score op is not loaded" + ) + assert hasattr(torch.ops.trtllm, "tri_attention_rotate_mean_score_coefficients"), ( + "TriAttention rotation op is not loaded" + ) + request_count = 3 + ( + group, + round_starts, + token_starts, + valid_seq_lens, + _, + mean_cos, + mean_sin, + _, + ) = _build_score_case( + request_count=request_count, + max_requests=4, + num_layers=2, + page_count=4, + prompt_len=5, + seed=20260723, + offsets=[1.0, 2.0, 4.0], + **case, + ) + device = group.output.device + sentinel = -54321.0 + + def run(mean_rotate_in_cta: bool) -> "tuple[torch.Tensor, torch.Tensor]": + group.output.fill_(sentinel) + valid_widths = torch.zeros(request_count, dtype=torch.int32, device=device) + scores = group.launch( + request_count, + valid_seq_lens, + valid_widths, + round_starts, + token_starts, + mean_cos, + mean_sin, + "mean", + mean_rotate_in_cta=mean_rotate_in_cta, + ).clone() + return scores, valid_widths + + # Reference leg: the kernel-round2 preparation (standalone rotation + # kernel + score kernels reading the global coefficient planes). + reference_scores, reference_widths = run(mean_rotate_in_cta=False) + in_cta_scores, in_cta_widths = run(mean_rotate_in_cta=True) + assert not reference_scores.eq(sentinel).all(), "reference leg scored nothing" + assert torch.equal(in_cta_scores, reference_scores) + assert torch.equal(in_cta_widths, reference_widths) From accceea9cf42ed40dac174d9e8c56dc3c0d309f1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 19:33:33 -0700 Subject: [PATCH 055/178] [None][chore] Move the tie-settlement reference kernel into its unit test The standalone tie-settlement kernel has no production caller since the fused settle-and-pack kernel became the only launched settle path; it only served as the bit-equality reference of the fused-kernel unit test. Keep it verbatim inside that test so the production module ships only kernels it launches. The standalone move-source packing kernel stays: the draft-model flow still launches it on pre-settled ordinals. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 89 +----------------- .../test_triattention_fused_settle_pack.py | 91 ++++++++++++++++++- 2 files changed, 92 insertions(+), 88 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 776b8d5bfc2c..8318020a87ea 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -1026,88 +1026,6 @@ def prepare_per_head_scores( ) -@triton.jit -def _settle_ties_after_topk_kernel( - scores, - seq_lens, - prompt_offsets, - provisional_indices, - output_indices, - WIDTH: tl.constexpr, - KEEP_COUNT: tl.constexpr, - OUTPUT_WIDTH: tl.constexpr, - BLOCK: tl.constexpr, -): - """Resolve boundary ties and emit increasing physical token indices. - - The production selectors launch the fused settle-and-pack kernel below; - this standalone version is kept as the reference the unit tests compare - the fused kernel against. - """ - row = tl.program_id(0) - row_scores = scores + row * WIDTH - row_selected = provisional_indices + row * KEEP_COUNT - row_output = output_indices + row * OUTPUT_WIDTH - # Scores are decode-relative; this row's pinned prompt length rebases the - # emitted ordinals to absolute positions (per row, so one launch may mix - # prompt lengths). - prompt_len = tl.load(prompt_offsets + row) - - 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, - ) - selected_score = tl.load( - row_scores + token_index, - mask=selected_mask, - other=float("inf"), - ).to(tl.float32) - threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) - - seq_len = tl.load(seq_lens + 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 < seq_len) - 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 < seq_len) - 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_len, - mask=selected, - ) - output_count += tl.sum(selected_i32) - ties_seen += tl.sum(tied_i32) - - # --------------------------------------------------------------------------- # # Compaction: pack the kept ordinals into per-request move indices. # # --------------------------------------------------------------------------- # @@ -1214,8 +1132,8 @@ def _settle_ties_and_pack_compaction_sources_kernel( ): """Settle one selection row's ties, then pack its compaction move sources. - One program per (request, selection row). The first half repeats - ``_settle_ties_after_topk_kernel`` verbatim: recover the top-k threshold + One program per (request, selection row). The first half settles the + provisional top-k: recover the top-k threshold from the provisional selection, count the strictly greater scores, then emit the kept ordinals in increasing order, rebased by the row's pinned prompt length. With ``HAS_PACK`` the same program then continues with the @@ -1225,7 +1143,8 @@ def _settle_ties_and_pack_compaction_sources_kernel( same conditions as the standalone kernel. Union selection has one row per request feeding every KV head's packed row, so that single program writes all of them. ``HAS_PACK=False`` compiles the second half away, leaving - exactly the standalone settle. Fusing the two launches was suggested by + exactly the settle stage (its pre-fusion standalone copy lives in the + fused-kernel unit test as the bit-equality reference). Fusing the two launches was suggested by Fanrong Li (torch-graph review 2026-07-20). """ request = tl.program_id(0) 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 index 816074f669ac..ad3c64f81245 100644 --- 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 @@ -4,8 +4,10 @@ """The fused settle-and-pack kernel must reproduce the original two-kernel sequence byte for byte. -The original tie settlement and move-source packing kernels stay in the -module as the reference here: every case launches them on one set of buffers +The reference legs are the pre-fusion kernels: the tie-settlement copy +kept in this file (its production original was deleted once the fused +kernel became the only launched settle path) and the move-source packing +kernel the module still ships for the draft flow: every case launches them on one set of buffers and the fused kernel on an identically initialized set, then requires ``torch.equal`` on the kept ordinals, the dense move sources, and the SWA move sources -- including the buffer regions neither path overwrites (rows @@ -15,6 +17,8 @@ import pytest import torch +import triton +import triton.language as tl from conftest import encode_block_offsets as _encode_block_offsets from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( @@ -25,10 +29,91 @@ ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _pack_compaction_sources_kernel, - _settle_ties_after_topk_kernel, _settle_ties_and_pack_compaction_sources_kernel, ) + +@triton.jit +def _settle_ties_after_topk_kernel( + scores, + seq_lens, + prompt_offsets, + provisional_indices, + output_indices, + WIDTH: tl.constexpr, + KEEP_COUNT: tl.constexpr, + OUTPUT_WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + """Resolve boundary ties and emit increasing physical token indices. + + Pre-fusion standalone kept verbatim as the fused kernel's bit-equality + reference; the production module ships only the fused launch. + """ + row = tl.program_id(0) + row_scores = scores + row * WIDTH + row_selected = provisional_indices + row * KEEP_COUNT + row_output = output_indices + row * OUTPUT_WIDTH + # Scores are decode-relative; this row's pinned prompt length rebases the + # emitted ordinals to absolute positions (per row, so one launch may mix + # prompt lengths). + prompt_len = tl.load(prompt_offsets + row) + + 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, + ) + selected_score = tl.load( + row_scores + token_index, + mask=selected_mask, + other=float("inf"), + ).to(tl.float32) + threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) + + seq_len = tl.load(seq_lens + 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 < seq_len) + 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 < seq_len) + 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_len, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) + + _BLOCK = 256 _NUM_WARPS = 4 From 6e2435596ab56126edb0a07f4b5151044ba15dba Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 19:52:25 -0700 Subject: [PATCH 056/178] [None][chore] Retire the register-staging sparse-KV compaction fallback The pipelined bf16 kernels won the A/B comparison everywhere and are the only path production dispatches, so the register-staging adapter behind them (layered buffer struct, launcher, dtype instantiations, and the half/float torch-op branches) is retired. Unsupported pool dtypes and geometries now fail loudly instead of silently degrading to a slower kernel. The unit tests move onto supported 32-token-page geometry and assert the rejection behavior. Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_float_float.cu | 1 - .../unfusedAttentionKernels_2_half_half.cu | 1 - .../unfusedAttentionKernels_2_template.h | 133 ++---------------- .../thop/sparseKvCacheCompactOp.cpp | 23 +-- .../triattention/compaction.py | 5 +- .../serial/test_sparse_kv_cache_compact.py | 54 +++---- 6 files changed, 46 insertions(+), 171 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu index 4b413712d787..4a150117540b 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu @@ -25,7 +25,6 @@ namespace kernels INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(float, float, KVBlockArray); INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(float, float, KVLinearBuffer); -INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(float); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu index 943207c2e589..a6fcb6dc36ad 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu @@ -25,7 +25,6 @@ namespace kernels INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(half, half, KVBlockArray); INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(half, half, KVLinearBuffer); -INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(half); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index 7482b0fc1882..a1a290549c20 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1756,70 +1756,6 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams //////////////////////////////////////////////////////////////////////////////////////////////////// -//! Layered KVCacheManagerV2 layout policy for the existing sparse-KV updater. -//! blockIdx.x selects an independent per-layer pool view. Layers in this launch -//! share one V2 block-offset table and the production FMHA updater copy loop. -struct KvCacheV2LayersBuffer -{ - int64_t const* poolPointers; - int32_t const* pageTable; - int32_t const* sourceLayerIndices; - int64_t sourceLayerStride; - // Head-plane stride of the packed move-source indices. The move buffers - // may be wider than one round's move count, so the stride comes from the - // host-side allocation width, not from the last source-offsets entry. - int64_t sourceHeadStride; - int64_t pageTableRequestStride; - int32_t tokensPerBlock; - int32_t const* destinationBases; - size_t bytesPerPage; - size_t bytesPerKvHalf; - - __device__ __forceinline__ uint8_t* getPool() const - { - return reinterpret_cast(static_cast(poolPointers[blockIdx.x])); - } - - // KVCacheManagerV2 block offsets encode page and K/V plane as - // 2*page + plane (K = 2p, V = 2p + 1); this table carries the K plane, - // so dividing by 2 recovers the page for both halves. - __device__ __forceinline__ void* getKBlockPtr(int32_t batchIdx, int32_t tokenIdx) const - { - int32_t const blockOffset = pageTable[batchIdx * pageTableRequestStride + tokenIdx / tokensPerBlock]; - int32_t const page = blockOffset / 2; - return getPool() + static_cast(page) * bytesPerPage; - } - - __device__ __forceinline__ void* getVBlockPtr(int32_t batchIdx, int32_t tokenIdx) const - { - int32_t const blockOffset = pageTable[batchIdx * pageTableRequestStride + tokenIdx / tokensPerBlock]; - int32_t const page = blockOffset / 2; - return getPool() + static_cast(page) * bytesPerPage + bytesPerKvHalf; - } - - __device__ __forceinline__ int32_t getKVLocalIdx( - int32_t tokenIdx, int32_t headIdx, int32_t valuesPerHead, int32_t headValueIdx) const - { - return (headIdx * tokensPerBlock + (tokenIdx % tokensPerBlock)) * valuesPerHead + headValueIdx; - } - - __device__ __forceinline__ int32_t getSparseKvSourceToken( - int32_t const* sourceIndices, int32_t headIdx, int32_t globalMove) const - { - int32_t const layer - = sourceLayerIndices == nullptr ? static_cast(blockIdx.x) : sourceLayerIndices[blockIdx.x]; - int64_t const offset = static_cast(layer) * sourceLayerStride - + static_cast(headIdx) * sourceHeadStride + globalMove; - return sourceIndices[offset]; - } - - __device__ __forceinline__ int32_t getSparseKvDestinationToken(int32_t batchIdx, int32_t requestMove) const - { - // Per-request landing positions: cohorts may mix prompt lengths. - return destinationBases[batchIdx] + requestMove; - } -}; - #ifdef ENABLE_BF16 // Optimized bf16 compaction fast path, ported from Fanrong Li's optimized @@ -1932,7 +1868,7 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) return; } - // Layer resolution matches KvCacheV2LayersBuffer::getSparseKvSourceToken: + // Layer resolution rule shared with the packed move-source layout: // without an explicit map, launch layer i reads source plane i (the flat // layout passes sourceLayerStride == 0, which collapses the term). int32_t const sourceLayer = params.sourceLayerIndices == nullptr ? layerIdx : params.sourceLayerIndices[layerIdx]; @@ -2432,28 +2368,6 @@ void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams } } -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -void launchSparseKvCacheCompactV2Layers( - QKVPreprocessingParams params, int32_t numLayers, cudaStream_t stream) -{ - constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / 16; - constexpr int32_t kDefaultSharedMemoryBytes = 48 * 1024; - constexpr int32_t kSharedBytesPerToken = 2 * kVectorsPerHead * sizeof(uint4); - constexpr bool kUseRegisterStaging = (HeadDim == 64 || HeadDim == 128) && sizeof(T) == 2; - constexpr int32_t kVectorThreads = kUseRegisterStaging ? kVectorsPerHead : 32; - constexpr int32_t kTokensPerTile - = kUseRegisterStaging ? 16 : (kSharedBytesPerToken * 32 <= kDefaultSharedMemoryBytes ? 32 : 16); - constexpr int32_t kBlockSize = kVectorThreads * kTokensPerTile; - static_assert(kSharedBytesPerToken * kTokensPerTile <= kDefaultSharedMemoryBytes); - dim3 const block(kVectorThreads, kTokensPerTile); - dim3 const grid(numLayers, params.kv_head_num, params.batch_size); - size_t const sharedBytes = kUseRegisterStaging ? 0 : 2 * block.y * kVectorsPerHead * sizeof(uint4); - updateSparseKvCacheAfterFmha - <<>>(params); -} - template void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, @@ -2464,13 +2378,11 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con #ifdef ENABLE_BF16 if constexpr (std::is_same_v) { - // Production path for bf16 pools with head size 64/128 and 32/128-token - // pages: the pipelined kernels won the A/B comparison against the - // register-staging kernel everywhere (verified 2026-07-20: 1.47x at - // batch 1, 1.09-1.30x at batch 32, 1.09-1.13x at batch 256, with - // byte-identical outputs), so they dispatch unconditionally here. The - // register-staging path below remains only as the fallback for other - // dtypes and geometries. + // The pipelined kernels are the only shipped path: they won the A/B + // comparison against the retired register-staging kernel everywhere + // (verified 2026-07-20: 1.47x at batch 1, 1.09-1.30x at batch 32, + // 1.09-1.13x at batch 256, byte-identical outputs). Unsupported pool + // dtypes and geometries fail the check below instead of falling back. if ((headDim == 64 || headDim == 128) && (tokensPerBlock == 32 || tokensPerBlock == 128)) { SparseKvCacheCompactV2Bf16Params fastParams{}; @@ -2509,35 +2421,10 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con } #endif // ENABLE_BF16 - KvCacheV2LayersBuffer buffer{}; - buffer.poolPointers = poolPointers; - buffer.pageTable = pageTable; - buffer.sourceLayerIndices = sourceLayerIndices; - buffer.sourceLayerStride = sourceLayerStride; - buffer.sourceHeadStride = sourceHeadStride; - buffer.pageTableRequestStride = pageTableRequestStride; - buffer.tokensPerBlock = tokensPerBlock; - buffer.destinationBases = destinationBases; - buffer.bytesPerKvHalf = static_cast(numKvHeads) * tokensPerBlock * headDim * sizeof(T); - buffer.bytesPerPage = 2 * buffer.bytesPerKvHalf; - - QKVPreprocessingParams params{}; - params.kv_cache_buffer = buffer; - params.sparse_kv_indices = sparseKvIndices; - params.sparse_kv_offsets = sparseKvOffsets; - params.batch_size = batchSize; - params.kv_head_num = numKvHeads; - params.size_per_head = headDim; - - switch (headDim) - { - case 16: launchSparseKvCacheCompactV2Layers<16>(params, numLayers, stream); break; - case 32: launchSparseKvCacheCompactV2Layers<32>(params, numLayers, stream); break; - case 64: launchSparseKvCacheCompactV2Layers<64>(params, numLayers, stream); break; - case 128: launchSparseKvCacheCompactV2Layers<128>(params, numLayers, stream); break; - case 256: launchSparseKvCacheCompactV2Layers<256>(params, numLayers, stream); break; - default: TLLM_CHECK_WITH_INFO(false, "Sparse KV compaction does not support head size %d", headDim); - } + TLLM_CHECK_WITH_INFO(false, + "Sparse KV compaction ships only the pipelined bf16 kernels (head size 64/128, page size 32/128 " + "tokens); got element size %zu, head size %d, %d tokens per page", + sizeof(T), headDim, tokensPerBlock); } //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index c8fac21f431d..4bb2c0a34121 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -77,8 +77,8 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor pageTable.get_device() == device, "sparse_kv_cache_compact_layers: block offsets must be on the pool device"); TORCH_CHECK(poolPointers.is_cuda() && poolPointers.get_device() == device - && poolPointers.scalar_type() == th::kInt64 && poolPointers.dim() == 1 - && poolPointers.size(0) == numLayers && poolPointers.is_contiguous(), + && poolPointers.scalar_type() == th::kInt64 && poolPointers.dim() == 1 && poolPointers.size(0) == numLayers + && poolPointers.is_contiguous(), "sparse_kv_cache_compact_layers: pool_pointers must be contiguous CUDA int64 [num_layers]"); TORCH_CHECK(sourceIndices.is_cuda() && sourceIndices.get_device() == device @@ -138,23 +138,12 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor sourceLayerPtr, sourceLayerStride, sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); } - else if (dtype == th::kHalf) - { - tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), - numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, - sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, - stream); - } - else if (dtype == th::kFloat) - { - tk::invokeSparseKvCacheCompactV2Layers(poolPointers.data_ptr(), pageTable.data_ptr(), - numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, - sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, - stream); - } else { - TORCH_CHECK(false, "sparse_kv_cache_compact_layers: unsupported pool dtype ", dtype); + TORCH_CHECK(false, + "sparse_kv_cache_compact_layers ships only the pipelined bf16 kernels (head size 64/128, page size " + "32/128 tokens); got pool dtype ", + dtype); } } diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 635842bef7d7..10dcfecefd15 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -31,7 +31,7 @@ import torch -_SUPPORTED_POOL_DTYPES = (torch.bfloat16, torch.float16, torch.float32) +_SUPPORTED_POOL_DTYPES = (torch.bfloat16,) class _CppCompactGroup(NamedTuple): @@ -112,8 +112,7 @@ def _validated_kv_head_count( for layer in layers ): raise ValueError( - f"{what} requires contiguous interleaved BF16/FP16/FP32 pools " - "with one common KV-head count" + f"{what} requires contiguous interleaved BF16 pools with one common KV-head count" ) return num_kv_heads diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index e601f1aa034a..9062e4a37cbf 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the layered V2 adapter over the existing sparse-KV updater.""" +"""Tests for the layered V2 sparse-KV compaction op (pipelined bf16 kernels).""" from typing import NamedTuple, Optional @@ -10,20 +10,19 @@ import tensorrt_llm # noqa: F401 # Register torch.ops.trtllm operators. -_TOKENS_PER_BLOCK = 4 +_TOKENS_PER_BLOCK = 32 _NUM_KV_HEADS = 2 _BATCH_SIZE = 2 _MAX_PAGES_PER_SEQUENCE = 3 _NUM_PAGES = _BATCH_SIZE * _MAX_PAGES_PER_SEQUENCE _PAGE_INDEX_DIVISOR = 2 -# Kernel-name substrings for the profiler probes below: the pipelined bf16 -# kernel dispatches unconditionally on eligible geometry, the register-staging -# kernel covers everything else. The byte-compare tests would pass no matter -# which kernel ran, so the probes pin down that the dispatch gate routes each -# case to the intended kernel. +# Kernel-name substrings for the profiler probes below. The pipelined bf16 +# kernels are the only shipped compaction path; the retired register-staging +# kernel (still present for the sparse-attention updater) must never appear +# in this op's launches. _FAST_KERNEL_NAME = "sparseKvCacheCompactV2Bf16PipelineKernel" -_EXISTING_KERNEL_NAME = "updateSparseKvCacheAfterFmha" +_RETIRED_KERNEL_NAME = "updateSparseKvCacheAfterFmha" def _encode_k_block_offsets( @@ -181,12 +180,8 @@ def _compact( @pytest.mark.parametrize( "dtype,head_dim", [ - (torch.float16, 16), - (torch.bfloat16, 32), - (torch.float32, 64), - (torch.float16, 128), - (torch.bfloat16, 256), - (torch.float32, 256), + (torch.bfloat16, 64), + (torch.bfloat16, 128), ], ) @pytest.mark.parametrize( @@ -538,18 +533,25 @@ def test_sparse_kv_cache_compact_layers_fast_geometry_per_layer_source(): (torch.bfloat16, 64, 16), # page size outside the gate ], ) -def test_sparse_kv_cache_compact_layers_fast_gate_fallback(dtype, head_dim, tokens_per_block): - # Near-miss geometries must fall back to the register-staging kernel and - # stay byte-correct. The profiler probe proves the gate actually rejected - # the case: a gate widened by accident would hand the pipelined kernel a - # geometry it was never built for. +def test_sparse_kv_cache_compact_layers_rejects_unsupported_geometry( + dtype, head_dim, tokens_per_block +): + # There is no fallback kernel: near-miss geometries must fail loudly + # instead of silently degrading, and the pools must stay untouched. case = _make_fast_geometry_case(head_dim, tokens_per_block, dtype=dtype) - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: - expected = _run_fast_geometry_case(case) - names = [event.name for event in profiler.events()] - assert any(_EXISTING_KERNEL_NAME in name for name in names) - assert not any(_FAST_KERNEL_NAME in name for name in names) - for actual, reference in zip(case.pools, expected): + arguments = _device_arguments( + case.pools, case.source_indices, case.source_offsets, case.source_layer_indices + ) + with pytest.raises((RuntimeError, ValueError), match="bf16|BF16"): + _compact( + case.pools, + case.page_tables, + arguments, + case.destination_bases, + batch_size=case.batch_size, + ) + torch.cuda.synchronize() + for actual, reference in zip(case.pools, case.pools_cpu): assert torch.equal(actual.cpu(), reference) @@ -565,4 +567,4 @@ def test_sparse_kv_cache_compact_layers_fast_path_actually_runs(head_dim, tokens _run_fast_geometry_case(case) names = [event.name for event in profiler.events()] assert any(_FAST_KERNEL_NAME in name for name in names) - assert not any(_EXISTING_KERNEL_NAME in name for name in names) + assert not any(_RETIRED_KERNEL_NAME in name for name in names) From 6eff0e96b3e7ef3e71b95009dda9f4394a77feb4 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 19:52:26 -0700 Subject: [PATCH 057/178] [None][chore] Fold the never-tuned offset ladder bound into a constant No caller ever passed a non-default offset_max_length, so the constructor knob becomes a module constant. Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 13dbab12da47..532ca1f74cf1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -76,6 +76,10 @@ # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) +# Upper bound of the geometric integration offset ladder [1, 2, 4, ...]; no +# caller ever tuned it, so it is a constant rather than a constructor knob. +_OFFSET_MAX_LENGTH = 65536 + def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Tensor: """Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <=max].""" @@ -1056,7 +1060,6 @@ def __init__( beta: int = 128, model_path: Optional[str] = None, calibration_path: Optional[str] = None, - offset_max_length: int = 65536, score_aggregation: str = "mean", eviction_mode: str = "union", normalize_scores: bool = True, @@ -1110,7 +1113,6 @@ def __init__( # Geometric integration offsets (built lazily on first eviction so the # device matches the cache pool). - self._offset_max_length = offset_max_length self._offsets: Optional[torch.Tensor] = None # Request presence records successful initialization. The record also @@ -1947,7 +1949,7 @@ def _fixed_resources_for( first_pool = layout.layer_pools[layout.dense_layers[0]] if self._offsets is None: - self._offsets = _build_geometric_offsets(self._offset_max_length, first_pool.device) + self._offsets = _build_geometric_offsets(_OFFSET_MAX_LENGTH, first_pool.device) q_real, q_imag, mlr_coef = self._local_score_calibration( layout.num_layers, layout.global_layers ) From 4de52275dd53741ecc2e7c5a638e82913be9e4a4 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 20:00:27 -0700 Subject: [PATCH 058/178] [None][chore] Move the cohort-tail test helper into the test conftest Production stages the per-family move offsets through the single round-metadata upload; only the unit tests drove the buffers standalone via set_protected_tails, so the helper moves off BatchedKVCacheCompaction into the kv_cache_compression conftest. Signed-off-by: tianruih --- .../triattention/compaction.py | 50 +------------------ .../_torch/kv_cache_compression/conftest.py | 48 ++++++++++++++++++ .../test_triattention_draft_cocompaction.py | 3 +- .../test_triattention_selection_compaction.py | 13 ++--- 4 files changed, 59 insertions(+), 55 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 10dcfecefd15..d2448754cba9 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -402,8 +402,8 @@ class BatchedKVCacheCompaction: object must support. A protected tail covers KV positions past the valid length reserved for a forward already in flight; each round's actual lengths arrive through the per-family move-offset - rows staged with the round metadata (`set_protected_tails` fills - the same rows for standalone use) and move with the kept tokens. + rows staged with the round metadata and move with the kept + tokens. `draft_*`: co-compressed draft-cache layout (union mode only); the draft reuses the target keep set and pins the same prompt. """ @@ -707,52 +707,6 @@ def _build_draft_compaction( destination_bases=self.prompt_offsets, ) - def _write_move_offsets(self, offsets: torch.Tensor, moves_per_request: List[int]) -> None: - cumulative = [0] - for count in moves_per_request: - cumulative.append(cumulative[-1] + count) - # Rows past the cohort are padding and contribute no moves. - cumulative.extend(cumulative[-1:] * (self.request_count - len(moves_per_request))) - offsets.copy_(torch.tensor(cumulative, dtype=torch.int32), non_blocking=True) - - def set_protected_tails( - self, - tail_lengths: List[int], - draft_tail_lengths: Optional[List[int]] = None, - ) -> None: - """Load this cohort's per-request protected tails into the move offsets. - - The pack kernel and the C++ compacts read every request's move range - from these offsets, so refreshing them retargets the fixed buffers to - the cohort at hand without any reallocation. - """ - if len(tail_lengths) > self.request_count: - raise ValueError("the cohort exceeds the compaction request capacity") - if any(tail < 0 or tail > self.protected_tail_capacity for tail in tail_lengths): - raise ValueError("a protected tail exceeds the configured capacity") - self._write_move_offsets( - self.target_dense_compaction.move_source_offsets, - [self.decode_keep_count + int(tail) for tail in tail_lengths], - ) - if self.target_swa_compaction is not None: - self._write_move_offsets( - self.target_swa_compaction.move_source_offsets, - [self.swa_window + int(tail) for tail in tail_lengths], - ) - if self.draft_compaction is not None: - if draft_tail_lengths is None: - draft_tail_lengths = [0] * len(tail_lengths) - if len(draft_tail_lengths) != len(tail_lengths): - raise ValueError("draft protected tails must match the cohort") - if any( - tail < 0 or tail > self.draft_protected_tail_capacity for tail in draft_tail_lengths - ): - raise ValueError("a draft protected tail exceeds the configured capacity") - self._write_move_offsets( - self.draft_compaction.move_source_offsets, - [self.decode_keep_count + int(tail) for tail in draft_tail_lengths], - ) - def compact(self) -> None: """Pack the move indices, then run every cache family's C++ compacts.""" if self.swa_destination_bases is not None: diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 69ffaf9b6e84..f8df2e35476f 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -41,6 +41,54 @@ def encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: return encoded +def _write_move_offsets(compaction, offsets, moves_per_request): + cumulative = [0] + for count in moves_per_request: + cumulative.append(cumulative[-1] + count) + # Rows past the cohort are padding and contribute no moves. + cumulative.extend(cumulative[-1:] * (compaction.request_count - len(moves_per_request))) + offsets.copy_(torch.tensor(cumulative, dtype=torch.int32), non_blocking=True) + + +def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): + """Load a cohort's per-request protected tails into the move offsets. + + Production stages these rows through the single round-metadata upload; + tests drive the fixed buffers directly through this helper (moved out of + BatchedKVCacheCompaction, whose production surface has no caller for it). + """ + if len(tail_lengths) > compaction.request_count: + raise ValueError("the cohort exceeds the compaction request capacity") + if any(tail < 0 or tail > compaction.protected_tail_capacity for tail in tail_lengths): + raise ValueError("a protected tail exceeds the configured capacity") + _write_move_offsets( + compaction, + compaction.target_dense_compaction.move_source_offsets, + [compaction.decode_keep_count + int(tail) for tail in tail_lengths], + ) + if compaction.target_swa_compaction is not None: + _write_move_offsets( + compaction, + compaction.target_swa_compaction.move_source_offsets, + [compaction.swa_window + int(tail) for tail in tail_lengths], + ) + if compaction.draft_compaction is not None: + if draft_tail_lengths is None: + draft_tail_lengths = [0] * len(tail_lengths) + if len(draft_tail_lengths) != len(tail_lengths): + raise ValueError("draft protected tails must match the cohort") + if any( + tail < 0 or tail > compaction.draft_protected_tail_capacity + for tail in draft_tail_lengths + ): + raise ValueError("a draft protected tail exceeds the configured capacity") + _write_move_offsets( + compaction, + compaction.draft_compaction.move_source_offsets, + [compaction.decode_keep_count + int(tail) for tail in draft_tail_lengths], + ) + + 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 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 index 4d7af4ea8ae4..65162b9e109a 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -23,6 +23,7 @@ from conftest import make_request as _make_request from conftest import make_triattention as _make_triattention from conftest import mocked_eviction_internals as _mocked_eviction_internals +from conftest import set_protected_tails as _set_protected_tails from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( BatchedKVCacheCompaction, @@ -108,7 +109,7 @@ def _launched_draft_compaction(draft_protected_tails): draft_kv_block_offsets=_encode_block_offsets(draft_tables), draft_page_table_slots={0: 0}, ) - compaction.set_protected_tails(target_protected_tails, draft_protected_tails) + _set_protected_tails(compaction, target_protected_tails, draft_protected_tails) compaction.compact() torch.cuda.synchronize(device) 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 index 8aa7be2995fa..0bd6606cd423 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -5,6 +5,7 @@ import pytest import torch from conftest import encode_block_offsets as _encode_block_offsets +from conftest import set_protected_tails as _set_protected_tails from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( BatchedKVCacheCompaction, @@ -442,7 +443,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) swa_window=None, protected_tail_capacity=max(protected_tails), ) - compaction.set_protected_tails(protected_tails) + _set_protected_tails(compaction, protected_tails) compaction.compact() torch.cuda.synchronize(device) @@ -549,7 +550,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): swa_window=None, protected_tail_capacity=max(protected_tails), ) - cohort_compaction.set_protected_tails(protected_tails) + _set_protected_tails(cohort_compaction, protected_tails) cohort_compaction.compact() expected_pools = [pool.clone() for pool in initial_pools] @@ -571,7 +572,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): swa_window=None, protected_tail_capacity=protected_tails[request], ) - single_compaction.set_protected_tails([protected_tails[request]]) + _set_protected_tails(single_compaction, [protected_tails[request]]) single_compaction.compact() torch.cuda.synchronize(device) @@ -692,7 +693,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): swa_window=None, protected_tail_capacity=0, ) - batched_compaction.set_protected_tails([0]) + _set_protected_tails(batched_compaction, [0]) batched_compaction.compact() torch.cuda.synchronize(device) @@ -850,7 +851,7 @@ def write_token(token: int, score: float) -> None: swa_window=None, protected_tail_capacity=protected_tail, ) - batched_compaction.set_protected_tails([protected_tail]) + _set_protected_tails(batched_compaction, [protected_tail]) def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) @@ -969,7 +970,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): swa_window=2, protected_tail_capacity=max(protected_tails), ) - compaction.set_protected_tails(protected_tails) + _set_protected_tails(compaction, protected_tails) compaction.compact() torch.cuda.synchronize(device) From c8c6bd7b092f3c572951cf5c55b258762475811e Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 20:39:51 -0700 Subject: [PATCH 059/178] [None][feat] Generalize the CuTe score kernel to 32-token pages The kernel was specialized to its validation geometry (128-token pages, 32 frequencies, GQA group 8). Promote the frequency count and page size to constructor parameters and let one 64-token compute tile span several pages: each TMA phase issues one box per page fragment into the same transaction barrier, the producer prefetches one page id per fragment, and the ragged last tile clamps its second fragment to the first valid page. The TMA descriptor geometry and swizzle now derive from the frequency count. The validated 128-token geometry keeps the single-fragment schedule. Unit oracle now also covers 32-token pages with a shuffled physical-page table and mid-tile valid lengths. Signed-off-by: tianruih --- .../triattention/triattention_cute_score.py | 315 ++++++++++++------ .../triattention/triattention_kernels.py | 4 +- .../test_triattention_cute_score.py | 49 ++- 3 files changed, 259 insertions(+), 109 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py index 19ec1f885885..d5f3303554c8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -55,16 +55,10 @@ def _cute_sqrt_keyword_mode() -> str: _CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() CTA_M = 64 -K = 96 N = 8 -NUM_FREQS = 32 THREADS = 128 -PAGE_TOKENS = 128 -RAW_K_HALF_ELEMENTS = CTA_M * 2 * NUM_FREQS RAW_K_VECTOR_ELEMENTS = 8 -RAW_K_SPLIT_PHASE_ELEMENTS = CTA_M * NUM_FREQS -RAW_K_SPLIT_TMA_COPY_BYTES = RAW_K_SPLIT_PHASE_ELEMENTS * (cutlass.BFloat16.width // 8) TMA_DESCRIPTOR_QWORDS = 16 @@ -155,13 +149,15 @@ def __init__( super().__init__() if pool_dtype is not cutlass.BFloat16: raise ValueError("TriAttention CuTe score requires BF16 K pages") - if num_freqs != NUM_FREQS: - raise ValueError("TriAttention CuTe score requires 32 frequencies") - if tokens_per_block != PAGE_TOKENS: - raise ValueError("TriAttention CuTe score requires 128-token pages") + 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 num_q_heads % num_kv_heads or num_q_heads // num_kv_heads != N: raise ValueError("TriAttention CuTe score requires GQA group 8") - if score_start % PAGE_TOKENS: + if score_start % tokens_per_block: raise ValueError("TriAttention CuTe score requires page-aligned score_start") if page_shards not in (2, 3): raise ValueError("TriAttention CuTe score requires two or three page shards") @@ -174,8 +170,21 @@ def __init__( self.num_tasks = num_segments * num_kv_heads self.page_shards = page_shards self.num_ctas = self.num_tasks * page_shards - self.max_pages = (seq_len + PAGE_TOKENS - 1) // PAGE_TOKENS - self.halves_per_page = PAGE_TOKENS // CTA_M + self.num_freqs = num_freqs + self.tokens_per_block = tokens_per_block + # cos/sin/mlr coefficient planes per frequency. + self.k_coeff = 3 * num_freqs + # One 64-token compute tile either sub-divides a page (128-token + # pages: two halves per page, one TMA box each) or spans several + # pages (32-token pages: two page fragments per phase). The + # validated 128-token geometry is the single-fragment case, so + # its schedule is unchanged. + self.box_tokens = min(CTA_M, tokens_per_block) + self.fragments_per_phase = CTA_M // self.box_tokens + self.pages_per_tile = self.fragments_per_phase + self.halves_per_page = max(1, tokens_per_block // CTA_M) + self.tile_tokens = self.halves_per_page * CTA_M + self.max_tiles = (seq_len + self.tile_tokens - 1) // self.tile_tokens # Measured final choices that still shape layouts or generated code. self.prefetch_depth = 4 @@ -184,8 +193,11 @@ def __init__( self.use_tma = True self.cpasync_schedule = "sync_each_half" self.split_raw_tma = True - self.raw_tma_feature_extent = NUM_FREQS - self.raw_tma_copy_bytes = RAW_K_SPLIT_TMA_COPY_BYTES + self.raw_tma_feature_extent = num_freqs + # Barrier transaction bytes for one phase: the full 64-token tile + # of one coefficient plane, regardless of how many page fragments + # deliver it. + self.raw_tma_copy_bytes = CTA_M * num_freqs * (cutlass.BFloat16.width // 8) self.raw_tma_pipeline_stages = 1 self.accumulator_pipeline_stages = 1 self.umma_accumulator_partitions = 1 @@ -203,11 +215,15 @@ def __init__( self.compact_token_loop = True self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape - if pool_kv_heads != num_kv_heads or pool_tokens != PAGE_TOKENS or pool_dim != 2 * NUM_FREQS: + if ( + pool_kv_heads != num_kv_heads + or pool_tokens != tokens_per_block + or 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_slot, self.s_dim = pool_strides - if self.s_slot != 2 * NUM_FREQS or self.s_dim != 1: - raise ValueError("K pages must be contiguous [128, 64]") + if self.s_slot != 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") @@ -233,7 +249,7 @@ def __call__( ): self.c_dtype = output.element_type self.c_layout = utils.LayoutEnum.COL_MAJOR - self.mma_tiler = (CTA_M, N, K) + self.mma_tiler = (CTA_M, N, self.k_coeff) self.cta_tile_shape_mnk = self.mma_tiler self.epi_tile = (CTA_M, N) @@ -254,20 +270,20 @@ def __call__( self.mma_tiler[:2], ) main_a_shape = ( - (CTA_M, N, NUM_FREQS) + (CTA_M, N, self.num_freqs) if self.main_operand_mode == "bf16_raw_three_term_weight" else self.mma_tiler ) a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, main_a_shape, cutlass.Float32, 1) raw_bf16_a_smem_layout = sm100_utils.make_smem_layout_a( raw_bf16_tiled_mma, - (CTA_M, N, 2 * NUM_FREQS), + (CTA_M, N, 2 * self.num_freqs), cutlass.BFloat16, 1, ) raw_bf16_split_a_smem_layout = sm100_utils.make_smem_layout_a( raw_bf16_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.BFloat16, 2, ) @@ -281,14 +297,14 @@ def __call__( raw_bf16_direct_a_smem_layout.inner, 0, cute.make_layout( - (self.raw_tma_feature_extent, CTA_M), + (self.raw_tma_feature_extent, self.box_tokens), stride=(1, self.raw_tma_feature_extent), ), ) raw_tma_source_layout = cute.make_layout( ( - 2 * NUM_FREQS, - PAGE_TOKENS, + 2 * self.num_freqs, + self.tokens_per_block, (self.num_kv_heads, self.num_physical_pages), ), stride=( @@ -305,19 +321,19 @@ def __call__( cpasync.CopyBulkTensorTileG2SOp(), raw_tma_source, raw_tma_smem_layout, - (self.raw_tma_feature_extent, CTA_M), + (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, N, 2 * NUM_FREQS), + (CTA_M, N, 2 * self.num_freqs), cutlass.BFloat16, 1, ) - # The magnitude residual has only K=32. A separate compact descriptor - # lets the producer issue its four UMMA steps before the first commit, - # rather than waiting for and overwriting the K=96 main A tile. + # The magnitude residual has only the frequency-count K. A separate + # compact descriptor lets the producer issue its UMMA steps before the + # first commit, rather than waiting for and overwriting the main A tile. magnitude_lo_smem_layout = sm100_utils.make_smem_layout_a( - tiled_mma, (CTA_M, N, NUM_FREQS), cutlass.Float32, 1 + tiled_mma, (CTA_M, N, self.num_freqs), cutlass.Float32, 1 ) magnitude_lo_tiled_mma = sm100_utils.make_trivial_tiled_mma( cutlass.Float16, @@ -329,30 +345,30 @@ def __call__( ) magnitude_lo_fp16_smem_layout = sm100_utils.make_smem_layout_a( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) magnitude_hi_fp16_smem_layout = sm100_utils.make_smem_layout_b( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) magnitude_fp16_b_smem_layout = sm100_utils.make_smem_layout_b( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) main_b_shape = ( - (CTA_M, N, NUM_FREQS) + (CTA_M, N, self.num_freqs) if self.main_operand_mode == "bf16_raw_three_term_weight" else self.mma_tiler ) @@ -386,12 +402,13 @@ def __call__( a_elements = cute.cosize(a_smem_layout.outer) * int( not self.shared_a_raw_alias and not self.fp16_magnitude_two_term ) - raw_k_elements = RAW_K_HALF_ELEMENTS * int( + raw_k_half_elements = CTA_M * 2 * self.num_freqs + raw_k_elements = raw_k_half_elements * int( self.k_staging_mode in ("half_page_cpasync", "half_page_tma") and not self.shared_a_raw_alias ) alias_a_elements = cute.cosize(a_smem_layout.outer) * int(self.shared_a_raw_alias) - alias_raw_k_elements = RAW_K_HALF_ELEMENTS * int(self.shared_a_raw_alias) + alias_raw_k_elements = raw_k_half_elements * int(self.shared_a_raw_alias) raw_bf16_a_elements = cute.cosize(raw_bf16_a_smem_layout.outer) * int( self.main_operand_mode == "bf16_raw_three_term_weight" and ( @@ -612,33 +629,45 @@ def kernel( # Each stage slice retains the K_SW64 pointer flags. Reuse only # the feature-first outer mapping for the corresponding TMA # destination so the swizzle is not applied twice. - raw_tma_shared_real = cute.make_tensor( - cpasync_raw_k_real.iterator, - raw_tma_smem_layout.outer, - ) - raw_tma_shared_imag = cute.make_tensor( - cpasync_raw_k_imag.iterator, - raw_tma_smem_layout.outer, - ) raw_tma_source_tiles = cute.local_tile( raw_tma_source, - (self.raw_tma_feature_extent, CTA_M), + (self.raw_tma_feature_extent, self.box_tokens), coord=(None, None, None), ) - raw_tma_shared_partition_real, raw_tma_global_partition = cpasync.tma_partition( - raw_tma_atom, - 0, - cute.make_layout(1), - cute.group_modes(raw_tma_shared_real, 0, 2), - cute.group_modes(raw_tma_source_tiles, 0, 2), - ) - raw_tma_shared_partition_imag, _ = cpasync.tma_partition( - raw_tma_atom, - 0, - cute.make_layout(1), - cute.group_modes(raw_tma_shared_imag, 0, 2), - cute.group_modes(raw_tma_source_tiles, 0, 2), - ) + # One smem view and TMA partition per page fragment of the + # 64-token tile. Fragment f lands box_tokens rows deeper in the + # same stage; the offset is a whole multiple of the swizzle + # period, so the descriptor swizzle stays phase-aligned. + raw_tma_shared_partition_real = [] + raw_tma_shared_partition_imag = [] + 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, + ) + 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), + ) + raw_tma_shared_partition_real.append(partition_real) + raw_tma_shared_partition_imag.append(partition_imag) + raw_tma_global_partition = global_partition raw_tensormap_manager = utils.TensorMapManager( utils.TensorMapUpdateMode.GMEM, 128, @@ -656,10 +685,11 @@ def kernel( swizzle=raw_bf16_b_smem_layout.inner, ) - page_index = self.score_start // PAGE_TOKENS + page_shard - page_start = page_index * PAGE_TOKENS + page_index = self.score_start // self.tile_tokens + page_shard + page_start = page_index * self.tile_tokens pages_processed = cutlass.Int32(0) producer_prefetched_page_id_lane0 = cutlass.Int32(0) + producer_prefetched_page_id_lane0_f1 = cutlass.Int32(0) if warp_idx == self.producer_warp_id: if lane_idx == 0: # ``page_ids`` is the flattened native block-offset staging @@ -671,8 +701,20 @@ def kernel( # ``encoded / kvFactor`` (triAttentionScoreKernels.cu), so # divide by two here as well. V-plane entries are never read. producer_prefetched_page_id_lane0 = ( - cutlass.Int32(page_ids[page_off + page_index]) // 2 + cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile]) // 2 ) + if cutlass.const_expr(self.fragments_per_phase == 2): + # The tail tile may not reach its second page; clamp + # to the first fragment (those scores lie past the + # valid width and are masked downstream) so the TMA + # never dereferences an unstaged block entry. + second_page_id = producer_prefetched_page_id_lane0 + if page_start + self.box_tokens < valid_seq_len: + second_page_id = ( + cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile + 1]) + // 2 + ) + producer_prefetched_page_id_lane0_f1 = second_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) @@ -713,15 +755,15 @@ def kernel( pipeline.PipelineUserType.Consumer, self.accumulator_pipeline_stages ) cute.arch.mbarrier_init_fence() - for weight_round in cutlass.range_constexpr(N * K // THREADS): + for weight_round in cutlass.range_constexpr(N * self.k_coeff // THREADS): linear_index = tidx + weight_round * THREADS - qg = linear_index // K - feature = linear_index % K - coefficient_kind = feature // NUM_FREQS - frequency = feature % NUM_FREQS - mean_offset = req_id * NUM_FREQS + frequency + 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 q_head = kv_head * self.group_size + qg - calib_offset = (layer_id * self.num_q_heads + q_head) * NUM_FREQS + frequency + 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]) @@ -779,13 +821,19 @@ def kernel( tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) thr_mma = tiled_mma.get_slice(0) - while page_start < valid_seq_len and pages_processed < self.max_pages: + while page_start < valid_seq_len and pages_processed < self.max_tiles: physical_page = cutlass.Int32(0) + physical_page_f1 = 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 == 2): + physical_page_f1 = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0_f1, + 0, + ) for page_half in cutlass.range_constexpr(self.halves_per_page): if warp_idx == self.producer_warp_id: # Phase 0 fills the packed 4-KiB K_SW64 real @@ -802,10 +850,27 @@ def kernel( (kv_head, physical_page), ) ], - raw_tma_shared_partition_real, + raw_tma_shared_partition_real[0], tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), tma_desc_ptr=raw_tma_descriptor_ptr, ) + if cutlass.const_expr(self.fragments_per_phase == 2): + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + page_half, + (kv_head, physical_page_f1), + ) + ], + raw_tma_shared_partition_real[1], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() raw_tma_pipeline.consumer_wait(raw_tma_consumer_state) raw_tma_pipeline.consumer_release(raw_tma_consumer_state) @@ -822,31 +887,71 @@ def kernel( (kv_head, physical_page), ) ], - raw_tma_shared_partition_imag, + raw_tma_shared_partition_imag[0], tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), tma_desc_ptr=raw_tma_descriptor_ptr, ) + if cutlass.const_expr(self.fragments_per_phase == 2): + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + page_half, + (kv_head, physical_page_f1), + ) + ], + raw_tma_shared_partition_imag[1], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() - if cutlass.const_expr(self.producer_page_id_prefetch and page_half == 1): + if cutlass.const_expr( + self.producer_page_id_prefetch and page_half == self.halves_per_page - 1 + ): next_page_id_lane0 = cutlass.Int32(0) + next_page_id_lane0_f1 = cutlass.Int32(0) if warp_idx == self.producer_warp_id: if lane_idx == 0: - next_page_start = page_start + PAGE_TOKENS * self.page_shards + next_page_start = page_start + self.tile_tokens * self.page_shards next_pages_processed = pages_processed + 1 if ( next_page_start < valid_seq_len - and next_pages_processed < self.max_pages + and next_pages_processed < self.max_tiles ): # Same K-plane decode as the initial # prefetch: entries are physical_page * 2. next_page_id_lane0 = ( cutlass.Int32( - page_ids[page_off + page_index + self.page_shards] + page_ids[ + page_off + + (page_index + self.page_shards) * self.pages_per_tile + ] ) // 2 ) + if cutlass.const_expr(self.fragments_per_phase == 2): + next_second_id = next_page_id_lane0 + if next_page_start + self.box_tokens < valid_seq_len: + next_second_id = ( + cutlass.Int32( + page_ids[ + page_off + + (page_index + self.page_shards) + * self.pages_per_tile + + 1 + ] + ) + // 2 + ) + next_page_id_lane0_f1 = next_second_id producer_prefetched_page_id_lane0 = next_page_id_lane0 + if cutlass.const_expr(self.fragments_per_phase == 2): + producer_prefetched_page_id_lane0_f1 = next_page_id_lane0_f1 # Submit B0-real while the imaginary TMA is in flight. tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] @@ -856,7 +961,7 @@ def kernel( tcgen05.Field.ACCUMULATE, False, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( raw_bf16_tiled_mma, tCtAcc, @@ -876,7 +981,7 @@ def kernel( # from global to the single raw shared buffer. for token_base in cutlass.range( 0, - CTA_M // 4, + CTA_M // (THREADS // 32), self.prefetch_depth, unroll_full=not self.compact_token_loop, ): @@ -884,7 +989,7 @@ def kernel( 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 * 4 + token = warp_idx + token_round * (THREADS // 32) staged_real[prefetch_index] = cutlass.Float32( cpasync_raw_k_0[ ( @@ -908,7 +1013,7 @@ def kernel( for prefetch_index in cutlass.range_constexpr(self.prefetch_depth): token_round = token_base + prefetch_index - token = warp_idx + token_round * 4 + token = warp_idx + token_round * (THREADS // 32) real = staged_real[prefetch_index] imag = staged_imag[prefetch_index] norm2 = real * real + imag * imag @@ -954,8 +1059,8 @@ def kernel( tcgen05.Field.ACCUMULATE, True, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): - imag_b_block = NUM_FREQS // 16 + raw_k_block + 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, @@ -963,7 +1068,7 @@ def kernel( tCrRawBf16B0[(None, None, imag_b_block, 0)], tCtAcc, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( raw_bf16_tiled_mma, tCtAcc, @@ -971,8 +1076,8 @@ def kernel( tCrRawBf16B1[(None, None, raw_k_block, 0)], tCtAcc, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): - imag_b_block = NUM_FREQS // 16 + raw_k_block + 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, @@ -988,7 +1093,7 @@ def kernel( tcgen05.Field.ACCUMULATE, True, ) - for magnitude_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( magnitude_lo_tiled_mma, tCtAcc, @@ -996,7 +1101,7 @@ def kernel( tCrMagnitudeFp16B0[(None, None, magnitude_k_block, 0)], tCtAcc, ) - for magnitude_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( magnitude_lo_tiled_mma, tCtAcc, @@ -1004,7 +1109,7 @@ def kernel( tCrMagnitudeFp16B1[(None, None, magnitude_k_block, 0)], tCtAcc, ) - for magnitude_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( magnitude_lo_tiled_mma, tCtAcc, @@ -1068,7 +1173,7 @@ def kernel( raw_tma_consumer_state.advance() page_index += self.page_shards - page_start += PAGE_TOKENS * self.page_shards + page_start += self.tile_tokens * self.page_shards pages_processed += 1 if warp_idx == self.producer_warp_id: raw_tma_pipeline.producer_tail(raw_tma_producer_state) @@ -1086,6 +1191,8 @@ def kernel( def _encode_tma_descriptors( layer_pools: list[torch.Tensor], layer_indices: list[int], + num_freqs: int, + tokens_per_block: int, ) -> torch.Tensor: """Encode one immutable feature-first TensorMap per layer index.""" anchor = layer_pools[layer_indices[0]] @@ -1099,14 +1206,17 @@ def _encode_tma_descriptors( 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, tokens_per_block, head_dim = pool.shape - if (kv_factor, tokens_per_block, head_dim) != (2, PAGE_TOKENS, 2 * NUM_FREQS): - raise ValueError("TriAttention CuTe score requires [page, 2, Hkv, 128, 64] pools") + _, kv_factor, num_kv_heads, pool_tokens, head_dim = pool.shape + if (kv_factor, pool_tokens, head_dim) != (2, 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, PAGE_TOKENS] + 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)) @@ -1115,7 +1225,7 @@ def _encode_tma_descriptors( global_dims.append(int(pool.shape[0])) global_strides_bytes.append(s_page * pool.element_size()) tensor_rank = len(global_dims) - box_dims = [NUM_FREQS, CTA_M] + [1] * (tensor_rank - 2) + 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), @@ -1125,7 +1235,14 @@ def _encode_tma_descriptors( [uint32(value) for value in box_dims], [uint32(1) for _ in range(tensor_rank)], cuda.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, - cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B, + # The swizzle must match the smem layout the TMA lands in; the + # sm100 helpers pick it from the inner-row byte count (one + # coefficient plane: num_freqs bf16 elements). + ( + 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, ) @@ -1197,7 +1314,9 @@ def __init__( self.max_requests = int(max_requests) self.num_layers = int(num_layers) self.num_kv_heads = int(num_kv_heads) - self.descriptors = _encode_tma_descriptors(layer_pools, layer_indices) + self.descriptors = _encode_tma_descriptors( + layer_pools, layer_indices, int(num_freqs), int(tokens_per_block) + ) self._torch_prefix = ( page_ids, seg_page_off, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 8318020a87ea..08aa2c46dc17 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -568,8 +568,8 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> torch.cuda.get_device_capability(anchor.device) == (10, 0) and anchor.dtype == torch.bfloat16 and kv_factor == 2 - and tokens_per_block == 128 - and num_freqs == 32 + and tokens_per_block in (32, 128) + and num_freqs in (32, 64) and num_q_heads == num_kv_heads * 8 and int(anchor.stride(-1)) == 1 # The kernel computes flat score offsets in 32-bit arithmetic. 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 index b61a0cbe4c8b..9349964b90e4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -10,7 +10,25 @@ not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), reason="TriAttention CuTe score kernel requires SM100", ) -def test_cute_score_matches_torch_mean_oracle(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize( + "tokens_per_block,page_permutation,valid_lens", + [ + # The originally validated geometry: 128-token pages, identity table. + (128, [0, 1], None), + # Production page size: a 64-token compute tile spans two pages, so a + # shuffled physical-page table catches any fragment/page mix-up. + (32, [3, 1, 4, 7, 5, 0, 2, 6], None), + # Ragged tails land mid-tile: the second page fragment of the last + # tile is clamped, and scores past the valid length are unspecified. + (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198]), + ], +) +def test_cute_score_matches_torch_mean_oracle( + monkeypatch: pytest.MonkeyPatch, + tokens_per_block: int, + page_permutation: list, + valid_lens: "list | None", +) -> None: pytest.importorskip("cutlass") monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_SCORE", "1") @@ -23,7 +41,11 @@ def test_cute_score_matches_torch_mean_oracle(monkeypatch: pytest.MonkeyPatch) - seq_len = 256 num_q_heads = 8 num_freqs = 32 - pool = (0.125 * torch.randn(2, 2, 1, 128, 64, device=device)).to(torch.bfloat16) + num_pages = seq_len // 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) @@ -37,9 +59,11 @@ def test_cute_score_matches_torch_mean_oracle(monkeypatch: pytest.MonkeyPatch) - # Native block-offset staging layout ([pool_slot, request, K/V plane, # block] int32): K-plane entries encode physical_page * kv_factor with - # kv_factor == 2. Both requests read pool pages [0, 1]. + # kv_factor == 2. Both requests read the same (permuted) page sequence. + k_plane = [2 * page for page in page_permutation] + v_plane = [2 * page + 1 for page in page_permutation] block_offsets = torch.tensor( - [[[[0, 2], [1, 3]], [[0, 2], [1, 3]]]], dtype=torch.int32, device=device + [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device ) group = _FixedScoreGroup( [pool], @@ -58,12 +82,18 @@ def test_cute_score_matches_torch_mean_oracle(monkeypatch: pytest.MonkeyPatch) - offsets, output_width=seq_len, ) - keys = pool[:, 0, 0].reshape(seq_len, 2 * num_freqs).float() + keys = ( + torch.cat([pool[page, 0, 0] for page in page_permutation], dim=0) + .reshape(seq_len, 2 * num_freqs) + .float() + ) k_real = keys[:, :num_freqs] k_imag = keys[:, num_freqs:] magnitude = torch.sqrt(k_real.square() + k_imag.square()) - valid_seq_lens = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) - valid_widths = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) + if valid_lens is None: + valid_lens = [seq_len, seq_len] + valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) + valid_widths = torch.tensor(valid_lens, dtype=torch.int32, device=device) round_starts_device = torch.tensor([seq_len, seq_len + 1], dtype=torch.int32, device=device) token_starts_device = torch.zeros(2, dtype=torch.int32, device=device) for request_count in (1, 2): @@ -87,9 +117,10 @@ def test_cute_score_matches_torch_mean_oracle(monkeypatch: pytest.MonkeyPatch) - + q_imag[0, :, None] * rotated_imag[None] + mlr_coef[0, :, None] * freq_scale_sq[None, None] * magnitude[None] ).sum(dim=-1) + valid = valid_lens[request] torch.testing.assert_close( - actual[request, 0], - expected, + actual[request, 0, :, :valid], + expected[:, :valid], rtol=5.0e-3, atol=5.0e-3, ) From a3c0053e009bf8ca89a62d9849f383a877804bab Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 20:51:20 -0700 Subject: [PATCH 060/178] [None][feat] Extend the CuTe score kernel to 64-frequency, GQA-group-4 heads Qwen3 heads carry 128-element K rows (64 frequencies) and 4 query heads per KV head. The 32 magnitude-staging lanes take one pass per 32 frequencies, and groups below the minimum tcgen05 MMA tile N=8 ride padded columns: the padded coefficients are zeroed (their score columns come out zero) and the adapter scratch is padded per KV head so the gather skips them. Group-8 geometries keep the previous schedule. The unit oracle gains both Qwen3-shaped legs. Signed-off-by: tianruih --- .../triattention/triattention_cute_score.py | 170 ++++++++++-------- .../triattention/triattention_kernels.py | 63 ++++--- .../test_triattention_cute_score.py | 20 ++- 3 files changed, 147 insertions(+), 106 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py index d5f3303554c8..cc8ca7a8fbe1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -155,8 +155,8 @@ def __init__( ) if tokens_per_block not in (32, 128): raise ValueError("TriAttention CuTe score requires 32- or 128-token pages") - if num_q_heads % num_kv_heads or num_q_heads // num_kv_heads != N: - raise ValueError("TriAttention CuTe score requires GQA group 8") + 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 score_start % tokens_per_block: raise ValueError("TriAttention CuTe score requires page-aligned score_start") if page_shards not in (2, 3): @@ -762,7 +762,16 @@ def kernel( coefficient_kind = feature // self.num_freqs frequency = feature % self.num_freqs mean_offset = req_id * self.num_freqs + frequency - q_head = kv_head * self.group_size + qg + # GQA groups below the minimum MMA tile N=8 ride padded + # columns: they read the group's first head (any valid + # address) and force zero coefficients, so the padded score + # columns come out zero and land in scratch rows the adapter + # never gathers. + qg_read = qg + if cutlass.const_expr(self.group_size < N): + 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]) @@ -776,6 +785,9 @@ def kernel( value = scale * (qr * msin + qi * mcos) else: value = scale * cutlass.Float32(mlr_coef[calib_offset]) + if cutlass.const_expr(self.group_size < N): + if qg >= self.group_size: + value = cutlass.Float32(0.0) raw_k_block = feature // 16 magnitude_k_block = frequency // 16 if coefficient_kind < 2: @@ -974,79 +986,84 @@ def kernel( True, ) raw_tma_pipeline.consumer_wait(raw_tma_consumer_state) - frequency = lane_idx - # Issue several independent token loads before consuming - # any of them. This bounded RMEM window is unchanged; the - # optional half-page staging only switches its K source - # from global to the single raw shared buffer. - for token_base in cutlass.range( - 0, - CTA_M // (THREADS // 32), - self.prefetch_depth, - unroll_full=not self.compact_token_loop, - ): - 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, - 0, + # Each of the 32 lanes 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 + # Issue several independent token loads before consuming + # any of them. This bounded RMEM window is unchanged; the + # optional half-page staging only switches its K source + # from global to the single raw shared buffer. + for token_base in cutlass.range( + 0, + CTA_M // (THREADS // 32), + self.prefetch_depth, + unroll_full=not self.compact_token_loop, + ): + 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, + 0, + ) + ] + ) + staged_imag[prefetch_index] = cutlass.Float32( + cpasync_raw_k_0[ + ( + (token, frequency % 16), + 0, + frequency // 16, + 1, + ) + ] + ) + + 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 + if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): + magnitude = cute.math.sqrt( + norm2, + approx=self.sqrt_mode == "approx", + ftz=self.magnitude_sqrt_ftz, ) - ] - ) - staged_imag[prefetch_index] = cutlass.Float32( - cpasync_raw_k_0[ - ( - (token, frequency % 16), - 0, - frequency // 16, - 1, + elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): + # cutlass 4.5 renamed the approximate-sqrt control + # to ``fastmath``; map the measured approx choice + # onto it to preserve the authored behavior. + magnitude = cute.math.sqrt( + norm2, fastmath=self.sqrt_mode == "approx" ) - ] - ) - - 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 - if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): - magnitude = cute.math.sqrt( - norm2, - approx=self.sqrt_mode == "approx", - ftz=self.magnitude_sqrt_ftz, + else: + # DSLs with neither spelling get the plain (IEEE) + # sqrt, which is strictly MORE accurate than the + # measured approx choice above; the unit test's + # 5e-3 oracle tolerance absorbs the difference. + magnitude = cute.math.sqrt(norm2) + magnitude_fp16_0 = cutlass.Float16(magnitude) + magnitude_fp16_1 = cutlass.Float16( + magnitude - cutlass.Float32(magnitude_fp16_0) ) - elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): - # cutlass 4.5 renamed the approximate-sqrt control - # to ``fastmath``; map the measured approx choice - # onto it to preserve the authored behavior. - magnitude = cute.math.sqrt(norm2, fastmath=self.sqrt_mode == "approx") - else: - # DSLs with neither spelling get the plain (IEEE) - # sqrt, which is strictly MORE accurate than the - # measured approx choice above; the unit test's - # 5e-3 oracle tolerance absorbs the difference. - magnitude = cute.math.sqrt(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 + 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() @@ -1124,10 +1141,7 @@ def kernel( cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) acc_pipeline.consumer_wait(acc_consumer_state) output_offset = ( - kv_head * self.group_size * self.sum_seq - + out_base - + page_start - + page_half * CTA_M + kv_head * N * self.sum_seq + out_base + page_start + page_half * CTA_M ) page_output = cute.make_tensor( output.iterator + output_offset, @@ -1136,7 +1150,7 @@ def kernel( stride=( 1, self.sum_seq, - self.group_size * self.sum_seq, + N * self.sum_seq, ), ), ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 08aa2c46dc17..876bd035b2e3 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -548,13 +548,10 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> any unsupported geometry this returns without importing the CuTe module, and every launch keeps using the compiled C++ score ops. - Geometry reality check: the supported contract below (SM100 exactly, - BF16 pools, 128-token pages, 64-element K rows / 32 frequencies, - 8 query heads per KV head) matches none of the current production - models (Qwen3 uses 128-element K rows, 32-token pages, and 4 query - heads per KV head; GPT-OSS uses 32-token pages), so today the kernel - fires only on the synthetic unit-test geometry. This wiring exists to - validate the kernel end to end while wider geometry support lands. + Supported contract: SM100 exactly, BF16 pools, 32- or 128-token + pages, 32 or 64 frequencies (head size 64/128), and 4 or 8 query + heads per KV head — this covers the Qwen3 and GPT-OSS production + geometries as well as the original validation shape. """ if self._cute_score_attempted: return @@ -570,10 +567,12 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> and kv_factor == 2 and tokens_per_block in (32, 128) and num_freqs in (32, 64) - and num_q_heads == num_kv_heads * 8 + and num_q_heads % num_kv_heads == 0 + and num_q_heads // num_kv_heads in (4, 8) and int(anchor.stride(-1)) == 1 - # The kernel computes flat score offsets in 32-bit arithmetic. - and num_q_heads * max_segments * self.seq_len < 2**31 + # The kernel computes flat score offsets in 32-bit arithmetic; + # group-4 geometries pad the head axis to the MMA tile N=8. + and num_kv_heads * 8 * max_segments * self.seq_len < 2**31 ) if not supported: return @@ -587,8 +586,11 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> # decode window from that scratch into ``self.output``. All # buffers below are persistent because the compiled kernel # captures their device pointers. + # The kernel writes one scratch row per padded head column + # (GQA group below 8 pads up to the MMA tile); the gather in + # ``launch`` reads only the real heads. scratch = torch.empty( - num_q_heads * max_segments * self.seq_len, + num_kv_heads * 8 * max_segments * self.seq_len, dtype=torch.float32, device=device, ) @@ -724,19 +726,40 @@ def launch( # past a request's valid width carry unscored scratch data, # matching the C++ op, whose consumers mask by valid width. num_q_heads = int(self.geometry_args[0]) + num_kv_heads = int(self.geometry_args[1]) + group_size = num_q_heads // num_kv_heads + # The scratch head axis is padded to the MMA tile N=8 per + # KV head; slicing the view to the real group size skips + # the zero padding columns. source = ( - self._cute_scratch[: num_q_heads * num_segments * self.seq_len] - .view(num_q_heads, request_count, self.num_layers, self.seq_len) - .permute(1, 2, 0, 3) - ) - columns = ( - token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1) - + self._cute_gather_columns + self._cute_scratch[: num_kv_heads * 8 * num_segments * self.seq_len] + .view(num_kv_heads, 8, request_count, self.num_layers, self.seq_len)[ + :, :group_size + ] + .permute(2, 3, 0, 1, 4) ) + columns = token_starts_device[:request_count].to(torch.int64).view( + -1, 1, 1, 1, 1 + ) + self._cute_gather_columns.view(1, 1, 1, 1, -1) columns = columns.clamp_(max=self.seq_len - 1).expand( - request_count, self.num_layers, num_q_heads, self.output_width + request_count, + self.num_layers, + num_kv_heads, + group_size, + self.output_width, + ) + torch.gather( + source, + 4, + columns, + out=output.view( + request_count, + self.num_layers, + num_kv_heads, + group_size, + self.output_width, + ), ) - torch.gather(source, 3, columns, out=output) return output _launch_tri_score_perhead( self, 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 index 9349964b90e4..3d9c339b1eb7 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -11,16 +11,20 @@ reason="TriAttention CuTe score kernel requires SM100", ) @pytest.mark.parametrize( - "tokens_per_block,page_permutation,valid_lens", + "tokens_per_block,page_permutation,valid_lens,num_freqs,num_q_heads", [ # The originally validated geometry: 128-token pages, identity table. - (128, [0, 1], None), - # Production page size: a 64-token compute tile spans two pages, so a - # shuffled physical-page table catches any fragment/page mix-up. - (32, [3, 1, 4, 7, 5, 0, 2, 6], None), + (128, [0, 1], None, 32, 8), + # GPT-OSS geometry: 32-token pages; a 64-token compute tile spans two + # pages, so a shuffled physical-page table catches fragment mix-ups. + (32, [3, 1, 4, 7, 5, 0, 2, 6], None, 32, 8), # Ragged tails land mid-tile: the second page fragment of the last # tile is clamped, and scores past the valid length are unspecified. - (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198]), + (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198], 32, 8), + # Qwen3 geometry: 128-element K rows (64 frequencies) and GQA group + # 4, which rides the MMA tile N=8 with zeroed padding columns. + (32, [3, 1, 4, 7, 5, 0, 2, 6], None, 64, 4), + (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198], 64, 4), ], ) def test_cute_score_matches_torch_mean_oracle( @@ -28,6 +32,8 @@ def test_cute_score_matches_torch_mean_oracle( tokens_per_block: int, page_permutation: list, valid_lens: "list | None", + num_freqs: int, + num_q_heads: int, ) -> None: pytest.importorskip("cutlass") monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_SCORE", "1") @@ -39,8 +45,6 @@ def test_cute_score_matches_torch_mean_oracle( torch.manual_seed(20260720) device = torch.device("cuda") seq_len = 256 - num_q_heads = 8 - num_freqs = 32 num_pages = seq_len // tokens_per_block assert sorted(page_permutation) == list(range(num_pages)) pool = ( From 104f1c1fadcc84eaf413e799420c05d9d9c45789 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 21:29:43 -0700 Subject: [PATCH 061/178] [None][chore] Log when the CuTe score path engages One INFO line with the accepted geometry, so end-to-end runs can prove which score kernel actually served them instead of inferring it. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 876bd035b2e3..581b083e1f69 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -630,6 +630,12 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> self._cute_scratch = scratch self._cute_seg_seq_len = seg_seq_len self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) + from tensorrt_llm.logger import logger + + logger.info( + f"TriAttention CuTe score enabled: {num_q_heads}q/{num_kv_heads}kv heads, " + f"{num_freqs} freqs, {tokens_per_block}-token pages" + ) except (ImportError, RuntimeError, ValueError, AssertionError) as error: warnings.warn( f"TriAttention CuTe score setup failed; using the C++ score ops: {error}", From 9c0cf959639f590e6ae671a04e7ba7380e6cb178 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 21:57:17 -0700 Subject: [PATCH 062/178] [None][test] Migrate compaction fixtures to the shipped bf16 geometry The retired register-staging fallback used to absorb fp16/fp32 pools and tiny page sizes in these fixtures; with the pipelined bf16 kernels as the only path they fail loudly, and one such failure poisoned the whole suite: pytest repr-ed the raised constructor frame after the test teardown had unmapped the V2 pool, and the dangling view read killed the CUDA context for every later test. Fixtures move to bf16 pools on 32-token pages with head size 64, with payloads and keep sets rescaled so every byte-preservation, protected-tail, capacity, and page-reuse assertion still discriminates. Suite scope back to green (194 passed). Signed-off-by: tianruih --- .../test_triattention_draft_cocompaction.py | 49 +++-- .../test_triattention_fused_settle_pack.py | 9 +- .../test_triattention_selection_compaction.py | 197 ++++++++++++------ 3 files changed, 172 insertions(+), 83 deletions(-) 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 index 65162b9e109a..ae6ab1ed799c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -44,15 +44,22 @@ def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: def _launched_draft_compaction(draft_protected_tails): - """Build target and draft pools with distinct head counts, then compact.""" + """Build target and draft pools with distinct head counts, then compact. + + The compact op ships only the pipelined bf16 kernels, so the pools use + the supported production geometry (bf16, 32-token pages, head_dim 64). + Pool payloads are a shifted ``arange % 251`` ramp: every value is exact + in bf16 and any wrong page/plane/head/token move lands on a different + byte pattern, so the equality checks below stay conclusive. + """ device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 target_kv_heads = 2 draft_kv_heads = 4 prompt_len = 2 decode_keep_count = 4 - tokens_per_block = 4 - head_dim = 16 + tokens_per_block = 32 + head_dim = 64 target_protected_tails = [2, 1] valid_seq_lens = [10, 9] @@ -60,22 +67,34 @@ def _launched_draft_compaction(draft_protected_tails): draft_tables = torch.tensor([[1, 0, 2], [5, 4, 3]], dtype=torch.int32, device=device) target_pools = [ ( - torch.arange( - 6 * 2 * target_kv_heads * tokens_per_block * head_dim, - dtype=torch.float32, - device=device, - ).view(6, 2, target_kv_heads, tokens_per_block, head_dim) - + layer * 100_000.0 + ( + torch.arange( + 6 * 2 * target_kv_heads * tokens_per_block * head_dim, + dtype=torch.int32, + device=device, + ) + + layer * 37 + ) + % 251 ) + .view(6, 2, target_kv_heads, tokens_per_block, head_dim) + .to(torch.bfloat16) for layer in range(2) ] draft_pool = ( - torch.arange( - 6 * 2 * draft_kv_heads * tokens_per_block * head_dim, - dtype=torch.float32, - device=device, - ).view(6, 2, draft_kv_heads, tokens_per_block, head_dim) - + 900_000.0 + ( + ( + torch.arange( + 6 * 2 * draft_kv_heads * tokens_per_block * head_dim, + dtype=torch.int32, + device=device, + ) + + 149 + ) + % 251 + ) + .view(6, 2, draft_kv_heads, tokens_per_block, head_dim) + .to(torch.bfloat16) ) assert target_pools[0].shape[2] != draft_pool.shape[2] initial_target = [pool.clone() for pool in target_pools] 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 index ad3c64f81245..0bbc6cfeb0b1 100644 --- 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 @@ -353,9 +353,14 @@ def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buff own keep buffer.""" device = torch.device("cuda", torch.cuda.current_device()) request_count, num_kv_heads, keep_count, width = 2, 2, 4, 16 - tokens_per_block, head_dim = 4, 8 + # BatchedKVCacheCompaction admits only bf16 pools in the compact op's + # supported geometry (32/128-token pages, head_dim 64/128). + tokens_per_block, head_dim = 32, 64 pools = [ - torch.zeros(6, 2, num_kv_heads, tokens_per_block, head_dim, device=device) for _ in range(2) + torch.zeros( + 6, 2, num_kv_heads, tokens_per_block, head_dim, dtype=torch.bfloat16, device=device + ) + for _ in range(2) ] page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) 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 index 0bd6606cd423..17e1f9dc0843 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -373,34 +373,46 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode): + # The compact op ships only the pipelined bf16 kernels: pools use the + # supported geometry (bf16, 32-token pages, head_dim 64), and the kept + # ordinals are spread across all three pages per request so the moves + # still cross page boundaries. The bf16-exact ``arange % 251`` payload + # keeps every wrong-move byte pattern distinguishable. device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 num_layers = 2 num_kv_heads = 2 prompt_len = 2 decode_keep_count = 4 - seq_len = 10 - tokens_per_block = 4 + seq_len = 80 + tokens_per_block = 32 pages_per_request = 3 - head_dim = 16 + head_dim = 64 protected_tails = [2, 1] page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) initial_pools = [ ( - torch.arange( - 6 * 2 * num_kv_heads * tokens_per_block * head_dim, - dtype=torch.float32, - device=device, - ).view(6, 2, num_kv_heads, tokens_per_block, head_dim) - + layer * 100_000.0 + ( + torch.arange( + 6 * 2 * num_kv_heads * tokens_per_block * head_dim, + dtype=torch.int32, + device=device, + ) + + layer * 37 + ) + % 251 ) + .view(6, 2, num_kv_heads, tokens_per_block, head_dim) + .to(torch.bfloat16) for layer in range(num_layers) ] pools = [pool.clone() for pool in initial_pools] # Kept ordinals are decode-only but hold absolute positions; the pinned # prompt tokens never appear in the selection rectangle. - union_decode = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) + union_decode = torch.tensor( + [[16, 32, 56, 72], [24, 40, 48, 64]], dtype=torch.int64, device=device + ) if eviction_mode == "union": keep = union_decode selection_rows = 1 @@ -418,7 +430,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) keep[request, row] = torch.tensor( sorted( { - 2 + ((request + row + offset * 2) % 8) + 2 + ((request + row + offset * 2) % 8) * 8 for offset in range(decode_keep_count) } ), @@ -490,12 +502,14 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): request_count = 2 num_layers = 2 num_kv_heads = 2 - seq_len = 10 + # bf16 pools in the compact op's supported geometry; three 32-token pages + # per request with a 80-token sequence keep the moves page-crossing. + seq_len = 80 decode_keep_count = 3 prompt_lens = [2, 5] protected_tails = [2, 1] - tokens_per_block = 4 - head_dim = 16 + tokens_per_block = 32 + head_dim = 64 decode_widths = [seq_len - prompt_len for prompt_len in prompt_lens] width = max(decode_widths) @@ -521,13 +535,18 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) initial_pools = [ ( - torch.arange( - 6 * 2 * num_kv_heads * tokens_per_block * head_dim, - dtype=torch.float32, - device=device, - ).view(6, 2, num_kv_heads, tokens_per_block, head_dim) - + layer * 100_000.0 + ( + torch.arange( + 6 * 2 * num_kv_heads * tokens_per_block * head_dim, + dtype=torch.int32, + device=device, + ) + + layer * 37 + ) + % 251 ) + .view(6, 2, num_kv_heads, tokens_per_block, head_dim) + .to(torch.bfloat16) for layer in range(num_layers) ] cohort_pools = [pool.clone() for pool in initial_pools] @@ -592,8 +611,12 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): num_layers = 3 seq_len = 8 keep_count = 2 - tokens_per_block = 4 - head_dim = 16 + # bf16 pools in the compact op's supported geometry. The scored tokens + # all live in each table's first entry, but the two tables still map the + # two storage groups onto different physical pages, which is what the + # layer-order alignment below depends on. + tokens_per_block = 32 + head_dim = 64 num_freqs = head_dim // 2 dense_layers = [0, 1, 2] dense_groups = [[0, 2], [1]] @@ -613,12 +636,19 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): pools = [] for layer, (table, values) in enumerate(zip(layer_tables, score_values)): pool = ( - torch.arange( - 2 * 2 * tokens_per_block * head_dim, - dtype=torch.float32, - device=device, - ).view(2, 2, 1, tokens_per_block, head_dim) - + layer * 10_000 + ( + ( + torch.arange( + 2 * 2 * tokens_per_block * head_dim, + dtype=torch.int32, + device=device, + ) + + layer * 37 + ) + % 251 + ) + .view(2, 2, 1, tokens_per_block, head_dim) + .to(torch.bfloat16) ) for token, value in enumerate(values): page = int(table[0, token // tokens_per_block]) @@ -701,14 +731,24 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): zip(initial_pools, pools, layer_tables) ): pages = table[0].to(torch.long) - before = before_pool[pages].permute(1, 2, 0, 3, 4).reshape(2, 1, seq_len, head_dim) + # 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)) def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): - """Run two real eviction rounds through one live V2 cache.""" + """Run two real eviction rounds through one live V2 cache. + + The cache uses the compact op's supported geometry (bf16, 32-token + pages, head_dim 64): the request spans three pages so that compacting to + two pages still releases one physical page for reuse. Token scores are + tracked in a host-side mirror and the expected keep sets are derived + from it, replacing the hand-written score tables of the old 4-token-page + fixture. + """ import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( @@ -721,12 +761,13 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): device = torch.device("cuda", torch.cuda.current_device()) request_id = 7 prompt_len = 2 - seq_len = 10 + seq_len = 66 protected_tail = 2 - compacted_capacity = 8 - tokens_per_block = 4 - head_dim = 16 + compacted_capacity = 34 + 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, @@ -742,7 +783,7 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): 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.HALF, + dtype=tensorrt_llm.bindings.DataType.BF16, vocab_size=128, ) @@ -780,18 +821,34 @@ 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 and distinct per token, so the + # byte comparisons below catch any wrong move. payload = ( - torch.arange(2 * head_dim, dtype=torch.float16, device=device) + ((torch.arange(2 * head_dim, dtype=torch.int32, device=device) + token * 37) % 251) .reshape(2, head_dim) - .add_(token * 64) + .to(torch.bfloat16) ) payload[0, 0] = score payload[0, num_freqs] = 0 pool[page, :, 0, offset].copy_(payload) - first_scores = [0, 0, 8, 1, 7, 2, 6, 3, 5, 4, 9, 0] - for token, score in enumerate(first_scores): - write_token(token, score) + # Host-side mirror of each physical position's score; expected keep + # sets are derived from it. Scores are distinct within the decode + # window (7 is invertible mod 64 and the window spans one residue + # cycle), so the selection is tie-free and deterministic. + 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, + ) q_real = torch.zeros(1, 1, num_freqs, dtype=torch.float32, device=device) q_imag = torch.zeros_like(q_real) @@ -822,7 +879,7 @@ def write_token(token: int, score: float) -> None: keep_set_selector = _BatchedUnionKeepSetSelector( rows=1, width=seq_len - prompt_len, - keep_count=compacted_capacity - prompt_len - protected_tail, + keep_count=keep_count, dtype=torch.float32, device=device, max_requests=1, @@ -847,7 +904,7 @@ def write_token(token: int, score: float) -> None: page_table_slots=score_staging.representative_slots, request_count=1, prompt_offsets=score_staging.token_starts_device[:1], - decode_keep_count=compacted_capacity - prompt_len - protected_tail, + decode_keep_count=keep_count, swa_window=None, protected_tail_capacity=protected_tail, ) @@ -888,11 +945,10 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: 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, - torch.tensor([2, 4, 6, 8], dtype=torch.long, device=device), - ) + 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:] @@ -909,18 +965,21 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: assert torch.equal(page_ids(request_id)[:2], retained_pages) assert torch.equal(page_ids(request_id)[2:], released_page) # The first protected tail becomes confirmed input to round two. Only - # later generated tokens and the next protected tail are written here. - write_token(8, 10) - write_token(9, 4.5) - write_token(10, 11) - write_token(11, 0.5) - assert torch.equal(snapshot(8), first_compacted) - + # later generated tokens and the next protected tail are written + # here. Mirror the physical relayout, then give the fresh tokens a + # disjoint higher score band (11 invertible mod 64 over a shorter + # window) so round two must select differently from round one. + 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, - torch.tensor([2, 3, 6, 8], dtype=torch.long, device=device), - ) + 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]) @@ -936,22 +995,26 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: def test_eager_compaction_rebases_masked_swa_window_and_tail(): + # bf16 pools in the compact op's supported geometry (32-token pages, + # head_dim 64); the kept ordinals and valid lengths span all three pages + # per request so the dense and SWA moves stay page-crossing. device = torch.device("cuda", torch.cuda.current_device()) dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) initial_pools = [ - torch.arange(6 * 2 * 1 * 4 * 16, dtype=torch.float32, device=device).view(6, 2, 1, 4, 16), - torch.arange(6 * 2 * 1 * 4 * 16, dtype=torch.float32, device=device).view(6, 2, 1, 4, 16) - + 1000.0, + ((torch.arange(6 * 2 * 1 * 32 * 64, dtype=torch.int32, device=device) + layer * 37) % 251) + .view(6, 2, 1, 32, 64) + .to(torch.bfloat16) + for layer in range(2) ] pools = [pool.clone() for pool in initial_pools] # Decode-only kept ordinals holding absolute positions past the prompt. keep = torch.tensor( - [[2, 4, 5, 7], [2, 3, 5, 6]], + [[16, 32, 40, 56], [16, 24, 40, 48]], dtype=torch.int64, device=device, ) - valid_seq_lens = torch.tensor([8, 7], dtype=torch.int32, device=device) + valid_seq_lens = torch.tensor([64, 56], dtype=torch.int32, device=device) protected_tails = [2, 1] compaction = BatchedKVCacheCompaction( eviction_mode="union", @@ -979,9 +1042,9 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): ): dense_pages = dense_tables[request].to(torch.long) swa_pages = swa_tables[request].to(torch.long) - dense_before = initial_pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 16) + dense_before = initial_pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) dense_after = pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape_as(dense_before) - swa_before = initial_pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 16) + swa_before = initial_pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) swa_after = pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape_as(swa_before) tail = torch.arange( valid_seq_len, @@ -1024,9 +1087,11 @@ def test_cache_families_read_the_staged_move_offsets_rows(): device = torch.device("cuda", torch.cuda.current_device()) dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) + # bf16 pools in the compact op's supported geometry; this test only + # constructs the compaction, so the contents stay zero. pools = [ - torch.zeros(6, 2, 1, 4, 16, dtype=torch.float32, device=device), - torch.zeros(6, 2, 1, 4, 16, dtype=torch.float32, device=device), + torch.zeros(6, 2, 1, 32, 64, dtype=torch.bfloat16, device=device), + torch.zeros(6, 2, 1, 32, 64, dtype=torch.bfloat16, device=device), ] keep = torch.tensor([[2, 4, 5, 7], [2, 3, 5, 6]], dtype=torch.int32, device=device) staged_rows = torch.zeros(2, 3, dtype=torch.int32, device=device) From 727ae54e3451812aab4bf1d2d14d52700626033f Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 21:59:39 -0700 Subject: [PATCH 063/178] [None][chore] Warn when the opt-in CuTe score cannot engage With the knob explicitly on, a silent geometry rejection or a compiled- variant miss at launch is indistinguishable from success; both now warn once with the concrete values, and the first engaged launch says so. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 581b083e1f69..1dc77c6676f9 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -575,6 +575,16 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> and num_kv_heads * 8 * max_segments * self.seq_len < 2**31 ) if not supported: + warnings.warn( + "TriAttention CuTe score not engaged despite " + "TRTLLM_TRIATTENTION_CUTE_SCORE=1: " + f"capability={torch.cuda.get_device_capability(anchor.device)}, " + f"dtype={anchor.dtype}, kv_factor={kv_factor}, " + f"tokens_per_block={tokens_per_block}, num_freqs={num_freqs}, " + f"heads={num_q_heads}q/{num_kv_heads}kv, " + f"stride={int(anchor.stride(-1))}, " + f"offset_audit={num_kv_heads * 8 * max_segments * self.seq_len}" + ) return device = anchor.device try: @@ -714,7 +724,22 @@ def launch( # outside CUDA graph capture. Default off: one attribute check. self.prepare_cute_score(mean_cos, mean_sin) runner = self._cute_score_runner + if runner is not None and not runner.supports(request_count): + missed = getattr(self, "_cute_supports_warned", set()) + if request_count not in missed: + missed.add(request_count) + self._cute_supports_warned = missed + warnings.warn( + f"TriAttention CuTe score compiled variants miss " + f"request_count={request_count}; falling back to the C++ " + f"score op for this round" + ) if runner is not None and runner.supports(request_count): + if not getattr(self, "_cute_engaged_logged", False): + self._cute_engaged_logged = True + warnings.warn( + f"TriAttention CuTe score engaged (request_count={request_count})" + ) # Stage per-segment valid lengths (segment = request x layer). torch.index_select( valid_seq_lens, From 2012bdee01385bbf487dd734048a55435efa0a1b Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 23:07:30 -0700 Subject: [PATCH 064/178] [None][feat] Make the CuTe DSL kernel the only TriAttention score path The C++ CUDA score stack (paged score, coefficient rotation, coefficient fold, their torch ops and phase-table plumbing) is deleted; the SM100 CuTe DSL kernel serves every supported geometry (bf16 pools, head size 64/128, 32/128-token pages, GQA group 4 or 8) and anything outside that contract raises at setup instead of routing to a slower kernel. Max aggregation and the quantized-pool coefficient path go with the stack: neither was reachable from the user-facing config. Two latent gaps surfaced by the swap are fixed and pinned by tests: the CuTe path now writes the per-request decode widths the selection kernels consume, and score bucket capacities are rounded up to the kernel compute tile whose unmasked tail stores would otherwise spill into the next segment. Score tests move onto the CuTe contract (SM100-only, matching CI). Signed-off-by: tianruih --- .../triAttentionScoreKernels.cu | 949 ------------------ .../triAttentionScoreKernels.h | 213 ---- cpp/tensorrt_llm/thop/CMakeLists.txt | 3 +- cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp | 452 --------- .../triattention/triattention.py | 66 +- .../triattention/triattention_cute_score.py | 20 +- .../triattention/triattention_kernels.py | 593 +++-------- .../_torch/kv_cache_compression/conftest.py | 28 +- .../test_triattention_cute_score.py | 10 +- .../test_triattention_phase_rotation.py | 236 ----- .../test_triattention_pipeline.py | 190 ++-- .../test_triattention_score_ops.py | 656 ++++-------- .../test_triattention_selection_compaction.py | 64 +- 13 files changed, 516 insertions(+), 2964 deletions(-) delete mode 100644 cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu delete mode 100644 cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h delete mode 100644 cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp delete mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu deleted file mode 100644 index 381f9bcb4a4f..000000000000 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.cu +++ /dev/null @@ -1,949 +0,0 @@ -/* - * 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. - */ - -// ============================================================================ -// TriAttention folded score kernels -// ============================================================================ -// -// Scores every cached decode token of every scored layer for KV eviction. -// The per-round trigonometry is folded into coefficient tables first -// (foldScoreCoefficientsLaunch), so the hot kernel is a pure fused -// multiply-add stream over paged KV: -// -// score(t, h) = sum_f K_re(t,f)*c_re + K_im(t,f)*c_im + |K(t,f)|*c_mlr -// -// One thread scores one token across all frequencies; a 128-thread CTA covers -// 128 consecutive tokens of one (request, layer, KV-head) segment. The hot -// loop has no shuffles and no barriers: each thread keeps one fused -// accumulator per query head of its GQA group ("mean" aggregation) or one -// partial sum per offset plane ("max" aggregation, where max over offsets -// does not commute through the frequency sum). Coefficient loads are -// lane-uniform 16-byte reads served by L1 broadcast (or by shared memory on -// the mean path's default in-CTA rotation, whose prologue is the one place a -// score CTA uses shared memory and a barrier); K loads are 16-byte chunks of -// 8 frequencies when the pool layout allows it, otherwise a fully strided -// scalar path runs the same math. -// -// The kernel accumulates the frequency reduction in sequential chunks (not a -// block-wide tree), so results are tolerance-equal, not bit-equal, against -// the unit tests' PyTorch oracle (the in-tree reference; the original Triton -// score kernel has been deleted). The valid-width side store IS exact: first -// tile, first head, first segment of each request, thread 0, before any -// early-out. -// -// This file must NOT be compiled with --use_fast_math: the fold kernel's -// cosf/sinf and the scalar-path precision are part of the accuracy contract. -// The approximate square root below is an explicit, scoped opt-in instead. -// ============================================================================ - -#include -#include -#include -#include -#include -#include - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::tri_attention_score -{ - -namespace -{ - -// |K| uses the hardware approximate square root (one MUFU op), gated by the -// unit suite's tolerance comparison against the PyTorch oracle (this matched -// the deleted Triton score kernel's tl.sqrt lowering). Define -// TRTLLM_TRI_ATTENTION_IEEE_SQRT to restore the IEEE sqrtf sequence if a -// future geometry needs the extra bits. -__device__ __forceinline__ float triSqrtApprox(float x) -{ -#ifdef TRTLLM_TRI_ATTENTION_IEEE_SQRT - return sqrtf(x); -#else - float y; - asm("sqrt.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); - return y; -#endif -} - -// The ONE mean-path coefficient rotation, compiled by BOTH the standalone -// triRotateMeanScoreCoefficientsKernel and the score kernels' in-CTA -// prologue. The operand order is load-bearing: a single source expression -// makes the compiler emit the same multiply/FMA contraction in every -// inlining site, so the in-CTA coefficients are bit-identical to the -// standalone kernel's output and end-to-end mean scores are bit-identical -// between the two preparation paths (the unit suite proves this with -// torch.equal). Do not reorder or "simplify" these expressions. -__device__ __forceinline__ float2 rotateCoefficientPair(float qre, float qim, float pc, float ps) -{ - float2 c; - c.x = qre * pc - qim * ps; - c.y = qim * pc + qre * ps; - return c; -} - -template -__device__ __forceinline__ float toFloat(T value); - -template <> -__device__ __forceinline__ float toFloat<__nv_bfloat16>(__nv_bfloat16 value) -{ - return __bfloat162float(value); -} - -template <> -__device__ __forceinline__ float toFloat(half value) -{ - return __half2float(value); -} - -template <> -__device__ __forceinline__ float toFloat(float value) -{ - return value; -} - -// Quantized pool elements are converted RAW (no scale applied): the per-layer -// dequantization scale is folded into the coefficient tables at fold time, so -// the hot loop stays a pure convert-and-FMA stream. -template <> -__device__ __forceinline__ float toFloat<__nv_fp8_e4m3>(__nv_fp8_e4m3 value) -{ - return static_cast(value); -} - -template <> -__device__ __forceinline__ float toFloat(int8_t value) -{ - return static_cast(value); -} - -// Unpack one 16-byte K chunk (8 consecutive frequencies) to fp32. -template -__device__ __forceinline__ void unpackChunk8(uint4 v, float* dst); - -template <> -__device__ __forceinline__ void unpackChunk8<__nv_bfloat16>(uint4 v, float* dst) -{ - auto const* p = reinterpret_cast<__nv_bfloat162 const*>(&v); -#pragma unroll - for (int i = 0; i < 4; ++i) - { - float2 f2 = __bfloat1622float2(p[i]); - dst[2 * i] = f2.x; - dst[2 * i + 1] = f2.y; - } -} - -template <> -__device__ __forceinline__ void unpackChunk8(uint4 v, float* dst) -{ - auto const* p = reinterpret_cast<__half2 const*>(&v); -#pragma unroll - for (int i = 0; i < 4; ++i) - { - float2 f2 = __half22float2(p[i]); - dst[2 * i] = f2.x; - dst[2 * i + 1] = f2.y; - } -} - -// Predicated 16-byte loads of one 8-frequency chunk of one token row. Row -// layout with unit frequency stride is [re f0..F-1 | im f0..F-1], so the -// imaginary half of chunk c sits sizeof(T) * numFreqs bytes further. -template -__device__ __forceinline__ void scoreLoadChunk(char const* row, bool valid, int numFreqs, int c, uint4& re4, uint4& im4) -{ - re4 = make_uint4(0u, 0u, 0u, 0u); - im4 = make_uint4(0u, 0u, 0u, 0u); - if (valid) - { - int const fByte = c * 8 * static_cast(sizeof(T)); - re4 = __ldg(reinterpret_cast(row + fByte)); - im4 = __ldg(reinterpret_cast(row + static_cast(sizeof(T)) * numFreqs + fByte)); - } -} - -// Mean-path in-CTA rotation prologue shared by the vectorized and scalar -// score kernels: rebuild the [coefCount = group * numFreqs] coefficient block -// this CTA will read — gather the request's phase-table row (L2-resident, -// shared by all of the request's CTAs) and the CTA's slice of the static -// pre-scaled calibration query, rotate with rotateCoefficientPair (the SAME -// arithmetic the standalone rotation kernel compiles, the bit-equality -// contract), and stage the results in shared memory. Thread-strided plain -// loads + FMAs only, one barrier; the prologue's temporaries die here, so -// ptxas can fold them into the main loop's register budget. -__device__ __forceinline__ void rotateCoefficientsIntoShared( - FoldedScoreParams const& a, int reqId, int layerId, int headBase, int coefCount, float* sCRe, float* sCIm) -{ - int64_t const phaseRow = static_cast(a.roundStarts[reqId]) * a.numFreqs; - int64_t const qBase = (static_cast(layerId) * a.numQueryHeads + headBase) * a.numFreqs; - for (int i = threadIdx.x; i < coefCount; i += kScoreBlockThreads) - { - int const f = i % a.numFreqs; - float const pc = a.phaseCos[phaseRow + f]; - float const ps = a.phaseSin[phaseRow + f]; - float2 const c = rotateCoefficientPair(a.qReS[qBase + i], a.qImS[qBase + i], pc, ps); - sCRe[i] = c.x; - sCIm[i] = c.y; - } - __syncthreads(); -} - -// Accumulate one 8-frequency chunk into this thread's per-head accumulators. -// cRe/cIm point at the coefficient source: the global per-round tables (with -// coff0 = flat index of (request, layer, first head of this block, chunk -// frequency 0)) or, under ROTATE_IN_CTA, this CTA's shared-memory block -// (with coff0 = the chunk's frequency offset — the head-block origin is the -// shared block itself). cMlr = the matching chunk pointer into the static -// request-independent [layer, head, freq] c_mlr table, pre-offset by the -// caller. Passing the resolved pointer (instead of a second flat offset) -// keeps this loop's live set at the tuned ~72-register baseline: one pointer -// register replaces the base + request-strided offset pair the c_mlr reads -// consumed before the table went static. All coefficient reads are -// lane-uniform 16-byte loads (__ldg is global-memory-only, so the shared -// block reads through plain loads — same values either way). |K| is computed -// once per (token, frequency) BEFORE the head loop so the GROUP heads share -// it from registers. -template -__device__ __forceinline__ void scoreComputeChunk(FoldedScoreParams const& a, float const* cRe, float const* cIm, - int64_t coff0, float const* cMlr, int64_t planeStride, uint4 re4, uint4 im4, float* accMean, float* accMlr, - float* accPos) -{ - static_assert(!(ROTATE_IN_CTA && USE_MAX), "in-CTA coefficient rotation exists only on the mean path"); - float kRe[8], kIm[8], kMag[8]; - unpackChunk8(re4, kRe); - unpackChunk8(im4, kIm); -#pragma unroll - for (int i = 0; i < 8; ++i) - { - kMag[i] = triSqrtApprox(kRe[i] * kRe[i] + kIm[i] * kIm[i]); - } -#pragma unroll - for (int hg = 0; hg < GROUP; ++hg) - { - int64_t const coff = coff0 + static_cast(hg) * a.numFreqs; - float const* cmp = cMlr + static_cast(hg) * a.numFreqs; - float4 const cm0 = __ldg(reinterpret_cast(cmp)); - float4 const cm1 = __ldg(reinterpret_cast(cmp + 4)); - float const cml[8] = {cm0.x, cm0.y, cm0.z, cm0.w, cm1.x, cm1.y, cm1.z, cm1.w}; - if constexpr (USE_MAX) - { - // The |K| term is offset independent: keep it in its own - // accumulator and add it after the offset max at store time. - float m = accMlr[hg]; -#pragma unroll - for (int i = 0; i < 8; ++i) - { - m = fmaf(kMag[i], cml[i], m); - } - accMlr[hg] = m; - // Unrolled with a live-plane guard so accPos indexing stays - // static (register-resident) despite the runtime offset count. -#pragma unroll - for (int o = 0; o < kMaxScoreOffsets; ++o) - { - if (o < a.numOffsets) - { - float const* crp = cRe + o * planeStride + coff; - float const* cip = cIm + o * planeStride + coff; - float4 const cr0 = __ldg(reinterpret_cast(crp)); - float4 const cr1 = __ldg(reinterpret_cast(crp + 4)); - float4 const ci0 = __ldg(reinterpret_cast(cip)); - float4 const ci1 = __ldg(reinterpret_cast(cip + 4)); - float const cre[8] = {cr0.x, cr0.y, cr0.z, cr0.w, cr1.x, cr1.y, cr1.z, cr1.w}; - float const cim[8] = {ci0.x, ci0.y, ci0.z, ci0.w, ci1.x, ci1.y, ci1.z, ci1.w}; - float p = accPos[hg * kMaxScoreOffsets + o]; -#pragma unroll - for (int i = 0; i < 8; ++i) - { - p = fmaf(kRe[i], cre[i], fmaf(kIm[i], cim[i], p)); - } - accPos[hg * kMaxScoreOffsets + o] = p; - } - } - } - else - { - float const* crp = cRe + coff; - float const* cip = cIm + coff; - float4 cr0, cr1, ci0, ci1; - if constexpr (ROTATE_IN_CTA) - { - // Shared-memory coefficient block written by this CTA's - // rotation prologue (16-byte aligned: numFreqs % 8 == 0 on - // the vectorized path). - cr0 = *reinterpret_cast(crp); - cr1 = *reinterpret_cast(crp + 4); - ci0 = *reinterpret_cast(cip); - ci1 = *reinterpret_cast(cip + 4); - } - else - { - cr0 = __ldg(reinterpret_cast(crp)); - cr1 = __ldg(reinterpret_cast(crp + 4)); - ci0 = __ldg(reinterpret_cast(cip)); - ci1 = __ldg(reinterpret_cast(cip + 4)); - } - float const cre[8] = {cr0.x, cr0.y, cr0.z, cr0.w, cr1.x, cr1.y, cr1.z, cr1.w}; - float const cim[8] = {ci0.x, ci0.y, ci0.z, ci0.w, ci1.x, ci1.y, ci1.z, ci1.w}; - float t = accMean[hg]; -#pragma unroll - for (int i = 0; i < 8; ++i) - { - // 3 fused multiply-adds per (token, frequency, head): the - // position and |K| terms share one accumulator chain, valid - // because the mean path has a single coefficient plane. - t = fmaf(kRe[i], cre[i], fmaf(kIm[i], cim[i], fmaf(kMag[i], cml[i], t))); - } - accMean[hg] = t; - } - } -} - -// Vectorized token-per-thread score kernel. Grid: x = 128-token tiles of the -// page-aligned decode span, y = (request, layer) segment, z = KV head (or one -// query head when a.zIsQueryHead covers GQA group sizes with no dedicated -// GROUP instantiation). STATIC_CHUNKS == 8 pins the production 64-frequency -// shape at compile time (fully unrolled chunk loop, the tuned register -// budget); STATIC_CHUNKS == 0 loops numFreqs / 8 chunks at runtime. -// -// minBlocksPerMultiprocessor = 7: tighter caps force ptxas into ~48-56 -// registers with stack spills in the fully unrolled inner loops; 7 CTAs/SM -// admits the ~72-register spill-free allocation this kernel was tuned at. -// -// ROTATE_IN_CTA (mean path only) prepends the shared-memory coefficient -// rotation prologue and points the mean coefficient reads at it; the -// standalone-rotation read path (ROTATE_IN_CTA == false) stays compiled for -// the max aggregation and for the unit tests' bit-equality reference leg. -template -__global__ void __launch_bounds__(kScoreBlockThreads, 7) triScoreVectorizedKernel(FoldedScoreParams a) -{ - static_assert(!(ROTATE_IN_CTA && USE_MAX), "in-CTA coefficient rotation exists only on the mean path"); - int const seg = blockIdx.y; - int const reqId = a.segRequestIds[seg]; - int const seqLen = a.requestSeqLens[reqId]; - int const tokenStart = a.requestTokenStarts[reqId]; - // Valid-width side store: evaluated before any early-out, so it fires - // exactly once per request even when the decode region is empty. - if (blockIdx.x == 0 && blockIdx.z == 0 && (seg % a.numLayers) == 0 && threadIdx.x == 0) - { - a.validWidthOut[reqId] = seqLen - tokenStart; - } - // Tiles start on the first page of the decode region: with 32-token pages - // every warp then covers exactly one page and its lane-identical page - // lookup collapses to one L1 broadcast. - int const alignedStart = (tokenStart / a.tokensPerBlock) * a.tokensPerBlock; - if (alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads >= seqLen) - { - return; // CTA-uniform: the whole tile is past this sequence - } - int const absT = alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads + static_cast(threadIdx.x); - bool const valid = absT >= tokenStart && absT < seqLen; - - int kvHead; - int headBase; - if (a.zIsQueryHead) - { - headBase = static_cast(blockIdx.z); - kvHead = headBase / (a.numQueryHeads / a.numKvHeads); - } - else - { - kvHead = static_cast(blockIdx.z); - headBase = kvHead * GROUP; - } - - int const layerId = a.segLayerIds[seg]; - // Mean-path in-CTA rotation: stage this CTA's GROUP x numFreqs - // coefficient block in shared memory (the early-out above is - // CTA-uniform, so the barrier inside is safe here and skipped CTAs do no - // rotation work). - float* sCRe = nullptr; - float* sCIm = nullptr; - if constexpr (ROTATE_IN_CTA) - { - extern __shared__ float coefficientSmem[]; - int const coefCount = GROUP * a.numFreqs; - sCRe = coefficientSmem; - sCIm = coefficientSmem + coefCount; - rotateCoefficientsIntoShared(a, reqId, layerId, headBase, coefCount, sCRe, sCIm); - } - int const page = absT / a.tokensPerBlock; - int const slot = absT - page * a.tokensPerBlock; - // Threads past the sequence tail must not touch the page table (their - // page ordinal may exceed the staged row); their K loads are predicated - // off below, so page 0 is a safe placeholder. - int encoded = 0; - if (absT < seqLen) - { - encoded = a.blockOffsets[a.segPageOffsets[seg] + page]; - } - // Page-table entries count K/V role pages; kvFactor converts to the - // layer pool page holding the K plane. - auto const physPage = static_cast(encoded / a.kvFactor); - auto const* layerBase = reinterpret_cast(a.layerBaseAddrs[layerId]); - char const* row = layerBase - + static_cast(sizeof(T)) - * (physPage * a.stridePage + static_cast(kvHead) * a.strideKvHead - + static_cast(slot) * a.strideSlot); - - // Coefficient source: the global per-round tables, or this CTA's shared - // block (whose head-block origin is index 0). - float const* coefRe = a.cRe; - float const* coefIm = a.cIm; - int64_t coff0 = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs - + static_cast(headBase) * a.numFreqs; - if constexpr (ROTATE_IN_CTA) - { - coefRe = sCRe; - coefIm = sCIm; - coff0 = 0; - } - // The MLR coefficient is position independent, so its table is folded once - // at initialization without a request axis: [layer, head, freq]. Hoist the - // row pointer here (no request stride) so the fully unrolled chunk loop - // sees a single pre-offset pointer, not an extra live flat offset. - float const* const cMlrRow = a.cMlr + (static_cast(layerId) * a.numQueryHeads + headBase) * a.numFreqs; - int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads - * static_cast(a.numFreqs); - - float accMean[GROUP]; - float accMlr[USE_MAX ? GROUP : 1]; - float accPos[USE_MAX ? GROUP * kMaxScoreOffsets : 1]; -#pragma unroll - for (int hg = 0; hg < GROUP; ++hg) - { - accMean[hg] = 0.0f; - } - if constexpr (USE_MAX) - { -#pragma unroll - for (int hg = 0; hg < GROUP; ++hg) - { - accMlr[hg] = 0.0f; - } -#pragma unroll - for (int i = 0; i < GROUP * kMaxScoreOffsets; ++i) - { - accPos[i] = 0.0f; - } - } - - if constexpr (STATIC_CHUNKS > 0) - { -#pragma unroll - for (int c = 0; c < STATIC_CHUNKS; ++c) - { - uint4 re4, im4; - scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); - scoreComputeChunk( - a, coefRe, coefIm, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); - } - } - else - { - int const chunkCount = a.numFreqs / 8; - for (int c = 0; c < chunkCount; ++c) - { - uint4 re4, im4; - scoreLoadChunk(row, valid, a.numFreqs, c, re4, im4); - scoreComputeChunk( - a, coefRe, coefIm, coff0 + c * 8, cMlrRow + c * 8, planeStride, re4, im4, accMean, accMlr, accPos); - } - } - - // Store: per (tile, head) the CTA writes one contiguous fp32 row; the - // predicate replicates the reference token mask and decode-region clip. - int const tDec = absT - tokenStart; - if (tDec >= 0 && absT < seqLen) - { -#pragma unroll - for (int hg = 0; hg < GROUP; ++hg) - { - float score; - if constexpr (USE_MAX) - { - float best = -INFINITY; -#pragma unroll - for (int o = 0; o < kMaxScoreOffsets; ++o) - { - if (o < a.numOffsets) - { - best = fmaxf(best, accPos[hg * kMaxScoreOffsets + o]); - } - } - score = best + accMlr[hg]; - } - else - { - score = accMean[hg]; - } - a.out[(static_cast(seg) * a.numQueryHeads + headBase + hg) * a.outputWidth + tDec] = score; - } - } -} - -// Strided scalar score kernel: same token-per-thread mapping and math as the -// vectorized kernel, but every K element is loaded through the full runtime -// stride set (any frequency count, any element stride, fp32 pools included). -// The GQA head loop runs at runtime, so any group size is covered. |K| is -// recomputed per head from the same loads — bit-identical to hoisting it. -// ROTATE_IN_CTA carries the same mean-path shared-memory rotation prologue -// as the vectorized kernel (here sized by the runtime GQA group). -template -__global__ void __launch_bounds__(kScoreBlockThreads) triScoreScalarKernel(FoldedScoreParams a) -{ - static_assert(!(ROTATE_IN_CTA && USE_MAX), "in-CTA coefficient rotation exists only on the mean path"); - int const seg = blockIdx.y; - int const reqId = a.segRequestIds[seg]; - int const seqLen = a.requestSeqLens[reqId]; - int const tokenStart = a.requestTokenStarts[reqId]; - // Valid-width side store: identical contract to the vectorized kernel. - if (blockIdx.x == 0 && blockIdx.z == 0 && (seg % a.numLayers) == 0 && threadIdx.x == 0) - { - a.validWidthOut[reqId] = seqLen - tokenStart; - } - int const alignedStart = (tokenStart / a.tokensPerBlock) * a.tokensPerBlock; - if (alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads >= seqLen) - { - return; - } - int const absT = alignedStart + static_cast(blockIdx.x) * kScoreBlockThreads + static_cast(threadIdx.x); - bool const valid = absT >= tokenStart && absT < seqLen; - - int const kvHead = blockIdx.z; - int const groupSize = a.numQueryHeads / a.numKvHeads; - int const layerId = a.segLayerIds[seg]; - // Mean-path in-CTA rotation: stage this CTA's groupSize x numFreqs - // coefficient block in shared memory (the early-out above is - // CTA-uniform, so the barrier inside is safe here). - float const* coefRe = a.cRe; - float const* coefIm = a.cIm; - if constexpr (ROTATE_IN_CTA) - { - extern __shared__ float coefficientSmem[]; - int const coefCount = groupSize * a.numFreqs; - float* sCRe = coefficientSmem; - float* sCIm = coefficientSmem + coefCount; - rotateCoefficientsIntoShared(a, reqId, layerId, kvHead * groupSize, coefCount, sCRe, sCIm); - coefRe = sCRe; - coefIm = sCIm; - } - int const page = absT / a.tokensPerBlock; - int const slot = absT - page * a.tokensPerBlock; - int encoded = 0; - if (absT < seqLen) - { - encoded = a.blockOffsets[a.segPageOffsets[seg] + page]; - } - auto const physPage = static_cast(encoded / a.kvFactor); - auto const* row = reinterpret_cast(a.layerBaseAddrs[layerId]) + physPage * a.stridePage - + static_cast(kvHead) * a.strideKvHead + static_cast(slot) * a.strideSlot; - - int64_t const planeStride = static_cast(a.numRequests) * a.numCalibratedLayers * a.numQueryHeads - * static_cast(a.numFreqs); - // The static request-independent [layer, head, freq] c_mlr table: hoist - // this block's first-head row pointer once (no request stride), mirroring - // the vectorized kernel; the head loop advances it by numFreqs per head. - float const* const cMlrRow - = a.cMlr + (static_cast(layerId) * a.numQueryHeads + kvHead * groupSize) * a.numFreqs; - int const tDec = absT - tokenStart; - bool const store = tDec >= 0 && absT < seqLen; - - for (int hg = 0; hg < groupSize; ++hg) - { - int const h = kvHead * groupSize + hg; - int64_t coff = (static_cast(reqId) * a.numCalibratedLayers + layerId) * a.numQueryHeads * a.numFreqs - + static_cast(h) * a.numFreqs; - if constexpr (ROTATE_IN_CTA) - { - // The shared coefficient block's origin is this CTA's head block. - coff = static_cast(hg) * a.numFreqs; - } - float const* const cml = cMlrRow + static_cast(hg) * a.numFreqs; - float acc = 0.0f; - float accMlr = 0.0f; - float accPos[kMaxScoreOffsets]; -#pragma unroll - for (int o = 0; o < kMaxScoreOffsets; ++o) - { - accPos[o] = 0.0f; - } - for (int f = 0; f < a.numFreqs; ++f) - { - float kRe = 0.0f; - float kIm = 0.0f; - if (valid) - { - kRe = toFloat(row[static_cast(f) * a.strideDim]); - kIm = toFloat(row[static_cast(a.numFreqs + f) * a.strideDim]); - } - float const kMag = triSqrtApprox(kRe * kRe + kIm * kIm); - if constexpr (USE_MAX) - { - accMlr = fmaf(kMag, cml[f], accMlr); -#pragma unroll - for (int o = 0; o < kMaxScoreOffsets; ++o) - { - if (o < a.numOffsets) - { - accPos[o] = fmaf(kRe, a.cRe[o * planeStride + coff + f], - fmaf(kIm, a.cIm[o * planeStride + coff + f], accPos[o])); - } - } - } - else - { - acc = fmaf(kRe, coefRe[coff + f], fmaf(kIm, coefIm[coff + f], fmaf(kMag, cml[f], acc))); - } - } - if (store) - { - float score; - if constexpr (USE_MAX) - { - float best = -INFINITY; -#pragma unroll - for (int o = 0; o < kMaxScoreOffsets; ++o) - { - if (o < a.numOffsets) - { - best = fmaxf(best, accPos[o]); - } - } - score = best + accMlr; - } - else - { - score = acc; - } - a.out[(static_cast(seg) * a.numQueryHeads + h) * a.outputWidth + tDec] = score; - } - } -} - -// Per-round coefficient fold: one thread per (request, layer, head, freq) -// element. On the max path each thread additionally writes one c_re/c_im -// value per offset plane (planes are `total` elements apart). The production -// mean path no longer runs this kernel (see the rotation kernel below); the -// c_mlr rows written here are request-identical, and the score kernels read -// only the leading [layer, head, freq] block of that buffer. -__global__ void triFoldScoreCoefficientsKernel(float const* __restrict__ qReal, float const* __restrict__ qImag, - float const* __restrict__ mlrCoef, float const* __restrict__ freqScaleSq, float const* __restrict__ meanCos, - float const* __restrict__ meanSin, float const* __restrict__ omega, float const* __restrict__ offsets, - int32_t const* __restrict__ roundStarts, float const* __restrict__ kvScales, float* __restrict__ cRe, - float* __restrict__ cIm, float* __restrict__ cMlr, int32_t numCalibratedLayers, int32_t numQueryHeads, - int32_t numFreqs, int32_t numOffsets, bool useMax, int64_t total) -{ - int64_t const idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx >= total) - { - return; - } - auto const f = static_cast(idx % numFreqs); - int64_t rest = idx / numFreqs; - auto const h = static_cast(rest % numQueryHeads); - rest /= numQueryHeads; - auto const l = static_cast(rest % numCalibratedLayers); - auto const req = static_cast(rest / numCalibratedLayers); - // Calibration tables are per (layer, head, freq); the request axis only - // enters through the phase terms below. - int64_t const cIdx = (static_cast(l) * numQueryHeads + h) * numFreqs + f; - float const qre = qReal[cIdx]; - float const qim = qImag[cIdx]; - float s = freqScaleSq[f]; - // Quantized-pool dequantization fold: K_real = scale_l * K_quant, so - // multiplying scale_l into ALL coefficient tables (c_mlr and every - // c_re/c_im plane below, since they all carry s) lets the score kernel - // read raw quantized elements. The |K| term relies on scale > 0 - // (validated host-side). Guarded (not "* 1.0f") so the float-pool path is - // instruction-identical to before this parameter existed. - if (kvScales != nullptr) - { - s *= kvScales[l]; - } - cMlr[idx] = mlrCoef[cIdx] * s; - if (!useMax) - { - float const mc = meanCos[static_cast(req) * numFreqs + f]; - float const ms = meanSin[static_cast(req) * numFreqs + f]; - cRe[idx] = s * (qre * mc - qim * ms); - cIm[idx] = s * (qim * mc + qre * ms); - } - else - { - float const om = omega[f]; - auto const rs = static_cast(roundStarts[req]); - for (int32_t o = 0; o < numOffsets; ++o) - { - float const phase = (rs + offsets[o]) * om; - float const cp = cosf(phase); - float const sp = sinf(phase); - cRe[o * total + idx] = s * (qre * cp - qim * sp); - cIm[o * total + idx] = s * (qim * cp + qre * sp); - } - } -} - -// Per-round mean-path replacement for the fold above: all trigonometry is -// tabulated once at initialization (RoPE-style position tables), so the round -// reduces to one table-row gather per request plus four multiplies and two -// adds per element. With -// phaseCos[pos, f] = freq_scale_sq[f] * (1/O) * sum_o cos((pos + offset_o) * omega_f) -// phaseSin[pos, f] = freq_scale_sq[f] * (1/O) * sum_o sin((pos + offset_o) * omega_f) -// qRealScaled / qImagScaled = kv_scale_l * q (identity for float pools) -// the rotation -// c_re = qRealScaled * phaseCos[rs] - qImagScaled * phaseSin[rs] -// c_im = qImagScaled * phaseCos[rs] + qRealScaled * phaseSin[rs] -// equals the mean fold's freq_scale_sq * kv_scale_l * (q rotated by the -// offset-mean phase at round start rs), because both scale factors distribute -// over the complex product. The position-independent c_mlr fold is fully -// static (folded once at initialization, request axis removed), so this -// kernel never writes it. Grid mirrors the fold kernel: one thread per -// (request, layer, head, freq) element. The host wrapper guarantees every -// round start indexes inside the tables. -// Design: Fanrong Li (torch-graph review, 2026-07-20). -// -// The production mean path now runs this rotation inside the score kernels' -// CTA prologue instead (rotateCoefficientsIntoShared, sharing -// rotateCoefficientPair with this kernel so the two paths stay bit-identical -// end to end). This standalone kernel remains as the unit tests' equality -// reference leg and as the fallback for coefficient blocks past the score -// launch's dynamic shared memory bound. -__global__ void triRotateMeanScoreCoefficientsKernel(float const* __restrict__ qRealScaled, - float const* __restrict__ qImagScaled, float const* __restrict__ phaseCos, float const* __restrict__ phaseSin, - int32_t const* __restrict__ roundStarts, float* __restrict__ cRe, float* __restrict__ cIm, - int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, int64_t total) -{ - int64_t const idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx >= total) - { - return; - } - auto const f = static_cast(idx % numFreqs); - int64_t const rest = idx / numFreqs; - int64_t const calibrationRows = static_cast(numCalibratedLayers) * numQueryHeads; - auto const req = static_cast(rest / calibrationRows); - // Calibration tables are per (layer, head, freq); the request-major flat - // index reduces onto them modulo the calibration extent. - int64_t const cIdx = (rest % calibrationRows) * numFreqs + f; - int64_t const phaseIdx = static_cast(roundStarts[req]) * numFreqs + f; - float const pc = phaseCos[phaseIdx]; - float const ps = phaseSin[phaseIdx]; - // rotateCoefficientPair is the ONE rotation expression shared with the - // score kernels' in-CTA prologue — the bit-equality contract between the - // two mean-path preparation flavors. - float2 const c = rotateCoefficientPair(qRealScaled[cIdx], qImagScaled[cIdx], pc, ps); - cRe[idx] = c.x; - cIm[idx] = c.y; -} - -// The in-CTA rotation flavor adds one instantiation per (T, GROUP, chunk -// mode) — mean-only, so the instantiation count per element type grows from -// 4 to 6 per GROUP (binary-size cost of keeping the standalone-rotation read -// path compiled for the max aggregation and the unit tests' equality leg). -template -void launchVectorized( - FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useMax, bool rotateInCta, cudaStream_t stream) -{ - bool const staticChunks = params.numFreqs == 64; - int32_t const effectiveGroup = params.zIsQueryHead ? 1 : groupSize; - // Dynamic shared memory for the rotation prologue's coefficient block - // (re + im); zero on the other paths. - size_t const smemBytes - = rotateInCta ? sizeof(float) * 2 * static_cast(effectiveGroup) * params.numFreqs : 0; -#define TRTLLM_TRI_SCORE_LAUNCH_GROUP(GROUP_V) \ - do \ - { \ - if (staticChunks) \ - { \ - if (useMax) \ - triScoreVectorizedKernel<<>>(params); \ - else if (rotateInCta) \ - triScoreVectorizedKernel \ - <<>>(params); \ - else \ - triScoreVectorizedKernel \ - <<>>(params); \ - } \ - else \ - { \ - if (useMax) \ - triScoreVectorizedKernel<<>>(params); \ - else if (rotateInCta) \ - triScoreVectorizedKernel \ - <<>>(params); \ - else \ - triScoreVectorizedKernel \ - <<>>(params); \ - } \ - } while (0) - switch (effectiveGroup) - { - case 1: TRTLLM_TRI_SCORE_LAUNCH_GROUP(1); break; - case 2: TRTLLM_TRI_SCORE_LAUNCH_GROUP(2); break; - case 4: TRTLLM_TRI_SCORE_LAUNCH_GROUP(4); break; - case 8: TRTLLM_TRI_SCORE_LAUNCH_GROUP(8); break; - default: - TLLM_CHECK_WITH_INFO(false, - "tri_attention_score: vectorized GQA group size must be 1/2/4/8 or use the per-query-head mapping (got " - "%d)", - effectiveGroup); - } -#undef TRTLLM_TRI_SCORE_LAUNCH_GROUP -} - -template -void launchScalar(FoldedScoreParams const& params, dim3 grid, bool useMax, bool rotateInCta, cudaStream_t stream) -{ - // The scalar rotation prologue sizes its shared coefficient block by the - // runtime GQA group (any group size is covered here). - size_t const smemBytes = rotateInCta - ? sizeof(float) * 2 * static_cast(params.numQueryHeads / params.numKvHeads) * params.numFreqs - : 0; - if (useMax) - { - triScoreScalarKernel<<>>(params); - } - else if (rotateInCta) - { - triScoreScalarKernel<<>>(params); - } - else - { - triScoreScalarKernel<<>>(params); - } -} - -// Launch flavor for bf16/fp16 pools, the only element types owning both load -// paths (the vectorized 16-byte chunk kernel and the strided scalar kernel). -template -void launchVectorizedOrScalar(FoldedScoreParams const& params, int32_t groupSize, dim3 grid, bool useVectorized, - bool useMax, bool rotateInCta, cudaStream_t stream) -{ - if (useVectorized) - { - launchVectorized(params, groupSize, grid, useMax, rotateInCta, stream); - } - else - { - launchScalar(params, grid, useMax, rotateInCta, stream); - } -} - -// Quantized pools are functional-only: no vectorized instantiation exists for -// them by design (their dequant scale is folded into the coefficients, so -// only the scalar load path knows how to read them). -template -void launchQuantizedScalar( - FoldedScoreParams const& params, dim3 grid, bool useVectorized, bool useMax, bool rotateInCta, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: quantized pools must use the scalar path"); - launchScalar(params, grid, useMax, rotateInCta, stream); -} - -} // namespace - -void foldScoreCoefficientsLaunch(float const* qReal, float const* qImag, float const* mlrCoef, float const* freqScaleSq, - float const* meanCos, float const* meanSin, float const* omega, float const* offsets, int32_t const* roundStarts, - float const* kvScales, float* cRe, float* cIm, float* cMlr, int32_t numRequests, int32_t numCalibratedLayers, - int32_t numQueryHeads, int32_t numFreqs, int32_t numOffsets, bool useMax, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(useMax ? (omega != nullptr && offsets != nullptr && roundStarts != nullptr) - : (meanCos != nullptr && meanSin != nullptr), - "tri_attention_score_fold: aggregation-path inputs are missing"); - int64_t const total = static_cast(numRequests) * numCalibratedLayers * numQueryHeads * numFreqs; - int32_t const threads = 256; - auto const blocks = static_cast((total + threads - 1) / threads); - triFoldScoreCoefficientsKernel<<>>(qReal, qImag, mlrCoef, freqScaleSq, meanCos, meanSin, - omega, offsets, roundStarts, kvScales, cRe, cIm, cMlr, numCalibratedLayers, numQueryHeads, numFreqs, numOffsets, - useMax, total); - TLLM_CUDA_CHECK(cudaGetLastError()); -} - -void rotateMeanScoreCoefficientsLaunch(float const* qRealScaled, float const* qImagScaled, float const* phaseCos, - float const* phaseSin, int32_t const* roundStarts, float* cRe, float* cIm, int32_t numRequests, - int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, cudaStream_t stream) -{ - int64_t const total = static_cast(numRequests) * numCalibratedLayers * numQueryHeads * numFreqs; - int32_t const threads = 256; - auto const blocks = static_cast((total + threads - 1) / threads); - triRotateMeanScoreCoefficientsKernel<<>>(qRealScaled, qImagScaled, phaseCos, phaseSin, - roundStarts, cRe, cIm, numCalibratedLayers, numQueryHeads, numFreqs, total); - TLLM_CUDA_CHECK(cudaGetLastError()); -} - -void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, - int32_t numSegments, bool useVectorized, bool useMax, bool rotateInCta, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(numSegments > 0 && numSegments <= 65535, - "tri_attention_score: request*layer segment count exceeds the CUDA grid limit"); - TLLM_CHECK_WITH_INFO(params.numOffsets >= 1 && params.numOffsets <= kMaxScoreOffsets, - "tri_attention_score: offset planes exceed the per-thread accumulator budget"); - TLLM_CHECK_WITH_INFO(!useVectorized || (params.numFreqs % 8 == 0 && params.strideDim == 1), - "tri_attention_score: vectorized path requires 8-frequency chunks with unit stride"); - TLLM_CHECK_WITH_INFO(!(rotateInCta && useMax), "tri_attention_score: in-CTA coefficient rotation is mean-only"); - TLLM_CHECK_WITH_INFO(!rotateInCta - || (params.phaseCos != nullptr && params.phaseSin != nullptr && params.qReS != nullptr - && params.qImS != nullptr && params.roundStarts != nullptr), - "tri_attention_score: in-CTA rotation inputs are missing"); - // Launches request dynamic shared memory without opting into the - // above-48KB attribute, so oversized coefficient blocks must fail loudly - // here; such geometries can still score through the standalone rotation - // launch (rotateInCta == false). The scalar prologue's runtime GQA group - // is the worst case (the vectorized per-query-head mapping uses 1). - TLLM_CHECK_WITH_INFO( - !rotateInCta || sizeof(float) * 2 * static_cast(groupSize) * params.numFreqs <= 48u * 1024u, - "tri_attention_score: in-CTA rotation coefficient block exceeds the 48KB dynamic shared memory bound"); - // Tile count covers the decode span plus the worst-case page-alignment - // slack (tokenStart may sit up to tokensPerBlock - 1 tokens into a page). - auto const tiles = static_cast( - (params.outputWidth + params.tokensPerBlock - 1 + kScoreBlockThreads - 1) / kScoreBlockThreads); - uint32_t const headBlocks = params.zIsQueryHead && useVectorized ? params.numQueryHeads : params.numKvHeads; - dim3 const grid(tiles, static_cast(numSegments), headBlocks); - switch (poolType) - { - case PoolElementType::kBFloat16: - launchVectorizedOrScalar<__nv_bfloat16>(params, groupSize, grid, useVectorized, useMax, rotateInCta, stream); - break; - case PoolElementType::kHalf: - launchVectorizedOrScalar(params, groupSize, grid, useVectorized, useMax, rotateInCta, stream); - break; - case PoolElementType::kFloat32: - // fp32 pools have 32-byte 8-frequency rows; the 16-byte chunk path - // does not apply, so they always take the strided scalar kernel. - TLLM_CHECK_WITH_INFO(!useVectorized, "tri_attention_score: fp32 pools must use the scalar path"); - launchScalar(params, grid, useMax, rotateInCta, stream); - break; - case PoolElementType::kFloat8E4M3: - launchQuantizedScalar<__nv_fp8_e4m3>(params, grid, useVectorized, useMax, rotateInCta, stream); - break; - case PoolElementType::kInt8: - launchQuantizedScalar(params, grid, useVectorized, useMax, rotateInCta, stream); - break; - } - TLLM_CUDA_CHECK(cudaGetLastError()); -} - -} // namespace kernels::tri_attention_score - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h b/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h deleted file mode 100644 index 43192f2f770d..000000000000 --- a/cpp/tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h +++ /dev/null @@ -1,213 +0,0 @@ -/* - * 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. - */ -#pragma once - -#include -#include - -#include "tensorrt_llm/common/config.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels::tri_attention_score -{ - -// The TriAttention trig score of one cached token t for query head h is -// score(t, h) = sum_f K_re(t, f) * c_re(h, f) -// + K_im(t, f) * c_im(h, f) -// + |K(t, f)| * c_mlr(h, f) -// where the c tables fold the per-round query calibration and phase terms so -// the per-token kernel touches no trigonometry. The fold runs once per -// eviction round; the score kernel then reads paged KV directly (one thread -// per token) and writes fp32 rows in the selector's layout. - -// Element type of the paged KV pools. The score kernel reads layers through -// raw per-layer base addresses (V2 exposes each layer as its own storage), so -// the caller passes the shared element type explicitly instead of a tensor. -// The quantized types (fp8/int8) are functional-only: they run the strided -// scalar kernel exclusively, and their per-layer dequantization scale is -// folded into the score coefficients (see foldScoreCoefficientsLaunch), so -// the score kernel reads raw quantized elements with zero hot-loop cost. -enum class PoolElementType : int32_t -{ - kBFloat16 = 0, - kHalf = 1, - kFloat32 = 2, - kFloat8E4M3 = 3, - kInt8 = 4, -}; - -// Upper bound on per-offset accumulators held by one score thread on the -// "max" aggregation path. The production "mean" path folds every offset into -// ONE coefficient plane, so this bound never constrains it; "max" needs one -// plane per offset, and the default geometric offset table exceeds 8, so a -// "max" run trips the fold op's TORCH_CHECK unless the offset budget is -// reduced. 8 keeps headroom without bloating the per-thread register budget. -inline constexpr int32_t kMaxScoreOffsets = 8; - -// Threads per score CTA; one thread scores one cached token. -inline constexpr int32_t kScoreBlockThreads = 128; - -// Fold the per-round score coefficients: -// c_re = fss * (q_re * cos - q_im * sin) -// c_im = fss * (q_im * cos + q_re * sin) -// c_mlr = mlr * fss -// Mean aggregation collapses all offsets into meanCos/meanSin beforehand and -// writes one plane (numOffsets == 1). Max aggregation cannot collapse (max -// does not commute through the frequency sum), so it writes one c_re/c_im -// plane per offset with cos/sin((round_start + offset) * omega); c_mlr is -// offset independent either way. Output layout per plane is -// [numRequests, numCalibratedLayers, numQueryHeads, numFreqs] fp32, indexed -// by ABSOLUTE layer id (matching the calibration tables). -// -// kvScales (nullable) carries one fp32 dequantization scale per ABSOLUTE -// layer id for quantized (fp8/int8) KV pools: K_real = scale_l * K_quant, so -// multiplying scale_l into all three coefficient tables (every per-offset -// plane included) lets the score kernel consume raw quantized elements. The -// |K| term relies on |scale * K_q| == scale * |K_q|, which only holds for -// scale > 0 — the host wrapper validates positivity before launch. -// -// The c_mlr rows this kernel writes are request-identical (the request axis -// only enters through the phase terms), and the score kernels consume c_mlr -// through a request-independent [L_cal, HQ, F] index, i.e. only the leading -// calibration block of the buffer. The production mean path skips this fold -// entirely (see rotateMeanScoreCoefficientsLaunch below); this launch remains -// the "max" aggregation path. -void foldScoreCoefficientsLaunch(float const* qReal, // [L_cal * HQ * F] - float const* qImag, // [L_cal * HQ * F] - float const* mlrCoef, // [L_cal * HQ * F] - float const* freqScaleSq, // [F] - float const* meanCos, // [numRequests * F] (mean path, else nullptr) - float const* meanSin, // [numRequests * F] (mean path, else nullptr) - float const* omega, // [F] (max path, else nullptr) - float const* offsets, // [numOffsets] (max path, else nullptr) - int32_t const* roundStarts, // [numRequests] (max path, else nullptr) - float const* kvScales, // [L_cal] per-layer dequant scale (quantized pools, else nullptr) - float* cRe, // [numOffsets, numRequests, L_cal, HQ, F] - float* cIm, // [numOffsets, numRequests, L_cal, HQ, F] - float* cMlr, // [numRequests, L_cal, HQ, F] - int32_t numRequests, int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, int32_t numOffsets, - bool useMax, cudaStream_t stream); - -// Mean-aggregation replacement for the per-round fold above: rotate the -// pre-scaled calibration query by tabulated phases instead of computing any -// trigonometry per round. phaseCos/phaseSin hold, for every possible round -// start position (built once at initialization, float64 accumulation), -// phaseCos[pos, f] = freq_scale_sq[f] * (1/O) * sum_o cos((pos + offset_o) * omega_f) -// phaseSin[pos, f] = freq_scale_sq[f] * (1/O) * sum_o sin((pos + offset_o) * omega_f) -// and qRealScaled/qImagScaled carry the per-layer KV dequantization scale -// (identity for float pools). Each thread gathers its request's table row and -// writes -// c_re = q_re_s * phaseCos[rs] - q_im_s * phaseSin[rs] -// c_im = q_im_s * phaseCos[rs] + q_re_s * phaseSin[rs] -// into the same [numRequests, L_cal, HQ, F] planes the fold produces, because -// freq_scale_sq and kv_scale distribute over the complex rotation. c_mlr has -// no position term, so on this path it is folded once at initialization into -// a static [L_cal, HQ, F] table and never rewritten per round. The "max" -// aggregation keeps foldScoreCoefficientsLaunch unchanged. Every round start -// must lie in [0, maxPosition); the host wrapper enforces that loudly before -// launch. Design: Fanrong Li (torch-graph review, 2026-07-20). -// -// The production mean path no longer runs this standalone launch either: the -// score kernels rebuild the same coefficients in an in-CTA shared-memory -// prologue (see FoldedScoreParams and foldedScoreLaunch's rotateInCta). The -// prologue compiles the SAME rotation arithmetic as this kernel, so the two -// preparation paths produce bit-identical scores; this launch stays as that -// equality proof's reference leg (the unit suite compares them with -// torch.equal) and as the fallback for coefficient blocks past the dynamic -// shared memory bound. -void rotateMeanScoreCoefficientsLaunch(float const* qRealScaled, // [L_cal * HQ * F] - float const* qImagScaled, // [L_cal * HQ * F] - float const* phaseCos, // [maxPosition * F] - float const* phaseSin, // [maxPosition * F] - int32_t const* roundStarts, // [numRequests], each in [0, maxPosition) - float* cRe, // [numRequests, L_cal, HQ, F] - float* cIm, // [numRequests, L_cal, HQ, F] - int32_t numRequests, int32_t numCalibratedLayers, int32_t numQueryHeads, int32_t numFreqs, cudaStream_t stream); - -// Everything one folded-score launch needs. One "segment" is one -// (request, scored layer) pair; segments are request-major so -// seg % numLayers == 0 identifies each request's first segment. -struct FoldedScoreParams -{ - int64_t const* layerBaseAddrs; // [num pools] absolute device addresses, ABSOLUTE layer id indexed - int32_t const* blockOffsets; // flattened native V2 page table - int64_t const* segPageOffsets; // [numSegments] offset of each segment's page row into blockOffsets - int32_t const* segRequestIds; // [numSegments] - int32_t const* segLayerIds; // [numSegments] ABSOLUTE layer ids - int32_t const* requestSeqLens; // [numRequests] - int32_t* validWidthOut; // [numRequests] side-store: seqLen - tokenStart, once per request - int32_t const* requestTokenStarts; // [numRequests] pinned prompt length = decode-region origin - // Per-round coefficient planes (see foldScoreCoefficientsLaunch / - // rotateMeanScoreCoefficientsLaunch); nullptr when the mean-path kernels - // rotate their coefficients in-CTA instead (rotateInCta below). - float const* cRe; - float const* cIm; - float const* cMlr; // static, request-independent [L_cal, HQ, F] MLR table - // Mean-path in-CTA rotation inputs, consumed only by rotateInCta launches - // (nullptr otherwise). Table contract and arithmetic are EXACTLY those of - // rotateMeanScoreCoefficientsLaunch; round starts must be host-validated - // against the tabulated position extent BEFORE launch (the kernels gather - // phase rows unguarded — a device-side bound would need its own check - // kernel). - float const* phaseCos; // [maxPosition, F] - float const* phaseSin; // [maxPosition, F] - float const* qReS; // [L_cal, HQ, F] pre-scaled calibration query, real part - float const* qImS; // [L_cal, HQ, F] pre-scaled calibration query, imaginary part - int32_t const* roundStarts; // [numRequests], each in [0, maxPosition) - float* out; // [segment, numQueryHeads, outputWidth] fp32 decode-only scores - int32_t outputWidth; - int32_t numLayers; // scored layers per request (the segment period) - int32_t numRequests; - int32_t numCalibratedLayers; - int32_t numQueryHeads; - int32_t numKvHeads; - int32_t numFreqs; - int32_t tokensPerBlock; - int32_t kvFactor; // page-table entries encode role pages; entry / kvFactor = pool page - int32_t numOffsets; // effective c_re/c_im planes (1 on the mean path) - // Grid mapping for GQA group sizes without a dedicated template - // instantiation: grid.z indexes single query heads instead of KV heads - // (the KV plane is derived per block; K traffic is repeated per head). - bool zIsQueryHead; - // HND pool element strides (elements, not bytes) shared by all layers. - int64_t stridePage; - int64_t strideKvHead; - int64_t strideSlot; - int64_t strideDim; -}; - -// Launch the folded score over paged KV. useVectorized selects 16-byte -// 8-frequency chunk loads (requires numFreqs % 8 == 0, strideDim == 1, -// bf16/fp16 pools, and 16-byte aligned bases/strides — the caller audits -// alignment); otherwise a fully strided scalar path runs the same math. -// groupSize = numQueryHeads / numKvHeads must be 1, 2, 4, or 8 unless -// params.zIsQueryHead maps grid.z to single query heads. -// -// rotateInCta (mean aggregation only) makes each score CTA rotate its own -// coefficient block from params.phaseCos/phaseSin/qReS/qImS/roundStarts into -// dynamic shared memory instead of reading pre-rotated global c_re/c_im -// planes, eliminating the standalone rotation launch and its per-round -// coefficient scratch. The prologue shares the standalone kernel's rotation -// arithmetic, so scores are bit-identical between the two preparation paths. -// The coefficient block (2 * group * numFreqs fp32) must fit the default -// 48KB dynamic shared memory bound — enforced loudly here. -void foldedScoreLaunch(FoldedScoreParams const& params, PoolElementType poolType, int32_t groupSize, - int32_t numSegments, bool useVectorized, bool useMax, bool rotateInCta, cudaStream_t stream); - -} // namespace kernels::tri_attention_score - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 02772cf9cf83..273ffa838a21 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -141,8 +141,7 @@ add_library( trtllmGenQKVProcessOp.cpp inplaceSliceCopyOp.cpp mhcOp.cpp - compressorOp.cpp - triAttentionScoreOp.cpp) + compressorOp.cpp) set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_libraries( th_common PRIVATE ${TORCH_LIBRARIES} th_utils ${Python3_LIBRARIES} diff --git a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp b/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp deleted file mode 100644 index 61a13dedd5e1..000000000000 --- a/cpp/tensorrt_llm/thop/triAttentionScoreOp.cpp +++ /dev/null @@ -1,452 +0,0 @@ -/* - * 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. - */ - -#include "tensorrt_llm/kernels/triAttentionScoreKernels/triAttentionScoreKernels.h" - -#include -#include - -namespace tk = tensorrt_llm::kernels::tri_attention_score; - -// One op per kernel launch, matching the file-level granularity of sibling -// kernel wrappers: the per-round coefficient preparation (the rotation op on -// the legacy mean path, the trigonometric fold on the max path) writes -// persistent buffers whose plane count depends on the aggregation mode, while -// the score launch consumes those buffers; callers time and re-plan them -// independently. The default mean path collapses to the score op alone (its -// kernels rotate coefficients in-CTA), so no preparation op runs there. - -namespace -{ - -// One dtype-parameterized operand validator; dtypeName keeps the message -// spelling (fp32 / int32 / int64) for each expected scalar type. -void checkContiguousCuda(torch::Tensor const& tensor, at::ScalarType dtype, char const* dtypeName, char const* name) -{ - TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == dtype, name, - " must be a contiguous ", dtypeName, " CUDA tensor"); -} - -// Validate the optional per-layer dequantization scales for quantized -// (fp8/int8) KV pools and return the device pointer (nullptr when absent). -// Scales are indexed by ABSOLUTE layer id, matching layer_base_addrs and the -// calibration tables, so the extent must cover every calibrated layer. The -// positivity check runs HOST-side (one small sync per eviction round on a -// functional-only path): the |K| coefficient fold assumes -// |scale * K_q| == scale * |K_q|, which silently corrupts scores for -// scale <= 0, so a loud host error here is required rather than a -// device-side assert. -float const* checkKvScales( - std::optional const& kv_scales, int64_t num_calibrated_layers, char const* op_name) -{ - if (!kv_scales.has_value()) - { - return nullptr; - } - checkContiguousCuda(*kv_scales, at::kFloat, "fp32", "kv_scales"); - TORCH_CHECK(kv_scales->numel() >= num_calibrated_layers, op_name, - ": kv_scales must carry one scale per calibrated layer (absolute layer id indexed)"); - TORCH_CHECK(kv_scales->min().item() > 0.0f, op_name, - ": kv_scales must be strictly positive (the |K| dequantization fold is only valid for positive scales)"); - return kv_scales->data_ptr(); -} - -// Fold the per-round TriAttention score coefficients into c_re/c_im/c_mlr -// (fp32, [num_offsets?, num_requests, num_calibrated_layers, heads, freqs]). -// The mean aggregation consumes offset-collapsed mean_cos/mean_sin and writes -// one plane; the max aggregation consumes omega/offsets/round_starts and -// writes one c_re/c_im plane per offset. kv_scales (quantized pools only) -// folds the per-layer dequantization scale into every coefficient table; the -// paired score op then reads raw quantized elements. This op cannot see the -// pool dtype, so presence-iff-quantized is enforced by the score op. The -// production mean path prepares its coefficients through the rotation op -// below instead; this fold remains the max-aggregation path (and keeps its -// mean branch for that kernel's documented contract). -void triAttentionFoldScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, torch::Tensor c_mlr, - torch::Tensor q_real, torch::Tensor q_imag, torch::Tensor mlr_coef, torch::Tensor freq_scale_sq, - std::optional mean_cos, std::optional mean_sin, std::optional omega, - std::optional offsets, std::optional round_starts, int64_t num_requests, - int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_freqs, int64_t num_offsets, bool use_max, - std::optional kv_scales) -{ - checkContiguousCuda(c_re, at::kFloat, "fp32", "c_re"); - checkContiguousCuda(c_im, at::kFloat, "fp32", "c_im"); - checkContiguousCuda(c_mlr, at::kFloat, "fp32", "c_mlr"); - checkContiguousCuda(q_real, at::kFloat, "fp32", "q_real"); - checkContiguousCuda(q_imag, at::kFloat, "fp32", "q_imag"); - checkContiguousCuda(mlr_coef, at::kFloat, "fp32", "mlr_coef"); - checkContiguousCuda(freq_scale_sq, at::kFloat, "fp32", "freq_scale_sq"); - TORCH_CHECK(num_requests > 0 && num_calibrated_layers > 0 && num_query_heads > 0 && num_freqs > 0, - "tri_attention_fold_score_coefficients: fold extents must be positive"); - TORCH_CHECK(num_offsets >= 1 && num_offsets <= tk::kMaxScoreOffsets, - "tri_attention_fold_score_coefficients: num_offsets must be in [1, ", tk::kMaxScoreOffsets, "], got ", - num_offsets); - int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; - int64_t const planes = use_max ? num_offsets : 1; - TORCH_CHECK(c_re.numel() >= planes * total && c_im.numel() >= planes * total && c_mlr.numel() >= total, - "tri_attention_fold_score_coefficients: coefficient buffers are undersized"); - int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; - TORCH_CHECK(q_real.numel() >= calibration && q_imag.numel() >= calibration && mlr_coef.numel() >= calibration - && freq_scale_sq.numel() >= num_freqs, - "tri_attention_fold_score_coefficients: calibration tensors are undersized for the fold extent"); - - float const* meanCosPtr = nullptr; - float const* meanSinPtr = nullptr; - float const* omegaPtr = nullptr; - float const* offsetsPtr = nullptr; - int32_t const* roundStartsPtr = nullptr; - if (use_max) - { - TORCH_CHECK(omega.has_value() && offsets.has_value() && round_starts.has_value(), - "tri_attention_fold_score_coefficients: max aggregation requires omega, offsets, and round_starts"); - checkContiguousCuda(*omega, at::kFloat, "fp32", "omega"); - checkContiguousCuda(*offsets, at::kFloat, "fp32", "offsets"); - checkContiguousCuda(*round_starts, at::kInt, "int32", "round_starts"); - TORCH_CHECK( - omega->numel() >= num_freqs && offsets->numel() >= num_offsets && round_starts->numel() >= num_requests, - "tri_attention_fold_score_coefficients: max-path inputs are undersized for the folded request count"); - omegaPtr = omega->data_ptr(); - offsetsPtr = offsets->data_ptr(); - roundStartsPtr = round_starts->data_ptr(); - } - else - { - TORCH_CHECK(mean_cos.has_value() && mean_sin.has_value(), - "tri_attention_fold_score_coefficients: mean aggregation requires mean_cos and mean_sin"); - checkContiguousCuda(*mean_cos, at::kFloat, "fp32", "mean_cos"); - checkContiguousCuda(*mean_sin, at::kFloat, "fp32", "mean_sin"); - TORCH_CHECK(mean_cos->numel() >= num_requests * num_freqs && mean_sin->numel() >= num_requests * num_freqs, - "tri_attention_fold_score_coefficients: mean_cos/mean_sin are undersized (the fold iterates one row per " - "folded request)"); - meanCosPtr = mean_cos->data_ptr(); - meanSinPtr = mean_sin->data_ptr(); - } - float const* kvScalesPtr = checkKvScales(kv_scales, num_calibrated_layers, "tri_attention_fold_score_coefficients"); - - auto stream = at::cuda::getCurrentCUDAStream(); - tk::foldScoreCoefficientsLaunch(q_real.data_ptr(), q_imag.data_ptr(), mlr_coef.data_ptr(), - freq_scale_sq.data_ptr(), meanCosPtr, meanSinPtr, omegaPtr, offsetsPtr, roundStartsPtr, kvScalesPtr, - c_re.data_ptr(), c_im.data_ptr(), c_mlr.data_ptr(), static_cast(num_requests), - static_cast(num_calibrated_layers), static_cast(num_query_heads), - static_cast(num_freqs), static_cast(num_offsets), use_max, stream); -} - -// Per-round mean-path coefficient rotation: gather each request's row of the -// initialization-time phase tables (freq_scale_sq, the offset mean, and — via -// the pre-scaled calibration query — the per-layer KV dequantization scale -// are already baked in) and rotate the query by it, writing the same -// c_re/c_im planes the mean fold produced, with zero trigonometry per round. -// c_mlr has no position term, so on this path it is folded once at -// initialization and this op never touches it. Round starts index the phase -// tables directly, so a start at or past max_position (the tabulated position -// count) must fail loudly here rather than gather out of range; like the -// kv_scales positivity check above, the bounds reduction costs one small -// host sync per eviction round. Design: Fanrong Li (torch-graph review, -// 2026-07-20). -// -// The production mean path no longer calls this op: the paged score op below -// performs the same rotation in its kernels' CTA prologue (bit-identically — -// the two paths share one device rotation expression). This op stays as the -// unit tests' equality reference leg and as the fallback for coefficient -// blocks past the score launch's dynamic shared memory bound. -void triAttentionRotateMeanScoreCoefficientsOp(torch::Tensor c_re, torch::Tensor c_im, torch::Tensor q_real_scaled, - torch::Tensor q_imag_scaled, torch::Tensor phase_cos, torch::Tensor phase_sin, torch::Tensor round_starts, - int64_t num_requests, int64_t num_calibrated_layers, int64_t num_query_heads, int64_t num_freqs, - int64_t max_position) -{ - checkContiguousCuda(c_re, at::kFloat, "fp32", "c_re"); - checkContiguousCuda(c_im, at::kFloat, "fp32", "c_im"); - checkContiguousCuda(q_real_scaled, at::kFloat, "fp32", "q_real_scaled"); - checkContiguousCuda(q_imag_scaled, at::kFloat, "fp32", "q_imag_scaled"); - checkContiguousCuda(phase_cos, at::kFloat, "fp32", "phase_cos"); - checkContiguousCuda(phase_sin, at::kFloat, "fp32", "phase_sin"); - checkContiguousCuda(round_starts, at::kInt, "int32", "round_starts"); - TORCH_CHECK( - num_requests > 0 && num_calibrated_layers > 0 && num_query_heads > 0 && num_freqs > 0 && max_position > 0, - "tri_attention_rotate_mean_score_coefficients: rotation extents must be positive"); - int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; - int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; - TORCH_CHECK(c_re.numel() >= total && c_im.numel() >= total, - "tri_attention_rotate_mean_score_coefficients: coefficient buffers are undersized"); - TORCH_CHECK(q_real_scaled.numel() >= calibration && q_imag_scaled.numel() >= calibration, - "tri_attention_rotate_mean_score_coefficients: scaled calibration tensors are undersized"); - TORCH_CHECK(phase_cos.numel() >= max_position * num_freqs && phase_sin.numel() >= max_position * num_freqs, - "tri_attention_rotate_mean_score_coefficients: phase tables do not cover max_position rows"); - TORCH_CHECK(round_starts.numel() >= num_requests, - "tri_attention_rotate_mean_score_coefficients: round_starts are undersized"); - auto const [minStart, maxStart] = round_starts.narrow(0, 0, num_requests).aminmax(); - TORCH_CHECK(minStart.item() >= 0 && maxStart.item() < max_position, - "tri_attention_rotate_mean_score_coefficients: a round start lies outside the tabulated position range " - "[0, ", - max_position, ")"); - - auto stream = at::cuda::getCurrentCUDAStream(); - tk::rotateMeanScoreCoefficientsLaunch(q_real_scaled.data_ptr(), q_imag_scaled.data_ptr(), - phase_cos.data_ptr(), phase_sin.data_ptr(), round_starts.data_ptr(), - c_re.data_ptr(), c_im.data_ptr(), static_cast(num_requests), - static_cast(num_calibrated_layers), static_cast(num_query_heads), - static_cast(num_freqs), stream); -} - -// Score every cached decode token of every (request, layer) segment against -// the folded coefficient tables, writing fp32 [segment, head, token] rows and -// each request's decode width. pool_anchor is one of the scored layer pools: -// the kernel reads all layers through layer_base_addrs (V2 exposes each layer -// as its own storage), and the anchor only supplies their common element type -// and the device; its data is never read through this argument. -// -// Mean-path coefficient preparation comes in two flavors. Passing phase_cos/ -// phase_sin/q_real_scaled/q_imag_scaled/round_starts (+ max_position) makes -// the score kernels rotate their own coefficient block in an in-CTA prologue -// — c_re/c_im must then be omitted (no per-round coefficient scratch exists -// at all). Passing c_re/c_im instead scores against pre-rotated global planes -// (the rotation op above, or the fold op on the max path). The two mean -// flavors share one rotation expression device-side, so their scores are -// bit-identical. The round-start bounds contract of the rotation op moves -// here on the in-CTA flavor: a start at or past max_position would gather -// past the phase tables, and a device-side guard would need its own check -// kernel, so the same host-side aminmax reduction (one small sync per -// eviction round) runs before the launch. -void triAttentionPagedScoreOp(torch::Tensor pool_anchor, torch::Tensor layer_base_addrs, torch::Tensor block_offsets, - torch::Tensor seg_page_offsets, torch::Tensor seg_request_ids, torch::Tensor seg_layer_ids, - torch::Tensor request_seq_lens, torch::Tensor valid_widths, torch::Tensor request_token_starts, - std::optional c_re, std::optional c_im, torch::Tensor c_mlr, torch::Tensor out, - int64_t output_width, int64_t num_layers, int64_t num_requests, int64_t num_calibrated_layers, - int64_t num_query_heads, int64_t num_kv_heads, int64_t num_freqs, int64_t tokens_per_block, int64_t kv_factor, - int64_t num_offsets, int64_t stride_page, int64_t stride_kv_head, int64_t stride_slot, int64_t stride_dim, - int64_t num_segments, bool use_max, bool use_vectorized, std::optional kv_scales, - std::optional phase_cos, std::optional phase_sin, - std::optional q_real_scaled, std::optional q_imag_scaled, - std::optional round_starts, int64_t max_position) -{ - TORCH_CHECK(use_max || num_offsets == 1, - "tri_attention_paged_score: mean aggregation consumes exactly one folded coefficient plane"); - checkContiguousCuda(layer_base_addrs, at::kLong, "int64", "layer_base_addrs"); - checkContiguousCuda(block_offsets, at::kInt, "int32", "block_offsets"); - checkContiguousCuda(seg_page_offsets, at::kLong, "int64", "seg_page_offsets"); - checkContiguousCuda(seg_request_ids, at::kInt, "int32", "seg_request_ids"); - checkContiguousCuda(seg_layer_ids, at::kInt, "int32", "seg_layer_ids"); - checkContiguousCuda(request_seq_lens, at::kInt, "int32", "request_seq_lens"); - checkContiguousCuda(valid_widths, at::kInt, "int32", "valid_widths"); - checkContiguousCuda(request_token_starts, at::kInt, "int32", "request_token_starts"); - checkContiguousCuda(c_mlr, at::kFloat, "fp32", "c_mlr"); - checkContiguousCuda(out, at::kFloat, "fp32", "out"); - TORCH_CHECK(pool_anchor.is_cuda(), "tri_attention_paged_score: pool anchor must be a CUDA tensor"); - - TORCH_CHECK(num_segments > 0 && num_segments <= 65535, - "tri_attention_paged_score: request*layer segment count exceeds the CUDA grid limit"); - TORCH_CHECK(output_width > 0 && num_layers > 0 && num_requests > 0 && num_calibrated_layers > 0 - && tokens_per_block > 0 && kv_factor > 0 && num_freqs > 0, - "tri_attention_paged_score: geometry extents must be positive"); - TORCH_CHECK(num_kv_heads > 0 && num_query_heads % num_kv_heads == 0, - "tri_attention_paged_score: query heads must be divisible by KV heads"); - // kMaxScoreOffsets is the per-thread accumulator budget baked into the - // kernels. It only constrains the "max" path (one coefficient plane per - // offset); the "mean" path always folds every offset into one plane, so - // the default geometric offset table (which is larger) still passes here. - TORCH_CHECK(num_offsets >= 1 && num_offsets <= tk::kMaxScoreOffsets, - "tri_attention_paged_score: num_offsets must be in [1, ", tk::kMaxScoreOffsets, "], got ", num_offsets); - TORCH_CHECK(seg_page_offsets.numel() >= num_segments && seg_request_ids.numel() >= num_segments - && seg_layer_ids.numel() >= num_segments, - "tri_attention_paged_score: segment metadata is undersized"); - TORCH_CHECK(request_seq_lens.numel() >= num_requests && valid_widths.numel() >= num_requests - && request_token_starts.numel() >= num_requests, - "tri_attention_paged_score: per-request metadata is undersized"); - int64_t const total = num_requests * num_calibrated_layers * num_query_heads * num_freqs; - // c_mlr is request independent: the score kernels index it as one static - // [layer, head, freq] table, so only the calibration extent is required. - int64_t const calibration = num_calibrated_layers * num_query_heads * num_freqs; - bool const rotateInCta = phase_cos.has_value() || phase_sin.has_value() || q_real_scaled.has_value() - || q_imag_scaled.has_value() || round_starts.has_value(); - if (rotateInCta) - { - TORCH_CHECK(!use_max, "tri_attention_paged_score: in-CTA coefficient rotation is mean-only"); - TORCH_CHECK(phase_cos.has_value() && phase_sin.has_value() && q_real_scaled.has_value() - && q_imag_scaled.has_value() && round_starts.has_value(), - "tri_attention_paged_score: in-CTA rotation requires phase_cos, phase_sin, q_real_scaled, " - "q_imag_scaled, and round_starts together"); - TORCH_CHECK(!c_re.has_value() && !c_im.has_value(), - "tri_attention_paged_score: in-CTA rotation reads no c_re/c_im planes; omit them"); - checkContiguousCuda(*phase_cos, at::kFloat, "fp32", "phase_cos"); - checkContiguousCuda(*phase_sin, at::kFloat, "fp32", "phase_sin"); - checkContiguousCuda(*q_real_scaled, at::kFloat, "fp32", "q_real_scaled"); - checkContiguousCuda(*q_imag_scaled, at::kFloat, "fp32", "q_imag_scaled"); - checkContiguousCuda(*round_starts, at::kInt, "int32", "round_starts"); - TORCH_CHECK(max_position > 0, "tri_attention_paged_score: in-CTA rotation requires a positive max_position"); - TORCH_CHECK(phase_cos->numel() >= max_position * num_freqs && phase_sin->numel() >= max_position * num_freqs, - "tri_attention_paged_score: phase tables do not cover max_position rows"); - TORCH_CHECK(q_real_scaled->numel() >= calibration && q_imag_scaled->numel() >= calibration, - "tri_attention_paged_score: scaled calibration tensors are undersized"); - TORCH_CHECK(round_starts->numel() >= num_requests, "tri_attention_paged_score: round_starts are undersized"); - // Same host-side bounds reduction as the standalone rotation op (one - // small sync per eviction round): the kernels gather phase rows - // unguarded, so an out-of-range start must fail loudly here. - auto const [minStart, maxStart] = round_starts->narrow(0, 0, num_requests).aminmax(); - TORCH_CHECK(minStart.item() >= 0 && maxStart.item() < max_position, - "tri_attention_paged_score: a round start lies outside the tabulated position range [0, ", max_position, - ")"); - } - else - { - TORCH_CHECK(c_re.has_value() && c_im.has_value(), - "tri_attention_paged_score: pre-rotated scoring requires c_re and c_im"); - checkContiguousCuda(*c_re, at::kFloat, "fp32", "c_re"); - checkContiguousCuda(*c_im, at::kFloat, "fp32", "c_im"); - TORCH_CHECK(c_re->numel() >= num_offsets * total && c_im->numel() >= num_offsets * total, - "tri_attention_paged_score: folded coefficient buffers are undersized"); - } - TORCH_CHECK(c_mlr.numel() >= calibration, "tri_attention_paged_score: folded coefficient buffers are undersized"); - TORCH_CHECK(out.numel() >= num_segments * num_query_heads * output_width, - "tri_attention_paged_score: score output buffer is undersized"); - - auto const dtype = pool_anchor.scalar_type(); - auto poolType = tk::PoolElementType::kBFloat16; - if (dtype == at::kBFloat16) - { - poolType = tk::PoolElementType::kBFloat16; - } - else if (dtype == at::kHalf) - { - poolType = tk::PoolElementType::kHalf; - } - else if (dtype == at::kFloat) - { - poolType = tk::PoolElementType::kFloat32; - } - else if (dtype == at::kFloat8_e4m3fn) - { - poolType = tk::PoolElementType::kFloat8E4M3; - } - else if (dtype == at::kChar) - { - poolType = tk::PoolElementType::kInt8; - } - else - { - TORCH_CHECK(false, "tri_attention_paged_score: unsupported KV pool dtype ", dtype, - " (supported: bf16, fp16, fp32, fp8_e4m3fn, int8)"); - } - // The score kernel never applies kv_scales itself (the fold op already - // multiplied them into the coefficient tables), but this op is the only - // one that sees the pool dtype, so it owns the presence contract: - // quantized elements without scales would be scored as raw integers, and - // scales alongside float pools would silently double-scale. - bool const quantizedPool = poolType == tk::PoolElementType::kFloat8E4M3 || poolType == tk::PoolElementType::kInt8; - TORCH_CHECK(!quantizedPool || kv_scales.has_value(), - "tri_attention_paged_score: quantized (fp8/int8) KV pools require per-layer kv_scales"); - TORCH_CHECK(quantizedPool || !kv_scales.has_value(), - "tri_attention_paged_score: kv_scales are only valid for quantized (fp8/int8) KV pools"); - checkKvScales(kv_scales, num_calibrated_layers, "tri_attention_paged_score"); - TORCH_CHECK(!use_vectorized || dtype == at::kBFloat16 || dtype == at::kHalf, - "tri_attention_paged_score: the vectorized path requires bf16 or fp16 pools"); - TORCH_CHECK(!use_vectorized || (num_freqs % 8 == 0 && stride_dim == 1), - "tri_attention_paged_score: the vectorized path requires num_freqs % 8 == 0 and a unit frequency stride"); - - auto const groupSize = static_cast(num_query_heads / num_kv_heads); - bool const vectorizedGroup = groupSize == 1 || groupSize == 2 || groupSize == 4 || groupSize == 8; - // Other GQA group sizes run the vectorized math one query head per grid.z - // block instead of one KV head (a runtime mapping, no extra template). - bool const zIsQueryHead = use_vectorized && !vectorizedGroup; - TORCH_CHECK((zIsQueryHead ? num_query_heads : num_kv_heads) <= 65535, - "tri_attention_paged_score: head count exceeds the CUDA grid limit"); - - tk::FoldedScoreParams params; - params.layerBaseAddrs = layer_base_addrs.data_ptr(); - params.blockOffsets = block_offsets.data_ptr(); - params.segPageOffsets = seg_page_offsets.data_ptr(); - params.segRequestIds = seg_request_ids.data_ptr(); - params.segLayerIds = seg_layer_ids.data_ptr(); - params.requestSeqLens = request_seq_lens.data_ptr(); - params.validWidthOut = valid_widths.data_ptr(); - params.requestTokenStarts = request_token_starts.data_ptr(); - params.cRe = rotateInCta ? nullptr : c_re->data_ptr(); - params.cIm = rotateInCta ? nullptr : c_im->data_ptr(); - params.cMlr = c_mlr.data_ptr(); - params.phaseCos = rotateInCta ? phase_cos->data_ptr() : nullptr; - params.phaseSin = rotateInCta ? phase_sin->data_ptr() : nullptr; - params.qReS = rotateInCta ? q_real_scaled->data_ptr() : nullptr; - params.qImS = rotateInCta ? q_imag_scaled->data_ptr() : nullptr; - params.roundStarts = rotateInCta ? round_starts->data_ptr() : nullptr; - params.out = out.data_ptr(); - params.outputWidth = static_cast(output_width); - params.numLayers = static_cast(num_layers); - params.numRequests = static_cast(num_requests); - params.numCalibratedLayers = static_cast(num_calibrated_layers); - params.numQueryHeads = static_cast(num_query_heads); - params.numKvHeads = static_cast(num_kv_heads); - params.numFreqs = static_cast(num_freqs); - params.tokensPerBlock = static_cast(tokens_per_block); - params.kvFactor = static_cast(kv_factor); - params.numOffsets = static_cast(num_offsets); - params.zIsQueryHead = zIsQueryHead; - params.stridePage = stride_page; - params.strideKvHead = stride_kv_head; - params.strideSlot = stride_slot; - params.strideDim = stride_dim; - - auto stream = at::cuda::getCurrentCUDAStream(); - tk::foldedScoreLaunch( - params, poolType, groupSize, static_cast(num_segments), use_vectorized, use_max, rotateInCta, stream); -} - -} // anonymous namespace - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def( - "tri_attention_fold_score_coefficients(" - "Tensor(a!) c_re, Tensor(b!) c_im, Tensor(c!) c_mlr, " - "Tensor q_real, Tensor q_imag, Tensor mlr_coef, Tensor freq_scale_sq, " - "Tensor? mean_cos, Tensor? mean_sin, " - "Tensor? omega, Tensor? offsets, Tensor? round_starts, " - "int num_requests, int num_calibrated_layers, " - "int num_query_heads, int num_freqs, " - "int num_offsets, bool use_max, Tensor? kv_scales=None) -> ()"); - - m.def( - "tri_attention_rotate_mean_score_coefficients(" - "Tensor(a!) c_re, Tensor(b!) c_im, " - "Tensor q_real_scaled, Tensor q_imag_scaled, " - "Tensor phase_cos, Tensor phase_sin, Tensor round_starts, " - "int num_requests, int num_calibrated_layers, " - "int num_query_heads, int num_freqs, int max_position) -> ()"); - - // c_re/c_im are optional: the mean path's default in-CTA coefficient - // rotation (phase_cos .. round_starts + max_position, all-or-nothing) - // replaces them entirely; the max path and the legacy mean path keep - // passing pre-rotated planes. - m.def( - "tri_attention_paged_score(" - "Tensor pool_anchor, Tensor layer_base_addrs, Tensor block_offsets, " - "Tensor seg_page_offsets, Tensor seg_request_ids, Tensor seg_layer_ids, " - "Tensor request_seq_lens, Tensor(a!) valid_widths, Tensor request_token_starts, " - "Tensor? c_re, Tensor? c_im, Tensor c_mlr, Tensor(b!) out, " - "int output_width, int num_layers, int num_requests, int num_calibrated_layers, " - "int num_query_heads, int num_kv_heads, int num_freqs, int tokens_per_block, " - "int kv_factor, int num_offsets, int stride_page, int stride_kv_head, " - "int stride_slot, int stride_dim, int num_segments, bool use_max, bool use_vectorized, " - "Tensor? kv_scales=None, " - "Tensor? phase_cos=None, Tensor? phase_sin=None, " - "Tensor? q_real_scaled=None, Tensor? q_imag_scaled=None, " - "Tensor? round_starts=None, int max_position=0) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("tri_attention_fold_score_coefficients", &triAttentionFoldScoreCoefficientsOp); - m.impl("tri_attention_rotate_mean_score_coefficients", &triAttentionRotateMeanScoreCoefficientsOp); - m.impl("tri_attention_paged_score", &triAttentionPagedScoreOp); -} diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 532ca1f74cf1..6c2c1fd8843f 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -744,10 +744,10 @@ def __init__( offsets, output_width=decode_width, ) - # Compile the optional SM100 CuTe score specialization here, outside - # any CUDA graph capture (compilation allocates and synchronizes). - # Default off: without TRTLLM_TRIATTENTION_CUTE_SCORE=1 this is a - # no-op and scoring stays on the compiled C++ score ops. + # Compile the SM100 CuTe score kernel (the only score implementation) + # here at workspace construction, outside any CUDA graph capture + # (compilation allocates and synchronizes). Unsupported geometry + # raises loudly right here rather than mid-round at the score launch. self.fused_group.prepare_cute_score(self.mean_cos, self.mean_sin) self.copy_done = torch.cuda.Event() # First record publishes constructor allocations to the V2 copy stream; @@ -759,22 +759,25 @@ def __init__( self.page_tables_active = False self.stream = None self._score_valid_widths: Optional[torch.Tensor] = None - self._score_aggregation: Optional[str] = None + self._score_launcher_bound = False - def bind_score_launcher(self, valid_widths: torch.Tensor, score_aggregation: str) -> None: - """Bind the per-row score widths and aggregation for these buffers.""" - if self._score_aggregation is not None: + def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> None: + """Bind the per-row score widths for these buffers (mean-only).""" + if self._score_launcher_bound: raise RuntimeError("TriAttention score launcher is already bound") - if score_aggregation not in ("mean", "max"): - raise ValueError(f"unsupported score aggregation: {score_aggregation}") + if aggregation != "mean": + raise ValueError( + f"unsupported score aggregation {aggregation!r}: max aggregation " + "was removed with the C++ score stack; only 'mean' exists" + ) self._score_valid_widths = valid_widths - self._score_aggregation = score_aggregation + self._score_launcher_bound = True def launch_prepared_score(self) -> torch.Tensor: """Launch the phase and score kernels over these buffers.""" from .triattention_kernels import prepare_mean_phase - if self._score_aggregation is None: + if not self._score_launcher_bound: raise RuntimeError("TriAttention score launcher is not bound") stream = torch.cuda.current_stream(self.device) if self.stream is None: @@ -786,30 +789,24 @@ def launch_prepared_score(self) -> torch.Tensor: raise _FixedScoreStreamMismatch( "TriAttention score launches must stay on the staging CUDA stream" ) - if self._score_aggregation == "mean" and self.fused_group._cute_score_runner is not None: - # mean_cos/mean_sin feed ONLY the opt-in CuTe score runner, whose - # compiled kernel captured their device pointers, so they must be - # refreshed before it launches. The default C++ mean path rotates - # init-time phase tables inside the score kernels' own CTA - # prologue instead, so production rounds launch zero phase or - # coefficient kernels. - prepare_mean_phase( - self.round_starts_device, - self.offsets, - self.omega, - self.mean_cos, - self.mean_sin, - self.max_requests, - ) + # mean_cos/mean_sin feed the CuTe score kernel, whose compiled launch + # captured their device pointers, so they must be refreshed from this + # round's staged round starts before it runs. + prepare_mean_phase( + self.round_starts_device, + self.offsets, + self.omega, + self.mean_cos, + self.mean_sin, + self.max_requests, + ) return self.fused_group.launch( self.max_requests, self.valid_seq_lens_device, self._score_valid_widths, - self.round_starts_device, self.token_starts_device, self.mean_cos, self.mean_sin, - self._score_aggregation, ) def stage( @@ -1060,7 +1057,6 @@ def __init__( beta: int = 128, model_path: Optional[str] = None, calibration_path: Optional[str] = None, - score_aggregation: str = "mean", eviction_mode: str = "union", normalize_scores: bool = True, pin_prefill: bool = True, @@ -1092,7 +1088,6 @@ def __init__( ) # All physical moves use the C++ V2 compaction operation. # No other compaction path exists. - self.score_aggregation = score_aggregation # Calibration is the OFFICIAL TriAttention .pt (passed via # calibration_path), resolved + converted on the first request # (on_request_init). TRT-LLM does NOT compute calibration; model_path is @@ -1927,6 +1922,13 @@ def _fixed_resources_for( self.top_B + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), ) seq_capacity = max(needed_page_tokens, int(mgr.max_seq_len)) + # The CuTe score kernel stores full compute tiles (64 tokens, or one + # page for 128-token pages) into a scratch strided by this bucket + # capacity, so the capacity must be tile-aligned (its geometry gate + # rejects anything else). Rounding up costs at most one tile of + # scratch per segment and never changes scoring semantics. + score_tile_tokens = max(64, int(mgr.tokens_per_block)) + seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) dense_groups = list(layout.storage_groups.values()) @@ -2003,7 +2005,7 @@ def _fixed_resources_for( provisional.zero_() score_staging.bind_score_launcher( keep_set_selector.valid_widths, - self.score_aggregation, + "mean", ) resources = _EvictionBuffers( score_staging=score_staging, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py index cc8ca7a8fbe1..8e38abed1b1f 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -2,15 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 """SM100 CuTe-DSL scorer for the TriAttention mean-score path. -This is the production specialization of the final workbench kernel. It -uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, and -producer-only page-ID lookahead. The public integration keeps the compiled -C++ score ops as the implementation for every geometry outside the exact -contract validated here. +This is the production specialization of the final workbench kernel — and +the ONLY score implementation. It uses split real/imag TMA loads, BF16 and +FP16 compensated UMMA, sqrt FTZ, and producer-only page-ID lookahead. +Geometries outside the exact contract validated here raise loudly at setup +(``_FixedScoreGroup.prepare_cute_score``); there is no fallback path. Page-table contract: ``page_ids`` is the flattened native block-offset -staging buffer ([pool_slot, request, K/V plane, block] int32) shared with -the C++ score op; K-plane entries encode ``physical_page * kv_factor`` and +staging buffer ([pool_slot, request, K/V plane, block] int32) produced by +the V2 manager; K-plane entries encode ``physical_page * kv_factor`` and are decoded inline (kv_factor == 2), so no per-round conversion pass is needed. """ @@ -696,10 +696,8 @@ def kernel( # buffer ([pool_slot, request, K/V plane, block] int32) and # ``page_off`` points at one request's K plane. K-plane # entries encode ``physical_page * kv_factor`` (kv_factor is - # 2 for the interleaved K/V pools this kernel requires); the - # C++ score op decodes the same buffer with - # ``encoded / kvFactor`` (triAttentionScoreKernels.cu), so - # divide by two here as well. V-plane entries are never read. + # 2 for the interleaved K/V pools this kernel requires), so + # divide by two here. V-plane entries are never read. producer_prefetched_page_id_lane0 = ( cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile]) // 2 ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 1dc77c6676f9..8860d175fbae 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -3,18 +3,15 @@ """GPU kernels for the TriAttention KV-eviction pipeline. The production path uses one fixed-shape trig-score launch across all dense -layers, CuTE-DSL TopK selection, and grouped C++ compaction. This module owns -the score launcher and its persistent metadata; scoring itself runs through -the compiled ``trtllm`` CUDA ops (on the mean path ONE folded paged-score -launch whose CTAs rotate their own coefficients from init-time phase tables; -on the max path a trigonometric coefficient fold, then the folded paged -score) for every supported geometry. The -original Triton score kernel has been deleted; -the unit tests validate the CUDA ops against an independent PyTorch oracle. -Selection and compaction live in their respective runtime modules. An -optional SM100 CuTe-DSL specialization (``triattention_cute_score.py``, -default off behind ``TRTLLM_TRIATTENTION_CUTE_SCORE=1``) can take over the -mean-aggregation score launch for one exactly-validated geometry. +layers, CuTE-DSL TopK selection, and grouped C++ compaction. Scoring runs +EXCLUSIVELY through the SM100 CuTe-DSL kernel +(``triattention_cute_score.py``): mean aggregation, BF16 KV pools, head size +64/128, 32/128-token pages, GQA group 4 or 8. There is deliberately no other +score path -- any geometry outside that contract raises loudly at setup +instead of routing to a slower kernel (the original Triton score kernel and +the C++ CUDA score stack have both been deleted). The unit tests validate +the CuTe kernel against an independent PyTorch oracle. Selection and +compaction live in their respective runtime modules. House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). @@ -25,8 +22,6 @@ from __future__ import annotations -import os -import warnings from typing import List import torch @@ -109,177 +104,6 @@ def prepare_mean_phase( ) -def _launch_tri_score_perhead( - group: "_FixedScoreGroup", - request_count: int, - num_segments: int, - valid_seq_lens: torch.Tensor, - valid_widths: torch.Tensor, - round_starts_device: torch.Tensor, - token_starts_device: torch.Tensor, - *, - score_aggregation: str, - mean_rotate_in_cta: bool = True, -) -> None: - """Prepare the per-round coefficients, then score paged KV via the C++ ops. - - The compiled ``trtllm`` score ops are THE implementation for every - geometry this launcher accepts; unsupported inputs fail loudly inside the - ops (TORCH_CHECK) instead of routing to another kernel. The unit tests - validate them against an independent PyTorch oracle. - - ``mean_rotate_in_cta`` (mean aggregation only, default on) hands the - per-round coefficient rotation to the score kernels' own CTA prologue: - one launch per round, zero coefficient scratch. The two-launch - preparation (standalone rotation kernel writing global tables the score - kernels then read) stays selectable because the unit tests prove the two - paths produce bit-identical scores — the kernels share one rotation - expression — and because geometries whose coefficient block exceeds the - score launch's shared-memory bound must fall back to it. - """ - if score_aggregation not in ("mean", "max"): - raise ValueError(f"unsupported score aggregation: {score_aggregation}") - if not ( - hasattr(torch.ops.trtllm, "tri_attention_fold_score_coefficients") - and hasattr(torch.ops.trtllm, "tri_attention_rotate_mean_score_coefficients") - and hasattr(torch.ops.trtllm, "tri_attention_paged_score") - ): - raise RuntimeError( - "this TensorRT-LLM build is missing the TriAttention score ops; rebuild the C++ " - "th_common extension (there is deliberately no Triton fallback: a loud failure " - "here beats silently scoring through a slower path)" - ) - use_max = score_aggregation == "max" - ( - num_q_heads, - num_kv_heads, - num_freqs, - tokens_per_block, - kv_factor, - num_offsets, - s_page, - s_kv_head, - s_slot, - s_dim, - ) = group.geometry_args - # Mean aggregation collapses all offsets into one coefficient plane; max - # keeps one c_re/c_im plane per offset because max does not commute - # through the frequency sum. - offset_planes = num_offsets if use_max else 1 - # Keyword operands of the score op that trigger its in-CTA coefficient - # rotation; empty when pre-rotated c_re/c_im planes are passed instead. - rotate_in_cta_kwargs: dict = {} - if use_max: - c_re, c_im, c_mlr = group._fold_coefficient_buffers(offset_planes, with_mlr=True) - q_real, q_imag, mlr_coef = group.pointer_middle - freq_scale_sq, omega, offsets = group.pointer_tail - torch.ops.trtllm.tri_attention_fold_score_coefficients( - c_re, - c_im, - c_mlr, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - None, - None, - omega, - offsets, - round_starts_device, - request_count, - group._num_calibrated_layers, - num_q_heads, - num_freqs, - offset_planes, - True, - # Per-layer dequant scales (quantized pools only, else None): the - # fold multiplies them into the coefficient tables so the score op - # below reads raw quantized elements at zero hot-loop cost. - group._kv_scales, - ) - elif mean_rotate_in_cta: - # Production mean path: the score kernels rotate the pre-scaled - # calibration query by the tabulated offset-mean phase of each - # request's round start in their own CTA prologue (tables built once - # at group construction; phase-table design: Fanrong Li, torch-graph - # review 2026-07-20). No rotation launch, no per-round coefficient - # scratch; c_mlr is the static init-time table. The prologue compiles - # the standalone rotation kernel's exact arithmetic, so scores are - # bit-identical to the two-launch path below. - c_re = None - c_im = None - c_mlr = group._mlr_fold - rotate_in_cta_kwargs = dict( - phase_cos=group._phase_cos, - phase_sin=group._phase_sin, - q_real_scaled=group._q_real_scaled, - q_imag_scaled=group._q_imag_scaled, - round_starts=round_starts_device, - max_position=group._max_position, - ) - else: - # Two-launch mean preparation: the standalone rotation kernel writes - # global c_re/c_im planes the score kernels then read. Kept callable - # as the unit tests' bit-equality reference leg and as the fallback - # for coefficient blocks past the score launch's shared-memory bound. - c_re, c_im, _ = group._fold_coefficient_buffers(offset_planes, with_mlr=False) - c_mlr = group._mlr_fold - torch.ops.trtllm.tri_attention_rotate_mean_score_coefficients( - c_re, - c_im, - group._q_real_scaled, - group._q_imag_scaled, - group._phase_cos, - group._phase_sin, - round_starts_device, - request_count, - group._num_calibrated_layers, - num_q_heads, - num_freqs, - group._max_position, - ) - pool_anchor, layer_base_addrs, block_offsets, seg_page_off, seg_req, seg_layer = ( - group.pointer_prefix - ) - torch.ops.trtllm.tri_attention_paged_score( - pool_anchor, - layer_base_addrs, - block_offsets, - seg_page_off, - seg_req, - seg_layer, - valid_seq_lens, - valid_widths, - token_starts_device, - c_re, - c_im, - c_mlr, - group.output, - group.output_width, - group.num_layers, - request_count, - group._num_calibrated_layers, - num_q_heads, - num_kv_heads, - num_freqs, - tokens_per_block, - kv_factor, - offset_planes, - s_page, - s_kv_head, - s_slot, - s_dim, - num_segments, - use_max, - group._use_vectorized, - # Validation-only here (presence must match the pool dtype; the max - # fold above or the init-time pre-scaled mean tables already consumed - # the values). - group._kv_scales, - **rotate_in_cta_kwargs, - ) - - class _FixedScoreGroup: """Persistent score metadata/output for one fixed geometry. @@ -288,12 +112,11 @@ class _FixedScoreGroup: uses the native TRT-LLM attention layout and ``page_table_slots`` maps each scored layer to its V2 pool slot. - LIFETIME CONTRACT: the group captures the scored layer pools as raw device - addresses (``layer_base_addrs``) and keeps a reference only to the anchor - pool (the score op's dtype witness). The caller owns ``layer_pools`` and must - keep every scored pool alive for as long as it launches through this group - (in production the V2 KV-cache manager does); a dropped pool leaves its - address dangling and scores read allocator-recycled memory. + LIFETIME: the group retains references to every scored layer pool -- the + SM100 CuTe score kernel encodes immutable TMA descriptors from their raw + device addresses at compile time, so the pools must stay alive (and stay + put) for as long as the group launches. In production the V2 KV-cache + manager owns them for the manager's lifetime. """ def __init__( @@ -313,7 +136,6 @@ def __init__( omega: torch.Tensor, offsets: torch.Tensor, output_width: int, - kv_scales: torch.Tensor | None = None, ) -> None: if not layer_indices or min(max_requests, page_count, seq_len) <= 0: raise ValueError("fixed score group requires non-empty positive geometry") @@ -347,16 +169,10 @@ def __init__( self.num_freqs, tokens_per_block, kv_factor, - int(offsets.numel()), - strides[0], - strides[2], - strides[3], - strides[4], ) # Per-layer ABSOLUTE base addresses. Layers may live in distinct # storages (V2 TensorWrapper-per-layer); only geometry must be uniform. element_size = p0.element_size() - bases_16b_aligned = True layer_base_addrs = torch.zeros(len(layer_pools), dtype=torch.int64, device=device) for layer in layer_indices: pool = layer_pools[layer] @@ -369,114 +185,14 @@ def __init__( address = int(pool.data_ptr()) if address % element_size: raise ValueError("fixed score layer base is not element-aligned") - bases_16b_aligned &= address % 16 == 0 layer_base_addrs[layer] = address - # The score op runs 16-byte 8-frequency K loads when the fixed layout - # guarantees aligned rows, and its strided scalar path otherwise. - # Audited ONCE here: bases and strides never change for this group. - strides_16b_aligned = all( - (element_size * stride) % 16 == 0 for stride in (strides[0], strides[2], strides[3]) - ) - self._use_vectorized = ( - p0.dtype in (torch.bfloat16, torch.float16) - and self.num_freqs % 8 == 0 - and strides[4] == 1 - and bases_16b_aligned - and strides_16b_aligned - ) - # Quantized (fp8/int8) pools are FUNCTIONAL-ONLY and scalar-path-only - # (the dtype gate above already excludes them from the vectorized - # path). Their per-layer dequantization scale is folded into the - # score coefficients at launch time, so it must be present up front; - # conversely, scales alongside a float pool would double-scale the - # coefficients, so that pairing is rejected just as loudly. - quantized_pool = p0.dtype in (torch.float8_e4m3fn, torch.int8) - if quantized_pool and kv_scales is None: - raise ValueError("quantized (fp8/int8) KV pools require per-layer kv_scales") - if not quantized_pool and kv_scales is not None: - raise ValueError("kv_scales are only valid for quantized (fp8/int8) KV pools") - self._kv_scales = ( - None - if kv_scales is None - else kv_scales.to(device=device, dtype=torch.float32).contiguous().view(-1) - ) # Calibration tables span every model layer; segments index them by - # ABSOLUTE layer id, so the fold covers the full calibrated extent. + # ABSOLUTE layer id, so the tables cover the full calibrated extent. self._num_calibrated_layers = q_real_LHF.numel() // (int(num_q_heads) * self.num_freqs) - # Segment layer ids index the fold tables ON DEVICE where they cannot - # be range-checked; validate the extent once here, loudly. + # Segment layer ids index the calibration tables ON DEVICE where they + # cannot be range-checked; validate the extent once here, loudly. if min(layer_indices) < 0 or max(layer_indices) >= self._num_calibrated_layers: raise ValueError("scored layer index exceeds the calibrated layer extent") - # Init-once tables for the mean-aggregation coefficient rotation - # (design: Fanrong Li, torch-graph review 2026-07-20). The - # offset-averaged phase of every possible round-start position is - # tabulated once, RoPE-style (float64 accumulation, stored fp32): - # phase_cos[pos, f] = freq_scale_sq[f] * mean_o cos((pos + offset_o) * omega_f) - # phase_sin[pos, f] = freq_scale_sq[f] * mean_o sin((pos + offset_o) * omega_f) - # and the per-layer KV dequantization scale is pre-multiplied into - # the calibration query (identity for float pools). The MLR - # coefficient has no position term, so its fold - # (kv_scale * freq_scale_sq * mlr) is fully static: the request axis - # disappears and nothing about it is recomputed per round. Every - # eviction round then gathers one table row per request and rotates - # the static query by it -- zero trigonometry at runtime. The score - # kernels perform that rotation themselves in an in-CTA prologue by - # default; the standalone rotation kernel remains the two-launch - # reference flavor (bit-identical scores, proven by the unit tests). - # - # Positions can legitimately reach the full sequence capacity, so the - # table covers [0, seq_len] inclusive. The 64-row floor keeps tiny - # synthetic unit-test geometries (whose logical round starts exceed - # their physical bucket) inside the table at negligible cost; every - # row is exact for its position, and production buckets (the - # manager's max sequence length) always exceed the floor. - self._max_position = max(int(seq_len), 63) + 1 - omega64 = omega.view(-1)[: self.num_freqs].to(torch.float64) - freq_scale_sq64 = freq_scale_sq.view(-1)[: self.num_freqs].to(torch.float64) - positions = torch.arange(self._max_position, dtype=torch.float64, device=device) - cos_acc = torch.zeros( - self._max_position, self.num_freqs, dtype=torch.float64, device=device - ) - sin_acc = torch.zeros_like(cos_acc) - for offset in offsets.to(torch.float64).tolist(): - angle = (positions + offset).unsqueeze(1) * omega64.unsqueeze(0) - cos_acc += torch.cos(angle) - sin_acc += torch.sin(angle) - phase_scale = freq_scale_sq64.unsqueeze(0) / float(offsets.numel()) - self._phase_cos = (cos_acc * phase_scale).to(torch.float32).contiguous() - self._phase_sin = (sin_acc * phase_scale).to(torch.float32).contiguous() - calibration_shape = (self._num_calibrated_layers, int(num_q_heads), self.num_freqs) - mlr_fold64 = mlr_coef_LHF.view(calibration_shape).to(torch.float64) * freq_scale_sq64 - if self._kv_scales is None: - # kv_scale is 1.0 for float pools: the pre-scaled query IS the - # calibration query (aliased, not copied). - self._q_real_scaled = q_real_LHF.view(-1) - self._q_imag_scaled = q_imag_LHF.view(-1) - else: - layer_scales64 = ( - self._kv_scales[: self._num_calibrated_layers].to(torch.float64).view(-1, 1, 1) - ) - self._q_real_scaled = ( - (q_real_LHF.view(calibration_shape).to(torch.float64) * layer_scales64) - .to(torch.float32) - .contiguous() - .view(-1) - ) - self._q_imag_scaled = ( - (q_imag_LHF.view(calibration_shape).to(torch.float64) * layer_scales64) - .to(torch.float32) - .contiguous() - .view(-1) - ) - mlr_fold64 = mlr_fold64 * layer_scales64 - self._mlr_fold = mlr_fold64.to(torch.float32).contiguous().view(-1) - # Folded per-round coefficient tables, allocated on first launch and - # keyed by plane count so switching aggregation (mean: one plane; - # max: one plane per offset) re-shapes without churn. - self._fold_buffers: dict = {} - # The anchor pool is passed to the CUDA score op ONLY as its dtype - # witness: the op recovers the pool element type from it and never - # reads data through it. seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( self.num_layers ) @@ -525,42 +241,44 @@ def __init__( mlr_coef_LHF.view(-1), ) self.pointer_tail = (freq_scale_sq, omega, offsets) - # Optional SM100 CuTe mean-score specialization (see - # triattention_cute_score.py), compiled by the first - # ``prepare_cute_score`` call and default OFF behind the - # TRTLLM_TRIATTENTION_CUTE_SCORE environment knob (read once here). - # The CuTe runner encodes TMA descriptors from the actual pool - # tensors, so pool references are retained ONLY when the knob is on; - # the default path keeps the raw-address-only lifetime contract - # documented above. + # The SM100 CuTe score kernel (see triattention_cute_score.py) is THE + # score implementation; it is compiled by the first + # ``prepare_cute_score`` call. The runner encodes TMA descriptors from + # the actual pool tensors, hence the pool references retained here + # (see the LIFETIME note in the class docstring). self.seq_len = int(seq_len) self._cute_score_runner = None self._cute_score_attempted = False - cute_score_enabled = os.environ.get("TRTLLM_TRIATTENTION_CUTE_SCORE", "0") == "1" - self._cute_layer_pools = list(layer_pools) if cute_score_enabled else None + self._cute_layer_pools = list(layer_pools) self._cute_layer_indices = [int(layer) for layer in layer_indices] def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> None: - """Compile the optional SM100 CuTe mean-score specialization once. + """Compile the SM100 CuTe score kernel once; raise loudly otherwise. Call this outside CUDA graph capture: compilation allocates memory - and synchronizes. With the environment knob unset (the default) or on - any unsupported geometry this returns without importing the CuTe - module, and every launch keeps using the compiled C++ score ops. + and synchronizes. The CuTe kernel is the ONLY score implementation, + so an unsupported geometry raises ValueError here and a runner + construction failure raises RuntimeError -- there is deliberately no + fallback path. Supported contract: SM100 exactly, BF16 pools, 32- or 128-token - pages, 32 or 64 frequencies (head size 64/128), and 4 or 8 query - heads per KV head — this covers the Qwen3 and GPT-OSS production + pages, 32 or 64 frequencies (head size 64/128), 4 or 8 query heads + per KV head, and a bucket capacity (``seq_len``) aligned to the + kernel's compute tile — this covers the Qwen3 and GPT-OSS production geometries as well as the original validation shape. """ if self._cute_score_attempted: return self._cute_score_attempted = True - if self._cute_layer_pools is None: - return anchor = self.pointer_prefix[0] - num_q_heads, num_kv_heads, num_freqs, tokens_per_block, kv_factor = self.geometry_args[:5] + num_q_heads, num_kv_heads, num_freqs, tokens_per_block, kv_factor = self.geometry_args max_segments = self.max_requests * self.num_layers + # The kernel's epilogue stores full compute tiles (64 tokens, or one + # page for 128-token pages) into a scratch whose per-segment stride + # is seq_len, without masking the ragged tail of the LAST tile; an + # unaligned bucket would silently spill scores into the next + # segment's region, so it is rejected here instead. + score_tile_tokens = max(64, int(tokens_per_block)) supported = ( torch.cuda.get_device_capability(anchor.device) == (10, 0) and anchor.dtype == torch.bfloat16 @@ -570,22 +288,24 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> and num_q_heads % num_kv_heads == 0 and num_q_heads // num_kv_heads in (4, 8) and int(anchor.stride(-1)) == 1 + and self.seq_len % score_tile_tokens == 0 # The kernel computes flat score offsets in 32-bit arithmetic; # group-4 geometries pad the head axis to the MMA tile N=8. and num_kv_heads * 8 * max_segments * self.seq_len < 2**31 ) if not supported: - warnings.warn( - "TriAttention CuTe score not engaged despite " - "TRTLLM_TRIATTENTION_CUTE_SCORE=1: " + raise ValueError( + "TriAttention score requires SM100, bf16 KV pools, head size " + "64/128, 32/128-token pages, GQA group 4 or 8, and a bucket " + "capacity aligned to the score compute tile; got " f"capability={torch.cuda.get_device_capability(anchor.device)}, " f"dtype={anchor.dtype}, kv_factor={kv_factor}, " f"tokens_per_block={tokens_per_block}, num_freqs={num_freqs}, " f"heads={num_q_heads}q/{num_kv_heads}kv, " f"stride={int(anchor.stride(-1))}, " + f"seq_len={self.seq_len} (tile {score_tile_tokens}), " f"offset_audit={num_kv_heads * 8 * max_segments * self.seq_len}" ) - return device = anchor.device try: from .triattention_cute_score import TriAttentionCuteScoreRunner @@ -637,72 +357,43 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> freq_scale_sq=self.pointer_tail[0], output=scratch, ) - self._cute_scratch = scratch - self._cute_seg_seq_len = seg_seq_len - self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) - from tensorrt_llm.logger import logger - - logger.info( - f"TriAttention CuTe score enabled: {num_q_heads}q/{num_kv_heads}kv heads, " - f"{num_freqs} freqs, {tokens_per_block}-token pages" - ) except (ImportError, RuntimeError, ValueError, AssertionError) as error: - warnings.warn( - f"TriAttention CuTe score setup failed; using the C++ score ops: {error}", - RuntimeWarning, - stacklevel=2, - ) - - def _fold_coefficient_buffers( - self, offset_planes: int, with_mlr: bool - ) -> "tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]": - """Return (c_re, c_im, c_mlr) per-round scratch for one plane count. - - Sized on ``max_requests`` so any launch's active ``request_count`` - fits without reallocation (each launch prepares only its active - rows). The production mean path allocates NOTHING here: its score - kernels rotate coefficients in-CTA, so only the max aggregation - (which also writes a per-round c_mlr) and the two-launch mean - reference path (static MLR table, c_re/c_im planes only) ever call - this. - """ - key = (offset_planes, with_mlr) - buffers = self._fold_buffers.get(key) - if buffers is None: - elements = ( - self.max_requests - * self._num_calibrated_layers - * int(self.geometry_args[0]) - * self.num_freqs - ) - device = self.output.device - c_re = torch.empty(offset_planes * elements, dtype=torch.float32, device=device) - c_im = torch.empty_like(c_re) - c_mlr = torch.empty(elements, dtype=torch.float32, device=device) if with_mlr else None - buffers = (c_re, c_im, c_mlr) - self._fold_buffers[key] = buffers - return buffers + raise RuntimeError( + "TriAttention CuTe score setup failed and no other score path exists" + ) from error + self._cute_scratch = scratch + self._cute_seg_seq_len = seg_seq_len + self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) + from tensorrt_llm.logger import logger + + logger.info( + f"TriAttention CuTe score enabled: {num_q_heads}q/{num_kv_heads}kv heads, " + f"{num_freqs} freqs, {tokens_per_block}-token pages" + ) def launch( self, request_count: int, valid_seq_lens: torch.Tensor, valid_widths: torch.Tensor, - round_starts_device: torch.Tensor, token_starts_device: torch.Tensor, mean_cos: torch.Tensor, mean_sin: torch.Tensor, - score_aggregation: str, - *, - mean_rotate_in_cta: bool = True, + aggregation: str = "mean", ) -> torch.Tensor: """Return decode-only scores as ``[request, layer, head, token]``. - ``mean_rotate_in_cta=False`` selects the two-launch mean coefficient - preparation (standalone rotation kernel + global-table score reads); - see ``_launch_tri_score_perhead``. Scores are bit-identical either - way — the unit tests compare the two with ``torch.equal``. + Runs the SM100 CuTe score kernel (the only score implementation) and + writes each request's decode width (``valid_seq_len - token_start``) + into ``valid_widths``, which the selection reduce kernels consume. + Only mean aggregation exists; ``request_count`` must be one of the + precompiled variants (1 or the group capacity). """ + if aggregation != "mean": + raise ValueError( + f"unsupported score aggregation {aggregation!r}: max aggregation " + "was removed with the C++ score stack; only 'mean' exists" + ) if request_count <= 0 or request_count > self.max_requests: raise ValueError("request count exceeds fixed score capacity") if ( @@ -713,95 +404,71 @@ def launch( ): raise ValueError("score output lengths do not fit the keep-set selector") num_segments = request_count * self.num_layers - if num_segments > 65535: - # Segments sit on the y grid axis (CUDA caps y/z at 65535) so the - # unbounded x axis can hold the token tiles of long sequences. - raise ValueError("request*layer segment count exceeds the CUDA grid limit") output = self.output[:request_count] - if score_aggregation == "mean": - # Lazy compile covers groups used without their owning workspace - # (unit tests); production compiles in the workspace constructor, - # outside CUDA graph capture. Default off: one attribute check. - self.prepare_cute_score(mean_cos, mean_sin) - runner = self._cute_score_runner - if runner is not None and not runner.supports(request_count): - missed = getattr(self, "_cute_supports_warned", set()) - if request_count not in missed: - missed.add(request_count) - self._cute_supports_warned = missed - warnings.warn( - f"TriAttention CuTe score compiled variants miss " - f"request_count={request_count}; falling back to the C++ " - f"score op for this round" - ) - if runner is not None and runner.supports(request_count): - if not getattr(self, "_cute_engaged_logged", False): - self._cute_engaged_logged = True - warnings.warn( - f"TriAttention CuTe score engaged (request_count={request_count})" - ) - # Stage per-segment valid lengths (segment = request x layer). - torch.index_select( - valid_seq_lens, - 0, - self.pointer_prefix[4][:num_segments], - out=self._cute_seg_seq_len[:num_segments], - ) - runner.launch(request_count, mean_cos, mean_sin) - # The kernel wrote full-sequence scores from physical token - # zero into its head-major scratch. Gather each request's - # decode window (starting at its pinned prompt length) into - # the group output so callers see exactly the layout the C++ - # score ops produce. This costs one extra read+write of the - # score volume per round, only on this opt-in path; columns - # past a request's valid width carry unscored scratch data, - # matching the C++ op, whose consumers mask by valid width. - num_q_heads = int(self.geometry_args[0]) - num_kv_heads = int(self.geometry_args[1]) - group_size = num_q_heads // num_kv_heads - # The scratch head axis is padded to the MMA tile N=8 per - # KV head; slicing the view to the real group size skips - # the zero padding columns. - source = ( - self._cute_scratch[: num_kv_heads * 8 * num_segments * self.seq_len] - .view(num_kv_heads, 8, request_count, self.num_layers, self.seq_len)[ - :, :group_size - ] - .permute(2, 3, 0, 1, 4) - ) - columns = token_starts_device[:request_count].to(torch.int64).view( - -1, 1, 1, 1, 1 - ) + self._cute_gather_columns.view(1, 1, 1, 1, -1) - columns = columns.clamp_(max=self.seq_len - 1).expand( - request_count, - self.num_layers, - num_kv_heads, - group_size, - self.output_width, - ) - torch.gather( - source, - 4, - columns, - out=output.view( - request_count, - self.num_layers, - num_kv_heads, - group_size, - self.output_width, - ), - ) - return output - _launch_tri_score_perhead( - self, - request_count, - num_segments, + # Lazy compile covers groups used without their owning workspace + # (unit tests); production compiles in the workspace constructor, + # outside CUDA graph capture. + self.prepare_cute_score(mean_cos, mean_sin) + runner = self._cute_score_runner + if runner is None or not runner.supports(request_count): + raise RuntimeError( + f"TriAttention CuTe score has no compiled variant for " + f"request_count={request_count} (capacity {self.max_requests}) " + "and no other score path exists" + ) + # Per-request decode widths for the selection reduce kernels; the + # deleted C++ score op used to write these (seq_len - token_start). + torch.sub( + valid_seq_lens[:request_count], + token_starts_device[:request_count], + out=valid_widths[:request_count], + ) + # Stage per-segment valid lengths (segment = request x layer). + torch.index_select( valid_seq_lens, - valid_widths, - round_starts_device, - token_starts_device, - score_aggregation=score_aggregation, - mean_rotate_in_cta=mean_rotate_in_cta, + 0, + self.pointer_prefix[4][:num_segments], + out=self._cute_seg_seq_len[:num_segments], + ) + runner.launch(request_count, mean_cos, mean_sin) + # The kernel wrote full-sequence scores from physical token zero into + # its head-major scratch. Gather each request's decode window + # (starting at its pinned prompt length) into the group output, the + # ``[request, layer, head, token]`` layout the selection kernels + # read. Columns past a request's valid width carry unscored scratch + # data; consumers mask by ``valid_widths``. + num_q_heads = int(self.geometry_args[0]) + num_kv_heads = int(self.geometry_args[1]) + group_size = num_q_heads // num_kv_heads + # The scratch head axis is padded to the MMA tile N=8 per KV head; + # slicing the view to the real group size skips the zero padding + # columns. + source = ( + self._cute_scratch[: num_kv_heads * 8 * num_segments * self.seq_len] + .view(num_kv_heads, 8, request_count, self.num_layers, self.seq_len)[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + columns = token_starts_device[:request_count].to(torch.int64).view( + -1, 1, 1, 1, 1 + ) + self._cute_gather_columns.view(1, 1, 1, 1, -1) + columns = columns.clamp_(max=self.seq_len - 1).expand( + request_count, + self.num_layers, + num_kv_heads, + group_size, + self.output_width, + ) + torch.gather( + source, + 4, + columns, + out=output.view( + request_count, + self.num_layers, + num_kv_heads, + group_size, + self.output_width, + ), ) return output diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index f8df2e35476f..16639539047a 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -194,12 +194,12 @@ def torch_tri_score_oracle( omega, offsets, layer_indices, - aggregation, ): - """Independent Torch implementation of the paged TriAttention score. + """Independent Torch implementation of the paged TriAttention mean score. - Covers both aggregations (mean and max), GQA head mapping via - ``head // group_size``, and the position-independent MLR term. + Covers GQA head mapping via ``head // group_size`` and the + position-independent MLR term. Mean aggregation only: it is the single + production aggregation (max was removed with the C++ score stack). """ scores = [] num_q_heads = int(q_real.shape[1]) @@ -228,23 +228,9 @@ def torch_tri_score_oracle( 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 - if aggregation == "mean": - position = ( - freq_scale_sq * (product_real * mean_cos - product_imag * mean_sin) - ).sum(dim=-1) - else: - position = ( - ( - freq_scale_sq[None, None, :] - * ( - product_real[None] * torch.cos(phase)[:, None, :] - - product_imag[None] * torch.sin(phase)[:, None, :] - ) - ) - .sum(dim=-1) - .max(dim=0) - .values - ) + 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] 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 index 3d9c339b1eb7..88f022640dfe 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Correctness coverage for the optional SM100 TriAttention CuTe scorer.""" +"""Correctness coverage for the SM100 TriAttention CuTe scorer (the only score path).""" import pytest import torch @@ -28,7 +28,6 @@ ], ) def test_cute_score_matches_torch_mean_oracle( - monkeypatch: pytest.MonkeyPatch, tokens_per_block: int, page_permutation: list, valid_lens: "list | None", @@ -36,7 +35,6 @@ def test_cute_score_matches_torch_mean_oracle( num_q_heads: int, ) -> None: pytest.importorskip("cutlass") - monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_SCORE", "1") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _FixedScoreGroup, @@ -98,7 +96,6 @@ def test_cute_score_matches_torch_mean_oracle( valid_lens = [seq_len, seq_len] valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) valid_widths = torch.tensor(valid_lens, dtype=torch.int32, device=device) - round_starts_device = torch.tensor([seq_len, seq_len + 1], dtype=torch.int32, device=device) token_starts_device = torch.zeros(2, dtype=torch.int32, device=device) for request_count in (1, 2): group.output.fill_(float("nan")) @@ -106,11 +103,9 @@ def test_cute_score_matches_torch_mean_oracle( request_count, valid_seq_lens, valid_widths, - round_starts_device, token_starts_device, mean_cos, mean_sin, - "mean", ) assert actual.shape == (request_count, 1, num_q_heads, seq_len) for request in range(request_count): @@ -130,6 +125,5 @@ def test_cute_score_matches_torch_mean_oracle( ) torch.cuda.synchronize() - # Fails loudly if setup silently fell back to the C++ score ops (whose - # scores would also match the oracle here). + # The CuTe runner is the only score path; prove setup actually built it. assert group._cute_score_runner is not None diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py deleted file mode 100644 index 3274a3336646..000000000000 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_phase_rotation.py +++ /dev/null @@ -1,236 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Mean-path phase-table rotation: numerical fidelity and path equality. - -The mean-aggregation coefficient preparation tabulates the offset-averaged -phase of every possible round-start position once at initialization (float64 -accumulation, stored fp32) and rotates the pre-scaled calibration query by -one gathered table row per request. The first test rebuilds the same -coefficients with direct float64 trigonometry at the exact round starts and -compares. The second proves the two places that rotation can run — the -standalone rotation kernel writing global coefficient planes, and the score -kernels' own in-CTA shared-memory prologue (the production default) — yield -BIT-IDENTICAL paged-score outputs, because both compile one shared rotation -expression. -""" - -import pytest -import torch -from test_triattention_score_ops import _build_case as _build_score_case - -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, -) - - -def test_phase_table_rotation_matches_direct_trig_fold(): - """Table + rotation == direct trig fold at positions {0, mid, last}. - - The last position exercises phases of tens of thousands of radians, - where fp32 runtime trigonometry would already have lost several digits - to argument reduction; the float64-built table must not. - """ - assert hasattr(torch.ops.trtllm, "tri_attention_rotate_mean_score_coefficients"), ( - "TriAttention rotation op is not loaded" - ) - device = torch.device("cuda", torch.cuda.current_device()) - torch.manual_seed(20260720) - num_layers = 2 - num_q_heads = 4 - head_dim = 16 - num_freqs = head_dim // 2 - # Sizes the phase table: the group tabulates [0, seq_len] inclusive. - seq_len = 32768 - pools = [ - torch.randn(2, 2, 2, 4, head_dim, device=device).to(torch.bfloat16) - for _ in range(num_layers) - ] - block_offsets = torch.zeros(1, 3, 2, 1, dtype=torch.int32, device=device) - q_real = torch.randn(num_layers, num_q_heads, num_freqs, device=device) - q_imag = torch.randn(num_layers, num_q_heads, num_freqs, device=device) - mlr_coef = torch.randn(num_layers, num_q_heads, num_freqs, device=device) - freq_scale_sq = torch.rand(num_freqs, device=device) + 0.5 - # RoPE-style inverse frequencies: omega[0] == 1.0 makes the last tabulated - # position a genuinely large trigonometric argument. - omega = 10000.0 ** (-torch.arange(num_freqs, device=device, dtype=torch.float32) / num_freqs) - offsets = torch.tensor([1.0, 2.0, 4.0, 8.0], dtype=torch.float32, device=device) - group = _FixedScoreGroup( - pools, - list(range(num_layers)), - 3, - 1, - seq_len, - num_q_heads, - block_offsets, - [0] * num_layers, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - output_width=4, - ) - max_position = group._max_position - assert max_position == seq_len + 1 - - positions = [0, max_position // 2, max_position - 1] - round_starts = torch.tensor(positions, dtype=torch.int32, device=device) - total = len(positions) * num_layers * num_q_heads * num_freqs - c_re = torch.full((total,), float("nan"), dtype=torch.float32, device=device) - c_im = torch.full_like(c_re, float("nan")) - torch.ops.trtllm.tri_attention_rotate_mean_score_coefficients( - c_re, - c_im, - group._q_real_scaled, - group._q_imag_scaled, - group._phase_cos, - group._phase_sin, - round_starts, - len(positions), - num_layers, - num_q_heads, - num_freqs, - max_position, - ) - - # Direct trigonometric fold: average the offset phases at each round - # start, then rotate and scale the calibration query (float64 throughout, - # cast to fp32 only at the end, exactly like the tables were built). - phase = ( - round_starts.to(torch.float64)[:, None, None] + offsets.to(torch.float64)[None, :, None] - ) * omega.to(torch.float64)[None, None, :] - mean_cos = torch.cos(phase).mean(dim=1)[:, None, None, :] - mean_sin = torch.sin(phase).mean(dim=1)[:, None, None, :] - fss = freq_scale_sq.to(torch.float64) - q_re = q_real.to(torch.float64)[None] - q_im = q_imag.to(torch.float64)[None] - reference_re = (fss * (q_re * mean_cos - q_im * mean_sin)).to(torch.float32) - reference_im = (fss * (q_im * mean_cos + q_re * mean_sin)).to(torch.float32) - - shape = (len(positions), num_layers, num_q_heads, num_freqs) - torch.testing.assert_close(c_re.view(shape), reference_re, rtol=1e-5, atol=1e-5) - torch.testing.assert_close(c_im.view(shape), reference_im, rtol=1e-5, atol=1e-5) - - -# One geometry per code path the in-CTA rotation prologue compiles into: -# both CUDA load paths (16-byte vectorized chunks and strided scalar), GQA -# groups 2/4/8 with dedicated template instantiations, group 3 through the -# per-query-head grid mapping, and every float pool dtype. The shared -# builder's per-request sequence lengths are ragged, so the write mask is -# compared too (sentinel-filled outputs). -_EQUALITY_CASES = [ - pytest.param( - dict( - head_dim=128, - tokens_per_block=32, - num_q_heads=8, - num_kv_heads=2, - dtype=torch.bfloat16, - ), - id="vectorized_bf16_group4", - ), - pytest.param( - dict( - head_dim=128, - tokens_per_block=32, - num_q_heads=8, - num_kv_heads=1, - dtype=torch.bfloat16, - ), - id="vectorized_bf16_group8", - ), - pytest.param( - dict( - head_dim=128, - tokens_per_block=32, - num_q_heads=6, - num_kv_heads=2, - dtype=torch.float16, - ), - id="vectorized_fp16_group3_per_query_head", - ), - pytest.param( - dict( - head_dim=8, - tokens_per_block=4, - num_q_heads=4, - num_kv_heads=2, - dtype=torch.bfloat16, - ), - id="scalar_bf16_group2", - ), - pytest.param( - dict( - head_dim=8, - tokens_per_block=4, - num_q_heads=6, - num_kv_heads=2, - dtype=torch.float32, - ), - id="scalar_fp32_runtime_group3", - ), -] - - -@pytest.mark.parametrize("case", _EQUALITY_CASES) -def test_in_cta_rotation_scores_bit_equal_to_standalone_rotation(case): - """Full mean-path score outputs are torch.equal between rotation flavors. - - Bit-equality (not a tolerance) is the contract: the score kernels' in-CTA - prologue and the standalone rotation kernel share one rotation - expression, and the score accumulation downstream is identical, so any - single-ulp drift means the arithmetic diverged and must fail here. - """ - assert hasattr(torch.ops.trtllm, "tri_attention_paged_score"), ( - "TriAttention paged score op is not loaded" - ) - assert hasattr(torch.ops.trtllm, "tri_attention_rotate_mean_score_coefficients"), ( - "TriAttention rotation op is not loaded" - ) - request_count = 3 - ( - group, - round_starts, - token_starts, - valid_seq_lens, - _, - mean_cos, - mean_sin, - _, - ) = _build_score_case( - request_count=request_count, - max_requests=4, - num_layers=2, - page_count=4, - prompt_len=5, - seed=20260723, - offsets=[1.0, 2.0, 4.0], - **case, - ) - device = group.output.device - sentinel = -54321.0 - - def run(mean_rotate_in_cta: bool) -> "tuple[torch.Tensor, torch.Tensor]": - group.output.fill_(sentinel) - valid_widths = torch.zeros(request_count, dtype=torch.int32, device=device) - scores = group.launch( - request_count, - valid_seq_lens, - valid_widths, - round_starts, - token_starts, - mean_cos, - mean_sin, - "mean", - mean_rotate_in_cta=mean_rotate_in_cta, - ).clone() - return scores, valid_widths - - # Reference leg: the kernel-round2 preparation (standalone rotation - # kernel + score kernels reading the global coefficient planes). - reference_scores, reference_widths = run(mean_rotate_in_cta=False) - in_cta_scores, in_cta_widths = run(mean_rotate_in_cta=True) - assert not reference_scores.eq(sentinel).all(), "reference leg scored nothing" - assert torch.equal(in_cta_scores, reference_scores) - assert torch.equal(in_cta_widths, reference_widths) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 39871d396015..0d874e5505e4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -57,6 +57,14 @@ _TORCH_TOPK_ORACLE = torch.topk +# The SM100 CuTe kernel is the only score path, so every test that actually +# launches scores (or builds the real staging buffers, whose constructor +# compiles the kernel) is 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 _set_request_state( manager, @@ -901,7 +909,7 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize score_staging.bind_score_launcher.assert_called_once_with( keep_set_selector.valid_widths, - manager.score_aggregation, + "mean", ) plan = build_selection.call_args.args[0] assert plan.eviction_mode == eviction_mode @@ -1157,8 +1165,10 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ) assert all(not hasattr(item, "page_ids") for item in prepared) + @requires_sm100 @pytest.mark.parametrize("request_count", [1, 7, 8]) def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, request_count): + pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _FixedScoreStagingBuffers, _FixedScoreStreamMismatch, @@ -1167,23 +1177,36 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re device = torch.device("cuda", torch.cuda.current_device()) max_requests = 8 page_count = 3 - seq_len = 7 - page_table_token_capacity = 11 - layer_elements = max_requests * page_count * 2 * 1 * 4 * 4 - shared = torch.randn(2 * layer_elements, device=device) + # The bucket capacity must be aligned to the score kernel's 64-token + # compute tile (the staging constructor compiles the kernel). + seq_len = 64 + page_table_token_capacity = 90 + tokens_per_block = 32 + head_dim = 64 + num_freqs = head_dim // 2 + num_q_heads = 8 + layer_elements = max_requests * page_count * 2 * 1 * tokens_per_block * head_dim + shared = torch.randn(2 * layer_elements, device=device).to(torch.bfloat16) + pool_shape = (max_requests * page_count, 2, 1, tokens_per_block, head_dim) pools = [ - shared[:layer_elements].view(max_requests * page_count, 2, 1, 4, 4), - shared[layer_elements:].view(max_requests * page_count, 2, 1, 4, 4), - torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device), - torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device), + shared[:layer_elements].view(pool_shape), + shared[layer_elements:].view(pool_shape), + torch.randn(pool_shape, device=device).to(torch.bfloat16), + torch.randn(pool_shape, device=device).to(torch.bfloat16), ] dense_groups = [[0, 1], [2]] representatives = [0, 2, 3] - q_real = torch.randn(4, 2, 4, dtype=torch.float64, device=device)[..., ::2] - q_imag = torch.randn(4, 2, 4, dtype=torch.float64, device=device)[..., ::2] - mlr = torch.randn(4, 2, 4, dtype=torch.float64, device=device)[..., ::2] - freq = torch.tensor([1.0, 0.0, 1.0, 0.0], dtype=torch.float64, device=device)[::2] - omega = torch.tensor([0.01, 0.0, 0.03, 0.0], dtype=torch.float64, device=device)[::2] + q_real = torch.randn(4, num_q_heads, 2 * num_freqs, dtype=torch.float64, device=device)[ + ..., ::2 + ] + q_imag = torch.randn(4, num_q_heads, 2 * num_freqs, dtype=torch.float64, device=device)[ + ..., ::2 + ] + mlr = torch.randn(4, num_q_heads, 2 * num_freqs, dtype=torch.float64, device=device)[ + ..., ::2 + ] + freq = (torch.rand(2 * num_freqs, dtype=torch.float64, device=device) + 0.5)[::2] + omega = (torch.rand(2 * num_freqs, dtype=torch.float64, device=device) * 0.05)[::2] offsets = torch.tensor([1.0, 0.0, 2.0, 0.0], dtype=torch.float64, device=device)[::2] assert not q_real.is_contiguous() assert not freq.is_contiguous() @@ -1196,8 +1219,8 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re representatives, max_requests, seq_len, - 2, - 2, + num_q_heads, + num_freqs, q_real, q_imag, mlr, @@ -1316,9 +1339,10 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): staging.stage(manager, request_ids, round_starts, token_starts) assert gather.call_count == calls - @pytest.mark.parametrize("request_count", [1, 7, 8]) - @pytest.mark.parametrize("aggregation", ["mean", "max"]) - def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, aggregation): + @requires_sm100 + @pytest.mark.parametrize("request_count", [1, 8]) + def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count): + pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _FixedScoreGroup, ) @@ -1327,24 +1351,29 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, torch.manual_seed(20260703 + request_count) max_requests = 8 page_count = 2 - seq_len = 7 + tokens_per_block = 32 + head_dim = 64 + num_freqs = head_dim // 2 + num_q_heads = 8 + seq_len = page_count * tokens_per_block prompt_len = 2 page_ids = torch.arange(max_requests * page_count, dtype=torch.int64, device=device).view( max_requests, page_count ) - layer_elements = max_requests * page_count * 2 * 1 * 4 * 4 - shared = torch.randn(2 * layer_elements, device=device) + layer_elements = max_requests * page_count * 2 * 1 * tokens_per_block * head_dim + shared = (0.125 * torch.randn(2 * layer_elements, device=device)).to(torch.bfloat16) + pool_shape = (max_requests * page_count, 2, 1, tokens_per_block, head_dim) pools = [ - shared[:layer_elements].view(max_requests * page_count, 2, 1, 4, 4), - shared[layer_elements:].view(max_requests * page_count, 2, 1, 4, 4), - torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device), + shared[:layer_elements].view(pool_shape), + shared[layer_elements:].view(pool_shape), + (0.125 * torch.randn(pool_shape, device=device)).to(torch.bfloat16), ] storage_groups = [[0, 1], [2]] - q_real = torch.randn(3, 2, 4, device=device)[..., ::2] - q_imag = torch.randn(3, 2, 4, device=device)[..., ::2] - mlr = torch.randn(3, 2, 4, device=device)[..., ::2] - freq = torch.tensor([0.7, 0.0, 1.3, 0.0], device=device)[::2] - omega = torch.tensor([0.013, 0.0, 0.071, 0.0], device=device)[::2] + q_real = (0.125 * torch.randn(3, num_q_heads, 2 * num_freqs, device=device))[..., ::2] + q_imag = (0.125 * torch.randn(3, num_q_heads, 2 * num_freqs, device=device))[..., ::2] + mlr = (0.125 * torch.randn(3, num_q_heads, 2 * num_freqs, device=device))[..., ::2] + freq = (torch.rand(2 * num_freqs, device=device) + 0.5)[::2] + omega = (torch.rand(2 * num_freqs, device=device) * 0.05)[::2] offsets = torch.tensor([1.0, 0.0, 2.0, 0.0, 4.0, 0.0], device=device)[::2] assert not q_real.is_contiguous() assert not q_imag.is_contiguous() @@ -1371,7 +1400,6 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, omega, offsets, [0, 1, 2], - aggregation, ) for layers in storage_groups: group = _FixedScoreGroup( @@ -1380,7 +1408,7 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, max_requests, page_count, seq_len, - 2, + num_q_heads, _encode_block_offsets(page_ids), [0] * len(layers), q_real, @@ -1391,39 +1419,36 @@ def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count, offsets, output_width=seq_len - prompt_len, ) - valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) + valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) fixed = group.launch( request_count, torch.tensor(seq_lens, dtype=torch.int32, device=device), valid_widths, - round_device, token_starts_device, torch.cos(phase).mean(dim=1), torch.sin(phase).mean(dim=1), - aggregation, ) - assert valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + assert valid_widths[:request_count].tolist() == [ + seq_len - prompt_len for seq_len in seq_lens + ] assert fixed.shape == ( request_count, len(layers), - 2, + num_q_heads, seq_len - prompt_len, ) for request in range(request_count): for layer_slot, layer in enumerate(layers): valid_width = seq_lens[request] - prompt_len segment = fixed[request, layer_slot, :, :valid_width] - expected = oracle[request * len(pools) + layer][:, prompt_len:] + expected = oracle[request * len(pools) + layer][ + :, prompt_len : prompt_len + valid_width + ] torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) - selected = torch.topk(segment.max(dim=0).values, 3).indices.sort().values - expected_selected = ( - torch.topk(expected.max(dim=0).values, 3).indices.sort().values - ) - assert torch.equal(selected, expected_selected) + @requires_sm100 @pytest.mark.parametrize("request_count", [1, 7, 8]) - @pytest.mark.parametrize("aggregation", ["mean", "max"]) - def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count, aggregation): + def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count): """ONE launch over layers in DISTINCT storages with DISTINCT block tables. This is the production V2 shape: get_buffers wraps every layer as its @@ -1431,6 +1456,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun the fused path must not assume a shared storage anchor or a shared per-request block table. """ + pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _FixedScoreStagingBuffers, @@ -1444,12 +1470,21 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun torch.manual_seed(20260707 + request_count) max_requests = request_count page_count = 2 - seq_len = 7 + tokens_per_block = 32 + head_dim = 64 + num_freqs = head_dim // 2 + num_q_heads = 8 + seq_len = page_count * tokens_per_block prompt_len = 1 num_layers = 3 # Three SEPARATE allocations (distinct storages, like V2 TensorWrapper). pools = [ - torch.randn(max_requests * page_count, 2, 1, 4, 4, device=device) + ( + 0.125 + * torch.randn( + max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device + ) + ).to(torch.bfloat16) for _ in range(num_layers) ] assert len({pool.untyped_storage().data_ptr() for pool in pools}) == num_layers @@ -1465,11 +1500,11 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun for _ in range(num_layers) ] ).contiguous() - q_real = torch.randn(num_layers, 2, 2, device=device) - q_imag = torch.randn(num_layers, 2, 2, device=device) - mlr = torch.randn(num_layers, 2, 2, device=device) - freq = torch.tensor([0.7, 1.3], device=device) - omega = torch.tensor([0.013, 0.071], device=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 = 0.125 * torch.randn(num_layers, num_q_heads, num_freqs, device=device) + freq = torch.rand(num_freqs, device=device) + 0.5 + omega = torch.rand(num_freqs, device=device) * 0.05 offsets = torch.tensor([1.0, 2.0, 4.0], device=device) round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 round_starts = round_device[:request_count].tolist() @@ -1483,7 +1518,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun max_requests, page_count, seq_len, - 2, + num_q_heads, block_offsets, layer_order, # slot i holds layer i's tables q_real, @@ -1496,28 +1531,25 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) - mean_cos = torch.empty(request_count, 2, dtype=torch.float32, device=device) + mean_cos = torch.empty(request_count, num_freqs, dtype=torch.float32, device=device) mean_sin = torch.empty_like(mean_cos) - if aggregation == "mean": - triattention_kernels.prepare_mean_phase( - round_device, - offsets, - omega, - mean_cos, - mean_sin, - request_count, - ) + triattention_kernels.prepare_mean_phase( + round_device, + offsets, + omega, + mean_cos, + mean_sin, + request_count, + ) score_sentinel = -12345.0 group.output.fill_(score_sentinel) checked = group.launch( request_count, valid_seq_lens, valid_widths, - round_device, token_starts, mean_cos, mean_sin, - aggregation, ).clone() staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) staging.device = group.output.device @@ -1532,8 +1564,8 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun staging.omega = omega staging.stream = None staging._score_valid_widths = None - staging._score_aggregation = None - staging.bind_score_launcher(valid_widths, aggregation) + staging._score_launcher_bound = False + staging.bind_score_launcher(valid_widths, "mean") group.output.fill_(score_sentinel) fixed = staging.launch_prepared_score().clone() torch.testing.assert_close(fixed, checked, rtol=0, atol=0) @@ -1553,13 +1585,14 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun omega, offsets, layer_order, - aggregation, ) for request in range(request_count): for layer_slot, layer in enumerate(layer_order): valid_width = seq_lens[request] - prompt_len segment = fixed[request, layer_slot, :, :valid_width] - expected = oracle[request * num_layers + layer][:, prompt_len:] + expected = oracle[request * num_layers + layer][ + :, prompt_len : prompt_len + valid_width + ] torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) round_device.add_(17) @@ -1571,15 +1604,14 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun device=device, ) ) - if aggregation == "mean": - triattention_kernels.prepare_mean_phase( - round_device, - offsets, - omega, - mean_cos, - mean_sin, - request_count, - ) + triattention_kernels.prepare_mean_phase( + round_device, + offsets, + omega, + mean_cos, + mean_sin, + request_count, + ) expected_second_widths = valid_seq_lens - prompt_len group.output.fill_(score_sentinel) valid_widths.fill_(-1) @@ -1587,11 +1619,9 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun request_count, valid_seq_lens, valid_widths, - round_device, token_starts, mean_cos, mean_sin, - aggregation, ).clone() group.output.fill_(score_sentinel) valid_widths.fill_(-1) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py index 2f2c39ce57c8..eda49607cbbd 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py @@ -1,16 +1,15 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""CUDA score ops vs an independent PyTorch oracle. - -`_launch_tri_score_perhead` calls the compiled `trtllm` fold + paged-score -ops unconditionally (no Triton fallback; the original Triton score kernel has -been deleted). These tests score the same paged pools through a pure-PyTorch -oracle and compare the two implementations across the geometry matrix the -launcher must cover: both CUDA load paths (vectorized 8-frequency chunks and -strided scalar), both aggregations, every supported pool dtype (bf16/fp16/ -fp32, plus functional-only fp8_e4m3fn/int8 with per-layer dequantization -scales), and GQA group sizes with and without a dedicated template -instantiation. +"""The CuTe score kernel vs an independent PyTorch oracle, through group.launch. + +The SM100 CuTe-DSL kernel (``triattention_cute_score.py``) is the ONLY score +implementation. These tests drive it through ``_FixedScoreGroup.launch`` -- +the exact production entry point -- across the supported production +geometries, with multi-layer segments, permuted page tables, ragged valid +lengths, and per-request prompt windows, and compare against a pure-PyTorch +oracle that recomputes everything independently. They also pin the +loud-failure contract: unsupported geometry, removed aggregations, and +uncompiled request counts raise instead of routing to another kernel. """ import pytest @@ -22,64 +21,14 @@ _FixedScoreGroup, ) - -def _require_score_ops() -> None: - """The compiled score ops are a hard prerequisite for these tests.""" - assert hasattr(torch.ops.trtllm, "tri_attention_fold_score_coefficients"), ( - "TriAttention fold op is not loaded" - ) - assert hasattr(torch.ops.trtllm, "tri_attention_paged_score"), ( - "TriAttention paged score op is not loaded" - ) - - -def _oracle_reference( - group: _FixedScoreGroup, - pools: list, - oracle_inputs: dict, - request_count: int, - seq_lens: list, - round_starts: torch.Tensor, - prompt_len: int, - aggregation: str, - sentinel: float, -) -> torch.Tensor: - """Sentinel-filled oracle scores in the ops' [request, layer, head, token] layout. - - The oracle scores every cached token in [0, seq_len); the score ops write - only the decode region [prompt_len, seq_len) at column origin 0. Slice the - prompt columns off each oracle row and leave every column the ops must not - touch at the sentinel, so the comparison covers the write MASK as well as - the values. - """ - num_layers = group.num_layers - oracle = _torch_tri_score_oracle( - pools, - oracle_inputs["page_ids"][:request_count], - seq_lens, - round_starts[:request_count].tolist(), - 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)), - aggregation, - ) - reference = torch.full_like(group.output, sentinel) - for request in range(request_count): - width = seq_lens[request] - prompt_len - for layer in range(num_layers): - reference[request, layer, :, :width] = oracle[request * num_layers + layer][ - :, prompt_len: - ] - return reference +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( *, - request_count: int, max_requests: int, num_layers: int, page_count: int, @@ -87,35 +36,33 @@ def _build_case( head_dim: int, num_q_heads: int, num_kv_heads: int, - dtype: torch.dtype, - offsets: list, prompt_len: int, seed: int, - pools: list | None = None, - kv_scales: torch.Tensor | None = None, + 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 - # Callers may inject prebuilt pools: the quantized tests build a quantized - # pool (plus its dequantized fp32 twin for the oracle reference leg) from - # ONE set of randoms and pass the quantized pool in here. - if pools is None: - pools = [ - torch.randn( + # The 0.125 scaling keeps the BF16 key/coefficient products small so the + # kernel-vs-oracle tolerance can stay tight across the frequency sum. + pools = [ + ( + 0.125 + * torch.randn( max_requests * page_count, 2, num_kv_heads, tokens_per_block, head_dim, device=device, - ).to(dtype) - for _ in range(num_layers) - ] + ) + ).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 = torch.randn(num_layers, num_q_heads, num_freqs, device=device) - q_imag = torch.randn(num_layers, num_q_heads, num_freqs, device=device) - mlr_coef = torch.randn(num_layers, num_q_heads, num_freqs, device=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) @@ -136,25 +83,15 @@ def _build_case( omega, offsets_t, output_width=capacity - prompt_len, - kv_scales=kv_scales, ) - # LIFETIME: the group records only raw device ADDRESSES of the scored - # layer pools (layer_base_addrs) and references just the anchor pool - # (pools[0], the kernel dtype witness). Production pools are owned by the - # KV-cache manager, so the group deliberately does not hold them; the - # test must keep the whole pool list alive itself. Dropping it here frees - # every non-anchor layer pool, whose blocks the caching allocator then - # recycles for the fold/output/reference tensors allocated later in the - # test — a use-after-free that reads back sentinel/coefficient bytes as - # K data (observed as layer>=1 inf/NaN score garbage). - group.test_pools_keepalive = pools round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) - seq_lens = [capacity - ((request * 3) % 5) for request in range(request_count)] + # Ragged valid lengths whose tails land mid-page and mid-compute-tile. + seq_lens = [capacity - ((request * 3) % 5) for request in range(max_requests)] valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) phase = (round_starts.float()[:, None, None] + offsets_t[None, :, None]) * omega[None, None, :] - mean_cos = torch.cos(phase).mean(dim=1) - mean_sin = torch.sin(phase).mean(dim=1) + mean_cos = torch.cos(phase).mean(dim=1).contiguous() + mean_sin = torch.sin(phase).mean(dim=1).contiguous() # Everything the PyTorch oracle needs to rebuild the reference leg # independently (it recomputes its own mean phases from these). oracle_inputs = dict( @@ -168,7 +105,6 @@ def _build_case( ) return ( group, - round_starts, token_starts, valid_seq_lens, seq_lens, @@ -178,440 +114,202 @@ def _build_case( ) -# --------------------------------------------------------------------------- -# Quantized (fp8_e4m3fn / int8) KV pools — FUNCTIONAL kernel-level coverage -# ONLY. End-to-end quantized-KV eviction is NOT validated here or anywhere -# else yet: nothing in the pipeline produces quantized pools, so these tests -# exercise the scalar score path + coefficient scale fold in isolation. The -# reference leg dequantizes the SAME quantized elements to fp32 pools and -# runs the PyTorch oracle on them, so quantization error cancels and only the -# scale-fold/loading math is under test. -# --------------------------------------------------------------------------- - -_QUANTIZED_GEOMETRY = dict( - request_count=3, - max_requests=4, - num_layers=2, - page_count=4, - tokens_per_block=4, - head_dim=8, - num_q_heads=4, - num_kv_heads=2, - prompt_len=5, - offsets=[1.0, 2.0, 4.0], -) - - -def _quantize_pools(raw_pools: list, dtype: torch.dtype): - """Per-layer amax quantization of fp32 pools. - - Returns (quantized pools, dequantized fp32 twin pools, per-layer scales). - The twin holds values dequantized FROM the quantized elements — NOT the - pre-quantization randoms — so quantization error is present in both legs - identically and the CUDA-vs-reference tolerance can stay tight. - """ - quant_max = 448.0 if dtype == torch.float8_e4m3fn else 127.0 - quantized, dequantized, scales = [], [], [] - for raw in raw_pools: - scale = raw.abs().amax().float() / quant_max - if dtype == torch.int8: - quant = torch.round(raw / scale).clamp(-127, 127).to(torch.int8) - else: - quant = (raw / scale).to(torch.float8_e4m3fn) - quantized.append(quant) - dequantized.append(quant.to(torch.float32) * scale) - scales.append(scale) - return quantized, dequantized, torch.stack(scales) - - -def _build_quantized_raw_pools(seed: int, device: torch.device) -> list: - """The shared fp32 randoms both quantized-test pools derive from.""" - torch.manual_seed(seed) - g = _QUANTIZED_GEOMETRY - return [ - torch.randn( - g["max_requests"] * g["page_count"], - 2, - g["num_kv_heads"], - g["tokens_per_block"], - g["head_dim"], - device=device, - ) - for _ in range(g["num_layers"]) - ] - - -# One entry per geometry class the launcher must cover. expected_vectorized -# white-boxes the launch-time path selection so the matrix provably exercises -# both CUDA load paths. +# One entry per supported production geometry: the Qwen3 shape (64 +# frequencies, GQA group 4 riding the padded MMA tile), the GPT-OSS shape +# (32 frequencies, group 8, 32-token pages spanning two page fragments per +# compute tile), and the originally validated 128-token-page shape. _CASES = [ pytest.param( dict( - head_dim=128, + max_requests=4, + num_layers=2, + page_count=4, tokens_per_block=32, - num_q_heads=8, - num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", - ), - True, - id="production_bf16_f64_group4_mean", - ), - pytest.param( - dict( head_dim=128, - tokens_per_block=32, num_q_heads=8, num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0, 8.0], - aggregation="max", - ), - True, - id="max_aggregation_four_offsets", - ), - pytest.param( - dict( - head_dim=128, - tokens_per_block=32, - num_q_heads=4, - num_kv_heads=2, - dtype=torch.float16, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", ), - True, - id="fp16_pool", + id="qwen3_f64_group4_tpb32", ), pytest.param( dict( - head_dim=128, - tokens_per_block=32, - num_q_heads=6, - num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", - ), - True, - id="group3_generic_head_mapping", - ), - pytest.param( - dict( - head_dim=32, + max_requests=2, + num_layers=3, + page_count=4, tokens_per_block=32, - num_q_heads=4, - num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", - ), - True, - id="f16_runtime_chunk_count", - ), - pytest.param( - dict( - head_dim=8, - tokens_per_block=4, - num_q_heads=4, - num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", - ), - False, - id="tiny_f4_scalar", - ), - pytest.param( - dict( - head_dim=8, - tokens_per_block=4, - num_q_heads=4, - num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0], - aggregation="max", - ), - False, - id="tiny_f4_scalar_max", - ), - pytest.param( - dict( - head_dim=12, - tokens_per_block=4, - num_q_heads=4, - num_kv_heads=2, - dtype=torch.bfloat16, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", + head_dim=64, + num_q_heads=8, + num_kv_heads=1, ), - False, - id="f6_nonpow2_masked_tail", + id="gptoss_f32_group8_tpb32", ), - # fp32 pools always take the scalar path; group=3 additionally exercises - # its runtime GQA loop against the reference (the other generic-group - # case is vectorized, the other scalar cases use a templated group size). pytest.param( dict( - head_dim=8, - tokens_per_block=4, - num_q_heads=6, - num_kv_heads=2, - dtype=torch.float32, - offsets=[1.0, 2.0, 4.0], - aggregation="mean", + max_requests=2, + num_layers=2, + page_count=2, + tokens_per_block=128, + head_dim=64, + num_q_heads=8, + num_kv_heads=1, ), - False, - id="fp32_scalar_generic_group3", + id="original_f32_group8_tpb128", ), ] +_QWEN3_CASE = dict(_CASES[0].values[0]) -class TestTriAttentionScoreOps: - @pytest.mark.parametrize("case,expected_vectorized", _CASES) - def test_cuda_ops_match_torch_oracle(self, case, expected_vectorized): - _require_score_ops() + +class TestTriAttentionScoreLaunch: + @requires_sm100 + @pytest.mark.parametrize("case", _CASES) + def test_cute_kernel_matches_torch_oracle(self, case): + pytest.importorskip("cutlass") case = dict(case) # parametrize reuses the dict across reruns - aggregation = case.pop("aggregation") - request_count = 3 prompt_len = 5 + max_requests = case["max_requests"] + num_layers = case["num_layers"] ( group, - round_starts, token_starts, valid_seq_lens, seq_lens, mean_cos, mean_sin, oracle_inputs, - ) = _build_case( - request_count=request_count, - max_requests=4, - num_layers=2, - page_count=4, - prompt_len=prompt_len, - seed=20260719, - **case, - ) - assert group._use_vectorized == expected_vectorized + ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) device = group.output.device - sentinel = -54321.0 - group.output.fill_(sentinel) - valid_widths_cuda = torch.empty(request_count, dtype=torch.int32, device=device) - cuda_scores = group.launch( - request_count, - valid_seq_lens, - valid_widths_cuda, - round_starts, - token_starts, - mean_cos, - mean_sin, - aggregation, - ).clone() - - # The oracle reads the same stored pool elements (up-cast to fp32, - # like the ops' loads), so the legs differ only by coefficient-fold - # association and reduction order. - reference = _oracle_reference( - group, - group.test_pools_keepalive, - oracle_inputs, - request_count, + oracle = _torch_tri_score_oracle( + group._cute_layer_pools, + oracle_inputs["page_ids"], seq_lens, - round_starts, - prompt_len, - aggregation, - sentinel, + [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)), ) - assert valid_widths_cuda.tolist() == [seq_len - prompt_len for seq_len in seq_lens] - # Sentinel-filled outputs make the comparison cover the write MASK as - # well as the values: any stray or missing store breaks equality. - # The ops' fp32 math tracks this oracle to ~2e-6 on these geometries, - # so 1e-4 is a tight gate with ample margin. - torch.testing.assert_close(cuda_scores, reference[:request_count], rtol=1e-4, atol=1e-4) - - @pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.int8], ids=["fp8_e4m3fn", "int8"]) - @pytest.mark.parametrize("aggregation", ["mean", "max"]) - def test_quantized_pool_matches_dequantized_reference(self, dtype, aggregation): - """Scale-folded scoring of quantized pools == dense scoring of their fp32 twin. - - The max aggregation additionally covers the per-offset coefficient - planes, which must all carry the folded per-layer scale. - """ - _require_score_ops() - device = torch.device("cuda", torch.cuda.current_device()) - seed = 20260721 - raw_pools = _build_quantized_raw_pools(seed, device) - quant_pools, dequant_pools, kv_scales = _quantize_pools(raw_pools, dtype) - ( - quant_group, - round_starts, - token_starts, - valid_seq_lens, - seq_lens, - mean_cos, - mean_sin, - oracle_inputs, - ) = _build_case( - dtype=dtype, - seed=seed, - pools=quant_pools, - kv_scales=kv_scales, - **_QUANTIZED_GEOMETRY, - ) - # Quantized pools must never select the vectorized load path. - assert quant_group._use_vectorized is False - request_count = _QUANTIZED_GEOMETRY["request_count"] - sentinel = -54321.0 - - quant_group.output.fill_(sentinel) - valid_widths_cuda = torch.empty(request_count, dtype=torch.int32, device=device) - cuda_scores = quant_group.launch( - request_count, - valid_seq_lens, - valid_widths_cuda, - round_starts, - token_starts, - mean_cos, - mean_sin, - aggregation, - ).clone() - - # Reference leg: the ORACLE over the dequantized-from-quantized fp32 - # twin pools (same page tables and calibration as the quantized group). - reference = _oracle_reference( - quant_group, - dequant_pools, - oracle_inputs, - request_count, - seq_lens, - round_starts, - _QUANTIZED_GEOMETRY["prompt_len"], - aggregation, - sentinel, - ) - - assert valid_widths_cuda.tolist() == [ - seq_len - _QUANTIZED_GEOMETRY["prompt_len"] for seq_len in seq_lens - ] - # Quantization error is identical in both legs (the reference pool is - # dequantized from the quantized values), so the only differences are - # scale-fold association (q*(s*c) vs (q*s)*c), the approximate sqrt, - # and reduction order — hence a tolerance close to the float cases'. - torch.testing.assert_close(cuda_scores, reference[:request_count], rtol=3e-3, atol=3e-3) - - def test_quantized_pool_missing_scales_raises(self): - with pytest.raises(ValueError, match="require per-layer kv_scales"): - _build_case(dtype=torch.int8, seed=20260722, **_QUANTIZED_GEOMETRY) - - def test_scales_with_float_pool_raises(self): - device = torch.device("cuda", torch.cuda.current_device()) - scales = torch.ones(_QUANTIZED_GEOMETRY["num_layers"], device=device) - with pytest.raises(ValueError, match="only valid for quantized"): - _build_case( - dtype=torch.bfloat16, seed=20260722, kv_scales=scales, **_QUANTIZED_GEOMETRY - ) - - def test_negative_scale_raises(self): - """Positivity is enforced host-side in the C++ op, at launch time. - - The |K| coefficient fold assumes |scale * K_q| == scale * |K_q|, which - breaks silently for non-positive scales, so the op must refuse them. - """ - _require_score_ops() - device = torch.device("cuda", torch.cuda.current_device()) - seed = 20260722 - raw_pools = _build_quantized_raw_pools(seed, device) - quant_pools, _, kv_scales = _quantize_pools(raw_pools, torch.int8) - bad_scales = kv_scales.clone() - bad_scales[0] = -bad_scales[0] - group, round_starts, token_starts, valid_seq_lens, _, mean_cos, mean_sin, _ = _build_case( - dtype=torch.int8, - seed=seed, - pools=quant_pools, - kv_scales=bad_scales, - **_QUANTIZED_GEOMETRY, - ) - request_count = _QUANTIZED_GEOMETRY["request_count"] - valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) - with pytest.raises(RuntimeError, match="strictly positive"): - group.launch( + # The runner precompiles exactly the request counts production + # launches: one and the full group capacity. + for request_count in dict.fromkeys((1, max_requests)): + group.output.fill_(float("nan")) + valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) + scores = group.launch( request_count, valid_seq_lens, valid_widths, - round_starts, token_starts, mean_cos, mean_sin, - "mean", ) - - def test_short_scales_raise(self): - """kv_scales must cover every calibrated layer, enforced at launch. - - The Python group only gates presence; the extent contract lives in the - C++ ops (segments index the fold tables by absolute layer id on - device, where a short scale tensor could not be range-checked). + assert scores.shape == ( + request_count, + num_layers, + case["num_q_heads"], + group.output_width, + ) + # The launch owns the per-request decode widths the selection + # reduce kernels consume (the deleted C++ op used to write them). + assert valid_widths[: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, + ) + + @requires_sm100 + def test_uncompiled_request_count_raises(self): + """A request count outside the precompiled variants fails loudly. + + There is no fallback kernel, so ``supports()`` misses must raise + instead of silently scoring through a slower path. """ - _require_score_ops() - device = torch.device("cuda", torch.cuda.current_device()) - seed = 20260722 - raw_pools = _build_quantized_raw_pools(seed, device) - quant_pools, _, kv_scales = _quantize_pools(raw_pools, torch.int8) - short_scales = kv_scales[:1] # the geometry calibrates two layers - group, round_starts, token_starts, valid_seq_lens, _, mean_cos, mean_sin, _ = _build_case( - dtype=torch.int8, - seed=seed, - pools=quant_pools, - kv_scales=short_scales, - **_QUANTIZED_GEOMETRY, + pytest.importorskip("cutlass") + ( + group, + token_starts, + valid_seq_lens, + _, + mean_cos, + mean_sin, + _, + ) = _build_case(prompt_len=5, seed=20260719, **_QWEN3_CASE) + valid_widths = torch.empty( + _QWEN3_CASE["max_requests"], dtype=torch.int32, device=group.output.device ) - request_count = _QUANTIZED_GEOMETRY["request_count"] - valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) - with pytest.raises(RuntimeError, match="one scale per calibrated layer"): + with pytest.raises(RuntimeError, match="no compiled variant"): group.launch( - request_count, + _QWEN3_CASE["max_requests"] - 1, valid_seq_lens, valid_widths, - round_starts, token_starts, mean_cos, mean_sin, - "mean", ) - def test_unsupported_pool_dtype_raises(self): - _require_score_ops() - # fp32 pools stay supported (the existing tiny-geometry unit suite - # drives them through this launcher); fp64 is genuinely outside the - # op's coverage and must fail loudly instead of routing elsewhere. - request_count = 2 - group, round_starts, token_starts, valid_seq_lens, _, mean_cos, mean_sin, _ = _build_case( - request_count=request_count, - max_requests=2, - num_layers=1, - page_count=2, - prompt_len=1, - seed=20260720, - head_dim=8, - tokens_per_block=4, - num_q_heads=2, - num_kv_heads=1, - dtype=torch.float64, - offsets=[1.0, 2.0], + def _tiny_unsupported_group(self, dtype: torch.dtype): + """A geometry far outside the CuTe contract (constructor accepts it).""" + 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 + pools = [ + torch.randn( + max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device + ).to(dtype) + for _ in range(num_layers) + ] + page_ids = ( + torch.arange(max_requests * page_count, device=device) + .view(max_requests, page_count) + .contiguous() ) - valid_widths = torch.empty(request_count, dtype=torch.int32, device=group.output.device) - with pytest.raises(RuntimeError, match="unsupported KV pool dtype"): - group.launch( - request_count, - valid_seq_lens, - valid_widths, - round_starts, - token_starts, - mean_cos, - mean_sin, - "mean", - ) + capacity = page_count * tokens_per_block + group = _FixedScoreGroup( + pools, + list(range(num_layers)), + max_requests, + page_count, + capacity, + 2, + _encode_block_offsets(page_ids), + [0] * num_layers, + torch.randn(num_layers, 2, num_freqs, device=device), + torch.randn(num_layers, 2, num_freqs, device=device), + torch.randn(num_layers, 2, num_freqs, device=device), + torch.rand(num_freqs, device=device) + 0.5, + torch.rand(num_freqs, device=device) * 0.05, + torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), + output_width=capacity - 1, + ) + device_args = dict(dtype=torch.int32, device=device) + return group, ( + torch.full((max_requests,), capacity, **device_args), + torch.empty(max_requests, **device_args), + torch.ones(max_requests, **device_args), + torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), + torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), + ) + + def test_unsupported_geometry_raises(self): + """Score setup outside the CuTe contract raises; nothing falls back.""" + group, launch_args = self._tiny_unsupported_group(torch.float32) + with pytest.raises(ValueError, match="TriAttention score requires SM100"): + group.launch(1, *launch_args) + + def test_max_aggregation_raises(self): + """Max aggregation was removed with the C++ score stack.""" + group, launch_args = self._tiny_unsupported_group(torch.bfloat16) + with pytest.raises(ValueError, match="max aggregation"): + group.launch(1, *launch_args, aggregation="max") 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 index 17e1f9dc0843..81bd19f6d4b3 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -23,6 +23,14 @@ def _require_cute_topk_op() -> None: ) +# 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 _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)) @@ -601,16 +609,25 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): assert torch.equal(cohort_pool, expected_pool) +@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") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _FixedScoreStagingBuffers, ) device = torch.device("cuda", torch.cuda.current_device()) num_layers = 3 + # The staged bucket capacity must be aligned to the score kernel's + # 64-token compute tile; the request itself stays 8 tokens long. + bucket_capacity = 64 seq_len = 8 keep_count = 2 + # GQA group 8 (the smallest CuTe-supported group with one KV head); all + # query heads share one zero calibration query and one MLR coefficient, + # so every head row carries the same |K|-driven score. + num_q_heads = 8 # bf16 pools in the compact op's supported geometry. The scored tokens # all live in each table's first entry, but the two tables still map the # two storage groups onto different physical pages, which is what the @@ -658,7 +675,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): pools.append(pool) initial_pools = [pool.clone() for pool in pools] - q_real = torch.zeros(num_layers, 1, num_freqs, dtype=torch.float32, device=device) + 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 @@ -670,8 +687,8 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): dense_layers, [0, 1], 1, - seq_len, - 1, + bucket_capacity, + num_q_heads, num_freqs, q_real, q_imag, @@ -692,9 +709,9 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): keep_set_selector = _BatchedPerHeadKeepSetSelector( eviction_mode="per_layer_perhead", dense_layers=tuple(dense_layers), - num_query_heads=1, + num_query_heads=num_q_heads, num_kv_heads=1, - width=seq_len, + width=bucket_capacity, keep_count=keep_count, dtype=torch.float32, device=device, @@ -739,16 +756,18 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): 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(): """Run two real eviction rounds through one live V2 cache. - The cache uses the compact op's supported geometry (bf16, 32-token - pages, head_dim 64): the request spans three pages so that compacting to - two pages still releases one physical page for reuse. Token scores are - tracked in a host-side mirror and the expected keep sets are derived - from it, replacing the hand-written score tables of the old 4-token-page - fixture. + The cache uses the score and compact kernels' supported geometry (bf16, + 32-token pages, head_dim 64): the request spans three pages so that + compacting to two pages still releases one physical page for reuse. + Token scores are tracked in a host-side mirror and the expected keep + sets are derived from it, replacing the hand-written score tables of the + old 4-token-page fixture. """ + pytest.importorskip("cutlass") import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( @@ -761,9 +780,12 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): device = torch.device("cuda", torch.cuda.current_device()) request_id = 7 prompt_len = 2 - seq_len = 66 + # The bucket capacity equals the confirmed length here and must be + # aligned to the score kernel's 64-token compute tile; the protected + # tail rides beyond it. + seq_len = 64 protected_tail = 2 - compacted_capacity = 34 + compacted_capacity = 36 tokens_per_block = 32 head_dim = 64 num_freqs = head_dim // 2 @@ -850,7 +872,11 @@ def expected_keep() -> torch.Tensor: device=device, ) - q_real = torch.zeros(1, 1, num_freqs, dtype=torch.float32, device=device) + # GQA group 8 (the smallest CuTe-supported group with one KV head); + # zero calibration query and a shared MLR coefficient give every + # query head the same |K|-driven score. + 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 @@ -863,7 +889,7 @@ def expected_keep() -> torch.Tensor: [0], 1, seq_len, - 1, + num_q_heads, num_freqs, q_real, q_imag, @@ -877,16 +903,18 @@ def expected_keep() -> torch.Tensor: page_table_token_capacity=seq_len + protected_tail, ) keep_set_selector = _BatchedUnionKeepSetSelector( - rows=1, + rows=num_q_heads, width=seq_len - prompt_len, keep_count=keep_count, dtype=torch.float32, device=device, max_requests=1, dense_layers=(0,), - num_query_heads=1, + num_query_heads=num_q_heads, num_kv_heads=1, - input_scores=score_staging.fused_group.output.view(1, 1, seq_len - prompt_len), + input_scores=score_staging.fused_group.output.view( + 1, num_q_heads, seq_len - prompt_len + ), normalize_scores=False, prompt_offsets_buffer=score_staging.token_starts_device, ) From dbaccb84ccab58fd6be1d556323a0f45ab794629 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 20 Jul 2026 23:14:24 -0700 Subject: [PATCH 065/178] [None][chore] Drop year-only diffs against main The retired-fallback commit left two dtype-instantiation files whose only remaining change against main was the copyright year bump; restore them byte-identical to main so the review diff carries no noise. Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_float_float.cu | 2 +- .../unfusedAttentionKernels_2_half_half.cu | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu index 4a150117540b..55e3e8756afe 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_float_float.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu index a6fcb6dc36ad..5abd544359d1 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_half_half.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); From 5013e9998421723a648d1254510631a71e4c5f16 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 00:45:02 -0700 Subject: [PATCH 066/178] [None][feat] Fuse settle+pack and gather mean phases from a position table Retire the two remaining per-round Triton launches on the eviction path: - Settle-ties and pack-compaction-sources now run as one fused kernel with HAS_SETTLE/HAS_PACK constexpr halves; the draft co-compaction flow packs pre-settled ordinals through the same kernel. - The per-round trig kernel is replaced by a MeanPhaseTable built once at setup (RoPE-style position rows of offset-averaged cos/sin) and a single vendored Triton gather into the CuTe-bound mean buffers. The table regrows host-side at staging time and the gather clamps stale rows instead of faulting. Signed-off-by: tianruih --- .../triattention/compaction.py | 65 +-- .../triattention/triattention.py | 61 ++- .../triattention/triattention_kernels.py | 386 +++++++++--------- .../test_triattention_fused_settle_pack.py | 75 +++- .../test_triattention_pipeline.py | 67 +-- 5 files changed, 382 insertions(+), 272 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index d2448754cba9..56bb83adf9c4 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -312,7 +312,7 @@ def _move_index_pack_launcher( swa_indices_arg = move_source_indices swa_total = 0 - from .triattention_kernels import _pack_compaction_sources_kernel + from .triattention_kernels import _settle_ties_and_pack_compaction_sources_kernel max_move = decode_keep_count + max_protected_tail if swa_total: @@ -336,39 +336,40 @@ def _move_index_pack_launcher( per_layer=per_layer, has_swa=swa_total > 0, ) - packed_row_count = num_dense_layers * num_kv_heads if per_layer else num_kv_heads - grid = ( - request_count, - packed_row_count, - (max_move + _PACK_BLOCK_TOKENS - 1) // _PACK_BLOCK_TOKENS, - ) - bound_tensors = ( - kept_token_ordinals, - valid_sequence_lengths, - move_source_offsets, - move_source_indices, - swa_offsets_arg, - swa_indices_arg, - ) - # Ordered to match the kernel's constexpr parameter declaration. - constexpr_values = dict( - DENSE_TOTAL=pack_arguments.dense_total, - SWA_TOTAL=pack_arguments.swa_total, - SELECTION_ROWS=pack_arguments.selection_rows, - SELECTION_STRIDE=pack_arguments.keep_count, - KEEP_COUNT=pack_arguments.keep_count, - NUM_KV_HEADS=pack_arguments.num_kv_heads, - SWA_WINDOW=pack_arguments.swa_window, - UNION=pack_arguments.union, - PER_LAYER=pack_arguments.per_layer, - HAS_SWA=pack_arguments.has_swa, - BLOCK=_PACK_BLOCK_TOKENS, - ) + # One program per (request, selection row); the settle half is compiled + # away because the ordinals arrive pre-settled (the draft flow reuses + # the target's keep set verbatim), so only the pack half runs. The + # settle-side pointer arguments are compiled away with it; any + # well-formed tensor stands in for them. + grid = (request_count, selection_rows) def launch_pack() -> None: - _pack_compaction_sources_kernel[grid]( - *bound_tensors, - **constexpr_values, + _settle_ties_and_pack_compaction_sources_kernel[grid]( + kept_token_ordinals, + valid_sequence_lengths, + move_source_offsets, + kept_token_ordinals, + kept_token_ordinals, + valid_sequence_lengths, + move_source_offsets, + move_source_indices, + swa_offsets_arg, + swa_indices_arg, + WIDTH=decode_keep_count, + KEEP_COUNT=decode_keep_count, + OUTPUT_WIDTH=decode_keep_count, + SELECTION_ROWS=selection_rows, + DENSE_TOTAL=pack_arguments.dense_total, + SWA_TOTAL=pack_arguments.swa_total, + MOVE_CAPACITY=pack_arguments.move_capacity, + NUM_KV_HEADS=pack_arguments.num_kv_heads, + SWA_WINDOW=pack_arguments.swa_window, + UNION=pack_arguments.union, + PER_LAYER=pack_arguments.per_layer, + HAS_SWA=pack_arguments.has_swa, + HAS_SETTLE=False, + HAS_PACK=True, + BLOCK=_PACK_BLOCK_TOKENS, num_warps=_PACK_NUM_WARPS, ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 6c2c1fd8843f..f108ad5f4d24 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -69,6 +69,9 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig if TYPE_CHECKING: + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + MeanPhaseTable, + ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests @@ -312,6 +315,7 @@ def _select_top_tokens(self) -> None: OUTPUT_WIDTH=self.keep_count, SELECTION_ROWS=self.selection_rows_per_request, **pack_shape, + HAS_SETTLE=True, BLOCK=256, num_warps=4, ) @@ -541,6 +545,7 @@ def __init__( freq_scale_sq: torch.Tensor, offsets: torch.Tensor, omega: torch.Tensor, + mean_phase_table: Optional["MeanPhaseTable"] = None, page_table_keys: Optional[List[object]] = None, num_page_table_slots: Optional[int] = None, decode_width: Optional[int] = None, @@ -551,7 +556,7 @@ def __init__( draft_num_page_table_slots: Optional[int] = None, draft_page_table_token_capacity: Optional[int] = None, ) -> None: - from .triattention_kernels import _FixedScoreGroup + from .triattention_kernels import MeanPhaseTable, _FixedScoreGroup if not dense_groups or not dense_layers or not page_representatives or max_requests <= 0: raise ValueError("fixed score metadata requires non-empty positive geometry") @@ -703,7 +708,9 @@ def __init__( dtype=torch.int32, device=self.device, ) - self.request_metadata_device = torch.empty( + # Zero-filled so an unstaged cohort gathers the phase table's row 0 + # instead of indexing it with uninitialized round starts. + self.request_metadata_device = torch.zeros( (6, max_requests + 1), dtype=torch.int32, device=self.device ) self.round_starts_device = self.request_metadata_device[0, :max_requests] @@ -716,12 +723,19 @@ def __init__( self.dense_move_offsets = self.request_metadata_device[3] self.swa_move_offsets = self.request_metadata_device[4] self.draft_move_offsets = self.request_metadata_device[5] - self.mean_cos = torch.empty( - (max_requests, num_freqs), dtype=torch.float32, device=self.device + # Stacked cos/sin planes: one index_select refreshes both, and each + # plane stays a contiguous tensor for the CuTe launch to capture. + self.mean_phases = torch.empty( + (2, max_requests, num_freqs), dtype=torch.float32, device=self.device ) - self.mean_sin = torch.empty_like(self.mean_cos) - self.offsets = offsets - self.omega = omega + self.mean_cos = self.mean_phases[0] + self.mean_sin = self.mean_phases[1] + # The phase table depends only on the shared calibration, so the + # manager passes one instance to every staging bucket; standalone + # construction (tests) builds a private one. + if mean_phase_table is None: + mean_phase_table = MeanPhaseTable(offsets, omega, initial_rows=seq_len) + self.mean_phase_table = mean_phase_table # ONE fused group across ALL dense layers: segments carry their own # layer base address and page-table slot, so distinct per-layer # storages/block tables no longer force one launch per storage group. @@ -774,9 +788,7 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> N self._score_launcher_bound = True def launch_prepared_score(self) -> torch.Tensor: - """Launch the phase and score kernels over these buffers.""" - from .triattention_kernels import prepare_mean_phase - + """Gather the phase means and launch the score kernel over these buffers.""" if not self._score_launcher_bound: raise RuntimeError("TriAttention score launcher is not bound") stream = torch.cuda.current_stream(self.device) @@ -790,14 +802,11 @@ def launch_prepared_score(self) -> torch.Tensor: "TriAttention score launches must stay on the staging CUDA stream" ) # mean_cos/mean_sin feed the CuTe score kernel, whose compiled launch - # captured their device pointers, so they must be refreshed from this - # round's staged round starts before it runs. - prepare_mean_phase( + # captured their device pointers, so they must be refreshed in place + # from this round's staged round starts before it runs. + self.mean_phase_table.gather( self.round_starts_device, - self.offsets, - self.omega, - self.mean_cos, - self.mean_sin, + self.mean_phases, self.max_requests, ) return self.fused_group.launch( @@ -864,6 +873,12 @@ def stage( ) except (OverflowError, RuntimeError, TypeError, ValueError): return False + if min(round_starts) < 0: + return False + # Grow the phase table while this cohort's round starts are still host + # integers: the gather clamps stale-capacity rows instead of faulting, + # so skipping this would silently mis-phase the cohort. + self.mean_phase_table.ensure(int(max(round_starts)) + 1) if not self._stage_page_tables_bulk( manager, request_ids, @@ -1109,6 +1124,7 @@ def __init__( # Geometric integration offsets (built lazily on first eviction so the # device matches the cache pool). self._offsets: Optional[torch.Tensor] = None + self._mean_phase_table: Optional["MeanPhaseTable"] = None # Request presence records successful initialization. The record also # owns the counters and physical length cleared at request finish. @@ -1952,6 +1968,16 @@ def _fixed_resources_for( first_pool = layout.layer_pools[layout.dense_layers[0]] if self._offsets is None: self._offsets = _build_geometric_offsets(_OFFSET_MAX_LENGTH, first_pool.device) + if self._mean_phase_table is None: + from .triattention_kernels import MeanPhaseTable + + self._mean_phase_table = MeanPhaseTable( + self._offsets, + self.calibration["omega"] + .to(device=first_pool.device, dtype=torch.float32) + .contiguous(), + initial_rows=seq_capacity, + ) q_real, q_imag, mlr_coef = self._local_score_calibration( layout.num_layers, layout.global_layers ) @@ -1970,6 +1996,7 @@ def _fixed_resources_for( freq_scale_sq=self._freq_scale_sq, offsets=self._offsets, omega=self.calibration["omega"], + mean_phase_table=self._mean_phase_table, page_table_keys=self._page_table_pool_keys(representatives, layout.global_layers), num_page_table_slots=layout.manager.num_pools, decode_width=decode_width, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 8860d175fbae..f5ea6019e985 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -22,7 +22,7 @@ from __future__ import annotations -from typing import List +from typing import List, Optional import torch import triton @@ -33,75 +33,130 @@ # --------------------------------------------------------------------------- # +# Positions past this row count are no longer exactly representable in fp32, +# so a larger table would silently degrade every downstream phase. +_MEAN_PHASE_MAX_ROWS = 1 << 24 + + @triton.jit -def _prepare_mean_phase_kernel( +def _gather_mean_phase_kernel( + table, round_starts, - offsets, - omega, - mean_cos, - mean_sin, + mean_phases, + table_rows, + table_plane_stride, + output_plane_stride, NUM_FREQS: tl.constexpr, - NUM_OFFSETS: tl.constexpr, F_BLOCK: tl.constexpr, ): - """Collapse all offset phases for one request into reusable frequency means.""" + """Copy one request's cos/sin table row into the stacked mean planes.""" request = tl.program_id(0) frequency = tl.arange(0, F_BLOCK) frequency_mask = frequency < NUM_FREQS - round_start = tl.load(round_starts + request) - angular_frequency = tl.load(omega + frequency, mask=frequency_mask, other=0.0) - cos_sum = tl.zeros((F_BLOCK,), tl.float32) - sin_sum = tl.zeros((F_BLOCK,), tl.float32) - for offset_index in tl.static_range(0, NUM_OFFSETS): - offset = tl.load(offsets + offset_index) - phase = (round_start + offset) * angular_frequency - cos_sum += tl.cos(phase) - sin_sum += tl.sin(phase) - output_offset = request * NUM_FREQS + frequency - scale = 1.0 / NUM_OFFSETS - tl.store(mean_cos + output_offset, cos_sum * scale, mask=frequency_mask) - tl.store(mean_sin + output_offset, sin_sum * scale, mask=frequency_mask) - - -def prepare_mean_phase( - round_starts: torch.Tensor, - offsets: torch.Tensor, - omega: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - request_count: int, -) -> None: - """Prepare mean score phases in one launch without intermediate tensors.""" - request_count = int(request_count) - if request_count <= 0 or request_count > round_starts.numel(): - raise ValueError("phase preparation request count is outside its fixed buffers") - num_freqs = int(omega.numel()) - num_offsets = int(offsets.numel()) - if ( - num_freqs <= 0 - or num_offsets <= 0 - or mean_cos.ndim != 2 - or mean_cos.shape[0] < request_count - or mean_cos.shape[1] != num_freqs - or mean_sin.shape != mean_cos.shape - or any( - tensor.device != round_starts.device for tensor in (offsets, omega, mean_cos, mean_sin) + row = tl.load(round_starts + request).to(tl.int64) + # Clamp instead of trusting the staged values: a stale-capacity row + # must degrade to a wrong phase, never to an out-of-bounds access. + row = tl.minimum(tl.maximum(row, 0), table_rows - 1) + source_offset = row * NUM_FREQS + frequency + mean_cos = tl.load(table + source_offset, mask=frequency_mask, other=0.0) + mean_sin = tl.load(table + table_plane_stride + source_offset, mask=frequency_mask, other=0.0) + output_offset = request.to(tl.int64) * NUM_FREQS + frequency + tl.store(mean_phases + output_offset, mean_cos, mask=frequency_mask) + tl.store(mean_phases + output_plane_stride + output_offset, mean_sin, mask=frequency_mask) + + +class MeanPhaseTable: + """RoPE-style position table of mean trig phases, gathered per round. + + The ``(2, rows, num_freqs)`` table stacks a cosine and a sine plane; + row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` over the + calibration offsets for every frequency, so refreshing a round's + stacked ``mean_cos``/``mean_sin`` planes is ONE vendored Triton gather + over the staged round starts instead of a per-round trig kernel. The + gather writes in place because the compiled CuTe score launch captured + the destination planes' device pointers. Eviction never runs under + CUDA graph capture, so the table itself may regrow; callers must + ``ensure`` capacity while the round starts are still host integers -- + the gather clamps a stale-capacity row to the last table row, which + cannot fault but yields that request a wrong phase. + """ + + def __init__(self, offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int) -> None: + if ( + offsets.numel() <= 0 + or omega.numel() <= 0 + or offsets.dtype != torch.float32 + or omega.dtype != torch.float32 + or offsets.device != omega.device + or omega.device.type != "cuda" + ): + raise ValueError("mean-phase tables require FP32 CUDA offsets and frequencies") + self.offsets = offsets.contiguous() + self.omega = omega.contiguous() + self._offset_values: List[float] = self.offsets.tolist() + self._table: Optional[torch.Tensor] = None + self._rows = 0 + self.ensure(max(int(initial_rows), 1)) + + @property + def rows(self) -> int: + return self._rows + + def ensure(self, rows: int) -> None: + """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" + 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 = 1 + while target < rows: + target *= 2 + target = min(max(target, 2 * self._rows), _MEAN_PHASE_MAX_ROWS) + positions = torch.arange(target, device=self.omega.device, dtype=torch.float32) + table = torch.zeros( + (2, target, self.omega.numel()), dtype=torch.float32, device=self.omega.device + ) + # Accumulate offset-by-offset in fp32, mirroring the retired + # per-round Triton kernel's summation order. + for offset in self._offset_values: + phase = torch.outer(positions + offset, self.omega) + table[0] += torch.cos(phase) + table[1] += torch.sin(phase) + self._table = table.mul_(1.0 / len(self._offset_values)) + self._rows = target + + def gather( + self, + round_starts: torch.Tensor, + mean_phases: torch.Tensor, + request_count: int, + ) -> None: + """Refresh the stacked cos/sin mean planes in place in one gather.""" + request_count = int(request_count) + if request_count <= 0 or request_count > round_starts.numel(): + raise ValueError("phase gather request count is outside its fixed buffers") + num_freqs = self.omega.numel() + if ( + mean_phases.shape != (2, request_count, num_freqs) + or not mean_phases.is_contiguous() + or mean_phases.dtype != torch.float32 + or round_starts.dtype != torch.int32 + or round_starts.device != self.omega.device + or mean_phases.device != self.omega.device + ): + raise ValueError("phase gather tensors do not share one valid FP32 geometry") + _gather_mean_phase_kernel[(request_count,)]( + self._table, + round_starts, + mean_phases, + self._rows, + self._rows * num_freqs, + request_count * num_freqs, + NUM_FREQS=num_freqs, + F_BLOCK=triton.next_power_of_2(num_freqs), + num_warps=1, ) - or round_starts.dtype != torch.int32 - or any(tensor.dtype != torch.float32 for tensor in (offsets, omega, mean_cos, mean_sin)) - ): - raise ValueError("phase preparation tensors do not share one valid FP32 geometry") - _prepare_mean_phase_kernel[(request_count,)]( - round_starts, - offsets, - omega, - mean_cos, - mean_sin, - NUM_FREQS=num_freqs, - NUM_OFFSETS=num_offsets, - F_BLOCK=triton.next_power_of_2(num_freqs), - num_warps=1, - ) class _FixedScoreGroup: @@ -752,78 +807,6 @@ def prepare_per_head_scores( # --------------------------------------------------------------------------- # -@triton.jit -def _pack_compaction_sources_kernel( - selected_indices, - valid_seq_lens, - dense_offsets, - dense_indices, - swa_offsets, - swa_indices, - DENSE_TOTAL: tl.constexpr, - SWA_TOTAL: tl.constexpr, - SELECTION_ROWS: tl.constexpr, - SELECTION_STRIDE: tl.constexpr, - KEEP_COUNT: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - SWA_WINDOW: tl.constexpr, - UNION: tl.constexpr, - PER_LAYER: tl.constexpr, - HAS_SWA: tl.constexpr, - BLOCK: tl.constexpr, -): - """Pack selected decode ordinals and protected tails for the C++ updater.""" - request = tl.program_id(0) - domain = tl.program_id(1) - move = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) - - dense_begin = tl.load(dense_offsets + request) - dense_end = tl.load(dense_offsets + request + 1) - dense_count = dense_end - dense_begin - seq_len = tl.load(valid_seq_lens + request) - - if UNION: - selection_domain = 0 - else: - selection_domain = domain - # Selection rows carry decode-only kept ordinals (already absolute), so - # rows are prompt-length independent and one cohort may mix prompt sizes. - selection_row = request * SELECTION_ROWS + selection_domain - selected = tl.load( - selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + move, - mask=move < KEEP_COUNT, - other=0, - ) - dense_source = tl.where(move < KEEP_COUNT, selected, seq_len + move - KEEP_COUNT) - dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move - tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) - - if HAS_SWA: - # Per-layer selection has one dense domain per (layer, head). SWA uses - # one shared source row per head, so only the first layer writes it. - if PER_LAYER: - write_swa = domain < NUM_KV_HEADS - else: - write_swa = move >= 0 - swa_begin = tl.load(swa_offsets + request) - swa_end = tl.load(swa_offsets + request + 1) - swa_count = swa_end - swa_begin - head = domain % NUM_KV_HEADS - swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move - swa_source = seq_len - SWA_WINDOW + move - tl.store( - swa_indices + swa_output, - swa_source, - mask=write_swa & (move < swa_count), - ) - - -# --------------------------------------------------------------------------- # -# Fused finalize: settle the top-k ties and pack the move indices in one # -# launch (fusion suggested by Fanrong Li, torch-graph review 2026-07-20). # -# --------------------------------------------------------------------------- # - - @triton.jit def _settle_ties_and_pack_compaction_sources_kernel( scores, @@ -848,6 +831,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( UNION: tl.constexpr, PER_LAYER: tl.constexpr, HAS_SWA: tl.constexpr, + HAS_SETTLE: tl.constexpr, HAS_PACK: tl.constexpr, BLOCK: tl.constexpr, ): @@ -864,9 +848,13 @@ def _settle_ties_and_pack_compaction_sources_kernel( same conditions as the standalone kernel. Union selection has one row per request feeding every KV head's packed row, so that single program writes all of them. ``HAS_PACK=False`` compiles the second half away, leaving - exactly the settle stage (its pre-fusion standalone copy lives in the - fused-kernel unit test as the bit-equality reference). Fusing the two launches was suggested by - Fanrong Li (torch-graph review 2026-07-20). + exactly the settle stage; ``HAS_SETTLE=False`` compiles the first half + away instead, packing pre-settled ordinals read from + ``output_indices`` -- the draft co-compaction flow, whose keep set is + the target's and needs no settling. The pre-fusion standalone copies + live in the fused-kernel unit test as the bit-equality references. + Fusing the launches was suggested by Fanrong Li (torch-graph review + 2026-07-20). """ request = tl.program_id(0) selection_domain = tl.program_id(1) @@ -874,70 +862,72 @@ def _settle_ties_and_pack_compaction_sources_kernel( row_scores = scores + row * WIDTH row_selected = provisional_indices + row * KEEP_COUNT row_output = output_indices + row * OUTPUT_WIDTH - # Scores are decode-relative; this row's pinned prompt length rebases the - # emitted ordinals to absolute positions (per row, so one launch may mix - # prompt lengths). - prompt_len = tl.load(prompt_offsets + row) - - 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, - ) - selected_score = tl.load( - row_scores + token_index, - mask=selected_mask, - other=float("inf"), - ).to(tl.float32) - threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) - - seq_len = tl.load(seq_lens + 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 < seq_len) - 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 < seq_len) - 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_len, - mask=selected, - ) - output_count += tl.sum(selected_i32) - ties_seen += tl.sum(tied_i32) + if HAS_SETTLE: + # Scores are decode-relative; this row's pinned prompt length rebases + # the emitted ordinals to absolute positions (per row, so one launch + # may mix prompt lengths). + prompt_len = tl.load(prompt_offsets + row) + + 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, + ) + selected_score = tl.load( + row_scores + token_index, + mask=selected_mask, + other=float("inf"), + ).to(tl.float32) + threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) + + seq_len = tl.load(seq_lens + 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 < seq_len) + 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 < seq_len) + 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_len, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) if HAS_PACK: - # The emission above scatters through other lanes of this program; - # make those global stores visible to every lane before the pack - # half reads the row back. - tl.debug_barrier() + if HAS_SETTLE: + # The emission above scatters through other lanes of this + # program; make those global stores visible to every lane + # before the pack half reads the row back. + tl.debug_barrier() dense_begin = tl.load(dense_offsets + request) dense_end = tl.load(dense_offsets + request + 1) dense_count = dense_end - dense_begin 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 index 0bbc6cfeb0b1..55aeabb4b5e1 100644 --- 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 @@ -28,11 +28,82 @@ _BatchedUnionKeepSetSelector, ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _pack_compaction_sources_kernel, _settle_ties_and_pack_compaction_sources_kernel, ) +@triton.jit +def _pack_compaction_sources_kernel( + selected_indices, + valid_seq_lens, + dense_offsets, + dense_indices, + swa_offsets, + swa_indices, + DENSE_TOTAL: tl.constexpr, + SWA_TOTAL: tl.constexpr, + SELECTION_ROWS: tl.constexpr, + SELECTION_STRIDE: tl.constexpr, + KEEP_COUNT: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + SWA_WINDOW: tl.constexpr, + UNION: tl.constexpr, + PER_LAYER: tl.constexpr, + HAS_SWA: tl.constexpr, + BLOCK: tl.constexpr, +): + """Pack selected decode ordinals and protected tails for the C++ updater.""" + request = tl.program_id(0) + domain = tl.program_id(1) + move = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + + dense_begin = tl.load(dense_offsets + request) + dense_end = tl.load(dense_offsets + request + 1) + dense_count = dense_end - dense_begin + seq_len = tl.load(valid_seq_lens + request) + + if UNION: + selection_domain = 0 + else: + selection_domain = domain + # Selection rows carry decode-only kept ordinals (already absolute), so + # rows are prompt-length independent and one cohort may mix prompt sizes. + selection_row = request * SELECTION_ROWS + selection_domain + selected = tl.load( + selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + move, + mask=move < KEEP_COUNT, + other=0, + ) + dense_source = tl.where(move < KEEP_COUNT, selected, seq_len + move - KEEP_COUNT) + dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) + + if HAS_SWA: + # Per-layer selection has one dense domain per (layer, head). SWA uses + # one shared source row per head, so only the first layer writes it. + if PER_LAYER: + write_swa = domain < NUM_KV_HEADS + else: + write_swa = move >= 0 + swa_begin = tl.load(swa_offsets + request) + swa_end = tl.load(swa_offsets + request + 1) + swa_count = swa_end - swa_begin + head = domain % NUM_KV_HEADS + swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move + swa_source = seq_len - SWA_WINDOW + move + tl.store( + swa_indices + swa_output, + swa_source, + mask=write_swa & (move < swa_count), + ) + + +# --------------------------------------------------------------------------- # +# Fused finalize: settle the top-k ties and pack the move indices in one # +# launch (fusion suggested by Fanrong Li, torch-graph review 2026-07-20). # +# --------------------------------------------------------------------------- # + + @triton.jit def _settle_ties_after_topk_kernel( scores, @@ -270,6 +341,7 @@ def test_fused_settle_pack_matches_two_kernel_sequence(eviction_mode, has_swa, w UNION=union, PER_LAYER=per_layer, HAS_SWA=has_swa, + HAS_SETTLE=True, HAS_PACK=True, BLOCK=_BLOCK, num_warps=_NUM_WARPS, @@ -338,6 +410,7 @@ def test_fused_kernel_without_pack_matches_standalone_settle(): UNION=False, PER_LAYER=False, HAS_SWA=False, + HAS_SETTLE=True, HAS_PACK=False, BLOCK=_BLOCK, num_warps=_NUM_WARPS, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 0d874e5505e4..d2b2799b8de3 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1232,10 +1232,10 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re assert staging.bucket_seq_len == seq_len assert staging.page_table_token_capacity == page_table_token_capacity assert staging.page_count == page_count - assert staging.offsets.dtype == torch.float32 - assert staging.offsets.is_contiguous() - assert staging.omega.dtype == torch.float32 - assert staging.omega.is_contiguous() + assert staging.mean_phase_table.offsets.dtype == torch.float32 + assert staging.mean_phase_table.offsets.is_contiguous() + assert staging.mean_phase_table.omega.dtype == torch.float32 + assert staging.mean_phase_table.omega.is_contiguous() fused = staging.fused_group for calibration in (*fused.pointer_middle[2:], *fused.pointer_tail): assert calibration.dtype == torch.float32 @@ -1531,16 +1531,14 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) - mean_cos = torch.empty(request_count, num_freqs, dtype=torch.float32, device=device) - mean_sin = torch.empty_like(mean_cos) - triattention_kernels.prepare_mean_phase( - round_device, - offsets, - omega, - mean_cos, - mean_sin, - request_count, - ) + mean_phases = torch.empty(2, request_count, num_freqs, dtype=torch.float32, device=device) + mean_cos = mean_phases[0] + mean_sin = mean_phases[1] + # Rows cover this launch and the +17 round-start advance below. + mean_phase_table = triattention_kernels.MeanPhaseTable( + offsets, omega, initial_rows=int(round_device.max()) + 18 + ) + mean_phase_table.gather(round_device, mean_phases, request_count) score_sentinel = -12345.0 group.output.fill_(score_sentinel) checked = group.launch( @@ -1558,10 +1556,10 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun staging.round_starts_device = round_device staging.valid_seq_lens_device = valid_seq_lens staging.token_starts_device = token_starts + staging.mean_phases = mean_phases staging.mean_cos = mean_cos staging.mean_sin = mean_sin - staging.offsets = offsets - staging.omega = omega + staging.mean_phase_table = mean_phase_table staging.stream = None staging._score_valid_widths = None staging._score_launcher_bound = False @@ -1604,14 +1602,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun device=device, ) ) - triattention_kernels.prepare_mean_phase( - round_device, - offsets, - omega, - mean_cos, - mean_sin, - request_count, - ) + mean_phase_table.gather(round_device, mean_phases, request_count) expected_second_widths = valid_seq_lens - prompt_len group.output.fill_(score_sentinel) valid_widths.fill_(-1) @@ -1637,6 +1628,34 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun class TestKernelMaskedSwa: + @pytest.mark.skipif(not torch.cuda.is_available(), reason="MeanPhaseTable is CUDA-only") + def test_mean_phase_table_regrowth_and_clamp(self): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + MeanPhaseTable, + ) + + device = torch.device("cuda", torch.cuda.current_device()) + omega = torch.rand(7, device=device) * 0.05 + offsets = torch.tensor([1.0, 2.0, 4.0], device=device) + rounds = torch.tensor([0, 3, 500, 65], dtype=torch.int32, device=device) + gathered = torch.empty(2, 4, 7, dtype=torch.float32, device=device) + small = MeanPhaseTable(offsets, omega, initial_rows=8) + assert small.rows == 8 + small.ensure(501) + assert small.rows == 512 + small.gather(rounds, gathered, 4) + # A regrown table is rebuilt from positions, so it must match a + # table built at the final size bit-for-bit. + fresh = MeanPhaseTable(offsets, omega, initial_rows=512) + expected = torch.empty_like(gathered) + fresh.gather(rounds, expected, 4) + torch.testing.assert_close(gathered, expected, rtol=0, atol=0) + # A stale-capacity row clamps to the last table row instead of + # reading out of bounds. + stale = torch.tensor([100_000, 511, 0, 1], dtype=torch.int32, device=device) + small.gather(stale, gathered, 4) + torch.testing.assert_close(gathered[:, 0], gathered[:, 1], rtol=0, atol=0) + def test_layer_partition_uses_local_model_config(self): mgr = _make_triattention() mgr.model_path = "/models/gpt-oss" From 1b3cd4512dccb84b515dd30a2de53a927cb0b89c Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 01:18:56 -0700 Subject: [PATCH 067/178] [None][fix] Clamp gathered phase rows and keep table construction device-agnostic The mean-phase gather kernel now clamps each staged round start into the table so padded or stale slots read row 0 or the last row instead of faulting; live cohorts are still host-validated at staging. Building the table is plain torch and no longer demands CUDA (CPU-mocked staging tests construct it), while the Triton gather remains CUDA-only. Also sort the staging import and reference the table type under TYPE_CHECKING. Signed-off-by: tianruih --- .../triattention/triattention.py | 21 ++-- .../triattention/triattention_kernels.py | 98 ++++++++++--------- .../test_triattention_pipeline.py | 38 +------ 3 files changed, 66 insertions(+), 91 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index f108ad5f4d24..5942250d41d8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -69,12 +69,11 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig if TYPE_CHECKING: - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - MeanPhaseTable, - ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests + from .triattention_kernels import MeanPhaseTable + # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) @@ -723,13 +722,10 @@ def __init__( self.dense_move_offsets = self.request_metadata_device[3] self.swa_move_offsets = self.request_metadata_device[4] self.draft_move_offsets = self.request_metadata_device[5] - # Stacked cos/sin planes: one index_select refreshes both, and each - # plane stays a contiguous tensor for the CuTe launch to capture. - self.mean_phases = torch.empty( - (2, max_requests, num_freqs), dtype=torch.float32, device=self.device + self.mean_cos = torch.empty( + (max_requests, num_freqs), dtype=torch.float32, device=self.device ) - self.mean_cos = self.mean_phases[0] - self.mean_sin = self.mean_phases[1] + self.mean_sin = torch.empty_like(self.mean_cos) # The phase table depends only on the shared calibration, so the # manager passes one instance to every staging bucket; standalone # construction (tests) builds a private one. @@ -806,7 +802,8 @@ def launch_prepared_score(self) -> torch.Tensor: # from this round's staged round starts before it runs. self.mean_phase_table.gather( self.round_starts_device, - self.mean_phases, + self.mean_cos, + self.mean_sin, self.max_requests, ) return self.fused_group.launch( @@ -876,8 +873,8 @@ def stage( if min(round_starts) < 0: return False # Grow the phase table while this cohort's round starts are still host - # integers: the gather clamps stale-capacity rows instead of faulting, - # so skipping this would silently mis-phase the cohort. + # integers: a stale-capacity gather is an out-of-bounds index_select + # on the device. self.mean_phase_table.ensure(int(max(round_starts)) + 1) if not self._stage_page_tables_bulk( manager, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index f5ea6019e985..662d43a5c6e8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -40,45 +40,45 @@ @triton.jit def _gather_mean_phase_kernel( - table, round_starts, - mean_phases, + table_cos, + table_sin, + mean_cos, + mean_sin, table_rows, - table_plane_stride, - output_plane_stride, NUM_FREQS: tl.constexpr, F_BLOCK: tl.constexpr, ): - """Copy one request's cos/sin table row into the stacked mean planes.""" + """Copy each request's precomputed phase-table row into the fixed buffers.""" request = tl.program_id(0) frequency = tl.arange(0, F_BLOCK) frequency_mask = frequency < NUM_FREQS - row = tl.load(round_starts + request).to(tl.int64) - # Clamp instead of trusting the staged values: a stale-capacity row - # must degrade to a wrong phase, never to an out-of-bounds access. - row = tl.minimum(tl.maximum(row, 0), table_rows - 1) - source_offset = row * NUM_FREQS + frequency - mean_cos = tl.load(table + source_offset, mask=frequency_mask, other=0.0) - mean_sin = tl.load(table + table_plane_stride + source_offset, mask=frequency_mask, other=0.0) - output_offset = request.to(tl.int64) * NUM_FREQS + frequency - tl.store(mean_phases + output_offset, mean_cos, mask=frequency_mask) - tl.store(mean_phases + output_plane_stride + output_offset, mean_sin, mask=frequency_mask) + table_row = tl.load(round_starts + request).to(tl.int64) + # Clamp stale or padded round starts into the table instead of faulting; + # staged cohorts are host-validated, so live rows are never clamped. + table_row = tl.minimum(tl.maximum(table_row, 0), table_rows - 1) + source_offset = table_row * NUM_FREQS + frequency + output_offset = request * NUM_FREQS + frequency + row_cos = tl.load(table_cos + source_offset, mask=frequency_mask, other=0.0) + row_sin = tl.load(table_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) class MeanPhaseTable: """RoPE-style position table of mean trig phases, gathered per round. - The ``(2, rows, num_freqs)`` table stacks a cosine and a sine plane; - row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` over the + Row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` over the calibration offsets for every frequency, so refreshing a round's - stacked ``mean_cos``/``mean_sin`` planes is ONE vendored Triton gather - over the staged round starts instead of a per-round trig kernel. The - gather writes in place because the compiled CuTe score launch captured - the destination planes' device pointers. Eviction never runs under - CUDA graph capture, so the table itself may regrow; callers must - ``ensure`` capacity while the round starts are still host integers -- - the gather clamps a stale-capacity row to the last table row, which - cannot fault but yields that request a wrong phase. + ``mean_cos``/``mean_sin`` is one pure-gather launch over the staged + round starts instead of a per-round trig kernel. The gather writes in + place because the compiled CuTe score launch captured the destination + buffers' device pointers. Eviction never runs under CUDA graph + capture, so the table itself may regrow; callers must ``ensure`` + capacity while the round starts are still host integers (the gather + clamps stale rows into the table rather than faulting). Building the + table is plain torch and works on any device; gathering launches the + Triton kernel and is CUDA-only. """ def __init__(self, offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int) -> None: @@ -88,13 +88,13 @@ def __init__(self, offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int or offsets.dtype != torch.float32 or omega.dtype != torch.float32 or offsets.device != omega.device - or omega.device.type != "cuda" ): - raise ValueError("mean-phase tables require FP32 CUDA offsets and frequencies") + raise ValueError("mean-phase tables require same-device FP32 offsets and frequencies") self.offsets = offsets.contiguous() self.omega = omega.contiguous() self._offset_values: List[float] = self.offsets.tolist() - self._table: Optional[torch.Tensor] = None + self._cos: Optional[torch.Tensor] = None + self._sin: Optional[torch.Tensor] = None self._rows = 0 self.ensure(max(int(initial_rows), 1)) @@ -114,45 +114,53 @@ def ensure(self, rows: int) -> None: target *= 2 target = min(max(target, 2 * self._rows), _MEAN_PHASE_MAX_ROWS) positions = torch.arange(target, device=self.omega.device, dtype=torch.float32) - table = torch.zeros( - (2, target, self.omega.numel()), dtype=torch.float32, device=self.omega.device + cos_table = torch.zeros( + (target, self.omega.numel()), dtype=torch.float32, device=self.omega.device ) + sin_table = torch.zeros_like(cos_table) # Accumulate offset-by-offset in fp32, mirroring the retired # per-round Triton kernel's summation order. for offset in self._offset_values: phase = torch.outer(positions + offset, self.omega) - table[0] += torch.cos(phase) - table[1] += torch.sin(phase) - self._table = table.mul_(1.0 / len(self._offset_values)) + cos_table += torch.cos(phase) + sin_table += torch.sin(phase) + scale = 1.0 / len(self._offset_values) + self._cos = cos_table.mul_(scale) + self._sin = sin_table.mul_(scale) self._rows = target def gather( self, round_starts: torch.Tensor, - mean_phases: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, request_count: int, ) -> None: - """Refresh the stacked cos/sin mean planes in place in one gather.""" + """Refresh the fixed mean buffers in place from staged round starts.""" request_count = int(request_count) if request_count <= 0 or request_count > round_starts.numel(): raise ValueError("phase gather request count is outside its fixed buffers") num_freqs = self.omega.numel() if ( - mean_phases.shape != (2, request_count, num_freqs) - or not mean_phases.is_contiguous() - or mean_phases.dtype != torch.float32 + mean_cos.ndim != 2 + or mean_cos.shape[0] < request_count + or mean_cos.shape[1] != num_freqs + or mean_sin.shape != mean_cos.shape or round_starts.dtype != torch.int32 - or round_starts.device != self.omega.device - or mean_phases.device != self.omega.device + or self.omega.device.type != "cuda" + or any( + tensor.device != self.omega.device for tensor in (round_starts, mean_cos, mean_sin) + ) + or any(tensor.dtype != torch.float32 for tensor in (mean_cos, mean_sin)) ): - raise ValueError("phase gather tensors do not share one valid FP32 geometry") + raise ValueError("phase gather tensors do not share one valid FP32 CUDA geometry") _gather_mean_phase_kernel[(request_count,)]( - self._table, round_starts, - mean_phases, + self._cos, + self._sin, + mean_cos, + mean_sin, self._rows, - self._rows * num_freqs, - request_count * num_freqs, NUM_FREQS=num_freqs, F_BLOCK=triton.next_power_of_2(num_freqs), num_warps=1, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index d2b2799b8de3..8e9f983eb198 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -1531,14 +1531,13 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) - mean_phases = torch.empty(2, request_count, num_freqs, dtype=torch.float32, device=device) - mean_cos = mean_phases[0] - mean_sin = mean_phases[1] + mean_cos = torch.empty(request_count, num_freqs, dtype=torch.float32, device=device) + mean_sin = torch.empty_like(mean_cos) # Rows cover this launch and the +17 round-start advance below. mean_phase_table = triattention_kernels.MeanPhaseTable( offsets, omega, initial_rows=int(round_device.max()) + 18 ) - mean_phase_table.gather(round_device, mean_phases, request_count) + mean_phase_table.gather(round_device, mean_cos, mean_sin, request_count) score_sentinel = -12345.0 group.output.fill_(score_sentinel) checked = group.launch( @@ -1556,7 +1555,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun staging.round_starts_device = round_device staging.valid_seq_lens_device = valid_seq_lens staging.token_starts_device = token_starts - staging.mean_phases = mean_phases staging.mean_cos = mean_cos staging.mean_sin = mean_sin staging.mean_phase_table = mean_phase_table @@ -1602,7 +1600,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun device=device, ) ) - mean_phase_table.gather(round_device, mean_phases, request_count) + mean_phase_table.gather(round_device, mean_cos, mean_sin, request_count) expected_second_widths = valid_seq_lens - prompt_len group.output.fill_(score_sentinel) valid_widths.fill_(-1) @@ -1628,34 +1626,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun class TestKernelMaskedSwa: - @pytest.mark.skipif(not torch.cuda.is_available(), reason="MeanPhaseTable is CUDA-only") - def test_mean_phase_table_regrowth_and_clamp(self): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - MeanPhaseTable, - ) - - device = torch.device("cuda", torch.cuda.current_device()) - omega = torch.rand(7, device=device) * 0.05 - offsets = torch.tensor([1.0, 2.0, 4.0], device=device) - rounds = torch.tensor([0, 3, 500, 65], dtype=torch.int32, device=device) - gathered = torch.empty(2, 4, 7, dtype=torch.float32, device=device) - small = MeanPhaseTable(offsets, omega, initial_rows=8) - assert small.rows == 8 - small.ensure(501) - assert small.rows == 512 - small.gather(rounds, gathered, 4) - # A regrown table is rebuilt from positions, so it must match a - # table built at the final size bit-for-bit. - fresh = MeanPhaseTable(offsets, omega, initial_rows=512) - expected = torch.empty_like(gathered) - fresh.gather(rounds, expected, 4) - torch.testing.assert_close(gathered, expected, rtol=0, atol=0) - # A stale-capacity row clamps to the last table row instead of - # reading out of bounds. - stale = torch.tensor([100_000, 511, 0, 1], dtype=torch.int32, device=device) - small.gather(stale, gathered, 4) - torch.testing.assert_close(gathered[:, 0], gathered[:, 1], rtol=0, atol=0) - def test_layer_partition_uses_local_model_config(self): mgr = _make_triattention() mgr.model_path = "/models/gpt-oss" From 2b3e973eac222280e673269d98a97b2d622aa82d Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 03:22:30 -0700 Subject: [PATCH 068/178] [None][chore] Drop a stale kernel-name reference from the fused pack docstring Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 662d43a5c6e8..f67fa042bb36 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -849,11 +849,11 @@ def _settle_ties_and_pack_compaction_sources_kernel( provisional top-k: recover the top-k threshold from the provisional selection, count the strictly greater scores, then emit the kept ordinals in increasing order, rebased by the row's pinned - prompt length. With ``HAS_PACK`` the same program then continues with the - ``_pack_compaction_sources_kernel`` work for the packed rows this - selection row feeds: the kept ordinals it just wrote, followed by the - request's protected tail, plus the SWA rows (latest window) under the - same conditions as the standalone kernel. Union selection has one row per + prompt length. With ``HAS_PACK`` the same program then packs the move + sources for the packed rows this selection row feeds: the kept ordinals + it just wrote, followed by the request's protected tail, plus the SWA + rows (latest window) under the + same conditions as the retired standalone kernel. Union selection has one row per request feeding every KV head's packed row, so that single program writes all of them. ``HAS_PACK=False`` compiles the second half away, leaving exactly the settle stage; ``HAS_SETTLE=False`` compiles the first half From 294176f5a602a6a6ffd8e04d7736942063e5ccf3 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 04:10:24 -0700 Subject: [PATCH 069/178] [None][feat] Integrate the fused score+stats+union pipeline (opt-in, off) Vendor Fanrong Li's two-kernel scheme: the SM100 CuTe score kernel gains per-page-shard partial row statistics, and a second CuTe kernel merges the partials, normalizes rows, and takes the cross-row union maximum directly into the selector's combined buffer, replacing the split row-stats/union/normalize launches when it engages. Wiring: the pipeline is opt-in via TRTLLM_TRIATTENTION_CUTE_UNION_FUSION and engages only for union mode with normalized scores on the fused contract (128-token pages, GQA group 8, 32 frequencies) when the cohort's pinned prompt lengths agree (the score start is a compile-time constant); anything else falls back to the split path. Fixes applied on top of the original patch: the head-plane stride reaches the kernel as Int64 (request*layer segments of seq_len columns exceed 2^31 at large batch), and the fast-sqrt spelling is probed for cutlass 4.5. KNOWN DEFECT (why the default stays off): under this environment's cutlass DSL the fused kernel corrupts every page after the first (page 0 matches the torch oracle to 1e-6; pages 1+ diverge by up to ~0.7 in both the plain and stats variants), which fingerprints the page-id prefetch pipeline. The equivalence test is committed as a strict xfail repro while the kernel/DSL mismatch is resolved with the author. Signed-off-by: tianruih --- .../triattention/triattention.py | 66 +- .../triattention_cute_score_fused.py | 1847 +++++++++++++++++ .../triattention_cute_selection.py | 516 +++++ .../triattention/triattention_kernels.py | 119 ++ .../test_triattention_cute_union_fusion.py | 231 +++ 5 files changed, 2777 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score_fused.py create mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py create mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_cute_union_fusion.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 5942250d41d8..b912f3576b62 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -387,6 +387,14 @@ def select_prepared_requests(self) -> None: ) self._select_top_tokens() + def select_prepared_union_scores(self) -> None: + """Select from normalized union rows already written into ``combined``. + + The fused score+stats+union pipeline produces the per-request union + rows on the device, so only the top-k settle-and-pack launch remains. + """ + self._select_top_tokens() + class _BatchedPerHeadKeepSetSelector(_BatchedKeepSetSelectorBase): """Fixed ``[request, ...]`` selector for both per-head modes.""" @@ -770,6 +778,7 @@ def __init__( self.stream = None self._score_valid_widths: Optional[torch.Tensor] = None self._score_launcher_bound = False + self.staged_uniform_token_start: Optional[int] = None def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> None: """Bind the per-row score widths for these buffers (mean-only).""" @@ -783,6 +792,42 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> N self._score_valid_widths = valid_widths self._score_launcher_bound = True + def launch_prepared_union_fusion(self, union_out: torch.Tensor) -> bool: + """Try the fused score+stats+union pipeline over these buffers. + + Returns False without side effects on the selection buffers when the + fused path cannot serve this cohort; the caller then runs the split + score and selection launches instead. + """ + if not self._score_launcher_bound: + raise RuntimeError("TriAttention score launcher is not bound") + stream = torch.cuda.current_stream(self.device) + if self.stream is None: + self.stream = stream + elif (stream.device, stream.cuda_stream) != ( + self.stream.device, + self.stream.cuda_stream, + ): + raise _FixedScoreStreamMismatch( + "TriAttention score launches must stay on the staging CUDA stream" + ) + self.mean_phase_table.gather( + self.round_starts_device, + self.mean_cos, + self.mean_sin, + self.max_requests, + ) + return self.fused_group.launch_cute_union_fusion( + self.max_requests, + self.valid_seq_lens_device, + self._score_valid_widths, + self.token_starts_device, + getattr(self, "staged_uniform_token_start", None), + self.mean_cos, + self.mean_sin, + union_out, + ) + def launch_prepared_score(self) -> torch.Tensor: """Gather the phase means and launch the score kernel over these buffers.""" if not self._score_launcher_bound: @@ -876,6 +921,13 @@ def stage( # integers: a stale-capacity gather is an out-of-bounds index_select # on the device. self.mean_phase_table.ensure(int(max(round_starts)) + 1) + # The fused score+stats+union kernel bakes one global score start at + # compile time, so it only serves cohorts whose pinned prompt lengths + # agree; padded rows are inert (zero valid length) regardless. + first_start = token_starts[0] + self.staged_uniform_token_start = ( + int(first_start) if all(start == first_start for start in token_starts) else None + ) if not self._stage_page_tables_bulk( manager, request_ids, @@ -2269,9 +2321,19 @@ def _evict_requests( try: with nvtx_range("triattention.score", color="blue"): - per_head = score_staging.launch_prepared_score() + # The fused pipeline covers score, row stats, normalization, + # and the cross-row union maximum in two kernels; when it + # declines this cohort the split launches run instead. + fused_union = ( + self.normalize_scores + and isinstance(keep_set_selector, _BatchedUnionKeepSetSelector) + and score_staging.launch_prepared_union_fusion(keep_set_selector.combined) + ) + per_head = None if fused_union else score_staging.launch_prepared_score() with nvtx_range("triattention.select", color="yellow"): - if isinstance(keep_set_selector, _BatchedUnionKeepSetSelector): + if fused_union: + keep_set_selector.select_prepared_union_scores() + elif isinstance(keep_set_selector, _BatchedUnionKeepSetSelector): keep_set_selector.select_prepared_requests() else: keep_set_selector.select_requests( 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..79137054befb --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score_fused.py @@ -0,0 +1,1847 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SM100 CuTe-DSL scorer for the TriAttention mean-score path. + +This is the production specialization of the final workbench kernel. It +uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, and +producer-only page-ID lookahead. The public integration keeps a Triton +fallback for every geometry outside the exact contract validated here. +""" + +from __future__ import annotations + +import threading + +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.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import from_dlpack + + +def _cute_sqrt_keyword_mode() -> str: + """Probe which fast-sqrt spelling this CuTe DSL's ``cute.math.sqrt`` takes. + + The approximate-sqrt control was renamed across DSL releases: some expose + ``approx``/``ftz`` keywords, cutlass 4.5 exposes a single ``fastmath`` + flag, and older releases expose a plain one-argument ``sqrt``. Passing an + unknown keyword raises TypeError at trace time (inside ``cute.compile``, + where it cannot be caught), so the capability is probed once at import + time via signature inspection and folded into a trace-time constant. + """ + import inspect + + try: + parameters = inspect.signature(cute.math.sqrt).parameters + except (TypeError, ValueError): + return "plain" + if "approx" in parameters and "ftz" in parameters: + return "approx_ftz" + if "fastmath" in parameters: + return "fastmath" + return "plain" + + +_CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() + +CTA_M = 128 +K = 96 +N = 8 +NUM_FREQS = 32 +THREADS = 256 +EPILOGUE_THREADS = 128 +RAW_PAGE_BUFFERS = 2 + +PAGE_TOKENS = 128 +RAW_K_HALF_ELEMENTS = CTA_M * 2 * NUM_FREQS * RAW_PAGE_BUFFERS +RAW_K_VECTOR_ELEMENTS = 8 +RAW_K_SPLIT_PHASE_ELEMENTS = CTA_M * NUM_FREQS +RAW_K_SPLIT_TMA_COPY_BYTES = RAW_K_SPLIT_PHASE_ELEMENTS * (cutlass.BFloat16.width // 8) +TMA_DESCRIPTOR_QWORDS = 16 +_SUPPORTED_PAGE_SHARDS = (2, 3) + + +class _TriScoreEpilogue: + """Minimal TMEM-to-global epilogue for the score specialization.""" + + def __init__(self) -> None: + self.acc_dtype = cutlass.Float32 + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + accumulator: cute.Tensor, + output: cute.Tensor, + epilogue_tile: cute.Tile, + use_2cta_instrs: bool, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + copy_atom = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epilogue_tile, + use_2cta_instrs, + ) + accumulator_epilogue = cute.flat_divide( + accumulator[((None, None), 0, 0)], + epilogue_tile, + ) + tiled_copy = tcgen05.make_tmem_copy( + copy_atom, + accumulator_epilogue[(None, None, 0, 0)], + ) + thread_copy = tiled_copy.get_slice(tidx) + thread_accumulator = thread_copy.partition_S(accumulator_epilogue) + output_epilogue = cute.flat_divide( + output[((None, None), 0, 0, None, None, None)], + epilogue_tile, + ) + thread_output = thread_copy.partition_D(output_epilogue) + register_accumulator = cute.make_rmem_tensor( + thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.acc_dtype, + ) + return tiled_copy, thread_accumulator, register_accumulator + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tiled_copy: cute.TiledCopy, + output: cute.Tensor, + epilogue_tile: cute.Tile, + _unused_smem: cute.Tensor, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + output_epilogue = cute.flat_divide( + output[((None, None), 0, 0, None, None, None)], + epilogue_tile, + ) + thread_copy = tiled_copy.get_slice(tidx) + thread_output = thread_copy.partition_D(output_epilogue) + register_output = cute.make_rmem_tensor( + thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.c_dtype, + ) + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return copy_atom, register_output, thread_output + + +class _TriAttentionScoreKernel(_TriScoreEpilogue): + """Assign one CTA to each segment/KV-head task and retain W across pages.""" + + def __init__( + self, + *, + num_layers: int, + seq_len: int, + score_start: int, + num_q_heads: int, + num_kv_heads: int, + num_freqs: int, + tokens_per_block: int, + pool_shape: tuple[int, int, int, int, int], + pool_strides: tuple[int, int, int, int, int], + pool_dtype: type[cutlass.Numeric], + page_shards: int, + write_partial_stats: bool = False, + ) -> None: + """Build the single validated production specialization.""" + super().__init__() + if pool_dtype is not cutlass.BFloat16: + raise ValueError("TriAttention CuTe score requires BF16 K pages") + if num_freqs != NUM_FREQS: + raise ValueError("TriAttention CuTe score requires 32 frequencies") + if tokens_per_block != PAGE_TOKENS: + raise ValueError("TriAttention CuTe score requires 128-token pages") + if num_q_heads % num_kv_heads or num_q_heads // num_kv_heads != N: + raise ValueError("TriAttention CuTe score requires GQA group 8") + if not 0 <= score_start < seq_len: + raise ValueError("TriAttention CuTe score_start is out of range") + if page_shards not in _SUPPORTED_PAGE_SHARDS: + raise ValueError("TriAttention CuTe score has unsupported page shards") + + self.score_start = score_start + self.seq_len = seq_len + 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.max_pages = (seq_len + PAGE_TOKENS - 1) // PAGE_TOKENS + self.halves_per_page = PAGE_TOKENS // CTA_M + + # Measured final choices that still shape layouts or generated code. + self.prefetch_depth = 4 + self.sqrt_mode = "approx" + self.k_staging_mode = "half_page_tma" + self.use_tma = True + self.cpasync_schedule = "sync_each_half" + self.split_raw_tma = True + self.raw_tma_feature_extent = NUM_FREQS + self.raw_tma_copy_bytes = RAW_K_SPLIT_TMA_COPY_BYTES + self.raw_tma_pipeline_stages = 2 * RAW_PAGE_BUFFERS if write_partial_stats else 1 + self.accumulator_pipeline_stages = 1 + self.umma_accumulator_partitions = 1 + self.raw_cpasync_direct_a = True + self.weight_builder_mode = "coefficient_scalar_bf16_two_term" + self.main_operand_mode = "bf16_raw_three_term_weight" + self.numerical_policy = "three_term" + self.magnitude_residual_mode = "fp16_mma_two_term_single_commit" + self.fp16_magnitude_two_term = True + self.magnitude_sqrt_ftz = True + self.producer_page_id_prefetch = True + self.producer_warp_id = 0 + self.physical_threads = THREADS + self.shared_a_raw_alias = False + self.compact_token_loop = True + + self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape + if pool_kv_heads != num_kv_heads or pool_tokens != PAGE_TOKENS or 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_slot, self.s_dim = pool_strides + if self.s_slot != 2 * NUM_FREQS or self.s_dim != 1: + raise ValueError("K pages must be contiguous [128, 64]") + 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 __call__( + self, + page_ids: cute.Tensor, + seg_page_off: cute.Tensor, + seg_req_id: cute.Tensor, + seg_layer_id: cute.Tensor, + seg_seq_len: cute.Tensor, + seg_out_offset: 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, + pool_template: 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, N, K) + self.cta_tile_shape_mnk = self.mma_tiler + self.epi_tile = (CTA_M, N) + + 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], + ) + main_a_shape = ( + (CTA_M, N, NUM_FREQS) + if self.main_operand_mode == "bf16_raw_three_term_weight" + else self.mma_tiler + ) + a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, main_a_shape, cutlass.Float32, 1) + raw_bf16_a_smem_layout = sm100_utils.make_smem_layout_a( + raw_bf16_tiled_mma, + (CTA_M, N, 2 * NUM_FREQS), + cutlass.BFloat16, + 1, + ) + raw_bf16_split_a_smem_layout = sm100_utils.make_smem_layout_a( + raw_bf16_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.BFloat16, + 2 * RAW_PAGE_BUFFERS, + ) + # The full transport uses one K_SW128 8-KiB tile. The split transport + # packs two K_SW64 4-KiB stages into that same allocation, one each for + # real and imaginary data; the compile-time schedule selects the view. + raw_bf16_direct_a_smem_layout = ( + raw_bf16_split_a_smem_layout if self.split_raw_tma else raw_bf16_a_smem_layout + ) + raw_tma_smem_layout = cute.make_composed_layout( + raw_bf16_direct_a_smem_layout.inner, + 0, + cute.make_layout( + (self.raw_tma_feature_extent, CTA_M), + stride=(1, self.raw_tma_feature_extent), + ), + ) + raw_tma_source_layout = cute.make_layout( + ( + 2 * NUM_FREQS, + PAGE_TOKENS, + (self.num_kv_heads, self.num_physical_pages), + ), + stride=( + self.s_dim, + self.s_slot, + (self.s_kv_head, self.s_page), + ), + ) + raw_tma_source = cute.make_tensor( + pool_template.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, CTA_M), + ) + raw_bf16_b_smem_layout = sm100_utils.make_smem_layout_b( + raw_bf16_tiled_mma, + (CTA_M, N, 2 * NUM_FREQS), + cutlass.BFloat16, + 1, + ) + # The magnitude residual has only K=32. A separate compact descriptor + # lets the producer issue its four UMMA steps before the first commit, + # rather than waiting for and overwriting the K=96 main A tile. + magnitude_lo_smem_layout = sm100_utils.make_smem_layout_a( + tiled_mma, (CTA_M, N, NUM_FREQS), cutlass.Float32, 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_lo_fp16_smem_layout = sm100_utils.make_smem_layout_a( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + magnitude_hi_fp16_smem_layout = sm100_utils.make_smem_layout_b( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + magnitude_fp16_b_smem_layout = sm100_utils.make_smem_layout_b( + magnitude_lo_tiled_mma, + (CTA_M, N, NUM_FREQS), + cutlass.Float16, + 1, + ) + main_b_shape = ( + (CTA_M, N, NUM_FREQS) + if self.main_operand_mode == "bf16_raw_three_term_weight" + else self.mma_tiler + ) + b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, main_b_shape, cutlass.Float32, 1) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # Keep an explicit stage mode even for the one-stage control. The + # singleton mode folds away in codegen and lets both specializations + # share the same producer/consumer slicing protocol. + self.num_accumulator_slots = ( + self.accumulator_pipeline_stages * self.umma_accumulator_partitions + ) + 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) + + b_hi_elements = cute.cosize(b_smem_layout.outer) * int(not self.fp16_magnitude_two_term) + b_lo_elements = cute.cosize(b_smem_layout.outer) * int( + self.numerical_policy == "three_term" and not self.fp16_magnitude_two_term + ) + magnitude_lo_elements = cute.cosize(magnitude_lo_fp16_smem_layout.outer) * int( + self.numerical_policy == "three_term" + and self.magnitude_residual_mode in ("fp16_smem", "fp16_mma_single_commit") + ) + magnitude_lo_fp32_elements = cute.cosize(magnitude_lo_smem_layout.outer) * int( + self.numerical_policy == "three_term" + and self.magnitude_residual_mode == "fp32_smem_single_commit" + ) + magnitude_hi_fp16_elements = cute.cosize(magnitude_hi_fp16_smem_layout.outer) * int( + self.numerical_policy == "three_term" + and self.magnitude_residual_mode == "fp16_mma_single_commit" + ) + a_elements = cute.cosize(a_smem_layout.outer) * int( + not self.shared_a_raw_alias and not self.fp16_magnitude_two_term + ) + raw_k_elements = RAW_K_HALF_ELEMENTS * int( + self.k_staging_mode in ("half_page_cpasync", "half_page_tma") + and not self.shared_a_raw_alias + ) + alias_a_elements = cute.cosize(a_smem_layout.outer) * int(self.shared_a_raw_alias) + alias_raw_k_elements = RAW_K_HALF_ELEMENTS * int(self.shared_a_raw_alias) + raw_bf16_a_elements = cute.cosize(raw_bf16_a_smem_layout.outer) * int( + self.main_operand_mode == "bf16_raw_three_term_weight" + and ( + not self.raw_cpasync_direct_a + or self.cpasync_schedule not in ("sync_each_half", "intra_half_overlap") + ) + ) + raw_bf16_b_elements = cute.cosize(raw_bf16_b_smem_layout.outer) * int( + self.main_operand_mode == "bf16_raw_three_term_weight" + ) + raw_bf16_b2_elements = raw_bf16_b_elements * int( + self.weight_builder_mode != "coefficient_scalar_bf16_two_term" + ) + magnitude_fp16_a_elements = cute.cosize(magnitude_fp16_a_smem_layout.outer) * int( + self.fp16_magnitude_two_term + ) + magnitude_fp16_b_elements = cute.cosize(magnitude_fp16_b_smem_layout.outer) * int( + self.fp16_magnitude_two_term + ) + stats_scratch_elements = 72 * int(self.write_partial_stats) + + @cute.union + class SharedARawAlias: + # The two descriptors are byte-identical in size (8 KiB) and have + # disjoint lifetimes in the alias specialization. + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, alias_a_elements], + 1024, + ] + sRawK: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, alias_raw_k_elements], + 16, + ] + + @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 * int(self.use_tma), + ] + tmem_holding_buf: cutlass.Int32 + sA: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, a_elements], + 1024, + ] + sB_hi: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, b_hi_elements], + 1024, + ] + sB_lo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, b_lo_elements], + 1024, + ] + sMagnitudeLo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_lo_elements], + 1024, + ] + sMagnitudeLoFp32: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, magnitude_lo_fp32_elements], + 1024, + ] + sMagnitudeHiFp16: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float16, magnitude_hi_fp16_elements], + 1024, + ] + sRawK: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_k_elements], + 1024, + ] + sRawBf16A: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_bf16_a_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, + ] + sRawBf16B2: cute.struct.Align[ + cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b2_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, + ] + sARawAlias: SharedARawAlias + + self.shared_storage = SharedStorage + # 64-bit: at large request counts this product exceeds 2^31 (the + # score scratch spans request*layer segments of seq_len columns), so + # the head-plane stride must reach the kernel as Int64. + sum_seq = cutlass.Int64(request_count * self.num_layers * self.seq_len) + 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, + page_ids, + seg_page_off, + seg_req_id, + seg_layer_id, + seg_seq_len, + seg_out_offset, + q_real, + q_imag, + mlr_coef, + mean_cos, + mean_sin, + freq_scale_sq, + output, + partial_stats, + sum_seq, + a_smem_layout, + raw_bf16_a_smem_layout, + raw_bf16_direct_a_smem_layout, + raw_bf16_split_a_smem_layout, + raw_tma_smem_layout, + raw_bf16_b_smem_layout, + magnitude_lo_smem_layout, + magnitude_lo_fp16_smem_layout, + magnitude_hi_fp16_smem_layout, + magnitude_fp16_a_smem_layout, + magnitude_fp16_b_smem_layout, + b_smem_layout, + ).launch( + grid=(num_ctas, 1, 1), + block=(self.physical_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, + page_ids: cute.Tensor, + seg_page_off: cute.Tensor, + seg_req_id: cute.Tensor, + seg_layer_id: cute.Tensor, + seg_seq_len: cute.Tensor, + seg_out_offset: 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, + sum_seq: cutlass.Int64, + a_smem_layout: cute.ComposedLayout, + raw_bf16_a_smem_layout: cute.ComposedLayout, + raw_bf16_direct_a_smem_layout: cute.ComposedLayout, + raw_bf16_split_a_smem_layout: cute.ComposedLayout, + raw_tma_smem_layout: cute.ComposedLayout, + raw_bf16_b_smem_layout: cute.ComposedLayout, + magnitude_lo_smem_layout: cute.ComposedLayout, + magnitude_lo_fp16_smem_layout: cute.ComposedLayout, + magnitude_hi_fp16_smem_layout: cute.ComposedLayout, + magnitude_fp16_a_smem_layout: cute.ComposedLayout, + magnitude_fp16_b_smem_layout: cute.ComposedLayout, + 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] + valid_seq_len = seg_seq_len[segment] + page_off = seg_page_off[segment] + out_base = seg_out_offset[segment] + + 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(72)) + 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)] + 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)] + # Each stage slice retains the K_SW64 pointer flags. Reuse only + # the feature-first outer mapping for the corresponding TMA + # destination so the swizzle is not applied twice. + raw_tma_shared_real = cute.make_tensor( + cpasync_raw_k_real.iterator, + raw_tma_smem_layout.outer, + ) + raw_tma_shared_imag = cute.make_tensor( + cpasync_raw_k_imag.iterator, + raw_tma_smem_layout.outer, + ) + raw_tma_shared_real_next = cute.make_tensor( + cpasync_raw_k_real_next.iterator, + raw_tma_smem_layout.outer, + ) + raw_tma_shared_imag_next = cute.make_tensor( + cpasync_raw_k_imag_next.iterator, + raw_tma_smem_layout.outer, + ) + raw_tma_source_tiles = cute.local_tile( + raw_tma_source, + (self.raw_tma_feature_extent, CTA_M), + coord=(None, None, None), + ) + raw_tma_shared_partition_real, raw_tma_global_partition = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_tma_shared_real, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + raw_tma_shared_partition_imag, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_tma_shared_imag, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + raw_tma_shared_partition_real_next, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_tma_shared_real_next, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + raw_tma_shared_partition_imag_next, _ = cpasync.tma_partition( + raw_tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(raw_tma_shared_imag_next, 0, 2), + cute.group_modes(raw_tma_source_tiles, 0, 2), + ) + 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, + ) + 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, + ) + + page_index = self.score_start // PAGE_TOKENS + page_shard + page_start = page_index * PAGE_TOKENS + shard_first_page_start = page_start + pages_processed = cutlass.Int32(0) + if cutlass.const_expr(self.write_partial_stats): + stats_page_scores_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) + stats_origins_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) + stats_sums_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) + stats_square_sums_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) + for stats_head in cutlass.range_constexpr(N): + 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) + shard_has_page = valid_seq_len > self.score_start and page_start < valid_seq_len + empty_shard = valid_seq_len <= self.score_start or page_start >= valid_seq_len + if cutlass.dynamic_expr(shard_has_page): + if warp_idx == self.producer_warp_id: + if lane_idx == 0: + producer_prefetched_page_id_lane0 = cutlass.Int32( + page_ids[page_off + page_index] + ) + 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() + for weight_round in cutlass.range_constexpr(N * K // THREADS): + linear_index = tidx + weight_round * THREADS + qg = linear_index // K + feature = linear_index % K + coefficient_kind = feature // NUM_FREQS + frequency = feature % NUM_FREQS + mean_offset = req_id * NUM_FREQS + frequency + q_head = kv_head * self.group_size + qg + calib_offset = (layer_id * self.num_q_heads + q_head) * 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]) + 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, + ) + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + 0, + (kv_head, prefetched_physical_page), + ) + ], + raw_tma_shared_partition_real, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + 0, + (kv_head, prefetched_physical_page), + ) + ], + raw_tma_shared_partition_imag, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + while ( + valid_seq_len > self.score_start + and page_start < valid_seq_len + and pages_processed < self.max_pages + ): + physical_page = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, + ) + for page_half in cutlass.range_constexpr(self.halves_per_page): + raw_page_buffer = cutlass.Int32(0) + if cutlass.const_expr(self.write_partial_stats): + raw_page_buffer = pages_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 4-KiB K_SW64 real + # view. Every producer-warp lane participates + # in the PipelineTmaAsync barrier election. + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + page_half, + (kv_head, physical_page), + ) + ], + raw_tma_shared_partition_real, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + 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) + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + page_half, + (kv_head, physical_page), + ) + ], + raw_tma_shared_partition_imag, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + + if cutlass.const_expr( + self.producer_page_id_prefetch and page_half == self.halves_per_page - 1 + ): + next_page_id_lane0 = cutlass.Int32(0) + if warp_idx == self.producer_warp_id: + if lane_idx == 0: + next_page_start = page_start + PAGE_TOKENS * self.page_shards + next_pages_processed = pages_processed + 1 + if ( + next_page_start < valid_seq_len + and next_pages_processed < self.max_pages + ): + next_page_id_lane0 = cutlass.Int32( + page_ids[page_off + page_index + self.page_shards] + ) + 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(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: + if cutlass.const_expr(page_half + 1 < self.halves_per_page): + prefetched_physical_page = physical_page + prefetch_next_raw = True + prefetched_page_half = page_half + 1 + else: + next_page_start = page_start + PAGE_TOKENS * self.page_shards + next_pages_processed = pages_processed + 1 + prefetch_next_raw = ( + next_page_start < valid_seq_len + and next_pages_processed < self.max_pages + ) + prefetched_physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, + ) + prefetched_page_half = 0 + if cutlass.dynamic_expr(prefetch_next_raw): + 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): + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + prefetched_page_half, + (kv_head, prefetched_physical_page), + ) + ], + raw_tma_shared_partition_real, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + else: + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + prefetched_page_half, + (kv_head, prefetched_physical_page), + ) + ], + raw_tma_shared_partition_real_next, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + if cutlass.dynamic_expr(next_raw_page_buffer == 0): + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + prefetched_page_half, + (kv_head, prefetched_physical_page), + ) + ], + raw_tma_shared_partition_imag, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + else: + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + prefetched_page_half, + (kv_head, prefetched_physical_page), + ) + ], + raw_tma_shared_partition_imag_next, + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + frequency = lane_idx + # Issue several independent token loads before consuming + # any of them. This bounded RMEM window is unchanged; the + # optional half-page staging only switches its K source + # from global to the single raw shared buffer. + for token_base in cutlass.range( + 0, + CTA_M // (THREADS // 32), + self.prefetch_depth, + unroll_full=not self.compact_token_loop, + ): + 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 + if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): + magnitude = cute.math.sqrt( + norm2, + approx=self.sqrt_mode == "approx", + ftz=self.magnitude_sqrt_ftz, + ) + elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): + # cutlass 4.5 renamed the approximate-sqrt control + # to ``fastmath``; map the measured approx choice + # onto it to preserve the authored behavior. + magnitude = cute.math.sqrt(norm2, fastmath=self.sqrt_mode == "approx") + else: + # DSLs with neither spelling get the plain (IEEE) + # sqrt, which is strictly MORE accurate than the + # measured approx choice; the equivalence test's + # tolerance absorbs the difference. + magnitude = cute.math.sqrt(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(NUM_FREQS // 16): + imag_b_block = 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(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(NUM_FREQS // 16): + imag_b_block = 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, + ) + # Keep the full four-product control unchanged. + # The independent omit_a1b1 mode drops only the + # second-order residual product; all other FP16 + # K16 products retain their original order. + magnitude_lo_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + True, + ) + for magnitude_k_block in cutlass.range_constexpr(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(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(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 pages_processed == 0: + if cutlass.const_expr(page_half == 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): + # The alternate raw-page buffer was filled while this + # page's UMMA completed. Release only the current imag + # phase after all of its asynchronous consumers finish. + raw_tma_pipeline.consumer_release(raw_tma_consumer_state) + raw_tma_consumer_state.advance() + # Every term multiplying sum_seq must stay 64-bit: with + # request*layer segments of seq_len columns the head-plane + # stride alone can exceed 2^31. + output_offset = ( + cutlass.Int64(kv_head * self.group_size) * sum_seq + + out_base + + page_start + + page_half * CTA_M + ) + page_output = cute.make_tensor( + output.iterator + output_offset, + cute.make_layout( + (CTA_M, N, 1), + stride=( + 1, + sum_seq, + self.group_size * sum_seq, + ), + ), + ) + 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 + tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( + epilogue_tidx, tCtAcc, tCgC, self.epi_tile, False + ) + simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( + epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile, None + ) + 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)) + if cutlass.const_expr( + self.score_start % CTA_M == 0 and self.seq_len % CTA_M == 0 + ): + cute.copy( + simt_atom, + tTR_rC, + tTR_gC[(None, None, None, subtile_idx)], + ) + elif cutlass.dynamic_expr( + page_start >= self.score_start and page_start + CTA_M <= self.seq_len + ): + cute.copy( + simt_atom, + tTR_rC, + tTR_gC[(None, None, None, subtile_idx)], + ) + else: + output_token = page_start + epilogue_tidx + if cutlass.dynamic_expr( + output_token >= self.score_start and output_token < self.seq_len + ): + 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 pages_processed == 0: + if cutlass.const_expr(page_half == 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 pages_processed == 0: + if cutlass.const_expr(page_half == 0): + for stats_head in cutlass.range_constexpr(N): + stats_origins_m128[stats_head] = sStats[stats_head] + stats_token = page_start + tidx + if tidx < EPILOGUE_THREADS: + if cutlass.dynamic_expr( + stats_token >= self.score_start and stats_token < valid_seq_len + ): + for stats_head in cutlass.range_constexpr(N): + 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 + ) + page_index += self.page_shards + page_start += PAGE_TOKENS * self.page_shards + pages_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(N): + 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 = 8 + (warp_idx * N + stats_head) * 2 + sStats[stats_scratch_base] = stats_sum + sStats[stats_scratch_base + 1] = stats_square_sum + cute.arch.barrier() + if warp_idx == 0: + if lane_idx < N: + 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 = 8 + (stats_warp * N + 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 = pages_processed * PAGE_TOKENS + if pages_processed > 0: + stats_invalid_prefix = self.score_start - shard_first_page_start + if cutlass.dynamic_expr(stats_invalid_prefix > 0): + stats_count_i32 = stats_count_i32 - stats_invalid_prefix + stats_last_page = page_start - PAGE_TOKENS * self.page_shards + stats_invalid_tail = stats_last_page + PAGE_TOKENS - valid_seq_len + 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 * N + lane_idx + stats_base = (stats_row * self.page_shards + page_shard) * 3 + partial_stats[stats_base] = stats_count + partial_stats[stats_base + 1] = stats_mean + partial_stats[stats_base + 2] = 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], +) -> torch.Tensor: + """Encode one immutable feature-first TensorMap per layer index.""" + 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, tokens_per_block, head_dim = pool.shape + if (kv_factor, tokens_per_block, head_dim) != (2, PAGE_TOKENS, 2 * NUM_FREQS): + raise ValueError("TriAttention CuTe score requires [page, 2, Hkv, 128, 64] 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, PAGE_TOKENS] + 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) + box_dims = [NUM_FREQS, CTA_M] + [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, + cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B, + 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) + + +class TriAttentionCuteScoreRunner: + """Compile and launch the exact SM100 mean-score specialization.""" + + def __init__( + self, + *, + layer_pools: list[torch.Tensor], + layer_indices: list[int], + max_requests: int, + num_layers: int, + seq_len: int, + score_start: int, + num_q_heads: int, + num_kv_heads: int, + num_freqs: int, + tokens_per_block: int, + page_ids: torch.Tensor, + seg_page_off: torch.Tensor, + seg_req_id: torch.Tensor, + seg_layer_id: torch.Tensor, + seg_seq_len: torch.Tensor, + seg_out_offset: torch.Tensor, + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + freq_scale_sq: torch.Tensor, + output: torch.Tensor, + enable_partial_stats: bool = False, + small_workload_page_shards: int = 3, + ) -> None: + self.max_requests = int(max_requests) + self.num_layers = int(num_layers) + self.width = int(seq_len - score_start) + self.num_q_heads = int(num_q_heads) + self.num_kv_heads = int(num_kv_heads) + self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) + self.enable_partial_stats = bool(enable_partial_stats) + if small_workload_page_shards not in _SUPPORTED_PAGE_SHARDS: + raise ValueError("TriAttention CuTe score has unsupported page shards") + self.small_workload_page_shards = int(small_workload_page_shards) + partial_stats_elements = ( + max_requests * num_layers * num_q_heads * self.small_workload_page_shards * 3 + if self.enable_partial_stats + else 1 + ) + self.partial_stats = torch.empty( + partial_stats_elements, + dtype=torch.float32, + device=output.device, + ) + self.descriptors = _encode_tma_descriptors(layer_pools, layer_indices) + self._torch_prefix = ( + page_ids, + seg_page_off, + seg_req_id, + seg_layer_id, + seg_seq_len, + seg_out_offset, + q_real, + q_imag, + mlr_coef, + ) + self._torch_tail = ( + freq_scale_sq, + output, + self.partial_stats, + layer_pools[layer_indices[0]], + self.descriptors, + ) + self._cute_prefix = tuple(_to_cute(tensor) for tensor in self._torch_prefix) + self._cute_tail = ( + _to_cute(freq_scale_sq), + _to_cute(output), + _to_cute(self.partial_stats), + _to_cute(layer_pools[layer_indices[0]]), + _to_cute(self.descriptors, assumed_align=128), + ) + self._compiled: dict[int, object] = {} + self._compiled_stats: dict[int, object] = {} + self._compiled_normalize_union: dict[int, object] = {} + self._page_shards: dict[int, int] = {} + compile_output_rows = max_requests if self.enable_partial_stats else 1 + self._normalize_union_compile_output = torch.empty( + (compile_output_rows, self.width), + dtype=torch.float32, + device=output.device, + ) + self._cute_selection_prefix = ( + _to_cute(output), + _to_cute(seg_seq_len), + _to_cute(seg_out_offset), + ) + static_geometry = ( + max_requests, + num_layers, + seq_len, + score_start, + num_q_heads, + num_kv_heads, + num_freqs, + tokens_per_block, + tuple(int(value) for value in layer_pools[layer_indices[0]].shape), + tuple(int(value) for value in layer_pools[layer_indices[0]].stride()), + ) + tensor_specs = tuple( + _tensor_spec(tensor) + for tensor in ( + *self._torch_prefix, + mean_cos.view(-1), + mean_sin.view(-1), + *self._torch_tail, + ) + ) + variants = [(1, self.small_workload_page_shards)] + if max_requests > 1: + variants.append((max_requests, 2)) + for request_count, page_shards in variants: + cache_key = ( + "triattention_cute_score", + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + kernel = _TriAttentionScoreKernel( + num_layers=num_layers, + seq_len=seq_len, + score_start=score_start, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + pool_shape=tuple( + int(value) for value in layer_pools[layer_indices[0]].shape + ), + pool_strides=tuple( + int(value) for value in layer_pools[layer_indices[0]].stride() + ), + pool_dtype=cutlass.BFloat16, + page_shards=page_shards, + ) + stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) + compiled = cute.compile( + kernel, + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled + self._compiled[request_count] = compiled + self._page_shards[request_count] = page_shards + if self.enable_partial_stats: + stats_cache_key = ( + "triattention_cute_score_stats", + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + with _COMPILE_LOCK: + compiled_stats = _COMPILED_KERNELS.get(stats_cache_key) + if compiled_stats is None: + stats_kernel = _TriAttentionScoreKernel( + num_layers=num_layers, + seq_len=seq_len, + score_start=score_start, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + pool_shape=tuple( + int(value) for value in layer_pools[layer_indices[0]].shape + ), + pool_strides=tuple( + int(value) for value in layer_pools[layer_indices[0]].stride() + ), + pool_dtype=cutlass.BFloat16, + page_shards=page_shards, + write_partial_stats=True, + ) + stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) + compiled_stats = cute.compile( + stats_kernel, + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[stats_cache_key] = compiled_stats + self._compiled_stats[request_count] = compiled_stats + + if max_requests > 1: + small_score = self._compiled[1] + large_score = self._compiled[max_requests] + small_stats = self._compiled_stats.get(1) + large_stats = self._compiled_stats.get(max_requests) + for request_count in range(1, max_requests + 1): + two_shard_ctas = request_count * num_layers * num_kv_heads * 2 + use_extra_score_shard = two_shard_ctas < 2 * self.sm_count + self._compiled[request_count] = ( + small_score if use_extra_score_shard else large_score + ) + self._page_shards[request_count] = ( + self.small_workload_page_shards if use_extra_score_shard else 2 + ) + if self.enable_partial_stats: + self._compiled_stats[request_count] = ( + small_stats if use_extra_score_shard else large_stats + ) + + if self.enable_partial_stats: + from .triattention_cute_selection import ( + _select_normalize_union_config, + _TriAttentionNormalizeUnionKernel, + ) + + compiled_configs: dict[tuple[int, int, int, int], object] = {} + for request_count in range(1, max_requests + 1): + page_shards = self._page_shards[request_count] + config = _select_normalize_union_config( + request_count, + self.width, + self.sm_count, + ) + config_key = (page_shards, *config) + compiled_selection = compiled_configs.get(config_key) + if compiled_selection is None: + cache_key = ( + "triattention_cute_normalize_union", + static_geometry, + tensor_specs, + config_key, + _tensor_spec(self._normalize_union_compile_output), + _tensor_spec(self.partial_stats), + ) + with _COMPILE_LOCK: + compiled_selection = _COMPILED_KERNELS.get(cache_key) + if compiled_selection is None: + tokens_per_lane, token_subtiles, row_cluster_ctas = config + kernel = _TriAttentionNormalizeUnionKernel( + num_layers=num_layers, + seq_len=seq_len, + score_start=score_start, + num_q_heads=num_q_heads, + page_shards=page_shards, + tokens_per_lane=tokens_per_lane, + token_subtiles=token_subtiles, + row_cluster_ctas=row_cluster_ctas, + ) + stream = cuda.CUstream( + torch.cuda.current_stream(output.device).cuda_stream + ) + compiled_selection = cute.compile( + kernel, + _to_cute(self.partial_stats), + *self._cute_selection_prefix, + _to_cute(self._normalize_union_compile_output.view(-1)), + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled_selection + compiled_configs[config_key] = compiled_selection + self._compiled_normalize_union[request_count] = compiled_selection + + def supports(self, request_count: int) -> bool: + """Return whether the dynamic specialization covers this request count.""" + return request_count in self._compiled + + def launch( + self, + request_count: int, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + ) -> None: + """Launch the CuTe score kernel on the current PyTorch stream.""" + stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) + self._compiled[request_count]( + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + request_count, + stream, + ) + + def launch_with_partial_stats( + self, + request_count: int, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + ) -> tuple[torch.Tensor, int]: + """Launch score and write deterministic partial row statistics.""" + stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) + self._compiled_stats[request_count]( + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + request_count, + stream, + ) + row_count = request_count * self.num_layers * self.num_q_heads + page_shards = self._page_shards[request_count] + stats = self.partial_stats[: row_count * page_shards * 3] + return stats.view(row_count, page_shards, 3), page_shards + + def supports_union_fusion(self, request_count: int) -> bool: + """Return whether the score/stats/union pipeline was precompiled.""" + return ( + request_count in self._compiled_stats + and request_count in self._compiled_normalize_union + ) + + def launch_union_fusion( + self, + request_count: int, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + union_scores: torch.Tensor, + ) -> None: + """Launch score plus stats followed by normalized union reduction.""" + stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) + self._compiled_stats[request_count]( + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + request_count, + stream, + ) + self._launch_union_finalize(request_count, union_scores, stream) + + def _launch_union_finalize( + self, + request_count: int, + union_scores: torch.Tensor, + stream: cuda.CUstream, + ) -> None: + if ( + union_scores.shape != (request_count, self.width) + or union_scores.dtype != torch.float32 + or union_scores.device != self.partial_stats.device + or not union_scores.is_contiguous() + ): + raise ValueError("TriAttention union output does not match the compiled geometry") + self._compiled_normalize_union[request_count]( + _to_cute(self.partial_stats), + *self._cute_selection_prefix, + _to_cute(union_scores.view(-1)), + request_count, + stream, + ) + + def launch_union_finalize( + self, + request_count: int, + union_scores: torch.Tensor, + ) -> None: + """Launch only the union finalizer for isolated validation and profiling.""" + stream = cuda.CUstream(torch.cuda.current_stream(union_scores.device).cuda_stream) + self._launch_union_finalize(request_count, union_scores, stream) 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..da08a08669b4 --- /dev/null +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -0,0 +1,516 @@ +# 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 + +_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 +_STATS_FIELDS = 3 +_STD_EPSILON = 1.0e-6 +_SUPPORTED_PAGE_SHARDS = (2, 3) + + +def _select_normalize_union_config( + request_count: int, + width: int, + sm_count: int, +) -> tuple[int, int, int]: + """Return tokens per lane, token subtiles, and 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, + ) + ) + + +@cute.jit +def _mapa_cluster(smem_ptr, peer_rank): + return _mapa_shared_cluster(smem_ptr, peer_rank) + + +@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, + ) + ) + + +@cute.jit +def _ld_cluster_f32(mapped_addr): + return _ld_shared_cluster_f32(mapped_addr) + + +class _TriAttentionNormalizeUnionKernel: + """Merge row moments, normalize scores, and reduce their elementwise maximum.""" + + def __init__( + self, + *, + num_layers: int, + seq_len: int, + score_start: int, + num_q_heads: int, + page_shards: int, + tokens_per_lane: int, + token_subtiles: int, + row_cluster_ctas: int, + ) -> None: + if min(num_layers, seq_len, num_q_heads) <= 0: + raise ValueError("TriAttention stats/union reduction requires positive geometry") + if not 0 <= score_start < seq_len: + raise ValueError("TriAttention stats/union score_start is out of range") + if page_shards not in _SUPPORTED_PAGE_SHARDS: + raise ValueError("TriAttention stats/union has unsupported page shards") + if tokens_per_lane not in (_SMALL_TOKENS_PER_LANE, _LARGE_TOKENS_PER_LANE): + raise ValueError("TriAttention stats/union has unsupported load width") + if token_subtiles not in (_SMALL_TOKEN_SUBTILES, _LARGE_TOKEN_SUBTILES): + raise ValueError("TriAttention stats/union has unsupported token subtiles") + if row_cluster_ctas not in (1, 2, 4): + raise ValueError("TriAttention stats/union has unsupported row cluster") + self.num_layers = num_layers + self.seq_len = seq_len + self.score_start = score_start + self.width = seq_len - score_start + 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 + self.num_token_tiles = (self.width + self.token_tile - 1) // self.token_tile + + @cute.jit + def __call__( + self, + partial_stats: cute.Tensor, + scores: cute.Tensor, + seg_seq_len: cute.Tensor, + seg_out_offset: cute.Tensor, + union_scores: cute.Tensor, + request_count: cutlass.Int32, + stream: cuda.CUstream, + ): + kernel = self.kernel( + partial_stats, + scores, + seg_seq_len, + seg_out_offset, + 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: + 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.kernel + def kernel( + self, + partial_stats: cute.Tensor, + scores: cute.Tensor, + seg_seq_len: cute.Tensor, + seg_out_offset: 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 + valid_width = seg_seq_len[first_segment] - self.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 + + 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 + 1] + m2_0 = partial_stats[stats_base + 2] + mean_1 = partial_stats[stats_base + 4] + m2_1 = partial_stats[stats_base + 5] + 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 + 1] + shard_m2 = partial_stats[stats_base + 2] + 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(q_head) * request_count * self.num_layers * self.seq_len + + seg_out_offset[segment] + + self.score_start + + subtile_first_token + ) + if cutlass.const_expr( + self.score_start % self.tokens_per_lane == 0 + and self.seq_len % self.tokens_per_lane == 0 + ) and cutlass.dynamic_expr( + subtile_first_token + self.tokens_per_lane <= valid_width + ): + score_tile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (scores.iterator + score_index).toint(), + AddressSpace.gmem, + assumed_align=self.tokens_per_lane * 4, + ), + cute.make_layout(self.tokens_per_lane), + ) + cute.copy( + score_copy_atom, + cute.coalesce(score_tile), + cute.coalesce(score_value_tiles[token_subtile]), + ) + else: + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + token = subtile_first_token + token_slot + if cutlass.dynamic_expr(token < valid_width): + score_value_tiles[token_subtile][token_slot] = scores[ + score_index + 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: + 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): + 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 + if cutlass.const_expr(self.width % self.tokens_per_lane == 0): + if cutlass.dynamic_expr( + subtile_first_token + self.tokens_per_lane <= self.width + ): + union_index = request_idx * self.width + subtile_first_token + union_tile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (union_scores.iterator + union_index).toint(), + AddressSpace.gmem, + assumed_align=self.tokens_per_lane * 4, + ), + cute.make_layout(self.tokens_per_lane), + ) + cute.copy( + score_copy_atom, + cute.coalesce(reduced_values), + cute.coalesce(union_tile), + ) + else: + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + token = subtile_first_token + token_slot + if cutlass.dynamic_expr(token < self.width): + union_scores[request_idx * self.width + token] = reduced_values[ + token_slot + ] + else: + # Warp 0 reduces its CTA's row partition, then CTA 0 combines the + # cluster's partial maxima through distributed shared memory. + 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: + 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): + union_value = warp_max[(0, token_subtile, token_slot, lane_idx)] + shared_offset = ( + token_subtile * self.tokens_per_lane * _WARP_SIZE + + token_slot * _WARP_SIZE + + lane_idx + ) + for peer_rank in cutlass.range_constexpr(1, self.row_cluster_ctas): + remote_addr = _mapa_cluster( + warp_max_ptr, + cutlass.Int32(peer_rank), + ) + union_value = cute.arch.fmax( + union_value, + _ld_cluster_f32(remote_addr + shared_offset * 4), + ) + reduced_values[token_slot] = union_value + subtile_first_token = first_token + token_subtile * self.subtile_token_tile + if cutlass.const_expr(self.width % self.tokens_per_lane == 0): + if cutlass.dynamic_expr( + subtile_first_token + self.tokens_per_lane <= self.width + ): + union_index = request_idx * self.width + subtile_first_token + union_tile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (union_scores.iterator + union_index).toint(), + AddressSpace.gmem, + assumed_align=self.tokens_per_lane * 4, + ), + cute.make_layout(self.tokens_per_lane), + ) + cute.copy( + score_copy_atom, + cute.coalesce(reduced_values), + cute.coalesce(union_tile), + ) + else: + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + token = subtile_first_token + token_slot + if cutlass.dynamic_expr(token < self.width): + union_scores[request_idx * self.width + token] = reduced_values[ + token_slot + ] + 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 index f67fa042bb36..d85d1d9defa7 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -22,6 +22,7 @@ from __future__ import annotations +import os from typing import List, Optional import torch @@ -427,6 +428,13 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> self._cute_scratch = scratch self._cute_seg_seq_len = seg_seq_len self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) + # Opt-in fused score+stats+union pipeline (Fanrong Li's two-kernel + # scheme). Runners are built lazily per uniform score start because + # the start is a compile-time constant of the fused kernel. + self._cute_union_fusion_enabled = ( + os.environ.get("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "0") == "1" + ) + self._union_fusion_runners = {} from tensorrt_llm.logger import logger logger.info( @@ -535,6 +543,117 @@ def launch( ) return output + def _union_fusion_runner_for( + self, score_start: int, mean_cos: torch.Tensor, mean_sin: torch.Tensor + ): + """Build (or reuse) the fused score/stats/union runner for one start. + + The fused kernel bakes ``score_start`` at compile time, so one runner + exists per distinct uniform prompt length; the cache is capped so a + prompt-length churn cannot trigger unbounded recompiles. ``None`` + entries record geometries the fused pipeline rejected. + """ + if score_start in self._union_fusion_runners: + return self._union_fusion_runners[score_start] + if len(self._union_fusion_runners) >= 4 or not 0 <= score_start < self.seq_len: + return None + num_q_heads, num_kv_heads, num_freqs, tokens_per_block, _ = self.geometry_args + entry = None + try: + from .triattention_cute_score_fused import ( + TriAttentionCuteScoreRunner as _FusedUnionScoreRunner, + ) + + device = self.output.device + seg_out_offset = ( + torch.arange(self.max_requests * self.num_layers, dtype=torch.int64, device=device) + * self.seq_len + ).to(torch.int32) + union_rows = torch.empty( + (self.max_requests, self.seq_len - score_start), + dtype=torch.float32, + device=device, + ) + runner = _FusedUnionScoreRunner( + layer_pools=self._cute_layer_pools, + layer_indices=self._cute_layer_indices, + max_requests=self.max_requests, + num_layers=self.num_layers, + seq_len=self.seq_len, + score_start=score_start, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + page_ids=self.pointer_prefix[2], + seg_page_off=self.pointer_prefix[3], + seg_req_id=self.pointer_prefix[4], + seg_layer_id=self.pointer_prefix[5], + seg_seq_len=self._cute_seg_seq_len, + seg_out_offset=seg_out_offset, + q_real=self.pointer_middle[0], + q_imag=self.pointer_middle[1], + mlr_coef=self.pointer_middle[2], + mean_cos=mean_cos, + mean_sin=mean_sin, + freq_scale_sq=self.pointer_tail[0], + output=self._cute_scratch, + enable_partial_stats=True, + ) + entry = (runner, union_rows) + except (ImportError, RuntimeError, ValueError, AssertionError) as error: + import warnings + + warnings.warn(f"TriAttention CuTe union fusion unavailable: {error}") + self._union_fusion_runners[score_start] = entry + return entry + + def launch_cute_union_fusion( + self, + request_count: int, + valid_seq_lens: torch.Tensor, + valid_widths: torch.Tensor, + token_starts_device: torch.Tensor, + score_start: Optional[int], + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + union_out: torch.Tensor, + ) -> bool: + """Run the fused score+stats+normalized-union pipeline when possible. + + Returns False (without launching) whenever the opt-in is off, the + cohort's prompt lengths are not uniform, or the geometry has no fused + specialization — the caller then falls back to the split score, + row-stats, and union path. + """ + self.prepare_cute_score(mean_cos, mean_sin) + if not self._cute_union_fusion_enabled or score_start is None: + return False + if request_count <= 0 or request_count > self.max_requests: + raise ValueError("request count exceeds fixed score capacity") + entry = self._union_fusion_runner_for(int(score_start), mean_cos, mean_sin) + if entry is None: + return False + runner, union_rows = entry + if not runner.supports_union_fusion(request_count): + return False + num_segments = request_count * self.num_layers + torch.sub( + valid_seq_lens[:request_count], + token_starts_device[:request_count], + out=valid_widths[:request_count], + ) + torch.index_select( + valid_seq_lens, + 0, + self.pointer_prefix[4][:num_segments], + out=self._cute_seg_seq_len[:num_segments], + ) + runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_rows[:request_count]) + columns = min(union_rows.shape[1], union_out.shape[1]) + union_out[:request_count, :columns].copy_(union_rows[:request_count, :columns]) + return True + # --------------------------------------------------------------------------- # # Selection: combine scores per mode, then finalize the top-k set. # 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..0a7a90c42ae7 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_union_fusion.py @@ -0,0 +1,231 @@ +# 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 (two CuTe kernels).""" + +import pytest +import torch + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="TriAttention CuTe kernels require SM100", +) +@pytest.mark.xfail( + reason="the fused score kernel misreads every page after the first under " + "this environment's cutlass DSL (page-id prefetch drift; page 0 matches " + "the oracle to 1e-6, pages 1+ are corrupt) — kept as the repro while the " + "kernel/DSL mismatch is resolved upstream", + strict=True, +) +@pytest.mark.parametrize( + "score_start,valid_lens", + [ + # Full-range scoring, both requests full length. + (0, None), + # Non-page-aligned uniform start with ragged valid lengths. + (37, [250, 198]), + # Page-aligned start. + (128, [250, 230]), + ], +) +def test_union_fusion_matches_split_pipeline( + score_start: int, + valid_lens: "list | None", + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The fused pipeline must reproduce the split score->stats->union rows. + + Geometry is the fused kernel's contract (128-token pages, GQA group 8, + 32 frequencies). The reference runs the production split path: the score + launch gathers each request's decode window, then ``prepare_union_scores`` + normalizes rows and takes the cross-row union maximum. + """ + pytest.importorskip("cutlass") + monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + prepare_union_scores, + ) + + torch.manual_seed(20260721) + device = torch.device("cuda") + seq_len = 256 + tokens_per_block = 128 + num_freqs = 32 + num_q_heads = 8 + num_pages = seq_len // tokens_per_block + 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) + round_starts = torch.tensor([float(seq_len), float(seq_len + 1)], device=device) + phase = (round_starts[:, 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() + + k_plane = [2 * page for page in range(num_pages)] + v_plane = [2 * page + 1 for page in range(num_pages)] + block_offsets = torch.tensor( + [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device + ) + group = _FixedScoreGroup( + [pool], + [0], + 2, + num_pages, + seq_len, + num_q_heads, + block_offsets, + [0], + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + output_width=seq_len, + ) + if valid_lens is None: + valid_lens = [seq_len, seq_len] + valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) + request_count = 2 + + # Reference: the split pipeline over the same decode windows. + split_widths = torch.empty(request_count, dtype=torch.int32, device=device) + token_starts = torch.full((request_count,), score_start, dtype=torch.int32, device=device) + per_head = group.launch( + request_count, + valid_seq_lens, + split_widths, + token_starts, + mean_cos, + mean_sin, + ) + rows = per_head.shape[1] * per_head.shape[2] + scores_rows = per_head.reshape(request_count, rows, seq_len).contiguous() + row_mean = torch.empty((request_count, rows, 1), dtype=torch.float32, device=device) + row_inv_std = torch.empty_like(row_mean) + expected = torch.empty((request_count, seq_len), dtype=torch.float32, device=device) + prepare_union_scores( + scores_rows, + split_widths, + row_mean, + row_inv_std, + expected, + request_count, + normalize_scores=True, + ) + + 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 + ) + launched = group.launch_cute_union_fusion( + request_count, + valid_seq_lens, + fused_widths, + token_starts, + score_start, + mean_cos, + mean_sin, + fused_out, + ) + assert launched, "fused union pipeline must engage on its contract geometry" + assert torch.equal(fused_widths, split_widths) + for request in range(request_count): + width = int(valid_lens[request]) - score_start + torch.testing.assert_close( + fused_out[request, :width], + expected[request, :width], + rtol=5.0e-3, + atol=5.0e-3, + ) + + # A cohort without one uniform prompt start must decline, not launch. + assert not group.launch_cute_union_fusion( + request_count, + valid_seq_lens, + fused_widths, + token_starts, + None, + mean_cos, + mean_sin, + fused_out, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="TriAttention CuTe kernels require SM100", +) +def test_union_fusion_declines_off_contract_geometry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """32-token pages sit outside the fused contract: decline, do not fault.""" + pytest.importorskip("cutlass") + monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + ) + + torch.manual_seed(20260721) + device = torch.device("cuda") + seq_len = 256 + tokens_per_block = 32 + num_freqs = 32 + num_q_heads = 8 + num_pages = seq_len // tokens_per_block + 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) + 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) + mean_cos = torch.cos(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) + mean_sin = torch.sin(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) + k_plane = [2 * page for page in range(num_pages)] + v_plane = [2 * page + 1 for page in range(num_pages)] + block_offsets = torch.tensor( + [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device + ) + group = _FixedScoreGroup( + [pool], + [0], + 2, + num_pages, + seq_len, + num_q_heads, + block_offsets, + [0], + q_real, + torch.randn_like(q_real) * 0.125, + torch.randn_like(q_real) * 0.125, + freq_scale_sq, + omega, + offsets, + output_width=seq_len, + ) + valid_seq_lens = torch.full((2,), seq_len, dtype=torch.int32, device=device) + widths = torch.empty(2, dtype=torch.int32, device=device) + token_starts = torch.zeros(2, dtype=torch.int32, device=device) + union_out = torch.empty((2, seq_len), dtype=torch.float32, device=device) + with pytest.warns(UserWarning, match="union fusion unavailable"): + launched = group.launch_cute_union_fusion( + 2, + valid_seq_lens, + widths, + token_starts, + 0, + mean_cos.contiguous(), + mean_sin.contiguous(), + union_out, + ) + assert not launched From c5795bb884b87dbedd04f407f4d4f096b84c7303 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 05:23:10 -0700 Subject: [PATCH 070/178] [None][fix] Decode staged page ids in the fused score kernel The fused pipeline's kernel read the staged K-plane block-offset entries as raw pool page indices, but the native staging layout encodes physical_page * kv_factor (2). Page 0 passed by accident (0 * 2 == 0) while every later page read out of the pool; the whole-pool probe now matches the torch oracle on every page and the fused-vs-split equivalence tests pass (the xfail repro marker is removed). Production-shaped kernel timing at the fused contract geometry (64q/8kv heads, 32 freqs, 128-token pages, 12 layers, seq 1280, capacity-256 launches): split 2.30/2.47/3.64 ms vs fused 0.93/1.09/2.17 ms per round at 1/32/256 active requests (2.48x/2.26x/1.67x). Signed-off-by: tianruih --- .../triattention/triattention_cute_score_fused.py | 14 ++++++++++---- .../test_triattention_cute_union_fusion.py | 7 ------- 2 files changed, 10 insertions(+), 11 deletions(-) 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 index 79137054befb..e39ee467375f 100644 --- 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 @@ -714,8 +714,11 @@ def kernel( if cutlass.dynamic_expr(shard_has_page): if warp_idx == self.producer_warp_id: if lane_idx == 0: - producer_prefetched_page_id_lane0 = cutlass.Int32( - page_ids[page_off + page_index] + # The staged K-plane entries encode physical_page * + # kv_factor (2); decode to the pool page index here, + # matching the production score kernel. + producer_prefetched_page_id_lane0 = ( + cutlass.Int32(page_ids[page_off + page_index]) // 2 ) tCrRawBf16ASplit = raw_bf16_tiled_mma.make_fragment_A(cpasync_raw_k_0) tCrRawBf16B0 = raw_bf16_tiled_mma.make_fragment_B(sRawBf16B0) @@ -945,8 +948,11 @@ def kernel( next_page_start < valid_seq_len and next_pages_processed < self.max_pages ): - next_page_id_lane0 = cutlass.Int32( - page_ids[page_off + page_index + self.page_shards] + next_page_id_lane0 = ( + cutlass.Int32( + page_ids[page_off + page_index + self.page_shards] + ) + // 2 ) producer_prefetched_page_id_lane0 = next_page_id_lane0 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 index 0a7a90c42ae7..229cf8524d75 100644 --- 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 @@ -10,13 +10,6 @@ not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), reason="TriAttention CuTe kernels require SM100", ) -@pytest.mark.xfail( - reason="the fused score kernel misreads every page after the first under " - "this environment's cutlass DSL (page-id prefetch drift; page 0 matches " - "the oracle to 1e-6, pages 1+ are corrupt) — kept as the repro while the " - "kernel/DSL mismatch is resolved upstream", - strict=True, -) @pytest.mark.parametrize( "score_start,valid_lens", [ From 448db5368ffe1b4c39106041faf14f4166ba0b20 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 05:46:03 -0700 Subject: [PATCH 071/178] [None][perf] Bucket the score scratch by cohort need and widen its base offset The score scratch was strided by max_seq_len per request*layer segment, which at large batch was both unindexable in 32 bits and tens of GiB (BS=256 x 8K bucket = 4.8e9 fp32 elements) for cohorts that never score past ~1K tokens per request. The bucket now follows what cohorts actually present (power-of-two, 1024 floor, growth reuses the existing rebuild path), the kernel's head-plane base offset is computed in Int64, and the setup audit guards only the remaining 32-bit plane product. First BS=256 CuTe capture on the production nsys recipe (the geometry both prior attempts fail-fasted on): score kernel 8.0 ms vs the C++-era 14.6 ms per round, scratch 19 GiB -> 4.8 GiB. Signed-off-by: tianruih --- .../triattention/triattention.py | 12 +++++++++++- .../triattention/triattention_cute_score.py | 8 +++++++- .../triattention/triattention_kernels.py | 8 +++++--- .../test_triattention_draft_cocompaction.py | 10 ++++++---- 4 files changed, 29 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b912f3576b62..098266647d12 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1986,7 +1986,17 @@ def _fixed_resources_for( needed_width, self.top_B + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), ) - seq_capacity = max(needed_page_tokens, int(mgr.max_seq_len)) + # Bucket the score scratch by what cohorts actually present instead of + # pinning it to max_seq_len: with pinned prompts the post-compaction + # length is bounded by prompt + budget + slack, so one power-of-two + # bucket serves the steady state, and a cohort that outgrows it simply + # rebuilds these resources through the capacity check above. The old + # max_seq_len floor made the scratch unindexable in 32 bits and tens + # of GiB at large batch (BS=256 x 8K bucket) for work that never + # scored past ~1K tokens per request. + seq_capacity = max(int(needed_page_tokens), 1024) + seq_capacity = 1 << (seq_capacity - 1).bit_length() + seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) # The CuTe score kernel stores full compute tiles (64 tokens, or one # page for 128-token pages) into a scratch strided by this bucket # capacity, so the capacity must be tile-aligned (its geometry gate diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py index 8e38abed1b1f..29785c702e71 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py @@ -1138,8 +1138,14 @@ def kernel( if cutlass.const_expr(page_half == 0): cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) acc_pipeline.consumer_wait(acc_consumer_state) + # 64-bit: the head-plane product exceeds 2^31 once the scratch + # spans many request*layer segments (large batch), so promote + # before the multiply reaches Int32 arithmetic. output_offset = ( - kv_head * N * self.sum_seq + out_base + page_start + page_half * CTA_M + cutlass.Int64(kv_head) * (N * self.sum_seq) + + out_base + + page_start + + page_half * CTA_M ) page_output = cute.make_tensor( output.iterator + output_offset, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index d85d1d9defa7..da1524680ec6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -353,9 +353,11 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> and num_q_heads // num_kv_heads in (4, 8) and int(anchor.stride(-1)) == 1 and self.seq_len % score_tile_tokens == 0 - # The kernel computes flat score offsets in 32-bit arithmetic; - # group-4 geometries pad the head axis to the MMA tile N=8. - and num_kv_heads * 8 * max_segments * self.seq_len < 2**31 + # The kernel's head-plane base offset is 64-bit; the widest + # 32-bit product left is one plane (N-1 head columns of one + # segment stride), which the score bucket keeps far below 2^31. + # Group-4 geometries pad the head axis to the MMA tile N=8. + and (8 - 1) * max_segments * self.seq_len < 2**31 ) if not supported: raise ValueError( 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 index ae6ab1ed799c..354a6ec790f9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -448,13 +448,15 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): ): resources = manager._fixed_resources_for(layout, prepared) - # The buffers follow the executor limits, not this one-request cohort. + # Request capacity follows the executor limits, while the score + # bucket follows what the cohort actually presents (power-of-two, + # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. assert resources.score_staging is score_staging assert score_cls.call_args.kwargs["max_requests"] == 8 assert score_cls.call_args.kwargs["decode_width"] == 4 + 2 * 128 - assert score_cls.call_args.kwargs["seq_len"] == 65536 - assert score_cls.call_args.kwargs["page_table_token_capacity"] == 65536 + 1 - assert score_cls.call_args.kwargs["draft_page_table_token_capacity"] == 65536 + 1 + assert score_cls.call_args.kwargs["seq_len"] == 1024 + assert score_cls.call_args.kwargs["page_table_token_capacity"] == 1024 + 1 + assert score_cls.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 # A second round with unchanged pools reuses the resident buffers and # keeps the cached compaction launches. From 78b7444da5a6573cd50a344c2c1f4eef4005e826 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 07:48:24 -0700 Subject: [PATCH 072/178] [None][feat] Generalize the fused score kernel to 32-token pages Port the production score kernel's page-fragment machinery onto the fused score+partial-stats kernel (CTA_M=128): one compute tile now spans fragments_per_phase pages, with per-fragment TMA views for every real/imag stage slice, per-fragment clamped page-id prefetch, page-id table indexing by page_index * pages_per_tile + fragment, and barrier transaction bytes kept at one full tile per phase. At 128-token pages fragments_per_phase == 1 const-folds every new block away, leaving the original single-copy schedule bit-identical. Equivalence vs the split pipeline passes at both 32- and 128-token pages across aligned, non-page-aligned, and ragged windows (shuffled physical-page tables catch fragment mix-ups); the full unit suite is green (131 passed). The GQA-group-8 requirement remains and now has its own decline test. Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 543 +++++++++++------- .../test_triattention_cute_union_fusion.py | 33 +- 2 files changed, 369 insertions(+), 207 deletions(-) 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 index e39ee467375f..7df5f7e9c0fd 100644 --- 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 @@ -56,7 +56,6 @@ def _cute_sqrt_keyword_mode() -> str: EPILOGUE_THREADS = 128 RAW_PAGE_BUFFERS = 2 -PAGE_TOKENS = 128 RAW_K_HALF_ELEMENTS = CTA_M * 2 * NUM_FREQS * RAW_PAGE_BUFFERS RAW_K_VECTOR_ELEMENTS = 8 RAW_K_SPLIT_PHASE_ELEMENTS = CTA_M * NUM_FREQS @@ -155,8 +154,8 @@ def __init__( raise ValueError("TriAttention CuTe score requires BF16 K pages") if num_freqs != NUM_FREQS: raise ValueError("TriAttention CuTe score requires 32 frequencies") - if tokens_per_block != PAGE_TOKENS: - raise ValueError("TriAttention CuTe score requires 128-token pages") + if tokens_per_block not in (32, 128): + raise ValueError("TriAttention CuTe score requires 32- or 128-token pages") if num_q_heads % num_kv_heads or num_q_heads // num_kv_heads != N: raise ValueError("TriAttention CuTe score requires GQA group 8") if not 0 <= score_start < seq_len: @@ -172,8 +171,18 @@ def __init__( self.group_size = num_q_heads // num_kv_heads self.page_shards = page_shards self.write_partial_stats = write_partial_stats - self.max_pages = (seq_len + PAGE_TOKENS - 1) // PAGE_TOKENS - self.halves_per_page = PAGE_TOKENS // CTA_M + self.tokens_per_block = tokens_per_block + # One 128-token compute tile either matches a page exactly (the + # validated 128-token geometry, one TMA box per phase) or spans + # several pages (32-token pages: four page fragments per phase, + # one TMA box each into the same transaction barrier). The + # single-fragment schedule of the validated geometry is unchanged. + self.box_tokens = min(CTA_M, tokens_per_block) + self.fragments_per_phase = CTA_M // self.box_tokens + self.pages_per_tile = self.fragments_per_phase + self.halves_per_page = max(1, tokens_per_block // CTA_M) + self.tile_tokens = self.halves_per_page * CTA_M + self.max_tiles = (seq_len + self.tile_tokens - 1) // self.tile_tokens # Measured final choices that still shape layouts or generated code. self.prefetch_depth = 4 @@ -183,6 +192,9 @@ def __init__( self.cpasync_schedule = "sync_each_half" self.split_raw_tma = True self.raw_tma_feature_extent = NUM_FREQS + # Barrier transaction bytes for one phase: the full 128-token tile + # of one coefficient plane, regardless of how many page fragments + # deliver it. self.raw_tma_copy_bytes = RAW_K_SPLIT_TMA_COPY_BYTES self.raw_tma_pipeline_stages = 2 * RAW_PAGE_BUFFERS if write_partial_stats else 1 self.accumulator_pipeline_stages = 1 @@ -201,11 +213,15 @@ def __init__( self.compact_token_loop = True self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape - if pool_kv_heads != num_kv_heads or pool_tokens != PAGE_TOKENS or pool_dim != 2 * NUM_FREQS: + if ( + pool_kv_heads != num_kv_heads + or pool_tokens != tokens_per_block + or 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_slot, self.s_dim = pool_strides if self.s_slot != 2 * NUM_FREQS or self.s_dim != 1: - raise ValueError("K pages must be contiguous [128, 64]") + 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") @@ -281,14 +297,14 @@ def __call__( raw_bf16_direct_a_smem_layout.inner, 0, cute.make_layout( - (self.raw_tma_feature_extent, CTA_M), + (self.raw_tma_feature_extent, self.box_tokens), stride=(1, self.raw_tma_feature_extent), ), ) raw_tma_source_layout = cute.make_layout( ( 2 * NUM_FREQS, - PAGE_TOKENS, + self.tokens_per_block, (self.num_kv_heads, self.num_physical_pages), ), stride=( @@ -305,7 +321,7 @@ def __call__( cpasync.CopyBulkTensorTileG2SOp(), raw_tma_source, raw_tma_smem_layout, - (self.raw_tma_feature_extent, CTA_M), + (self.raw_tma_feature_extent, self.box_tokens), ) raw_bf16_b_smem_layout = sm100_utils.make_smem_layout_b( raw_bf16_tiled_mma, @@ -630,55 +646,72 @@ def kernel( # Each stage slice retains the K_SW64 pointer flags. Reuse only # the feature-first outer mapping for the corresponding TMA # destination so the swizzle is not applied twice. - raw_tma_shared_real = cute.make_tensor( - cpasync_raw_k_real.iterator, - raw_tma_smem_layout.outer, - ) - raw_tma_shared_imag = cute.make_tensor( - cpasync_raw_k_imag.iterator, - raw_tma_smem_layout.outer, - ) - raw_tma_shared_real_next = cute.make_tensor( - cpasync_raw_k_real_next.iterator, - raw_tma_smem_layout.outer, - ) - raw_tma_shared_imag_next = cute.make_tensor( - cpasync_raw_k_imag_next.iterator, - raw_tma_smem_layout.outer, - ) raw_tma_source_tiles = cute.local_tile( raw_tma_source, - (self.raw_tma_feature_extent, CTA_M), + (self.raw_tma_feature_extent, self.box_tokens), coord=(None, None, None), ) - raw_tma_shared_partition_real, raw_tma_global_partition = cpasync.tma_partition( - raw_tma_atom, - 0, - cute.make_layout(1), - cute.group_modes(raw_tma_shared_real, 0, 2), - cute.group_modes(raw_tma_source_tiles, 0, 2), - ) - raw_tma_shared_partition_imag, _ = cpasync.tma_partition( - raw_tma_atom, - 0, - cute.make_layout(1), - cute.group_modes(raw_tma_shared_imag, 0, 2), - cute.group_modes(raw_tma_source_tiles, 0, 2), - ) - raw_tma_shared_partition_real_next, _ = cpasync.tma_partition( - raw_tma_atom, - 0, - cute.make_layout(1), - cute.group_modes(raw_tma_shared_real_next, 0, 2), - cute.group_modes(raw_tma_source_tiles, 0, 2), - ) - raw_tma_shared_partition_imag_next, _ = cpasync.tma_partition( - raw_tma_atom, - 0, - cute.make_layout(1), - cute.group_modes(raw_tma_shared_imag_next, 0, 2), - cute.group_modes(raw_tma_source_tiles, 0, 2), - ) + # One smem view and TMA partition per page fragment of the + # 128-token tile, for each of the four stage slices. Fragment f + # lands box_tokens rows deeper in the same stage; the offset is a + # whole multiple of the swizzle period, so the descriptor swizzle + # stays phase-aligned. + 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, + ) + 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), + ) + 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) + 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, @@ -696,8 +729,8 @@ def kernel( swizzle=raw_bf16_b_smem_layout.inner, ) - page_index = self.score_start // PAGE_TOKENS + page_shard - page_start = page_index * PAGE_TOKENS + page_index = self.score_start // self.tile_tokens + page_shard + page_start = page_index * self.tile_tokens shard_first_page_start = page_start pages_processed = cutlass.Int32(0) if cutlass.const_expr(self.write_partial_stats): @@ -709,6 +742,15 @@ def kernel( 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) + if cutlass.const_expr(self.fragments_per_phase > 1): + # Per-fragment page-id registers for multi-page compute tiles. + # Slot 0 is unused: fragment 0 keeps the scalar broadcast + # registers of the validated single-fragment schedule. + producer_prefetched_page_ids_lane0 = cute.make_rmem_tensor( + (self.pages_per_tile,), cutlass.Int32 + ) + physical_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) + prefetched_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) shard_has_page = valid_seq_len > self.score_start and page_start < valid_seq_len empty_shard = valid_seq_len <= self.score_start or page_start >= valid_seq_len if cutlass.dynamic_expr(shard_has_page): @@ -718,8 +760,26 @@ def kernel( # kv_factor (2); decode to the pool page index here, # matching the production score kernel. producer_prefetched_page_id_lane0 = ( - cutlass.Int32(page_ids[page_off + page_index]) // 2 + cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile]) // 2 ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + # The tail tile may not reach this fragment's + # page; clamp to the first fragment (those + # scores lie past the valid width and are + # masked downstream) so the TMA never + # dereferences an unstaged block entry. + fragment_page_id = producer_prefetched_page_id_lane0 + if page_start + fragment * self.box_tokens < valid_seq_len: + fragment_page_id = ( + cutlass.Int32( + page_ids[ + page_off + page_index * self.pages_per_tile + fragment + ] + ) + // 2 + ) + 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) @@ -840,42 +900,60 @@ def kernel( producer_prefetched_page_id_lane0, 0, ) - raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + prefetched_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], 0, - 0, - (kv_head, prefetched_physical_page), ) - ], - raw_tma_shared_partition_real, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - 0, - (kv_head, prefetched_physical_page), - ) - ], - raw_tma_shared_partition_imag, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() while ( valid_seq_len > self.score_start and page_start < valid_seq_len - and pages_processed < self.max_pages + and pages_processed < self.max_tiles ): physical_page = cutlass.Int32(0) if warp_idx == self.producer_warp_id: @@ -883,6 +961,12 @@ def kernel( producer_prefetched_page_id_lane0, 0, ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + physical_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], + 0, + ) for page_half in cutlass.range_constexpr(self.halves_per_page): raw_page_buffer = cutlass.Int32(0) if cutlass.const_expr(self.write_partial_stats): @@ -895,22 +979,26 @@ def kernel( # view. Every producer-warp lane participates # in the PipelineTmaAsync barrier election. raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - page_half, - (kv_head, physical_page), - ) - ], - raw_tma_shared_partition_real, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = physical_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + page_half, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() raw_tma_pipeline.consumer_wait(raw_tma_consumer_state) raw_tma_pipeline.consumer_release(raw_tma_consumer_state) @@ -918,22 +1006,26 @@ def kernel( 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) - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - page_half, - (kv_head, physical_page), - ) - ], - raw_tma_shared_partition_imag, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = physical_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + page_half, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() if cutlass.const_expr( @@ -942,18 +1034,45 @@ def kernel( next_page_id_lane0 = cutlass.Int32(0) if warp_idx == self.producer_warp_id: if lane_idx == 0: - next_page_start = page_start + PAGE_TOKENS * self.page_shards + next_page_start = page_start + self.tile_tokens * self.page_shards next_pages_processed = pages_processed + 1 if ( next_page_start < valid_seq_len - and next_pages_processed < self.max_pages + and next_pages_processed < self.max_tiles ): next_page_id_lane0 = ( cutlass.Int32( - page_ids[page_off + page_index + self.page_shards] + page_ids[ + page_off + + (page_index + self.page_shards) * self.pages_per_tile + ] ) // 2 ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + # Same tail-tile clamp as the initial + # prefetch: fall back to the first + # fragment's page. + next_fragment_page_id = next_page_id_lane0 + if ( + next_page_start + fragment * self.box_tokens + < valid_seq_len + ): + next_fragment_page_id = ( + cutlass.Int32( + page_ids[ + page_off + + (page_index + self.page_shards) + * self.pages_per_tile + + fragment + ] + ) + // 2 + ) + 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. @@ -984,90 +1103,112 @@ def kernel( prefetch_next_raw = True prefetched_page_half = page_half + 1 else: - next_page_start = page_start + PAGE_TOKENS * self.page_shards + next_page_start = page_start + self.tile_tokens * self.page_shards next_pages_processed = pages_processed + 1 prefetch_next_raw = ( next_page_start < valid_seq_len - and next_pages_processed < self.max_pages + 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.pages_per_tile): + prefetched_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], + 0, + ) prefetched_page_half = 0 if cutlass.dynamic_expr(prefetch_next_raw): 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): - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - prefetched_page_half, - (kv_head, prefetched_physical_page), - ) - ], - raw_tma_shared_partition_real, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + prefetched_page_half, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) else: - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - prefetched_page_half, - (kv_head, prefetched_physical_page), - ) - ], - raw_tma_shared_partition_real_next, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + prefetched_page_half, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real_next[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() raw_tma_pipeline.producer_acquire(raw_tma_producer_state) if cutlass.dynamic_expr(next_raw_page_buffer == 0): - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - prefetched_page_half, - (kv_head, prefetched_physical_page), - ) - ], - raw_tma_shared_partition_imag, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + prefetched_page_half, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) else: - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - prefetched_page_half, - (kv_head, prefetched_physical_page), - ) - ], - raw_tma_shared_partition_imag_next, - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + prefetched_page_half, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag_next[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) raw_tma_producer_state.advance() frequency = lane_idx # Issue several independent token loads before consuming @@ -1338,7 +1479,7 @@ def kernel( stats_square_sums_m128[stats_head] + stats_delta * stats_delta ) page_index += self.page_shards - page_start += PAGE_TOKENS * self.page_shards + page_start += self.tile_tokens * self.page_shards pages_processed += 1 if warp_idx == self.producer_warp_id: raw_tma_pipeline.producer_tail(raw_tma_producer_state) @@ -1365,13 +1506,13 @@ def kernel( stats_scratch_base = 8 + (stats_warp * N + 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 = pages_processed * PAGE_TOKENS + stats_count_i32 = pages_processed * self.tile_tokens if pages_processed > 0: stats_invalid_prefix = self.score_start - shard_first_page_start if cutlass.dynamic_expr(stats_invalid_prefix > 0): stats_count_i32 = stats_count_i32 - stats_invalid_prefix - stats_last_page = page_start - PAGE_TOKENS * self.page_shards - stats_invalid_tail = stats_last_page + PAGE_TOKENS - valid_seq_len + stats_last_page = page_start - self.tile_tokens * self.page_shards + stats_invalid_tail = stats_last_page + self.tile_tokens - valid_seq_len if cutlass.dynamic_expr(stats_invalid_tail > 0): stats_count_i32 = stats_count_i32 - stats_invalid_tail stats_count = cutlass.Float32(stats_count_i32) @@ -1402,6 +1543,7 @@ def kernel( def _encode_tma_descriptors( layer_pools: list[torch.Tensor], layer_indices: list[int], + tokens_per_block: int, ) -> torch.Tensor: """Encode one immutable feature-first TensorMap per layer index.""" anchor = layer_pools[layer_indices[0]] @@ -1415,14 +1557,17 @@ def _encode_tma_descriptors( 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, tokens_per_block, head_dim = pool.shape - if (kv_factor, tokens_per_block, head_dim) != (2, PAGE_TOKENS, 2 * NUM_FREQS): - raise ValueError("TriAttention CuTe score requires [page, 2, Hkv, 128, 64] pools") + _, kv_factor, num_kv_heads, pool_tokens, head_dim = pool.shape + if (kv_factor, pool_tokens, head_dim) != (2, 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, PAGE_TOKENS] + 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)) @@ -1431,7 +1576,9 @@ def _encode_tma_descriptors( global_dims.append(int(pool.shape[0])) global_strides_bytes.append(s_page * pool.element_size()) tensor_rank = len(global_dims) - box_dims = [NUM_FREQS, CTA_M] + [1] * (tensor_rank - 2) + # One TMA box covers one coefficient plane of one page fragment + # (the whole page for the validated 128-token geometry). + 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), @@ -1532,7 +1679,9 @@ def __init__( dtype=torch.float32, device=output.device, ) - self.descriptors = _encode_tma_descriptors(layer_pools, layer_indices) + self.descriptors = _encode_tma_descriptors( + layer_pools, layer_indices, int(tokens_per_block) + ) self._torch_prefix = ( page_ids, seg_page_off, 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 index 229cf8524d75..2a5c86241c76 100644 --- 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 @@ -10,6 +10,7 @@ not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), reason="TriAttention CuTe kernels require SM100", ) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) @pytest.mark.parametrize( "score_start,valid_lens", [ @@ -22,16 +23,18 @@ ], ) def test_union_fusion_matches_split_pipeline( + tokens_per_block: int, score_start: int, valid_lens: "list | None", monkeypatch: pytest.MonkeyPatch, ) -> None: """The fused pipeline must reproduce the split score->stats->union rows. - Geometry is the fused kernel's contract (128-token pages, GQA group 8, - 32 frequencies). The reference runs the production split path: the score - launch gathers each request's decode window, then ``prepare_union_scores`` - normalizes rows and takes the cross-row union maximum. + Geometry is the fused kernel's contract (32- or 128-token pages, GQA + group 8, 32 frequencies). The reference runs the production split path: + the score launch gathers each request's decode window, then + ``prepare_union_scores`` normalizes rows and takes the cross-row union + maximum. """ pytest.importorskip("cutlass") monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") @@ -44,10 +47,15 @@ def test_union_fusion_matches_split_pipeline( torch.manual_seed(20260721) device = torch.device("cuda") seq_len = 256 - tokens_per_block = 128 num_freqs = 32 num_q_heads = 8 num_pages = seq_len // tokens_per_block + # 32-token pages: one 128-token compute tile spans four pages, so a + # shuffled physical-page table catches any fragment/page mix-up. The + # ragged valid lengths above land mid-tile, exercising the clamped + # tail fragments. + 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) @@ -62,8 +70,8 @@ def test_union_fusion_matches_split_pipeline( mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() - k_plane = [2 * page for page in range(num_pages)] - v_plane = [2 * page + 1 for page in range(num_pages)] + k_plane = [2 * page for page in page_permutation] + v_plane = [2 * page + 1 for page in page_permutation] block_offsets = torch.tensor( [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device ) @@ -160,7 +168,12 @@ def test_union_fusion_matches_split_pipeline( def test_union_fusion_declines_off_contract_geometry( monkeypatch: pytest.MonkeyPatch, ) -> None: - """32-token pages sit outside the fused contract: decline, do not fault.""" + """GQA group 4 sits outside the fused contract: decline, do not fault. + + The split score path pads group-4 head columns up to the MMA tile, but + the fused stats epilogue requires the full GQA group 8, so the runner + must warn and fall back instead of engaging. + """ pytest.importorskip("cutlass") monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") @@ -171,9 +184,9 @@ def test_union_fusion_declines_off_contract_geometry( torch.manual_seed(20260721) device = torch.device("cuda") seq_len = 256 - tokens_per_block = 32 + tokens_per_block = 128 num_freqs = 32 - num_q_heads = 8 + num_q_heads = 4 num_pages = seq_len // tokens_per_block pool = ( 0.125 * torch.randn(num_pages, 2, 1, tokens_per_block, 2 * num_freqs, device=device) From b86594e70585af5afdb64ca061468077b3e55681 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 09:13:15 -0700 Subject: [PATCH 073/178] [None][feat] Extend the fused score kernel to 64 frequencies and GQA group 4 Turn the fused kernel's frequency count, coefficient-plane depth, and GQA group into instance parameters: group-4 heads ride the MMA tile N=8 with zero-padded columns exactly as the production score kernel does, 64-frequency geometries double the coefficient planes and the raw/magnitude K-block loops, and the partial-stats epilogue and union finalize keep padded columns out of the merged statistics. Together with the 32-token-page support this puts the qwen3-8B production geometry (32q/8kv heads, 64 freqs) inside the fused contract. Fused-vs-split equivalence passes across the geometry matrix (32/128-token pages x (32 freqs, group 8)/(64 freqs, group 4) x aligned/non-aligned/ragged windows); the full unit suite is green (136 passed). Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 300 ++++++++++-------- .../triattention_cute_selection.py | 29 +- .../test_triattention_cute_union_fusion.py | 159 ++++------ 3 files changed, 266 insertions(+), 222 deletions(-) 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 index 7df5f7e9c0fd..317d2e9d26ff 100644 --- 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 @@ -49,17 +49,14 @@ def _cute_sqrt_keyword_mode() -> str: _CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() CTA_M = 128 -K = 96 +# Minimum tcgen05 MMA tile N: GQA groups below 8 ride zero-padded head +# columns (see the weight-builder loop and the partial-stats epilogue). N = 8 -NUM_FREQS = 32 THREADS = 256 EPILOGUE_THREADS = 128 RAW_PAGE_BUFFERS = 2 -RAW_K_HALF_ELEMENTS = CTA_M * 2 * NUM_FREQS * RAW_PAGE_BUFFERS RAW_K_VECTOR_ELEMENTS = 8 -RAW_K_SPLIT_PHASE_ELEMENTS = CTA_M * NUM_FREQS -RAW_K_SPLIT_TMA_COPY_BYTES = RAW_K_SPLIT_PHASE_ELEMENTS * (cutlass.BFloat16.width // 8) TMA_DESCRIPTOR_QWORDS = 16 _SUPPORTED_PAGE_SHARDS = (2, 3) @@ -152,12 +149,14 @@ def __init__( super().__init__() if pool_dtype is not cutlass.BFloat16: raise ValueError("TriAttention CuTe score requires BF16 K pages") - if num_freqs != NUM_FREQS: - raise ValueError("TriAttention CuTe score requires 32 frequencies") + 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 num_q_heads % num_kv_heads or num_q_heads // num_kv_heads != N: - raise ValueError("TriAttention CuTe score requires GQA group 8") + 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 not 0 <= score_start < seq_len: raise ValueError("TriAttention CuTe score_start is out of range") if page_shards not in _SUPPORTED_PAGE_SHARDS: @@ -171,6 +170,9 @@ def __init__( 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 128-token compute tile either matches a page exactly (the # validated 128-token geometry, one TMA box per phase) or spans @@ -191,11 +193,11 @@ def __init__( self.use_tma = True self.cpasync_schedule = "sync_each_half" self.split_raw_tma = True - self.raw_tma_feature_extent = NUM_FREQS + self.raw_tma_feature_extent = num_freqs # Barrier transaction bytes for one phase: the full 128-token tile # of one coefficient plane, regardless of how many page fragments # deliver it. - self.raw_tma_copy_bytes = RAW_K_SPLIT_TMA_COPY_BYTES + 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 self.accumulator_pipeline_stages = 1 self.umma_accumulator_partitions = 1 @@ -216,12 +218,12 @@ def __init__( if ( pool_kv_heads != num_kv_heads or pool_tokens != tokens_per_block - or pool_dim != 2 * NUM_FREQS + or 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_slot, self.s_dim = pool_strides - if self.s_slot != 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_slot != 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") @@ -249,7 +251,7 @@ def __call__( ): self.c_dtype = output.element_type self.c_layout = utils.LayoutEnum.COL_MAJOR - self.mma_tiler = (CTA_M, N, K) + self.mma_tiler = (CTA_M, N, self.k_coeff) self.cta_tile_shape_mnk = self.mma_tiler self.epi_tile = (CTA_M, N) @@ -270,20 +272,20 @@ def __call__( self.mma_tiler[:2], ) main_a_shape = ( - (CTA_M, N, NUM_FREQS) + (CTA_M, N, self.num_freqs) if self.main_operand_mode == "bf16_raw_three_term_weight" else self.mma_tiler ) a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, main_a_shape, cutlass.Float32, 1) raw_bf16_a_smem_layout = sm100_utils.make_smem_layout_a( raw_bf16_tiled_mma, - (CTA_M, N, 2 * NUM_FREQS), + (CTA_M, N, 2 * self.num_freqs), cutlass.BFloat16, 1, ) raw_bf16_split_a_smem_layout = sm100_utils.make_smem_layout_a( raw_bf16_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.BFloat16, 2 * RAW_PAGE_BUFFERS, ) @@ -303,7 +305,7 @@ def __call__( ) raw_tma_source_layout = cute.make_layout( ( - 2 * NUM_FREQS, + 2 * self.num_freqs, self.tokens_per_block, (self.num_kv_heads, self.num_physical_pages), ), @@ -325,15 +327,15 @@ def __call__( ) raw_bf16_b_smem_layout = sm100_utils.make_smem_layout_b( raw_bf16_tiled_mma, - (CTA_M, N, 2 * NUM_FREQS), + (CTA_M, N, 2 * self.num_freqs), cutlass.BFloat16, 1, ) - # The magnitude residual has only K=32. A separate compact descriptor - # lets the producer issue its four UMMA steps before the first commit, - # rather than waiting for and overwriting the K=96 main A tile. + # The magnitude residual has only the frequency-count K. A separate + # compact descriptor lets the producer issue its UMMA steps before the + # first commit, rather than waiting for and overwriting the main A tile. magnitude_lo_smem_layout = sm100_utils.make_smem_layout_a( - tiled_mma, (CTA_M, N, NUM_FREQS), cutlass.Float32, 1 + tiled_mma, (CTA_M, N, self.num_freqs), cutlass.Float32, 1 ) magnitude_lo_tiled_mma = sm100_utils.make_trivial_tiled_mma( cutlass.Float16, @@ -345,30 +347,30 @@ def __call__( ) magnitude_lo_fp16_smem_layout = sm100_utils.make_smem_layout_a( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) magnitude_hi_fp16_smem_layout = sm100_utils.make_smem_layout_b( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) magnitude_fp16_b_smem_layout = sm100_utils.make_smem_layout_b( magnitude_lo_tiled_mma, - (CTA_M, N, NUM_FREQS), + (CTA_M, N, self.num_freqs), cutlass.Float16, 1, ) main_b_shape = ( - (CTA_M, N, NUM_FREQS) + (CTA_M, N, self.num_freqs) if self.main_operand_mode == "bf16_raw_three_term_weight" else self.mma_tiler ) @@ -402,12 +404,13 @@ def __call__( a_elements = cute.cosize(a_smem_layout.outer) * int( not self.shared_a_raw_alias and not self.fp16_magnitude_two_term ) - raw_k_elements = RAW_K_HALF_ELEMENTS * int( + raw_k_half_elements = CTA_M * 2 * self.num_freqs * RAW_PAGE_BUFFERS + raw_k_elements = raw_k_half_elements * int( self.k_staging_mode in ("half_page_cpasync", "half_page_tma") and not self.shared_a_raw_alias ) alias_a_elements = cute.cosize(a_smem_layout.outer) * int(self.shared_a_raw_alias) - alias_raw_k_elements = RAW_K_HALF_ELEMENTS * int(self.shared_a_raw_alias) + alias_raw_k_elements = raw_k_half_elements * int(self.shared_a_raw_alias) raw_bf16_a_elements = cute.cosize(raw_bf16_a_smem_layout.outer) * int( self.main_operand_mode == "bf16_raw_three_term_weight" and ( @@ -824,15 +827,24 @@ def kernel( num_threads=EPILOGUE_THREADS, ) cute.arch.mbarrier_init_fence() - for weight_round in cutlass.range_constexpr(N * K // THREADS): + for weight_round in cutlass.range_constexpr(N * self.k_coeff // THREADS): linear_index = tidx + weight_round * THREADS - qg = linear_index // K - feature = linear_index % K - coefficient_kind = feature // NUM_FREQS - frequency = feature % NUM_FREQS - mean_offset = req_id * NUM_FREQS + frequency - q_head = kv_head * self.group_size + qg - calib_offset = (layer_id * self.num_q_heads + q_head) * NUM_FREQS + frequency + 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 + # GQA groups below the minimum MMA tile N=8 ride padded + # columns: they read the group's first head (any valid + # address) and force zero coefficients, so the padded score + # columns come out zero and land in scratch rows the union + # finalizer never reads. + qg_read = qg + if cutlass.const_expr(self.group_size < N): + 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]) @@ -845,6 +857,9 @@ def kernel( value = scale * (qr * msin + qi * mcos) else: value = scale * cutlass.Float32(mlr_coef[calib_offset]) + if cutlass.const_expr(self.group_size < N): + if qg >= self.group_size: + value = cutlass.Float32(0.0) raw_k_block = feature // 16 magnitude_k_block = frequency // 16 if coefficient_kind < 2: @@ -1083,7 +1098,7 @@ def kernel( tcgen05.Field.ACCUMULATE, False, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( raw_bf16_tiled_mma, tCtAcc, @@ -1210,79 +1225,84 @@ def kernel( tma_desc_ptr=raw_tma_descriptor_ptr, ) raw_tma_producer_state.advance() - frequency = lane_idx - # Issue several independent token loads before consuming - # any of them. This bounded RMEM window is unchanged; the - # optional half-page staging only switches its K source - # from global to the single raw shared buffer. - for token_base in cutlass.range( - 0, - CTA_M // (THREADS // 32), - self.prefetch_depth, - unroll_full=not self.compact_token_loop, - ): - 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, + # Each of the 32 lanes 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 + # Issue several independent token loads before consuming + # any of them. This bounded RMEM window is unchanged; the + # optional half-page staging only switches its K source + # from global to the single raw shared buffer. + for token_base in cutlass.range( + 0, + CTA_M // (THREADS // 32), + self.prefetch_depth, + unroll_full=not self.compact_token_loop, + ): + 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 + if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): + magnitude = cute.math.sqrt( + norm2, + approx=self.sqrt_mode == "approx", + ftz=self.magnitude_sqrt_ftz, ) - ] - ) - staged_imag[prefetch_index] = cutlass.Float32( - cpasync_raw_k_0[ - ( - (token, frequency % 16), - 0, - frequency // 16, - raw_imag_stage, + elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): + # cutlass 4.5 renamed the approximate-sqrt control + # to ``fastmath``; map the measured approx choice + # onto it to preserve the authored behavior. + magnitude = cute.math.sqrt( + norm2, fastmath=self.sqrt_mode == "approx" ) - ] - ) - - 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 - if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): - magnitude = cute.math.sqrt( - norm2, - approx=self.sqrt_mode == "approx", - ftz=self.magnitude_sqrt_ftz, + else: + # DSLs with neither spelling get the plain (IEEE) + # sqrt, which is strictly MORE accurate than the + # measured approx choice; the equivalence test's + # tolerance absorbs the difference. + magnitude = cute.math.sqrt(norm2) + magnitude_fp16_0 = cutlass.Float16(magnitude) + magnitude_fp16_1 = cutlass.Float16( + magnitude - cutlass.Float32(magnitude_fp16_0) ) - elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): - # cutlass 4.5 renamed the approximate-sqrt control - # to ``fastmath``; map the measured approx choice - # onto it to preserve the authored behavior. - magnitude = cute.math.sqrt(norm2, fastmath=self.sqrt_mode == "approx") - else: - # DSLs with neither spelling get the plain (IEEE) - # sqrt, which is strictly MORE accurate than the - # measured approx choice; the equivalence test's - # tolerance absorbs the difference. - magnitude = cute.math.sqrt(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 + 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() @@ -1295,8 +1315,8 @@ def kernel( tcgen05.Field.ACCUMULATE, True, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): - imag_b_block = NUM_FREQS // 16 + raw_k_block + 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, @@ -1304,7 +1324,7 @@ def kernel( tCrRawBf16B0[(None, None, imag_b_block, 0)], tCtAcc, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for raw_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( raw_bf16_tiled_mma, tCtAcc, @@ -1312,8 +1332,8 @@ def kernel( tCrRawBf16B1[(None, None, raw_k_block, 0)], tCtAcc, ) - for raw_k_block in cutlass.range_constexpr(NUM_FREQS // 16): - imag_b_block = NUM_FREQS // 16 + raw_k_block + 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, @@ -1329,7 +1349,7 @@ def kernel( tcgen05.Field.ACCUMULATE, True, ) - for magnitude_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( magnitude_lo_tiled_mma, tCtAcc, @@ -1337,7 +1357,7 @@ def kernel( tCrMagnitudeFp16B0[(None, None, magnitude_k_block, 0)], tCtAcc, ) - for magnitude_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( magnitude_lo_tiled_mma, tCtAcc, @@ -1345,7 +1365,7 @@ def kernel( tCrMagnitudeFp16B1[(None, None, magnitude_k_block, 0)], tCtAcc, ) - for magnitude_k_block in cutlass.range_constexpr(NUM_FREQS // 16): + for magnitude_k_block in cutlass.range_constexpr(self.num_freqs // 16): cute.gemm( magnitude_lo_tiled_mma, tCtAcc, @@ -1367,12 +1387,11 @@ def kernel( raw_tma_consumer_state.advance() # Every term multiplying sum_seq must stay 64-bit: with # request*layer segments of seq_len columns the head-plane - # stride alone can exceed 2^31. + # stride alone can exceed 2^31. The scratch head axis is + # padded to the MMA tile N=8 per KV head (group-4 columns + # 4..7 land in padded planes holding zero scores). output_offset = ( - cutlass.Int64(kv_head * self.group_size) * sum_seq - + out_base - + page_start - + page_half * CTA_M + cutlass.Int64(kv_head * N) * sum_seq + out_base + page_start + page_half * CTA_M ) page_output = cute.make_tensor( output.iterator + output_offset, @@ -1381,7 +1400,7 @@ def kernel( stride=( 1, sum_seq, - self.group_size * sum_seq, + N * sum_seq, ), ), ) @@ -1499,7 +1518,11 @@ def kernel( sStats[stats_scratch_base + 1] = stats_square_sum cute.arch.barrier() if warp_idx == 0: - if lane_idx < N: + # Padded head columns (GQA group below the MMA tile N=8) + # carry zero scores; only the real heads' statistics are + # merged and written, in the compact row layout the union + # finalizer reads (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): @@ -1526,7 +1549,7 @@ def kernel( stats_square_sum - stats_sum * stats_sum * inverse_count, cutlass.Float32(0.0), ) - stats_row = task * N + lane_idx + stats_row = task * self.group_size + lane_idx stats_base = (stats_row * self.page_shards + page_shard) * 3 partial_stats[stats_base] = stats_count partial_stats[stats_base + 1] = stats_mean @@ -1543,6 +1566,7 @@ def kernel( def _encode_tma_descriptors( layer_pools: list[torch.Tensor], layer_indices: list[int], + num_freqs: int, tokens_per_block: int, ) -> torch.Tensor: """Encode one immutable feature-first TensorMap per layer index.""" @@ -1558,16 +1582,16 @@ def _encode_tma_descriptors( 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) != (2, tokens_per_block, 2 * NUM_FREQS): + if (kv_factor, pool_tokens, head_dim) != (2, tokens_per_block, 2 * num_freqs): raise ValueError( f"TriAttention CuTe score requires [page, 2, Hkv, {tokens_per_block}, " - f"{2 * NUM_FREQS}] pools" + 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_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)) @@ -1578,7 +1602,7 @@ def _encode_tma_descriptors( tensor_rank = len(global_dims) # One TMA box covers one coefficient plane of one page fragment # (the whole page for the validated 128-token geometry). - box_dims = [NUM_FREQS, min(CTA_M, tokens_per_block)] + [1] * (tensor_rank - 2) + 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), @@ -1588,7 +1612,14 @@ def _encode_tma_descriptors( [uint32(value) for value in box_dims], [uint32(1) for _ in range(tensor_rank)], cuda.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE, - cuda.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B, + # The swizzle must match the smem layout the TMA lands in; the + # sm100 helpers pick it from the inner-row byte count (one + # coefficient plane: num_freqs bf16 elements). + ( + 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, ) @@ -1680,7 +1711,7 @@ def __init__( device=output.device, ) self.descriptors = _encode_tma_descriptors( - layer_pools, layer_indices, int(tokens_per_block) + layer_pools, layer_indices, int(num_freqs), int(tokens_per_block) ) self._torch_prefix = ( page_ids, @@ -1883,6 +1914,11 @@ def __init__( seq_len=seq_len, score_start=score_start, num_q_heads=num_q_heads, + # The score scratch pads each KV head's group + # of head planes to the MMA tile N=8; the + # finalizer maps real head rows onto those + # padded planes (identity for GQA group 8). + num_kv_heads=num_kv_heads, page_shards=page_shards, tokens_per_lane=tokens_per_lane, token_subtiles=token_subtiles, 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 index da08a08669b4..7cadd8575035 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -37,6 +37,9 @@ _STATS_FIELDS = 3 _STD_EPSILON = 1.0e-6 _SUPPORTED_PAGE_SHARDS = (2, 3) +# The fused score kernel pads each KV head's group of score planes up to the +# minimum tcgen05 MMA tile (GQA groups below 8 ride zero-padded columns). +_PADDED_HEAD_COLUMNS = 8 def _select_normalize_union_config( @@ -120,6 +123,7 @@ def __init__( seq_len: int, score_start: int, num_q_heads: int, + num_kv_heads: int | None = None, page_shards: int, tokens_per_lane: int, token_subtiles: int, @@ -131,6 +135,24 @@ def __init__( raise ValueError("TriAttention stats/union score_start is out of range") if page_shards not in _SUPPORTED_PAGE_SHARDS: raise ValueError("TriAttention stats/union has unsupported page shards") + # ``num_kv_heads`` declares that the score scratch pads each KV + # head's group of head planes up to the MMA tile: real head row + # ``q_head`` then lives in scratch plane ``kv * 8 + qg``. Omitting + # it keeps the compact plane-per-head layout (GQA group 8 is + # identical either way). The partial-stats rows are always compact. + if num_kv_heads is None: + self.score_group_size = num_q_heads + self.score_head_pad = 0 + else: + if num_kv_heads <= 0 or num_q_heads % num_kv_heads: + raise ValueError("TriAttention stats/union requires uniform GQA groups") + self.score_group_size = num_q_heads // num_kv_heads + if self.score_group_size > _PADDED_HEAD_COLUMNS: + raise ValueError( + "TriAttention stats/union supports GQA groups up to the " + f"padded head tile ({_PADDED_HEAD_COLUMNS})" + ) + self.score_head_pad = _PADDED_HEAD_COLUMNS - self.score_group_size if tokens_per_lane not in (_SMALL_TOKENS_PER_LANE, _LARGE_TOKENS_PER_LANE): raise ValueError("TriAttention stats/union has unsupported load width") if token_subtiles not in (_SMALL_TOKEN_SUBTILES, _LARGE_TOKEN_SUBTILES): @@ -295,6 +317,11 @@ def kernel( 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 (possibly padded) score plane; + # the padded planes carry zero scores and 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) @@ -330,7 +357,7 @@ def kernel( 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(q_head) * request_count * self.num_layers * self.seq_len + cutlass.Int64(score_plane) * request_count * self.num_layers * self.seq_len + seg_out_offset[segment] + self.score_start + subtile_first_token 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 index 2a5c86241c76..b4a05ce799cf 100644 --- 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 @@ -5,36 +5,25 @@ import pytest import torch - -@pytest.mark.skipif( +_SM100_ONLY = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), reason="TriAttention CuTe kernels require SM100", ) -@pytest.mark.parametrize("tokens_per_block", [32, 128]) -@pytest.mark.parametrize( - "score_start,valid_lens", - [ - # Full-range scoring, both requests full length. - (0, None), - # Non-page-aligned uniform start with ragged valid lengths. - (37, [250, 198]), - # Page-aligned start. - (128, [250, 230]), - ], -) -def test_union_fusion_matches_split_pipeline( + + +def _check_union_fusion_matches_split_pipeline( tokens_per_block: int, + num_freqs: int, + num_q_heads: int, score_start: int, valid_lens: "list | None", monkeypatch: pytest.MonkeyPatch, ) -> None: """The fused pipeline must reproduce the split score->stats->union rows. - Geometry is the fused kernel's contract (32- or 128-token pages, GQA - group 8, 32 frequencies). The reference runs the production split path: - the score launch gathers each request's decode window, then - ``prepare_union_scores`` normalizes rows and takes the cross-row union - maximum. + The reference runs the production split path: the score launch gathers + each request's decode window, then ``prepare_union_scores`` normalizes + rows and takes the cross-row union maximum. """ pytest.importorskip("cutlass") monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") @@ -47,13 +36,11 @@ def test_union_fusion_matches_split_pipeline( torch.manual_seed(20260721) device = torch.device("cuda") seq_len = 256 - num_freqs = 32 - num_q_heads = 8 num_pages = seq_len // tokens_per_block # 32-token pages: one 128-token compute tile spans four pages, so a # shuffled physical-page table catches any fragment/page mix-up. The - # ragged valid lengths above land mid-tile, exercising the clamped - # tail fragments. + # ragged valid lengths land mid-tile, exercising the clamped tail + # fragments. 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 = ( @@ -161,77 +148,71 @@ def test_union_fusion_matches_split_pipeline( ) -@pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), - reason="TriAttention CuTe kernels require SM100", +@_SM100_ONLY +@pytest.mark.parametrize( + "tokens_per_block,num_freqs,num_q_heads,score_start,valid_lens", + [ + # The originally validated geometry: 32 frequencies (64-element K + # rows), GQA group 8, across full-range, non-page-aligned ragged, + # and page-aligned score starts. + (32, 32, 8, 0, None), + (32, 32, 8, 37, [250, 198]), + (32, 32, 8, 128, [250, 230]), + (128, 32, 8, 0, None), + (128, 32, 8, 37, [250, 198]), + (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, 0, None), + (32, 64, 4, 37, [250, 198]), + (128, 64, 4, 0, None), + (128, 64, 4, 128, [250, 230]), + ], ) -def test_union_fusion_declines_off_contract_geometry( +def test_union_fusion_matches_split_pipeline( + tokens_per_block: int, + num_freqs: int, + num_q_heads: int, + score_start: int, + valid_lens: "list | None", monkeypatch: pytest.MonkeyPatch, ) -> None: - """GQA group 4 sits outside the fused contract: decline, do not fault. + _check_union_fusion_matches_split_pipeline( + tokens_per_block, num_freqs, num_q_heads, score_start, valid_lens, monkeypatch + ) - The split score path pads group-4 head columns up to the MMA tile, but - the fused stats epilogue requires the full GQA group 8, so the runner - must warn and fall back instead of engaging. + +@_SM100_ONLY +def test_union_fusion_engages_gqa4_narrow_heads(monkeypatch: pytest.MonkeyPatch) -> None: + """GQA group 4 with 32 frequencies (the formerly declined geometry) engages. + + The fused kernel pads group-4 head columns 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 — so this mixed geometry must launch and match the split path. """ - pytest.importorskip("cutlass") - monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") + _check_union_fusion_matches_split_pipeline(128, 32, 4, 0, None, monkeypatch) - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) - torch.manual_seed(20260721) - device = torch.device("cuda") - seq_len = 256 - tokens_per_block = 128 - num_freqs = 32 - num_q_heads = 4 - num_pages = seq_len // tokens_per_block - 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) - 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) - mean_cos = torch.cos(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) - mean_sin = torch.sin(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) - k_plane = [2 * page for page in range(num_pages)] - v_plane = [2 * page + 1 for page in range(num_pages)] - block_offsets = torch.tensor( - [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device - ) - group = _FixedScoreGroup( - [pool], - [0], - 2, - num_pages, - seq_len, - num_q_heads, - block_offsets, - [0], - q_real, - torch.randn_like(q_real) * 0.125, - torch.randn_like(q_real) * 0.125, - freq_scale_sq, - omega, - offsets, - output_width=seq_len, +def test_union_fusion_rejects_unsupported_frequency_count() -> None: + """16 frequencies (head size 32) sit outside the fused kernel contract.""" + cutlass = pytest.importorskip("cutlass") + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused import ( # noqa: E501 + _TriAttentionScoreKernel, ) - valid_seq_lens = torch.full((2,), seq_len, dtype=torch.int32, device=device) - widths = torch.empty(2, dtype=torch.int32, device=device) - token_starts = torch.zeros(2, dtype=torch.int32, device=device) - union_out = torch.empty((2, seq_len), dtype=torch.float32, device=device) - with pytest.warns(UserWarning, match="union fusion unavailable"): - launched = group.launch_cute_union_fusion( - 2, - valid_seq_lens, - widths, - token_starts, - 0, - mean_cos.contiguous(), - mean_sin.contiguous(), - union_out, + + with pytest.raises(ValueError, match="frequencies"): + _TriAttentionScoreKernel( + num_layers=1, + seq_len=256, + score_start=0, + num_q_heads=8, + num_kv_heads=1, + num_freqs=16, + tokens_per_block=128, + pool_shape=(2, 2, 1, 128, 32), + pool_strides=(8192, 4096, 4096, 32, 1), + pool_dtype=cutlass.BFloat16, + page_shards=3, ) - assert not launched From 3936dafc9a31da8a391d874671b07a5d951c3226 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 09:56:02 -0700 Subject: [PATCH 074/178] [None][feat] Read the fused score window start from per-request metadata The fused pipeline's window start was a compile-time constant, so cohorts with mixed prompt lengths had to fall back to the split launches. The start now rides the per-request metadata exactly like the valid lengths: the kernels load it per segment, the scoring loop, stats domain, and union normalization become per-request, and one compiled runner per geometry serves every cohort. The per-start runner cache and the uniform-prompt decline are gone. Fused-vs-split equivalence passes a 14-case matrix including mixed-prompt cohorts across both page sizes and both geometry families; the full unit suite is green (140 passed). Signed-off-by: tianruih --- .../triattention/triattention.py | 9 -- .../triattention_cute_score_fused.py | 56 ++++----- .../triattention_cute_selection.py | 114 ++++++++++-------- .../triattention/triattention_kernels.py | 58 ++++----- .../test_triattention_cute_union_fusion.py | 46 ++++--- 5 files changed, 146 insertions(+), 137 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 098266647d12..b30090e0b46c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -778,7 +778,6 @@ def __init__( self.stream = None self._score_valid_widths: Optional[torch.Tensor] = None self._score_launcher_bound = False - self.staged_uniform_token_start: Optional[int] = None def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> None: """Bind the per-row score widths for these buffers (mean-only).""" @@ -822,7 +821,6 @@ def launch_prepared_union_fusion(self, union_out: torch.Tensor) -> bool: self.valid_seq_lens_device, self._score_valid_widths, self.token_starts_device, - getattr(self, "staged_uniform_token_start", None), self.mean_cos, self.mean_sin, union_out, @@ -921,13 +919,6 @@ def stage( # integers: a stale-capacity gather is an out-of-bounds index_select # on the device. self.mean_phase_table.ensure(int(max(round_starts)) + 1) - # The fused score+stats+union kernel bakes one global score start at - # compile time, so it only serves cohorts whose pinned prompt lengths - # agree; padded rows are inert (zero valid length) regardless. - first_start = token_starts[0] - self.staged_uniform_token_start = ( - int(first_start) if all(start == first_start for start in token_starts) else None - ) if not self._stage_page_tables_bulk( manager, request_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 index 317d2e9d26ff..5af12cdf13c7 100644 --- 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 @@ -134,7 +134,6 @@ def __init__( *, num_layers: int, seq_len: int, - score_start: int, num_q_heads: int, num_kv_heads: int, num_freqs: int, @@ -157,12 +156,9 @@ def __init__( raise ValueError("TriAttention CuTe score requires 32- or 128-token pages") 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 not 0 <= score_start < seq_len: - raise ValueError("TriAttention CuTe score_start is out of range") if page_shards not in _SUPPORTED_PAGE_SHARDS: raise ValueError("TriAttention CuTe score has unsupported page shards") - self.score_start = score_start self.seq_len = seq_len self.num_layers = num_layers self.num_q_heads = num_q_heads @@ -236,6 +232,7 @@ def __call__( seg_layer_id: cute.Tensor, seg_seq_len: cute.Tensor, seg_out_offset: cute.Tensor, + token_starts: cute.Tensor, q_real: cute.Tensor, q_imag: cute.Tensor, mlr_coef: cute.Tensor, @@ -539,6 +536,7 @@ class SharedStorage: seg_layer_id, seg_seq_len, seg_out_offset, + token_starts, q_real, q_imag, mlr_coef, @@ -581,6 +579,7 @@ def kernel( seg_layer_id: cute.Tensor, seg_seq_len: cute.Tensor, seg_out_offset: cute.Tensor, + token_starts: cute.Tensor, q_real: cute.Tensor, q_imag: cute.Tensor, mlr_coef: cute.Tensor, @@ -616,6 +615,12 @@ def kernel( valid_seq_len = seg_seq_len[segment] page_off = seg_page_off[segment] out_base = seg_out_offset[segment] + # Per-request score window start (the request's pinned prompt + # length), loaded like the other per-segment metadata: each CTA + # owns one segment, so its whole schedule derives from one start. + # Scratch writes stay absolute; only the scoring/stats domain and + # the first scored page move per request. + score_start = cutlass.Int32(token_starts[req_id]) smem = utils.SmemAllocator() storage = smem.allocate(self.shared_storage) @@ -732,7 +737,7 @@ def kernel( swizzle=raw_bf16_b_smem_layout.inner, ) - page_index = self.score_start // self.tile_tokens + page_shard + page_index = score_start // self.tile_tokens + page_shard page_start = page_index * self.tile_tokens shard_first_page_start = page_start pages_processed = cutlass.Int32(0) @@ -754,8 +759,8 @@ def kernel( ) physical_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) prefetched_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) - shard_has_page = valid_seq_len > self.score_start and page_start < valid_seq_len - empty_shard = valid_seq_len <= self.score_start or page_start >= valid_seq_len + shard_has_page = valid_seq_len > score_start and page_start < valid_seq_len + empty_shard = valid_seq_len <= score_start or page_start >= valid_seq_len if cutlass.dynamic_expr(shard_has_page): if warp_idx == self.producer_warp_id: if lane_idx == 0: @@ -966,7 +971,7 @@ def kernel( ) raw_tma_producer_state.advance() while ( - valid_seq_len > self.score_start + valid_seq_len > score_start and page_start < valid_seq_len and pages_processed < self.max_tiles ): @@ -1424,16 +1429,12 @@ def kernel( tTR_rAcc, ) tTR_rC.store(tTR_rAcc.load().to(self.c_dtype)) - if cutlass.const_expr( - self.score_start % CTA_M == 0 and self.seq_len % CTA_M == 0 - ): - cute.copy( - simt_atom, - tTR_rC, - tTR_gC[(None, None, None, subtile_idx)], - ) - elif cutlass.dynamic_expr( - page_start >= self.score_start and page_start + CTA_M <= self.seq_len + # The window start is a per-request runtime value, so + # the tile-interior fast path is a dynamic predicate; + # only the straddling first tile takes the per-token + # branch. + if cutlass.dynamic_expr( + page_start >= score_start and page_start + CTA_M <= self.seq_len ): cute.copy( simt_atom, @@ -1443,7 +1444,7 @@ def kernel( else: output_token = page_start + epilogue_tidx if cutlass.dynamic_expr( - output_token >= self.score_start and output_token < self.seq_len + output_token >= score_start and output_token < self.seq_len ): cute.copy( simt_atom, @@ -1484,7 +1485,7 @@ def kernel( stats_token = page_start + tidx if tidx < EPILOGUE_THREADS: if cutlass.dynamic_expr( - stats_token >= self.score_start and stats_token < valid_seq_len + stats_token >= score_start and stats_token < valid_seq_len ): for stats_head in cutlass.range_constexpr(N): stats_delta = ( @@ -1531,7 +1532,7 @@ def kernel( stats_square_sum = stats_square_sum + sStats[stats_scratch_base + 1] stats_count_i32 = pages_processed * self.tile_tokens if pages_processed > 0: - stats_invalid_prefix = self.score_start - shard_first_page_start + stats_invalid_prefix = score_start - shard_first_page_start if cutlass.dynamic_expr(stats_invalid_prefix > 0): stats_count_i32 = stats_count_i32 - stats_invalid_prefix stats_last_page = page_start - self.tile_tokens * self.page_shards @@ -1669,7 +1670,6 @@ def __init__( max_requests: int, num_layers: int, seq_len: int, - score_start: int, num_q_heads: int, num_kv_heads: int, num_freqs: int, @@ -1680,6 +1680,7 @@ def __init__( seg_layer_id: torch.Tensor, seg_seq_len: torch.Tensor, seg_out_offset: torch.Tensor, + token_starts: torch.Tensor, q_real: torch.Tensor, q_imag: torch.Tensor, mlr_coef: torch.Tensor, @@ -1692,7 +1693,10 @@ def __init__( ) -> None: self.max_requests = int(max_requests) self.num_layers = int(num_layers) - self.width = int(seq_len - score_start) + # The score window start is a per-request runtime input + # (``token_starts``), so the widest window — the whole bucket — + # sizes every start-dependent buffer. + self.width = int(seq_len) self.num_q_heads = int(num_q_heads) self.num_kv_heads = int(num_kv_heads) self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) @@ -1720,6 +1724,7 @@ def __init__( seg_layer_id, seg_seq_len, seg_out_offset, + token_starts, q_real, q_imag, mlr_coef, @@ -1753,12 +1758,12 @@ def __init__( _to_cute(output), _to_cute(seg_seq_len), _to_cute(seg_out_offset), + _to_cute(token_starts), ) static_geometry = ( max_requests, num_layers, seq_len, - score_start, num_q_heads, num_kv_heads, num_freqs, @@ -1792,7 +1797,6 @@ def __init__( kernel = _TriAttentionScoreKernel( num_layers=num_layers, seq_len=seq_len, - score_start=score_start, num_q_heads=num_q_heads, num_kv_heads=num_kv_heads, num_freqs=num_freqs, @@ -1833,7 +1837,6 @@ def __init__( stats_kernel = _TriAttentionScoreKernel( num_layers=num_layers, seq_len=seq_len, - score_start=score_start, num_q_heads=num_q_heads, num_kv_heads=num_kv_heads, num_freqs=num_freqs, @@ -1912,7 +1915,6 @@ def __init__( kernel = _TriAttentionNormalizeUnionKernel( num_layers=num_layers, seq_len=seq_len, - score_start=score_start, num_q_heads=num_q_heads, # The score scratch pads each KV head's group # of head planes to the MMA tile N=8; the 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 index 7cadd8575035..3bbc38761b4e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -121,7 +121,6 @@ def __init__( *, num_layers: int, seq_len: int, - score_start: int, num_q_heads: int, num_kv_heads: int | None = None, page_shards: int, @@ -131,8 +130,6 @@ def __init__( ) -> None: if min(num_layers, seq_len, num_q_heads) <= 0: raise ValueError("TriAttention stats/union reduction requires positive geometry") - if not 0 <= score_start < seq_len: - raise ValueError("TriAttention stats/union score_start is out of range") if page_shards not in _SUPPORTED_PAGE_SHARDS: raise ValueError("TriAttention stats/union has unsupported page shards") # ``num_kv_heads`` declares that the score scratch pads each KV @@ -161,8 +158,10 @@ def __init__( raise ValueError("TriAttention stats/union has unsupported row cluster") self.num_layers = num_layers self.seq_len = seq_len - self.score_start = score_start - self.width = seq_len - score_start + # The score window start is per-request runtime metadata + # (``token_starts``); the widest window — the whole bucket — + # sizes the output rows and the token-tile grid. + self.width = seq_len self.num_q_heads = num_q_heads self.num_rows = num_layers * num_q_heads self.page_shards = page_shards @@ -182,6 +181,7 @@ def __call__( scores: cute.Tensor, seg_seq_len: cute.Tensor, seg_out_offset: cute.Tensor, + token_starts: cute.Tensor, union_scores: cute.Tensor, request_count: cutlass.Int32, stream: cuda.CUstream, @@ -191,6 +191,7 @@ def __call__( scores, seg_seq_len, seg_out_offset, + token_starts, union_scores, request_count, ) @@ -219,6 +220,7 @@ def kernel( scores: cute.Tensor, seg_seq_len: cute.Tensor, seg_out_offset: cute.Tensor, + token_starts: cute.Tensor, union_scores: cute.Tensor, request_count: cutlass.Int32, ): @@ -237,7 +239,10 @@ def kernel( 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 - valid_width = seg_seq_len[first_segment] - self.score_start + # Per-request score window start: the normalization domain and the + # union output row both cover [0, valid - start) for this request. + score_start = cutlass.Int32(token_starts[request_idx]) + valid_width = seg_seq_len[first_segment] - score_start warp_max_ptr = cute.arch.alloc_smem( cutlass.Float32, self.reduce_threads * self.tokens_per_lane * self.token_subtiles, @@ -359,14 +364,17 @@ def kernel( score_index = ( cutlass.Int64(score_plane) * request_count * self.num_layers * self.seq_len + seg_out_offset[segment] - + self.score_start + + score_start + subtile_first_token ) + # The vectorized load also needs the runtime start aligned to + # the lane width (subtile_first_token and the segment stride + # are aligned whenever seq_len is). if cutlass.const_expr( - self.score_start % self.tokens_per_lane == 0 - and self.seq_len % self.tokens_per_lane == 0 + self.seq_len % self.tokens_per_lane == 0 ) and cutlass.dynamic_expr( - subtile_first_token + self.tokens_per_lane <= valid_width + score_start % self.tokens_per_lane == 0 + and subtile_first_token + self.tokens_per_lane <= valid_width ): score_tile = cute.make_tensor( cute.make_ptr( @@ -450,29 +458,33 @@ def kernel( ) reduced_values[token_slot] = union_value subtile_first_token = first_token + token_subtile * self.subtile_token_tile - if cutlass.const_expr(self.width % self.tokens_per_lane == 0): - if cutlass.dynamic_expr( - subtile_first_token + self.tokens_per_lane <= self.width - ): - union_index = request_idx * self.width + subtile_first_token - union_tile = cute.make_tensor( - cute.make_ptr( - cutlass.Float32, - (union_scores.iterator + union_index).toint(), - AddressSpace.gmem, - assumed_align=self.tokens_per_lane * 4, - ), - cute.make_layout(self.tokens_per_lane), - ) - cute.copy( - score_copy_atom, - cute.coalesce(reduced_values), - cute.coalesce(union_tile), - ) + # The output row covers this request's own window, + # [0, valid - start); the straddling subtile falls back + # to the per-token stores. + if cutlass.const_expr( + self.width % self.tokens_per_lane == 0 + ) and cutlass.dynamic_expr( + subtile_first_token + self.tokens_per_lane <= valid_width + ): + union_index = request_idx * self.width + subtile_first_token + union_tile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (union_scores.iterator + union_index).toint(), + AddressSpace.gmem, + assumed_align=self.tokens_per_lane * 4, + ), + cute.make_layout(self.tokens_per_lane), + ) + cute.copy( + score_copy_atom, + cute.coalesce(reduced_values), + cute.coalesce(union_tile), + ) else: for token_slot in cutlass.range_constexpr(self.tokens_per_lane): token = subtile_first_token + token_slot - if cutlass.dynamic_expr(token < self.width): + if cutlass.dynamic_expr(token < valid_width): union_scores[request_idx * self.width + token] = reduced_values[ token_slot ] @@ -513,29 +525,31 @@ def kernel( ) reduced_values[token_slot] = union_value subtile_first_token = first_token + token_subtile * self.subtile_token_tile - if cutlass.const_expr(self.width % self.tokens_per_lane == 0): - if cutlass.dynamic_expr( - subtile_first_token + self.tokens_per_lane <= self.width - ): - union_index = request_idx * self.width + subtile_first_token - union_tile = cute.make_tensor( - cute.make_ptr( - cutlass.Float32, - (union_scores.iterator + union_index).toint(), - AddressSpace.gmem, - assumed_align=self.tokens_per_lane * 4, - ), - cute.make_layout(self.tokens_per_lane), - ) - cute.copy( - score_copy_atom, - cute.coalesce(reduced_values), - cute.coalesce(union_tile), - ) + # Same per-request output domain as the single-CTA path. + if cutlass.const_expr( + self.width % self.tokens_per_lane == 0 + ) and cutlass.dynamic_expr( + subtile_first_token + self.tokens_per_lane <= valid_width + ): + union_index = request_idx * self.width + subtile_first_token + union_tile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (union_scores.iterator + union_index).toint(), + AddressSpace.gmem, + assumed_align=self.tokens_per_lane * 4, + ), + cute.make_layout(self.tokens_per_lane), + ) + cute.copy( + score_copy_atom, + cute.coalesce(reduced_values), + cute.coalesce(union_tile), + ) else: for token_slot in cutlass.range_constexpr(self.tokens_per_lane): token = subtile_first_token + token_slot - if cutlass.dynamic_expr(token < self.width): + if cutlass.dynamic_expr(token < valid_width): union_scores[request_idx * self.width + token] = reduced_values[ token_slot ] diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index da1524680ec6..a9898a0051dc 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -431,12 +431,13 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> self._cute_seg_seq_len = seg_seq_len self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) # Opt-in fused score+stats+union pipeline (Fanrong Li's two-kernel - # scheme). Runners are built lazily per uniform score start because - # the start is a compile-time constant of the fused kernel. + # scheme). ONE runner serves every cohort: the score window start is + # per-request runtime metadata, not a compile-time constant. self._cute_union_fusion_enabled = ( os.environ.get("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "0") == "1" ) - self._union_fusion_runners = {} + self._union_fusion_runner_built = False + self._union_fusion_runner_entry = None from tensorrt_llm.logger import logger logger.info( @@ -545,20 +546,17 @@ def launch( ) return output - def _union_fusion_runner_for( - self, score_start: int, mean_cos: torch.Tensor, mean_sin: torch.Tensor - ): - """Build (or reuse) the fused score/stats/union runner for one start. + def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): + """Build (or reuse) the ONE fused score/stats/union runner. - The fused kernel bakes ``score_start`` at compile time, so one runner - exists per distinct uniform prompt length; the cache is capped so a - prompt-length churn cannot trigger unbounded recompiles. ``None`` - entries record geometries the fused pipeline rejected. + The score window start is per-request runtime metadata staged into a + persistent device buffer, so a single compiled runner serves every + cohort. A ``None`` entry records that the fused pipeline rejected the + geometry. """ - if score_start in self._union_fusion_runners: - return self._union_fusion_runners[score_start] - if len(self._union_fusion_runners) >= 4 or not 0 <= score_start < self.seq_len: - return None + if self._union_fusion_runner_built: + return self._union_fusion_runner_entry + self._union_fusion_runner_built = True num_q_heads, num_kv_heads, num_freqs, tokens_per_block, _ = self.geometry_args entry = None try: @@ -571,18 +569,22 @@ def _union_fusion_runner_for( torch.arange(self.max_requests * self.num_layers, dtype=torch.int64, device=device) * self.seq_len ).to(torch.int32) + # The union output rows are sized by the whole bucket (the widest + # possible window); consumers mask by the per-request widths. union_rows = torch.empty( - (self.max_requests, self.seq_len - score_start), + (self.max_requests, self.seq_len), dtype=torch.float32, device=device, ) + # Per-request score window starts, staged before each launch; the + # compiled kernels capture this buffer's device pointer. + token_starts = torch.zeros(self.max_requests, dtype=torch.int32, device=device) runner = _FusedUnionScoreRunner( layer_pools=self._cute_layer_pools, layer_indices=self._cute_layer_indices, max_requests=self.max_requests, num_layers=self.num_layers, seq_len=self.seq_len, - score_start=score_start, num_q_heads=num_q_heads, num_kv_heads=num_kv_heads, num_freqs=num_freqs, @@ -593,6 +595,7 @@ def _union_fusion_runner_for( seg_layer_id=self.pointer_prefix[5], seg_seq_len=self._cute_seg_seq_len, seg_out_offset=seg_out_offset, + token_starts=token_starts, q_real=self.pointer_middle[0], q_imag=self.pointer_middle[1], mlr_coef=self.pointer_middle[2], @@ -602,12 +605,12 @@ def _union_fusion_runner_for( output=self._cute_scratch, enable_partial_stats=True, ) - entry = (runner, union_rows) + entry = (runner, union_rows, token_starts) except (ImportError, RuntimeError, ValueError, AssertionError) as error: import warnings warnings.warn(f"TriAttention CuTe union fusion unavailable: {error}") - self._union_fusion_runners[score_start] = entry + self._union_fusion_runner_entry = entry return entry def launch_cute_union_fusion( @@ -616,27 +619,27 @@ def launch_cute_union_fusion( valid_seq_lens: torch.Tensor, valid_widths: torch.Tensor, token_starts_device: torch.Tensor, - score_start: Optional[int], mean_cos: torch.Tensor, mean_sin: torch.Tensor, union_out: torch.Tensor, ) -> bool: """Run the fused score+stats+normalized-union pipeline when possible. - Returns False (without launching) whenever the opt-in is off, the - cohort's prompt lengths are not uniform, or the geometry has no fused - specialization — the caller then falls back to the split score, - row-stats, and union path. + Each request scores its own window (``token_starts_device`` carries + the per-request pinned prompt lengths), so mixed-prompt cohorts are + served directly. Returns False (without launching) only when the + opt-in is off or the geometry has no fused specialization — the + caller then falls back to the split score, row-stats, and union path. """ self.prepare_cute_score(mean_cos, mean_sin) - if not self._cute_union_fusion_enabled or score_start is None: + if not self._cute_union_fusion_enabled: return False if request_count <= 0 or request_count > self.max_requests: raise ValueError("request count exceeds fixed score capacity") - entry = self._union_fusion_runner_for(int(score_start), mean_cos, mean_sin) + entry = self._union_fusion_runner(mean_cos, mean_sin) if entry is None: return False - runner, union_rows = entry + runner, union_rows, staged_token_starts = entry if not runner.supports_union_fusion(request_count): return False num_segments = request_count * self.num_layers @@ -651,6 +654,7 @@ def launch_cute_union_fusion( self.pointer_prefix[4][:num_segments], out=self._cute_seg_seq_len[:num_segments], ) + staged_token_starts[:request_count].copy_(token_starts_device[:request_count]) runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_rows[:request_count]) columns = min(union_rows.shape[1], union_out.shape[1]) union_out[:request_count, :columns].copy_(union_rows[:request_count, :columns]) 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 index b4a05ce799cf..654413f55d3a 100644 --- 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 @@ -15,15 +15,17 @@ def _check_union_fusion_matches_split_pipeline( tokens_per_block: int, num_freqs: int, num_q_heads: int, - score_start: int, + score_starts: "int | list", valid_lens: "list | None", monkeypatch: pytest.MonkeyPatch, ) -> None: """The fused pipeline must reproduce the split score->stats->union rows. - The reference runs the production split path: the score launch gathers - each request's decode window, then ``prepare_union_scores`` normalizes - rows and takes the cross-row union maximum. + ``score_starts`` is either one uniform window start or a per-request + list (the fused kernels read the start per request at runtime). The + reference runs the production split path: the score launch gathers each + request's decode window, then ``prepare_union_scores`` normalizes rows + and takes the cross-row union maximum. """ pytest.importorskip("cutlass") monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") @@ -83,10 +85,13 @@ def _check_union_fusion_matches_split_pipeline( valid_lens = [seq_len, seq_len] valid_seq_lens = 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 split pipeline over the same decode windows. split_widths = torch.empty(request_count, dtype=torch.int32, device=device) - token_starts = torch.full((request_count,), score_start, dtype=torch.int32, device=device) + token_starts = torch.tensor(score_starts, dtype=torch.int32, device=device) per_head = group.launch( request_count, valid_seq_lens, @@ -119,7 +124,6 @@ def _check_union_fusion_matches_split_pipeline( valid_seq_lens, fused_widths, token_starts, - score_start, mean_cos, mean_sin, fused_out, @@ -127,7 +131,7 @@ def _check_union_fusion_matches_split_pipeline( assert launched, "fused union pipeline must engage on its contract geometry" assert torch.equal(fused_widths, split_widths) for request in range(request_count): - width = int(valid_lens[request]) - score_start + width = int(valid_lens[request]) - int(score_starts[request]) torch.testing.assert_close( fused_out[request, :width], expected[request, :width], @@ -135,26 +139,14 @@ def _check_union_fusion_matches_split_pipeline( atol=5.0e-3, ) - # A cohort without one uniform prompt start must decline, not launch. - assert not group.launch_cute_union_fusion( - request_count, - valid_seq_lens, - fused_widths, - token_starts, - None, - mean_cos, - mean_sin, - fused_out, - ) - @_SM100_ONLY @pytest.mark.parametrize( - "tokens_per_block,num_freqs,num_q_heads,score_start,valid_lens", + "tokens_per_block,num_freqs,num_q_heads,score_starts,valid_lens", [ # The originally validated geometry: 32 frequencies (64-element K # rows), GQA group 8, across full-range, non-page-aligned ragged, - # and page-aligned score starts. + # and page-aligned uniform window starts. (32, 32, 8, 0, None), (32, 32, 8, 37, [250, 198]), (32, 32, 8, 128, [250, 230]), @@ -167,18 +159,25 @@ def _check_union_fusion_matches_split_pipeline( (32, 64, 4, 37, [250, 198]), (128, 64, 4, 0, None), (128, 64, 4, 128, [250, 230]), + # Mixed-prompt cohorts: each request scores its own window (one + # start mid-tile, one page-aligned) — the case the fused pipeline + # previously declined. + (32, 32, 8, [37, 128], [250, 198]), + (128, 32, 8, [37, 128], None), + (32, 64, 4, [37, 128], None), + (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_start: int, + score_starts: "int | list", valid_lens: "list | None", monkeypatch: pytest.MonkeyPatch, ) -> None: _check_union_fusion_matches_split_pipeline( - tokens_per_block, num_freqs, num_q_heads, score_start, valid_lens, monkeypatch + tokens_per_block, num_freqs, num_q_heads, score_starts, valid_lens, monkeypatch ) @@ -206,7 +205,6 @@ def test_union_fusion_rejects_unsupported_frequency_count() -> None: _TriAttentionScoreKernel( num_layers=1, seq_len=256, - score_start=0, num_q_heads=8, num_kv_heads=1, num_freqs=16, From 623df3d4c22f01fe30fb2a5c27c38fde33835ab8 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 10:28:28 -0700 Subject: [PATCH 075/178] [None][feat] Retire the split union scoring path The fused score+stats+union pipeline is now the only union path, completing the single-path contract: the split union reduction (_score_union_kernel, prepare_union_scores), the union selector's split branch and its normalization buffers, and the opt-in environment gate are deleted; a fused setup failure now raises instead of warning and falling back. _score_row_stats_kernel stays: the per-head and per-layer-per-head selection paths still normalize with it. Union with normalize_scores=False is rejected at construction (the fused pipeline always z-normalizes; the un-normalized variant died with the split path). The retired kernels live on as reference copies in the fused-pipeline equivalence tests, which pass a 16-case matrix (mixed prompts, both page sizes, both geometry families) plus a torch-oracle check of the references themselves; the full unit suite is green (140 passed). Signed-off-by: tianruih --- .../triattention/triattention.py | 85 +++--- .../triattention/triattention_kernels.py | 153 +++-------- tensorrt_llm/llmapi/llm_args.py | 3 +- .../test_triattention_cute_union_fusion.py | 256 +++++++++++++++++- .../test_triattention_fused_settle_pack.py | 2 - .../test_triattention_pipeline.py | 18 +- .../test_triattention_selection_compaction.py | 101 ++----- 7 files changed, 347 insertions(+), 271 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b30090e0b46c..4e14d5105b6a 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -321,7 +321,15 @@ def _select_top_tokens(self) -> None: class _BatchedUnionKeepSetSelector(_BatchedKeepSetSelectorBase): - """Persistent ``[request, ...]`` buffers for union selection.""" + """Persistent ``[request, ...]`` buffers for union selection. + + The fused score+stats+union CuTe pipeline is THE union score producer: + it writes the normalized per-request union rows straight into + ``combined``, so this selector owns only the top-k settle-and-pack + stage. The split row-stats/union-reduce Triton launches were retired + with their buffers; their standalone copies live in the fused-pipeline + unit test as references. + """ def __init__( self, @@ -335,8 +343,6 @@ def __init__( dense_layers: Tuple[int, ...] = (), num_query_heads: int = 0, num_kv_heads: int = 0, - input_scores: Optional[torch.Tensor] = None, - normalize_scores: bool = True, prompt_offsets_buffer: Optional[torch.Tensor] = None, ) -> None: if rows <= 0: @@ -353,11 +359,6 @@ def __init__( device=device, max_requests=max_requests, ) - if input_scores is None: - raise ValueError("union selection requires its fixed score input") - - self.row_mean = torch.empty((max_requests, rows, 1), dtype=dtype, device=self.device) - self.row_std = torch.empty_like(self.row_mean) self.combined = torch.empty((max_requests, width), dtype=dtype, device=self.device) self.final_indices = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=self.device @@ -368,24 +369,6 @@ def __init__( (max_requests, self.keep_count), dtype=torch.int32, device=self.device ) self._bind_selection_rows(self.combined, self.valid_widths, self.final_indices, self.keep) - # Callers select from exactly this tensor with exactly this flag. - self.input_scores = input_scores - self.normalize_scores = bool(normalize_scores) - - def select_prepared_requests(self) -> None: - """Select from the CUDA score tensor bound to this fixed selector.""" - from .triattention_kernels import prepare_union_scores - - prepare_union_scores( - self.input_scores, - self.valid_widths, - self.row_mean, - self.row_std, - self.combined, - self.max_requests, - normalize_scores=self.normalize_scores, - ) - self._select_top_tokens() def select_prepared_union_scores(self) -> None: """Select from normalized union rows already written into ``combined``. @@ -791,12 +774,12 @@ def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> N self._score_valid_widths = valid_widths self._score_launcher_bound = True - def launch_prepared_union_fusion(self, union_out: torch.Tensor) -> bool: - """Try the fused score+stats+union pipeline over these buffers. + def launch_prepared_union_fusion(self, union_out: torch.Tensor) -> None: + """Run the fused score+stats+union pipeline over these buffers. - Returns False without side effects on the selection buffers when the - fused path cannot serve this cohort; the caller then runs the split - score and selection launches instead. + This is THE union score path; there is deliberately no fallback. A + geometry or capacity the fused pipeline cannot serve raises loudly + instead of routing to the retired split launches. """ if not self._score_launcher_bound: raise RuntimeError("TriAttention score launcher is not bound") @@ -1132,6 +1115,14 @@ def __init__( "'union', 'per_head', 'per_layer_perhead'" ) self.normalize_scores = bool(normalize_scores) + if self.eviction_mode == "union" and not self.normalize_scores: + # The fused score+stats+union CuTe pipeline is THE union path and + # always z-normalizes; the split un-normalized union launches + # were retired with the Triton/C++ score stacks. + raise ValueError( + "TriAttention union eviction requires normalize_scores=True: " + "the fused union pipeline always z-normalizes score rows" + ) self.pin_prefill = bool(pin_prefill) # cpt=False (default): budget counts DECODE tokens only (pinned prompt is # extra). cpt=True: budget INCLUDES the pinned prompt. @@ -1573,8 +1564,6 @@ def _configured_protected_tail_capacity(self) -> int: def _build_cross_request_keep_set_selector( plan: _CrossRequestSelectionPlan, *, - input_scores: Optional[torch.Tensor] = None, - normalize_scores: bool = True, prompt_offsets_buffer: Optional[torch.Tensor] = None, ) -> Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector]: """Allocate one fixed ``[request, ...]`` keep-set selector.""" @@ -1589,8 +1578,6 @@ def _build_cross_request_keep_set_selector( dense_layers=plan.dense_layers, num_query_heads=plan.num_query_heads, num_kv_heads=plan.num_kv_heads, - input_scores=input_scores, - normalize_scores=normalize_scores, prompt_offsets_buffer=prompt_offsets_buffer, ) return _BatchedPerHeadKeepSetSelector( @@ -2066,12 +2053,6 @@ def _fixed_resources_for( device=first_pool.device, max_requests=request_capacity, ), - input_scores=score_staging.fused_group.output.view( - request_capacity, - len(layout.dense_layers) * int(self._H), - decode_width, - ), - normalize_scores=self.normalize_scores, prompt_offsets_buffer=score_staging.token_starts_device, ) # Padded rows carry zero valid width; their provisional TopK entries @@ -2321,21 +2302,21 @@ def _evict_requests( keep_set_selector.refresh_row_prompt_offsets() try: + fused_union = isinstance(keep_set_selector, _BatchedUnionKeepSetSelector) with nvtx_range("triattention.score", color="blue"): - # The fused pipeline covers score, row stats, normalization, - # and the cross-row union maximum in two kernels; when it - # declines this cohort the split launches run instead. - fused_union = ( - self.normalize_scores - and isinstance(keep_set_selector, _BatchedUnionKeepSetSelector) - and score_staging.launch_prepared_union_fusion(keep_set_selector.combined) - ) - per_head = None if fused_union else score_staging.launch_prepared_score() + # Union rounds run the fused pipeline (score, row stats, + # normalization, and the cross-row union maximum in two CuTe + # kernels) — THE only union path; it raises loudly if it + # cannot serve this cohort. The per-head modes run the split + # score launch and their own reduction kernels. + if fused_union: + score_staging.launch_prepared_union_fusion(keep_set_selector.combined) + per_head = None + else: + per_head = score_staging.launch_prepared_score() with nvtx_range("triattention.select", color="yellow"): if fused_union: keep_set_selector.select_prepared_union_scores() - elif isinstance(keep_set_selector, _BatchedUnionKeepSetSelector): - keep_set_selector.select_prepared_requests() else: keep_set_selector.select_requests( per_head, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index a9898a0051dc..59a035a07867 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -9,9 +9,15 @@ 64/128, 32/128-token pages, GQA group 4 or 8. There is deliberately no other score path -- any geometry outside that contract raises loudly at setup instead of routing to a slower kernel (the original Triton score kernel and -the C++ CUDA score stack have both been deleted). The unit tests validate -the CuTe kernel against an independent PyTorch oracle. Selection and -compaction live in their respective runtime modules. +the C++ CUDA score stack have both been deleted). Union eviction likewise +runs EXCLUSIVELY through the fused score+stats+union CuTe pipeline +(``triattention_cute_score_fused.py`` + ``triattention_cute_selection.py``); +the split Triton row-stats/union-reduce launches were retired with it, and +their standalone copies live in the fused-pipeline unit test as references +(the row-stats kernel itself remains: the per-head modes still normalize +with it). The unit tests validate the CuTe kernels against independent +PyTorch oracles. Selection and compaction live in their respective runtime +modules. House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). @@ -22,7 +28,6 @@ from __future__ import annotations -import os from typing import List, Optional import torch @@ -430,13 +435,10 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> self._cute_scratch = scratch self._cute_seg_seq_len = seg_seq_len self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) - # Opt-in fused score+stats+union pipeline (Fanrong Li's two-kernel - # scheme). ONE runner serves every cohort: the score window start is - # per-request runtime metadata, not a compile-time constant. - self._cute_union_fusion_enabled = ( - os.environ.get("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "0") == "1" - ) - self._union_fusion_runner_built = False + # Fused score+stats+union pipeline (Fanrong Li's two-kernel scheme): + # THE union path, built lazily on the first union launch. ONE runner + # serves every cohort: the score window start is per-request runtime + # metadata, not a compile-time constant. self._union_fusion_runner_entry = None from tensorrt_llm.logger import logger @@ -551,14 +553,12 @@ def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): The score window start is per-request runtime metadata staged into a persistent device buffer, so a single compiled runner serves every - cohort. A ``None`` entry records that the fused pipeline rejected the - geometry. + cohort. This is THE union path: a construction failure raises loudly + instead of recording a fallback. """ - if self._union_fusion_runner_built: + if self._union_fusion_runner_entry is not None: return self._union_fusion_runner_entry - self._union_fusion_runner_built = True num_q_heads, num_kv_heads, num_freqs, tokens_per_block, _ = self.geometry_args - entry = None try: from .triattention_cute_score_fused import ( TriAttentionCuteScoreRunner as _FusedUnionScoreRunner, @@ -605,13 +605,12 @@ def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): output=self._cute_scratch, enable_partial_stats=True, ) - entry = (runner, union_rows, token_starts) except (ImportError, RuntimeError, ValueError, AssertionError) as error: - import warnings - - warnings.warn(f"TriAttention CuTe union fusion unavailable: {error}") - self._union_fusion_runner_entry = entry - return entry + raise RuntimeError( + "TriAttention CuTe union fusion setup failed and no other union path exists" + ) from error + self._union_fusion_runner_entry = (runner, union_rows, token_starts) + return self._union_fusion_runner_entry def launch_cute_union_fusion( self, @@ -622,26 +621,25 @@ def launch_cute_union_fusion( mean_cos: torch.Tensor, mean_sin: torch.Tensor, union_out: torch.Tensor, - ) -> bool: - """Run the fused score+stats+normalized-union pipeline when possible. + ) -> None: + """Run the fused score+stats+normalized-union pipeline (THE union path). Each request scores its own window (``token_starts_device`` carries the per-request pinned prompt lengths), so mixed-prompt cohorts are - served directly. Returns False (without launching) only when the - opt-in is off or the geometry has no fused specialization — the - caller then falls back to the split score, row-stats, and union path. + served directly. There is deliberately no fallback: an unsupported + geometry or request count raises loudly instead of routing to the + retired split score/row-stats/union launches. """ self.prepare_cute_score(mean_cos, mean_sin) - if not self._cute_union_fusion_enabled: - return False if request_count <= 0 or request_count > self.max_requests: raise ValueError("request count exceeds fixed score capacity") - entry = self._union_fusion_runner(mean_cos, mean_sin) - if entry is None: - return False - runner, union_rows, staged_token_starts = entry + runner, union_rows, staged_token_starts = self._union_fusion_runner(mean_cos, mean_sin) if not runner.supports_union_fusion(request_count): - return False + raise RuntimeError( + f"TriAttention CuTe union fusion has no compiled variant for " + f"request_count={request_count} (capacity {self.max_requests}) " + "and no other union path exists" + ) num_segments = request_count * self.num_layers torch.sub( valid_seq_lens[:request_count], @@ -658,7 +656,6 @@ def launch_cute_union_fusion( runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_rows[:request_count]) columns = min(union_rows.shape[1], union_out.shape[1]) union_out[:request_count, :columns].copy_(union_rows[:request_count, :columns]) - return True # --------------------------------------------------------------------------- # @@ -701,92 +698,6 @@ def _score_row_stats_kernel( tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) -@triton.jit -def _score_union_kernel( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - ROWS: tl.constexpr, - WIDTH: tl.constexpr, - NORMALIZE: tl.constexpr, - BLOCK: tl.constexpr, -): - """Normalize score rows and reduce them directly to one request-level union.""" - request = tl.program_id(0) - token = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - valid_width = tl.load(valid_widths + request) - valid_token = token < valid_width - union_max = tl.full((BLOCK,), -float("inf"), tl.float32) - for row in tl.range(0, ROWS): - flat_row = request * ROWS + row - value = tl.load( - scores + flat_row * WIDTH + 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")) - union_max = tl.maximum(union_max, value) - tl.store(combined + request * WIDTH + token, union_max, mask=token < WIDTH) - - -def prepare_union_scores( - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - combined: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, -) -> None: - """Mask, normalize, and union-reduce score rows in two or three launches.""" - request_count = int(request_count) - if not scores.is_cuda or scores.ndim != 3 or scores.dtype != torch.float32: - raise ValueError("union score preparation requires contiguous CUDA FP32 rows") - if not scores.is_contiguous() or request_count != scores.shape[0]: - raise ValueError("union score preparation request geometry does not match") - _, rows, width = scores.shape - if ( - valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.numel() < request_count * rows - or row_inv_std.shape != row_mean.shape - or combined.shape != (request_count, width) - ): - raise ValueError("union score preparation buffers do not match") - stats_block = 256 - if normalize_scores: - _score_row_stats_kernel[(request_count * rows,)]( - scores, - valid_widths, - row_mean, - row_inv_std, - ROWS=rows, - WIDTH=width, - BLOCK=stats_block, - num_warps=4, - ) - union_block = 32 - _score_union_kernel[(request_count, triton.cdiv(width, union_block))]( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - ROWS=rows, - WIDTH=width, - NORMALIZE=normalize_scores, - BLOCK=union_block, - num_warps=1, - ) - - @triton.jit def _score_per_head_reduce_kernel( scores, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9bd95c2f57ee..c241d123861a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3450,7 +3450,8 @@ class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): normalize_scores: bool = Field( default=True, description="Z-normalize each head's scores over the decode region " - "before selection (upstream default).") + "before selection (upstream default). `union` eviction requires True: " + "its fused score+stats+union pipeline always normalizes.") pin_prefill: bool = Field( default=True, description="Always preserve the prompt (prefill) tokens; only decode " 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 index 654413f55d3a..79d23da42225 100644 --- 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 @@ -4,12 +4,132 @@ import pytest import torch +import triton +import triton.language as tl _SM100_ONLY = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), reason="TriAttention CuTe kernels require SM100", ) +# --------------------------------------------------------------------------- # +# Standalone reference copies of the RETIRED split-union launches. The fused +# score+stats+union CuTe pipeline is THE production union path; these +# pre-retirement Triton copies exist only as the equivalence references here +# (precedent: the standalone settle/pack copies in +# test_triattention_fused_settle_pack.py). +# --------------------------------------------------------------------------- # + + +@triton.jit +def _reference_score_row_stats_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + """Compute one valid-prefix mean and inverse standard deviation per score row.""" + flat_row = tl.program_id(0) + request = flat_row // ROWS + valid_width = tl.load(valid_widths + request) + score_row = scores + flat_row * WIDTH + lane = tl.arange(0, BLOCK) + score_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < valid_width + value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) + score_sum += tl.sum(value, axis=0) + mean = score_sum / valid_width + square_sum = 0.0 + for start in tl.static_range(0, WIDTH, BLOCK): + token = start + lane + valid = token < valid_width + 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 / valid_width) + tl.store(row_mean + flat_row, mean) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) + + +@triton.jit +def _reference_score_union_kernel( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, + NORMALIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + """Normalize score rows and reduce them directly to one request-level union.""" + request = tl.program_id(0) + token = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + valid_width = tl.load(valid_widths + request) + valid_token = token < valid_width + union_max = tl.full((BLOCK,), -float("inf"), tl.float32) + for row in tl.range(0, ROWS): + flat_row = request * ROWS + row + value = tl.load( + scores + flat_row * WIDTH + 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")) + union_max = tl.maximum(union_max, value) + tl.store(combined + request * WIDTH + token, union_max, mask=token < WIDTH) + + +def _reference_prepare_union_scores( + scores: torch.Tensor, + valid_widths: torch.Tensor, + row_mean: torch.Tensor, + row_inv_std: torch.Tensor, + combined: torch.Tensor, + request_count: int, + *, + normalize_scores: bool, +) -> None: + """Mask, normalize, and union-reduce score rows in one or two launches.""" + request_count = int(request_count) + assert scores.is_cuda and scores.ndim == 3 and scores.dtype == torch.float32 + assert scores.is_contiguous() and request_count == scores.shape[0] + _, rows, width = scores.shape + stats_block = 256 + if normalize_scores: + _reference_score_row_stats_kernel[(request_count * rows,)]( + scores, + valid_widths, + row_mean, + row_inv_std, + ROWS=rows, + WIDTH=width, + BLOCK=stats_block, + num_warps=4, + ) + union_block = 32 + _reference_score_union_kernel[(request_count, triton.cdiv(width, union_block))]( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + ROWS=rows, + WIDTH=width, + NORMALIZE=normalize_scores, + BLOCK=union_block, + num_warps=1, + ) + def _check_union_fusion_matches_split_pipeline( tokens_per_block: int, @@ -17,22 +137,20 @@ def _check_union_fusion_matches_split_pipeline( num_q_heads: int, score_starts: "int | list", valid_lens: "list | None", - monkeypatch: pytest.MonkeyPatch, ) -> None: """The fused pipeline must reproduce the split score->stats->union rows. ``score_starts`` is either one uniform window start or a per-request list (the fused kernels read the start per request at runtime). The - reference runs the production split path: the score launch gathers each - request's decode window, then ``prepare_union_scores`` normalizes rows - and takes the cross-row union maximum. + reference runs the retired split path: the score launch gathers each + request's decode window, then the standalone + ``_reference_prepare_union_scores`` copy normalizes rows and takes the + cross-row union maximum. """ pytest.importorskip("cutlass") - monkeypatch.setenv("TRTLLM_TRIATTENTION_CUTE_UNION_FUSION", "1") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _FixedScoreGroup, - prepare_union_scores, ) torch.manual_seed(20260721) @@ -105,7 +223,7 @@ def _check_union_fusion_matches_split_pipeline( row_mean = torch.empty((request_count, rows, 1), dtype=torch.float32, device=device) row_inv_std = torch.empty_like(row_mean) expected = torch.empty((request_count, seq_len), dtype=torch.float32, device=device) - prepare_union_scores( + _reference_prepare_union_scores( scores_rows, split_widths, row_mean, @@ -119,7 +237,7 @@ def _check_union_fusion_matches_split_pipeline( fused_out = torch.full( (request_count, seq_len), float("nan"), dtype=torch.float32, device=device ) - launched = group.launch_cute_union_fusion( + group.launch_cute_union_fusion( request_count, valid_seq_lens, fused_widths, @@ -128,7 +246,6 @@ def _check_union_fusion_matches_split_pipeline( mean_sin, fused_out, ) - assert launched, "fused union pipeline must engage on its contract geometry" assert torch.equal(fused_widths, split_widths) for request in range(request_count): width = int(valid_lens[request]) - int(score_starts[request]) @@ -174,15 +291,14 @@ def test_union_fusion_matches_split_pipeline( num_q_heads: int, score_starts: "int | list", valid_lens: "list | None", - monkeypatch: pytest.MonkeyPatch, ) -> None: _check_union_fusion_matches_split_pipeline( - tokens_per_block, num_freqs, num_q_heads, score_starts, valid_lens, monkeypatch + tokens_per_block, num_freqs, num_q_heads, score_starts, valid_lens ) @_SM100_ONLY -def test_union_fusion_engages_gqa4_narrow_heads(monkeypatch: pytest.MonkeyPatch) -> None: +def test_union_fusion_engages_gqa4_narrow_heads() -> None: """GQA group 4 with 32 frequencies (the formerly declined geometry) engages. The fused kernel pads group-4 head columns up to the MMA tile N=8 with @@ -190,7 +306,121 @@ def test_union_fusion_engages_gqa4_narrow_heads(monkeypatch: pytest.MonkeyPatch) rows, and the union finalizer maps head rows onto the padded score planes — so this mixed geometry must launch and match the split path. """ - _check_union_fusion_matches_split_pipeline(128, 32, 4, 0, None, monkeypatch) + _check_union_fusion_matches_split_pipeline(128, 32, 4, 0, None) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_reference_union_preparation_matches_ragged_torch_reference() -> None: + """The standalone reference copy must match a pure-torch oracle. + + Moved from the selection/compaction suite when the production + ``prepare_union_scores`` was retired: this keeps the reference copy the + fused-pipeline equivalence tests compare against honest. + """ + device = torch.device("cuda", torch.cuda.current_device()) + request_count, rows, width = 2, 7, 97 + generator = torch.Generator(device=device).manual_seed(17) + scores = torch.randn( + request_count, + rows, + width, + generator=generator, + dtype=torch.float32, + device=device, + ) + valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) + row_mean = torch.empty(request_count, rows, 1, dtype=torch.float32, device=device) + row_inv_std = torch.empty_like(row_mean) + combined = torch.empty(request_count, width, device=device) + + _reference_prepare_union_scores( + scores, + valid_widths, + row_mean, + row_inv_std, + combined, + request_count, + normalize_scores=True, + ) + torch.cuda.synchronize(device) + + expected = torch.full_like(combined, float("-inf")) + for request, valid_width in enumerate(valid_widths.tolist()): + valid_scores = scores[request, :, :valid_width] + mean = valid_scores.mean(dim=1, keepdim=True) + std = torch.linalg.vector_norm(valid_scores - mean, dim=1, keepdim=True) + std = (std / valid_width**0.5).clamp_min(1e-6) + expected[request, :valid_width] = ((valid_scores - mean) / std).amax(dim=0) + assert torch.allclose(combined, expected, rtol=2e-5, atol=2e-5) + + +@_SM100_ONLY +def test_union_fusion_setup_failure_raises(monkeypatch: pytest.MonkeyPatch) -> None: + """A fused-runner construction failure raises loudly: no fallback remains.""" + pytest.importorskip("cutlass") + + import tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused as fused_module # noqa: E501 + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + ) + + torch.manual_seed(20260721) + device = torch.device("cuda") + seq_len = 256 + tokens_per_block = 128 + num_freqs = 32 + num_q_heads = 8 + num_pages = seq_len // tokens_per_block + 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) + 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) + mean_cos = torch.cos(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) + mean_sin = torch.sin(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) + k_plane = [2 * page for page in range(num_pages)] + v_plane = [2 * page + 1 for page in range(num_pages)] + block_offsets = torch.tensor( + [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device + ) + group = _FixedScoreGroup( + [pool], + [0], + 2, + num_pages, + seq_len, + num_q_heads, + block_offsets, + [0], + q_real, + torch.randn_like(q_real) * 0.125, + torch.randn_like(q_real) * 0.125, + freq_scale_sq, + omega, + offsets, + output_width=seq_len, + ) + + def _refuse_construction(**_kwargs): + raise ValueError("synthetic fused-runner construction failure") + + monkeypatch.setattr(fused_module, "TriAttentionCuteScoreRunner", _refuse_construction) + valid_seq_lens = torch.full((2,), seq_len, dtype=torch.int32, device=device) + widths = torch.empty(2, dtype=torch.int32, device=device) + token_starts = torch.zeros(2, dtype=torch.int32, device=device) + union_out = torch.empty((2, seq_len), dtype=torch.float32, device=device) + with pytest.raises(RuntimeError, match="no other union path exists"): + group.launch_cute_union_fusion( + 2, + valid_seq_lens, + widths, + token_starts, + mean_cos.contiguous(), + mean_sin.contiguous(), + union_out, + ) def test_union_fusion_rejects_unsupported_frequency_count() -> None: 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 index 55aeabb4b5e1..5a0a365d1c1b 100644 --- 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 @@ -444,8 +444,6 @@ def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buff dtype=torch.float32, device=device, max_requests=request_count, - input_scores=torch.zeros(request_count, 3, width, device=device), - normalize_scores=False, ) def build_compaction(kept_token_ordinals): diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 8e9f983eb198..7f02596cff33 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -834,10 +834,11 @@ def test_cross_request_union_matches_oracle_at_high_keep_counts(self, keep_count dtype=scores.dtype, device=device, max_requests=len(request_scores), - input_scores=scores, - normalize_scores=False, ) - selector.select_prepared_requests() + # The fused CuTe pipeline is the production union-row producer; this + # top-k routing test stages the prepared rows into ``combined``. + selector.combined.copy_(scores.amax(dim=1)) + selector.select_prepared_union_scores() selected = selector.keep.cpu() for actual, expected_keep in zip(selected, expected): @@ -850,6 +851,16 @@ class TestFixedScoreMetadata: def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize_scores): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module + if eviction_mode == "union" and not normalize_scores: + # The fused pipeline (THE union path) always z-normalizes, so + # this combination is rejected loudly at construction. + with pytest.raises(ValueError, match="normalize_scores=True"): + _make_triattention( + top_B=4, + eviction_mode=eviction_mode, + normalize_scores=normalize_scores, + ) + return manager = _make_triattention( top_B=4, eviction_mode=eviction_mode, @@ -913,7 +924,6 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize ) plan = build_selection.call_args.args[0] assert plan.eviction_mode == eviction_mode - assert build_selection.call_args.kwargs["normalize_scores"] is normalize_scores assert resources.score_staging is score_staging assert resources.keep_set_selector is keep_set_selector 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 index 81bd19f6d4b3..de9e9035a467 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -152,6 +152,8 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( def test_union_eager_runs_the_registered_cute_op(): _require_cute_topk_op() device = torch.device("cuda", torch.cuda.current_device()) + # The fused CuTe pipeline is the production union-row producer; these + # selector tests stage prepared union rows into ``combined`` directly. scores = torch.randn(2, 4, 96, dtype=torch.float32, device=device) selector = _BatchedUnionKeepSetSelector( rows=4, @@ -160,62 +162,42 @@ def test_union_eager_runs_the_registered_cute_op(): dtype=torch.float32, device=device, max_requests=2, - input_scores=scores, - normalize_scores=True, ) - selector.select_prepared_requests() + selector.combined.copy_(scores.amax(dim=1)) + selector.select_prepared_union_scores() torch.cuda.synchronize(device) assert torch.all(selector.keep[:, 1:] >= selector.keep[:, :-1]) -@pytest.mark.parametrize("normalize_scores", [False, True]) -def test_prepared_union_scores_match_checked_launch_and_exact_indices(normalize_scores): +def test_prepared_union_rows_select_exact_indices(): _require_cute_topk_op() - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels - device = torch.device("cuda", torch.cuda.current_device()) - request_count, rows, width, keep_count = 2, 7, 97, 64 + request_count, width, keep_count = 2, 97, 64 generator = torch.Generator(device=device).manual_seed(53) - scores = torch.randn( + combined_rows = torch.randn( request_count, - rows, width, generator=generator, dtype=torch.float32, device=device, ) valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) - reference_mean = torch.empty(request_count, rows, 1, dtype=torch.float32, device=device) - reference_inv_std = torch.empty_like(reference_mean) - reference_combined = torch.empty(request_count, width, dtype=torch.float32, device=device) - triattention_kernels.prepare_union_scores( - scores, - valid_widths, - reference_mean, - reference_inv_std, - reference_combined, - request_count, - normalize_scores=normalize_scores, - ) selector = _BatchedUnionKeepSetSelector( - rows=rows, + rows=7, width=width, keep_count=keep_count, dtype=torch.float32, device=device, max_requests=request_count, - input_scores=scores, - normalize_scores=normalize_scores, ) selector.valid_widths.copy_(valid_widths) - selector.select_prepared_requests() - actual_combined = selector.combined.cpu() + selector.combined.copy_(combined_rows) + selector.select_prepared_union_scores() actual_keep = selector.keep.cpu() - expected_combined = reference_combined.cpu() + expected_combined = combined_rows.cpu() torch.cuda.synchronize(device) - assert torch.equal(actual_combined, expected_combined) for request, valid_width in enumerate(valid_widths.cpu().tolist()): expected_keep = torch.sort( _stable_topk(expected_combined[request], valid_width, keep_count) @@ -246,8 +228,6 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid dtype=torch.float32, device=device, max_requests=request_count, - input_scores=scores, - normalize_scores=False, ) selector.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) # Write the shared per-request prompt lengths the way production staging @@ -256,7 +236,8 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) selector.refresh_row_prompt_offsets() - selector.select_prepared_requests() + selector.combined.copy_(scores.amax(dim=1)) + selector.select_prepared_union_scores() actual = selector.keep.cpu() combined = scores.amax(dim=1).cpu() @@ -267,46 +248,10 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid assert torch.equal(actual[request], expected_decode) -def test_fused_union_preparation_matches_ragged_torch_reference(): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - prepare_union_scores, - ) - - device = torch.device("cuda", torch.cuda.current_device()) - request_count, rows, width = 2, 7, 97 - generator = torch.Generator(device=device).manual_seed(17) - scores = torch.randn( - request_count, - rows, - width, - generator=generator, - dtype=torch.float32, - device=device, - ) - valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) - row_mean = torch.empty(request_count, rows, 1, dtype=torch.float32, device=device) - row_inv_std = torch.empty_like(row_mean) - combined = torch.empty(request_count, width, device=device) - - prepare_union_scores( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - request_count, - normalize_scores=True, - ) - torch.cuda.synchronize(device) - - expected = torch.full_like(combined, float("-inf")) - for request, valid_width in enumerate(valid_widths.tolist()): - valid_scores = scores[request, :, :valid_width] - mean = valid_scores.mean(dim=1, keepdim=True) - std = torch.linalg.vector_norm(valid_scores - mean, dim=1, keepdim=True) - std = (std / valid_width**0.5).clamp_min(1e-6) - expected[request, :valid_width] = ((valid_scores - mean) / std).amax(dim=0) - assert torch.allclose(combined, expected, rtol=2e-5, atol=2e-5) +# The retired split-union preparation (``prepare_union_scores``) is covered +# by its standalone reference copy in test_triattention_cute_union_fusion.py, +# which validates it against a pure-torch oracle and uses it as the +# fused-pipeline equivalence reference. @pytest.mark.parametrize("per_layer", [False, True]) @@ -912,10 +857,6 @@ def expected_keep() -> torch.Tensor: dense_layers=(0,), num_query_heads=num_q_heads, num_kv_heads=1, - input_scores=score_staging.fused_group.output.view( - 1, num_q_heads, seq_len - prompt_len - ), - normalize_scores=False, prompt_offsets_buffer=score_staging.token_starts_device, ) score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") @@ -949,8 +890,12 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: [seq_len + protected_tail], ) keep_set_selector.refresh_row_prompt_offsets() - score_staging.launch_prepared_score() - keep_set_selector.select_prepared_requests() + # THE union path: the fused pipeline writes normalized union rows + # into ``combined``. Z-normalization is monotonic per row and all + # query heads carry identical scores here, so the expected keep + # set (derived from raw scores) is unchanged. + score_staging.launch_prepared_union_fusion(keep_set_selector.combined) + keep_set_selector.select_prepared_union_scores() selected = keep_set_selector.keep[0].clone().to(torch.long) batched_compaction.compact() score_staging.mark_page_tables_consumed(manager._stream) From 56498dfdbc634dbd61a58dab0268f0c870820e94 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 21 Jul 2026 23:18:21 -0700 Subject: [PATCH 076/178] [None][feat] Retire the single-shot CuTe score kernel The fused score pack is now the only score implementation: the per-head and per-layer-per-head selection paths launch its score-only entry, and the single-shot runner (triattention_cute_score.py) is deleted. The fixed score group stages per-request window starts for the score launch exactly like the union path, so the per-head modes drop the legacy uniform-score_start limitation (each request scores its own decode window instead of the full sequence from token zero) and the runner dispatches every request count up to the group capacity instead of only one and the capacity. The union fusion runner now shares the group's staged window starts and segment offsets. The score-only entry writes the same head-major scratch layout the group gather already consumes, and its per-token score math is unchanged (same UMMA sequence, fp32 accumulation), so generations are token-identical across the swap. The full unit suite is green (140 passed). Against an unmodified 623df3d tree the qwen3-8B per_head and per_layer_perhead e2e digests are token-identical, as are the qwen3-8B and gpt-oss-20b union digests and the eagle3(one-model) overlap+graph leg's eviction generations. Signed-off-by: tianruih --- .../triattention/triattention_cute_score.py | 1444 ----------------- .../triattention_cute_score_fused.py | 11 +- .../triattention/triattention_kernels.py | 125 +- .../test_triattention_cute_union_fusion.py | 9 +- .../test_triattention_score_ops.py | 37 +- 5 files changed, 102 insertions(+), 1524 deletions(-) delete mode 100644 tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py deleted file mode 100644 index 29785c702e71..000000000000 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_score.py +++ /dev/null @@ -1,1444 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""SM100 CuTe-DSL scorer for the TriAttention mean-score path. - -This is the production specialization of the final workbench kernel — and -the ONLY score implementation. It uses split real/imag TMA loads, BF16 and -FP16 compensated UMMA, sqrt FTZ, and producer-only page-ID lookahead. -Geometries outside the exact contract validated here raise loudly at setup -(``_FixedScoreGroup.prepare_cute_score``); there is no fallback path. - -Page-table contract: ``page_ids`` is the flattened native block-offset -staging buffer ([pool_slot, request, K/V plane, block] int32) produced by -the V2 manager; K-plane entries encode ``physical_page * kv_factor`` and -are decoded inline (kv_factor == 2), so no per-round conversion pass is -needed. -""" - -from __future__ import annotations - -import inspect -import threading - -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.cute.nvgpu import cpasync, tcgen05 -from cutlass.cute.runtime import from_dlpack - - -def _cute_sqrt_keyword_mode() -> str: - """Probe which fast-sqrt spelling this CuTe DSL's ``cute.math.sqrt`` takes. - - The approximate-sqrt control was renamed across DSL releases: some expose - ``approx``/``ftz`` keywords, cutlass 4.5 exposes a single ``fastmath`` - flag, and older releases expose a plain one-argument ``sqrt``. Passing an - unknown keyword raises TypeError at trace time (inside ``cute.compile``, - where it cannot be caught), so the capability is probed once at import - time via signature inspection and folded into a trace-time constant. - """ - try: - parameters = inspect.signature(cute.math.sqrt).parameters - except (TypeError, ValueError): - return "plain" - if "approx" in parameters and "ftz" in parameters: - return "approx_ftz" - if "fastmath" in parameters: - return "fastmath" - return "plain" - - -_CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() - -CTA_M = 64 -N = 8 -THREADS = 128 - -RAW_K_VECTOR_ELEMENTS = 8 -TMA_DESCRIPTOR_QWORDS = 16 - - -class _TriScoreEpilogue: - """Minimal TMEM-to-global epilogue for the score specialization.""" - - def __init__(self) -> None: - self.acc_dtype = cutlass.Float32 - - def epilog_tmem_copy_and_partition( - self, - tidx: cutlass.Int32, - accumulator: cute.Tensor, - output: cute.Tensor, - epilogue_tile: cute.Tile, - use_2cta_instrs: bool, - ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: - copy_atom = sm100_utils.get_tmem_load_op( - self.cta_tile_shape_mnk, - self.c_layout, - self.c_dtype, - self.acc_dtype, - epilogue_tile, - use_2cta_instrs, - ) - accumulator_epilogue = cute.flat_divide( - accumulator[((None, None), 0, 0)], - epilogue_tile, - ) - tiled_copy = tcgen05.make_tmem_copy( - copy_atom, - accumulator_epilogue[(None, None, 0, 0)], - ) - thread_copy = tiled_copy.get_slice(tidx) - thread_accumulator = thread_copy.partition_S(accumulator_epilogue) - output_epilogue = cute.flat_divide( - output[((None, None), 0, 0, None, None, None)], - epilogue_tile, - ) - thread_output = thread_copy.partition_D(output_epilogue) - register_accumulator = cute.make_rmem_tensor( - thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.acc_dtype, - ) - return tiled_copy, thread_accumulator, register_accumulator - - def epilog_gmem_copy_and_partition( - self, - tidx: cutlass.Int32, - tiled_copy: cute.TiledCopy, - output: cute.Tensor, - epilogue_tile: cute.Tile, - _unused_smem: cute.Tensor, - ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: - output_epilogue = cute.flat_divide( - output[((None, None), 0, 0, None, None, None)], - epilogue_tile, - ) - thread_copy = tiled_copy.get_slice(tidx) - thread_output = thread_copy.partition_D(output_epilogue) - register_output = cute.make_rmem_tensor( - thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.c_dtype, - ) - copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - return copy_atom, register_output, thread_output - - -class _TriAttentionScoreKernel(_TriScoreEpilogue): - """Assign one CTA to each segment/KV-head task and retain W across pages.""" - - def __init__( - self, - *, - num_segments: int, - seq_len: int, - score_start: int, - num_q_heads: int, - num_kv_heads: int, - num_freqs: int, - tokens_per_block: int, - pool_shape: tuple[int, int, int, int, int], - pool_strides: tuple[int, int, int, int, int], - pool_dtype: type[cutlass.Numeric], - page_shards: int, - ) -> None: - """Build the single validated production specialization.""" - super().__init__() - if pool_dtype is not cutlass.BFloat16: - raise ValueError("TriAttention CuTe score requires BF16 K pages") - 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 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 score_start % tokens_per_block: - raise ValueError("TriAttention CuTe score requires page-aligned score_start") - if page_shards not in (2, 3): - raise ValueError("TriAttention CuTe score requires two or three page shards") - - self.score_start = score_start - self.num_q_heads = num_q_heads - self.num_kv_heads = num_kv_heads - self.group_size = num_q_heads // num_kv_heads - self.sum_seq = num_segments * seq_len - self.num_tasks = num_segments * num_kv_heads - self.page_shards = page_shards - self.num_ctas = self.num_tasks * page_shards - self.num_freqs = num_freqs - self.tokens_per_block = tokens_per_block - # cos/sin/mlr coefficient planes per frequency. - self.k_coeff = 3 * num_freqs - # One 64-token compute tile either sub-divides a page (128-token - # pages: two halves per page, one TMA box each) or spans several - # pages (32-token pages: two page fragments per phase). The - # validated 128-token geometry is the single-fragment case, so - # its schedule is unchanged. - self.box_tokens = min(CTA_M, tokens_per_block) - self.fragments_per_phase = CTA_M // self.box_tokens - self.pages_per_tile = self.fragments_per_phase - self.halves_per_page = max(1, tokens_per_block // CTA_M) - self.tile_tokens = self.halves_per_page * CTA_M - self.max_tiles = (seq_len + self.tile_tokens - 1) // self.tile_tokens - - # Measured final choices that still shape layouts or generated code. - self.prefetch_depth = 4 - self.sqrt_mode = "approx" - self.k_staging_mode = "half_page_tma" - self.use_tma = True - self.cpasync_schedule = "sync_each_half" - self.split_raw_tma = True - self.raw_tma_feature_extent = num_freqs - # Barrier transaction bytes for one phase: the full 64-token tile - # of one coefficient plane, regardless of how many page fragments - # deliver it. - self.raw_tma_copy_bytes = CTA_M * num_freqs * (cutlass.BFloat16.width // 8) - self.raw_tma_pipeline_stages = 1 - self.accumulator_pipeline_stages = 1 - self.umma_accumulator_partitions = 1 - self.raw_cpasync_direct_a = True - self.weight_builder_mode = "coefficient_scalar_bf16_two_term" - self.main_operand_mode = "bf16_raw_three_term_weight" - self.numerical_policy = "three_term" - self.magnitude_residual_mode = "fp16_mma_two_term_single_commit" - self.fp16_magnitude_two_term = True - self.magnitude_sqrt_ftz = True - self.producer_page_id_prefetch = True - self.producer_warp_id = 0 - self.physical_threads = THREADS - self.shared_a_raw_alias = False - self.compact_token_loop = True - - self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape - if ( - pool_kv_heads != num_kv_heads - or pool_tokens != tokens_per_block - or 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_slot, self.s_dim = pool_strides - if self.s_slot != 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 __call__( - self, - page_ids: cute.Tensor, - seg_page_off: cute.Tensor, - seg_req_id: cute.Tensor, - seg_layer_id: cute.Tensor, - seg_seq_len: cute.Tensor, - seg_out_offset: 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, - pool_template: cute.Tensor, - raw_tma_descriptors: cute.Tensor, - stream: cuda.CUstream, - ): - self.c_dtype = output.element_type - self.c_layout = utils.LayoutEnum.COL_MAJOR - self.mma_tiler = (CTA_M, N, self.k_coeff) - self.cta_tile_shape_mnk = self.mma_tiler - self.epi_tile = (CTA_M, N) - - 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], - ) - main_a_shape = ( - (CTA_M, N, self.num_freqs) - if self.main_operand_mode == "bf16_raw_three_term_weight" - else self.mma_tiler - ) - a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, main_a_shape, cutlass.Float32, 1) - raw_bf16_a_smem_layout = sm100_utils.make_smem_layout_a( - raw_bf16_tiled_mma, - (CTA_M, N, 2 * self.num_freqs), - cutlass.BFloat16, - 1, - ) - raw_bf16_split_a_smem_layout = sm100_utils.make_smem_layout_a( - raw_bf16_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.BFloat16, - 2, - ) - # The full transport uses one K_SW128 8-KiB tile. The split transport - # packs two K_SW64 4-KiB stages into that same allocation, one each for - # real and imaginary data; the compile-time schedule selects the view. - raw_bf16_direct_a_smem_layout = ( - raw_bf16_split_a_smem_layout if self.split_raw_tma else raw_bf16_a_smem_layout - ) - 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_slot, - (self.s_kv_head, self.s_page), - ), - ) - raw_tma_source = cute.make_tensor( - pool_template.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, N, 2 * self.num_freqs), - cutlass.BFloat16, - 1, - ) - # The magnitude residual has only the frequency-count K. A separate - # compact descriptor lets the producer issue its UMMA steps before the - # first commit, rather than waiting for and overwriting the main A tile. - magnitude_lo_smem_layout = sm100_utils.make_smem_layout_a( - tiled_mma, (CTA_M, N, self.num_freqs), cutlass.Float32, 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_lo_fp16_smem_layout = sm100_utils.make_smem_layout_a( - magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.Float16, - 1, - ) - magnitude_hi_fp16_smem_layout = sm100_utils.make_smem_layout_b( - magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.Float16, - 1, - ) - magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( - magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.Float16, - 1, - ) - magnitude_fp16_b_smem_layout = sm100_utils.make_smem_layout_b( - magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.Float16, - 1, - ) - main_b_shape = ( - (CTA_M, N, self.num_freqs) - if self.main_operand_mode == "bf16_raw_three_term_weight" - else self.mma_tiler - ) - b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, main_b_shape, cutlass.Float32, 1) - acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - # Keep an explicit stage mode even for the one-stage control. The - # singleton mode folds away in codegen and lets both specializations - # share the same producer/consumer slicing protocol. - self.num_accumulator_slots = ( - self.accumulator_pipeline_stages * self.umma_accumulator_partitions - ) - 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) - - b_hi_elements = cute.cosize(b_smem_layout.outer) * int(not self.fp16_magnitude_two_term) - b_lo_elements = cute.cosize(b_smem_layout.outer) * int( - self.numerical_policy == "three_term" and not self.fp16_magnitude_two_term - ) - magnitude_lo_elements = cute.cosize(magnitude_lo_fp16_smem_layout.outer) * int( - self.numerical_policy == "three_term" - and self.magnitude_residual_mode in ("fp16_smem", "fp16_mma_single_commit") - ) - magnitude_lo_fp32_elements = cute.cosize(magnitude_lo_smem_layout.outer) * int( - self.numerical_policy == "three_term" - and self.magnitude_residual_mode == "fp32_smem_single_commit" - ) - magnitude_hi_fp16_elements = cute.cosize(magnitude_hi_fp16_smem_layout.outer) * int( - self.numerical_policy == "three_term" - and self.magnitude_residual_mode == "fp16_mma_single_commit" - ) - a_elements = cute.cosize(a_smem_layout.outer) * int( - not self.shared_a_raw_alias and not self.fp16_magnitude_two_term - ) - raw_k_half_elements = CTA_M * 2 * self.num_freqs - raw_k_elements = raw_k_half_elements * int( - self.k_staging_mode in ("half_page_cpasync", "half_page_tma") - and not self.shared_a_raw_alias - ) - alias_a_elements = cute.cosize(a_smem_layout.outer) * int(self.shared_a_raw_alias) - alias_raw_k_elements = raw_k_half_elements * int(self.shared_a_raw_alias) - raw_bf16_a_elements = cute.cosize(raw_bf16_a_smem_layout.outer) * int( - self.main_operand_mode == "bf16_raw_three_term_weight" - and ( - not self.raw_cpasync_direct_a - or self.cpasync_schedule not in ("sync_each_half", "intra_half_overlap") - ) - ) - raw_bf16_b_elements = cute.cosize(raw_bf16_b_smem_layout.outer) * int( - self.main_operand_mode == "bf16_raw_three_term_weight" - ) - raw_bf16_b2_elements = raw_bf16_b_elements * int( - self.weight_builder_mode != "coefficient_scalar_bf16_two_term" - ) - magnitude_fp16_a_elements = cute.cosize(magnitude_fp16_a_smem_layout.outer) * int( - self.fp16_magnitude_two_term - ) - magnitude_fp16_b_elements = cute.cosize(magnitude_fp16_b_smem_layout.outer) * int( - self.fp16_magnitude_two_term - ) - - @cute.union - class SharedARawAlias: - # The two descriptors are byte-identical in size (8 KiB) and have - # disjoint lifetimes in the alias specialization. - sA: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, alias_a_elements], - 1024, - ] - sRawK: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, alias_raw_k_elements], - 16, - ] - - @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 * int(self.use_tma), - ] - tmem_holding_buf: cutlass.Int32 - sA: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, a_elements], - 1024, - ] - sB_hi: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, b_hi_elements], - 1024, - ] - sB_lo: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, b_lo_elements], - 1024, - ] - sMagnitudeLo: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float16, magnitude_lo_elements], - 1024, - ] - sMagnitudeLoFp32: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, magnitude_lo_fp32_elements], - 1024, - ] - sMagnitudeHiFp16: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float16, magnitude_hi_fp16_elements], - 1024, - ] - sRawK: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, raw_k_elements], - 1024, - ] - sRawBf16A: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, raw_bf16_a_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, - ] - sRawBf16B2: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b2_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, - ] - sARawAlias: SharedARawAlias - - self.shared_storage = SharedStorage - self.kernel( - tiled_mma, - raw_bf16_tiled_mma, - magnitude_lo_tiled_mma, - raw_tma_atom, - raw_tma_tensor, - raw_tma_descriptors, - page_ids, - seg_page_off, - seg_req_id, - seg_layer_id, - seg_seq_len, - seg_out_offset, - q_real, - q_imag, - mlr_coef, - mean_cos, - mean_sin, - freq_scale_sq, - output, - a_smem_layout, - raw_bf16_a_smem_layout, - raw_bf16_direct_a_smem_layout, - raw_bf16_split_a_smem_layout, - raw_tma_smem_layout, - raw_bf16_b_smem_layout, - magnitude_lo_smem_layout, - magnitude_lo_fp16_smem_layout, - magnitude_hi_fp16_smem_layout, - magnitude_fp16_a_smem_layout, - magnitude_fp16_b_smem_layout, - b_smem_layout, - ).launch( - grid=(self.num_ctas, 1, 1), - block=(self.physical_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, - page_ids: cute.Tensor, - seg_page_off: cute.Tensor, - seg_req_id: cute.Tensor, - seg_layer_id: cute.Tensor, - seg_seq_len: cute.Tensor, - seg_out_offset: 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, - a_smem_layout: cute.ComposedLayout, - raw_bf16_a_smem_layout: cute.ComposedLayout, - raw_bf16_direct_a_smem_layout: cute.ComposedLayout, - raw_bf16_split_a_smem_layout: cute.ComposedLayout, - raw_tma_smem_layout: cute.ComposedLayout, - raw_bf16_b_smem_layout: cute.ComposedLayout, - magnitude_lo_smem_layout: cute.ComposedLayout, - magnitude_lo_fp16_smem_layout: cute.ComposedLayout, - magnitude_hi_fp16_smem_layout: cute.ComposedLayout, - magnitude_fp16_a_smem_layout: cute.ComposedLayout, - magnitude_fp16_b_smem_layout: cute.ComposedLayout, - 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] - valid_seq_len = seg_seq_len[segment] - page_off = seg_page_off[segment] - out_base = seg_out_offset[segment] - - 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, - ) - 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)] - # Each stage slice retains the K_SW64 pointer flags. Reuse only - # the feature-first outer mapping for the corresponding 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 and TMA partition per page fragment of the - # 64-token tile. Fragment f lands box_tokens rows deeper in the - # same stage; the offset is a whole multiple of the swizzle - # period, so the descriptor swizzle stays phase-aligned. - raw_tma_shared_partition_real = [] - raw_tma_shared_partition_imag = [] - 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, - ) - 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), - ) - raw_tma_shared_partition_real.append(partition_real) - raw_tma_shared_partition_imag.append(partition_imag) - 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, - ) - 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, - ) - - page_index = self.score_start // self.tile_tokens + page_shard - page_start = page_index * self.tile_tokens - pages_processed = cutlass.Int32(0) - producer_prefetched_page_id_lane0 = cutlass.Int32(0) - producer_prefetched_page_id_lane0_f1 = cutlass.Int32(0) - if warp_idx == self.producer_warp_id: - if lane_idx == 0: - # ``page_ids`` is the flattened native block-offset staging - # buffer ([pool_slot, request, K/V plane, block] int32) and - # ``page_off`` points at one request's K plane. K-plane - # entries encode ``physical_page * kv_factor`` (kv_factor is - # 2 for the interleaved K/V pools this kernel requires), so - # divide by two here. V-plane entries are never read. - producer_prefetched_page_id_lane0 = ( - cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile]) // 2 - ) - if cutlass.const_expr(self.fragments_per_phase == 2): - # The tail tile may not reach its second page; clamp - # to the first fragment (those scores lie past the - # valid width and are masked downstream) so the TMA - # never dereferences an unstaged block entry. - second_page_id = producer_prefetched_page_id_lane0 - if page_start + self.box_tokens < valid_seq_len: - second_page_id = ( - cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile + 1]) - // 2 - ) - producer_prefetched_page_id_lane0_f1 = second_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 - ) - cute.arch.mbarrier_init_fence() - for weight_round in cutlass.range_constexpr(N * 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 - # GQA groups below the minimum MMA tile N=8 ride padded - # columns: they read the group's first head (any valid - # address) and force zero coefficients, so the padded score - # columns come out zero and land in scratch rows the adapter - # never gathers. - qg_read = qg - if cutlass.const_expr(self.group_size < N): - 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 < N): - 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) - - thr_mma = tiled_mma.get_slice(0) - while page_start < valid_seq_len and pages_processed < self.max_tiles: - physical_page = cutlass.Int32(0) - physical_page_f1 = 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 == 2): - physical_page_f1 = cute.arch.shuffle_sync( - producer_prefetched_page_id_lane0_f1, - 0, - ) - for page_half in cutlass.range_constexpr(self.halves_per_page): - if warp_idx == self.producer_warp_id: - # Phase 0 fills the packed 4-KiB K_SW64 real - # view. Every producer-warp lane participates - # in the PipelineTmaAsync barrier election. - raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - page_half, - (kv_head, physical_page), - ) - ], - raw_tma_shared_partition_real[0], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - if cutlass.const_expr(self.fragments_per_phase == 2): - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - page_half, - (kv_head, physical_page_f1), - ) - ], - raw_tma_shared_partition_real[1], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - 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 warp_idx == self.producer_warp_id: - raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - page_half, - (kv_head, physical_page), - ) - ], - raw_tma_shared_partition_imag[0], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier(raw_tma_producer_state), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - if cutlass.const_expr(self.fragments_per_phase == 2): - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - page_half, - (kv_head, physical_page_f1), - ) - ], - raw_tma_shared_partition_imag[1], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - raw_tma_producer_state.advance() - - if cutlass.const_expr( - self.producer_page_id_prefetch and page_half == self.halves_per_page - 1 - ): - next_page_id_lane0 = cutlass.Int32(0) - next_page_id_lane0_f1 = cutlass.Int32(0) - if warp_idx == self.producer_warp_id: - if lane_idx == 0: - next_page_start = page_start + self.tile_tokens * self.page_shards - next_pages_processed = pages_processed + 1 - if ( - next_page_start < valid_seq_len - and next_pages_processed < self.max_tiles - ): - # Same K-plane decode as the initial - # prefetch: entries are physical_page * 2. - next_page_id_lane0 = ( - cutlass.Int32( - page_ids[ - page_off - + (page_index + self.page_shards) * self.pages_per_tile - ] - ) - // 2 - ) - if cutlass.const_expr(self.fragments_per_phase == 2): - next_second_id = next_page_id_lane0 - if next_page_start + self.box_tokens < valid_seq_len: - next_second_id = ( - cutlass.Int32( - page_ids[ - page_off - + (page_index + self.page_shards) - * self.pages_per_tile - + 1 - ] - ) - // 2 - ) - next_page_id_lane0_f1 = next_second_id - producer_prefetched_page_id_lane0 = next_page_id_lane0 - if cutlass.const_expr(self.fragments_per_phase == 2): - producer_prefetched_page_id_lane0_f1 = next_page_id_lane0_f1 - - # 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, 0)], - 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) - # Each of the 32 lanes 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 - # Issue several independent token loads before consuming - # any of them. This bounded RMEM window is unchanged; the - # optional half-page staging only switches its K source - # from global to the single raw shared buffer. - for token_base in cutlass.range( - 0, - CTA_M // (THREADS // 32), - self.prefetch_depth, - unroll_full=not self.compact_token_loop, - ): - 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, - 0, - ) - ] - ) - staged_imag[prefetch_index] = cutlass.Float32( - cpasync_raw_k_0[ - ( - (token, frequency % 16), - 0, - frequency // 16, - 1, - ) - ] - ) - - 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 - if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): - magnitude = cute.math.sqrt( - norm2, - approx=self.sqrt_mode == "approx", - ftz=self.magnitude_sqrt_ftz, - ) - elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): - # cutlass 4.5 renamed the approximate-sqrt control - # to ``fastmath``; map the measured approx choice - # onto it to preserve the authored behavior. - magnitude = cute.math.sqrt( - norm2, fastmath=self.sqrt_mode == "approx" - ) - else: - # DSLs with neither spelling get the plain (IEEE) - # sqrt, which is strictly MORE accurate than the - # measured approx choice above; the unit test's - # 5e-3 oracle tolerance absorbs the difference. - magnitude = cute.math.sqrt(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, 1)], - 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, 0)], - 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, 1)], - tCrRawBf16B1[(None, None, imag_b_block, 0)], - tCtAcc, - ) - # Keep the full four-product control unchanged. - # The independent omit_a1b1 mode drops only the - # second-order residual product; all other FP16 - # K16 products retain their original order. - 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 pages_processed == 0: - if cutlass.const_expr(page_half == 0): - cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) - acc_pipeline.consumer_wait(acc_consumer_state) - # 64-bit: the head-plane product exceeds 2^31 once the scratch - # spans many request*layer segments (large batch), so promote - # before the multiply reaches Int32 arithmetic. - output_offset = ( - cutlass.Int64(kv_head) * (N * self.sum_seq) - + out_base - + page_start - + page_half * CTA_M - ) - page_output = cute.make_tensor( - output.iterator + output_offset, - cute.make_layout( - (CTA_M, N, 1), - stride=( - 1, - self.sum_seq, - N * self.sum_seq, - ), - ), - ) - gC_mnl = cute.local_tile(page_output, self.epi_tile, (None, None, None)) - tCgC = thr_mma.partition_C(gC_mnl) - tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( - tidx, tCtAcc, tCgC, self.epi_tile, False - ) - simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( - tidx, tiled_copy_t2r, tCgC, self.epi_tile, None - ) - 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)) - for subtile_idx in range(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)) - cute.copy( - simt_atom, - tTR_rC, - tTR_gC[(None, None, None, subtile_idx)], - ) - - cute.arch.fence_view_async_tmem_load() - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() - cute.arch.barrier() - raw_tma_pipeline.consumer_release(raw_tma_consumer_state) - raw_tma_consumer_state.advance() - - page_index += self.page_shards - page_start += self.tile_tokens * self.page_shards - pages_processed += 1 - if warp_idx == self.producer_warp_id: - raw_tma_pipeline.producer_tail(raw_tma_producer_state) - if warp_idx == self.producer_warp_id: - acc_pipeline.producer_tail(acc_producer_state) - 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: - """Encode one immutable feature-first TensorMap per layer index.""" - 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) != (2, 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) - 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 smem layout the TMA lands in; the - # sm100 helpers pick it from the inner-row byte count (one - # coefficient plane: num_freqs bf16 elements). - ( - 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) - - -class TriAttentionCuteScoreRunner: - """Compile and launch the exact SM100 mean-score specialization.""" - - def __init__( - self, - *, - layer_pools: list[torch.Tensor], - layer_indices: list[int], - max_requests: int, - num_layers: int, - seq_len: int, - score_start: int, - num_q_heads: int, - num_kv_heads: int, - num_freqs: int, - tokens_per_block: int, - page_ids: torch.Tensor, - seg_page_off: torch.Tensor, - seg_req_id: torch.Tensor, - seg_layer_id: torch.Tensor, - seg_seq_len: torch.Tensor, - seg_out_offset: torch.Tensor, - q_real: torch.Tensor, - q_imag: torch.Tensor, - mlr_coef: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - freq_scale_sq: torch.Tensor, - output: torch.Tensor, - ) -> None: - self.max_requests = int(max_requests) - self.num_layers = int(num_layers) - self.num_kv_heads = int(num_kv_heads) - self.descriptors = _encode_tma_descriptors( - layer_pools, layer_indices, int(num_freqs), int(tokens_per_block) - ) - self._torch_prefix = ( - page_ids, - seg_page_off, - seg_req_id, - seg_layer_id, - seg_seq_len, - seg_out_offset, - q_real, - q_imag, - mlr_coef, - ) - self._torch_tail = ( - freq_scale_sq, - output, - layer_pools[layer_indices[0]], - self.descriptors, - ) - self._cute_prefix = tuple(_to_cute(tensor) for tensor in self._torch_prefix) - self._cute_tail = ( - _to_cute(freq_scale_sq), - _to_cute(output), - _to_cute(layer_pools[layer_indices[0]]), - _to_cute(self.descriptors, assumed_align=128), - ) - self._compiled: dict[int, object] = {} - static_geometry = ( - max_requests * num_layers, - seq_len, - score_start, - num_q_heads, - num_kv_heads, - num_freqs, - tokens_per_block, - tuple(int(value) for value in layer_pools[layer_indices[0]].shape), - tuple(int(value) for value in layer_pools[layer_indices[0]].stride()), - ) - tensor_specs = tuple( - _tensor_spec(tensor) - for tensor in ( - *self._torch_prefix, - mean_cos.view(-1), - mean_sin.view(-1), - *self._torch_tail, - ) - ) - variants = [(1, 3)] - if max_requests > 1: - variants.append((max_requests, 2)) - for request_count, page_shards in variants: - cache_key = ( - "triattention_cute_score", - static_geometry, - tensor_specs, - request_count, - page_shards, - ) - with _COMPILE_LOCK: - compiled = _COMPILED_KERNELS.get(cache_key) - if compiled is None: - kernel = _TriAttentionScoreKernel( - num_segments=request_count * num_layers, - seq_len=seq_len, - score_start=score_start, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - pool_shape=tuple( - int(value) for value in layer_pools[layer_indices[0]].shape - ), - pool_strides=tuple( - int(value) for value in layer_pools[layer_indices[0]].stride() - ), - pool_dtype=cutlass.BFloat16, - page_shards=page_shards, - ) - stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) - compiled = cute.compile( - kernel, - *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), - *self._cute_tail, - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled - self._compiled[request_count] = compiled - - def supports(self, request_count: int) -> bool: - """Return whether an exact static specialization was precompiled.""" - return request_count in self._compiled - - def launch( - self, - request_count: int, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - ) -> None: - """Launch the CuTe score kernel on the current PyTorch stream.""" - stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) - self._compiled[request_count]( - *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), - *self._cute_tail, - stream, - ) 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 index 5af12cdf13c7..9ef401676e77 100644 --- 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 @@ -2,10 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """SM100 CuTe-DSL scorer for the TriAttention mean-score path. -This is the production specialization of the final workbench kernel. It -uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, and -producer-only page-ID lookahead. The public integration keeps a Triton -fallback for every geometry outside the exact contract validated here. +This is the production specialization of the final workbench kernel — and +the ONLY score implementation: the per-head modes launch its score-only +entry and union eviction launches its fused score+stats+union pipeline. It +uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, +and producer-only page-ID lookahead. Geometries outside the exact contract +validated here raise loudly at setup +(``_FixedScoreGroup.prepare_cute_score``); there is no fallback path. """ from __future__ import annotations diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 59a035a07867..42172f036b66 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -4,15 +4,16 @@ The production path uses one fixed-shape trig-score launch across all dense layers, CuTE-DSL TopK selection, and grouped C++ compaction. Scoring runs -EXCLUSIVELY through the SM100 CuTe-DSL kernel -(``triattention_cute_score.py``): mean aggregation, BF16 KV pools, head size -64/128, 32/128-token pages, GQA group 4 or 8. There is deliberately no other -score path -- any geometry outside that contract raises loudly at setup -instead of routing to a slower kernel (the original Triton score kernel and -the C++ CUDA score stack have both been deleted). Union eviction likewise -runs EXCLUSIVELY through the fused score+stats+union CuTe pipeline -(``triattention_cute_score_fused.py`` + ``triattention_cute_selection.py``); -the split Triton row-stats/union-reduce launches were retired with it, and +EXCLUSIVELY through the SM100 CuTe-DSL fused score pack +(``triattention_cute_score_fused.py``): mean aggregation, BF16 KV pools, +head size 64/128, 32/128-token pages, GQA group 4 or 8, per-request score +window starts. There is deliberately no other score path -- any geometry +outside that contract raises loudly at setup instead of routing to a slower +kernel (the original Triton score kernel, the C++ CUDA score stack, and the +single-shot CuTe score kernel have all been deleted). The per-head modes use +the pack's score-only entry; union eviction runs its fused score+stats+union +pipeline (with ``triattention_cute_selection.py``). The split Triton +row-stats/union-reduce launches were retired with the union fusion, and their standalone copies live in the fused-pipeline unit test as references (the row-stats kernel itself remains: the per-head modes still normalize with it). The unit tests validate the CuTe kernels against independent @@ -310,11 +311,12 @@ def __init__( mlr_coef_LHF.view(-1), ) self.pointer_tail = (freq_scale_sq, omega, offsets) - # The SM100 CuTe score kernel (see triattention_cute_score.py) is THE - # score implementation; it is compiled by the first - # ``prepare_cute_score`` call. The runner encodes TMA descriptors from - # the actual pool tensors, hence the pool references retained here - # (see the LIFETIME note in the class docstring). + # The SM100 CuTe fused score pack (see + # triattention_cute_score_fused.py) is THE score implementation; its + # score-only entry is compiled by the first ``prepare_cute_score`` + # call. The runner encodes TMA descriptors from the actual pool + # tensors, hence the pool references retained here (see the LIFETIME + # note in the class docstring). self.seq_len = int(seq_len) self._cute_score_runner = None self._cute_score_attempted = False @@ -322,18 +324,18 @@ def __init__( self._cute_layer_indices = [int(layer) for layer in layer_indices] def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> None: - """Compile the SM100 CuTe score kernel once; raise loudly otherwise. + """Compile the fused CuTe runner's score-only entry; raise loudly otherwise. Call this outside CUDA graph capture: compilation allocates memory - and synchronizes. The CuTe kernel is the ONLY score implementation, - so an unsupported geometry raises ValueError here and a runner - construction failure raises RuntimeError -- there is deliberately no - fallback path. + and synchronizes. The fused score pack is the ONLY score + implementation, so an unsupported geometry raises ValueError here + and a runner construction failure raises RuntimeError -- there is + deliberately no fallback path. Supported contract: SM100 exactly, BF16 pools, 32- or 128-token pages, 32 or 64 frequencies (head size 64/128), 4 or 8 query heads per KV head, and a bucket capacity (``seq_len``) aligned to the - kernel's compute tile — this covers the Qwen3 and GPT-OSS production + historical score tile — this covers the Qwen3 and GPT-OSS production geometries as well as the original validation shape. """ if self._cute_score_attempted: @@ -342,11 +344,11 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> anchor = self.pointer_prefix[0] num_q_heads, num_kv_heads, num_freqs, tokens_per_block, kv_factor = self.geometry_args max_segments = self.max_requests * self.num_layers - # The kernel's epilogue stores full compute tiles (64 tokens, or one - # page for 128-token pages) into a scratch whose per-segment stride - # is seq_len, without masking the ragged tail of the LAST tile; an - # unaligned bucket would silently spill scores into the next - # segment's region, so it is rejected here instead. + # The retired single-shot kernel stored full unmasked compute tiles + # (64 tokens, or one page for 128-token pages), which forced + # tile-aligned buckets. The fused kernel masks its ragged tail, but + # the bucket contract is kept unchanged so the kernel swap cannot + # silently admit new geometry (production pow2 buckets satisfy it). score_tile_tokens = max(64, int(tokens_per_block)) supported = ( torch.cuda.get_device_capability(anchor.device) == (10, 0) @@ -379,14 +381,14 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> ) device = anchor.device try: - from .triattention_cute_score import TriAttentionCuteScoreRunner + from .triattention_cute_score_fused import TriAttentionCuteScoreRunner - # The kernel scores the FULL sequence from physical token zero - # into its own head-major scratch (row = query head, column = - # segment * seq_len + token); ``launch`` gathers each request's - # decode window from that scratch into ``self.output``. All - # buffers below are persistent because the compiled kernel - # captures their device pointers. + # The kernel scores each request's window (from its staged + # per-request start) into its own head-major scratch (row = + # query head, column = segment * seq_len + token); ``launch`` + # gathers each request's decode window from that scratch into + # ``self.output``. All buffers below are persistent because the + # compiled kernel captures their device pointers. # The kernel writes one scratch row per padded head column # (GQA group below 8 pads up to the MMA tile); the gather in # ``launch`` reads only the real heads. @@ -399,6 +401,10 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> seg_out_offset = ( torch.arange(max_segments, dtype=torch.int64, device=device) * self.seq_len ).to(torch.int32) + # Per-request score window starts, staged before each launch; + # the compiled kernels capture this buffer's device pointer. + # The union fusion runner shares it (and the segment buffers). + token_starts = torch.zeros(self.max_requests, dtype=torch.int32, device=device) gather_columns = torch.arange(self.output_width, dtype=torch.int64, device=device) self._cute_score_runner = TriAttentionCuteScoreRunner( layer_pools=self._cute_layer_pools, @@ -406,10 +412,6 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> max_requests=self.max_requests, num_layers=self.num_layers, seq_len=self.seq_len, - # Always score from physical token zero: per-request prompt - # windows are applied by the gather in ``launch`` instead of - # one global page-aligned start. - score_start=0, num_q_heads=num_q_heads, num_kv_heads=num_kv_heads, num_freqs=num_freqs, @@ -420,6 +422,7 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> seg_layer_id=self.pointer_prefix[5], seg_seq_len=seg_seq_len, seg_out_offset=seg_out_offset, + token_starts=token_starts, q_real=self.pointer_middle[0], q_imag=self.pointer_middle[1], mlr_coef=self.pointer_middle[2], @@ -427,6 +430,10 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> mean_sin=mean_sin, freq_scale_sq=self.pointer_tail[0], output=scratch, + # Score-only mode: the stats and union-finalize kernels are + # compiled lazily by ``_union_fusion_runner`` when (and only + # when) union eviction actually launches. + enable_partial_stats=False, ) except (ImportError, RuntimeError, ValueError, AssertionError) as error: raise RuntimeError( @@ -434,6 +441,8 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> ) from error self._cute_scratch = scratch self._cute_seg_seq_len = seg_seq_len + self._cute_seg_out_offset = seg_out_offset + self._cute_token_starts = token_starts self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) # Fused score+stats+union pipeline (Fanrong Li's two-kernel scheme): # THE union path, built lazily on the first union launch. ONE runner @@ -459,11 +468,13 @@ def launch( ) -> torch.Tensor: """Return decode-only scores as ``[request, layer, head, token]``. - Runs the SM100 CuTe score kernel (the only score implementation) and - writes each request's decode width (``valid_seq_len - token_start``) - into ``valid_widths``, which the selection reduce kernels consume. - Only mean aggregation exists; ``request_count`` must be one of the - precompiled variants (1 or the group capacity). + Runs the fused CuTe runner's score-only entry (the only score + implementation) and writes each request's decode width + (``valid_seq_len - token_start``) into ``valid_widths``, which the + selection reduce kernels consume. Each request scores its own window + from its staged start, so one cohort may mix prompt lengths. Only + mean aggregation exists; the runner dispatches every request count + up to the group capacity. """ if aggregation != "mean": raise ValueError( @@ -506,13 +517,16 @@ def launch( self.pointer_prefix[4][:num_segments], out=self._cute_seg_seq_len[:num_segments], ) + # Stage the per-request score window starts: the compiled kernel + # captured this buffer's pointer and reads one start per request. + self._cute_token_starts[:request_count].copy_(token_starts_device[:request_count]) runner.launch(request_count, mean_cos, mean_sin) - # The kernel wrote full-sequence scores from physical token zero into - # its head-major scratch. Gather each request's decode window - # (starting at its pinned prompt length) into the group output, the - # ``[request, layer, head, token]`` layout the selection kernels - # read. Columns past a request's valid width carry unscored scratch - # data; consumers mask by ``valid_widths``. + # The kernel wrote each request's window scores (from its pinned + # prompt length) into its head-major scratch. Gather each request's + # decode window into the group output, the ``[request, layer, head, + # token]`` layout the selection kernels read. Columns past a + # request's valid width carry unscored scratch data; consumers mask + # by ``valid_widths``. num_q_heads = int(self.geometry_args[0]) num_kv_heads = int(self.geometry_args[1]) group_size = num_q_heads // num_kv_heads @@ -565,10 +579,6 @@ def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): ) device = self.output.device - seg_out_offset = ( - torch.arange(self.max_requests * self.num_layers, dtype=torch.int64, device=device) - * self.seq_len - ).to(torch.int32) # The union output rows are sized by the whole bucket (the widest # possible window); consumers mask by the per-request widths. union_rows = torch.empty( @@ -576,9 +586,10 @@ def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): dtype=torch.float32, device=device, ) - # Per-request score window starts, staged before each launch; the - # compiled kernels capture this buffer's device pointer. - token_starts = torch.zeros(self.max_requests, dtype=torch.int32, device=device) + # The scratch, segment buffers, and staged per-request window + # starts are shared with the score-only runner built by + # ``prepare_cute_score``; the compiled kernels capture their + # device pointers. runner = _FusedUnionScoreRunner( layer_pools=self._cute_layer_pools, layer_indices=self._cute_layer_indices, @@ -594,8 +605,8 @@ def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): seg_req_id=self.pointer_prefix[4], seg_layer_id=self.pointer_prefix[5], seg_seq_len=self._cute_seg_seq_len, - seg_out_offset=seg_out_offset, - token_starts=token_starts, + seg_out_offset=self._cute_seg_out_offset, + token_starts=self._cute_token_starts, q_real=self.pointer_middle[0], q_imag=self.pointer_middle[1], mlr_coef=self.pointer_middle[2], @@ -609,7 +620,7 @@ def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): raise RuntimeError( "TriAttention CuTe union fusion setup failed and no other union path exists" ) from error - self._union_fusion_runner_entry = (runner, union_rows, token_starts) + self._union_fusion_runner_entry = (runner, union_rows, self._cute_token_starts) return self._union_fusion_runner_entry def launch_cute_union_fusion( 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 index 79d23da42225..bf22373277a5 100644 --- 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 @@ -356,7 +356,12 @@ def test_reference_union_preparation_matches_ragged_torch_reference() -> None: @_SM100_ONLY def test_union_fusion_setup_failure_raises(monkeypatch: pytest.MonkeyPatch) -> None: - """A fused-runner construction failure raises loudly: no fallback remains.""" + """A fused-runner construction failure raises loudly: no fallback remains. + + The single fused pack serves both the score-only and union entries, so + the construction failure surfaces at score setup (``prepare_cute_score`` + runs before the union runner is built). + """ pytest.importorskip("cutlass") import tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused as fused_module # noqa: E501 @@ -411,7 +416,7 @@ def _refuse_construction(**_kwargs): widths = torch.empty(2, dtype=torch.int32, device=device) token_starts = torch.zeros(2, dtype=torch.int32, device=device) union_out = torch.empty((2, seq_len), dtype=torch.float32, device=device) - with pytest.raises(RuntimeError, match="no other union path exists"): + with pytest.raises(RuntimeError, match="no other score path exists"): group.launch_cute_union_fusion( 2, valid_seq_lens, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py index eda49607cbbd..1d4993a1de7a 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py @@ -2,14 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 """The CuTe score kernel vs an independent PyTorch oracle, through group.launch. -The SM100 CuTe-DSL kernel (``triattention_cute_score.py``) is the ONLY score -implementation. These tests drive it through ``_FixedScoreGroup.launch`` -- -the exact production entry point -- across the supported production -geometries, with multi-layer segments, permuted page tables, ragged valid -lengths, and per-request prompt windows, and compare against a pure-PyTorch -oracle that recomputes everything independently. They also pin the -loud-failure contract: unsupported geometry, removed aggregations, and -uncompiled request counts raise instead of routing to another kernel. +The SM100 CuTe-DSL fused score pack (``triattention_cute_score_fused.py``, +score-only entry) is the ONLY score implementation. These tests drive it +through ``_FixedScoreGroup.launch`` -- the exact production entry point -- +across the supported production geometries, with multi-layer segments, +permuted page tables, ragged valid lengths, and per-request prompt windows, +and compare against a pure-PyTorch oracle that recomputes everything +independently. They also pin the loud-failure contract: unsupported +geometry, removed aggregations, and request counts beyond the group +capacity raise instead of routing to another kernel. """ import pytest @@ -194,9 +195,9 @@ def test_cute_kernel_matches_torch_oracle(self, case): list(range(num_layers)), ) - # The runner precompiles exactly the request counts production - # launches: one and the full group capacity. - for request_count in dict.fromkeys((1, max_requests)): + # The fused runner dispatches every request count up to the group + # capacity; cover one, an intermediate count, and the capacity. + for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): group.output.fill_(float("nan")) valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) scores = group.launch( @@ -229,11 +230,13 @@ def test_cute_kernel_matches_torch_oracle(self, case): ) @requires_sm100 - def test_uncompiled_request_count_raises(self): - """A request count outside the precompiled variants fails loudly. + def test_request_count_beyond_capacity_raises(self): + """A request count beyond the group capacity fails loudly. - There is no fallback kernel, so ``supports()`` misses must raise - instead of silently scoring through a slower path. + The fused runner dispatches every request count up to the group + capacity (covered by the oracle matrix above); there is no fallback + kernel, so anything beyond it must raise instead of silently + scoring through a slower path. """ pytest.importorskip("cutlass") ( @@ -248,9 +251,9 @@ def test_uncompiled_request_count_raises(self): valid_widths = torch.empty( _QWEN3_CASE["max_requests"], dtype=torch.int32, device=group.output.device ) - with pytest.raises(RuntimeError, match="no compiled variant"): + with pytest.raises(ValueError, match="exceeds fixed score capacity"): group.launch( - _QWEN3_CASE["max_requests"] - 1, + _QWEN3_CASE["max_requests"] + 1, valid_seq_lens, valid_widths, token_starts, From be073edd22e0cd402d27a5e932d6a596994df621 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 01:32:23 -0700 Subject: [PATCH 077/178] [None][fix] Fold 64-bit score offsets into the union finalizer's tail loads The union finalizer's per-token fallback loads indexed the score scratch with a flat Int64 that lowers through the DSL's 32-bit dynamic coordinate. At the qwen3-8B serve geometry (request capacity 64 with the 16384-token score bucket) the scratch spans 2,415,919,104 elements, so the last KV head's group planes sit past 2^31 and the wrapped loads land about 8 GiB below the scratch: an illegal memory access whenever any request's pinned prompt length is not a multiple of the lane width, which real serve cohorts hit on every eviction round. The vectorized branch already folds the same offset into an explicit 64-bit pointer; the fallback now builds the identical pointer and indexes only within the lane, so both branches read exactly the same elements and in-range behavior is unchanged. This was the reproducible trtllm-serve qwen3-8B mbs=64 eviction crash. The async traceback blamed the following top-k launch, but CUDA_LAUNCH_BLOCKING pinned the union finalizer, and a standalone one-round repro crashed and cleared on exactly the window-start alignment bit at the same geometry. Bench and offline runs never hit the branch past 2^31 because their synthetic window starts are lane-aligned, and the poisson serve arm ran the 8192-token bucket, below 2^31. The new giant-scratch regression legs drive the fused pipeline at the production shape (capacity 64, 36 layers, bucket 16384, unaligned per-request window starts, full-capacity launch with zero-length tail rows) plus a capacity-32 control below 2^31, and check both against the split score gather plus the standalone union reference. The full unit suite is green (142 passed, including the two new legs). The serve repro that previously died at its first eviction bands (natural AIME pool-32, mbs=64) ran past every band on this fix with zero failures and 80 natural completions before the watchdog cap. Signed-off-by: tianruih --- .../triattention_cute_selection.py | 22 ++- .../test_triattention_cute_union_fusion.py | 153 ++++++++++++++++++ 2 files changed, 172 insertions(+), 3 deletions(-) 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 index 3bbc38761b4e..304778f061c4 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -391,12 +391,28 @@ def kernel( cute.coalesce(score_value_tiles[token_subtile]), ) else: + # Fold the 64-bit flat index into the pointer BEFORE the + # per-element loads, exactly like the vectorized branch: + # the score scratch exceeds 2^31 elements at large + # request counts, and indexing ``scores`` with the flat + # Int64 goes through the DSL's 32-bit dynamic coordinate, + # which wraps. This branch runs whenever a request's + # window start is not lane-aligned (any pinned prompt + # length not divisible by ``tokens_per_lane``), so real + # serve cohorts hit it on every eviction round. + score_tail = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + (scores.iterator + score_index).toint(), + AddressSpace.gmem, + assumed_align=4, + ), + cute.make_layout(self.tokens_per_lane), + ) for token_slot in cutlass.range_constexpr(self.tokens_per_lane): token = subtile_first_token + token_slot if cutlass.dynamic_expr(token < valid_width): - score_value_tiles[token_subtile][token_slot] = scores[ - score_index + token_slot - ] + score_value_tiles[token_subtile][token_slot] = score_tail[token_slot] else: score_value_tiles[token_subtile][token_slot] = cutlass.Float32( float("-inf") 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 index bf22373277a5..39d00a70dbe5 100644 --- 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 @@ -449,3 +449,156 @@ def test_union_fusion_rejects_unsupported_frequency_count() -> None: pool_dtype=cutlass.BFloat16, 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 standalone union reference. + """ + pytest.importorskip("cutlass") + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, + ) + + 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) + round_starts = torch.arange(max_requests, dtype=torch.float32, device=device) + seq_len + phase = (round_starts[:, 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() + + page_ids = torch.arange(num_pages, dtype=torch.int32, device=device) + block_offsets = torch.zeros(1, max_requests, 2, num_pages, dtype=torch.int32, device=device) + block_offsets[0, :request_count, 0] = 2 * page_ids + block_offsets[0, :request_count, 1] = 2 * page_ids + 1 + group = _FixedScoreGroup( + layer_pools, + list(range(num_layers)), + max_requests, + num_pages, + seq_len, + num_q_heads, + block_offsets, + [0] * num_layers, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + output_width=decode_window, + ) + assert (num_kv_heads * 8 * max_requests * num_layers * seq_len > 2**31) == (max_requests == 64) + + valid_seq_lens = torch.zeros(max_requests, dtype=torch.int32, device=device) + token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) + valid_seq_lens[:request_count] = torch.tensor(valid_lens, dtype=torch.int32, device=device) + token_starts[:request_count] = torch.tensor(score_starts, dtype=torch.int32, device=device) + + # Reference: the split score gather over the same decode windows, then + # the standalone normalize-and-union copy. + split_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) + per_head = group.launch( + request_count, + valid_seq_lens, + split_widths, + token_starts, + mean_cos, + mean_sin, + ) + rows = per_head.shape[1] * per_head.shape[2] + scores_rows = per_head.reshape(request_count, rows, decode_window).contiguous() + row_mean = torch.empty((request_count, rows, 1), dtype=torch.float32, device=device) + row_inv_std = torch.empty_like(row_mean) + expected = torch.empty((request_count, decode_window), dtype=torch.float32, device=device) + _reference_prepare_union_scores( + scores_rows, + split_widths, + row_mean, + row_inv_std, + expected, + request_count, + normalize_scores=True, + ) + + # 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 + ) + group.launch_cute_union_fusion( + max_requests, + valid_seq_lens, + fused_widths, + token_starts, + 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, + ) From 98bab8b7b98242e3f4e1572a6bc84d006b08ae25 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 01:38:59 -0700 Subject: [PATCH 078/178] [None][fix] Skip top-k sentinel lanes in the settle threshold gather The top-k pads any row shorter than KEEP_COUNT with -1 sentinels (a zero-width padded row is all sentinels), and the settle's threshold pass gathered scores through those indices unguarded, reading the element before the row -- for row zero, four bytes before the score buffer. The read never changed an outcome: a sentinel lane can only drag the threshold at or below the row minimum, which still emits every real ordinal of such a short row, and full rows never carry sentinels. Guard the gather on ``token_index >= 0`` so sentinel lanes stop dereferencing and stop joining the threshold; rows without sentinels take exactly the same loads as before, bit for bit. The new unit drives the settle with production-shaped sentinel padding (zero-width, short, boundary, and full rows) and requires the emitted ordinals plus the untouched tail slots to match the reference exactly. The full unit suite is green (143 passed, including the new sentinel leg). The token-identity digest gate (qwen3-8B and gpt-oss-20b union plus qwen3-8B per_head and per_layer_perhead, against an unmodified 56498df baseline) runs on this bit-neutral stack before any numerics-affecting commit lands on top of it. Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 8 +- .../test_triattention_fused_settle_pack.py | 81 +++++++++++++++++++ 2 files changed, 88 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 42172f036b66..73a644b17a28 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -932,9 +932,15 @@ def _settle_ties_and_pack_compaction_sources_kernel( mask=selected_mask, other=0, ) + # Rows shorter than KEEP_COUNT arrive padded with -1 sentinels + # from the top-k's short-row path (zero-width padded rows are all + # sentinels). Mask those lanes out of the gather so no lane + # dereferences ``row_scores - 1`` and a sentinel never joins the + # threshold; rows without sentinels load exactly as before. + selected_valid = selected_mask & (token_index >= 0) selected_score = tl.load( row_scores + token_index, - mask=selected_mask, + mask=selected_valid, other=float("inf"), ).to(tl.float32) threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) 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 index 5a0a365d1c1b..8d945d689d6a 100644 --- 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 @@ -420,6 +420,87 @@ def test_fused_kernel_without_pack_matches_standalone_settle(): assert torch.equal(output_fused, output_reference) +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) + placeholder = row_lengths + _settle_ties_and_pack_compaction_sources_kernel[(rows_total, 1)]( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output, + placeholder, + placeholder, + placeholder, + placeholder, + placeholder, + WIDTH=width, + KEEP_COUNT=keep_count, + OUTPUT_WIDTH=keep_count, + SELECTION_ROWS=1, + DENSE_TOTAL=0, + SWA_TOTAL=0, + MOVE_CAPACITY=0, + NUM_KV_HEADS=1, + SWA_WINDOW=0, + UNION=False, + PER_LAYER=False, + HAS_SWA=False, + HAS_SETTLE=True, + HAS_PACK=False, + BLOCK=_BLOCK, + num_warps=_NUM_WARPS, + ) + 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" + + def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buffers(): """The handoff exports the live move buffers, drops the compaction-time dense pack launch, and the selector only accepts a packing that reads its From 2f03cd9ffe9d389345a17afef76e2ec0a1b85e75 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 01:52:23 -0700 Subject: [PATCH 079/178] [None][chore] Carry the due count on an ungated eviction-round NVTX range The eviction round's NVTX range was behind the TLLM_NVTX_DEBUG gate, so production nsys captures showed the round's kernels with no marker for how many requests the cohort carried. Emit the range unconditionally with the due count in the message: the eviction path runs outside CUDA-graph capture, so a dynamic message is safe, and the cost is one host-side f-string per eviction round. Any nsys capture now shows the per-round cohort size directly. The pipeline lifecycle test tracks eviction/resize ordering through nvtx_range, so its expected timeline gains the new range around the eviction dispatch. The full unit suite is green (143 passed). The change writes one host-side string and cannot affect generations; the token-identity digest gate runs on the two fix commits below it. Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 9 ++++++++- .../kv_cache_compression/test_triattention_pipeline.py | 2 ++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 4e14d5105b6a..92520ef5d986 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1462,7 +1462,14 @@ def _periodic_evict( if not due_requests: return num_layers = self._num_layers_from_manager() - with nvtx_range_debug("triattention.evict_request_group", color="purple"): + # Ungated NVTX with the due count in the message, so any nsys capture + # shows how many requests each eviction round carries. This path runs + # outside CUDA-graph capture, so the dynamic message is safe; the cost + # is one host-side f-string per eviction round. + with nvtx_range( + f"triattention.evict_request_group reqs={len(due_requests)}", + color="purple", + ): capacity_targets = self._evict_requests( due_requests, num_layers, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 7f02596cff33..e5e06af64224 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -440,7 +440,9 @@ def track_range(name, **kwargs): mgr.kv_cache_manager._stream.wait_event.assert_not_called() cache.resize.assert_called_once_with(1024 + 4096, None) assert timeline == [ + "enter:triattention.evict_request_group reqs=1", "compact_dispatch", + "exit:triattention.evict_request_group reqs=1", "enter:triattention.resize", "exit:triattention.resize", ] From f0d756d0ec5635f939c09afefc0ea3330f33274c Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 02:01:56 -0700 Subject: [PATCH 080/178] [None][perf] Emit the approximate FTZ square root in the fused score kernel CuTe DSL 4.5 renamed the approximate-sqrt control to ``fastmath``, but in 4.5.0 that keyword only attaches an MLIR fast-math flag and does not select the PTX approximate square root, so the fused score kernel's magnitude term lowered to full-precision RSQ sequences with range checks, branches, and Newton refinement. Emit ``sqrt.approx.ftz.f32`` directly through inline PTX, restoring the approximate sqrt the kernel was authored and validated with. The kernel change follows fanrongl's cutlass 4.5 sqrt patch (kvcache-v2-triattention-no-engine-cutlass45-sqrt.patch, applied verbatim; import order and one blank line adjusted for the formatter). This commit is numerics-affecting by design, so its acceptance basis differs from the bit-neutral commits below it: end-to-end digests are expected to drift in ulps rather than match token-for-token. Acceptance recorded instead as (a) canonical-cell keep-set invariance: qwen bs32 at window 1040, top-512 keep sets identical on all 32 rows, normalized union rows within 3.1e-6 absolute; (b) isolated K1 latency at that cell 1647.8us -> 1299.4us (-21.1%), matching fanrongl's formal candidate measurements (-39% at his 9216-token window, both against the same exact-sqrt baseline); (c) coherent end-to-end legs with eviction engaged rather than digest-matched runs; (d) the full unit suite green (143 passed) including the score-oracle tolerances. Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) 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 index 9ef401676e77..d289c632c592 100644 --- 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 @@ -22,8 +22,10 @@ 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 def _cute_sqrt_keyword_mode() -> str: @@ -51,6 +53,25 @@ def _cute_sqrt_keyword_mode() -> str: _CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() + +@dsl_user_op +def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: + """Emit the approximate FTZ square root missing from CuTe DSL 4.5.""" + 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 # Minimum tcgen05 MMA tile N: GQA groups below 8 ride zero-padded head # columns (see the weight-builder loop and the partial-stats epilogue). @@ -1286,12 +1307,7 @@ def kernel( ftz=self.magnitude_sqrt_ftz, ) elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): - # cutlass 4.5 renamed the approximate-sqrt control - # to ``fastmath``; map the measured approx choice - # onto it to preserve the authored behavior. - magnitude = cute.math.sqrt( - norm2, fastmath=self.sqrt_mode == "approx" - ) + magnitude = _sqrt_approx_ftz(norm2) else: # DSLs with neither spelling get the plain (IEEE) # sqrt, which is strictly MORE accurate than the From 6da91b63d5faf68361180e5d8c9672c7e7a4a4d2 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 06:36:07 -0700 Subject: [PATCH 081/178] [None][chore] Trim narration comments and drop dead score-runner wrappers Comment-only trims across the TriAttention module: deleted-kernel history, review attributions, and lineage notes; every geometry/stream/int64 contract stays. Removes two zero-reference TriAttentionCuteScoreRunner methods (launch_with_partial_stats, launch_union_finalize) kept only for isolated profiling; the private union finalizer they wrapped remains on the production launch path. Net -48 lines, no behavior change. Signed-off-by: tianruih --- .../triattention/compaction.py | 5 +- .../triattention/triattention.py | 17 ++--- .../triattention_cute_score_fused.py | 33 +-------- .../triattention/triattention_kernels.py | 67 ++++++++----------- .../test_triattention_fused_settle_pack.py | 2 +- 5 files changed, 38 insertions(+), 86 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 56bb83adf9c4..9576293b916f 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -224,9 +224,8 @@ def _compact_groups( class _MoveSourcePackArguments(NamedTuple): """One cache family's move-index packing, described as plain launch data. - The selection-side fused settle-and-pack launch (design suggested by - Fanrong Li, torch-graph review 2026-07-20) consumes this to pack the - dense/SWA move sources in the same kernel that finalizes the kept + The selection-side fused settle-and-pack launch consumes this to pack + the dense/SWA move sources in the same kernel that finalizes the kept ordinals, instead of a second launch at compaction time. """ diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 92520ef5d986..04c63c32327c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -223,8 +223,7 @@ def fuse_move_source_pack(self, pack_arguments) -> None: """Pack compaction move sources inside this selector's settle launch. ``pack_arguments`` is the dense/SWA packing description exported by - ``BatchedKVCacheCompaction.hand_move_source_pack_to_selection`` - (fusion suggested by Fanrong Li, torch-graph review 2026-07-20). The + ``BatchedKVCacheCompaction.hand_move_source_pack_to_selection``. The fused kernel reads back the kept ordinals it just wrote, so the packing must read this selector's own keep buffer, and the packing geometry must match the selection rows this selector settles. @@ -326,9 +325,7 @@ class _BatchedUnionKeepSetSelector(_BatchedKeepSetSelectorBase): The fused score+stats+union CuTe pipeline is THE union score producer: it writes the normalized per-request union rows straight into ``combined``, so this selector owns only the top-k settle-and-pack - stage. The split row-stats/union-reduce Triton launches were retired - with their buffers; their standalone copies live in the fused-pipeline - unit test as references. + stage. """ def __init__( @@ -725,7 +722,7 @@ def __init__( self.mean_phase_table = mean_phase_table # ONE fused group across ALL dense layers: segments carry their own # layer base address and page-table slot, so distinct per-layer - # storages/block tables no longer force one launch per storage group. + # storages/block tables share a single launch. _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} _page_table_slots = [self.representative_slots[_rep_of[layer]] for layer in dense_layers] self.fused_group = _FixedScoreGroup( @@ -778,8 +775,7 @@ def launch_prepared_union_fusion(self, union_out: torch.Tensor) -> None: """Run the fused score+stats+union pipeline over these buffers. This is THE union score path; there is deliberately no fallback. A - geometry or capacity the fused pipeline cannot serve raises loudly - instead of routing to the retired split launches. + geometry or capacity the fused pipeline cannot serve raises loudly. """ if not self._score_launcher_bound: raise RuntimeError("TriAttention score launcher is not bound") @@ -1116,9 +1112,8 @@ def __init__( ) self.normalize_scores = bool(normalize_scores) if self.eviction_mode == "union" and not self.normalize_scores: - # The fused score+stats+union CuTe pipeline is THE union path and - # always z-normalizes; the split un-normalized union launches - # were retired with the Triton/C++ score stacks. + # The fused score+stats+union CuTe pipeline is THE union path + # and always z-normalizes. raise ValueError( "TriAttention union eviction requires normalize_scores=True: " "the fused union pipeline always z-normalizes score rows" 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 index d289c632c592..3bf397b73282 100644 --- 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 @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """SM100 CuTe-DSL scorer for the TriAttention mean-score path. -This is the production specialization of the final workbench kernel — and -the ONLY score implementation: the per-head modes launch its score-only +This is the ONLY score implementation: the per-head modes launch its score-only entry and union eviction launches its fused score+stats+union pipeline. It uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, and producer-only page-ID lookahead. Geometries outside the exact contract @@ -1981,27 +1980,6 @@ def launch( stream, ) - def launch_with_partial_stats( - self, - request_count: int, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - ) -> tuple[torch.Tensor, int]: - """Launch score and write deterministic partial row statistics.""" - stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) - self._compiled_stats[request_count]( - *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), - *self._cute_tail, - request_count, - stream, - ) - row_count = request_count * self.num_layers * self.num_q_heads - page_shards = self._page_shards[request_count] - stats = self.partial_stats[: row_count * page_shards * 3] - return stats.view(row_count, page_shards, 3), page_shards - def supports_union_fusion(self, request_count: int) -> bool: """Return whether the score/stats/union pipeline was precompiled.""" return ( @@ -2048,12 +2026,3 @@ def _launch_union_finalize( request_count, stream, ) - - def launch_union_finalize( - self, - request_count: int, - union_scores: torch.Tensor, - ) -> None: - """Launch only the union finalizer for isolated validation and profiling.""" - stream = cuda.CUstream(torch.cuda.current_stream(union_scores.device).cuda_stream) - self._launch_union_finalize(request_count, union_scores, stream) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 73a644b17a28..5ae607986cdf 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -8,17 +8,12 @@ (``triattention_cute_score_fused.py``): mean aggregation, BF16 KV pools, head size 64/128, 32/128-token pages, GQA group 4 or 8, per-request score window starts. There is deliberately no other score path -- any geometry -outside that contract raises loudly at setup instead of routing to a slower -kernel (the original Triton score kernel, the C++ CUDA score stack, and the -single-shot CuTe score kernel have all been deleted). The per-head modes use -the pack's score-only entry; union eviction runs its fused score+stats+union -pipeline (with ``triattention_cute_selection.py``). The split Triton -row-stats/union-reduce launches were retired with the union fusion, and -their standalone copies live in the fused-pipeline unit test as references -(the row-stats kernel itself remains: the per-head modes still normalize -with it). The unit tests validate the CuTe kernels against independent -PyTorch oracles. Selection and compaction live in their respective runtime -modules. +outside that contract raises loudly at setup. The per-head modes use the +pack's score-only entry plus the row-stats kernel for normalization; union +eviction runs the fused score+stats+union pipeline (with +``triattention_cute_selection.py``). The unit tests validate the CuTe +kernels against independent PyTorch oracles. Selection and compaction live +in their respective runtime modules. House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). @@ -125,8 +120,8 @@ def ensure(self, rows: int) -> None: (target, self.omega.numel()), dtype=torch.float32, device=self.omega.device ) sin_table = torch.zeros_like(cos_table) - # Accumulate offset-by-offset in fp32, mirroring the retired - # per-round Triton kernel's summation order. + # Accumulate offset-by-offset in fp32 (fixed summation order keeps + # the table bit-stable across rebuilds). for offset in self._offset_values: phase = torch.outer(positions + offset, self.omega) cos_table += torch.cos(phase) @@ -344,10 +339,9 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> anchor = self.pointer_prefix[0] num_q_heads, num_kv_heads, num_freqs, tokens_per_block, kv_factor = self.geometry_args max_segments = self.max_requests * self.num_layers - # The retired single-shot kernel stored full unmasked compute tiles - # (64 tokens, or one page for 128-token pages), which forced - # tile-aligned buckets. The fused kernel masks its ragged tail, but - # the bucket contract is kept unchanged so the kernel swap cannot + # The fused kernel masks its ragged tail, but the scratch bucket + # contract stays tile-aligned (64 tokens, or one page for + # 128-token pages) so bucket geometry cannot # silently admit new geometry (production pow2 buckets satisfy it). score_tile_tokens = max(64, int(tokens_per_block)) supported = ( @@ -444,7 +438,7 @@ def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> self._cute_seg_out_offset = seg_out_offset self._cute_token_starts = token_starts self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) - # Fused score+stats+union pipeline (Fanrong Li's two-kernel scheme): + # Fused score+stats+union pipeline (two launches): # THE union path, built lazily on the first union launch. ONE runner # serves every cohort: the score window start is per-request runtime # metadata, not a compile-time constant. @@ -503,8 +497,8 @@ def launch( f"request_count={request_count} (capacity {self.max_requests}) " "and no other score path exists" ) - # Per-request decode widths for the selection reduce kernels; the - # deleted C++ score op used to write these (seq_len - token_start). + # Per-request decode widths (seq_len - token_start) for the + # selection reduce kernels. torch.sub( valid_seq_lens[:request_count], token_starts_device[:request_count], @@ -638,8 +632,7 @@ def launch_cute_union_fusion( Each request scores its own window (``token_starts_device`` carries the per-request pinned prompt lengths), so mixed-prompt cohorts are served directly. There is deliberately no fallback: an unsupported - geometry or request count raises loudly instead of routing to the - retired split score/row-stats/union launches. + geometry or request count raises loudly. """ self.prepare_cute_score(mean_cos, mean_sin) if request_count <= 0 or request_count > self.max_requests: @@ -893,23 +886,19 @@ def _settle_ties_and_pack_compaction_sources_kernel( """Settle one selection row's ties, then pack its compaction move sources. One program per (request, selection row). The first half settles the - provisional top-k: recover the top-k threshold - from the provisional selection, count the strictly greater scores, then - emit the kept ordinals in increasing order, rebased by the row's pinned - prompt length. With ``HAS_PACK`` the same program then packs the move - sources for the packed rows this selection row feeds: the kept ordinals - it just wrote, followed by the request's protected tail, plus the SWA - rows (latest window) under the - same conditions as the retired standalone kernel. Union selection has one row per - request feeding every KV head's packed row, so that single program writes - all of them. ``HAS_PACK=False`` compiles the second half away, leaving - exactly the settle stage; ``HAS_SETTLE=False`` compiles the first half - away instead, packing pre-settled ordinals read from - ``output_indices`` -- the draft co-compaction flow, whose keep set is - the target's and needs no settling. The pre-fusion standalone copies - live in the fused-kernel unit test as the bit-equality references. - Fusing the launches was suggested by Fanrong Li (torch-graph review - 2026-07-20). + provisional top-k: recover the threshold from the provisional + selection, count the strictly greater scores, then emit the kept + ordinals in increasing order, rebased by the row's pinned prompt + length. With ``HAS_PACK`` the same program then packs the move sources + for the packed rows this selection row feeds: the kept ordinals it + just wrote, the request's protected tail, plus the SWA rows (latest + window). Union selection has one row per request feeding every KV + head's packed row, so that single program writes all of them. + ``HAS_PACK=False`` compiles the second half away, leaving exactly the + settle stage; ``HAS_SETTLE=False`` compiles the first half away + instead, packing pre-settled ordinals read from ``output_indices`` -- + the draft co-compaction flow, whose keep set is the target's and needs + no settling. """ request = tl.program_id(0) selection_domain = tl.program_id(1) 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 index 8d945d689d6a..c2c9ac20f390 100644 --- 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 @@ -100,7 +100,7 @@ def _pack_compaction_sources_kernel( # --------------------------------------------------------------------------- # # Fused finalize: settle the top-k ties and pack the move indices in one # -# launch (fusion suggested by Fanrong Li, torch-graph review 2026-07-20). # +# launch. # # --------------------------------------------------------------------------- # From 346eb1a76984fb2a0a261b65ea383e57bb5532cb Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 06:57:44 -0700 Subject: [PATCH 082/178] [None][test] Slim TriAttention unit suite: torch oracles, table-driven cases, guard dedup Corner-test slimming tranche 1 (-1048 LOC, 5809 -> 4761 across the TriAttention unit surface; test callables 63 -> 40 in the consolidated files, all parametrized case coverage preserved unless noted): - Retired in-test Triton reference-kernel replicas: the pre-fusion settle/pack copies and the split row-stats/union-reduce references now have pure-torch oracles (settle/pack comparisons stay exact on integer outputs; union comparisons keep the existing tolerance). - Deleted the reference-vs-reference validation test the replicas needed. - Merged test_triattention_score_ops.py into test_triattention_cute_score.py with unified staging helpers. - Table-driven consolidation: draft-contract guards into one parametrized admission table; page-table copy/consumer-order pair; request-init, request-finish, prepare-growth, hook, and masked-SWA pairs; compact-op contract tests over a shared builder. - Guard dedup: one representative per reject family (compact-op geometry, union-fusion setup failure, frequency count). - Deleted tests subsumed by stronger peers: standalone top-k routing (heavy-ties selection oracle covers k in {4,4096,8192}), sorted-output and exact-indices smokes, oracle-self-consistency (tested test-helper algebra), two-group score oracle (grouped-pool staging still covered by the per-layer alignment test), bookkeeping-cumulative (superset in the co-compaction monotone test). - Incident anchors stay runnable: giant-scratch unaligned start (capacity 64 and 32), settle sentinel padding, per-request score-start rows, scratch-bucket asserts, overlap-tail exclusion, identity-compaction gate pair, two-rounds byte-exactness, mixed-prompt cohort, draft admission fail-closed gates; per_head/per_layer_perhead coverage stays runnable throughout. Signed-off-by: tianruih --- .../_torch/kv_cache_compression/conftest.py | 134 +++ .../test_triattention_cute_score.py | 278 +++++- .../test_triattention_cute_union_fusion.py | 333 ++----- .../test_triattention_draft_cocompaction.py | 164 ++-- .../test_triattention_fused_settle_pack.py | 310 +++---- .../test_triattention_pipeline.py | 816 +++--------------- .../test_triattention_score_ops.py | 318 ------- .../test_triattention_selection_compaction.py | 222 +---- .../serial/test_sparse_kv_cache_compact.py | 266 +++--- 9 files changed, 901 insertions(+), 1940 deletions(-) delete mode 100644 tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 16639539047a..015302f3aa78 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -89,6 +89,140 @@ def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): ) +def make_ramp_pools( + count, + *, + num_kv_heads=2, + pages=6, + tokens_per_block=32, + head_dim=64, + layer_stride=37, + base=0, + device=None, +): + """bf16 pools carrying a shifted ``arange % 251`` ramp payload. + + Every value is exact in bf16 and any wrong page/plane/head/token move + lands on a different byte pattern, so pool-equality checks in the + compaction tests stay conclusive. Geometry defaults to the compact op's + supported production shape (32-token pages, head_dim 64). + """ + return [ + ( + ( + torch.arange( + pages * 2 * num_kv_heads * tokens_per_block * head_dim, + dtype=torch.int32, + device=device, + ) + + base + + layer * layer_stride + ) + % 251 + ) + .view(pages, 2, num_kv_heads, tokens_per_block, head_dim) + .to(torch.bfloat16) + for layer in range(count) + ] + + +def build_compaction(**overrides): + """``BatchedKVCacheCompaction`` with the suite's default 2-layer geometry.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( + BatchedKVCacheCompaction, + ) + + args = dict( + eviction_mode="union", + dense_layers=[0, 1], + swa_layers=[], + layer_group_representative={0: 0, 1: 1}, + layer_pool_keys=[("dense", 0), ("dense", 0)], + page_table_slots={0: 0, 1: 0}, + request_count=2, + decode_keep_count=4, + swa_window=None, + ) + args.update(overrides) + return BatchedKVCacheCompaction(**args) + + +def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None): + """A ``__new__``-built staging shell for the bulk page-table copy tests.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _FixedScoreStagingBuffers, + ) + + staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + staging.device = device + staging.max_requests = max_requests + staging.copy_block_count = copy_block_count + if page_count is not None: + staging.page_count = page_count + staging.bulk_copy_done = torch.cuda.Event() + staging.bulk_consume_done = torch.cuda.Event() + staging.page_tables_active = False + staging.copy_done = torch.cuda.Event() + staging.copy_pending = False + staging._bulk_offsets_src = torch.empty( + 1, max_requests, 2, copy_block_count, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging._bulk_copy_idx_src = torch.arange( + max_requests, dtype=torch.int32, device="cpu", pin_memory=True + ) + staging.block_offsets_device = torch.empty( + 1, max_requests, 2, copy_block_count, dtype=torch.int32, device=device + ) + return staging + + +def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): + """The manager surface ``_stage_page_tables_bulk``/``stage`` consume.""" + 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_fixed_resources_stubs(manager, *, decode_width=260): + """Stub the calibration/staging surfaces around ``_fixed_resources_for``.""" + manager._H = 2 + manager._F = 2 + manager._freq_scale_sq = torch.ones(2) + manager._offsets = torch.ones(2) + manager.calibration = {"omega": torch.ones(2)} + manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) + manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) + pool = torch.empty(8, 2, 1, 4, 4) + layout = SimpleNamespace( + manager=SimpleNamespace(num_pools=1), + num_layers=2, + global_layers=[0, 1], + layer_pools=[pool, pool], + dense_layers=[0, 1], + swa_layers=[], + storage_groups={0: [0, 1]}, + pool_view_fingerprint=(("fixed",),), + ) + score_staging = SimpleNamespace( + fused_group=SimpleNamespace(output=torch.empty(8, 4, decode_width)), + bind_score_launcher=mock.Mock(), + token_starts_device=torch.zeros(8, dtype=torch.int32), + decode_width=decode_width, + page_table_token_capacity=65537, + max_requests=8, + ) + keep_set_selector = SimpleNamespace( + valid_widths=torch.empty(8, dtype=torch.int32), + top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), + ) + return layout, score_staging, keep_set_selector + + 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 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 index 88f022640dfe..82e71ab40b33 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -1,15 +1,37 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Correctness coverage for the SM100 TriAttention CuTe scorer (the only score path).""" +"""The SM100 TriAttention CuTe scorer (the only score path) vs PyTorch oracles. + +Two layers of coverage over ``_FixedScoreGroup.launch`` -- the exact +production entry point. The kernel-numerics matrix drives a single-layer +group across the supported page geometries (permuted physical pages, ragged +valid lengths, GQA group 4 riding the padded MMA tile) against inline +oracle math. The launch-path matrix drives multi-layer groups across the +named production geometries (Qwen3, GPT-OSS, the originally validated +128-token-page shape) against the shared pure-PyTorch oracle, sweeps +request counts up to the group capacity, and checks the per-request +decode-width metadata the selection reduce kernels consume. The contract +tests pin the loud-failure behavior: unsupported geometry, removed +aggregations, and request counts beyond capacity raise -- there is no +fallback score kernel. +""" import pytest import torch +from conftest import encode_block_offsets as _encode_block_offsets +from conftest import torch_tri_score_oracle as _torch_tri_score_oracle +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _FixedScoreGroup, +) -@pytest.mark.skipif( +requires_sm100 = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), - reason="TriAttention CuTe score kernel requires SM100", + reason="TriAttention score requires SM100", ) + + +@requires_sm100 @pytest.mark.parametrize( "tokens_per_block,page_permutation,valid_lens,num_freqs,num_q_heads", [ @@ -36,10 +58,6 @@ def test_cute_score_matches_torch_mean_oracle( ) -> None: pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) - torch.manual_seed(20260720) device = torch.device("cuda") seq_len = 256 @@ -127,3 +145,249 @@ def test_cute_score_matches_torch_mean_oracle( torch.cuda.synchronize() # The CuTe runner is the only score path; prove setup actually built it. assert group._cute_score_runner is not None + + +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 + # The 0.125 scaling keeps the BF16 key/coefficient products small so the + # kernel-vs-oracle tolerance can stay tight across the frequency sum. + 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 + group = _FixedScoreGroup( + pools, + list(range(num_layers)), + max_requests, + page_count, + capacity, + num_q_heads, + _encode_block_offsets(page_ids), + [0] * num_layers, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets_t, + output_width=capacity - prompt_len, + ) + round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() + token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) + # Ragged valid lengths whose tails land mid-page and mid-compute-tile. + seq_lens = [capacity - ((request * 3) % 5) for request in range(max_requests)] + valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + phase = (round_starts.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() + # Everything the PyTorch oracle needs to rebuild the reference leg + # independently (it recomputes its own mean phases from these). + 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 ( + group, + token_starts, + valid_seq_lens, + 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: the Qwen3 shape (64 +# frequencies, GQA group 4 riding the padded MMA tile), the GPT-OSS shape +# (32 frequencies, group 8, 32-token pages spanning two page fragments per +# compute tile), and the originally validated 128-token-page shape. +_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"] + ( + group, + token_starts, + valid_seq_lens, + seq_lens, + mean_cos, + mean_sin, + oracle_inputs, + ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) + device = group.output.device + + oracle = _torch_tri_score_oracle( + group._cute_layer_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)), + ) + + # The fused runner dispatches every request count up to the group + # capacity; cover one, an intermediate count, and the capacity. + for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): + group.output.fill_(float("nan")) + valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) + scores = group.launch( + request_count, + valid_seq_lens, + valid_widths, + token_starts, + mean_cos, + mean_sin, + ) + assert scores.shape == ( + request_count, + num_layers, + case["num_q_heads"], + group.output_width, + ) + # The launch owns the per-request decode widths the selection + # reduce kernels consume. + assert valid_widths[: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 _tiny_unsupported_group(dtype: torch.dtype): + """A geometry far outside the CuTe contract (constructor accepts it).""" + 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 + pools = [ + torch.randn(max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device).to( + dtype + ) + for _ in range(num_layers) + ] + page_ids = ( + torch.arange(max_requests * page_count, device=device) + .view(max_requests, page_count) + .contiguous() + ) + capacity = page_count * tokens_per_block + group = _FixedScoreGroup( + pools, + list(range(num_layers)), + max_requests, + page_count, + capacity, + 2, + _encode_block_offsets(page_ids), + [0] * num_layers, + torch.randn(num_layers, 2, num_freqs, device=device), + torch.randn(num_layers, 2, num_freqs, device=device), + torch.randn(num_layers, 2, num_freqs, device=device), + torch.rand(num_freqs, device=device) + 0.5, + torch.rand(num_freqs, device=device) * 0.05, + torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), + output_width=capacity - 1, + ) + device_args = dict(dtype=torch.int32, device=device) + return group, ( + torch.full((max_requests,), capacity, **device_args), + torch.empty(max_requests, **device_args), + torch.ones(max_requests, **device_args), + torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), + torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), + ) + + +# The loud-failure contract, one representative per guard family. All three +# guards fire before any kernel work, in this order: removed aggregation, +# request count beyond the group capacity (previously exercised on a +# production-shaped SM100 group; the check is layered before compilation, so +# the tiny group covers the same code path), unsupported geometry at setup. +@pytest.mark.parametrize( + "dtype,request_count,launch_kwargs,match", + [ + pytest.param( + torch.float32, 1, {}, "TriAttention score requires SM100", id="unsupported_geometry" + ), + pytest.param( + torch.bfloat16, 1, {"aggregation": "max"}, "max aggregation", id="max_aggregation" + ), + pytest.param(torch.bfloat16, 3, {}, "exceeds fixed score capacity", id="beyond_capacity"), + ], +) +def test_score_launch_contract_raises(dtype, request_count, launch_kwargs, match): + group, launch_args = _tiny_unsupported_group(dtype) + with pytest.raises(ValueError, match=match): + group.launch(request_count, *launch_args, **launch_kwargs) 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 index 39d00a70dbe5..2996344d697f 100644 --- 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 @@ -1,134 +1,43 @@ # 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 (two CuTe kernels).""" +"""Equivalence coverage for the fused score+stats+union pipeline (two CuTe kernels). + +The reference side gathers the SAME production score rows (the fused pack's +score-only entry) and normalizes + union-reduces them with a pure-torch +float32 oracle. The fused-vs-reference comparison was always tolerance-based +(the fused pipeline's reduction order differs from any reference); the +tolerances are unchanged from the retired Triton reference copies. +""" import pytest import torch -import triton -import triton.language as tl _SM100_ONLY = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), reason="TriAttention CuTe kernels require SM100", ) -# --------------------------------------------------------------------------- # -# Standalone reference copies of the RETIRED split-union launches. The fused -# score+stats+union CuTe pipeline is THE production union path; these -# pre-retirement Triton copies exist only as the equivalence references here -# (precedent: the standalone settle/pack copies in -# test_triattention_fused_settle_pack.py). -# --------------------------------------------------------------------------- # - - -@triton.jit -def _reference_score_row_stats_kernel( - scores, - valid_widths, - row_mean, - row_inv_std, - ROWS: tl.constexpr, - WIDTH: tl.constexpr, - BLOCK: tl.constexpr, -): - """Compute one valid-prefix mean and inverse standard deviation per score row.""" - flat_row = tl.program_id(0) - request = flat_row // ROWS - valid_width = tl.load(valid_widths + request) - score_row = scores + flat_row * WIDTH - lane = tl.arange(0, BLOCK) - score_sum = 0.0 - for start in tl.static_range(0, WIDTH, BLOCK): - token = start + lane - valid = token < valid_width - value = tl.load(score_row + token, mask=valid, other=0.0).to(tl.float32) - score_sum += tl.sum(value, axis=0) - mean = score_sum / valid_width - square_sum = 0.0 - for start in tl.static_range(0, WIDTH, BLOCK): - token = start + lane - valid = token < valid_width - 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 / valid_width) - tl.store(row_mean + flat_row, mean) - tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) - - -@triton.jit -def _reference_score_union_kernel( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - ROWS: tl.constexpr, - WIDTH: tl.constexpr, - NORMALIZE: tl.constexpr, - BLOCK: tl.constexpr, -): - """Normalize score rows and reduce them directly to one request-level union.""" - request = tl.program_id(0) - token = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - valid_width = tl.load(valid_widths + request) - valid_token = token < valid_width - union_max = tl.full((BLOCK,), -float("inf"), tl.float32) - for row in tl.range(0, ROWS): - flat_row = request * ROWS + row - value = tl.load( - scores + flat_row * WIDTH + 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")) - union_max = tl.maximum(union_max, value) - tl.store(combined + request * WIDTH + token, union_max, mask=token < WIDTH) - - -def _reference_prepare_union_scores( - scores: torch.Tensor, - valid_widths: torch.Tensor, - row_mean: torch.Tensor, - row_inv_std: torch.Tensor, - combined: torch.Tensor, - request_count: int, - *, - normalize_scores: bool, -) -> None: - """Mask, normalize, and union-reduce score rows in one or two launches.""" - request_count = int(request_count) - assert scores.is_cuda and scores.ndim == 3 and scores.dtype == torch.float32 - assert scores.is_contiguous() and request_count == scores.shape[0] - _, rows, width = scores.shape - stats_block = 256 - if normalize_scores: - _reference_score_row_stats_kernel[(request_count * rows,)]( - scores, - valid_widths, - row_mean, - row_inv_std, - ROWS=rows, - WIDTH=width, - BLOCK=stats_block, - num_warps=4, - ) - union_block = 32 - _reference_score_union_kernel[(request_count, triton.cdiv(width, union_block))]( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - ROWS=rows, - WIDTH=width, - NORMALIZE=normalize_scores, - BLOCK=union_block, - num_warps=1, + +def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: + """Pure-torch union oracle: z-normalize each row's valid prefix, union-max. + + Mirrors the production union semantics: per-row mean and biased std over + the valid prefix (std clamped at 1e-6), then the per-token maximum across + the request's rows; tokens past the valid width stay ``-inf``. + """ + 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): + valid_width = int(valid_widths[request]) + if valid_width <= 0: + continue + valid = scores_rows[request, :, :valid_width].to(torch.float32) + mean = valid.mean(dim=1, keepdim=True) + std = ((valid - mean).square().sum(dim=1, keepdim=True) / valid_width).sqrt() + combined[request, :valid_width] = ((valid - mean) / std.clamp_min(1e-6)).amax(dim=0) + return combined def _check_union_fusion_matches_split_pipeline( @@ -138,14 +47,12 @@ def _check_union_fusion_matches_split_pipeline( score_starts: "int | list", valid_lens: "list | None", ) -> None: - """The fused pipeline must reproduce the split score->stats->union rows. + """The fused pipeline must reproduce the split score->normalize->union rows. ``score_starts`` is either one uniform window start or a per-request list (the fused kernels read the start per request at runtime). The - reference runs the retired split path: the score launch gathers each - request's decode window, then the standalone - ``_reference_prepare_union_scores`` copy normalizes rows and takes the - cross-row union maximum. + reference leg runs the production score-only launch over the same decode + windows, then the pure-torch union oracle. """ pytest.importorskip("cutlass") @@ -207,7 +114,8 @@ def _check_union_fusion_matches_split_pipeline( score_starts = [score_starts] * request_count assert len(score_starts) == request_count - # Reference: the split pipeline over the same decode windows. + # 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) token_starts = torch.tensor(score_starts, dtype=torch.int32, device=device) per_head = group.launch( @@ -220,18 +128,7 @@ def _check_union_fusion_matches_split_pipeline( ) rows = per_head.shape[1] * per_head.shape[2] scores_rows = per_head.reshape(request_count, rows, seq_len).contiguous() - row_mean = torch.empty((request_count, rows, 1), dtype=torch.float32, device=device) - row_inv_std = torch.empty_like(row_mean) - expected = torch.empty((request_count, seq_len), dtype=torch.float32, device=device) - _reference_prepare_union_scores( - scores_rows, - split_widths, - row_mean, - row_inv_std, - expected, - request_count, - normalize_scores=True, - ) + expected = _reference_union_scores(scores_rows, split_widths) fused_widths = torch.empty(request_count, dtype=torch.int32, device=device) fused_out = torch.full( @@ -276,6 +173,11 @@ def _check_union_fusion_matches_split_pipeline( (32, 64, 4, 37, [250, 198]), (128, 64, 4, 0, None), (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 cohorts: each request scores its own window (one # start mid-tile, one page-aligned) — the case the fused pipeline # previously declined. @@ -298,71 +200,37 @@ def test_union_fusion_matches_split_pipeline( @_SM100_ONLY -def test_union_fusion_engages_gqa4_narrow_heads() -> None: - """GQA group 4 with 32 frequencies (the formerly declined geometry) engages. - - The fused kernel pads group-4 head columns 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 — so this mixed geometry must launch and match the split path. +@pytest.mark.parametrize("guard", ["runner_construction_failure", "frequency_count"]) +def test_union_fusion_guards_raise(guard: str, monkeypatch: pytest.MonkeyPatch) -> None: + """One representative per fused-pipeline guard family raises loudly. + + ``runner_construction_failure``: a fused-runner construction failure + surfaces as the no-fallback RuntimeError at score setup + (``prepare_cute_score`` runs before the union runner is built). + ``frequency_count``: 16 frequencies (head size 32) sit outside the fused + kernel contract and are rejected at kernel construction. """ - _check_union_fusion_matches_split_pipeline(128, 32, 4, 0, None) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def test_reference_union_preparation_matches_ragged_torch_reference() -> None: - """The standalone reference copy must match a pure-torch oracle. - - Moved from the selection/compaction suite when the production - ``prepare_union_scores`` was retired: this keeps the reference copy the - fused-pipeline equivalence tests compare against honest. - """ - device = torch.device("cuda", torch.cuda.current_device()) - request_count, rows, width = 2, 7, 97 - generator = torch.Generator(device=device).manual_seed(17) - scores = torch.randn( - request_count, - rows, - width, - generator=generator, - dtype=torch.float32, - device=device, - ) - valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) - row_mean = torch.empty(request_count, rows, 1, dtype=torch.float32, device=device) - row_inv_std = torch.empty_like(row_mean) - combined = torch.empty(request_count, width, device=device) - - _reference_prepare_union_scores( - scores, - valid_widths, - row_mean, - row_inv_std, - combined, - request_count, - normalize_scores=True, - ) - torch.cuda.synchronize(device) - - expected = torch.full_like(combined, float("-inf")) - for request, valid_width in enumerate(valid_widths.tolist()): - valid_scores = scores[request, :, :valid_width] - mean = valid_scores.mean(dim=1, keepdim=True) - std = torch.linalg.vector_norm(valid_scores - mean, dim=1, keepdim=True) - std = (std / valid_width**0.5).clamp_min(1e-6) - expected[request, :valid_width] = ((valid_scores - mean) / std).amax(dim=0) - assert torch.allclose(combined, expected, rtol=2e-5, atol=2e-5) - + cutlass = pytest.importorskip("cutlass") -@_SM100_ONLY -def test_union_fusion_setup_failure_raises(monkeypatch: pytest.MonkeyPatch) -> None: - """A fused-runner construction failure raises loudly: no fallback remains. + if guard == "frequency_count": + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused import ( # noqa: E501 + _TriAttentionScoreKernel, + ) - The single fused pack serves both the score-only and union entries, so - the construction failure surfaces at score setup (``prepare_cute_score`` - runs before the union runner is built). - """ - pytest.importorskip("cutlass") + with pytest.raises(ValueError, match="frequencies"): + _TriAttentionScoreKernel( + num_layers=1, + seq_len=256, + num_q_heads=8, + num_kv_heads=1, + num_freqs=16, + tokens_per_block=128, + pool_shape=(2, 2, 1, 128, 32), + pool_strides=(8192, 4096, 4096, 32, 1), + pool_dtype=cutlass.BFloat16, + page_shards=3, + ) + return import tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused as fused_module # noqa: E501 from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( @@ -371,20 +239,13 @@ def test_union_fusion_setup_failure_raises(monkeypatch: pytest.MonkeyPatch) -> N torch.manual_seed(20260721) device = torch.device("cuda") - seq_len = 256 - tokens_per_block = 128 - num_freqs = 32 - num_q_heads = 8 + seq_len, tokens_per_block, num_freqs, num_q_heads = 256, 128, 32, 8 num_pages = seq_len // tokens_per_block 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) - 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) - mean_cos = torch.cos(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) - mean_sin = torch.sin(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) k_plane = [2 * page for page in range(num_pages)] v_plane = [2 * page + 1 for page in range(num_pages)] block_offsets = torch.tensor( @@ -402,9 +263,9 @@ def test_union_fusion_setup_failure_raises(monkeypatch: pytest.MonkeyPatch) -> N q_real, torch.randn_like(q_real) * 0.125, torch.randn_like(q_real) * 0.125, - freq_scale_sq, + torch.linspace(0.5, 1.5, num_freqs, device=device), omega, - offsets, + torch.tensor([1.0, 2.0, 4.0], device=device), output_width=seq_len, ) @@ -412,42 +273,17 @@ def _refuse_construction(**_kwargs): raise ValueError("synthetic fused-runner construction failure") monkeypatch.setattr(fused_module, "TriAttentionCuteScoreRunner", _refuse_construction) - valid_seq_lens = torch.full((2,), seq_len, dtype=torch.int32, device=device) - widths = torch.empty(2, dtype=torch.int32, device=device) - token_starts = torch.zeros(2, dtype=torch.int32, device=device) - union_out = torch.empty((2, seq_len), dtype=torch.float32, device=device) + mean_cos = torch.cos(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) + mean_sin = torch.sin(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) with pytest.raises(RuntimeError, match="no other score path exists"): group.launch_cute_union_fusion( 2, - valid_seq_lens, - widths, - token_starts, + torch.full((2,), seq_len, dtype=torch.int32, device=device), + torch.empty(2, dtype=torch.int32, device=device), + torch.zeros(2, dtype=torch.int32, device=device), mean_cos.contiguous(), mean_sin.contiguous(), - union_out, - ) - - -def test_union_fusion_rejects_unsupported_frequency_count() -> None: - """16 frequencies (head size 32) sit outside the fused kernel contract.""" - cutlass = 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, - seq_len=256, - num_q_heads=8, - num_kv_heads=1, - num_freqs=16, - tokens_per_block=128, - pool_shape=(2, 2, 1, 128, 32), - pool_strides=(8192, 4096, 4096, 32, 1), - pool_dtype=cutlass.BFloat16, - page_shards=3, + torch.empty((2, seq_len), dtype=torch.float32, device=device), ) @@ -481,7 +317,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: 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 standalone union reference. + score-gather plus the pure-torch union oracle. """ pytest.importorskip("cutlass") @@ -553,7 +389,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: token_starts[:request_count] = torch.tensor(score_starts, dtype=torch.int32, device=device) # Reference: the split score gather over the same decode windows, then - # the standalone normalize-and-union copy. + # the pure-torch union oracle. split_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) per_head = group.launch( request_count, @@ -565,18 +401,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: ) rows = per_head.shape[1] * per_head.shape[2] scores_rows = per_head.reshape(request_count, rows, decode_window).contiguous() - row_mean = torch.empty((request_count, rows, 1), dtype=torch.float32, device=device) - row_inv_std = torch.empty_like(row_mean) - expected = torch.empty((request_count, decode_window), dtype=torch.float32, device=device) - _reference_prepare_union_scores( - scores_rows, - split_widths, - row_mean, - row_inv_std, - expected, - request_count, - normalize_scores=True, - ) + 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) 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 index 354a6ec790f9..0d8e8ca1de37 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -9,8 +9,9 @@ appended as ordinals ``valid_seq_len + 0..tail-1``, and both caches land at ``destination_base = prompt_len``. These tests cover the physical draft moves, the packed move indices, stream ordering across both cache managers, the -speculative admission gates, the published compressed-token invariant, and -prepared-compaction cache invalidation. +speculative admission gates (one representative per guard family), the +published compressed-token invariant, and prepared-compaction cache +invalidation. """ from types import SimpleNamespace @@ -18,16 +19,16 @@ import pytest import torch +from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_fixed_resources_stubs as _make_fixed_resources_stubs +from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_request as _make_request from conftest import make_triattention as _make_triattention from conftest import mocked_eviction_internals as _mocked_eviction_internals from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - BatchedKVCacheCompaction, -) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( TriAttention, _FixedScoreStagingBuffers, @@ -47,55 +48,19 @@ def _launched_draft_compaction(draft_protected_tails): """Build target and draft pools with distinct head counts, then compact. The compact op ships only the pipelined bf16 kernels, so the pools use - the supported production geometry (bf16, 32-token pages, head_dim 64). - Pool payloads are a shifted ``arange % 251`` ramp: every value is exact - in bf16 and any wrong page/plane/head/token move lands on a different - byte pattern, so the equality checks below stay conclusive. + the supported production geometry (bf16, 32-token pages, head_dim 64) and + the conclusive shifted ``arange % 251`` ramp payload. """ device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 - target_kv_heads = 2 - draft_kv_heads = 4 prompt_len = 2 - decode_keep_count = 4 - tokens_per_block = 32 - head_dim = 64 target_protected_tails = [2, 1] valid_seq_lens = [10, 9] target_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) draft_tables = torch.tensor([[1, 0, 2], [5, 4, 3]], dtype=torch.int32, device=device) - target_pools = [ - ( - ( - torch.arange( - 6 * 2 * target_kv_heads * tokens_per_block * head_dim, - dtype=torch.int32, - device=device, - ) - + layer * 37 - ) - % 251 - ) - .view(6, 2, target_kv_heads, tokens_per_block, head_dim) - .to(torch.bfloat16) - for layer in range(2) - ] - draft_pool = ( - ( - ( - torch.arange( - 6 * 2 * draft_kv_heads * tokens_per_block * head_dim, - dtype=torch.int32, - device=device, - ) - + 149 - ) - % 251 - ) - .view(6, 2, draft_kv_heads, tokens_per_block, head_dim) - .to(torch.bfloat16) - ) + target_pools = _make_ramp_pools(2, num_kv_heads=2, device=device) + draft_pool = _make_ramp_pools(1, num_kv_heads=4, base=149, device=device)[0] assert target_pools[0].shape[2] != draft_pool.shape[2] initial_target = [pool.clone() for pool in target_pools] initial_draft = draft_pool.clone() @@ -104,21 +69,13 @@ def _launched_draft_compaction(draft_protected_tails): # never appear in the selection rectangle. keep = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) - compaction = BatchedKVCacheCompaction( - eviction_mode="union", + compaction = _build_compaction( layer_pools=target_pools, - dense_layers=[0, 1], - swa_layers=[], - layer_group_representative={0: 0, 1: 1}, layer_pool_keys=[("pool", 0), ("pool", 0)], kept_token_ordinals=keep.to(torch.int32), valid_sequence_lengths=torch.tensor(valid_seq_lens, dtype=torch.int32, device=device), kv_block_offsets=_encode_block_offsets(target_tables), - page_table_slots={0: 0, 1: 0}, - request_count=request_count, prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), - decode_keep_count=decode_keep_count, - swa_window=None, protected_tail_capacity=max(target_protected_tails), draft_layer_pools=[draft_pool], draft_layers=[0], @@ -150,11 +107,14 @@ def _launched_draft_compaction(draft_protected_tails): ) -def test_draft_pools_receive_target_union_keep_set_and_own_tail(): - built = _launched_draft_compaction(draft_protected_tails=[1, 2]) +@pytest.mark.parametrize("draft_protected_tails", [[1, 1], [1, 2]]) +def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protected_tails): + built = _launched_draft_compaction(draft_protected_tails=draft_protected_tails) device = built.device prompt_len = built.prompt_len + expected_offsets = [0] + expected_moves = [] for request in range(built.request_count): valid = built.valid_seq_lens[request] # Target dense layers compact the union keep set plus the target tail. @@ -207,27 +167,13 @@ def test_draft_pools_receive_target_union_keep_set_and_own_tail(): before[:, head].index_select(1, draft_source), ) + expected_moves.append(draft_source.to(torch.int32)) + expected_offsets.append(expected_offsets[-1] + int(draft_source.numel())) -@pytest.mark.parametrize("draft_protected_tails", [[1, 1], [1, 2]]) -def test_draft_pack_matches_keep_broadcast_and_tail_ordinal_oracle(draft_protected_tails): - built = _launched_draft_compaction(draft_protected_tails=draft_protected_tails) + # The packed draft move indices must match the same broadcast-plus-tail + # oracle the physical moves followed. draft_compaction = built.compaction.draft_compaction - - expected_offsets = [0] - expected_moves = [] - for request in range(built.request_count): - decode = built.keep[request].to(torch.int32) - tail = torch.arange( - built.valid_seq_lens[request], - built.valid_seq_lens[request] + draft_protected_tails[request], - dtype=torch.int32, - device=built.device, - ) - moves = torch.cat((decode, tail)) - expected_moves.append(moves) - expected_offsets.append(expected_offsets[-1] + int(moves.numel())) expected_row = torch.cat(expected_moves) - assert draft_compaction.move_source_offsets.cpu().tolist() == expected_offsets draft_indices = draft_compaction.move_source_indices # The index buffer is sized for the widest tail (the capacity); this @@ -268,32 +214,47 @@ def test_mark_page_tables_consumed_orders_both_manager_streams(): @pytest.mark.parametrize( "gate,match", [ - ("union_only", "union"), + # One representative per guard family (the per-mode/per-config + # variants raise through the same checks). + ("union_only_per_head", "union"), + ("union_only_per_layer", "union"), ("draft_kv_factor", "standard key/value cache"), ("full_attention_draft", "full-attention draft"), - ("dflash", "standard paged cache compacted together"), + ("callsite_dflash", "standard paged cache compacted together"), + ("callsite_draft_target", "standard paged cache compacted together"), + ("callsite_pard", "standard paged cache compacted together"), ], ) def test_draft_admission_gates_raise(gate, match): draft_manager = _make_fake_v2(is_draft=True) if gate == "full_attention_draft": draft_manager.max_attention_window_vec = [128] - if gate == "dflash": - # DFlash reads cross-attention context buffers, not a paged KV cache; - # the call-site speculative gate rejects before any manager is - # created. + if gate.startswith("callsite_"): + # These draft contracts read cross-attention buffers or unvalidated + # paged tails; the call-site speculative gate rejects every one of + # them before any manager is created. from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, + DraftTargetDecodingConfig, + PARDDecodingConfig, TriAttentionKvCacheCompressionConfig, ) + spec_config = { + "callsite_dflash": lambda: DFlashDecodingConfig(max_draft_len=3), + "callsite_draft_target": lambda: DraftTargetDecodingConfig( + max_draft_len=3, + speculative_model="/tmp/draft-target-model", + ), + "callsite_pard": lambda: PARDDecodingConfig(max_draft_len=3), + }[gate]() with pytest.raises(ValueError, match=match): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 ), - DFlashDecodingConfig(max_draft_len=3), + spec_config, draft_manager, ) return @@ -301,7 +262,10 @@ def test_draft_admission_gates_raise(gate, match): _make_fake_v2(), top_B=8, model_path="/models/test", - eviction_mode="per_head" if gate == "union_only" else "union", + eviction_mode={ + "union_only_per_head": "per_head", + "union_only_per_layer": "per_layer_perhead", + }.get(gate, "union"), draft_kv_cache_manager=draft_manager, ) if gate == "draft_kv_factor": @@ -359,6 +323,10 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): confirmed = state.confirmed_kv_length cache.capacity = confirmed assert confirmed == 2 + 4 + # The staged logical position restores the uncompressed + # length: physical confirmed plus everything evicted so far. + prepared = internals.attach.call_args.args[0] + assert prepared[0].round_start == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. assert request.py_num_compressed_tokens == uncompressed - confirmed @@ -379,13 +347,7 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention(top_B=4) - manager._H = 2 - manager._F = 2 - manager._freq_scale_sq = torch.ones(2) - manager._offsets = torch.ones(2) - manager.calibration = {"omega": torch.ones(2)} - manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) - manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) + layout, score_staging, keep_set_selector = _make_fixed_resources_stubs(manager) draft_manager = _make_fake_v2(is_draft=True) draft_manager.num_pools = 1 manager.draft_kv_cache_manager = draft_manager @@ -398,30 +360,6 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): pool_view_fingerprint=(), ) ) - pool = torch.empty(8, 2, 1, 4, 4) - layout = SimpleNamespace( - manager=SimpleNamespace(num_pools=1), - num_layers=2, - global_layers=[0, 1], - layer_pools=[pool, pool], - dense_layers=[0, 1], - swa_layers=[], - storage_groups={0: [0, 1]}, - pool_view_fingerprint=(("fixed",),), - ) - - score_staging = SimpleNamespace( - fused_group=SimpleNamespace(output=torch.empty(8, 4, 260)), - bind_score_launcher=mock.Mock(), - token_starts_device=torch.zeros(8, dtype=torch.int32), - decode_width=260, - page_table_token_capacity=65537, - max_requests=8, - ) - keep_set_selector = SimpleNamespace( - valid_widths=torch.empty(8, dtype=torch.int32), - top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), - ) prepared = [ _PreparedEviction( request=_make_request(7), 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 index c2c9ac20f390..e5e35df1b007 100644 --- 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 @@ -1,24 +1,21 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The fused settle-and-pack kernel must reproduce the original two-kernel -sequence byte for byte. - -The reference legs are the pre-fusion kernels: the tie-settlement copy -kept in this file (its production original was deleted once the fused -kernel became the only launched settle path) and the move-source packing -kernel the module still ships for the draft flow: every case launches them on one set of buffers -and the fused kernel on an identically initialized set, then requires -``torch.equal`` on the kept ordinals, the dense move sources, and the SWA -move sources -- including the buffer regions neither path overwrites (rows -shorter than the keep count leave stale entries behind, and the packing -forwards those stale entries the same way in both paths). +"""The fused settle-and-pack kernel must reproduce the settle/pack semantics +exactly. + +The reference is a pure-torch oracle implementing the integer settle and +pack semantics (threshold recovery with sentinel-skip, strictly-greater +count, lowest-index tie quota, ascending prompt-rebased emission, then the +dense/SWA move-source packing). Every output is integer-valued, so the +comparisons remain ``torch.equal`` — including the buffer regions neither +path overwrites (rows shorter than the keep count leave stale entries +behind, and the packing forwards those stale entries the same way in both +paths). """ import pytest import torch -import triton -import triton.language as tl from conftest import encode_block_offsets as _encode_block_offsets from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( @@ -31,162 +28,84 @@ _settle_ties_and_pack_compaction_sources_kernel, ) +_BLOCK = 256 +_NUM_WARPS = 4 + + +def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, keep_count): + """Settle each row's provisional top-k in place (exact integer semantics). -@triton.jit -def _pack_compaction_sources_kernel( - selected_indices, + Threshold = min score over the provisional lanes (``-1`` sentinel lanes + are skipped, contributing +inf, so an all-sentinel row settles inertly); + keep every strictly-greater score in the valid width, then fill the + remaining quota with threshold ties in increasing index order; emit the + kept ordinals ascending, rebased by the row's pinned prompt length. + Output entries past the emitted count keep their previous (stale) value. + """ + 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 _pack_oracle( + settled, valid_seq_lens, dense_offsets, - dense_indices, + dense_out, swa_offsets, - swa_indices, - DENSE_TOTAL: tl.constexpr, - SWA_TOTAL: tl.constexpr, - SELECTION_ROWS: tl.constexpr, - SELECTION_STRIDE: tl.constexpr, - KEEP_COUNT: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - SWA_WINDOW: tl.constexpr, - UNION: tl.constexpr, - PER_LAYER: tl.constexpr, - HAS_SWA: tl.constexpr, - BLOCK: tl.constexpr, -): - """Pack selected decode ordinals and protected tails for the C++ updater.""" - request = tl.program_id(0) - domain = tl.program_id(1) - move = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) - - dense_begin = tl.load(dense_offsets + request) - dense_end = tl.load(dense_offsets + request + 1) - dense_count = dense_end - dense_begin - seq_len = tl.load(valid_seq_lens + request) - - if UNION: - selection_domain = 0 - else: - selection_domain = domain - # Selection rows carry decode-only kept ordinals (already absolute), so - # rows are prompt-length independent and one cohort may mix prompt sizes. - selection_row = request * SELECTION_ROWS + selection_domain - selected = tl.load( - selected_indices + selection_row.to(tl.int64) * SELECTION_STRIDE + move, - mask=move < KEEP_COUNT, - other=0, - ) - dense_source = tl.where(move < KEEP_COUNT, selected, seq_len + move - KEEP_COUNT) - dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move - tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) - - if HAS_SWA: - # Per-layer selection has one dense domain per (layer, head). SWA uses - # one shared source row per head, so only the first layer writes it. - if PER_LAYER: - write_swa = domain < NUM_KV_HEADS - else: - write_swa = move >= 0 - swa_begin = tl.load(swa_offsets + request) - swa_end = tl.load(swa_offsets + request + 1) - swa_count = swa_end - swa_begin - head = domain % NUM_KV_HEADS - swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move - swa_source = seq_len - SWA_WINDOW + move - tl.store( - swa_indices + swa_output, - swa_source, - mask=write_swa & (move < swa_count), - ) - - -# --------------------------------------------------------------------------- # -# Fused finalize: settle the top-k ties and pack the move indices in one # -# launch. # -# --------------------------------------------------------------------------- # - - -@triton.jit -def _settle_ties_after_topk_kernel( - scores, - seq_lens, - prompt_offsets, - provisional_indices, - output_indices, - WIDTH: tl.constexpr, - KEEP_COUNT: tl.constexpr, - OUTPUT_WIDTH: tl.constexpr, - BLOCK: tl.constexpr, + swa_out, + *, + selection_rows, + keep_count, + num_kv_heads, + swa_window, + union, + per_layer, + has_swa, ): - """Resolve boundary ties and emit increasing physical token indices. - - Pre-fusion standalone kept verbatim as the fused kernel's bit-equality - reference; the production module ships only the fused launch. + """Pack dense/SWA move sources from the settled ordinals (in place). + + Dense rows forward the settled output content verbatim for the first + ``keep_count`` moves (stale entries included, exactly like the kernel's + unconditional gather) and append the protected tail + ``seq_len + move - keep_count``; SWA rows write the latest-window + ordinals once per KV head (per-layer packs share one SWA row per head, + written only by the first layer's domains). """ - row = tl.program_id(0) - row_scores = scores + row * WIDTH - row_selected = provisional_indices + row * KEEP_COUNT - row_output = output_indices + row * OUTPUT_WIDTH - # Scores are decode-relative; this row's pinned prompt length rebases the - # emitted ordinals to absolute positions (per row, so one launch may mix - # prompt lengths). - prompt_len = tl.load(prompt_offsets + row) - - 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, - ) - selected_score = tl.load( - row_scores + token_index, - mask=selected_mask, - other=float("inf"), - ).to(tl.float32) - threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) - - seq_len = tl.load(seq_lens + 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 < seq_len) - 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 < seq_len) - 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_len, - mask=selected, - ) - output_count += tl.sum(selected_i32) - ties_seen += tl.sum(tied_i32) - - -_BLOCK = 256 -_NUM_WARPS = 4 + request_count = int(valid_seq_lens.shape[0]) + packed_rows = int(dense_out.shape[0]) + dense_total = int(dense_out.shape[1]) + swa_total = int(swa_out.shape[1]) + for request in range(request_count): + seq_len = int(valid_seq_lens[request]) + dense_begin = int(dense_offsets[request]) + dense_count = int(dense_offsets[request + 1]) - dense_begin + for domain in range(packed_rows): + selection_domain = 0 if union else domain + settled_row = settled[request * selection_rows + selection_domain] + for move in range(dense_count): + value = int(settled_row[move]) if move < keep_count else seq_len + move - keep_count + dense_out.view(-1)[domain * dense_total + dense_begin + move] = value + if has_swa and (domain < num_kv_heads if per_layer else True): + swa_begin = int(swa_offsets[request]) + swa_count = int(swa_offsets[request + 1]) - swa_begin + head = domain % num_kv_heads + for move in range(swa_count): + swa_out.view(-1)[head * swa_total + swa_begin + move] = ( + seq_len - swa_window + move + ) def _selection_rows_for(eviction_mode: str, num_layers: int, num_kv_heads: int) -> int: @@ -215,7 +134,7 @@ def _staged_offsets(counts, device): (350, 300), ], ) -def test_fused_settle_pack_matches_two_kernel_sequence(eviction_mode, has_swa, width, keep_count): +def test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, keep_count): device = torch.device("cuda", torch.cuda.current_device()) request_count, num_layers, num_kv_heads = 3, 2, 2 union = eviction_mode == "union" @@ -277,41 +196,23 @@ def test_fused_settle_pack_matches_two_kernel_sequence(eviction_mode, has_swa, w output_reference = output_stale.clone() dense_reference = dense_stale.clone() swa_reference = swa_stale.clone() - _settle_ties_after_topk_kernel[(rows_total,)]( - scores, - row_lengths, - row_prompt_offsets, - provisional, - output_reference, - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, + _settle_oracle( + scores, row_lengths, row_prompt_offsets, provisional, output_reference, keep_count ) - swa_offsets_arg = swa_offsets if has_swa else dense_offsets - swa_reference_arg = swa_reference if has_swa else dense_reference - _pack_compaction_sources_kernel[ - (request_count, packed_rows, (move_capacity + _BLOCK - 1) // _BLOCK) - ]( + _pack_oracle( output_reference, valid_seq_lens, dense_offsets, dense_reference, - swa_offsets_arg, - swa_reference_arg, - DENSE_TOTAL=dense_total, - SWA_TOTAL=swa_total if has_swa else 0, - SELECTION_ROWS=selection_rows, - SELECTION_STRIDE=keep_count, - KEEP_COUNT=keep_count, - NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window if has_swa else 0, - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=has_swa, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, + swa_offsets if has_swa else dense_offsets, + swa_reference if has_swa else dense_reference, + selection_rows=selection_rows, + keep_count=keep_count, + num_kv_heads=num_kv_heads, + swa_window=swa_window if has_swa else 0, + union=union, + per_layer=per_layer, + has_swa=has_swa, ) output_fused = output_stale.clone() @@ -327,7 +228,7 @@ def test_fused_settle_pack_matches_two_kernel_sequence(eviction_mode, has_swa, w valid_seq_lens, dense_offsets, dense_fused, - swa_offsets_arg, + swa_offsets if has_swa else dense_offsets, swa_fused_arg, WIDTH=width, KEEP_COUNT=keep_count, @@ -353,8 +254,8 @@ def test_fused_settle_pack_matches_two_kernel_sequence(eviction_mode, has_swa, w assert torch.equal(swa_fused, swa_reference), f"SWA moves differ (seed {seed})" -def test_fused_kernel_without_pack_matches_standalone_settle(): - """``HAS_PACK=False`` must leave exactly the standalone settle kernel.""" +def test_fused_kernel_without_pack_matches_settle_oracle(): + """``HAS_PACK=False`` must leave exactly the settle stage.""" device = torch.device("cuda", torch.cuda.current_device()) rows_total, width, keep_count = 6, 33, 7 generator = torch.Generator(device=device).manual_seed(11) @@ -372,17 +273,8 @@ def test_fused_kernel_without_pack_matches_standalone_settle(): ) output_reference = output_stale.clone() - _settle_ties_after_topk_kernel[(rows_total,)]( - scores, - row_lengths, - row_prompt_offsets, - provisional, - output_reference, - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, + _settle_oracle( + scores, row_lengths, row_prompt_offsets, provisional, output_reference, keep_count ) output_fused = output_stale.clone() diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index e5e06af64224..51d37c511ca0 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -20,9 +20,10 @@ the model's standard attention over the compacted cache; the manager publishes the cumulative evicted count on ``LlmRequest.py_num_compressed_tokens`` and the model engine subtracts it where it builds ``num_cached_tokens_per_seq``. These -tests cover the config, construction, compressed-count publication, eager -selection, page-table staging, bounded request chunks, and request lifecycle. -Model-level correctness is covered by separate end-to-end tests. +tests cover the config, construction, eviction lifecycle, page-table staging, +and the fixed score buffers. Draft co-compression contracts live in +``test_triattention_draft_cocompaction.py``; model-level correctness is covered +by separate end-to-end tests. """ from types import SimpleNamespace @@ -31,8 +32,11 @@ import pytest import torch from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_bare_staging as _make_bare_staging from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_fixed_resources_stubs as _make_fixed_resources_stubs from conftest import make_request as _make_request +from conftest import make_staging_manager as _make_staging_manager from conftest import make_triattention as _make_triattention from conftest import mocked_eviction_internals as _mocked_eviction_internals from conftest import torch_tri_score_oracle as _torch_tri_score_oracle @@ -42,7 +46,6 @@ # compression manager -- no attention classes or KV-cache-manager subclass. from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( TriAttention, - _BatchedUnionKeepSetSelector, _PreparedEviction, _PreparedGenerationBatch, _RequestCompressionState, @@ -55,8 +58,6 @@ from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig -_TORCH_TOPK_ORACLE = torch.topk - # The SM100 CuTe kernel is the only score path, so every test that actually # launches scores (or builds the real staging buffers, whose constructor # compiles the kernel) is SM100-only, like the production feature itself. @@ -104,51 +105,6 @@ def _prepared_eviction( ) -def _union_oracle(scores: torch.Tensor, keep_count: int) -> torch.Tensor: - """Independent expected-result implementation of union selection.""" - combined = scores.max(dim=0).values - row_top = _TORCH_TOPK_ORACLE( - scores, - keep_count, - dim=1, - sorted=False, - ).indices - union_mask = torch.zeros(scores.shape[1], dtype=torch.bool, device=scores.device) - union_mask.scatter_(0, row_top.reshape(-1), True) - union_indices = torch.nonzero(union_mask, as_tuple=False).flatten() - if union_indices.numel() >= keep_count: - candidates = combined.index_select(0, union_indices) - relative = _TORCH_TOPK_ORACLE( - candidates, - keep_count, - sorted=False, - ).indices - return torch.sort(union_indices.index_select(0, relative)).values - - remaining = keep_count - int(union_indices.numel()) - residual = combined.clone() - residual[union_mask] = float("-inf") - extra = _TORCH_TOPK_ORACLE( - residual, - remaining, - sorted=False, - ).indices - return torch.sort(torch.cat((union_indices, extra))).values - - -def _distinct_topk_scores(width: int, rows: int = 2) -> torch.Tensor: - """Create deterministic finite rows without top-k boundary ties.""" - token = torch.arange(width, dtype=torch.float32) - return torch.stack( - [ - torch.sin(token * (0.0017 + row * 0.0003)) - + token * (0.00011 + row * 0.000013) - + row * 0.000001 - for row in range(rows) - ] - ) - - @pytest.fixture def flat_calibration_pt(tmp_path): """Build a minimal valid calibration ``.pt`` in our flat runtime schema.""" @@ -169,8 +125,8 @@ def _make_hf_config(**values): return SimpleNamespace(get_text_config=lambda: text_config) -class TestKvCacheCompressionConfig: - def test_llm_args_dispatches_concrete_and_unknown_algorithms(self): +class TestConfigAndFactory: + def test_llm_args_dispatch_and_validation(self): from tensorrt_llm.llmapi.llm_args import TorchLlmArgs tri_args = TorchLlmArgs( @@ -195,11 +151,28 @@ def test_llm_args_dispatches_concrete_and_unknown_algorithms(self): model="dummy", kv_cache_compression_config={"algorithm": "future_method"}, ) - - def test_eviction_mode_validated(self): with pytest.raises(ValidationError): TriAttentionKvCacheCompressionConfig(eviction_mode="made_up_mode") + def test_factory_returns_triattention_and_propagates_config_fields(self): + # A plain V2 manager (block reuse off) yields a TriAttention instance. + # Calibration is deferred to the first request, so construction needs + # no calibration file or CUDA. + fake_v2 = _make_fake_v2(enable_block_reuse=False) + cfg = TriAttentionKvCacheCompressionConfig( + top_B=32, + beta=16, + eviction_mode="per_head", + model_path="/models/test", + calibration_path="/calib/test.pt", + ) + mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) + assert isinstance(mgr, TriAttention) + assert mgr.top_B == 32 + assert mgr.beta == 16 + assert mgr.eviction_mode == "per_head" + assert mgr.kv_cache_manager is fake_v2 + class TestTriAttentionClass: def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): @@ -247,35 +220,27 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): mock.call(101, Role.KEY), ] - def test_triattention_enables_capacity_only_on_target_manager(self): + @pytest.mark.parametrize("num_extra_kv_tokens,reserved_draft", [(0, 0), (4, 4)]) + def test_request_init_marks_capacity_only_and_tracks_state( + self, num_extra_kv_tokens, reserved_draft + ): + # Speculative capacity (extra KV tokens / reserved draft width) is + # accepted at request init; the target manager is marked so V2 sizing + # keeps logical max_seq_len while capacity is reclaimed and reused. manager = _make_fake_v2() + manager.num_extra_kv_tokens = num_extra_kv_tokens + manager._kv_reserve_draft_tokens = reserved_draft triattention = TriAttention(manager, top_B=8, model_path="/models/test") triattention._attention_layer_partition_cache = ([], [], None) triattention._calibrated = True - first = _make_request(11) - second = _make_request(12) - triattention.on_request_init(first) - triattention.on_request_init(second) + triattention.on_request_init(_make_request(11)) + triattention.on_request_init(_make_request(12)) assert triattention.adjusts_generation_kv_length is True assert manager.kv_compression_manages_history assert set(triattention._request_states) == {11, 12} - def test_request_init_accepts_speculative_capacity(self): - manager = _make_fake_v2() - manager.num_extra_kv_tokens = 4 - manager._kv_reserve_draft_tokens = 4 - triattention = TriAttention(manager, top_B=8, model_path="/models/test") - triattention._attention_layer_partition_cache = ([], [], None) - triattention._calibrated = True - request = _make_request(11) - - triattention.on_request_init(request) - - assert manager.kv_compression_manages_history - assert set(triattention._request_states) == {11} - def test_resolve_accepts_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() mgr.calibration_path = flat_calibration_pt @@ -285,56 +250,11 @@ def test_resolve_accepts_flat_pt(self, flat_calibration_pt): assert key in loaded -# --------------------------------------------------------------------------- -# Eviction publishes the cumulative evicted count on the request; the model -# engine reads it back where it builds num_cached_tokens_per_seq. -# --------------------------------------------------------------------------- - - class TestCompressedTokenPublication: - def test_manager_is_marked_capacity_only_and_requests_default_to_zero(self): - mgr = _make_triattention() - manager = mgr.kv_cache_manager - # The compression-manager base marks the target manager so V2 sizing - # keeps logical max_seq_len while capacity is reclaimed and reused. - assert manager.kv_compression_manages_history - request = _make_request(7) - # Default 0 keeps the engine's num_cached subtraction a no-op until - # the first eviction publishes a count. - assert request.py_num_compressed_tokens == 0 - - def test_eviction_bookkeeping_publishes_cumulative_count(self): - # The eviction bookkeeping writes the cumulative evicted count on the - # request in the same step that compacts the cache; this is the - # channel's only producer. - manager = _make_triattention(top_B=4) - manager.kv_cache_manager._stream = mock.Mock() - request = _make_request(7, py_prompt_len=2) - _set_request_state(manager, 7, confirmed_kv_length=10) - - with _mocked_eviction_internals(manager) as internals: - first = manager._evict_requests([(request, 7)], 2) - - assert first == [(7, 6)] - # 10 confirmed - (2 pinned prompt + 4 decode budget) = 4 evicted. - assert request.py_num_compressed_tokens == 4 - assert manager._request_states[7].confirmed_kv_length == 6 - internals.batched_compaction.compact.assert_called_once_with() - internals.score_staging.mark_page_tables_consumed.assert_called_once_with( - manager.kv_cache_manager._stream - ) - - # Round two: 6 retained + 8 newly confirmed decode tokens. - manager._request_states[7].confirmed_kv_length = 14 - with _mocked_eviction_internals(manager) as internals: - second = manager._evict_requests([(request, 7)], 2) - - assert second == [(7, 6)] - # The count is cumulative and never decreases: 4 + (14 - 6) = 12. - assert request.py_num_compressed_tokens == 12 - # The staged logical position restores the uncompressed length. - prepared = internals.attach.call_args.args[0] - assert prepared[0].round_start == 14 + 4 + # The cumulative/monotone publication contract itself (including the + # uncompressed round_start restoration) is covered end to end by + # test_triattention_draft_cocompaction.py:: + # test_compressed_count_is_monotone_and_tracks_confirmed_length. def test_identity_compaction_is_rejected_instead_of_published(self): manager = _make_triattention(top_B=4) @@ -350,20 +270,19 @@ def test_identity_compaction_is_rejected_instead_of_published(self): class TestEvictionLifecycle: - def test_triattention_prepare_only_snapshots_and_update_uses_final_hook(self): + def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): + # Structural: prepare is a snapshot-only override; the eviction runs + # from the framework's final on_generation_step_end hook. assert "prepare_resources" in TriAttention.__dict__ assert "update_resources" not in TriAttention.__dict__ assert "on_generation_step_end" in TriAttention.__dict__ - def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): manager = _make_triattention() - request = _make_request(7) batch = SimpleNamespace( context_requests=[], context_requests_last_chunk=[], - generation_requests=[request], + generation_requests=[_make_request(7)], ) - with mock.patch.object(manager, "_periodic_evict") as periodic_evict: manager.prepare_resources(batch) periodic_evict.assert_not_called() @@ -372,10 +291,6 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): periodic_evict.assert_called_once_with(batch) - def test_non_v2_manager_is_always_rejected(self): - with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - TriAttention(SimpleNamespace(), top_B=8) - @staticmethod def _make_due_decode_request(seq_len): request = _make_request( @@ -407,29 +322,14 @@ def _make_due_decode_request(seq_len): return mgr, request, batch def test_identity_gate_preserves_real_eviction_round(self): - import contextlib - - import tensorrt_llm._torch.kv_cache_compression.triattention.triattention as tri_module - mgr, request, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) - timeline = [] cache = mgr.kv_cache_manager.kv_cache_map[7] def compact(*args, protected_tail_lengths, **_kwargs): assert protected_tail_lengths == {7: 0} - timeline.append("compact_dispatch") return [(7, 1024 + 4096)] - @contextlib.contextmanager - def track_range(name, **kwargs): - timeline.append(f"enter:{name}") - yield - timeline.append(f"exit:{name}") - - with ( - mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict, - mock.patch.object(tri_module, "nvtx_range", side_effect=track_range), - ): + with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: mgr._periodic_evict(batch) evict.assert_called_once_with( @@ -439,13 +339,6 @@ def track_range(name, **kwargs): ) mgr.kv_cache_manager._stream.wait_event.assert_not_called() cache.resize.assert_called_once_with(1024 + 4096, None) - assert timeline == [ - "enter:triattention.evict_request_group reqs=1", - "compact_dispatch", - "exit:triattention.evict_request_group reqs=1", - "enter:triattention.resize", - "exit:triattention.resize", - ] def test_suspended_cache_rejects_batch_before_cadence_mutation(self): manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) @@ -500,17 +393,28 @@ def test_eager_eviction_runs_large_due_cohort_in_one_round(self): assert [len(call.args[0]) for call in evict.call_args_list] == [65] assert resize.call_count == 1 - def test_request_finish_keeps_eviction_buffers_resident(self): + def test_request_finish_clears_state_but_keeps_buffers_resident(self): manager = _make_triattention() - request = _make_request(7) - _set_request_state(manager, 7) + _set_request_state( + manager, + 7, + generation_steps=1, + evicted_tokens=127, + confirmed_kv_length=128, + ) buffers = object() compaction = object() manager._eviction_resources = buffers manager._batched_compaction = compaction + manager._prepared_generation_batch = _PreparedGenerationBatch( + batch=SimpleNamespace(), + growth_by_request={7: 1}, + ) - manager.on_request_finish(request) + manager.on_request_finish(_make_request(7)) + assert manager._request_states == {} + assert manager._prepared_generation_batch.growth_by_request == {} # The buffers are sized for the executor limits, not one cohort, so # they stay resident for the next generation batch. assert manager._eviction_resources is buffers @@ -598,25 +502,11 @@ def test_mla_selfkonly_cache_is_rejected(self): with pytest.raises(ValueError, match="standard key/value KV cache"): manager._validate_v2_compatibility() - def test_one_model_mtp_co_compression_contract_is_accepted(self): - draft_manager = _make_fake_v2(is_draft=True) - manager = TriAttention( - _make_fake_v2(), - top_B=8, - model_path="/models/test", - draft_kv_cache_manager=draft_manager, - ) - - manager._validate_v2_compatibility() - assert manager.kv_cache_manager.kv_compression_manages_history is True - # The draft cache is compacted together with the target, so its - # physical length diverges from the logical length the same way. - assert draft_manager.kv_compression_manages_history is True - - def test_draft_co_compression_accepts_smaller_draft_max_seq_len(self): + def test_one_model_draft_co_compression_contract_is_accepted(self): # Co-compression keeps the draft's physical length equal to the - # target's, so the draft does not have to cover the target's logical - # maximum sequence length. + # target's, so a draft with a smaller max_seq_len than the target's + # logical maximum is accepted, and both managers are marked as + # diverging from the logical length. draft_manager = _make_fake_v2(is_draft=True) draft_manager.max_seq_len = 8192 manager = TriAttention( @@ -627,47 +517,8 @@ def test_draft_co_compression_accepts_smaller_draft_max_seq_len(self): ) manager._validate_v2_compatibility() - - @pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) - def test_draft_co_compression_requires_union_mode(self, eviction_mode): - manager = TriAttention( - _make_fake_v2(), - top_B=8, - model_path="/models/test", - eviction_mode=eviction_mode, - draft_kv_cache_manager=_make_fake_v2(is_draft=True), - ) - - with pytest.raises(ValueError, match="union"): - manager._validate_v2_compatibility() - - def test_draft_co_compression_rejects_mla_draft_cache(self): - draft_manager = _make_fake_v2(is_draft=True) - manager = TriAttention( - _make_fake_v2(), - top_B=8, - model_path="/models/test", - draft_kv_cache_manager=draft_manager, - ) - # The base class already rejects a mismatched draft kv_factor at - # construction; this guards against later runtime divergence too. - draft_manager.kv_factor = 1 - - with pytest.raises(ValueError, match="standard key/value cache"): - manager._validate_v2_compatibility() - - def test_draft_co_compression_requires_full_attention_draft(self): - draft_manager = _make_fake_v2(is_draft=True) - draft_manager.max_attention_window_vec = [128] - manager = TriAttention( - _make_fake_v2(), - top_B=8, - model_path="/models/test", - draft_kv_cache_manager=draft_manager, - ) - - with pytest.raises(ValueError, match="full-attention draft"): - manager._validate_v2_compatibility() + assert manager.kv_cache_manager.kv_compression_manages_history is True + assert draft_manager.kv_compression_manages_history is True def test_resize_shrinks_draft_cache_with_its_own_protected_tail(self): retained = 1024 + 4096 @@ -721,130 +572,34 @@ def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): manager._validate_v2_compatibility() - @pytest.mark.parametrize("mode", ["draft_target", "pard"]) - def test_unvalidated_paged_draft_tail_contracts_remain_fail_closed(self, mode): - from tensorrt_llm.llmapi.llm_args import DraftTargetDecodingConfig, PARDDecodingConfig - - if mode == "draft_target": - spec_config = DraftTargetDecodingConfig( - max_draft_len=3, - speculative_model="/tmp/draft-target-model", - ) - else: - spec_config = PARDDecodingConfig(max_draft_len=3) - # The call-site speculative gate rejects before any manager is created. - from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec - from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig - - with pytest.raises(ValueError, match="standard paged cache compacted together"): - validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 - ), - spec_config, - _make_fake_v2(is_draft=True), - ) - - def test_dflash_spec_mode_is_rejected(self): - # Policy: the DFlash draft reads cross-attention context buffers, not - # a paged KV cache, so compression cannot cover it. The call-site - # speculative gate rejects before any manager is created. - from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec - from tensorrt_llm.llmapi.llm_args import ( - DFlashDecodingConfig, - TriAttentionKvCacheCompressionConfig, - ) - - with pytest.raises(ValueError, match="standard paged cache compacted together"): - validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 - ), - DFlashDecodingConfig(max_draft_len=3), - _make_fake_v2(is_draft=True), - ) - - def test_prepare_snapshots_fixed_linear_generation_growth(self): + @pytest.mark.parametrize( + "num_extra_kv_tokens,reserved_draft,draft_tokens,expected_growth", + [ + (2, 0, [1, 2, 3], 4), + # The reserved draft width protects capacity even when this step's + # actual draft is shorter. + (0, 6, [1, 2], 7), + ], + ) + def test_prepare_snapshots_fixed_linear_generation_growth( + self, num_extra_kv_tokens, reserved_draft, draft_tokens, expected_growth + ): manager = _make_fake_v2() - manager.num_extra_kv_tokens = 2 + manager.num_extra_kv_tokens = num_extra_kv_tokens + manager._kv_reserve_draft_tokens = reserved_draft manager.kv_cache_map = { 7: SimpleNamespace(capacity=106, is_active=True), } triattention = TriAttention(manager, top_B=8, model_path="/models/test") batch = SimpleNamespace( context_requests=[], - generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], + generation_requests=[_make_request(7, py_draft_tokens=draft_tokens)], ) triattention.prepare_resources(batch) assert triattention._prepared_generation_batch.batch is batch - assert triattention._prepared_generation_batch.growth_by_request == {7: 4} - - def test_prepare_protects_reserved_draft_width(self): - manager = _make_fake_v2() - manager._kv_reserve_draft_tokens = 6 - manager.kv_cache_map = { - 7: SimpleNamespace(capacity=106, is_active=True), - } - triattention = TriAttention(manager, top_B=8, model_path="/models/test") - batch = SimpleNamespace( - context_requests=[], - generation_requests=[_make_request(7, py_draft_tokens=[1, 2])], - ) - - triattention.prepare_resources(batch) - - assert triattention._prepared_generation_batch.growth_by_request == {7: 7} - - def test_request_finish_clears_compression_state(self): - request = SimpleNamespace(py_request_id=7) - mgr = _make_triattention() - _set_request_state( - mgr, - 7, - generation_steps=1, - evicted_tokens=127, - confirmed_kv_length=128, - ) - mgr._prepared_generation_batch = _PreparedGenerationBatch( - batch=SimpleNamespace(), - growth_by_request={7: 1}, - ) - - mgr.on_request_finish(request) - - assert mgr._request_states == {} - assert mgr._prepared_generation_batch.growth_by_request == {} - - -class TestTopKRouting: - @pytest.mark.parametrize("keep_count", [4096, 8192]) - def test_cross_request_union_matches_oracle_at_high_keep_counts(self, keep_count): - width = keep_count + 64 - request_scores = [ - _distinct_topk_scores(width), - _distinct_topk_scores(width).roll(17, dims=1) + 0.000007, - ] - expected = [_union_oracle(scores, keep_count) for scores in request_scores] - device = torch.device("cuda", torch.cuda.current_device()) - scores = torch.stack(request_scores).to(device) - selector = _BatchedUnionKeepSetSelector( - request_scores[0].shape[0], - width, - keep_count, - dtype=scores.dtype, - device=device, - max_requests=len(request_scores), - ) - # The fused CuTe pipeline is the production union-row producer; this - # top-k routing test stages the prepared rows into ``combined``. - selector.combined.copy_(scores.amax(dim=1)) - selector.select_prepared_union_scores() - selected = selector.keep.cpu() - - for actual, expected_keep in zip(selected, expected): - assert torch.equal(actual, expected_keep.to(torch.int32)) + assert triattention._prepared_generation_batch.growth_by_request == {7: expected_growth} class TestFixedScoreMetadata: @@ -868,35 +623,9 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize eviction_mode=eviction_mode, normalize_scores=normalize_scores, ) - manager._H = 2 - manager._F = 2 - manager._freq_scale_sq = torch.ones(2) - manager._offsets = torch.ones(2) - manager.calibration = {"omega": torch.ones(2)} - manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) - manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) - pool = torch.empty(8, 2, 1, 4, 4) - layout = SimpleNamespace( - manager=SimpleNamespace(num_pools=1), - num_layers=2, - global_layers=[0, 1], - layer_pools=[pool, pool], - dense_layers=[0, 1], - swa_layers=[], - storage_groups={0: [0, 1]}, - pool_view_fingerprint=(("fixed",),), - ) # The buffers follow the executor limits: eight requests (max batch # size) by 260 decode tokens (top_B plus two eviction periods). - score_staging = SimpleNamespace( - fused_group=SimpleNamespace(output=torch.empty(8, 4, 260)), - bind_score_launcher=mock.Mock(), - token_starts_device=torch.zeros(8, dtype=torch.int32), - ) - keep_set_selector = SimpleNamespace( - valid_widths=torch.empty(8, dtype=torch.int32), - top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), - ) + layout, score_staging, keep_set_selector = _make_fixed_resources_stubs(manager) prepared = [ _prepared_eviction( _make_request(7), @@ -929,11 +658,16 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize assert resources.score_staging is score_staging assert resources.keep_set_selector is keep_set_selector - def test_bulk_page_table_copy_uses_immutable_host_snapshots(self): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, - ) + def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): + """The bulk copy stages immutable host snapshots, and the next copy + waits for the previous cohort's consumers. + Both persistent V2 host inputs (the block-offset table and the + index-mapper slot assignment) may mutate as soon as ``stage`` returns; + the staged device tables must reflect the values at staging time. A + subsequent bulk copy must also wait until the previous round's + consumers (recorded by ``mark_page_tables_consumed``) are done. + """ device = torch.device("cuda", torch.cuda.current_device()) current_stream = torch.cuda.current_stream(device) manager_stream = torch.cuda.Stream(device=device) @@ -957,44 +691,11 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): ) gather = mock.Mock(side_effect=gather_k_block_offsets) - - staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = device - staging.max_requests = 1 - staging.page_count = 5 - staging.copy_block_count = 8 - staging.bulk_copy_done = torch.cuda.Event() - staging.bulk_consume_done = torch.cuda.Event() - staging.page_tables_active = False - staging.copy_done = torch.cuda.Event() - staging.copy_pending = False - staging._bulk_offsets_src = torch.empty( - 1, 1, 2, 8, dtype=torch.int32, device="cpu", pin_memory=True - ) - staging._bulk_copy_idx_src = torch.arange( - 1, dtype=torch.int32, device="cpu", pin_memory=True - ) - staging.block_offsets_device = torch.empty(1, 1, 2, 8, dtype=torch.int32, device=device) + staging = _make_bare_staging(device, max_requests=1, copy_block_count=8, page_count=5) staging.copy_done.record(current_stream) + manager = _make_staging_manager(host_table, gather, manager_stream) - manager = SimpleNamespace( - host_kv_cache_block_offsets=host_table, - kv_factor=2, - layer_offsets={10: 0}, - layer_to_pool_mapping_dict={0: 0}, - index_mapper=SimpleNamespace(gather_k_block_offsets=gather), - index_scales=torch.tensor([2], dtype=torch.int32, device="cpu", pin_memory=True), - kv_offset=torch.tensor([1], dtype=torch.int32, device="cpu", pin_memory=True), - _stream=manager_stream, - ) - - with torch.cuda.stream(manager_stream): - torch.cuda._sleep(50_000_000) - with mock.patch.object( - torch, - "index_select", - side_effect=AssertionError("page-table staging used torch.index_select"), - ): + def stage_once(): assert staging._stage_page_tables_bulk( manager, [7], @@ -1003,11 +704,19 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): staging.block_offsets_device, staging.copy_block_count, ) - assert staging._bulk_offsets_src.shape[-1] == 8 - assert staging._bulk_offsets_src.shape[1] == 1 - # Mutate both persistent V2 host inputs before the delayed kernel reads. - # The staged result must still reflect row 0 values [3, 4, 5, 6, 7]. + # Round 1: mutate the host table and the slot assignment right after + # staging, with the manager stream artificially delayed. The staged + # result must still reflect the values at staging time. + with torch.cuda.stream(manager_stream): + torch.cuda._sleep(50_000_000) + with mock.patch.object( + torch, + "index_select", + side_effect=AssertionError("page-table staging used torch.index_select"), + ): + stage_once() + assert staging._bulk_offsets_src.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() @@ -1015,18 +724,12 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): 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 cohort. host_table[0, 0, 0, :5] = torch.tensor([18, 19, 20, 21, 22], dtype=torch.int32) selected_slot[0] = 0 with torch.cuda.stream(manager_stream): torch.cuda._sleep(50_000_000) - assert staging._stage_page_tables_bulk( - manager, - [7], - current_stream, - staging._bulk_offsets_src, - staging.block_offsets_device, - staging.copy_block_count, - ) + 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() @@ -1034,57 +737,10 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): 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] - def test_next_bulk_copy_waits_for_page_table_consumers(self): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, - ) - - device = torch.device("cuda", torch.cuda.current_device()) - current_stream = torch.cuda.current_stream(device) - manager_stream = torch.cuda.Stream(device=device) - host_table = torch.zeros(1, 1, 2, 4, dtype=torch.int32, device="cpu", pin_memory=True) - host_table[0, 0, 0] = torch.tensor([1, 2, 3, 4], dtype=torch.int32) - - def gather_k_block_offsets(source, destination, request_ids, num_blocks): - assert request_ids == [7] - destination[:, :1, 0, :num_blocks].copy_(source[:, :1, 0, :num_blocks]) - - staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = device - staging.max_requests = 1 - staging.copy_block_count = 4 - staging.bulk_copy_done = torch.cuda.Event() - staging.bulk_consume_done = torch.cuda.Event() - staging.page_tables_active = False - staging.copy_done = torch.cuda.Event() - staging.copy_pending = False - staging._bulk_offsets_src = torch.empty( - 1, 1, 2, 4, dtype=torch.int32, device="cpu", pin_memory=True - ) - staging._bulk_copy_idx_src = torch.arange( - 1, dtype=torch.int32, device="cpu", pin_memory=True - ) - staging.block_offsets_device = torch.empty(1, 1, 2, 4, dtype=torch.int32, device=device) - staging.copy_done.record(current_stream) - manager = SimpleNamespace( - host_kv_cache_block_offsets=host_table, - kv_factor=2, - index_mapper=SimpleNamespace( - gather_k_block_offsets=mock.Mock(side_effect=gather_k_block_offsets) - ), - index_scales=torch.tensor([2], dtype=torch.int32, device="cpu", pin_memory=True), - kv_offset=torch.tensor([1], dtype=torch.int32, device="cpu", pin_memory=True), - _stream=manager_stream, - ) - - assert staging._stage_page_tables_bulk( - manager, - [7], - current_stream, - staging._bulk_offsets_src, - staging.block_offsets_device, - staging.copy_block_count, - ) + # Round 3: a delayed consumer read (snapshot) queued before + # ``mark_page_tables_consumed`` must complete before the next bulk + # copy overwrites the device tables. + selected_slot[0] = 0 manager_stream.synchronize() snapshot = torch.empty_like(staging.block_offsets_device) torch.cuda._sleep(20_000_000) @@ -1092,42 +748,11 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): staging.page_tables_active = True staging.mark_page_tables_consumed(manager_stream) - host_table[0, 0, 0] = torch.tensor([5, 6, 7, 8], dtype=torch.int32) - assert staging._stage_page_tables_bulk( - manager, - [7], - current_stream, - staging._bulk_offsets_src, - staging.block_offsets_device, - staging.copy_block_count, - ) + stage_once() current_stream.synchronize() - assert snapshot[0, 0, 0].tolist() == [2, 4, 6, 8] - assert staging.block_offsets_device[0, 0, 0].tolist() == [10, 12, 14, 16] - - def test_cross_stream_staging_is_rejected_before_page_table_query(self): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, - _FixedScoreStreamMismatch, - ) - - staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = torch.device("cuda", torch.cuda.current_device()) - staging.max_requests = 8 - staging.stream = SimpleNamespace(device=torch.device("cuda:0"), cuda_stream=4) - staging.page_tables_active = False - staging.copy_pending = False - staging.copy_done = SimpleNamespace(query=mock.Mock(), synchronize=mock.Mock()) - staging.draft_block_offsets_device = None - manager = mock.Mock() - - other_stream = SimpleNamespace(device=torch.device("cuda:0"), cuda_stream=5) - with mock.patch.object(torch.cuda, "current_stream", return_value=other_stream): - with pytest.raises(_FixedScoreStreamMismatch, match="first CUDA stream"): - staging.stage(manager, [1], [8.0], [0]) - staging.copy_done.query.assert_not_called() - staging.copy_done.synchronize.assert_not_called() + 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] def test_staged_page_tables_bypass_per_request_cuda_materialization(self): manager = _make_triattention() @@ -1290,15 +915,10 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): ) gather = mock.Mock(side_effect=gather_k_block_offsets) - manager = SimpleNamespace( - enable_swa_scratch_reuse=False, - host_kv_cache_block_offsets=host_table, - kv_factor=2, - index_mapper=SimpleNamespace(gather_k_block_offsets=gather), - index_scales=torch.full((3,), 2, dtype=torch.int32, pin_memory=True), - kv_offset=torch.ones(3, dtype=torch.int32, pin_memory=True), - _stream=torch.cuda.Stream(device=device), + manager = _make_staging_manager( + host_table, gather, torch.cuda.Stream(device=device), num_slots=3 ) + manager.enable_swa_scratch_reuse = False assert not staging.stage( manager, @@ -1351,113 +971,6 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): staging.stage(manager, request_ids, round_starts, token_starts) assert gather.call_count == calls - @requires_sm100 - @pytest.mark.parametrize("request_count", [1, 8]) - def test_fixed_score_matches_torch_oracle_across_two_groups(self, request_count): - pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) - - device = torch.device("cuda", torch.cuda.current_device()) - torch.manual_seed(20260703 + request_count) - max_requests = 8 - page_count = 2 - tokens_per_block = 32 - head_dim = 64 - num_freqs = head_dim // 2 - num_q_heads = 8 - seq_len = page_count * tokens_per_block - prompt_len = 2 - page_ids = torch.arange(max_requests * page_count, dtype=torch.int64, device=device).view( - max_requests, page_count - ) - layer_elements = max_requests * page_count * 2 * 1 * tokens_per_block * head_dim - shared = (0.125 * torch.randn(2 * layer_elements, device=device)).to(torch.bfloat16) - pool_shape = (max_requests * page_count, 2, 1, tokens_per_block, head_dim) - pools = [ - shared[:layer_elements].view(pool_shape), - shared[layer_elements:].view(pool_shape), - (0.125 * torch.randn(pool_shape, device=device)).to(torch.bfloat16), - ] - storage_groups = [[0, 1], [2]] - q_real = (0.125 * torch.randn(3, num_q_heads, 2 * num_freqs, device=device))[..., ::2] - q_imag = (0.125 * torch.randn(3, num_q_heads, 2 * num_freqs, device=device))[..., ::2] - mlr = (0.125 * torch.randn(3, num_q_heads, 2 * num_freqs, device=device))[..., ::2] - freq = (torch.rand(2 * num_freqs, device=device) + 0.5)[::2] - omega = (torch.rand(2 * num_freqs, device=device) * 0.05)[::2] - offsets = torch.tensor([1.0, 0.0, 2.0, 0.0, 4.0, 0.0], device=device)[::2] - assert not q_real.is_contiguous() - assert not q_imag.is_contiguous() - assert not mlr.is_contiguous() - assert not freq.is_contiguous() - assert not omega.is_contiguous() - assert not offsets.is_contiguous() - round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 - round_starts = round_device[:request_count].tolist() - token_starts_device = torch.full( - (max_requests,), prompt_len, dtype=torch.int32, device=device - ) - seq_lens = [seq_len - request % 2 for request in range(request_count)] - phase = (round_device[:, None, None] + offsets[None, :, None]) * omega[None, None] - oracle = _torch_tri_score_oracle( - pools, - page_ids[:request_count], - seq_lens, - round_starts, - q_real, - q_imag, - mlr, - freq, - omega, - offsets, - [0, 1, 2], - ) - for layers in storage_groups: - group = _FixedScoreGroup( - pools, - layers, - max_requests, - page_count, - seq_len, - num_q_heads, - _encode_block_offsets(page_ids), - [0] * len(layers), - q_real, - q_imag, - mlr, - freq, - omega, - offsets, - output_width=seq_len - prompt_len, - ) - valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) - fixed = group.launch( - request_count, - torch.tensor(seq_lens, dtype=torch.int32, device=device), - valid_widths, - token_starts_device, - torch.cos(phase).mean(dim=1), - torch.sin(phase).mean(dim=1), - ) - assert valid_widths[:request_count].tolist() == [ - seq_len - prompt_len for seq_len in seq_lens - ] - assert fixed.shape == ( - request_count, - len(layers), - num_q_heads, - seq_len - prompt_len, - ) - for request in range(request_count): - for layer_slot, layer in enumerate(layers): - valid_width = seq_lens[request] - prompt_len - segment = fixed[request, layer_slot, :, :valid_width] - expected = oracle[request * len(pools) + layer][ - :, prompt_len : prompt_len + valid_width - ] - torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) - @requires_sm100 @pytest.mark.parametrize("request_count", [1, 7, 8]) def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count): @@ -1466,14 +979,12 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun This is the production V2 shape: get_buffers wraps every layer as its own TensorWrapper storage and every layer allocates its own pages, so the fused path must not assume a shared storage anchor or a shared - per-request block table. + per-request block table. The launch is checked against the independent + Torch oracle, then relaunched after a round-start advance and a block + table rebind. """ pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, - _FixedScoreStreamMismatch, - ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _FixedScoreGroup, ) @@ -1552,7 +1063,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun mean_phase_table.gather(round_device, mean_cos, mean_sin, request_count) score_sentinel = -12345.0 group.output.fill_(score_sentinel) - checked = group.launch( + fixed = group.launch( request_count, valid_seq_lens, valid_widths, @@ -1560,23 +1071,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun mean_cos, mean_sin, ).clone() - staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = group.output.device - staging.max_requests = request_count - staging.fused_group = group - staging.round_starts_device = round_device - staging.valid_seq_lens_device = valid_seq_lens - staging.token_starts_device = token_starts - staging.mean_cos = mean_cos - staging.mean_sin = mean_sin - staging.mean_phase_table = mean_phase_table - staging.stream = None - staging._score_valid_widths = None - staging._score_launcher_bound = False - staging.bind_score_launcher(valid_widths, "mean") - group.output.fill_(score_sentinel) - fixed = staging.launch_prepared_score().clone() - torch.testing.assert_close(fixed, checked, rtol=0, atol=0) assert valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] # The deployed fused score must agree with the independent Torch oracle @@ -1616,7 +1110,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun expected_second_widths = valid_seq_lens - prompt_len group.output.fill_(score_sentinel) valid_widths.fill_(-1) - checked_second = group.launch( + second_launch = group.launch( request_count, valid_seq_lens, valid_widths, @@ -1624,24 +1118,16 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun mean_cos, mean_sin, ).clone() - group.output.fill_(score_sentinel) - valid_widths.fill_(-1) - second_launch = staging.launch_prepared_score().clone() - torch.testing.assert_close(second_launch, checked_second, rtol=0, atol=0) assert torch.equal(valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) - other_stream = torch.cuda.Stream(device=device) - with torch.cuda.stream(other_stream): - with pytest.raises(_FixedScoreStreamMismatch, match="staging CUDA stream"): - staging.launch_prepared_score() - class TestKernelMaskedSwa: - def test_layer_partition_uses_local_model_config(self): + @pytest.mark.parametrize("top_B,fits_window", [(128, True), (127, False)]) + def test_layer_partition_uses_local_config_and_validates_window(self, top_B, fits_window): mgr = _make_triattention() mgr.model_path = "/models/gpt-oss" - mgr.top_B = 128 + mgr.top_B = top_B mgr.kv_cache_manager = SimpleNamespace(pp_layers=[0, 1, 2, 3]) config = _make_hf_config( layer_types=[ @@ -1654,6 +1140,11 @@ def test_layer_partition_uses_local_model_config(self): ) 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="decode budget top_B=127"): + mgr._attention_layer_partition(4) + return dense, sliding, window = mgr._attention_layer_partition(4) load.assert_called_once_with( @@ -1662,40 +1153,3 @@ def test_layer_partition_uses_local_model_config(self): assert dense == [1, 3] assert sliding == [0, 2] assert window == 128 - - def test_layer_partition_rejects_decode_budget_smaller_than_window(self): - mgr = _make_triattention() - mgr.model_path = "/models/gpt-oss" - mgr.top_B = 127 - mgr.kv_cache_manager = SimpleNamespace(pp_layers=[0, 1]) - config = _make_hf_config( - layer_types=["sliding_attention", "full_attention"], - sliding_window=128, - ) - - with ( - mock.patch("transformers.AutoConfig.from_pretrained", return_value=config), - pytest.raises(ValueError, match="decode budget top_B=127"), - ): - mgr._attention_layer_partition(2) - - -class TestFactory: - def test_returns_triattention_instance_and_propagates_config_fields(self): - # A plain V2 manager (block reuse off) yields a TriAttention instance. - # Calibration is deferred to the first request, so construction needs - # no calibration file or CUDA. - fake_v2 = _make_fake_v2(enable_block_reuse=False) - cfg = TriAttentionKvCacheCompressionConfig( - top_B=32, - beta=16, - eviction_mode="per_head", - model_path="/models/test", - calibration_path="/calib/test.pt", - ) - mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) - assert isinstance(mgr, TriAttention) - assert mgr.top_B == 32 - assert mgr.beta == 16 - assert mgr.eviction_mode == "per_head" - assert mgr.kv_cache_manager is fake_v2 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py deleted file mode 100644 index 1d4993a1de7a..000000000000 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_score_ops.py +++ /dev/null @@ -1,318 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The CuTe score kernel vs an independent PyTorch oracle, through group.launch. - -The SM100 CuTe-DSL fused score pack (``triattention_cute_score_fused.py``, -score-only entry) is the ONLY score implementation. These tests drive it -through ``_FixedScoreGroup.launch`` -- the exact production entry point -- -across the supported production geometries, with multi-layer segments, -permuted page tables, ragged valid lengths, and per-request prompt windows, -and compare against a pure-PyTorch oracle that recomputes everything -independently. They also pin the loud-failure contract: unsupported -geometry, removed aggregations, and request counts beyond the group -capacity raise instead of routing to another kernel. -""" - -import pytest -import torch -from conftest import encode_block_offsets as _encode_block_offsets -from conftest import torch_tri_score_oracle as _torch_tri_score_oracle - -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, -) - -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 - # The 0.125 scaling keeps the BF16 key/coefficient products small so the - # kernel-vs-oracle tolerance can stay tight across the frequency sum. - 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 - group = _FixedScoreGroup( - pools, - list(range(num_layers)), - max_requests, - page_count, - capacity, - num_q_heads, - _encode_block_offsets(page_ids), - [0] * num_layers, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets_t, - output_width=capacity - prompt_len, - ) - round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() - token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) - # Ragged valid lengths whose tails land mid-page and mid-compute-tile. - seq_lens = [capacity - ((request * 3) % 5) for request in range(max_requests)] - valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) - phase = (round_starts.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() - # Everything the PyTorch oracle needs to rebuild the reference leg - # independently (it recomputes its own mean phases from these). - 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 ( - group, - token_starts, - valid_seq_lens, - seq_lens, - mean_cos, - mean_sin, - oracle_inputs, - ) - - -# One entry per supported production geometry: the Qwen3 shape (64 -# frequencies, GQA group 4 riding the padded MMA tile), the GPT-OSS shape -# (32 frequencies, group 8, 32-token pages spanning two page fragments per -# compute tile), and the originally validated 128-token-page shape. -_CASES = [ - pytest.param( - dict( - max_requests=4, - num_layers=2, - page_count=4, - tokens_per_block=32, - head_dim=128, - num_q_heads=8, - num_kv_heads=2, - ), - id="qwen3_f64_group4_tpb32", - ), - pytest.param( - dict( - max_requests=2, - num_layers=3, - page_count=4, - tokens_per_block=32, - head_dim=64, - num_q_heads=8, - num_kv_heads=1, - ), - id="gptoss_f32_group8_tpb32", - ), - pytest.param( - dict( - max_requests=2, - num_layers=2, - page_count=2, - tokens_per_block=128, - head_dim=64, - num_q_heads=8, - num_kv_heads=1, - ), - id="original_f32_group8_tpb128", - ), -] - -_QWEN3_CASE = dict(_CASES[0].values[0]) - - -class TestTriAttentionScoreLaunch: - @requires_sm100 - @pytest.mark.parametrize("case", _CASES) - def test_cute_kernel_matches_torch_oracle(self, 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"] - ( - group, - token_starts, - valid_seq_lens, - seq_lens, - mean_cos, - mean_sin, - oracle_inputs, - ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) - device = group.output.device - - oracle = _torch_tri_score_oracle( - group._cute_layer_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)), - ) - - # The fused runner dispatches every request count up to the group - # capacity; cover one, an intermediate count, and the capacity. - for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): - group.output.fill_(float("nan")) - valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) - scores = group.launch( - request_count, - valid_seq_lens, - valid_widths, - token_starts, - mean_cos, - mean_sin, - ) - assert scores.shape == ( - request_count, - num_layers, - case["num_q_heads"], - group.output_width, - ) - # The launch owns the per-request decode widths the selection - # reduce kernels consume (the deleted C++ op used to write them). - assert valid_widths[: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, - ) - - @requires_sm100 - def test_request_count_beyond_capacity_raises(self): - """A request count beyond the group capacity fails loudly. - - The fused runner dispatches every request count up to the group - capacity (covered by the oracle matrix above); there is no fallback - kernel, so anything beyond it must raise instead of silently - scoring through a slower path. - """ - pytest.importorskip("cutlass") - ( - group, - token_starts, - valid_seq_lens, - _, - mean_cos, - mean_sin, - _, - ) = _build_case(prompt_len=5, seed=20260719, **_QWEN3_CASE) - valid_widths = torch.empty( - _QWEN3_CASE["max_requests"], dtype=torch.int32, device=group.output.device - ) - with pytest.raises(ValueError, match="exceeds fixed score capacity"): - group.launch( - _QWEN3_CASE["max_requests"] + 1, - valid_seq_lens, - valid_widths, - token_starts, - mean_cos, - mean_sin, - ) - - def _tiny_unsupported_group(self, dtype: torch.dtype): - """A geometry far outside the CuTe contract (constructor accepts it).""" - 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 - pools = [ - torch.randn( - max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device - ).to(dtype) - for _ in range(num_layers) - ] - page_ids = ( - torch.arange(max_requests * page_count, device=device) - .view(max_requests, page_count) - .contiguous() - ) - capacity = page_count * tokens_per_block - group = _FixedScoreGroup( - pools, - list(range(num_layers)), - max_requests, - page_count, - capacity, - 2, - _encode_block_offsets(page_ids), - [0] * num_layers, - torch.randn(num_layers, 2, num_freqs, device=device), - torch.randn(num_layers, 2, num_freqs, device=device), - torch.randn(num_layers, 2, num_freqs, device=device), - torch.rand(num_freqs, device=device) + 0.5, - torch.rand(num_freqs, device=device) * 0.05, - torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), - output_width=capacity - 1, - ) - device_args = dict(dtype=torch.int32, device=device) - return group, ( - torch.full((max_requests,), capacity, **device_args), - torch.empty(max_requests, **device_args), - torch.ones(max_requests, **device_args), - torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), - torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), - ) - - def test_unsupported_geometry_raises(self): - """Score setup outside the CuTe contract raises; nothing falls back.""" - group, launch_args = self._tiny_unsupported_group(torch.float32) - with pytest.raises(ValueError, match="TriAttention score requires SM100"): - group.launch(1, *launch_args) - - def test_max_aggregation_raises(self): - """Max aggregation was removed with the C++ score stack.""" - group, launch_args = self._tiny_unsupported_group(torch.bfloat16) - with pytest.raises(ValueError, match="max aggregation"): - group.launch(1, *launch_args, aggregation="max") 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 index de9e9035a467..7791857876b1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -4,12 +4,11 @@ import pytest import torch +from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_ramp_pools as _make_ramp_pools from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - BatchedKVCacheCompaction, -) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _BatchedPerHeadKeepSetSelector, _BatchedUnionKeepSetSelector, @@ -71,36 +70,6 @@ def _per_head_keep_oracle( return torch.stack(rows) -def _legacy_union(scores: torch.Tensor, keep_count: int) -> torch.Tensor: - row_top = { - int(index) - for row in scores - for index in _stable_topk(row, row.numel(), keep_count).tolist() - } - combined = scores.max(dim=0).values - ordered = sorted( - row_top, - key=lambda index: (-float(combined[index]), index), - ) - return torch.tensor(ordered[:keep_count], dtype=torch.long) - - -@pytest.mark.parametrize("rows,width,keep_count", [(2, 8, 4), (5, 17, 7)]) -def test_direct_union_topk_matches_legacy_union_with_heavy_ties(rows, width, keep_count): - for seed in range(25): - generator = torch.Generator().manual_seed(seed) - scores = torch.randint( - -2, - 3, - (rows, width), - generator=generator, - dtype=torch.int32, - ).to(torch.float32) - combined = scores.max(dim=0).values - direct = _stable_topk(combined, width, keep_count).to(torch.long) - assert torch.equal(direct, _legacy_union(scores, keep_count)) - - @pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) @pytest.mark.parametrize("normalize_scores", [False, True]) def test_per_head_selection_matches_torch_oracle_on_selector_stream( @@ -149,64 +118,11 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( assert torch.equal(second, expected) -def test_union_eager_runs_the_registered_cute_op(): - _require_cute_topk_op() - device = torch.device("cuda", torch.cuda.current_device()) - # The fused CuTe pipeline is the production union-row producer; these - # selector tests stage prepared union rows into ``combined`` directly. - scores = torch.randn(2, 4, 96, dtype=torch.float32, device=device) - selector = _BatchedUnionKeepSetSelector( - rows=4, - width=96, - keep_count=64, - dtype=torch.float32, - device=device, - max_requests=2, - ) - selector.combined.copy_(scores.amax(dim=1)) - selector.select_prepared_union_scores() - torch.cuda.synchronize(device) - assert torch.all(selector.keep[:, 1:] >= selector.keep[:, :-1]) - - -def test_prepared_union_rows_select_exact_indices(): - _require_cute_topk_op() - device = torch.device("cuda", torch.cuda.current_device()) - request_count, width, keep_count = 2, 97, 64 - generator = torch.Generator(device=device).manual_seed(53) - combined_rows = torch.randn( - request_count, - width, - generator=generator, - dtype=torch.float32, - device=device, - ) - valid_widths = torch.tensor([83, 91], dtype=torch.int32, device=device) - - selector = _BatchedUnionKeepSetSelector( - rows=7, - width=width, - keep_count=keep_count, - dtype=torch.float32, - device=device, - max_requests=request_count, - ) - selector.valid_widths.copy_(valid_widths) - selector.combined.copy_(combined_rows) - selector.select_prepared_union_scores() - actual_keep = selector.keep.cpu() - expected_combined = combined_rows.cpu() - torch.cuda.synchronize(device) - - for request, valid_width in enumerate(valid_widths.cpu().tolist()): - expected_keep = torch.sort( - _stable_topk(expected_combined[request], valid_width, keep_count) - ).values - assert torch.equal(actual_keep[request], expected_keep) - - @pytest.mark.parametrize("keep_count,width", [(4, 64), (4096, 4224), (8192, 9216)]) def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): + # Heavily tied integer scores with ragged valid widths and per-request + # prompt rebase: the strongest oracle over the direct union top-k path + # (it subsumes the sorted-output and exact-indices smoke variants). _require_cute_topk_op() device = torch.device("cuda", torch.cuda.current_device()) prompt_len = 17 @@ -248,12 +164,6 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid assert torch.equal(actual[request], expected_decode) -# The retired split-union preparation (``prepare_union_scores``) is covered -# by its standalone reference copy in test_triattention_cute_union_fusion.py, -# which validates it against a pure-torch oracle and uses it as the -# fused-pipeline equivalence reference. - - @pytest.mark.parametrize("per_layer", [False, True]) @pytest.mark.parametrize("normalize_scores", [False, True]) def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, normalize_scores): @@ -329,8 +239,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) # The compact op ships only the pipelined bf16 kernels: pools use the # supported geometry (bf16, 32-token pages, head_dim 64), and the kept # ordinals are spread across all three pages per request so the moves - # still cross page boundaries. The bf16-exact ``arange % 251`` payload - # keeps every wrong-move byte pattern distinguishable. + # still cross page boundaries. device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 num_layers = 2 @@ -343,22 +252,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) head_dim = 64 protected_tails = [2, 1] page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - initial_pools = [ - ( - ( - torch.arange( - 6 * 2 * num_kv_heads * tokens_per_block * head_dim, - dtype=torch.int32, - device=device, - ) - + layer * 37 - ) - % 251 - ) - .view(6, 2, num_kv_heads, tokens_per_block, head_dim) - .to(torch.bfloat16) - for layer in range(num_layers) - ] + initial_pools = _make_ramp_pools(num_layers, device=device) pools = [pool.clone() for pool in initial_pools] # Kept ordinals are decode-only but hold absolute positions; the pinned @@ -391,21 +285,13 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) device=device, ) - compaction = BatchedKVCacheCompaction( + compaction = _build_compaction( eviction_mode=eviction_mode, layer_pools=pools, - dense_layers=[0, 1], - swa_layers=[], - layer_group_representative={0: 0, 1: 1}, - layer_pool_keys=[("dense", 0), ("dense", 0)], kept_token_ordinals=keep.to(torch.int32), valid_sequence_lengths=torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device), kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), - page_table_slots={0: 0, 1: 0}, - request_count=request_count, prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), - decode_keep_count=decode_keep_count, - swa_window=None, protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) @@ -454,15 +340,12 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 num_layers = 2 - num_kv_heads = 2 # bf16 pools in the compact op's supported geometry; three 32-token pages # per request with a 80-token sequence keep the moves page-crossing. seq_len = 80 decode_keep_count = 3 prompt_lens = [2, 5] protected_tails = [2, 1] - tokens_per_block = 32 - head_dim = 64 decode_widths = [seq_len - prompt_len for prompt_len in prompt_lens] width = max(decode_widths) @@ -486,40 +369,17 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): ) page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - initial_pools = [ - ( - ( - torch.arange( - 6 * 2 * num_kv_heads * tokens_per_block * head_dim, - dtype=torch.int32, - device=device, - ) - + layer * 37 - ) - % 251 - ) - .view(6, 2, num_kv_heads, tokens_per_block, head_dim) - .to(torch.bfloat16) - for layer in range(num_layers) - ] + initial_pools = _make_ramp_pools(num_layers, device=device) cohort_pools = [pool.clone() for pool in initial_pools] keep_cuda = keep.to(device) valid_seq_lens = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) - cohort_compaction = BatchedKVCacheCompaction( - eviction_mode="union", + cohort_compaction = _build_compaction( layer_pools=cohort_pools, - dense_layers=[0, 1], - swa_layers=[], - layer_group_representative={0: 0, 1: 1}, - layer_pool_keys=[("dense", 0), ("dense", 0)], kept_token_ordinals=keep_cuda, valid_sequence_lengths=valid_seq_lens, kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), - page_table_slots={0: 0, 1: 0}, - request_count=request_count, prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, - swa_window=None, protected_tail_capacity=max(protected_tails), ) _set_protected_tails(cohort_compaction, protected_tails) @@ -527,21 +387,14 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): expected_pools = [pool.clone() for pool in initial_pools] for request in range(request_count): - single_compaction = BatchedKVCacheCompaction( - eviction_mode="union", + single_compaction = _build_compaction( layer_pools=expected_pools, - dense_layers=[0, 1], - swa_layers=[], - layer_group_representative={0: 0, 1: 1}, - layer_pool_keys=[("dense", 0), ("dense", 0)], kept_token_ordinals=keep_cuda[request : request + 1], valid_sequence_lengths=valid_seq_lens[request : request + 1], kv_block_offsets=_encode_block_offsets(page_tables[request : request + 1].unsqueeze(0)), - page_table_slots={0: 0, 1: 0}, request_count=1, prompt_offsets=torch.tensor([prompt_lens[request]], dtype=torch.int32, device=device), decode_keep_count=decode_keep_count, - swa_window=None, protected_tail_capacity=protected_tails[request], ) _set_protected_tails(single_compaction, [protected_tails[request]]) @@ -595,29 +448,13 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): ) expected_keep = torch.tensor([[[1, 6], [2, 5], [3, 7]]], dtype=torch.int32, device=device) - pools = [] - for layer, (table, values) in enumerate(zip(layer_tables, score_values)): - pool = ( - ( - ( - torch.arange( - 2 * 2 * tokens_per_block * head_dim, - dtype=torch.int32, - device=device, - ) - + layer * 37 - ) - % 251 - ) - .view(2, 2, 1, tokens_per_block, head_dim) - .to(torch.bfloat16) - ) + 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 - pools.append(pool) initial_pools = [pool.clone() for pool in pools] q_real = torch.zeros(num_layers, num_q_heads, num_freqs, dtype=torch.float32, device=device) @@ -668,11 +505,10 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): keep_set_selector.select_requests(scores, normalize_scores=False) assert torch.equal(keep_set_selector.keep, expected_keep) - batched_compaction = BatchedKVCacheCompaction( + batched_compaction = _build_compaction( eviction_mode="per_layer_perhead", layer_pools=pools, dense_layers=dense_layers, - swa_layers=[], layer_group_representative=layer_group_representative, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], kept_token_ordinals=keep_set_selector.keep[:1], @@ -682,15 +518,14 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): request_count=1, prompt_offsets=torch.zeros(1, dtype=torch.int32, device=device), decode_keep_count=keep_count, - swa_window=None, protected_tail_capacity=0, ) _set_protected_tails(batched_compaction, [0]) batched_compaction.compact() torch.cuda.synchronize(device) - for layer, (before_pool, after_pool, table) in enumerate( - zip(initial_pools, pools, layer_tables) + 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 @@ -709,8 +544,7 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): 32-token pages, head_dim 64): the request spans three pages so that compacting to two pages still releases one physical page for reuse. Token scores are tracked in a host-side mirror and the expected keep - sets are derived from it, replacing the hand-written score tables of the - old 4-token-page fixture. + sets are derived from it. """ pytest.importorskip("cutlass") import tensorrt_llm @@ -860,11 +694,9 @@ def expected_keep() -> torch.Tensor: prompt_offsets_buffer=score_staging.token_starts_device, ) score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") - batched_compaction = BatchedKVCacheCompaction( - eviction_mode="union", + batched_compaction = _build_compaction( layer_pools=[pool], dense_layers=[0], - swa_layers=[], layer_group_representative={0: 0}, layer_pool_keys=[("pool", 0)], kept_token_ordinals=keep_set_selector.keep[:1], @@ -874,7 +706,6 @@ def expected_keep() -> torch.Tensor: request_count=1, prompt_offsets=score_staging.token_starts_device[:1], decode_keep_count=keep_count, - swa_window=None, protected_tail_capacity=protected_tail, ) _set_protected_tails(batched_compaction, [protected_tail]) @@ -974,12 +805,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): device = torch.device("cuda", torch.cuda.current_device()) dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) - initial_pools = [ - ((torch.arange(6 * 2 * 1 * 32 * 64, dtype=torch.int32, device=device) + layer * 37) % 251) - .view(6, 2, 1, 32, 64) - .to(torch.bfloat16) - for layer in range(2) - ] + initial_pools = _make_ramp_pools(2, num_kv_heads=1, device=device) pools = [pool.clone() for pool in initial_pools] # Decode-only kept ordinals holding absolute positions past the prompt. keep = torch.tensor( @@ -989,8 +815,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): ) valid_seq_lens = torch.tensor([64, 56], dtype=torch.int32, device=device) protected_tails = [2, 1] - compaction = BatchedKVCacheCompaction( - eviction_mode="union", + compaction = _build_compaction( layer_pools=pools, dense_layers=[0], swa_layers=[1], @@ -1000,9 +825,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): valid_sequence_lengths=valid_seq_lens, kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), page_table_slots={0: 0, 1: 1}, - request_count=2, prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), - decode_keep_count=4, swa_window=2, protected_tail_capacity=max(protected_tails), ) @@ -1070,8 +893,7 @@ def test_cache_families_read_the_staged_move_offsets_rows(): staged_rows = torch.zeros(2, 3, dtype=torch.int32, device=device) dense_offsets_row = staged_rows[0] swa_offsets_row = staged_rows[1] - compaction = BatchedKVCacheCompaction( - eviction_mode="union", + compaction = _build_compaction( layer_pools=pools, dense_layers=[0], swa_layers=[1], @@ -1081,9 +903,7 @@ def test_cache_families_read_the_staged_move_offsets_rows(): valid_sequence_lengths=torch.tensor([8, 7], dtype=torch.int32, device=device), kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), page_table_slots={0: 0, 1: 1}, - request_count=2, prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), - decode_keep_count=4, swa_window=2, protected_tail_capacity=2, dense_move_offsets=dense_offsets_row, diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 9062e4a37cbf..5ece2dd60765 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -13,8 +13,6 @@ _TOKENS_PER_BLOCK = 32 _NUM_KV_HEADS = 2 _BATCH_SIZE = 2 -_MAX_PAGES_PER_SEQUENCE = 3 -_NUM_PAGES = _BATCH_SIZE * _MAX_PAGES_PER_SEQUENCE _PAGE_INDEX_DIVISOR = 2 # Kernel-name substrings for the profiler probes below. The pipelined bf16 @@ -52,8 +50,10 @@ def _make_pools( dtype: torch.dtype, head_dim: int, page_index_scale: int = _PAGE_INDEX_DIVISOR, + pages_per_seq: int = 3, + sequential_pages: bool = False, ) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: - num_pages = _NUM_PAGES * page_index_scale // _PAGE_INDEX_DIVISOR + num_pages = _BATCH_SIZE * pages_per_seq * page_index_scale // _PAGE_INDEX_DIVISOR shape = ( num_pages, 2, @@ -67,9 +67,15 @@ def _make_pools( for layer in range(num_layers) ] pools = [pool.cuda() for pool in pools_cpu] - raw_pages = [[4, 1, 5], [2, 0, 3]] - assert set(raw_pages[0]).isdisjoint(raw_pages[1]) - raw_page_table = torch.tensor(raw_pages, dtype=torch.int32, device="cuda") + if sequential_pages: + raw_page_table = torch.arange( + _BATCH_SIZE * pages_per_seq, dtype=torch.int32, device="cuda" + ).reshape(_BATCH_SIZE, pages_per_seq) + else: + assert pages_per_seq == 3 + raw_pages = [[4, 1, 5], [2, 0, 3]] + assert set(raw_pages[0]).isdisjoint(raw_pages[1]) + raw_page_table = torch.tensor(raw_pages, dtype=torch.int32, device="cuda") page_table = _encode_k_block_offsets(raw_page_table, page_index_scale) page_tables = [page_table] * num_layers assert page_tables[0].stride(0) == 2 * page_tables[0].shape[1] @@ -177,78 +183,72 @@ def _compact( ) -@pytest.mark.parametrize( - "dtype,head_dim", - [ - (torch.bfloat16, 64), - (torch.bfloat16, 128), - ], -) -@pytest.mark.parametrize( - "destination_base,page_index_scale", - [(0, 2), (2, 2), (2, 4)], -) -def test_sparse_kv_cache_compact_layers(dtype, head_dim, destination_base, page_index_scale): - pools_cpu, pools, page_tables = _make_pools(3, dtype, head_dim, page_index_scale) - page_tables_cpu = [page_table.cpu() for page_table in page_tables] - source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) - source_row = torch.tensor([2, 5, 8, 3, 7, 10], dtype=torch.int32) - source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() - expected = _reference_compact( - pools_cpu, - page_tables_cpu, - source_indices, - source_offsets, - destination_base, +_SMALL_ROW = [2, 5, 8, 3, 7, 10] +# One byte-equality case per launch shape the op serves: dtype/head_dim and +# destination/page-scale sweeps, per-request destination bases (one launch +# mixing pinned-prompt lengths), a 3-D per-layer source with layer routing, +# and a multi-tile launch (40/35 moves across 24-page sequences). +_LAYER_CASES = [ + pytest.param( + dict(head_dim=head_dim, dest=dest, scale=scale), + id=f"bf16_h{head_dim}_dest{dest}_scale{scale}", ) - arguments = _device_arguments(pools, source_indices, source_offsets) - - _compact(pools, page_tables, arguments, destination_base) - torch.cuda.synchronize() - - for actual, reference in zip(pools, expected): - assert torch.equal(actual.cpu(), reference) - - -def test_sparse_kv_cache_compact_layers_per_request_destination_bases(): - # One launch may mix pinned-prompt lengths: each request lands at its own - # destination base. - pools_cpu, pools, page_tables = _make_pools(2, torch.bfloat16, 64) - page_tables_cpu = [page_table.cpu() for page_table in page_tables] - source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) - source_row = torch.tensor([3, 5, 8, 6, 7, 10], dtype=torch.int32) - source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() - destination_bases = [2, 5] - expected = _reference_compact( - pools_cpu, - page_tables_cpu, - source_indices, - source_offsets, - destination_bases, + for head_dim in (64, 128) + for dest, scale in ((0, 2), (2, 2), (2, 4)) +] + [ + pytest.param( + dict(head_dim=64, num_layers=2, row=[3, 5, 8, 6, 7, 10], dest=[2, 5]), + id="per_request_destination_bases", + ), + pytest.param( + dict( + head_dim=64, + num_layers=2, + dest=2, + indices3d=[ + [[2, 5, 8, 3, 7, 10], [3, 6, 9, 2, 5, 8]], + [[3, 7, 10, 2, 6, 9], [2, 5, 9, 3, 6, 10]], + [[4, 7, 9, 3, 6, 8], [3, 5, 8, 4, 7, 10]], + ], + layer_indices=[2, 0], + ), + id="per_layer_source", + ), + pytest.param( + dict( + head_dim=64, + num_layers=2, + dest=2, + pages_per_seq=24, + sequential_pages=True, + offsets=(0, 40, 75), + row=list(range(40, 80)) + list(range(36, 71)), + ), + id="multiple_tiles", + ), +] + + +@pytest.mark.parametrize("case", _LAYER_CASES) +def test_sparse_kv_cache_compact_layers(case): + pools_cpu, pools, page_tables = _make_pools( + case.get("num_layers", 3), + torch.bfloat16, + case["head_dim"], + case.get("scale", _PAGE_INDEX_DIVISOR), + pages_per_seq=case.get("pages_per_seq", 3), + sequential_pages=case.get("sequential_pages", False), ) - arguments = _device_arguments(pools, source_indices, source_offsets) - - _compact(pools, page_tables, arguments, destination_bases) - torch.cuda.synchronize() - - for actual, reference in zip(pools, expected): - assert torch.equal(actual.cpu(), reference) - - -def test_sparse_kv_cache_compact_layers_per_layer_source(): - pools_cpu, pools, page_tables = _make_pools(2, torch.bfloat16, 64) page_tables_cpu = [page_table.cpu() for page_table in page_tables] - source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) - source_indices = torch.tensor( - [ - [[2, 5, 8, 3, 7, 10], [3, 6, 9, 2, 5, 8]], - [[3, 7, 10, 2, 6, 9], [2, 5, 9, 3, 6, 10]], - [[4, 7, 9, 3, 6, 8], [3, 5, 8, 4, 7, 10]], - ], - dtype=torch.int32, - ) - source_layer_indices = torch.tensor([2, 0], dtype=torch.int32) - destination_base = 2 + source_offsets = torch.tensor(case.get("offsets", (0, 3, 6)), dtype=torch.int32) + if "indices3d" in case: + source_indices = torch.tensor(case["indices3d"], dtype=torch.int32) + source_layer_indices = torch.tensor(case["layer_indices"], dtype=torch.int32) + else: + source_row = torch.tensor(case.get("row", _SMALL_ROW), dtype=torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + source_layer_indices = None + destination_base = case.get("dest", 0) expected = _reference_compact( pools_cpu, page_tables_cpu, @@ -257,67 +257,8 @@ def test_sparse_kv_cache_compact_layers_per_layer_source(): destination_base, source_layer_indices, ) - arguments = _device_arguments( - pools, - source_indices, - source_offsets, - source_layer_indices, - ) - - _compact(pools, page_tables, arguments, destination_base) - torch.cuda.synchronize() - - for actual, reference in zip(pools, expected): - assert torch.equal(actual.cpu(), reference) - - -def test_sparse_kv_cache_compact_layers_rejects_flat_source_with_layer_indices(): - # A flat [kv_heads, total] source with per-layer indices would silently - # read layer 0 for every launch; the op rejects the combination instead. - _, pools, page_tables = _make_pools(2, torch.bfloat16, 64) - source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) - source_row = torch.tensor([2, 5, 8, 3, 7, 10], dtype=torch.int32) - source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() - source_layer_indices = torch.tensor([0, 0], dtype=torch.int32) arguments = _device_arguments(pools, source_indices, source_offsets, source_layer_indices) - with pytest.raises(RuntimeError, match="require 3-D per-layer source_indices"): - _compact(pools, page_tables, arguments, 0) - - -def test_sparse_kv_cache_compact_layers_multiple_tiles(): - num_layers = 2 - max_pages_per_sequence = 24 - num_pages = _BATCH_SIZE * max_pages_per_sequence - shape = (num_pages, 2, _NUM_KV_HEADS, _TOKENS_PER_BLOCK, 64) - numel = torch.Size(shape).numel() - pools_cpu = [ - ((torch.arange(numel, dtype=torch.int32) + layer * 37) % 251) - .reshape(shape) - .to(torch.bfloat16) - for layer in range(num_layers) - ] - pools = [pool.cuda() for pool in pools_cpu] - raw_page_table = torch.arange(num_pages, dtype=torch.int32, device="cuda").reshape( - _BATCH_SIZE, max_pages_per_sequence - ) - page_table = _encode_k_block_offsets(raw_page_table) - page_tables = [page_table] * num_layers - assert page_tables[0].stride(0) == 2 * max_pages_per_sequence - page_tables_cpu = [table.cpu() for table in page_tables] - source_offsets = torch.tensor([0, 40, 75], dtype=torch.int32) - source_row = torch.cat((torch.arange(40, 80), torch.arange(36, 71))).to(torch.int32) - source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() - destination_base = 2 - expected = _reference_compact( - pools_cpu, - page_tables_cpu, - source_indices, - source_offsets, - destination_base, - ) - arguments = _device_arguments(pools, source_indices, source_offsets) - _compact(pools, page_tables, arguments, destination_base) torch.cuda.synchronize() @@ -330,7 +271,7 @@ def test_sparse_kv_cache_compact_layers_cuda_graph_replay(): pools_cpu, pools, page_tables = _make_pools(3, torch.bfloat16, 64) page_tables_cpu = [page_table.cpu() for page_table in page_tables] source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) - source_row = torch.tensor([2, 5, 8, 3, 7, 10], dtype=torch.int32) + source_row = torch.tensor(_SMALL_ROW, dtype=torch.int32) source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() replay_row = torch.tensor([3, 6, 9, 2, 5, 8], dtype=torch.int32) replay_indices = replay_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() @@ -362,7 +303,7 @@ def test_sparse_kv_cache_compact_layers_cuda_graph_replay(): # # The fast path only dispatches for bf16 pools with head_dim 64/128 and # 32/128-token pages, so the cases below use their own builder instead of the -# 4-token-page fixtures above. +# 3-page fixtures above. _FAST_BATCH_SIZE = 3 # Per-request move counts: two ragged tiles plus a full one (pipeline steady @@ -504,45 +445,56 @@ def _run_fast_geometry_case(case: _FastGeometryCase) -> list[torch.Tensor]: # The full fast-path gate matrix. 128-token pages are the geometry the ported # kernel was written for; 32-token pages and head_dim 128 are this tree's -# production configuration. +# production configuration. The per-layer-source row keeps the 3-D routing +# path runnable through the fast kernel. _FAST_GEOMETRY_MATRIX = [(64, 32), (128, 32), (64, 128), (128, 128)] -@pytest.mark.parametrize("head_dim,tokens_per_block", _FAST_GEOMETRY_MATRIX) -def test_sparse_kv_cache_compact_layers_fast_geometry(head_dim, tokens_per_block): +@pytest.mark.parametrize( + "head_dim,tokens_per_block,per_layer", + [(h, t, False) for h, t in _FAST_GEOMETRY_MATRIX] + [(64, 32, True)], +) +def test_sparse_kv_cache_compact_layers_fast_geometry(head_dim, tokens_per_block, per_layer): # Eligible geometry always dispatches the pipelined kernel; the # byte-compare against the CPU reference is the correctness net. - case = _make_fast_geometry_case(head_dim, tokens_per_block) - expected = _run_fast_geometry_case(case) - for actual, reference in zip(case.pools, expected): - assert torch.equal(actual.cpu(), reference) - - -def test_sparse_kv_cache_compact_layers_fast_geometry_per_layer_source(): - case = _make_fast_geometry_case(64, 32, per_layer_sources=True) + case = _make_fast_geometry_case(head_dim, tokens_per_block, per_layer_sources=per_layer) expected = _run_fast_geometry_case(case) for actual, reference in zip(case.pools, expected): assert torch.equal(actual.cpu(), reference) +# One representative per reject family: dtype outside the bf16-only gate, +# head_dim outside the gate, page size outside the gate, and a flat 2-D +# source combined with per-layer indices (which would silently read layer 0 +# for every launch). There is no fallback kernel: every reject must fail +# loudly and leave the pools untouched. @pytest.mark.parametrize( - "dtype,head_dim,tokens_per_block", + "dtype,head_dim,tokens_per_block,flat_with_layer_indices,match", [ - (torch.float16, 64, 32), # dtype outside the bf16-only gate - (torch.bfloat16, 256, 32), # head_dim outside the gate - (torch.bfloat16, 64, 16), # page size outside the gate + pytest.param(torch.float16, 64, 32, False, "bf16|BF16", id="dtype_outside_gate"), + pytest.param(torch.bfloat16, 256, 32, False, "bf16|BF16", id="head_dim_outside_gate"), + pytest.param(torch.bfloat16, 64, 16, False, "bf16|BF16", id="page_size_outside_gate"), + pytest.param( + torch.bfloat16, + 64, + 32, + True, + "require 3-D per-layer source_indices", + id="flat_source_with_layer_indices", + ), ], ) -def test_sparse_kv_cache_compact_layers_rejects_unsupported_geometry( - dtype, head_dim, tokens_per_block +def test_sparse_kv_cache_compact_layers_rejects_invalid_launch( + dtype, head_dim, tokens_per_block, flat_with_layer_indices, match ): - # There is no fallback kernel: near-miss geometries must fail loudly - # instead of silently degrading, and the pools must stay untouched. case = _make_fast_geometry_case(head_dim, tokens_per_block, dtype=dtype) + source_layer_indices = case.source_layer_indices + if flat_with_layer_indices: + source_layer_indices = torch.tensor([0, 0], dtype=torch.int32) arguments = _device_arguments( - case.pools, case.source_indices, case.source_offsets, case.source_layer_indices + case.pools, case.source_indices, case.source_offsets, source_layer_indices ) - with pytest.raises((RuntimeError, ValueError), match="bf16|BF16"): + with pytest.raises((RuntimeError, ValueError), match=match): _compact( case.pools, case.page_tables, From f21f1ddd602b1307686dff20c0f2f128ebc9de51 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 07:01:36 -0700 Subject: [PATCH 083/178] [None][refactor] Flatten keep-set selector hierarchy into direct kernel calls Behavior-identical de-OOP pass over the TriAttention selection plumbing, mirroring the flat style of the other sparse backends: - _BatchedKeepSetSelectorBase + _BatchedUnionKeepSetSelector + _BatchedPerHeadKeepSetSelector (3-class hierarchy) collapse into one plain _BatchedKeepSetSelector whose mode differences are if-branches; _bind_selection_rows indirection inlined into the constructor branches. - _CrossRequestSelectionPlan NamedTuple deleted; its fields are passed directly at the sole construction site. - The _build_cross_request_keep_set_selector factory staticmethod is replaced by direct construction in _fixed_resources_for. - isinstance-based union dispatch becomes an eviction_mode check; the getattr provisional-buffer probe becomes a plain mode branch. Same kernel launches with the same arguments, buffer allocation order, validation order, and error messages; union, per_head, and per_layer_perhead all verified token-identical A/B against the pre-refactor tree (Qwen3-8B digest legs) with the unit suite green. Net -117 product lines; tests updated for the merged class name. Signed-off-by: tianruih --- .../triattention/triattention.py | 319 ++++++------------ .../_torch/kv_cache_compression/conftest.py | 8 + .../test_triattention_draft_cocompaction.py | 4 +- .../test_triattention_fused_settle_pack.py | 5 +- .../test_triattention_pipeline.py | 7 +- .../test_triattention_selection_compaction.py | 13 +- 6 files changed, 124 insertions(+), 232 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 04c63c32327c..b71e1d69c7ef 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -54,7 +54,7 @@ """ from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple import torch @@ -99,21 +99,6 @@ class _FixedScoreStreamMismatch(RuntimeError): """Raised when fixed score staging buffers are used from another CUDA stream.""" -class _CrossRequestSelectionPlan(NamedTuple): - """Selection dimensions used to allocate reusable fixed buffers.""" - - eviction_mode: str - dense_layers: Tuple[int, ...] - num_query_heads: int - num_kv_heads: int - rows: int - width: int - keep_count: int - dtype: torch.dtype - device: torch.device - max_requests: int - - class _RuntimeKVLayout(NamedTuple): """Manager-lifetime layer and pool views used by every eviction.""" @@ -132,24 +117,48 @@ class _RuntimeKVLayout(NamedTuple): pool_view_fingerprint: Tuple[tuple, ...] -class _BatchedKeepSetSelectorBase: - """Shared fixed buffers and row views for keep-set selectors.""" +class _BatchedKeepSetSelector: + """Fixed ``[request, ...]`` keep-set selection buffers for every eviction mode. + + Union mode: the fused score+stats+union CuTe pipeline is THE union score + producer -- it writes the normalized per-request union rows straight into + ``combined``, so the selector owns only the top-k settle-and-pack stage. + The per-head modes (``per_head``, ``per_layer_perhead``) own the per-head + score preparation buffers and settle one selection row per head or per + (layer, head). + """ def __init__( self, *, eviction_mode: str, - dense_layers: Tuple[int, ...], - num_query_heads: int, - num_kv_heads: int, width: int, keep_count: int, - selection_rows_per_request: int = 1, - prompt_offsets_buffer: Optional[torch.Tensor] = None, dtype: torch.dtype, device: torch.device, max_requests: int, + dense_layers: Tuple[int, ...] = (), + num_query_heads: int = 0, + num_kv_heads: int = 0, + rows: Optional[int] = None, + prompt_offsets_buffer: Optional[torch.Tensor] = None, ) -> None: + if eviction_mode == "union": + if rows is None or rows <= 0: + raise ValueError("cross-request selection requires rows > 0") + selection_rows = 1 + elif eviction_mode in ("per_head", "per_layer_perhead"): + if not dense_layers or min(num_query_heads, num_kv_heads, max_requests) <= 0: + raise ValueError( + "per-head selection requires positive layer, head, and request counts" + ) + if num_query_heads % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + selection_rows = ( + num_kv_heads if eviction_mode == "per_head" else len(dense_layers) * num_kv_heads + ) + else: + raise ValueError(f"unsupported per-head eviction mode: {eviction_mode}") if width <= keep_count or keep_count <= 0: raise ValueError("keep-set selection requires width > keep_count > 0") if max_requests <= 0: @@ -169,7 +178,7 @@ def __init__( # decode-relative and these offsets rebase emitted ordinals, so one # cohort may mix prompt lengths. ``row_prompt_offsets`` is the # row-major expansion consumed by the finalizer. - self.selection_rows_per_request = int(selection_rows_per_request) + self.selection_rows_per_request = int(selection_rows) # Optional compaction move packing fused into the settle launch; set # once the compaction buffers exist (see ``fuse_move_source_pack``). self._move_source_pack = None @@ -195,6 +204,56 @@ def __init__( dtype=torch.int32, device=self.device, ) + if eviction_mode == "union": + self.combined = torch.empty((max_requests, width), dtype=dtype, device=self.device) + self.final_indices = torch.empty( + (max_requests, keep_count), dtype=torch.int32, device=self.device + ) + # Kept decode ordinals only: rows are prompt-length independent, so + # one selector serves cohorts with mixed prompt lengths. + self.keep = torch.empty( + (max_requests, self.keep_count), dtype=torch.int32, device=self.device + ) + # Row-major views consumed by the top-k settle launch. + self._selection_scores_rows = self.combined + self._selection_row_lengths = self.valid_widths + self._provisional_rows = self.final_indices + self._keep_rows = self.keep + else: + self.num_layers = len(self.dense_layers) + self.selection_rows = selection_rows + score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) + self.row_mean = torch.empty(score_shape[:-1] + (1,), dtype=dtype, device=self.device) + self.row_std = torch.empty_like(self.row_mean) + self.selection_scores = torch.empty( + (self.max_requests, self.selection_rows, self.width), + dtype=dtype, + device=self.device, + ) + self.row_seq_lens = torch.full( + (self.max_requests, self.selection_rows), + self.width, + dtype=torch.int32, + device=self.device, + ) + selection_shape = (self.max_requests, self.selection_rows, self.keep_count) + self.top_indices_i32 = torch.empty( + selection_shape, dtype=torch.int32, device=self.device + ) + # Kept decode ordinals only: rows are prompt-length independent, so + # one selector serves cohorts with mixed prompt lengths. + self.keep = torch.empty(selection_shape, dtype=torch.int32, device=self.device) + self.selection_scores_flat = self.selection_scores.view( + self.max_requests * self.selection_rows, self.width + ) + self.row_seq_lens_flat = self.row_seq_lens.view(-1) + self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) + self.keep_flat = self.keep.view(-1, self.keep_count) + # Row-major views consumed by the top-k settle launch. + self._selection_scores_rows = self.selection_scores_flat + self._selection_row_lengths = self.row_seq_lens_flat + self._provisional_rows = self.top_indices_i32_flat + self._keep_rows = self.keep_flat def refresh_row_prompt_offsets(self) -> None: """Re-expand the per-request prompt offsets into their row-major view. @@ -206,19 +265,6 @@ def refresh_row_prompt_offsets(self) -> None: self.prompt_offsets.unsqueeze(1).expand(-1, self.selection_rows_per_request) ) - def _bind_selection_rows( - self, - scores_rows: torch.Tensor, - row_lengths: torch.Tensor, - provisional_indices: torch.Tensor, - keep_rows: torch.Tensor, - ) -> None: - """Keep row-major views of the buffers the top-k selection reads.""" - self._selection_scores_rows = scores_rows - self._selection_row_lengths = row_lengths - self._provisional_rows = provisional_indices - self._keep_rows = keep_rows - def fuse_move_source_pack(self, pack_arguments) -> None: """Pack compaction move sources inside this selector's settle launch. @@ -318,55 +364,6 @@ def _select_top_tokens(self) -> None: num_warps=4, ) - -class _BatchedUnionKeepSetSelector(_BatchedKeepSetSelectorBase): - """Persistent ``[request, ...]`` buffers for union selection. - - The fused score+stats+union CuTe pipeline is THE union score producer: - it writes the normalized per-request union rows straight into - ``combined``, so this selector owns only the top-k settle-and-pack - stage. - """ - - def __init__( - self, - rows: int, - width: int, - keep_count: int, - *, - dtype: torch.dtype, - device: torch.device, - max_requests: int, - dense_layers: Tuple[int, ...] = (), - num_query_heads: int = 0, - num_kv_heads: int = 0, - prompt_offsets_buffer: Optional[torch.Tensor] = None, - ) -> None: - if rows <= 0: - raise ValueError("cross-request selection requires rows > 0") - super().__init__( - eviction_mode="union", - dense_layers=dense_layers, - num_query_heads=num_query_heads, - num_kv_heads=num_kv_heads, - width=width, - keep_count=keep_count, - prompt_offsets_buffer=prompt_offsets_buffer, - dtype=dtype, - device=device, - max_requests=max_requests, - ) - self.combined = torch.empty((max_requests, width), dtype=dtype, device=self.device) - self.final_indices = torch.empty( - (max_requests, keep_count), dtype=torch.int32, device=self.device - ) - # Kept decode ordinals only: rows are prompt-length independent, so - # one selector serves cohorts with mixed prompt lengths. - self.keep = torch.empty( - (max_requests, self.keep_count), dtype=torch.int32, device=self.device - ) - self._bind_selection_rows(self.combined, self.valid_widths, self.final_indices, self.keep) - def select_prepared_union_scores(self) -> None: """Select from normalized union rows already written into ``combined``. @@ -375,85 +372,6 @@ def select_prepared_union_scores(self) -> None: """ self._select_top_tokens() - -class _BatchedPerHeadKeepSetSelector(_BatchedKeepSetSelectorBase): - """Fixed ``[request, ...]`` selector for both per-head modes.""" - - def __init__( - self, - *, - eviction_mode: str, - dense_layers: Tuple[int, ...], - num_query_heads: int, - num_kv_heads: int, - width: int, - keep_count: int, - dtype: torch.dtype, - device: torch.device, - max_requests: int, - prompt_offsets_buffer: Optional[torch.Tensor] = None, - ) -> None: - if eviction_mode not in ("per_head", "per_layer_perhead"): - raise ValueError(f"unsupported per-head eviction mode: {eviction_mode}") - if not dense_layers or min(num_query_heads, num_kv_heads, max_requests) <= 0: - raise ValueError("per-head selection requires positive layer, head, and request counts") - if num_query_heads % num_kv_heads: - raise ValueError("query heads must be divisible by KV heads") - selection_rows = ( - num_kv_heads if eviction_mode == "per_head" else len(dense_layers) * num_kv_heads - ) - super().__init__( - eviction_mode=eviction_mode, - dense_layers=dense_layers, - num_query_heads=num_query_heads, - num_kv_heads=num_kv_heads, - width=width, - keep_count=keep_count, - selection_rows_per_request=selection_rows, - prompt_offsets_buffer=prompt_offsets_buffer, - dtype=dtype, - device=device, - max_requests=max_requests, - ) - self.num_layers = len(self.dense_layers) - self.selection_rows = selection_rows - - score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) - self.row_mean = torch.empty(score_shape[:-1] + (1,), dtype=dtype, device=self.device) - self.row_std = torch.empty_like(self.row_mean) - self.selection_scores = torch.empty( - (self.max_requests, self.selection_rows, self.width), - dtype=dtype, - device=self.device, - ) - self.row_seq_lens = torch.full( - (self.max_requests, self.selection_rows), - self.width, - dtype=torch.int32, - device=self.device, - ) - selection_shape = (self.max_requests, self.selection_rows, self.keep_count) - self.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=self.device) - # Kept decode ordinals only: rows are prompt-length independent, so - # one selector serves cohorts with mixed prompt lengths. - self.keep = torch.empty( - (self.max_requests, self.selection_rows, self.keep_count), - dtype=torch.int32, - device=self.device, - ) - self.selection_scores_flat = self.selection_scores.view( - self.max_requests * self.selection_rows, self.width - ) - self.row_seq_lens_flat = self.row_seq_lens.view(-1) - self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) - self.keep_flat = self.keep.view(-1, self.keep_count) - self._bind_selection_rows( - self.selection_scores_flat, - self.row_seq_lens_flat, - self.top_indices_i32_flat, - self.keep_flat, - ) - def select_requests( self, scores: torch.Tensor, @@ -1067,7 +985,7 @@ class _EvictionBuffers: """Reusable fixed score and selection buffers for one runtime shape.""" score_staging: _FixedScoreStagingBuffers - keep_set_selector: Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector] + keep_set_selector: _BatchedKeepSetSelector class TriAttention(BaseKVCacheCompressionManager): @@ -1562,39 +1480,6 @@ def _configured_protected_tail_capacity(self) -> int: raise RuntimeError("KVCacheManagerV2 exposes an invalid protected-tail capacity") return capacity - @staticmethod - def _build_cross_request_keep_set_selector( - plan: _CrossRequestSelectionPlan, - *, - prompt_offsets_buffer: Optional[torch.Tensor] = None, - ) -> Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector]: - """Allocate one fixed ``[request, ...]`` keep-set selector.""" - if plan.eviction_mode == "union": - return _BatchedUnionKeepSetSelector( - plan.rows, - plan.width, - plan.keep_count, - dtype=plan.dtype, - device=plan.device, - max_requests=plan.max_requests, - dense_layers=plan.dense_layers, - num_query_heads=plan.num_query_heads, - num_kv_heads=plan.num_kv_heads, - prompt_offsets_buffer=prompt_offsets_buffer, - ) - return _BatchedPerHeadKeepSetSelector( - eviction_mode=plan.eviction_mode, - dense_layers=plan.dense_layers, - num_query_heads=plan.num_query_heads, - num_kv_heads=plan.num_kv_heads, - width=plan.width, - keep_count=plan.keep_count, - dtype=plan.dtype, - device=plan.device, - max_requests=plan.max_requests, - prompt_offsets_buffer=prompt_offsets_buffer, - ) - def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: """Drop this request's per-request length and eviction state.""" request_id = request.py_request_id @@ -2042,27 +1927,25 @@ def _fixed_resources_for( page_table_token_capacity=page_table_token_capacity, **draft_kwargs, ) - keep_set_selector = self._build_cross_request_keep_set_selector( - _CrossRequestSelectionPlan( - eviction_mode=self.eviction_mode, - dense_layers=tuple(layout.dense_layers), - num_query_heads=int(self._H), - num_kv_heads=int(first_pool.shape[2]), - rows=len(layout.dense_layers) * int(self._H), - width=decode_width, - keep_count=self.top_B, - dtype=torch.float32, - device=first_pool.device, - max_requests=request_capacity, - ), + keep_set_selector = _BatchedKeepSetSelector( + eviction_mode=self.eviction_mode, + dense_layers=tuple(layout.dense_layers), + num_query_heads=int(self._H), + num_kv_heads=int(first_pool.shape[2]), + rows=len(layout.dense_layers) * int(self._H), + width=decode_width, + keep_count=self.top_B, + dtype=torch.float32, + device=first_pool.device, + max_requests=request_capacity, prompt_offsets_buffer=score_staging.token_starts_device, ) # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. - provisional = getattr(keep_set_selector, "final_indices", None) - if provisional is None: - provisional = keep_set_selector.top_indices_i32 - provisional.zero_() + if self.eviction_mode == "union": + keep_set_selector.final_indices.zero_() + else: + keep_set_selector.top_indices_i32.zero_() score_staging.bind_score_launcher( keep_set_selector.valid_widths, "mean", @@ -2081,7 +1964,7 @@ def _batched_compaction_for( layout: _RuntimeKVLayout, prepared: Sequence[_PreparedEviction], score_staging: _FixedScoreStagingBuffers, - keep_set_selector: Union[_BatchedUnionKeepSetSelector, _BatchedPerHeadKeepSetSelector], + keep_set_selector: _BatchedKeepSetSelector, ): """Build or reuse the C++ compaction launches for one cohort.""" from .compaction import BatchedKVCacheCompaction @@ -2304,7 +2187,7 @@ def _evict_requests( keep_set_selector.refresh_row_prompt_offsets() try: - fused_union = isinstance(keep_set_selector, _BatchedUnionKeepSetSelector) + fused_union = keep_set_selector.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): # Union rounds run the fused pipeline (score, row stats, # normalization, and the cross-row union maximum in two CuTe diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 015302f3aa78..51836d691a5b 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -218,7 +218,9 @@ def make_fixed_resources_stubs(manager, *, decode_width=260): ) keep_set_selector = SimpleNamespace( valid_widths=torch.empty(8, dtype=torch.int32), + # The builder zero-fills the mode's provisional top-k buffer. top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), + final_indices=torch.zeros(8, 4, dtype=torch.int32), ) return layout, score_staging, keep_set_selector @@ -287,10 +289,16 @@ def mocked_eviction_internals(manager): """Run the real ``_evict_requests`` body around mocked GPU launches.""" score_staging = SimpleNamespace( launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), + launch_prepared_union_fusion=mock.Mock(), mark_page_tables_consumed=mock.Mock(), ) + # The dispatch reads the selector's own eviction mode, so the stub + # mirrors the manager's and carries both mode paths' launch surfaces. keep_set_selector = SimpleNamespace( + eviction_mode=manager.eviction_mode, + combined=torch.zeros(1), select_requests=mock.Mock(), + select_prepared_union_scores=mock.Mock(), refresh_row_prompt_offsets=mock.Mock(), ) resources = SimpleNamespace( 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 index 0d8e8ca1de37..29b9a3b95d72 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -379,8 +379,8 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): return_value=score_staging, ) as score_cls, mock.patch.object( - manager, - "_build_cross_request_keep_set_selector", + module, + "_BatchedKeepSetSelector", return_value=keep_set_selector, ), ): 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 index e5e35df1b007..75cc5b837d7f 100644 --- 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 @@ -22,7 +22,7 @@ BatchedKVCacheCompaction, ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _BatchedUnionKeepSetSelector, + _BatchedKeepSetSelector, ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _settle_ties_and_pack_compaction_sources_kernel, @@ -410,7 +410,8 @@ def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buff ] page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - selector = _BatchedUnionKeepSetSelector( + selector = _BatchedKeepSetSelector( + eviction_mode="union", rows=3, width=width, keep_count=keep_count, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 51d37c511ca0..97c2d4374c78 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -642,8 +642,8 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize return_value=score_staging, ), mock.patch.object( - manager, - "_build_cross_request_keep_set_selector", + module, + "_BatchedKeepSetSelector", return_value=keep_set_selector, ) as build_selection, ): @@ -653,8 +653,7 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize keep_set_selector.valid_widths, "mean", ) - plan = build_selection.call_args.args[0] - assert plan.eviction_mode == eviction_mode + assert build_selection.call_args.kwargs["eviction_mode"] == eviction_mode assert resources.score_staging is score_staging assert resources.keep_set_selector is keep_set_selector 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 index 7791857876b1..d6740ea43d67 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -10,8 +10,7 @@ from conftest import set_protected_tails as _set_protected_tails from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _BatchedPerHeadKeepSetSelector, - _BatchedUnionKeepSetSelector, + _BatchedKeepSetSelector, ) @@ -95,7 +94,7 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( device = torch.device("cuda", torch.cuda.current_device()) stream = torch.cuda.Stream(device=device) with torch.cuda.stream(stream): - selector = _BatchedPerHeadKeepSetSelector( + selector = _BatchedKeepSetSelector( eviction_mode=eviction_mode, dense_layers=tuple(range(layers)), num_query_heads=query_heads, @@ -137,7 +136,8 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid device=device, ).to(torch.float32) valid_widths = (width, width - 32) - selector = _BatchedUnionKeepSetSelector( + selector = _BatchedKeepSetSelector( + eviction_mode="union", rows=rows, width=width, keep_count=keep_count, @@ -488,7 +488,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): score_staging.round_starts_device.fill_(0) score_staging.valid_seq_lens_device.fill_(seq_len) score_staging.token_starts_device.fill_(0) - keep_set_selector = _BatchedPerHeadKeepSetSelector( + keep_set_selector = _BatchedKeepSetSelector( eviction_mode="per_layer_perhead", dense_layers=tuple(dense_layers), num_query_heads=num_q_heads, @@ -681,7 +681,8 @@ def expected_keep() -> torch.Tensor: decode_width=seq_len - prompt_len, page_table_token_capacity=seq_len + protected_tail, ) - keep_set_selector = _BatchedUnionKeepSetSelector( + keep_set_selector = _BatchedKeepSetSelector( + eviction_mode="union", rows=num_q_heads, width=seq_len - prompt_len, keep_count=keep_count, From a05be27f622f54e6f2bee387c4dc61abb87fae63 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 08:24:28 -0700 Subject: [PATCH 084/178] [None][refactor] Straight-line the eviction round flow One prepare, one round: prepare_eviction_workspace() is the single one-time constructor for the whole eviction stack (grouped-assert validation, buffer staging, eager CuTe compilation of exactly the mode-launched entries, C++ compaction launch data, fused settle-and-pack wiring) and returns a plain namespace of tensors, events, and compiled runners. stage_eviction_cohort() and run_eviction_round() are straight-line module functions that feed the kernels directly: build metadata -> phase gather -> score (fused union K1+K2, or score-only + decode-window gather + stats/reduce) -> top-k op -> settle-and-pack kernel -> C++ compacts. The layered launch chains and their owners are gone: _FixedScoreStagingBuffers, _BatchedKeepSetSelector, _FixedScoreGroup, MeanPhaseTable (now a plain table dict + module functions), and BatchedKVCacheCompaction (now stateless builders returning plain launch data plus run_cache_compactions). Classes remain only at the framework boundary (TriAttention) and for compiled CuTe artifacts. Records became plain dicts/tuples; the staging bool-reject protocol became typed errors; the lazy-compile state machines and one-shot binder are gone. Per-round re-validation of manager-enforced invariants is deleted in favor of one-time gates (host block-offset table shape, SWA scratch remapping, draft wiring parity); lifecycle guards, the 2^31 offset audits, and the CuTe geometry contract raise stay. All Triton kernel bodies are AST-identical to the previous revision; kernel launch sequences and arguments are unchanged, so all three eviction modes stay token-identical. Signed-off-by: tianruih --- .../triattention/compaction.py | 722 +++--- .../triattention/triattention.py | 2185 ++++++++--------- .../triattention_cute_score_fused.py | 2 +- .../triattention/triattention_kernels.py | 726 +----- .../_torch/kv_cache_compression/conftest.py | 122 +- .../test_triattention_cute_score.py | 308 ++- .../test_triattention_cute_union_fusion.py | 265 +- .../test_triattention_draft_cocompaction.py | 153 +- .../test_triattention_fused_settle_pack.py | 71 +- .../test_triattention_pipeline.py | 404 ++- .../test_triattention_selection_compaction.py | 338 +-- 11 files changed, 2345 insertions(+), 2951 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 9576293b916f..34ce298998f6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -19,69 +19,23 @@ staged V2 block offsets, this module packs per-request move indices with one Triton launch per compacted cache (one launch covers the target's dense and SWA families; a co-compressed draft adds a second) and then moves the -surviving KV in place with batched C++ compact launches. Inputs are plain -tensors, so any eviction method that produces a kept-token set per request -can drive it. A driver that finalizes the keep set in its own GPU launch can -take the target's dense/SWA packing over into that launch instead (see -``hand_move_source_pack_to_selection``); the draft always packs here. +surviving KV in place with batched C++ compact launches. Everything is plain +tensors and dicts: ``build_cache_compactions`` allocates the launch data once +per geometry (called by ``triattention.prepare_eviction_workspace``) and +``run_cache_compactions`` fires the kernels directly each round. A driver +that finalizes the keep set in its own GPU launch takes the target's +dense/SWA packing over into that launch (``fuse_dense_pack_into_selection``); +the draft always packs here. """ from collections import OrderedDict -from typing import Callable, Dict, List, NamedTuple, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple import torch _SUPPORTED_POOL_DTYPES = (torch.bfloat16,) -class _CppCompactGroup(NamedTuple): - """One layered sparse-KV updater launch over pools sharing a block table.""" - - pools: Tuple[torch.Tensor, ...] - page_table: torch.Tensor - pool_pointers: torch.Tensor - source_layer_indices: Optional[torch.Tensor] - - def compact( - self, source: torch.Tensor, offsets: torch.Tensor, destination_bases: torch.Tensor - ) -> None: - torch.ops.trtllm.sparse_kv_cache_compact_layers( - list(self.pools), - self.pool_pointers, - self.page_table, - source, - offsets, - destination_bases, - self.source_layer_indices, - ) - - -class _SingleCacheCompaction(NamedTuple): - """One compacted cache family (target dense, target SWA, or draft). - - Holds the launch that packs this family's move indices (None - when an earlier family's pack call fills them in the same run), the - C++ compact groups that consume them, and the destination base the moved - tokens land at. - """ - - move_index_pack: Optional[Callable[[], None]] - cpp_compact_groups: Tuple[_CppCompactGroup, ...] - move_source_indices: torch.Tensor - move_source_offsets: torch.Tensor - # Per-request landing positions; may alias staged prompt lengths so the - # values track the current round without a refresh. - destination_bases: torch.Tensor - - def compact(self) -> None: - if self.move_index_pack is not None: - self.move_index_pack() - for group in self.cpp_compact_groups: - group.compact( - self.move_source_indices, self.move_source_offsets, self.destination_bases - ) - - def _validated_kv_head_count( pools: List[torch.Tensor], layers: Tuple[int, ...], @@ -167,11 +121,13 @@ def _compact_groups( pool_keys: Tuple[object, ...], device: torch.device, per_layer_slots: Optional[Dict[int, int]] = None, -) -> Tuple[_CppCompactGroup, ...]: +) -> Tuple[Dict[str, object], ...]: """Batch layers into one C++ launch per uniform V2 pool. - ``per_layer_slots`` maps each layer to its selection row; it is only set - when every dense layer keeps its own token set (per-layer eviction). + Each returned dict is the plain launch data for one + ``sparse_kv_cache_compact_layers`` call. ``per_layer_slots`` maps each + layer to its selection row; it is only set when every dense layer keeps + its own token set (per-layer eviction). """ grouped = OrderedDict() for layer, pool, page_table in entries: @@ -201,8 +157,8 @@ def _compact_groups( device=device, ) result.append( - _CppCompactGroup( - pools=pools, + dict( + pools=list(pools), page_table=page_tables[0], pool_pointers=torch.tensor( [pool.data_ptr() for pool in pools], @@ -221,38 +177,7 @@ def _compact_groups( _PACK_NUM_WARPS = 4 -class _MoveSourcePackArguments(NamedTuple): - """One cache family's move-index packing, described as plain launch data. - - The selection-side fused settle-and-pack launch consumes this to pack - the dense/SWA move sources in the same kernel that finalizes the kept - ordinals, instead of a second launch at compaction time. - """ - - kept_token_ordinals: torch.Tensor - valid_sequence_lengths: torch.Tensor - dense_offsets: torch.Tensor - dense_indices: torch.Tensor - # With no SWA family these alias the dense tensors and ``has_swa`` - # specializes every SWA load and store away. - swa_offsets: torch.Tensor - swa_indices: torch.Tensor - dense_total: int - swa_total: int - selection_rows: int - keep_count: int - request_count: int - num_kv_heads: int - swa_window: int - # Widest per-request move count any staged offsets may express; the - # packing loop covers exactly this many move slots per packed row. - move_capacity: int - union: bool - per_layer: bool - has_swa: bool - - -def _move_index_pack_launcher( +def build_move_pack_arguments( kept_token_ordinals: torch.Tensor, valid_sequence_lengths: torch.Tensor, move_source_offsets: torch.Tensor, @@ -266,15 +191,13 @@ def _move_index_pack_launcher( swa_window: int, swa_move_source_offsets: Optional[torch.Tensor], swa_move_source_indices: Optional[torch.Tensor], -) -> Tuple[Callable[[], None], _MoveSourcePackArguments]: - """Build one launch of the move-index packing kernel. - - The kernel reads the kept-token ordinals and each request's valid length - and writes the packed per-(layer, head) move source indices consumed by - the C++ compact launches. Only the caller-provided selection tensors are - validated here; the move buffers are allocated by this module. The - returned arguments bundle describes the same packing so a fused - selection-side launch can take it over. +) -> Dict[str, object]: + """Describe one move-index packing as plain kernel launch data. + + The packing kernel reads the kept-token ordinals and each request's valid + length and writes the packed per-(layer, head) move source indices + consumed by the C++ compact launches. ``launch_move_pack`` fires it + standalone; a fused selection-side settle launch consumes the same dict. """ per_layer = eviction_mode == "per_layer_perhead" union = eviction_mode == "union" @@ -311,12 +234,10 @@ def _move_index_pack_launcher( swa_indices_arg = move_source_indices swa_total = 0 - from .triattention_kernels import _settle_ties_and_pack_compaction_sources_kernel - max_move = decode_keep_count + max_protected_tail if swa_total: max_move = max(max_move, swa_window + max_protected_tail) - pack_arguments = _MoveSourcePackArguments( + return dict( kept_token_ordinals=kept_token_ordinals, valid_sequence_lengths=valid_sequence_lengths, dense_offsets=move_source_offsets, @@ -330,314 +251,236 @@ def _move_index_pack_launcher( request_count=request_count, num_kv_heads=num_kv_heads, swa_window=swa_window, + # Widest per-request move count any staged offsets may express; the + # packing loop covers exactly this many move slots per packed row. move_capacity=max_move, union=union, per_layer=per_layer, has_swa=swa_total > 0, ) - # One program per (request, selection row); the settle half is compiled - # away because the ordinals arrive pre-settled (the draft flow reuses - # the target's keep set verbatim), so only the pack half runs. The - # settle-side pointer arguments are compiled away with it; any - # well-formed tensor stands in for them. - grid = (request_count, selection_rows) - - def launch_pack() -> None: - _settle_ties_and_pack_compaction_sources_kernel[grid]( - kept_token_ordinals, - valid_sequence_lengths, - move_source_offsets, - kept_token_ordinals, - kept_token_ordinals, - valid_sequence_lengths, - move_source_offsets, - move_source_indices, - swa_offsets_arg, - swa_indices_arg, - WIDTH=decode_keep_count, - KEEP_COUNT=decode_keep_count, - OUTPUT_WIDTH=decode_keep_count, - SELECTION_ROWS=selection_rows, - DENSE_TOTAL=pack_arguments.dense_total, - SWA_TOTAL=pack_arguments.swa_total, - MOVE_CAPACITY=pack_arguments.move_capacity, - NUM_KV_HEADS=pack_arguments.num_kv_heads, - SWA_WINDOW=pack_arguments.swa_window, - UNION=pack_arguments.union, - PER_LAYER=pack_arguments.per_layer, - HAS_SWA=pack_arguments.has_swa, - HAS_SETTLE=False, - HAS_PACK=True, - BLOCK=_PACK_BLOCK_TOKENS, - num_warps=_PACK_NUM_WARPS, - ) - return launch_pack, pack_arguments + +def launch_move_pack(pack: Dict[str, object]) -> None: + """Fire one standalone move-index packing launch. + + One program per (request, selection row); the settle half is compiled + away because the ordinals arrive pre-settled (the draft flow reuses the + target's keep set verbatim), so only the pack half runs. The settle-side + pointer arguments are compiled away with it; any well-formed tensor + stands in for them. + """ + from .triattention_kernels import _settle_ties_and_pack_compaction_sources_kernel + + kept = pack["kept_token_ordinals"] + _settle_ties_and_pack_compaction_sources_kernel[ + (pack["request_count"], pack["selection_rows"]) + ]( + kept, + pack["valid_sequence_lengths"], + pack["dense_offsets"], + kept, + kept, + pack["valid_sequence_lengths"], + pack["dense_offsets"], + pack["dense_indices"], + pack["swa_offsets"], + pack["swa_indices"], + WIDTH=pack["keep_count"], + KEEP_COUNT=pack["keep_count"], + OUTPUT_WIDTH=pack["keep_count"], + SELECTION_ROWS=pack["selection_rows"], + DENSE_TOTAL=pack["dense_total"], + SWA_TOTAL=pack["swa_total"], + MOVE_CAPACITY=pack["move_capacity"], + NUM_KV_HEADS=pack["num_kv_heads"], + SWA_WINDOW=pack["swa_window"], + UNION=pack["union"], + PER_LAYER=pack["per_layer"], + HAS_SWA=pack["has_swa"], + HAS_SETTLE=False, + HAS_PACK=True, + BLOCK=_PACK_BLOCK_TOKENS, + num_warps=_PACK_NUM_WARPS, + ) -class BatchedKVCacheCompaction: - """Batched physical compaction of the KV caches for one fixed geometry. +def build_cache_compactions( + *, + eviction_mode: str, + layer_pools: List[torch.Tensor], + dense_layers: List[int], + swa_layers: List[int], + layer_group_representative: Dict[int, int], + kept_token_ordinals: torch.Tensor, + valid_sequence_lengths: torch.Tensor, + kv_block_offsets: torch.Tensor, + page_table_slots: Dict[int, int], + request_count: int, + prompt_offsets: torch.Tensor, + decode_keep_count: int, + swa_window: Optional[int], + layer_pool_keys: List[object], + protected_tail_capacity: int = 0, + fuse_dense_pack_into_selection: bool = False, + draft_layer_pools: Optional[List[torch.Tensor]] = None, + draft_layers: Optional[List[int]] = None, + draft_layer_group_representative: Optional[Dict[int, int]] = None, + draft_layer_pool_keys: Optional[List[object]] = None, + draft_protected_tail_capacity: Optional[int] = None, + draft_kv_block_offsets: Optional[torch.Tensor] = None, + draft_page_table_slots: Optional[Dict[int, int]] = None, + dense_move_offsets: Optional[torch.Tensor] = None, + swa_move_offsets: Optional[torch.Tensor] = None, + draft_move_offsets: Optional[torch.Tensor] = None, +) -> Dict[str, object]: + """Allocate the per-geometry compaction launch data as one plain dict. Dense layers keep the prompt in place and compact the selected decode tokens plus any target KV reserved for the next overlapped forward; kernel-masked SWA layers keep the latest window plus the same protected tail. A co-compressed draft cache reuses the target's kept token ordinals - (broadcast over the draft's own KV-head count) plus the draft's own - protected tail, landing at the same destination base. - - Key constructor inputs: - `kept_token_ordinals`: increasing kept decode ordinals (absolute - positions) per request; shape `[requests, keep]` for `union`, - with a selection-row dimension in between for the per-head - modes. Prompt tokens never move, so the rectangle is - prompt-length independent and one cohort may mix prompt sizes; - `prompt_offsets` carries each request's pinned prompt length. - `kv_block_offsets`: the staged V2 block-offset snapshot laid out as - `[slot, request, K/V, block]`, where a block offset encodes - page and K/V plane as `2*page + plane`. - `page_table_slots` / `layer_group_representative`: map each layer's - group representative to its snapshot slot; layers that share a - slot must share one block-offset table. - `protected_tail_capacity`: widest per-request protected tail this - object must support. A protected tail covers KV positions past - the valid length reserved for a forward already in flight; each - round's actual lengths arrive through the per-family move-offset - rows staged with the round metadata and move with the kept - tokens. - `draft_*`: co-compressed draft-cache layout (union mode only); the - draft reuses the target keep set and pins the same prompt. + (broadcast over the draft's own KV-head count, union mode only) plus the + draft's own protected tail, landing at the same destination base. + + ``kept_token_ordinals`` carries increasing kept decode ordinals (absolute + positions) per request; prompt tokens never move, so the rectangle is + prompt-length independent and ``prompt_offsets`` carries each request's + pinned prompt length. ``kv_block_offsets`` is the staged V2 snapshot + ``[slot, request, K/V, block]`` (offset = ``2*page + plane``); + ``protected_tail_capacity`` is the widest per-request tail this geometry + must support -- actual per-round lengths arrive through the staged + move-offset rows. With ``fuse_dense_pack_into_selection`` the target's + dense/SWA packing is left to the caller's fused settle launch (the + returned ``dense_pack`` dict describes it) and only the C++ moves run + here; the draft always keeps its own pack launch. + + Returns ``{"families": [...], "dense_pack": ..., "num_kv_heads": ..., + "swa_window": ..., "swa_destination_bases": ...}`` where each family is + ``{"pack": dict|None, "groups": (...), "source": t, "offsets": t, + "destination_bases": t}``. """ + if eviction_mode not in ("union", "per_head", "per_layer_perhead"): + raise ValueError(f"unsupported compaction mode: {eviction_mode}") + if request_count <= 0 or decode_keep_count <= 0: + raise ValueError("batched compaction requires requests and retained tokens") + if not dense_layers: + raise ValueError("batched compaction requires at least one dense layer") + if draft_layers and eviction_mode != "union": + raise ValueError("draft co-compaction supports only union eviction") + if draft_layer_pools is not None and not draft_layers: + raise ValueError("draft pools were given without any draft layers") + if not swa_layers and swa_window: + raise ValueError("swa_window was given without any SWA layers") + + device = layer_pools[dense_layers[0]].device + # The move buffers are allocated on the pool device, so the selection + # tensors feeding the pack kernel must already live there. + if kept_token_ordinals.device != device: + raise ValueError("kept-token ordinals must live on the pool device") + request_count = int(request_count) + if ( + prompt_offsets.shape != (request_count,) + or prompt_offsets.dtype != torch.int32 + or prompt_offsets.device != device + or not prompt_offsets.is_contiguous() + ): + raise ValueError("per-request prompt offsets do not match the cohort") + decode_keep_count = int(decode_keep_count) + if protected_tail_capacity < 0: + raise ValueError("the protected-tail capacity must be non-negative") + protected_tail_capacity = int(protected_tail_capacity) + dense_layers = tuple(int(layer) for layer in dense_layers) + swa_layers = tuple(int(layer) for layer in swa_layers) + if len(layer_pool_keys) != len(layer_pools): + raise ValueError("pool keys must match the layer-pool count") + layer_pool_keys = tuple(layer_pool_keys) - def __init__( - self, - *, - eviction_mode: str, - layer_pools: List[torch.Tensor], - dense_layers: List[int], - swa_layers: List[int], - layer_group_representative: Dict[int, int], - kept_token_ordinals: torch.Tensor, - valid_sequence_lengths: torch.Tensor, - kv_block_offsets: torch.Tensor, - page_table_slots: Dict[int, int], - request_count: int, - prompt_offsets: torch.Tensor, - decode_keep_count: int, - swa_window: Optional[int], - layer_pool_keys: List[object], - protected_tail_capacity: int = 0, - draft_layer_pools: Optional[List[torch.Tensor]] = None, - draft_layers: Optional[List[int]] = None, - draft_layer_group_representative: Optional[Dict[int, int]] = None, - draft_layer_pool_keys: Optional[List[object]] = None, - draft_protected_tail_capacity: Optional[int] = None, - draft_kv_block_offsets: Optional[torch.Tensor] = None, - draft_page_table_slots: Optional[Dict[int, int]] = None, - dense_move_offsets: Optional[torch.Tensor] = None, - swa_move_offsets: Optional[torch.Tensor] = None, - draft_move_offsets: Optional[torch.Tensor] = None, - ) -> None: - if eviction_mode not in ("union", "per_head", "per_layer_perhead"): - raise ValueError(f"unsupported compaction mode: {eviction_mode}") - if request_count <= 0 or decode_keep_count <= 0: - raise ValueError("batched compaction requires requests and retained tokens") - if not dense_layers: - raise ValueError("batched compaction requires at least one dense layer") - if draft_layers and eviction_mode != "union": - raise ValueError("draft co-compaction supports only union eviction") - if draft_layer_pools is not None and not draft_layers: - raise ValueError("draft pools were given without any draft layers") - if not swa_layers and swa_window: - raise ValueError("swa_window was given without any SWA layers") - - self.eviction_mode = eviction_mode - self.device = layer_pools[dense_layers[0]].device - # The move buffers are allocated on the pool device, so the selection - # tensors feeding the pack kernel must already live there. - if kept_token_ordinals.device != self.device: - raise ValueError("kept-token ordinals must live on the pool device") - self.request_count = int(request_count) - # Per-request pinned prompt lengths; this usually aliases the staged - # prompt buffer, so the values track the current round. Only the - # geometry is validated here. - if ( - prompt_offsets.shape != (self.request_count,) - or prompt_offsets.dtype != torch.int32 - or prompt_offsets.device != self.device - or not prompt_offsets.is_contiguous() - ): - raise ValueError("per-request prompt offsets do not match the cohort") - self.prompt_offsets = prompt_offsets - self.decode_keep_count = int(decode_keep_count) - if protected_tail_capacity < 0: - raise ValueError("the protected-tail capacity must be non-negative") - self.protected_tail_capacity = int(protected_tail_capacity) - self.dense_layers = tuple(int(layer) for layer in dense_layers) - self.swa_layers = tuple(int(layer) for layer in swa_layers) - if len(layer_pool_keys) != len(layer_pools): - raise ValueError("pool keys must match the layer-pool count") - self.layer_pool_keys = tuple(layer_pool_keys) - - per_layer = self.eviction_mode == "per_layer_perhead" - self.num_kv_heads = _validated_kv_head_count( - layer_pools, - (*self.dense_layers, *self.swa_layers), - self.device, - "batched compaction", - ) - dense_index_prefix = ( - (len(self.dense_layers), self.num_kv_heads) if per_layer else (self.num_kv_heads,) - ) - dense_move_indices, dense_move_offsets = _make_move_buffers( - dense_index_prefix, - [self.decode_keep_count + self.protected_tail_capacity] * self.request_count, - self.device, - external_offsets=dense_move_offsets, - ) - page_table_for = _page_table_provider( - page_table_slots, - kv_block_offsets, - self.device, - self.request_count, - "compaction", - ) - dense_entries = [ - (layer, layer_pools[layer], page_table_for(layer_group_representative[layer])) - for layer in self.dense_layers - ] - - self.swa_window = 0 - self.swa_destination_bases = None - swa_move_indices = None - swa_entries = [] - if not self.swa_layers: - # No SWA family: drop the unused offsets row. With SWA layers the - # constructor argument must stay live so the family reads the - # per-round staged offsets instead of its construction-time sizes. - swa_move_offsets = None - if self.swa_layers: - if swa_window is None or swa_window <= 0: - raise ValueError("SWA compaction requires a valid retained window") - # Per-request window validity (prompt + decode keep >= window) is - # prompt-dependent and checked by the caller each round. - self.swa_window = int(swa_window) - self.swa_destination_bases = torch.empty_like(self.prompt_offsets) - swa_move_indices, swa_move_offsets = _make_move_buffers( - (self.num_kv_heads,), - [self.swa_window + self.protected_tail_capacity] * self.request_count, - self.device, - external_offsets=swa_move_offsets, - ) - # SWA layers are staged as their own page-table representatives. - swa_entries = [ - (layer, layer_pools[layer], page_table_for(layer)) for layer in self.swa_layers - ] - - dense_slots = ( - {layer: slot for slot, layer in enumerate(self.dense_layers)} if per_layer else None - ) - dense_pack, self._dense_pack_arguments = _move_index_pack_launcher( - kept_token_ordinals, - valid_sequence_lengths, - dense_move_offsets, - dense_move_indices, - eviction_mode=self.eviction_mode, - decode_keep_count=self.decode_keep_count, - num_dense_layers=len(self.dense_layers), - num_kv_heads=self.num_kv_heads, - max_protected_tail=self.protected_tail_capacity, - swa_window=self.swa_window, - swa_move_source_offsets=swa_move_offsets, - swa_move_source_indices=swa_move_indices, + per_layer = eviction_mode == "per_layer_perhead" + num_kv_heads = _validated_kv_head_count( + layer_pools, (*dense_layers, *swa_layers), device, "batched compaction" + ) + dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer else (num_kv_heads,) + dense_move_indices, dense_move_offsets = _make_move_buffers( + dense_index_prefix, + [decode_keep_count + protected_tail_capacity] * request_count, + device, + external_offsets=dense_move_offsets, + ) + page_table_for = _page_table_provider( + page_table_slots, kv_block_offsets, device, request_count, "compaction" + ) + dense_entries = [ + (layer, layer_pools[layer], page_table_for(layer_group_representative[layer])) + for layer in dense_layers + ] + + swa_destination_bases = None + swa_move_indices = None + swa_entries = [] + if not swa_layers: + # No SWA family: drop the unused offsets row. With SWA layers the + # argument must stay live so the family reads the per-round staged + # offsets instead of its construction-time sizes. + swa_move_offsets = None + swa_window = 0 + else: + if swa_window is None or swa_window <= 0: + raise ValueError("SWA compaction requires a valid retained window") + # Per-request window validity (prompt + decode keep >= window) is + # prompt-dependent and checked by the caller each round. + swa_window = int(swa_window) + swa_destination_bases = torch.empty_like(prompt_offsets) + swa_move_indices, swa_move_offsets = _make_move_buffers( + (num_kv_heads,), + [swa_window + protected_tail_capacity] * request_count, + device, + external_offsets=swa_move_offsets, ) - self.target_dense_compaction = _SingleCacheCompaction( - move_index_pack=dense_pack, - cpp_compact_groups=_compact_groups( - dense_entries, self.layer_pool_keys, self.device, dense_slots - ), - move_source_indices=dense_move_indices, - move_source_offsets=dense_move_offsets, - destination_bases=self.prompt_offsets, + # SWA layers are staged as their own page-table representatives. + swa_entries = [(layer, layer_pools[layer], page_table_for(layer)) for layer in swa_layers] + + dense_slots = {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer else None + dense_pack = build_move_pack_arguments( + kept_token_ordinals, + valid_sequence_lengths, + dense_move_offsets, + dense_move_indices, + eviction_mode=eviction_mode, + decode_keep_count=decode_keep_count, + num_dense_layers=len(dense_layers), + num_kv_heads=num_kv_heads, + max_protected_tail=protected_tail_capacity, + swa_window=swa_window, + swa_move_source_offsets=swa_move_offsets, + swa_move_source_indices=swa_move_indices, + ) + families = [ + dict( + name="dense", + # A fused selection-side settle launch packs the dense/SWA move + # sources when it finalizes the kept ordinals; only the C++ moves + # stay here. Each round then packs exactly once. + pack=None if fuse_dense_pack_into_selection else dense_pack, + groups=_compact_groups(dense_entries, layer_pool_keys, device, dense_slots), + source=dense_move_indices, + offsets=dense_move_offsets, + destination_bases=prompt_offsets, ) + ] + if swa_layers: # The dense pack call fills the SWA move buffers in the same run. - self.target_swa_compaction = None - if self.swa_layers: - self.target_swa_compaction = _SingleCacheCompaction( - move_index_pack=None, - cpp_compact_groups=_compact_groups(swa_entries, self.layer_pool_keys, self.device), - move_source_indices=swa_move_indices, - move_source_offsets=swa_move_offsets, - destination_bases=self.swa_destination_bases, + families.append( + dict( + name="swa", + pack=None, + groups=_compact_groups(swa_entries, layer_pool_keys, device), + source=swa_move_indices, + offsets=swa_move_offsets, + destination_bases=swa_destination_bases, ) - - self.draft_compaction = None - self.draft_protected_tail_capacity = 0 - if draft_layers: - self.draft_compaction = self._build_draft_compaction( - kept_token_ordinals, - valid_sequence_lengths, - draft_layer_pools=draft_layer_pools, - draft_layers=draft_layers, - draft_layer_group_representative=draft_layer_group_representative, - draft_layer_pool_keys=draft_layer_pool_keys, - draft_protected_tail_capacity=draft_protected_tail_capacity, - draft_kv_block_offsets=draft_kv_block_offsets, - draft_page_table_slots=draft_page_table_slots, - draft_move_offsets=draft_move_offsets, - ) - - self.cache_compactions = tuple( - compaction - for compaction in ( - self.target_dense_compaction, - self.target_swa_compaction, - self.draft_compaction, - ) - if compaction is not None ) - def hand_move_source_pack_to_selection(self) -> _MoveSourcePackArguments: - """Hand the dense/SWA move packing over to the selection launch. - - Returns the packing description and drops this object's own dense - pack launch, so each round packs exactly once: the caller's fused - settle-and-pack kernel fills the move buffers when it finalizes the - kept ordinals, and ``compact`` then only runs the C++ moves. The - co-compressed draft keeps its own pack launch because it broadcasts - the finalized keep set over the draft's own KV-head layout. - """ - self.target_dense_compaction = self.target_dense_compaction._replace(move_index_pack=None) - self.cache_compactions = tuple( - compaction - for compaction in ( - self.target_dense_compaction, - self.target_swa_compaction, - self.draft_compaction, - ) - if compaction is not None - ) - return self._dense_pack_arguments - - def _build_draft_compaction( - self, - kept_token_ordinals: torch.Tensor, - valid_sequence_lengths: torch.Tensor, - *, - draft_layer_pools: Optional[List[torch.Tensor]], - draft_layers: List[int], - draft_layer_group_representative: Optional[Dict[int, int]], - draft_layer_pool_keys: Optional[List[object]], - draft_protected_tail_capacity: Optional[int], - draft_kv_block_offsets: Optional[torch.Tensor], - draft_page_table_slots: Optional[Dict[int, int]], - draft_move_offsets: Optional[torch.Tensor] = None, - ) -> _SingleCacheCompaction: - """Build the co-compressed draft cache's own pack and launch groups. - - The draft forms its own launch groups so it may use a different - KV-head count than the target. Union-only eviction is enforced by - the constructor before dense groups are built. - """ + if draft_layers: if ( draft_layer_pools is None or draft_layer_group_representative is None @@ -648,28 +491,23 @@ def _build_draft_compaction( raise ValueError("draft co-compaction requires staged draft page tables") if draft_protected_tail_capacity is not None and draft_protected_tail_capacity < 0: raise ValueError("the draft protected-tail capacity must be non-negative") - self.draft_protected_tail_capacity = int(draft_protected_tail_capacity or 0) + draft_tail = int(draft_protected_tail_capacity or 0) if len(draft_layer_pool_keys) != len(draft_layer_pools): raise ValueError("draft pool keys must match the draft layer-pool count") draft_layers = tuple(int(layer) for layer in draft_layers) + # The draft forms its own launch groups so it may use a different + # KV-head count than the target. draft_num_kv_heads = _validated_kv_head_count( - draft_layer_pools, - draft_layers, - self.device, - "draft co-compaction", + draft_layer_pools, draft_layers, device, "draft co-compaction" ) draft_move_indices, draft_move_offsets = _make_move_buffers( (draft_num_kv_heads,), - [self.decode_keep_count + self.draft_protected_tail_capacity] * self.request_count, - self.device, + [decode_keep_count + draft_tail] * request_count, + device, external_offsets=draft_move_offsets, ) draft_page_table_for = _page_table_provider( - draft_page_table_slots, - draft_kv_block_offsets, - self.device, - self.request_count, - "draft", + draft_page_table_slots, draft_kv_block_offsets, device, request_count, "draft" ) draft_entries = [ ( @@ -679,43 +517,71 @@ def _build_draft_compaction( ) for layer in draft_layers ] - # In union mode the pack kernel reads selection row 0 for every - # packed row, so one more pack launch broadcasts the target keep - # set over the draft KV heads and appends the draft's own tail - # ordinals (valid_seq_len + 0..tail-1). - draft_pack, _ = _move_index_pack_launcher( + # In union mode the pack kernel reads selection row 0 for every packed + # row, so one more pack launch broadcasts the target keep set over the + # draft KV heads and appends the draft's own tail ordinals. + draft_pack = build_move_pack_arguments( kept_token_ordinals, valid_sequence_lengths, draft_move_offsets, draft_move_indices, eviction_mode="union", - decode_keep_count=self.decode_keep_count, + decode_keep_count=decode_keep_count, num_dense_layers=1, num_kv_heads=draft_num_kv_heads, - max_protected_tail=self.draft_protected_tail_capacity, + max_protected_tail=draft_tail, swa_window=0, swa_move_source_offsets=None, swa_move_source_indices=None, ) - return _SingleCacheCompaction( - move_index_pack=draft_pack, - cpp_compact_groups=_compact_groups( - draft_entries, tuple(draft_layer_pool_keys), self.device - ), - move_source_indices=draft_move_indices, - move_source_offsets=draft_move_offsets, - destination_bases=self.prompt_offsets, + families.append( + dict( + name="draft", + pack=draft_pack, + groups=_compact_groups(draft_entries, tuple(draft_layer_pool_keys), device), + source=draft_move_indices, + offsets=draft_move_offsets, + destination_bases=prompt_offsets, + ) ) - def compact(self) -> None: - """Pack the move indices, then run every cache family's C++ compacts.""" - if self.swa_destination_bases is not None: - # The prompt offsets may have been re-staged since construction; - # rebase the SWA landing positions for this round. - torch.add( - self.prompt_offsets, - self.decode_keep_count - self.swa_window, - out=self.swa_destination_bases, + return dict( + families=families, + dense_pack=dense_pack, + num_kv_heads=num_kv_heads, + swa_window=swa_window, + swa_destination_bases=swa_destination_bases, + prompt_offsets=prompt_offsets, + decode_keep_count=decode_keep_count, + request_count=request_count, + protected_tail_capacity=protected_tail_capacity, + draft_protected_tail_capacity=( + int(draft_protected_tail_capacity or 0) if draft_layers else 0 + ), + ) + + +def run_cache_compactions(compaction: Dict[str, object]) -> None: + """Pack the move indices, then run every cache family's C++ compacts.""" + swa_destination_bases = compaction["swa_destination_bases"] + if swa_destination_bases is not None: + # The prompt offsets may have been re-staged since construction; + # rebase the SWA landing positions for this round. + torch.add( + compaction["prompt_offsets"], + compaction["decode_keep_count"] - compaction["swa_window"], + out=swa_destination_bases, + ) + for family in compaction["families"]: + if family["pack"] is not None: + launch_move_pack(family["pack"]) + for group in family["groups"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers( + group["pools"], + group["pool_pointers"], + group["page_table"], + family["source"], + family["offsets"], + family["destination_bases"], + group["source_layer_indices"], ) - for compaction in self.cache_compactions: - compaction.compact() diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b71e1d69c7ef..72787ea0778b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -35,6 +35,13 @@ in the same round with the target's kept token set (union mode only), so target and draft always share one physical KV length. +Structure: ``prepare_eviction_workspace`` is the ONE one-time constructor for +the whole eviction stack -- it validates the geometry, allocates every buffer, +compiles the mode-needed CuTe entries eagerly, and builds the C++ compaction +launch data. The result is a plain namespace of tensors, events, and ints. +``stage_eviction_cohort`` and ``run_eviction_round`` are the per-round flow: +straight-line module functions that feed the kernels directly. + KV layout: the decode kernel stores keys in HND layout ``[num_pages, kv_factor, num_kv_heads, tokens_per_block, head_dim]``. The Python gather / score / compact code MUST read ``get_buffers`` with ``kv_layout="HND"``; @@ -53,8 +60,8 @@ The scoring math follows the same upstream reference (``methods/pruning_utils.py``). """ -from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple +from types import SimpleNamespace +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple import torch @@ -68,12 +75,19 @@ from tensorrt_llm.logger import logger from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig +from .compaction import build_cache_compactions, run_cache_compactions +from .triattention_kernels import ( + _settle_ties_and_pack_compaction_sources_kernel, + build_mean_phase_table, + gather_mean_phases, + grow_mean_phase_table, + prepare_per_head_scores, +) + if TYPE_CHECKING: from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests - from .triattention_kernels import MeanPhaseTable - # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) @@ -82,6 +96,10 @@ # caller ever tuned it, so it is a constant rather than a constructor knob. _OFFSET_MAX_LENGTH = 65536 +# Stream-affinity contract: the staged buffers and compiled launches are bound +# to the first CUDA stream that uses them. +_STREAM_MISMATCH = "TriAttention eviction launches must stay on the staging CUDA stream" + def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Tensor: """Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <=max].""" @@ -95,897 +113,870 @@ def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Ten return torch.tensor(offsets, device=device, dtype=torch.float32) -class _FixedScoreStreamMismatch(RuntimeError): - """Raised when fixed score staging buffers are used from another CUDA stream.""" - - -class _RuntimeKVLayout(NamedTuple): - """Manager-lifetime layer and pool views used by every eviction.""" - - manager: object - num_layers: int - global_layers: List[int] - layer_pools: List[torch.Tensor] - dense_layers: List[int] - swa_layers: List[int] - swa_window: Optional[int] - storage_groups: Dict[object, List[int]] - layer_group_representative: Dict[int, int] - layer_pool_keys: Tuple[object, ...] - pool_representatives: Tuple[int, ...] - pool_page_counts: Tuple[int, ...] - pool_view_fingerprint: Tuple[tuple, ...] - - -class _BatchedKeepSetSelector: - """Fixed ``[request, ...]`` keep-set selection buffers for every eviction mode. - - Union mode: the fused score+stats+union CuTe pipeline is THE union score - producer -- it writes the normalized per-request union rows straight into - ``combined``, so the selector owns only the top-k settle-and-pack stage. - The per-head modes (``per_head``, ``per_layer_perhead``) own the per-head - score preparation buffers and settle one selection row per head or per - (layer, head). +def _page_table_slot_layout( + page_representatives: List[int], + page_table_keys: List[object], +) -> Tuple[Dict[int, int], int]: + """Map representative layers to page-table snapshot slots.""" + if len(page_table_keys) != len(page_representatives): + raise ValueError("page-table keys must match the representative count") + use_pool_ids = all( + isinstance(key, tuple) + and len(key) == 2 + and key[0] == "pool" + and isinstance(key[1], int) + and key[1] >= 0 + for key in page_table_keys + ) + unique_slots = [] + key_to_slot = {} + representative_slots = {} + for representative, key in zip(page_representatives, page_table_keys): + slot = key_to_slot.get(key) + if slot is None: + slot = int(key[1]) if use_pool_ids else len(key_to_slot) + key_to_slot[key] = slot + unique_slots.append(slot) + representative_slots[representative] = slot + slot_count = max(unique_slots, default=-1) + 1 + return representative_slots, slot_count + + +def _bind_workspace_stream(ws: SimpleNamespace) -> torch.cuda.Stream: + """Bind the workspace to the current stream on first use, then enforce it.""" + stream = torch.cuda.current_stream(ws.device) + if ws.stream is None: + ws.stream = stream + elif (stream.device, stream.cuda_stream) != (ws.stream.device, ws.stream.cuda_stream): + raise RuntimeError(_STREAM_MISMATCH) + return stream + + +def _allocate_page_table_plane( + layer_pools: List[torch.Tensor], + page_representatives: List[int], + page_table_keys: List[object], + num_page_table_slots: Optional[int], + token_capacity: int, + max_requests: int, + device: torch.device, + what: str, +) -> Tuple[Dict[int, int], int, torch.Tensor, torch.Tensor]: + """Allocate one staged block-offset plane (host pinned + device).""" + representative_slots, minimum_slots = _page_table_slot_layout( + page_representatives, page_table_keys + ) + if num_page_table_slots is None: + num_page_table_slots = minimum_slots + if num_page_table_slots < minimum_slots: + raise ValueError(f"{what}page-table slot capacity does not cover every V2 pool") + tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) + if int(layer_pools[page_representatives[0]].shape[1]) != 2: + raise ValueError(f"{what}page-table staging requires an interleaved K/V pool") + page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block + if any( + (token_capacity + int(layer_pools[layer].shape[3]) - 1) // int(layer_pools[layer].shape[3]) + != page_count + for layer in page_representatives + ): + raise ValueError(f"{what}page-table staging requires a uniform page count") + copy_block_count = (page_count + 3) // 4 * 4 + plane_shape = (num_page_table_slots, max_requests, 2, copy_block_count) + host = torch.empty(plane_shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) + dev = torch.empty(plane_shape, dtype=torch.int32, device=device) + return representative_slots, copy_block_count, host, dev + + +def prepare_eviction_workspace( + *, + eviction_mode: str, + layer_pools: List[torch.Tensor], + dense_groups: List[List[int]], + dense_layers: List[int], + swa_layers: Sequence[int] = (), + swa_window: Optional[int] = None, + layer_group_representative: Optional[Dict[int, int]] = None, + layer_pool_keys: Optional[List[object]] = None, + page_representatives: List[int], + max_requests: int, + seq_len: int, + num_q_heads: int, + num_freqs: int, + keep_count: int, + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + freq_scale_sq: torch.Tensor, + offsets: torch.Tensor, + omega: torch.Tensor, + phase: Optional[Dict[str, object]] = None, + page_table_keys: Optional[List[object]] = None, + num_page_table_slots: Optional[int] = None, + decode_width: Optional[int] = None, + page_table_token_capacity: Optional[int] = None, + protected_tail_capacity: int = 0, + build_compaction: bool = True, + draft_layer_pools: Optional[List[torch.Tensor]] = None, + draft_layers: Optional[List[int]] = None, + draft_layer_group_representative: Optional[Dict[int, int]] = None, + draft_layer_pool_keys: Optional[List[object]] = None, + draft_page_representatives: Optional[List[int]] = None, + draft_page_table_keys: Optional[List[object]] = None, + draft_num_page_table_slots: Optional[int] = None, + draft_page_table_token_capacity: Optional[int] = None, + draft_protected_tail_capacity: int = 0, +) -> SimpleNamespace: + """Build the ONE plain-namespace workspace for the whole eviction stack. + + This is the single one-time constructor: geometry validation, buffer + staging, eager CuTe compilation for exactly the entries the eviction mode + launches, selection buffers, and the C++ compaction launch data. It runs + outside CUDA graph capture (compilation allocates and synchronizes) and + raises loudly on any unsupported geometry -- there is deliberately no + fallback path. The returned namespace holds tensors, events, streams, + compiled runners, and ints; all flow logic lives in + ``stage_eviction_cohort`` and ``run_eviction_round``. + + The workspace retains references to every scored layer pool: the SM100 + CuTe score kernel encodes immutable TMA descriptors from their raw device + addresses at compile time, so the pools must stay alive and stay put for + as long as the workspace launches (the V2 manager owns them for its own + lifetime in production). """ - - def __init__( - self, - *, - eviction_mode: str, - width: int, - keep_count: int, - dtype: torch.dtype, - device: torch.device, - max_requests: int, - dense_layers: Tuple[int, ...] = (), - num_query_heads: int = 0, - num_kv_heads: int = 0, - rows: Optional[int] = None, - prompt_offsets_buffer: Optional[torch.Tensor] = None, - ) -> None: - if eviction_mode == "union": - if rows is None or rows <= 0: - raise ValueError("cross-request selection requires rows > 0") - selection_rows = 1 - elif eviction_mode in ("per_head", "per_layer_perhead"): - if not dense_layers or min(num_query_heads, num_kv_heads, max_requests) <= 0: - raise ValueError( - "per-head selection requires positive layer, head, and request counts" - ) - if num_query_heads % num_kv_heads: - raise ValueError("query heads must be divisible by KV heads") - selection_rows = ( - num_kv_heads if eviction_mode == "per_head" else len(dense_layers) * num_kv_heads - ) - else: - raise ValueError(f"unsupported per-head eviction mode: {eviction_mode}") - if width <= keep_count or keep_count <= 0: - raise ValueError("keep-set selection requires width > keep_count > 0") - if max_requests <= 0: - raise ValueError("keep-set selection requires a positive request capacity") - self.eviction_mode = eviction_mode - self.dense_layers = tuple(int(layer) for layer in dense_layers) - self.num_query_heads = int(num_query_heads) - self.num_kv_heads = int(num_kv_heads) - self.width = int(width) - self.keep_count = int(keep_count) - self.device = device - self.max_requests = int(max_requests) - self.valid_widths = torch.full( - (self.max_requests,), self.width, dtype=torch.int32, device=self.device - ) - # Per-request pinned prompt lengths, refreshed each round: scores are - # decode-relative and these offsets rebase emitted ordinals, so one - # cohort may mix prompt lengths. ``row_prompt_offsets`` is the - # row-major expansion consumed by the finalizer. - self.selection_rows_per_request = int(selection_rows) - # Optional compaction move packing fused into the settle launch; set - # once the compaction buffers exist (see ``fuse_move_source_pack``). - self._move_source_pack = None - if prompt_offsets_buffer is not None: - # Share the staging buffers' per-request prompt lengths so the - # values are written once per round. - if ( - prompt_offsets_buffer.shape != (self.max_requests,) - or prompt_offsets_buffer.dtype != torch.int32 - or prompt_offsets_buffer.device != self.device - ): - raise ValueError("prompt offsets buffer does not match the selector geometry") - self.prompt_offsets = prompt_offsets_buffer - else: - self.prompt_offsets = torch.zeros( - (self.max_requests,), dtype=torch.int32, device=self.device - ) - if self.selection_rows_per_request == 1: - self.row_prompt_offsets = self.prompt_offsets - else: - self.row_prompt_offsets = torch.zeros( - (self.max_requests * self.selection_rows_per_request,), - dtype=torch.int32, - device=self.device, - ) - if eviction_mode == "union": - self.combined = torch.empty((max_requests, width), dtype=dtype, device=self.device) - self.final_indices = torch.empty( - (max_requests, keep_count), dtype=torch.int32, device=self.device - ) - # Kept decode ordinals only: rows are prompt-length independent, so - # one selector serves cohorts with mixed prompt lengths. - self.keep = torch.empty( - (max_requests, self.keep_count), dtype=torch.int32, device=self.device - ) - # Row-major views consumed by the top-k settle launch. - self._selection_scores_rows = self.combined - self._selection_row_lengths = self.valid_widths - self._provisional_rows = self.final_indices - self._keep_rows = self.keep - else: - self.num_layers = len(self.dense_layers) - self.selection_rows = selection_rows - score_shape = (self.max_requests, self.num_layers, self.num_query_heads, self.width) - self.row_mean = torch.empty(score_shape[:-1] + (1,), dtype=dtype, device=self.device) - self.row_std = torch.empty_like(self.row_mean) - self.selection_scores = torch.empty( - (self.max_requests, self.selection_rows, self.width), - dtype=dtype, - device=self.device, - ) - self.row_seq_lens = torch.full( - (self.max_requests, self.selection_rows), - self.width, - dtype=torch.int32, - device=self.device, - ) - selection_shape = (self.max_requests, self.selection_rows, self.keep_count) - self.top_indices_i32 = torch.empty( - selection_shape, dtype=torch.int32, device=self.device - ) - # Kept decode ordinals only: rows are prompt-length independent, so - # one selector serves cohorts with mixed prompt lengths. - self.keep = torch.empty(selection_shape, dtype=torch.int32, device=self.device) - self.selection_scores_flat = self.selection_scores.view( - self.max_requests * self.selection_rows, self.width - ) - self.row_seq_lens_flat = self.row_seq_lens.view(-1) - self.top_indices_i32_flat = self.top_indices_i32.view(-1, self.keep_count) - self.keep_flat = self.keep.view(-1, self.keep_count) - # Row-major views consumed by the top-k settle launch. - self._selection_scores_rows = self.selection_scores_flat - self._selection_row_lengths = self.row_seq_lens_flat - self._provisional_rows = self.top_indices_i32_flat - self._keep_rows = self.keep_flat - - def refresh_row_prompt_offsets(self) -> None: - """Re-expand the per-request prompt offsets into their row-major view. - - Called after the shared per-request buffer was staged externally. - """ - if self.row_prompt_offsets is not self.prompt_offsets: - self.row_prompt_offsets.view(self.max_requests, self.selection_rows_per_request).copy_( - self.prompt_offsets.unsqueeze(1).expand(-1, self.selection_rows_per_request) - ) - - def fuse_move_source_pack(self, pack_arguments) -> None: - """Pack compaction move sources inside this selector's settle launch. - - ``pack_arguments`` is the dense/SWA packing description exported by - ``BatchedKVCacheCompaction.hand_move_source_pack_to_selection``. The - fused kernel reads back the kept ordinals it just wrote, so the - packing must read this selector's own keep buffer, and the packing - geometry must match the selection rows this selector settles. - """ - if ( - pack_arguments.kept_token_ordinals.data_ptr() != self._keep_rows.data_ptr() - or pack_arguments.kept_token_ordinals.numel() != self._keep_rows.numel() - ): - raise ValueError("fused move packing must read this selector's keep buffer") + from .triattention_cute_score_fused import TriAttentionCuteScoreRunner + + if eviction_mode not in ("union", "per_head", "per_layer_perhead"): + raise ValueError(f"unsupported eviction mode: {eviction_mode}") + if not dense_groups or not dense_layers or not page_representatives or max_requests <= 0: + raise ValueError("fixed score metadata requires non-empty positive geometry") + grouped_layers = [layer for layers in dense_groups for layer in layers] + if ( + len(grouped_layers) != len(dense_layers) + or len(set(grouped_layers)) != len(grouped_layers) + or len(set(dense_layers)) != len(dense_layers) + or set(dense_layers) != set(grouped_layers) + ): + raise ValueError("dense layer order must cover every grouped layer exactly once") + device = layer_pools[page_representatives[0]].device + if device.type != "cuda": + raise ValueError("fixed score metadata is CUDA-only") + max_requests = int(max_requests) + seq_len = int(seq_len) + if page_table_token_capacity is None: + page_table_token_capacity = seq_len + if page_table_token_capacity < seq_len: + raise ValueError("page-table capacity cannot be smaller than the score bucket") + page_table_token_capacity = int(page_table_token_capacity) + # Decode-width capacity of the score buffers; per-request prompt lengths + # are staged runtime metadata. + if decode_width is None: + decode_width = seq_len + if decode_width <= 0 or decode_width > seq_len: + raise ValueError("fixed score decode width exceeds the sequence capacity") + decode_width = int(decode_width) + keep_count = int(keep_count) + if decode_width <= keep_count or keep_count <= 0: + raise ValueError("keep-set selection requires width > keep_count > 0") + + q_real = q_real.to(device=device, dtype=torch.float32).contiguous() + q_imag = q_imag.to(device=device, dtype=torch.float32).contiguous() + mlr_coef = mlr_coef.to(device=device, dtype=torch.float32).contiguous() + freq_scale_sq = freq_scale_sq.to(device=device, dtype=torch.float32).contiguous() + offsets = offsets.to(device=device, dtype=torch.float32).contiguous() + omega = omega.to(device=device, dtype=torch.float32).contiguous() + if page_table_keys is None: + page_table_keys = list(range(len(page_representatives))) + + ws = SimpleNamespace() + ws.eviction_mode = eviction_mode + ws.device = device + ws.max_requests = max_requests + ws.bucket_seq_len = seq_len + ws.decode_width = decode_width + ws.keep_count = keep_count + ws.page_table_token_capacity = page_table_token_capacity + + # ---- staged page-table planes (target, plus the co-compressed draft) --- + ( + ws.representative_slots, + ws.copy_block_count, + ws._bulk_offsets_src, + ws.block_offsets_device, + ) = _allocate_page_table_plane( + layer_pools, + page_representatives, + page_table_keys, + num_page_table_slots, + page_table_token_capacity, + max_requests, + device, + "", + ) + # The draft is never scored: these offsets feed only the draft compacts. + ws.draft_block_offsets_device = None + ws._draft_bulk_offsets_src = None + ws.draft_representative_slots = {} + ws.draft_copy_block_count = 0 + if draft_layer_pools is not None: if ( - pack_arguments.selection_rows != self.selection_rows_per_request - or pack_arguments.keep_count != self.keep_count - or pack_arguments.request_count * self.selection_rows_per_request - != int(self._keep_rows.shape[0]) + not draft_page_representatives + or draft_page_table_keys is None + or draft_page_table_token_capacity is None + or draft_page_table_token_capacity <= 0 ): - raise ValueError("fused move packing does not match the selector geometry") - self._move_source_pack = pack_arguments - - def _select_top_tokens(self) -> None: - """Pick the top-k with the CuTE selector, then settle its output. - - The CuTE top-k is fast but breaks score ties arbitrarily and emits - indices in arbitrary order; the settle kernel recomputes the threshold - membership with lowest-index-wins ties, rebases each row by its prompt - offset, and writes sorted ordinals. When a compaction move packing is - fused in, the same launch also packs each request's dense/SWA move - source indices from the ordinals it just settled. - """ - from .triattention_kernels import _settle_ties_and_pack_compaction_sources_kernel - - rows = int(self._selection_scores_rows.shape[0]) - # 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, - self._selection_row_lengths, - self._provisional_rows, - self.keep_count, - 1, - ) - pack = self._move_source_pack - if pack is None: - # Settle only: the pack half is compiled away, so its tensor - # parameters are never read; any resident tensor stands in. - placeholder = self._selection_row_lengths - pack_tensors = (placeholder,) * 5 - pack_shape = dict( - DENSE_TOTAL=0, - SWA_TOTAL=0, - MOVE_CAPACITY=0, - NUM_KV_HEADS=1, - SWA_WINDOW=0, - UNION=False, - PER_LAYER=False, - HAS_SWA=False, - HAS_PACK=False, - ) - else: - pack_tensors = ( - pack.valid_sequence_lengths, - pack.dense_offsets, - pack.dense_indices, - pack.swa_offsets, - pack.swa_indices, - ) - pack_shape = dict( - DENSE_TOTAL=pack.dense_total, - SWA_TOTAL=pack.swa_total, - MOVE_CAPACITY=pack.move_capacity, - NUM_KV_HEADS=pack.num_kv_heads, - SWA_WINDOW=pack.swa_window, - UNION=pack.union, - PER_LAYER=pack.per_layer, - HAS_SWA=pack.has_swa, - HAS_PACK=True, + raise ValueError( + "draft page-table staging requires representatives, keys, and capacity" ) - _settle_ties_and_pack_compaction_sources_kernel[ - (rows // self.selection_rows_per_request, self.selection_rows_per_request) - ]( - self._selection_scores_rows, - self._selection_row_lengths, - self.row_prompt_offsets, - self._provisional_rows, - self._keep_rows, - *pack_tensors, - WIDTH=self.width, - KEEP_COUNT=self.keep_count, - OUTPUT_WIDTH=self.keep_count, - SELECTION_ROWS=self.selection_rows_per_request, - **pack_shape, - HAS_SETTLE=True, - BLOCK=256, - num_warps=4, - ) - - def select_prepared_union_scores(self) -> None: - """Select from normalized union rows already written into ``combined``. - - The fused score+stats+union pipeline produces the per-request union - rows on the device, so only the top-k settle-and-pack launch remains. - """ - self._select_top_tokens() - - def select_requests( - self, - scores: torch.Tensor, - *, - normalize_scores: bool, - ) -> None: - from .triattention_kernels import prepare_per_head_scores - - expected_shape = ( - self.max_requests, - self.num_layers, - self.num_query_heads, - self.width, - ) - if tuple(scores.shape) != expected_shape or not scores.is_contiguous(): - raise ValueError("per-head scores do not match the selector geometry") - prepare_per_head_scores( - scores, - self.valid_widths, - self.row_mean, - self.row_std, - self.selection_scores, - self.row_seq_lens, - self.max_requests, - num_kv_heads=self.num_kv_heads, - per_layer=self.eviction_mode == "per_layer_perhead", - normalize_scores=normalize_scores, - ) - self._select_top_tokens() - - -class _FixedScoreStagingBuffers: - """Pool-bound fixed score metadata with one nonblocking page-table upload.""" - - @staticmethod - def _page_table_slot_layout( - page_representatives: List[int], - page_table_keys: List[object], - ) -> Tuple[Dict[int, int], int]: - if len(page_table_keys) != len(page_representatives): - raise ValueError("page-table keys must match the representative count") - use_pool_ids = all( - isinstance(key, tuple) - and len(key) == 2 - and key[0] == "pool" - and isinstance(key[1], int) - and key[1] >= 0 - for key in page_table_keys + ( + ws.draft_representative_slots, + ws.draft_copy_block_count, + ws._draft_bulk_offsets_src, + ws.draft_block_offsets_device, + ) = _allocate_page_table_plane( + draft_layer_pools, + draft_page_representatives, + draft_page_table_keys, + draft_num_page_table_slots, + int(draft_page_table_token_capacity), + max_requests, + device, + "draft ", ) - unique_slots = [] - key_to_slot = {} - representative_slots = {} - for representative, key in zip(page_representatives, page_table_keys): - slot = key_to_slot.get(key) - if slot is None: - slot = int(key[1]) if use_pool_ids else len(key_to_slot) - key_to_slot[key] = slot - unique_slots.append(slot) - representative_slots[representative] = slot - slot_count = max(unique_slots, default=-1) + 1 - return representative_slots, slot_count - - def __init__( - self, - layer_pools: List[torch.Tensor], - dense_groups: List[List[int]], - dense_layers: List[int], - page_representatives: List[int], - max_requests: int, - seq_len: int, - num_q_heads: int, - num_freqs: int, - q_real: torch.Tensor, - q_imag: torch.Tensor, - mlr_coef: torch.Tensor, - freq_scale_sq: torch.Tensor, - offsets: torch.Tensor, - omega: torch.Tensor, - mean_phase_table: Optional["MeanPhaseTable"] = None, - page_table_keys: Optional[List[object]] = None, - num_page_table_slots: Optional[int] = None, - decode_width: Optional[int] = None, - page_table_token_capacity: Optional[int] = None, - draft_layer_pools: Optional[List[torch.Tensor]] = None, - draft_page_representatives: Optional[List[int]] = None, - draft_page_table_keys: Optional[List[object]] = None, - draft_num_page_table_slots: Optional[int] = None, - draft_page_table_token_capacity: Optional[int] = None, - ) -> None: - from .triattention_kernels import MeanPhaseTable, _FixedScoreGroup - if not dense_groups or not dense_layers or not page_representatives or max_requests <= 0: - raise ValueError("fixed score metadata requires non-empty positive geometry") - grouped_layers = [layer for layers in dense_groups for layer in layers] + # ---- per-round metadata table: ONE host-to-device copy per round ------- + # Three metadata rows (logical position, valid length, prompt length) plus + # one move-offsets row per compacted cache family; offsets rows have + # request_capacity + 1 entries, hence the extra column. + ws.request_metadata_host = torch.empty( + (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + ) + ws._bulk_copy_idx_src = torch.arange( + max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + ) + # Zero-filled so an unstaged cohort gathers the phase table's row 0 + # instead of indexing it with uninitialized round starts. + ws.request_metadata_device = torch.zeros( + (6, max_requests + 1), dtype=torch.int32, device=device + ) + ws.round_starts_device = ws.request_metadata_device[0, :max_requests] + ws.valid_seq_lens_device = ws.request_metadata_device[1, :max_requests] + # Per-request pinned prompt lengths: the score kernel starts each + # request's decode window here, so one bucket may mix prompt lengths. + ws.token_starts_device = ws.request_metadata_device[2, :max_requests] + ws.dense_move_offsets = ws.request_metadata_device[3] + ws.swa_move_offsets = ws.request_metadata_device[4] + ws.draft_move_offsets = ws.request_metadata_device[5] + ws.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) + ws.mean_sin = torch.empty_like(ws.mean_cos) + # The phase table depends only on the shared calibration; the manager + # shares one dict with every workspace (tests may pass None for a + # private table). + if phase is None: + phase = build_mean_phase_table(offsets, omega, initial_rows=seq_len) + ws.phase = phase + + # ---- score state: ONE fused group across ALL dense layers -------------- + # Segments carry their own page-table slot, so distinct per-layer + # storages/block tables share a single launch. + p0 = layer_pools[dense_layers[0]] + if p0.ndim != 5: + raise ValueError("fixed score group requires HND pools") + _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape + if num_q_heads % num_kv_heads: + raise ValueError("query heads must be divisible by KV heads") + if int(num_freqs) != head_dim // 2: + raise ValueError("calibration frequency count must match half the head dim") + strides = tuple(int(value) for value in p0.stride()) + element_size = p0.element_size() + for layer in dense_layers: + pool = layer_pools[layer] if ( - len(grouped_layers) != len(dense_layers) - or len(set(grouped_layers)) != len(grouped_layers) - or len(set(dense_layers)) != len(dense_layers) - or set(dense_layers) != set(grouped_layers) + tuple(pool.shape[1:]) != tuple(p0.shape[1:]) + or tuple(pool.stride()) != strides + or pool.dtype != p0.dtype ): - raise ValueError("dense layer order must cover every grouped layer exactly once") - self.device = layer_pools[page_representatives[0]].device - if self.device.type != "cuda": - raise ValueError("fixed score metadata is CUDA-only") - self.max_requests = max_requests - self.bucket_seq_len = seq_len - if page_table_token_capacity is None: - page_table_token_capacity = seq_len - if page_table_token_capacity < seq_len: - raise ValueError("page-table capacity cannot be smaller than the score bucket") - self.page_table_token_capacity = int(page_table_token_capacity) - # Decode-width capacity of the score buffers; per-request prompt - # lengths are staged runtime metadata. Default: the whole sequence - # capacity is scorable. - if decode_width is None: - decode_width = int(seq_len) - if decode_width <= 0 or decode_width > seq_len: - raise ValueError("fixed score decode width exceeds the sequence capacity") - self.decode_width = int(decode_width) - q_real = q_real.to(device=self.device, dtype=torch.float32).contiguous() - q_imag = q_imag.to(device=self.device, dtype=torch.float32).contiguous() - mlr_coef = mlr_coef.to(device=self.device, dtype=torch.float32).contiguous() - freq_scale_sq = freq_scale_sq.to(device=self.device, dtype=torch.float32).contiguous() - offsets = offsets.to(device=self.device, dtype=torch.float32).contiguous() - omega = omega.to(device=self.device, dtype=torch.float32).contiguous() - if page_table_keys is None: - page_table_keys = list(range(len(page_representatives))) - self.representative_slots, minimum_page_table_slots = self._page_table_slot_layout( - page_representatives, page_table_keys + raise ValueError("fixed score layers must share one uniform geometry") + if int(pool.data_ptr()) % element_size: + raise ValueError("fixed score layer base is not element-aligned") + ws.num_layers = len(dense_layers) + ws.num_q_heads = int(num_q_heads) + ws.num_kv_heads = int(num_kv_heads) + ws.num_freqs = int(num_freqs) + ws.tokens_per_block = int(tokens_per_block) + # Calibration tables span every model layer; segments index them by + # ABSOLUTE layer id ON DEVICE where they cannot be range-checked, so + # validate the extent once here, loudly. + num_calibrated_layers = q_real.numel() // (ws.num_q_heads * ws.num_freqs) + if min(dense_layers) < 0 or max(dense_layers) >= num_calibrated_layers: + raise ValueError("scored layer index exceeds the calibrated layer extent") + _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} + page_table_slots = [ws.representative_slots[_rep_of[layer]] for layer in dense_layers] + ws.seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( + ws.num_layers + ) + seg_layer = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( + max_requests + ) + block_offsets = ws.block_offsets_device + slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) + if int(slots_t.max()) >= int(block_offsets.shape[0]): + raise ValueError("page table slot exceeds staged page-id planes") + req_idx = torch.arange(max_requests, dtype=torch.int64, device=device).repeat_interleave( + ws.num_layers + ) + slot_idx = slots_t.repeat(max_requests) + seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) + + max_segments = max_requests * ws.num_layers + # The fused kernel masks its ragged tail, but the scratch bucket contract + # stays tile-aligned (64 tokens, or one page for 128-token pages) so + # bucket geometry cannot silently admit new geometry. + score_tile_tokens = max(64, ws.tokens_per_block) + supported = ( + torch.cuda.get_device_capability(device) == (10, 0) + and p0.dtype == torch.bfloat16 + and kv_factor == 2 + and ws.tokens_per_block in (32, 128) + and ws.num_freqs in (32, 64) + and ws.num_q_heads % ws.num_kv_heads == 0 + and ws.num_q_heads // ws.num_kv_heads in (4, 8) + and int(p0.stride(-1)) == 1 + and seq_len % score_tile_tokens == 0 + # The kernel's head-plane base offset is 64-bit; the widest 32-bit + # product left is one plane (N-1 head columns of one segment stride), + # which the score bucket keeps far below 2^31. Group-4 geometries pad + # the head axis to the MMA tile N=8. + and (8 - 1) * max_segments * seq_len < 2**31 + ) + if not supported: + raise ValueError( + "TriAttention score requires SM100, bf16 KV pools, head size " + "64/128, 32/128-token pages, GQA group 4 or 8, and a bucket " + "capacity aligned to the score compute tile; got " + f"capability={torch.cuda.get_device_capability(device)}, " + f"dtype={p0.dtype}, kv_factor={kv_factor}, " + f"tokens_per_block={ws.tokens_per_block}, num_freqs={ws.num_freqs}, " + f"heads={ws.num_q_heads}q/{ws.num_kv_heads}kv, " + f"stride={int(p0.stride(-1))}, " + f"seq_len={seq_len} (tile {score_tile_tokens}), " + f"offset_audit={ws.num_kv_heads * 8 * max_segments * seq_len}" ) - if num_page_table_slots is None: - num_page_table_slots = minimum_page_table_slots - if num_page_table_slots < minimum_page_table_slots: - raise ValueError("page-table slot capacity does not cover every V2 pool") - tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) - if int(layer_pools[page_representatives[0]].shape[1]) != 2: - raise ValueError("fixed score metadata requires an interleaved K/V pool") - self.page_count = ( - self.page_table_token_capacity + tokens_per_block - 1 - ) // tokens_per_block - self.copy_block_count = (self.page_count + 3) // 4 * 4 - if any( - (self.page_table_token_capacity + int(layer_pools[layer].shape[3]) - 1) - // int(layer_pools[layer].shape[3]) - != self.page_count - for layer in page_representatives - ): - raise ValueError("fixed score metadata requires a uniform page count") - device_page_shape = ( - num_page_table_slots, - max_requests, - 2, - self.copy_block_count, + # The kernel scores each request's window (from its staged per-request + # start) into a head-major scratch padded to the MMA tile N=8 per KV + # head; consumers read only the real heads. All buffers below are + # persistent because the compiled kernels capture their device pointers. + ws.cute_scratch = torch.empty( + ws.num_kv_heads * 8 * max_segments * seq_len, dtype=torch.float32, device=device + ) + ws.seg_seq_len = torch.zeros(max_segments, dtype=torch.int32, device=device) + seg_out_offset = (torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len).to( + torch.int32 + ) + ws.cute_token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) + ws.gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( + 1, 1, 1, 1, -1 + ) + # Compile the SM100 CuTe entries this mode launches -- HERE at workspace + # construction, outside any CUDA graph capture (compilation allocates and + # synchronizes). Union rounds run the fused score+stats+union pipeline; + # the per-head modes run the score-only entry. There is deliberately no + # other score path and no fallback. + union = eviction_mode == "union" + ws.union_rows = None + if union: + # Union output rows are sized by the whole bucket (the widest + # possible window); consumers mask by the per-request widths. + ws.union_rows = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) + try: + ws.runner = TriAttentionCuteScoreRunner( + layer_pools=list(layer_pools), + layer_indices=[int(layer) for layer in dense_layers], + max_requests=max_requests, + num_layers=ws.num_layers, + seq_len=seq_len, + num_q_heads=ws.num_q_heads, + num_kv_heads=ws.num_kv_heads, + num_freqs=ws.num_freqs, + tokens_per_block=ws.tokens_per_block, + page_ids=block_offsets.view(-1), + seg_page_off=seg_page_off, + seg_req_id=ws.seg_req, + seg_layer_id=seg_layer, + seg_seq_len=ws.seg_seq_len, + seg_out_offset=seg_out_offset, + token_starts=ws.cute_token_starts, + q_real=q_real.view(-1), + q_imag=q_imag.view(-1), + mlr_coef=mlr_coef.view(-1), + mean_cos=ws.mean_cos, + mean_sin=ws.mean_sin, + freq_scale_sq=freq_scale_sq, + output=ws.cute_scratch, + enable_partial_stats=union, ) - # One table carries every per-round host value: three metadata rows - # (logical position, valid length, prompt length) plus one move-offsets - # row per compacted cache family, so each round pays exactly one - # host-to-device copy. Offsets rows have request_capacity + 1 entries, - # hence the extra column. - self.request_metadata_host = torch.empty( - (6, max_requests + 1), - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), + except (ImportError, RuntimeError, ValueError, AssertionError) as error: + raise RuntimeError( + "TriAttention CuTe score setup failed and no other score path exists" + ) from error + logger.info( + f"TriAttention CuTe score enabled: {ws.num_q_heads}q/{ws.num_kv_heads}kv heads, " + f"{ws.num_freqs} freqs, {ws.tokens_per_block}-token pages" + ) + + # ---- selection buffers -------------------------------------------------- + # Per-request valid decode widths, refreshed each round from the staged + # lengths; prompt offsets alias the staged per-request prompt lengths so + # the values are written once per round. + ws.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) + ws.prompt_offsets = ws.token_starts_device + if union: + ws.selection_rows_per_request = 1 + ws.row_prompt_offsets = ws.prompt_offsets + # The fused pipeline writes normalized per-request union rows straight + # into ``combined``; only the top-k settle-and-pack stage remains. + ws.combined = torch.empty((max_requests, decode_width), dtype=torch.float32, device=device) + ws.final_indices = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) + # Kept decode ordinals only: rows are prompt-length independent, so + # one workspace serves cohorts with mixed prompt lengths. + ws.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) + # Row-major views consumed by the top-k settle launch. + ws.selection_scores_rows = ws.combined + ws.selection_row_lengths = ws.valid_widths + ws.provisional_rows = ws.final_indices + ws.keep_rows = ws.keep + # Padded rows carry zero valid width; their provisional TopK entries + # must still be in-range ordinals for the finalizer's score gather. + ws.final_indices.zero_() + ws.score_output = None + else: + selection_rows = ( + ws.num_kv_heads if eviction_mode == "per_head" else ws.num_layers * ws.num_kv_heads + ) + ws.selection_rows_per_request = selection_rows + ws.row_prompt_offsets = torch.zeros( + (max_requests * selection_rows,), dtype=torch.int32, device=device ) - self._bulk_copy_idx_src = torch.arange( + # Decode-only per-head scores gathered from the CuTe scratch, the + # ``[request, layer, head, token]`` layout the reduce kernels read. + ws.score_output = torch.empty( max_requests, - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), + ws.num_layers, + ws.num_q_heads, + decode_width, + dtype=torch.float32, + device=device, ) - self._bulk_offsets_src = torch.empty( - device_page_shape, - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), + score_shape = (max_requests, ws.num_layers, ws.num_q_heads, 1) + ws.row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) + ws.row_std = torch.empty_like(ws.row_mean) + ws.selection_scores = torch.empty( + (max_requests, selection_rows, decode_width), dtype=torch.float32, device=device ) - self.block_offsets_device = torch.empty( - device_page_shape, - dtype=torch.int32, - device=self.device, + ws.row_seq_lens = torch.full( + (max_requests, selection_rows), decode_width, dtype=torch.int32, device=device ) - # Optional second block-offset staging plane for a co-compressed draft - # KV cache. The draft is never scored: these offsets feed only the - # draft compact launches. - self.draft_block_offsets_device: Optional[torch.Tensor] = None - self._draft_bulk_offsets_src: Optional[torch.Tensor] = None - self.draft_representative_slots: Dict[int, int] = {} - self.draft_copy_block_count = 0 - self.draft_page_table_token_capacity = 0 - if draft_layer_pools is not None: - if ( - not draft_page_representatives - or draft_page_table_keys is None - or draft_page_table_token_capacity is None - or draft_page_table_token_capacity <= 0 - ): - raise ValueError( - "draft page-table staging requires representatives, keys, and capacity" - ) - ( - self.draft_representative_slots, - minimum_draft_slots, - ) = self._page_table_slot_layout(draft_page_representatives, draft_page_table_keys) - if draft_num_page_table_slots is None: - draft_num_page_table_slots = minimum_draft_slots - if draft_num_page_table_slots < minimum_draft_slots: - raise ValueError("draft page-table slot capacity does not cover every V2 pool") - draft_tokens_per_block = int(draft_layer_pools[draft_page_representatives[0]].shape[3]) - if int(draft_layer_pools[draft_page_representatives[0]].shape[1]) != 2: - raise ValueError("draft page-table staging requires an interleaved K/V pool") - self.draft_page_table_token_capacity = int(draft_page_table_token_capacity) - draft_capacity = self.draft_page_table_token_capacity - draft_page_count = ( - draft_capacity + draft_tokens_per_block - 1 - ) // draft_tokens_per_block - if any( - (draft_capacity + int(draft_layer_pools[layer].shape[3]) - 1) - // int(draft_layer_pools[layer].shape[3]) - != draft_page_count - for layer in draft_page_representatives - ): - raise ValueError("draft page-table staging requires a uniform page count") - self.draft_copy_block_count = (draft_page_count + 3) // 4 * 4 - draft_page_shape = ( - draft_num_page_table_slots, - max_requests, - 2, - self.draft_copy_block_count, - ) - self._draft_bulk_offsets_src = torch.empty( - draft_page_shape, - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), - ) - self.draft_block_offsets_device = torch.empty( - draft_page_shape, - dtype=torch.int32, - device=self.device, + selection_shape = (max_requests, selection_rows, keep_count) + ws.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=device) + ws.keep = torch.empty(selection_shape, dtype=torch.int32, device=device) + ws.selection_scores_rows = ws.selection_scores.view( + max_requests * selection_rows, decode_width + ) + ws.selection_row_lengths = ws.row_seq_lens.view(-1) + ws.provisional_rows = ws.top_indices_i32.view(-1, keep_count) + ws.keep_rows = ws.keep.view(-1, keep_count) + ws.top_indices_i32.zero_() + + # ---- compaction launch data + settle/pack fusion ------------------------ + ws.compaction = None + # One settle program per (request, selection row). + ws.settle_grid = (max_requests, ws.selection_rows_per_request) + if build_compaction: + if layer_group_representative is None or layer_pool_keys is None: + raise ValueError("compaction requires the layer grouping and pool keys") + draft_kwargs = {} + if draft_layers: + draft_kwargs = dict( + draft_layer_pools=draft_layer_pools, + draft_layers=list(draft_layers), + draft_layer_group_representative=draft_layer_group_representative, + draft_layer_pool_keys=draft_layer_pool_keys, + draft_protected_tail_capacity=int(draft_protected_tail_capacity), + draft_kv_block_offsets=ws.draft_block_offsets_device, + draft_page_table_slots=ws.draft_representative_slots, + draft_move_offsets=ws.draft_move_offsets, ) - # Zero-filled so an unstaged cohort gathers the phase table's row 0 - # instead of indexing it with uninitialized round starts. - self.request_metadata_device = torch.zeros( - (6, max_requests + 1), dtype=torch.int32, device=self.device + ws.compaction = build_cache_compactions( + eviction_mode=eviction_mode, + layer_pools=layer_pools, + dense_layers=list(dense_layers), + swa_layers=list(swa_layers), + layer_group_representative=layer_group_representative, + kept_token_ordinals=ws.keep, + valid_sequence_lengths=ws.valid_seq_lens_device, + kv_block_offsets=ws.block_offsets_device, + page_table_slots=ws.representative_slots, + request_count=max_requests, + prompt_offsets=ws.token_starts_device, + decode_keep_count=keep_count, + swa_window=swa_window, + layer_pool_keys=list(layer_pool_keys), + protected_tail_capacity=int(protected_tail_capacity), + # Tails vary per round (in-flight growth), so the per-family move + # offsets ride the staged metadata rows each round. + dense_move_offsets=ws.dense_move_offsets, + swa_move_offsets=ws.swa_move_offsets, + # ONE launch settles the kept ordinals and packs the dense/SWA + # move sources; ``run_cache_compactions`` then only runs the C++ + # moves (plus the draft's own pack). + fuse_dense_pack_into_selection=True, + **draft_kwargs, ) - self.round_starts_device = self.request_metadata_device[0, :max_requests] - self.valid_seq_lens_device = self.request_metadata_device[1, :max_requests] - # Per-request pinned prompt lengths: the score kernel starts each - # request's decode window here, so one bucket may mix prompt lengths. - self.token_starts_device = self.request_metadata_device[2, :max_requests] - # Per-family move offsets consumed by the compaction pack kernel and - # the C++ compact launches; refreshed with the metadata each round. - self.dense_move_offsets = self.request_metadata_device[3] - self.swa_move_offsets = self.request_metadata_device[4] - self.draft_move_offsets = self.request_metadata_device[5] - self.mean_cos = torch.empty( - (max_requests, num_freqs), dtype=torch.float32, device=self.device + pack = ws.compaction["dense_pack"] + # The fused kernel reads back the kept ordinals it just wrote, so the + # packing must read this workspace's own keep buffer. + assert ( + pack["kept_token_ordinals"].data_ptr() == ws.keep_rows.data_ptr() + and pack["selection_rows"] == ws.selection_rows_per_request + and pack["keep_count"] == keep_count + ), "fused move packing must match the selection geometry" + ws.settle_pack_tensors = ( + pack["valid_sequence_lengths"], + pack["dense_offsets"], + pack["dense_indices"], + pack["swa_offsets"], + pack["swa_indices"], ) - self.mean_sin = torch.empty_like(self.mean_cos) - # The phase table depends only on the shared calibration, so the - # manager passes one instance to every staging bucket; standalone - # construction (tests) builds a private one. - if mean_phase_table is None: - mean_phase_table = MeanPhaseTable(offsets, omega, initial_rows=seq_len) - self.mean_phase_table = mean_phase_table - # ONE fused group across ALL dense layers: segments carry their own - # layer base address and page-table slot, so distinct per-layer - # storages/block tables share a single launch. - _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} - _page_table_slots = [self.representative_slots[_rep_of[layer]] for layer in dense_layers] - self.fused_group = _FixedScoreGroup( - layer_pools, - dense_layers, - max_requests, - self.page_count, - seq_len, - num_q_heads, - self.block_offsets_device, - _page_table_slots, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - output_width=decode_width, + ws.settle_pack_shape = dict( + DENSE_TOTAL=pack["dense_total"], + SWA_TOTAL=pack["swa_total"], + MOVE_CAPACITY=pack["move_capacity"], + NUM_KV_HEADS=pack["num_kv_heads"], + SWA_WINDOW=pack["swa_window"], + UNION=pack["union"], + PER_LAYER=pack["per_layer"], + HAS_SWA=pack["has_swa"], + HAS_PACK=True, + ) + else: + # Settle-only workspaces (unit tests): the pack half is compiled + # away, so its tensor parameters are never read. + placeholder = ws.selection_row_lengths + ws.settle_pack_tensors = (placeholder,) * 5 + ws.settle_pack_shape = dict( + DENSE_TOTAL=0, + SWA_TOTAL=0, + MOVE_CAPACITY=0, + NUM_KV_HEADS=1, + SWA_WINDOW=0, + UNION=False, + PER_LAYER=False, + HAS_SWA=False, + HAS_PACK=False, ) - # Compile the SM100 CuTe score kernel (the only score implementation) - # here at workspace construction, outside any CUDA graph capture - # (compilation allocates and synchronizes). Unsupported geometry - # raises loudly right here rather than mid-round at the score launch. - self.fused_group.prepare_cute_score(self.mean_cos, self.mean_sin) - self.copy_done = torch.cuda.Event() - # First record publishes constructor allocations to the V2 copy stream; - # later records protect pinned metadata before the next cohort reuses it. - self.copy_done.record(torch.cuda.current_stream(self.device)) - self.bulk_copy_done = torch.cuda.Event() - self.bulk_consume_done = torch.cuda.Event() - self.copy_pending = False - self.page_tables_active = False - self.stream = None - self._score_valid_widths: Optional[torch.Tensor] = None - self._score_launcher_bound = False - - def bind_score_launcher(self, valid_widths: torch.Tensor, aggregation: str) -> None: - """Bind the per-row score widths for these buffers (mean-only).""" - if self._score_launcher_bound: - raise RuntimeError("TriAttention score launcher is already bound") - if aggregation != "mean": - raise ValueError( - f"unsupported score aggregation {aggregation!r}: max aggregation " - "was removed with the C++ score stack; only 'mean' exists" - ) - self._score_valid_widths = valid_widths - self._score_launcher_bound = True - - def launch_prepared_union_fusion(self, union_out: torch.Tensor) -> None: - """Run the fused score+stats+union pipeline over these buffers. - This is THE union score path; there is deliberately no fallback. A - geometry or capacity the fused pipeline cannot serve raises loudly. - """ - if not self._score_launcher_bound: - raise RuntimeError("TriAttention score launcher is not bound") - stream = torch.cuda.current_stream(self.device) - if self.stream is None: - self.stream = stream - elif (stream.device, stream.cuda_stream) != ( - self.stream.device, - self.stream.cuda_stream, - ): - raise _FixedScoreStreamMismatch( - "TriAttention score launches must stay on the staging CUDA stream" - ) - self.mean_phase_table.gather( - self.round_starts_device, - self.mean_cos, - self.mean_sin, - self.max_requests, + # ---- round-ordering events ---------------------------------------------- + ws.copy_done = torch.cuda.Event() + # First record publishes constructor allocations to the V2 copy stream; + # later records protect pinned metadata before the next cohort reuses it. + ws.copy_done.record(torch.cuda.current_stream(device)) + ws.bulk_copy_done = torch.cuda.Event() + ws.bulk_consume_done = torch.cuda.Event() + ws.copy_pending = False + ws.page_tables_active = False + ws.stream = None + return ws + + +def _stage_block_offsets( + ws: SimpleNamespace, + manager: KVCacheManagerV2, + request_ids: List[int], + current_stream: torch.cuda.Stream, + source: torch.Tensor, + destination: torch.Tensor, + copy_block_count: int, +) -> None: + """Copy one request group's V2 block offsets before live compaction. + + Uses the V2 block-offset kernel with an immutable pinned snapshot of the + selected host-table rows: this enqueues asynchronous host-memory reads, + and TriAttention later resizes the same cache, which mutates the + manager's table in place. The IndexMapper synchronously resolves request + slots and gathers only their beam-0 K block offsets, decoupling both live + inputs before the native asynchronous copy consumes the snapshot with + identity indices. ``dst[pool, r, 0(K), :]`` holds ``base_page * + index_scales``; score and compact decode that K plane inline. + """ + if ws.copy_pending and not ws.copy_done.query(): + ws.copy_done.synchronize() + # The native device copy reads only K and derives V with kv_offset. + manager.index_mapper.gather_k_block_offsets( + manager.host_kv_cache_block_offsets, + source, + request_ids, + copy_block_count, + ) + manager._stream.wait_event(ws.copy_done) + copy_batch_block_offsets_to_device( + source, + destination, + ws._bulk_copy_idx_src[: len(request_ids)], + manager.index_scales, + manager.kv_offset, + manager._stream.cuda_stream, + ) + ws.bulk_copy_done.record(manager._stream) + current_stream.wait_event(ws.bulk_copy_done) + + +def stage_eviction_cohort( + ws: SimpleNamespace, + manager: KVCacheManagerV2, + request_ids: List[int], + round_starts: List[int], + token_starts: List[int], + seq_lens: Optional[List[int]] = None, + page_table_seq_lens: Optional[List[int]] = None, + draft_manager: Optional[KVCacheManagerV2] = None, + dense_move_offsets: Optional[List[int]] = None, + swa_move_offsets: Optional[List[int]] = None, + draft_move_offsets: Optional[List[int]] = None, +) -> None: + """Copy one eviction cohort into the reusable device buffers. + + ``token_starts`` carries each request's pinned prompt length; the score + kernel starts that request's decode window there, so the cohort may mix + prompt lengths. Raises on any invalid cohort -- there is no fallback. + """ + request_count = len(request_ids) + if ( + request_count == 0 + or request_count > ws.max_requests + or len(round_starts) != request_count + or len(token_starts) != request_count + ): + raise ValueError("eviction cohort does not fit the workspace request capacity") + stream = _bind_workspace_stream(ws) + if ws.page_tables_active: + raise RuntimeError("previous page-table cohort is still active") + if seq_lens is None: + seq_lens = [ws.bucket_seq_len] * request_count + if page_table_seq_lens is None: + page_table_seq_lens = seq_lens + if len(seq_lens) != request_count or len(page_table_seq_lens) != request_count: + raise ValueError("eviction cohort length lists do not match the request count") + if (draft_manager is None) != (ws.draft_block_offsets_device is None): + raise RuntimeError("draft staging requires the workspace built with the draft cache") + request_metadata = torch.as_tensor((round_starts, seq_lens, token_starts), dtype=torch.int32) + if min(round_starts) < 0: + raise ValueError("eviction round starts must be non-negative") + # Grow the phase table while this cohort's round starts are still host + # integers: a stale-capacity gather is an out-of-bounds index_select on + # the device. + grow_mean_phase_table(ws.phase, int(max(round_starts)) + 1) + _stage_block_offsets( + ws, + manager, + request_ids, + stream, + ws._bulk_offsets_src, + ws.block_offsets_device, + ws.copy_block_count, + ) + if draft_manager is not None: + _stage_block_offsets( + ws, + draft_manager, + request_ids, + stream, + ws._draft_bulk_offsets_src, + ws.draft_block_offsets_device, + ws.draft_copy_block_count, ) - return self.fused_group.launch_cute_union_fusion( - self.max_requests, - self.valid_seq_lens_device, - self._score_valid_widths, - self.token_starts_device, - self.mean_cos, - self.mean_sin, - union_out, + ws.request_metadata_host[:3, :request_count].copy_(request_metadata) + # Rows past this cohort are padding: zero lengths keep the score kernel + # and selection inert for them. + ws.request_metadata_host[:3, request_count:].zero_() + # This round's per-family move offsets ride the same table, so the single + # device copy below carries them too. + for row, family_offsets in ( + (3, dense_move_offsets), + (4, swa_move_offsets), + (5, draft_move_offsets), + ): + if family_offsets is not None: + ws.request_metadata_host[row, : len(family_offsets)].copy_( + torch.as_tensor(family_offsets, dtype=torch.int32) + ) + try: + # Copy the fixed backing once. Only the first ``request_count`` + # columns are consumed by this cohort. + ws.request_metadata_device.copy_(ws.request_metadata_host, non_blocking=True) + finally: + # Guard the pinned metadata until its asynchronous copies complete. + # Page-table device-buffer reuse is guarded separately after compact. + ws.copy_done.record(stream) + ws.copy_pending = True + ws.page_tables_active = True + # The staged per-request prompt lengths are shared with the selection; + # per-head modes re-expand them into their row-major view here. + if ws.row_prompt_offsets is not ws.prompt_offsets: + ws.row_prompt_offsets.view(ws.max_requests, ws.selection_rows_per_request).copy_( + ws.prompt_offsets.unsqueeze(1).expand(-1, ws.selection_rows_per_request) ) - def launch_prepared_score(self) -> torch.Tensor: - """Gather the phase means and launch the score kernel over these buffers.""" - if not self._score_launcher_bound: - raise RuntimeError("TriAttention score launcher is not bound") - stream = torch.cuda.current_stream(self.device) - if self.stream is None: - self.stream = stream - elif (stream.device, stream.cuda_stream) != ( - self.stream.device, - self.stream.cuda_stream, - ): - raise _FixedScoreStreamMismatch( - "TriAttention score launches must stay on the staging CUDA stream" - ) - # mean_cos/mean_sin feed the CuTe score kernel, whose compiled launch - # captured their device pointers, so they must be refreshed in place - # from this round's staged round starts before it runs. - self.mean_phase_table.gather( - self.round_starts_device, - self.mean_cos, - self.mean_sin, - self.max_requests, + +def mark_page_tables_consumed(ws: SimpleNamespace, *manager_streams: torch.cuda.Stream) -> None: + """Order V2 page-table reuse and resize after this cohort's compact. + + Every passed manager stream (target, and the draft when co-compressed) + waits on one event recorded after the compact launches, so neither cache + can free or reallocate pages this cohort is still reading. + """ + if not ws.page_tables_active: + raise RuntimeError("TriAttention page tables were not staged") + ws.bulk_consume_done.record(torch.cuda.current_stream(ws.device)) + for manager_stream in manager_streams: + manager_stream.wait_event(ws.bulk_consume_done) + ws.page_tables_active = False + + +def settle_top_tokens(ws: SimpleNamespace) -> None: + """Pick the top-k with the CuTE selector, then settle its output. + + The CuTE top-k is fast but breaks score ties arbitrarily and emits + indices in arbitrary order; the settle kernel recomputes the threshold + membership with lowest-index-wins ties, rebases each row by its prompt + offset, and writes sorted ordinals. The same launch packs each request's + dense/SWA compaction move sources from the ordinals it just settled + (workspaces built without compaction compile the pack half away). + """ + # The trailing 1 is next_n: decode scores one query token per request. + torch.ops.trtllm.cute_dsl_indexer_topk_decode( + ws.selection_scores_rows, + ws.selection_row_lengths, + ws.provisional_rows, + ws.keep_count, + 1, + ) + _settle_ties_and_pack_compaction_sources_kernel[ws.settle_grid]( + ws.selection_scores_rows, + ws.selection_row_lengths, + ws.row_prompt_offsets, + ws.provisional_rows, + ws.keep_rows, + *ws.settle_pack_tensors, + WIDTH=ws.decode_width, + KEEP_COUNT=ws.keep_count, + OUTPUT_WIDTH=ws.keep_count, + SELECTION_ROWS=ws.selection_rows_per_request, + **ws.settle_pack_shape, + HAS_SETTLE=True, + BLOCK=256, + num_warps=4, + ) + + +def run_eviction_round(ws: SimpleNamespace, normalize_scores: bool) -> None: + """One staged eviction round, kernels fired directly in sequence. + + Union: phase gather, fused score+stats+union (two CuTe launches), top-k, + settle-and-pack, C++ compacts. Per-head modes: phase gather, score-only + CuTe launch, decode-window gather, stats+reduce kernels, top-k, + settle-and-pack, C++ compacts. Every launch covers the full request + capacity; padded rows past the staged cohort carry zero lengths and stay + inert. + """ + _bind_workspace_stream(ws) + request_count = ws.max_requests + num_segments = request_count * ws.num_layers + union = ws.eviction_mode == "union" + with nvtx_range("triattention.score", color="blue"): + # mean_cos/mean_sin feed the compiled score launches, which captured + # their device pointers: refresh them in place from this round's + # staged round starts. The same holds for the per-request decode + # widths, per-segment valid lengths, and staged window starts. + gather_mean_phases( + ws.phase, ws.round_starts_device, ws.mean_cos, ws.mean_sin, request_count ) - return self.fused_group.launch( - self.max_requests, - self.valid_seq_lens_device, - self._score_valid_widths, - self.token_starts_device, - self.mean_cos, - self.mean_sin, + torch.sub( + ws.valid_seq_lens_device[:request_count], + ws.token_starts_device[:request_count], + out=ws.valid_widths[:request_count], ) - - def stage( - self, - manager: KVCacheManagerV2, - request_ids: List[int], - round_starts: List[int], - token_starts: List[int], - seq_lens: Optional[List[int]] = None, - page_table_seq_lens: Optional[List[int]] = None, - draft_manager: Optional[KVCacheManagerV2] = None, - dense_move_offsets: Optional[List[int]] = None, - swa_move_offsets: Optional[List[int]] = None, - draft_move_offsets: Optional[List[int]] = None, - ) -> bool: - """Copy one eviction cohort into reusable device buffers. - - ``token_starts`` carries each request's pinned prompt length; the - score kernel starts that request's decode window there, so the cohort - may mix prompt lengths. - """ - request_count = len(request_ids) - if ( - request_count == 0 - or request_count > self.max_requests - or len(round_starts) != request_count - or len(token_starts) != request_count - ): - return False - if (draft_manager is None) != (self.draft_block_offsets_device is None): - return False - stream = torch.cuda.current_stream(self.device) - if self.stream is None: - self.stream = stream - elif (stream.device, stream.cuda_stream) != ( - self.stream.device, - self.stream.cuda_stream, - ): - raise _FixedScoreStreamMismatch( - "TriAttention fixed score metadata is bound to its first CUDA stream" - ) - if self.page_tables_active: - raise RuntimeError("previous page-table cohort is still active") - if seq_lens is None: - seq_lens = [self.bucket_seq_len] * request_count - if page_table_seq_lens is None: - page_table_seq_lens = seq_lens - if len(seq_lens) != request_count or len(page_table_seq_lens) != request_count: - return False - if manager.enable_swa_scratch_reuse: - raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") - try: - request_metadata = torch.as_tensor( - (round_starts, seq_lens, token_starts), dtype=torch.int32 - ) - except (OverflowError, RuntimeError, TypeError, ValueError): - return False - if min(round_starts) < 0: - return False - # Grow the phase table while this cohort's round starts are still host - # integers: a stale-capacity gather is an out-of-bounds index_select - # on the device. - self.mean_phase_table.ensure(int(max(round_starts)) + 1) - if not self._stage_page_tables_bulk( - manager, - request_ids, - stream, - self._bulk_offsets_src, - self.block_offsets_device, - self.copy_block_count, - ): - return False - if draft_manager is not None: - if draft_manager.enable_swa_scratch_reuse: + torch.index_select( + ws.valid_seq_lens_device, + 0, + ws.seg_req[:num_segments], + out=ws.seg_seq_len[:num_segments], + ) + ws.cute_token_starts[:request_count].copy_(ws.token_starts_device[:request_count]) + if union: + if not ws.runner.supports_union_fusion(request_count): raise RuntimeError( - "TriAttention does not support V2 SWA scratch page-table remapping" + f"TriAttention CuTe union fusion has no compiled variant for " + f"request_count={request_count} (capacity {ws.max_requests}) " + "and no other union path exists" ) - assert self._draft_bulk_offsets_src is not None - assert self.draft_block_offsets_device is not None - if not self._stage_page_tables_bulk( - draft_manager, - request_ids, - stream, - self._draft_bulk_offsets_src, - self.draft_block_offsets_device, - self.draft_copy_block_count, - ): - return False - self.request_metadata_host[:3, :request_count].copy_(request_metadata) - # Rows past this cohort are padding: zero lengths keep the score - # kernel and selection inert for them. - self.request_metadata_host[:3, request_count:].zero_() - # This round's per-family move offsets ride the same table, so the - # single device copy below carries them too. - for row, family_offsets in ( - (3, dense_move_offsets), - (4, swa_move_offsets), - (5, draft_move_offsets), - ): - if family_offsets is not None: - self.request_metadata_host[row, : len(family_offsets)].copy_( - torch.as_tensor(family_offsets, dtype=torch.int32) + ws.runner.launch_union_fusion( + request_count, ws.mean_cos, ws.mean_sin, ws.union_rows[:request_count] + ) + columns = min(ws.union_rows.shape[1], ws.combined.shape[1]) + ws.combined[:request_count, :columns].copy_(ws.union_rows[:request_count, :columns]) + else: + if not ws.runner.supports(request_count): + raise RuntimeError( + f"TriAttention CuTe score has no compiled variant for " + f"request_count={request_count} (capacity {ws.max_requests}) " + "and no other score path exists" ) - try: - # Copy the fixed backing once. Only the first ``request_count`` - # columns are consumed by this cohort. - self.request_metadata_device.copy_(self.request_metadata_host, non_blocking=True) - finally: - # Guard the pinned metadata until its asynchronous copies complete. - # Page-table device-buffer reuse is guarded separately after compact. - self.copy_done.record(stream) - self.copy_pending = True - self.page_tables_active = True - return True - - def _stage_page_tables_bulk( - self, - manager: KVCacheManagerV2, - request_ids: List[int], - current_stream: torch.cuda.Stream, - source: torch.Tensor, - destination: torch.Tensor, - copy_block_count: int, - ) -> bool: - """Copy one request group's V2 block offsets before live compaction. - - Uses the V2 block-offset kernel with an immutable pinned snapshot of the - selected host-table rows. The snapshot is required because - this method enqueues asynchronous host-memory reads; TriAttention later - resizes the same cache, which mutates the manager's table in place. - The IndexMapper synchronously resolves request slots and gathers only - their beam-0 K block offsets, decoupling both live inputs before the - native asynchronous copy consumes the snapshot with identity indices. - ``dst[pool, r, 0(K), :]`` holds ``base_page * index_scales``. Score and - compact decode that K plane inline, avoiding any conversion kernel. - """ - if not request_ids or len(request_ids) > self.max_requests: - return False - - host_table = manager.host_kv_cache_block_offsets - num_pools, _, kv_planes, max_blocks = host_table.shape - if ( - host_table.dtype != torch.int32 - or kv_planes != 2 - or copy_block_count > max_blocks - or int(manager.kv_factor) != 2 - or num_pools != destination.shape[0] - ): - return False - request_count = len(request_ids) - submitted = False - try: - if self.copy_pending and not self.copy_done.query(): - self.copy_done.synchronize() - # The native device copy reads only K and derives V with kv_offset. - manager.index_mapper.gather_k_block_offsets( - host_table, - source, - request_ids, - copy_block_count, + ws.runner.launch(request_count, ws.mean_cos, ws.mean_sin) + # The kernel wrote each request's window scores (from its pinned + # prompt length) into the head-major scratch, padded to the MMA + # tile N=8 per KV head. Gather each request's decode window into + # the [request, layer, head, token] layout the reduce kernels + # read; columns past a request's valid width carry unscored + # scratch data masked by ``valid_widths``. + group_size = ws.num_q_heads // ws.num_kv_heads + source = ( + ws.cute_scratch[: ws.num_kv_heads * 8 * num_segments * ws.bucket_seq_len] + .view(ws.num_kv_heads, 8, request_count, ws.num_layers, ws.bucket_seq_len)[ + :, :group_size + ] + .permute(2, 3, 0, 1, 4) + ) + columns = ( + ws.token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + + ws.gather_columns + ) + columns = columns.clamp_(max=ws.bucket_seq_len - 1).expand( + request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width ) - manager._stream.wait_event(self.copy_done) - copy_batch_block_offsets_to_device( + torch.gather( source, - destination, - self._bulk_copy_idx_src[:request_count], - manager.index_scales, - manager.kv_offset, - manager._stream.cuda_stream, + 4, + columns, + out=ws.score_output[:request_count].view( + request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + ), ) - submitted = True - self.bulk_copy_done.record(manager._stream) - current_stream.wait_event(self.bulk_copy_done) - except (AttributeError, IndexError, KeyError, RuntimeError, TypeError, ValueError) as exc: - if submitted: - raise RuntimeError( - "TriAttention bulk page-table copy failed after GPU submission" - ) from exc - logger.warning(f"TriAttention bulk page-table staging failed: {exc}") - return False - return True - - def mark_page_tables_consumed(self, *manager_streams: torch.cuda.Stream) -> None: - """Order V2 page-table reuse and resize after this cohort's compact. - - Every passed manager stream (target, and the draft when co-compressed) - waits on one event recorded after the compact launches, so neither - cache can free or reallocate pages this cohort is still reading. - """ - if not self.page_tables_active: - raise RuntimeError("TriAttention page tables were not staged") - self.bulk_consume_done.record(torch.cuda.current_stream(self.device)) - for manager_stream in manager_streams: - manager_stream.wait_event(self.bulk_consume_done) - self.page_tables_active = False - - -@dataclass(frozen=True, kw_only=True, slots=True) -class _PreparedEviction: - """Request metadata validated before score, select, and compact.""" - - request: "LlmRequest" - request_id: int - seq_len: int - round_start: int - prompt_len: int - expected_keep_count: int - protected_tail: int - - -@dataclass(kw_only=True, slots=True) -class _RequestCompressionState: - """Mutable compression state owned by one live request.""" - - generation_steps: int = 0 - evicted_tokens: int = 0 - confirmed_kv_length: Optional[int] = None - - -@dataclass(kw_only=True, slots=True) -class _PreparedGenerationBatch: - """Target growth reserved by the most recently prepared generation batch.""" - - batch: "ScheduledRequests" - growth_by_request: Dict[int, int] - - -@dataclass(kw_only=True, slots=True) -class _EvictionBuffers: - """Reusable fixed score and selection buffers for one runtime shape.""" - - score_staging: _FixedScoreStagingBuffers - keep_set_selector: _BatchedKeepSetSelector + with nvtx_range("triattention.select", color="yellow"): + if not union: + prepare_per_head_scores( + ws.score_output[:request_count], + ws.valid_widths, + ws.row_mean, + ws.row_std, + ws.selection_scores, + ws.row_seq_lens, + request_count, + num_kv_heads=ws.num_kv_heads, + per_layer=ws.eviction_mode == "per_layer_perhead", + normalize_scores=normalize_scores, + ) + settle_top_tokens(ws) + with nvtx_range("triattention.compact", color="purple"): + run_cache_compactions(ws.compaction) class TriAttention(BaseKVCacheCompressionManager): @@ -1068,26 +1059,28 @@ def __init__( # Geometric integration offsets (built lazily on first eviction so the # device matches the cache pool). self._offsets: Optional[torch.Tensor] = None - self._mean_phase_table: Optional["MeanPhaseTable"] = None + # Mean-phase table dict, shared by reference with every workspace so + # it persists across workspace rebuilds. + self._phase: Optional[Dict[str, object]] = None - # Request presence records successful initialization. The record also - # owns the counters and physical length cleared at request finish. - self._request_states: Dict[int, _RequestCompressionState] = {} + # Request presence records successful initialization. Each value is a + # plain dict {generation_steps, evicted_tokens, confirmed_kv_length}. + self._request_states: Dict[int, Dict[str, object]] = {} # The overlap executor prepares B(n) before finalizing B(n-1). Keep the - # exact fixed-linear generation width for that currently in-flight batch; - # the final hook treats those slots as an opaque suffix. - self._prepared_generation_batch: Optional[_PreparedGenerationBatch] = None - # Eviction buffers are built once at the first eviction, sized to + # exact fixed-linear generation width for that currently in-flight + # batch as ``(batch, {request_id: growth})``; the final hook treats + # those slots as an opaque suffix. + self._prepared_generation_batch: Optional[Tuple[object, Dict[int, int]]] = None + # The eviction workspace is built once at the first eviction, sized to # capacity bounds, and reused for the manager's lifetime. - self._eviction_resources: Optional[_EvictionBuffers] = None - self._eviction_pool_fingerprint: Optional[tuple] = None - self._batched_compaction = None + self._workspace: Optional[SimpleNamespace] = None + self._workspace_fingerprint: Optional[tuple] = None self._local_to_global_layers_cache: Optional[List[int]] = None self._attention_layer_partition_cache: Optional[ Tuple[List[int], List[int], Optional[int]] ] = None - self._runtime_kv_layout_cache: Optional[_RuntimeKVLayout] = None - self._draft_runtime_kv_layout_cache: Optional[_RuntimeKVLayout] = None + self._runtime_kv_layout_cache: Optional[Dict[str, object]] = None + self._draft_runtime_kv_layout_cache: Optional[Dict[str, object]] = None def on_request_init(self, request: "LlmRequest", **kwargs) -> None: """Mark capacity-only decode and resolve calibration once. @@ -1101,7 +1094,11 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: self._validate_request_capacity(request) num_layers = self._num_layers_from_manager() self._attention_layer_partition(num_layers) - self._request_states[request_id] = _RequestCompressionState() + self._request_states[request_id] = { + "generation_steps": 0, + "evicted_tokens": 0, + "confirmed_kv_length": None, + } self._ensure_calibrated() def _validate_request_capacity(self, request: "LlmRequest") -> None: @@ -1203,6 +1200,8 @@ def _validate_v2_compatibility(self) -> None: raise ValueError("TriAttention does not support disaggregated serving") if manager.max_beam_width != 1: raise ValueError("TriAttention requires beam-width-one decoding") + if manager.enable_swa_scratch_reuse: + raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") # Speculative feature gates (resolved draft length, linear drafting, # mode whitelist) run in the factory, where spec_config lives. The # draft cache itself is validated here whenever one is attached. @@ -1219,6 +1218,10 @@ def _validate_v2_compatibility(self) -> None: "the target, so the draft cache must be a standard " "key/value cache" ) + if draft_manager.enable_swa_scratch_reuse: + raise RuntimeError( + "TriAttention does not support V2 SWA scratch page-table remapping" + ) if self.eviction_mode != "union": raise ValueError( "TriAttention draft KV co-compression supports only " @@ -1274,19 +1277,16 @@ def prepare_resources(self, scheduled_batch: "ScheduledRequests") -> None: self.kv_cache_manager._kv_reserve_draft_tokens, ) generation_growth[request_id] = growth - self._prepared_generation_batch = _PreparedGenerationBatch( - batch=scheduled_batch, - growth_by_request=generation_growth, - ) + self._prepared_generation_batch = (scheduled_batch, generation_growth) def _inflight_generation_growth( self, scheduled_batch: "ScheduledRequests", request_id: int ) -> int: """Return exact newer target allocation width under overlap scheduling.""" prepared = self._prepared_generation_batch - if prepared is None or scheduled_batch is prepared.batch: + if prepared is None or scheduled_batch is prepared[0]: return 0 - return prepared.growth_by_request.get(request_id, 0) + return prepared[1].get(request_id, 0) def _periodic_evict( self, @@ -1347,11 +1347,11 @@ def _periodic_evict( f"history {kv_cache.history_length}" ) request_state = self._request_states[request_id] - request_state.confirmed_kv_length = seq_len - previous_step = request_state.generation_steps + request_state["confirmed_kv_length"] = seq_len + previous_step = request_state["generation_steps"] confirmed_delta = 1 + int(request.py_num_accepted_draft_tokens) step = previous_step + confirmed_delta - request_state.generation_steps = step + request_state["generation_steps"] = step if previous_step // self.beta >= step // self.beta: continue if seq_len <= self._minimum_evictable_length(request, seq_len): @@ -1367,7 +1367,7 @@ def _periodic_evict( protected_tails[request_id] = protected_tail due_requests.append((request, request_id)) - # (2) Compact all affected dense and kernel-masked SWA layers, then release + # Compact all affected dense and kernel-masked SWA layers, then release # the unreachable tail directly through V2's public resize primitive. # Prompt lengths and tails are per-request metadata, so the whole due # cohort runs as one batched round (the workspace holds max_batch_size @@ -1486,7 +1486,7 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: self._request_states.pop(request_id, None) prepared = self._prepared_generation_batch if prepared is not None: - prepared.growth_by_request.pop(request_id, None) + prepared[1].pop(request_id, None) # The workspace stays resident across idle periods: its memory is a # deliberate one-time cost and rebuilding it per burst would reintroduce # allocation on the decode hot path. @@ -1495,18 +1495,6 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: # Helpers (eviction / scoring / V2 cache access / calibration) # # ================================================================== # - # --- Upstream-faithful eviction modes (per_head / per_layer_perhead / union) --- - # - # These reproduce github.com/WeianMao/triattention's selection: scores are NOT - # averaged over heads (each KV head keeps its own token set), they are - # z-normalized per head over the decode region, the prompt (prefill) tokens are - # pinned, and there is no recency window. The kept COUNT stays uniform (= top_B) - # so paged attention - # and the num_cached bookkeeping are unchanged; only the kept SET differs per - # head. Kept K keeps its original RoPE rotation (scored post-RoPE), so a head - # holding a different token set still scores the correct relative distance - # and no per-head position tracking is needed. - def _local_to_global_layers(self, num_layers: int) -> List[int]: """Return V2's global layer id for every local TriAttention layer slot.""" cached = self._local_to_global_layers_cache @@ -1619,30 +1607,29 @@ def _attention_layer_partition( self._attention_layer_partition_cache = result return result - def _runtime_kv_layout(self, num_layers: int) -> _RuntimeKVLayout: + def _runtime_kv_layout(self, num_layers: int) -> Dict[str, object]: """Return stable V2 pool views and layer groups for eviction. KVCacheManagerV2 keeps GPU virtual addresses and layer geometry stable, while opt-in pool rebalance can change the page dimension. Cache all layer views, then query the live page count for one representative per - physical pool before reuse. This avoids rebuilding TensorWrapper views - on every eviction while retaining the same fail-closed rebalance check. + physical pool before reuse (fail-closed rebalance check). """ cached = self._runtime_kv_layout_cache manager = self.kv_cache_manager if cached is not None: - if cached.num_layers != num_layers: + if cached["num_layers"] != num_layers: raise ValueError( - f"TriAttention layer count changed from {cached.num_layers} to {num_layers}" + f"TriAttention layer count changed from {cached['num_layers']} to {num_layers}" ) - if cached.manager is not manager: + if cached["manager"] is not manager: raise RuntimeError("TriAttention target KV cache manager changed at runtime") current_page_counts = self._pool_page_counts( manager, - cached.global_layers, - cached.pool_representatives, + cached["global_layers"], + cached["pool_representatives"], ) - if current_page_counts != cached.pool_page_counts: + if current_page_counts != cached["pool_page_counts"]: raise RuntimeError( "TriAttention V2 pool layout changed after the layout was built; " "KV pool rebalance is not supported" @@ -1675,7 +1662,7 @@ def _build_runtime_kv_layout( swa_window: Optional[int], dense_storage_groups: Optional[Dict[object, List[int]]], what: str, - ) -> _RuntimeKVLayout: + ) -> Dict[str, object]: """Build the manager-lifetime layer and pool views one eviction reads. ``dense_storage_groups`` restricts the compaction groups to the dense @@ -1704,7 +1691,7 @@ def _build_runtime_kv_layout( layer: layers[0] for layers in storage_groups.values() for layer in layers } pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) - return _RuntimeKVLayout( + return dict( manager=manager, num_layers=num_layers, global_layers=global_layers, @@ -1724,7 +1711,7 @@ def _build_runtime_kv_layout( ), ) - def _draft_runtime_kv_layout(self) -> _RuntimeKVLayout: + def _draft_runtime_kv_layout(self) -> Dict[str, object]: """Return stable draft V2 pool views, mirroring ``_runtime_kv_layout``. The draft cache is compacted with the target's kept token set, so its @@ -1736,14 +1723,14 @@ def _draft_runtime_kv_layout(self) -> _RuntimeKVLayout: raise RuntimeError("TriAttention has no draft KV cache manager to lay out") cached = self._draft_runtime_kv_layout_cache if cached is not None: - if cached.manager is not manager: + if cached["manager"] is not manager: raise RuntimeError("TriAttention draft KV cache manager changed at runtime") current_page_counts = self._pool_page_counts( manager, - cached.global_layers, - cached.pool_representatives, + cached["global_layers"], + cached["pool_representatives"], ) - if current_page_counts != cached.pool_page_counts: + if current_page_counts != cached["pool_page_counts"]: raise RuntimeError( "TriAttention draft V2 pool layout changed after the layout " "was built; KV pool rebalance is not supported" @@ -1797,52 +1784,50 @@ def _pool_view_fingerprint(pools: List[torch.Tensor]) -> Tuple[tuple, ...]: for pool in pools ) - def _fixed_resources_for( + def _workspace_for( self, - layout: _RuntimeKVLayout, - prepared: Sequence[_PreparedEviction], - ) -> _EvictionBuffers: - """Return the eviction buffers, building them once at first use. + layout: Dict[str, object], + prepared: Sequence[Dict[str, object]], + ) -> SimpleNamespace: + """Return the eviction workspace, building it once at first use. The request capacity follows the executor's max batch size (memory scales linearly with it) and the decode-width capacity follows the eviction bound (compaction keeps the scored decode region near - ``top_B`` plus one period of growth), so one set of buffers serves - every round. They are rebuilt only when the pool views change or a - round outgrows them. + ``top_B`` plus one period of growth), so one workspace serves every + round. It is rebuilt only when the pool views change or a round + outgrows it. """ if not prepared: raise ValueError("TriAttention eviction requires at least one request") - needed_width = max(item.seq_len - item.prompt_len for item in prepared) - needed_page_tokens = max(item.seq_len + item.protected_tail for item in prepared) + needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) + needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) needed_requests = len(prepared) draft_fingerprint = None if self.draft_kv_cache_manager is not None: draft_layout = self._draft_runtime_kv_layout() draft_fingerprint = ( - draft_layout.pool_page_counts, - draft_layout.pool_view_fingerprint, + draft_layout["pool_page_counts"], + draft_layout["pool_view_fingerprint"], ) fingerprint = ( self.eviction_mode, self.top_B, - tuple(layout.dense_layers), - layout.pool_view_fingerprint, + tuple(layout["dense_layers"]), + layout["pool_view_fingerprint"], draft_fingerprint, ) - resources = self._eviction_resources - if resources is not None: - staging = resources.score_staging + ws = self._workspace + if ws is not None: if ( - self._eviction_pool_fingerprint == fingerprint - and needed_width <= staging.decode_width - and needed_page_tokens <= staging.page_table_token_capacity - and needed_requests <= staging.max_requests + self._workspace_fingerprint == fingerprint + and needed_width <= ws.decode_width + and needed_page_tokens <= ws.page_table_token_capacity + and needed_requests <= ws.max_requests ): - return resources + return ws # Pools changed or this round outgrew the buffers: rebuild. - self._eviction_resources = None - self._batched_compaction = None + self._workspace = None mgr = self.kv_cache_manager tail_capacity = self._configured_protected_tail_capacity() @@ -1855,47 +1840,62 @@ def _fixed_resources_for( # pinning it to max_seq_len: with pinned prompts the post-compaction # length is bounded by prompt + budget + slack, so one power-of-two # bucket serves the steady state, and a cohort that outgrows it simply - # rebuilds these resources through the capacity check above. The old - # max_seq_len floor made the scratch unindexable in 32 bits and tens - # of GiB at large batch (BS=256 x 8K bucket) for work that never - # scored past ~1K tokens per request. + # rebuilds the workspace through the capacity check above. A + # max_seq_len floor would make the scratch unindexable in 32 bits and + # tens of GiB at large batch for work that never scores past ~1K + # tokens per request. seq_capacity = max(int(needed_page_tokens), 1024) seq_capacity = 1 << (seq_capacity - 1).bit_length() seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) - # The CuTe score kernel stores full compute tiles (64 tokens, or one - # page for 128-token pages) into a scratch strided by this bucket - # capacity, so the capacity must be tile-aligned (its geometry gate - # rejects anything else). Rounding up costs at most one tile of - # scratch per segment and never changes scoring semantics. + # The CuTe score kernel stores full compute tiles into a scratch + # strided by this bucket capacity, so the capacity must be + # tile-aligned (the geometry gate rejects anything else). Rounding up + # costs at most one tile of scratch per segment. score_tile_tokens = max(64, int(mgr.tokens_per_block)) seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) - dense_groups = list(layout.storage_groups.values()) + dense_groups = list(layout["storage_groups"].values()) representatives = [group[0] for group in dense_groups] - representatives.extend(layer for layer in layout.swa_layers if layer not in representatives) + representatives.extend( + layer for layer in layout["swa_layers"] if layer not in representatives + ) draft_kwargs = {} if self.draft_kv_cache_manager is not None: draft_layout = self._draft_runtime_kv_layout() draft_tail_capacity = self._draft_protected_tail_capacity() - draft_representatives = list(draft_layout.pool_representatives) + draft_representatives = list(draft_layout["pool_representatives"]) draft_kwargs = dict( - draft_layer_pools=draft_layout.layer_pools, + draft_layer_pools=draft_layout["layer_pools"], + draft_layers=draft_layout["dense_layers"], + draft_layer_group_representative=draft_layout["layer_group_representative"], + draft_layer_pool_keys=list(draft_layout["layer_pool_keys"]), draft_page_representatives=draft_representatives, draft_page_table_keys=[ - draft_layout.layer_pool_keys[layer] for layer in draft_representatives + draft_layout["layer_pool_keys"][layer] for layer in draft_representatives ], draft_num_page_table_slots=self.draft_kv_cache_manager.num_pools, draft_page_table_token_capacity=seq_capacity + draft_tail_capacity, + draft_protected_tail_capacity=draft_tail_capacity, ) - first_pool = layout.layer_pools[layout.dense_layers[0]] + # One-time shape gate on the V2-allocated host block-offset tables the + # bulk staging reads every round (int32 [pools, slots, K/V, blocks]). + checked_managers = [("", mgr)] + if self.draft_kv_cache_manager is not None: + checked_managers.append(("draft ", self.draft_kv_cache_manager)) + for what, manager in checked_managers: + table = manager.host_kv_cache_block_offsets + if table.dtype != torch.int32 or table.ndim != 4 or table.shape[2] != 2: + raise RuntimeError( + f"KVCacheManagerV2 exposes an invalid {what}host block-offset table" + ) + + first_pool = layout["layer_pools"][layout["dense_layers"][0]] if self._offsets is None: self._offsets = _build_geometric_offsets(_OFFSET_MAX_LENGTH, first_pool.device) - if self._mean_phase_table is None: - from .triattention_kernels import MeanPhaseTable - - self._mean_phase_table = MeanPhaseTable( + if self._phase is None: + self._phase = build_mean_phase_table( self._offsets, self.calibration["omega"] .to(device=first_pool.device, dtype=torch.float32) @@ -1903,133 +1903,45 @@ def _fixed_resources_for( initial_rows=seq_capacity, ) q_real, q_imag, mlr_coef = self._local_score_calibration( - layout.num_layers, layout.global_layers + layout["num_layers"], layout["global_layers"] ) - score_staging = _FixedScoreStagingBuffers( - layout.layer_pools, + ws = prepare_eviction_workspace( + eviction_mode=self.eviction_mode, + layer_pools=layout["layer_pools"], dense_groups=dense_groups, - dense_layers=layout.dense_layers, + dense_layers=layout["dense_layers"], + swa_layers=layout["swa_layers"], + swa_window=layout["swa_window"], + layer_group_representative=layout["layer_group_representative"], + layer_pool_keys=list(layout["layer_pool_keys"]), page_representatives=representatives, max_requests=request_capacity, seq_len=seq_capacity, num_q_heads=int(self._H), num_freqs=int(self._F), + keep_count=self.top_B, q_real=q_real, q_imag=q_imag, mlr_coef=mlr_coef, freq_scale_sq=self._freq_scale_sq, offsets=self._offsets, omega=self.calibration["omega"], - mean_phase_table=self._mean_phase_table, - page_table_keys=self._page_table_pool_keys(representatives, layout.global_layers), - num_page_table_slots=layout.manager.num_pools, + phase=self._phase, + page_table_keys=self._page_table_pool_keys(representatives, layout["global_layers"]), + num_page_table_slots=layout["manager"].num_pools, decode_width=decode_width, page_table_token_capacity=page_table_token_capacity, + protected_tail_capacity=tail_capacity, **draft_kwargs, ) - keep_set_selector = _BatchedKeepSetSelector( - eviction_mode=self.eviction_mode, - dense_layers=tuple(layout.dense_layers), - num_query_heads=int(self._H), - num_kv_heads=int(first_pool.shape[2]), - rows=len(layout.dense_layers) * int(self._H), - width=decode_width, - keep_count=self.top_B, - dtype=torch.float32, - device=first_pool.device, - max_requests=request_capacity, - prompt_offsets_buffer=score_staging.token_starts_device, - ) - # Padded rows carry zero valid width; their provisional TopK entries - # must still be in-range ordinals for the finalizer's score gather. - if self.eviction_mode == "union": - keep_set_selector.final_indices.zero_() - else: - keep_set_selector.top_indices_i32.zero_() - score_staging.bind_score_launcher( - keep_set_selector.valid_widths, - "mean", - ) - resources = _EvictionBuffers( - score_staging=score_staging, - keep_set_selector=keep_set_selector, - ) - self._eviction_resources = resources - self._eviction_pool_fingerprint = fingerprint - return resources - - def _batched_compaction_for( - self, - *, - layout: _RuntimeKVLayout, - prepared: Sequence[_PreparedEviction], - score_staging: _FixedScoreStagingBuffers, - keep_set_selector: _BatchedKeepSetSelector, - ): - """Build or reuse the C++ compaction launches for one cohort.""" - from .compaction import BatchedKVCacheCompaction - - if layout.swa_layers and layout.swa_window: - # SWA landing positions are prompt-dependent; reject a request - # whose retained span cannot cover the model window this round. - for item in prepared: - if item.prompt_len + self.top_B < int(layout.swa_window): - raise ValueError( - f"Request {item.request_id} retains " - f"{item.prompt_len + self.top_B} tokens, below the " - f"sliding window {layout.swa_window}" - ) - batched_compaction = self._batched_compaction - if batched_compaction is None: - draft_kwargs = {} - if self.draft_kv_cache_manager is not None: - draft_layout = self._draft_runtime_kv_layout() - draft_kwargs = dict( - draft_layer_pools=draft_layout.layer_pools, - draft_layers=draft_layout.dense_layers, - draft_layer_group_representative=draft_layout.layer_group_representative, - draft_layer_pool_keys=list(draft_layout.layer_pool_keys), - draft_protected_tail_capacity=self._draft_protected_tail_capacity(), - draft_kv_block_offsets=score_staging.draft_block_offsets_device, - draft_page_table_slots=score_staging.draft_representative_slots, - ) - batched_compaction = BatchedKVCacheCompaction( - eviction_mode=self.eviction_mode, - layer_pools=layout.layer_pools, - dense_layers=layout.dense_layers, - swa_layers=layout.swa_layers, - layer_group_representative=layout.layer_group_representative, - layer_pool_keys=list(layout.layer_pool_keys), - kept_token_ordinals=keep_set_selector.keep, - valid_sequence_lengths=score_staging.valid_seq_lens_device, - kv_block_offsets=score_staging.block_offsets_device, - page_table_slots=score_staging.representative_slots, - request_count=score_staging.max_requests, - prompt_offsets=score_staging.token_starts_device, - decode_keep_count=self.top_B, - swa_window=layout.swa_window, - protected_tail_capacity=self._configured_protected_tail_capacity(), - dense_move_offsets=score_staging.dense_move_offsets, - swa_move_offsets=score_staging.swa_move_offsets, - draft_move_offsets=score_staging.draft_move_offsets, - **draft_kwargs, - ) - # One launch settles the kept ordinals and packs the dense/SWA - # move sources; the compaction keeps only its C++ moves (plus the - # draft's own pack). Both caches are invalidated together, so the - # fused packing always points at the live compaction buffers. - keep_set_selector.fuse_move_source_pack( - batched_compaction.hand_move_source_pack_to_selection() - ) - self._batched_compaction = batched_compaction - # Tails vary per round (in-flight growth), so the per-family move - # offsets ride the staged metadata table each round. - return batched_compaction + self._workspace = ws + self._workspace_fingerprint = fingerprint + return ws def _move_offsets_for( self, - layout: _RuntimeKVLayout, - prepared: Sequence[_PreparedEviction], + layout: Dict[str, object], + prepared: Sequence[Dict[str, object]], capacity: int, ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: """Build this round's per-family move offsets, padded to the capacity. @@ -2045,11 +1957,11 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets.extend(offsets[-1:] * (capacity - len(moves_per_request))) return offsets - tails = [item.protected_tail for item in prepared] + tails = [item["protected_tail"] for item in prepared] dense = padded_offsets([self.top_B + tail for tail in tails]) swa = None - if layout.swa_layers and layout.swa_window: - swa = padded_offsets([int(layout.swa_window) + tail for tail in tails]) + if layout["swa_layers"] and layout["swa_window"]: + swa = padded_offsets([int(layout["swa_window"]) + tail for tail in tails]) draft = None if self.draft_kv_cache_manager is not None: draft_tail = self._draft_protected_tail_capacity() @@ -2089,35 +2001,6 @@ def _dense_layer_pool_groups( groups.setdefault(pool_key, []).append(layer) return groups - def _attach_page_ids( - self, - prepared: Sequence[_PreparedEviction], - staging: _FixedScoreStagingBuffers, - layout: _RuntimeKVLayout, - ) -> None: - dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( - layout, prepared, staging.max_requests - ) - try: - staged = staging.stage( - self.kv_cache_manager, - [item.request_id for item in prepared], - [item.round_start for item in prepared], - [item.prompt_len for item in prepared], - [item.seq_len for item in prepared], - [item.seq_len + item.protected_tail for item in prepared], - draft_manager=self.draft_kv_cache_manager, - dense_move_offsets=dense_offsets, - swa_move_offsets=swa_offsets, - draft_move_offsets=draft_offsets, - ) - except _FixedScoreStreamMismatch: - raise - except Exception as exc: - raise RuntimeError("TriAttention score staging failed") from exc - if not staged: - raise RuntimeError("TriAttention page-table staging rejected the cohort") - def _evict_requests( self, evict_reqs, @@ -2137,20 +2020,19 @@ def _evict_requests( layout = self._runtime_kv_layout(num_layers) # Resolve request length and page metadata before mutating any layer. - prepared: List[_PreparedEviction] = [] + prepared: List[Dict[str, object]] = [] with nvtx_range("triattention.metadata", color="cyan"): for request, rid in evict_reqs: request_state = self._request_states.get(rid) - seq_len = None if request_state is None else request_state.confirmed_kv_length + seq_len = None if request_state is None else request_state["confirmed_kv_length"] if seq_len is None: raise RuntimeError(f"Missing confirmed KV length for request {rid}") # Restore the uncompressed confirmed logical position from the # physical prefix and cumulative eviction count. - round_start = seq_len + request_state.evicted_tokens + round_start = seq_len + request_state["evicted_tokens"] minimum_evictable_length = self._minimum_evictable_length(request, seq_len) if seq_len <= minimum_evictable_length: continue - expected_keep_count = minimum_evictable_length protected_tail = int(protected_tail_lengths.get(rid, 0)) if protected_tail < 0 or protected_tail > protected_tail_capacity: raise RuntimeError( @@ -2158,78 +2040,71 @@ def _evict_requests( f"configured capacity {protected_tail_capacity}" ) prepared.append( - _PreparedEviction( - request=request, - request_id=rid, - seq_len=int(seq_len), - round_start=int(round_start), - prompt_len=min(int(request.py_prompt_len), int(seq_len)), - expected_keep_count=expected_keep_count, - protected_tail=protected_tail, - ) + { + "request": request, + "request_id": rid, + "seq_len": int(seq_len), + "round_start": int(round_start), + "prompt_len": min(int(request.py_prompt_len), int(seq_len)), + "expected_keep_count": minimum_evictable_length, + "protected_tail": protected_tail, + } ) if not prepared: return [] with nvtx_range_debug("triattention.staging_lookup", color="blue"): - resources = self._fixed_resources_for(layout, prepared) - score_staging = resources.score_staging - keep_set_selector = resources.keep_set_selector - batched_compaction = self._batched_compaction_for( - layout=layout, - prepared=prepared, - score_staging=score_staging, - keep_set_selector=keep_set_selector, - ) + ws = self._workspace_for(layout, prepared) + if layout["swa_layers"] and layout["swa_window"]: + # SWA landing positions are prompt-dependent; reject a request + # whose retained span cannot cover the model window this round. + for item in prepared: + if item["prompt_len"] + self.top_B < int(layout["swa_window"]): + raise ValueError( + f"Request {item['request_id']} retains " + f"{item['prompt_len'] + self.top_B} tokens, below the " + f"sliding window {layout['swa_window']}" + ) with nvtx_range_debug("triattention.page_table_stage", color="orange"): - self._attach_page_ids(prepared, score_staging, layout) - # The staged per-request prompt lengths are shared with the - # selector; per-head modes re-expand them to selection rows here. - keep_set_selector.refresh_row_prompt_offsets() + dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( + layout, prepared, ws.max_requests + ) + stage_eviction_cohort( + ws, + self.kv_cache_manager, + [item["request_id"] for item in prepared], + [item["round_start"] for item in prepared], + [item["prompt_len"] for item in prepared], + [item["seq_len"] for item in prepared], + [item["seq_len"] + item["protected_tail"] for item in prepared], + draft_manager=self.draft_kv_cache_manager, + dense_move_offsets=dense_offsets, + swa_move_offsets=swa_offsets, + draft_move_offsets=draft_offsets, + ) try: - fused_union = keep_set_selector.eviction_mode == "union" - with nvtx_range("triattention.score", color="blue"): - # Union rounds run the fused pipeline (score, row stats, - # normalization, and the cross-row union maximum in two CuTe - # kernels) — THE only union path; it raises loudly if it - # cannot serve this cohort. The per-head modes run the split - # score launch and their own reduction kernels. - if fused_union: - score_staging.launch_prepared_union_fusion(keep_set_selector.combined) - per_head = None - else: - per_head = score_staging.launch_prepared_score() - with nvtx_range("triattention.select", color="yellow"): - if fused_union: - keep_set_selector.select_prepared_union_scores() - else: - keep_set_selector.select_requests( - per_head, - normalize_scores=self.normalize_scores, - ) - with nvtx_range("triattention.compact", color="purple"): - batched_compaction.compact() + run_eviction_round(ws, self.normalize_scores) finally: consumer_streams = [self.kv_cache_manager._stream] if self.draft_kv_cache_manager is not None: consumer_streams.append(self.draft_kv_cache_manager._stream) - score_staging.mark_page_tables_consumed(*consumer_streams) + mark_page_tables_consumed(ws, *consumer_streams) capacity_targets = [] for item in prepared: - keep_count = item.expected_keep_count - evicted = item.seq_len - keep_count + keep_count = item["expected_keep_count"] + evicted = item["seq_len"] - keep_count if evicted <= 0: raise RuntimeError("TriAttention attempted an identity compaction") - request_state = self._request_states[item.request_id] - request_state.evicted_tokens += evicted - request_state.confirmed_kv_length = keep_count + request_state = self._request_states[item["request_id"]] + request_state["evicted_tokens"] += evicted + request_state["confirmed_kv_length"] = keep_count # Publish the cumulative count on the request: this is the # manager's only channel to the runtime. The model engine # reads it back where it builds num_cached_tokens_per_seq, # so the kernels see the compacted KV length next step. - item.request.py_num_compressed_tokens = request_state.evicted_tokens - capacity_targets.append((item.request_id, keep_count)) + item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] + capacity_targets.append((item["request_id"], keep_count)) return capacity_targets def _num_layers_from_manager(self) -> int: 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 index 3bf397b73282..a5fa066b93b4 100644 --- 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 @@ -7,7 +7,7 @@ uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, and producer-only page-ID lookahead. Geometries outside the exact contract validated here raise loudly at setup -(``_FixedScoreGroup.prepare_cute_score``); there is no fallback path. +(``triattention.prepare_eviction_workspace``); there is no fallback path. """ from __future__ import annotations diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 5ae607986cdf..69e0c5185053 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -2,18 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 """GPU kernels for the TriAttention KV-eviction pipeline. -The production path uses one fixed-shape trig-score launch across all dense -layers, CuTE-DSL TopK selection, and grouped C++ compaction. Scoring runs -EXCLUSIVELY through the SM100 CuTe-DSL fused score pack +Scoring runs EXCLUSIVELY through the SM100 CuTe-DSL fused score pack (``triattention_cute_score_fused.py``): mean aggregation, BF16 KV pools, head size 64/128, 32/128-token pages, GQA group 4 or 8, per-request score -window starts. There is deliberately no other score path -- any geometry -outside that contract raises loudly at setup. The per-head modes use the -pack's score-only entry plus the row-stats kernel for normalization; union -eviction runs the fused score+stats+union pipeline (with -``triattention_cute_selection.py``). The unit tests validate the CuTe -kernels against independent PyTorch oracles. Selection and compaction live -in their respective runtime modules. +window starts. Any geometry outside that contract raises loudly at +workspace construction. The per-head modes use the pack's score-only entry +plus the row-stats kernel; union eviction runs the fused +score+stats+union pipeline. One-time buffer staging and runner compilation +live in ``triattention.prepare_eviction_workspace``; this module keeps the +Triton kernels, their launch helpers, and the mean-phase table builders. House rules honored throughout: * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). @@ -24,14 +21,14 @@ from __future__ import annotations -from typing import List, Optional +from typing import Dict import torch import triton import triton.language as tl # --------------------------------------------------------------------------- # -# Scoring: trig-score every cached token across all dense layers. # +# Mean-phase table: RoPE-style position table of mean trig phases. # # --------------------------------------------------------------------------- # @@ -67,599 +64,94 @@ def _gather_mean_phase_kernel( tl.store(mean_sin + output_offset, row_sin, mask=frequency_mask) -class MeanPhaseTable: - """RoPE-style position table of mean trig phases, gathered per round. +def build_mean_phase_table( + offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int +) -> Dict[str, object]: + """Build the plain-dict mean-phase table shared by every workspace. Row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` over the calibration offsets for every frequency, so refreshing a round's - ``mean_cos``/``mean_sin`` is one pure-gather launch over the staged - round starts instead of a per-round trig kernel. The gather writes in - place because the compiled CuTe score launch captured the destination - buffers' device pointers. Eviction never runs under CUDA graph - capture, so the table itself may regrow; callers must ``ensure`` - capacity while the round starts are still host integers (the gather - clamps stale rows into the table rather than faulting). Building the - table is plain torch and works on any device; gathering launches the - Triton kernel and is CUDA-only. + ``mean_cos``/``mean_sin`` is one pure-gather launch over the staged round + starts. The dict is shared BY REFERENCE between the manager and its + workspace, so ``grow_mean_phase_table`` reaches both. Grow the table while + the round starts are still host integers; the gather kernel clamps stale + rows instead of faulting. """ + if ( + offsets.numel() <= 0 + or omega.numel() <= 0 + or offsets.dtype != torch.float32 + or omega.dtype != torch.float32 + or offsets.device != omega.device + ): + raise ValueError("mean-phase tables require same-device FP32 offsets and frequencies") + phase: Dict[str, object] = { + "offsets": offsets.contiguous(), + "omega": omega.contiguous(), + "offset_values": offsets.tolist(), + "cos": None, + "sin": None, + "rows": 0, + } + grow_mean_phase_table(phase, max(int(initial_rows), 1)) + return phase + + +def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: + """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" + rows = int(rows) + if rows <= phase["rows"]: + return + if rows > _MEAN_PHASE_MAX_ROWS: + raise ValueError(f"a {rows}-row mean-phase table exceeds the exact-FP32 position range") + target = 1 + while target < rows: + target *= 2 + target = min(max(target, 2 * phase["rows"]), _MEAN_PHASE_MAX_ROWS) + omega = phase["omega"] + positions = torch.arange(target, device=omega.device, dtype=torch.float32) + cos_table = torch.zeros((target, omega.numel()), dtype=torch.float32, device=omega.device) + sin_table = torch.zeros_like(cos_table) + # Accumulate offset-by-offset in fp32 (fixed summation order keeps the + # table bit-stable across rebuilds). + for offset in phase["offset_values"]: + angle = torch.outer(positions + offset, omega) + cos_table += torch.cos(angle) + sin_table += torch.sin(angle) + scale = 1.0 / len(phase["offset_values"]) + phase["cos"] = cos_table.mul_(scale) + phase["sin"] = sin_table.mul_(scale) + phase["rows"] = target + + +def gather_mean_phases( + phase: Dict[str, object], + round_starts: torch.Tensor, + mean_cos: torch.Tensor, + mean_sin: torch.Tensor, + request_count: int, +) -> None: + """Refresh the fixed mean buffers in place from staged round starts. - def __init__(self, offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int) -> None: - if ( - offsets.numel() <= 0 - or omega.numel() <= 0 - or offsets.dtype != torch.float32 - or omega.dtype != torch.float32 - or offsets.device != omega.device - ): - raise ValueError("mean-phase tables require same-device FP32 offsets and frequencies") - self.offsets = offsets.contiguous() - self.omega = omega.contiguous() - self._offset_values: List[float] = self.offsets.tolist() - self._cos: Optional[torch.Tensor] = None - self._sin: Optional[torch.Tensor] = None - self._rows = 0 - self.ensure(max(int(initial_rows), 1)) - - @property - def rows(self) -> int: - return self._rows - - def ensure(self, rows: int) -> None: - """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" - 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 = 1 - while target < rows: - target *= 2 - 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.omega.numel()), dtype=torch.float32, device=self.omega.device - ) - sin_table = torch.zeros_like(cos_table) - # Accumulate offset-by-offset in fp32 (fixed summation order keeps - # the table bit-stable across rebuilds). - for offset in self._offset_values: - phase = torch.outer(positions + offset, self.omega) - cos_table += torch.cos(phase) - sin_table += torch.sin(phase) - scale = 1.0 / len(self._offset_values) - self._cos = cos_table.mul_(scale) - self._sin = sin_table.mul_(scale) - self._rows = target - - def gather( - self, - round_starts: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - request_count: int, - ) -> None: - """Refresh the fixed mean buffers in place from staged round starts.""" - request_count = int(request_count) - if request_count <= 0 or request_count > round_starts.numel(): - raise ValueError("phase gather request count is outside its fixed buffers") - num_freqs = self.omega.numel() - if ( - mean_cos.ndim != 2 - or mean_cos.shape[0] < request_count - or mean_cos.shape[1] != num_freqs - or mean_sin.shape != mean_cos.shape - or round_starts.dtype != torch.int32 - or self.omega.device.type != "cuda" - or any( - tensor.device != self.omega.device for tensor in (round_starts, mean_cos, mean_sin) - ) - or any(tensor.dtype != torch.float32 for tensor in (mean_cos, mean_sin)) - ): - raise ValueError("phase gather tensors do not share one valid FP32 CUDA geometry") - _gather_mean_phase_kernel[(request_count,)]( - round_starts, - self._cos, - self._sin, - mean_cos, - mean_sin, - self._rows, - NUM_FREQS=num_freqs, - F_BLOCK=triton.next_power_of_2(num_freqs), - num_warps=1, - ) - - -class _FixedScoreGroup: - """Persistent score metadata/output for one fixed geometry. - - Since the per-layer absolute-address ABI, ONE group can span dense layers - living in DISTINCT storages with DISTINCT block tables. ``block_offsets`` - uses the native TRT-LLM attention layout and ``page_table_slots`` maps each - scored layer to its V2 pool slot. - - LIFETIME: the group retains references to every scored layer pool -- the - SM100 CuTe score kernel encodes immutable TMA descriptors from their raw - device addresses at compile time, so the pools must stay alive (and stay - put) for as long as the group launches. In production the V2 KV-cache - manager owns them for the manager's lifetime. + Writes in place because the compiled CuTe score launch captured the + destination buffers' device pointers. CUDA-only; eviction never runs + under CUDA graph capture. """ - - def __init__( - self, - layer_pools: List[torch.Tensor], - layer_indices: List[int], - max_requests: int, - page_count: int, - seq_len: int, - num_q_heads: int, - block_offsets: torch.Tensor, # [num_pools, max_requests, 2, copied_blocks] int32 - page_table_slots: List[int], # per scored layer: pool slot into block_offsets - q_real_LHF: torch.Tensor, - q_imag_LHF: torch.Tensor, - mlr_coef_LHF: torch.Tensor, - freq_scale_sq: torch.Tensor, - omega: torch.Tensor, - offsets: torch.Tensor, - output_width: int, - ) -> None: - if not layer_indices or min(max_requests, page_count, seq_len) <= 0: - raise ValueError("fixed score group requires non-empty positive geometry") - if output_width <= 0 or output_width > seq_len: - raise ValueError("fixed score group requires a decode width within its capacity") - if len(page_table_slots) != len(layer_indices): - raise ValueError("page_table_slots must align with layer_indices") - self.max_requests = max_requests - # Prompt lengths are per-request kernel inputs; this capacity only - # sizes the widest possible decode window of the output buffer. - self.output_width = int(output_width) - self.num_layers = len(layer_indices) - p0 = layer_pools[layer_indices[0]] - if p0.ndim != 5: - raise ValueError("fixed score group requires HND pools") - device = p0.device - q_real_LHF = q_real_LHF.to(device=device, dtype=torch.float32).contiguous() - q_imag_LHF = q_imag_LHF.to(device=device, dtype=torch.float32).contiguous() - mlr_coef_LHF = mlr_coef_LHF.to(device=device, dtype=torch.float32).contiguous() - freq_scale_sq = freq_scale_sq.to(device=device, dtype=torch.float32).contiguous() - omega = omega.to(device=device, dtype=torch.float32).contiguous() - offsets = offsets.to(device=device, dtype=torch.float32).contiguous() - _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape - if num_q_heads % num_kv_heads: - raise ValueError("query heads must be divisible by KV heads") - self.num_freqs = head_dim // 2 - strides = tuple(int(value) for value in p0.stride()) - self.geometry_args = ( - num_q_heads, - num_kv_heads, - self.num_freqs, - tokens_per_block, - kv_factor, - ) - # Per-layer ABSOLUTE base addresses. Layers may live in distinct - # storages (V2 TensorWrapper-per-layer); only geometry must be uniform. - element_size = p0.element_size() - layer_base_addrs = torch.zeros(len(layer_pools), dtype=torch.int64, device=device) - for layer in layer_indices: - pool = layer_pools[layer] - if ( - tuple(pool.shape[1:]) != tuple(p0.shape[1:]) - or tuple(pool.stride()) != strides - or pool.dtype != p0.dtype - ): - raise ValueError("fixed score layers must share one uniform geometry") - address = int(pool.data_ptr()) - if address % element_size: - raise ValueError("fixed score layer base is not element-aligned") - layer_base_addrs[layer] = address - # Calibration tables span every model layer; segments index them by - # ABSOLUTE layer id, so the tables cover the full calibrated extent. - self._num_calibrated_layers = q_real_LHF.numel() // (int(num_q_heads) * self.num_freqs) - # Segment layer ids index the calibration tables ON DEVICE where they - # cannot be range-checked; validate the extent once here, loudly. - if min(layer_indices) < 0 or max(layer_indices) >= self._num_calibrated_layers: - raise ValueError("scored layer index exceeds the calibrated layer extent") - seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( - self.num_layers - ) - seg_layer = torch.tensor(layer_indices, dtype=torch.int32, device=device).repeat( - max_requests - ) - # Each segment reads the K plane for one request from the same native - # block-offset buffer used by TRT-LLM attention metadata preparation. - if ( - block_offsets.ndim != 4 - or tuple(block_offsets.shape[1:3]) != (max_requests, 2) - or block_offsets.shape[3] < page_count - or block_offsets.dtype != torch.int32 - or block_offsets.device != device - ): - raise ValueError("block offsets do not match fixed score geometry") - if not block_offsets.is_contiguous(): - raise ValueError("fixed score block offsets must be contiguous") - slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) - if int(slots_t.max()) >= int(block_offsets.shape[0]): - raise ValueError("page table slot exceeds staged page-id planes") - req_idx = torch.arange(max_requests, dtype=torch.int64, device=device).repeat_interleave( - self.num_layers - ) - slot_idx = slots_t.repeat(max_requests) - seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - self.output = torch.empty( - max_requests, - self.num_layers, - num_q_heads, - self.output_width, - dtype=torch.float32, - device=device, - ) - self.pointer_prefix = ( - p0, - layer_base_addrs, - block_offsets.view(-1), - seg_page_off, - seg_req, - seg_layer, - ) - self.pointer_middle = ( - q_real_LHF.view(-1), - q_imag_LHF.view(-1), - mlr_coef_LHF.view(-1), - ) - self.pointer_tail = (freq_scale_sq, omega, offsets) - # The SM100 CuTe fused score pack (see - # triattention_cute_score_fused.py) is THE score implementation; its - # score-only entry is compiled by the first ``prepare_cute_score`` - # call. The runner encodes TMA descriptors from the actual pool - # tensors, hence the pool references retained here (see the LIFETIME - # note in the class docstring). - self.seq_len = int(seq_len) - self._cute_score_runner = None - self._cute_score_attempted = False - self._cute_layer_pools = list(layer_pools) - self._cute_layer_indices = [int(layer) for layer in layer_indices] - - def prepare_cute_score(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor) -> None: - """Compile the fused CuTe runner's score-only entry; raise loudly otherwise. - - Call this outside CUDA graph capture: compilation allocates memory - and synchronizes. The fused score pack is the ONLY score - implementation, so an unsupported geometry raises ValueError here - and a runner construction failure raises RuntimeError -- there is - deliberately no fallback path. - - Supported contract: SM100 exactly, BF16 pools, 32- or 128-token - pages, 32 or 64 frequencies (head size 64/128), 4 or 8 query heads - per KV head, and a bucket capacity (``seq_len``) aligned to the - historical score tile — this covers the Qwen3 and GPT-OSS production - geometries as well as the original validation shape. - """ - if self._cute_score_attempted: - return - self._cute_score_attempted = True - anchor = self.pointer_prefix[0] - num_q_heads, num_kv_heads, num_freqs, tokens_per_block, kv_factor = self.geometry_args - max_segments = self.max_requests * self.num_layers - # The fused kernel masks its ragged tail, but the scratch bucket - # contract stays tile-aligned (64 tokens, or one page for - # 128-token pages) so bucket geometry cannot - # silently admit new geometry (production pow2 buckets satisfy it). - score_tile_tokens = max(64, int(tokens_per_block)) - supported = ( - torch.cuda.get_device_capability(anchor.device) == (10, 0) - and anchor.dtype == torch.bfloat16 - and kv_factor == 2 - and tokens_per_block in (32, 128) - and num_freqs in (32, 64) - and num_q_heads % num_kv_heads == 0 - and num_q_heads // num_kv_heads in (4, 8) - and int(anchor.stride(-1)) == 1 - and self.seq_len % score_tile_tokens == 0 - # The kernel's head-plane base offset is 64-bit; the widest - # 32-bit product left is one plane (N-1 head columns of one - # segment stride), which the score bucket keeps far below 2^31. - # Group-4 geometries pad the head axis to the MMA tile N=8. - and (8 - 1) * max_segments * self.seq_len < 2**31 - ) - if not supported: - raise ValueError( - "TriAttention score requires SM100, bf16 KV pools, head size " - "64/128, 32/128-token pages, GQA group 4 or 8, and a bucket " - "capacity aligned to the score compute tile; got " - f"capability={torch.cuda.get_device_capability(anchor.device)}, " - f"dtype={anchor.dtype}, kv_factor={kv_factor}, " - f"tokens_per_block={tokens_per_block}, num_freqs={num_freqs}, " - f"heads={num_q_heads}q/{num_kv_heads}kv, " - f"stride={int(anchor.stride(-1))}, " - f"seq_len={self.seq_len} (tile {score_tile_tokens}), " - f"offset_audit={num_kv_heads * 8 * max_segments * self.seq_len}" - ) - device = anchor.device - try: - from .triattention_cute_score_fused import TriAttentionCuteScoreRunner - - # The kernel scores each request's window (from its staged - # per-request start) into its own head-major scratch (row = - # query head, column = segment * seq_len + token); ``launch`` - # gathers each request's decode window from that scratch into - # ``self.output``. All buffers below are persistent because the - # compiled kernel captures their device pointers. - # The kernel writes one scratch row per padded head column - # (GQA group below 8 pads up to the MMA tile); the gather in - # ``launch`` reads only the real heads. - scratch = torch.empty( - num_kv_heads * 8 * max_segments * self.seq_len, - dtype=torch.float32, - device=device, - ) - seg_seq_len = torch.zeros(max_segments, dtype=torch.int32, device=device) - seg_out_offset = ( - torch.arange(max_segments, dtype=torch.int64, device=device) * self.seq_len - ).to(torch.int32) - # Per-request score window starts, staged before each launch; - # the compiled kernels capture this buffer's device pointer. - # The union fusion runner shares it (and the segment buffers). - token_starts = torch.zeros(self.max_requests, dtype=torch.int32, device=device) - gather_columns = torch.arange(self.output_width, dtype=torch.int64, device=device) - self._cute_score_runner = TriAttentionCuteScoreRunner( - layer_pools=self._cute_layer_pools, - layer_indices=self._cute_layer_indices, - max_requests=self.max_requests, - num_layers=self.num_layers, - seq_len=self.seq_len, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - page_ids=self.pointer_prefix[2], - seg_page_off=self.pointer_prefix[3], - seg_req_id=self.pointer_prefix[4], - seg_layer_id=self.pointer_prefix[5], - seg_seq_len=seg_seq_len, - seg_out_offset=seg_out_offset, - token_starts=token_starts, - q_real=self.pointer_middle[0], - q_imag=self.pointer_middle[1], - mlr_coef=self.pointer_middle[2], - mean_cos=mean_cos, - mean_sin=mean_sin, - freq_scale_sq=self.pointer_tail[0], - output=scratch, - # Score-only mode: the stats and union-finalize kernels are - # compiled lazily by ``_union_fusion_runner`` when (and only - # when) union eviction actually launches. - enable_partial_stats=False, - ) - except (ImportError, RuntimeError, ValueError, AssertionError) as error: - raise RuntimeError( - "TriAttention CuTe score setup failed and no other score path exists" - ) from error - self._cute_scratch = scratch - self._cute_seg_seq_len = seg_seq_len - self._cute_seg_out_offset = seg_out_offset - self._cute_token_starts = token_starts - self._cute_gather_columns = gather_columns.view(1, 1, 1, -1) - # Fused score+stats+union pipeline (two launches): - # THE union path, built lazily on the first union launch. ONE runner - # serves every cohort: the score window start is per-request runtime - # metadata, not a compile-time constant. - self._union_fusion_runner_entry = None - from tensorrt_llm.logger import logger - - logger.info( - f"TriAttention CuTe score enabled: {num_q_heads}q/{num_kv_heads}kv heads, " - f"{num_freqs} freqs, {tokens_per_block}-token pages" - ) - - def launch( - self, - request_count: int, - valid_seq_lens: torch.Tensor, - valid_widths: torch.Tensor, - token_starts_device: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - aggregation: str = "mean", - ) -> torch.Tensor: - """Return decode-only scores as ``[request, layer, head, token]``. - - Runs the fused CuTe runner's score-only entry (the only score - implementation) and writes each request's decode width - (``valid_seq_len - token_start``) into ``valid_widths``, which the - selection reduce kernels consume. Each request scores its own window - from its staged start, so one cohort may mix prompt lengths. Only - mean aggregation exists; the runner dispatches every request count - up to the group capacity. - """ - if aggregation != "mean": - raise ValueError( - f"unsupported score aggregation {aggregation!r}: max aggregation " - "was removed with the C++ score stack; only 'mean' exists" - ) - if request_count <= 0 or request_count > self.max_requests: - raise ValueError("request count exceeds fixed score capacity") - if ( - valid_widths.ndim != 1 - or valid_widths.numel() < request_count - or valid_widths.dtype != torch.int32 - or valid_widths.device != self.output.device - ): - raise ValueError("score output lengths do not fit the keep-set selector") - num_segments = request_count * self.num_layers - output = self.output[:request_count] - # Lazy compile covers groups used without their owning workspace - # (unit tests); production compiles in the workspace constructor, - # outside CUDA graph capture. - self.prepare_cute_score(mean_cos, mean_sin) - runner = self._cute_score_runner - if runner is None or not runner.supports(request_count): - raise RuntimeError( - f"TriAttention CuTe score has no compiled variant for " - f"request_count={request_count} (capacity {self.max_requests}) " - "and no other score path exists" - ) - # Per-request decode widths (seq_len - token_start) for the - # selection reduce kernels. - torch.sub( - valid_seq_lens[:request_count], - token_starts_device[:request_count], - out=valid_widths[:request_count], - ) - # Stage per-segment valid lengths (segment = request x layer). - torch.index_select( - valid_seq_lens, - 0, - self.pointer_prefix[4][:num_segments], - out=self._cute_seg_seq_len[:num_segments], - ) - # Stage the per-request score window starts: the compiled kernel - # captured this buffer's pointer and reads one start per request. - self._cute_token_starts[:request_count].copy_(token_starts_device[:request_count]) - runner.launch(request_count, mean_cos, mean_sin) - # The kernel wrote each request's window scores (from its pinned - # prompt length) into its head-major scratch. Gather each request's - # decode window into the group output, the ``[request, layer, head, - # token]`` layout the selection kernels read. Columns past a - # request's valid width carry unscored scratch data; consumers mask - # by ``valid_widths``. - num_q_heads = int(self.geometry_args[0]) - num_kv_heads = int(self.geometry_args[1]) - group_size = num_q_heads // num_kv_heads - # The scratch head axis is padded to the MMA tile N=8 per KV head; - # slicing the view to the real group size skips the zero padding - # columns. - source = ( - self._cute_scratch[: num_kv_heads * 8 * num_segments * self.seq_len] - .view(num_kv_heads, 8, request_count, self.num_layers, self.seq_len)[:, :group_size] - .permute(2, 3, 0, 1, 4) - ) - columns = token_starts_device[:request_count].to(torch.int64).view( - -1, 1, 1, 1, 1 - ) + self._cute_gather_columns.view(1, 1, 1, 1, -1) - columns = columns.clamp_(max=self.seq_len - 1).expand( - request_count, - self.num_layers, - num_kv_heads, - group_size, - self.output_width, - ) - torch.gather( - source, - 4, - columns, - out=output.view( - request_count, - self.num_layers, - num_kv_heads, - group_size, - self.output_width, - ), - ) - return output - - def _union_fusion_runner(self, mean_cos: torch.Tensor, mean_sin: torch.Tensor): - """Build (or reuse) the ONE fused score/stats/union runner. - - The score window start is per-request runtime metadata staged into a - persistent device buffer, so a single compiled runner serves every - cohort. This is THE union path: a construction failure raises loudly - instead of recording a fallback. - """ - if self._union_fusion_runner_entry is not None: - return self._union_fusion_runner_entry - num_q_heads, num_kv_heads, num_freqs, tokens_per_block, _ = self.geometry_args - try: - from .triattention_cute_score_fused import ( - TriAttentionCuteScoreRunner as _FusedUnionScoreRunner, - ) - - device = self.output.device - # The union output rows are sized by the whole bucket (the widest - # possible window); consumers mask by the per-request widths. - union_rows = torch.empty( - (self.max_requests, self.seq_len), - dtype=torch.float32, - device=device, - ) - # The scratch, segment buffers, and staged per-request window - # starts are shared with the score-only runner built by - # ``prepare_cute_score``; the compiled kernels capture their - # device pointers. - runner = _FusedUnionScoreRunner( - layer_pools=self._cute_layer_pools, - layer_indices=self._cute_layer_indices, - max_requests=self.max_requests, - num_layers=self.num_layers, - seq_len=self.seq_len, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - page_ids=self.pointer_prefix[2], - seg_page_off=self.pointer_prefix[3], - seg_req_id=self.pointer_prefix[4], - seg_layer_id=self.pointer_prefix[5], - seg_seq_len=self._cute_seg_seq_len, - seg_out_offset=self._cute_seg_out_offset, - token_starts=self._cute_token_starts, - q_real=self.pointer_middle[0], - q_imag=self.pointer_middle[1], - mlr_coef=self.pointer_middle[2], - mean_cos=mean_cos, - mean_sin=mean_sin, - freq_scale_sq=self.pointer_tail[0], - output=self._cute_scratch, - enable_partial_stats=True, - ) - except (ImportError, RuntimeError, ValueError, AssertionError) as error: - raise RuntimeError( - "TriAttention CuTe union fusion setup failed and no other union path exists" - ) from error - self._union_fusion_runner_entry = (runner, union_rows, self._cute_token_starts) - return self._union_fusion_runner_entry - - def launch_cute_union_fusion( - self, - request_count: int, - valid_seq_lens: torch.Tensor, - valid_widths: torch.Tensor, - token_starts_device: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - union_out: torch.Tensor, - ) -> None: - """Run the fused score+stats+normalized-union pipeline (THE union path). - - Each request scores its own window (``token_starts_device`` carries - the per-request pinned prompt lengths), so mixed-prompt cohorts are - served directly. There is deliberately no fallback: an unsupported - geometry or request count raises loudly. - """ - self.prepare_cute_score(mean_cos, mean_sin) - if request_count <= 0 or request_count > self.max_requests: - raise ValueError("request count exceeds fixed score capacity") - runner, union_rows, staged_token_starts = self._union_fusion_runner(mean_cos, mean_sin) - if not runner.supports_union_fusion(request_count): - raise RuntimeError( - f"TriAttention CuTe union fusion has no compiled variant for " - f"request_count={request_count} (capacity {self.max_requests}) " - "and no other union path exists" - ) - num_segments = request_count * self.num_layers - torch.sub( - valid_seq_lens[:request_count], - token_starts_device[:request_count], - out=valid_widths[:request_count], - ) - torch.index_select( - valid_seq_lens, - 0, - self.pointer_prefix[4][:num_segments], - out=self._cute_seg_seq_len[:num_segments], - ) - staged_token_starts[:request_count].copy_(token_starts_device[:request_count]) - runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_rows[:request_count]) - columns = min(union_rows.shape[1], union_out.shape[1]) - union_out[:request_count, :columns].copy_(union_rows[:request_count, :columns]) + request_count = int(request_count) + if request_count <= 0 or request_count > round_starts.numel(): + raise ValueError("phase gather request count is outside its fixed buffers") + num_freqs = phase["omega"].numel() + _gather_mean_phase_kernel[(request_count,)]( + round_starts, + phase["cos"], + phase["sin"], + mean_cos, + mean_sin, + phase["rows"], + NUM_FREQS=num_freqs, + F_BLOCK=triton.next_power_of_2(num_freqs), + num_warps=1, + ) # --------------------------------------------------------------------------- # @@ -791,28 +283,24 @@ def prepare_per_head_scores( """Normalize and reduce score rows for either per-head eviction mode.""" request_count = int(request_count) num_kv_heads = int(num_kv_heads) - if not scores.is_cuda or scores.ndim != 4 or scores.dtype != torch.float32: - raise ValueError("per-head score preparation requires CUDA FP32 rows") - if not scores.is_contiguous() or request_count != scores.shape[0]: - raise ValueError("per-head score preparation request geometry does not match") + # Essentials only: the buffers were allocated together at workspace + # construction with matching geometry; dtype/layout is the kernel contract. + assert scores.is_cuda and scores.ndim == 4 and scores.dtype == torch.float32, ( + "per-head score preparation requires CUDA FP32 [request, layer, head, token] rows" + ) + assert scores.is_contiguous() and request_count == scores.shape[0], ( + "per-head score preparation request geometry does not match" + ) _, num_layers, num_query_heads, width = scores.shape - if num_kv_heads <= 0 or num_query_heads % num_kv_heads: - raise ValueError("per-head score preparation requires valid GQA geometry") + assert num_kv_heads > 0 and num_query_heads % num_kv_heads == 0, ( + "per-head score preparation requires valid GQA geometry" + ) selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - if ( - valid_widths.shape != (request_count,) - or valid_widths.dtype != torch.int32 - or valid_widths.device != scores.device - or row_mean.numel() < request_count * num_layers * num_query_heads - or row_inv_std.shape != row_mean.shape - or selection_scores.shape != (request_count, selection_rows, width) - or selection_scores.dtype != torch.float32 - or selection_scores.device != scores.device - or selection_seq_lens.shape != (request_count, selection_rows) - or selection_seq_lens.dtype != torch.int32 - or selection_seq_lens.device != scores.device - ): - raise ValueError("per-head score preparation buffers do not match") + assert ( + selection_scores.shape == (request_count, selection_rows, width) + and selection_seq_lens.shape == (request_count, selection_rows) + and row_mean.numel() >= request_count * num_layers * num_query_heads + ), "per-head score preparation buffers do not match" stats_block = 256 rows = num_layers * num_query_heads diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 51836d691a5b..85b3947d2769 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -41,12 +41,20 @@ def encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: return encoded +def compaction_family(compaction, name): + """Return one cache family dict ("dense", "swa", or "draft") or None.""" + for family in compaction["families"]: + if family["name"] == name: + return family + return None + + def _write_move_offsets(compaction, offsets, moves_per_request): cumulative = [0] for count in moves_per_request: cumulative.append(cumulative[-1] + count) # Rows past the cohort are padding and contribute no moves. - cumulative.extend(cumulative[-1:] * (compaction.request_count - len(moves_per_request))) + cumulative.extend(cumulative[-1:] * (compaction["request_count"] - len(moves_per_request))) offsets.copy_(torch.tensor(cumulative, dtype=torch.int32), non_blocking=True) @@ -54,38 +62,39 @@ def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): """Load a cohort's per-request protected tails into the move offsets. Production stages these rows through the single round-metadata upload; - tests drive the fixed buffers directly through this helper (moved out of - BatchedKVCacheCompaction, whose production surface has no caller for it). + tests drive the fixed buffers directly through this helper. """ - if len(tail_lengths) > compaction.request_count: + if len(tail_lengths) > compaction["request_count"]: raise ValueError("the cohort exceeds the compaction request capacity") - if any(tail < 0 or tail > compaction.protected_tail_capacity for tail in tail_lengths): + if any(tail < 0 or tail > compaction["protected_tail_capacity"] for tail in tail_lengths): raise ValueError("a protected tail exceeds the configured capacity") _write_move_offsets( compaction, - compaction.target_dense_compaction.move_source_offsets, - [compaction.decode_keep_count + int(tail) for tail in tail_lengths], + compaction_family(compaction, "dense")["offsets"], + [compaction["decode_keep_count"] + int(tail) for tail in tail_lengths], ) - if compaction.target_swa_compaction is not None: + swa_family = compaction_family(compaction, "swa") + if swa_family is not None: _write_move_offsets( compaction, - compaction.target_swa_compaction.move_source_offsets, - [compaction.swa_window + int(tail) for tail in tail_lengths], + swa_family["offsets"], + [compaction["swa_window"] + int(tail) for tail in tail_lengths], ) - if compaction.draft_compaction is not None: + draft_family = compaction_family(compaction, "draft") + if draft_family is not None: if draft_tail_lengths is None: draft_tail_lengths = [0] * len(tail_lengths) if len(draft_tail_lengths) != len(tail_lengths): raise ValueError("draft protected tails must match the cohort") if any( - tail < 0 or tail > compaction.draft_protected_tail_capacity + tail < 0 or tail > compaction["draft_protected_tail_capacity"] for tail in draft_tail_lengths ): raise ValueError("a draft protected tail exceeds the configured capacity") _write_move_offsets( compaction, - compaction.draft_compaction.move_source_offsets, - [compaction.decode_keep_count + int(tail) for tail in draft_tail_lengths], + draft_family["offsets"], + [compaction["decode_keep_count"] + int(tail) for tail in draft_tail_lengths], ) @@ -127,9 +136,9 @@ def make_ramp_pools( def build_compaction(**overrides): - """``BatchedKVCacheCompaction`` with the suite's default 2-layer geometry.""" + """``build_cache_compactions`` with the suite's default 2-layer geometry.""" from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - BatchedKVCacheCompaction, + build_cache_compactions, ) args = dict( @@ -144,21 +153,18 @@ def build_compaction(**overrides): swa_window=None, ) args.update(overrides) - return BatchedKVCacheCompaction(**args) + return build_cache_compactions(**args) def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None): - """A ``__new__``-built staging shell for the bulk page-table copy tests.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, - ) - - staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) + """A bare workspace namespace for the bulk page-table copy tests.""" + staging = SimpleNamespace() staging.device = device staging.max_requests = max_requests staging.copy_block_count = copy_block_count if page_count is not None: staging.page_count = page_count + staging.stream = None staging.bulk_copy_done = torch.cuda.Event() staging.bulk_consume_done = torch.cuda.Event() staging.page_tables_active = False @@ -188,41 +194,38 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): ) -def make_fixed_resources_stubs(manager, *, decode_width=260): - """Stub the calibration/staging surfaces around ``_fixed_resources_for``.""" +def make_workspace_stubs(manager, *, decode_width=260): + """Stub the calibration/layout surfaces around ``_workspace_for``.""" manager._H = 2 manager._F = 2 manager._freq_scale_sq = torch.ones(2) manager._offsets = torch.ones(2) + manager._phase = {"rows": 8} manager.calibration = {"omega": torch.ones(2)} manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) pool = torch.empty(8, 2, 1, 4, 4) - layout = SimpleNamespace( + layout = dict( manager=SimpleNamespace(num_pools=1), num_layers=2, global_layers=[0, 1], layer_pools=[pool, pool], dense_layers=[0, 1], swa_layers=[], + swa_window=None, storage_groups={0: [0, 1]}, + layer_group_representative={0: 0, 1: 0}, + layer_pool_keys=(("pool", 0), ("pool", 0)), pool_view_fingerprint=(("fixed",),), ) - score_staging = SimpleNamespace( - fused_group=SimpleNamespace(output=torch.empty(8, 4, decode_width)), - bind_score_launcher=mock.Mock(), - token_starts_device=torch.zeros(8, dtype=torch.int32), + workspace = SimpleNamespace( decode_width=decode_width, page_table_token_capacity=65537, max_requests=8, - ) - keep_set_selector = SimpleNamespace( + token_starts_device=torch.zeros(8, dtype=torch.int32), valid_widths=torch.empty(8, dtype=torch.int32), - # The builder zero-fills the mode's provisional top-k buffer. - top_indices_i32=torch.zeros(8, 4, dtype=torch.int32), - final_indices=torch.zeros(8, 4, dtype=torch.int32), ) - return layout, score_staging, keep_set_selector + return layout, workspace def make_fake_v2(enable_block_reuse=False, *, is_draft=False): @@ -231,6 +234,7 @@ def make_fake_v2(enable_block_reuse=False, *, is_draft=False): fake_v2 = KVCacheManagerV2.__new__(KVCacheManagerV2) fake_v2.enable_block_reuse = enable_block_reuse + fake_v2.enable_swa_scratch_reuse = False fake_v2.is_draft = is_draft fake_v2.kv_compression_manages_history = False fake_v2.kv_factor = 2 @@ -250,7 +254,7 @@ def make_fake_v2(enable_block_reuse=False, *, is_draft=False): 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.empty(1, dtype=torch.int64) + 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 = {} @@ -287,40 +291,22 @@ def make_request(request_id, **overrides): @contextmanager def mocked_eviction_internals(manager): """Run the real ``_evict_requests`` body around mocked GPU launches.""" - score_staging = SimpleNamespace( - launch_prepared_score=mock.Mock(return_value=torch.zeros(1)), - launch_prepared_union_fusion=mock.Mock(), - mark_page_tables_consumed=mock.Mock(), - ) - # The dispatch reads the selector's own eviction mode, so the stub - # mirrors the manager's and carries both mode paths' launch surfaces. - keep_set_selector = SimpleNamespace( - eviction_mode=manager.eviction_mode, - combined=torch.zeros(1), - select_requests=mock.Mock(), - select_prepared_union_scores=mock.Mock(), - refresh_row_prompt_offsets=mock.Mock(), - ) - resources = SimpleNamespace( - score_staging=score_staging, - keep_set_selector=keep_set_selector, - ) - batched_compaction = SimpleNamespace(compact=mock.Mock()) + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module + + workspace = SimpleNamespace(max_requests=8) + layout = dict(swa_layers=[], swa_window=None) with ( - mock.patch.object(manager, "_runtime_kv_layout", return_value=SimpleNamespace()), - mock.patch.object(manager, "_fixed_resources_for", return_value=resources), - mock.patch.object( - manager, - "_batched_compaction_for", - return_value=batched_compaction, - ), - mock.patch.object(manager, "_attach_page_ids") as attach, + mock.patch.object(manager, "_runtime_kv_layout", return_value=layout), + mock.patch.object(manager, "_workspace_for", return_value=workspace), + mock.patch.object(module, "stage_eviction_cohort") as stage, + mock.patch.object(module, "run_eviction_round") as run_round, + mock.patch.object(module, "mark_page_tables_consumed") as consumed, ): yield SimpleNamespace( - score_staging=score_staging, - keep_set_selector=keep_set_selector, - batched_compaction=batched_compaction, - attach=attach, + workspace=workspace, + stage=stage, + run_round=run_round, + consumed=consumed, ) 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 index 82e71ab40b33..0ce12b5de609 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -2,35 +2,126 @@ # SPDX-License-Identifier: Apache-2.0 """The SM100 TriAttention CuTe scorer (the only score path) vs PyTorch oracles. -Two layers of coverage over ``_FixedScoreGroup.launch`` -- the exact -production entry point. The kernel-numerics matrix drives a single-layer -group across the supported page geometries (permuted physical pages, ragged -valid lengths, GQA group 4 riding the padded MMA tile) against inline -oracle math. The launch-path matrix drives multi-layer groups across the -named production geometries (Qwen3, GPT-OSS, the originally validated -128-token-page shape) against the shared pure-PyTorch oracle, sweeps -request counts up to the group capacity, and checks the per-request -decode-width metadata the selection reduce kernels consume. The contract -tests pin the loud-failure behavior: unsupported geometry, removed -aggregations, and request counts beyond capacity raise -- there is no -fallback score kernel. +Two layers of coverage over the production score leg (workspace metadata +staging, the compiled runner launch, and the decode-window gather -- the same +sequence ``run_eviction_round`` fires). The kernel-numerics matrix drives a +single-layer workspace across the supported page geometries (permuted +physical pages, ragged valid lengths, GQA group 4 riding the padded MMA tile) +against inline oracle math. The launch-path matrix drives multi-layer +workspaces across the named production geometries (Qwen3, GPT-OSS, the +originally validated 128-token-page shape) against the shared pure-PyTorch +oracle, sweeps request counts up to the workspace capacity, and checks the +per-request decode-width metadata the selection reduce kernels consume. The +contract tests pin the loud-failure behavior: unsupported geometry raises at +workspace construction and an oversized cohort is rejected at staging -- +there is no fallback score kernel. """ +from types import SimpleNamespace + import pytest import torch from conftest import encode_block_offsets as _encode_block_offsets from conftest import torch_tri_score_oracle as _torch_tri_score_oracle -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, -) - 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_score_workspace( + *, + layer_pools, + max_requests, + seq_len, + num_q_heads, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + decode_width=None, + eviction_mode="per_head", +): + """A score-only workspace over one shared page-table slot (no compaction).""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + prepare_eviction_workspace, + ) + + num_layers = len(layer_pools) + return prepare_eviction_workspace( + eviction_mode=eviction_mode, + layer_pools=layer_pools, + dense_groups=[list(range(num_layers))], + dense_layers=list(range(num_layers)), + page_representatives=[0], + max_requests=max_requests, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_freqs=int(q_real.shape[-1]), + keep_count=1, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + offsets=offsets, + omega=omega, + decode_width=decode_width, + build_compaction=False, + ) + + +def _write_block_offsets(ws, encoded): + """Load a test page table into the workspace's staged block-offset plane.""" + ws.block_offsets_device.zero_() + ws.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) + + +def _launch_split_scores( + ws, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin +): + """The production score-only leg: stage metadata, fire the compiled + runner, gather each request's decode window (``run_eviction_round``'s + per-head sequence, parameterized by the request count).""" + num_segments = request_count * ws.num_layers + torch.sub( + valid_seq_lens[:request_count], + token_starts[:request_count], + out=valid_widths[:request_count], + ) + torch.index_select( + valid_seq_lens, 0, ws.seg_req[:num_segments], out=ws.seg_seq_len[:num_segments] + ) + ws.cute_token_starts[:request_count].copy_(token_starts[:request_count]) + assert ws.runner.supports(request_count) + ws.runner.launch(request_count, mean_cos, mean_sin) + group_size = ws.num_q_heads // ws.num_kv_heads + source = ( + ws.cute_scratch[: ws.num_kv_heads * 8 * num_segments * ws.bucket_seq_len] + .view(ws.num_kv_heads, 8, request_count, ws.num_layers, ws.bucket_seq_len)[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + columns = token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + ws.gather_columns + columns = columns.clamp_(max=ws.bucket_seq_len - 1).expand( + request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + ) + output = torch.full( + (request_count, ws.num_layers, ws.num_q_heads, ws.decode_width), + float("nan"), + dtype=torch.float32, + device=ws.device, + ) + torch.gather( + source, + 4, + columns, + out=output.view(request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width), + ) + return output + + @requires_sm100 @pytest.mark.parametrize( "tokens_per_block,page_permutation,valid_lens,num_freqs,num_q_heads", @@ -77,30 +168,26 @@ def test_cute_score_matches_torch_mean_oracle( mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() + ws = _make_score_workspace( + 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, + ) # Native block-offset staging layout ([pool_slot, request, K/V plane, # block] int32): K-plane entries encode physical_page * kv_factor with # kv_factor == 2. Both requests read the same (permuted) page sequence. k_plane = [2 * page for page in page_permutation] v_plane = [2 * page + 1 for page in page_permutation] - block_offsets = torch.tensor( - [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device - ) - group = _FixedScoreGroup( - [pool], - [0], - 2, - 2, - seq_len, - num_q_heads, - block_offsets, - [0], - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - output_width=seq_len, + _write_block_offsets( + ws, + torch.tensor([[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device), ) keys = ( torch.cat([pool[page, 0, 0] for page in page_permutation], dim=0) @@ -116,8 +203,8 @@ def test_cute_score_matches_torch_mean_oracle( valid_widths = torch.tensor(valid_lens, dtype=torch.int32, device=device) token_starts_device = torch.zeros(2, dtype=torch.int32, device=device) for request_count in (1, 2): - group.output.fill_(float("nan")) - actual = group.launch( + actual = _launch_split_scores( + ws, request_count, valid_seq_lens, valid_widths, @@ -143,8 +230,9 @@ def test_cute_score_matches_torch_mean_oracle( ) torch.cuda.synchronize() - # The CuTe runner is the only score path; prove setup actually built it. - assert group._cute_score_runner is not None + # The CuTe runner is the only score path, compiled eagerly at workspace + # construction; prove setup actually built it. + assert ws.runner is not None def _build_case( @@ -187,23 +275,20 @@ def _build_case( 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 - group = _FixedScoreGroup( - pools, - list(range(num_layers)), - max_requests, - page_count, - capacity, - num_q_heads, - _encode_block_offsets(page_ids), - [0] * num_layers, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets_t, - output_width=capacity - prompt_len, + ws = _make_score_workspace( + 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(ws, _encode_block_offsets(page_ids)) round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) # Ragged valid lengths whose tails land mid-page and mid-compute-tile. @@ -224,7 +309,8 @@ def _build_case( offsets=offsets_t, ) return ( - group, + ws, + pools, token_starts, valid_seq_lens, seq_lens, @@ -266,7 +352,8 @@ def test_cute_kernel_matches_torch_oracle(case): max_requests = case["max_requests"] num_layers = case["num_layers"] ( - group, + ws, + pools, token_starts, valid_seq_lens, seq_lens, @@ -274,10 +361,10 @@ def test_cute_kernel_matches_torch_oracle(case): mean_sin, oracle_inputs, ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) - device = group.output.device + device = ws.device oracle = _torch_tri_score_oracle( - group._cute_layer_pools, + pools, oracle_inputs["page_ids"], seq_lens, [int(start) for start in range(9, 9 + max_requests)], @@ -290,12 +377,14 @@ def test_cute_kernel_matches_torch_oracle(case): list(range(num_layers)), ) - # The fused runner dispatches every request count up to the group - # capacity; cover one, an intermediate count, and the capacity. + # The compiled runner serves every request count up to the workspace + # capacity and nothing beyond it; cover one, an intermediate count, and + # the capacity. + assert not ws.runner.supports(max_requests + 1) for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): - group.output.fill_(float("nan")) valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) - scores = group.launch( + scores = _launch_split_scores( + ws, request_count, valid_seq_lens, valid_widths, @@ -307,9 +396,9 @@ def test_cute_kernel_matches_torch_oracle(case): request_count, num_layers, case["num_q_heads"], - group.output_width, + ws.decode_width, ) - # The launch owns the per-request decode widths the selection + # The score leg owns the per-request decode widths the selection # reduce kernels consume. assert valid_widths[:request_count].tolist() == [ seq_lens[request] - prompt_len for request in range(request_count) @@ -325,69 +414,44 @@ def test_cute_kernel_matches_torch_oracle(case): ) -def _tiny_unsupported_group(dtype: torch.dtype): - """A geometry far outside the CuTe contract (constructor accepts it).""" +# The loud-failure contract, one representative per guard family: unsupported +# geometry raises at workspace construction (the only score path compiles +# eagerly there -- no fallback), and an oversized cohort is rejected by the +# host-side staging gate before any GPU work. +def test_unsupported_geometry_raises_at_workspace_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 pools with a 4-token page and 4 frequencies sit far outside the + # CuTe contract on every device. pools = [ - torch.randn(max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device).to( - dtype - ) + torch.randn(max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device) for _ in range(num_layers) ] - page_ids = ( - torch.arange(max_requests * page_count, device=device) - .view(max_requests, page_count) - .contiguous() - ) - capacity = page_count * tokens_per_block - group = _FixedScoreGroup( - pools, - list(range(num_layers)), - max_requests, - page_count, - capacity, - 2, - _encode_block_offsets(page_ids), - [0] * num_layers, - torch.randn(num_layers, 2, num_freqs, device=device), - torch.randn(num_layers, 2, num_freqs, device=device), - torch.randn(num_layers, 2, num_freqs, device=device), - torch.rand(num_freqs, device=device) + 0.5, - torch.rand(num_freqs, device=device) * 0.05, - torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), - output_width=capacity - 1, - ) - device_args = dict(dtype=torch.int32, device=device) - return group, ( - torch.full((max_requests,), capacity, **device_args), - torch.empty(max_requests, **device_args), - torch.ones(max_requests, **device_args), - torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), - torch.zeros(max_requests, num_freqs, dtype=torch.float32, device=device), - ) + calib = torch.randn(num_layers, 2, num_freqs, device=device) + with pytest.raises(ValueError, match="TriAttention score requires SM100"): + _make_score_workspace( + 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, + ) -# The loud-failure contract, one representative per guard family. All three -# guards fire before any kernel work, in this order: removed aggregation, -# request count beyond the group capacity (previously exercised on a -# production-shaped SM100 group; the check is layered before compilation, so -# the tiny group covers the same code path), unsupported geometry at setup. -@pytest.mark.parametrize( - "dtype,request_count,launch_kwargs,match", - [ - pytest.param( - torch.float32, 1, {}, "TriAttention score requires SM100", id="unsupported_geometry" - ), - pytest.param( - torch.bfloat16, 1, {"aggregation": "max"}, "max aggregation", id="max_aggregation" - ), - pytest.param(torch.bfloat16, 3, {}, "exceeds fixed score capacity", id="beyond_capacity"), - ], -) -def test_score_launch_contract_raises(dtype, request_count, launch_kwargs, match): - group, launch_args = _tiny_unsupported_group(dtype) - with pytest.raises(ValueError, match=match): - group.launch(request_count, *launch_args, **launch_kwargs) +def test_oversized_cohort_is_rejected_at_staging(): + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + stage_eviction_cohort, + ) + + ws = SimpleNamespace(max_requests=2) + with pytest.raises(ValueError, match="does not fit the workspace request capacity"): + stage_eviction_cohort(ws, None, [1, 2, 3], [0, 0, 0], [0, 0, 0]) 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 index 2996344d697f..e69180f26e0f 100644 --- 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 @@ -3,10 +3,11 @@ """Equivalence coverage for the fused score+stats+union pipeline (two CuTe kernels). The reference side gathers the SAME production score rows (the fused pack's -score-only entry) and normalizes + union-reduces them with a pure-torch -float32 oracle. The fused-vs-reference comparison was always tolerance-based -(the fused pipeline's reduction order differs from any reference); the -tolerances are unchanged from the retired Triton reference copies. +score-only entry, which every workspace runner compiles) and normalizes + +union-reduces them with a pure-torch float32 oracle. The fused-vs-reference +comparison was always tolerance-based (the fused pipeline's reduction order +differs from any reference); the tolerances are unchanged from the retired +Triton reference copies. """ import pytest @@ -18,6 +19,114 @@ ) +def _make_union_workspace( + *, + layer_pools, + max_requests, + seq_len, + num_q_heads, + q_real, + q_imag, + mlr_coef, + freq_scale_sq, + omega, + offsets, + decode_width=None, +): + """A union-mode workspace over one shared page-table slot (no compaction). + + The union runner also compiles the score-only entries, so one workspace + serves both the fused pipeline and the split reference leg. + """ + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + prepare_eviction_workspace, + ) + + num_layers = len(layer_pools) + return prepare_eviction_workspace( + eviction_mode="union", + layer_pools=layer_pools, + dense_groups=[list(range(num_layers))], + dense_layers=list(range(num_layers)), + page_representatives=[0], + max_requests=max_requests, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_freqs=int(q_real.shape[-1]), + keep_count=1, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + offsets=offsets, + omega=omega, + decode_width=decode_width, + build_compaction=False, + ) + + +def _write_block_offsets(ws, encoded): + """Load a test page table into the workspace's staged block-offset plane.""" + ws.block_offsets_device.zero_() + ws.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) + + +def _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts): + """Stage the per-round score metadata exactly like ``run_eviction_round``.""" + num_segments = request_count * ws.num_layers + torch.sub( + valid_seq_lens[:request_count], + token_starts[:request_count], + out=valid_widths[:request_count], + ) + torch.index_select( + valid_seq_lens, 0, ws.seg_req[:num_segments], out=ws.seg_seq_len[:num_segments] + ) + ws.cute_token_starts[:request_count].copy_(token_starts[:request_count]) + + +def _launch_split_scores( + ws, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin +): + """The production score-only leg plus the decode-window gather.""" + _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts) + assert ws.runner.supports(request_count) + ws.runner.launch(request_count, mean_cos, mean_sin) + num_segments = request_count * ws.num_layers + group_size = ws.num_q_heads // ws.num_kv_heads + source = ( + ws.cute_scratch[: ws.num_kv_heads * 8 * num_segments * ws.bucket_seq_len] + .view(ws.num_kv_heads, 8, request_count, ws.num_layers, ws.bucket_seq_len)[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + columns = token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + ws.gather_columns + columns = columns.clamp_(max=ws.bucket_seq_len - 1).expand( + request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + ) + output = torch.full( + (request_count, ws.num_layers, ws.num_q_heads, ws.decode_width), + float("nan"), + dtype=torch.float32, + device=ws.device, + ) + torch.gather( + source, + 4, + columns, + out=output.view(request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width), + ) + return output + + +def _launch_union_fusion( + ws, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out +): + """The fused score+stats+normalized-union pipeline (THE union path).""" + _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts) + assert ws.runner.supports_union_fusion(request_count) + ws.runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_out[:request_count]) + + def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: """Pure-torch union oracle: z-normalize each row's valid prefix, union-max. @@ -56,10 +165,6 @@ def _check_union_fusion_matches_split_pipeline( """ pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) - torch.manual_seed(20260721) device = torch.device("cuda") seq_len = 256 @@ -84,27 +189,23 @@ def _check_union_fusion_matches_split_pipeline( mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() + ws = _make_union_workspace( + 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, + ) k_plane = [2 * page for page in page_permutation] v_plane = [2 * page + 1 for page in page_permutation] - block_offsets = torch.tensor( - [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device - ) - group = _FixedScoreGroup( - [pool], - [0], - 2, - num_pages, - seq_len, - num_q_heads, - block_offsets, - [0], - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - output_width=seq_len, + _write_block_offsets( + ws, + torch.tensor([[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device), ) if valid_lens is None: valid_lens = [seq_len, seq_len] @@ -118,7 +219,8 @@ def _check_union_fusion_matches_split_pipeline( # then the pure-torch union oracle. split_widths = torch.empty(request_count, dtype=torch.int32, device=device) token_starts = torch.tensor(score_starts, dtype=torch.int32, device=device) - per_head = group.launch( + per_head = _launch_split_scores( + ws, request_count, valid_seq_lens, split_widths, @@ -134,7 +236,8 @@ def _check_union_fusion_matches_split_pipeline( fused_out = torch.full( (request_count, seq_len), float("nan"), dtype=torch.float32, device=device ) - group.launch_cute_union_fusion( + _launch_union_fusion( + ws, request_count, valid_seq_lens, fused_widths, @@ -204,11 +307,11 @@ def test_union_fusion_matches_split_pipeline( def test_union_fusion_guards_raise(guard: str, monkeypatch: pytest.MonkeyPatch) -> None: """One representative per fused-pipeline guard family raises loudly. - ``runner_construction_failure``: a fused-runner construction failure - surfaces as the no-fallback RuntimeError at score setup - (``prepare_cute_score`` runs before the union runner is built). - ``frequency_count``: 16 frequencies (head size 32) sit outside the fused - kernel contract and are rejected at kernel construction. + ``runner_construction_failure``: a runner construction failure surfaces + as the no-fallback RuntimeError at workspace construction (the only + place the CuTe entries compile). ``frequency_count``: 16 frequencies + (head size 32) sit outside the fused kernel contract and are rejected at + kernel construction. """ cutlass = pytest.importorskip("cutlass") @@ -233,9 +336,6 @@ def test_union_fusion_guards_raise(guard: str, monkeypatch: pytest.MonkeyPatch) return import tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused as fused_module # noqa: E501 - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) torch.manual_seed(20260721) device = torch.device("cuda") @@ -246,44 +346,23 @@ def test_union_fusion_guards_raise(guard: str, monkeypatch: pytest.MonkeyPatch) ).to(torch.bfloat16) q_real = 0.125 * torch.randn(1, num_q_heads, num_freqs, device=device) omega = torch.linspace(0.01, 0.03, num_freqs, device=device) - k_plane = [2 * page for page in range(num_pages)] - v_plane = [2 * page + 1 for page in range(num_pages)] - block_offsets = torch.tensor( - [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device - ) - group = _FixedScoreGroup( - [pool], - [0], - 2, - num_pages, - seq_len, - num_q_heads, - block_offsets, - [0], - q_real, - torch.randn_like(q_real) * 0.125, - torch.randn_like(q_real) * 0.125, - torch.linspace(0.5, 1.5, num_freqs, device=device), - omega, - torch.tensor([1.0, 2.0, 4.0], device=device), - output_width=seq_len, - ) def _refuse_construction(**_kwargs): raise ValueError("synthetic fused-runner construction failure") monkeypatch.setattr(fused_module, "TriAttentionCuteScoreRunner", _refuse_construction) - mean_cos = torch.cos(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) - mean_sin = torch.sin(torch.outer(torch.tensor([256.0, 257.0], device=device), omega)) with pytest.raises(RuntimeError, match="no other score path exists"): - group.launch_cute_union_fusion( - 2, - torch.full((2,), seq_len, dtype=torch.int32, device=device), - torch.empty(2, dtype=torch.int32, device=device), - torch.zeros(2, dtype=torch.int32, device=device), - mean_cos.contiguous(), - mean_sin.contiguous(), - torch.empty((2, seq_len), dtype=torch.float32, device=device), + _make_union_workspace( + layer_pools=[pool], + max_requests=2, + seq_len=seq_len, + num_q_heads=num_q_heads, + q_real=q_real, + q_imag=torch.randn_like(q_real) * 0.125, + mlr_coef=torch.randn_like(q_real) * 0.125, + freq_scale_sq=torch.linspace(0.5, 1.5, num_freqs, device=device), + omega=omega, + offsets=torch.tensor([1.0, 2.0, 4.0], device=device), ) @@ -321,10 +400,6 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: """ pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) - torch.manual_seed(20260722) device = torch.device("cuda") num_layers = 36 @@ -360,28 +435,24 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() - page_ids = torch.arange(num_pages, dtype=torch.int32, device=device) - block_offsets = torch.zeros(1, max_requests, 2, num_pages, dtype=torch.int32, device=device) - block_offsets[0, :request_count, 0] = 2 * page_ids - block_offsets[0, :request_count, 1] = 2 * page_ids + 1 - group = _FixedScoreGroup( - layer_pools, - list(range(num_layers)), - max_requests, - num_pages, - seq_len, - num_q_heads, - block_offsets, - [0] * num_layers, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - output_width=decode_window, + ws = _make_union_workspace( + layer_pools=layer_pools, + max_requests=max_requests, + 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, ) - assert (num_kv_heads * 8 * max_requests * num_layers * seq_len > 2**31) == (max_requests == 64) + assert (ws.cute_scratch.numel() > 2**31) == (max_requests == 64) + page_ids = torch.arange(num_pages, dtype=torch.int32, device=device) + ws.block_offsets_device.zero_() + ws.block_offsets_device[0, :request_count, 0, :num_pages] = 2 * page_ids + ws.block_offsets_device[0, :request_count, 1, :num_pages] = 2 * page_ids + 1 valid_seq_lens = torch.zeros(max_requests, dtype=torch.int32, device=device) token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) @@ -391,7 +462,8 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: # 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 = group.launch( + per_head = _launch_split_scores( + ws, request_count, valid_seq_lens, split_widths, @@ -408,7 +480,8 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: fused_out = torch.full( (max_requests, seq_len), float("nan"), dtype=torch.float32, device=device ) - group.launch_cute_union_fusion( + _launch_union_fusion( + ws, max_requests, valid_seq_lens, fused_widths, 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 index 29b9a3b95d72..0838aa0ff730 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -10,8 +10,7 @@ ``destination_base = prompt_len``. These tests cover the physical draft moves, the packed move indices, stream ordering across both cache managers, the speculative admission gates (one representative per guard family), the -published compressed-token invariant, and prepared-compaction cache -invalidation. +published compressed-token invariant, and workspace rebuild/invalidation. """ from types import SimpleNamespace @@ -20,23 +19,28 @@ import pytest import torch from conftest import build_compaction as _build_compaction +from conftest import compaction_family as _compaction_family from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_fake_v2 as _make_fake_v2 -from conftest import make_fixed_resources_stubs as _make_fixed_resources_stubs from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_request as _make_request from conftest import make_triattention as _make_triattention +from conftest import make_workspace_stubs as _make_workspace_stubs from conftest import mocked_eviction_internals as _mocked_eviction_internals from conftest import set_protected_tails as _set_protected_tails +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import run_cache_compactions from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( TriAttention, - _FixedScoreStagingBuffers, - _PreparedEviction, - _RequestCompressionState, + mark_page_tables_consumed, ) +def _fresh_request_state(): + """One request's compression ledger, as the manager initializes it.""" + return {"generation_steps": 0, "evicted_tokens": 0, "confirmed_kv_length": None} + + def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: """Gather one request's pages into [K/V, head, token, dim] order.""" num_kv_heads = int(pool.shape[2]) @@ -86,7 +90,7 @@ def _launched_draft_compaction(draft_protected_tails): draft_page_table_slots={0: 0}, ) _set_protected_tails(compaction, target_protected_tails, draft_protected_tails) - compaction.compact() + run_cache_compactions(compaction) torch.cuda.synchronize(device) return SimpleNamespace( @@ -172,10 +176,10 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect # The packed draft move indices must match the same broadcast-plus-tail # oracle the physical moves followed. - draft_compaction = built.compaction.draft_compaction + draft_family = _compaction_family(built.compaction, "draft") expected_row = torch.cat(expected_moves) - assert draft_compaction.move_source_offsets.cpu().tolist() == expected_offsets - draft_indices = draft_compaction.move_source_indices + assert draft_family["offsets"].cpu().tolist() == expected_offsets + draft_indices = draft_family["source"] # The index buffer is sized for the widest tail (the capacity); this # round's moves are packed at the front, where the offsets point. capacity_total = built.request_count * ( @@ -188,27 +192,28 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect def test_mark_page_tables_consumed_orders_both_manager_streams(): - staging = _FixedScoreStagingBuffers.__new__(_FixedScoreStagingBuffers) - staging.device = torch.device("cuda", torch.cuda.current_device()) - staging.page_tables_active = True event = mock.Mock() - staging.bulk_consume_done = event + workspace = SimpleNamespace( + device=torch.device("cuda", torch.cuda.current_device()), + page_tables_active=True, + bulk_consume_done=event, + ) target_stream = mock.Mock() draft_stream = mock.Mock() compute_stream = SimpleNamespace() with mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream): - staging.mark_page_tables_consumed(target_stream, draft_stream) + mark_page_tables_consumed(workspace, target_stream, draft_stream) # One event records the compact launches; BOTH cache managers wait on it, # so neither can free or reallocate pages this cohort is still reading. event.record.assert_called_once_with(compute_stream) target_stream.wait_event.assert_called_once_with(event) draft_stream.wait_event.assert_called_once_with(event) - assert staging.page_tables_active is False + assert workspace.page_tables_active is False with pytest.raises(RuntimeError, match="not staged"): - staging.mark_page_tables_consumed(target_stream, draft_stream) + mark_page_tables_consumed(workspace, target_stream, draft_stream) @pytest.mark.parametrize( @@ -298,7 +303,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager.draft_kv_cache_manager = draft_manager request = _make_request(7, py_prompt_len=2, py_num_accepted_draft_tokens=1) - manager._request_states[7] = _RequestCompressionState() + manager._request_states[7] = _fresh_request_state() batch = SimpleNamespace(generation_requests=[request]) # Every step confirms one sampled token plus one accepted draft token. @@ -308,7 +313,6 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): previous_published = 0 eviction_rounds = 0 with _mocked_eviction_internals(manager) as internals: - score_staging = internals.score_staging for _ in range(6): uncompressed += 2 confirmed += 2 @@ -317,16 +321,18 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager._periodic_evict(batch) state = manager._request_states[7] - if state.confirmed_kv_length < confirmed: + if state["confirmed_kv_length"] < confirmed: # An eviction round compacted the cache to prompt + budget. eviction_rounds += 1 - confirmed = state.confirmed_kv_length + confirmed = state["confirmed_kv_length"] cache.capacity = confirmed assert confirmed == 2 + 4 # The staged logical position restores the uncompressed - # length: physical confirmed plus everything evicted so far. - prepared = internals.attach.call_args.args[0] - assert prepared[0].round_start == uncompressed + # length: physical confirmed plus everything evicted so far + # (stage_eviction_cohort args: ws, manager, ids, round_starts, + # prompt lengths, seq lens, page-table lens). + round_starts = internals.stage.call_args.args[3] + assert round_starts[0] == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. assert request.py_num_compressed_tokens == uncompressed - confirmed @@ -338,8 +344,8 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): # Each round the draft cache shrinks with the target and both manager # streams are ordered after the compact launches. assert draft_cache.resize.call_args_list == [mock.call(7, None)] * eviction_rounds - assert score_staging.mark_page_tables_consumed.call_args_list == ( - [mock.call(target._stream, draft_manager._stream)] * eviction_rounds + assert internals.consumed.call_args_list == ( + [mock.call(internals.workspace, target._stream, draft_manager._stream)] * eviction_rounds ) @@ -347,13 +353,21 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention(top_B=4) - layout, score_staging, keep_set_selector = _make_fixed_resources_stubs(manager) + layout, workspace = _make_workspace_stubs(manager) + # The one-time host block-offset table shape gate reads the real manager + # tables (int32 [pools, slots, K/V, blocks]). + manager.kv_cache_manager.host_kv_cache_block_offsets = torch.zeros( + 1, 8, 2, 4, dtype=torch.int32 + ) draft_manager = _make_fake_v2(is_draft=True) draft_manager.num_pools = 1 + draft_manager.host_kv_cache_block_offsets = torch.zeros(1, 8, 2, 4, dtype=torch.int32) manager.draft_kv_cache_manager = draft_manager manager._draft_runtime_kv_layout = mock.Mock( - return_value=SimpleNamespace( + return_value=dict( layer_pools=[], + dense_layers=[], + layer_group_representative={}, pool_representatives=(), layer_pool_keys=(), pool_page_counts=(4,), @@ -361,53 +375,50 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): ) ) prepared = [ - _PreparedEviction( - request=_make_request(7), - request_id=7, - seq_len=8, - round_start=8, - prompt_len=0, - expected_keep_count=4, - protected_tail=0, - ) + { + "request": _make_request(7), + "request_id": 7, + "seq_len": 8, + "round_start": 8, + "prompt_len": 0, + "expected_keep_count": 4, + "protected_tail": 0, + } ] - with ( - mock.patch.object( - module, - "_FixedScoreStagingBuffers", - return_value=score_staging, - ) as score_cls, - mock.patch.object( - module, - "_BatchedKeepSetSelector", - return_value=keep_set_selector, - ), - ): - resources = manager._fixed_resources_for(layout, prepared) + with mock.patch.object( + module, + "prepare_eviction_workspace", + return_value=workspace, + ) as prepare: + resources = manager._workspace_for(layout, prepared) # Request capacity follows the executor limits, while the score # bucket follows what the cohort actually presents (power-of-two, # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. - assert resources.score_staging is score_staging - assert score_cls.call_args.kwargs["max_requests"] == 8 - assert score_cls.call_args.kwargs["decode_width"] == 4 + 2 * 128 - assert score_cls.call_args.kwargs["seq_len"] == 1024 - assert score_cls.call_args.kwargs["page_table_token_capacity"] == 1024 + 1 - assert score_cls.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 - - # A second round with unchanged pools reuses the resident buffers and - # keeps the cached compaction launches. - cached_compaction = object() - manager._batched_compaction = cached_compaction - assert manager._fixed_resources_for(layout, prepared) is resources - assert score_cls.call_count == 1 - assert manager._batched_compaction is cached_compaction - - # A pool change invalidates both the buffers and the compaction - # launches that alias them. - layout.pool_view_fingerprint = (("moved",),) - rebuilt = manager._fixed_resources_for(layout, prepared) + assert resources is workspace + assert prepare.call_args.kwargs["eviction_mode"] == "union" + assert prepare.call_args.kwargs["max_requests"] == 8 + assert prepare.call_args.kwargs["decode_width"] == 4 + 2 * 128 + assert prepare.call_args.kwargs["seq_len"] == 1024 + assert prepare.call_args.kwargs["page_table_token_capacity"] == 1024 + 1 + assert prepare.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 + + # A second round with unchanged pools reuses the resident workspace + # (and with it the compaction launch data it carries). + assert manager._workspace_for(layout, prepared) is resources + assert prepare.call_count == 1 + + # A pool change invalidates the whole workspace, compaction included. + layout["pool_view_fingerprint"] = (("moved",),) + rebuilt_workspace = SimpleNamespace( + decode_width=workspace.decode_width, + page_table_token_capacity=workspace.page_table_token_capacity, + max_requests=workspace.max_requests, + ) + prepare.return_value = rebuilt_workspace + rebuilt = manager._workspace_for(layout, prepared) assert rebuilt is not resources - assert score_cls.call_count == 2 - assert manager._batched_compaction is None + assert rebuilt is rebuilt_workspace + assert prepare.call_count == 2 + assert manager._workspace is rebuilt_workspace 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 index 75cc5b837d7f..17ce57b990ce 100644 --- 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 @@ -16,14 +16,10 @@ import pytest import torch +from conftest import compaction_family as _compaction_family from conftest import encode_block_offsets as _encode_block_offsets -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - BatchedKVCacheCompaction, -) -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _BatchedKeepSetSelector, -) +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import build_cache_compactions from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _settle_ties_and_pack_compaction_sources_kernel, ) @@ -393,13 +389,14 @@ def test_settle_handles_topk_sentinel_padding(): assert (output[row, emitted:] == stale).all(), f"row {row} tail" -def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buffers(): - """The handoff exports the live move buffers, drops the compaction-time - dense pack launch, and the selector only accepts a packing that reads its - own keep buffer.""" +def test_pack_fusion_drops_the_compaction_dense_pack_and_exports_live_buffers(): + """Fused packing must remove the compaction-side dense pack launch (each + round packs exactly once, in the selection settle), and the exported pack + description must point at the very move buffers and keep ordinals the C++ + compacts consume.""" device = torch.device("cuda", torch.cuda.current_device()) - request_count, num_kv_heads, keep_count, width = 2, 2, 4, 16 - # BatchedKVCacheCompaction admits only bf16 pools in the compact op's + request_count, num_kv_heads, keep_count = 2, 2, 4 + # The compaction builder admits only bf16 pools in the compact op's # supported geometry (32/128-token pages, head_dim 64/128). tokens_per_block, head_dim = 32, 64 pools = [ @@ -409,26 +406,17 @@ def test_pack_handoff_disables_compaction_dense_pack_and_selector_validates_buff for _ in range(2) ] page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + keep = torch.zeros(request_count, keep_count, dtype=torch.int32, device=device) - selector = _BatchedKeepSetSelector( - eviction_mode="union", - rows=3, - width=width, - keep_count=keep_count, - dtype=torch.float32, - device=device, - max_requests=request_count, - ) - - def build_compaction(kept_token_ordinals): - return BatchedKVCacheCompaction( + def build(fuse_dense_pack_into_selection): + return build_cache_compactions( eviction_mode="union", layer_pools=pools, dense_layers=[0, 1], swa_layers=[], layer_group_representative={0: 0, 1: 1}, layer_pool_keys=[("dense", 0), ("dense", 0)], - kept_token_ordinals=kept_token_ordinals, + kept_token_ordinals=keep, valid_sequence_lengths=torch.full( (request_count,), 10, dtype=torch.int32, device=device ), @@ -439,23 +427,20 @@ def build_compaction(kept_token_ordinals): decode_keep_count=keep_count, swa_window=None, protected_tail_capacity=1, + fuse_dense_pack_into_selection=fuse_dense_pack_into_selection, ) - compaction = build_compaction(selector.keep) - assert compaction.target_dense_compaction.move_index_pack is not None - - pack_arguments = compaction.hand_move_source_pack_to_selection() - assert compaction.target_dense_compaction.move_index_pack is None - assert len(compaction.cache_compactions) == 1 - assert compaction.cache_compactions[0] is compaction.target_dense_compaction - assert pack_arguments.dense_indices is compaction.target_dense_compaction.move_source_indices - assert pack_arguments.dense_offsets is compaction.target_dense_compaction.move_source_offsets - - selector.fuse_move_source_pack(pack_arguments) - assert selector._move_source_pack is pack_arguments - - # A packing built over any other keep buffer must be rejected: the fused - # kernel reads back the ordinals it just wrote. - foreign = build_compaction(torch.zeros_like(selector.keep)) - with pytest.raises(ValueError, match="keep buffer"): - selector.fuse_move_source_pack(foreign.hand_move_source_pack_to_selection()) + standalone = build(fuse_dense_pack_into_selection=False) + standalone_dense = _compaction_family(standalone, "dense") + assert standalone_dense["pack"] is standalone["dense_pack"] + + fused = build(fuse_dense_pack_into_selection=True) + dense = _compaction_family(fused, "dense") + assert dense["pack"] is None + assert len(fused["families"]) == 1 + pack = fused["dense_pack"] + assert pack["dense_indices"] is dense["source"] + assert pack["dense_offsets"] is dense["offsets"] + # The fused settle kernel reads back the ordinals it just wrote, so the + # packing must describe the caller's own keep buffer. + assert pack["kept_token_ordinals"] is keep diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 97c2d4374c78..6338b6cd5dd9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -34,23 +34,17 @@ from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_bare_staging as _make_bare_staging from conftest import make_fake_v2 as _make_fake_v2 -from conftest import make_fixed_resources_stubs as _make_fixed_resources_stubs from conftest import make_request as _make_request from conftest import make_staging_manager as _make_staging_manager from conftest import make_triattention as _make_triattention +from conftest import make_workspace_stubs as _make_workspace_stubs from conftest import mocked_eviction_internals as _mocked_eviction_internals from conftest import torch_tri_score_oracle as _torch_tri_score_oracle from pydantic import ValidationError # 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 ( - TriAttention, - _PreparedEviction, - _PreparedGenerationBatch, - _RequestCompressionState, - _RuntimeKVLayout, -) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention # Framework base class lives in pyexecutor.resource_manager; the factory lives # in pyexecutor._util (next to _create_kv_cache_manager), matching #15106. @@ -75,11 +69,11 @@ def _set_request_state( evicted_tokens=0, confirmed_kv_length=None, ): - state = _RequestCompressionState( - generation_steps=generation_steps, - evicted_tokens=evicted_tokens, - confirmed_kv_length=confirmed_kv_length, - ) + state = { + "generation_steps": generation_steps, + "evicted_tokens": evicted_tokens, + "confirmed_kv_length": confirmed_kv_length, + } manager._request_states[request_id] = state return state @@ -94,15 +88,15 @@ def _prepared_eviction( round_start=None, prompt_len=0, ): - return _PreparedEviction( - request=request, - request_id=request_id, - seq_len=seq_len, - round_start=int(seq_len if round_start is None else round_start), - prompt_len=prompt_len, - expected_keep_count=expected_keep_count, - protected_tail=protected_tail, - ) + return { + "request": request, + "request_id": request_id, + "seq_len": seq_len, + "round_start": int(seq_len if round_start is None else round_start), + "prompt_len": prompt_len, + "expected_keep_count": expected_keep_count, + "protected_tail": protected_tail, + } @pytest.fixture @@ -185,7 +179,7 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): ) triattention = _make_triattention() triattention.kv_cache_manager = manager - cached = _RuntimeKVLayout( + cached = dict( manager=manager, num_layers=3, global_layers=[10, 11, 12], @@ -351,22 +345,22 @@ def test_suspended_cache_rejects_batch_before_cadence_mutation(self): with pytest.raises(RuntimeError, match="request 8 must be resumed"): manager._periodic_evict(batch) - assert first_state.generation_steps == 127 - assert first_state.confirmed_kv_length is None - assert second_state.generation_steps == 127 - assert second_state.confirmed_kv_length is None + assert first_state["generation_steps"] == 127 + assert first_state["confirmed_kv_length"] is None + assert second_state["generation_steps"] == 127 + assert second_state["confirmed_kv_length"] is None def test_non_boundary_step_skips_eviction_geometry(self): manager, _, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) state = manager._request_states[7] - state.generation_steps = 126 + state["generation_steps"] = 126 with mock.patch.object(manager, "_minimum_evictable_length") as keep_count: manager._periodic_evict(batch) keep_count.assert_not_called() - assert state.generation_steps == 127 - assert state.confirmed_kv_length == 1024 + 4096 + 1 + assert state["generation_steps"] == 127 + assert state["confirmed_kv_length"] == 1024 + 4096 + 1 def test_eager_eviction_runs_large_due_cohort_in_one_round(self): manager, _, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) @@ -402,23 +396,17 @@ def test_request_finish_clears_state_but_keeps_buffers_resident(self): evicted_tokens=127, confirmed_kv_length=128, ) - buffers = object() - compaction = object() - manager._eviction_resources = buffers - manager._batched_compaction = compaction - manager._prepared_generation_batch = _PreparedGenerationBatch( - batch=SimpleNamespace(), - growth_by_request={7: 1}, - ) + workspace = object() + manager._workspace = workspace + manager._prepared_generation_batch = (SimpleNamespace(), {7: 1}) manager.on_request_finish(_make_request(7)) assert manager._request_states == {} - assert manager._prepared_generation_batch.growth_by_request == {} - # The buffers are sized for the executor limits, not one cohort, so - # they stay resident for the next generation batch. - assert manager._eviction_resources is buffers - assert manager._batched_compaction is compaction + assert manager._prepared_generation_batch[1] == {} + # The workspace is sized for the executor limits, not one cohort, so + # it stays resident for the next generation batch. + assert manager._workspace is workspace @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): @@ -432,9 +420,9 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): cache = mgr.kv_cache_manager.kv_cache_map[7] cache.capacity = confirmed + tail mgr.kv_cache_manager.num_extra_kv_tokens = reserve - mgr._prepared_generation_batch = _PreparedGenerationBatch( - batch=SimpleNamespace(generation_requests=[request]), - growth_by_request={7: current_growth}, + mgr._prepared_generation_batch = ( + SimpleNamespace(generation_requests=[request]), + {7: current_growth}, ) draft_manager = _make_fake_v2(is_draft=True) draft_cache = SimpleNamespace(is_active=True, resize=mock.Mock(return_value=True)) @@ -443,7 +431,7 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): mgr.draft_kv_cache_manager = draft_manager def compact(*_args, **_kwargs): - mgr._request_states[7].confirmed_kv_length = retained + mgr._request_states[7]["confirmed_kv_length"] = retained return [(7, retained)] with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: @@ -454,7 +442,7 @@ def compact(*_args, **_kwargs): 2, protected_tail_lengths={7: tail}, ) - assert mgr._request_states[7].confirmed_kv_length == retained + assert mgr._request_states[7]["confirmed_kv_length"] == retained cache.resize.assert_called_once_with(retained + tail, None) # The draft cache shrinks in the same round, to the same retained # length plus the draft's own protected tail. @@ -492,7 +480,7 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): manager._periodic_evict(SimpleNamespace(generation_requests=[request])) - assert manager._request_states[7].confirmed_kv_length == physical_confirmed + assert manager._request_states[7]["confirmed_kv_length"] == physical_confirmed cache.resize.assert_not_called() def test_mla_selfkonly_cache_is_rejected(self): @@ -598,14 +586,16 @@ def test_prepare_snapshots_fixed_linear_generation_growth( triattention.prepare_resources(batch) - assert triattention._prepared_generation_batch.batch is batch - assert triattention._prepared_generation_batch.growth_by_request == {7: expected_growth} + assert triattention._prepared_generation_batch[0] is batch + assert triattention._prepared_generation_batch[1] == {7: expected_growth} class TestFixedScoreMetadata: @pytest.mark.parametrize("normalize_scores", [False, True]) @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) - def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize_scores): + def test_workspace_build_receives_mode_and_capacity_kwargs( + self, eviction_mode, normalize_scores + ): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module if eviction_mode == "union" and not normalize_scores: @@ -623,9 +613,9 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize eviction_mode=eviction_mode, normalize_scores=normalize_scores, ) - # The buffers follow the executor limits: eight requests (max batch + # The workspace follows the executor limits: eight requests (max batch # size) by 260 decode tokens (top_B plus two eviction periods). - layout, score_staging, keep_set_selector = _make_fixed_resources_stubs(manager) + layout, workspace = _make_workspace_stubs(manager) prepared = [ _prepared_eviction( _make_request(7), @@ -635,27 +625,28 @@ def test_fixed_buffers_bind_score_after_selection(self, eviction_mode, normalize ) ] - with ( - mock.patch.object( - module, - "_FixedScoreStagingBuffers", - return_value=score_staging, - ), - mock.patch.object( - module, - "_BatchedKeepSetSelector", - return_value=keep_set_selector, - ) as build_selection, + with mock.patch.object( + module, + "prepare_eviction_workspace", + return_value=workspace, + ) as build_workspace: + resources = manager._workspace_for(layout, prepared) + + assert resources is workspace + kwargs = build_workspace.call_args.kwargs + assert kwargs["eviction_mode"] == eviction_mode + assert kwargs["keep_count"] == 4 + assert kwargs["max_requests"] == 8 + assert kwargs["decode_width"] == 260 + assert kwargs["phase"] is manager._phase + assert kwargs["layer_pool_keys"] == list(layout["layer_pool_keys"]) + # The cached workspace serves later rounds without rebuilding. + with mock.patch.object( + module, + "prepare_eviction_workspace", + side_effect=AssertionError("workspace was rebuilt"), ): - resources = manager._fixed_resources_for(layout, prepared) - - score_staging.bind_score_launcher.assert_called_once_with( - keep_set_selector.valid_widths, - "mean", - ) - assert build_selection.call_args.kwargs["eviction_mode"] == eviction_mode - assert resources.score_staging is score_staging - assert resources.keep_set_selector is keep_set_selector + assert manager._workspace_for(layout, prepared) is workspace def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): """The bulk copy stages immutable host snapshots, and the next copy @@ -667,6 +658,11 @@ def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): subsequent bulk copy must also wait until the previous round's consumers (recorded by ``mark_page_tables_consumed``) are done. """ + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _stage_block_offsets, + mark_page_tables_consumed, + ) + device = torch.device("cuda", torch.cuda.current_device()) current_stream = torch.cuda.current_stream(device) manager_stream = torch.cuda.Stream(device=device) @@ -695,7 +691,9 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): manager = _make_staging_manager(host_table, gather, manager_stream) def stage_once(): - assert staging._stage_page_tables_bulk( + # Raises on any staging failure; success returns None. + _stage_block_offsets( + staging, manager, [7], current_stream, @@ -745,7 +743,7 @@ def stage_once(): torch.cuda._sleep(20_000_000) snapshot.copy_(staging.block_offsets_device) staging.page_tables_active = True - staging.mark_page_tables_consumed(manager_stream) + mark_page_tables_consumed(staging, manager_stream) stage_once() current_stream.synchronize() @@ -754,60 +752,48 @@ def stage_once(): assert staging.block_offsets_device[0, 0, 0, :5].tolist() == [46, 48, 50, 52, 54] def test_staged_page_tables_bypass_per_request_cuda_materialization(self): - manager = _make_triattention() - get_batch = mock.Mock() - manager.kv_cache_manager = SimpleNamespace(get_batch_cache_indices=get_batch) - staging = SimpleNamespace( - stage=mock.Mock(return_value=True), - max_requests=8, + manager = _make_triattention(top_B=4) + manager.kv_cache_manager.num_extra_kv_tokens = 3 + manager.kv_cache_manager._stream = mock.Mock() + manager.kv_cache_manager.get_batch_cache_indices = mock.Mock( + side_effect=AssertionError("eviction staged page tables per request") ) - prepared = [ - _prepared_eviction( - SimpleNamespace(), - request_id=7, - round_start=8, - seq_len=8, - prompt_len=3, - expected_keep_count=6, - protected_tail=2, - ), - _prepared_eviction( - SimpleNamespace(), - request_id=8, - round_start=9, - seq_len=9, - prompt_len=5, - expected_keep_count=6, - protected_tail=3, - ), - ] + first = _make_request(7, py_prompt_len=3) + second = _make_request(8, py_prompt_len=5) + _set_request_state(manager, 7, confirmed_kv_length=8) + _set_request_state(manager, 8, confirmed_kv_length=10) + + with _mocked_eviction_internals(manager) as internals: + manager._evict_requests( + [(first, 7), (second, 8)], + 2, + protected_tail_lengths={7: 2, 8: 3}, + ) - layout = SimpleNamespace(swa_layers=[], swa_window=None) - manager._attach_page_ids(prepared, staging, layout) - - # top_B=8: per-request moves are keep + tail = [10, 11]; padded rows - # repeat the final offset out to the request capacity. - staging.stage.assert_called_once_with( - manager.kv_cache_manager, - [7, 8], - [8, 9], - [3, 5], - [8, 9], - [10, 12], - draft_manager=None, - dense_move_offsets=[0, 10, 21, 21, 21, 21, 21, 21, 21], - swa_move_offsets=None, - draft_move_offsets=None, + # One batched staging call carries the whole cohort: request ids, + # round starts, pinned prompt lengths, valid lengths, and page-table + # lengths (valid + protected tail). top_B=4: per-request moves are + # keep + tail = [6, 7]; padded rows repeat the final offset out to + # the request capacity. + args = internals.stage.call_args + assert args.args[0] is internals.workspace + assert args.args[1] is manager.kv_cache_manager + assert args.args[2:7] == ([7, 8], [8, 10], [3, 5], [8, 10], [10, 13]) + assert args.kwargs["draft_manager"] is None + assert args.kwargs["dense_move_offsets"] == [0, 6, 13, 13, 13, 13, 13, 13, 13] + assert args.kwargs["swa_move_offsets"] is None + assert args.kwargs["draft_move_offsets"] is None + internals.consumed.assert_called_once_with( + internals.workspace, manager.kv_cache_manager._stream ) - assert all(not hasattr(item, "page_ids") for item in prepared) @requires_sm100 @pytest.mark.parametrize("request_count", [1, 7, 8]) def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, request_count): pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, - _FixedScoreStreamMismatch, + prepare_eviction_workspace, + stage_eviction_cohort, ) device = torch.device("cuda", torch.cuda.current_device()) @@ -848,34 +834,37 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re assert not freq.is_contiguous() assert not omega.is_contiguous() assert not offsets.is_contiguous() - staging = _FixedScoreStagingBuffers( - pools, - dense_groups, - [0, 1, 2], - representatives, - max_requests, - seq_len, - num_q_heads, - num_freqs, - q_real, - q_imag, - mlr, - freq, - offsets, - omega, + # The workspace constructor converts every non-contiguous fp64 + # calibration input to contiguous fp32 (the runner's flat views would + # otherwise fail) and compiles the score kernel here. + staging = prepare_eviction_workspace( + eviction_mode="per_head", + layer_pools=pools, + dense_groups=dense_groups, + dense_layers=[0, 1, 2], + page_representatives=representatives, + max_requests=max_requests, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_freqs=num_freqs, + keep_count=4, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr, + freq_scale_sq=freq, + offsets=offsets, + omega=omega, page_table_token_capacity=page_table_token_capacity, + build_compaction=False, ) assert staging.bucket_seq_len == seq_len assert staging.page_table_token_capacity == page_table_token_capacity - assert staging.page_count == page_count - assert staging.mean_phase_table.offsets.dtype == torch.float32 - assert staging.mean_phase_table.offsets.is_contiguous() - assert staging.mean_phase_table.omega.dtype == torch.float32 - assert staging.mean_phase_table.omega.is_contiguous() - fused = staging.fused_group - for calibration in (*fused.pointer_middle[2:], *fused.pointer_tail): - assert calibration.dtype == torch.float32 - assert calibration.is_contiguous() + # page_count 3 rounds up to the 4-block copy granule. + assert staging.copy_block_count == (page_count + 3) // 4 * 4 + assert staging.phase["offsets"].dtype == torch.float32 + assert staging.phase["offsets"].is_contiguous() + assert staging.phase["omega"].dtype == torch.float32 + assert staging.phase["omega"].is_contiguous() tables = { 10: [ [3 * request, 3 * request + 1, 3 * request + 2] for request in range(request_count) @@ -917,23 +906,27 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): manager = _make_staging_manager( host_table, gather, torch.cuda.Stream(device=device), num_slots=3 ) - manager.enable_swa_scratch_reuse = False - assert not staging.stage( - manager, - request_ids, - [2**31] * request_count, - token_starts, - [seq_len] * request_count, - [10] * request_count, - ) + # Round starts past the int32 metadata range fail loudly before any + # GPU work is enqueued. + with pytest.raises((RuntimeError, OverflowError, ValueError)): + stage_eviction_cohort( + staging, + manager, + request_ids, + [2**31] * request_count, + token_starts, + [seq_len] * request_count, + [10] * request_count, + ) assert gather.call_count == 0 with mock.patch.object( torch, "index_select", side_effect=AssertionError("page-table staging used torch.index_select"), ): - assert staging.stage( + stage_eviction_cohort( + staging, manager, request_ids, round_starts, @@ -966,8 +959,8 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): calls = gather.call_count other_stream = torch.cuda.Stream(device=device) with torch.cuda.stream(other_stream): - with pytest.raises(_FixedScoreStreamMismatch, match="first CUDA stream"): - staging.stage(manager, request_ids, round_starts, token_starts) + with pytest.raises(RuntimeError, match="staging CUDA stream"): + stage_eviction_cohort(staging, manager, request_ids, round_starts, token_starts) assert gather.call_count == calls @requires_sm100 @@ -983,10 +976,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun table rebind. """ pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention_kernels - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _FixedScoreGroup, - ) + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(20260707 + request_count) @@ -1030,47 +1020,46 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun offsets = torch.tensor([1.0, 2.0, 4.0], device=device) round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 round_starts = round_device[:request_count].tolist() - token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) - block_offsets = _encode_block_offsets(page_ids_3d) - group = _FixedScoreGroup( - pools, - layer_order, - max_requests, - page_count, - seq_len, - num_q_heads, - block_offsets, - layer_order, # slot i holds layer i's tables - q_real, - q_imag, - mlr, - freq, - omega, - offsets, - output_width=seq_len - prompt_len, + # One group per layer: slot i holds layer i's own block table. + ws = module.prepare_eviction_workspace( + eviction_mode="per_head", + layer_pools=pools, + dense_groups=[[layer] for layer in layer_order], + dense_layers=layer_order, + page_representatives=layer_order, + max_requests=max_requests, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_freqs=num_freqs, + keep_count=4, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr, + freq_scale_sq=freq, + offsets=offsets, + omega=omega, + decode_width=seq_len - prompt_len, + build_compaction=False, ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) - valid_widths = torch.empty(request_count, dtype=torch.int32, device=device) - mean_cos = torch.empty(request_count, num_freqs, dtype=torch.float32, device=device) - mean_sin = torch.empty_like(mean_cos) - # Rows cover this launch and the +17 round-start advance below. - mean_phase_table = triattention_kernels.MeanPhaseTable( - offsets, omega, initial_rows=int(round_device.max()) + 18 - ) - mean_phase_table.gather(round_device, mean_cos, mean_sin, request_count) + + def stage_round(): + # Stage the round metadata straight into the fixed device rows + # (the cohort staging path is covered by the staging test above). + ws.round_starts_device.copy_(round_device) + ws.valid_seq_lens_device[:request_count].copy_(valid_seq_lens) + ws.token_starts_device.fill_(prompt_len) + ws.block_offsets_device[:, :, :, :page_count].copy_(_encode_block_offsets(page_ids_3d)) + score_sentinel = -12345.0 - group.output.fill_(score_sentinel) - fixed = group.launch( - request_count, - valid_seq_lens, - valid_widths, - token_starts, - mean_cos, - mean_sin, - ).clone() - assert valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + stage_round() + ws.score_output.fill_(score_sentinel) + with mock.patch.object(module, "run_cache_compactions"): + module.run_eviction_round(ws, normalize_scores=False) + fixed = ws.score_output.clone() + assert ws.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] # The deployed fused score must agree with the independent Torch oracle # when every layer owns a distinct V2 block table. @@ -1097,7 +1086,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) round_device.add_(17) - block_offsets.copy_(_encode_block_offsets(page_ids_3d.roll(1, dims=2))) + page_ids_3d = page_ids_3d.roll(1, dims=2) valid_seq_lens.copy_( torch.tensor( [seq_len - (request + 1) % 2 for request in range(request_count)], @@ -1105,19 +1094,14 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun device=device, ) ) - mean_phase_table.gather(round_device, mean_cos, mean_sin, request_count) expected_second_widths = valid_seq_lens - prompt_len - group.output.fill_(score_sentinel) - valid_widths.fill_(-1) - second_launch = group.launch( - request_count, - valid_seq_lens, - valid_widths, - token_starts, - mean_cos, - mean_sin, - ).clone() - assert torch.equal(valid_widths, expected_second_widths) + stage_round() + ws.score_output.fill_(score_sentinel) + ws.valid_widths.fill_(-1) + with mock.patch.object(module, "run_cache_compactions"): + module.run_eviction_round(ws, normalize_scores=False) + second_launch = ws.score_output.clone() + assert torch.equal(ws.valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index d6740ea43d67..c4b80b26d1c9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -2,15 +2,20 @@ # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace + import pytest import torch from conftest import build_compaction as _build_compaction +from conftest import compaction_family as _compaction_family from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_ramp_pools as _make_ramp_pools from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _BatchedKeepSetSelector, +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import run_cache_compactions +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + prepare_per_head_scores, ) @@ -29,6 +34,105 @@ def _require_cute_topk_op() -> None: ) +def _make_selection_ws( + *, + eviction_mode, + width, + keep_count, + device, + max_requests, + num_layers=1, + num_query_heads=1, + num_kv_heads=1, +): + """Selection-only workspace: exactly the buffers the one prepare allocates + for the mode, without the CuTe score state or compaction (the settle's + pack half is compiled away).""" + ws = SimpleNamespace( + eviction_mode=eviction_mode, + device=device, + max_requests=max_requests, + decode_width=width, + keep_count=keep_count, + num_layers=num_layers, + num_q_heads=num_query_heads, + num_kv_heads=num_kv_heads, + stream=None, + ) + ws.valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) + ws.prompt_offsets = torch.zeros(max_requests, dtype=torch.int32, device=device) + if eviction_mode == "union": + ws.selection_rows_per_request = 1 + ws.row_prompt_offsets = ws.prompt_offsets + ws.combined = torch.empty((max_requests, width), dtype=torch.float32, device=device) + # Padded rows carry zero valid width; their provisional TopK entries + # must still be in-range ordinals for the finalizer's score gather. + ws.final_indices = torch.zeros((max_requests, keep_count), dtype=torch.int32, device=device) + ws.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) + ws.selection_scores_rows = ws.combined + ws.selection_row_lengths = ws.valid_widths + ws.provisional_rows = ws.final_indices + ws.keep_rows = ws.keep + else: + selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads + ws.selection_rows_per_request = selection_rows + ws.row_prompt_offsets = torch.zeros( + max_requests * selection_rows, dtype=torch.int32, device=device + ) + ws.row_mean = torch.empty( + max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device + ) + ws.row_std = torch.empty_like(ws.row_mean) + ws.selection_scores = torch.empty( + (max_requests, selection_rows, width), dtype=torch.float32, device=device + ) + ws.row_seq_lens = torch.full( + (max_requests, selection_rows), width, dtype=torch.int32, device=device + ) + ws.top_indices_i32 = torch.zeros( + (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device + ) + ws.keep = torch.empty( + (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device + ) + ws.selection_scores_rows = ws.selection_scores.view(max_requests * selection_rows, width) + ws.selection_row_lengths = ws.row_seq_lens.view(-1) + ws.provisional_rows = ws.top_indices_i32.view(-1, keep_count) + ws.keep_rows = ws.keep.view(-1, keep_count) + ws.settle_grid = (max_requests, ws.selection_rows_per_request) + placeholder = ws.selection_row_lengths + ws.settle_pack_tensors = (placeholder,) * 5 + ws.settle_pack_shape = dict( + DENSE_TOTAL=0, + SWA_TOTAL=0, + MOVE_CAPACITY=0, + NUM_KV_HEADS=1, + SWA_WINDOW=0, + UNION=False, + PER_LAYER=False, + HAS_SWA=False, + HAS_PACK=False, + ) + return ws + + +def _select_per_head(ws, scores, *, normalize_scores): + """The per-head selection flow: reduce kernels, then top-k settle.""" + prepare_per_head_scores( + scores, + ws.valid_widths, + ws.row_mean, + ws.row_std, + ws.selection_scores, + ws.row_seq_lens, + ws.max_requests, + num_kv_heads=ws.num_kv_heads, + per_layer=ws.eviction_mode == "per_layer_perhead", + normalize_scores=normalize_scores, + ) + settle_top_tokens(ws) + + 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)) @@ -94,23 +198,22 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( device = torch.device("cuda", torch.cuda.current_device()) stream = torch.cuda.Stream(device=device) with torch.cuda.stream(stream): - selector = _BatchedKeepSetSelector( + ws = _make_selection_ws( eviction_mode=eviction_mode, - dense_layers=tuple(range(layers)), - num_query_heads=query_heads, - num_kv_heads=kv_heads, width=width, keep_count=keep_count, - dtype=torch.float32, device=device, max_requests=request_count, + num_layers=layers, + num_query_heads=query_heads, + num_kv_heads=kv_heads, ) - selector.valid_widths.copy_(valid_widths.to(device)) + ws.valid_widths.copy_(valid_widths.to(device)) scores = scores_cpu.to(device) - selector.select_requests(scores, normalize_scores=normalize_scores) - first = selector.keep.cpu() - selector.select_requests(scores, normalize_scores=normalize_scores) - second = selector.keep.cpu() + _select_per_head(ws, scores, normalize_scores=normalize_scores) + first = ws.keep.cpu() + _select_per_head(ws, scores, normalize_scores=normalize_scores) + second = ws.keep.cpu() stream.synchronize() assert torch.equal(first, expected) @@ -136,25 +239,23 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid device=device, ).to(torch.float32) valid_widths = (width, width - 32) - selector = _BatchedKeepSetSelector( + ws = _make_selection_ws( eviction_mode="union", - rows=rows, width=width, keep_count=keep_count, - dtype=torch.float32, device=device, max_requests=request_count, ) - selector.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) + ws.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) # Write the shared per-request prompt lengths the way production staging - # does: fill the bound buffer, then re-expand the row-major view. - selector.prompt_offsets[:request_count].copy_( + # does: the union row-major view aliases the per-request buffer. + ws.prompt_offsets[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) - selector.refresh_row_prompt_offsets() - selector.combined.copy_(scores.amax(dim=1)) - selector.select_prepared_union_scores() - actual = selector.keep.cpu() + assert ws.row_prompt_offsets is ws.prompt_offsets + ws.combined.copy_(scores.amax(dim=1)) + settle_top_tokens(ws) + actual = ws.keep.cpu() combined = scores.amax(dim=1).cpu() for request, valid_width in enumerate(valid_widths): @@ -167,10 +268,6 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid @pytest.mark.parametrize("per_layer", [False, True]) @pytest.mark.parametrize("normalize_scores", [False, True]) def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, normalize_scores): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - prepare_per_head_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) @@ -295,7 +392,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) - compaction.compact() + run_cache_compactions(compaction) torch.cuda.synchronize(device) for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): @@ -383,7 +480,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): protected_tail_capacity=max(protected_tails), ) _set_protected_tails(cohort_compaction, protected_tails) - cohort_compaction.compact() + run_cache_compactions(cohort_compaction) expected_pools = [pool.clone() for pool in initial_pools] for request in range(request_count): @@ -398,7 +495,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): protected_tail_capacity=protected_tails[request], ) _set_protected_tails(single_compaction, [protected_tails[request]]) - single_compaction.compact() + run_cache_compactions(single_compaction) torch.cuda.synchronize(device) # The two requests own disjoint pages, so whole-pool equality proves the @@ -412,7 +509,8 @@ 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") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, + prepare_eviction_workspace, + run_eviction_round, ) device = torch.device("cuda", torch.cuda.current_device()) @@ -463,65 +561,53 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): mlr_coef[:, :, 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - score_staging = _FixedScoreStagingBuffers( - pools, - dense_groups, - dense_layers, - [0, 1], - 1, - bucket_capacity, - num_q_heads, - num_freqs, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - torch.zeros(1, dtype=torch.float32, device=device), - torch.zeros(num_freqs, dtype=torch.float32, device=device), - page_table_keys=[("pool", 0), ("pool", 1)], - num_page_table_slots=2, - ) - score_staging.block_offsets_device.zero_() - score_staging.block_offsets_device[..., :2].copy_( - _encode_block_offsets(torch.stack(page_tables)) - ) - score_staging.round_starts_device.fill_(0) - score_staging.valid_seq_lens_device.fill_(seq_len) - score_staging.token_starts_device.fill_(0) - keep_set_selector = _BatchedKeepSetSelector( + ws = prepare_eviction_workspace( eviction_mode="per_layer_perhead", - dense_layers=tuple(dense_layers), - num_query_heads=num_q_heads, - num_kv_heads=1, - width=bucket_capacity, - keep_count=keep_count, - dtype=torch.float32, - device=device, + layer_pools=pools, + dense_groups=dense_groups, + dense_layers=dense_layers, + page_representatives=[0, 1], max_requests=1, + seq_len=bucket_capacity, + num_q_heads=num_q_heads, + num_freqs=num_freqs, + keep_count=keep_count, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + offsets=torch.zeros(1, dtype=torch.float32, device=device), + omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), + page_table_keys=[("pool", 0), ("pool", 1)], + num_page_table_slots=2, + build_compaction=False, ) - score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") - - scores = score_staging.launch_prepared_score() - keep_set_selector.select_requests(scores, normalize_scores=False) - assert torch.equal(keep_set_selector.keep, expected_keep) - - batched_compaction = _build_compaction( + ws.block_offsets_device.zero_() + ws.block_offsets_device[..., :2].copy_(_encode_block_offsets(torch.stack(page_tables))) + ws.round_starts_device.fill_(0) + ws.valid_seq_lens_device.fill_(seq_len) + ws.token_starts_device.fill_(0) + + # The compaction packs its own move indices here (the settle's pack half + # is compiled away in a workspace built without compaction). + ws.compaction = _build_compaction( eviction_mode="per_layer_perhead", layer_pools=pools, dense_layers=dense_layers, layer_group_representative=layer_group_representative, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], - kept_token_ordinals=keep_set_selector.keep[:1], - valid_sequence_lengths=score_staging.valid_seq_lens_device[:1], - kv_block_offsets=score_staging.block_offsets_device, - page_table_slots=score_staging.representative_slots, + kept_token_ordinals=ws.keep[:1], + valid_sequence_lengths=ws.valid_seq_lens_device[:1], + kv_block_offsets=ws.block_offsets_device, + page_table_slots=ws.representative_slots, request_count=1, prompt_offsets=torch.zeros(1, dtype=torch.int32, device=device), decode_keep_count=keep_count, protected_tail_capacity=0, ) - _set_protected_tails(batched_compaction, [0]) - batched_compaction.compact() + _set_protected_tails(ws.compaction, [0]) + run_eviction_round(ws, normalize_scores=False) + assert torch.equal(ws.keep, expected_keep) torch.cuda.synchronize(device) for before_pool, after_pool, table, layer in zip( @@ -550,7 +636,10 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _FixedScoreStagingBuffers, + mark_page_tables_consumed, + prepare_eviction_workspace, + run_eviction_round, + stage_eviction_cohort, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -661,76 +750,53 @@ def expected_keep() -> torch.Tensor: mlr_coef[..., 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - score_staging = _FixedScoreStagingBuffers( - [pool], - [[0]], - [0], - [0], - 1, - seq_len, - num_q_heads, - num_freqs, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - torch.zeros(1, dtype=torch.float32, device=device), - torch.zeros(num_freqs, dtype=torch.float32, device=device), - page_table_keys=[("pool", 0)], - num_page_table_slots=1, - decode_width=seq_len - prompt_len, - page_table_token_capacity=seq_len + protected_tail, - ) - keep_set_selector = _BatchedKeepSetSelector( + ws = prepare_eviction_workspace( eviction_mode="union", - rows=num_q_heads, - width=seq_len - prompt_len, - keep_count=keep_count, - dtype=torch.float32, - device=device, - max_requests=1, - dense_layers=(0,), - num_query_heads=num_q_heads, - num_kv_heads=1, - prompt_offsets_buffer=score_staging.token_starts_device, - ) - score_staging.bind_score_launcher(keep_set_selector.valid_widths, "mean") - batched_compaction = _build_compaction( layer_pools=[pool], + dense_groups=[[0]], dense_layers=[0], layer_group_representative={0: 0}, layer_pool_keys=[("pool", 0)], - kept_token_ordinals=keep_set_selector.keep[:1], - valid_sequence_lengths=score_staging.valid_seq_lens_device[:1], - kv_block_offsets=score_staging.block_offsets_device, - page_table_slots=score_staging.representative_slots, - request_count=1, - prompt_offsets=score_staging.token_starts_device[:1], - decode_keep_count=keep_count, + page_representatives=[0], + max_requests=1, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_freqs=num_freqs, + keep_count=keep_count, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + offsets=torch.zeros(1, dtype=torch.float32, device=device), + omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), + page_table_keys=[("pool", 0)], + num_page_table_slots=1, + decode_width=seq_len - prompt_len, + page_table_token_capacity=seq_len + protected_tail, protected_tail_capacity=protected_tail, ) - _set_protected_tails(batched_compaction, [protected_tail]) def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) - assert score_staging.stage( + stage_eviction_cohort( + ws, manager, [request_id], [0], [prompt_len], [seq_len], [seq_len + protected_tail], + dense_move_offsets=[0, keep_count + protected_tail], ) - keep_set_selector.refresh_row_prompt_offsets() # THE union path: the fused pipeline writes normalized union rows # into ``combined``. Z-normalization is monotonic per row and all # query heads carry identical scores here, so the expected keep - # set (derived from raw scores) is unchanged. - score_staging.launch_prepared_union_fusion(keep_set_selector.combined) - keep_set_selector.select_prepared_union_scores() - selected = keep_set_selector.keep[0].clone().to(torch.long) - batched_compaction.compact() - score_staging.mark_page_tables_consumed(manager._stream) + # set (derived from raw scores) is unchanged. The settle launch + # packs the move sources and the C++ compacts run in the same + # round call; the kept ordinals stay readable afterwards. + run_eviction_round(ws, normalize_scores=True) + selected = ws.keep[0].clone().to(torch.long) + mark_page_tables_consumed(ws, manager._stream) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) after = snapshot(compacted_capacity) @@ -831,7 +897,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) - compaction.compact() + run_cache_compactions(compaction) torch.cuda.synchronize(device) for request, (valid_seq_len, tail_length) in enumerate( @@ -879,7 +945,7 @@ def test_cache_families_read_the_staged_move_offsets_rows(): compacts request slots that are not in the staged cohort; on sliding-window models the padded slots then produce negative source ordinals and an illegal memory access. Binding the staged row by - reference is part of the constructor contract. + reference is part of the builder contract. """ device = torch.device("cuda", torch.cuda.current_device()) dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) @@ -910,12 +976,8 @@ def test_cache_families_read_the_staged_move_offsets_rows(): dense_move_offsets=dense_offsets_row, swa_move_offsets=swa_offsets_row, ) - assert ( - compaction.target_dense_compaction.move_source_offsets.data_ptr() - == dense_offsets_row.data_ptr() - ) - assert compaction.target_swa_compaction is not None - assert ( - compaction.target_swa_compaction.move_source_offsets.data_ptr() - == swa_offsets_row.data_ptr() - ) + dense_family = _compaction_family(compaction, "dense") + assert dense_family["offsets"].data_ptr() == dense_offsets_row.data_ptr() + swa_family = _compaction_family(compaction, "swa") + assert swa_family is not None + assert swa_family["offsets"].data_ptr() == swa_offsets_row.data_ptr() From 90664ff8c769c15821b6a1ffa9532f6e3aa8200a Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 08:53:18 -0700 Subject: [PATCH 085/178] [None][refactor] Trust the pipeline: delete in-flow re-validation and always-true kernel flags Validation philosophy: a check whose violation already fails loudly downstream, or whose preconditions hold by construction on the straight-line flow, is deleted outright. The CuTe runner's own contract raises are the single validation layer for score geometry (bf16/page size/GQA group/pool layout); construction-guaranteed workspace, staging, and compaction pre-checks are gone, as is the bool-era cohort gauntlet. Survivors are exactly the silent-corruption class: the 2^31 scratch-plane audit, the device-indexed calibration-extent audit, the page-table reuse/consume lifecycle guards, the phase-table exact-FP32 row bound, the bucket-alignment guarantee at the bucket builder, and the config-time _validate_v2_compatibility gates. The settle-and-pack kernel drops HAS_PACK (always true in production: the selection settle always packs via the fused dense pack, and the draft pack launch packs with HAS_SETTLE=False) and OUTPUT_WIDTH (always KEEP_COUNT); the settle-only placeholder machinery and the build_compaction test knob die with it. HAS_SETTLE/UNION/PER_LAYER/ HAS_SWA remain: they vary in production. Stream-pinning machinery is deleted (kernels launch on the implicit current stream; the cross-stream staging event ordering remains as plain code). The never-read page_table_seq_lens staging parameter is deleted. Hook conformance: the fixed-linear growth snapshot moves from a prepare_resources override into the on_generation_step_begin hook (ordering-equivalent: the base template fires that hook last). Signed-off-by: tianruih --- .../triattention/compaction.py | 118 +----- .../triattention/triattention.py | 344 +++++------------- .../triattention/triattention_kernels.py | 157 ++++---- .../_torch/kv_cache_compression/conftest.py | 1 - .../test_triattention_cute_score.py | 32 +- .../test_triattention_cute_union_fusion.py | 5 +- .../test_triattention_fused_settle_pack.py | 79 +--- .../test_triattention_pipeline.py | 33 +- .../test_triattention_selection_compaction.py | 19 +- 9 files changed, 200 insertions(+), 588 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 34ce298998f6..321326fbd4b8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -33,43 +33,6 @@ import torch -_SUPPORTED_POOL_DTYPES = (torch.bfloat16,) - - -def _validated_kv_head_count( - pools: List[torch.Tensor], - layers: Tuple[int, ...], - device: torch.device, - what: str, -) -> int: - """Return the common KV-head count of one launch side's pools. - - The C++ compact op reads the interleaved V2 layout - ``[page, K/V, head, token, dim]`` and takes the KV-head count from each - launch's pool shape, so every layer on one side must agree on it. - """ - first = pools[layers[0]] - if first.ndim != 5 or first.shape[2] <= 0: - raise ValueError( - f"{what} pools must be 5-D interleaved V2 pools " - f"[pages, K/V, heads, tokens, dim]; layer {layers[0]} has shape " - f"{tuple(first.shape)}" - ) - num_kv_heads = int(first.shape[2]) - if not all( - pools[layer].ndim == 5 - and pools[layer].shape[1] == 2 - and pools[layer].device == device - and int(pools[layer].shape[2]) == num_kv_heads - and pools[layer].is_contiguous() - and pools[layer].dtype in _SUPPORTED_POOL_DTYPES - for layer in layers - ): - raise ValueError( - f"{what} requires contiguous interleaved BF16 pools with one common KV-head count" - ) - return num_kv_heads - def _make_move_buffers( index_prefix: Tuple[int, ...], @@ -99,18 +62,13 @@ def _page_table_provider( request_count: int, what: str, ) -> Callable[[int], torch.Tensor]: - """Return validated per-slot K block-offset views, cached per slot.""" + """Return per-slot K block-offset views, cached per slot.""" tables: Dict[int, torch.Tensor] = {} def page_table_for(representative: int) -> torch.Tensor: slot = page_table_slots[representative] if slot not in tables: - block_offsets = kv_block_offsets[slot, :request_count, 0] - if block_offsets.device != device or block_offsets.dtype != torch.int32: - raise ValueError(f"{what} block offsets must be int32 tensors on the pool device") - if block_offsets.ndim != 2 or block_offsets.stride(1) != 1: - raise ValueError(f"{what} K block offsets must have a contiguous block dimension") - tables[slot] = block_offsets + tables[slot] = kv_block_offsets[slot, :request_count, 0] return tables[slot] return page_table_for @@ -145,10 +103,6 @@ def _compact_groups( layers = tuple(entry[0] for entry in group_entries) pools = tuple(entry[1] for entry in group_entries) page_tables = tuple(entry[2] for entry in group_entries) - if len({int(pool.data_ptr()) for pool in pools}) != len(pools): - raise ValueError("layered compaction requires a distinct pool view for every layer") - if len({int(page_table.data_ptr()) for page_table in page_tables}) != 1: - raise ValueError("layers in one V2 pool must share one block-offset table") source_layer_indices = None if per_layer_slots is not None: source_layer_indices = torch.tensor( @@ -208,22 +162,8 @@ def build_move_pack_arguments( selection_rows = num_dense_layers * num_kv_heads else: selection_rows = num_kv_heads - selection_prefix = (request_count,) if union else (request_count, selection_rows) # Selection rows carry decode-only kept ordinals (already absolute), so # the rectangle is prompt-length independent. - expected_selection = (*selection_prefix, decode_keep_count) - if ( - request_count <= 0 - or tuple(kept_token_ordinals.shape) != expected_selection - or valid_sequence_lengths.shape != (request_count,) - ): - raise ValueError( - f"prepared compaction packing expects kept ordinals of shape " - f"{expected_selection} and one valid length per request; got " - f"{tuple(kept_token_ordinals.shape)} and " - f"{tuple(valid_sequence_lengths.shape)}" - ) - if swa_move_source_indices is not None: swa_offsets_arg = swa_move_source_offsets swa_indices_arg = swa_move_source_indices @@ -287,7 +227,6 @@ def launch_move_pack(pack: Dict[str, object]) -> None: pack["swa_indices"], WIDTH=pack["keep_count"], KEEP_COUNT=pack["keep_count"], - OUTPUT_WIDTH=pack["keep_count"], SELECTION_ROWS=pack["selection_rows"], DENSE_TOTAL=pack["dense_total"], SWA_TOTAL=pack["swa_total"], @@ -298,7 +237,6 @@ def launch_move_pack(pack: Dict[str, object]) -> None: PER_LAYER=pack["per_layer"], HAS_SWA=pack["has_swa"], HAS_SETTLE=False, - HAS_PACK=True, BLOCK=_PACK_BLOCK_TOKENS, num_warps=_PACK_NUM_WARPS, ) @@ -359,46 +297,18 @@ def build_cache_compactions( ``{"pack": dict|None, "groups": (...), "source": t, "offsets": t, "destination_bases": t}``. """ - if eviction_mode not in ("union", "per_head", "per_layer_perhead"): - raise ValueError(f"unsupported compaction mode: {eviction_mode}") - if request_count <= 0 or decode_keep_count <= 0: - raise ValueError("batched compaction requires requests and retained tokens") - if not dense_layers: - raise ValueError("batched compaction requires at least one dense layer") - if draft_layers and eviction_mode != "union": - raise ValueError("draft co-compaction supports only union eviction") - if draft_layer_pools is not None and not draft_layers: - raise ValueError("draft pools were given without any draft layers") - if not swa_layers and swa_window: - raise ValueError("swa_window was given without any SWA layers") - device = layer_pools[dense_layers[0]].device - # The move buffers are allocated on the pool device, so the selection - # tensors feeding the pack kernel must already live there. - if kept_token_ordinals.device != device: - raise ValueError("kept-token ordinals must live on the pool device") request_count = int(request_count) - if ( - prompt_offsets.shape != (request_count,) - or prompt_offsets.dtype != torch.int32 - or prompt_offsets.device != device - or not prompt_offsets.is_contiguous() - ): - raise ValueError("per-request prompt offsets do not match the cohort") decode_keep_count = int(decode_keep_count) - if protected_tail_capacity < 0: - raise ValueError("the protected-tail capacity must be non-negative") protected_tail_capacity = int(protected_tail_capacity) dense_layers = tuple(int(layer) for layer in dense_layers) swa_layers = tuple(int(layer) for layer in swa_layers) - if len(layer_pool_keys) != len(layer_pools): - raise ValueError("pool keys must match the layer-pool count") layer_pool_keys = tuple(layer_pool_keys) per_layer = eviction_mode == "per_layer_perhead" - num_kv_heads = _validated_kv_head_count( - layer_pools, (*dense_layers, *swa_layers), device, "batched compaction" - ) + # The C++ compact op takes the KV-head count from each launch's pool + # shape [pages, K/V, heads, tokens, dim]. + num_kv_heads = int(layer_pools[dense_layers[0]].shape[2]) dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer else (num_kv_heads,) dense_move_indices, dense_move_offsets = _make_move_buffers( dense_index_prefix, @@ -424,8 +334,6 @@ def build_cache_compactions( swa_move_offsets = None swa_window = 0 else: - if swa_window is None or swa_window <= 0: - raise ValueError("SWA compaction requires a valid retained window") # Per-request window validity (prompt + decode keep >= window) is # prompt-dependent and checked by the caller each round. swa_window = int(swa_window) @@ -481,25 +389,11 @@ def build_cache_compactions( ) if draft_layers: - if ( - draft_layer_pools is None - or draft_layer_group_representative is None - or draft_layer_pool_keys is None - ): - raise ValueError("draft co-compaction requires the full draft layout") - if draft_kv_block_offsets is None or draft_page_table_slots is None: - raise ValueError("draft co-compaction requires staged draft page tables") - if draft_protected_tail_capacity is not None and draft_protected_tail_capacity < 0: - raise ValueError("the draft protected-tail capacity must be non-negative") draft_tail = int(draft_protected_tail_capacity or 0) - if len(draft_layer_pool_keys) != len(draft_layer_pools): - raise ValueError("draft pool keys must match the draft layer-pool count") draft_layers = tuple(int(layer) for layer in draft_layers) # The draft forms its own launch groups so it may use a different # KV-head count than the target. - draft_num_kv_heads = _validated_kv_head_count( - draft_layer_pools, draft_layers, device, "draft co-compaction" - ) + draft_num_kv_heads = int(draft_layer_pools[draft_layers[0]].shape[2]) draft_move_indices, draft_move_offsets = _make_move_buffers( (draft_num_kv_heads,), [decode_keep_count + draft_tail] * request_count, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 72787ea0778b..eaae0a62dc39 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -96,15 +96,10 @@ # caller ever tuned it, so it is a constant rather than a constructor knob. _OFFSET_MAX_LENGTH = 65536 -# Stream-affinity contract: the staged buffers and compiled launches are bound -# to the first CUDA stream that uses them. -_STREAM_MISMATCH = "TriAttention eviction launches must stay on the staging CUDA stream" - +# Stream-affinity contract: the staged buffers and compiled launches are bound def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Tensor: """Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <=max].""" - if max_length < 1: - raise ValueError("offset_max_length must be >= 1") offsets: List[float] = [] value = 1 while value <= max_length: @@ -118,8 +113,6 @@ def _page_table_slot_layout( page_table_keys: List[object], ) -> Tuple[Dict[int, int], int]: """Map representative layers to page-table snapshot slots.""" - if len(page_table_keys) != len(page_representatives): - raise ValueError("page-table keys must match the representative count") use_pool_ids = all( isinstance(key, tuple) and len(key) == 2 @@ -142,16 +135,6 @@ def _page_table_slot_layout( return representative_slots, slot_count -def _bind_workspace_stream(ws: SimpleNamespace) -> torch.cuda.Stream: - """Bind the workspace to the current stream on first use, then enforce it.""" - stream = torch.cuda.current_stream(ws.device) - if ws.stream is None: - ws.stream = stream - elif (stream.device, stream.cuda_stream) != (ws.stream.device, ws.stream.cuda_stream): - raise RuntimeError(_STREAM_MISMATCH) - return stream - - def _allocate_page_table_plane( layer_pools: List[torch.Tensor], page_representatives: List[int], @@ -168,18 +151,8 @@ def _allocate_page_table_plane( ) if num_page_table_slots is None: num_page_table_slots = minimum_slots - if num_page_table_slots < minimum_slots: - raise ValueError(f"{what}page-table slot capacity does not cover every V2 pool") tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) - if int(layer_pools[page_representatives[0]].shape[1]) != 2: - raise ValueError(f"{what}page-table staging requires an interleaved K/V pool") page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block - if any( - (token_capacity + int(layer_pools[layer].shape[3]) - 1) // int(layer_pools[layer].shape[3]) - != page_count - for layer in page_representatives - ): - raise ValueError(f"{what}page-table staging requires a uniform page count") copy_block_count = (page_count + 3) // 4 * 4 plane_shape = (num_page_table_slots, max_requests, 2, copy_block_count) host = torch.empty(plane_shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) @@ -215,7 +188,6 @@ def prepare_eviction_workspace( decode_width: Optional[int] = None, page_table_token_capacity: Optional[int] = None, protected_tail_capacity: int = 0, - build_compaction: bool = True, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_layers: Optional[List[int]] = None, draft_layer_group_representative: Optional[Dict[int, int]] = None, @@ -228,55 +200,31 @@ def prepare_eviction_workspace( ) -> SimpleNamespace: """Build the ONE plain-namespace workspace for the whole eviction stack. - This is the single one-time constructor: geometry validation, buffer - staging, eager CuTe compilation for exactly the entries the eviction mode - launches, selection buffers, and the C++ compaction launch data. It runs - outside CUDA graph capture (compilation allocates and synchronizes) and - raises loudly on any unsupported geometry -- there is deliberately no - fallback path. The returned namespace holds tensors, events, streams, - compiled runners, and ints; all flow logic lives in - ``stage_eviction_cohort`` and ``run_eviction_round``. - - The workspace retains references to every scored layer pool: the SM100 - CuTe score kernel encodes immutable TMA descriptors from their raw device - addresses at compile time, so the pools must stay alive and stay put for - as long as the workspace launches (the V2 manager owns them for its own - lifetime in production). + The single one-time constructor: buffer staging, eager CuTe compilation + for exactly the entries the eviction mode launches, selection buffers, + and the C++ compaction launch data. Runs outside CUDA graph capture + (compilation allocates and synchronizes); the CuTe runner validates its + own geometry contract and raises loudly -- there is no fallback path. + The returned namespace holds tensors, events, compiled runners, and + ints; all flow logic lives in ``stage_eviction_cohort`` and + ``run_eviction_round``. The workspace retains references to every scored + layer pool: the compiled kernels encode immutable TMA descriptors from + their raw device addresses, so the pools must stay alive and stay put. """ from .triattention_cute_score_fused import TriAttentionCuteScoreRunner - if eviction_mode not in ("union", "per_head", "per_layer_perhead"): - raise ValueError(f"unsupported eviction mode: {eviction_mode}") - if not dense_groups or not dense_layers or not page_representatives or max_requests <= 0: - raise ValueError("fixed score metadata requires non-empty positive geometry") - grouped_layers = [layer for layers in dense_groups for layer in layers] - if ( - len(grouped_layers) != len(dense_layers) - or len(set(grouped_layers)) != len(grouped_layers) - or len(set(dense_layers)) != len(dense_layers) - or set(dense_layers) != set(grouped_layers) - ): - raise ValueError("dense layer order must cover every grouped layer exactly once") device = layer_pools[page_representatives[0]].device - if device.type != "cuda": - raise ValueError("fixed score metadata is CUDA-only") max_requests = int(max_requests) seq_len = int(seq_len) if page_table_token_capacity is None: page_table_token_capacity = seq_len - if page_table_token_capacity < seq_len: - raise ValueError("page-table capacity cannot be smaller than the score bucket") page_table_token_capacity = int(page_table_token_capacity) # Decode-width capacity of the score buffers; per-request prompt lengths # are staged runtime metadata. if decode_width is None: decode_width = seq_len - if decode_width <= 0 or decode_width > seq_len: - raise ValueError("fixed score decode width exceeds the sequence capacity") decode_width = int(decode_width) keep_count = int(keep_count) - if decode_width <= keep_count or keep_count <= 0: - raise ValueError("keep-set selection requires width > keep_count > 0") q_real = q_real.to(device=device, dtype=torch.float32).contiguous() q_imag = q_imag.to(device=device, dtype=torch.float32).contiguous() @@ -318,15 +266,6 @@ def prepare_eviction_workspace( ws.draft_representative_slots = {} ws.draft_copy_block_count = 0 if draft_layer_pools is not None: - if ( - not draft_page_representatives - or draft_page_table_keys is None - or draft_page_table_token_capacity is None - or draft_page_table_token_capacity <= 0 - ): - raise ValueError( - "draft page-table staging requires representatives, keys, and capacity" - ) ( ws.draft_representative_slots, ws.draft_copy_block_count, @@ -377,27 +316,10 @@ def prepare_eviction_workspace( # ---- score state: ONE fused group across ALL dense layers -------------- # Segments carry their own page-table slot, so distinct per-layer - # storages/block tables share a single launch. + # storages/block tables share a single launch. Pool geometry, dtype, and + # layout are validated by the CuTe runner itself below. p0 = layer_pools[dense_layers[0]] - if p0.ndim != 5: - raise ValueError("fixed score group requires HND pools") _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape - if num_q_heads % num_kv_heads: - raise ValueError("query heads must be divisible by KV heads") - if int(num_freqs) != head_dim // 2: - raise ValueError("calibration frequency count must match half the head dim") - strides = tuple(int(value) for value in p0.stride()) - element_size = p0.element_size() - for layer in dense_layers: - pool = layer_pools[layer] - if ( - tuple(pool.shape[1:]) != tuple(p0.shape[1:]) - or tuple(pool.stride()) != strides - or pool.dtype != p0.dtype - ): - raise ValueError("fixed score layers must share one uniform geometry") - if int(pool.data_ptr()) % element_size: - raise ValueError("fixed score layer base is not element-aligned") ws.num_layers = len(dense_layers) ws.num_q_heads = int(num_q_heads) ws.num_kv_heads = int(num_kv_heads) @@ -419,8 +341,6 @@ def prepare_eviction_workspace( ) block_offsets = ws.block_offsets_device slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) - if int(slots_t.max()) >= int(block_offsets.shape[0]): - raise ValueError("page table slot exceeds staged page-id planes") req_idx = torch.arange(max_requests, dtype=torch.int64, device=device).repeat_interleave( ws.num_layers ) @@ -428,43 +348,15 @@ def prepare_eviction_workspace( seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) max_segments = max_requests * ws.num_layers - # The fused kernel masks its ragged tail, but the scratch bucket contract - # stays tile-aligned (64 tokens, or one page for 128-token pages) so - # bucket geometry cannot silently admit new geometry. - score_tile_tokens = max(64, ws.tokens_per_block) - supported = ( - torch.cuda.get_device_capability(device) == (10, 0) - and p0.dtype == torch.bfloat16 - and kv_factor == 2 - and ws.tokens_per_block in (32, 128) - and ws.num_freqs in (32, 64) - and ws.num_q_heads % ws.num_kv_heads == 0 - and ws.num_q_heads // ws.num_kv_heads in (4, 8) - and int(p0.stride(-1)) == 1 - and seq_len % score_tile_tokens == 0 - # The kernel's head-plane base offset is 64-bit; the widest 32-bit - # product left is one plane (N-1 head columns of one segment stride), - # which the score bucket keeps far below 2^31. Group-4 geometries pad - # the head axis to the MMA tile N=8. - and (8 - 1) * max_segments * seq_len < 2**31 - ) - if not supported: + # Head axis pads to the MMA tile N=8; one padded scratch plane must stay + # 32-bit indexable (wraparound = silent wild reads, not a clean error). + if (8 - 1) * max_segments * seq_len >= 2**31: raise ValueError( - "TriAttention score requires SM100, bf16 KV pools, head size " - "64/128, 32/128-token pages, GQA group 4 or 8, and a bucket " - "capacity aligned to the score compute tile; got " - f"capability={torch.cuda.get_device_capability(device)}, " - f"dtype={p0.dtype}, kv_factor={kv_factor}, " - f"tokens_per_block={ws.tokens_per_block}, num_freqs={ws.num_freqs}, " - f"heads={ws.num_q_heads}q/{ws.num_kv_heads}kv, " - f"stride={int(p0.stride(-1))}, " - f"seq_len={seq_len} (tile {score_tile_tokens}), " - f"offset_audit={ws.num_kv_heads * 8 * max_segments * seq_len}" + f"score bucket overflows the 32-bit scratch plane: {(8 - 1) * max_segments * seq_len}" ) - # The kernel scores each request's window (from its staged per-request - # start) into a head-major scratch padded to the MMA tile N=8 per KV - # head; consumers read only the real heads. All buffers below are - # persistent because the compiled kernels capture their device pointers. + # The kernel scores each request's window into a head-major scratch; + # all buffers below are persistent because the compiled kernels capture + # their device pointers. ws.cute_scratch = torch.empty( ws.num_kv_heads * 8 * max_segments * seq_len, dtype=torch.float32, device=device ) @@ -514,7 +406,7 @@ def prepare_eviction_workspace( output=ws.cute_scratch, enable_partial_stats=union, ) - except (ImportError, RuntimeError, ValueError, AssertionError) as error: + except (ImportError, RuntimeError, TypeError, ValueError, AssertionError) as error: raise RuntimeError( "TriAttention CuTe score setup failed and no other score path exists" ) from error @@ -587,92 +479,64 @@ def prepare_eviction_workspace( ws.top_indices_i32.zero_() # ---- compaction launch data + settle/pack fusion ------------------------ - ws.compaction = None # One settle program per (request, selection row). ws.settle_grid = (max_requests, ws.selection_rows_per_request) - if build_compaction: - if layer_group_representative is None or layer_pool_keys is None: - raise ValueError("compaction requires the layer grouping and pool keys") - draft_kwargs = {} - if draft_layers: - draft_kwargs = dict( - draft_layer_pools=draft_layer_pools, - draft_layers=list(draft_layers), - draft_layer_group_representative=draft_layer_group_representative, - draft_layer_pool_keys=draft_layer_pool_keys, - draft_protected_tail_capacity=int(draft_protected_tail_capacity), - draft_kv_block_offsets=ws.draft_block_offsets_device, - draft_page_table_slots=ws.draft_representative_slots, - draft_move_offsets=ws.draft_move_offsets, - ) - ws.compaction = build_cache_compactions( - eviction_mode=eviction_mode, - layer_pools=layer_pools, - dense_layers=list(dense_layers), - swa_layers=list(swa_layers), - layer_group_representative=layer_group_representative, - kept_token_ordinals=ws.keep, - valid_sequence_lengths=ws.valid_seq_lens_device, - kv_block_offsets=ws.block_offsets_device, - page_table_slots=ws.representative_slots, - request_count=max_requests, - prompt_offsets=ws.token_starts_device, - decode_keep_count=keep_count, - swa_window=swa_window, - layer_pool_keys=list(layer_pool_keys), - protected_tail_capacity=int(protected_tail_capacity), - # Tails vary per round (in-flight growth), so the per-family move - # offsets ride the staged metadata rows each round. - dense_move_offsets=ws.dense_move_offsets, - swa_move_offsets=ws.swa_move_offsets, - # ONE launch settles the kept ordinals and packs the dense/SWA - # move sources; ``run_cache_compactions`` then only runs the C++ - # moves (plus the draft's own pack). - fuse_dense_pack_into_selection=True, - **draft_kwargs, - ) - pack = ws.compaction["dense_pack"] - # The fused kernel reads back the kept ordinals it just wrote, so the - # packing must read this workspace's own keep buffer. - assert ( - pack["kept_token_ordinals"].data_ptr() == ws.keep_rows.data_ptr() - and pack["selection_rows"] == ws.selection_rows_per_request - and pack["keep_count"] == keep_count - ), "fused move packing must match the selection geometry" - ws.settle_pack_tensors = ( - pack["valid_sequence_lengths"], - pack["dense_offsets"], - pack["dense_indices"], - pack["swa_offsets"], - pack["swa_indices"], - ) - ws.settle_pack_shape = dict( - DENSE_TOTAL=pack["dense_total"], - SWA_TOTAL=pack["swa_total"], - MOVE_CAPACITY=pack["move_capacity"], - NUM_KV_HEADS=pack["num_kv_heads"], - SWA_WINDOW=pack["swa_window"], - UNION=pack["union"], - PER_LAYER=pack["per_layer"], - HAS_SWA=pack["has_swa"], - HAS_PACK=True, - ) - else: - # Settle-only workspaces (unit tests): the pack half is compiled - # away, so its tensor parameters are never read. - placeholder = ws.selection_row_lengths - ws.settle_pack_tensors = (placeholder,) * 5 - ws.settle_pack_shape = dict( - DENSE_TOTAL=0, - SWA_TOTAL=0, - MOVE_CAPACITY=0, - NUM_KV_HEADS=1, - SWA_WINDOW=0, - UNION=False, - PER_LAYER=False, - HAS_SWA=False, - HAS_PACK=False, + draft_kwargs = {} + if draft_layers: + draft_kwargs = dict( + draft_layer_pools=draft_layer_pools, + draft_layers=list(draft_layers), + draft_layer_group_representative=draft_layer_group_representative, + draft_layer_pool_keys=draft_layer_pool_keys, + draft_protected_tail_capacity=int(draft_protected_tail_capacity), + draft_kv_block_offsets=ws.draft_block_offsets_device, + draft_page_table_slots=ws.draft_representative_slots, + draft_move_offsets=ws.draft_move_offsets, ) + ws.compaction = build_cache_compactions( + eviction_mode=eviction_mode, + layer_pools=layer_pools, + dense_layers=list(dense_layers), + swa_layers=list(swa_layers), + layer_group_representative=layer_group_representative, + kept_token_ordinals=ws.keep, + valid_sequence_lengths=ws.valid_seq_lens_device, + kv_block_offsets=ws.block_offsets_device, + page_table_slots=ws.representative_slots, + request_count=max_requests, + prompt_offsets=ws.token_starts_device, + decode_keep_count=keep_count, + swa_window=swa_window, + layer_pool_keys=list(layer_pool_keys), + protected_tail_capacity=int(protected_tail_capacity), + # Tails vary per round (in-flight growth), so the per-family move + # offsets ride the staged metadata rows each round. + dense_move_offsets=ws.dense_move_offsets, + swa_move_offsets=ws.swa_move_offsets, + # ONE launch settles the kept ordinals and packs the dense/SWA + # move sources; ``run_cache_compactions`` then only runs the C++ + # moves (plus the draft's own pack). + fuse_dense_pack_into_selection=True, + **draft_kwargs, + ) + pack = ws.compaction["dense_pack"] + ws.settle_pack_tensors = ( + pack["valid_sequence_lengths"], + pack["dense_offsets"], + pack["dense_indices"], + pack["swa_offsets"], + pack["swa_indices"], + ) + ws.settle_pack_shape = dict( + DENSE_TOTAL=pack["dense_total"], + SWA_TOTAL=pack["swa_total"], + MOVE_CAPACITY=pack["move_capacity"], + NUM_KV_HEADS=pack["num_kv_heads"], + SWA_WINDOW=pack["swa_window"], + UNION=pack["union"], + PER_LAYER=pack["per_layer"], + HAS_SWA=pack["has_swa"], + ) # ---- round-ordering events ---------------------------------------------- ws.copy_done = torch.cuda.Event() @@ -683,7 +547,6 @@ def prepare_eviction_workspace( ws.bulk_consume_done = torch.cuda.Event() ws.copy_pending = False ws.page_tables_active = False - ws.stream = None return ws @@ -736,7 +599,6 @@ def stage_eviction_cohort( round_starts: List[int], token_starts: List[int], seq_lens: Optional[List[int]] = None, - page_table_seq_lens: Optional[List[int]] = None, draft_manager: Optional[KVCacheManagerV2] = None, dense_move_offsets: Optional[List[int]] = None, swa_move_offsets: Optional[List[int]] = None, @@ -746,30 +608,17 @@ def stage_eviction_cohort( ``token_starts`` carries each request's pinned prompt length; the score kernel starts that request's decode window there, so the cohort may mix - prompt lengths. Raises on any invalid cohort -- there is no fallback. + prompt lengths. """ request_count = len(request_ids) - if ( - request_count == 0 - or request_count > ws.max_requests - or len(round_starts) != request_count - or len(token_starts) != request_count - ): - raise ValueError("eviction cohort does not fit the workspace request capacity") - stream = _bind_workspace_stream(ws) + stream = torch.cuda.current_stream(ws.device) + # Reuse guard: staging over a cohort whose pages are still being read + # would silently corrupt the in-flight compaction. if ws.page_tables_active: raise RuntimeError("previous page-table cohort is still active") if seq_lens is None: seq_lens = [ws.bucket_seq_len] * request_count - if page_table_seq_lens is None: - page_table_seq_lens = seq_lens - if len(seq_lens) != request_count or len(page_table_seq_lens) != request_count: - raise ValueError("eviction cohort length lists do not match the request count") - if (draft_manager is None) != (ws.draft_block_offsets_device is None): - raise RuntimeError("draft staging requires the workspace built with the draft cache") request_metadata = torch.as_tensor((round_starts, seq_lens, token_starts), dtype=torch.int32) - if min(round_starts) < 0: - raise ValueError("eviction round starts must be non-negative") # Grow the phase table while this cohort's round starts are still host # integers: a stale-capacity gather is an out-of-bounds index_select on # the device. @@ -868,7 +717,6 @@ def settle_top_tokens(ws: SimpleNamespace) -> None: *ws.settle_pack_tensors, WIDTH=ws.decode_width, KEEP_COUNT=ws.keep_count, - OUTPUT_WIDTH=ws.keep_count, SELECTION_ROWS=ws.selection_rows_per_request, **ws.settle_pack_shape, HAS_SETTLE=True, @@ -887,7 +735,6 @@ def run_eviction_round(ws: SimpleNamespace, normalize_scores: bool) -> None: capacity; padded rows past the staged cohort carry zero lengths and stay inert. """ - _bind_workspace_stream(ws) request_count = ws.max_requests num_segments = request_count * ws.num_layers union = ws.eviction_mode == "union" @@ -912,24 +759,12 @@ def run_eviction_round(ws: SimpleNamespace, normalize_scores: bool) -> None: ) ws.cute_token_starts[:request_count].copy_(ws.token_starts_device[:request_count]) if union: - if not ws.runner.supports_union_fusion(request_count): - raise RuntimeError( - f"TriAttention CuTe union fusion has no compiled variant for " - f"request_count={request_count} (capacity {ws.max_requests}) " - "and no other union path exists" - ) ws.runner.launch_union_fusion( request_count, ws.mean_cos, ws.mean_sin, ws.union_rows[:request_count] ) columns = min(ws.union_rows.shape[1], ws.combined.shape[1]) ws.combined[:request_count, :columns].copy_(ws.union_rows[:request_count, :columns]) else: - if not ws.runner.supports(request_count): - raise RuntimeError( - f"TriAttention CuTe score has no compiled variant for " - f"request_count={request_count} (capacity {ws.max_requests}) " - "and no other score path exists" - ) ws.runner.launch(request_count, ws.mean_cos, ws.mean_sin) # The kernel wrote each request's window scores (from its pinned # prompt length) into the head-major scratch, padded to the MMA @@ -1266,9 +1101,8 @@ def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) with nvtx_range_debug("triattention.generation_step_end", color="blue"): self._periodic_evict(scheduled_batch) - def prepare_resources(self, scheduled_batch: "ScheduledRequests") -> None: + def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Snapshot fixed-linear target growth; mutation remains in final update.""" - super().prepare_resources(scheduled_batch) generation_growth = {} for request in scheduled_batch.generation_requests: request_id = request.py_request_id @@ -1853,6 +1687,9 @@ def _workspace_for( # costs at most one tile of scratch per segment. score_tile_tokens = max(64, int(mgr.tokens_per_block)) seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens + # Mis-tiled buckets stripe the score scratch silently; the builder + # guarantees alignment right here. + assert seq_capacity % score_tile_tokens == 0 page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) dense_groups = list(layout["storage_groups"].values()) @@ -1879,18 +1716,6 @@ def _workspace_for( draft_protected_tail_capacity=draft_tail_capacity, ) - # One-time shape gate on the V2-allocated host block-offset tables the - # bulk staging reads every round (int32 [pools, slots, K/V, blocks]). - checked_managers = [("", mgr)] - if self.draft_kv_cache_manager is not None: - checked_managers.append(("draft ", self.draft_kv_cache_manager)) - for what, manager in checked_managers: - table = manager.host_kv_cache_block_offsets - if table.dtype != torch.int32 or table.ndim != 4 or table.shape[2] != 2: - raise RuntimeError( - f"KVCacheManagerV2 exposes an invalid {what}host block-offset table" - ) - first_pool = layout["layer_pools"][layout["dense_layers"][0]] if self._offsets is None: self._offsets = _build_geometric_offsets(_OFFSET_MAX_LENGTH, first_pool.device) @@ -2075,7 +1900,6 @@ def _evict_requests( [item["round_start"] for item in prepared], [item["prompt_len"] for item in prepared], [item["seq_len"] for item in prepared], - [item["seq_len"] + item["protected_tail"] for item in prepared], draft_manager=self.draft_kv_cache_manager, dense_move_offsets=dense_offsets, swa_move_offsets=swa_offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 69e0c5185053..7e4b6bc0a0ad 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -77,14 +77,6 @@ def build_mean_phase_table( the round starts are still host integers; the gather kernel clamps stale rows instead of faulting. """ - if ( - offsets.numel() <= 0 - or omega.numel() <= 0 - or offsets.dtype != torch.float32 - or omega.dtype != torch.float32 - or offsets.device != omega.device - ): - raise ValueError("mean-phase tables require same-device FP32 offsets and frequencies") phase: Dict[str, object] = { "offsets": offsets.contiguous(), "omega": omega.contiguous(), @@ -137,9 +129,6 @@ def gather_mean_phases( destination buffers' device pointers. CUDA-only; eviction never runs under CUDA graph capture. """ - request_count = int(request_count) - if request_count <= 0 or request_count > round_starts.numel(): - raise ValueError("phase gather request count is outside its fixed buffers") num_freqs = phase["omega"].numel() _gather_mean_phase_kernel[(request_count,)]( round_starts, @@ -283,25 +272,8 @@ def prepare_per_head_scores( """Normalize and reduce score rows for either per-head eviction mode.""" request_count = int(request_count) num_kv_heads = int(num_kv_heads) - # Essentials only: the buffers were allocated together at workspace - # construction with matching geometry; dtype/layout is the kernel contract. - assert scores.is_cuda and scores.ndim == 4 and scores.dtype == torch.float32, ( - "per-head score preparation requires CUDA FP32 [request, layer, head, token] rows" - ) - assert scores.is_contiguous() and request_count == scores.shape[0], ( - "per-head score preparation request geometry does not match" - ) _, num_layers, num_query_heads, width = scores.shape - assert num_kv_heads > 0 and num_query_heads % num_kv_heads == 0, ( - "per-head score preparation requires valid GQA geometry" - ) selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - assert ( - selection_scores.shape == (request_count, selection_rows, width) - and selection_seq_lens.shape == (request_count, selection_rows) - and row_mean.numel() >= request_count * num_layers * num_query_heads - ), "per-head score preparation buffers do not match" - stats_block = 256 rows = num_layers * num_query_heads if normalize_scores: @@ -357,7 +329,6 @@ def _settle_ties_and_pack_compaction_sources_kernel( swa_indices, WIDTH: tl.constexpr, KEEP_COUNT: tl.constexpr, - OUTPUT_WIDTH: tl.constexpr, SELECTION_ROWS: tl.constexpr, DENSE_TOTAL: tl.constexpr, SWA_TOTAL: tl.constexpr, @@ -368,23 +339,20 @@ def _settle_ties_and_pack_compaction_sources_kernel( PER_LAYER: tl.constexpr, HAS_SWA: tl.constexpr, HAS_SETTLE: tl.constexpr, - HAS_PACK: tl.constexpr, BLOCK: tl.constexpr, ): """Settle one selection row's ties, then pack its compaction move sources. - One program per (request, selection row). The first half settles the - provisional top-k: recover the threshold from the provisional - selection, count the strictly greater scores, then emit the kept - ordinals in increasing order, rebased by the row's pinned prompt - length. With ``HAS_PACK`` the same program then packs the move sources - for the packed rows this selection row feeds: the kept ordinals it - just wrote, the request's protected tail, plus the SWA rows (latest - window). Union selection has one row per request feeding every KV - head's packed row, so that single program writes all of them. - ``HAS_PACK=False`` compiles the second half away, leaving exactly the - settle stage; ``HAS_SETTLE=False`` compiles the first half away - instead, packing pre-settled ordinals read from ``output_indices`` -- + One program per (request, selection row). The settle half recovers the + threshold from the provisional top-k, counts the strictly greater + scores, then emits the kept ordinals in increasing order + (lowest-index-wins ties), rebased by the row's pinned prompt length. + The pack half then writes the move sources for the packed rows this + selection row feeds: the kept ordinals it just wrote, the request's + protected tail, plus the SWA rows (latest window). Union selection has + one row per request feeding every KV head's packed row, so that single + program writes all of them. ``HAS_SETTLE=False`` compiles the settle + half away, packing pre-settled ordinals read from ``output_indices`` -- the draft co-compaction flow, whose keep set is the target's and needs no settling. """ @@ -393,7 +361,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( row = request * SELECTION_ROWS + selection_domain row_scores = scores + row * WIDTH row_selected = provisional_indices + row * KEEP_COUNT - row_output = output_indices + row * OUTPUT_WIDTH + row_output = output_indices + row * KEEP_COUNT if HAS_SETTLE: # Scores are decode-relative; this row's pinned prompt length rebases # the emitted ordinals to absolute positions (per row, so one launch @@ -460,63 +428,62 @@ def _settle_ties_and_pack_compaction_sources_kernel( output_count += tl.sum(selected_i32) ties_seen += tl.sum(tied_i32) - if HAS_PACK: - if HAS_SETTLE: - # The emission above scatters through other lanes of this - # program; make those global stores visible to every lane - # before the pack half reads the row back. - tl.debug_barrier() - dense_begin = tl.load(dense_offsets + request) - dense_end = tl.load(dense_offsets + request + 1) - dense_count = dense_end - dense_begin - valid_len = tl.load(valid_seq_lens + request) + if HAS_SETTLE: + # The emission above scatters through other lanes of this program; + # make those global stores visible to every lane before the pack + # half reads the row back. + tl.debug_barrier() + dense_begin = tl.load(dense_offsets + request) + dense_end = tl.load(dense_offsets + request + 1) + dense_count = dense_end - dense_begin + valid_len = tl.load(valid_seq_lens + request) + if HAS_SWA: + swa_begin = tl.load(swa_offsets + request) + swa_end = tl.load(swa_offsets + request + 1) + swa_count = swa_end - swa_begin + for move_start in tl.static_range(0, MOVE_CAPACITY, BLOCK): + move = move_start + tl.arange(0, BLOCK) + selected = tl.load( + row_output + move, + mask=move < KEEP_COUNT, + other=0, + ) + dense_source = tl.where(move < KEEP_COUNT, selected, valid_len + move - KEEP_COUNT) + if UNION: + # The one union row per request feeds every KV head's packed + # row with the same move sources. + for head in tl.static_range(0, NUM_KV_HEADS): + tl.store( + dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, + dense_source, + mask=move < dense_count, + ) + else: + domain = tl.program_id(1) + dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) if HAS_SWA: - swa_begin = tl.load(swa_offsets + request) - swa_end = tl.load(swa_offsets + request + 1) - swa_count = swa_end - swa_begin - for move_start in tl.static_range(0, MOVE_CAPACITY, BLOCK): - move = move_start + tl.arange(0, BLOCK) - selected = tl.load( - row_output + move, - mask=move < KEEP_COUNT, - other=0, - ) - dense_source = tl.where(move < KEEP_COUNT, selected, valid_len + move - KEEP_COUNT) + swa_source = valid_len - SWA_WINDOW + move if UNION: - # The one union row per request feeds every KV head's packed - # row with the same move sources. for head in tl.static_range(0, NUM_KV_HEADS): tl.store( - dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, - dense_source, - mask=move < dense_count, + swa_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, + swa_source, + mask=move < swa_count, ) else: domain = tl.program_id(1) - dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move - tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) - if HAS_SWA: - swa_source = valid_len - SWA_WINDOW + move - if UNION: - for head in tl.static_range(0, NUM_KV_HEADS): - tl.store( - swa_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, - swa_source, - mask=move < swa_count, - ) + # Per-layer selection has one dense domain per (layer, + # head). SWA uses one shared source row per head, so only + # the first layer writes it. + if PER_LAYER: + write_swa = domain < NUM_KV_HEADS else: - domain = tl.program_id(1) - # Per-layer selection has one dense domain per (layer, - # head). SWA uses one shared source row per head, so only - # the first layer writes it. - if PER_LAYER: - write_swa = domain < NUM_KV_HEADS - else: - write_swa = move >= 0 - head = domain % NUM_KV_HEADS - swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move - tl.store( - swa_indices + swa_output, - swa_source, - mask=write_swa & (move < swa_count), - ) + write_swa = move >= 0 + head = domain % NUM_KV_HEADS + swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move + tl.store( + swa_indices + swa_output, + swa_source, + mask=write_swa & (move < swa_count), + ) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 85b3947d2769..74ce740feaa8 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -164,7 +164,6 @@ def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None staging.copy_block_count = copy_block_count if page_count is not None: staging.page_count = page_count - staging.stream = None staging.bulk_copy_done = torch.cuda.Event() staging.bulk_consume_done = torch.cuda.Event() staging.page_tables_active = 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 index 0ce12b5de609..950c15259283 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -12,13 +12,11 @@ originally validated 128-token-page shape) against the shared pure-PyTorch oracle, sweeps request counts up to the workspace capacity, and checks the per-request decode-width metadata the selection reduce kernels consume. The -contract tests pin the loud-failure behavior: unsupported geometry raises at -workspace construction and an oversized cohort is rejected at staging -- -there is no fallback score kernel. +contract test pins the loud-failure behavior: unsupported geometry raises +from the CuTe runner's own validation at workspace construction -- there is +no fallback score kernel. """ -from types import SimpleNamespace - import pytest import torch from conftest import encode_block_offsets as _encode_block_offsets @@ -45,7 +43,7 @@ def _make_score_workspace( decode_width=None, eviction_mode="per_head", ): - """A score-only workspace over one shared page-table slot (no compaction).""" + """A score-only workspace over one shared page-table slot.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( prepare_eviction_workspace, ) @@ -69,7 +67,8 @@ def _make_score_workspace( offsets=offsets, omega=omega, decode_width=decode_width, - build_compaction=False, + layer_group_representative={layer: 0 for layer in range(num_layers)}, + layer_pool_keys=[("pool", 0)] * num_layers, ) @@ -414,10 +413,9 @@ def test_cute_kernel_matches_torch_oracle(case): ) -# The loud-failure contract, one representative per guard family: unsupported -# geometry raises at workspace construction (the only score path compiles -# eagerly there -- no fallback), and an oversized cohort is rejected by the -# host-side staging gate before any GPU work. +# The loud-failure contract: unsupported geometry raises from the CuTe +# runner's own validation during the eager compile at workspace +# construction, surfaced as the no-fallback RuntimeError. def test_unsupported_geometry_raises_at_workspace_construction(): pytest.importorskip("cutlass") device = torch.device("cuda", torch.cuda.current_device()) @@ -431,7 +429,7 @@ def test_unsupported_geometry_raises_at_workspace_construction(): for _ in range(num_layers) ] calib = torch.randn(num_layers, 2, num_freqs, device=device) - with pytest.raises(ValueError, match="TriAttention score requires SM100"): + with pytest.raises(RuntimeError, match="no other score path exists"): _make_score_workspace( layer_pools=pools, max_requests=max_requests, @@ -445,13 +443,3 @@ def test_unsupported_geometry_raises_at_workspace_construction(): offsets=torch.tensor([1.0, 2.0], dtype=torch.float32, device=device), decode_width=page_count * tokens_per_block - 1, ) - - -def test_oversized_cohort_is_rejected_at_staging(): - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - stage_eviction_cohort, - ) - - ws = SimpleNamespace(max_requests=2) - with pytest.raises(ValueError, match="does not fit the workspace request capacity"): - stage_eviction_cohort(ws, None, [1, 2, 3], [0, 0, 0], [0, 0, 0]) 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 index e69180f26e0f..33a04fca5853 100644 --- 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 @@ -33,7 +33,7 @@ def _make_union_workspace( offsets, decode_width=None, ): - """A union-mode workspace over one shared page-table slot (no compaction). + """A union-mode workspace over one shared page-table slot. The union runner also compiles the score-only entries, so one workspace serves both the fused pipeline and the split reference leg. @@ -61,7 +61,8 @@ def _make_union_workspace( offsets=offsets, omega=omega, decode_width=decode_width, - build_compaction=False, + layer_group_representative={layer: 0 for layer in range(num_layers)}, + layer_pool_keys=[("pool", 0)] * num_layers, ) 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 index 17ce57b990ce..9bd6faa50305 100644 --- 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 @@ -228,7 +228,6 @@ def test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, k swa_fused_arg, WIDTH=width, KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, SELECTION_ROWS=selection_rows, DENSE_TOTAL=dense_total, SWA_TOTAL=swa_total if has_swa else 0, @@ -239,7 +238,6 @@ def test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, k PER_LAYER=per_layer, HAS_SWA=has_swa, HAS_SETTLE=True, - HAS_PACK=True, BLOCK=_BLOCK, num_warps=_NUM_WARPS, ) @@ -250,64 +248,6 @@ def test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, k assert torch.equal(swa_fused, swa_reference), f"SWA moves differ (seed {seed})" -def test_fused_kernel_without_pack_matches_settle_oracle(): - """``HAS_PACK=False`` must leave exactly the settle stage.""" - device = torch.device("cuda", torch.cuda.current_device()) - rows_total, width, keep_count = 6, 33, 7 - generator = torch.Generator(device=device).manual_seed(11) - scores = torch.randint( - -2, 3, (rows_total, width), generator=generator, dtype=torch.int32, device=device - ).to(torch.float32) - row_lengths = torch.tensor([0, 3, 9, 17, 33, 33], dtype=torch.int32, device=device) - row_prompt_offsets = torch.tensor([5, 0, 2, 0, 1, 4], dtype=torch.int32, device=device) - 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() - 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_fused = output_stale.clone() - placeholder = row_lengths - _settle_ties_and_pack_compaction_sources_kernel[(rows_total, 1)]( - scores, - row_lengths, - row_prompt_offsets, - provisional, - output_fused, - placeholder, - placeholder, - placeholder, - placeholder, - placeholder, - WIDTH=width, - KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, - SELECTION_ROWS=1, - DENSE_TOTAL=0, - SWA_TOTAL=0, - MOVE_CAPACITY=0, - NUM_KV_HEADS=1, - SWA_WINDOW=0, - UNION=False, - PER_LAYER=False, - HAS_SWA=False, - HAS_SETTLE=True, - HAS_PACK=False, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, - ) - torch.cuda.synchronize(device) - - assert torch.equal(output_fused, output_reference) - - def test_settle_handles_topk_sentinel_padding(): """Rows shorter than KEEP_COUNT arrive -1-padded and must settle inertly. @@ -341,32 +281,33 @@ def test_settle_handles_topk_sentinel_padding(): stale = 0x5EED output = torch.full((rows_total, keep_count), stale, dtype=torch.int32, device=device) - placeholder = row_lengths + # The pack half always runs now; zero per-request move counts mask every + # pack store off, so the settle assertions below stay byte-exact. + dense_offsets = torch.zeros(rows_total + 1, dtype=torch.int32, device=device) + dense_indices = torch.zeros(1, dtype=torch.int32, device=device) _settle_ties_and_pack_compaction_sources_kernel[(rows_total, 1)]( scores, row_lengths, row_prompt_offsets, provisional, output, - placeholder, - placeholder, - placeholder, - placeholder, - placeholder, + row_lengths, + dense_offsets, + dense_indices, + dense_offsets, + dense_indices, WIDTH=width, KEEP_COUNT=keep_count, - OUTPUT_WIDTH=keep_count, SELECTION_ROWS=1, DENSE_TOTAL=0, SWA_TOTAL=0, - MOVE_CAPACITY=0, + MOVE_CAPACITY=keep_count, NUM_KV_HEADS=1, SWA_WINDOW=0, UNION=False, PER_LAYER=False, HAS_SWA=False, HAS_SETTLE=True, - HAS_PACK=False, BLOCK=_BLOCK, num_warps=_NUM_WARPS, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 6338b6cd5dd9..a7b62fa60537 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -265,10 +265,12 @@ def test_identity_compaction_is_rejected_instead_of_published(self): class TestEvictionLifecycle: def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): - # Structural: prepare is a snapshot-only override; the eviction runs - # from the framework's final on_generation_step_end hook. - assert "prepare_resources" in TriAttention.__dict__ + # Structural: TriAttention implements hooks only, never the base + # template methods; the growth snapshot rides the step-begin hook and + # the eviction runs from the final on_generation_step_end hook. + assert "prepare_resources" not in TriAttention.__dict__ assert "update_resources" not in TriAttention.__dict__ + assert "on_generation_step_begin" in TriAttention.__dict__ assert "on_generation_step_end" in TriAttention.__dict__ manager = _make_triattention() @@ -771,14 +773,13 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ) # One batched staging call carries the whole cohort: request ids, - # round starts, pinned prompt lengths, valid lengths, and page-table - # lengths (valid + protected tail). top_B=4: per-request moves are - # keep + tail = [6, 7]; padded rows repeat the final offset out to - # the request capacity. + # round starts, pinned prompt lengths, and valid lengths. top_B=4: + # per-request moves are keep + tail = [6, 7]; padded rows repeat the + # final offset out to the request capacity. args = internals.stage.call_args assert args.args[0] is internals.workspace assert args.args[1] is manager.kv_cache_manager - assert args.args[2:7] == ([7, 8], [8, 10], [3, 5], [8, 10], [10, 13]) + assert args.args[2:6] == ([7, 8], [8, 10], [3, 5], [8, 10]) assert args.kwargs["draft_manager"] is None assert args.kwargs["dense_move_offsets"] == [0, 6, 13, 13, 13, 13, 13, 13, 13] assert args.kwargs["swa_move_offsets"] is None @@ -789,7 +790,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): @requires_sm100 @pytest.mark.parametrize("request_count", [1, 7, 8]) - def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, request_count): + def test_staging_stages_dense_and_swa_tables(self, request_count): pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( prepare_eviction_workspace, @@ -855,7 +856,8 @@ def test_staging_stages_dense_and_swa_tables_and_rejects_stream_changes(self, re offsets=offsets, omega=omega, page_table_token_capacity=page_table_token_capacity, - build_compaction=False, + layer_group_representative={0: 0, 1: 0, 2: 2}, + layer_pool_keys=[("pool", 0), ("pool", 0), ("pool", 2), ("pool", 3)], ) assert staging.bucket_seq_len == seq_len assert staging.page_table_token_capacity == page_table_token_capacity @@ -917,7 +919,6 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): [2**31] * request_count, token_starts, [seq_len] * request_count, - [10] * request_count, ) assert gather.call_count == 0 with mock.patch.object( @@ -932,7 +933,6 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): round_starts, token_starts, [seq_len] * request_count, - [10] * request_count, ) torch.cuda.current_stream(device).synchronize() assert staging.round_starts_device.untyped_storage().data_ptr() == ( @@ -956,12 +956,6 @@ def gather_k_block_offsets(source, destination, requested_ids, num_blocks): staging.block_offsets_device[slot, :request_count, 0, :page_count], expected, ) - calls = gather.call_count - other_stream = torch.cuda.Stream(device=device) - with torch.cuda.stream(other_stream): - with pytest.raises(RuntimeError, match="staging CUDA stream"): - stage_eviction_cohort(staging, manager, request_ids, round_starts, token_starts) - assert gather.call_count == calls @requires_sm100 @pytest.mark.parametrize("request_count", [1, 7, 8]) @@ -1041,7 +1035,8 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun offsets=offsets, omega=omega, decode_width=seq_len - prompt_len, - build_compaction=False, + layer_group_representative={layer: layer for layer in layer_order}, + layer_pool_keys=[("pool", layer) for layer in layer_order], ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) 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 index c4b80b26d1c9..32ba340ab148 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -100,18 +100,21 @@ def _make_selection_ws( ws.provisional_rows = ws.top_indices_i32.view(-1, keep_count) ws.keep_rows = ws.keep.view(-1, keep_count) ws.settle_grid = (max_requests, ws.selection_rows_per_request) - placeholder = ws.selection_row_lengths - ws.settle_pack_tensors = (placeholder,) * 5 + # The settle launch always packs now; zero per-request move counts mask + # every pack store off, so these workspaces stay selection-only. + zero_offsets = torch.zeros(max_requests + 1, dtype=torch.int32, device=device) + zero_lengths = torch.zeros(max_requests, dtype=torch.int32, device=device) + zero_indices = torch.zeros(1, dtype=torch.int32, device=device) + ws.settle_pack_tensors = (zero_lengths, zero_offsets, zero_indices, zero_offsets, zero_indices) ws.settle_pack_shape = dict( DENSE_TOTAL=0, SWA_TOTAL=0, - MOVE_CAPACITY=0, + MOVE_CAPACITY=keep_count, NUM_KV_HEADS=1, SWA_WINDOW=0, UNION=False, PER_LAYER=False, HAS_SWA=False, - HAS_PACK=False, ) return ws @@ -580,7 +583,8 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), page_table_keys=[("pool", 0), ("pool", 1)], num_page_table_slots=2, - build_compaction=False, + layer_group_representative=layer_group_representative, + layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], ) ws.block_offsets_device.zero_() ws.block_offsets_device[..., :2].copy_(_encode_block_offsets(torch.stack(page_tables))) @@ -588,8 +592,8 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): ws.valid_seq_lens_device.fill_(seq_len) ws.token_starts_device.fill_(0) - # The compaction packs its own move indices here (the settle's pack half - # is compiled away in a workspace built without compaction). + # This compaction packs its own move indices; the workspace's fused pack + # stays masked off (its staged move-offsets row is all zeros). ws.compaction = _build_compaction( eviction_mode="per_layer_perhead", layer_pools=pools, @@ -785,7 +789,6 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: [0], [prompt_len], [seq_len], - [seq_len + protected_tail], dense_move_offsets=[0, keep_count + protected_tail], ) # THE union path: the fused pipeline writes normalized union rows From 66fe91da0f5164a70cf22e2878aee73cc354ab07 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 09:18:10 -0700 Subject: [PATCH 086/178] [None][chore] Kernel comment clarity + one-value scaffolding removal The fused score pack drops its sixteen one-value tuning flags and everything they carried: zero-size shared-storage fields, the SharedARawAlias union, dead smem layout constructions, unreachable ternary arms, and seven unused kernel signature parameters. Every deletion is Python/const_expr-side; the generated SM100 kernels are behaviorally identical (storage read census re-verified against the kernel body). The selection kernel loses its dead compact-layout branch and construction-guaranteed ctor checks; the page-shard knob that was never overridden becomes a constant. The TMA descriptor encoder's geometry checks stay: they are the designated single validation layer. Comments across the kernel files now state plainly what each block computes, for a reader who knows CUDA but not the project history; index-layout maps, 2^31 pointer-fold notes, and alignment requirements stay as one-line contracts. Also: the provably no-op calibration key re-check is gone (both call sites construct or gate the exact key set), the twin protected-tail derivations merged into one helper, and the geometric-offset ladder builder folded into its only call site. Signed-off-by: tianruih --- .../triattention/triattention.py | 56 ++-- .../triattention_cute_score_fused.py | 256 +++--------------- .../triattention_cute_selection.py | 36 +-- .../triattention/triattention_kernels.py | 8 +- 4 files changed, 69 insertions(+), 287 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index eaae0a62dc39..233979cca1d3 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -98,16 +98,6 @@ # Stream-affinity contract: the staged buffers and compiled launches are bound -def _build_geometric_offsets(max_length: int, device: torch.device) -> torch.Tensor: - """Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <=max].""" - offsets: List[float] = [] - value = 1 - while value <= max_length: - offsets.append(float(value)) - value *= 2 - return torch.tensor(offsets, device=device, dtype=torch.float32) - - def _page_table_slot_layout( page_representatives: List[int], page_table_keys: List[object], @@ -135,6 +125,14 @@ def _page_table_slot_layout( return representative_slots, slot_count +def _protected_tail_capacity(manager: KVCacheManagerV2, what: str) -> int: + """The V2 tail (extra KV + draft reserve + 1) moved with every compaction.""" + capacity = int(manager.num_extra_kv_tokens) + int(manager._kv_reserve_draft_tokens) + 1 + if capacity <= 0: + raise RuntimeError(f"{what}KVCacheManagerV2 exposes an invalid protected-tail capacity") + return capacity + + def _allocate_page_table_plane( layer_pools: List[torch.Tensor], page_representatives: List[int], @@ -992,13 +990,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: def _draft_protected_tail_capacity(self) -> int: """Return the draft tail moved and re-reserved by every co-compression.""" - draft_manager = self.draft_kv_cache_manager - capacity = ( - int(draft_manager.num_extra_kv_tokens) + int(draft_manager._kv_reserve_draft_tokens) + 1 - ) - if capacity <= 0: - raise RuntimeError("draft KVCacheManagerV2 exposes an invalid protected-tail capacity") - return capacity + return _protected_tail_capacity(self.draft_kv_cache_manager, "draft ") def _ensure_calibrated(self) -> None: """Resolve calibration once for the first request.""" @@ -1305,14 +1297,7 @@ def _local_score_calibration( def _configured_protected_tail_capacity(self) -> int: """Return the largest target tail reserved by the native V2 lifecycle.""" - capacity = ( - int(self.kv_cache_manager.num_extra_kv_tokens) - + int(self.kv_cache_manager._kv_reserve_draft_tokens) - + 1 - ) - if capacity <= 0: - raise RuntimeError("KVCacheManagerV2 exposes an invalid protected-tail capacity") - return capacity + return _protected_tail_capacity(self.kv_cache_manager, "") def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: """Drop this request's per-request length and eviction state.""" @@ -1718,7 +1703,12 @@ def _workspace_for( first_pool = layout["layer_pools"][layout["dense_layers"][0]] if self._offsets is None: - self._offsets = _build_geometric_offsets(_OFFSET_MAX_LENGTH, first_pool.device) + # Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <= max]. + self._offsets = torch.tensor( + [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], + device=first_pool.device, + dtype=torch.float32, + ) if self._phase is None: self._phase = build_mean_phase_table( self._offsets, @@ -1956,9 +1946,7 @@ def _resolve_calibration(self) -> Dict[str, torch.Tensor]: ) raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): - calib = {k: (v.to("cuda") if torch.is_tensor(v) else v) for k, v in raw.items()} - self._validate_calibration(calib) - return calib + return {k: (v.to("cuda") if torch.is_tensor(v) else v) for k, v in raw.items()} if isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): return self._convert_official_calibration(raw) got = sorted(raw.keys()) if isinstance(raw, dict) else type(raw).__name__ @@ -2001,7 +1989,6 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: "omega": omega.to("cuda"), "freq_scale_sq": freq_scale_sq.to("cuda"), } - self._validate_calibration(calib) logger.info( f"TriAttention: converted official calibration {self.calibration_path}" f" -> E_q[L={num_layers}, H={num_heads}, F={freq_count}]" @@ -2033,12 +2020,3 @@ def _rope_tables(self, freq_count: int): omega = (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() scale_sq = 1.0 return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) - - def _validate_calibration(self, calibration: Dict[str, torch.Tensor]) -> None: - """Verify the calibration dict has the expected keys.""" - missing = _REQUIRED_CALIBRATION_KEYS - set(calibration.keys()) - if missing: - raise ValueError( - f"TriAttention calibration is missing keys: {sorted(missing)}; " - f"got {sorted(calibration.keys())}." - ) 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 index a5fa066b93b4..7e9a76aaafba 100644 --- 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 @@ -5,9 +5,8 @@ This is the ONLY score implementation: the per-head modes launch its score-only entry and union eviction launches its fused score+stats+union pipeline. It uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, -and producer-only page-ID lookahead. Geometries outside the exact contract -validated here raise loudly at setup -(``triattention.prepare_eviction_workspace``); there is no fallback path. +and producer-only page-ID lookahead. Geometry outside the exact contract +raises loudly at kernel construction; there is no fallback path. """ from __future__ import annotations @@ -82,6 +81,8 @@ def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Fl 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 _TriScoreEpilogue: @@ -133,7 +134,6 @@ def epilog_gmem_copy_and_partition( tiled_copy: cute.TiledCopy, output: cute.Tensor, epilogue_tile: cute.Tile, - _unused_smem: cute.Tensor, ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: output_epilogue = cute.flat_divide( output[((None, None), 0, 0, None, None, None)], @@ -193,11 +193,10 @@ def __init__( # cos/sin/mlr coefficient planes per frequency. self.k_coeff = 3 * num_freqs self.tokens_per_block = tokens_per_block - # One 128-token compute tile either matches a page exactly (the - # validated 128-token geometry, one TMA box per phase) or spans - # several pages (32-token pages: four page fragments per phase, - # one TMA box each into the same transaction barrier). The - # single-fragment schedule of the validated geometry is unchanged. + # One 128-token compute tile either matches a page exactly + # (128-token pages: one TMA box per phase) or spans several pages + # (32-token pages: four page fragments per phase, one TMA box each + # into the same transaction barrier). self.box_tokens = min(CTA_M, tokens_per_block) self.fragments_per_phase = CTA_M // self.box_tokens self.pages_per_tile = self.fragments_per_phase @@ -205,13 +204,8 @@ def __init__( self.tile_tokens = self.halves_per_page * CTA_M self.max_tiles = (seq_len + self.tile_tokens - 1) // self.tile_tokens - # Measured final choices that still shape layouts or generated code. + # Producer staging constants baked into the generated code. self.prefetch_depth = 4 - self.sqrt_mode = "approx" - self.k_staging_mode = "half_page_tma" - self.use_tma = True - self.cpasync_schedule = "sync_each_half" - self.split_raw_tma = True self.raw_tma_feature_extent = num_freqs # Barrier transaction bytes for one phase: the full 128-token tile # of one coefficient plane, regardless of how many page fragments @@ -219,19 +213,8 @@ def __init__( 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 self.accumulator_pipeline_stages = 1 - self.umma_accumulator_partitions = 1 - self.raw_cpasync_direct_a = True - self.weight_builder_mode = "coefficient_scalar_bf16_two_term" - self.main_operand_mode = "bf16_raw_three_term_weight" - self.numerical_policy = "three_term" - self.magnitude_residual_mode = "fp16_mma_two_term_single_commit" - self.fp16_magnitude_two_term = True - self.magnitude_sqrt_ftz = True - self.producer_page_id_prefetch = True self.producer_warp_id = 0 self.physical_threads = THREADS - self.shared_a_raw_alias = False - self.compact_token_loop = True self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape if ( @@ -291,30 +274,14 @@ def __call__( tcgen05.CtaGroup.ONE, self.mma_tiler[:2], ) - main_a_shape = ( - (CTA_M, N, self.num_freqs) - if self.main_operand_mode == "bf16_raw_three_term_weight" - else self.mma_tiler - ) - a_smem_layout = sm100_utils.make_smem_layout_a(tiled_mma, main_a_shape, cutlass.Float32, 1) - raw_bf16_a_smem_layout = sm100_utils.make_smem_layout_a( - raw_bf16_tiled_mma, - (CTA_M, N, 2 * self.num_freqs), - cutlass.BFloat16, - 1, - ) - raw_bf16_split_a_smem_layout = sm100_utils.make_smem_layout_a( + # Split raw-K transport: two K_SW64 4-KiB stages per raw-page buffer, + # one each for the real and imaginary K bands. + raw_bf16_direct_a_smem_layout = sm100_utils.make_smem_layout_a( raw_bf16_tiled_mma, (CTA_M, N, self.num_freqs), cutlass.BFloat16, 2 * RAW_PAGE_BUFFERS, ) - # The full transport uses one K_SW128 8-KiB tile. The split transport - # packs two K_SW64 4-KiB stages into that same allocation, one each for - # real and imaginary data; the compile-time schedule selects the view. - raw_bf16_direct_a_smem_layout = ( - raw_bf16_split_a_smem_layout if self.split_raw_tma else raw_bf16_a_smem_layout - ) raw_tma_smem_layout = cute.make_composed_layout( raw_bf16_direct_a_smem_layout.inner, 0, @@ -351,12 +318,6 @@ def __call__( cutlass.BFloat16, 1, ) - # The magnitude residual has only the frequency-count K. A separate - # compact descriptor lets the producer issue its UMMA steps before the - # first commit, rather than waiting for and overwriting the main A tile. - magnitude_lo_smem_layout = sm100_utils.make_smem_layout_a( - tiled_mma, (CTA_M, N, self.num_freqs), cutlass.Float32, 1 - ) magnitude_lo_tiled_mma = sm100_utils.make_trivial_tiled_mma( cutlass.Float16, tcgen05.OperandMajorMode.K, @@ -365,18 +326,6 @@ def __call__( tcgen05.CtaGroup.ONE, self.mma_tiler[:2], ) - magnitude_lo_fp16_smem_layout = sm100_utils.make_smem_layout_a( - magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.Float16, - 1, - ) - magnitude_hi_fp16_smem_layout = sm100_utils.make_smem_layout_b( - magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), - cutlass.Float16, - 1, - ) magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( magnitude_lo_tiled_mma, (CTA_M, N, self.num_freqs), @@ -389,123 +338,33 @@ def __call__( cutlass.Float16, 1, ) - main_b_shape = ( - (CTA_M, N, self.num_freqs) - if self.main_operand_mode == "bf16_raw_three_term_weight" - else self.mma_tiler - ) - b_smem_layout = sm100_utils.make_smem_layout_b(tiled_mma, main_b_shape, cutlass.Float32, 1) acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - # Keep an explicit stage mode even for the one-stage control. The - # singleton mode folds away in codegen and lets both specializations - # share the same producer/consumer slicing protocol. - self.num_accumulator_slots = ( - self.accumulator_pipeline_stages * self.umma_accumulator_partitions - ) + # One accumulator slot; the explicit slot mode keeps the shared + # producer/consumer slicing protocol and folds away in codegen. + 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) - b_hi_elements = cute.cosize(b_smem_layout.outer) * int(not self.fp16_magnitude_two_term) - b_lo_elements = cute.cosize(b_smem_layout.outer) * int( - self.numerical_policy == "three_term" and not self.fp16_magnitude_two_term - ) - magnitude_lo_elements = cute.cosize(magnitude_lo_fp16_smem_layout.outer) * int( - self.numerical_policy == "three_term" - and self.magnitude_residual_mode in ("fp16_smem", "fp16_mma_single_commit") - ) - magnitude_lo_fp32_elements = cute.cosize(magnitude_lo_smem_layout.outer) * int( - self.numerical_policy == "three_term" - and self.magnitude_residual_mode == "fp32_smem_single_commit" - ) - magnitude_hi_fp16_elements = cute.cosize(magnitude_hi_fp16_smem_layout.outer) * int( - self.numerical_policy == "three_term" - and self.magnitude_residual_mode == "fp16_mma_single_commit" - ) - a_elements = cute.cosize(a_smem_layout.outer) * int( - not self.shared_a_raw_alias and not self.fp16_magnitude_two_term - ) - raw_k_half_elements = CTA_M * 2 * self.num_freqs * RAW_PAGE_BUFFERS - raw_k_elements = raw_k_half_elements * int( - self.k_staging_mode in ("half_page_cpasync", "half_page_tma") - and not self.shared_a_raw_alias - ) - alias_a_elements = cute.cosize(a_smem_layout.outer) * int(self.shared_a_raw_alias) - alias_raw_k_elements = raw_k_half_elements * int(self.shared_a_raw_alias) - raw_bf16_a_elements = cute.cosize(raw_bf16_a_smem_layout.outer) * int( - self.main_operand_mode == "bf16_raw_three_term_weight" - and ( - not self.raw_cpasync_direct_a - or self.cpasync_schedule not in ("sync_each_half", "intra_half_overlap") - ) - ) - raw_bf16_b_elements = cute.cosize(raw_bf16_b_smem_layout.outer) * int( - self.main_operand_mode == "bf16_raw_three_term_weight" - ) - raw_bf16_b2_elements = raw_bf16_b_elements * int( - self.weight_builder_mode != "coefficient_scalar_bf16_two_term" - ) - magnitude_fp16_a_elements = cute.cosize(magnitude_fp16_a_smem_layout.outer) * int( - self.fp16_magnitude_two_term - ) - magnitude_fp16_b_elements = cute.cosize(magnitude_fp16_b_smem_layout.outer) * int( - self.fp16_magnitude_two_term - ) + # Two double-buffered raw-K stages (real+imag halves per page buffer). + raw_k_elements = CTA_M * 2 * self.num_freqs * 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 = 72 * int(self.write_partial_stats) - @cute.union - class SharedARawAlias: - # The two descriptors are byte-identical in size (8 KiB) and have - # disjoint lifetimes in the alias specialization. - sA: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, alias_a_elements], - 1024, - ] - sRawK: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, alias_raw_k_elements], - 16, - ] - @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 * int(self.use_tma), + 2 * self.raw_tma_pipeline_stages, ] tmem_holding_buf: cutlass.Int32 - sA: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, a_elements], - 1024, - ] - sB_hi: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, b_hi_elements], - 1024, - ] - sB_lo: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, b_lo_elements], - 1024, - ] - sMagnitudeLo: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float16, magnitude_lo_elements], - 1024, - ] - sMagnitudeLoFp32: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float32, magnitude_lo_fp32_elements], - 1024, - ] - sMagnitudeHiFp16: cute.struct.Align[ - cute.struct.MemRange[cutlass.Float16, magnitude_hi_fp16_elements], - 1024, - ] sRawK: cute.struct.Align[ cute.struct.MemRange[cutlass.BFloat16, raw_k_elements], 1024, ] - sRawBf16A: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, raw_bf16_a_elements], - 1024, - ] sRawBf16B0: cute.struct.Align[ cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b_elements], 1024, @@ -514,10 +373,6 @@ class SharedStorage: cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b_elements], 1024, ] - sRawBf16B2: cute.struct.Align[ - cute.struct.MemRange[cutlass.BFloat16, raw_bf16_b2_elements], - 1024, - ] sMagnitudeFp16A0: cute.struct.Align[ cute.struct.MemRange[cutlass.Float16, magnitude_fp16_a_elements], 1024, @@ -538,7 +393,6 @@ class SharedStorage: cute.struct.MemRange[cutlass.Float32, stats_scratch_elements], 16, ] - sARawAlias: SharedARawAlias self.shared_storage = SharedStorage # 64-bit: at large request counts this product exceeds 2^31 (the @@ -569,18 +423,11 @@ class SharedStorage: output, partial_stats, sum_seq, - a_smem_layout, - raw_bf16_a_smem_layout, raw_bf16_direct_a_smem_layout, - raw_bf16_split_a_smem_layout, raw_tma_smem_layout, raw_bf16_b_smem_layout, - magnitude_lo_smem_layout, - magnitude_lo_fp16_smem_layout, - magnitude_hi_fp16_smem_layout, magnitude_fp16_a_smem_layout, magnitude_fp16_b_smem_layout, - b_smem_layout, ).launch( grid=(num_ctas, 1, 1), block=(self.physical_threads, 1, 1), @@ -612,18 +459,11 @@ def kernel( output: cute.Tensor, partial_stats: cute.Tensor, sum_seq: cutlass.Int64, - a_smem_layout: cute.ComposedLayout, - raw_bf16_a_smem_layout: cute.ComposedLayout, raw_bf16_direct_a_smem_layout: cute.ComposedLayout, - raw_bf16_split_a_smem_layout: cute.ComposedLayout, raw_tma_smem_layout: cute.ComposedLayout, raw_bf16_b_smem_layout: cute.ComposedLayout, - magnitude_lo_smem_layout: cute.ComposedLayout, - magnitude_lo_fp16_smem_layout: cute.ComposedLayout, - magnitude_hi_fp16_smem_layout: cute.ComposedLayout, magnitude_fp16_a_smem_layout: cute.ComposedLayout, magnitude_fp16_b_smem_layout: cute.ComposedLayout, - b_smem_layout: cute.ComposedLayout, ): tidx, _, _ = cute.arch.thread_idx() cta_index, _, _ = cute.arch.block_idx() @@ -855,6 +695,10 @@ def kernel( num_threads=EPILOGUE_THREADS, ) cute.arch.mbarrier_init_fence() + # Build the per-(head, frequency) score coefficients: the cos/sin + # bands rotated by the mean future phase (scale*(qr*C - qi*S), + # scale*(qr*S + qi*C)) split into bf16 value+residual pairs, and the + # MLR magnitude coefficient split into fp16 pairs. for weight_round in cutlass.range_constexpr(N * self.k_coeff // THREADS): linear_index = tidx + weight_round * THREADS qg = linear_index // self.k_coeff @@ -1071,9 +915,7 @@ def kernel( ) raw_tma_producer_state.advance() - if cutlass.const_expr( - self.producer_page_id_prefetch and page_half == self.halves_per_page - 1 - ): + if cutlass.const_expr(page_half == self.halves_per_page - 1): next_page_id_lane0 = cutlass.Int32(0) if warp_idx == self.producer_warp_id: if lane_idx == 0: @@ -1257,15 +1099,13 @@ def kernel( # frequency heads take two passes. for freq_rep in cutlass.range_constexpr(self.num_freqs // 32): frequency = lane_idx + 32 * freq_rep - # Issue several independent token loads before consuming - # any of them. This bounded RMEM window is unchanged; the - # optional half-page staging only switches its K source - # from global to the single raw shared buffer. + # Stage prefetch_depth independent token loads from the + # raw-K shared buffer before consuming any of them. for token_base in cutlass.range( 0, CTA_M // (THREADS // 32), self.prefetch_depth, - unroll_full=not self.compact_token_loop, + 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) @@ -1300,18 +1140,12 @@ def kernel( imag = staged_imag[prefetch_index] norm2 = real * real + imag * imag if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): - magnitude = cute.math.sqrt( - norm2, - approx=self.sqrt_mode == "approx", - ftz=self.magnitude_sqrt_ftz, - ) + magnitude = cute.math.sqrt(norm2, approx=True, ftz=True) elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): magnitude = _sqrt_approx_ftz(norm2) else: - # DSLs with neither spelling get the plain (IEEE) - # sqrt, which is strictly MORE accurate than the - # measured approx choice; the equivalence test's - # tolerance absorbs the difference. + # Plain IEEE sqrt (more accurate than approx); + # the equivalence tolerance covers the difference. magnitude = cute.math.sqrt(norm2) magnitude_fp16_0 = cutlass.Float16(magnitude) magnitude_fp16_1 = cutlass.Float16( @@ -1364,10 +1198,9 @@ def kernel( tCrRawBf16B1[(None, None, imag_b_block, 0)], tCtAcc, ) - # Keep the full four-product control unchanged. - # The independent omit_a1b1 mode drops only the - # second-order residual product; all other FP16 - # K16 products retain their original order. + # Compensated FP16 magnitude accumulation: + # |K|*coeff = A0*B0 + A0*B1 + A1*B0 (the A1*B1 term is + # below fp32 accumulation resolution and dropped). magnitude_lo_tiled_mma.set( tcgen05.Field.ACCUMULATE, True, @@ -1434,7 +1267,7 @@ def kernel( epilogue_tidx, tCtAcc, tCgC, self.epi_tile, False ) simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( - epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile, None + epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile ) tTR_gC = tTR_gC[(None, None, None, None, None, 0, 0, 0)] tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) @@ -1707,7 +1540,6 @@ def __init__( freq_scale_sq: torch.Tensor, output: torch.Tensor, enable_partial_stats: bool = False, - small_workload_page_shards: int = 3, ) -> None: self.max_requests = int(max_requests) self.num_layers = int(num_layers) @@ -1719,11 +1551,8 @@ def __init__( self.num_kv_heads = int(num_kv_heads) self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) self.enable_partial_stats = bool(enable_partial_stats) - if small_workload_page_shards not in _SUPPORTED_PAGE_SHARDS: - raise ValueError("TriAttention CuTe score has unsupported page shards") - self.small_workload_page_shards = int(small_workload_page_shards) partial_stats_elements = ( - max_requests * num_layers * num_q_heads * self.small_workload_page_shards * 3 + max_requests * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * 3 if self.enable_partial_stats else 1 ) @@ -1798,7 +1627,7 @@ def __init__( *self._torch_tail, ) ) - variants = [(1, self.small_workload_page_shards)] + variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] if max_requests > 1: variants.append((max_requests, 2)) for request_count, page_shards in variants: @@ -1894,7 +1723,7 @@ def __init__( small_score if use_extra_score_shard else large_score ) self._page_shards[request_count] = ( - self.small_workload_page_shards if use_extra_score_shard else 2 + SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 ) if self.enable_partial_stats: self._compiled_stats[request_count] = ( @@ -2012,13 +1841,6 @@ def _launch_union_finalize( union_scores: torch.Tensor, stream: cuda.CUstream, ) -> None: - if ( - union_scores.shape != (request_count, self.width) - or union_scores.dtype != torch.float32 - or union_scores.device != self.partial_stats.device - or not union_scores.is_contiguous() - ): - raise ValueError("TriAttention union output does not match the compiled geometry") self._compiled_normalize_union[request_count]( _to_cute(self.partial_stats), *self._cute_selection_prefix, 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 index 304778f061c4..bf6cff2ae77f 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -36,7 +36,6 @@ _SMALL_TILE_RESIDENT_CTAS_PER_SM = 6 _STATS_FIELDS = 3 _STD_EPSILON = 1.0e-6 -_SUPPORTED_PAGE_SHARDS = (2, 3) # The fused score kernel pads each KV head's group of score planes up to the # minimum tcgen05 MMA tile (GQA groups below 8 ride zero-padded columns). _PADDED_HEAD_COLUMNS = 8 @@ -122,40 +121,17 @@ def __init__( num_layers: int, seq_len: int, num_q_heads: int, - num_kv_heads: int | None = None, + num_kv_heads: int, page_shards: int, tokens_per_lane: int, token_subtiles: int, row_cluster_ctas: int, ) -> None: - if min(num_layers, seq_len, num_q_heads) <= 0: - raise ValueError("TriAttention stats/union reduction requires positive geometry") - if page_shards not in _SUPPORTED_PAGE_SHARDS: - raise ValueError("TriAttention stats/union has unsupported page shards") - # ``num_kv_heads`` declares that the score scratch pads each KV - # head's group of head planes up to the MMA tile: real head row - # ``q_head`` then lives in scratch plane ``kv * 8 + qg``. Omitting - # it keeps the compact plane-per-head layout (GQA group 8 is - # identical either way). The partial-stats rows are always compact. - if num_kv_heads is None: - self.score_group_size = num_q_heads - self.score_head_pad = 0 - else: - if num_kv_heads <= 0 or num_q_heads % num_kv_heads: - raise ValueError("TriAttention stats/union requires uniform GQA groups") - self.score_group_size = num_q_heads // num_kv_heads - if self.score_group_size > _PADDED_HEAD_COLUMNS: - raise ValueError( - "TriAttention stats/union supports GQA groups up to the " - f"padded head tile ({_PADDED_HEAD_COLUMNS})" - ) - self.score_head_pad = _PADDED_HEAD_COLUMNS - self.score_group_size - if tokens_per_lane not in (_SMALL_TOKENS_PER_LANE, _LARGE_TOKENS_PER_LANE): - raise ValueError("TriAttention stats/union has unsupported load width") - if token_subtiles not in (_SMALL_TOKEN_SUBTILES, _LARGE_TOKEN_SUBTILES): - raise ValueError("TriAttention stats/union has unsupported token subtiles") - if row_cluster_ctas not in (1, 2, 4): - raise ValueError("TriAttention stats/union has unsupported row cluster") + # The score scratch pads each KV head's group of head planes up to + # the MMA tile: real head row ``q_head`` lives in scratch plane + # ``kv * 8 + qg``. The partial-stats rows are always 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.seq_len = seq_len # The score window start is per-request runtime metadata diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 7e4b6bc0a0ad..ed67d5530272 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -201,7 +201,13 @@ def _score_per_head_reduce_kernel( NORMALIZE: tl.constexpr, BLOCK: tl.constexpr, ): - """Reduce query-head score rows into one selector row per KV-head domain.""" + """Reduce query-head score rows into one selector row per KV-head domain. + + per_layer: row (layer, kv_head) = max over the KV head's query group. + Otherwise: row kv_head = mean over layers of that per-layer group max. + Optionally z-normalizes each query-head row with the precomputed + mean/inv-std before reducing. + """ request = tl.program_id(0) selection_row = tl.program_id(1) token_block = tl.program_id(2) From c80d71441083fec762e55bb052ad91053ab39de5 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 09:30:30 -0700 Subject: [PATCH 087/178] [None][test] Collapse corner grids and mock-theater to representative coverage Pipeline tests lose the flows whose assertions are implied by stronger end-to-end neighbors (staged-tables semantics, one-batched-call duplicates, identity-gate and cadence permutations, per-family resize duplicates) and the workspace-kwarg matrix collapses its redundant axes. Corner grids keep boundary plus production-shaped rows: the fused-vs-split union matrix 15 -> 8 (production Qwen3/GPT-OSS shapes and the group-4 padding special kept), the CuTe runner matrix 5 -> 3, the heavy-ties sweep keeps the smallest and the 8k incident row. Raise-path coverage thins to one test per surviving guard family. All incident anchors stay: giant-scratch with cap-32 control, verbatim sentinel settle, per-request score-start rows, bucket asserts, CUDA-graph replay, draft byte-compares, the live two-round page-reuse test, and the per-head/per-layer oracles. Signed-off-by: tianruih --- .../test_triattention_cute_score.py | 2 - .../test_triattention_cute_union_fusion.py | 13 +- .../test_triattention_draft_cocompaction.py | 19 +- .../test_triattention_fused_settle_pack.py | 60 --- .../test_triattention_pipeline.py | 343 ++---------------- .../test_triattention_selection_compaction.py | 54 +-- 6 files changed, 45 insertions(+), 446 deletions(-) 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 index 950c15259283..6dd917f93bdc 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -129,13 +129,11 @@ def _launch_split_scores( (128, [0, 1], None, 32, 8), # GPT-OSS geometry: 32-token pages; a 64-token compute tile spans two # pages, so a shuffled physical-page table catches fragment mix-ups. - (32, [3, 1, 4, 7, 5, 0, 2, 6], None, 32, 8), # Ragged tails land mid-tile: the second page fragment of the last # tile is clamped, and scores past the valid length are unspecified. (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198], 32, 8), # Qwen3 geometry: 128-element K rows (64 frequencies) and GQA group # 4, which rides the MMA tile N=8 with zeroed padding columns. - (32, [3, 1, 4, 7, 5, 0, 2, 6], None, 64, 4), (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198], 64, 4), ], ) 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 index 33a04fca5853..eaddfd382bf5 100644 --- 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 @@ -262,20 +262,14 @@ def _check_union_fusion_matches_split_pipeline( @pytest.mark.parametrize( "tokens_per_block,num_freqs,num_q_heads,score_starts,valid_lens", [ - # The originally validated geometry: 32 frequencies (64-element K - # rows), GQA group 8, across full-range, non-page-aligned ragged, - # and page-aligned uniform window starts. + # 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), - (32, 32, 8, 37, [250, 198]), - (32, 32, 8, 128, [250, 230]), - (128, 32, 8, 0, None), - (128, 32, 8, 37, [250, 198]), (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, 0, None), (32, 64, 4, 37, [250, 198]), - (128, 64, 4, 0, None), (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 @@ -286,7 +280,6 @@ def _check_union_fusion_matches_split_pipeline( # start mid-tile, one page-aligned) — the case the fused pipeline # previously declined. (32, 32, 8, [37, 128], [250, 198]), - (128, 32, 8, [37, 128], None), (32, 64, 4, [37, 128], None), (128, 64, 4, [37, 128], [250, 230]), ], 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 index 0838aa0ff730..8518800a5dc7 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -222,12 +222,9 @@ def test_mark_page_tables_consumed_orders_both_manager_streams(): # One representative per guard family (the per-mode/per-config # variants raise through the same checks). ("union_only_per_head", "union"), - ("union_only_per_layer", "union"), ("draft_kv_factor", "standard key/value cache"), ("full_attention_draft", "full-attention draft"), ("callsite_dflash", "standard paged cache compacted together"), - ("callsite_draft_target", "standard paged cache compacted together"), - ("callsite_pard", "standard paged cache compacted together"), ], ) def test_draft_admission_gates_raise(gate, match): @@ -241,19 +238,10 @@ def test_draft_admission_gates_raise(gate, match): from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, - DraftTargetDecodingConfig, - PARDDecodingConfig, TriAttentionKvCacheCompressionConfig, ) - spec_config = { - "callsite_dflash": lambda: DFlashDecodingConfig(max_draft_len=3), - "callsite_draft_target": lambda: DraftTargetDecodingConfig( - max_draft_len=3, - speculative_model="/tmp/draft-target-model", - ), - "callsite_pard": lambda: PARDDecodingConfig(max_draft_len=3), - }[gate]() + spec_config = DFlashDecodingConfig(max_draft_len=3) with pytest.raises(ValueError, match=match): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( @@ -267,10 +255,7 @@ def test_draft_admission_gates_raise(gate, match): _make_fake_v2(), top_B=8, model_path="/models/test", - eviction_mode={ - "union_only_per_head": "per_head", - "union_only_per_layer": "per_layer_perhead", - }.get(gate, "union"), + eviction_mode="per_head" if gate == "union_only_per_head" else "union", draft_kv_cache_manager=draft_manager, ) if gate == "draft_kv_factor": 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 index 9bd6faa50305..8737bfaff447 100644 --- 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 @@ -16,10 +16,7 @@ import pytest import torch -from conftest import compaction_family as _compaction_family -from conftest import encode_block_offsets as _encode_block_offsets -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import build_cache_compactions from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _settle_ties_and_pack_compaction_sources_kernel, ) @@ -328,60 +325,3 @@ def test_settle_handles_topk_sentinel_padding(): ) assert torch.equal(output[row, :emitted], expected_row), f"row {row}" assert (output[row, emitted:] == stale).all(), f"row {row} tail" - - -def test_pack_fusion_drops_the_compaction_dense_pack_and_exports_live_buffers(): - """Fused packing must remove the compaction-side dense pack launch (each - round packs exactly once, in the selection settle), and the exported pack - description must point at the very move buffers and keep ordinals the C++ - compacts consume.""" - device = torch.device("cuda", torch.cuda.current_device()) - request_count, num_kv_heads, keep_count = 2, 2, 4 - # The compaction builder admits only bf16 pools in the compact op's - # supported geometry (32/128-token pages, head_dim 64/128). - tokens_per_block, head_dim = 32, 64 - pools = [ - torch.zeros( - 6, 2, num_kv_heads, tokens_per_block, head_dim, dtype=torch.bfloat16, device=device - ) - for _ in range(2) - ] - page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - keep = torch.zeros(request_count, keep_count, dtype=torch.int32, device=device) - - def build(fuse_dense_pack_into_selection): - return build_cache_compactions( - eviction_mode="union", - layer_pools=pools, - dense_layers=[0, 1], - swa_layers=[], - layer_group_representative={0: 0, 1: 1}, - layer_pool_keys=[("dense", 0), ("dense", 0)], - kept_token_ordinals=keep, - valid_sequence_lengths=torch.full( - (request_count,), 10, dtype=torch.int32, device=device - ), - kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), - page_table_slots={0: 0, 1: 0}, - request_count=request_count, - prompt_offsets=torch.zeros(request_count, dtype=torch.int32, device=device), - decode_keep_count=keep_count, - swa_window=None, - protected_tail_capacity=1, - fuse_dense_pack_into_selection=fuse_dense_pack_into_selection, - ) - - standalone = build(fuse_dense_pack_into_selection=False) - standalone_dense = _compaction_family(standalone, "dense") - assert standalone_dense["pack"] is standalone["dense_pack"] - - fused = build(fuse_dense_pack_into_selection=True) - dense = _compaction_family(fused, "dense") - assert dense["pack"] is None - assert len(fused["families"]) == 1 - pack = fused["dense_pack"] - assert pack["dense_indices"] is dense["source"] - assert pack["dense_offsets"] is dense["offsets"] - # The fused settle kernel reads back the ordinals it just wrote, so the - # packing must describe the caller's own keep buffer. - assert pack["kept_token_ordinals"] is keep diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index a7b62fa60537..41822ca07c7c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -149,7 +149,6 @@ def test_llm_args_dispatch_and_validation(self): TriAttentionKvCacheCompressionConfig(eviction_mode="made_up_mode") def test_factory_returns_triattention_and_propagates_config_fields(self): - # A plain V2 manager (block reuse off) yields a TriAttention instance. # Calibration is deferred to the first request, so construction needs # no calibration file or CUDA. fake_v2 = _make_fake_v2(enable_block_reuse=False) @@ -214,16 +213,13 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): mock.call(101, Role.KEY), ] - @pytest.mark.parametrize("num_extra_kv_tokens,reserved_draft", [(0, 0), (4, 4)]) - def test_request_init_marks_capacity_only_and_tracks_state( - self, num_extra_kv_tokens, reserved_draft - ): + def test_request_init_marks_capacity_only_and_tracks_state(self): # Speculative capacity (extra KV tokens / reserved draft width) is # accepted at request init; the target manager is marked so V2 sizing # keeps logical max_seq_len while capacity is reclaimed and reused. manager = _make_fake_v2() - manager.num_extra_kv_tokens = num_extra_kv_tokens - manager._kv_reserve_draft_tokens = reserved_draft + manager.num_extra_kv_tokens = 4 + manager._kv_reserve_draft_tokens = 4 triattention = TriAttention(manager, top_B=8, model_path="/models/test") triattention._attention_layer_partition_cache = ([], [], None) triattention._calibrated = True @@ -317,25 +313,6 @@ def _make_due_decode_request(seq_len): mgr.top_B = 4096 return mgr, request, batch - def test_identity_gate_preserves_real_eviction_round(self): - mgr, request, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) - cache = mgr.kv_cache_manager.kv_cache_map[7] - - def compact(*args, protected_tail_lengths, **_kwargs): - assert protected_tail_lengths == {7: 0} - return [(7, 1024 + 4096)] - - with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: - mgr._periodic_evict(batch) - - evict.assert_called_once_with( - [(request, 7)], - 2, - protected_tail_lengths={7: 0}, - ) - mgr.kv_cache_manager._stream.wait_event.assert_not_called() - cache.resize.assert_called_once_with(1024 + 4096, None) - def test_suspended_cache_rejects_batch_before_cadence_mutation(self): manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) second_request = _make_request(8, py_prompt_len=1024) @@ -352,43 +329,6 @@ def test_suspended_cache_rejects_batch_before_cadence_mutation(self): assert second_state["generation_steps"] == 127 assert second_state["confirmed_kv_length"] is None - def test_non_boundary_step_skips_eviction_geometry(self): - manager, _, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) - state = manager._request_states[7] - state["generation_steps"] = 126 - - with mock.patch.object(manager, "_minimum_evictable_length") as keep_count: - manager._periodic_evict(batch) - - keep_count.assert_not_called() - assert state["generation_steps"] == 127 - assert state["confirmed_kv_length"] == 1024 + 4096 + 1 - - def test_eager_eviction_runs_large_due_cohort_in_one_round(self): - manager, _, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) - requests = [] - caches = {} - for request_id in range(65): - request = _make_request(request_id, py_prompt_len=1024) - requests.append(request) - caches[request_id] = SimpleNamespace( - capacity=1024 + 4096 + 1, - history_length=1024, - is_active=True, - ) - _set_request_state(manager, request_id, generation_steps=127) - manager.kv_cache_manager.kv_cache_map = caches - batch = SimpleNamespace(generation_requests=requests) - - with ( - mock.patch.object(manager, "_evict_requests", return_value=[]) as evict, - mock.patch.object(manager, "_resize_compacted_requests") as resize, - ): - manager._periodic_evict(batch) - - assert [len(call.args[0]) for call in evict.call_args_list] == [65] - assert resize.call_count == 1 - def test_request_finish_clears_state_but_keeps_buffers_resident(self): manager = _make_triattention() _set_request_state( @@ -450,15 +390,6 @@ def compact(*_args, **_kwargs): # length plus the draft's own protected tail. draft_cache.resize.assert_called_once_with(retained + 1, None) - def test_missing_draft_cache_fails_the_due_eviction_round(self): - mgr, request, batch = self._make_due_decode_request(seq_len=1024 + 4096 + 1) - mgr.draft_kv_cache_manager = _make_fake_v2(is_draft=True) - - with mock.patch.object(mgr, "_evict_requests") as evict: - with pytest.raises(RuntimeError, match="missing or.*suspended draft KV cache"): - mgr._periodic_evict(batch) - evict.assert_not_called() - def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): physical_confirmed = 6100 manager = _make_triattention(beta=128) @@ -509,59 +440,18 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): manager._validate_v2_compatibility() assert manager.kv_cache_manager.kv_compression_manages_history is True assert draft_manager.kv_compression_manages_history is True + # A one-model MTP contract also passes the call-site speculative gate. + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig - def test_resize_shrinks_draft_cache_with_its_own_protected_tail(self): - retained = 1024 + 4096 - manager = _make_triattention() - target_cache = SimpleNamespace( - capacity=retained + 10, - is_active=True, - resize=mock.Mock(return_value=True), - ) - manager.kv_cache_manager = SimpleNamespace( - kv_cache_map={7: target_cache}, - ) - draft_manager = _make_fake_v2(is_draft=True) - draft_manager.num_extra_kv_tokens = 2 - draft_manager._kv_reserve_draft_tokens = 3 - draft_cache = SimpleNamespace(is_active=True, resize=mock.Mock(return_value=True)) - draft_manager.kv_cache_map = {7: draft_cache} - manager.draft_kv_cache_manager = draft_manager - - manager._resize_compacted_requests([(7, retained)], {7: 4}) - - target_cache.resize.assert_called_once_with(retained + 4, None) - # Draft protected tail = num_extra_kv_tokens + reserved draft width + 1. - draft_cache.resize.assert_called_once_with(retained + 6, None) - - def test_mtp_eagle_paged_draft_length_contract_is_accepted(self): - # A one-model MTP contract passes the call-site speculative gate, and - # the factory then builds a manager that validates cleanly. - from tensorrt_llm._torch.pyexecutor._util import ( - create_kv_cache_compression_manager, - validate_kv_cache_compression_with_spec, - ) - from tensorrt_llm.llmapi.llm_args import ( - MTPDecodingConfig, - TriAttentionKvCacheCompressionConfig, - ) - - config = TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 - ) - assert config.kv_cache_compression_mode.is_eviction_method() is True - draft_manager = _make_fake_v2(is_draft=True) validate_kv_cache_compression_with_spec( - config, MTPDecodingConfig(max_draft_len=1), draft_manager - ) - manager = create_kv_cache_compression_manager( - config, - _make_fake_v2(), - draft_kv_cache_manager=draft_manager, + TriAttentionKvCacheCompressionConfig( + model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + ), + MTPDecodingConfig(max_draft_len=1), + draft_manager, ) - manager._validate_v2_compatibility() - @pytest.mark.parametrize( "num_extra_kv_tokens,reserved_draft,draft_tokens,expected_growth", [ @@ -593,28 +483,16 @@ def test_prepare_snapshots_fixed_linear_generation_growth( class TestFixedScoreMetadata: - @pytest.mark.parametrize("normalize_scores", [False, True]) @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) - def test_workspace_build_receives_mode_and_capacity_kwargs( - self, eviction_mode, normalize_scores - ): + def test_workspace_build_receives_mode_and_capacity_kwargs(self, eviction_mode): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - if eviction_mode == "union" and not normalize_scores: + if eviction_mode == "union": # The fused pipeline (THE union path) always z-normalizes, so - # this combination is rejected loudly at construction. + # normalize_scores=False is rejected loudly at construction. with pytest.raises(ValueError, match="normalize_scores=True"): - _make_triattention( - top_B=4, - eviction_mode=eviction_mode, - normalize_scores=normalize_scores, - ) - return - manager = _make_triattention( - top_B=4, - eviction_mode=eviction_mode, - normalize_scores=normalize_scores, - ) + _make_triattention(top_B=4, eviction_mode="union", normalize_scores=False) + manager = _make_triattention(top_B=4, eviction_mode=eviction_mode) # The workspace follows the executor limits: eight requests (max batch # size) by 260 decode tokens (top_B plus two eviction periods). layout, workspace = _make_workspace_stubs(manager) @@ -650,6 +528,24 @@ def test_workspace_build_receives_mode_and_capacity_kwargs( ): assert manager._workspace_for(layout, prepared) is workspace + def test_stage_rejects_int32_overflowing_round_starts(self): + # Round starts past the int32 metadata range fail loudly (in the host + # metadata build) before any GPU work is enqueued. + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + stage_eviction_cohort, + ) + + device = torch.device("cuda", torch.cuda.current_device()) + staging = _make_bare_staging(device, max_requests=1, copy_block_count=8) + gather = mock.Mock() + manager = _make_staging_manager( + torch.zeros(1, 2, 2, 12, dtype=torch.int32), gather, torch.cuda.Stream(device=device) + ) + + with pytest.raises((RuntimeError, OverflowError, ValueError)): + stage_eviction_cohort(staging, manager, [7], [2**31], [0], [64]) + assert gather.call_count == 0 + def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): """The bulk copy stages immutable host snapshots, and the next copy waits for the previous cohort's consumers. @@ -789,176 +685,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ) @requires_sm100 - @pytest.mark.parametrize("request_count", [1, 7, 8]) - def test_staging_stages_dense_and_swa_tables(self, request_count): - pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - prepare_eviction_workspace, - stage_eviction_cohort, - ) - - device = torch.device("cuda", torch.cuda.current_device()) - max_requests = 8 - page_count = 3 - # The bucket capacity must be aligned to the score kernel's 64-token - # compute tile (the staging constructor compiles the kernel). - seq_len = 64 - page_table_token_capacity = 90 - tokens_per_block = 32 - head_dim = 64 - num_freqs = head_dim // 2 - num_q_heads = 8 - layer_elements = max_requests * page_count * 2 * 1 * tokens_per_block * head_dim - shared = torch.randn(2 * layer_elements, device=device).to(torch.bfloat16) - pool_shape = (max_requests * page_count, 2, 1, tokens_per_block, head_dim) - pools = [ - shared[:layer_elements].view(pool_shape), - shared[layer_elements:].view(pool_shape), - torch.randn(pool_shape, device=device).to(torch.bfloat16), - torch.randn(pool_shape, device=device).to(torch.bfloat16), - ] - dense_groups = [[0, 1], [2]] - representatives = [0, 2, 3] - q_real = torch.randn(4, num_q_heads, 2 * num_freqs, dtype=torch.float64, device=device)[ - ..., ::2 - ] - q_imag = torch.randn(4, num_q_heads, 2 * num_freqs, dtype=torch.float64, device=device)[ - ..., ::2 - ] - mlr = torch.randn(4, num_q_heads, 2 * num_freqs, dtype=torch.float64, device=device)[ - ..., ::2 - ] - freq = (torch.rand(2 * num_freqs, dtype=torch.float64, device=device) + 0.5)[::2] - omega = (torch.rand(2 * num_freqs, dtype=torch.float64, device=device) * 0.05)[::2] - offsets = torch.tensor([1.0, 0.0, 2.0, 0.0], dtype=torch.float64, device=device)[::2] - assert not q_real.is_contiguous() - assert not freq.is_contiguous() - assert not omega.is_contiguous() - assert not offsets.is_contiguous() - # The workspace constructor converts every non-contiguous fp64 - # calibration input to contiguous fp32 (the runner's flat views would - # otherwise fail) and compiles the score kernel here. - staging = prepare_eviction_workspace( - eviction_mode="per_head", - layer_pools=pools, - dense_groups=dense_groups, - dense_layers=[0, 1, 2], - page_representatives=representatives, - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_freqs=num_freqs, - keep_count=4, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr, - freq_scale_sq=freq, - offsets=offsets, - omega=omega, - page_table_token_capacity=page_table_token_capacity, - layer_group_representative={0: 0, 1: 0, 2: 2}, - layer_pool_keys=[("pool", 0), ("pool", 0), ("pool", 2), ("pool", 3)], - ) - assert staging.bucket_seq_len == seq_len - assert staging.page_table_token_capacity == page_table_token_capacity - # page_count 3 rounds up to the 4-block copy granule. - assert staging.copy_block_count == (page_count + 3) // 4 * 4 - assert staging.phase["offsets"].dtype == torch.float32 - assert staging.phase["offsets"].is_contiguous() - assert staging.phase["omega"].dtype == torch.float32 - assert staging.phase["omega"].is_contiguous() - tables = { - 10: [ - [3 * request, 3 * request + 1, 3 * request + 2] for request in range(request_count) - ], - 12: [ - [3 * request + 2, 3 * request + 1, 3 * request] for request in range(request_count) - ], - 13: [ - [23 - 3 * request, 22 - 3 * request, 21 - 3 * request] - for request in range(request_count) - ], - } - - request_ids = list(range(request_count)) - round_starts = [131_071 + request for request in request_ids] - # Per-request pinned prompt lengths: one cohort may mix them. - token_starts = list(request_ids) - host_table = torch.zeros( - 3, - max_requests, - 2, - staging.copy_block_count, - dtype=torch.int32, - device="cpu", - pin_memory=True, - ) - for slot, global_layer in enumerate((10, 12, 13)): - host_table[slot, :request_count, 0, :page_count].copy_( - torch.tensor(tables[global_layer], dtype=torch.int32) - ) - - def gather_k_block_offsets(source, destination, requested_ids, num_blocks): - for destination_row, request_id in enumerate(requested_ids): - destination[:, destination_row, 0, :num_blocks].copy_( - source[:, request_id, 0, :num_blocks] - ) - - gather = mock.Mock(side_effect=gather_k_block_offsets) - manager = _make_staging_manager( - host_table, gather, torch.cuda.Stream(device=device), num_slots=3 - ) - - # Round starts past the int32 metadata range fail loudly before any - # GPU work is enqueued. - with pytest.raises((RuntimeError, OverflowError, ValueError)): - stage_eviction_cohort( - staging, - manager, - request_ids, - [2**31] * request_count, - token_starts, - [seq_len] * request_count, - ) - assert gather.call_count == 0 - with mock.patch.object( - torch, - "index_select", - side_effect=AssertionError("page-table staging used torch.index_select"), - ): - stage_eviction_cohort( - staging, - manager, - request_ids, - round_starts, - token_starts, - [seq_len] * request_count, - ) - torch.cuda.current_stream(device).synchronize() - assert staging.round_starts_device.untyped_storage().data_ptr() == ( - staging.valid_seq_lens_device.untyped_storage().data_ptr() - ) - assert torch.equal( - staging.round_starts_device[:request_count], - torch.tensor(round_starts, dtype=torch.int32, device=device), - ) - assert torch.equal( - staging.valid_seq_lens_device[:request_count], - torch.full((request_count,), seq_len, dtype=torch.int32, device=device), - ) - assert torch.equal( - staging.token_starts_device[:request_count], - torch.tensor(token_starts, dtype=torch.int32, device=device), - ) - for slot, global_layer in enumerate((10, 12, 13)): - expected = torch.tensor(tables[global_layer], dtype=torch.int32, device=device) * 2 - assert torch.equal( - staging.block_offsets_device[slot, :request_count, 0, :page_count], - expected, - ) - - @requires_sm100 - @pytest.mark.parametrize("request_count", [1, 7, 8]) + @pytest.mark.parametrize("request_count", [1, 8]) def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count): """ONE launch over layers in DISTINCT storages with DISTINCT block tables. 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 index 32ba340ab148..fd3a6460c00c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -7,7 +7,6 @@ import pytest import torch from conftest import build_compaction as _build_compaction -from conftest import compaction_family as _compaction_family from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_ramp_pools as _make_ramp_pools from conftest import set_protected_tails as _set_protected_tails @@ -176,8 +175,10 @@ def _per_head_keep_oracle( return torch.stack(rows) -@pytest.mark.parametrize("eviction_mode", ["per_head", "per_layer_perhead"]) -@pytest.mark.parametrize("normalize_scores", [False, True]) +@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 ): @@ -223,7 +224,7 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( assert torch.equal(second, expected) -@pytest.mark.parametrize("keep_count,width", [(4, 64), (4096, 4224), (8192, 9216)]) +@pytest.mark.parametrize("keep_count,width", [(4, 64), (8192, 9216)]) def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): # Heavily tied integer scores with ragged valid widths and per-request # prompt rebase: the strongest oracle over the direct union top-k path @@ -939,48 +940,3 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): swa_after.index_select(2, swa_destination), swa_before.index_select(2, swa_source), ) - - -def test_cache_families_read_the_staged_move_offsets_rows(): - """Every cache family must consume the caller-staged offsets row. - - A family that silently falls back to its construction-time offsets - compacts request slots that are not in the staged cohort; on - sliding-window models the padded slots then produce negative source - ordinals and an illegal memory access. Binding the staged row by - reference is part of the builder contract. - """ - device = torch.device("cuda", torch.cuda.current_device()) - dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) - swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) - # bf16 pools in the compact op's supported geometry; this test only - # constructs the compaction, so the contents stay zero. - pools = [ - torch.zeros(6, 2, 1, 32, 64, dtype=torch.bfloat16, device=device), - torch.zeros(6, 2, 1, 32, 64, dtype=torch.bfloat16, device=device), - ] - keep = torch.tensor([[2, 4, 5, 7], [2, 3, 5, 6]], dtype=torch.int32, device=device) - staged_rows = torch.zeros(2, 3, dtype=torch.int32, device=device) - dense_offsets_row = staged_rows[0] - swa_offsets_row = staged_rows[1] - compaction = _build_compaction( - layer_pools=pools, - dense_layers=[0], - swa_layers=[1], - layer_group_representative={0: 0}, - layer_pool_keys=[("dense", 0), ("swa", 0)], - kept_token_ordinals=keep, - valid_sequence_lengths=torch.tensor([8, 7], dtype=torch.int32, device=device), - kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), - page_table_slots={0: 0, 1: 1}, - prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), - swa_window=2, - protected_tail_capacity=2, - dense_move_offsets=dense_offsets_row, - swa_move_offsets=swa_offsets_row, - ) - dense_family = _compaction_family(compaction, "dense") - assert dense_family["offsets"].data_ptr() == dense_offsets_row.data_ptr() - swa_family = _compaction_family(compaction, "swa") - assert swa_family is not None - assert swa_family["offsets"].data_ptr() == swa_offsets_row.data_ptr() From bb506ce3884f3ad16f0d09cff48984459654c612 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:00:50 -0700 Subject: [PATCH 088/178] [None][fix] Reject generation requests that skipped on_request_init A request reaching the generation hook without its init hook means the framework lifecycle was violated; silently initializing it here hid the breakage and made the calibration re-check downstream permanently unreachable. The late-init fallback becomes a hard RuntimeError and the dead not-calibrated return leg is deleted (on_request_init always ran). The _rope_tables except-narrowing originally paired with this knife is parked: its pre-edit runtime receipt exposed a transformers-5.5.4 rope-schema regression (rope_scaling -> rope_parameters) that silently routes both production models through the analytic fallback with a wrong theta; the fix changes scoring numerics on the official- calibration path and needs its own ruling and receipt. Signed-off-by: tianruih --- .../triattention/triattention.py | 6 ++++-- .../test_triattention_pipeline.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 233979cca1d3..b5f0c10207e6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1142,9 +1142,11 @@ def _periodic_evict( "the final update hook" ) if request_id not in self._request_states: - self.on_request_init(request) + raise RuntimeError( + f"request {request_id} reached generation without on_request_init" + ) resolved_requests.append((request, request_id, kv_cache)) - if not resolved_requests or not self._calibrated: + if not resolved_requests: return protected_tails: Dict[int, int] = {} due_requests = [] diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 41822ca07c7c..852ad592f445 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -283,6 +283,18 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): periodic_evict.assert_called_once_with(batch) + def test_unregistered_generation_request_is_rejected(self): + # A generation request whose on_request_init never ran is a framework + # ordering bug: it fails loudly instead of late-initializing here. + manager = _make_triattention() + manager._calibrated = True + cache = SimpleNamespace(capacity=8, history_length=0, is_active=True) + manager.kv_cache_manager.kv_cache_map = {7: cache} + request = _make_request(7, py_prompt_len=2) + + with pytest.raises(RuntimeError, match="without on_request_init"): + manager._periodic_evict(SimpleNamespace(generation_requests=[request])) + @staticmethod def _make_due_decode_request(seq_len): request = _make_request( From aa026ffac0f15083ca938d8a99ed5e329dbd6ef1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:00:50 -0700 Subject: [PATCH 089/178] [None][refactor] Build the eviction cohort where its values are born _periodic_evict already resolves every per-request quantity (confirmed length, protected tail, due-ness, keep target); it now assembles the prepared per-request records right there, inside the same triattention.metadata NVTX range, and hands the list to _evict_requests. The downstream re-derivation loop, the due-ness re-filter, the protected_tail_lengths parameter, and the dead missing-confirmed-length raise are gone; the tail-capacity accounting raise runs once at the resolution site and the evicted<=0 identity-compaction raise remains the contract at the bookkeeping loop. Direct _evict_requests callers now speak trust-the-pipeline: the identity-compaction incident anchor re-anchors at _periodic_evict level (due-filter proves no publication) plus a direct contract-raise probe. Signed-off-by: tianruih --- .../triattention/triattention.py | 164 ++++++++---------- .../test_triattention_pipeline.py | 74 +++++++- 2 files changed, 140 insertions(+), 98 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b5f0c10207e6..bb7db73aac84 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1149,58 +1149,82 @@ def _periodic_evict( if not resolved_requests: return protected_tails: Dict[int, int] = {} - due_requests = [] - - # Resolve every active target cache before changing cadence state. The - # captured cache objects also avoid repeating the V2 map lookup here. - for request, request_id, kv_cache in resolved_requests: - raw_capacity = int(kv_cache.capacity) - # One-engine speculative decoding keeps a fixed reserve E. Under - # overlap, B(n) is allocated/enqueued before finalizing B(n-1), so - # its exact scheduler growth Q is also opaque. Both spans are - # contiguous after the stable target prefix and move byte-for-byte. - protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( - scheduled_batch, request_id - ) - seq_len = raw_capacity - protected_tail - if seq_len < 0 or protected_tail < 0: - raise RuntimeError( - f"Request {request_id} has an inconsistent protected target tail: " - f"confirmed={seq_len}, capacity={raw_capacity}, " - f"protected_tail={protected_tail}" - ) - if seq_len < kv_cache.history_length: - raise RuntimeError( - f"Request {request_id} KV length {seq_len} is below finalized " - f"history {kv_cache.history_length}" + prepared: List[Dict[str, object]] = [] + protected_tail_capacity = self._configured_protected_tail_capacity() + + # Resolve every active target cache before changing cadence state (the + # captured cache objects also avoid repeating the V2 map lookup here), + # building the due cohort's per-request eviction metadata in the same + # pass -- ``_evict_requests`` trusts it as-is. + with nvtx_range("triattention.metadata", color="cyan"): + for request, request_id, kv_cache in resolved_requests: + raw_capacity = int(kv_cache.capacity) + # One-engine speculative decoding keeps a fixed reserve E. + # Under overlap, B(n) is allocated/enqueued before finalizing + # B(n-1), so its exact scheduler growth Q is also opaque. Both + # spans are contiguous after the stable target prefix and move + # byte-for-byte. + protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( + scheduled_batch, request_id ) - request_state = self._request_states[request_id] - request_state["confirmed_kv_length"] = seq_len - previous_step = request_state["generation_steps"] - confirmed_delta = 1 + int(request.py_num_accepted_draft_tokens) - step = previous_step + confirmed_delta - request_state["generation_steps"] = step - if previous_step // self.beta >= step // self.beta: - continue - if seq_len <= self._minimum_evictable_length(request, seq_len): - continue - if self.draft_kv_cache_manager is not None: - draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) - if draft_kv_cache is None or not draft_kv_cache.is_active: + seq_len = raw_capacity - protected_tail + if seq_len < 0 or protected_tail < 0: + raise RuntimeError( + f"Request {request_id} has an inconsistent protected target tail: " + f"confirmed={seq_len}, capacity={raw_capacity}, " + f"protected_tail={protected_tail}" + ) + if protected_tail > protected_tail_capacity: + raise RuntimeError( + f"Request {request_id} protected tail {protected_tail} exceeds " + f"configured capacity {protected_tail_capacity}" + ) + if seq_len < kv_cache.history_length: raise RuntimeError( - "TriAttention cannot co-compress a missing or " - f"suspended draft KV cache; request {request_id} must " - "be resumed before the final update hook" + f"Request {request_id} KV length {seq_len} is below finalized " + f"history {kv_cache.history_length}" ) - protected_tails[request_id] = protected_tail - due_requests.append((request, request_id)) + request_state = self._request_states[request_id] + request_state["confirmed_kv_length"] = seq_len + previous_step = request_state["generation_steps"] + confirmed_delta = 1 + int(request.py_num_accepted_draft_tokens) + step = previous_step + confirmed_delta + request_state["generation_steps"] = step + if previous_step // self.beta >= step // self.beta: + continue + expected_keep_count = self._minimum_evictable_length(request, seq_len) + if seq_len <= expected_keep_count: + continue + if self.draft_kv_cache_manager is not None: + draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) + if draft_kv_cache is None or not draft_kv_cache.is_active: + raise RuntimeError( + "TriAttention cannot co-compress a missing or " + f"suspended draft KV cache; request {request_id} must " + "be resumed before the final update hook" + ) + protected_tails[request_id] = protected_tail + prepared.append( + { + "request": request, + "request_id": request_id, + "seq_len": int(seq_len), + # Restore the uncompressed confirmed logical position + # from the physical prefix and cumulative eviction + # count. + "round_start": int(seq_len + request_state["evicted_tokens"]), + "prompt_len": min(int(request.py_prompt_len), int(seq_len)), + "expected_keep_count": expected_keep_count, + "protected_tail": protected_tail, + } + ) # Compact all affected dense and kernel-masked SWA layers, then release # the unreachable tail directly through V2's public resize primitive. # Prompt lengths and tails are per-request metadata, so the whole due # cohort runs as one batched round (the workspace holds max_batch_size # requests, which bounds any generation batch). - if not due_requests: + if not prepared: return num_layers = self._num_layers_from_manager() # Ungated NVTX with the due count in the message, so any nsys capture @@ -1208,14 +1232,10 @@ def _periodic_evict( # outside CUDA-graph capture, so the dynamic message is safe; the cost # is one host-side f-string per eviction round. with nvtx_range( - f"triattention.evict_request_group reqs={len(due_requests)}", + f"triattention.evict_request_group reqs={len(prepared)}", color="purple", ): - capacity_targets = self._evict_requests( - due_requests, - num_layers, - protected_tail_lengths=protected_tails, - ) + capacity_targets = self._evict_requests(prepared, num_layers) self._resize_compacted_requests(capacity_targets, protected_tails) def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: @@ -1820,55 +1840,19 @@ def _dense_layer_pool_groups( def _evict_requests( self, - evict_reqs, + prepared: List[Dict[str, object]], num_layers: int, - protected_tail_lengths: Optional[Dict[int, int]] = None, ) -> List[Tuple[int, int]]: - """Score and compact requests, returning ``(request_id, capacity)`` targets. + """Score and compact a prepared cohort, returning ``(request_id, capacity)`` targets. - Only full-attention layers participate in scoring. For kernel-masked SWA + ``prepared`` carries the per-request eviction metadata resolved by + ``_periodic_evict`` (every entry is due and evictable). Only + full-attention layers participate in scoring. For kernel-masked SWA layers, the latest model window is rebased to the tail of the common compacted prefix before the request-wide capacity is reduced. """ - if protected_tail_lengths is None: - protected_tail_lengths = {} - protected_tail_capacity = self._configured_protected_tail_capacity() with nvtx_range_debug("triattention.resolve_layout", color="blue"): layout = self._runtime_kv_layout(num_layers) - - # Resolve request length and page metadata before mutating any layer. - prepared: List[Dict[str, object]] = [] - with nvtx_range("triattention.metadata", color="cyan"): - for request, rid in evict_reqs: - request_state = self._request_states.get(rid) - seq_len = None if request_state is None else request_state["confirmed_kv_length"] - if seq_len is None: - raise RuntimeError(f"Missing confirmed KV length for request {rid}") - # Restore the uncompressed confirmed logical position from the - # physical prefix and cumulative eviction count. - round_start = seq_len + request_state["evicted_tokens"] - minimum_evictable_length = self._minimum_evictable_length(request, seq_len) - if seq_len <= minimum_evictable_length: - continue - protected_tail = int(protected_tail_lengths.get(rid, 0)) - if protected_tail < 0 or protected_tail > protected_tail_capacity: - raise RuntimeError( - f"Request {rid} protected tail {protected_tail} exceeds " - f"configured capacity {protected_tail_capacity}" - ) - prepared.append( - { - "request": request, - "request_id": rid, - "seq_len": int(seq_len), - "round_start": int(round_start), - "prompt_len": min(int(request.py_prompt_len), int(seq_len)), - "expected_keep_count": minimum_evictable_length, - "protected_tail": protected_tail, - } - ) - if not prepared: - return [] with nvtx_range_debug("triattention.staging_lookup", color="blue"): ws = self._workspace_for(layout, prepared) if layout["swa_layers"] and layout["swa_window"]: diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 852ad592f445..15e97b331c6d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -250,12 +250,39 @@ def test_identity_compaction_is_rejected_instead_of_published(self): manager = _make_triattention(top_B=4) manager.kv_cache_manager._stream = mock.Mock() request = _make_request(7, py_prompt_len=2) - # Selection keeps every token: seq_len == prompt + budget + 1 evicts - # one token; seq_len == prompt + budget must never publish. - _set_request_state(manager, 7, confirmed_kv_length=6) + # Selection keeps every token at seq_len == prompt + budget: the due + # filter must drop the request before any eviction work or publication. + 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} + manager._calibrated = True + state = _set_request_state(manager, 7, generation_steps=127) + + with _mocked_eviction_internals(manager) as internals: + manager._periodic_evict(SimpleNamespace(generation_requests=[request])) + internals.run_round.assert_not_called() + assert request.py_num_compressed_tokens == 0 + assert state["evicted_tokens"] == 0 + cache.resize.assert_not_called() + + # A prepared entry that violates the keep contract fails loudly in the + # bookkeeping loop instead of publishing an identity compaction. with _mocked_eviction_internals(manager): - assert manager._evict_requests([(request, 7)], 2) == [] + with pytest.raises(RuntimeError, match="identity compaction"): + manager._evict_requests( + [ + _prepared_eviction( + request, + request_id=7, + seq_len=6, + expected_keep_count=6, + prompt_len=2, + ) + ], + 2, + ) assert request.py_num_compressed_tokens == 0 @@ -317,6 +344,7 @@ def _make_due_decode_request(seq_len): pp_layers=[0, 1], _stream=mock.Mock(), num_extra_kv_tokens=0, + _kv_reserve_draft_tokens=0, ) mgr._L = 2 mgr._request_states = {} @@ -374,6 +402,9 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): cache = mgr.kv_cache_manager.kv_cache_map[7] cache.capacity = confirmed + tail mgr.kv_cache_manager.num_extra_kv_tokens = reserve + # The configured tail capacity (reserve + draft reserve + 1) must + # cover this round's actual tail, as in production. + mgr.kv_cache_manager._kv_reserve_draft_tokens = current_growth mgr._prepared_generation_batch = ( SimpleNamespace(generation_requests=[request]), {7: current_growth}, @@ -391,10 +422,21 @@ def compact(*_args, **_kwargs): with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: mgr._periodic_evict(batch) + # The resolved per-request metadata is threaded in as-is: the tail is + # excluded from seq_len and the keep target is prompt + budget. evict.assert_called_once_with( - [(request, 7)], + [ + { + "request": request, + "request_id": 7, + "seq_len": confirmed, + "round_start": confirmed, + "prompt_len": 1024, + "expected_keep_count": retained, + "protected_tail": tail, + } + ], 2, - protected_tail_lengths={7: tail}, ) assert mgr._request_states[7]["confirmed_kv_length"] == retained cache.resize.assert_called_once_with(retained + tail, None) @@ -675,9 +717,25 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): with _mocked_eviction_internals(manager) as internals: manager._evict_requests( - [(first, 7), (second, 8)], + [ + _prepared_eviction( + first, + request_id=7, + seq_len=8, + prompt_len=3, + expected_keep_count=7, + protected_tail=2, + ), + _prepared_eviction( + second, + request_id=8, + seq_len=10, + prompt_len=5, + expected_keep_count=9, + protected_tail=3, + ), + ], 2, - protected_tail_lengths={7: 2, 8: 3}, ) # One batched staging call carries the whole cohort: request ids, From 0bee387a8d69a0c49edcf12eb2988cc73b8829e6 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:01:27 -0700 Subject: [PATCH 090/178] [None][refactor] Lock the fused dense pack on and drop the compiled-state predicates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit build_cache_compactions loses fuse_dense_pack_into_selection: the fused settle-and-pack is production's only shape, so the dense family is always built pack-free and the pack description is always returned for the selection launch. run_cache_compactions' family-pack branch stays: the draft family packs through it every co-compaction round. The test scaffolding (conftest build_compaction) re-attaches the dense pack so standalone compaction tests keep packing their move buffers before the C++ moves — without it they would compact from uninitialized indices; tests that already fire launch_move_pack explicitly stay correct (the pack launch is idempotent). The CuTe runner's supports()/supports_union_fusion() predicates had no product caller left (the round checks died with the trusted-pipeline knife); tests assert the compiled-state dicts directly, preserving each assert's meaning including the over-capacity negative. Signed-off-by: tianruih --- .../triattention/compaction.py | 20 +++++++++---------- .../triattention/triattention.py | 1 - .../triattention_cute_score_fused.py | 11 ---------- .../_torch/kv_cache_compression/conftest.py | 7 ++++++- .../test_triattention_cute_score.py | 4 ++-- .../test_triattention_cute_union_fusion.py | 7 +++++-- .../test_triattention_selection_compaction.py | 18 ++++++++++++++--- 7 files changed, 37 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py index 321326fbd4b8..52d233cf16ef 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py @@ -22,10 +22,10 @@ surviving KV in place with batched C++ compact launches. Everything is plain tensors and dicts: ``build_cache_compactions`` allocates the launch data once per geometry (called by ``triattention.prepare_eviction_workspace``) and -``run_cache_compactions`` fires the kernels directly each round. A driver -that finalizes the keep set in its own GPU launch takes the target's -dense/SWA packing over into that launch (``fuse_dense_pack_into_selection``); -the draft always packs here. +``run_cache_compactions`` fires the kernels directly each round. The +target's dense/SWA packing always rides the driver's fused settle launch +(the returned ``dense_pack`` dict describes it); the draft always packs +here with its own launch. """ from collections import OrderedDict @@ -259,7 +259,6 @@ def build_cache_compactions( swa_window: Optional[int], layer_pool_keys: List[object], protected_tail_capacity: int = 0, - fuse_dense_pack_into_selection: bool = False, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_layers: Optional[List[int]] = None, draft_layer_group_representative: Optional[Dict[int, int]] = None, @@ -287,10 +286,9 @@ def build_cache_compactions( ``[slot, request, K/V, block]`` (offset = ``2*page + plane``); ``protected_tail_capacity`` is the widest per-request tail this geometry must support -- actual per-round lengths arrive through the staged - move-offset rows. With ``fuse_dense_pack_into_selection`` the target's - dense/SWA packing is left to the caller's fused settle launch (the - returned ``dense_pack`` dict describes it) and only the C++ moves run - here; the draft always keeps its own pack launch. + move-offset rows. The target's dense/SWA packing is left to the caller's + fused settle launch (the returned ``dense_pack`` dict describes it) and + only the C++ moves run here; the draft always keeps its own pack launch. Returns ``{"families": [...], "dense_pack": ..., "num_kv_heads": ..., "swa_window": ..., "swa_destination_bases": ...}`` where each family is @@ -365,10 +363,10 @@ def build_cache_compactions( families = [ dict( name="dense", - # A fused selection-side settle launch packs the dense/SWA move + # The fused selection-side settle launch packs the dense/SWA move # sources when it finalizes the kept ordinals; only the C++ moves # stay here. Each round then packs exactly once. - pack=None if fuse_dense_pack_into_selection else dense_pack, + pack=None, groups=_compact_groups(dense_entries, layer_pool_keys, device, dense_slots), source=dense_move_indices, offsets=dense_move_offsets, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index bb7db73aac84..c4c97f0658cf 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -514,7 +514,6 @@ def prepare_eviction_workspace( # ONE launch settles the kept ordinals and packs the dense/SWA # move sources; ``run_cache_compactions`` then only runs the C++ # moves (plus the draft's own pack). - fuse_dense_pack_into_selection=True, **draft_kwargs, ) pack = ws.compaction["dense_pack"] 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 index 7e9a76aaafba..831ae677a4b3 100644 --- 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 @@ -1788,10 +1788,6 @@ def __init__( compiled_configs[config_key] = compiled_selection self._compiled_normalize_union[request_count] = compiled_selection - def supports(self, request_count: int) -> bool: - """Return whether the dynamic specialization covers this request count.""" - return request_count in self._compiled - def launch( self, request_count: int, @@ -1809,13 +1805,6 @@ def launch( stream, ) - def supports_union_fusion(self, request_count: int) -> bool: - """Return whether the score/stats/union pipeline was precompiled.""" - return ( - request_count in self._compiled_stats - and request_count in self._compiled_normalize_union - ) - def launch_union_fusion( self, request_count: int, diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 74ce740feaa8..1c5157da42f2 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -153,7 +153,12 @@ def build_compaction(**overrides): swa_window=None, ) args.update(overrides) - return build_cache_compactions(**args) + compaction = build_cache_compactions(**args) + # Production fuses the dense pack into the settle launch; standalone + # compaction tests re-attach it so run_cache_compactions packs the dense + # move buffers before the C++ moves (firing the pack twice is idempotent). + compaction_family(compaction, "dense")["pack"] = compaction["dense_pack"] + return compaction def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None): 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 index 6dd917f93bdc..895c66a6275a 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -94,7 +94,7 @@ def _launch_split_scores( valid_seq_lens, 0, ws.seg_req[:num_segments], out=ws.seg_seq_len[:num_segments] ) ws.cute_token_starts[:request_count].copy_(token_starts[:request_count]) - assert ws.runner.supports(request_count) + assert request_count in ws.runner._compiled ws.runner.launch(request_count, mean_cos, mean_sin) group_size = ws.num_q_heads // ws.num_kv_heads source = ( @@ -377,7 +377,7 @@ def test_cute_kernel_matches_torch_oracle(case): # The compiled runner serves every request count up to the workspace # capacity and nothing beyond it; cover one, an intermediate count, and # the capacity. - assert not ws.runner.supports(max_requests + 1) + assert max_requests + 1 not in ws.runner._compiled for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) scores = _launch_split_scores( 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 index eaddfd382bf5..addc96647302 100644 --- 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 @@ -91,7 +91,7 @@ def _launch_split_scores( ): """The production score-only leg plus the decode-window gather.""" _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts) - assert ws.runner.supports(request_count) + assert request_count in ws.runner._compiled ws.runner.launch(request_count, mean_cos, mean_sin) num_segments = request_count * ws.num_layers group_size = ws.num_q_heads // ws.num_kv_heads @@ -124,7 +124,10 @@ def _launch_union_fusion( ): """The fused score+stats+normalized-union pipeline (THE union path).""" _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts) - assert ws.runner.supports_union_fusion(request_count) + assert ( + request_count in ws.runner._compiled_stats + and request_count in ws.runner._compiled_normalize_union + ) ws.runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_out[:request_count]) 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 index fd3a6460c00c..71fcc702fbf4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -11,7 +11,10 @@ from conftest import make_ramp_pools as _make_ramp_pools from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import run_cache_compactions +from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( + launch_move_pack, + run_cache_compactions, +) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( prepare_per_head_scores, @@ -396,6 +399,9 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) + # Production packs these buffers inside the fused settle launch; with + # pre-settled ordinals the standalone pack call is its exact analog. + launch_move_pack(compaction["dense_pack"]) run_cache_compactions(compaction) torch.cuda.synchronize(device) @@ -484,6 +490,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): protected_tail_capacity=max(protected_tails), ) _set_protected_tails(cohort_compaction, protected_tails) + launch_move_pack(cohort_compaction["dense_pack"]) run_cache_compactions(cohort_compaction) expected_pools = [pool.clone() for pool in initial_pools] @@ -499,6 +506,7 @@ def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): protected_tail_capacity=protected_tails[request], ) _set_protected_tails(single_compaction, [protected_tails[request]]) + launch_move_pack(single_compaction["dense_pack"]) run_cache_compactions(single_compaction) torch.cuda.synchronize(device) @@ -593,8 +601,10 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): ws.valid_seq_lens_device.fill_(seq_len) ws.token_starts_device.fill_(0) - # This compaction packs its own move indices; the workspace's fused pack - # stays masked off (its staged move-offsets row is all zeros). + # This compaction packs its own move indices (the ordinals settle + # mid-round, so the pack rides the family slot exactly like the draft's + # own pack does in production); the workspace's fused pack stays masked + # off (its staged move-offsets row is all zeros). ws.compaction = _build_compaction( eviction_mode="per_layer_perhead", layer_pools=pools, @@ -611,6 +621,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): protected_tail_capacity=0, ) _set_protected_tails(ws.compaction, [0]) + ws.compaction["families"][0]["pack"] = ws.compaction["dense_pack"] run_eviction_round(ws, normalize_scores=False) assert torch.equal(ws.keep, expected_keep) torch.cuda.synchronize(device) @@ -901,6 +912,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) + launch_move_pack(compaction["dense_pack"]) run_cache_compactions(compaction) torch.cuda.synchronize(device) From 6c4381d1ff4efbe7f477aa04bc70127087df6220 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:01:27 -0700 Subject: [PATCH 091/178] [None][doc] Align the README with the shipped defaults and contracts top_B's documented default becomes 2048 (matching the code example and YAML in the same file); model_path's description gains its second real consumer (kernel-masked SWA layer classification from the model config); normalize_scores documents that union eviction requires True. The module notes on the calibration-path docstring that calibration resolves lazily on the first request. Signed-off-by: tianruih --- examples/triattention/README.md | 6 +++--- .../kv_cache_compression/triattention/triattention.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index 97deea79eb81..3dbdbd8ed384 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -114,13 +114,13 @@ trtllm-eval --model --config config.yaml longbench_v2 --max_outp `TriAttentionKvCacheCompressionConfig` controls the compression ratio and the eviction algorithm: -* **`top_B`** (int, default=1024): Tokens kept at each eviction (the upstream `budget`). Prompt tokens are always preserved on top of this. Smaller `top_B` → more compression. +* **`top_B`** (int, default=2048): Tokens kept at each eviction (the upstream `budget`). Prompt tokens are always preserved on top of this. Smaller `top_B` → 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). +* **`normalize_scores`** (bool, default=True): Z-normalize each head's scores over the decode region before selection (upstream default). `union` eviction requires `True` (the fused union pipeline always z-normalizes; construction rejects `False`). * **`pin_prefill`** (bool, default=True): Always preserve the prompt (prefill) tokens; only decode tokens compete for the budget (upstream behaviour). * **`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 only to derive the model's RoPE tables when converting the official calibration file. +* **`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/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index c4c97f0658cf..d1c1ab5e0cba 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1921,7 +1921,8 @@ def _resolve_calibration(self) -> Dict[str, torch.Tensor]: ``calibration_path``; we only run inference. Both the official R-KV layout (``{metadata, stats{"layerLL_headHH": {q_mean_real, q_mean_imag, q_abs_mean}}}``) and our already-converted flat layout are accepted -- the - official one is converted here.""" + official one is converted here. Calibration resolves lazily on the + first request (``on_request_init``), not at manager construction.""" if self.calibration_path is None: raise ValueError( "TriAttention requires `calibration_path`: a calibration .pt from " From 563d1cd884e19108bc3d3aca3e926eb4c2ae27e9 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:01:45 -0700 Subject: [PATCH 092/178] [None][test] Retire mock legs and duplicate grids ruled out with the seams D1: the runner-construction-failure mock leg goes; the frequency-count guard survives as a plain test. D2: the standalone torch-mean oracle test folds into the launch-path matrix, whose seq_lens formula now includes deep-ragged tails and a capacity-58 request so the fully-invalid-page-fragment clamp stays covered through the shared oracle. D3: the cohort-vs-singles equivalence test retires; heterogeneous prompts fold into the byte-exact eager-compaction oracle (per-request prompt offsets, rebased keep sets, unweakened equalities). D4: the workspace-build mock plumbing collapses to one case, the normalize_scores rejection stands alone, the duplicated reuse block is gone, and the unique kwargs asserts live on in the draft file's superset test. Signed-off-by: tianruih --- .../test_triattention_cute_score.py | 118 +----------------- .../test_triattention_cute_union_fusion.py | 75 +++-------- .../test_triattention_draft_cocompaction.py | 5 + .../test_triattention_pipeline.py | 31 ++--- .../test_triattention_selection_compaction.py | 84 +------------ 5 files changed, 46 insertions(+), 267 deletions(-) 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 index 895c66a6275a..9e3358d8baf9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -121,117 +121,6 @@ def _launch_split_scores( return output -@requires_sm100 -@pytest.mark.parametrize( - "tokens_per_block,page_permutation,valid_lens,num_freqs,num_q_heads", - [ - # The originally validated geometry: 128-token pages, identity table. - (128, [0, 1], None, 32, 8), - # GPT-OSS geometry: 32-token pages; a 64-token compute tile spans two - # pages, so a shuffled physical-page table catches fragment mix-ups. - # Ragged tails land mid-tile: the second page fragment of the last - # tile is clamped, and scores past the valid length are unspecified. - (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198], 32, 8), - # Qwen3 geometry: 128-element K rows (64 frequencies) and GQA group - # 4, which rides the MMA tile N=8 with zeroed padding columns. - (32, [3, 1, 4, 7, 5, 0, 2, 6], [250, 198], 64, 4), - ], -) -def test_cute_score_matches_torch_mean_oracle( - tokens_per_block: int, - page_permutation: list, - valid_lens: "list | None", - num_freqs: int, - num_q_heads: int, -) -> None: - pytest.importorskip("cutlass") - - torch.manual_seed(20260720) - device = torch.device("cuda") - seq_len = 256 - num_pages = seq_len // 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) - round_starts = torch.tensor([float(seq_len), float(seq_len + 1)], device=device) - phase = (round_starts[:, 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() - - ws = _make_score_workspace( - 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, - ) - # Native block-offset staging layout ([pool_slot, request, K/V plane, - # block] int32): K-plane entries encode physical_page * kv_factor with - # kv_factor == 2. Both requests read the same (permuted) page sequence. - k_plane = [2 * page for page in page_permutation] - v_plane = [2 * page + 1 for page in page_permutation] - _write_block_offsets( - ws, - torch.tensor([[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device), - ) - keys = ( - torch.cat([pool[page, 0, 0] for page in page_permutation], dim=0) - .reshape(seq_len, 2 * num_freqs) - .float() - ) - k_real = keys[:, :num_freqs] - k_imag = keys[:, num_freqs:] - magnitude = torch.sqrt(k_real.square() + k_imag.square()) - if valid_lens is None: - valid_lens = [seq_len, seq_len] - valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) - valid_widths = torch.tensor(valid_lens, dtype=torch.int32, device=device) - token_starts_device = torch.zeros(2, dtype=torch.int32, device=device) - for request_count in (1, 2): - actual = _launch_split_scores( - ws, - request_count, - valid_seq_lens, - valid_widths, - token_starts_device, - mean_cos, - mean_sin, - ) - assert actual.shape == (request_count, 1, num_q_heads, seq_len) - for request in range(request_count): - rotated_real = freq_scale_sq * (k_real * mean_cos[request] + k_imag * mean_sin[request]) - rotated_imag = freq_scale_sq * (k_imag * mean_cos[request] - k_real * mean_sin[request]) - expected = ( - q_real[0, :, None] * rotated_real[None] - + q_imag[0, :, None] * rotated_imag[None] - + mlr_coef[0, :, None] * freq_scale_sq[None, None] * magnitude[None] - ).sum(dim=-1) - valid = valid_lens[request] - torch.testing.assert_close( - actual[request, 0, :, :valid], - expected[:, :valid], - rtol=5.0e-3, - atol=5.0e-3, - ) - - torch.cuda.synchronize() - # The CuTe runner is the only score path, compiled eagerly at workspace - # construction; prove setup actually built it. - assert ws.runner is not None - - def _build_case( *, max_requests: int, @@ -288,8 +177,11 @@ def _build_case( _write_block_offsets(ws, _encode_block_offsets(page_ids)) round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) - # Ragged valid lengths whose tails land mid-page and mid-compute-tile. - seq_lens = [capacity - ((request * 3) % 5) for request in range(max_requests)] + # Ragged valid lengths: shallow tails land mid-page and mid-compute-tile; + # the 58-deep tail leaves a whole trailing 32-token page fragment past the + # valid length, so the fully-invalid-fragment clamp stays covered. + tail_cuts = (0, 58, 3, 33) + seq_lens = [capacity - tail_cuts[request % len(tail_cuts)] for request in range(max_requests)] valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) phase = (round_starts.float()[:, None, None] + offsets_t[None, :, None]) * omega[None, None, :] mean_cos = torch.cos(phase).mean(dim=1).contiguous() 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 index addc96647302..ec992beb2464 100644 --- 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 @@ -300,66 +300,27 @@ def test_union_fusion_matches_split_pipeline( @_SM100_ONLY -@pytest.mark.parametrize("guard", ["runner_construction_failure", "frequency_count"]) -def test_union_fusion_guards_raise(guard: str, monkeypatch: pytest.MonkeyPatch) -> None: - """One representative per fused-pipeline guard family raises loudly. - - ``runner_construction_failure``: a runner construction failure surfaces - as the no-fallback RuntimeError at workspace construction (the only - place the CuTe entries compile). ``frequency_count``: 16 frequencies - (head size 32) sit outside the fused kernel contract and are rejected at - kernel construction. - """ +def test_union_fusion_frequency_count_guard_raises() -> None: + """16 frequencies (head size 32) sit outside the fused kernel contract + and are rejected at kernel construction.""" cutlass = pytest.importorskip("cutlass") - if guard == "frequency_count": - 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, - seq_len=256, - num_q_heads=8, - num_kv_heads=1, - num_freqs=16, - tokens_per_block=128, - pool_shape=(2, 2, 1, 128, 32), - pool_strides=(8192, 4096, 4096, 32, 1), - pool_dtype=cutlass.BFloat16, - page_shards=3, - ) - return - - import tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused as fused_module # noqa: E501 - - torch.manual_seed(20260721) - device = torch.device("cuda") - seq_len, tokens_per_block, num_freqs, num_q_heads = 256, 128, 32, 8 - num_pages = seq_len // tokens_per_block - 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) - omega = torch.linspace(0.01, 0.03, num_freqs, device=device) + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused import ( # noqa: E501 + _TriAttentionScoreKernel, + ) - def _refuse_construction(**_kwargs): - raise ValueError("synthetic fused-runner construction failure") - - monkeypatch.setattr(fused_module, "TriAttentionCuteScoreRunner", _refuse_construction) - with pytest.raises(RuntimeError, match="no other score path exists"): - _make_union_workspace( - layer_pools=[pool], - max_requests=2, - seq_len=seq_len, - num_q_heads=num_q_heads, - q_real=q_real, - q_imag=torch.randn_like(q_real) * 0.125, - mlr_coef=torch.randn_like(q_real) * 0.125, - freq_scale_sq=torch.linspace(0.5, 1.5, num_freqs, device=device), - omega=omega, - offsets=torch.tensor([1.0, 2.0, 4.0], device=device), + with pytest.raises(ValueError, match="frequencies"): + _TriAttentionScoreKernel( + num_layers=1, + seq_len=256, + num_q_heads=8, + num_kv_heads=1, + num_freqs=16, + tokens_per_block=128, + pool_shape=(2, 2, 1, 128, 32), + pool_strides=(8192, 4096, 4096, 32, 1), + pool_dtype=cutlass.BFloat16, + page_shards=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 index 8518800a5dc7..2be6e28fbc9d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -388,6 +388,11 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): assert prepare.call_args.kwargs["seq_len"] == 1024 assert prepare.call_args.kwargs["page_table_token_capacity"] == 1024 + 1 assert prepare.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 + # Migrated from the pipeline workspace-kwargs test: the budget, the + # shared phase-table dict, and the pool keys thread through unchanged. + assert prepare.call_args.kwargs["keep_count"] == manager.top_B + assert prepare.call_args.kwargs["phase"] is manager._phase + assert prepare.call_args.kwargs["layer_pool_keys"] == list(layout["layer_pool_keys"]) # A second round with unchanged pools reuses the resident workspace # (and with it the compaction launch data it carries). diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 15e97b331c6d..c9c189b5baf1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -537,16 +537,19 @@ def test_prepare_snapshots_fixed_linear_generation_growth( class TestFixedScoreMetadata: - @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) - def test_workspace_build_receives_mode_and_capacity_kwargs(self, eviction_mode): + def test_union_rejects_unnormalized_scores(self): + # The fused pipeline (THE union path) always z-normalizes, so + # normalize_scores=False is rejected loudly at construction. + with pytest.raises(ValueError, match="normalize_scores=True"): + _make_triattention(top_B=4, eviction_mode="union", normalize_scores=False) + + def test_workspace_build_receives_mode_and_capacity_kwargs(self): + # Mode/phase/pool-key threading and cached reuse are covered by the + # workspace-rebuild superset test in + # test_triattention_draft_cocompaction.py. from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - if eviction_mode == "union": - # The fused pipeline (THE union path) always z-normalizes, so - # normalize_scores=False is rejected loudly at construction. - with pytest.raises(ValueError, match="normalize_scores=True"): - _make_triattention(top_B=4, eviction_mode="union", normalize_scores=False) - manager = _make_triattention(top_B=4, eviction_mode=eviction_mode) + manager = _make_triattention(top_B=4) # The workspace follows the executor limits: eight requests (max batch # size) by 260 decode tokens (top_B plus two eviction periods). layout, workspace = _make_workspace_stubs(manager) @@ -568,19 +571,9 @@ def test_workspace_build_receives_mode_and_capacity_kwargs(self, eviction_mode): assert resources is workspace kwargs = build_workspace.call_args.kwargs - assert kwargs["eviction_mode"] == eviction_mode - assert kwargs["keep_count"] == 4 + assert kwargs["eviction_mode"] == "union" assert kwargs["max_requests"] == 8 assert kwargs["decode_width"] == 260 - assert kwargs["phase"] is manager._phase - assert kwargs["layer_pool_keys"] == list(layout["layer_pool_keys"]) - # The cached workspace serves later rounds without rebuilding. - with mock.patch.object( - module, - "prepare_eviction_workspace", - side_effect=AssertionError("workspace was rebuilt"), - ): - assert manager._workspace_for(layout, prepared) is workspace def test_stage_rejects_int32_overflowing_round_starts(self): # Round starts past the int32 metadata range fail loudly (in the host 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 index 71fcc702fbf4..c96e68bf83b7 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -348,7 +348,9 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) request_count = 2 num_layers = 2 num_kv_heads = 2 - prompt_len = 2 + # Per-request pinned prompts: one cohort mixes prompt lengths, so the + # byte-exact oracle also proves per-request destination rebasing. + prompt_lens = [2, 5] decode_keep_count = 4 seq_len = 80 tokens_per_block = 32 @@ -381,7 +383,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) keep[request, row] = torch.tensor( sorted( { - 2 + ((request + row + offset * 2) % 8) * 8 + prompt_lens[request] + ((request + row + offset * 2) % 8) * 8 for offset in range(decode_keep_count) } ), @@ -395,7 +397,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) kept_token_ordinals=keep.to(torch.int32), valid_sequence_lengths=torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device), kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), - prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), + prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) @@ -407,6 +409,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): for request in range(request_count): + prompt_len = prompt_lens[request] pages = page_tables[request].to(torch.long) before = ( before_pool[pages] @@ -441,81 +444,6 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) ) -def test_union_mixed_prompt_lengths_cohort_matches_single_request_compactions(): - """One union cohort mixing prompt lengths compacts byte-identically to - running the same two requests as two single-request compactions.""" - device = torch.device("cuda", torch.cuda.current_device()) - request_count = 2 - num_layers = 2 - # bf16 pools in the compact op's supported geometry; three 32-token pages - # per request with a 80-token sequence keep the moves page-crossing. - seq_len = 80 - decode_keep_count = 3 - prompt_lens = [2, 5] - protected_tails = [2, 1] - decode_widths = [seq_len - prompt_len for prompt_len in prompt_lens] - width = max(decode_widths) - - # Oracle selection: decode-relative scores per request, rebased to - # absolute ordinals by each request's own prompt offset. - generator = torch.Generator().manual_seed(11) - scores = torch.randint( - -8, - 9, - (request_count, 1, width), - generator=generator, - dtype=torch.int32, - ).to(torch.float32) - keep = torch.stack( - [ - torch.sort( - _stable_topk(scores[request, 0], decode_width, decode_keep_count) + prompt_len - ).values - for request, (prompt_len, decode_width) in enumerate(zip(prompt_lens, decode_widths)) - ] - ) - - page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - initial_pools = _make_ramp_pools(num_layers, device=device) - cohort_pools = [pool.clone() for pool in initial_pools] - keep_cuda = keep.to(device) - valid_seq_lens = torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device) - cohort_compaction = _build_compaction( - layer_pools=cohort_pools, - kept_token_ordinals=keep_cuda, - valid_sequence_lengths=valid_seq_lens, - kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), - prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), - decode_keep_count=decode_keep_count, - protected_tail_capacity=max(protected_tails), - ) - _set_protected_tails(cohort_compaction, protected_tails) - launch_move_pack(cohort_compaction["dense_pack"]) - run_cache_compactions(cohort_compaction) - - expected_pools = [pool.clone() for pool in initial_pools] - for request in range(request_count): - single_compaction = _build_compaction( - layer_pools=expected_pools, - kept_token_ordinals=keep_cuda[request : request + 1], - valid_sequence_lengths=valid_seq_lens[request : request + 1], - kv_block_offsets=_encode_block_offsets(page_tables[request : request + 1].unsqueeze(0)), - request_count=1, - prompt_offsets=torch.tensor([prompt_lens[request]], dtype=torch.int32, device=device), - decode_keep_count=decode_keep_count, - protected_tail_capacity=protected_tails[request], - ) - _set_protected_tails(single_compaction, [protected_tails[request]]) - launch_move_pack(single_compaction["dense_pack"]) - run_cache_compactions(single_compaction) - torch.cuda.synchronize(device) - - # The two requests own disjoint pages, so whole-pool equality proves the - # cohort produced exactly the two single-request results. - for cohort_pool, expected_pool in zip(cohort_pools, expected_pools): - assert torch.equal(cohort_pool, expected_pool) - - @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.""" From 95d79f75a40aa16344de9cd90a59f23a166e3d0f Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:13:58 -0700 Subject: [PATCH 093/178] [None][refactor] Rename TriAttention top_B to budget (official name, snake_case) The config field, manager attribute, constructor kwarg, error strings, tests, and README all follow; the golden llm-args manifest entry renames and re-sorts within its group. The description drops the 'upstream budget' parenthetical: budget now IS the name. Signed-off-by: tianruih --- examples/triattention/README.md | 10 ++--- .../triattention/triattention.py | 40 +++++++++---------- tensorrt_llm/llmapi/llm_args.py | 8 ++-- .../usage/llm_args_golden_manifest.json | 14 +++---- .../_torch/kv_cache_compression/conftest.py | 2 +- .../test_triattention_draft_cocompaction.py | 10 ++--- .../test_triattention_pipeline.py | 36 ++++++++--------- 7 files changed, 60 insertions(+), 60 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index 3dbdbd8ed384..0389f38332ad 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -2,7 +2,7 @@ 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 `top_B` tokens, and physically compacts the cache — reducing KV-cache memory so more sequences fit on a GPU at once. +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). @@ -11,7 +11,7 @@ For technical details see the paper [TriAttention](https://arxiv.org/abs/2604.04 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 it and converts it to its runtime schema at the first request. -2. **Periodic eviction (Stage during generation).** Every `beta` confirmed generation tokens, once a sequence is over budget, TriAttention scores the whole cache, selects `top_B` tokens to keep (the prompt tokens are preserved on top of the budget), and physically compacts the KV cache down to the kept set. A speculative iteration may confirm multiple tokens; crossing multiple periods in one update is coalesced into one eviction. +2. **Periodic eviction (Stage during generation).** Every `beta` confirmed generation tokens, once a sequence is over budget, TriAttention scores the whole cache, selects `budget` tokens to keep (the prompt tokens are preserved on top of the budget), and physically compacts the KV cache down to the kept 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`. The scoring and compaction kernels are implemented in **Triton**. @@ -65,7 +65,7 @@ from tensorrt_llm.llmapi import (KvCacheConfig, # 1. Configure the eviction manager + point it at the calibration file. compression_config = TriAttentionKvCacheCompressionConfig( - top_B=2048, # tokens kept at each eviction (prompt is kept on top) + 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 @@ -96,7 +96,7 @@ Pass the configs via `--config config.yaml`. The field names match the Python co backend: pytorch kv_cache_compression_config: algorithm: triattention - top_B: 2048 + budget: 2048 beta: 64 eviction_mode: union calibration_path: /path/to/qwen3-8b-calibration.pt @@ -114,7 +114,7 @@ trtllm-eval --model --config config.yaml longbench_v2 --max_outp `TriAttentionKvCacheCompressionConfig` controls the compression ratio and the eviction algorithm: -* **`top_B`** (int, default=2048): Tokens kept at each eviction (the upstream `budget`). Prompt tokens are always preserved on top of this. Smaller `top_B` → more compression. +* **`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. diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d1c1ab5e0cba..1e7027fc67dc 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -827,7 +827,7 @@ class TriAttention(BaseKVCacheCompressionManager): def __init__( self, kv_cache_manager: KVCacheManagerV2, - top_B: int, + budget: int, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, beta: int = 128, model_path: Optional[str] = None, @@ -838,10 +838,10 @@ def __init__( count_prompt_tokens: bool = False, ): super().__init__(kv_cache_manager, draft_kv_cache_manager) - self.top_B = top_B + self.budget = budget self.beta = beta - if self.top_B <= 0 or self.beta <= 0: - raise ValueError("TriAttention top_B and beta must both be positive") + if self.budget <= 0 or self.beta <= 0: + raise ValueError("TriAttention budget and beta must both be positive") # Which token set each eviction round keeps. The user-facing meaning of # each mode is documented on TriAttentionKvCacheCompressionConfig # (llm_args); implementation notes live above the selection helpers. @@ -939,7 +939,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: # V2 mirrors the resolved speculative draft length (0 without spec). speculative_overshoot = int(manager.max_draft_len) first_eviction_decode_length = ( - self.top_B // self.beta + 1 + self.budget // self.beta + 1 ) * self.beta + speculative_overshoot decode_capacity = min(int(request.py_max_new_tokens), first_eviction_decode_length) confirmed_capacity = int(request.py_prompt_len) + decode_capacity @@ -954,7 +954,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: raise ValueError( "TriAttention target KV capacity is too small to reach its first " f"eviction: request requires {required_capacity} tokens " - f"(prompt={request.py_prompt_len}, budget={self.top_B}, " + f"(prompt={request.py_prompt_len}, budget={self.budget}, " f"beta={self.beta}, decode before eviction or completion=" f"{decode_capacity}, speculative overshoot=" f"{speculative_overshoot}, protected tail=" @@ -979,7 +979,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: raise ValueError( "TriAttention draft KV capacity is too small to reach the first " f"co-compression: request requires {draft_required_capacity} " - f"tokens (prompt={request.py_prompt_len}, budget={self.top_B}, " + f"tokens (prompt={request.py_prompt_len}, budget={self.budget}, " f"beta={self.beta}, decode before eviction or completion=" f"{decode_capacity}, draft protected tail={draft_protected_tail}), " f"but the draft V2 pool covers " @@ -1279,13 +1279,13 @@ def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: """Return the largest cache length for which selection is an identity. - With a decode-only budget, pinned prompt tokens do not consume ``top_B``. + With a decode-only budget, pinned prompt tokens do not consume ``budget``. Selection therefore keeps every token until the cache exceeds - ``prompt_len + top_B``. The constructor guarantees the decode-only + ``prompt_len + budget``. The constructor guarantees the decode-only budget (``pin_prefill=True``, ``count_prompt_tokens=False``). """ prompt_len = min(int(request.py_prompt_len), seq_len) - return prompt_len + self.top_B + return prompt_len + self.budget def _local_score_calibration( self, @@ -1437,9 +1437,9 @@ def _attention_layer_partition( "TriAttention requires a positive integer model sliding_window " "when layer_types contains sliding attention" ) - if self.top_B < raw_window: + if self.budget < raw_window: raise ValueError( - f"TriAttention decode budget top_B={self.top_B} must be at least " + f"TriAttention budget={self.budget} must be at least " f"the kernel-masked SWA window size {raw_window}" ) window_size = raw_window @@ -1634,7 +1634,7 @@ def _workspace_for( The request capacity follows the executor's max batch size (memory scales linearly with it) and the decode-width capacity follows the eviction bound (compaction keeps the scored decode region near - ``top_B`` plus one period of growth), so one workspace serves every + ``budget`` plus one period of growth), so one workspace serves every round. It is rebuilt only when the pool views change or a round outgrows it. """ @@ -1652,7 +1652,7 @@ def _workspace_for( ) fingerprint = ( self.eviction_mode, - self.top_B, + self.budget, tuple(layout["dense_layers"]), layout["pool_view_fingerprint"], draft_fingerprint, @@ -1674,7 +1674,7 @@ def _workspace_for( request_capacity = max(needed_requests, int(mgr.max_batch_size)) decode_width = max( needed_width, - self.top_B + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), + self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), ) # Bucket the score scratch by what cohorts actually present instead of # pinning it to max_seq_len: with pinned prompts the post-compaction @@ -1755,7 +1755,7 @@ def _workspace_for( seq_len=seq_capacity, num_q_heads=int(self._H), num_freqs=int(self._F), - keep_count=self.top_B, + keep_count=self.budget, q_real=q_real, q_imag=q_imag, mlr_coef=mlr_coef, @@ -1794,14 +1794,14 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: return offsets tails = [item["protected_tail"] for item in prepared] - dense = padded_offsets([self.top_B + tail for tail in tails]) + dense = padded_offsets([self.budget + tail for tail in tails]) swa = None if layout["swa_layers"] and layout["swa_window"]: swa = padded_offsets([int(layout["swa_window"]) + tail for tail in tails]) draft = None if self.draft_kv_cache_manager is not None: draft_tail = self._draft_protected_tail_capacity() - draft = padded_offsets([self.top_B + draft_tail] * len(prepared)) + draft = padded_offsets([self.budget + draft_tail] * len(prepared)) return dense, swa, draft def _page_table_pool_keys( @@ -1858,10 +1858,10 @@ def _evict_requests( # SWA landing positions are prompt-dependent; reject a request # whose retained span cannot cover the model window this round. for item in prepared: - if item["prompt_len"] + self.top_B < int(layout["swa_window"]): + if item["prompt_len"] + self.budget < int(layout["swa_window"]): raise ValueError( f"Request {item['request_id']} retains " - f"{item['prompt_len'] + self.top_B} tokens, below the " + f"{item['prompt_len'] + self.budget} tokens, below the " f"sliding window {layout['swa_window']}" ) with nvtx_range_debug("triattention.page_table_stage", color="orange"): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 8c28829b31a0..cf2ba9ec9b75 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3456,11 +3456,11 @@ class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): default=True, description="Always preserve the prompt (prefill) tokens; only decode " "tokens compete for the budget (upstream behaviour).") - top_B: int = Field( + budget: int = Field( default=2048, gt=0, - description="Tokens kept at each periodic eviction (upstream `budget`; " - "prompt tokens are always preserved on top).") + description="Tokens kept at each periodic eviction; prompt tokens are " + "always preserved on top.") beta: int = Field( default=128, gt=0, @@ -3499,7 +3499,7 @@ def _require_calibration_inputs(self): def to_manager_kwargs(self) -> dict: """Constructor kwargs for the TriAttention manager.""" return { - "top_B": self.top_B, + "budget": self.budget, "beta": self.beta, "model_path": self.model_path, "calibration_path": self.calibration_path, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index bc19de3efec7..047aaa3bbbb5 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -560,6 +560,13 @@ "kind": "value", "path": "kv_cache_compression_config.beta" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "kv_cache_compression_config.budget" + }, { "allowed_values": [], "annotation": "", @@ -592,13 +599,6 @@ "kind": "value", "path": "kv_cache_compression_config.pin_prefill" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_cache_compression_config.top_B" - }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 1c5157da42f2..af512e8446ad 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -269,7 +269,7 @@ def make_triattention(**overrides): """Construct a fully initialized manager for method-level unit tests.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention - options = {"top_B": 8, "model_path": "/models/test"} + options = {"budget": 8, "model_path": "/models/test"} options.update(overrides) return TriAttention(make_fake_v2(), **options) 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 index 2be6e28fbc9d..4a28f77a9cd5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -245,7 +245,7 @@ def test_draft_admission_gates_raise(gate, match): with pytest.raises(ValueError, match=match): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + model_path="/models/test", calibration_path="/calib/test.pt", budget=8 ), spec_config, draft_manager, @@ -253,7 +253,7 @@ def test_draft_admission_gates_raise(gate, match): return manager = TriAttention( _make_fake_v2(), - top_B=8, + budget=8, model_path="/models/test", eviction_mode="per_head" if gate == "union_only_per_head" else "union", draft_kv_cache_manager=draft_manager, @@ -268,7 +268,7 @@ def test_draft_admission_gates_raise(gate, match): def test_compressed_count_is_monotone_and_tracks_confirmed_length(): - manager = _make_triattention(top_B=4, beta=4) + manager = _make_triattention(budget=4, beta=4) manager._calibrated = True manager._attention_layer_partition_cache = ([0, 1], [], None) target = manager.kv_cache_manager @@ -337,7 +337,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - manager = _make_triattention(top_B=4) + manager = _make_triattention(budget=4) layout, workspace = _make_workspace_stubs(manager) # The one-time host block-offset table shape gate reads the real manager # tables (int32 [pools, slots, K/V, blocks]). @@ -390,7 +390,7 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): assert prepare.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 # Migrated from the pipeline workspace-kwargs test: the budget, the # shared phase-table dict, and the pool keys thread through unchanged. - assert prepare.call_args.kwargs["keep_count"] == manager.top_B + assert prepare.call_args.kwargs["keep_count"] == manager.budget assert prepare.call_args.kwargs["phase"] is manager._phase assert prepare.call_args.kwargs["layer_pool_keys"] == list(layout["layer_pool_keys"]) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index c9c189b5baf1..da164a95f1a4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -135,7 +135,7 @@ def test_llm_args_dispatch_and_validation(self): tri_args.kv_cache_compression_config, TriAttentionKvCacheCompressionConfig, ) - assert tri_args.kv_cache_compression_config.top_B == 2048 + assert tri_args.kv_cache_compression_config.budget == 2048 assert tri_args.kv_cache_compression_config.beta == 128 # The union dispatches on the algorithm tag, so an unknown algorithm @@ -153,7 +153,7 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): # no calibration file or CUDA. fake_v2 = _make_fake_v2(enable_block_reuse=False) cfg = TriAttentionKvCacheCompressionConfig( - top_B=32, + budget=32, beta=16, eviction_mode="per_head", model_path="/models/test", @@ -161,7 +161,7 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): ) mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) - assert mgr.top_B == 32 + assert mgr.budget == 32 assert mgr.beta == 16 assert mgr.eviction_mode == "per_head" assert mgr.kv_cache_manager is fake_v2 @@ -220,7 +220,7 @@ def test_request_init_marks_capacity_only_and_tracks_state(self): manager = _make_fake_v2() manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 - triattention = TriAttention(manager, top_B=8, model_path="/models/test") + triattention = TriAttention(manager, budget=8, model_path="/models/test") triattention._attention_layer_partition_cache = ([], [], None) triattention._calibrated = True @@ -247,7 +247,7 @@ class TestCompressedTokenPublication: # test_compressed_count_is_monotone_and_tracks_confirmed_length. def test_identity_compaction_is_rejected_instead_of_published(self): - manager = _make_triattention(top_B=4) + manager = _make_triattention(budget=4) manager.kv_cache_manager._stream = mock.Mock() request = _make_request(7, py_prompt_len=2) # Selection keeps every token at seq_len == prompt + budget: the due @@ -350,7 +350,7 @@ def _make_due_decode_request(seq_len): mgr._request_states = {} _set_request_state(mgr, 7, generation_steps=127) mgr.beta = 128 - mgr.top_B = 4096 + mgr.budget = 4096 return mgr, request, batch def test_suspended_cache_rejects_batch_before_cadence_mutation(self): @@ -486,7 +486,7 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): draft_manager.max_seq_len = 8192 manager = TriAttention( _make_fake_v2(), - top_B=8, + budget=8, model_path="/models/test", draft_kv_cache_manager=draft_manager, ) @@ -500,7 +500,7 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): validate_kv_cache_compression_with_spec( TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", top_B=8 + model_path="/models/test", calibration_path="/calib/test.pt", budget=8 ), MTPDecodingConfig(max_draft_len=1), draft_manager, @@ -524,7 +524,7 @@ def test_prepare_snapshots_fixed_linear_generation_growth( manager.kv_cache_map = { 7: SimpleNamespace(capacity=106, is_active=True), } - triattention = TriAttention(manager, top_B=8, model_path="/models/test") + triattention = TriAttention(manager, budget=8, model_path="/models/test") batch = SimpleNamespace( context_requests=[], generation_requests=[_make_request(7, py_draft_tokens=draft_tokens)], @@ -541,7 +541,7 @@ def test_union_rejects_unnormalized_scores(self): # The fused pipeline (THE union path) always z-normalizes, so # normalize_scores=False is rejected loudly at construction. with pytest.raises(ValueError, match="normalize_scores=True"): - _make_triattention(top_B=4, eviction_mode="union", normalize_scores=False) + _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) def test_workspace_build_receives_mode_and_capacity_kwargs(self): # Mode/phase/pool-key threading and cached reuse are covered by the @@ -549,9 +549,9 @@ def test_workspace_build_receives_mode_and_capacity_kwargs(self): # test_triattention_draft_cocompaction.py. from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - manager = _make_triattention(top_B=4) + manager = _make_triattention(budget=4) # The workspace follows the executor limits: eight requests (max batch - # size) by 260 decode tokens (top_B plus two eviction periods). + # size) by 260 decode tokens (budget plus two eviction periods). layout, workspace = _make_workspace_stubs(manager) prepared = [ _prepared_eviction( @@ -697,7 +697,7 @@ def stage_once(): assert staging.block_offsets_device[0, 0, 0, :5].tolist() == [46, 48, 50, 52, 54] def test_staged_page_tables_bypass_per_request_cuda_materialization(self): - manager = _make_triattention(top_B=4) + manager = _make_triattention(budget=4) manager.kv_cache_manager.num_extra_kv_tokens = 3 manager.kv_cache_manager._stream = mock.Mock() manager.kv_cache_manager.get_batch_cache_indices = mock.Mock( @@ -732,7 +732,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ) # One batched staging call carries the whole cohort: request ids, - # round starts, pinned prompt lengths, and valid lengths. top_B=4: + # round starts, pinned prompt lengths, and valid lengths. budget=4: # per-request moves are keep + tail = [6, 7]; padded rows repeat the # final offset out to the request capacity. args = internals.stage.call_args @@ -891,11 +891,11 @@ def stage_round(): class TestKernelMaskedSwa: - @pytest.mark.parametrize("top_B,fits_window", [(128, True), (127, False)]) - def test_layer_partition_uses_local_config_and_validates_window(self, top_B, fits_window): + @pytest.mark.parametrize("budget,fits_window", [(128, True), (127, False)]) + def test_layer_partition_uses_local_config_and_validates_window(self, budget, fits_window): mgr = _make_triattention() mgr.model_path = "/models/gpt-oss" - mgr.top_B = top_B + mgr.budget = budget mgr.kv_cache_manager = SimpleNamespace(pp_layers=[0, 1, 2, 3]) config = _make_hf_config( layer_types=[ @@ -910,7 +910,7 @@ def test_layer_partition_uses_local_config_and_validates_window(self, top_B, fit 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="decode budget top_B=127"): + with pytest.raises(ValueError, match="budget=127"): mgr._attention_layer_partition(4) return dense, sliding, window = mgr._attention_layer_partition(4) From 30b4cb3655d78fd947d9f0ff8f75c0035e33d950 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 20:44:10 -0700 Subject: [PATCH 094/178] [None][fix] Resolve RoPE tables via the transformers 5.5 rope_parameters API transformers 5.5 moved rope_theta/rope_scaling into config.rope_parameters and dropped the "default" key from ROPE_INIT_FUNCTIONS; the old lookup threw inside the bare except and silently fell back to the analytic inv_freq with base=10000 for BOTH production models (qwen3: true theta 1,000,000; gpt-oss: true theta 150,000 plus a lost yarn attention_factor 1.3466). Only the official {metadata,stats} calibration conversion path was affected; flat calibration files carry omega/freq_scale_sq in-file. _rope_tables now reads both config generations (rope_parameters first, legacy rope_scaling/rope_theta fallback), resolves theta from the right location, computes plain RoPE with the standard formula, and routes scaled variants through transformers' rope-init so attention_factor is honored. The analytic fallback survives only for ImportError; unknown rope_type, layer-type-keyed schemas, and init failures raise loudly naming the model path and the rope config seen. Signed-off-by: tianruih --- .../triattention/triattention.py | 63 +++++++++++++++---- .../test_triattention_pipeline.py | 54 ++++++++++++++++ 2 files changed, 104 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 1e7027fc67dc..b5db030288b1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -1984,25 +1984,62 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: def _rope_tables(self, freq_count: int): """RoPE ``omega`` (inv_freq) + ``freq_scale_sq`` (squared position-0 amplitude) from the model config -- model-intrinsic, corpus-independent - (the official file does not store them). transformers' rope-init handles - plain and scaled RoPE; plain RoPE has attention_factor 1 so freq_scale_sq - is all ones. Falls back to the analytic inv_freq if rope-init is absent.""" + (the official file does not store them). Reads both config generations: + transformers>=5.5 ``rope_parameters`` (rope_theta folded inside) and the + legacy top-level ``rope_scaling``/``rope_theta``. Plain RoPE uses the + standard formula with the resolved theta (attention_factor 1); scaled + variants (yarn, llama3, ...) go through transformers' rope-init so their + attention_factor is honored. The analytic fallback survives ONLY for + ImportError (rope-init module absent); every other failure raises.""" + import transformers from transformers import AutoConfig cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() config_values = cfg.to_dict() head_dim = freq_count * 2 - base = float(config_values.get("rope_theta", 10000.0)) - try: - from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + rope_params = ( + config_values.get("rope_parameters") or config_values.get("rope_scaling") or {} + ) + if rope_params and all(isinstance(v, dict) for v in rope_params.values()): + raise ValueError( + f"TriAttention: layer-type-keyed rope_parameters are not supported for " + f"calibration conversion (model {self.model_path}); got {rope_params!r}." + ) + rope_type = rope_params.get("rope_type") or rope_params.get("type") or "default" + theta_seen = rope_params.get("rope_theta", config_values.get("rope_theta")) + base = float(theta_seen) if theta_seen is not None else 10000.0 + + def analytic_inv_freq(): + idx = torch.arange(0, head_dim, 2, dtype=torch.float32) + return (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() - scaling = config_values.get("rope_scaling") or {} - rope_type = scaling.get("rope_type") or scaling.get("type") or "default" - inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") + if rope_type == "default": + # The original RoPE formula (transformers>=5.5 computes it per-model + # and no longer keys "default" in ROPE_INIT_FUNCTIONS). + omega, scale_sq = analytic_inv_freq(), 1.0 + else: + try: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + except ImportError: + logger.warning( + f"TriAttention: transformers rope-init unavailable; using the analytic " + f"inv_freq with theta={base} for {self.model_path} and IGNORING " + f"rope_type={rope_type!r} scaling corrections." + ) + return analytic_inv_freq(), torch.ones(freq_count, dtype=torch.float32) + if rope_type not in ROPE_INIT_FUNCTIONS: + raise ValueError( + f"TriAttention: unknown rope_type {rope_type!r} for {self.model_path} " + f"(transformers {transformers.__version__} provides " + f"{sorted(ROPE_INIT_FUNCTIONS)}); rope config seen: {rope_params!r}." + ) + try: + inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") + except Exception as exc: + raise ValueError( + f"TriAttention: rope-init {rope_type!r} failed for {self.model_path}; " + f"rope config seen: {rope_params!r}." + ) from exc omega = inv_freq.to(torch.float32)[:freq_count].clone() scale_sq = float(attention_factor) ** 2 - except Exception: - idx = torch.arange(0, head_dim, 2, dtype=torch.float32) - omega = (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() - scale_sq = 1.0 return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index da164a95f1a4..574430bea322 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -239,6 +239,60 @@ def test_resolve_accepts_flat_pt(self, flat_calibration_pt): for key in ("E_q", "E_q_norm", "omega", "freq_scale_sq"): assert key in loaded + 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 + + mgr.model_path = plain + omega, freq_scale_sq = mgr._rope_tables(freq_count) + 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 + omega_yarn, freq_scale_sq_yarn = mgr._rope_tables(freq_count) + # 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 cumulative/monotone publication contract itself (including the From d475bf985abc3aa10660ab302281a82ddb96521b Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 21:26:14 -0700 Subject: [PATCH 095/178] [None][refactor] Hoist compaction to package level and inline the round's launches Knife 16: compaction.py moves to kv_cache_compression/ (it is a general post-eviction compaction component, not TriAttention-specific); prepare_eviction_workspace becomes init_eviction_buffers and the compaction builder init_compaction_buffers (symmetric constructors); the SimpleNamespace is now bufs everywhere and the word workspace dies. run_cache_compactions / launch_move_pack / build_move_pack_arguments dissolve: run_eviction_round fires the SWA rebase add, the draft-family pack kernel, and the per-family C++ compact ops directly with explicit arguments; the pack geometry (selection_rows / move_capacity / dense_total / settle launch data) is computed once in init_compaction_buffers and flattened onto plain bufs fields. Per-round host re-marshaling dies with it (perf item #18): the valid-widths subtraction folds into the phase-gather kernel, the score and union-finalize kernels read valid_seq_lens[req_id] / token_starts[req_id] as per-CTA scalars straight from the staged metadata rows (pointer capture, assumed_align=4 for the 4-byte-aligned row views), and the per-segment seg_seq_len index_select plus the cute_token_starts copy and both persistent buffers are deleted. Tests follow: conftest builds bundles via init_compaction_buffers and replicates the round's compact stage for standalone bundles (run_compaction/launch_family_pack); the per-layer order test attaches its bundle wholesale so the fused settle packs it production-shaped; the cute stage helpers write the staged metadata rows the runners now capture. Signed-off-by: tianruih --- .../{triattention => }/compaction.py | 268 +++----- .../triattention/triattention.py | 599 ++++++++++-------- .../triattention_cute_score_fused.py | 25 +- .../triattention_cute_selection.py | 8 +- .../triattention/triattention_kernels.py | 21 +- .../_torch/kv_cache_compression/conftest.py | 114 +++- .../test_triattention_cute_score.py | 91 +-- .../test_triattention_cute_union_fusion.py | 105 +-- .../test_triattention_draft_cocompaction.py | 57 +- .../test_triattention_pipeline.py | 71 ++- .../test_triattention_selection_compaction.py | 192 +++--- 11 files changed, 803 insertions(+), 748 deletions(-) rename tensorrt_llm/_torch/kv_cache_compression/{triattention => }/compaction.py (61%) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py similarity index 61% rename from tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py rename to tensorrt_llm/_torch/kv_cache_compression/compaction.py index 52d233cf16ef..052e35607645 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,17 +15,17 @@ """Batched physical KV-cache compaction for eviction-based compression. -Given each request's kept-token ordinals, its valid sequence length, and the -staged V2 block offsets, this module packs per-request move indices with one -Triton launch per compacted cache (one launch covers the target's dense and -SWA families; a co-compressed draft adds a second) and then moves the -surviving KV in place with batched C++ compact launches. Everything is plain -tensors and dicts: ``build_cache_compactions`` allocates the launch data once -per geometry (called by ``triattention.prepare_eviction_workspace``) and -``run_cache_compactions`` fires the kernels directly each round. The -target's dense/SWA packing always rides the driver's fused settle launch -(the returned ``dense_pack`` dict describes it); the draft always packs -here with its own launch. +A general post-eviction component: given each request's kept-token ordinals, +its valid sequence length, and the staged V2 block offsets, the surviving KV +moves in place through batched C++ compact launches, fed by per-request move +indices packed on device. Everything is plain tensors and dicts: +``init_compaction_buffers`` allocates the launch data once per geometry +(called by ``triattention.init_eviction_buffers``) and returns one bundle +whose fields the eviction driver fires directly each round -- the target's +dense/SWA packing rides the driver's fused settle launch (described by the +bundle's ``settle_pack_tensors``/``settle_pack_shape``), the co-compressed +draft's own pack launch and every family's C++ moves are inlined in +``triattention.run_eviction_round``. """ from collections import OrderedDict @@ -125,124 +125,7 @@ def _compact_groups( return tuple(result) -# Launch shape of the move-index packing kernel: tokens per program along the -# move axis, and its warp count. -_PACK_BLOCK_TOKENS = 256 -_PACK_NUM_WARPS = 4 - - -def build_move_pack_arguments( - kept_token_ordinals: torch.Tensor, - valid_sequence_lengths: torch.Tensor, - move_source_offsets: torch.Tensor, - move_source_indices: torch.Tensor, - *, - eviction_mode: str, - decode_keep_count: int, - num_dense_layers: int, - num_kv_heads: int, - max_protected_tail: int, - swa_window: int, - swa_move_source_offsets: Optional[torch.Tensor], - swa_move_source_indices: Optional[torch.Tensor], -) -> Dict[str, object]: - """Describe one move-index packing as plain kernel launch data. - - The packing kernel reads the kept-token ordinals and each request's valid - length and writes the packed per-(layer, head) move source indices - consumed by the C++ compact launches. ``launch_move_pack`` fires it - standalone; a fused selection-side settle launch consumes the same dict. - """ - per_layer = eviction_mode == "per_layer_perhead" - union = eviction_mode == "union" - request_count = int(kept_token_ordinals.shape[0]) if kept_token_ordinals.ndim else 0 - if union: - selection_rows = 1 - elif per_layer: - selection_rows = num_dense_layers * num_kv_heads - else: - selection_rows = num_kv_heads - # Selection rows carry decode-only kept ordinals (already absolute), so - # the rectangle is prompt-length independent. - if swa_move_source_indices is not None: - swa_offsets_arg = swa_move_source_offsets - swa_indices_arg = swa_move_source_indices - swa_total = int(swa_move_source_indices.shape[-1]) - else: - # HAS_SWA specializes all corresponding loads and stores away. - swa_offsets_arg = move_source_offsets - swa_indices_arg = move_source_indices - swa_total = 0 - - max_move = decode_keep_count + max_protected_tail - if swa_total: - max_move = max(max_move, swa_window + max_protected_tail) - return dict( - kept_token_ordinals=kept_token_ordinals, - valid_sequence_lengths=valid_sequence_lengths, - dense_offsets=move_source_offsets, - dense_indices=move_source_indices, - swa_offsets=swa_offsets_arg, - swa_indices=swa_indices_arg, - dense_total=int(move_source_indices.shape[-1]), - swa_total=swa_total, - selection_rows=selection_rows, - keep_count=decode_keep_count, - request_count=request_count, - num_kv_heads=num_kv_heads, - swa_window=swa_window, - # Widest per-request move count any staged offsets may express; the - # packing loop covers exactly this many move slots per packed row. - move_capacity=max_move, - union=union, - per_layer=per_layer, - has_swa=swa_total > 0, - ) - - -def launch_move_pack(pack: Dict[str, object]) -> None: - """Fire one standalone move-index packing launch. - - One program per (request, selection row); the settle half is compiled - away because the ordinals arrive pre-settled (the draft flow reuses the - target's keep set verbatim), so only the pack half runs. The settle-side - pointer arguments are compiled away with it; any well-formed tensor - stands in for them. - """ - from .triattention_kernels import _settle_ties_and_pack_compaction_sources_kernel - - kept = pack["kept_token_ordinals"] - _settle_ties_and_pack_compaction_sources_kernel[ - (pack["request_count"], pack["selection_rows"]) - ]( - kept, - pack["valid_sequence_lengths"], - pack["dense_offsets"], - kept, - kept, - pack["valid_sequence_lengths"], - pack["dense_offsets"], - pack["dense_indices"], - pack["swa_offsets"], - pack["swa_indices"], - WIDTH=pack["keep_count"], - KEEP_COUNT=pack["keep_count"], - SELECTION_ROWS=pack["selection_rows"], - DENSE_TOTAL=pack["dense_total"], - SWA_TOTAL=pack["swa_total"], - MOVE_CAPACITY=pack["move_capacity"], - NUM_KV_HEADS=pack["num_kv_heads"], - SWA_WINDOW=pack["swa_window"], - UNION=pack["union"], - PER_LAYER=pack["per_layer"], - HAS_SWA=pack["has_swa"], - HAS_SETTLE=False, - BLOCK=_PACK_BLOCK_TOKENS, - num_warps=_PACK_NUM_WARPS, - ) - - -def build_cache_compactions( +def init_compaction_buffers( *, eviction_mode: str, layer_pools: List[torch.Tensor], @@ -286,14 +169,17 @@ def build_cache_compactions( ``[slot, request, K/V, block]`` (offset = ``2*page + plane``); ``protected_tail_capacity`` is the widest per-request tail this geometry must support -- actual per-round lengths arrive through the staged - move-offset rows. The target's dense/SWA packing is left to the caller's - fused settle launch (the returned ``dense_pack`` dict describes it) and - only the C++ moves run here; the draft always keeps its own pack launch. - - Returns ``{"families": [...], "dense_pack": ..., "num_kv_heads": ..., - "swa_window": ..., "swa_destination_bases": ...}`` where each family is - ``{"pack": dict|None, "groups": (...), "source": t, "offsets": t, - "destination_bases": t}``. + move-offset rows. Nothing launches here: the target's dense/SWA packing + rides the driver's fused settle launch (``settle_pack_tensors`` + + ``settle_pack_shape`` describe it) and the driver's round function fires + the draft pack (``draft_pack``) and every family's C++ moves directly. + + Returns one plain bundle dict: ``families`` (each ``{"name", "groups", + "source", "offsets", "destination_bases"}``), the settle/pack launch data + above, ``draft_pack`` (``None`` or ``{"indices", "offsets", + "dense_total", "move_capacity", "num_kv_heads"}``), ``swa_rebase_delta`` + (per-round SWA destination rebase), and the geometry scalars + (``selection_rows``, ``request_count``, ``decode_keep_count``, ...). """ device = layer_pools[dense_layers[0]].device request_count = int(request_count) @@ -346,27 +232,47 @@ def build_cache_compactions( swa_entries = [(layer, layer_pools[layer], page_table_for(layer)) for layer in swa_layers] dense_slots = {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer else None - dense_pack = build_move_pack_arguments( - kept_token_ordinals, + union = eviction_mode == "union" + if union: + selection_rows = 1 + elif per_layer: + selection_rows = len(dense_layers) * num_kv_heads + else: + selection_rows = num_kv_heads + # Selection rows carry decode-only kept ordinals (already absolute), so + # the rectangle is prompt-length independent. HAS_SWA specializes the SWA + # loads and stores away, so without SWA layers the dense buffers stand in + # for the compiled-away SWA pointer arguments. + has_swa = swa_move_indices is not None + swa_total = int(swa_move_indices.shape[-1]) if has_swa else 0 + # Widest per-request move count any staged offsets may express; the + # packing loop covers exactly this many move slots per packed row. + move_capacity = decode_keep_count + protected_tail_capacity + if has_swa: + move_capacity = max(move_capacity, swa_window + protected_tail_capacity) + settle_pack_tensors = ( valid_sequence_lengths, dense_move_offsets, dense_move_indices, - eviction_mode=eviction_mode, - decode_keep_count=decode_keep_count, - num_dense_layers=len(dense_layers), - num_kv_heads=num_kv_heads, - max_protected_tail=protected_tail_capacity, - swa_window=swa_window, - swa_move_source_offsets=swa_move_offsets, - swa_move_source_indices=swa_move_indices, + swa_move_offsets if has_swa else dense_move_offsets, + swa_move_indices if has_swa else dense_move_indices, + ) + settle_pack_shape = dict( + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=swa_total, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window, + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=has_swa, ) families = [ dict( name="dense", # The fused selection-side settle launch packs the dense/SWA move # sources when it finalizes the kept ordinals; only the C++ moves - # stay here. Each round then packs exactly once. - pack=None, + # consume these buffers. Each round then packs exactly once. groups=_compact_groups(dense_entries, layer_pool_keys, device, dense_slots), source=dense_move_indices, offsets=dense_move_offsets, @@ -374,11 +280,10 @@ def build_cache_compactions( ) ] if swa_layers: - # The dense pack call fills the SWA move buffers in the same run. + # The fused dense pack fills the SWA move buffers in the same launch. families.append( dict( name="swa", - pack=None, groups=_compact_groups(swa_entries, layer_pool_keys, device), source=swa_move_indices, offsets=swa_move_offsets, @@ -410,39 +315,42 @@ def build_cache_compactions( for layer in draft_layers ] # In union mode the pack kernel reads selection row 0 for every packed - # row, so one more pack launch broadcasts the target keep set over the - # draft KV heads and appends the draft's own tail ordinals. - draft_pack = build_move_pack_arguments( - kept_token_ordinals, - valid_sequence_lengths, - draft_move_offsets, - draft_move_indices, - eviction_mode="union", - decode_keep_count=decode_keep_count, - num_dense_layers=1, + # row, so the driver fires one more pack launch broadcasting the + # target keep set over the draft KV heads and appending the draft's + # own tail ordinals; these are that launch's geometry constants. + draft_pack = dict( + indices=draft_move_indices, + offsets=draft_move_offsets, + dense_total=int(draft_move_indices.shape[-1]), + move_capacity=decode_keep_count + draft_tail, num_kv_heads=draft_num_kv_heads, - max_protected_tail=draft_tail, - swa_window=0, - swa_move_source_offsets=None, - swa_move_source_indices=None, ) families.append( dict( name="draft", - pack=draft_pack, groups=_compact_groups(draft_entries, tuple(draft_layer_pool_keys), device), source=draft_move_indices, offsets=draft_move_offsets, destination_bases=prompt_offsets, ) ) + else: + draft_pack = None return dict( families=families, - dense_pack=dense_pack, + kept_token_ordinals=kept_token_ordinals, + valid_sequence_lengths=valid_sequence_lengths, + selection_rows=selection_rows, + settle_pack_tensors=settle_pack_tensors, + settle_pack_shape=settle_pack_shape, + draft_pack=draft_pack, num_kv_heads=num_kv_heads, swa_window=swa_window, swa_destination_bases=swa_destination_bases, + # The prompt offsets may be re-staged each round; the driver rebases + # the SWA landing positions with this delta before the moves. + swa_rebase_delta=decode_keep_count - swa_window, prompt_offsets=prompt_offsets, decode_keep_count=decode_keep_count, request_count=request_count, @@ -451,29 +359,3 @@ def build_cache_compactions( int(draft_protected_tail_capacity or 0) if draft_layers else 0 ), ) - - -def run_cache_compactions(compaction: Dict[str, object]) -> None: - """Pack the move indices, then run every cache family's C++ compacts.""" - swa_destination_bases = compaction["swa_destination_bases"] - if swa_destination_bases is not None: - # The prompt offsets may have been re-staged since construction; - # rebase the SWA landing positions for this round. - torch.add( - compaction["prompt_offsets"], - compaction["decode_keep_count"] - compaction["swa_window"], - out=swa_destination_bases, - ) - for family in compaction["families"]: - if family["pack"] is not None: - launch_move_pack(family["pack"]) - for group in family["groups"]: - torch.ops.trtllm.sparse_kv_cache_compact_layers( - group["pools"], - group["pool_pointers"], - group["page_table"], - family["source"], - family["offsets"], - family["destination_bases"], - group["source_layer_indices"], - ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b5db030288b1..17bf80fecf11 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -35,12 +35,13 @@ in the same round with the target's kept token set (union mode only), so target and draft always share one physical KV length. -Structure: ``prepare_eviction_workspace`` is the ONE one-time constructor for -the whole eviction stack -- it validates the geometry, allocates every buffer, +Structure: ``init_eviction_buffers`` is the ONE one-time constructor for the +whole eviction stack -- it validates the geometry, allocates every buffer, compiles the mode-needed CuTe entries eagerly, and builds the C++ compaction -launch data. The result is a plain namespace of tensors, events, and ints. -``stage_eviction_cohort`` and ``run_eviction_round`` are the per-round flow: -straight-line module functions that feed the kernels directly. +launch data (``init_compaction_buffers``). The result is a plain namespace of +tensors, events, and ints. ``stage_eviction_cohort`` and +``run_eviction_round`` are the per-round flow: straight-line module functions +that feed the kernels directly. KV layout: the decode kernel stores keys in HND layout ``[num_pages, kv_factor, num_kv_heads, tokens_per_block, head_dim]``. The Python @@ -75,7 +76,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig -from .compaction import build_cache_compactions, run_cache_compactions +from ..compaction import init_compaction_buffers from .triattention_kernels import ( _settle_ties_and_pack_compaction_sources_kernel, build_mean_phase_table, @@ -158,7 +159,7 @@ def _allocate_page_table_plane( return representative_slots, copy_block_count, host, dev -def prepare_eviction_workspace( +def init_eviction_buffers( *, eviction_mode: str, layer_pools: List[torch.Tensor], @@ -196,7 +197,7 @@ def prepare_eviction_workspace( draft_page_table_token_capacity: Optional[int] = None, draft_protected_tail_capacity: int = 0, ) -> SimpleNamespace: - """Build the ONE plain-namespace workspace for the whole eviction stack. + """Build the ONE plain namespace of buffers for the whole eviction stack. The single one-time constructor: buffer staging, eager CuTe compilation for exactly the entries the eviction mode launches, selection buffers, @@ -205,7 +206,7 @@ def prepare_eviction_workspace( own geometry contract and raises loudly -- there is no fallback path. The returned namespace holds tensors, events, compiled runners, and ints; all flow logic lives in ``stage_eviction_cohort`` and - ``run_eviction_round``. The workspace retains references to every scored + ``run_eviction_round``. The buffers retain references to every scored layer pool: the compiled kernels encode immutable TMA descriptors from their raw device addresses, so the pools must stay alive and stay put. """ @@ -233,21 +234,21 @@ def prepare_eviction_workspace( if page_table_keys is None: page_table_keys = list(range(len(page_representatives))) - ws = SimpleNamespace() - ws.eviction_mode = eviction_mode - ws.device = device - ws.max_requests = max_requests - ws.bucket_seq_len = seq_len - ws.decode_width = decode_width - ws.keep_count = keep_count - ws.page_table_token_capacity = page_table_token_capacity + bufs = SimpleNamespace() + bufs.eviction_mode = eviction_mode + bufs.device = device + bufs.max_requests = max_requests + bufs.bucket_seq_len = seq_len + bufs.decode_width = decode_width + bufs.keep_count = keep_count + bufs.page_table_token_capacity = page_table_token_capacity # ---- staged page-table planes (target, plus the co-compressed draft) --- ( - ws.representative_slots, - ws.copy_block_count, - ws._bulk_offsets_src, - ws.block_offsets_device, + bufs.representative_slots, + bufs.copy_block_count, + bufs._bulk_offsets_src, + bufs.block_offsets_device, ) = _allocate_page_table_plane( layer_pools, page_representatives, @@ -259,16 +260,16 @@ def prepare_eviction_workspace( "", ) # The draft is never scored: these offsets feed only the draft compacts. - ws.draft_block_offsets_device = None - ws._draft_bulk_offsets_src = None - ws.draft_representative_slots = {} - ws.draft_copy_block_count = 0 + bufs.draft_block_offsets_device = None + bufs._draft_bulk_offsets_src = None + bufs.draft_representative_slots = {} + bufs.draft_copy_block_count = 0 if draft_layer_pools is not None: ( - ws.draft_representative_slots, - ws.draft_copy_block_count, - ws._draft_bulk_offsets_src, - ws.draft_block_offsets_device, + bufs.draft_representative_slots, + bufs.draft_copy_block_count, + bufs._draft_bulk_offsets_src, + bufs.draft_block_offsets_device, ) = _allocate_page_table_plane( draft_layer_pools, draft_page_representatives, @@ -284,33 +285,33 @@ def prepare_eviction_workspace( # Three metadata rows (logical position, valid length, prompt length) plus # one move-offsets row per compacted cache family; offsets rows have # request_capacity + 1 entries, hence the extra column. - ws.request_metadata_host = torch.empty( + bufs.request_metadata_host = torch.empty( (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) - ws._bulk_copy_idx_src = torch.arange( + bufs._bulk_copy_idx_src = torch.arange( max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) # Zero-filled so an unstaged cohort gathers the phase table's row 0 # instead of indexing it with uninitialized round starts. - ws.request_metadata_device = torch.zeros( + bufs.request_metadata_device = torch.zeros( (6, max_requests + 1), dtype=torch.int32, device=device ) - ws.round_starts_device = ws.request_metadata_device[0, :max_requests] - ws.valid_seq_lens_device = ws.request_metadata_device[1, :max_requests] + bufs.round_starts_device = bufs.request_metadata_device[0, :max_requests] + bufs.valid_seq_lens_device = bufs.request_metadata_device[1, :max_requests] # Per-request pinned prompt lengths: the score kernel starts each # request's decode window here, so one bucket may mix prompt lengths. - ws.token_starts_device = ws.request_metadata_device[2, :max_requests] - ws.dense_move_offsets = ws.request_metadata_device[3] - ws.swa_move_offsets = ws.request_metadata_device[4] - ws.draft_move_offsets = ws.request_metadata_device[5] - ws.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) - ws.mean_sin = torch.empty_like(ws.mean_cos) + bufs.token_starts_device = bufs.request_metadata_device[2, :max_requests] + bufs.dense_move_offsets = bufs.request_metadata_device[3] + bufs.swa_move_offsets = bufs.request_metadata_device[4] + bufs.draft_move_offsets = bufs.request_metadata_device[5] + bufs.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) + bufs.mean_sin = torch.empty_like(bufs.mean_cos) # The phase table depends only on the shared calibration; the manager - # shares one dict with every workspace (tests may pass None for a + # shares one dict with every buffer namespace (tests may pass None for a # private table). if phase is None: phase = build_mean_phase_table(offsets, omega, initial_rows=seq_len) - ws.phase = phase + bufs.phase = phase # ---- score state: ONE fused group across ALL dense layers -------------- # Segments carry their own page-table slot, so distinct per-layer @@ -318,34 +319,34 @@ def prepare_eviction_workspace( # layout are validated by the CuTe runner itself below. p0 = layer_pools[dense_layers[0]] _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape - ws.num_layers = len(dense_layers) - ws.num_q_heads = int(num_q_heads) - ws.num_kv_heads = int(num_kv_heads) - ws.num_freqs = int(num_freqs) - ws.tokens_per_block = int(tokens_per_block) + bufs.num_layers = len(dense_layers) + bufs.num_q_heads = int(num_q_heads) + bufs.num_kv_heads = int(num_kv_heads) + bufs.num_freqs = int(num_freqs) + bufs.tokens_per_block = int(tokens_per_block) # Calibration tables span every model layer; segments index them by # ABSOLUTE layer id ON DEVICE where they cannot be range-checked, so # validate the extent once here, loudly. - num_calibrated_layers = q_real.numel() // (ws.num_q_heads * ws.num_freqs) + num_calibrated_layers = q_real.numel() // (bufs.num_q_heads * bufs.num_freqs) if min(dense_layers) < 0 or max(dense_layers) >= num_calibrated_layers: raise ValueError("scored layer index exceeds the calibrated layer extent") _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} - page_table_slots = [ws.representative_slots[_rep_of[layer]] for layer in dense_layers] - ws.seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( - ws.num_layers + page_table_slots = [bufs.representative_slots[_rep_of[layer]] for layer in dense_layers] + bufs.seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( + bufs.num_layers ) seg_layer = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( max_requests ) - block_offsets = ws.block_offsets_device + block_offsets = bufs.block_offsets_device slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) req_idx = torch.arange(max_requests, dtype=torch.int64, device=device).repeat_interleave( - ws.num_layers + bufs.num_layers ) slot_idx = slots_t.repeat(max_requests) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - max_segments = max_requests * ws.num_layers + max_segments = max_requests * bufs.num_layers # Head axis pads to the MMA tile N=8; one padded scratch plane must stay # 32-bit indexable (wraparound = silent wild reads, not a clean error). if (8 - 1) * max_segments * seq_len >= 2**31: @@ -355,53 +356,55 @@ def prepare_eviction_workspace( # The kernel scores each request's window into a head-major scratch; # all buffers below are persistent because the compiled kernels capture # their device pointers. - ws.cute_scratch = torch.empty( - ws.num_kv_heads * 8 * max_segments * seq_len, dtype=torch.float32, device=device + bufs.cute_scratch = torch.empty( + bufs.num_kv_heads * 8 * max_segments * seq_len, dtype=torch.float32, device=device ) - ws.seg_seq_len = torch.zeros(max_segments, dtype=torch.int32, device=device) seg_out_offset = (torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len).to( torch.int32 ) - ws.cute_token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) - ws.gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( + bufs.gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( 1, 1, 1, 1, -1 ) - # Compile the SM100 CuTe entries this mode launches -- HERE at workspace + # Compile the SM100 CuTe entries this mode launches -- HERE at buffer # construction, outside any CUDA graph capture (compilation allocates and # synchronizes). Union rounds run the fused score+stats+union pipeline; # the per-head modes run the score-only entry. There is deliberately no # other score path and no fallback. union = eviction_mode == "union" - ws.union_rows = None + bufs.union_rows = None if union: # Union output rows are sized by the whole bucket (the widest # possible window); consumers mask by the per-request widths. - ws.union_rows = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) + bufs.union_rows = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) try: - ws.runner = TriAttentionCuteScoreRunner( + bufs.runner = TriAttentionCuteScoreRunner( layer_pools=list(layer_pools), layer_indices=[int(layer) for layer in dense_layers], max_requests=max_requests, - num_layers=ws.num_layers, + num_layers=bufs.num_layers, seq_len=seq_len, - num_q_heads=ws.num_q_heads, - num_kv_heads=ws.num_kv_heads, - num_freqs=ws.num_freqs, - tokens_per_block=ws.tokens_per_block, + num_q_heads=bufs.num_q_heads, + num_kv_heads=bufs.num_kv_heads, + num_freqs=bufs.num_freqs, + tokens_per_block=bufs.tokens_per_block, page_ids=block_offsets.view(-1), seg_page_off=seg_page_off, - seg_req_id=ws.seg_req, + seg_req_id=bufs.seg_req, seg_layer_id=seg_layer, - seg_seq_len=ws.seg_seq_len, + # The kernels read per-request lengths and window starts straight + # from the staged metadata rows (pointer capture; the rows are + # int32 views into the persistent metadata table, so the round + # needs no per-round re-marshaling copies for them). + valid_seq_lens=bufs.valid_seq_lens_device, seg_out_offset=seg_out_offset, - token_starts=ws.cute_token_starts, + token_starts=bufs.token_starts_device, q_real=q_real.view(-1), q_imag=q_imag.view(-1), mlr_coef=mlr_coef.view(-1), - mean_cos=ws.mean_cos, - mean_sin=ws.mean_sin, + mean_cos=bufs.mean_cos, + mean_sin=bufs.mean_sin, freq_scale_sq=freq_scale_sq, - output=ws.cute_scratch, + output=bufs.cute_scratch, enable_partial_stats=union, ) except (ImportError, RuntimeError, TypeError, ValueError, AssertionError) as error: @@ -409,76 +412,82 @@ def prepare_eviction_workspace( "TriAttention CuTe score setup failed and no other score path exists" ) from error logger.info( - f"TriAttention CuTe score enabled: {ws.num_q_heads}q/{ws.num_kv_heads}kv heads, " - f"{ws.num_freqs} freqs, {ws.tokens_per_block}-token pages" + f"TriAttention CuTe score enabled: {bufs.num_q_heads}q/{bufs.num_kv_heads}kv heads, " + f"{bufs.num_freqs} freqs, {bufs.tokens_per_block}-token pages" ) # ---- selection buffers -------------------------------------------------- # Per-request valid decode widths, refreshed each round from the staged # lengths; prompt offsets alias the staged per-request prompt lengths so # the values are written once per round. - ws.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) - ws.prompt_offsets = ws.token_starts_device + bufs.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) + bufs.prompt_offsets = bufs.token_starts_device if union: - ws.selection_rows_per_request = 1 - ws.row_prompt_offsets = ws.prompt_offsets + bufs.selection_rows_per_request = 1 + bufs.row_prompt_offsets = bufs.prompt_offsets # The fused pipeline writes normalized per-request union rows straight # into ``combined``; only the top-k settle-and-pack stage remains. - ws.combined = torch.empty((max_requests, decode_width), dtype=torch.float32, device=device) - ws.final_indices = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) + bufs.combined = torch.empty( + (max_requests, decode_width), dtype=torch.float32, device=device + ) + bufs.final_indices = torch.empty( + (max_requests, keep_count), dtype=torch.int32, device=device + ) # Kept decode ordinals only: rows are prompt-length independent, so - # one workspace serves cohorts with mixed prompt lengths. - ws.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) + # one buffer namespace serves cohorts with mixed prompt lengths. + bufs.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) # Row-major views consumed by the top-k settle launch. - ws.selection_scores_rows = ws.combined - ws.selection_row_lengths = ws.valid_widths - ws.provisional_rows = ws.final_indices - ws.keep_rows = ws.keep + bufs.selection_scores_rows = bufs.combined + bufs.selection_row_lengths = bufs.valid_widths + bufs.provisional_rows = bufs.final_indices + bufs.keep_rows = bufs.keep # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. - ws.final_indices.zero_() - ws.score_output = None + bufs.final_indices.zero_() + bufs.score_output = None else: selection_rows = ( - ws.num_kv_heads if eviction_mode == "per_head" else ws.num_layers * ws.num_kv_heads + bufs.num_kv_heads + if eviction_mode == "per_head" + else bufs.num_layers * bufs.num_kv_heads ) - ws.selection_rows_per_request = selection_rows - ws.row_prompt_offsets = torch.zeros( + bufs.selection_rows_per_request = selection_rows + bufs.row_prompt_offsets = torch.zeros( (max_requests * selection_rows,), dtype=torch.int32, device=device ) # Decode-only per-head scores gathered from the CuTe scratch, the # ``[request, layer, head, token]`` layout the reduce kernels read. - ws.score_output = torch.empty( + bufs.score_output = torch.empty( max_requests, - ws.num_layers, - ws.num_q_heads, + bufs.num_layers, + bufs.num_q_heads, decode_width, dtype=torch.float32, device=device, ) - score_shape = (max_requests, ws.num_layers, ws.num_q_heads, 1) - ws.row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) - ws.row_std = torch.empty_like(ws.row_mean) - ws.selection_scores = torch.empty( + score_shape = (max_requests, bufs.num_layers, bufs.num_q_heads, 1) + bufs.row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) + bufs.row_std = torch.empty_like(bufs.row_mean) + bufs.selection_scores = torch.empty( (max_requests, selection_rows, decode_width), dtype=torch.float32, device=device ) - ws.row_seq_lens = torch.full( + bufs.row_seq_lens = torch.full( (max_requests, selection_rows), decode_width, dtype=torch.int32, device=device ) selection_shape = (max_requests, selection_rows, keep_count) - ws.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=device) - ws.keep = torch.empty(selection_shape, dtype=torch.int32, device=device) - ws.selection_scores_rows = ws.selection_scores.view( + bufs.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=device) + bufs.keep = torch.empty(selection_shape, dtype=torch.int32, device=device) + bufs.selection_scores_rows = bufs.selection_scores.view( max_requests * selection_rows, decode_width ) - ws.selection_row_lengths = ws.row_seq_lens.view(-1) - ws.provisional_rows = ws.top_indices_i32.view(-1, keep_count) - ws.keep_rows = ws.keep.view(-1, keep_count) - ws.top_indices_i32.zero_() + bufs.selection_row_lengths = bufs.row_seq_lens.view(-1) + bufs.provisional_rows = bufs.top_indices_i32.view(-1, keep_count) + bufs.keep_rows = bufs.keep.view(-1, keep_count) + bufs.top_indices_i32.zero_() # ---- compaction launch data + settle/pack fusion ------------------------ # One settle program per (request, selection row). - ws.settle_grid = (max_requests, ws.selection_rows_per_request) + bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) draft_kwargs = {} if draft_layers: draft_kwargs = dict( @@ -487,68 +496,57 @@ def prepare_eviction_workspace( draft_layer_group_representative=draft_layer_group_representative, draft_layer_pool_keys=draft_layer_pool_keys, draft_protected_tail_capacity=int(draft_protected_tail_capacity), - draft_kv_block_offsets=ws.draft_block_offsets_device, - draft_page_table_slots=ws.draft_representative_slots, - draft_move_offsets=ws.draft_move_offsets, + draft_kv_block_offsets=bufs.draft_block_offsets_device, + draft_page_table_slots=bufs.draft_representative_slots, + draft_move_offsets=bufs.draft_move_offsets, ) - ws.compaction = build_cache_compactions( + compaction = init_compaction_buffers( eviction_mode=eviction_mode, layer_pools=layer_pools, dense_layers=list(dense_layers), swa_layers=list(swa_layers), layer_group_representative=layer_group_representative, - kept_token_ordinals=ws.keep, - valid_sequence_lengths=ws.valid_seq_lens_device, - kv_block_offsets=ws.block_offsets_device, - page_table_slots=ws.representative_slots, + kept_token_ordinals=bufs.keep, + valid_sequence_lengths=bufs.valid_seq_lens_device, + kv_block_offsets=bufs.block_offsets_device, + page_table_slots=bufs.representative_slots, request_count=max_requests, - prompt_offsets=ws.token_starts_device, + prompt_offsets=bufs.token_starts_device, decode_keep_count=keep_count, swa_window=swa_window, layer_pool_keys=list(layer_pool_keys), protected_tail_capacity=int(protected_tail_capacity), # Tails vary per round (in-flight growth), so the per-family move # offsets ride the staged metadata rows each round. - dense_move_offsets=ws.dense_move_offsets, - swa_move_offsets=ws.swa_move_offsets, - # ONE launch settles the kept ordinals and packs the dense/SWA - # move sources; ``run_cache_compactions`` then only runs the C++ - # moves (plus the draft's own pack). + dense_move_offsets=bufs.dense_move_offsets, + swa_move_offsets=bufs.swa_move_offsets, **draft_kwargs, ) - pack = ws.compaction["dense_pack"] - ws.settle_pack_tensors = ( - pack["valid_sequence_lengths"], - pack["dense_offsets"], - pack["dense_indices"], - pack["swa_offsets"], - pack["swa_indices"], - ) - ws.settle_pack_shape = dict( - DENSE_TOTAL=pack["dense_total"], - SWA_TOTAL=pack["swa_total"], - MOVE_CAPACITY=pack["move_capacity"], - NUM_KV_HEADS=pack["num_kv_heads"], - SWA_WINDOW=pack["swa_window"], - UNION=pack["union"], - PER_LAYER=pack["per_layer"], - HAS_SWA=pack["has_swa"], - ) + # Flatten the launch data to plain fields: ONE fused launch settles the + # kept ordinals and packs the dense/SWA move sources + # (``settle_top_tokens``); ``run_eviction_round`` fires the draft pack + # and every family's C++ moves directly. + bufs.compaction_families = compaction["families"] + bufs.settle_pack_tensors = compaction["settle_pack_tensors"] + bufs.settle_pack_shape = compaction["settle_pack_shape"] + bufs.draft_pack = compaction["draft_pack"] + bufs.swa_destination_bases = compaction["swa_destination_bases"] + bufs.swa_rebase_delta = compaction["swa_rebase_delta"] # ---- round-ordering events ---------------------------------------------- - ws.copy_done = torch.cuda.Event() + bufs.copy_done = torch.cuda.Event() # First record publishes constructor allocations to the V2 copy stream; # later records protect pinned metadata before the next cohort reuses it. - ws.copy_done.record(torch.cuda.current_stream(device)) - ws.bulk_copy_done = torch.cuda.Event() - ws.bulk_consume_done = torch.cuda.Event() - ws.copy_pending = False - ws.page_tables_active = False - return ws + bufs.copy_done.record(torch.cuda.current_stream(device)) + bufs.bulk_copy_done = torch.cuda.Event() + bufs.bulk_consume_done = torch.cuda.Event() + bufs.copy_pending = False + bufs.page_tables_active = False + return bufs def _stage_block_offsets( - ws: SimpleNamespace, + bufs: SimpleNamespace, manager: KVCacheManagerV2, request_ids: List[int], current_stream: torch.cuda.Stream, @@ -567,8 +565,8 @@ def _stage_block_offsets( identity indices. ``dst[pool, r, 0(K), :]`` holds ``base_page * index_scales``; score and compact decode that K plane inline. """ - if ws.copy_pending and not ws.copy_done.query(): - ws.copy_done.synchronize() + if bufs.copy_pending and not bufs.copy_done.query(): + bufs.copy_done.synchronize() # The native device copy reads only K and derives V with kv_offset. manager.index_mapper.gather_k_block_offsets( manager.host_kv_cache_block_offsets, @@ -576,21 +574,21 @@ def _stage_block_offsets( request_ids, copy_block_count, ) - manager._stream.wait_event(ws.copy_done) + manager._stream.wait_event(bufs.copy_done) copy_batch_block_offsets_to_device( source, destination, - ws._bulk_copy_idx_src[: len(request_ids)], + bufs._bulk_copy_idx_src[: len(request_ids)], manager.index_scales, manager.kv_offset, manager._stream.cuda_stream, ) - ws.bulk_copy_done.record(manager._stream) - current_stream.wait_event(ws.bulk_copy_done) + bufs.bulk_copy_done.record(manager._stream) + current_stream.wait_event(bufs.bulk_copy_done) def stage_eviction_cohort( - ws: SimpleNamespace, + bufs: SimpleNamespace, manager: KVCacheManagerV2, request_ids: List[int], round_starts: List[int], @@ -608,41 +606,41 @@ def stage_eviction_cohort( prompt lengths. """ request_count = len(request_ids) - stream = torch.cuda.current_stream(ws.device) + stream = torch.cuda.current_stream(bufs.device) # Reuse guard: staging over a cohort whose pages are still being read # would silently corrupt the in-flight compaction. - if ws.page_tables_active: + if bufs.page_tables_active: raise RuntimeError("previous page-table cohort is still active") if seq_lens is None: - seq_lens = [ws.bucket_seq_len] * request_count + seq_lens = [bufs.bucket_seq_len] * request_count request_metadata = torch.as_tensor((round_starts, seq_lens, token_starts), dtype=torch.int32) # Grow the phase table while this cohort's round starts are still host # integers: a stale-capacity gather is an out-of-bounds index_select on # the device. - grow_mean_phase_table(ws.phase, int(max(round_starts)) + 1) + grow_mean_phase_table(bufs.phase, int(max(round_starts)) + 1) _stage_block_offsets( - ws, + bufs, manager, request_ids, stream, - ws._bulk_offsets_src, - ws.block_offsets_device, - ws.copy_block_count, + bufs._bulk_offsets_src, + bufs.block_offsets_device, + bufs.copy_block_count, ) if draft_manager is not None: _stage_block_offsets( - ws, + bufs, draft_manager, request_ids, stream, - ws._draft_bulk_offsets_src, - ws.draft_block_offsets_device, - ws.draft_copy_block_count, + bufs._draft_bulk_offsets_src, + bufs.draft_block_offsets_device, + bufs.draft_copy_block_count, ) - ws.request_metadata_host[:3, :request_count].copy_(request_metadata) + bufs.request_metadata_host[:3, :request_count].copy_(request_metadata) # Rows past this cohort are padding: zero lengths keep the score kernel # and selection inert for them. - ws.request_metadata_host[:3, request_count:].zero_() + bufs.request_metadata_host[:3, request_count:].zero_() # This round's per-family move offsets ride the same table, so the single # device copy below carries them too. for row, family_offsets in ( @@ -651,43 +649,43 @@ def stage_eviction_cohort( (5, draft_move_offsets), ): if family_offsets is not None: - ws.request_metadata_host[row, : len(family_offsets)].copy_( + bufs.request_metadata_host[row, : len(family_offsets)].copy_( torch.as_tensor(family_offsets, dtype=torch.int32) ) try: # Copy the fixed backing once. Only the first ``request_count`` # columns are consumed by this cohort. - ws.request_metadata_device.copy_(ws.request_metadata_host, non_blocking=True) + bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) finally: # Guard the pinned metadata until its asynchronous copies complete. # Page-table device-buffer reuse is guarded separately after compact. - ws.copy_done.record(stream) - ws.copy_pending = True - ws.page_tables_active = True + bufs.copy_done.record(stream) + bufs.copy_pending = True + bufs.page_tables_active = True # The staged per-request prompt lengths are shared with the selection; # per-head modes re-expand them into their row-major view here. - if ws.row_prompt_offsets is not ws.prompt_offsets: - ws.row_prompt_offsets.view(ws.max_requests, ws.selection_rows_per_request).copy_( - ws.prompt_offsets.unsqueeze(1).expand(-1, ws.selection_rows_per_request) + if bufs.row_prompt_offsets is not bufs.prompt_offsets: + bufs.row_prompt_offsets.view(bufs.max_requests, bufs.selection_rows_per_request).copy_( + bufs.prompt_offsets.unsqueeze(1).expand(-1, bufs.selection_rows_per_request) ) -def mark_page_tables_consumed(ws: SimpleNamespace, *manager_streams: torch.cuda.Stream) -> None: +def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cuda.Stream) -> None: """Order V2 page-table reuse and resize after this cohort's compact. Every passed manager stream (target, and the draft when co-compressed) waits on one event recorded after the compact launches, so neither cache can free or reallocate pages this cohort is still reading. """ - if not ws.page_tables_active: + if not bufs.page_tables_active: raise RuntimeError("TriAttention page tables were not staged") - ws.bulk_consume_done.record(torch.cuda.current_stream(ws.device)) + bufs.bulk_consume_done.record(torch.cuda.current_stream(bufs.device)) for manager_stream in manager_streams: - manager_stream.wait_event(ws.bulk_consume_done) - ws.page_tables_active = False + manager_stream.wait_event(bufs.bulk_consume_done) + bufs.page_tables_active = False -def settle_top_tokens(ws: SimpleNamespace) -> None: +def settle_top_tokens(bufs: SimpleNamespace) -> None: """Pick the top-k with the CuTE selector, then settle its output. The CuTE top-k is fast but breaks score ties arbitrarily and emits @@ -695,120 +693,163 @@ def settle_top_tokens(ws: SimpleNamespace) -> None: membership with lowest-index-wins ties, rebases each row by its prompt offset, and writes sorted ordinals. The same launch packs each request's dense/SWA compaction move sources from the ordinals it just settled - (workspaces built without compaction compile the pack half away). + (buffers built without compaction compile the pack half away). """ # The trailing 1 is next_n: decode scores one query token per request. torch.ops.trtllm.cute_dsl_indexer_topk_decode( - ws.selection_scores_rows, - ws.selection_row_lengths, - ws.provisional_rows, - ws.keep_count, + bufs.selection_scores_rows, + bufs.selection_row_lengths, + bufs.provisional_rows, + bufs.keep_count, 1, ) - _settle_ties_and_pack_compaction_sources_kernel[ws.settle_grid]( - ws.selection_scores_rows, - ws.selection_row_lengths, - ws.row_prompt_offsets, - ws.provisional_rows, - ws.keep_rows, - *ws.settle_pack_tensors, - WIDTH=ws.decode_width, - KEEP_COUNT=ws.keep_count, - SELECTION_ROWS=ws.selection_rows_per_request, - **ws.settle_pack_shape, + _settle_ties_and_pack_compaction_sources_kernel[bufs.settle_grid]( + bufs.selection_scores_rows, + bufs.selection_row_lengths, + bufs.row_prompt_offsets, + bufs.provisional_rows, + bufs.keep_rows, + *bufs.settle_pack_tensors, + WIDTH=bufs.decode_width, + KEEP_COUNT=bufs.keep_count, + SELECTION_ROWS=bufs.selection_rows_per_request, + **bufs.settle_pack_shape, HAS_SETTLE=True, BLOCK=256, num_warps=4, ) -def run_eviction_round(ws: SimpleNamespace, normalize_scores: bool) -> None: +def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: """One staged eviction round, kernels fired directly in sequence. - Union: phase gather, fused score+stats+union (two CuTe launches), top-k, - settle-and-pack, C++ compacts. Per-head modes: phase gather, score-only - CuTe launch, decode-window gather, stats+reduce kernels, top-k, - settle-and-pack, C++ compacts. Every launch covers the full request - capacity; padded rows past the staged cohort carry zero lengths and stay - inert. + Union: phase gather (also derives valid widths), fused score+stats+union + (two CuTe launches), top-k, settle-and-pack, C++ compacts. Per-head + modes: phase gather, score-only CuTe launch, decode-window gather, + stats+reduce kernels, top-k, settle-and-pack, C++ compacts. A + co-compressed draft adds its own pack launch before its C++ moves. + Every launch covers the full request capacity; padded rows past the + staged cohort carry zero lengths and stay inert. """ - request_count = ws.max_requests - num_segments = request_count * ws.num_layers - union = ws.eviction_mode == "union" + request_count = bufs.max_requests + union = bufs.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): # mean_cos/mean_sin feed the compiled score launches, which captured # their device pointers: refresh them in place from this round's - # staged round starts. The same holds for the per-request decode - # widths, per-segment valid lengths, and staged window starts. + # staged round starts. The same launch derives the per-request valid + # decode widths; the compiled kernels read valid lengths and window + # starts straight from the staged metadata rows (pointer capture). gather_mean_phases( - ws.phase, ws.round_starts_device, ws.mean_cos, ws.mean_sin, request_count + bufs.phase, + bufs.round_starts_device, + bufs.mean_cos, + bufs.mean_sin, + bufs.valid_seq_lens_device, + bufs.token_starts_device, + bufs.valid_widths, + request_count, ) - torch.sub( - ws.valid_seq_lens_device[:request_count], - ws.token_starts_device[:request_count], - out=ws.valid_widths[:request_count], - ) - torch.index_select( - ws.valid_seq_lens_device, - 0, - ws.seg_req[:num_segments], - out=ws.seg_seq_len[:num_segments], - ) - ws.cute_token_starts[:request_count].copy_(ws.token_starts_device[:request_count]) if union: - ws.runner.launch_union_fusion( - request_count, ws.mean_cos, ws.mean_sin, ws.union_rows[:request_count] + bufs.runner.launch_union_fusion( + request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_rows[:request_count] ) - columns = min(ws.union_rows.shape[1], ws.combined.shape[1]) - ws.combined[:request_count, :columns].copy_(ws.union_rows[:request_count, :columns]) + columns = min(bufs.union_rows.shape[1], bufs.combined.shape[1]) + bufs.combined[:request_count, :columns].copy_(bufs.union_rows[:request_count, :columns]) else: - ws.runner.launch(request_count, ws.mean_cos, ws.mean_sin) + bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) # The kernel wrote each request's window scores (from its pinned # prompt length) into the head-major scratch, padded to the MMA # tile N=8 per KV head. Gather each request's decode window into # the [request, layer, head, token] layout the reduce kernels # read; columns past a request's valid width carry unscored # scratch data masked by ``valid_widths``. - group_size = ws.num_q_heads // ws.num_kv_heads + group_size = bufs.num_q_heads // bufs.num_kv_heads + num_segments = request_count * bufs.num_layers source = ( - ws.cute_scratch[: ws.num_kv_heads * 8 * num_segments * ws.bucket_seq_len] - .view(ws.num_kv_heads, 8, request_count, ws.num_layers, ws.bucket_seq_len)[ + bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] + .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ :, :group_size ] .permute(2, 3, 0, 1, 4) ) columns = ( - ws.token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) - + ws.gather_columns + bufs.token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + + bufs.gather_columns ) - columns = columns.clamp_(max=ws.bucket_seq_len - 1).expand( - request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width ) torch.gather( source, 4, columns, - out=ws.score_output[:request_count].view( - request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + out=bufs.score_output[:request_count].view( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width ), ) with nvtx_range("triattention.select", color="yellow"): if not union: prepare_per_head_scores( - ws.score_output[:request_count], - ws.valid_widths, - ws.row_mean, - ws.row_std, - ws.selection_scores, - ws.row_seq_lens, + bufs.score_output[:request_count], + bufs.valid_widths, + bufs.row_mean, + bufs.row_std, + bufs.selection_scores, + bufs.row_seq_lens, request_count, - num_kv_heads=ws.num_kv_heads, - per_layer=ws.eviction_mode == "per_layer_perhead", + num_kv_heads=bufs.num_kv_heads, + per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - settle_top_tokens(ws) + settle_top_tokens(bufs) with nvtx_range("triattention.compact", color="purple"): - run_cache_compactions(ws.compaction) + if bufs.swa_destination_bases is not None: + # The prompt offsets may have been re-staged since construction; + # rebase the SWA landing positions for this round. + torch.add(bufs.prompt_offsets, bufs.swa_rebase_delta, out=bufs.swa_destination_bases) + for family in bufs.compaction_families: + if family["name"] == "draft": + # One more pack launch broadcasts the target keep set over + # the draft KV heads and appends the draft's own tail + # ordinals. HAS_SETTLE=False compiles the settle half away + # (the ordinals arrive pre-settled), so any well-formed + # tensor stands in for the settle-side pointer arguments. + _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( + bufs.keep, + bufs.valid_seq_lens_device, + bufs.draft_pack["offsets"], + bufs.keep, + bufs.keep, + bufs.valid_seq_lens_device, + bufs.draft_pack["offsets"], + bufs.draft_pack["indices"], + bufs.draft_pack["offsets"], + bufs.draft_pack["indices"], + WIDTH=bufs.keep_count, + KEEP_COUNT=bufs.keep_count, + SELECTION_ROWS=1, + DENSE_TOTAL=bufs.draft_pack["dense_total"], + SWA_TOTAL=0, + MOVE_CAPACITY=bufs.draft_pack["move_capacity"], + NUM_KV_HEADS=bufs.draft_pack["num_kv_heads"], + SWA_WINDOW=0, + UNION=True, + PER_LAYER=False, + HAS_SWA=False, + HAS_SETTLE=False, + BLOCK=256, + num_warps=4, + ) + for group in family["groups"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers( + group["pools"], + group["pool_pointers"], + group["page_table"], + family["source"], + family["offsets"], + family["destination_bases"], + group["source_layer_indices"], + ) class TriAttention(BaseKVCacheCompressionManager): @@ -891,8 +932,8 @@ def __init__( # Geometric integration offsets (built lazily on first eviction so the # device matches the cache pool). self._offsets: Optional[torch.Tensor] = None - # Mean-phase table dict, shared by reference with every workspace so - # it persists across workspace rebuilds. + # Mean-phase table dict, shared by reference with every buffer + # namespace so it persists across buffer rebuilds. self._phase: Optional[Dict[str, object]] = None # Request presence records successful initialization. Each value is a @@ -903,10 +944,10 @@ def __init__( # batch as ``(batch, {request_id: growth})``; the final hook treats # those slots as an opaque suffix. self._prepared_generation_batch: Optional[Tuple[object, Dict[int, int]]] = None - # The eviction workspace is built once at the first eviction, sized to + # The eviction buffers are built once at the first eviction, sized to # capacity bounds, and reused for the manager's lifetime. - self._workspace: Optional[SimpleNamespace] = None - self._workspace_fingerprint: Optional[tuple] = None + self._buffers: Optional[SimpleNamespace] = None + self._buffers_fingerprint: Optional[tuple] = None self._local_to_global_layers_cache: Optional[List[int]] = None self._attention_layer_partition_cache: Optional[ Tuple[List[int], List[int], Optional[int]] @@ -1221,7 +1262,7 @@ def _periodic_evict( # Compact all affected dense and kernel-masked SWA layers, then release # the unreachable tail directly through V2's public resize primitive. # Prompt lengths and tails are per-request metadata, so the whole due - # cohort runs as one batched round (the workspace holds max_batch_size + # cohort runs as one batched round (the buffers hold max_batch_size # requests, which bounds any generation batch). if not prepared: return @@ -1327,7 +1368,7 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: prepared = self._prepared_generation_batch if prepared is not None: prepared[1].pop(request_id, None) - # The workspace stays resident across idle periods: its memory is a + # The buffers stay resident across idle periods: their memory is a # deliberate one-time cost and rebuilding it per burst would reintroduce # allocation on the decode hot path. @@ -1624,19 +1665,19 @@ def _pool_view_fingerprint(pools: List[torch.Tensor]) -> Tuple[tuple, ...]: for pool in pools ) - def _workspace_for( + def _buffers_for( self, layout: Dict[str, object], prepared: Sequence[Dict[str, object]], ) -> SimpleNamespace: - """Return the eviction workspace, building it once at first use. + """Return the eviction buffers, building them once at first use. The request capacity follows the executor's max batch size (memory scales linearly with it) and the decode-width capacity follows the eviction bound (compaction keeps the scored decode region near - ``budget`` plus one period of growth), so one workspace serves every - round. It is rebuilt only when the pool views change or a round - outgrows it. + ``budget`` plus one period of growth), so one set of buffers serves + every round. It is rebuilt only when the pool views change or a + round outgrows it. """ if not prepared: raise ValueError("TriAttention eviction requires at least one request") @@ -1657,17 +1698,17 @@ def _workspace_for( layout["pool_view_fingerprint"], draft_fingerprint, ) - ws = self._workspace - if ws is not None: + bufs = self._buffers + if bufs is not None: if ( - self._workspace_fingerprint == fingerprint - and needed_width <= ws.decode_width - and needed_page_tokens <= ws.page_table_token_capacity - and needed_requests <= ws.max_requests + self._buffers_fingerprint == fingerprint + and needed_width <= bufs.decode_width + and needed_page_tokens <= bufs.page_table_token_capacity + and needed_requests <= bufs.max_requests ): - return ws + return bufs # Pools changed or this round outgrew the buffers: rebuild. - self._workspace = None + self._buffers = None mgr = self.kv_cache_manager tail_capacity = self._configured_protected_tail_capacity() @@ -1680,7 +1721,7 @@ def _workspace_for( # pinning it to max_seq_len: with pinned prompts the post-compaction # length is bounded by prompt + budget + slack, so one power-of-two # bucket serves the steady state, and a cohort that outgrows it simply - # rebuilds the workspace through the capacity check above. A + # rebuilds the buffers through the capacity check above. A # max_seq_len floor would make the scratch unindexable in 32 bits and # tens of GiB at large batch for work that never scores past ~1K # tokens per request. @@ -1741,7 +1782,7 @@ def _workspace_for( q_real, q_imag, mlr_coef = self._local_score_calibration( layout["num_layers"], layout["global_layers"] ) - ws = prepare_eviction_workspace( + bufs = init_eviction_buffers( eviction_mode=self.eviction_mode, layer_pools=layout["layer_pools"], dense_groups=dense_groups, @@ -1770,9 +1811,9 @@ def _workspace_for( protected_tail_capacity=tail_capacity, **draft_kwargs, ) - self._workspace = ws - self._workspace_fingerprint = fingerprint - return ws + self._buffers = bufs + self._buffers_fingerprint = fingerprint + return bufs def _move_offsets_for( self, @@ -1853,7 +1894,7 @@ def _evict_requests( with nvtx_range_debug("triattention.resolve_layout", color="blue"): layout = self._runtime_kv_layout(num_layers) with nvtx_range_debug("triattention.staging_lookup", color="blue"): - ws = self._workspace_for(layout, prepared) + bufs = self._buffers_for(layout, prepared) if layout["swa_layers"] and layout["swa_window"]: # SWA landing positions are prompt-dependent; reject a request # whose retained span cannot cover the model window this round. @@ -1866,10 +1907,10 @@ def _evict_requests( ) with nvtx_range_debug("triattention.page_table_stage", color="orange"): dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( - layout, prepared, ws.max_requests + layout, prepared, bufs.max_requests ) stage_eviction_cohort( - ws, + bufs, self.kv_cache_manager, [item["request_id"] for item in prepared], [item["round_start"] for item in prepared], @@ -1882,12 +1923,12 @@ def _evict_requests( ) try: - run_eviction_round(ws, self.normalize_scores) + run_eviction_round(bufs, self.normalize_scores) finally: consumer_streams = [self.kv_cache_manager._stream] if self.draft_kv_cache_manager is not None: consumer_streams.append(self.draft_kv_cache_manager._stream) - mark_page_tables_consumed(ws, *consumer_streams) + mark_page_tables_consumed(bufs, *consumer_streams) capacity_targets = [] for item in prepared: 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 index 831ae677a4b3..dc1d6b8f40ba 100644 --- 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 @@ -236,7 +236,7 @@ def __call__( seg_page_off: cute.Tensor, seg_req_id: cute.Tensor, seg_layer_id: cute.Tensor, - seg_seq_len: cute.Tensor, + valid_seq_lens: cute.Tensor, seg_out_offset: cute.Tensor, token_starts: cute.Tensor, q_real: cute.Tensor, @@ -411,7 +411,7 @@ class SharedStorage: seg_page_off, seg_req_id, seg_layer_id, - seg_seq_len, + valid_seq_lens, seg_out_offset, token_starts, q_real, @@ -447,7 +447,7 @@ def kernel( seg_page_off: cute.Tensor, seg_req_id: cute.Tensor, seg_layer_id: cute.Tensor, - seg_seq_len: cute.Tensor, + valid_seq_lens: cute.Tensor, seg_out_offset: cute.Tensor, token_starts: cute.Tensor, q_real: cute.Tensor, @@ -475,7 +475,7 @@ def kernel( kv_head = task % self.num_kv_heads req_id = seg_req_id[segment] layer_id = seg_layer_id[segment] - valid_seq_len = seg_seq_len[segment] + valid_seq_len = valid_seq_lens[req_id] page_off = seg_page_off[segment] out_base = seg_out_offset[segment] # Per-request score window start (the request's pinned prompt @@ -1529,7 +1529,7 @@ def __init__( seg_page_off: torch.Tensor, seg_req_id: torch.Tensor, seg_layer_id: torch.Tensor, - seg_seq_len: torch.Tensor, + valid_seq_lens: torch.Tensor, seg_out_offset: torch.Tensor, token_starts: torch.Tensor, q_real: torch.Tensor, @@ -1569,7 +1569,7 @@ def __init__( seg_page_off, seg_req_id, seg_layer_id, - seg_seq_len, + valid_seq_lens, seg_out_offset, token_starts, q_real, @@ -1583,7 +1583,14 @@ def __init__( layer_pools[layer_indices[0]], self.descriptors, ) - self._cute_prefix = tuple(_to_cute(tensor) for tensor in self._torch_prefix) + # valid_seq_lens/token_starts are row views into the staged metadata + # table (byte offset 4*(max_requests+1)*row): only 4-byte aligned, + # and only ever read as per-CTA scalars. + prefix_aligns = (16, 16, 16, 16, 4, 16, 4, 16, 16, 16) + self._cute_prefix = tuple( + _to_cute(tensor, assumed_align=align) + for tensor, align in zip(self._torch_prefix, prefix_aligns) + ) self._cute_tail = ( _to_cute(freq_scale_sq), _to_cute(output), @@ -1603,9 +1610,9 @@ def __init__( ) self._cute_selection_prefix = ( _to_cute(output), - _to_cute(seg_seq_len), + _to_cute(valid_seq_lens, assumed_align=4), _to_cute(seg_out_offset), - _to_cute(token_starts), + _to_cute(token_starts, assumed_align=4), ) static_geometry = ( max_requests, 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 index bf6cff2ae77f..ea5aae2b1209 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -155,7 +155,7 @@ def __call__( self, partial_stats: cute.Tensor, scores: cute.Tensor, - seg_seq_len: cute.Tensor, + valid_seq_lens: cute.Tensor, seg_out_offset: cute.Tensor, token_starts: cute.Tensor, union_scores: cute.Tensor, @@ -165,7 +165,7 @@ def __call__( kernel = self.kernel( partial_stats, scores, - seg_seq_len, + valid_seq_lens, seg_out_offset, token_starts, union_scores, @@ -194,7 +194,7 @@ def kernel( self, partial_stats: cute.Tensor, scores: cute.Tensor, - seg_seq_len: cute.Tensor, + valid_seq_lens: cute.Tensor, seg_out_offset: cute.Tensor, token_starts: cute.Tensor, union_scores: cute.Tensor, @@ -218,7 +218,7 @@ def kernel( # Per-request score window start: the normalization domain and the # union output row both cover [0, valid - start) for this request. score_start = cutlass.Int32(token_starts[request_idx]) - valid_width = seg_seq_len[first_segment] - score_start + valid_width = valid_seq_lens[request_idx] - score_start warp_max_ptr = cute.arch.alloc_smem( cutlass.Float32, self.reduce_threads * self.tokens_per_lane * self.token_subtiles, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index ed67d5530272..e476242a5deb 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -9,7 +9,7 @@ workspace construction. The per-head modes use the pack's score-only entry plus the row-stats kernel; union eviction runs the fused score+stats+union pipeline. One-time buffer staging and runner compilation -live in ``triattention.prepare_eviction_workspace``; this module keeps the +live in ``triattention.init_eviction_buffers``; this module keeps the Triton kernels, their launch helpers, and the mean-phase table builders. House rules honored throughout: @@ -44,11 +44,18 @@ def _gather_mean_phase_kernel( table_sin, mean_cos, mean_sin, + valid_seq_lens, + token_starts, + valid_widths, table_rows, NUM_FREQS: tl.constexpr, F_BLOCK: tl.constexpr, ): - """Copy each request's precomputed phase-table row into the fixed buffers.""" + """Copy each request's precomputed phase-table row into the fixed buffers. + + The same launch derives the request's valid decode width (valid length + minus window start) so the round needs no separate subtraction launch. + """ request = tl.program_id(0) frequency = tl.arange(0, F_BLOCK) frequency_mask = frequency < NUM_FREQS @@ -62,6 +69,8 @@ def _gather_mean_phase_kernel( row_sin = tl.load(table_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) + width = tl.load(valid_seq_lens + request) - tl.load(token_starts + request) + tl.store(valid_widths + request, width) def build_mean_phase_table( @@ -121,9 +130,12 @@ def gather_mean_phases( round_starts: torch.Tensor, mean_cos: torch.Tensor, mean_sin: torch.Tensor, + valid_seq_lens: torch.Tensor, + token_starts: torch.Tensor, + valid_widths: torch.Tensor, request_count: int, ) -> None: - """Refresh the fixed mean buffers in place from staged round starts. + """Refresh the fixed mean buffers and valid widths from staged metadata. Writes in place because the compiled CuTe score launch captured the destination buffers' device pointers. CUDA-only; eviction never runs @@ -136,6 +148,9 @@ def gather_mean_phases( phase["sin"], mean_cos, mean_sin, + valid_seq_lens, + token_starts, + valid_widths, phase["rows"], NUM_FREQS=num_freqs, F_BLOCK=triton.next_power_of_2(num_freqs), diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index af512e8446ad..32ae62a432ac 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -136,10 +136,8 @@ def make_ramp_pools( def build_compaction(**overrides): - """``build_cache_compactions`` with the suite's default 2-layer geometry.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - build_cache_compactions, - ) + """``init_compaction_buffers`` with the suite's default 2-layer geometry.""" + from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers args = dict( eviction_mode="union", @@ -153,16 +151,96 @@ def build_compaction(**overrides): swa_window=None, ) args.update(overrides) - compaction = build_cache_compactions(**args) - # Production fuses the dense pack into the settle launch; standalone - # compaction tests re-attach it so run_cache_compactions packs the dense - # move buffers before the C++ moves (firing the pack twice is idempotent). - compaction_family(compaction, "dense")["pack"] = compaction["dense_pack"] - return compaction + return init_compaction_buffers(**args) + + +def launch_family_pack(compaction, name): + """Standalone HAS_SETTLE=False pack launch for one family of a bundle. + + Production packs dense/SWA inside the fused settle launch and fires the + draft pack inline in ``run_eviction_round``; standalone compaction tests + pack here from the bundle's own geometry so the C++ moves read + initialized indices. The settle-side pointer arguments are compiled away + (any well-formed tensor stands in). + """ + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _settle_ties_and_pack_compaction_sources_kernel, + ) + + family = compaction_family(compaction, name) + kept = compaction["kept_token_ordinals"] + valid = compaction["valid_sequence_lengths"] + keep_count = compaction["decode_keep_count"] + if name == "draft": + draft = compaction["draft_pack"] + rows = 1 + shape = dict( + DENSE_TOTAL=draft["dense_total"], + SWA_TOTAL=0, + MOVE_CAPACITY=draft["move_capacity"], + NUM_KV_HEADS=draft["num_kv_heads"], + SWA_WINDOW=0, + UNION=True, + PER_LAYER=False, + HAS_SWA=False, + ) + swa_offsets, swa_indices = family["offsets"], family["source"] + else: + rows = compaction["selection_rows"] + shape = compaction["settle_pack_shape"] + swa_family = compaction_family(compaction, "swa") + swa_offsets = swa_family["offsets"] if swa_family else family["offsets"] + swa_indices = swa_family["source"] if swa_family else family["source"] + _settle_ties_and_pack_compaction_sources_kernel[(compaction["request_count"], rows)]( + kept, + valid, + family["offsets"], + kept, + kept, + valid, + family["offsets"], + family["source"], + swa_offsets, + swa_indices, + WIDTH=keep_count, + KEEP_COUNT=keep_count, + SELECTION_ROWS=rows, + **shape, + HAS_SETTLE=False, + BLOCK=256, + num_warps=4, + ) + + +def run_compaction(compaction, pack=("dense", "draft")): + """Test-side replica of ``run_eviction_round``'s compact stage. + + Optional standalone family packs, the SWA destination rebase, then every + family's C++ moves -- the same sequence production fires inline. + """ + if compaction["swa_destination_bases"] is not None: + torch.add( + compaction["prompt_offsets"], + compaction["swa_rebase_delta"], + out=compaction["swa_destination_bases"], + ) + for family in compaction["families"]: + if family["name"] in pack and family["name"] != "swa": + launch_family_pack(compaction, family["name"]) + for group in family["groups"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers( + group["pools"], + group["pool_pointers"], + group["page_table"], + family["source"], + family["offsets"], + family["destination_bases"], + group["source_layer_indices"], + ) def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None): - """A bare workspace namespace for the bulk page-table copy tests.""" + """A bare buffer namespace for the bulk page-table copy tests.""" staging = SimpleNamespace() staging.device = device staging.max_requests = max_requests @@ -198,8 +276,8 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): ) -def make_workspace_stubs(manager, *, decode_width=260): - """Stub the calibration/layout surfaces around ``_workspace_for``.""" +def make_buffer_stubs(manager, *, decode_width=260): + """Stub the calibration/layout surfaces around ``_buffers_for``.""" manager._H = 2 manager._F = 2 manager._freq_scale_sq = torch.ones(2) @@ -222,14 +300,14 @@ def make_workspace_stubs(manager, *, decode_width=260): layer_pool_keys=(("pool", 0), ("pool", 0)), pool_view_fingerprint=(("fixed",),), ) - workspace = SimpleNamespace( + buffers = SimpleNamespace( decode_width=decode_width, page_table_token_capacity=65537, max_requests=8, token_starts_device=torch.zeros(8, dtype=torch.int32), valid_widths=torch.empty(8, dtype=torch.int32), ) - return layout, workspace + return layout, buffers def make_fake_v2(enable_block_reuse=False, *, is_draft=False): @@ -297,17 +375,17 @@ def mocked_eviction_internals(manager): """Run the real ``_evict_requests`` body around mocked GPU launches.""" from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - workspace = SimpleNamespace(max_requests=8) + buffers = SimpleNamespace(max_requests=8) layout = dict(swa_layers=[], swa_window=None) with ( mock.patch.object(manager, "_runtime_kv_layout", return_value=layout), - mock.patch.object(manager, "_workspace_for", return_value=workspace), + mock.patch.object(manager, "_buffers_for", return_value=buffers), mock.patch.object(module, "stage_eviction_cohort") as stage, mock.patch.object(module, "run_eviction_round") as run_round, mock.patch.object(module, "mark_page_tables_consumed") as consumed, ): yield SimpleNamespace( - workspace=workspace, + buffers=buffers, stage=stage, run_round=run_round, consumed=consumed, 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 index 9e3358d8baf9..2444f16fbd86 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -2,18 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 """The SM100 TriAttention CuTe scorer (the only score path) vs PyTorch oracles. -Two layers of coverage over the production score leg (workspace metadata +Two layers of coverage over the production score leg (buffer metadata staging, the compiled runner launch, and the decode-window gather -- the same sequence ``run_eviction_round`` fires). The kernel-numerics matrix drives a -single-layer workspace across the supported page geometries (permuted +single-layer buffers across the supported page geometries (permuted physical pages, ragged valid lengths, GQA group 4 riding the padded MMA tile) against inline oracle math. The launch-path matrix drives multi-layer -workspaces across the named production geometries (Qwen3, GPT-OSS, the +buffers across the named production geometries (Qwen3, GPT-OSS, the originally validated 128-token-page shape) against the shared pure-PyTorch -oracle, sweeps request counts up to the workspace capacity, and checks the +oracle, sweeps request counts up to the buffer capacity, and checks the per-request decode-width metadata the selection reduce kernels consume. The contract test pins the loud-failure behavior: unsupported geometry raises -from the CuTe runner's own validation at workspace construction -- there is +from the CuTe runner's own validation at buffer construction -- there is no fallback score kernel. """ @@ -28,7 +28,7 @@ ) -def _make_score_workspace( +def _make_score_buffers( *, layer_pools, max_requests, @@ -43,13 +43,13 @@ def _make_score_workspace( decode_width=None, eviction_mode="per_head", ): - """A score-only workspace over one shared page-table slot.""" + """Score-only buffers over one shared page-table slot.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - prepare_eviction_workspace, + init_eviction_buffers, ) num_layers = len(layer_pools) - return prepare_eviction_workspace( + return init_eviction_buffers( eviction_mode=eviction_mode, layer_pools=layer_pools, dense_groups=[list(range(num_layers))], @@ -72,51 +72,58 @@ def _make_score_workspace( ) -def _write_block_offsets(ws, encoded): - """Load a test page table into the workspace's staged block-offset plane.""" - ws.block_offsets_device.zero_() - ws.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) +def _write_block_offsets(bufs, encoded): + """Load a test page table into the staged block-offset plane.""" + bufs.block_offsets_device.zero_() + bufs.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) def _launch_split_scores( - ws, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin + bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin ): """The production score-only leg: stage metadata, fire the compiled runner, gather each request's decode window (``run_eviction_round``'s per-head sequence, parameterized by the request count).""" - num_segments = request_count * ws.num_layers + num_segments = request_count * bufs.num_layers torch.sub( valid_seq_lens[:request_count], token_starts[:request_count], out=valid_widths[:request_count], ) - torch.index_select( - valid_seq_lens, 0, ws.seg_req[:num_segments], out=ws.seg_seq_len[:num_segments] - ) - ws.cute_token_starts[:request_count].copy_(token_starts[:request_count]) - assert request_count in ws.runner._compiled - ws.runner.launch(request_count, mean_cos, mean_sin) - group_size = ws.num_q_heads // ws.num_kv_heads + # The compiled runner reads valid lengths and window starts straight from + # the staged metadata rows (pointer capture), so stage them exactly like + # ``stage_eviction_cohort`` does. + bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) + bufs.token_starts_device[:request_count].copy_(token_starts[:request_count]) + assert request_count in bufs.runner._compiled + bufs.runner.launch(request_count, mean_cos, mean_sin) + group_size = bufs.num_q_heads // bufs.num_kv_heads source = ( - ws.cute_scratch[: ws.num_kv_heads * 8 * num_segments * ws.bucket_seq_len] - .view(ws.num_kv_heads, 8, request_count, ws.num_layers, ws.bucket_seq_len)[:, :group_size] + bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] + .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ + :, :group_size + ] .permute(2, 3, 0, 1, 4) ) - columns = token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + ws.gather_columns - columns = columns.clamp_(max=ws.bucket_seq_len - 1).expand( - request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + columns = ( + token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + bufs.gather_columns + ) + columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width ) output = torch.full( - (request_count, ws.num_layers, ws.num_q_heads, ws.decode_width), + (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), float("nan"), dtype=torch.float32, - device=ws.device, + device=bufs.device, ) torch.gather( source, 4, columns, - out=output.view(request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width), + out=output.view( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + ), ) return output @@ -161,7 +168,7 @@ def _build_case( 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 - ws = _make_score_workspace( + bufs = _make_score_buffers( layer_pools=pools, max_requests=max_requests, seq_len=capacity, @@ -174,7 +181,7 @@ def _build_case( offsets=offsets_t, decode_width=capacity - prompt_len, ) - _write_block_offsets(ws, _encode_block_offsets(page_ids)) + _write_block_offsets(bufs, _encode_block_offsets(page_ids)) round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) # Ragged valid lengths: shallow tails land mid-page and mid-compute-tile; @@ -198,7 +205,7 @@ def _build_case( offsets=offsets_t, ) return ( - ws, + bufs, pools, token_starts, valid_seq_lens, @@ -241,7 +248,7 @@ def test_cute_kernel_matches_torch_oracle(case): max_requests = case["max_requests"] num_layers = case["num_layers"] ( - ws, + bufs, pools, token_starts, valid_seq_lens, @@ -250,7 +257,7 @@ def test_cute_kernel_matches_torch_oracle(case): mean_sin, oracle_inputs, ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) - device = ws.device + device = bufs.device oracle = _torch_tri_score_oracle( pools, @@ -266,14 +273,14 @@ def test_cute_kernel_matches_torch_oracle(case): list(range(num_layers)), ) - # The compiled runner serves every request count up to the workspace + # The compiled runner serves every request count up to the buffer # capacity and nothing beyond it; cover one, an intermediate count, and # the capacity. - assert max_requests + 1 not in ws.runner._compiled + assert max_requests + 1 not in bufs.runner._compiled for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) scores = _launch_split_scores( - ws, + bufs, request_count, valid_seq_lens, valid_widths, @@ -285,7 +292,7 @@ def test_cute_kernel_matches_torch_oracle(case): request_count, num_layers, case["num_q_heads"], - ws.decode_width, + bufs.decode_width, ) # The score leg owns the per-request decode widths the selection # reduce kernels consume. @@ -304,9 +311,9 @@ def test_cute_kernel_matches_torch_oracle(case): # The loud-failure contract: unsupported geometry raises from the CuTe -# runner's own validation during the eager compile at workspace +# runner's own validation during the eager compile at buffer # construction, surfaced as the no-fallback RuntimeError. -def test_unsupported_geometry_raises_at_workspace_construction(): +def test_unsupported_geometry_raises_at_buffer_construction(): pytest.importorskip("cutlass") device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(20260722) @@ -320,7 +327,7 @@ def test_unsupported_geometry_raises_at_workspace_construction(): ] calib = torch.randn(num_layers, 2, num_freqs, device=device) with pytest.raises(RuntimeError, match="no other score path exists"): - _make_score_workspace( + _make_score_buffers( layer_pools=pools, max_requests=max_requests, seq_len=page_count * tokens_per_block, 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 index ec992beb2464..03da8a048684 100644 --- 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 @@ -3,7 +3,7 @@ """Equivalence coverage for the fused score+stats+union pipeline (two CuTe kernels). The reference side gathers the SAME production score rows (the fused pack's -score-only entry, which every workspace runner compiles) and normalizes + +score-only entry, which every buffer-namespace runner compiles) and normalizes + union-reduces them with a pure-torch float32 oracle. The fused-vs-reference comparison was always tolerance-based (the fused pipeline's reduction order differs from any reference); the tolerances are unchanged from the retired @@ -19,7 +19,7 @@ ) -def _make_union_workspace( +def _make_union_buffers( *, layer_pools, max_requests, @@ -33,17 +33,17 @@ def _make_union_workspace( offsets, decode_width=None, ): - """A union-mode workspace over one shared page-table slot. + """Union-mode buffers over one shared page-table slot. - The union runner also compiles the score-only entries, so one workspace - serves both the fused pipeline and the split reference leg. + The union runner also compiles the score-only entries, so one buffer + namespace serves both the fused pipeline and the split reference leg. """ from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - prepare_eviction_workspace, + init_eviction_buffers, ) num_layers = len(layer_pools) - return prepare_eviction_workspace( + return init_eviction_buffers( eviction_mode="union", layer_pools=layer_pools, dense_groups=[list(range(num_layers))], @@ -66,69 +66,78 @@ def _make_union_workspace( ) -def _write_block_offsets(ws, encoded): - """Load a test page table into the workspace's staged block-offset plane.""" - ws.block_offsets_device.zero_() - ws.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) +def _write_block_offsets(bufs, encoded): + """Load a test page table into the staged block-offset plane.""" + bufs.block_offsets_device.zero_() + bufs.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) -def _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts): - """Stage the per-round score metadata exactly like ``run_eviction_round``.""" - num_segments = request_count * ws.num_layers +def _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts): + """Stage the per-round score metadata exactly like production. + + The compiled runner reads valid lengths and window starts straight from + the staged metadata rows (pointer capture), so stage them like + ``stage_eviction_cohort`` does; the width subtraction mirrors what the + production phase-gather launch derives on device. + """ torch.sub( valid_seq_lens[:request_count], token_starts[:request_count], out=valid_widths[:request_count], ) - torch.index_select( - valid_seq_lens, 0, ws.seg_req[:num_segments], out=ws.seg_seq_len[:num_segments] - ) - ws.cute_token_starts[:request_count].copy_(token_starts[:request_count]) + bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) + bufs.token_starts_device[:request_count].copy_(token_starts[:request_count]) def _launch_split_scores( - ws, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin + bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin ): """The production score-only leg plus the decode-window gather.""" - _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts) - assert request_count in ws.runner._compiled - ws.runner.launch(request_count, mean_cos, mean_sin) - num_segments = request_count * ws.num_layers - group_size = ws.num_q_heads // ws.num_kv_heads + _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) + assert request_count in bufs.runner._compiled + bufs.runner.launch(request_count, mean_cos, mean_sin) + num_segments = request_count * bufs.num_layers + group_size = bufs.num_q_heads // bufs.num_kv_heads source = ( - ws.cute_scratch[: ws.num_kv_heads * 8 * num_segments * ws.bucket_seq_len] - .view(ws.num_kv_heads, 8, request_count, ws.num_layers, ws.bucket_seq_len)[:, :group_size] + bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] + .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ + :, :group_size + ] .permute(2, 3, 0, 1, 4) ) - columns = token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + ws.gather_columns - columns = columns.clamp_(max=ws.bucket_seq_len - 1).expand( - request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width + columns = ( + token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + bufs.gather_columns + ) + columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width ) output = torch.full( - (request_count, ws.num_layers, ws.num_q_heads, ws.decode_width), + (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), float("nan"), dtype=torch.float32, - device=ws.device, + device=bufs.device, ) torch.gather( source, 4, columns, - out=output.view(request_count, ws.num_layers, ws.num_kv_heads, group_size, ws.decode_width), + out=output.view( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + ), ) return output def _launch_union_fusion( - ws, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out + bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out ): """The fused score+stats+normalized-union pipeline (THE union path).""" - _stage_score_metadata(ws, request_count, valid_seq_lens, valid_widths, token_starts) + _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) assert ( - request_count in ws.runner._compiled_stats - and request_count in ws.runner._compiled_normalize_union + request_count in bufs.runner._compiled_stats + and request_count in bufs.runner._compiled_normalize_union ) - ws.runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_out[:request_count]) + bufs.runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_out[:request_count]) def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: @@ -193,7 +202,7 @@ def _check_union_fusion_matches_split_pipeline( mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() - ws = _make_union_workspace( + bufs = _make_union_buffers( layer_pools=[pool], max_requests=2, seq_len=seq_len, @@ -208,7 +217,7 @@ def _check_union_fusion_matches_split_pipeline( k_plane = [2 * page for page in page_permutation] v_plane = [2 * page + 1 for page in page_permutation] _write_block_offsets( - ws, + bufs, torch.tensor([[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device), ) if valid_lens is None: @@ -224,7 +233,7 @@ def _check_union_fusion_matches_split_pipeline( split_widths = torch.empty(request_count, dtype=torch.int32, device=device) token_starts = torch.tensor(score_starts, dtype=torch.int32, device=device) per_head = _launch_split_scores( - ws, + bufs, request_count, valid_seq_lens, split_widths, @@ -241,7 +250,7 @@ def _check_union_fusion_matches_split_pipeline( (request_count, seq_len), float("nan"), dtype=torch.float32, device=device ) _launch_union_fusion( - ws, + bufs, request_count, valid_seq_lens, fused_widths, @@ -393,7 +402,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() - ws = _make_union_workspace( + bufs = _make_union_buffers( layer_pools=layer_pools, max_requests=max_requests, seq_len=seq_len, @@ -406,11 +415,11 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: offsets=offsets, decode_width=decode_window, ) - assert (ws.cute_scratch.numel() > 2**31) == (max_requests == 64) + assert (bufs.cute_scratch.numel() > 2**31) == (max_requests == 64) page_ids = torch.arange(num_pages, dtype=torch.int32, device=device) - ws.block_offsets_device.zero_() - ws.block_offsets_device[0, :request_count, 0, :num_pages] = 2 * page_ids - ws.block_offsets_device[0, :request_count, 1, :num_pages] = 2 * page_ids + 1 + bufs.block_offsets_device.zero_() + bufs.block_offsets_device[0, :request_count, 0, :num_pages] = 2 * page_ids + bufs.block_offsets_device[0, :request_count, 1, :num_pages] = 2 * page_ids + 1 valid_seq_lens = torch.zeros(max_requests, dtype=torch.int32, device=device) token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) @@ -421,7 +430,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: # the pure-torch union oracle. split_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) per_head = _launch_split_scores( - ws, + bufs, request_count, valid_seq_lens, split_widths, @@ -439,7 +448,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: (max_requests, seq_len), float("nan"), dtype=torch.float32, device=device ) _launch_union_fusion( - ws, + bufs, max_requests, valid_seq_lens, fused_widths, 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 index 4a28f77a9cd5..c91cadd20ba1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -10,7 +10,7 @@ ``destination_base = prompt_len``. These tests cover the physical draft moves, the packed move indices, stream ordering across both cache managers, the speculative admission gates (one representative per guard family), the -published compressed-token invariant, and workspace rebuild/invalidation. +published compressed-token invariant, and buffer rebuild/invalidation. """ from types import SimpleNamespace @@ -21,15 +21,15 @@ from conftest import build_compaction as _build_compaction from conftest import compaction_family as _compaction_family from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_buffer_stubs as _make_buffer_stubs from conftest import make_fake_v2 as _make_fake_v2 from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_request as _make_request from conftest import make_triattention as _make_triattention -from conftest import make_workspace_stubs as _make_workspace_stubs from conftest import mocked_eviction_internals as _mocked_eviction_internals +from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import run_cache_compactions from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( TriAttention, mark_page_tables_consumed, @@ -90,7 +90,7 @@ def _launched_draft_compaction(draft_protected_tails): draft_page_table_slots={0: 0}, ) _set_protected_tails(compaction, target_protected_tails, draft_protected_tails) - run_cache_compactions(compaction) + _run_compaction(compaction) torch.cuda.synchronize(device) return SimpleNamespace( @@ -193,7 +193,7 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect def test_mark_page_tables_consumed_orders_both_manager_streams(): event = mock.Mock() - workspace = SimpleNamespace( + buffers = SimpleNamespace( device=torch.device("cuda", torch.cuda.current_device()), page_tables_active=True, bulk_consume_done=event, @@ -203,17 +203,17 @@ def test_mark_page_tables_consumed_orders_both_manager_streams(): compute_stream = SimpleNamespace() with mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream): - mark_page_tables_consumed(workspace, target_stream, draft_stream) + mark_page_tables_consumed(buffers, target_stream, draft_stream) # One event records the compact launches; BOTH cache managers wait on it, # so neither can free or reallocate pages this cohort is still reading. event.record.assert_called_once_with(compute_stream) target_stream.wait_event.assert_called_once_with(event) draft_stream.wait_event.assert_called_once_with(event) - assert workspace.page_tables_active is False + assert buffers.page_tables_active is False with pytest.raises(RuntimeError, match="not staged"): - mark_page_tables_consumed(workspace, target_stream, draft_stream) + mark_page_tables_consumed(buffers, target_stream, draft_stream) @pytest.mark.parametrize( @@ -314,7 +314,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert confirmed == 2 + 4 # The staged logical position restores the uncompressed # length: physical confirmed plus everything evicted so far - # (stage_eviction_cohort args: ws, manager, ids, round_starts, + # (stage_eviction_cohort args: bufs, manager, ids, round_starts, # prompt lengths, seq lens, page-table lens). round_starts = internals.stage.call_args.args[3] assert round_starts[0] == uncompressed @@ -330,7 +330,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): # streams are ordered after the compact launches. assert draft_cache.resize.call_args_list == [mock.call(7, None)] * eviction_rounds assert internals.consumed.call_args_list == ( - [mock.call(internals.workspace, target._stream, draft_manager._stream)] * eviction_rounds + [mock.call(internals.buffers, target._stream, draft_manager._stream)] * eviction_rounds ) @@ -338,7 +338,7 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention(budget=4) - layout, workspace = _make_workspace_stubs(manager) + layout, buffers = _make_buffer_stubs(manager) # The one-time host block-offset table shape gate reads the real manager # tables (int32 [pools, slots, K/V, blocks]). manager.kv_cache_manager.host_kv_cache_block_offsets = torch.zeros( @@ -373,42 +373,43 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): with mock.patch.object( module, - "prepare_eviction_workspace", - return_value=workspace, + "init_eviction_buffers", + return_value=buffers, ) as prepare: - resources = manager._workspace_for(layout, prepared) + resources = manager._buffers_for(layout, prepared) # Request capacity follows the executor limits, while the score # bucket follows what the cohort actually presents (power-of-two, # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. - assert resources is workspace + assert resources is buffers assert prepare.call_args.kwargs["eviction_mode"] == "union" assert prepare.call_args.kwargs["max_requests"] == 8 assert prepare.call_args.kwargs["decode_width"] == 4 + 2 * 128 assert prepare.call_args.kwargs["seq_len"] == 1024 assert prepare.call_args.kwargs["page_table_token_capacity"] == 1024 + 1 assert prepare.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 - # Migrated from the pipeline workspace-kwargs test: the budget, the + # Migrated from the pipeline buffer-kwargs test: the budget, the # shared phase-table dict, and the pool keys thread through unchanged. assert prepare.call_args.kwargs["keep_count"] == manager.budget assert prepare.call_args.kwargs["phase"] is manager._phase assert prepare.call_args.kwargs["layer_pool_keys"] == list(layout["layer_pool_keys"]) - # A second round with unchanged pools reuses the resident workspace - # (and with it the compaction launch data it carries). - assert manager._workspace_for(layout, prepared) is resources + # A second round with unchanged pools reuses the resident buffers + # (and with them the compaction launch data they carry). + assert manager._buffers_for(layout, prepared) is resources assert prepare.call_count == 1 - # A pool change invalidates the whole workspace, compaction included. + # A pool change invalidates the whole buffer namespace, compaction + # included. layout["pool_view_fingerprint"] = (("moved",),) - rebuilt_workspace = SimpleNamespace( - decode_width=workspace.decode_width, - page_table_token_capacity=workspace.page_table_token_capacity, - max_requests=workspace.max_requests, + rebuilt_buffers = SimpleNamespace( + decode_width=buffers.decode_width, + page_table_token_capacity=buffers.page_table_token_capacity, + max_requests=buffers.max_requests, ) - prepare.return_value = rebuilt_workspace - rebuilt = manager._workspace_for(layout, prepared) + prepare.return_value = rebuilt_buffers + rebuilt = manager._buffers_for(layout, prepared) assert rebuilt is not resources - assert rebuilt is rebuilt_workspace + assert rebuilt is rebuilt_buffers assert prepare.call_count == 2 - assert manager._workspace is rebuilt_workspace + assert manager._buffers is rebuilt_buffers diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 574430bea322..54c4bc7dd666 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -33,11 +33,11 @@ import torch from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_bare_staging as _make_bare_staging +from conftest import make_buffer_stubs as _make_buffer_stubs 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_triattention as _make_triattention -from conftest import make_workspace_stubs as _make_workspace_stubs from conftest import mocked_eviction_internals as _mocked_eviction_internals from conftest import torch_tri_score_oracle as _torch_tri_score_oracle from pydantic import ValidationError @@ -432,17 +432,17 @@ def test_request_finish_clears_state_but_keeps_buffers_resident(self): evicted_tokens=127, confirmed_kv_length=128, ) - workspace = object() - manager._workspace = workspace + buffers = object() + manager._buffers = buffers manager._prepared_generation_batch = (SimpleNamespace(), {7: 1}) manager.on_request_finish(_make_request(7)) assert manager._request_states == {} assert manager._prepared_generation_batch[1] == {} - # The workspace is sized for the executor limits, not one cohort, so - # it stays resident for the next generation batch. - assert manager._workspace is workspace + # The buffers are sized for the executor limits, not one cohort, so + # they stay resident for the next generation batch. + assert manager._buffers is buffers @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): @@ -597,16 +597,16 @@ def test_union_rejects_unnormalized_scores(self): with pytest.raises(ValueError, match="normalize_scores=True"): _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) - def test_workspace_build_receives_mode_and_capacity_kwargs(self): + def test_buffer_build_receives_mode_and_capacity_kwargs(self): # Mode/phase/pool-key threading and cached reuse are covered by the - # workspace-rebuild superset test in + # buffer-rebuild superset test in # test_triattention_draft_cocompaction.py. from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention(budget=4) - # The workspace follows the executor limits: eight requests (max batch + # The buffers follow the executor limits: eight requests (max batch # size) by 260 decode tokens (budget plus two eviction periods). - layout, workspace = _make_workspace_stubs(manager) + layout, buffers = _make_buffer_stubs(manager) prepared = [ _prepared_eviction( _make_request(7), @@ -618,13 +618,13 @@ def test_workspace_build_receives_mode_and_capacity_kwargs(self): with mock.patch.object( module, - "prepare_eviction_workspace", - return_value=workspace, - ) as build_workspace: - resources = manager._workspace_for(layout, prepared) + "init_eviction_buffers", + return_value=buffers, + ) as build_buffers: + resources = manager._buffers_for(layout, prepared) - assert resources is workspace - kwargs = build_workspace.call_args.kwargs + assert resources is buffers + kwargs = build_buffers.call_args.kwargs assert kwargs["eviction_mode"] == "union" assert kwargs["max_requests"] == 8 assert kwargs["decode_width"] == 260 @@ -790,7 +790,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): # per-request moves are keep + tail = [6, 7]; padded rows repeat the # final offset out to the request capacity. args = internals.stage.call_args - assert args.args[0] is internals.workspace + assert args.args[0] is internals.buffers assert args.args[1] is manager.kv_cache_manager assert args.args[2:6] == ([7, 8], [8, 10], [3, 5], [8, 10]) assert args.kwargs["draft_manager"] is None @@ -798,7 +798,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): assert args.kwargs["swa_move_offsets"] is None assert args.kwargs["draft_move_offsets"] is None internals.consumed.assert_called_once_with( - internals.workspace, manager.kv_cache_manager._stream + internals.buffers, manager.kv_cache_manager._stream ) @requires_sm100 @@ -861,7 +861,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) # One group per layer: slot i holds layer i's own block table. - ws = module.prepare_eviction_workspace( + bufs = module.init_eviction_buffers( eviction_mode="per_head", layer_pools=pools, dense_groups=[[layer] for layer in layer_order], @@ -887,18 +887,22 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun def stage_round(): # Stage the round metadata straight into the fixed device rows # (the cohort staging path is covered by the staging test above). - ws.round_starts_device.copy_(round_device) - ws.valid_seq_lens_device[:request_count].copy_(valid_seq_lens) - ws.token_starts_device.fill_(prompt_len) - ws.block_offsets_device[:, :, :, :page_count].copy_(_encode_block_offsets(page_ids_3d)) + bufs.round_starts_device.copy_(round_device) + bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens) + bufs.token_starts_device.fill_(prompt_len) + bufs.block_offsets_device[:, :, :, :page_count].copy_( + _encode_block_offsets(page_ids_3d) + ) score_sentinel = -12345.0 + # Isolate the score stage: an empty family list makes the compact + # stage a no-op (this synthetic round never stages move offsets). + bufs.compaction_families = [] stage_round() - ws.score_output.fill_(score_sentinel) - with mock.patch.object(module, "run_cache_compactions"): - module.run_eviction_round(ws, normalize_scores=False) - fixed = ws.score_output.clone() - assert ws.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + bufs.score_output.fill_(score_sentinel) + module.run_eviction_round(bufs, normalize_scores=False) + fixed = bufs.score_output.clone() + assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] # The deployed fused score must agree with the independent Torch oracle # when every layer owns a distinct V2 block table. @@ -935,12 +939,11 @@ def stage_round(): ) expected_second_widths = valid_seq_lens - prompt_len stage_round() - ws.score_output.fill_(score_sentinel) - ws.valid_widths.fill_(-1) - with mock.patch.object(module, "run_cache_compactions"): - module.run_eviction_round(ws, normalize_scores=False) - second_launch = ws.score_output.clone() - assert torch.equal(ws.valid_widths, expected_second_widths) + bufs.score_output.fill_(score_sentinel) + bufs.valid_widths.fill_(-1) + module.run_eviction_round(bufs, normalize_scores=False) + second_launch = bufs.score_output.clone() + assert torch.equal(bufs.valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index c96e68bf83b7..a028aa68785f 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -9,12 +9,9 @@ from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_ramp_pools as _make_ramp_pools +from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.compaction import ( - launch_move_pack, - run_cache_compactions, -) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( prepare_per_head_scores, @@ -36,7 +33,7 @@ def _require_cute_topk_op() -> None: ) -def _make_selection_ws( +def _make_selection_buffers( *, eviction_mode, width, @@ -47,10 +44,10 @@ def _make_selection_ws( num_query_heads=1, num_kv_heads=1, ): - """Selection-only workspace: exactly the buffers the one prepare allocates + """Selection-only buffers: exactly what the one constructor allocates for the mode, without the CuTe score state or compaction (the settle's pack half is compiled away).""" - ws = SimpleNamespace( + bufs = SimpleNamespace( eviction_mode=eviction_mode, device=device, max_requests=max_requests, @@ -61,54 +58,64 @@ def _make_selection_ws( num_kv_heads=num_kv_heads, stream=None, ) - ws.valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) - ws.prompt_offsets = torch.zeros(max_requests, dtype=torch.int32, device=device) + bufs.valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) + bufs.prompt_offsets = torch.zeros(max_requests, dtype=torch.int32, device=device) if eviction_mode == "union": - ws.selection_rows_per_request = 1 - ws.row_prompt_offsets = ws.prompt_offsets - ws.combined = torch.empty((max_requests, width), dtype=torch.float32, device=device) + bufs.selection_rows_per_request = 1 + bufs.row_prompt_offsets = bufs.prompt_offsets + bufs.combined = torch.empty((max_requests, width), dtype=torch.float32, device=device) # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. - ws.final_indices = torch.zeros((max_requests, keep_count), dtype=torch.int32, device=device) - ws.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) - ws.selection_scores_rows = ws.combined - ws.selection_row_lengths = ws.valid_widths - ws.provisional_rows = ws.final_indices - ws.keep_rows = ws.keep + bufs.final_indices = torch.zeros( + (max_requests, keep_count), dtype=torch.int32, device=device + ) + bufs.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) + bufs.selection_scores_rows = bufs.combined + bufs.selection_row_lengths = bufs.valid_widths + bufs.provisional_rows = bufs.final_indices + bufs.keep_rows = bufs.keep else: selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads - ws.selection_rows_per_request = selection_rows - ws.row_prompt_offsets = torch.zeros( + bufs.selection_rows_per_request = selection_rows + bufs.row_prompt_offsets = torch.zeros( max_requests * selection_rows, dtype=torch.int32, device=device ) - ws.row_mean = torch.empty( + bufs.row_mean = torch.empty( max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device ) - ws.row_std = torch.empty_like(ws.row_mean) - ws.selection_scores = torch.empty( + bufs.row_std = torch.empty_like(bufs.row_mean) + bufs.selection_scores = torch.empty( (max_requests, selection_rows, width), dtype=torch.float32, device=device ) - ws.row_seq_lens = torch.full( + bufs.row_seq_lens = torch.full( (max_requests, selection_rows), width, dtype=torch.int32, device=device ) - ws.top_indices_i32 = torch.zeros( + bufs.top_indices_i32 = torch.zeros( (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device ) - ws.keep = torch.empty( + bufs.keep = torch.empty( (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device ) - ws.selection_scores_rows = ws.selection_scores.view(max_requests * selection_rows, width) - ws.selection_row_lengths = ws.row_seq_lens.view(-1) - ws.provisional_rows = ws.top_indices_i32.view(-1, keep_count) - ws.keep_rows = ws.keep.view(-1, keep_count) - ws.settle_grid = (max_requests, ws.selection_rows_per_request) + bufs.selection_scores_rows = bufs.selection_scores.view( + max_requests * selection_rows, width + ) + bufs.selection_row_lengths = bufs.row_seq_lens.view(-1) + bufs.provisional_rows = bufs.top_indices_i32.view(-1, keep_count) + bufs.keep_rows = bufs.keep.view(-1, keep_count) + bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) # The settle launch always packs now; zero per-request move counts mask - # every pack store off, so these workspaces stay selection-only. + # every pack store off, so these buffers stay selection-only. zero_offsets = torch.zeros(max_requests + 1, dtype=torch.int32, device=device) zero_lengths = torch.zeros(max_requests, dtype=torch.int32, device=device) zero_indices = torch.zeros(1, dtype=torch.int32, device=device) - ws.settle_pack_tensors = (zero_lengths, zero_offsets, zero_indices, zero_offsets, zero_indices) - ws.settle_pack_shape = dict( + bufs.settle_pack_tensors = ( + zero_lengths, + zero_offsets, + zero_indices, + zero_offsets, + zero_indices, + ) + bufs.settle_pack_shape = dict( DENSE_TOTAL=0, SWA_TOTAL=0, MOVE_CAPACITY=keep_count, @@ -118,24 +125,24 @@ def _make_selection_ws( PER_LAYER=False, HAS_SWA=False, ) - return ws + return bufs -def _select_per_head(ws, scores, *, normalize_scores): +def _select_per_head(bufs, scores, *, normalize_scores): """The per-head selection flow: reduce kernels, then top-k settle.""" prepare_per_head_scores( scores, - ws.valid_widths, - ws.row_mean, - ws.row_std, - ws.selection_scores, - ws.row_seq_lens, - ws.max_requests, - num_kv_heads=ws.num_kv_heads, - per_layer=ws.eviction_mode == "per_layer_perhead", + bufs.valid_widths, + bufs.row_mean, + bufs.row_std, + bufs.selection_scores, + bufs.row_seq_lens, + bufs.max_requests, + num_kv_heads=bufs.num_kv_heads, + per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - settle_top_tokens(ws) + settle_top_tokens(bufs) def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: @@ -205,7 +212,7 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( device = torch.device("cuda", torch.cuda.current_device()) stream = torch.cuda.Stream(device=device) with torch.cuda.stream(stream): - ws = _make_selection_ws( + bufs = _make_selection_buffers( eviction_mode=eviction_mode, width=width, keep_count=keep_count, @@ -215,12 +222,12 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( num_query_heads=query_heads, num_kv_heads=kv_heads, ) - ws.valid_widths.copy_(valid_widths.to(device)) + bufs.valid_widths.copy_(valid_widths.to(device)) scores = scores_cpu.to(device) - _select_per_head(ws, scores, normalize_scores=normalize_scores) - first = ws.keep.cpu() - _select_per_head(ws, scores, normalize_scores=normalize_scores) - second = ws.keep.cpu() + _select_per_head(bufs, scores, normalize_scores=normalize_scores) + first = bufs.keep.cpu() + _select_per_head(bufs, scores, normalize_scores=normalize_scores) + second = bufs.keep.cpu() stream.synchronize() assert torch.equal(first, expected) @@ -246,23 +253,23 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid device=device, ).to(torch.float32) valid_widths = (width, width - 32) - ws = _make_selection_ws( + bufs = _make_selection_buffers( eviction_mode="union", width=width, keep_count=keep_count, device=device, max_requests=request_count, ) - ws.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) + bufs.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) # Write the shared per-request prompt lengths the way production staging # does: the union row-major view aliases the per-request buffer. - ws.prompt_offsets[:request_count].copy_( + bufs.prompt_offsets[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) - assert ws.row_prompt_offsets is ws.prompt_offsets - ws.combined.copy_(scores.amax(dim=1)) - settle_top_tokens(ws) - actual = ws.keep.cpu() + assert bufs.row_prompt_offsets is bufs.prompt_offsets + bufs.combined.copy_(scores.amax(dim=1)) + settle_top_tokens(bufs) + actual = bufs.keep.cpu() combined = scores.amax(dim=1).cpu() for request, valid_width in enumerate(valid_widths): @@ -402,9 +409,9 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) ) _set_protected_tails(compaction, protected_tails) # Production packs these buffers inside the fused settle launch; with - # pre-settled ordinals the standalone pack call is its exact analog. - launch_move_pack(compaction["dense_pack"]) - run_cache_compactions(compaction) + # pre-settled ordinals the standalone pack in run_compaction is its + # exact analog. + _run_compaction(compaction) torch.cuda.synchronize(device) for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): @@ -449,7 +456,7 @@ 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") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - prepare_eviction_workspace, + init_eviction_buffers, run_eviction_round, ) @@ -501,7 +508,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): mlr_coef[:, :, 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - ws = prepare_eviction_workspace( + bufs = init_eviction_buffers( eviction_mode="per_layer_perhead", layer_pools=pools, dense_groups=dense_groups, @@ -523,35 +530,41 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): layer_group_representative=layer_group_representative, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], ) - ws.block_offsets_device.zero_() - ws.block_offsets_device[..., :2].copy_(_encode_block_offsets(torch.stack(page_tables))) - ws.round_starts_device.fill_(0) - ws.valid_seq_lens_device.fill_(seq_len) - ws.token_starts_device.fill_(0) - - # This compaction packs its own move indices (the ordinals settle - # mid-round, so the pack rides the family slot exactly like the draft's - # own pack does in production); the workspace's fused pack stays masked - # off (its staged move-offsets row is all zeros). - ws.compaction = _build_compaction( + bufs.block_offsets_device.zero_() + bufs.block_offsets_device[..., :2].copy_(_encode_block_offsets(torch.stack(page_tables))) + bufs.round_starts_device.fill_(0) + bufs.valid_seq_lens_device.fill_(seq_len) + bufs.token_starts_device.fill_(0) + + # Attach a standalone bundle wholesale (families AND the fused settle + # launch data), replacing the constructor-built one: the fused settle + # launch then packs this bundle's construction-time move offsets exactly + # like production packs the staged rows, and the round's inline C++ + # moves consume the same buffers. + compaction = _build_compaction( eviction_mode="per_layer_perhead", layer_pools=pools, dense_layers=dense_layers, layer_group_representative=layer_group_representative, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], - kept_token_ordinals=ws.keep[:1], - valid_sequence_lengths=ws.valid_seq_lens_device[:1], - kv_block_offsets=ws.block_offsets_device, - page_table_slots=ws.representative_slots, + kept_token_ordinals=bufs.keep[:1], + valid_sequence_lengths=bufs.valid_seq_lens_device[:1], + kv_block_offsets=bufs.block_offsets_device, + page_table_slots=bufs.representative_slots, request_count=1, prompt_offsets=torch.zeros(1, dtype=torch.int32, device=device), decode_keep_count=keep_count, protected_tail_capacity=0, ) - _set_protected_tails(ws.compaction, [0]) - ws.compaction["families"][0]["pack"] = ws.compaction["dense_pack"] - run_eviction_round(ws, normalize_scores=False) - assert torch.equal(ws.keep, expected_keep) + _set_protected_tails(compaction, [0]) + bufs.compaction_families = compaction["families"] + bufs.settle_pack_tensors = compaction["settle_pack_tensors"] + bufs.settle_pack_shape = compaction["settle_pack_shape"] + bufs.swa_destination_bases = compaction["swa_destination_bases"] + bufs.swa_rebase_delta = compaction["swa_rebase_delta"] + bufs.draft_pack = compaction["draft_pack"] + run_eviction_round(bufs, normalize_scores=False) + assert torch.equal(bufs.keep, expected_keep) torch.cuda.synchronize(device) for before_pool, after_pool, table, layer in zip( @@ -580,8 +593,8 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + init_eviction_buffers, mark_page_tables_consumed, - prepare_eviction_workspace, run_eviction_round, stage_eviction_cohort, ) @@ -694,7 +707,7 @@ def expected_keep() -> torch.Tensor: mlr_coef[..., 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - ws = prepare_eviction_workspace( + bufs = init_eviction_buffers( eviction_mode="union", layer_pools=[pool], dense_groups=[[0]], @@ -723,7 +736,7 @@ def expected_keep() -> torch.Tensor: def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) stage_eviction_cohort( - ws, + bufs, manager, [request_id], [0], @@ -737,9 +750,9 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: # set (derived from raw scores) is unchanged. The settle launch # packs the move sources and the C++ compacts run in the same # round call; the kept ordinals stay readable afterwards. - run_eviction_round(ws, normalize_scores=True) - selected = ws.keep[0].clone().to(torch.long) - mark_page_tables_consumed(ws, manager._stream) + run_eviction_round(bufs, normalize_scores=True) + selected = bufs.keep[0].clone().to(torch.long) + mark_page_tables_consumed(bufs, manager._stream) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) after = snapshot(compacted_capacity) @@ -840,8 +853,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) - launch_move_pack(compaction["dense_pack"]) - run_cache_compactions(compaction) + _run_compaction(compaction) torch.cuda.synchronize(device) for request, (valid_seq_len, tail_length) in enumerate( From a0b90d41cfff107fdac78681e40b0a5f0a9c0472 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 21:41:36 -0700 Subject: [PATCH 096/178] [None][chore] Zero-risk clarity pass from the findings ledger (knife 17A) Ledger items #2,4,6,8,9,12,33,34 (comment/rename only) + #13 (init audit) + cpp #23,26,27,28 (comment/rename + one provably-identical check-loop bound; inert without a rebuild): - #2 tile/token units: kernel-body 'page' names that lied about units become tile_index/tile_start_token/tiles_processed/... (physical_page and fragment_page stay: true pool-page indices). - #4/#6 stale K_SW64/4-KiB and 'matching the production score kernel' comments rewritten in derived quantities. - #8 the 2^31 scratch audit names the guarded quantity (7 padded score planes x plane stride, the K1 epilogue's i32 head-axis fold) and the seg_out_offset int32 cast points at it; cross-ref added at the fused epilogue fold. - #9 union-store i32 bound pinned at both store arms; #12 grid/cluster and shard-pick heuristics get intent comments (2-wave threshold). - #13 per-head score/selection rectangles get the same contract-mode init audit as the scratch plane (Triton kernels fold row offsets in int32; contract-legal per-head geometry could overflow silently). - #33 'plane' terminology unified (score plane vs K/V plane vs coefficient plane; cpp head-plane stride -> head-row stride). - #34 the C++ in-place-copy precondition is now documented on both producers (settle kernel, init_compaction_buffers). - cpp: #23 op/header doc blocks describe the shipped pipelined-bf16 design; #26 V2/Fast-qualified names align on the op-name stem (invokeSparseKvCacheCompactLayers, dispatchSparseKvCacheCompactGeometry, kSparseKvCompactTokensPerTile, SparseKvCacheCompactBf16Params, INSTANTIATE_SPARSE_KV_CACHE_COMPACT_LAYERS); #27 'Tip-ABI' labels point at the port note; #28 the pools validation loop starts at layer 1 (layer 0 defined the reference geometry). Signed-off-by: tianruih --- .../kernels/unfusedAttentionKernels.h | 19 +-- .../unfusedAttentionKernels_2_bf16_bf16.cu | 2 +- .../unfusedAttentionKernels_2_template.h | 58 ++++----- .../thop/sparseKvCacheCompactOp.cpp | 22 ++-- .../_torch/kv_cache_compression/compaction.py | 5 +- .../triattention/triattention.py | 21 +++- .../triattention_cute_score_fused.py | 117 ++++++++++-------- .../triattention_cute_selection.py | 25 +++- .../triattention/triattention_kernels.py | 6 + 9 files changed, 169 insertions(+), 106 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 1cf98efa94b5..06d58bdbe9cf 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -413,16 +413,17 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams template void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams params, cudaStream_t stream); -//! Adapt a uniform group of KVCacheManagerV2 layer pools to the existing -//! sparse-KV updater. Device pointer arrays allow one layered launch while the -//! updater remains the only implementation of the in-place copy loop. Within -//! each request and head, source ordinals must increase strictly and satisfy -//! destinationBases[request] + move <= source[move]; the per-request bases -//! let one launch cover a cohort with mixed pinned-prompt lengths. -//! sourceHeadStride is the head-plane stride of sparseKvIndices: the index -//! buffers may be wider than one round's total move count. +//! Adapt a uniform group of KVCacheManagerV2 layer pools to the batched +//! compaction kernels (double-buffered cp.async pipelined bf16 copies, one +//! CTA per layer/KV-head/request). Device pointer arrays allow one layered +//! launch. Within each request and head, source ordinals must increase +//! strictly and satisfy destinationBases[request] + move <= source[move]; +//! the per-request bases let one launch cover a cohort with mixed +//! pinned-prompt lengths. sourceHeadStride is the head-row stride of +//! sparseKvIndices: the index buffers may be wider than one round's total +//! move count. template -void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, +void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, int64_t sourceLayerStride, int64_t sourceHeadStride, int32_t const* sparseKvOffsets, int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu index ec1b92aad4fa..023702c4a72f 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu @@ -26,7 +26,7 @@ namespace kernels #ifdef ENABLE_BF16 INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(__nv_bfloat16, __nv_bfloat16, KVBlockArray); INSTANTIATE_ATTENTION_INPUT_OUTPUT_PROCESSING(__nv_bfloat16, __nv_bfloat16, KVLinearBuffer); -INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(__nv_bfloat16); +INSTANTIATE_SPARSE_KV_CACHE_COMPACT_LAYERS(__nv_bfloat16); #endif } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index a1a290549c20..e7b1ec5642a4 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1758,23 +1758,24 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams #ifdef ENABLE_BF16 -// Optimized bf16 compaction fast path, ported from Fanrong Li's optimized +// Pipelined bf16 compaction kernels, ported from Fanrong Li's optimized // compact kernels (snapshot 2026-07-19). The port keeps the double-buffered -// cp.async pipeline intact and adapts only the addressing to this tree's ABI: +// cp.async pipeline intact and adapts only the addressing to this +// repository's KVCacheManagerV2 ABI: // (a) the per-layer page-table pointer array became one flat int32 V2 // K-plane block-offset table shared by all layers (entries encode // 2 * page + plane with plane == 0, so >> 1 recovers the page), strided // per request; // (b) the host-scalar destination base became per-request device bases, read // once per CTA (one launch covers a cohort with mixed prompt lengths); -// (c) the head-plane stride of the move-source indices is an explicit +// (c) the head-row stride of the move-source indices is an explicit // parameter instead of being derived from sourceOffsets[batchSize] on // device: the move buffers are allocation-wide, so a device-derived // stride would silently read the wrong plane for every KV head above // head 0. // The original kernels were written for 128-token pages and Dh = 64; the port // additionally parameterizes the page and head-vector math so 32-token pages -// and Dh = 128 (this tree's production geometry) take the same pipeline. +// and Dh = 128 (the production geometry here) take the same pipeline. namespace compact_detail { @@ -1813,10 +1814,10 @@ __device__ __forceinline__ void waitGroup() } // namespace compact_detail // One pipeline stage moves a 32-token tile regardless of the page geometry. -constexpr int32_t kSparseKvCompactFastTokensPerTile = 32; +constexpr int32_t kSparseKvCompactTokensPerTile = 32; //! Launch parameters for the pipelined bf16 fast-path compaction kernels. -struct SparseKvCacheCompactV2Bf16Params +struct SparseKvCacheCompactBf16Params { int64_t const* poolPointers; int32_t const* pageTable; @@ -1841,7 +1842,7 @@ struct SparseKvCacheCompactV2Bf16Params //! threadIdx.y walks the tokens of a tile. template __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) - * kSparseKvCompactFastTokensPerTile) void sparseKvCacheCompactV2Bf16PipelineKernel(SparseKvCacheCompactV2Bf16Params + * kSparseKvCompactTokensPerTile) void sparseKvCacheCompactV2Bf16PipelineKernel(SparseKvCacheCompactBf16Params params) { static_assert(std::is_same_v); @@ -1852,7 +1853,7 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) // 16B vectors per head: Dh64 -> 8 lanes (block 8x32 = 256 threads), // Dh128 -> 16 lanes (block 16x32 = 512 threads). constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); - constexpr int32_t kTokensPerTile = kSparseKvCompactFastTokensPerTile; + constexpr int32_t kTokensPerTile = kSparseKvCompactTokensPerTile; constexpr int32_t kVectorsPerTile = kTokensPerTile * kVectorsPerHead; // A buffer holds one K tile plus one V tile; two buffers ping-pong. constexpr int32_t kVectorsPerBuffer = 2 * kVectorsPerTile; @@ -1872,15 +1873,16 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) // without an explicit map, launch layer i reads source plane i (the flat // layout passes sourceLayerStride == 0, which collapses the term). int32_t const sourceLayer = params.sourceLayerIndices == nullptr ? layerIdx : params.sourceLayerIndices[layerIdx]; - // Tip-ABI adaptation (c): head planes are strided by the allocation width + // ABI adaptation (c) -- see the port note above the compact_detail + // namespace: head rows are strided by the allocation width // of the move buffers; this request's range within a plane starts at // moveBegin. int64_t const sourceMoveBase = static_cast(sourceLayer) * params.sourceLayerStride + static_cast(kvHeadIdx) * params.sourceHeadStride + moveBegin; - // Tip-ABI adaptation (b): per-request landing position. + // ABI adaptation (b) -- see the port note above the compact_detail namespace: per-request landing position. int32_t const destinationBase = params.destinationBases[batchIdx]; auto* const pool = reinterpret_cast(static_cast(params.poolPointers[layerIdx])); - // Tip-ABI adaptation (a): flat V2 K-plane block-offset table; each lookup + // ABI adaptation (a) -- see the port note above the compact_detail namespace: flat V2 K-plane block-offset table; each lookup // below decodes an entry to a page with >> 1. TokensPerBlock is a // compile-time power of two, so / and % lower to shifts and masks. int32_t const* const pageTable = params.pageTable + static_cast(batchIdx) * params.pageTableRequestStride; @@ -1996,7 +1998,7 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) //! helper below). template __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) - * kSparseKvCompactFastTokensPerTile) void sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel(SparseKvCacheCompactV2Bf16Params + * kSparseKvCompactTokensPerTile) void sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel(SparseKvCacheCompactBf16Params params) { static_assert(std::is_same_v); @@ -2004,7 +2006,7 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) static_assert(TokensPerBlock == 32 || TokensPerBlock == 128); constexpr int32_t kBufferCount = 2; constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); - constexpr int32_t kTokensPerTile = kSparseKvCompactFastTokensPerTile; + constexpr int32_t kTokensPerTile = kSparseKvCompactTokensPerTile; constexpr int32_t kVectorsPerTile = kTokensPerTile * kVectorsPerHead; constexpr int32_t kVectorsPerBuffer = 2 * kVectorsPerTile; @@ -2143,10 +2145,10 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) } template -void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactV2Bf16Params const& params, cudaStream_t stream) +void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) { constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); - dim3 const block(kVectorsPerHead, kSparseKvCompactFastTokensPerTile); + dim3 const block(kVectorsPerHead, kSparseKvCompactTokensPerTile); dim3 const grid(params.numLayers, params.numKvHeads, params.batchSize); // Two ping-pong buffers x (K tile + V tile) of 32 tokens x kVectorsPerHead // 16B vectors: @@ -2154,28 +2156,28 @@ void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactV2Bf16Params c // Dh128: 4 * 32 * 16 * 16 B = 32 KiB // Both fit the 48 KiB per-CTA dynamic shared memory default, so no // cudaFuncSetAttribute opt-in is required. - size_t const sharedBytes = 4 * kSparseKvCompactFastTokensPerTile * kVectorsPerHead * sizeof(uint4); + size_t const sharedBytes = 4 * kSparseKvCompactTokensPerTile * kVectorsPerHead * sizeof(uint4); sparseKvCacheCompactV2Bf16PipelineKernel<<>>(params); } template void launchSparseKvCacheCompactV2Bf16DestinationPagePipeline( - SparseKvCacheCompactV2Bf16Params const& params, cudaStream_t stream) + SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) { constexpr int32_t kDestinationPageBuffers = 2; constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); - dim3 const block(kVectorsPerHead, kSparseKvCompactFastTokensPerTile); + dim3 const block(kVectorsPerHead, kSparseKvCompactTokensPerTile); dim3 const grid(params.numLayers, params.numKvHeads, params.batchSize); // Same 16/32 KiB tile buffers as the plain pipeline plus two staged // destination page indices; still far below the 48 KiB default. - size_t const sharedBytes = 4 * kSparseKvCompactFastTokensPerTile * kVectorsPerHead * sizeof(uint4) + size_t const sharedBytes = 4 * kSparseKvCompactTokensPerTile * kVectorsPerHead * sizeof(uint4) + kDestinationPageBuffers * sizeof(int32_t); sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel <<>>(params); } template -void dispatchSparseKvCacheCompactV2FastGeometry(SparseKvCacheCompactV2Bf16Params const& params, cudaStream_t stream) +void dispatchSparseKvCacheCompactGeometry(SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) { // The destination-page variant requires every request's destination base // to be 32-token tile aligned so each tile lands in one page, but the @@ -2369,7 +2371,7 @@ void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams } template -void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, +void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, int64_t sourceLayerStride, int64_t sourceHeadStride, int32_t const* sparseKvOffsets, int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, @@ -2385,7 +2387,7 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con // dtypes and geometries fail the check below instead of falling back. if ((headDim == 64 || headDim == 128) && (tokensPerBlock == 32 || tokensPerBlock == 128)) { - SparseKvCacheCompactV2Bf16Params fastParams{}; + SparseKvCacheCompactBf16Params fastParams{}; fastParams.poolPointers = poolPointers; fastParams.pageTable = pageTable; fastParams.sourceIndices = sparseKvIndices; @@ -2402,19 +2404,19 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con fastParams.bytesPerPage = 2 * fastParams.bytesPerKvHalf; if (headDim == 64 && tokensPerBlock == 32) { - dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + dispatchSparseKvCacheCompactGeometry(fastParams, stream); } else if (headDim == 64 && tokensPerBlock == 128) { - dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + dispatchSparseKvCacheCompactGeometry(fastParams, stream); } else if (headDim == 128 && tokensPerBlock == 32) { - dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + dispatchSparseKvCacheCompactGeometry(fastParams, stream); } else { - dispatchSparseKvCacheCompactV2FastGeometry(fastParams, stream); + dispatchSparseKvCacheCompactGeometry(fastParams, stream); } return; } @@ -2442,8 +2444,8 @@ void invokeSparseKvCacheCompactV2Layers(int64_t const* poolPointers, int32_t con QKVPreprocessingParams params, cudaStream_t stream); \ //////////////////////////////////////////////////////////////////////////////////////////////////// -#define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_V2_LAYERS(T) \ - template void invokeSparseKvCacheCompactV2Layers(int64_t const*, int32_t const*, int32_t, int64_t, \ +#define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_LAYERS(T) \ + template void invokeSparseKvCacheCompactLayers(int64_t const*, int32_t const*, int32_t, int64_t, \ int32_t const*, int32_t const*, int64_t, int64_t, int32_t const*, int32_t const*, int32_t, int32_t, int32_t, \ int32_t, cudaStream_t); diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index 4bb2c0a34121..b63845156c9d 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -33,13 +33,14 @@ namespace torch_ext { //! Adapt one uniform group of KVCacheManagerV2 HND layer pools to the -//! existing sparse-KV post-FMHA updater. Layers share one V2 block-offset -//! table; per-request destinationBases replace the former arbitrary -//! destination tensor because every compaction move targets one contiguous -//! interval per request. Within each request and KV head, the caller must -//! supply increasing source ordinals with destinationBases[request] + move -//! <= source[move], which makes the updater's forward tiled in-place copy -//! safe. +//! batched compaction kernels: dedicated double-buffered cp.async pipelined +//! bf16 kernels, one CTA per (layer, KV head, request), addressed through a +//! flat V2 K-plane block-offset table shared by all layers. Per-request +//! destinationBases replace the former arbitrary destination tensor because +//! every compaction move targets one contiguous interval per request. +//! Within each request and KV head, the caller must supply increasing +//! source ordinals with destinationBases[request] + move <= source[move], +//! which makes the forward tiled in-place copy safe. void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor const& poolPointers, th::Tensor const& pageTable, th::Tensor const& sourceIndices, th::Tensor const& sourceOffsets, th::Tensor const& destinationBases, std::optional const& sourceLayerIndices) @@ -64,7 +65,8 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor auto const batchSize = static_cast(pageTable.size(0)); auto const pageTableRequestStride = pageTable.stride(0); - for (int32_t layer = 0; layer < numLayers; ++layer) + // Layer 0 defined the reference geometry in the firstPool checks above. + for (int32_t layer = 1; layer < numLayers; ++layer) { auto const& pool = pools[layer]; TORCH_CHECK(pool.is_cuda() && pool.get_device() == device && pool.scalar_type() == dtype && pool.dim() == 5 @@ -115,7 +117,7 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor // source_offsets carve each request's move range out of source_indices; // the values live on device and the kernel trusts them. The index buffer - // may be wider than one round's total move count, so the head-plane + // may be wider than one round's total move count, so the head-row // stride passed below comes from the tensor shape, not from the offsets. TORCH_CHECK(sourceOffsets.is_cuda() && sourceOffsets.get_device() == device && sourceOffsets.scalar_type() == th::kInt32 && sourceOffsets.is_contiguous() && sourceOffsets.dim() == 1 @@ -133,7 +135,7 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor auto const sourceHeadStride = sourceIndices.size(-1); if (dtype == th::kBFloat16) { - tk::invokeSparseKvCacheCompactV2Layers<__nv_bfloat16>(poolPointers.data_ptr(), + tk::invokeSparseKvCacheCompactLayers<__nv_bfloat16>(poolPointers.data_ptr(), pageTable.data_ptr(), numLayers, pageTableRequestStride, sourceIndices.data_ptr(), sourceLayerPtr, sourceLayerStride, sourceHeadStride, sourceOffsets.data_ptr(), bases, batchSize, numKvHeads, tokensPerBlock, headDim, stream); diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 052e35607645..e752c97d47b7 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -165,7 +165,10 @@ def init_compaction_buffers( ``kept_token_ordinals`` carries increasing kept decode ordinals (absolute positions) per request; prompt tokens never move, so the rectangle is prompt-length independent and ``prompt_offsets`` carries each request's - pinned prompt length. ``kv_block_offsets`` is the staged V2 snapshot + pinned prompt length. The increasing order is LOAD-BEARING: + sparseKvCacheCompactOp.cpp's forward tiled in-place copy requires + increasing source ordinals with ``destinationBases[request] + move <= + source[move]``. ``kv_block_offsets`` is the staged V2 snapshot ``[slot, request, K/V, block]`` (offset = ``2*page + plane``); ``protected_tail_capacity`` is the widest per-request tail this geometry must support -- actual per-round lengths arrive through the staged diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 17bf80fecf11..a095a6f5b611 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -347,11 +347,15 @@ def init_eviction_buffers( seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) max_segments = max_requests * bufs.num_layers - # Head axis pads to the MMA tile N=8; one padded scratch plane must stay - # 32-bit indexable (wraparound = silent wild reads, not a clean error). + # The K1 epilogue folds the head axis as a DSL dynamic coordinate + # (kv_head * N <= 7 score planes of max_segments * seq_len columns each), + # so the largest 32-bit-folded offset is 7 planes x the plane stride; + # every downstream i32 offset (including the seg_out_offset cast below) + # is bounded by it. Wraparound would be a silent wild read, not a clean + # error, hence the loud audit. if (8 - 1) * max_segments * seq_len >= 2**31: raise ValueError( - f"score bucket overflows the 32-bit scratch plane: {(8 - 1) * max_segments * seq_len}" + f"score bucket overflows the 32-bit score plane: {(8 - 1) * max_segments * seq_len}" ) # The kernel scores each request's window into a head-major scratch; # all buffers below are persistent because the compiled kernels capture @@ -359,6 +363,7 @@ def init_eviction_buffers( bufs.cute_scratch = torch.empty( bufs.num_kv_heads * 8 * max_segments * seq_len, dtype=torch.float32, device=device ) + # int32 is safe here: covered by the 2^31 score-plane audit above. seg_out_offset = (torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len).to( torch.int32 ) @@ -451,6 +456,16 @@ def init_eviction_buffers( if eviction_mode == "per_head" else bufs.num_layers * bufs.num_kv_heads ) + # The Triton stats/reduce/settle kernels fold score-row offsets in + # int32; both per-head rectangles must stay 32-bit indexable + # (wraparound = silent wild reads, not a clean error). + score_rect = max_requests * bufs.num_layers * bufs.num_q_heads * decode_width + selection_rect = max_requests * selection_rows * max(decode_width, keep_count) + if max(score_rect, selection_rect) >= 2**31: + raise ValueError( + f"per-head score rectangles overflow 32-bit indexing: " + f"scores {score_rect}, selection {selection_rect}" + ) bufs.selection_rows_per_request = selection_rows bufs.row_prompt_offsets = torch.zeros( (max_requests * selection_rows,), dtype=torch.int32, device=device 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 index dc1d6b8f40ba..53ba4cc96ce6 100644 --- 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 @@ -274,8 +274,10 @@ def __call__( tcgen05.CtaGroup.ONE, self.mma_tiler[:2], ) - # Split raw-K transport: two K_SW64 4-KiB stages per raw-page buffer, - # one each for the real and imaginary K bands. + # Split raw-K transport: one real and one imaginary stage per + # raw-page buffer, each a CTA_M x num_freqs bf16 tile + # (raw_tma_copy_bytes per phase); the descriptor encoder picks the + # SW64/SW128 swizzle from the num_freqs row width. raw_bf16_direct_a_smem_layout = sm100_utils.make_smem_layout_a( raw_bf16_tiled_mma, (CTA_M, N, self.num_freqs), @@ -397,7 +399,7 @@ class SharedStorage: self.shared_storage = SharedStorage # 64-bit: at large request counts this product exceeds 2^31 (the # score scratch spans request*layer segments of seq_len columns), so - # the head-plane stride must reach the kernel as Int64. + # the score-plane stride must reach the kernel as Int64. sum_seq = cutlass.Int64(request_count * self.num_layers * self.seq_len) num_ctas = request_count * self.num_layers * self.num_kv_heads * self.page_shards self.kernel( @@ -514,8 +516,8 @@ def kernel( cpasync_raw_k_imag = cpasync_raw_k_0[(None, None, None, 1)] 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)] - # Each stage slice retains the K_SW64 pointer flags. Reuse only - # the feature-first outer mapping for the corresponding TMA + # Each stage slice retains the swizzled smem pointer flags. Reuse + # only the feature-first outer mapping for the corresponding TMA # destination so the swizzle is not applied twice. raw_tma_source_tiles = cute.local_tile( raw_tma_source, @@ -600,10 +602,10 @@ def kernel( swizzle=raw_bf16_b_smem_layout.inner, ) - page_index = score_start // self.tile_tokens + page_shard - page_start = page_index * self.tile_tokens - shard_first_page_start = page_start - pages_processed = cutlass.Int32(0) + tile_index = score_start // self.tile_tokens + page_shard + tile_start_token = tile_index * self.tile_tokens + 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((N,), cutlass.Float32) stats_origins_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) @@ -622,16 +624,15 @@ def kernel( ) physical_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) prefetched_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) - shard_has_page = valid_seq_len > score_start and page_start < valid_seq_len - empty_shard = valid_seq_len <= score_start or page_start >= valid_seq_len + shard_has_page = valid_seq_len > score_start and tile_start_token < valid_seq_len + empty_shard = valid_seq_len <= score_start or tile_start_token >= valid_seq_len if cutlass.dynamic_expr(shard_has_page): if warp_idx == self.producer_warp_id: if lane_idx == 0: # The staged K-plane entries encode physical_page * - # kv_factor (2); decode to the pool page index here, - # matching the production score kernel. + # kv_factor (2); decode to the pool page index here. producer_prefetched_page_id_lane0 = ( - cutlass.Int32(page_ids[page_off + page_index * self.pages_per_tile]) // 2 + cutlass.Int32(page_ids[page_off + tile_index * self.pages_per_tile]) // 2 ) if cutlass.const_expr(self.fragments_per_phase > 1): for fragment in cutlass.range_constexpr(1, self.pages_per_tile): @@ -641,11 +642,11 @@ def kernel( # masked downstream) so the TMA never # dereferences an unstaged block entry. fragment_page_id = producer_prefetched_page_id_lane0 - if page_start + fragment * self.box_tokens < valid_seq_len: + if tile_start_token + fragment * self.box_tokens < valid_seq_len: fragment_page_id = ( cutlass.Int32( page_ids[ - page_off + page_index * self.pages_per_tile + fragment + page_off + tile_index * self.pages_per_tile + fragment ] ) // 2 @@ -839,8 +840,8 @@ def kernel( raw_tma_producer_state.advance() while ( valid_seq_len > score_start - and page_start < valid_seq_len - and pages_processed < self.max_tiles + and tile_start_token < valid_seq_len + and tiles_processed < self.max_tiles ): physical_page = cutlass.Int32(0) if warp_idx == self.producer_warp_id: @@ -857,14 +858,15 @@ def kernel( for page_half in cutlass.range_constexpr(self.halves_per_page): raw_page_buffer = cutlass.Int32(0) if cutlass.const_expr(self.write_partial_stats): - raw_page_buffer = pages_processed % RAW_PAGE_BUFFERS + 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 4-KiB K_SW64 real - # view. Every producer-warp lane participates - # in the PipelineTmaAsync barrier election. + # Phase 0 fills the packed real-band stage view + # (CTA_M x num_freqs bf16, raw_tma_copy_bytes). + # Every producer-warp lane participates in the + # PipelineTmaAsync barrier election. raw_tma_pipeline.producer_acquire(raw_tma_producer_state) for fragment in cutlass.range_constexpr(self.fragments_per_phase): fragment_page = physical_page @@ -919,17 +921,19 @@ def kernel( next_page_id_lane0 = cutlass.Int32(0) if warp_idx == self.producer_warp_id: if lane_idx == 0: - next_page_start = page_start + self.tile_tokens * self.page_shards - next_pages_processed = pages_processed + 1 + next_tile_start_token = ( + tile_start_token + self.tile_tokens * self.page_shards + ) + next_pages_processed = tiles_processed + 1 if ( - next_page_start < valid_seq_len + next_tile_start_token < valid_seq_len and next_pages_processed < self.max_tiles ): next_page_id_lane0 = ( cutlass.Int32( page_ids[ page_off - + (page_index + self.page_shards) * self.pages_per_tile + + (tile_index + self.page_shards) * self.pages_per_tile ] ) // 2 @@ -941,14 +945,14 @@ def kernel( # fragment's page. next_fragment_page_id = next_page_id_lane0 if ( - next_page_start + fragment * self.box_tokens + next_tile_start_token + fragment * self.box_tokens < valid_seq_len ): next_fragment_page_id = ( cutlass.Int32( page_ids[ page_off - + (page_index + self.page_shards) + + (tile_index + self.page_shards) * self.pages_per_tile + fragment ] @@ -988,10 +992,12 @@ def kernel( prefetch_next_raw = True prefetched_page_half = page_half + 1 else: - next_page_start = page_start + self.tile_tokens * self.page_shards - next_pages_processed = pages_processed + 1 + next_tile_start_token = ( + tile_start_token + self.tile_tokens * self.page_shards + ) + next_pages_processed = tiles_processed + 1 prefetch_next_raw = ( - next_page_start < valid_seq_len + next_tile_start_token < valid_seq_len and next_pages_processed < self.max_tiles ) prefetched_physical_page = cute.arch.shuffle_sync( @@ -1231,7 +1237,7 @@ def kernel( ) acc_pipeline.producer_commit(acc_producer_state) acc_producer_state.advance() - if pages_processed == 0: + if tiles_processed == 0: if cutlass.const_expr(page_half == 0): cute.arch.relinquish_tmem_alloc_permit(is_two_cta=False) acc_pipeline.consumer_wait(acc_consumer_state) @@ -1242,12 +1248,17 @@ def kernel( raw_tma_pipeline.consumer_release(raw_tma_consumer_state) raw_tma_consumer_state.advance() # Every term multiplying sum_seq must stay 64-bit: with - # request*layer segments of seq_len columns the head-plane + # request*layer segments of seq_len columns the score-plane # stride alone can exceed 2^31. The scratch head axis is # padded to the MMA tile N=8 per KV head (group-4 columns - # 4..7 land in padded planes holding zero scores). + # 4..7 land in padded score planes holding zero scores). The + # host audit in triattention.init_eviction_buffers bounds the + # 32-bit head-axis fold (kv_head * N <= 7 planes) here. output_offset = ( - cutlass.Int64(kv_head * N) * sum_seq + out_base + page_start + page_half * CTA_M + cutlass.Int64(kv_head * N) * sum_seq + + out_base + + tile_start_token + + page_half * CTA_M ) page_output = cute.make_tensor( output.iterator + output_offset, @@ -1285,7 +1296,8 @@ def kernel( # only the straddling first tile takes the per-token # branch. if cutlass.dynamic_expr( - page_start >= score_start and page_start + CTA_M <= self.seq_len + tile_start_token >= score_start + and tile_start_token + CTA_M <= self.seq_len ): cute.copy( simt_atom, @@ -1293,7 +1305,7 @@ def kernel( tTR_gC[(None, None, None, subtile_idx)], ) else: - output_token = page_start + epilogue_tidx + output_token = tile_start_token + epilogue_tidx if cutlass.dynamic_expr( output_token >= score_start and output_token < self.seq_len ): @@ -1310,7 +1322,7 @@ def kernel( stats_page_scores_m128[stats_head] = cutlass.Float32( stats_output[stats_value] ) - if pages_processed == 0: + if tiles_processed == 0: if cutlass.const_expr(page_half == 0): if tidx == 0: sStats[stats_head] = cutlass.Float32( @@ -1329,11 +1341,11 @@ def kernel( raw_tma_pipeline.consumer_release(raw_tma_consumer_state) raw_tma_consumer_state.advance() if cutlass.const_expr(self.write_partial_stats): - if pages_processed == 0: + if tiles_processed == 0: if cutlass.const_expr(page_half == 0): for stats_head in cutlass.range_constexpr(N): stats_origins_m128[stats_head] = sStats[stats_head] - stats_token = page_start + tidx + stats_token = tile_start_token + tidx if tidx < EPILOGUE_THREADS: if cutlass.dynamic_expr( stats_token >= score_start and stats_token < valid_seq_len @@ -1349,9 +1361,9 @@ def kernel( stats_square_sums_m128[stats_head] = ( stats_square_sums_m128[stats_head] + stats_delta * stats_delta ) - page_index += self.page_shards - page_start += self.tile_tokens * self.page_shards - pages_processed += 1 + tile_index += self.page_shards + tile_start_token += self.tile_tokens * 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) @@ -1381,13 +1393,17 @@ def kernel( stats_scratch_base = 8 + (stats_warp * N + 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 = pages_processed * self.tile_tokens - if pages_processed > 0: - stats_invalid_prefix = score_start - shard_first_page_start + stats_count_i32 = tiles_processed * self.tile_tokens + 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_page = page_start - self.tile_tokens * self.page_shards - stats_invalid_tail = stats_last_page + self.tile_tokens - valid_seq_len + stats_last_tile_start_token = ( + tile_start_token - self.tile_tokens * self.page_shards + ) + stats_invalid_tail = ( + stats_last_tile_start_token + self.tile_tokens - valid_seq_len + ) if cutlass.dynamic_expr(stats_invalid_tail > 0): stats_count_i32 = stats_count_i32 - stats_invalid_tail stats_count = cutlass.Float32(stats_count_i32) @@ -1724,6 +1740,9 @@ def __init__( small_stats = self._compiled_stats.get(1) large_stats = self._compiled_stats.get(max_requests) for request_count in range(1, max_requests + 1): + # Shard-pick heuristic: give small cohorts the extra page + # shard while the 2-shard grid stays under two waves + # (2 * sm_count CTAs); larger cohorts already fill the GPU. two_shard_ctas = request_count * num_layers * num_kv_heads * 2 use_extra_score_shard = two_shard_ctas < 2 * self.sm_count self._compiled[request_count] = ( @@ -1771,7 +1790,7 @@ def __init__( seq_len=seq_len, num_q_heads=num_q_heads, # The score scratch pads each KV head's group - # of head planes to the MMA tile N=8; the + # of score planes to the MMA tile N=8; the # finalizer maps real head rows onto those # padded planes (identity for GQA group 8). num_kv_heads=num_kv_heads, 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 index ea5aae2b1209..bb07036ee6c2 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -46,7 +46,13 @@ def _select_normalize_union_config( width: int, sm_count: int, ) -> tuple[int, int, int]: - """Return tokens per lane, token subtiles, and row-cluster CTAs.""" + """Return tokens per lane, token subtiles, and row-cluster CTAs. + + Heuristic: prefer the small token tile (more CTAs, cluster-widened per + row) while the whole grid still fits the residency bound below; past it, + fall back to the large tile with no clustering so each CTA carries more + tokens instead of oversubscribing the SMs. + """ 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 @@ -127,8 +133,8 @@ def __init__( token_subtiles: int, row_cluster_ctas: int, ) -> None: - # The score scratch pads each KV head's group of head planes up to - # the MMA tile: real head row ``q_head`` lives in scratch plane + # The score scratch pads each KV head's group of score planes up to + # the MMA tile: real head row ``q_head`` lives in score plane # ``kv * 8 + qg``. The partial-stats rows are always compact. self.score_group_size = num_q_heads // num_kv_heads self.score_head_pad = _PADDED_HEAD_COLUMNS - self.score_group_size @@ -178,6 +184,10 @@ def __call__( stream=stream, ) else: + # 1D-flattened cluster grid: the X extent is divisible by the + # cluster size and cluster peers are consecutive CTAs, so the + # kernel's (request, token-tile, rank) decode must keep exactly + # this factor order. kernel.launch( grid=( request_count * self.num_token_tiles * self.row_cluster_ctas, @@ -452,7 +462,11 @@ def kernel( subtile_first_token = first_token + token_subtile * self.subtile_token_tile # The output row covers this request's own window, # [0, valid - start); the straddling subtile falls back - # to the per-token stores. + # to the per-token stores. union_scores spans + # request_count * width < 2^31 by the host score-plane + # audit (triattention.init_eviction_buffers), so the i32 + # union_index cannot wrap here -- unlike the scratch, + # whose offsets fold through Int64. if cutlass.const_expr( self.width % self.tokens_per_lane == 0 ) and cutlass.dynamic_expr( @@ -517,7 +531,8 @@ def kernel( ) reduced_values[token_slot] = union_value subtile_first_token = first_token + token_subtile * self.subtile_token_tile - # Same per-request output domain as the single-CTA path. + # Same per-request output domain (and the same host-audit + # bound on the i32 union_index) as the single-CTA path. if cutlass.const_expr( self.width % self.tokens_per_lane == 0 ) and cutlass.dynamic_expr( diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index e476242a5deb..d7e6d3646af1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -376,6 +376,12 @@ def _settle_ties_and_pack_compaction_sources_kernel( half away, packing pre-settled ordinals read from ``output_indices`` -- the draft co-compaction flow, whose keep set is the target's and needs no settling. + + The increasing-ordinal emission and the tail placement are LOAD-BEARING + for the consumer: sparseKvCacheCompactOp.cpp's forward tiled in-place + copy requires increasing source ordinals with + ``destinationBases[request] + move <= source[move]`` (the SWA window + inequality is the SWA-family instance of the same invariant). """ request = tl.program_id(0) selection_domain = tl.program_id(1) From 78aab7d7408d0f1472bbd669f5dece2da2b97811 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 21:55:18 -0700 Subject: [PATCH 097/178] [None][refactor] Codegen-identical dedup batch + scaffolding kills (knife 17B + E2-E4) Ledger items #3,5,7,10,11,14,15,16,24,25,29,30,31,32 plus the unparked E2/E3/E4, one packed equivalence gate: - #3 halves_per_page is provably 1 under the geometry contract: the loop wrapper, page_half coords, and the dead prefetch arm die; the contract pins tokens_per_block <= CTA_M at construction. - #5/#10/#29/#30 single-source constants: STATS_ORIGIN_SLOTS/ STATS_SCRATCH_ELEMENTS/K_PLANES_PER_POOL_PAGE named in the fused file; N + STATS_FIELDS/STATS_MEAN/STATS_M2 exported to selection (private copies deleted, the shard-1 +4/+5 literals rewritten) and to the host (scratch alloc/audit/gather literals). - #7 the K2 gmem lane-tile idiom folds into one trace-time helper (4 sites); #11 the DSM peer-read offset derives from warp_max.layout via crd2idx with the *4 scale spelled Float32.width//8. - #14 Triton optional-pointer convention: settle-side pointer arithmetic moves inside HAS_SETTLE; the draft pack and the no-SWA settle tensors pass None (the has_swa=False ARM stays, per the verified caveat). - #15/#31 SETTLE_PACK_BLOCK/NUM_WARPS homed next to the kernel and imported by every launch site incl. tests; #16 the pack half reads selection_domain and composes the SWA mask (always-true trick gone). - #32 STD_EPSILON shared between the Triton and CuTe normalizers. - E2 the selection reduce+store tail is one maintenance point across the single-CTA/cluster arms (+19 LOC, honest cost for the fold). - E3 the _CUTE_SQRT_KWARG_MODE probe collapses to the unconditional inline-asm sqrt (digest-proven bit-identical at f0d756d). - E4 score-only CuTe entries compile only for the per-head modes; the union-fusion test builds its own score-only reference buffers. - cpp (inert without rebuild): #24 the uncompilable Layered arm and register-staging scaffolding die; #25 the never-launched destination-page pipeline kernel (~180 LOC) dies with the dispatch stub collapsed (parked on branch tr-parked-destination-page-kernel; equivalence comparison must be name-normalized, the mangled symbol changed). Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 292 +---- .../_torch/kv_cache_compression/compaction.py | 8 +- .../triattention/triattention.py | 43 +- .../triattention_cute_score_fused.py | 1095 ++++++++--------- .../triattention_cute_selection.py | 261 ++-- .../triattention/triattention_kernels.py | 46 +- .../_torch/kv_cache_compression/conftest.py | 20 +- .../test_triattention_cute_union_fusion.py | 57 +- .../test_triattention_fused_settle_pack.py | 9 +- 9 files changed, 824 insertions(+), 1007 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index e7b1ec5642a4..ceab9ac55c42 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1848,7 +1848,7 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) static_assert(std::is_same_v); static_assert(HeadDim == 64 || HeadDim == 128); // 128-token pages are the geometry the kernel was written for; 32-token - // pages cover this tree's production configuration (one tile == one page). + // pages cover the supported production configuration (one tile == one page). static_assert(TokensPerBlock == 32 || TokensPerBlock == 128); // 16B vectors per head: Dh64 -> 8 lanes (block 8x32 = 256 threads), // Dh128 -> 16 lanes (block 16x32 = 512 threads). @@ -1989,161 +1989,9 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) } } -//! Variant of the pipeline kernel that stages the destination page index in -//! shared memory once per tile instead of having every thread look it up. -//! Only valid when each request's destination base is 32-token tile aligned, -//! so a whole tile lands in one destination page (both supported page sizes -//! are multiples of the tile). Compiled but not yet dispatched: the bases -//! live on device, so the host cannot prove alignment (see the dispatch -//! helper below). -template -__global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) - * kSparseKvCompactTokensPerTile) void sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel(SparseKvCacheCompactBf16Params - params) -{ - static_assert(std::is_same_v); - static_assert(HeadDim == 64 || HeadDim == 128); - static_assert(TokensPerBlock == 32 || TokensPerBlock == 128); - constexpr int32_t kBufferCount = 2; - constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); - constexpr int32_t kTokensPerTile = kSparseKvCompactTokensPerTile; - constexpr int32_t kVectorsPerTile = kTokensPerTile * kVectorsPerHead; - constexpr int32_t kVectorsPerBuffer = 2 * kVectorsPerTile; - - int32_t const layerIdx = static_cast(blockIdx.x); - int32_t const kvHeadIdx = static_cast(blockIdx.y); - int32_t const batchIdx = static_cast(blockIdx.z); - int32_t const moveBegin = params.sourceOffsets[batchIdx]; - int32_t const moveEnd = params.sourceOffsets[batchIdx + 1]; - int32_t const moveCount = moveEnd - moveBegin; - if (moveCount <= 0) - { - return; - } - - int32_t const sourceLayer = params.sourceLayerIndices == nullptr ? layerIdx : params.sourceLayerIndices[layerIdx]; - int64_t const sourceMoveBase = static_cast(sourceLayer) * params.sourceLayerStride - + static_cast(kvHeadIdx) * params.sourceHeadStride + moveBegin; - int32_t const destinationBase = params.destinationBases[batchIdx]; - auto* const pool = reinterpret_cast(static_cast(params.poolPointers[layerIdx])); - int32_t const* const pageTable = params.pageTable + static_cast(batchIdx) * params.pageTableRequestStride; - - extern __shared__ uint4 sharedVectors[]; - auto* const sharedDestinationPages = reinterpret_cast(sharedVectors + kBufferCount * kVectorsPerBuffer); - int32_t const sharedVector - = static_cast(threadIdx.y) * kVectorsPerHead + static_cast(threadIdx.x); - int32_t currentBuffer = 0; - int32_t currentRequestMove = static_cast(threadIdx.y); - bool currentValid = currentRequestMove < moveCount; - int32_t currentSourceToken = currentValid ? params.sourceIndices[sourceMoveBase + currentRequestMove] : -1; - uint4* currentSharedK = sharedVectors; - uint4* currentSharedV = currentSharedK + kVectorsPerTile; - - // The dispatch guard tile-aligns the destination base, so every 32-token - // tile lies within one destination page. One producer stages the decoded - // page beside buffer 0; the existing prologue barrier publishes both - // products. - if (threadIdx.x == 0 && threadIdx.y == 0) - { - sharedDestinationPages[currentBuffer] = pageTable[destinationBase / TokensPerBlock] >> 1; - } - uint4 const* currentSourceKVector = nullptr; - uint4 const* currentSourceVVector = nullptr; - if (currentValid) - { - int32_t const sourcePage = pageTable[currentSourceToken / TokensPerBlock] >> 1; - auto* const sourcePageBase = pool + static_cast(sourcePage) * params.bytesPerPage; - auto const* const sourceK = reinterpret_cast(sourcePageBase); - auto const* const sourceV = reinterpret_cast(sourcePageBase + params.bytesPerKvHalf); - int32_t const localVector = (kvHeadIdx * TokensPerBlock + currentSourceToken % TokensPerBlock) * kVectorsPerHead - + static_cast(threadIdx.x); - currentSourceKVector = &sourceK[localVector]; - currentSourceVVector = &sourceV[localVector]; - } - uint32_t const currentSourceBytes = currentValid ? sizeof(uint4) : 0U; - compact_detail::copyAsync(¤tSharedK[sharedVector], currentSourceKVector, currentSourceBytes); - compact_detail::copyAsync(¤tSharedV[sharedVector], currentSourceVVector, currentSourceBytes); - compact_detail::commitGroup(); - compact_detail::waitGroup<0>(); - __syncthreads(); - - for (int32_t nextTileBegin = kTokensPerTile; nextTileBegin < moveCount; nextTileBegin += kTokensPerTile) - { - int32_t const nextRequestMove = nextTileBegin + static_cast(threadIdx.y); - bool const nextValid = nextRequestMove < moveCount; - int32_t const nextSourceToken = nextValid ? params.sourceIndices[sourceMoveBase + nextRequestMove] : -1; - int32_t const nextBuffer = (nextTileBegin / kTokensPerTile) & 1; - uint4* const nextSharedK = sharedVectors + nextBuffer * kVectorsPerBuffer; - uint4* const nextSharedV = nextSharedK + kVectorsPerTile; - - // Produce the page paired with the next ping-pong buffer. The loop's existing final barrier publishes it - // before the buffer becomes current, so destination staging does not add a CTA synchronization. - if (threadIdx.x == 0 && threadIdx.y == 0) - { - int32_t const nextDestinationToken = destinationBase + nextTileBegin; - sharedDestinationPages[nextBuffer] = pageTable[nextDestinationToken / TokensPerBlock] >> 1; - } - uint4 const* nextSourceKVector = nullptr; - uint4 const* nextSourceVVector = nullptr; - if (nextValid) - { - int32_t const sourcePage = pageTable[nextSourceToken / TokensPerBlock] >> 1; - auto* const sourcePageBase = pool + static_cast(sourcePage) * params.bytesPerPage; - auto const* const sourceK = reinterpret_cast(sourcePageBase); - auto const* const sourceV = reinterpret_cast(sourcePageBase + params.bytesPerKvHalf); - int32_t const localVector - = (kvHeadIdx * TokensPerBlock + nextSourceToken % TokensPerBlock) * kVectorsPerHead - + static_cast(threadIdx.x); - nextSourceKVector = &sourceK[localVector]; - nextSourceVVector = &sourceV[localVector]; - } - uint32_t const nextSourceBytes = nextValid ? sizeof(uint4) : 0U; - compact_detail::copyAsync(&nextSharedK[sharedVector], nextSourceKVector, nextSourceBytes); - compact_detail::copyAsync(&nextSharedV[sharedVector], nextSourceVVector, nextSourceBytes); - compact_detail::commitGroup(); - - // The compaction contract provides strictly increasing sources per request/head and - // dst(i) = destinationBase + i <= src(i). For current i and future j, i < j implies - // dst(i) < destinationBase + j <= src(j), so current stores cannot alias future prefetch sources. - int32_t const destinationToken = destinationBase + currentRequestMove; - if (currentValid && currentSourceToken != destinationToken) - { - int32_t const destinationPage = sharedDestinationPages[currentBuffer]; - auto* const destinationPageBase = pool + static_cast(destinationPage) * params.bytesPerPage; - auto* const destinationK = reinterpret_cast(destinationPageBase); - auto* const destinationV = reinterpret_cast(destinationPageBase + params.bytesPerKvHalf); - int32_t const localVector - = (kvHeadIdx * TokensPerBlock + destinationToken % TokensPerBlock) * kVectorsPerHead - + static_cast(threadIdx.x); - destinationK[localVector] = currentSharedK[sharedVector]; - destinationV[localVector] = currentSharedV[sharedVector]; - } - - compact_detail::waitGroup<0>(); - __syncthreads(); - currentBuffer = nextBuffer; - currentRequestMove = nextRequestMove; - currentValid = nextValid; - currentSourceToken = nextSourceToken; - currentSharedK = nextSharedK; - currentSharedV = nextSharedV; - } - - // The final tile and its destination page scalar already completed the existing wait and CTA barrier. - int32_t const destinationToken = destinationBase + currentRequestMove; - if (currentValid && currentSourceToken != destinationToken) - { - int32_t const destinationPage = sharedDestinationPages[currentBuffer]; - auto* const destinationPageBase = pool + static_cast(destinationPage) * params.bytesPerPage; - auto* const destinationK = reinterpret_cast(destinationPageBase); - auto* const destinationV = reinterpret_cast(destinationPageBase + params.bytesPerKvHalf); - int32_t const localVector = (kvHeadIdx * TokensPerBlock + destinationToken % TokensPerBlock) * kVectorsPerHead - + static_cast(threadIdx.x); - destinationK[localVector] = currentSharedK[sharedVector]; - destinationV[localVector] = currentSharedV[sharedVector]; - } -} - +// A destination-page-staging variant (needs a host-side proof that every +// destination base is tile-aligned) is parked on branch +// tr-parked-destination-page-kernel pending compaction.py alignment-flag plumbing. template void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) { @@ -2160,60 +2008,22 @@ void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactBf16Params con sparseKvCacheCompactV2Bf16PipelineKernel<<>>(params); } -template -void launchSparseKvCacheCompactV2Bf16DestinationPagePipeline( - SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) -{ - constexpr int32_t kDestinationPageBuffers = 2; - constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); - dim3 const block(kVectorsPerHead, kSparseKvCompactTokensPerTile); - dim3 const grid(params.numLayers, params.numKvHeads, params.batchSize); - // Same 16/32 KiB tile buffers as the plain pipeline plus two staged - // destination page indices; still far below the 48 KiB default. - size_t const sharedBytes = 4 * kSparseKvCompactTokensPerTile * kVectorsPerHead * sizeof(uint4) - + kDestinationPageBuffers * sizeof(int32_t); - sparseKvCacheCompactV2Bf16DestinationPagePipelineKernel - <<>>(params); -} - -template -void dispatchSparseKvCacheCompactGeometry(SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) -{ - // The destination-page variant requires every request's destination base - // to be 32-token tile aligned so each tile lands in one page, but the - // bases live on device where the host cannot check them. Until a host-side - // alignment flag is plumbed through compaction.py, always launch the plain - // pipeline. A plain `if` (not `if constexpr`) keeps the destination-page - // kernel instantiated and compiling. - constexpr bool kDestinationPageVariantEnabled = false; - if (kDestinationPageVariantEnabled) - { - launchSparseKvCacheCompactV2Bf16DestinationPagePipeline(params, stream); - return; - } - launchSparseKvCacheCompactV2Bf16Pipeline(params, stream); -} #endif // ENABLE_BF16 -template +template __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( QKVPreprocessingParams params) { // The number of 16B vectors per head size in the kv cache. constexpr int VECS_PER_HEAD = Dh * sizeof(TCache) / 16; static_assert(BLOCK_SIZE % VECS_PER_HEAD == 0, "Kernel block should be able to handle entire heads."); - // D64 and D128 map one complete K/V vector to each x lane, so registers can - // preserve the read-before-write value across the in-place ordering barrier. - constexpr bool use_register_staging = Layered && (Dh == 64 || Dh == 128) && sizeof(TCache) == 2; int const batch_idx = blockIdx.z; int const kv_head_idx = blockIdx.y; - // Head-plane stride of the packed indices for the non-layered layout; the - // layered layout carries its stride on the buffer (its index buffers may - // be wider than one round's move count). - [[maybe_unused]] int const total_num_sparse_kv_tokens = params.sparse_kv_offsets[params.batch_size]; + // Head-row stride of the packed indices. + int const total_num_sparse_kv_tokens = params.sparse_kv_offsets[params.batch_size]; int const sparse_start_idx = params.sparse_kv_offsets[batch_idx]; int const sparse_end_idx = params.sparse_kv_offsets[batch_idx + 1]; @@ -2228,51 +2038,27 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( for (int token_block_offset = 0; token_block_offset < num_sparse_tokens; token_block_offset += tokens_per_block) { - uint4 key_vector; - uint4 value_vector; int const sparse_token_offset = token_block_offset + threadIdx.y; if (sparse_token_offset < num_sparse_tokens) { int const global_sparse_idx = sparse_start_idx + sparse_token_offset; - int src_token_idx; - if constexpr (Layered) - { - src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( - params.sparse_kv_indices, kv_head_idx, global_sparse_idx); - } - else - { - int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; - src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; - } + int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; + int const src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; void* src_k_ptr = params.kv_cache_buffer.getKBlockPtr(batch_idx, src_token_idx); void* src_v_ptr = params.kv_cache_buffer.getVBlockPtr(batch_idx, src_token_idx); auto const src_k_block_ptr = reinterpret_cast(src_k_ptr); auto const src_v_block_ptr = reinterpret_cast(src_v_ptr); - if constexpr (use_register_staging) + for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) { - int const head_vec_idx = threadIdx.x; auto const src_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); auto const src_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - key_vector = src_k_block_ptr[src_k_vec_idx]; - value_vector = src_v_block_ptr[src_v_vec_idx]; - } - else - { - for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) - { - auto const src_k_vec_idx - = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - auto const src_v_vec_idx - = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_k_block_ptr[src_k_vec_idx]; - v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_v_block_ptr[src_v_vec_idx]; - } + k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_k_block_ptr[src_k_vec_idx]; + v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_v_block_ptr[src_v_vec_idx]; } } __syncthreads(); @@ -2280,20 +2066,9 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( if (sparse_token_offset < num_sparse_tokens) { int const global_sparse_idx = sparse_start_idx + sparse_token_offset; - int src_token_idx; - int dst_token_idx; - if constexpr (Layered) - { - src_token_idx = params.kv_cache_buffer.getSparseKvSourceToken( - params.sparse_kv_indices, kv_head_idx, global_sparse_idx); - dst_token_idx = params.kv_cache_buffer.getSparseKvDestinationToken(batch_idx, sparse_token_offset); - } - else - { - int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; - src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; - dst_token_idx = sparse_token_offset; - } + int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; + int const src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; + int const dst_token_idx = sparse_token_offset; if (src_token_idx != dst_token_idx) { @@ -2302,27 +2077,14 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( auto const dst_k_block_ptr = reinterpret_cast(dst_k_ptr); auto const dst_v_block_ptr = reinterpret_cast(dst_v_ptr); - if constexpr (use_register_staging) - { - int const head_vec_idx = threadIdx.x; - auto const dst_k_vec_idx - = params.kv_cache_buffer.getKVLocalIdx(dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - auto const dst_v_vec_idx - = params.kv_cache_buffer.getKVLocalIdx(dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - dst_k_block_ptr[dst_k_vec_idx] = key_vector; - dst_v_block_ptr[dst_v_vec_idx] = value_vector; - } - else + for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) { - for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) - { - auto const dst_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx( - dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - auto const dst_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx( - dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - dst_k_block_ptr[dst_k_vec_idx] = k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; - dst_v_block_ptr[dst_v_vec_idx] = v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; - } + auto const dst_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx( + dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + auto const dst_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx( + dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + dst_k_block_ptr[dst_k_vec_idx] = k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; + dst_v_block_ptr[dst_v_vec_idx] = v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; } } } @@ -2344,7 +2106,7 @@ void kernelSparseDispatchHeadSize(QKVPreprocessingParams param // grid.x is always 1 to avoid data races dim3 grid(1, params.kv_head_num, params.batch_size); - updateSparseKvCacheAfterFmha + updateSparseKvCacheAfterFmha <<>>(params); } @@ -2404,19 +2166,19 @@ void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const fastParams.bytesPerPage = 2 * fastParams.bytesPerKvHalf; if (headDim == 64 && tokensPerBlock == 32) { - dispatchSparseKvCacheCompactGeometry(fastParams, stream); + launchSparseKvCacheCompactV2Bf16Pipeline(fastParams, stream); } else if (headDim == 64 && tokensPerBlock == 128) { - dispatchSparseKvCacheCompactGeometry(fastParams, stream); + launchSparseKvCacheCompactV2Bf16Pipeline(fastParams, stream); } else if (headDim == 128 && tokensPerBlock == 32) { - dispatchSparseKvCacheCompactGeometry(fastParams, stream); + launchSparseKvCacheCompactV2Bf16Pipeline(fastParams, stream); } else { - dispatchSparseKvCacheCompactGeometry(fastParams, stream); + launchSparseKvCacheCompactV2Bf16Pipeline(fastParams, stream); } return; } diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index e752c97d47b7..3d159366a759 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -244,8 +244,8 @@ def init_compaction_buffers( selection_rows = num_kv_heads # Selection rows carry decode-only kept ordinals (already absolute), so # the rectangle is prompt-length independent. HAS_SWA specializes the SWA - # loads and stores away, so without SWA layers the dense buffers stand in - # for the compiled-away SWA pointer arguments. + # loads and stores away, so without SWA layers the SWA pointer arguments + # are None (the Triton optional-pointer convention). has_swa = swa_move_indices is not None swa_total = int(swa_move_indices.shape[-1]) if has_swa else 0 # Widest per-request move count any staged offsets may express; the @@ -257,8 +257,8 @@ def init_compaction_buffers( valid_sequence_lengths, dense_move_offsets, dense_move_indices, - swa_move_offsets if has_swa else dense_move_offsets, - swa_move_indices if has_swa else dense_move_indices, + swa_move_offsets if has_swa else None, + swa_move_indices if has_swa else None, ) settle_pack_shape = dict( DENSE_TOTAL=int(dense_move_indices.shape[-1]), diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index a095a6f5b611..bbdf2b643777 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -78,6 +78,8 @@ from ..compaction import init_compaction_buffers from .triattention_kernels import ( + SETTLE_PACK_BLOCK, + SETTLE_PACK_NUM_WARPS, _settle_ties_and_pack_compaction_sources_kernel, build_mean_phase_table, gather_mean_phases, @@ -210,6 +212,7 @@ def init_eviction_buffers( layer pool: the compiled kernels encode immutable TMA descriptors from their raw device addresses, so the pools must stay alive and stay put. """ + from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS from .triattention_cute_score_fused import TriAttentionCuteScoreRunner device = layer_pools[page_representatives[0]].device @@ -353,15 +356,19 @@ def init_eviction_buffers( # every downstream i32 offset (including the seg_out_offset cast below) # is bounded by it. Wraparound would be a silent wild read, not a clean # error, hence the loud audit. - if (8 - 1) * max_segments * seq_len >= 2**31: + if (PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len >= 2**31: raise ValueError( - f"score bucket overflows the 32-bit score plane: {(8 - 1) * max_segments * seq_len}" + "score bucket overflows the 32-bit score plane: " + f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len}" ) # The kernel scores each request's window into a head-major scratch; # all buffers below are persistent because the compiled kernels capture # their device pointers. + bufs.padded_head_columns = PADDED_HEAD_COLUMNS bufs.cute_scratch = torch.empty( - bufs.num_kv_heads * 8 * max_segments * seq_len, dtype=torch.float32, device=device + bufs.num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * seq_len, + dtype=torch.float32, + device=device, ) # int32 is safe here: covered by the 2^31 score-plane audit above. seg_out_offset = (torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len).to( @@ -730,8 +737,8 @@ def settle_top_tokens(bufs: SimpleNamespace) -> None: SELECTION_ROWS=bufs.selection_rows_per_request, **bufs.settle_pack_shape, HAS_SETTLE=True, - BLOCK=256, - num_warps=4, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, ) @@ -780,9 +787,10 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: # scratch data masked by ``valid_widths``. group_size = bufs.num_q_heads // bufs.num_kv_heads num_segments = request_count * bufs.num_layers + pad = bufs.padded_head_columns source = ( - bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] - .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ + bufs.cute_scratch[: bufs.num_kv_heads * pad * num_segments * bufs.bucket_seq_len] + .view(bufs.num_kv_heads, pad, request_count, bufs.num_layers, bufs.bucket_seq_len)[ :, :group_size ] .permute(2, 3, 0, 1, 4) @@ -827,19 +835,20 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: # One more pack launch broadcasts the target keep set over # the draft KV heads and appends the draft's own tail # ordinals. HAS_SETTLE=False compiles the settle half away - # (the ordinals arrive pre-settled), so any well-formed - # tensor stands in for the settle-side pointer arguments. + # (the ordinals arrive pre-settled), so the settle-side + # pointer arguments are None; the pack half reads the + # settled ordinals through output_indices. _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( + None, + None, + None, + None, bufs.keep, bufs.valid_seq_lens_device, bufs.draft_pack["offsets"], - bufs.keep, - bufs.keep, - bufs.valid_seq_lens_device, - bufs.draft_pack["offsets"], - bufs.draft_pack["indices"], - bufs.draft_pack["offsets"], bufs.draft_pack["indices"], + None, + None, WIDTH=bufs.keep_count, KEEP_COUNT=bufs.keep_count, SELECTION_ROWS=1, @@ -852,8 +861,8 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: PER_LAYER=False, HAS_SWA=False, HAS_SETTLE=False, - BLOCK=256, - num_warps=4, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, ) for group in family["groups"]: torch.ops.trtllm.sparse_kv_cache_compact_layers( 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 index 53ba4cc96ce6..326f62a527cd 100644 --- 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 @@ -26,35 +26,14 @@ from cutlass.cutlass_dsl import T, dsl_user_op -def _cute_sqrt_keyword_mode() -> str: - """Probe which fast-sqrt spelling this CuTe DSL's ``cute.math.sqrt`` takes. - - The approximate-sqrt control was renamed across DSL releases: some expose - ``approx``/``ftz`` keywords, cutlass 4.5 exposes a single ``fastmath`` - flag, and older releases expose a plain one-argument ``sqrt``. Passing an - unknown keyword raises TypeError at trace time (inside ``cute.compile``, - where it cannot be caught), so the capability is probed once at import - time via signature inspection and folded into a trace-time constant. - """ - import inspect - - try: - parameters = inspect.signature(cute.math.sqrt).parameters - except (TypeError, ValueError): - return "plain" - if "approx" in parameters and "ftz" in parameters: - return "approx_ftz" - if "fastmath" in parameters: - return "fastmath" - return "plain" - - -_CUTE_SQRT_KWARG_MODE = _cute_sqrt_keyword_mode() - - @dsl_user_op def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: - """Emit the approximate FTZ square root missing from CuTe DSL 4.5.""" + """Emit the approximate FTZ square root as inline PTX. + + ``cute.math.sqrt``'s fast-sqrt keyword spelling varies across CuTe DSL + releases; the inline-asm form is release-independent and matches the + ``sqrt.approx.ftz.f32`` the fused score path has always used. + """ return cutlass.Float32( llvm.inline_asm( T.f32(), @@ -77,6 +56,19 @@ def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Fl THREADS = 256 EPILOGUE_THREADS = 128 RAW_PAGE_BUFFERS = 2 +# partial_stats is a flat [stats_row, page_shard, {count, mean, m2}] record +# array (stats_row = segment * num_q_heads + q_head); the selection finalizer +# imports these field constants. +STATS_FIELDS = 3 +STATS_MEAN = 1 +STATS_M2 = 2 +# Stats smem scratch: N first-tile score origins, then one (sum, square-sum) +# pair per (epilogue warp, padded head column). +STATS_ORIGIN_SLOTS = N +STATS_SCRATCH_ELEMENTS = STATS_ORIGIN_SLOTS + (EPILOGUE_THREADS // 32) * N * 2 +# HND K pools interleave K/V planes: 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 @@ -177,6 +169,11 @@ def __init__( ) 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: + # The schedule assumes one page never spans multiple compute + # tiles (the retired page_half loop generalized this; both + # supported page sizes make it a single iteration). + 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: @@ -200,8 +197,7 @@ def __init__( self.box_tokens = min(CTA_M, tokens_per_block) self.fragments_per_phase = CTA_M // self.box_tokens self.pages_per_tile = self.fragments_per_phase - self.halves_per_page = max(1, tokens_per_block // CTA_M) - self.tile_tokens = self.halves_per_page * CTA_M + self.tile_tokens = CTA_M self.max_tiles = (seq_len + self.tile_tokens - 1) // self.tile_tokens # Producer staging constants baked into the generated code. @@ -352,7 +348,7 @@ def __call__( 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 = 72 * int(self.write_partial_stats) + stats_scratch_elements = STATS_SCRATCH_ELEMENTS * int(self.write_partial_stats) @cute.struct class SharedStorage: @@ -506,7 +502,7 @@ def kernel( swizzle=magnitude_fp16_b_smem_layout.inner, ) if cutlass.const_expr(self.write_partial_stats): - sStats = storage.sStats.get_tensor(cute.make_layout(72)) + 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, @@ -632,7 +628,8 @@ def kernel( # The staged K-plane entries encode physical_page * # kv_factor (2); decode to the pool page index here. producer_prefetched_page_id_lane0 = ( - cutlass.Int32(page_ids[page_off + tile_index * self.pages_per_tile]) // 2 + cutlass.Int32(page_ids[page_off + tile_index * self.pages_per_tile]) + // K_PLANES_PER_POOL_PAGE ) if cutlass.const_expr(self.fragments_per_phase > 1): for fragment in cutlass.range_constexpr(1, self.pages_per_tile): @@ -649,7 +646,7 @@ def kernel( page_off + tile_index * self.pages_per_tile + fragment ] ) - // 2 + // 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) @@ -855,512 +852,476 @@ def kernel( producer_prefetched_page_ids_lane0[fragment], 0, ) - for page_half in cutlass.range_constexpr(self.halves_per_page): - 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 - # (CTA_M x num_freqs bf16, raw_tma_copy_bytes). - # Every producer-warp lane participates in the - # PipelineTmaAsync barrier election. - raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = physical_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - page_half, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - 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) - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = physical_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - page_half, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - raw_tma_producer_state.advance() - - if cutlass.const_expr(page_half == self.halves_per_page - 1): - 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 + self.tile_tokens * self.page_shards - ) - next_pages_processed = tiles_processed + 1 - if ( - next_tile_start_token < valid_seq_len - and next_pages_processed < self.max_tiles - ): - next_page_id_lane0 = ( - cutlass.Int32( - page_ids[ - page_off - + (tile_index + self.page_shards) * self.pages_per_tile - ] - ) - // 2 - ) - if cutlass.const_expr(self.fragments_per_phase > 1): - for fragment in cutlass.range_constexpr(1, self.pages_per_tile): - # Same tail-tile clamp as the initial - # prefetch: 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 - < valid_seq_len - ): - next_fragment_page_id = ( - cutlass.Int32( - page_ids[ - page_off - + (tile_index + self.page_shards) - * self.pages_per_tile - + fragment - ] - ) - // 2 - ) - 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)] + 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: - 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, + # Phase 0 fills the packed real-band stage view + # (CTA_M x num_freqs bf16, raw_tma_copy_bytes). + # Every producer-warp lane participates in the + # PipelineTmaAsync barrier election. + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = physical_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, ) - raw_bf16_tiled_mma.set( - tcgen05.Field.ACCUMULATE, - True, + 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) + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = physical_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, ) - raw_tma_pipeline.consumer_wait(raw_tma_consumer_state) - if cutlass.const_expr(self.write_partial_stats): - if warp_idx == self.producer_warp_id: - if cutlass.const_expr(page_half + 1 < self.halves_per_page): - prefetched_physical_page = physical_page - prefetch_next_raw = True - prefetched_page_half = page_half + 1 - else: - next_tile_start_token = ( - tile_start_token + self.tile_tokens * self.page_shards - ) - next_pages_processed = tiles_processed + 1 - prefetch_next_raw = ( - next_tile_start_token < valid_seq_len - 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.pages_per_tile): - prefetched_page_fragments[fragment] = cute.arch.shuffle_sync( - producer_prefetched_page_ids_lane0[fragment], - 0, - ) - prefetched_page_half = 0 - if cutlass.dynamic_expr(prefetch_next_raw): - 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): - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - prefetched_page_half, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - else: - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - prefetched_page_half, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real_next[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - raw_tma_producer_state.advance() - raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - if cutlass.dynamic_expr(next_raw_page_buffer == 0): - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - prefetched_page_half, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - else: - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - prefetched_page_half, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag_next[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) - raw_tma_producer_state.advance() - # Each of the 32 lanes 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 from the - # raw-K shared buffer before consuming any of them. - for token_base in cutlass.range( - 0, - CTA_M // (THREADS // 32), - self.prefetch_depth, - unroll_full=False, + 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 + self.tile_tokens * self.page_shards + next_pages_processed = tiles_processed + 1 + if ( + next_tile_start_token < valid_seq_len + and next_pages_processed < self.max_tiles ): - 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, - ) + next_page_id_lane0 = ( + cutlass.Int32( + page_ids[ + page_off + (tile_index + self.page_shards) * self.pages_per_tile ] ) - staged_imag[prefetch_index] = cutlass.Float32( - cpasync_raw_k_0[ - ( - (token, frequency % 16), - 0, - frequency // 16, - raw_imag_stage, + // K_PLANES_PER_POOL_PAGE + ) + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + # Same tail-tile clamp as the initial + # prefetch: 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 + < valid_seq_len + ): + next_fragment_page_id = ( + cutlass.Int32( + page_ids[ + page_off + + (tile_index + self.page_shards) + * self.pages_per_tile + + fragment + ] + ) + // K_PLANES_PER_POOL_PAGE ) - ] - ) - - 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 - if cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "approx_ftz"): - magnitude = cute.math.sqrt(norm2, approx=True, ftz=True) - elif cutlass.const_expr(_CUTE_SQRT_KWARG_MODE == "fastmath"): - magnitude = _sqrt_approx_ftz(norm2) - else: - # Plain IEEE sqrt (more accurate than approx); - # the equivalence tolerance covers the difference. - magnitude = cute.math.sqrt(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)] + producer_prefetched_page_ids_lane0[fragment] = next_fragment_page_id + producer_prefetched_page_id_lane0 = next_page_id_lane0 - if warp_idx == self.producer_warp_id: - # Finish B0-imag, then issue B1-real and B1-imag. + # 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, ) - 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 accumulation: - # |K|*coeff = A0*B0 + A0*B1 + A1*B0 (the A1*B1 term is - # below fp32 accumulation resolution and dropped). - magnitude_lo_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 + self.tile_tokens * self.page_shards + next_pages_processed = tiles_processed + 1 + prefetch_next_raw = ( + next_tile_start_token < valid_seq_len + and next_pages_processed < self.max_tiles + ) + prefetched_physical_page = cute.arch.shuffle_sync( + producer_prefetched_page_id_lane0, + 0, ) - 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, + if cutlass.const_expr(self.fragments_per_phase > 1): + for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + prefetched_page_fragments[fragment] = cute.arch.shuffle_sync( + producer_prefetched_page_ids_lane0[fragment], + 0, + ) + if cutlass.dynamic_expr(prefetch_next_raw): + 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): + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + else: + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 0, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_real_next[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + raw_tma_pipeline.producer_acquire(raw_tma_producer_state) + if cutlass.dynamic_expr(next_raw_page_buffer == 0): + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + else: + for fragment in cutlass.range_constexpr(self.fragments_per_phase): + fragment_page = prefetched_physical_page + if cutlass.const_expr(fragment > 0): + fragment_page = prefetched_page_fragments[fragment] + cute.copy( + raw_tma_atom, + raw_tma_global_partition[ + ( + None, + 1, + 0, + (kv_head, fragment_page), + ) + ], + raw_tma_shared_partition_imag_next[fragment], + tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( + raw_tma_producer_state + ), + tma_desc_ptr=raw_tma_descriptor_ptr, + ) + raw_tma_producer_state.advance() + # Each of the 32 lanes 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 from the + # raw-K shared buffer 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, + ) + ] ) - 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, + staged_imag[prefetch_index] = cutlass.Float32( + cpasync_raw_k_0[ + ( + (token, frequency % 16), + 0, + frequency // 16, + raw_imag_stage, + ) + ] ) - 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, + + 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) ) - acc_pipeline.producer_commit(acc_producer_state) - acc_producer_state.advance() - if tiles_processed == 0: - if cutlass.const_expr(page_half == 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): - # The alternate raw-page buffer was filled while this - # page's UMMA completed. Release only the current imag - # phase after all of its asynchronous consumers finish. - raw_tma_pipeline.consumer_release(raw_tma_consumer_state) - raw_tma_consumer_state.advance() - # Every term multiplying sum_seq must stay 64-bit: with - # request*layer segments of seq_len columns the score-plane - # stride alone can exceed 2^31. The scratch head axis is - # padded to the MMA tile N=8 per KV head (group-4 columns - # 4..7 land in padded score planes holding zero scores). The - # host audit in triattention.init_eviction_buffers bounds the - # 32-bit head-axis fold (kv_head * N <= 7 planes) here. - output_offset = ( - cutlass.Int64(kv_head * N) * sum_seq - + out_base - + tile_start_token - + page_half * CTA_M - ) - page_output = cute.make_tensor( - output.iterator + output_offset, - cute.make_layout( - (CTA_M, N, 1), - stride=( - 1, - sum_seq, - N * sum_seq, - ), - ), - ) - 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 - tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epilogue_tidx, tCtAcc, tCgC, self.epi_tile, False + 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, ) - simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( - epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile + 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 accumulation: + # |K|*coeff = A0*B0 + A0*B1 + A1*B0 (the A1*B1 term is + # below fp32 accumulation resolution and dropped). + magnitude_lo_tiled_mma.set( + tcgen05.Field.ACCUMULATE, + True, ) - 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])): + 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): + # The alternate raw-page buffer was filled while this + # page's UMMA completed. Release only the current imag + # phase after all of its asynchronous consumers finish. + raw_tma_pipeline.consumer_release(raw_tma_consumer_state) + raw_tma_consumer_state.advance() + # Every term multiplying sum_seq must stay 64-bit: with + # request*layer segments of seq_len columns the score-plane + # stride alone can exceed 2^31. The scratch head axis is + # padded to the MMA tile N=8 per KV head (group-4 columns + # 4..7 land in padded score planes holding zero scores). The + # host audit in triattention.init_eviction_buffers bounds the + # 32-bit head-axis fold (kv_head * N <= 7 planes) here. + output_offset = cutlass.Int64(kv_head * N) * sum_seq + out_base + tile_start_token + page_output = cute.make_tensor( + output.iterator + output_offset, + cute.make_layout( + (CTA_M, N, 1), + stride=( + 1, + sum_seq, + N * sum_seq, + ), + ), + ) + 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 + tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( + epilogue_tidx, tCtAcc, tCgC, self.epi_tile, False + ) + simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( + epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile + ) + 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)) + # The window start is a per-request runtime value, so + # the tile-interior fast path is a dynamic predicate; + # 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.seq_len + ): cute.copy( - tiled_copy_t2r, - tTR_tAcc[(None, None, None, subtile_idx)], - tTR_rAcc, + simt_atom, + tTR_rC, + tTR_gC[(None, None, None, subtile_idx)], ) - tTR_rC.store(tTR_rAcc.load().to(self.c_dtype)) - # The window start is a per-request runtime value, so - # the tile-interior fast path is a dynamic predicate; - # only the straddling first tile takes the per-token - # branch. + else: + output_token = tile_start_token + epilogue_tidx if cutlass.dynamic_expr( - tile_start_token >= score_start - and tile_start_token + CTA_M <= self.seq_len + output_token >= score_start and output_token < self.seq_len ): 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.seq_len - ): - 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 cutlass.const_expr(page_half == 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: - if cutlass.const_expr(page_half == 0): - for stats_head in cutlass.range_constexpr(N): - 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 < valid_seq_len - ): - for stats_head in cutlass.range_constexpr(N): - 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 - ) + 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(N): + 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 < valid_seq_len + ): + for stats_head in cutlass.range_constexpr(N): + 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 += self.tile_tokens * self.page_shards tiles_processed += 1 @@ -1377,7 +1338,7 @@ def kernel( stats_square_sum, stats_offset ) if lane_idx == 0 and warp_idx < EPILOGUE_THREADS // 32: - stats_scratch_base = 8 + (warp_idx * N + stats_head) * 2 + stats_scratch_base = STATS_ORIGIN_SLOTS + (warp_idx * N + stats_head) * 2 sStats[stats_scratch_base] = stats_sum sStats[stats_scratch_base + 1] = stats_square_sum cute.arch.barrier() @@ -1390,7 +1351,7 @@ def kernel( 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 = 8 + (stats_warp * N + lane_idx) * 2 + stats_scratch_base = STATS_ORIGIN_SLOTS + (stats_warp * N + 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 * self.tile_tokens @@ -1418,10 +1379,10 @@ def kernel( cutlass.Float32(0.0), ) stats_row = task * self.group_size + lane_idx - stats_base = (stats_row * self.page_shards + page_shard) * 3 + stats_base = (stats_row * self.page_shards + page_shard) * STATS_FIELDS partial_stats[stats_base] = stats_count - partial_stats[stats_base + 1] = stats_mean - partial_stats[stats_base + 2] = stats_m2 + 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) @@ -1450,7 +1411,11 @@ def _encode_tma_descriptors( 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) != (2, tokens_per_block, 2 * num_freqs): + 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" @@ -1567,8 +1532,10 @@ def __init__( self.num_kv_heads = int(num_kv_heads) self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) self.enable_partial_stats = bool(enable_partial_stats) + # One [stats_row, page_shard, {count, mean, m2}] record array (see the + # STATS_FIELDS constants above). partial_stats_elements = ( - max_requests * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * 3 + max_requests * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * STATS_FIELDS if self.enable_partial_stats else 1 ) @@ -1654,44 +1621,47 @@ def __init__( if max_requests > 1: variants.append((max_requests, 2)) for request_count, page_shards in variants: - cache_key = ( - "triattention_cute_score", - static_geometry, - tensor_specs, - request_count, - page_shards, - ) - with _COMPILE_LOCK: - compiled = _COMPILED_KERNELS.get(cache_key) - if compiled is None: - kernel = _TriAttentionScoreKernel( - num_layers=num_layers, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - pool_shape=tuple( - int(value) for value in layer_pools[layer_indices[0]].shape - ), - pool_strides=tuple( - int(value) for value in layer_pools[layer_indices[0]].stride() - ), - pool_dtype=cutlass.BFloat16, - page_shards=page_shards, - ) - stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) - compiled = cute.compile( - kernel, - *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), - *self._cute_tail, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled - self._compiled[request_count] = compiled + if not self.enable_partial_stats: + # Per-head modes launch the score-only entry; the union + # runner compiles ONLY its fused stats+union pipeline below. + cache_key = ( + "triattention_cute_score", + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + kernel = _TriAttentionScoreKernel( + num_layers=num_layers, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + pool_shape=tuple( + int(value) for value in layer_pools[layer_indices[0]].shape + ), + pool_strides=tuple( + int(value) for value in layer_pools[layer_indices[0]].stride() + ), + pool_dtype=cutlass.BFloat16, + page_shards=page_shards, + ) + stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) + compiled = cute.compile( + kernel, + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled + self._compiled[request_count] = compiled self._page_shards[request_count] = page_shards if self.enable_partial_stats: stats_cache_key = ( @@ -1735,8 +1705,8 @@ def __init__( self._compiled_stats[request_count] = compiled_stats if max_requests > 1: - small_score = self._compiled[1] - large_score = self._compiled[max_requests] + small_score = self._compiled.get(1) + large_score = self._compiled.get(max_requests) small_stats = self._compiled_stats.get(1) large_stats = self._compiled_stats.get(max_requests) for request_count in range(1, max_requests + 1): @@ -1745,9 +1715,10 @@ def __init__( # (2 * sm_count CTAs); larger cohorts already fill the GPU. two_shard_ctas = request_count * num_layers * num_kv_heads * 2 use_extra_score_shard = two_shard_ctas < 2 * self.sm_count - self._compiled[request_count] = ( - small_score if use_extra_score_shard else large_score - ) + if not self.enable_partial_stats: + self._compiled[request_count] = ( + small_score if use_extra_score_shard else large_score + ) self._page_shards[request_count] = ( SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 ) 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 index bb07036ee6c2..7423ac70b205 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -25,6 +25,15 @@ from cutlass.cute.typing import Pointer as CutePointer from cutlass.cutlass_dsl import T, dsl_user_op +from .triattention_cute_score_fused import STATS_FIELDS as _STATS_FIELDS +from .triattention_cute_score_fused import STATS_M2, STATS_MEAN + +# Single-sourced constants: the fused score file owns the padded-head tile +# (N) and the partial-stats record layout it writes; the Triton kernels +# module owns the z-normalization epsilon shared by both selection paths. +from .triattention_cute_score_fused import N as _PADDED_HEAD_COLUMNS +from .triattention_kernels import STD_EPSILON as _STD_EPSILON + _REDUCE_THREADS = 256 _WARP_SIZE = 32 _REDUCE_WARPS = _REDUCE_THREADS // _WARP_SIZE @@ -34,11 +43,6 @@ _SMALL_TOKEN_SUBTILES = 1 _MAX_ROW_CLUSTER_CTAS = 4 _SMALL_TILE_RESIDENT_CTAS_PER_SM = 6 -_STATS_FIELDS = 3 -_STD_EPSILON = 1.0e-6 -# The fused score kernel pads each KV head's group of score planes up to the -# minimum tcgen05 MMA tile (GQA groups below 8 ride zero-padded columns). -_PADDED_HEAD_COLUMNS = 8 def _select_normalize_union_config( @@ -118,6 +122,25 @@ def _ld_cluster_f32(mapped_addr): return _ld_shared_cluster_f32(mapped_addr) +def _gmem_lane_tile(iterator, flat_index, tokens_per_lane, assumed_align): + """One lane's gmem tile of ``tokens_per_lane`` fp32 values. + + Folds the (possibly 64-bit) flat index into the pointer BEFORE any + element access, so nothing routes through the DSL's 32-bit dynamic + coordinate; ``assumed_align`` is the only per-site difference between + the vectorized and fallback arms. + """ + 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.""" @@ -199,6 +222,81 @@ def __call__( 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, + valid_width: cutlass.Int32, + first_token: cutlass.Int32, + lane_idx: cutlass.Int32, + from_cluster_peers: cutlass.Constexpr, + ): + """Final peer reduce + union-row store, shared by both arms. + + The peer-reduction source is the ONLY difference between the + single-CTA and cluster arms (other warps' smem rows vs DSM cluster + peers), selected at trace time by ``from_cluster_peers``. + """ + 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 object the smem + # tensor was built with (no hand-derived strides). + 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_cluster(warp_max_ptr, cutlass.Int32(peer_rank)) + union_value = cute.arch.fmax( + union_value, + _ld_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 + # The output row covers this request's own window, + # [0, valid - start); the straddling subtile falls back to the + # per-token stores. union_scores spans request_count * width + # < 2^31 by the host score-plane audit + # (triattention.init_eviction_buffers), so the i32 union_index + # cannot wrap here -- unlike the scratch, whose offsets fold + # through Int64. + if cutlass.const_expr(self.width % self.tokens_per_lane == 0) and cutlass.dynamic_expr( + subtile_first_token + self.tokens_per_lane <= valid_width + ): + union_index = request_idx * self.width + 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: + for token_slot in cutlass.range_constexpr(self.tokens_per_lane): + token = subtile_first_token + token_slot + if cutlass.dynamic_expr(token < valid_width): + union_scores[request_idx * self.width + token] = reduced_values[token_slot] + @cute.kernel def kernel( self, @@ -320,10 +418,10 @@ def kernel( 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 + 1] - m2_0 = partial_stats[stats_base + 2] - mean_1 = partial_stats[stats_base + 4] - m2_1 = partial_stats[stats_base + 5] + 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 @@ -332,8 +430,8 @@ def kernel( 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 + 1] - shard_m2 = partial_stats[stats_base + 2] + 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] @@ -362,14 +460,11 @@ def kernel( score_start % self.tokens_per_lane == 0 and subtile_first_token + self.tokens_per_lane <= valid_width ): - score_tile = cute.make_tensor( - cute.make_ptr( - cutlass.Float32, - (scores.iterator + score_index).toint(), - AddressSpace.gmem, - assumed_align=self.tokens_per_lane * 4, - ), - cute.make_layout(self.tokens_per_lane), + score_tile = _gmem_lane_tile( + scores.iterator, + score_index, + self.tokens_per_lane, + self.tokens_per_lane * 4, ) cute.copy( score_copy_atom, @@ -386,14 +481,8 @@ def kernel( # window start is not lane-aligned (any pinned prompt # length not divisible by ``tokens_per_lane``), so real # serve cohorts hit it on every eviction round. - score_tail = cute.make_tensor( - cute.make_ptr( - cutlass.Float32, - (scores.iterator + score_index).toint(), - AddressSpace.gmem, - assumed_align=4, - ), - cute.make_layout(self.tokens_per_lane), + 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 @@ -449,51 +538,18 @@ def kernel( cute.arch.sync_threads() if cutlass.const_expr(self.row_cluster_ctas == 1): if warp_idx == 0: - 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): - 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 - # The output row covers this request's own window, - # [0, valid - start); the straddling subtile falls back - # to the per-token stores. union_scores spans - # request_count * width < 2^31 by the host score-plane - # audit (triattention.init_eviction_buffers), so the i32 - # union_index cannot wrap here -- unlike the scratch, - # whose offsets fold through Int64. - if cutlass.const_expr( - self.width % self.tokens_per_lane == 0 - ) and cutlass.dynamic_expr( - subtile_first_token + self.tokens_per_lane <= valid_width - ): - union_index = request_idx * self.width + subtile_first_token - union_tile = cute.make_tensor( - cute.make_ptr( - cutlass.Float32, - (union_scores.iterator + union_index).toint(), - AddressSpace.gmem, - assumed_align=self.tokens_per_lane * 4, - ), - cute.make_layout(self.tokens_per_lane), - ) - cute.copy( - score_copy_atom, - cute.coalesce(reduced_values), - cute.coalesce(union_tile), - ) - else: - for token_slot in cutlass.range_constexpr(self.tokens_per_lane): - token = subtile_first_token + token_slot - if cutlass.dynamic_expr(token < valid_width): - union_scores[request_idx * self.width + token] = reduced_values[ - token_slot - ] + self._reduce_and_store_union_rows( + union_scores, + union_values, + warp_max, + warp_max_ptr, + score_copy_atom, + request_idx, + valid_width, + first_token, + lane_idx, + False, + ) else: # Warp 0 reduces its CTA's row partition, then CTA 0 combines the # cluster's partial maxima through distributed shared memory. @@ -511,54 +567,17 @@ def kernel( cute.arch.cluster_arrive_relaxed() cute.arch.cluster_wait() if cta_rank == 0 and warp_idx == 0: - 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): - union_value = warp_max[(0, token_subtile, token_slot, lane_idx)] - shared_offset = ( - token_subtile * self.tokens_per_lane * _WARP_SIZE - + token_slot * _WARP_SIZE - + lane_idx - ) - for peer_rank in cutlass.range_constexpr(1, self.row_cluster_ctas): - remote_addr = _mapa_cluster( - warp_max_ptr, - cutlass.Int32(peer_rank), - ) - union_value = cute.arch.fmax( - union_value, - _ld_cluster_f32(remote_addr + shared_offset * 4), - ) - reduced_values[token_slot] = union_value - subtile_first_token = first_token + token_subtile * self.subtile_token_tile - # Same per-request output domain (and the same host-audit - # bound on the i32 union_index) as the single-CTA path. - if cutlass.const_expr( - self.width % self.tokens_per_lane == 0 - ) and cutlass.dynamic_expr( - subtile_first_token + self.tokens_per_lane <= valid_width - ): - union_index = request_idx * self.width + subtile_first_token - union_tile = cute.make_tensor( - cute.make_ptr( - cutlass.Float32, - (union_scores.iterator + union_index).toint(), - AddressSpace.gmem, - assumed_align=self.tokens_per_lane * 4, - ), - cute.make_layout(self.tokens_per_lane), - ) - cute.copy( - score_copy_atom, - cute.coalesce(reduced_values), - cute.coalesce(union_tile), - ) - else: - for token_slot in cutlass.range_constexpr(self.tokens_per_lane): - token = subtile_first_token + token_slot - if cutlass.dynamic_expr(token < valid_width): - union_scores[request_idx * self.width + token] = reduced_values[ - token_slot - ] + self._reduce_and_store_union_rows( + union_scores, + union_values, + warp_max, + warp_max_ptr, + score_copy_atom, + request_idx, + valid_width, + 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 index d7e6d3646af1..33602d57bdf8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -36,6 +36,13 @@ # so a larger table would silently degrade every downstream phase. _MEAN_PHASE_MAX_ROWS = 1 << 24 +# Score z-normalization epsilon, shared with the CuTe union pipeline +# (triattention_cute_selection.py imports it): both kernels bake the same +# literal, keeping the eviction modes' normalization consistent. Plain float +# (the CuTe DSL traces it directly); the Triton stats kernel receives it as +# an explicit constexpr parameter. +STD_EPSILON = 1e-6 + @triton.jit def _gather_mean_phase_kernel( @@ -172,6 +179,7 @@ def _score_row_stats_kernel( ROWS: tl.constexpr, WIDTH: tl.constexpr, BLOCK: tl.constexpr, + EPSILON: tl.constexpr, ): """Compute one valid-prefix mean and inverse standard deviation per score row.""" flat_row = tl.program_id(0) @@ -195,7 +203,7 @@ def _score_row_stats_kernel( square_sum += tl.sum(centered * centered, axis=0) std = tl.sqrt(square_sum / valid_width) tl.store(row_mean + flat_row, mean) - tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, 1e-6)) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, EPSILON)) @triton.jit @@ -295,6 +303,8 @@ def prepare_per_head_scores( num_kv_heads = int(num_kv_heads) _, num_layers, num_query_heads, width = scores.shape selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads + # 256 lanes / 4 warps: one program spans the row in a few static loop + # trips without starving occupancy; matches the settle/pack shape. stats_block = 256 rows = num_layers * num_query_heads if normalize_scores: @@ -306,6 +316,7 @@ def prepare_per_head_scores( ROWS=rows, WIDTH=width, BLOCK=stats_block, + EPSILON=STD_EPSILON, num_warps=4, ) reduction_block = 256 @@ -336,6 +347,14 @@ def prepare_per_head_scores( # --------------------------------------------------------------------------- # +# Launch shape of the settle/pack kernel: tokens per program along the +# width/move axis, and its warp count. One pair for every launch site (the +# fused settle, the draft pack, and the standalone test packs) so a retune +# cannot diverge silently. +SETTLE_PACK_BLOCK = 256 +SETTLE_PACK_NUM_WARPS = 4 + + @triton.jit def _settle_ties_and_pack_compaction_sources_kernel( scores, @@ -386,10 +405,10 @@ def _settle_ties_and_pack_compaction_sources_kernel( request = tl.program_id(0) selection_domain = tl.program_id(1) row = request * SELECTION_ROWS + selection_domain - row_scores = scores + row * WIDTH - row_selected = provisional_indices + row * KEEP_COUNT row_output = output_indices + row * KEEP_COUNT if HAS_SETTLE: + row_scores = scores + row * WIDTH + row_selected = provisional_indices + row * KEEP_COUNT # Scores are decode-relative; this row's pinned prompt length rebases # the emitted ordinals to absolute positions (per row, so one launch # may mix prompt lengths). @@ -486,8 +505,9 @@ def _settle_ties_and_pack_compaction_sources_kernel( mask=move < dense_count, ) else: - domain = tl.program_id(1) - dense_output = domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + dense_output = ( + selection_domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + ) tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) if HAS_SWA: swa_source = valid_len - SWA_WINDOW + move @@ -499,18 +519,16 @@ def _settle_ties_and_pack_compaction_sources_kernel( mask=move < swa_count, ) else: - domain = tl.program_id(1) - # Per-layer selection has one dense domain per (layer, - # head). SWA uses one shared source row per head, so only - # the first layer writes it. + swa_mask = move < swa_count if PER_LAYER: - write_swa = domain < NUM_KV_HEADS - else: - write_swa = move >= 0 - head = domain % NUM_KV_HEADS + # Per-layer selection has one dense domain per (layer, + # head). SWA uses one shared source row per head, so + # only the first layer's domains write it. + swa_mask = swa_mask & (selection_domain < NUM_KV_HEADS) + head = selection_domain % NUM_KV_HEADS swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move tl.store( swa_indices + swa_output, swa_source, - mask=write_swa & (move < swa_count), + mask=swa_mask, ) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 32ae62a432ac..0d6c5cb158ef 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -164,6 +164,8 @@ def launch_family_pack(compaction, name): (any well-formed tensor stands in). """ from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + SETTLE_PACK_BLOCK, + SETTLE_PACK_NUM_WARPS, _settle_ties_and_pack_compaction_sources_kernel, ) @@ -184,18 +186,18 @@ def launch_family_pack(compaction, name): PER_LAYER=False, HAS_SWA=False, ) - swa_offsets, swa_indices = family["offsets"], family["source"] + swa_offsets, swa_indices = None, None else: rows = compaction["selection_rows"] shape = compaction["settle_pack_shape"] swa_family = compaction_family(compaction, "swa") - swa_offsets = swa_family["offsets"] if swa_family else family["offsets"] - swa_indices = swa_family["source"] if swa_family else family["source"] + swa_offsets = swa_family["offsets"] if swa_family else None + swa_indices = swa_family["source"] if swa_family else None _settle_ties_and_pack_compaction_sources_kernel[(compaction["request_count"], rows)]( - kept, - valid, - family["offsets"], - kept, + None, + None, + None, + None, kept, valid, family["offsets"], @@ -207,8 +209,8 @@ def launch_family_pack(compaction, name): SELECTION_ROWS=rows, **shape, HAS_SETTLE=False, - BLOCK=256, - num_warps=4, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, ) 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 index 03da8a048684..264d677ee2a7 100644 --- 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 @@ -32,11 +32,13 @@ def _make_union_buffers( omega, offsets, decode_width=None, + eviction_mode="union", ): - """Union-mode buffers over one shared page-table slot. + """Buffers over one shared page-table slot (union mode by default). - The union runner also compiles the score-only entries, so one buffer - namespace serves both the fused pipeline and the split reference leg. + Union runners compile only the fused pipeline, so the split reference + leg builds its own score-only buffers with ``eviction_mode="per_head"`` + over the same pools. """ from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( init_eviction_buffers, @@ -44,7 +46,7 @@ def _make_union_buffers( num_layers = len(layer_pools) return init_eviction_buffers( - eviction_mode="union", + eviction_mode=eviction_mode, layer_pools=layer_pools, dense_groups=[list(range(num_layers))], dense_layers=list(range(num_layers)), @@ -214,12 +216,26 @@ def _check_union_fusion_matches_split_pipeline( omega=omega, offsets=offsets, ) + ref_bufs = _make_union_buffers( + 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, + eviction_mode="per_head", + ) k_plane = [2 * page for page in page_permutation] v_plane = [2 * page + 1 for page in page_permutation] - _write_block_offsets( - bufs, - torch.tensor([[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device), + encoded = torch.tensor( + [[[k_plane, v_plane], [k_plane, v_plane]]], dtype=torch.int32, device=device ) + _write_block_offsets(bufs, encoded) + _write_block_offsets(ref_bufs, encoded) if valid_lens is None: valid_lens = [seq_len, seq_len] valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) @@ -233,7 +249,7 @@ def _check_union_fusion_matches_split_pipeline( split_widths = torch.empty(request_count, dtype=torch.int32, device=device) token_starts = torch.tensor(score_starts, dtype=torch.int32, device=device) per_head = _launch_split_scores( - bufs, + ref_bufs, request_count, valid_seq_lens, split_widths, @@ -416,10 +432,27 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: decode_width=decode_window, ) assert (bufs.cute_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_bufs = _make_union_buffers( + layer_pools=layer_pools, + max_requests=request_count, + 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, + eviction_mode="per_head", + ) page_ids = torch.arange(num_pages, dtype=torch.int32, device=device) - bufs.block_offsets_device.zero_() - bufs.block_offsets_device[0, :request_count, 0, :num_pages] = 2 * page_ids - bufs.block_offsets_device[0, :request_count, 1, :num_pages] = 2 * page_ids + 1 + for staged in (bufs, ref_bufs): + 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 valid_seq_lens = torch.zeros(max_requests, dtype=torch.int32, device=device) token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) @@ -430,7 +463,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: # the pure-torch union oracle. split_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) per_head = _launch_split_scores( - bufs, + ref_bufs, request_count, valid_seq_lens, split_widths, 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 index 8737bfaff447..40703f34a476 100644 --- 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 @@ -17,13 +17,16 @@ import pytest import torch +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + SETTLE_PACK_BLOCK as _BLOCK, +) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + SETTLE_PACK_NUM_WARPS as _NUM_WARPS, +) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _settle_ties_and_pack_compaction_sources_kernel, ) -_BLOCK = 256 -_NUM_WARPS = 4 - def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, keep_count): """Settle each row's provisional top-k in place (exact integer semantics). From 4f22db1d30702cb8d6e8c95ed0edb51a7c273e15 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 22:19:21 -0700 Subject: [PATCH 098/178] [None][refactor] Framework census F-batch: compaction contract + hook/config honesty (knife 22) User-approved F1-F7: - F1 _page_table_provider (closure factory, dead 'what'/'device' params, worthless identity cache) dies; call sites take the slot views inline. - F2 _make_move_buffers collapses to a scalar-count index alloc; the per-request move offsets are always caller-owned rows (production stages them; conftest allocates the capacity cumsum the old internal fallback produced). - F3 the bundle is launch data, not an input mirror: only the six fields the driver fires survive (families, settle_pack_tensors/shape, draft_pack, swa_destination_bases, swa_rebase_delta); conftest mirrors the geometry inputs itself; the now-dead kept_token_ordinals param is gone. - F4 init_compaction_buffers takes the derived selection facts (union/per_layer) instead of the TriAttention config literal; module docstring rewritten as a method-agnostic contract. - F5 the generation-step hooks document their true firing conditions (once per executor iteration, generation_requests may be empty) and **kwargs as reserved. - F6 the union/normalize_scores cross-field contract surfaces at config validation (manager raise stays as the op-boundary backstop); field description in contract language; golden manifest regenerated = byte-identical (receipt: GENERATE_RC=0 MANIFEST_UNCHANGED). - F7 has_independent_draft_kv_cache (zero production consumers) dies; the executor test asserts the real field. Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 193 ++++++++---------- .../triattention/triattention.py | 4 +- .../_torch/pyexecutor/resource_manager.py | 18 +- tensorrt_llm/llmapi/llm_args.py | 11 +- .../test_kv_cache_compression_manager.py | 2 +- .../_torch/kv_cache_compression/conftest.py | 49 ++++- .../test_triattention_pipeline.py | 9 + 7 files changed, 162 insertions(+), 124 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 3d159366a759..34ba40c6d059 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,63 +15,40 @@ """Batched physical KV-cache compaction for eviction-based compression. -A general post-eviction component: given each request's kept-token ordinals, -its valid sequence length, and the staged V2 block offsets, the surviving KV -moves in place through batched C++ compact launches, fed by per-request move -indices packed on device. Everything is plain tensors and dicts: -``init_compaction_buffers`` allocates the launch data once per geometry -(called by ``triattention.init_eviction_buffers``) and returns one bundle -whose fields the eviction driver fires directly each round -- the target's -dense/SWA packing rides the driver's fused settle launch (described by the -bundle's ``settle_pack_tensors``/``settle_pack_shape``), the co-compressed -draft's own pack launch and every family's C++ moves are inlined in -``triattention.run_eviction_round``. +Contract: the caller supplies increasing kept decode ordinals per selection +row (absolute positions), each request's valid length and pinned prompt +length, and the staged V2 block offsets; the surviving KV then moves in +place through batched C++ compact launches, fed by per-request move indices +packed on device. Everything is plain tensors and dicts: +``init_compaction_buffers`` allocates the launch data once per geometry and +returns one bundle whose fields the eviction driver fires directly each +round -- the target's dense/SWA packing rides the driver's fused settle +launch (described by the bundle's ``settle_pack_tensors``/ +``settle_pack_shape``); the driver also fires the co-compressed draft's own +pack launch (``draft_pack``) and every family's C++ moves. """ from collections import OrderedDict -from typing import Callable, Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple import torch -def _make_move_buffers( +def _make_move_indices( index_prefix: Tuple[int, ...], - moves_per_request: List[int], + moves_per_request: int, + request_count: int, device: torch.device, - external_offsets: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Allocate the packed source-index buffer and its per-request offsets. +) -> torch.Tensor: + """Packed source-index buffer sized for the widest per-request moves. - ``external_offsets`` shares a caller-owned device row (refreshed together - with the round metadata in one copy) instead of allocating one here; the - index buffer is always sized for the widest per-request move counts. + The per-request move offsets are always caller-owned rows (refreshed + together with the round metadata in one copy), so only the index + rectangle is allocated here. """ - offsets = [0] - for count in moves_per_request: - offsets.append(offsets[-1] + count) - indices = torch.empty((*index_prefix, offsets[-1]), dtype=torch.int32, device=device) - if external_offsets is not None: - return indices, external_offsets - return indices, torch.tensor(offsets, dtype=torch.int32, device=device) - - -def _page_table_provider( - page_table_slots: Dict[int, int], - kv_block_offsets: torch.Tensor, - device: torch.device, - request_count: int, - what: str, -) -> Callable[[int], torch.Tensor]: - """Return per-slot K block-offset views, cached per slot.""" - tables: Dict[int, torch.Tensor] = {} - - def page_table_for(representative: int) -> torch.Tensor: - slot = page_table_slots[representative] - if slot not in tables: - tables[slot] = kv_block_offsets[slot, :request_count, 0] - return tables[slot] - - return page_table_for + return torch.empty( + (*index_prefix, moves_per_request * request_count), dtype=torch.int32, device=device + ) def _compact_groups( @@ -127,12 +104,12 @@ def _compact_groups( def init_compaction_buffers( *, - eviction_mode: str, + union: bool, + per_layer: bool, layer_pools: List[torch.Tensor], dense_layers: List[int], swa_layers: List[int], layer_group_representative: Dict[int, int], - kept_token_ordinals: torch.Tensor, valid_sequence_lengths: torch.Tensor, kv_block_offsets: torch.Tensor, page_table_slots: Dict[int, int], @@ -149,40 +126,45 @@ def init_compaction_buffers( draft_protected_tail_capacity: Optional[int] = None, draft_kv_block_offsets: Optional[torch.Tensor] = None, draft_page_table_slots: Optional[Dict[int, int]] = None, - dense_move_offsets: Optional[torch.Tensor] = None, + dense_move_offsets: torch.Tensor, swa_move_offsets: Optional[torch.Tensor] = None, draft_move_offsets: Optional[torch.Tensor] = None, ) -> Dict[str, object]: """Allocate the per-geometry compaction launch data as one plain dict. - Dense layers keep the prompt in place and compact the selected decode - tokens plus any target KV reserved for the next overlapped forward; - kernel-masked SWA layers keep the latest window plus the same protected - tail. A co-compressed draft cache reuses the target's kept token ordinals - (broadcast over the draft's own KV-head count, union mode only) plus the - draft's own protected tail, landing at the same destination base. + ``union``/``per_layer`` are the two selection-geometry facts consumed + here: one shared selection row per request (union), or one per + (layer, KV head) (per_layer), else one per KV head. Dense layers keep + the prompt in place and compact the selected decode tokens plus any + target KV reserved for the next overlapped forward; kernel-masked SWA + layers keep the latest window plus the same protected tail. A + co-compressed draft cache reuses the target's single-row kept ordinals + (broadcast over the draft's own KV-head count) plus the draft's own + protected tail, landing at the same destination base. - ``kept_token_ordinals`` carries increasing kept decode ordinals (absolute - positions) per request; prompt tokens never move, so the rectangle is - prompt-length independent and ``prompt_offsets`` carries each request's - pinned prompt length. The increasing order is LOAD-BEARING: - sparseKvCacheCompactOp.cpp's forward tiled in-place copy requires - increasing source ordinals with ``destinationBases[request] + move <= - source[move]``. ``kv_block_offsets`` is the staged V2 snapshot + The driver's settle launch packs increasing kept decode ordinals + (absolute positions) per selection row into the index buffers here; + prompt tokens never move, so the rectangle is prompt-length independent + and ``prompt_offsets`` carries each request's pinned prompt length. The + increasing order is LOAD-BEARING: sparseKvCacheCompactOp.cpp's forward + tiled in-place copy requires increasing source ordinals with + ``destinationBases[request] + move <= source[move]``. ``kv_block_offsets`` is the staged V2 snapshot ``[slot, request, K/V, block]`` (offset = ``2*page + plane``); ``protected_tail_capacity`` is the widest per-request tail this geometry must support -- actual per-round lengths arrive through the staged - move-offset rows. Nothing launches here: the target's dense/SWA packing - rides the driver's fused settle launch (``settle_pack_tensors`` + - ``settle_pack_shape`` describe it) and the driver's round function fires - the draft pack (``draft_pack``) and every family's C++ moves directly. + move-offset rows, which are caller-owned. Nothing launches here: the + target's dense/SWA packing rides the driver's fused settle launch + (``settle_pack_tensors`` + ``settle_pack_shape`` describe it) and the + driver's round function fires the draft pack (``draft_pack``) and every + family's C++ moves directly. - Returns one plain bundle dict: ``families`` (each ``{"name", "groups", - "source", "offsets", "destination_bases"}``), the settle/pack launch data - above, ``draft_pack`` (``None`` or ``{"indices", "offsets", - "dense_total", "move_capacity", "num_kv_heads"}``), ``swa_rebase_delta`` - (per-round SWA destination rebase), and the geometry scalars - (``selection_rows``, ``request_count``, ``decode_keep_count``, ...). + Returns EXACTLY the launch data the driver fires: ``families`` (each + ``{"name", "groups", "source", "offsets", "destination_bases"}``), + ``settle_pack_tensors``/``settle_pack_shape``, ``draft_pack`` (``None`` + or ``{"indices", "offsets", "dense_total", "move_capacity", + "num_kv_heads"}``), ``swa_destination_bases``, and ``swa_rebase_delta`` + (per-round SWA destination rebase). The bundle is launch data, not an + input mirror. """ device = layer_pools[dense_layers[0]].device request_count = int(request_count) @@ -192,22 +174,24 @@ def init_compaction_buffers( swa_layers = tuple(int(layer) for layer in swa_layers) layer_pool_keys = tuple(layer_pool_keys) - per_layer = eviction_mode == "per_layer_perhead" # The C++ compact op takes the KV-head count from each launch's pool # shape [pages, K/V, heads, tokens, dim]. num_kv_heads = int(layer_pools[dense_layers[0]].shape[2]) dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer else (num_kv_heads,) - dense_move_indices, dense_move_offsets = _make_move_buffers( + dense_move_indices = _make_move_indices( dense_index_prefix, - [decode_keep_count + protected_tail_capacity] * request_count, + decode_keep_count + protected_tail_capacity, + request_count, device, - external_offsets=dense_move_offsets, - ) - page_table_for = _page_table_provider( - page_table_slots, kv_block_offsets, device, request_count, "compaction" ) dense_entries = [ - (layer, layer_pools[layer], page_table_for(layer_group_representative[layer])) + ( + layer, + layer_pools[layer], + kv_block_offsets[ + page_table_slots[layer_group_representative[layer]], :request_count, 0 + ], + ) for layer in dense_layers ] @@ -225,23 +209,23 @@ def init_compaction_buffers( # prompt-dependent and checked by the caller each round. swa_window = int(swa_window) swa_destination_bases = torch.empty_like(prompt_offsets) - swa_move_indices, swa_move_offsets = _make_move_buffers( + swa_move_indices = _make_move_indices( (num_kv_heads,), - [swa_window + protected_tail_capacity] * request_count, + swa_window + protected_tail_capacity, + request_count, device, - external_offsets=swa_move_offsets, ) # SWA layers are staged as their own page-table representatives. - swa_entries = [(layer, layer_pools[layer], page_table_for(layer)) for layer in swa_layers] + swa_entries = [ + ( + layer, + layer_pools[layer], + kv_block_offsets[page_table_slots[layer], :request_count, 0], + ) + for layer in swa_layers + ] dense_slots = {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer else None - union = eviction_mode == "union" - if union: - selection_rows = 1 - elif per_layer: - selection_rows = len(dense_layers) * num_kv_heads - else: - selection_rows = num_kv_heads # Selection rows carry decode-only kept ordinals (already absolute), so # the rectangle is prompt-length independent. HAS_SWA specializes the SWA # loads and stores away, so without SWA layers the SWA pointer arguments @@ -300,20 +284,21 @@ def init_compaction_buffers( # The draft forms its own launch groups so it may use a different # KV-head count than the target. draft_num_kv_heads = int(draft_layer_pools[draft_layers[0]].shape[2]) - draft_move_indices, draft_move_offsets = _make_move_buffers( + draft_move_indices = _make_move_indices( (draft_num_kv_heads,), - [decode_keep_count + draft_tail] * request_count, + decode_keep_count + draft_tail, + request_count, device, - external_offsets=draft_move_offsets, - ) - draft_page_table_for = _page_table_provider( - draft_page_table_slots, draft_kv_block_offsets, device, request_count, "draft" ) draft_entries = [ ( layer, draft_layer_pools[layer], - draft_page_table_for(draft_layer_group_representative[layer]), + draft_kv_block_offsets[ + draft_page_table_slots[draft_layer_group_representative[layer]], + :request_count, + 0, + ], ) for layer in draft_layers ] @@ -342,23 +327,11 @@ def init_compaction_buffers( return dict( families=families, - kept_token_ordinals=kept_token_ordinals, - valid_sequence_lengths=valid_sequence_lengths, - selection_rows=selection_rows, settle_pack_tensors=settle_pack_tensors, settle_pack_shape=settle_pack_shape, draft_pack=draft_pack, - num_kv_heads=num_kv_heads, - swa_window=swa_window, swa_destination_bases=swa_destination_bases, # The prompt offsets may be re-staged each round; the driver rebases # the SWA landing positions with this delta before the moves. swa_rebase_delta=decode_keep_count - swa_window, - prompt_offsets=prompt_offsets, - decode_keep_count=decode_keep_count, - request_count=request_count, - protected_tail_capacity=protected_tail_capacity, - draft_protected_tail_capacity=( - int(draft_protected_tail_capacity or 0) if draft_layers else 0 - ), ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index bbdf2b643777..9c770126dd81 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -523,12 +523,12 @@ def init_eviction_buffers( draft_move_offsets=bufs.draft_move_offsets, ) compaction = init_compaction_buffers( - eviction_mode=eviction_mode, + union=union, + per_layer=eviction_mode == "per_layer_perhead", layer_pools=layer_pools, dense_layers=list(dense_layers), swa_layers=list(swa_layers), layer_group_representative=layer_group_representative, - kept_token_ordinals=bufs.keep, valid_sequence_lengths=bufs.valid_seq_lens_device, kv_block_offsets=bufs.block_offsets_device, page_table_slots=bufs.representative_slots, diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 157e2d7d720a..d983613fc995 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2462,10 +2462,6 @@ def __init__( draft_kv_cache_manager.kv_compression_manages_history = ( self.adjusts_generation_kv_length) - @property - def has_independent_draft_kv_cache(self) -> bool: - return self.draft_kv_cache_manager is not None - # ================================================================== # # KV-cache lifecycle hooks (5, in temporal order). # # Subclasses override what they need; all default to no-op. # @@ -2490,15 +2486,23 @@ def on_generation_step_begin( scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per generation step before this step's forward.""" + """Fired once per executor iteration, before this step's forward. + + The batch may contain only context requests and + ``scheduled_batch.generation_requests`` may be empty. ``**kwargs`` + is reserved for forward compatibility (never populated today). + """ def on_generation_step_end( self, scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per generation step, after every layer's forward - completes. Override for periodic or budget-triggered eviction. + """Fired once per executor iteration, after every layer's forward + completes. The batch may contain only context requests and + ``scheduled_batch.generation_requests`` may be empty; ``**kwargs`` + is reserved for forward compatibility. Override for periodic or + budget-triggered eviction. """ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cf2ba9ec9b75..462f49d7486b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3450,8 +3450,9 @@ class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): 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.") + "before selection (upstream default). `union` eviction always " + "normalizes; False is only valid with `per_head`/`per_layer_perhead`." + ) pin_prefill: bool = Field( default=True, description="Always preserve the prompt (prefill) tokens; only decode " @@ -3494,6 +3495,12 @@ def _require_calibration_inputs(self): "TriAttention requires both model_path and calibration_path; " "TRT-LLM consumes an official calibration file and does not " "compute one.") + # Same config-validation-time surfacing for the cross-field contract + # (the manager re-raises at construction as the op-boundary backstop). + if self.eviction_mode == "union" and not self.normalize_scores: + raise ValueError( + "union eviction always normalizes scores; normalize_scores=" + "False is only valid with per_head/per_layer_perhead.") return self def to_manager_kwargs(self) -> dict: 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 4e1c70f758dd..a074504f3dd8 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -149,7 +149,7 @@ def test_length_adjustment_marks_target_and_draft_v2(self): assert manager.kv_cache_manager is target assert manager.draft_kv_cache_manager is draft - assert manager.has_independent_draft_kv_cache + assert manager.draft_kv_cache_manager is not None assert target.kv_compression_manages_history is True assert draft.kv_compression_manages_history is True diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 0d6c5cb158ef..3ec78b351e07 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -136,7 +136,14 @@ def make_ramp_pools( def build_compaction(**overrides): - """``init_compaction_buffers`` with the suite's default 2-layer geometry.""" + """``init_compaction_buffers`` with the suite's default 2-layer geometry. + + Accepts ``eviction_mode`` for test ergonomics (translated to the + builder's derived selection facts), allocates the caller-owned move + offset rows the production driver would stage (capacity cumsum), and + mirrors the geometry inputs onto the bundle for the standalone helpers + here -- production reads only the launch fields. + """ from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers args = dict( @@ -151,7 +158,45 @@ def build_compaction(**overrides): swa_window=None, ) args.update(overrides) - return init_compaction_buffers(**args) + mode = args.pop("eviction_mode") + union = mode == "union" + per_layer = mode == "per_layer_perhead" + kept = args.pop("kept_token_ordinals") + request_count = args["request_count"] + keep_count = args["decode_keep_count"] + tail = int(args.get("protected_tail_capacity", 0)) + draft_tail = int(args.get("draft_protected_tail_capacity") or 0) + device = args["layer_pools"][args["dense_layers"][0]].device + swa_window = int(args["swa_window"] or 0) if args["swa_layers"] else 0 + + def capacity_offsets(count): + return torch.arange(0, (request_count + 1) * count, count, dtype=torch.int32, device=device) + + args.setdefault("dense_move_offsets", capacity_offsets(keep_count + tail)) + if args["swa_layers"]: + args.setdefault("swa_move_offsets", capacity_offsets(swa_window + tail)) + if args.get("draft_layers"): + args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) + num_kv_heads = int(args["layer_pools"][args["dense_layers"][0]].shape[2]) + compaction = init_compaction_buffers(union=union, per_layer=per_layer, **args) + # Test-side mirror of the construction inputs: production reads only the + # launch fields; the standalone helpers here need the geometry back. + compaction.update( + kept_token_ordinals=kept, + valid_sequence_lengths=args["valid_sequence_lengths"], + prompt_offsets=args["prompt_offsets"], + request_count=request_count, + decode_keep_count=keep_count, + protected_tail_capacity=tail, + draft_protected_tail_capacity=draft_tail if args.get("draft_layers") else 0, + swa_window=swa_window, + selection_rows=( + 1 + if union + else (len(args["dense_layers"]) * num_kv_heads if per_layer else num_kv_heads) + ), + ) + return compaction def launch_family_pack(compaction, name): diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 54c4bc7dd666..4e6af7ba308c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -147,6 +147,15 @@ def test_llm_args_dispatch_and_validation(self): ) with pytest.raises(ValidationError): TriAttentionKvCacheCompressionConfig(eviction_mode="made_up_mode") + # Cross-field contract surfaces at config validation (the manager + # re-raises at construction as the op-boundary backstop). + with pytest.raises(ValidationError, match="normalize"): + TriAttentionKvCacheCompressionConfig( + model_path="/models/test", + calibration_path="/calib/test.pt", + eviction_mode="union", + normalize_scores=False, + ) def test_factory_returns_triattention_and_propagates_config_fields(self): # Calibration is deferred to the first request, so construction needs From c136453d4f780128c8179487c349df9891241f1b Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 22:31:09 -0700 Subject: [PATCH 099/178] [None][test] Test tranche 3: dedupe scaffolding and fold duplicate guards (knife 19) Net -165 LOC (honest shortfall vs the -1300 target: after four prior tranches the remaining surface is protected material -- incident anchors, Triton-constexpr-distinct settle rows, the user-protected per_head/per_layer oracles, and the official-calibration contract anchors; measured file-by-file in the lane report). - The CuTe test scaffolding duplicated across the score and union-fusion files (buffer builders, block-offset staging, metadata staging, the split score leg) is single-sourced in conftest. - test_buffer_build_receives_mode_and_capacity_kwargs dies as a strict subset of the draft file's rebuild/invalidation superset test. - test_mla_selfkonly_cache_is_rejected folds into the draft guard matrix as a target_kv_factor row (same raise site; case union grew). - The single-caller union-vs-split check helper folds into its parametrized test (8-row matrix byte-identical). All protected anchors verified present post-cut (grep receipts in the lane report); official-calibration coverage untouched per the product contract. Signed-off-by: tianruih --- .../_torch/kv_cache_compression/conftest.py | 112 +++++++++ .../test_triattention_cute_score.py | 123 +--------- .../test_triattention_cute_union_fusion.py | 218 +++--------------- .../test_triattention_draft_cocompaction.py | 6 +- .../test_triattention_pipeline.py | 40 ---- 5 files changed, 167 insertions(+), 332 deletions(-) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 3ec78b351e07..edaeabff8b3b 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -496,3 +496,115 @@ def torch_tri_score_oracle( head_scores.append(position + mlr) scores.append(torch.stack(head_scores)) return scores + + +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, +): + """Real eviction buffers over one shared page-table slot. + + Shared by the CuTe score and union-fusion tests. Union runners compile + only the fused pipeline, so split-reference legs build their own + score-only buffers with ``eviction_mode="per_head"`` over the same pools. + """ + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + init_eviction_buffers, + ) + + num_layers = len(layer_pools) + return init_eviction_buffers( + eviction_mode=eviction_mode, + layer_pools=layer_pools, + dense_groups=[list(range(num_layers))], + dense_layers=list(range(num_layers)), + page_representatives=[0], + max_requests=max_requests, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_freqs=int(q_real.shape[-1]), + keep_count=1, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + offsets=offsets, + omega=omega, + decode_width=decode_width, + layer_group_representative={layer: 0 for layer in range(num_layers)}, + layer_pool_keys=[("pool", 0)] * num_layers, + ) + + +def write_block_offsets(bufs, encoded): + """Load a test page table into the staged block-offset plane.""" + bufs.block_offsets_device.zero_() + bufs.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) + + +def stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts): + """Stage the per-round score metadata exactly like production. + + The compiled runner reads valid lengths and window starts straight from + the staged metadata rows (pointer capture), so stage them like + ``stage_eviction_cohort`` does; the width subtraction mirrors what the + production phase-gather launch derives on device. + """ + torch.sub( + valid_seq_lens[:request_count], + token_starts[:request_count], + out=valid_widths[:request_count], + ) + bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) + bufs.token_starts_device[:request_count].copy_(token_starts[:request_count]) + + +def launch_split_scores( + bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin +): + """The production score-only leg plus the decode-window gather + (``run_eviction_round``'s per-head sequence, parameterized by count).""" + stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) + assert request_count in bufs.runner._compiled + bufs.runner.launch(request_count, mean_cos, mean_sin) + num_segments = request_count * bufs.num_layers + group_size = bufs.num_q_heads // bufs.num_kv_heads + source = ( + bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] + .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ + :, :group_size + ] + .permute(2, 3, 0, 1, 4) + ) + columns = ( + token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + bufs.gather_columns + ) + columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + ) + output = torch.full( + (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), + float("nan"), + dtype=torch.float32, + device=bufs.device, + ) + torch.gather( + source, + 4, + columns, + out=output.view( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + ), + ) + return output 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 index 2444f16fbd86..808c0eee9459 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -2,15 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 """The SM100 TriAttention CuTe scorer (the only score path) vs PyTorch oracles. -Two layers of coverage over the production score leg (buffer metadata -staging, the compiled runner launch, and the decode-window gather -- the same -sequence ``run_eviction_round`` fires). The kernel-numerics matrix drives a -single-layer buffers across the supported page geometries (permuted -physical pages, ragged valid lengths, GQA group 4 riding the padded MMA tile) -against inline oracle math. The launch-path matrix drives multi-layer -buffers across the named production geometries (Qwen3, GPT-OSS, the -originally validated 128-token-page shape) against the shared pure-PyTorch -oracle, sweeps request counts up to the buffer capacity, and checks the +The launch-path matrix drives multi-layer buffers across the named +production geometries (Qwen3, GPT-OSS, the originally validated +128-token-page shape) against the shared pure-PyTorch oracle -- permuted +physical pages, ragged valid lengths, GQA group 4 riding the padded MMA +tile -- sweeps request counts up to the buffer capacity, and checks the per-request decode-width metadata the selection reduce kernels consume. The contract test pins the loud-failure behavior: unsupported geometry raises from the CuTe runner's own validation at buffer construction -- there is @@ -20,7 +16,10 @@ 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), @@ -28,106 +27,6 @@ ) -def _make_score_buffers( - *, - layer_pools, - max_requests, - seq_len, - num_q_heads, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - decode_width=None, - eviction_mode="per_head", -): - """Score-only buffers over one shared page-table slot.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - init_eviction_buffers, - ) - - num_layers = len(layer_pools) - return init_eviction_buffers( - eviction_mode=eviction_mode, - layer_pools=layer_pools, - dense_groups=[list(range(num_layers))], - dense_layers=list(range(num_layers)), - page_representatives=[0], - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_freqs=int(q_real.shape[-1]), - keep_count=1, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=freq_scale_sq, - offsets=offsets, - omega=omega, - decode_width=decode_width, - layer_group_representative={layer: 0 for layer in range(num_layers)}, - layer_pool_keys=[("pool", 0)] * num_layers, - ) - - -def _write_block_offsets(bufs, encoded): - """Load a test page table into the staged block-offset plane.""" - bufs.block_offsets_device.zero_() - bufs.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) - - -def _launch_split_scores( - bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin -): - """The production score-only leg: stage metadata, fire the compiled - runner, gather each request's decode window (``run_eviction_round``'s - per-head sequence, parameterized by the request count).""" - num_segments = request_count * bufs.num_layers - torch.sub( - valid_seq_lens[:request_count], - token_starts[:request_count], - out=valid_widths[:request_count], - ) - # The compiled runner reads valid lengths and window starts straight from - # the staged metadata rows (pointer capture), so stage them exactly like - # ``stage_eviction_cohort`` does. - bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) - bufs.token_starts_device[:request_count].copy_(token_starts[:request_count]) - assert request_count in bufs.runner._compiled - bufs.runner.launch(request_count, mean_cos, mean_sin) - group_size = bufs.num_q_heads // bufs.num_kv_heads - source = ( - bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] - .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ - :, :group_size - ] - .permute(2, 3, 0, 1, 4) - ) - columns = ( - token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + bufs.gather_columns - ) - columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width - ) - output = torch.full( - (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), - float("nan"), - dtype=torch.float32, - device=bufs.device, - ) - torch.gather( - source, - 4, - columns, - out=output.view( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width - ), - ) - return output - - def _build_case( *, max_requests: int, @@ -168,7 +67,8 @@ def _build_case( 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 - bufs = _make_score_buffers( + bufs = _make_cute_buffers( + eviction_mode="per_head", layer_pools=pools, max_requests=max_requests, seq_len=capacity, @@ -327,7 +227,8 @@ def test_unsupported_geometry_raises_at_buffer_construction(): ] calib = torch.randn(num_layers, 2, num_freqs, device=device) with pytest.raises(RuntimeError, match="no other score path exists"): - _make_score_buffers( + _make_cute_buffers( + eviction_mode="per_head", layer_pools=pools, max_requests=max_requests, seq_len=page_count * tokens_per_block, 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 index 264d677ee2a7..94ed83ca4574 100644 --- 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 @@ -12,6 +12,10 @@ 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), @@ -19,117 +23,6 @@ ) -def _make_union_buffers( - *, - layer_pools, - max_requests, - seq_len, - num_q_heads, - q_real, - q_imag, - mlr_coef, - freq_scale_sq, - omega, - offsets, - decode_width=None, - eviction_mode="union", -): - """Buffers over one shared page-table slot (union mode by default). - - Union runners compile only the fused pipeline, so the split reference - leg builds its own score-only buffers with ``eviction_mode="per_head"`` - over the same pools. - """ - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - init_eviction_buffers, - ) - - num_layers = len(layer_pools) - return init_eviction_buffers( - eviction_mode=eviction_mode, - layer_pools=layer_pools, - dense_groups=[list(range(num_layers))], - dense_layers=list(range(num_layers)), - page_representatives=[0], - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_freqs=int(q_real.shape[-1]), - keep_count=1, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=freq_scale_sq, - offsets=offsets, - omega=omega, - decode_width=decode_width, - layer_group_representative={layer: 0 for layer in range(num_layers)}, - layer_pool_keys=[("pool", 0)] * num_layers, - ) - - -def _write_block_offsets(bufs, encoded): - """Load a test page table into the staged block-offset plane.""" - bufs.block_offsets_device.zero_() - bufs.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) - - -def _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts): - """Stage the per-round score metadata exactly like production. - - The compiled runner reads valid lengths and window starts straight from - the staged metadata rows (pointer capture), so stage them like - ``stage_eviction_cohort`` does; the width subtraction mirrors what the - production phase-gather launch derives on device. - """ - torch.sub( - valid_seq_lens[:request_count], - token_starts[:request_count], - out=valid_widths[:request_count], - ) - bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) - bufs.token_starts_device[:request_count].copy_(token_starts[:request_count]) - - -def _launch_split_scores( - bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin -): - """The production score-only leg plus the decode-window gather.""" - _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) - assert request_count in bufs.runner._compiled - bufs.runner.launch(request_count, mean_cos, mean_sin) - num_segments = request_count * bufs.num_layers - group_size = bufs.num_q_heads // bufs.num_kv_heads - source = ( - bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] - .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ - :, :group_size - ] - .permute(2, 3, 0, 1, 4) - ) - columns = ( - token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + bufs.gather_columns - ) - columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width - ) - output = torch.full( - (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), - float("nan"), - dtype=torch.float32, - device=bufs.device, - ) - torch.gather( - source, - 4, - columns, - out=output.view( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width - ), - ) - return output - - def _launch_union_fusion( bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out ): @@ -164,7 +57,33 @@ def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tenso return combined -def _check_union_fusion_matches_split_pipeline( +@_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 cohorts: each request scores its own window (one + # start mid-tile, one page-aligned) — the case the fused pipeline + # previously declined. + (32, 32, 8, [37, 128], [250, 198]), + (32, 64, 4, [37, 128], None), + (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, @@ -204,7 +123,7 @@ def _check_union_fusion_matches_split_pipeline( mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() - bufs = _make_union_buffers( + common = dict( layer_pools=[pool], max_requests=2, seq_len=seq_len, @@ -216,19 +135,9 @@ def _check_union_fusion_matches_split_pipeline( omega=omega, offsets=offsets, ) - ref_bufs = _make_union_buffers( - 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, - eviction_mode="per_head", - ) + bufs = _make_cute_buffers(eviction_mode="union", **common) + # The split reference leg runs on its own score-only buffers. + ref_bufs = _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( @@ -286,44 +195,6 @@ def _check_union_fusion_matches_split_pipeline( ) -@_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 cohorts: each request scores its own window (one - # start mid-tile, one page-aligned) — the case the fused pipeline - # previously declined. - (32, 32, 8, [37, 128], [250, 198]), - (32, 64, 4, [37, 128], None), - (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_union_fusion_matches_split_pipeline( - tokens_per_block, num_freqs, num_q_heads, score_starts, valid_lens - ) - - @_SM100_ONLY def test_union_fusion_frequency_count_guard_raises() -> None: """16 frequencies (head size 32) sit outside the fused kernel contract @@ -418,9 +289,8 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: mean_cos = torch.cos(phase).mean(dim=1).contiguous() mean_sin = torch.sin(phase).mean(dim=1).contiguous() - bufs = _make_union_buffers( + common = dict( layer_pools=layer_pools, - max_requests=max_requests, seq_len=seq_len, num_q_heads=num_q_heads, q_real=q_real, @@ -431,23 +301,11 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: offsets=offsets, decode_width=decode_window, ) + bufs = _make_cute_buffers(eviction_mode="union", max_requests=max_requests, **common) assert (bufs.cute_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_bufs = _make_union_buffers( - layer_pools=layer_pools, - max_requests=request_count, - 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, - eviction_mode="per_head", - ) + ref_bufs = _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 (bufs, ref_bufs): staged.block_offsets_device.zero_() 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 index c91cadd20ba1..a1d392da03f5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -223,6 +223,8 @@ def test_mark_page_tables_consumed_orders_both_manager_streams(): # variants raise through the same checks). ("union_only_per_head", "union"), ("draft_kv_factor", "standard key/value cache"), + # Same check family on the TARGET cache (MLA SELFKONLY, kv_factor 1). + ("target_kv_factor", "standard key/value KV cache"), ("full_attention_draft", "full-attention draft"), ("callsite_dflash", "standard paged cache compacted together"), ], @@ -256,12 +258,14 @@ def test_draft_admission_gates_raise(gate, match): budget=8, model_path="/models/test", eviction_mode="per_head" if gate == "union_only_per_head" else "union", - draft_kv_cache_manager=draft_manager, + draft_kv_cache_manager=None if gate == "target_kv_factor" else draft_manager, ) if gate == "draft_kv_factor": # Flipping kv_factor after construction exercises TriAttention's own # runtime gate. draft_manager.kv_factor = 1 + if gate == "target_kv_factor": + manager.kv_cache_manager.kv_factor = 1 with pytest.raises(ValueError, match=match): manager._validate_v2_compatibility() diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 4e6af7ba308c..886746b95a9c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -33,7 +33,6 @@ import torch from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_bare_staging as _make_bare_staging -from conftest import make_buffer_stubs as _make_buffer_stubs 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 @@ -533,13 +532,6 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): assert manager._request_states[7]["confirmed_kv_length"] == physical_confirmed cache.resize.assert_not_called() - def test_mla_selfkonly_cache_is_rejected(self): - manager = _make_triattention() - manager.kv_cache_manager.kv_factor = 1 - - with pytest.raises(ValueError, match="standard key/value KV cache"): - manager._validate_v2_compatibility() - def test_one_model_draft_co_compression_contract_is_accepted(self): # Co-compression keeps the draft's physical length equal to the # target's, so a draft with a smaller max_seq_len than the target's @@ -606,38 +598,6 @@ def test_union_rejects_unnormalized_scores(self): with pytest.raises(ValueError, match="normalize_scores=True"): _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) - def test_buffer_build_receives_mode_and_capacity_kwargs(self): - # Mode/phase/pool-key threading and cached reuse are covered by the - # buffer-rebuild superset test in - # test_triattention_draft_cocompaction.py. - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - - manager = _make_triattention(budget=4) - # The buffers follow the executor limits: eight requests (max batch - # size) by 260 decode tokens (budget plus two eviction periods). - layout, buffers = _make_buffer_stubs(manager) - prepared = [ - _prepared_eviction( - _make_request(7), - request_id=7, - seq_len=8, - expected_keep_count=4, - ) - ] - - with mock.patch.object( - module, - "init_eviction_buffers", - return_value=buffers, - ) as build_buffers: - resources = manager._buffers_for(layout, prepared) - - assert resources is buffers - kwargs = build_buffers.call_args.kwargs - assert kwargs["eviction_mode"] == "union" - assert kwargs["max_requests"] == 8 - assert kwargs["decode_width"] == 260 - def test_stage_rejects_int32_overflowing_round_starts(self): # Round starts past the int32 metadata range fail loudly (in the host # metadata build) before any GPU work is enqueued. From 4542815558258ebaf777c3e47f6a42e958cf5701 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 22:33:04 -0700 Subject: [PATCH 100/178] [None][refactor] Straighten init_eviction_buffers: one construction path (knife 20) Census #10 executed with the verifier corrections: the test-only None-derive path dies -- page_table_keys / num_page_table_slots / decode_width / page_table_token_capacity are required (callers say what they mean; the ONE documented test seam left is the shared phase dict), the isinstance-sniffing _page_table_slot_layout collapses to direct ('pool', id)-to-slot mapping, stage_eviction_cohort's dead seq_lens default and _allocate_page_table_plane's dead 'what' label die. Field hygiene (grep-audited write-vs-read over product+tests): the metadata-row aliases dense/swa/draft_move_offsets and draft_representative_slots were init-internal wiring -- now locals; the six calibration-tensor casts fold to one assignment; seg_req reuse replaces a duplicate arange chain (identical values, one fewer init launch). Audits, NVTX, per-head branches, and the official-calibration paths untouched. Test fallout folded here (the four sites the required params touch): conftest's shared CuTe builder and the three direct constructor call sites pass the capacities explicitly. Signed-off-by: tianruih --- .../triattention/triattention.py | 105 ++++++------------ .../_torch/kv_cache_compression/conftest.py | 7 ++ .../test_triattention_pipeline.py | 3 + .../test_triattention_selection_compaction.py | 2 + 4 files changed, 46 insertions(+), 71 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 9c770126dd81..b281af413e02 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -101,33 +101,6 @@ # Stream-affinity contract: the staged buffers and compiled launches are bound -def _page_table_slot_layout( - page_representatives: List[int], - page_table_keys: List[object], -) -> Tuple[Dict[int, int], int]: - """Map representative layers to page-table snapshot slots.""" - use_pool_ids = all( - isinstance(key, tuple) - and len(key) == 2 - and key[0] == "pool" - and isinstance(key[1], int) - and key[1] >= 0 - for key in page_table_keys - ) - unique_slots = [] - key_to_slot = {} - representative_slots = {} - for representative, key in zip(page_representatives, page_table_keys): - slot = key_to_slot.get(key) - if slot is None: - slot = int(key[1]) if use_pool_ids else len(key_to_slot) - key_to_slot[key] = slot - unique_slots.append(slot) - representative_slots[representative] = slot - slot_count = max(unique_slots, default=-1) + 1 - return representative_slots, slot_count - - def _protected_tail_capacity(manager: KVCacheManagerV2, what: str) -> int: """The V2 tail (extra KV + draft reserve + 1) moved with every compaction.""" capacity = int(manager.num_extra_kv_tokens) + int(manager._kv_reserve_draft_tokens) + 1 @@ -140,18 +113,20 @@ def _allocate_page_table_plane( layer_pools: List[torch.Tensor], page_representatives: List[int], page_table_keys: List[object], - num_page_table_slots: Optional[int], + num_page_table_slots: int, token_capacity: int, max_requests: int, device: torch.device, - what: str, ) -> Tuple[Dict[int, int], int, torch.Tensor, torch.Tensor]: - """Allocate one staged block-offset plane (host pinned + device).""" - representative_slots, minimum_slots = _page_table_slot_layout( - page_representatives, page_table_keys - ) - if num_page_table_slots is None: - num_page_table_slots = minimum_slots + """Allocate one staged block-offset plane (host pinned + device). + + The ``("pool", id)`` page-table keys ARE the snapshot slot numbering: + one scheme, one constructor contract. + """ + representative_slots = { + representative: int(key[1]) + for representative, key in zip(page_representatives, page_table_keys) + } tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block copy_block_count = (page_count + 3) // 4 * 4 @@ -184,10 +159,10 @@ def init_eviction_buffers( offsets: torch.Tensor, omega: torch.Tensor, phase: Optional[Dict[str, object]] = None, - page_table_keys: Optional[List[object]] = None, - num_page_table_slots: Optional[int] = None, - decode_width: Optional[int] = None, - page_table_token_capacity: Optional[int] = None, + page_table_keys: List[object], + num_page_table_slots: int, + decode_width: int, + page_table_token_capacity: int, protected_tail_capacity: int = 0, draft_layer_pools: Optional[List[torch.Tensor]] = None, draft_layers: Optional[List[int]] = None, @@ -218,24 +193,16 @@ def init_eviction_buffers( device = layer_pools[page_representatives[0]].device max_requests = int(max_requests) seq_len = int(seq_len) - if page_table_token_capacity is None: - page_table_token_capacity = seq_len page_table_token_capacity = int(page_table_token_capacity) # Decode-width capacity of the score buffers; per-request prompt lengths # are staged runtime metadata. - if decode_width is None: - decode_width = seq_len decode_width = int(decode_width) keep_count = int(keep_count) - q_real = q_real.to(device=device, dtype=torch.float32).contiguous() - q_imag = q_imag.to(device=device, dtype=torch.float32).contiguous() - mlr_coef = mlr_coef.to(device=device, dtype=torch.float32).contiguous() - freq_scale_sq = freq_scale_sq.to(device=device, dtype=torch.float32).contiguous() - offsets = offsets.to(device=device, dtype=torch.float32).contiguous() - omega = omega.to(device=device, dtype=torch.float32).contiguous() - if page_table_keys is None: - page_table_keys = list(range(len(page_representatives))) + q_real, q_imag, mlr_coef, freq_scale_sq, offsets, omega = ( + tensor.to(device=device, dtype=torch.float32).contiguous() + for tensor in (q_real, q_imag, mlr_coef, freq_scale_sq, offsets, omega) + ) bufs = SimpleNamespace() bufs.eviction_mode = eviction_mode @@ -260,16 +227,15 @@ def init_eviction_buffers( page_table_token_capacity, max_requests, device, - "", ) # The draft is never scored: these offsets feed only the draft compacts. bufs.draft_block_offsets_device = None bufs._draft_bulk_offsets_src = None - bufs.draft_representative_slots = {} bufs.draft_copy_block_count = 0 + draft_page_slots: Dict[int, int] = {} if draft_layer_pools is not None: ( - bufs.draft_representative_slots, + draft_page_slots, bufs.draft_copy_block_count, bufs._draft_bulk_offsets_src, bufs.draft_block_offsets_device, @@ -277,11 +243,10 @@ def init_eviction_buffers( draft_layer_pools, draft_page_representatives, draft_page_table_keys, - draft_num_page_table_slots, + int(draft_num_page_table_slots), int(draft_page_table_token_capacity), max_requests, device, - "draft ", ) # ---- per-round metadata table: ONE host-to-device copy per round ------- @@ -304,14 +269,16 @@ def init_eviction_buffers( # Per-request pinned prompt lengths: the score kernel starts each # request's decode window here, so one bucket may mix prompt lengths. bufs.token_starts_device = bufs.request_metadata_device[2, :max_requests] - bufs.dense_move_offsets = bufs.request_metadata_device[3] - bufs.swa_move_offsets = bufs.request_metadata_device[4] - bufs.draft_move_offsets = bufs.request_metadata_device[5] + # Rows 3-5 carry the per-family move offsets; the compaction bundle holds + # these views as each family's ``offsets``, so they are wiring, not state. + dense_move_offsets_row = bufs.request_metadata_device[3] + swa_move_offsets_row = bufs.request_metadata_device[4] + draft_move_offsets_row = bufs.request_metadata_device[5] bufs.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) bufs.mean_sin = torch.empty_like(bufs.mean_cos) # The phase table depends only on the shared calibration; the manager - # shares one dict with every buffer namespace (tests may pass None for a - # private table). + # shares one dict with every buffer namespace. ``phase=None`` is the ONE + # documented test seam: unit fixtures get a private table built here. if phase is None: phase = build_mean_phase_table(offsets, omega, initial_rows=seq_len) bufs.phase = phase @@ -343,9 +310,7 @@ def init_eviction_buffers( ) block_offsets = bufs.block_offsets_device slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) - req_idx = torch.arange(max_requests, dtype=torch.int64, device=device).repeat_interleave( - bufs.num_layers - ) + req_idx = bufs.seg_req.to(torch.int64) slot_idx = slots_t.repeat(max_requests) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) @@ -519,8 +484,8 @@ def init_eviction_buffers( draft_layer_pool_keys=draft_layer_pool_keys, draft_protected_tail_capacity=int(draft_protected_tail_capacity), draft_kv_block_offsets=bufs.draft_block_offsets_device, - draft_page_table_slots=bufs.draft_representative_slots, - draft_move_offsets=bufs.draft_move_offsets, + draft_page_table_slots=draft_page_slots, + draft_move_offsets=draft_move_offsets_row, ) compaction = init_compaction_buffers( union=union, @@ -540,8 +505,8 @@ def init_eviction_buffers( protected_tail_capacity=int(protected_tail_capacity), # Tails vary per round (in-flight growth), so the per-family move # offsets ride the staged metadata rows each round. - dense_move_offsets=bufs.dense_move_offsets, - swa_move_offsets=bufs.swa_move_offsets, + dense_move_offsets=dense_move_offsets_row, + swa_move_offsets=swa_move_offsets_row, **draft_kwargs, ) # Flatten the launch data to plain fields: ONE fused launch settles the @@ -615,7 +580,7 @@ def stage_eviction_cohort( request_ids: List[int], round_starts: List[int], token_starts: List[int], - seq_lens: Optional[List[int]] = None, + seq_lens: List[int], draft_manager: Optional[KVCacheManagerV2] = None, dense_move_offsets: Optional[List[int]] = None, swa_move_offsets: Optional[List[int]] = None, @@ -633,8 +598,6 @@ def stage_eviction_cohort( # would silently corrupt the in-flight compaction. if bufs.page_tables_active: raise RuntimeError("previous page-table cohort is still active") - if seq_lens is None: - seq_lens = [bufs.bucket_seq_len] * request_count request_metadata = torch.as_tensor((round_starts, seq_lens, token_starts), dtype=torch.int32) # Grow the phase table while this cohort's round starts are still host # integers: a stale-capacity gather is an out-of-bounds index_select on diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index edaeabff8b3b..57763703dad5 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -524,6 +524,10 @@ def make_cute_buffers( ) num_layers = len(layer_pools) + # The constructor takes every capacity explicitly (no test-only + # None-derive path); the widest window defaults keep old call sites. + if decode_width is None: + decode_width = seq_len return init_eviction_buffers( eviction_mode=eviction_mode, layer_pools=layer_pools, @@ -542,6 +546,9 @@ def make_cute_buffers( offsets=offsets, omega=omega, decode_width=decode_width, + page_table_keys=[("pool", 0)], + num_page_table_slots=1, + page_table_token_capacity=seq_len, layer_group_representative={layer: 0 for layer in range(num_layers)}, layer_pool_keys=[("pool", 0)] * num_layers, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 886746b95a9c..3ec78a871ffc 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -848,6 +848,9 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun offsets=offsets, omega=omega, decode_width=seq_len - prompt_len, + page_table_keys=[("pool", layer) for layer in layer_order], + num_page_table_slots=num_layers, + page_table_token_capacity=seq_len, layer_group_representative={layer: layer for layer in layer_order}, layer_pool_keys=[("pool", layer) for layer in layer_order], ) 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 index a028aa68785f..f2825691a052 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -525,8 +525,10 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): freq_scale_sq=freq_scale_sq, offsets=torch.zeros(1, dtype=torch.float32, device=device), omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), + decode_width=bucket_capacity, page_table_keys=[("pool", 0), ("pool", 1)], num_page_table_slots=2, + page_table_token_capacity=bucket_capacity, layer_group_representative=layer_group_representative, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], ) From 14d622ecc4c5ed316e1bcc926d3b131e07ccf255 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 22:48:18 -0700 Subject: [PATCH 101/178] [None][refactor] Fold runner compile ritual and single-heir helpers Census #11/#15/#16 on the CuTe score runner: - ONE compile ritual replaces the two duplicated ~35-line score/stats blocks: enable_partial_stats selects the variant key and target dict once, ctor kwargs are hoisted, and the small/large tail rebind twins collapse onto the selected dict. Cache keys and compile order are string- and sequence-identical; the E4 gate (score-only entries exist only for per-head runners) rides the variant slice. - The _TriScoreEpilogue single-heir mixin folds into _TriAttentionScoreKernel (acc_dtype class attribute; the two epilogue methods move verbatim; the super() indirection dies). - _launch_union_finalize (single caller) inlines into launch_union_fusion: the union round reads as two consecutive compiled-kernel calls. Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 280 +++++++----------- 1 file changed, 113 insertions(+), 167 deletions(-) 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 index 326f62a527cd..07c68e4896d8 100644 --- 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 @@ -77,73 +77,12 @@ def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Fl SMALL_WORKLOAD_PAGE_SHARDS = 3 -class _TriScoreEpilogue: - """Minimal TMEM-to-global epilogue for the score specialization.""" - - def __init__(self) -> None: - self.acc_dtype = cutlass.Float32 - - def epilog_tmem_copy_and_partition( - self, - tidx: cutlass.Int32, - accumulator: cute.Tensor, - output: cute.Tensor, - epilogue_tile: cute.Tile, - use_2cta_instrs: bool, - ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: - copy_atom = sm100_utils.get_tmem_load_op( - self.cta_tile_shape_mnk, - self.c_layout, - self.c_dtype, - self.acc_dtype, - epilogue_tile, - use_2cta_instrs, - ) - accumulator_epilogue = cute.flat_divide( - accumulator[((None, None), 0, 0)], - epilogue_tile, - ) - tiled_copy = tcgen05.make_tmem_copy( - copy_atom, - accumulator_epilogue[(None, None, 0, 0)], - ) - thread_copy = tiled_copy.get_slice(tidx) - thread_accumulator = thread_copy.partition_S(accumulator_epilogue) - output_epilogue = cute.flat_divide( - output[((None, None), 0, 0, None, None, None)], - epilogue_tile, - ) - thread_output = thread_copy.partition_D(output_epilogue) - register_accumulator = cute.make_rmem_tensor( - thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.acc_dtype, - ) - return tiled_copy, thread_accumulator, register_accumulator - - def epilog_gmem_copy_and_partition( - self, - tidx: cutlass.Int32, - tiled_copy: cute.TiledCopy, - output: cute.Tensor, - epilogue_tile: cute.Tile, - ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: - output_epilogue = cute.flat_divide( - output[((None, None), 0, 0, None, None, None)], - epilogue_tile, - ) - thread_copy = tiled_copy.get_slice(tidx) - thread_output = thread_copy.partition_D(output_epilogue) - register_output = cute.make_rmem_tensor( - thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.c_dtype, - ) - copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - return copy_atom, register_output, thread_output - - -class _TriAttentionScoreKernel(_TriScoreEpilogue): +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, *, @@ -160,7 +99,6 @@ def __init__( write_partial_stats: bool = False, ) -> None: """Build the single validated production specialization.""" - super().__init__() if pool_dtype is not cutlass.BFloat16: raise ValueError("TriAttention CuTe score requires BF16 K pages") if num_freqs not in (32, 64): @@ -225,6 +163,63 @@ def __init__( 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") + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + accumulator: cute.Tensor, + output: cute.Tensor, + epilogue_tile: cute.Tile, + use_2cta_instrs: bool, + ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + copy_atom = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epilogue_tile, + use_2cta_instrs, + ) + accumulator_epilogue = cute.flat_divide( + accumulator[((None, None), 0, 0)], + epilogue_tile, + ) + tiled_copy = tcgen05.make_tmem_copy( + copy_atom, + accumulator_epilogue[(None, None, 0, 0)], + ) + thread_copy = tiled_copy.get_slice(tidx) + thread_accumulator = thread_copy.partition_S(accumulator_epilogue) + output_epilogue = cute.flat_divide( + output[((None, None), 0, 0, None, None, None)], + epilogue_tile, + ) + thread_output = thread_copy.partition_D(output_epilogue) + register_accumulator = cute.make_rmem_tensor( + thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.acc_dtype, + ) + return tiled_copy, thread_accumulator, register_accumulator + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tiled_copy: cute.TiledCopy, + output: cute.Tensor, + epilogue_tile: cute.Tile, + ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + output_epilogue = cute.flat_divide( + output[((None, None), 0, 0, None, None, None)], + epilogue_tile, + ) + thread_copy = tiled_copy.get_slice(tidx) + thread_output = thread_copy.partition_D(output_epilogue) + register_output = cute.make_rmem_tensor( + thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, + self.c_dtype, + ) + copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + return copy_atom, register_output, thread_output + @cute.jit def __call__( self, @@ -1620,112 +1615,71 @@ def __init__( variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] if max_requests > 1: variants.append((max_requests, 2)) + # ONE compile ritual: per-head runners compile the score-only entry, + # the union runner ONLY its fused stats+union pipeline (plus the + # normalize finalizer below). Same cache keys and compile order as + # the former per-variant blocks; write_partial_stats=False is the + # kernel's default, so passing it explicitly is a no-op. + kernel_kwargs = dict( + num_layers=num_layers, + seq_len=seq_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + num_freqs=num_freqs, + tokens_per_block=tokens_per_block, + pool_shape=tuple(int(value) for value in layer_pools[layer_indices[0]].shape), + pool_strides=tuple(int(value) for value in layer_pools[layer_indices[0]].stride()), + pool_dtype=cutlass.BFloat16, + ) + if self.enable_partial_stats: + variant_key = "triattention_cute_score_stats" + compiled_entries = self._compiled_stats + else: + variant_key = "triattention_cute_score" + compiled_entries = self._compiled for request_count, page_shards in variants: - if not self.enable_partial_stats: - # Per-head modes launch the score-only entry; the union - # runner compiles ONLY its fused stats+union pipeline below. - cache_key = ( - "triattention_cute_score", - static_geometry, - tensor_specs, - request_count, - page_shards, - ) - with _COMPILE_LOCK: - compiled = _COMPILED_KERNELS.get(cache_key) - if compiled is None: - kernel = _TriAttentionScoreKernel( - num_layers=num_layers, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - pool_shape=tuple( - int(value) for value in layer_pools[layer_indices[0]].shape - ), - pool_strides=tuple( - int(value) for value in layer_pools[layer_indices[0]].stride() - ), - pool_dtype=cutlass.BFloat16, - page_shards=page_shards, - ) - stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) - compiled = cute.compile( - kernel, - *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), - *self._cute_tail, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled - self._compiled[request_count] = compiled + cache_key = ( + variant_key, + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + kernel = _TriAttentionScoreKernel( + **kernel_kwargs, + page_shards=page_shards, + write_partial_stats=self.enable_partial_stats, + ) + stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) + compiled = cute.compile( + kernel, + *self._cute_prefix, + _to_cute(mean_cos.view(-1)), + _to_cute(mean_sin.view(-1)), + *self._cute_tail, + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled + compiled_entries[request_count] = compiled self._page_shards[request_count] = page_shards - if self.enable_partial_stats: - stats_cache_key = ( - "triattention_cute_score_stats", - static_geometry, - tensor_specs, - request_count, - page_shards, - ) - with _COMPILE_LOCK: - compiled_stats = _COMPILED_KERNELS.get(stats_cache_key) - if compiled_stats is None: - stats_kernel = _TriAttentionScoreKernel( - num_layers=num_layers, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, - num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - pool_shape=tuple( - int(value) for value in layer_pools[layer_indices[0]].shape - ), - pool_strides=tuple( - int(value) for value in layer_pools[layer_indices[0]].stride() - ), - pool_dtype=cutlass.BFloat16, - page_shards=page_shards, - write_partial_stats=True, - ) - stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) - compiled_stats = cute.compile( - stats_kernel, - *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), - *self._cute_tail, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[stats_cache_key] = compiled_stats - self._compiled_stats[request_count] = compiled_stats if max_requests > 1: - small_score = self._compiled.get(1) - large_score = self._compiled.get(max_requests) - small_stats = self._compiled_stats.get(1) - large_stats = self._compiled_stats.get(max_requests) + small = compiled_entries.get(1) + large = compiled_entries.get(max_requests) for request_count in range(1, max_requests + 1): # Shard-pick heuristic: give small cohorts the extra page # shard while the 2-shard grid stays under two waves # (2 * sm_count CTAs); larger cohorts already fill the GPU. two_shard_ctas = request_count * num_layers * num_kv_heads * 2 use_extra_score_shard = two_shard_ctas < 2 * self.sm_count - if not self.enable_partial_stats: - self._compiled[request_count] = ( - small_score if use_extra_score_shard else large_score - ) + compiled_entries[request_count] = small if use_extra_score_shard else large self._page_shards[request_count] = ( SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 ) - if self.enable_partial_stats: - self._compiled_stats[request_count] = ( - small_stats if use_extra_score_shard else large_stats - ) if self.enable_partial_stats: from .triattention_cute_selection import ( @@ -1819,14 +1773,6 @@ def launch_union_fusion( request_count, stream, ) - self._launch_union_finalize(request_count, union_scores, stream) - - def _launch_union_finalize( - self, - request_count: int, - union_scores: torch.Tensor, - stream: cuda.CUstream, - ) -> None: self._compiled_normalize_union[request_count]( _to_cute(self.partial_stats), *self._cute_selection_prefix, From 16ddb74b78609d4b7b1a92397ea2d8d2cac88442 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:03:54 -0700 Subject: [PATCH 102/178] [None][perf] Host-work census batch: per-step and per-round dispatch cuts (knife 21) Census #1-9, every verifier correction honored: - #1 the begin hook stores ONE batch reference; the growth constant (1 + _kv_reserve_draft_tokens) is cached at init and the membership id-set builds once per prepared batch (overlap-scheduler correction), not per step or per round. - #2+#7 _periodic_evict runs step accounting and the beta cadence gate FIRST; capacity/protected-tail/seq_len math and the consistency raises move INTO the due branch (guards relocated, never weakened). The write-only confirmed_kv_length ledger field dies repo-wide. - #3 the SWA destination rebase rides the phase-gather kernel behind a constexpr HAS_SWA branch (pointer captured at init; the no-SWA specialization compiles it away); the per-round torch.add dies. - #4 compact-launch args prebind at init (attach_compaction_bundle); the round loop is bare op calls. Honest note: the win is Python dict-chasing only. - #5+#9 prepared entries thread the resolved kv_cache/draft_kv_cache objects; _resize_compacted_requests calls .resize directly with only the is_active guard (both re-.get()+silent-skip pairs die). - #6 _validate_v2_compatibility memoizes behind _v2_validated (the loud raise still fires on the first request). - #8 the per-round SWA retained-span loop dies; the init-time budget >= window raise is the single enforcement point. Test fallout folded here (12 sites): the ledger helper and its consumers drop confirmed_kv_length (the draft co-compaction round detector now tracks the evicted_tokens delta with equivalent arithmetic), the growth test asserts the cached-constant contract, the overlap-tail test threads the resolved objects, the capacity-ledger test observes the due-branch seq_len, and the SWA e2e test attaches its standalone bundle through attach_compaction_bundle so the round consumes the bundle's prebinds. Signed-off-by: tianruih --- .../triattention/triattention.py | 310 +++++++++++------- .../triattention/triattention_kernels.py | 29 +- .../test_triattention_draft_cocompaction.py | 8 +- .../test_triattention_pipeline.py | 88 +++-- .../test_triattention_selection_compaction.py | 10 +- 5 files changed, 261 insertions(+), 184 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b281af413e02..1bb102b22124 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -67,7 +67,7 @@ import torch from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState, get_draft_token_length +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import BaseKVCacheCompressionManager from tensorrt_llm._utils import nvtx_range, nvtx_range_debug, prefer_pinned from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( @@ -95,6 +95,13 @@ # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) +# Generation requests skipped by every eviction step (hoisted: the resolve +# loop runs per request per decode step). +_SKIP_REQUEST_STATES = ( + LlmRequestState.GENERATION_COMPLETE, + LlmRequestState.CONTEXT_INIT, +) + # Upper bound of the geometric integration offset ladder [1, 2, 4, ...]; no # caller ever tuned it, so it is a constant rather than a constructor knob. _OFFSET_MAX_LENGTH = 65536 @@ -136,6 +143,87 @@ def _allocate_page_table_plane( return representative_slots, copy_block_count, host, dev +def attach_compaction_bundle(bufs: SimpleNamespace, compaction: Dict[str, object]) -> None: + """Flatten one compaction bundle into prebound launch fields. + + The bundle's families/draft_pack dicts stay as the build-time + description; the round loop fires bare op calls over the tuples + prebound here (the args are init-frozen, so per-round dict-chasing is + pure overhead). Tests attach standalone bundles through the same + helper. + """ + bufs.compaction_families = compaction["families"] + bufs.settle_pack_tensors = compaction["settle_pack_tensors"] + bufs.settle_pack_shape = compaction["settle_pack_shape"] + bufs.draft_pack = compaction["draft_pack"] + bufs.swa_destination_bases = compaction["swa_destination_bases"] + bufs.swa_rebase_delta = compaction["swa_rebase_delta"] + + def group_args(family): + return tuple( + ( + group["pools"], + group["pool_pointers"], + group["page_table"], + family["source"], + family["offsets"], + family["destination_bases"], + group["source_layer_indices"], + ) + for group in family["groups"] + ) + + target_args: List[tuple] = [] + draft_args: Tuple[tuple, ...] = () + for family in compaction["families"]: + if family["name"] == "draft": + # The draft's own pack launch must precede its moves; keeping the + # draft family separate preserves the pack-before-moves order. + draft_args = group_args(family) + else: + target_args.extend(group_args(family)) + bufs.compact_launch_args = tuple(target_args) + bufs.draft_compact_launch_args = draft_args + draft_pack = compaction["draft_pack"] + if draft_pack is None: + bufs.draft_pack_args = None + bufs.draft_pack_kwargs = None + return + # One more pack launch broadcasts the target keep set over the draft KV + # heads and appends the draft's own tail ordinals. HAS_SETTLE=False + # compiles the settle half away (the ordinals arrive pre-settled), so + # the settle-side pointer arguments are None; the pack half reads the + # settled ordinals through output_indices. + bufs.draft_pack_args = ( + None, + None, + None, + None, + bufs.keep, + bufs.valid_seq_lens_device, + draft_pack["offsets"], + draft_pack["indices"], + None, + None, + ) + bufs.draft_pack_kwargs = dict( + WIDTH=bufs.keep_count, + KEEP_COUNT=bufs.keep_count, + SELECTION_ROWS=1, + DENSE_TOTAL=draft_pack["dense_total"], + SWA_TOTAL=0, + MOVE_CAPACITY=draft_pack["move_capacity"], + NUM_KV_HEADS=draft_pack["num_kv_heads"], + SWA_WINDOW=0, + UNION=True, + PER_LAYER=False, + HAS_SWA=False, + HAS_SETTLE=False, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, + ) + + def init_eviction_buffers( *, eviction_mode: str, @@ -509,16 +597,11 @@ def init_eviction_buffers( swa_move_offsets=swa_move_offsets_row, **draft_kwargs, ) - # Flatten the launch data to plain fields: ONE fused launch settles the - # kept ordinals and packs the dense/SWA move sources + # Flatten the launch data to plain prebound fields: ONE fused launch + # settles the kept ordinals and packs the dense/SWA move sources # (``settle_top_tokens``); ``run_eviction_round`` fires the draft pack - # and every family's C++ moves directly. - bufs.compaction_families = compaction["families"] - bufs.settle_pack_tensors = compaction["settle_pack_tensors"] - bufs.settle_pack_shape = compaction["settle_pack_shape"] - bufs.draft_pack = compaction["draft_pack"] - bufs.swa_destination_bases = compaction["swa_destination_bases"] - bufs.swa_rebase_delta = compaction["swa_rebase_delta"] + # and every family's C++ moves as bare prebound calls. + attach_compaction_bundle(bufs, compaction) # ---- round-ordering events ---------------------------------------------- bufs.copy_done = torch.cuda.Event() @@ -733,6 +816,10 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: bufs.token_starts_device, bufs.valid_widths, request_count, + # The prompt offsets may have been re-staged since construction; + # the same launch rebases this round's SWA landing positions. + swa_destination_bases=bufs.swa_destination_bases, + rebase_delta=bufs.swa_rebase_delta, ) if union: bufs.runner.launch_union_fusion( @@ -789,54 +876,16 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: ) settle_top_tokens(bufs) with nvtx_range("triattention.compact", color="purple"): - if bufs.swa_destination_bases is not None: - # The prompt offsets may have been re-staged since construction; - # rebase the SWA landing positions for this round. - torch.add(bufs.prompt_offsets, bufs.swa_rebase_delta, out=bufs.swa_destination_bases) - for family in bufs.compaction_families: - if family["name"] == "draft": - # One more pack launch broadcasts the target keep set over - # the draft KV heads and appends the draft's own tail - # ordinals. HAS_SETTLE=False compiles the settle half away - # (the ordinals arrive pre-settled), so the settle-side - # pointer arguments are None; the pack half reads the - # settled ordinals through output_indices. - _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( - None, - None, - None, - None, - bufs.keep, - bufs.valid_seq_lens_device, - bufs.draft_pack["offsets"], - bufs.draft_pack["indices"], - None, - None, - WIDTH=bufs.keep_count, - KEEP_COUNT=bufs.keep_count, - SELECTION_ROWS=1, - DENSE_TOTAL=bufs.draft_pack["dense_total"], - SWA_TOTAL=0, - MOVE_CAPACITY=bufs.draft_pack["move_capacity"], - NUM_KV_HEADS=bufs.draft_pack["num_kv_heads"], - SWA_WINDOW=0, - UNION=True, - PER_LAYER=False, - HAS_SWA=False, - HAS_SETTLE=False, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, - ) - for group in family["groups"]: - torch.ops.trtllm.sparse_kv_cache_compact_layers( - group["pools"], - group["pool_pointers"], - group["page_table"], - family["source"], - family["offsets"], - family["destination_bases"], - group["source_layer_indices"], - ) + # Bare prebound calls (``attach_compaction_bundle``); the draft's own + # pack launch precedes its moves. + for args in bufs.compact_launch_args: + torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) + if bufs.draft_pack_args is not None: + _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( + *bufs.draft_pack_args, **bufs.draft_pack_kwargs + ) + for args in bufs.draft_compact_launch_args: + torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) class TriAttention(BaseKVCacheCompressionManager): @@ -924,13 +973,21 @@ def __init__( self._phase: Optional[Dict[str, object]] = None # Request presence records successful initialization. Each value is a - # plain dict {generation_steps, evicted_tokens, confirmed_kv_length}. + # plain dict {generation_steps, evicted_tokens}. self._request_states: Dict[int, Dict[str, object]] = {} - # The overlap executor prepares B(n) before finalizing B(n-1). Keep the - # exact fixed-linear generation width for that currently in-flight - # batch as ``(batch, {request_id: growth})``; the final hook treats - # those slots as an opaque suffix. - self._prepared_generation_batch: Optional[Tuple[object, Dict[int, int]]] = None + # The overlap executor prepares B(n) before finalizing B(n-1). Keep a + # bare reference to that in-flight batch; its membership id-set and + # the fixed-linear growth constant (1 + the reserved draft width, + # which bounds every step's actual draft) are resolved lazily on the + # first overlap miss. The final hook treats those slots as an opaque + # suffix. + self._prepared_generation_batch: Optional[object] = None + self._prepared_generation_ids: Optional[set] = None + self._generation_growth: Optional[int] = None + # Manager-invariant validation and tail capacity are memoized (the + # loud raises still fire on the first request). + self._v2_validated = False + self._protected_tail_cache: Optional[int] = None # The eviction buffers are built once at the first eviction, sized to # capacity bounds, and reused for the manager's lifetime. self._buffers: Optional[SimpleNamespace] = None @@ -957,7 +1014,6 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: self._request_states[request_id] = { "generation_steps": 0, "evicted_tokens": 0, - "confirmed_kv_length": None, } self._ensure_calibrated() @@ -1039,7 +1095,13 @@ def _ensure_calibrated(self) -> None: self._calibrated = True def _validate_v2_compatibility(self) -> None: - """Reject runtime modes outside the V2 physical-compaction contract.""" + """Reject runtime modes outside the V2 physical-compaction contract. + + Manager-invariant, so the O(num_layers) scans run once; the loud + raise still fires on the first request. + """ + if self._v2_validated: + return manager = self.kv_cache_manager if not isinstance(manager, KVCacheManagerV2): raise ValueError("TriAttention physical eviction requires KVCacheManagerV2") @@ -1098,6 +1160,7 @@ def _validate_v2_compatibility(self) -> None: "TriAttention requires full-attention V2 lifecycles; native SWA, " "VSWA, and SSM pools are not supported" ) + self._v2_validated = True # The framework drives all request-lifecycle hooks. TriAttention resolves # calibration on request init, evicts periodically at generation-step end, @@ -1121,25 +1184,33 @@ def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) self._periodic_evict(scheduled_batch) def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Snapshot fixed-linear target growth; mutation remains in final update.""" - generation_growth = {} - for request in scheduled_batch.generation_requests: - request_id = request.py_request_id - growth = 1 + max( - get_draft_token_length(request), - self.kv_cache_manager._kv_reserve_draft_tokens, - ) - generation_growth[request_id] = growth - self._prepared_generation_batch = (scheduled_batch, generation_growth) + """Snapshot the prepared batch; mutation remains in final update.""" + self._prepared_generation_batch = scheduled_batch + self._prepared_generation_ids = None def _inflight_generation_growth( self, scheduled_batch: "ScheduledRequests", request_id: int ) -> int: - """Return exact newer target allocation width under overlap scheduling.""" + """Return exact newer target allocation width under overlap scheduling. + + The width is the fixed-linear constant ``1 + reserved draft`` for + members of the prepared batch (the reserve bounds every step's + actual draft; the configured-tail guard enforces the bound). + """ prepared = self._prepared_generation_batch - if prepared is None or scheduled_batch is prepared[0]: + if prepared is None or scheduled_batch is prepared: return 0 - return prepared[1].get(request_id, 0) + member_ids = self._prepared_generation_ids + if member_ids is None: + member_ids = {request.py_request_id for request in prepared.generation_requests} + self._prepared_generation_ids = member_ids + if request_id not in member_ids: + return 0 + growth = self._generation_growth + if growth is None: + growth = 1 + int(self.kv_cache_manager._kv_reserve_draft_tokens) + self._generation_growth = growth + return growth def _periodic_evict( self, @@ -1153,10 +1224,7 @@ def _periodic_evict( mgr = self.kv_cache_manager resolved_requests = [] for request in gen_requests: - if request.is_dummy or request.state in ( - LlmRequestState.GENERATION_COMPLETE, - LlmRequestState.CONTEXT_INIT, - ): + if request.is_dummy or request.state in _SKIP_REQUEST_STATES: continue request_id = request.py_request_id kv_cache = mgr.kv_cache_map.get(request_id) @@ -1180,11 +1248,22 @@ def _periodic_evict( protected_tail_capacity = self._configured_protected_tail_capacity() # Resolve every active target cache before changing cadence state (the - # captured cache objects also avoid repeating the V2 map lookup here), - # building the due cohort's per-request eviction metadata in the same - # pass -- ``_evict_requests`` trusts it as-is. + # captured cache objects thread all the way to resize -- V2 map + # lookups happen once per step), building the due cohort's + # per-request eviction metadata in the same pass -- + # ``_evict_requests`` trusts it as-is. with nvtx_range("triattention.metadata", color="cyan"): for request, request_id, kv_cache in resolved_requests: + # Step accounting and the beta cadence gate come first: a + # non-due request costs two dict ops and two int ops. The + # capacity/tail math and the consistency raises run in the + # due branch, at the point their values are consumed. + request_state = self._request_states[request_id] + previous_step = request_state["generation_steps"] + step = previous_step + 1 + int(request.py_num_accepted_draft_tokens) + request_state["generation_steps"] = step + if previous_step // self.beta >= step // self.beta: + continue raw_capacity = int(kv_cache.capacity) # One-engine speculative decoding keeps a fixed reserve E. # Under overlap, B(n) is allocated/enqueued before finalizing @@ -1211,17 +1290,10 @@ def _periodic_evict( f"Request {request_id} KV length {seq_len} is below finalized " f"history {kv_cache.history_length}" ) - request_state = self._request_states[request_id] - request_state["confirmed_kv_length"] = seq_len - previous_step = request_state["generation_steps"] - confirmed_delta = 1 + int(request.py_num_accepted_draft_tokens) - step = previous_step + confirmed_delta - request_state["generation_steps"] = step - if previous_step // self.beta >= step // self.beta: - continue expected_keep_count = self._minimum_evictable_length(request, seq_len) if seq_len <= expected_keep_count: continue + draft_kv_cache = None if self.draft_kv_cache_manager is not None: draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) if draft_kv_cache is None or not draft_kv_cache.is_active: @@ -1235,6 +1307,8 @@ def _periodic_evict( { "request": request, "request_id": request_id, + "kv_cache": kv_cache, + "draft_kv_cache": draft_kv_cache, "seq_len": int(seq_len), # Restore the uncompressed confirmed logical position # from the physical prefix and cumulative eviction @@ -1266,15 +1340,19 @@ def _periodic_evict( self._resize_compacted_requests(capacity_targets, protected_tails) def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: + """Release each compacted tail through the caches resolved this hook. + + ``capacity_targets`` threads the ``(request_id, kv_cache, + draft_kv_cache, keep_count)`` tuples resolved by ``_periodic_evict`` + within this same synchronous hook, so no V2 map re-lookup happens + here; only the cheap is_active guard remains. + """ if not capacity_targets: return - mgr = self.kv_cache_manager - draft_manager = self.draft_kv_cache_manager with nvtx_range("triattention.resize", color="red"): with nvtx_range_debug("triattention.v2_resize", color="red"): - for rid, target_capacity in capacity_targets: - kv_cache = mgr.kv_cache_map.get(rid) - if kv_cache is None or not kv_cache.is_active: + for rid, kv_cache, _, target_capacity in capacity_targets: + if not kv_cache.is_active: continue if target_capacity > kv_cache.capacity: raise RuntimeError( @@ -1288,14 +1366,13 @@ def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: f"Failed to resize compacted KV cache for request {rid} " f"to {resized_capacity} tokens" ) - if draft_manager is not None: + if self.draft_kv_cache_manager is not None: # The draft cache was compacted with the same kept token # set, so it shrinks to the same retained length plus its # own protected tail. draft_protected_tail = self._draft_protected_tail_capacity() - for rid, target_capacity in capacity_targets: - draft_kv_cache = draft_manager.kv_cache_map.get(rid) - if draft_kv_cache is None or not draft_kv_cache.is_active: + for rid, _, draft_kv_cache, target_capacity in capacity_targets: + if not draft_kv_cache.is_active: continue draft_capacity = target_capacity + draft_protected_tail if not draft_kv_cache.resize(draft_capacity, None): @@ -1346,15 +1423,13 @@ def _local_score_calibration( def _configured_protected_tail_capacity(self) -> int: """Return the largest target tail reserved by the native V2 lifecycle.""" - return _protected_tail_capacity(self.kv_cache_manager, "") + if self._protected_tail_cache is None: + self._protected_tail_cache = _protected_tail_capacity(self.kv_cache_manager, "") + return self._protected_tail_cache def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: """Drop this request's per-request length and eviction state.""" - request_id = request.py_request_id - self._request_states.pop(request_id, None) - prepared = self._prepared_generation_batch - if prepared is not None: - prepared[1].pop(request_id, None) + self._request_states.pop(request.py_request_id, None) # The buffers stay resident across idle periods: their memory is a # deliberate one-time cost and rebuilding it per burst would reintroduce # allocation on the decode hot path. @@ -1881,17 +1956,9 @@ def _evict_requests( with nvtx_range_debug("triattention.resolve_layout", color="blue"): layout = self._runtime_kv_layout(num_layers) with nvtx_range_debug("triattention.staging_lookup", color="blue"): + # Retained spans always cover the model window: construction + # rejects budget < window, and the pinned prompt only adds. bufs = self._buffers_for(layout, prepared) - if layout["swa_layers"] and layout["swa_window"]: - # SWA landing positions are prompt-dependent; reject a request - # whose retained span cannot cover the model window this round. - for item in prepared: - if item["prompt_len"] + self.budget < int(layout["swa_window"]): - raise ValueError( - f"Request {item['request_id']} retains " - f"{item['prompt_len'] + self.budget} tokens, below the " - f"sliding window {layout['swa_window']}" - ) with nvtx_range_debug("triattention.page_table_stage", color="orange"): dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( layout, prepared, bufs.max_requests @@ -1925,13 +1992,14 @@ def _evict_requests( raise RuntimeError("TriAttention attempted an identity compaction") request_state = self._request_states[item["request_id"]] request_state["evicted_tokens"] += evicted - request_state["confirmed_kv_length"] = keep_count # Publish the cumulative count on the request: this is the # manager's only channel to the runtime. The model engine # reads it back where it builds num_cached_tokens_per_seq, # so the kernels see the compacted KV length next step. item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] - capacity_targets.append((item["request_id"], keep_count)) + capacity_targets.append( + (item["request_id"], item["kv_cache"], item["draft_kv_cache"], keep_count) + ) return capacity_targets def _num_layers_from_manager(self) -> int: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 33602d57bdf8..1d25857f1b2e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -21,7 +21,7 @@ from __future__ import annotations -from typing import Dict +from typing import Dict, Optional import torch import triton @@ -54,14 +54,19 @@ def _gather_mean_phase_kernel( valid_seq_lens, token_starts, valid_widths, + swa_destination_bases, table_rows, + rebase_delta, NUM_FREQS: tl.constexpr, F_BLOCK: tl.constexpr, + HAS_SWA: tl.constexpr, ): """Copy each request's precomputed phase-table row into the fixed buffers. The same launch derives the request's valid decode width (valid length - minus window start) so the round needs no separate subtraction launch. + minus window start), and with SWA layers also rebases the round's SWA + landing positions (window start plus the init-frozen delta), so the + round needs no separate subtraction or rebase launches. """ request = tl.program_id(0) frequency = tl.arange(0, F_BLOCK) @@ -76,8 +81,10 @@ def _gather_mean_phase_kernel( row_sin = tl.load(table_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) - width = tl.load(valid_seq_lens + request) - tl.load(token_starts + request) - tl.store(valid_widths + request, width) + token_start = tl.load(token_starts + request) + tl.store(valid_widths + request, tl.load(valid_seq_lens + request) - token_start) + if HAS_SWA: + tl.store(swa_destination_bases + request, token_start + rebase_delta) def build_mean_phase_table( @@ -141,12 +148,17 @@ def gather_mean_phases( token_starts: torch.Tensor, valid_widths: torch.Tensor, request_count: int, + *, + swa_destination_bases: Optional[torch.Tensor] = None, + rebase_delta: int = 0, ) -> None: """Refresh the fixed mean buffers and valid widths from staged metadata. - Writes in place because the compiled CuTe score launch captured the - destination buffers' device pointers. CUDA-only; eviction never runs - under CUDA graph capture. + With SWA layers the same launch rebases ``swa_destination_bases`` + (window start + ``rebase_delta``); without them the HAS_SWA branch is + compiled away and the pointer argument is None. Writes in place because + the compiled CuTe score launch captured the destination buffers' device + pointers. CUDA-only; eviction never runs under CUDA graph capture. """ num_freqs = phase["omega"].numel() _gather_mean_phase_kernel[(request_count,)]( @@ -158,9 +170,12 @@ def gather_mean_phases( valid_seq_lens, token_starts, valid_widths, + swa_destination_bases, phase["rows"], + 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, ) 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 index a1d392da03f5..cb30b27c87a4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -38,7 +38,7 @@ def _fresh_request_state(): """One request's compression ledger, as the manager initializes it.""" - return {"generation_steps": 0, "evicted_tokens": 0, "confirmed_kv_length": None} + return {"generation_steps": 0, "evicted_tokens": 0} def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: @@ -300,6 +300,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): confirmed = uncompressed cache.capacity = confirmed previous_published = 0 + previous_evicted = 0 eviction_rounds = 0 with _mocked_eviction_internals(manager) as internals: for _ in range(6): @@ -310,10 +311,11 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager._periodic_evict(batch) state = manager._request_states[7] - if state["confirmed_kv_length"] < confirmed: + if state["evicted_tokens"] > previous_evicted: # An eviction round compacted the cache to prompt + budget. eviction_rounds += 1 - confirmed = state["confirmed_kv_length"] + confirmed -= state["evicted_tokens"] - previous_evicted + previous_evicted = state["evicted_tokens"] cache.capacity = confirmed assert confirmed == 2 + 4 # The staged logical position restores the uncompressed diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 3ec78a871ffc..7758810a1e1e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -60,18 +60,10 @@ ) -def _set_request_state( - manager, - request_id, - *, - generation_steps=0, - evicted_tokens=0, - confirmed_kv_length=None, -): +def _set_request_state(manager, request_id, *, generation_steps=0, evicted_tokens=0): state = { "generation_steps": generation_steps, "evicted_tokens": evicted_tokens, - "confirmed_kv_length": confirmed_kv_length, } manager._request_states[request_id] = state return state @@ -90,6 +82,8 @@ def _prepared_eviction( return { "request": request, "request_id": request_id, + "kv_cache": None, + "draft_kv_cache": None, "seq_len": seq_len, "round_start": int(seq_len if round_start is None else round_start), "prompt_len": prompt_len, @@ -427,27 +421,22 @@ def test_suspended_cache_rejects_batch_before_cadence_mutation(self): manager._periodic_evict(batch) assert first_state["generation_steps"] == 127 - assert first_state["confirmed_kv_length"] is None assert second_state["generation_steps"] == 127 - assert second_state["confirmed_kv_length"] is None def test_request_finish_clears_state_but_keeps_buffers_resident(self): manager = _make_triattention() - _set_request_state( - manager, - 7, - generation_steps=1, - evicted_tokens=127, - confirmed_kv_length=128, - ) + _set_request_state(manager, 7, generation_steps=1, evicted_tokens=127) buffers = object() manager._buffers = buffers - manager._prepared_generation_batch = (SimpleNamespace(), {7: 1}) + batch = SimpleNamespace() + manager._prepared_generation_batch = batch manager.on_request_finish(_make_request(7)) assert manager._request_states == {} - assert manager._prepared_generation_batch[1] == {} + # The batch reference survives finish (it belongs to the step, not + # the request). + assert manager._prepared_generation_batch is batch # The buffers are sized for the executor limits, not one cohort, so # they stay resident for the next generation batch. assert manager._buffers is buffers @@ -465,12 +454,10 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): cache.capacity = confirmed + tail mgr.kv_cache_manager.num_extra_kv_tokens = reserve # The configured tail capacity (reserve + draft reserve + 1) must - # cover this round's actual tail, as in production. - mgr.kv_cache_manager._kv_reserve_draft_tokens = current_growth - mgr._prepared_generation_batch = ( - SimpleNamespace(generation_requests=[request]), - {7: current_growth}, - ) + # cover this round's actual tail, as in production; the growth + # constant is 1 + _kv_reserve_draft_tokens for batch members. + mgr.kv_cache_manager._kv_reserve_draft_tokens = current_growth - 1 + mgr._prepared_generation_batch = 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} @@ -478,8 +465,7 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): mgr.draft_kv_cache_manager = draft_manager def compact(*_args, **_kwargs): - mgr._request_states[7]["confirmed_kv_length"] = retained - return [(7, retained)] + return [(7, cache, draft_cache, retained)] with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: mgr._periodic_evict(batch) @@ -491,6 +477,8 @@ def compact(*_args, **_kwargs): { "request": request, "request_id": 7, + "kv_cache": cache, + "draft_kv_cache": draft_cache, "seq_len": confirmed, "round_start": confirmed, "prompt_len": 1024, @@ -500,17 +488,18 @@ def compact(*_args, **_kwargs): ], 2, ) - assert mgr._request_states[7]["confirmed_kv_length"] == retained cache.resize.assert_called_once_with(retained + tail, None) # The draft cache shrinks in the same round, to the same retained # length plus the draft's own protected tail. 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 seq_len must come from the physical capacity ledger + # (capacity minus the protected tail), never the logical length. physical_confirmed = 6100 manager = _make_triattention(beta=128) manager._calibrated = True - _set_request_state(manager, 7, evicted_tokens=100) + _set_request_state(manager, 7, generation_steps=127, evicted_tokens=100) cache = SimpleNamespace( capacity=physical_confirmed, history_length=1024, @@ -527,9 +516,11 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): py_draft_tokens=[1, 2, 3, 4], ) - manager._periodic_evict(SimpleNamespace(generation_requests=[request])) + with mock.patch.object(manager, "_evict_requests", return_value=[]) as evict: + manager._periodic_evict(SimpleNamespace(generation_requests=[request])) - assert manager._request_states[7]["confirmed_kv_length"] == physical_confirmed + prepared = evict.call_args.args[0] + assert prepared[0]["seq_len"] == physical_confirmed cache.resize.assert_not_called() def test_one_model_draft_co_compression_contract_is_accepted(self): @@ -562,19 +553,18 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): ) @pytest.mark.parametrize( - "num_extra_kv_tokens,reserved_draft,draft_tokens,expected_growth", + "reserved_draft,expected_growth", [ - (2, 0, [1, 2, 3], 4), - # The reserved draft width protects capacity even when this step's - # actual draft is shorter. - (0, 6, [1, 2], 7), + (0, 1), + # The reserved draft width protects capacity regardless of any + # step's actual draft length: growth is the cached constant. + (6, 7), ], ) def test_prepare_snapshots_fixed_linear_generation_growth( - self, num_extra_kv_tokens, reserved_draft, draft_tokens, expected_growth + self, reserved_draft, expected_growth ): manager = _make_fake_v2() - manager.num_extra_kv_tokens = num_extra_kv_tokens manager._kv_reserve_draft_tokens = reserved_draft manager.kv_cache_map = { 7: SimpleNamespace(capacity=106, is_active=True), @@ -582,13 +572,17 @@ def test_prepare_snapshots_fixed_linear_generation_growth( triattention = TriAttention(manager, budget=8, model_path="/models/test") batch = SimpleNamespace( context_requests=[], - generation_requests=[_make_request(7, py_draft_tokens=draft_tokens)], + generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], ) triattention.prepare_resources(batch) - assert triattention._prepared_generation_batch[0] is batch - assert triattention._prepared_generation_batch[1] == {7: expected_growth} + assert triattention._prepared_generation_batch is batch + # Members of the prepared batch grow by the cached constant; others + # by zero. The prepared batch itself is the identity early-out. + assert triattention._inflight_generation_growth(SimpleNamespace(), 7) == expected_growth + assert triattention._inflight_generation_growth(SimpleNamespace(), 99) == 0 + assert triattention._inflight_generation_growth(batch, 7) == 0 class TestFixedScoreMetadata: @@ -728,8 +722,8 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ) first = _make_request(7, py_prompt_len=3) second = _make_request(8, py_prompt_len=5) - _set_request_state(manager, 7, confirmed_kv_length=8) - _set_request_state(manager, 8, confirmed_kv_length=10) + _set_request_state(manager, 7) + _set_request_state(manager, 8) with _mocked_eviction_internals(manager) as internals: manager._evict_requests( @@ -867,9 +861,9 @@ def stage_round(): ) score_sentinel = -12345.0 - # Isolate the score stage: an empty family list makes the compact - # stage a no-op (this synthetic round never stages move offsets). - bufs.compaction_families = [] + # Isolate the score stage: empty prebound launch args make the + # compact stage a no-op (this synthetic round never stages offsets). + bufs.compact_launch_args = () stage_round() bufs.score_output.fill_(score_sentinel) module.run_eviction_round(bufs, normalize_scores=False) 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 index f2825691a052..84ce4fe2316c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -456,6 +456,7 @@ 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") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + attach_compaction_bundle, init_eviction_buffers, run_eviction_round, ) @@ -559,12 +560,9 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): protected_tail_capacity=0, ) _set_protected_tails(compaction, [0]) - bufs.compaction_families = compaction["families"] - bufs.settle_pack_tensors = compaction["settle_pack_tensors"] - bufs.settle_pack_shape = compaction["settle_pack_shape"] - bufs.swa_destination_bases = compaction["swa_destination_bases"] - bufs.swa_rebase_delta = compaction["swa_rebase_delta"] - bufs.draft_pack = compaction["draft_pack"] + # Attach wholesale: the round consumes the PREBOUND launch args, so the + # bundle must land through the same helper production uses. + attach_compaction_bundle(bufs, compaction) run_eviction_round(bufs, normalize_scores=False) assert torch.equal(bufs.keep, expected_keep) torch.cuda.synchronize(device) From d069359b3fd3fdb8f4491ea03f057831b66f8b7c Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:08:25 -0700 Subject: [PATCH 103/178] [None][refactor] Fold the eight raw-band TMA copy rituals into one trace-time helper (E1) The hot producer loop repeated the same TMA copy ritual eight times (real/imag bands x current/next page x both dynamic_expr arms); the blocks fold into _stage_raw_band_copies, a trace-time helper parameterizing the byte offsets and stage slices. The pipeline handshake is deliberately NOT folded: every producer_acquire / producer_commit / advance stays inline at its original site (the double-buffer state is shared across the dynamic_expr arms), so the acquire/advance sequence per specialization is auditable as unchanged -- the diff contains zero handshake lines. Net -39 (honest shortfall vs the -90 estimate: no nested closures in @cute.jit and no handshake folding, both refusals per the deadlock caveat). Perf receipt rides the knife-18 captures. Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 282 ++++++++---------- 1 file changed, 122 insertions(+), 160 deletions(-) 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 index 07c68e4896d8..6db07534b07a 100644 --- 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 @@ -220,6 +220,46 @@ def epilog_gmem_copy_and_partition( copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) return copy_atom, register_output, thread_output + @cute.jit + def _stage_raw_band_copies( + self, + raw_tma_pipeline, + raw_tma_producer_state, + band, + first_page, + page_fragments, + shared_partition, + stage_args, + ): + """One raw-K band's fragment copies into an acquired pipeline stage. + + The CALLER owns the pipeline handshake (producer_acquire before, + state.advance after): the double-buffered prefetch shares ONE + acquire/advance across its dynamic destination arms, so the + handshake cannot live here. ``band`` is the trace-time real/imag + plane index; ``page_fragments`` is only read for the multi-fragment + specialization (None otherwise). + """ + 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, @@ -584,6 +624,9 @@ def kernel( (raw_tma_descriptors.iterator + layer_id * TMA_DESCRIPTOR_QWORDS).align(128), cute.AddressSpace.generic, ) + # Trace-time invariants of every raw-band stage copy (see + # _stage_raw_band_copies); 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, @@ -606,6 +649,8 @@ def kernel( 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 for multi-page compute tiles. # Slot 0 is unused: fragment 0 keeps the scalar broadcast @@ -615,6 +660,8 @@ def kernel( ) physical_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) prefetched_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) + physical_fragments_arg = physical_page_fragments + prefetched_fragments_arg = prefetched_page_fragments shard_has_page = valid_seq_len > score_start and tile_start_token < valid_seq_len empty_shard = valid_seq_len <= score_start or tile_start_token >= valid_seq_len if cutlass.dynamic_expr(shard_has_page): @@ -787,48 +834,26 @@ def kernel( 0, ) raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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) - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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 ( valid_seq_len > score_start @@ -859,26 +884,15 @@ def kernel( # Every producer-warp lane participates in the # PipelineTmaAsync barrier election. raw_tma_pipeline.producer_acquire(raw_tma_producer_state) - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = physical_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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) @@ -886,26 +900,15 @@ def kernel( 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) - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = physical_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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) @@ -989,94 +992,53 @@ def kernel( 0, ) if cutlass.dynamic_expr(prefetch_next_raw): + # ONE acquire/advance per band is SHARED across the + # dynamic destination arms (current vs next buffer); + # only the smem destination differs per arm. 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): - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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: - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 0, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_real_next[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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): - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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: - for fragment in cutlass.range_constexpr(self.fragments_per_phase): - fragment_page = prefetched_physical_page - if cutlass.const_expr(fragment > 0): - fragment_page = prefetched_page_fragments[fragment] - cute.copy( - raw_tma_atom, - raw_tma_global_partition[ - ( - None, - 1, - 0, - (kv_head, fragment_page), - ) - ], - raw_tma_shared_partition_imag_next[fragment], - tma_bar_ptr=raw_tma_pipeline.producer_get_barrier( - raw_tma_producer_state - ), - tma_desc_ptr=raw_tma_descriptor_ptr, - ) + 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 of the 32 lanes stages one frequency per pass; 64- # frequency heads take two passes. From 4c39809fdbdcee4847cc5176b3fbe0fcd634e25e Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:28:37 -0700 Subject: [PATCH 104/178] [None][chore] Restore the draft-cache property and short hook docstrings Signed-off-by: tianruih --- .../_torch/pyexecutor/resource_manager.py | 18 +++++++----------- .../test_kv_cache_compression_manager.py | 2 +- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index d983613fc995..157e2d7d720a 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2462,6 +2462,10 @@ def __init__( draft_kv_cache_manager.kv_compression_manages_history = ( self.adjusts_generation_kv_length) + @property + def has_independent_draft_kv_cache(self) -> bool: + return self.draft_kv_cache_manager is not None + # ================================================================== # # KV-cache lifecycle hooks (5, in temporal order). # # Subclasses override what they need; all default to no-op. # @@ -2486,23 +2490,15 @@ def on_generation_step_begin( scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per executor iteration, before this step's forward. - - The batch may contain only context requests and - ``scheduled_batch.generation_requests`` may be empty. ``**kwargs`` - is reserved for forward compatibility (never populated today). - """ + """Fired once per generation step before this step's forward.""" def on_generation_step_end( self, scheduled_batch: "ScheduledRequests", **kwargs, ) -> None: - """Fired once per executor iteration, after every layer's forward - completes. The batch may contain only context requests and - ``scheduled_batch.generation_requests`` may be empty; ``**kwargs`` - is reserved for forward compatibility. Override for periodic or - budget-triggered eviction. + """Fired once per generation step, after every layer's forward + completes. Override for periodic or budget-triggered eviction. """ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: 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 a074504f3dd8..4e1c70f758dd 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -149,7 +149,7 @@ def test_length_adjustment_marks_target_and_draft_v2(self): assert manager.kv_cache_manager is target assert manager.draft_kv_cache_manager is draft - assert manager.draft_kv_cache_manager is not None + assert manager.has_independent_draft_kv_cache assert target.kv_compression_manages_history is True assert draft.kv_compression_manages_history is True From 3f28019a45f25ce9ec1d30a694e740a924b0fc4c Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:31:31 -0700 Subject: [PATCH 105/178] [None][chore] Revert the config-time normalize_scores validator to the manager raise Signed-off-by: tianruih --- tensorrt_llm/llmapi/llm_args.py | 11 +- .../test_triattention_pipeline.py | 108 ++++++++++-------- 2 files changed, 62 insertions(+), 57 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 462f49d7486b..cf2ba9ec9b75 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3450,9 +3450,8 @@ class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): normalize_scores: bool = Field( default=True, description="Z-normalize each head's scores over the decode region " - "before selection (upstream default). `union` eviction always " - "normalizes; False is only valid with `per_head`/`per_layer_perhead`." - ) + "before selection (upstream default). `union` eviction requires True: " + "its fused score+stats+union pipeline always normalizes.") pin_prefill: bool = Field( default=True, description="Always preserve the prompt (prefill) tokens; only decode " @@ -3495,12 +3494,6 @@ def _require_calibration_inputs(self): "TriAttention requires both model_path and calibration_path; " "TRT-LLM consumes an official calibration file and does not " "compute one.") - # Same config-validation-time surfacing for the cross-field contract - # (the manager re-raises at construction as the op-boundary backstop). - if self.eviction_mode == "union" and not self.normalize_scores: - raise ValueError( - "union eviction always normalizes scores; normalize_scores=" - "False is only valid with per_head/per_layer_perhead.") return self def to_manager_kwargs(self) -> dict: diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 7758810a1e1e..f3bcd5557304 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -15,16 +15,10 @@ """Unit tests for the TriAttention compression-manager pipeline. -TriAttention is a pure KV-cache compression method on the PR-15106 framework: it -has NO sparse-attention config and NO attention backend of its own. Decode runs -the model's standard attention over the compacted cache; the manager publishes -the cumulative evicted count on ``LlmRequest.py_num_compressed_tokens`` and the -model engine subtracts it where it builds ``num_cached_tokens_per_seq``. These -tests cover the config, construction, eviction lifecycle, page-table staging, -and the fixed score buffers. Draft co-compression contracts live in -``test_triattention_draft_cocompaction.py``; model-level correctness is covered -by separate end-to-end tests. -""" +Config, construction, eviction lifecycle, page-table staging, and the fixed +score buffers; 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 @@ -140,15 +134,6 @@ def test_llm_args_dispatch_and_validation(self): ) with pytest.raises(ValidationError): TriAttentionKvCacheCompressionConfig(eviction_mode="made_up_mode") - # Cross-field contract surfaces at config validation (the manager - # re-raises at construction as the op-boundary backstop). - with pytest.raises(ValidationError, match="normalize"): - TriAttentionKvCacheCompressionConfig( - model_path="/models/test", - calibration_path="/calib/test.pt", - eviction_mode="union", - normalize_scores=False, - ) def test_factory_returns_triattention_and_propagates_config_fields(self): # Calibration is deferred to the first request, so construction needs @@ -215,10 +200,9 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): mock.call(101, Role.KEY), ] - def test_request_init_marks_capacity_only_and_tracks_state(self): - # Speculative capacity (extra KV tokens / reserved draft width) is - # accepted at request init; the target manager is marked so V2 sizing - # keeps logical max_seq_len while capacity is reclaimed and reused. + def test_request_init_and_finish_lifecycle(self): + # Init: speculative capacity accepted, manager marked, state tracked. + # Finish: state cleared; buffers and the step's batch stay resident. manager = _make_fake_v2() manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 @@ -233,6 +217,16 @@ def test_request_init_marks_capacity_only_and_tracks_state(self): assert manager.kv_compression_manages_history assert set(triattention._request_states) == {11, 12} + buffers = object() + triattention._buffers = buffers + batch = SimpleNamespace() + triattention._prepared_generation_batch = batch + triattention.on_request_finish(_make_request(11)) + triattention.on_request_finish(_make_request(12)) + assert triattention._request_states == {} + assert triattention._prepared_generation_batch is batch + assert triattention._buffers is buffers + def test_resolve_accepts_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() mgr.calibration_path = flat_calibration_pt @@ -241,6 +235,46 @@ def test_resolve_accepts_flat_pt(self, flat_calibration_pt): for key in ("E_q", "E_q_norm", "omega", "freq_scale_sq"): assert key in loaded + def test_resolve_converts_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_theta=10000.0) + + with mock.patch("transformers.AutoConfig.from_pretrained", return_value=config): + converted = mgr._resolve_calibration() + + assert set(converted) == {"E_q", "E_q_norm", "omega", "freq_scale_sq"} + assert converted["E_q"].shape == (num_layers, num_heads, freq_count) + torch.testing.assert_close( + converted["E_q"][1, 0].cpu(), + torch.complex(torch.full((freq_count,), 10.0), torch.full((freq_count,), 1.0)), + ) + torch.testing.assert_close( + converted["E_q_norm"][1, 1].cpu(), torch.full((freq_count,), 3.0) + ) + assert converted["omega"].numel() == freq_count + idx = torch.arange(0, 2 * freq_count, 2, dtype=torch.float32) + torch.testing.assert_close( + converted["omega"].cpu(), 1.0 / (10000.0 ** (idx / (2 * freq_count))) + ) + assert torch.equal(converted["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 @@ -306,8 +340,7 @@ def test_identity_compaction_is_rejected_instead_of_published(self): manager = _make_triattention(budget=4) manager.kv_cache_manager._stream = mock.Mock() request = _make_request(7, py_prompt_len=2) - # Selection keeps every token at seq_len == prompt + budget: the due - # filter must drop the request before any eviction work or publication. + # 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) ) @@ -344,9 +377,7 @@ def test_identity_compaction_is_rejected_instead_of_published(self): class TestEvictionLifecycle: def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): - # Structural: TriAttention implements hooks only, never the base - # template methods; the growth snapshot rides the step-begin hook and - # the eviction runs from the final on_generation_step_end hook. + # Hooks only, never the base template methods. assert "prepare_resources" not in TriAttention.__dict__ assert "update_resources" not in TriAttention.__dict__ assert "on_generation_step_begin" in TriAttention.__dict__ @@ -367,8 +398,7 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): periodic_evict.assert_called_once_with(batch) def test_unregistered_generation_request_is_rejected(self): - # A generation request whose on_request_init never ran is a framework - # ordering bug: it fails loudly instead of late-initializing here. + # Missing on_request_init = framework ordering bug: fail loudly. manager = _make_triattention() manager._calibrated = True cache = SimpleNamespace(capacity=8, history_length=0, is_active=True) @@ -423,24 +453,6 @@ def test_suspended_cache_rejects_batch_before_cadence_mutation(self): assert first_state["generation_steps"] == 127 assert second_state["generation_steps"] == 127 - def test_request_finish_clears_state_but_keeps_buffers_resident(self): - manager = _make_triattention() - _set_request_state(manager, 7, generation_steps=1, evicted_tokens=127) - buffers = object() - manager._buffers = buffers - batch = SimpleNamespace() - manager._prepared_generation_batch = batch - - manager.on_request_finish(_make_request(7)) - - assert manager._request_states == {} - # The batch reference survives finish (it belongs to the step, not - # the request). - assert manager._prepared_generation_batch is batch - # The buffers are sized for the executor limits, not one cohort, so - # they stay resident for the next generation batch. - assert manager._buffers is buffers - @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): confirmed = 1024 + 4096 + 1 + accepted From af1909ea75d59311b2d87772aca740b0ba21dd2d Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:42:02 -0700 Subject: [PATCH 106/178] [None][test] Test tranche 3 wave 2: merges, narration trims, official-layout conversion test Signed-off-by: tianruih --- .../_torch/kv_cache_compression/conftest.py | 73 ++++---------- .../test_triattention_cute_score.py | 38 ++------ .../test_triattention_cute_union_fusion.py | 35 ++----- .../test_triattention_draft_cocompaction.py | 40 +++----- .../test_triattention_fused_settle_pack.py | 40 ++------ .../test_triattention_pipeline.py | 53 +++-------- .../test_triattention_selection_compaction.py | 94 ++++++------------- .../serial/test_sparse_kv_cache_compact.py | 70 +++++--------- 8 files changed, 117 insertions(+), 326 deletions(-) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 57763703dad5..850ebe93b1e8 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -21,11 +21,7 @@ def encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: - """Build the native V2 [pool, request, K/V, block] layout. - - Accepts ``[request, block]`` page ids (one pool) or ``[pool, request, - block]``; K offsets encode as ``2*page`` and V as ``2*page + 1``. - """ + """Native V2 [pool, request, K/V, block] layout: K = 2*page, V = K+1.""" if page_ids.ndim == 2: page_ids = page_ids.unsqueeze(0) encoded = torch.empty( @@ -59,11 +55,7 @@ def _write_move_offsets(compaction, offsets, moves_per_request): def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): - """Load a cohort's per-request protected tails into the move offsets. - - Production stages these rows through the single round-metadata upload; - tests drive the fixed buffers directly through this helper. - """ + """Load per-request protected tails into the staged move offsets.""" if len(tail_lengths) > compaction["request_count"]: raise ValueError("the cohort exceeds the compaction request capacity") if any(tail < 0 or tail > compaction["protected_tail_capacity"] for tail in tail_lengths): @@ -109,13 +101,8 @@ def make_ramp_pools( base=0, device=None, ): - """bf16 pools carrying a shifted ``arange % 251`` ramp payload. - - Every value is exact in bf16 and any wrong page/plane/head/token move - lands on a different byte pattern, so pool-equality checks in the - compaction tests stay conclusive. Geometry defaults to the compact op's - supported production shape (32-token pages, head_dim 64). - """ + """bf16 pools with a shifted ``arange % 251`` ramp: every wrong move + lands on a different byte pattern (supported geometry defaults).""" return [ ( ( @@ -136,14 +123,9 @@ def make_ramp_pools( def build_compaction(**overrides): - """``init_compaction_buffers`` with the suite's default 2-layer geometry. - - Accepts ``eviction_mode`` for test ergonomics (translated to the - builder's derived selection facts), allocates the caller-owned move - offset rows the production driver would stage (capacity cumsum), and - mirrors the geometry inputs onto the bundle for the standalone helpers - here -- production reads only the launch fields. - """ + """``init_compaction_buffers`` with the suite's 2-layer defaults: + translates ``eviction_mode``, allocates the caller-owned offset rows + (capacity cumsum), and mirrors geometry inputs onto the bundle.""" from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers args = dict( @@ -200,14 +182,8 @@ def capacity_offsets(count): def launch_family_pack(compaction, name): - """Standalone HAS_SETTLE=False pack launch for one family of a bundle. - - Production packs dense/SWA inside the fused settle launch and fires the - draft pack inline in ``run_eviction_round``; standalone compaction tests - pack here from the bundle's own geometry so the C++ moves read - initialized indices. The settle-side pointer arguments are compiled away - (any well-formed tensor stands in). - """ + """Standalone HAS_SETTLE=False pack for one family so the C++ moves + read initialized indices (settle-side pointers are compiled away).""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( SETTLE_PACK_BLOCK, SETTLE_PACK_NUM_WARPS, @@ -260,11 +236,7 @@ def launch_family_pack(compaction, name): def run_compaction(compaction, pack=("dense", "draft")): - """Test-side replica of ``run_eviction_round``'s compact stage. - - Optional standalone family packs, the SWA destination rebase, then every - family's C++ moves -- the same sequence production fires inline. - """ + """Replica of the round's compact stage: packs, SWA rebase, C++ moves.""" if compaction["swa_destination_bases"] is not None: torch.add( compaction["prompt_offsets"], @@ -452,12 +424,8 @@ def torch_tri_score_oracle( offsets, layer_indices, ): - """Independent Torch implementation of the paged TriAttention mean score. - - Covers GQA head mapping via ``head // group_size`` and the - position-independent MLR term. Mean aggregation only: it is the single - production aggregation (max was removed with the C++ score stack). - """ + """Independent Torch oracle of the paged mean score (GQA mapping via + ``head // group_size`` plus the position-independent MLR term).""" scores = [] num_q_heads = int(q_real.shape[1]) for request, seq_len in enumerate(seq_lens): @@ -513,12 +481,8 @@ def make_cute_buffers( offsets, decode_width=None, ): - """Real eviction buffers over one shared page-table slot. - - Shared by the CuTe score and union-fusion tests. Union runners compile - only the fused pipeline, so split-reference legs build their own - score-only buffers with ``eviction_mode="per_head"`` over the same pools. - """ + """Real eviction buffers over one shared page-table slot; split + reference legs use ``eviction_mode="per_head"`` over the same pools.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( init_eviction_buffers, ) @@ -561,13 +525,8 @@ def write_block_offsets(bufs, encoded): def stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts): - """Stage the per-round score metadata exactly like production. - - The compiled runner reads valid lengths and window starts straight from - the staged metadata rows (pointer capture), so stage them like - ``stage_eviction_cohort`` does; the width subtraction mirrors what the - production phase-gather launch derives on device. - """ + """Stage the per-round score metadata exactly like production (the + compiled runner reads the staged rows via pointer capture).""" torch.sub( valid_seq_lens[:request_count], token_starts[:request_count], 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 index 808c0eee9459..391a705c6fbc 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -1,17 +1,9 @@ # 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 PyTorch oracles. +"""The SM100 TriAttention CuTe scorer (the only score path) vs oracles. -The launch-path matrix drives multi-layer buffers across the named -production geometries (Qwen3, GPT-OSS, the originally validated -128-token-page shape) against the shared pure-PyTorch oracle -- permuted -physical pages, ragged valid lengths, GQA group 4 riding the padded MMA -tile -- sweeps request counts up to the buffer capacity, and checks the -per-request decode-width metadata the selection reduce kernels consume. The -contract test pins the loud-failure behavior: unsupported geometry raises -from the CuTe runner's own validation at buffer construction -- there is -no fallback score kernel. -""" +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 @@ -43,8 +35,7 @@ def _build_case( device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(seed) num_freqs = head_dim // 2 - # The 0.125 scaling keeps the BF16 key/coefficient products small so the - # kernel-vs-oracle tolerance can stay tight across the frequency sum. + # 0.125 scaling keeps the kernel-vs-oracle tolerance tight. pools = [ ( 0.125 @@ -84,17 +75,13 @@ def _build_case( _write_block_offsets(bufs, _encode_block_offsets(page_ids)) round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) - # Ragged valid lengths: shallow tails land mid-page and mid-compute-tile; - # the 58-deep tail leaves a whole trailing 32-token page fragment past the - # valid length, so the fully-invalid-fragment clamp stays covered. + # 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)] valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) phase = (round_starts.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() - # Everything the PyTorch oracle needs to rebuild the reference leg - # independently (it recomputes its own mean phases from these). oracle_inputs = dict( page_ids=page_ids, q_real=q_real, @@ -128,10 +115,7 @@ def _geometry(max_requests, num_layers, page_count, tokens_per_block, head_dim, ) -# One entry per supported production geometry: the Qwen3 shape (64 -# frequencies, GQA group 4 riding the padded MMA tile), the GPT-OSS shape -# (32 frequencies, group 8, 32-token pages spanning two page fragments per -# compute tile), and the originally validated 128-token-page shape. +# 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"), @@ -173,9 +157,7 @@ def test_cute_kernel_matches_torch_oracle(case): list(range(num_layers)), ) - # The compiled runner serves every request count up to the buffer - # capacity and nothing beyond it; cover one, an intermediate count, and - # the capacity. + # Every count up to capacity is served, nothing beyond. assert max_requests + 1 not in bufs.runner._compiled for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) @@ -210,17 +192,13 @@ def test_cute_kernel_matches_torch_oracle(case): ) -# The loud-failure contract: unsupported geometry raises from the CuTe -# runner's own validation during the eager compile at buffer -# construction, surfaced as the no-fallback RuntimeError. 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 pools with a 4-token page and 4 frequencies sit far outside the - # CuTe contract on every device. + # 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) 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 index 94ed83ca4574..64f5b8912a11 100644 --- 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 @@ -1,14 +1,10 @@ # 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 (two CuTe kernels). +"""Equivalence coverage for the fused score+stats+union pipeline. -The reference side gathers the SAME production score rows (the fused pack's -score-only entry, which every buffer-namespace runner compiles) and normalizes + -union-reduces them with a pure-torch float32 oracle. The fused-vs-reference -comparison was always tolerance-based (the fused pipeline's reduction order -differs from any reference); the tolerances are unchanged from the retired -Triton reference copies. -""" +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 @@ -36,12 +32,8 @@ def _launch_union_fusion( def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: - """Pure-torch union oracle: z-normalize each row's valid prefix, union-max. - - Mirrors the production union semantics: per-row mean and biased std over - the valid prefix (std clamped at 1e-6), then the per-token maximum across - the request's rows; tokens past the valid width stay ``-inf``. - """ + """Union oracle: per-row mean/biased-std z-norm over the valid prefix + (std clamped at 1e-6), union-max across rows, ``-inf`` past the width.""" request_count, _, width = scores_rows.shape combined = torch.full( (request_count, width), float("-inf"), dtype=torch.float32, device=scores_rows.device @@ -90,23 +82,16 @@ def test_union_fusion_matches_split_pipeline( score_starts: "int | list", valid_lens: "list | None", ) -> None: - """The fused pipeline must reproduce the split score->normalize->union rows. - - ``score_starts`` is either one uniform window start or a per-request - list (the fused kernels read the start per request at runtime). The - reference leg runs the production score-only launch over the same decode - windows, then the pure-torch union oracle. - """ + """Fused rows must reproduce split score->normalize->union rows; + ``score_starts`` is uniform or per-request (read at runtime).""" pytest.importorskip("cutlass") torch.manual_seed(20260721) device = torch.device("cuda") seq_len = 256 num_pages = seq_len // tokens_per_block - # 32-token pages: one 128-token compute tile spans four pages, so a - # shuffled physical-page table catches any fragment/page mix-up. The - # ragged valid lengths land mid-tile, exercising the clamped tail - # fragments. + # 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 = ( 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 index cb30b27c87a4..40bc6f0b8751 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -1,17 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Draft KV co-compression tests for TriAttention eviction. - -With one-model speculative decoding, TriAttention compacts the separate draft -KV cache in the same round as the target: the target's union keep set is -broadcast over the draft's own KV heads, the draft's own protected tail is -appended as ordinals ``valid_seq_len + 0..tail-1``, and both caches land at -``destination_base = prompt_len``. These tests cover the physical draft moves, -the packed move indices, stream ordering across both cache managers, the -speculative admission gates (one representative per guard family), the -published compressed-token invariant, and buffer rebuild/invalidation. -""" +"""Draft KV co-compression: the target's union keep set broadcasts over +the draft's own KV heads, the draft's tail appends as ordinals, both land +at ``destination_base = prompt_len``. Covers the physical moves, packed +indices, stream ordering, admission gates, the published compressed-token +invariant, and buffer rebuild/invalidation.""" from types import SimpleNamespace from unittest import mock @@ -49,12 +43,8 @@ def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: def _launched_draft_compaction(draft_protected_tails): - """Build target and draft pools with distinct head counts, then compact. - - The compact op ships only the pipelined bf16 kernels, so the pools use - the supported production geometry (bf16, 32-token pages, head_dim 64) and - the conclusive shifted ``arange % 251`` ramp payload. - """ + """Target and draft pools with distinct head counts (supported bf16 + geometry, mod-251 ramp payload), compacted in one round.""" device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 prompt_len = 2 @@ -69,8 +59,6 @@ def _launched_draft_compaction(draft_protected_tails): initial_target = [pool.clone() for pool in target_pools] initial_draft = draft_pool.clone() - # Kept ordinals are decode-only but absolute; the pinned prompt tokens - # never appear in the selection rectangle. keep = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) compaction = _build_compaction( @@ -145,9 +133,7 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect before.index_select(2, target_source), ) - # The draft compacts the SAME kept ordinals through its OWN page - # table, over its own head count, with its own protected tail, at - # destination_base = prompt_len. + # Same kept ordinals through the draft's OWN table/heads/tail. draft_pages = built.draft_tables[request].to(torch.long) draft_tail = torch.arange( valid, @@ -174,14 +160,12 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect expected_moves.append(draft_source.to(torch.int32)) expected_offsets.append(expected_offsets[-1] + int(draft_source.numel())) - # The packed draft move indices must match the same broadcast-plus-tail - # oracle the physical moves followed. + # Packed indices must match the same broadcast-plus-tail oracle. draft_family = _compaction_family(built.compaction, "draft") expected_row = torch.cat(expected_moves) assert draft_family["offsets"].cpu().tolist() == expected_offsets draft_indices = draft_family["source"] - # The index buffer is sized for the widest tail (the capacity); this - # round's moves are packed at the front, where the offsets point. + # Capacity-sized buffer; this round's moves pack at the front. capacity_total = built.request_count * ( int(built.keep.shape[1]) + max(built.draft_protected_tails) ) @@ -234,9 +218,7 @@ def test_draft_admission_gates_raise(gate, match): if gate == "full_attention_draft": draft_manager.max_attention_window_vec = [128] if gate.startswith("callsite_"): - # These draft contracts read cross-attention buffers or unvalidated - # paged tails; the call-site speculative gate rejects every one of - # them before any manager is created. + # Call-site speculative gate: rejected before any manager exists. from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import ( DFlashDecodingConfig, 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 index 40703f34a476..cb1acd5f402d 100644 --- 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 @@ -1,18 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The fused settle-and-pack kernel must reproduce the settle/pack semantics -exactly. - -The reference is a pure-torch oracle implementing the integer settle and -pack semantics (threshold recovery with sentinel-skip, strictly-greater -count, lowest-index tie quota, ascending prompt-rebased emission, then the -dense/SWA move-source packing). Every output is integer-valued, so the -comparisons remain ``torch.equal`` — including the buffer regions neither -path overwrites (rows shorter than the keep count leave stale entries -behind, and the packing forwards those stale entries the same way in both -paths). -""" +"""Fused settle-and-pack vs a pure-torch integer oracle (threshold +recovery with sentinel-skip, strictly-greater count, lowest-index tie +quota, ascending prompt-rebased emission, dense/SWA packing). All outputs +are integers, so comparisons are ``torch.equal`` including stale regions.""" import pytest import torch @@ -29,15 +21,9 @@ def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, keep_count): - """Settle each row's provisional top-k in place (exact integer semantics). - - Threshold = min score over the provisional lanes (``-1`` sentinel lanes - are skipped, contributing +inf, so an all-sentinel row settles inertly); - keep every strictly-greater score in the valid width, then fill the - remaining quota with threshold ties in increasing index order; emit the - kept ordinals ascending, rebased by the row's pinned prompt length. - Output entries past the emitted count keep their previous (stale) value. - """ + """Settle in place: threshold = min over non-sentinel provisional + lanes; keep strictly-greater, fill quota with lowest-index ties, emit + ascending rebased by prompt; entries past the emitted count stay.""" 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] @@ -71,15 +57,9 @@ def _pack_oracle( per_layer, has_swa, ): - """Pack dense/SWA move sources from the settled ordinals (in place). - - Dense rows forward the settled output content verbatim for the first - ``keep_count`` moves (stale entries included, exactly like the kernel's - unconditional gather) and append the protected tail - ``seq_len + move - keep_count``; SWA rows write the latest-window - ordinals once per KV head (per-layer packs share one SWA row per head, - written only by the first layer's domains). - """ + """Pack in place: dense rows forward settled content verbatim (stale + included) then append the tail ``seq_len + move - keep_count``; SWA + rows write latest-window ordinals once per KV head.""" request_count = int(valid_seq_lens.shape[0]) packed_rows = int(dense_out.shape[0]) dense_total = int(dense_out.shape[1]) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index f3bcd5557304..b2fb3a633a6f 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -331,10 +331,8 @@ def config_dir(name, body): class TestCompressedTokenPublication: - # The cumulative/monotone publication contract itself (including the - # uncompressed round_start restoration) is covered end to end by - # test_triattention_draft_cocompaction.py:: - # test_compressed_count_is_monotone_and_tracks_confirmed_length. + # The monotone publication contract is covered end to end in + # test_triattention_draft_cocompaction.py. def test_identity_compaction_is_rejected_instead_of_published(self): manager = _make_triattention(budget=4) @@ -356,8 +354,7 @@ def test_identity_compaction_is_rejected_instead_of_published(self): assert state["evicted_tokens"] == 0 cache.resize.assert_not_called() - # A prepared entry that violates the keep contract fails loudly in the - # bookkeeping loop instead of publishing an identity compaction. + # A keep-contract violation fails loudly, never publishes. with _mocked_eviction_internals(manager): with pytest.raises(RuntimeError, match="identity compaction"): manager._evict_requests( @@ -465,9 +462,7 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): cache = mgr.kv_cache_manager.kv_cache_map[7] cache.capacity = confirmed + tail mgr.kv_cache_manager.num_extra_kv_tokens = reserve - # The configured tail capacity (reserve + draft reserve + 1) must - # cover this round's actual tail, as in production; the growth - # constant is 1 + _kv_reserve_draft_tokens for batch members. + # Growth constant = 1 + _kv_reserve_draft_tokens for batch members. mgr.kv_cache_manager._kv_reserve_draft_tokens = current_growth - 1 mgr._prepared_generation_batch = SimpleNamespace(generation_requests=[request]) draft_manager = _make_fake_v2(is_draft=True) @@ -482,8 +477,7 @@ def compact(*_args, **_kwargs): with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: mgr._periodic_evict(batch) - # The resolved per-request metadata is threaded in as-is: the tail is - # excluded from seq_len and the keep target is prompt + budget. + # Tail excluded from seq_len; keep target = prompt + budget. evict.assert_called_once_with( [ { @@ -501,8 +495,6 @@ def compact(*_args, **_kwargs): 2, ) cache.resize.assert_called_once_with(retained + tail, None) - # The draft cache shrinks in the same round, to the same retained - # length plus the draft's own protected tail. draft_cache.resize.assert_called_once_with(retained + 1, None) def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): @@ -536,10 +528,8 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): cache.resize.assert_not_called() def test_one_model_draft_co_compression_contract_is_accepted(self): - # Co-compression keeps the draft's physical length equal to the - # target's, so a draft with a smaller max_seq_len than the target's - # logical maximum is accepted, and both managers are marked as - # diverging from the logical length. + # Draft physical length tracks the target's: smaller draft + # max_seq_len is accepted and both managers are marked. draft_manager = _make_fake_v2(is_draft=True) draft_manager.max_seq_len = 8192 manager = TriAttention( @@ -552,7 +542,6 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): manager._validate_v2_compatibility() assert manager.kv_cache_manager.kv_compression_manages_history is True assert draft_manager.kv_compression_manages_history is True - # A one-model MTP contract also passes the call-site speculative gate. from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig @@ -760,10 +749,8 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): 2, ) - # One batched staging call carries the whole cohort: request ids, - # round starts, pinned prompt lengths, and valid lengths. budget=4: - # per-request moves are keep + tail = [6, 7]; padded rows repeat the - # final offset out to the request capacity. + # One batched staging call carries the whole cohort; budget=4 moves + # are keep + tail = [6, 7], padded rows repeat the final offset. args = internals.stage.call_args assert args.args[0] is internals.buffers assert args.args[1] is manager.kv_cache_manager @@ -779,15 +766,9 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): @requires_sm100 @pytest.mark.parametrize("request_count", [1, 8]) def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count): - """ONE launch over layers in DISTINCT storages with DISTINCT block tables. - - This is the production V2 shape: get_buffers wraps every layer as its - own TensorWrapper storage and every layer allocates its own pages, so - the fused path must not assume a shared storage anchor or a shared - per-request block table. The launch is checked against the independent - Torch oracle, then relaunched after a round-start advance and a block - table rebind. - """ + """ONE launch over layers in DISTINCT storages with DISTINCT block + tables (the production V2 shape), checked against the Torch oracle, + then relaunched after a round-start advance and a table rebind.""" pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module @@ -802,7 +783,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun seq_len = page_count * tokens_per_block prompt_len = 1 num_layers = 3 - # Three SEPARATE allocations (distinct storages, like V2 TensorWrapper). pools = [ ( 0.125 @@ -813,7 +793,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun for _ in range(num_layers) ] assert len({pool.untyped_storage().data_ptr() for pool in pools}) == num_layers - # A DIFFERENT block table per layer (per-layer page allocation). generator = torch.Generator(device="cpu").manual_seed(7 + request_count) page_ids_3d = torch.stack( [ @@ -835,7 +814,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun round_starts = round_device[:request_count].tolist() seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) - # One group per layer: slot i holds layer i's own block table. bufs = module.init_eviction_buffers( eviction_mode="per_head", layer_pools=pools, @@ -863,8 +841,6 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) def stage_round(): - # Stage the round metadata straight into the fixed device rows - # (the cohort staging path is covered by the staging test above). bufs.round_starts_device.copy_(round_device) bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens) bufs.token_starts_device.fill_(prompt_len) @@ -873,8 +849,7 @@ def stage_round(): ) score_sentinel = -12345.0 - # Isolate the score stage: empty prebound launch args make the - # compact stage a no-op (this synthetic round never stages offsets). + # Empty prebound launch args: the compact stage is a no-op. bufs.compact_launch_args = () stage_round() bufs.score_output.fill_(score_sentinel) @@ -882,8 +857,6 @@ def stage_round(): fixed = bufs.score_output.clone() assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] - # The deployed fused score must agree with the independent Torch oracle - # when every layer owns a distinct V2 block table. oracle = _torch_tri_score_oracle( pools, {layer: page_ids_3d[layer, :request_count] for layer in layer_order}, 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 index 84ce4fe2316c..f5d615a709a0 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -44,9 +44,8 @@ def _make_selection_buffers( num_query_heads=1, num_kv_heads=1, ): - """Selection-only buffers: exactly what the one constructor allocates - for the mode, without the CuTe score state or compaction (the settle's - pack half is compiled away).""" + """Selection-only buffers for the mode, without CuTe score state or + compaction (the settle's pack half is masked off).""" bufs = SimpleNamespace( eviction_mode=eviction_mode, device=device, @@ -236,9 +235,8 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( @pytest.mark.parametrize("keep_count,width", [(4, 64), (8192, 9216)]) def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, width): - # Heavily tied integer scores with ragged valid widths and per-request - # prompt rebase: the strongest oracle over the direct union top-k path - # (it subsumes the sorted-output and exact-indices smoke variants). + # 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 @@ -261,8 +259,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid max_requests=request_count, ) bufs.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) - # Write the shared per-request prompt lengths the way production staging - # does: the union row-major view aliases the per-request buffer. + # The union row-major view aliases the per-request buffer. bufs.prompt_offsets[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) @@ -347,16 +344,13 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode): - # The compact op ships only the pipelined bf16 kernels: pools use the - # supported geometry (bf16, 32-token pages, head_dim 64), and the kept - # ordinals are spread across all three pages per request so the moves - # still cross page boundaries. + # Supported bf16 geometry; kept ordinals span all three pages so moves + # cross page boundaries. device = torch.device("cuda", torch.cuda.current_device()) request_count = 2 num_layers = 2 num_kv_heads = 2 - # Per-request pinned prompts: one cohort mixes prompt lengths, so the - # byte-exact oracle also proves per-request destination rebasing. + # Mixed prompt lengths prove per-request destination rebasing. prompt_lens = [2, 5] decode_keep_count = 4 seq_len = 80 @@ -368,8 +362,7 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) initial_pools = _make_ramp_pools(num_layers, device=device) pools = [pool.clone() for pool in initial_pools] - # Kept ordinals are decode-only but hold absolute positions; the pinned - # prompt tokens never appear in the selection rectangle. + # Decode-only kept ordinals holding absolute positions. union_decode = torch.tensor( [[16, 32, 56, 72], [24, 40, 48, 64]], dtype=torch.int64, device=device ) @@ -463,19 +456,15 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): device = torch.device("cuda", torch.cuda.current_device()) num_layers = 3 - # The staged bucket capacity must be aligned to the score kernel's - # 64-token compute tile; the request itself stays 8 tokens long. + # Bucket capacity aligned to the 64-token compute tile; request stays 8. bucket_capacity = 64 seq_len = 8 keep_count = 2 - # GQA group 8 (the smallest CuTe-supported group with one KV head); all - # query heads share one zero calibration query and one MLR coefficient, - # so every head row carries the same |K|-driven score. + # GQA group 8, zero calibration query, shared MLR: every head row + # carries the same |K|-driven score. num_q_heads = 8 - # bf16 pools in the compact op's supported geometry. The scored tokens - # all live in each table's first entry, but the two tables still map the - # two storage groups onto different physical pages, which is what the - # layer-order alignment below depends on. + # 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 @@ -539,11 +528,8 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): bufs.valid_seq_lens_device.fill_(seq_len) bufs.token_starts_device.fill_(0) - # Attach a standalone bundle wholesale (families AND the fused settle - # launch data), replacing the constructor-built one: the fused settle - # launch then packs this bundle's construction-time move offsets exactly - # like production packs the staged rows, and the round's inline C++ - # moves consume the same buffers. + # Replace the constructor-built bundle wholesale; the settle launch + # packs this bundle's construction-time move offsets. compaction = _build_compaction( eviction_mode="per_layer_perhead", layer_pools=pools, @@ -560,8 +546,6 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): protected_tail_capacity=0, ) _set_protected_tails(compaction, [0]) - # Attach wholesale: the round consumes the PREBOUND launch args, so the - # bundle must land through the same helper production uses. attach_compaction_bundle(bufs, compaction) run_eviction_round(bufs, normalize_scores=False) assert torch.equal(bufs.keep, expected_keep) @@ -581,14 +565,9 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): @requires_sm100 def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): - """Run two real eviction rounds through one live V2 cache. - - The cache uses the score and compact kernels' supported geometry (bf16, - 32-token pages, head_dim 64): the request spans three pages so that - compacting to two pages still releases one physical page for reuse. - Token scores are tracked in a host-side mirror and the expected keep - sets are derived from it. - """ + """Two real eviction rounds through one live V2 cache: three pages + compact to two, releasing one physical page for reuse; expected keep + sets derive from a host-side score mirror.""" pytest.importorskip("cutlass") import tensorrt_llm import tensorrt_llm.bindings @@ -605,9 +584,7 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): device = torch.device("cuda", torch.cuda.current_device()) request_id = 7 prompt_len = 2 - # The bucket capacity equals the confirmed length here and must be - # aligned to the score kernel's 64-token compute tile; the protected - # tail rides beyond it. + # Bucket == confirmed length, 64-token-tile aligned; tail rides beyond. seq_len = 64 protected_tail = 2 compacted_capacity = 36 @@ -668,8 +645,7 @@ 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 and distinct per token, so the - # byte comparisons below catch any wrong move. + # 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) @@ -679,10 +655,7 @@ def write_token(token: int, score: float) -> None: payload[0, num_freqs] = 0 pool[page, :, 0, offset].copy_(payload) - # Host-side mirror of each physical position's score; expected keep - # sets are derived from it. Scores are distinct within the decode - # window (7 is invertible mod 64 and the window spans one residue - # cycle), so the selection is tie-free and deterministic. + # 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 @@ -697,9 +670,7 @@ def expected_keep() -> torch.Tensor: device=device, ) - # GQA group 8 (the smallest CuTe-supported group with one KV head); - # zero calibration query and a shared MLR coefficient give every - # query head the same |K|-driven score. + # 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) @@ -744,12 +715,8 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: [seq_len], dense_move_offsets=[0, keep_count + protected_tail], ) - # THE union path: the fused pipeline writes normalized union rows - # into ``combined``. Z-normalization is monotonic per row and all - # query heads carry identical scores here, so the expected keep - # set (derived from raw scores) is unchanged. The settle launch - # packs the move sources and the C++ compacts run in the same - # round call; the kept ordinals stay readable afterwards. + # THE union path (fused pipeline). Z-normalization is monotonic + # per row, so the raw-score keep set is unchanged. run_eviction_round(bufs, normalize_scores=True) selected = bufs.keep[0].clone().to(torch.long) mark_page_tables_consumed(bufs, manager._stream) @@ -791,11 +758,8 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: 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) - # The first protected tail becomes confirmed input to round two. Only - # later generated tokens and the next protected tail are written - # here. Mirror the physical relayout, then give the fresh tokens a - # disjoint higher score band (11 invertible mod 64 over a shorter - # window) so round two must select differently from round one. + # 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): @@ -822,9 +786,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: def test_eager_compaction_rebases_masked_swa_window_and_tail(): - # bf16 pools in the compact op's supported geometry (32-token pages, - # head_dim 64); the kept ordinals and valid lengths span all three pages - # per request so the dense and SWA moves stay page-crossing. + # Supported bf16 geometry; dense and SWA moves stay page-crossing. device = torch.device("cuda", torch.cuda.current_device()) dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 5ece2dd60765..81bb0aa72d3e 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -15,10 +15,8 @@ _BATCH_SIZE = 2 _PAGE_INDEX_DIVISOR = 2 -# Kernel-name substrings for the profiler probes below. The pipelined bf16 -# kernels are the only shipped compaction path; the retired register-staging -# kernel (still present for the sparse-attention updater) must never appear -# in this op's launches. +# Profiler probe names: the pipelined bf16 kernels are the only shipped +# path; the retired register-staging kernel must never appear. _FAST_KERNEL_NAME = "sparseKvCacheCompactV2Bf16PipelineKernel" _RETIRED_KERNEL_NAME = "updateSparseKvCacheAfterFmha" @@ -116,9 +114,8 @@ def _reference_compact( for group_layer, (source_pool, destination_pool, page_table) in enumerate( zip(original, expected, page_tables) ): - # The kernel decodes K offsets as offset // 2 (the V2 2*page+plane - # encoding) regardless of the scale the table was built with; the - # reference mirrors the kernel, not the encoder. + # The kernel decodes K offsets as offset // 2 regardless of the + # encoder's scale; the reference mirrors the kernel. raw_page_table = page_table // _PAGE_INDEX_DIVISOR if source_indices.ndim == 2: layer_sources = source_indices @@ -184,10 +181,8 @@ def _compact( _SMALL_ROW = [2, 5, 8, 3, 7, 10] -# One byte-equality case per launch shape the op serves: dtype/head_dim and -# destination/page-scale sweeps, per-request destination bases (one launch -# mixing pinned-prompt lengths), a 3-D per-layer source with layer routing, -# and a multi-tile launch (40/35 moves across 24-page sequences). +# One byte-equality case per launch shape: head_dim/destination/page-scale +# sweep, per-request bases, 3-D per-layer routing, multi-tile. _LAYER_CASES = [ pytest.param( dict(head_dim=head_dim, dest=dest, scale=scale), @@ -300,22 +295,14 @@ def test_sparse_kv_cache_compact_layers_cuda_graph_replay(): # --- Production-shaped geometry for the pipelined bf16 fast path ---------- -# -# The fast path only dispatches for bf16 pools with head_dim 64/128 and -# 32/128-token pages, so the cases below use their own builder instead of the -# 3-page fixtures above. _FAST_BATCH_SIZE = 3 -# Per-request move counts: two ragged tiles plus a full one (pipeline steady -# state), an empty request (kernel early-return), and a single ragged tile -# (prologue/epilogue only). +# Ragged+full tiles, an empty request, and a prologue/epilogue-only tile. _FAST_MOVE_COUNTS = (71, 0, 29) # Mixed prompt lengths: none tile- or page-aligned. _FAST_DESTINATION_BASES = [3, 9, 17] -# The production move-index buffers are allocation-wide: pad the head-plane -# width beyond this round's total move count so a kernel that derives the -# stride on device (instead of honoring the explicit sourceHeadStride) reads -# padding for heads > 0 and fails the byte-compare. +# Allocation-wide index buffers: padding past the round's total move count +# makes a device-derived stride read padding and fail the byte-compare. _FAST_SOURCE_PAD = 37 _FAST_IDENTITY_MOVES = 5 @@ -443,10 +430,8 @@ def _run_fast_geometry_case(case: _FastGeometryCase) -> list[torch.Tensor]: return expected -# The full fast-path gate matrix. 128-token pages are the geometry the ported -# kernel was written for; 32-token pages and head_dim 128 are this tree's -# production configuration. The per-layer-source row keeps the 3-D routing -# path runnable through the fast kernel. +# The full fast-path gate matrix; the per-layer-source row keeps the 3-D +# routing path runnable through the fast kernel. _FAST_GEOMETRY_MATRIX = [(64, 32), (128, 32), (64, 128), (128, 128)] @@ -455,19 +440,21 @@ def _run_fast_geometry_case(case: _FastGeometryCase) -> list[torch.Tensor]: [(h, t, False) for h, t in _FAST_GEOMETRY_MATRIX] + [(64, 32, True)], ) def test_sparse_kv_cache_compact_layers_fast_geometry(head_dim, tokens_per_block, per_layer): - # Eligible geometry always dispatches the pipelined kernel; the - # byte-compare against the CPU reference is the correctness net. + # Byte-compare against the CPU reference, plus the dispatch probe: the + # pipelined kernel must actually run (and the retired one must not) -- + # every byte-equality here would still pass on a silent fallback. case = _make_fast_geometry_case(head_dim, tokens_per_block, per_layer_sources=per_layer) - expected = _run_fast_geometry_case(case) + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: + expected = _run_fast_geometry_case(case) + names = [event.name for event in profiler.events()] + assert any(_FAST_KERNEL_NAME in name for name in names) + assert not any(_RETIRED_KERNEL_NAME in name for name in names) for actual, reference in zip(case.pools, expected): assert torch.equal(actual.cpu(), reference) -# One representative per reject family: dtype outside the bf16-only gate, -# head_dim outside the gate, page size outside the gate, and a flat 2-D -# source combined with per-layer indices (which would silently read layer 0 -# for every launch). There is no fallback kernel: every reject must fail -# loudly and leave the pools untouched. +# One representative per reject family; no fallback kernel exists, so every +# reject must fail loudly and leave the pools untouched. @pytest.mark.parametrize( "dtype,head_dim,tokens_per_block,flat_with_layer_indices,match", [ @@ -505,18 +492,3 @@ def test_sparse_kv_cache_compact_layers_rejects_invalid_launch( torch.cuda.synchronize() for actual, reference in zip(case.pools, case.pools_cpu): assert torch.equal(actual.cpu(), reference) - - -@pytest.mark.parametrize("head_dim,tokens_per_block", _FAST_GEOMETRY_MATRIX) -def test_sparse_kv_cache_compact_layers_fast_path_actually_runs(head_dim, tokens_per_block): - # Guard against the fast-path gate silently never firing: every - # byte-equality test in this module would still pass if the dispatcher - # fell through to the register-staging kernel. Assert via the profiler - # that the pipelined kernel ran (and the fallback did not) for each of - # the four static geometry dispatch branches. - case = _make_fast_geometry_case(head_dim, tokens_per_block) - with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profiler: - _run_fast_geometry_case(case) - names = [event.name for event in profiler.events()] - assert any(_FAST_KERNEL_NAME in name for name in names) - assert not any(_RETIRED_KERNEL_NAME in name for name in names) From 0b959b24bb7a40e4d54826022537dbcf3a97850f Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:46:51 -0700 Subject: [PATCH 107/178] [None][chore] Trim docstrings and comments to upstream brevity Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 101 +--- .../triattention/triattention.py | 499 ++++-------------- .../triattention_cute_score_fused.py | 88 +-- .../triattention_cute_selection.py | 51 +- .../triattention/triattention_kernels.py | 128 ++--- 5 files changed, 204 insertions(+), 663 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 34ba40c6d059..d753aacf0091 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,17 +15,8 @@ """Batched physical KV-cache compaction for eviction-based compression. -Contract: the caller supplies increasing kept decode ordinals per selection -row (absolute positions), each request's valid length and pinned prompt -length, and the staged V2 block offsets; the surviving KV then moves in -place through batched C++ compact launches, fed by per-request move indices -packed on device. Everything is plain tensors and dicts: -``init_compaction_buffers`` allocates the launch data once per geometry and -returns one bundle whose fields the eviction driver fires directly each -round -- the target's dense/SWA packing rides the driver's fused settle -launch (described by the bundle's ``settle_pack_tensors``/ -``settle_pack_shape``); the driver also fires the co-compressed draft's own -pack launch (``draft_pack``) and every family's C++ moves. +``init_compaction_buffers`` allocates plain-dict launch data once per +geometry; the eviction driver fires the bundle's fields directly each round. """ from collections import OrderedDict @@ -40,12 +31,8 @@ def _make_move_indices( request_count: int, device: torch.device, ) -> torch.Tensor: - """Packed source-index buffer sized for the widest per-request moves. - - The per-request move offsets are always caller-owned rows (refreshed - together with the round metadata in one copy), so only the index - rectangle is allocated here. - """ + """Packed source-index buffer sized for the widest per-request moves + (the move offsets are caller-owned rows).""" return torch.empty( (*index_prefix, moves_per_request * request_count), dtype=torch.int32, device=device ) @@ -57,13 +44,9 @@ def _compact_groups( device: torch.device, per_layer_slots: Optional[Dict[int, int]] = None, ) -> Tuple[Dict[str, object], ...]: - """Batch layers into one C++ launch per uniform V2 pool. - - Each returned dict is the plain launch data for one - ``sparse_kv_cache_compact_layers`` call. ``per_layer_slots`` maps each - layer to its selection row; it is only set when every dense layer keeps - its own token set (per-layer eviction). - """ + """Batch layers into one ``sparse_kv_cache_compact_layers`` launch per + uniform V2 pool. ``per_layer_slots`` maps layers to selection rows + (per-layer eviction only).""" grouped = OrderedDict() for layer, pool, page_table in entries: key = ( @@ -130,41 +113,12 @@ def init_compaction_buffers( swa_move_offsets: Optional[torch.Tensor] = None, draft_move_offsets: Optional[torch.Tensor] = None, ) -> Dict[str, object]: - """Allocate the per-geometry compaction launch data as one plain dict. - - ``union``/``per_layer`` are the two selection-geometry facts consumed - here: one shared selection row per request (union), or one per - (layer, KV head) (per_layer), else one per KV head. Dense layers keep - the prompt in place and compact the selected decode tokens plus any - target KV reserved for the next overlapped forward; kernel-masked SWA - layers keep the latest window plus the same protected tail. A - co-compressed draft cache reuses the target's single-row kept ordinals - (broadcast over the draft's own KV-head count) plus the draft's own - protected tail, landing at the same destination base. - - The driver's settle launch packs increasing kept decode ordinals - (absolute positions) per selection row into the index buffers here; - prompt tokens never move, so the rectangle is prompt-length independent - and ``prompt_offsets`` carries each request's pinned prompt length. The - increasing order is LOAD-BEARING: sparseKvCacheCompactOp.cpp's forward - tiled in-place copy requires increasing source ordinals with - ``destinationBases[request] + move <= source[move]``. ``kv_block_offsets`` is the staged V2 snapshot - ``[slot, request, K/V, block]`` (offset = ``2*page + plane``); - ``protected_tail_capacity`` is the widest per-request tail this geometry - must support -- actual per-round lengths arrive through the staged - move-offset rows, which are caller-owned. Nothing launches here: the - target's dense/SWA packing rides the driver's fused settle launch - (``settle_pack_tensors`` + ``settle_pack_shape`` describe it) and the - driver's round function fires the draft pack (``draft_pack``) and every - family's C++ moves directly. + """Allocate per-geometry compaction launch data for the driver's round. - Returns EXACTLY the launch data the driver fires: ``families`` (each - ``{"name", "groups", "source", "offsets", "destination_bases"}``), - ``settle_pack_tensors``/``settle_pack_shape``, ``draft_pack`` (``None`` - or ``{"indices", "offsets", "dense_total", "move_capacity", - "num_kv_heads"}``), ``swa_destination_bases``, and ``swa_rebase_delta`` - (per-round SWA destination rebase). The bundle is launch data, not an - input mirror. + Selection rows: union = one per request; per_layer = one per (layer, KV + head); else one per KV head. Move sources must be increasing kept ordinals + with destination_bases[request] + move <= source[move] (C++ in-place copy + contract). Returns the launch bundle the round function fires directly. """ device = layer_pools[dense_layers[0]].device request_count = int(request_count) @@ -174,8 +128,7 @@ def init_compaction_buffers( swa_layers = tuple(int(layer) for layer in swa_layers) layer_pool_keys = tuple(layer_pool_keys) - # The C++ compact op takes the KV-head count from each launch's pool - # shape [pages, K/V, heads, tokens, dim]. + # Pool shape [pages, K/V, heads, tokens, dim]. num_kv_heads = int(layer_pools[dense_layers[0]].shape[2]) dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer else (num_kv_heads,) dense_move_indices = _make_move_indices( @@ -199,14 +152,9 @@ def init_compaction_buffers( swa_move_indices = None swa_entries = [] if not swa_layers: - # No SWA family: drop the unused offsets row. With SWA layers the - # argument must stay live so the family reads the per-round staged - # offsets instead of its construction-time sizes. swa_move_offsets = None swa_window = 0 else: - # Per-request window validity (prompt + decode keep >= window) is - # prompt-dependent and checked by the caller each round. swa_window = int(swa_window) swa_destination_bases = torch.empty_like(prompt_offsets) swa_move_indices = _make_move_indices( @@ -226,14 +174,10 @@ def init_compaction_buffers( ] dense_slots = {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer else None - # Selection rows carry decode-only kept ordinals (already absolute), so - # the rectangle is prompt-length independent. HAS_SWA specializes the SWA - # loads and stores away, so without SWA layers the SWA pointer arguments - # are None (the Triton optional-pointer convention). + # Without SWA layers the SWA pointer args are None (HAS_SWA=False). has_swa = swa_move_indices is not None swa_total = int(swa_move_indices.shape[-1]) if has_swa else 0 - # Widest per-request move count any staged offsets may express; the - # packing loop covers exactly this many move slots per packed row. + # Widest per-request move count any staged offsets may express. move_capacity = decode_keep_count + protected_tail_capacity if has_swa: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) @@ -257,9 +201,6 @@ def init_compaction_buffers( families = [ dict( name="dense", - # The fused selection-side settle launch packs the dense/SWA move - # sources when it finalizes the kept ordinals; only the C++ moves - # consume these buffers. Each round then packs exactly once. groups=_compact_groups(dense_entries, layer_pool_keys, device, dense_slots), source=dense_move_indices, offsets=dense_move_offsets, @@ -267,7 +208,6 @@ def init_compaction_buffers( ) ] if swa_layers: - # The fused dense pack fills the SWA move buffers in the same launch. families.append( dict( name="swa", @@ -281,8 +221,7 @@ def init_compaction_buffers( if draft_layers: draft_tail = int(draft_protected_tail_capacity or 0) draft_layers = tuple(int(layer) for layer in draft_layers) - # The draft forms its own launch groups so it may use a different - # KV-head count than the target. + # Own launch groups: the draft may use a different KV-head count. draft_num_kv_heads = int(draft_layer_pools[draft_layers[0]].shape[2]) draft_move_indices = _make_move_indices( (draft_num_kv_heads,), @@ -302,10 +241,7 @@ def init_compaction_buffers( ) for layer in draft_layers ] - # In union mode the pack kernel reads selection row 0 for every packed - # row, so the driver fires one more pack launch broadcasting the - # target keep set over the draft KV heads and appending the draft's - # own tail ordinals; these are that launch's geometry constants. + # Geometry constants for the draft's own pack launch. draft_pack = dict( indices=draft_move_indices, offsets=draft_move_offsets, @@ -331,7 +267,6 @@ def init_compaction_buffers( settle_pack_shape=settle_pack_shape, draft_pack=draft_pack, swa_destination_bases=swa_destination_bases, - # The prompt offsets may be re-staged each round; the driver rebases - # the SWA landing positions with this delta before the moves. + # Per-round SWA destination rebase delta. swa_rebase_delta=decode_keep_count - swa_window, ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 1bb102b22124..22612a2426d6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -13,52 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TriAttention KV-cache compression: periodic physical KV eviction. - -Every ``beta`` confirmed generation tokens TriAttention scores each cached token with a -trigonometric importance score (computed from offline-calibrated statistics of -the model's pre-RoPE query vectors) and physically deletes the tokens below the -top-B keep set. There is no context-phase work and no per-step attention mask: -the eviction runs in the compression manager's final -``on_generation_step_end`` hook. - -TriAttention is a :class:`BaseKVCacheCompressionManager` and nothing more -- it -has no attention backend of its own; decode runs the model's standard dense -kernel over the compacted cache. TriAttention derives each request's effective -confirmed physical length after V2's native update/rewind and publishes the -cumulative evicted count on ``LlmRequest.py_num_compressed_tokens``; attention -metadata reads it back through the KV cache manager on the next step. -Physical reclaim uses V2's existing resize path directly after compaction. An -already-enqueued speculative suffix is excluded from scoring and appended -unchanged to the retained prefix by the same per-layer compact operation. -With one-model speculative decoding, the separate draft KV cache is compacted -in the same round with the target's kept token set (union mode only), so -target and draft always share one physical KV length. - -Structure: ``init_eviction_buffers`` is the ONE one-time constructor for the -whole eviction stack -- it validates the geometry, allocates every buffer, -compiles the mode-needed CuTe entries eagerly, and builds the C++ compaction -launch data (``init_compaction_buffers``). The result is a plain namespace of -tensors, events, and ints. ``stage_eviction_cohort`` and -``run_eviction_round`` are the per-round flow: straight-line module functions -that feed the kernels directly. - -KV layout: the decode kernel stores keys in HND layout -``[num_pages, kv_factor, num_kv_heads, tokens_per_block, head_dim]``. The Python -gather / score / compact code MUST read ``get_buffers`` with ``kv_layout="HND"``; -reading the default NHD silently swaps the token and head axes and scrambles the -cache. - -Position handling: kept keys retain their original RoPE rotation (no re-RoPE on -compaction). The model engine keeps the decode query at its true absolute -position while the attention metadata uses the compacted physical length, so a -query against a kept key at its original rotation still yields the correct -relative distance. - -Calibration is NOT computed here: the user calibrates with the official tool -(github.com/WeianMao/triattention) and passes that .pt via ``calibration_path``; -the manager converts it to our runtime schema at load (see _resolve_calibration). -The scoring math follows the same upstream reference (``methods/pruning_utils.py``). +"""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 types import SimpleNamespace @@ -95,19 +57,16 @@ # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) -# Generation requests skipped by every eviction step (hoisted: the resolve -# loop runs per request per decode step). +# Generation requests skipped by every eviction step. _SKIP_REQUEST_STATES = ( LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT, ) -# Upper bound of the geometric integration offset ladder [1, 2, 4, ...]; no -# caller ever tuned it, so it is a constant rather than a constructor knob. +# Upper bound of the geometric integration offset ladder [1, 2, 4, ...]. _OFFSET_MAX_LENGTH = 65536 -# Stream-affinity contract: the staged buffers and compiled launches are bound def _protected_tail_capacity(manager: KVCacheManagerV2, what: str) -> int: """The V2 tail (extra KV + draft reserve + 1) moved with every compaction.""" capacity = int(manager.num_extra_kv_tokens) + int(manager._kv_reserve_draft_tokens) + 1 @@ -127,8 +86,7 @@ def _allocate_page_table_plane( ) -> Tuple[Dict[int, int], int, torch.Tensor, torch.Tensor]: """Allocate one staged block-offset plane (host pinned + device). - The ``("pool", id)`` page-table keys ARE the snapshot slot numbering: - one scheme, one constructor contract. + The ``("pool", id)`` page-table keys are the snapshot slot numbering. """ representative_slots = { representative: int(key[1]) @@ -144,14 +102,7 @@ def _allocate_page_table_plane( def attach_compaction_bundle(bufs: SimpleNamespace, compaction: Dict[str, object]) -> None: - """Flatten one compaction bundle into prebound launch fields. - - The bundle's families/draft_pack dicts stay as the build-time - description; the round loop fires bare op calls over the tuples - prebound here (the args are init-frozen, so per-round dict-chasing is - pure overhead). Tests attach standalone bundles through the same - helper. - """ + """Flatten one compaction bundle into the prebound launch fields the round fires.""" bufs.compaction_families = compaction["families"] bufs.settle_pack_tensors = compaction["settle_pack_tensors"] bufs.settle_pack_shape = compaction["settle_pack_shape"] @@ -177,8 +128,7 @@ def group_args(family): draft_args: Tuple[tuple, ...] = () for family in compaction["families"]: if family["name"] == "draft": - # The draft's own pack launch must precede its moves; keeping the - # draft family separate preserves the pack-before-moves order. + # Kept separate: the draft pack launch must precede the draft moves. draft_args = group_args(family) else: target_args.extend(group_args(family)) @@ -189,11 +139,7 @@ def group_args(family): bufs.draft_pack_args = None bufs.draft_pack_kwargs = None return - # One more pack launch broadcasts the target keep set over the draft KV - # heads and appends the draft's own tail ordinals. HAS_SETTLE=False - # compiles the settle half away (the ordinals arrive pre-settled), so - # the settle-side pointer arguments are None; the pack half reads the - # settled ordinals through output_indices. + # Pack-only launch (HAS_SETTLE=False): settle-side pointers are None. bufs.draft_pack_args = ( None, None, @@ -262,18 +208,10 @@ def init_eviction_buffers( draft_page_table_token_capacity: Optional[int] = None, draft_protected_tail_capacity: int = 0, ) -> SimpleNamespace: - """Build the ONE plain namespace of buffers for the whole eviction stack. - - The single one-time constructor: buffer staging, eager CuTe compilation - for exactly the entries the eviction mode launches, selection buffers, - and the C++ compaction launch data. Runs outside CUDA graph capture - (compilation allocates and synchronizes); the CuTe runner validates its - own geometry contract and raises loudly -- there is no fallback path. - The returned namespace holds tensors, events, compiled runners, and - ints; all flow logic lives in ``stage_eviction_cohort`` and - ``run_eviction_round``. The buffers retain references to every scored - layer pool: the compiled kernels encode immutable TMA descriptors from - their raw device addresses, so the pools must stay alive and stay put. + """Build the one namespace of buffers, compiled launches, and compaction data. + + Runs once per geometry, outside CUDA graph capture. The compiled kernels + capture raw pool addresses, so the scored pools must stay alive and stay put. """ from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS from .triattention_cute_score_fused import TriAttentionCuteScoreRunner @@ -282,8 +220,6 @@ def init_eviction_buffers( max_requests = int(max_requests) seq_len = int(seq_len) page_table_token_capacity = int(page_table_token_capacity) - # Decode-width capacity of the score buffers; per-request prompt lengths - # are staged runtime metadata. decode_width = int(decode_width) keep_count = int(keep_count) @@ -337,44 +273,34 @@ def init_eviction_buffers( device, ) - # ---- per-round metadata table: ONE host-to-device copy per round ------- - # Three metadata rows (logical position, valid length, prompt length) plus - # one move-offsets row per compacted cache family; offsets rows have - # request_capacity + 1 entries, hence the extra column. + # ---- per-round metadata table: one host-to-device copy per round ------- + # Rows: logical position, valid length, prompt length, then one + # move-offsets row per family (offsets rows need the +1 column). bufs.request_metadata_host = torch.empty( (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) bufs._bulk_copy_idx_src = torch.arange( max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) - # Zero-filled so an unstaged cohort gathers the phase table's row 0 - # instead of indexing it with uninitialized round starts. + # Zero-filled: an unstaged cohort must gather the phase table's row 0. bufs.request_metadata_device = torch.zeros( (6, max_requests + 1), dtype=torch.int32, device=device ) bufs.round_starts_device = bufs.request_metadata_device[0, :max_requests] bufs.valid_seq_lens_device = bufs.request_metadata_device[1, :max_requests] - # Per-request pinned prompt lengths: the score kernel starts each - # request's decode window here, so one bucket may mix prompt lengths. + # Per-request pinned prompt lengths (per-request decode window starts). bufs.token_starts_device = bufs.request_metadata_device[2, :max_requests] - # Rows 3-5 carry the per-family move offsets; the compaction bundle holds - # these views as each family's ``offsets``, so they are wiring, not state. dense_move_offsets_row = bufs.request_metadata_device[3] swa_move_offsets_row = bufs.request_metadata_device[4] draft_move_offsets_row = bufs.request_metadata_device[5] bufs.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) bufs.mean_sin = torch.empty_like(bufs.mean_cos) - # The phase table depends only on the shared calibration; the manager - # shares one dict with every buffer namespace. ``phase=None`` is the ONE - # documented test seam: unit fixtures get a private table built here. + # ``phase=None`` is the one documented test seam (private table built here). if phase is None: phase = build_mean_phase_table(offsets, omega, initial_rows=seq_len) bufs.phase = phase - # ---- score state: ONE fused group across ALL dense layers -------------- - # Segments carry their own page-table slot, so distinct per-layer - # storages/block tables share a single launch. Pool geometry, dtype, and - # layout are validated by the CuTe runner itself below. + # ---- score state: one fused group across all dense layers -------------- p0 = layer_pools[dense_layers[0]] _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape bufs.num_layers = len(dense_layers) @@ -382,9 +308,8 @@ def init_eviction_buffers( bufs.num_kv_heads = int(num_kv_heads) bufs.num_freqs = int(num_freqs) bufs.tokens_per_block = int(tokens_per_block) - # Calibration tables span every model layer; segments index them by - # ABSOLUTE layer id ON DEVICE where they cannot be range-checked, so - # validate the extent once here, loudly. + # Segments index calibration by absolute layer id on device, where it + # cannot be range-checked: validate the extent here. num_calibrated_layers = q_real.numel() // (bufs.num_q_heads * bufs.num_freqs) if min(dense_layers) < 0 or max(dense_layers) >= num_calibrated_layers: raise ValueError("scored layer index exceeds the calibrated layer extent") @@ -403,20 +328,13 @@ def init_eviction_buffers( seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) max_segments = max_requests * bufs.num_layers - # The K1 epilogue folds the head axis as a DSL dynamic coordinate - # (kv_head * N <= 7 score planes of max_segments * seq_len columns each), - # so the largest 32-bit-folded offset is 7 planes x the plane stride; - # every downstream i32 offset (including the seg_out_offset cast below) - # is bounded by it. Wraparound would be a silent wild read, not a clean - # error, hence the loud audit. + # The score plane must stay 32-bit indexable (wraparound = silent wild read). if (PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len >= 2**31: raise ValueError( "score bucket overflows the 32-bit score plane: " f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len}" ) - # The kernel scores each request's window into a head-major scratch; - # all buffers below are persistent because the compiled kernels capture - # their device pointers. + # Persistent buffers: the compiled kernels capture their device pointers. bufs.padded_head_columns = PADDED_HEAD_COLUMNS bufs.cute_scratch = torch.empty( bufs.num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * seq_len, @@ -430,16 +348,11 @@ def init_eviction_buffers( bufs.gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( 1, 1, 1, 1, -1 ) - # Compile the SM100 CuTe entries this mode launches -- HERE at buffer - # construction, outside any CUDA graph capture (compilation allocates and - # synchronizes). Union rounds run the fused score+stats+union pipeline; - # the per-head modes run the score-only entry. There is deliberately no - # other score path and no fallback. + # Compile the mode's SM100 CuTe entries; no other score path, no fallback. union = eviction_mode == "union" bufs.union_rows = None if union: - # Union output rows are sized by the whole bucket (the widest - # possible window); consumers mask by the per-request widths. + # Bucket-wide rows; consumers mask by the per-request widths. bufs.union_rows = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) try: bufs.runner = TriAttentionCuteScoreRunner( @@ -456,10 +369,7 @@ def init_eviction_buffers( seg_page_off=seg_page_off, seg_req_id=bufs.seg_req, seg_layer_id=seg_layer, - # The kernels read per-request lengths and window starts straight - # from the staged metadata rows (pointer capture; the rows are - # int32 views into the persistent metadata table, so the round - # needs no per-round re-marshaling copies for them). + # Pointer capture of the staged metadata rows. valid_seq_lens=bufs.valid_seq_lens_device, seg_out_offset=seg_out_offset, token_starts=bufs.token_starts_device, @@ -482,32 +392,25 @@ def init_eviction_buffers( ) # ---- selection buffers -------------------------------------------------- - # Per-request valid decode widths, refreshed each round from the staged - # lengths; prompt offsets alias the staged per-request prompt lengths so - # the values are written once per round. bufs.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) bufs.prompt_offsets = bufs.token_starts_device if union: bufs.selection_rows_per_request = 1 bufs.row_prompt_offsets = bufs.prompt_offsets - # The fused pipeline writes normalized per-request union rows straight - # into ``combined``; only the top-k settle-and-pack stage remains. bufs.combined = torch.empty( (max_requests, decode_width), dtype=torch.float32, device=device ) bufs.final_indices = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=device ) - # Kept decode ordinals only: rows are prompt-length independent, so - # one buffer namespace serves cohorts with mixed prompt lengths. + # Kept decode ordinals only (prompt-length independent rows). bufs.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) # Row-major views consumed by the top-k settle launch. bufs.selection_scores_rows = bufs.combined bufs.selection_row_lengths = bufs.valid_widths bufs.provisional_rows = bufs.final_indices bufs.keep_rows = bufs.keep - # Padded rows carry zero valid width; their provisional TopK entries - # must still be in-range ordinals for the finalizer's score gather. + # Padded rows still need in-range ordinals for the finalizer's gather. bufs.final_indices.zero_() bufs.score_output = None else: @@ -516,9 +419,7 @@ def init_eviction_buffers( if eviction_mode == "per_head" else bufs.num_layers * bufs.num_kv_heads ) - # The Triton stats/reduce/settle kernels fold score-row offsets in - # int32; both per-head rectangles must stay 32-bit indexable - # (wraparound = silent wild reads, not a clean error). + # Both rectangles must stay 32-bit indexable (wraparound = wild reads). score_rect = max_requests * bufs.num_layers * bufs.num_q_heads * decode_width selection_rect = max_requests * selection_rows * max(decode_width, keep_count) if max(score_rect, selection_rect) >= 2**31: @@ -530,8 +431,7 @@ def init_eviction_buffers( bufs.row_prompt_offsets = torch.zeros( (max_requests * selection_rows,), dtype=torch.int32, device=device ) - # Decode-only per-head scores gathered from the CuTe scratch, the - # ``[request, layer, head, token]`` layout the reduce kernels read. + # [request, layer, head, token] layout read by the reduce kernels. bufs.score_output = torch.empty( max_requests, bufs.num_layers, @@ -561,7 +461,6 @@ def init_eviction_buffers( bufs.top_indices_i32.zero_() # ---- compaction launch data + settle/pack fusion ------------------------ - # One settle program per (request, selection row). bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) draft_kwargs = {} if draft_layers: @@ -591,22 +490,15 @@ def init_eviction_buffers( swa_window=swa_window, layer_pool_keys=list(layer_pool_keys), protected_tail_capacity=int(protected_tail_capacity), - # Tails vary per round (in-flight growth), so the per-family move - # offsets ride the staged metadata rows each round. + # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, swa_move_offsets=swa_move_offsets_row, **draft_kwargs, ) - # Flatten the launch data to plain prebound fields: ONE fused launch - # settles the kept ordinals and packs the dense/SWA move sources - # (``settle_top_tokens``); ``run_eviction_round`` fires the draft pack - # and every family's C++ moves as bare prebound calls. attach_compaction_bundle(bufs, compaction) # ---- round-ordering events ---------------------------------------------- bufs.copy_done = torch.cuda.Event() - # First record publishes constructor allocations to the V2 copy stream; - # later records protect pinned metadata before the next cohort reuses it. bufs.copy_done.record(torch.cuda.current_stream(device)) bufs.bulk_copy_done = torch.cuda.Event() bufs.bulk_consume_done = torch.cuda.Event() @@ -626,18 +518,11 @@ def _stage_block_offsets( ) -> None: """Copy one request group's V2 block offsets before live compaction. - Uses the V2 block-offset kernel with an immutable pinned snapshot of the - selected host-table rows: this enqueues asynchronous host-memory reads, - and TriAttention later resizes the same cache, which mutates the - manager's table in place. The IndexMapper synchronously resolves request - slots and gathers only their beam-0 K block offsets, decoupling both live - inputs before the native asynchronous copy consumes the snapshot with - identity indices. ``dst[pool, r, 0(K), :]`` holds ``base_page * - index_scales``; score and compact decode that K plane inline. + Gathers an immutable pinned snapshot of the beam-0 K block offsets before + the asynchronous device copy (resize later mutates the live host table). """ if bufs.copy_pending and not bufs.copy_done.query(): bufs.copy_done.synchronize() - # The native device copy reads only K and derives V with kv_offset. manager.index_mapper.gather_k_block_offsets( manager.host_kv_cache_block_offsets, source, @@ -671,20 +556,15 @@ def stage_eviction_cohort( ) -> None: """Copy one eviction cohort into the reusable device buffers. - ``token_starts`` carries each request's pinned prompt length; the score - kernel starts that request's decode window there, so the cohort may mix - prompt lengths. + ``token_starts`` carries each request's pinned prompt length (per-request + decode window start), so one cohort may mix prompt lengths. """ request_count = len(request_ids) stream = torch.cuda.current_stream(bufs.device) - # Reuse guard: staging over a cohort whose pages are still being read - # would silently corrupt the in-flight compaction. if bufs.page_tables_active: raise RuntimeError("previous page-table cohort is still active") request_metadata = torch.as_tensor((round_starts, seq_lens, token_starts), dtype=torch.int32) - # Grow the phase table while this cohort's round starts are still host - # integers: a stale-capacity gather is an out-of-bounds index_select on - # the device. + # Must grow while the round starts are still host integers. grow_mean_phase_table(bufs.phase, int(max(round_starts)) + 1) _stage_block_offsets( bufs, @@ -706,11 +586,8 @@ def stage_eviction_cohort( bufs.draft_copy_block_count, ) bufs.request_metadata_host[:3, :request_count].copy_(request_metadata) - # Rows past this cohort are padding: zero lengths keep the score kernel - # and selection inert for them. + # Zero lengths keep the score kernel and selection inert for padded rows. bufs.request_metadata_host[:3, request_count:].zero_() - # This round's per-family move offsets ride the same table, so the single - # device copy below carries them too. for row, family_offsets in ( (3, dense_move_offsets), (4, swa_move_offsets), @@ -721,17 +598,13 @@ def stage_eviction_cohort( torch.as_tensor(family_offsets, dtype=torch.int32) ) try: - # Copy the fixed backing once. Only the first ``request_count`` - # columns are consumed by this cohort. bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) finally: - # Guard the pinned metadata until its asynchronous copies complete. - # Page-table device-buffer reuse is guarded separately after compact. + # Guards the pinned metadata until the asynchronous copies complete. bufs.copy_done.record(stream) bufs.copy_pending = True bufs.page_tables_active = True - # The staged per-request prompt lengths are shared with the selection; - # per-head modes re-expand them into their row-major view here. + # Per-head modes re-expand the prompt lengths into their row-major view. if bufs.row_prompt_offsets is not bufs.prompt_offsets: bufs.row_prompt_offsets.view(bufs.max_requests, bufs.selection_rows_per_request).copy_( bufs.prompt_offsets.unsqueeze(1).expand(-1, bufs.selection_rows_per_request) @@ -739,12 +612,7 @@ def stage_eviction_cohort( def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cuda.Stream) -> None: - """Order V2 page-table reuse and resize after this cohort's compact. - - Every passed manager stream (target, and the draft when co-compressed) - waits on one event recorded after the compact launches, so neither cache - can free or reallocate pages this cohort is still reading. - """ + """Order V2 page-table reuse and resize after this cohort's compact.""" if not bufs.page_tables_active: raise RuntimeError("TriAttention page tables were not staged") bufs.bulk_consume_done.record(torch.cuda.current_stream(bufs.device)) @@ -754,14 +622,10 @@ def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cud def settle_top_tokens(bufs: SimpleNamespace) -> None: - """Pick the top-k with the CuTE selector, then settle its output. - - The CuTE top-k is fast but breaks score ties arbitrarily and emits - indices in arbitrary order; the settle kernel recomputes the threshold - membership with lowest-index-wins ties, rebases each row by its prompt - offset, and writes sorted ordinals. The same launch packs each request's - dense/SWA compaction move sources from the ordinals it just settled - (buffers built without compaction compile the pack half away). + """Pick the top-k, settle ties to sorted ordinals, and pack the move sources. + + The settle kernel resolves the selector's arbitrary tie-breaks with + lowest-index-wins and rebases each row by its prompt offset. """ # The trailing 1 is next_n: decode scores one query token per request. torch.ops.trtllm.cute_dsl_indexer_topk_decode( @@ -789,24 +653,15 @@ def settle_top_tokens(bufs: SimpleNamespace) -> None: def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: - """One staged eviction round, kernels fired directly in sequence. - - Union: phase gather (also derives valid widths), fused score+stats+union - (two CuTe launches), top-k, settle-and-pack, C++ compacts. Per-head - modes: phase gather, score-only CuTe launch, decode-window gather, - stats+reduce kernels, top-k, settle-and-pack, C++ compacts. A - co-compressed draft adds its own pack launch before its C++ moves. - Every launch covers the full request capacity; padded rows past the - staged cohort carry zero lengths and stay inert. + """Fire one staged eviction round: score, select, settle, compact. + + Every launch covers the full request capacity; padded rows carry zero + lengths and stay inert. """ request_count = bufs.max_requests union = bufs.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): - # mean_cos/mean_sin feed the compiled score launches, which captured - # their device pointers: refresh them in place from this round's - # staged round starts. The same launch derives the per-request valid - # decode widths; the compiled kernels read valid lengths and window - # starts straight from the staged metadata rows (pointer capture). + # In-place refresh: the compiled score launches captured these pointers. gather_mean_phases( bufs.phase, bufs.round_starts_device, @@ -816,8 +671,6 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: bufs.token_starts_device, bufs.valid_widths, request_count, - # The prompt offsets may have been re-staged since construction; - # the same launch rebases this round's SWA landing positions. swa_destination_bases=bufs.swa_destination_bases, rebase_delta=bufs.swa_rebase_delta, ) @@ -829,12 +682,9 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: bufs.combined[:request_count, :columns].copy_(bufs.union_rows[:request_count, :columns]) else: bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) - # The kernel wrote each request's window scores (from its pinned - # prompt length) into the head-major scratch, padded to the MMA - # tile N=8 per KV head. Gather each request's decode window into - # the [request, layer, head, token] layout the reduce kernels - # read; columns past a request's valid width carry unscored - # scratch data masked by ``valid_widths``. + # Gather each decode window from the head-major scratch into the + # [request, layer, head, token] layout the reduce kernels read; + # columns past a request's valid width are masked downstream. group_size = bufs.num_q_heads // bufs.num_kv_heads num_segments = request_count * bufs.num_layers pad = bufs.padded_head_columns @@ -876,8 +726,7 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: ) settle_top_tokens(bufs) with nvtx_range("triattention.compact", color="purple"): - # Bare prebound calls (``attach_compaction_bundle``); the draft's own - # pack launch precedes its moves. + # Prebound calls; the draft pack launch precedes the draft moves. for args in bufs.compact_launch_args: torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) if bufs.draft_pack_args is not None: @@ -891,12 +740,9 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: class TriAttention(BaseKVCacheCompressionManager): """Periodic physical KV eviction driven by trigonometric importance scoring. - Overrides ``on_generation_step_end``: every ``beta`` confirmed generation tokens it - reads the cached keys through the ``KVCacheManagerV2``, scores each token - with offline-calibrated stats, and physically evicts the tokens below the - keep set. Full-attention layers are scored; kernel-masked SWA layers preserve - their latest window in the same compacted prefix. Every layer ends with the - same request-wide cached length. + Scores full-attention layers every ``beta`` confirmed tokens and evicts + below the keep set; kernel-masked SWA layers keep their latest window. + Every layer ends with the same request-wide cached length. """ adjusts_generation_kv_length = True @@ -919,9 +765,6 @@ def __init__( self.beta = beta if self.budget <= 0 or self.beta <= 0: raise ValueError("TriAttention budget and beta must both be positive") - # Which token set each eviction round keeps. The user-facing meaning of - # each mode is documented on TriAttentionKvCacheCompressionConfig - # (llm_args); implementation notes live above the selection helpers. self.eviction_mode = eviction_mode if self.eviction_mode not in ("union", "per_head", "per_layer_perhead"): raise ValueError( @@ -930,27 +773,20 @@ def __init__( ) self.normalize_scores = bool(normalize_scores) if self.eviction_mode == "union" and not self.normalize_scores: - # The fused score+stats+union CuTe pipeline is THE union path - # and always z-normalizes. raise ValueError( "TriAttention union eviction requires normalize_scores=True: " "the fused union pipeline always z-normalizes score rows" ) self.pin_prefill = bool(pin_prefill) - # cpt=False (default): budget counts DECODE tokens only (pinned prompt is - # extra). cpt=True: budget INCLUDES the pinned prompt. + # False (default): the budget counts decode tokens only. self.count_prompt_tokens = bool(count_prompt_tokens) if not self.pin_prefill or self.count_prompt_tokens: raise ValueError( "TriAttention physical KV reclaim requires pin_prefill=True and " "count_prompt_tokens=False so finalized prompt KV is preserved" ) - # All physical moves use the C++ V2 compaction operation. - # No other compaction path exists. - # Calibration is the OFFICIAL TriAttention .pt (passed via - # calibration_path), resolved + converted on the first request - # (on_request_init). TRT-LLM does NOT compute calibration; model_path is - # used for RoPE tables and local layer_types/sliding_window metadata. + # Calibration is the official TriAttention .pt, converted on the + # first request; TRT-LLM does not compute calibration. self.model_path = model_path if self.model_path is None: raise ValueError( @@ -965,31 +801,22 @@ def __init__( self._F: Optional[int] = None self._freq_scale_sq: Optional[torch.Tensor] = None - # Geometric integration offsets (built lazily on first eviction so the - # device matches the cache pool). + # Geometric integration offsets, built lazily on first eviction. self._offsets: Optional[torch.Tensor] = None - # Mean-phase table dict, shared by reference with every buffer - # namespace so it persists across buffer rebuilds. + # Mean-phase table dict, shared by reference with every buffer namespace. self._phase: Optional[Dict[str, object]] = None - # Request presence records successful initialization. Each value is a - # plain dict {generation_steps, evicted_tokens}. + # Per-request {generation_steps, evicted_tokens}. self._request_states: Dict[int, Dict[str, object]] = {} - # The overlap executor prepares B(n) before finalizing B(n-1). Keep a - # bare reference to that in-flight batch; its membership id-set and - # the fixed-linear growth constant (1 + the reserved draft width, - # which bounds every step's actual draft) are resolved lazily on the - # first overlap miss. The final hook treats those slots as an opaque - # suffix. + # In-flight overlap batch reference; membership id-set and the growth + # constant (1 + reserved draft width) resolve lazily. self._prepared_generation_batch: Optional[object] = None self._prepared_generation_ids: Optional[set] = None self._generation_growth: Optional[int] = None - # Manager-invariant validation and tail capacity are memoized (the - # loud raises still fire on the first request). + # Memoized manager invariants. self._v2_validated = False self._protected_tail_cache: Optional[int] = None - # The eviction buffers are built once at the first eviction, sized to - # capacity bounds, and reused for the manager's lifetime. + # Built once at the first eviction, reused for the manager's lifetime. self._buffers: Optional[SimpleNamespace] = None self._buffers_fingerprint: Optional[tuple] = None self._local_to_global_layers_cache: Optional[List[int]] = None @@ -1000,11 +827,7 @@ def __init__( self._draft_runtime_kv_layout_cache: Optional[Dict[str, object]] = None def on_request_init(self, request: "LlmRequest", **kwargs) -> None: - """Mark capacity-only decode and resolve calibration once. - - Loads the user-supplied OFFICIAL calibration .pt and converts it to our - runtime schema (see _resolve_calibration). TRT-LLM does not calibrate. - """ + """Validate once and resolve the official calibration for the first request.""" request_id = request.py_request_id if request_id not in self._request_states: self._validate_v2_compatibility() @@ -1020,7 +843,6 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: def _validate_request_capacity(self, request: "LlmRequest") -> None: """Require enough target page-table capacity to reach first eviction.""" manager = self.kv_cache_manager - # V2 mirrors the resolved speculative draft length (0 without spec). speculative_overshoot = int(manager.max_draft_len) first_eviction_decode_length = ( self.budget // self.beta + 1 @@ -1082,10 +904,8 @@ def _ensure_calibrated(self) -> None: self.calibration = self._resolve_calibration() self._H = int(self.calibration["E_q"].shape[1]) self._F = int(self.calibration["E_q"].shape[2]) - # Squared per-frequency RoPE scaling factor (required calibration key). self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) - # Pre-split query stats + MLR coefficient for the score kernel so - # it doesn't recompute (E_q_norm - |E_q|) per call. Shapes [L, H, F]. + # Pre-split query stats + MLR coefficient, shapes [L, H, F]. _Eq = self.calibration["E_q"] self._triattn_q_real = _Eq.real.to(torch.float32).contiguous() self._triattn_q_imag = _Eq.imag.to(torch.float32).contiguous() @@ -1095,11 +915,7 @@ def _ensure_calibrated(self) -> None: self._calibrated = True def _validate_v2_compatibility(self) -> None: - """Reject runtime modes outside the V2 physical-compaction contract. - - Manager-invariant, so the O(num_layers) scans run once; the loud - raise still fires on the first request. - """ + """Reject runtime modes outside the V2 physical-compaction contract (memoized).""" if self._v2_validated: return manager = self.kv_cache_manager @@ -1118,9 +934,8 @@ def _validate_v2_compatibility(self) -> None: raise ValueError("TriAttention requires beam-width-one decoding") if manager.enable_swa_scratch_reuse: raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") - # Speculative feature gates (resolved draft length, linear drafting, - # mode whitelist) run in the factory, where spec_config lives. The - # draft cache itself is validated here whenever one is attached. + # Speculative feature gates run in the factory; the draft cache itself + # is validated here. draft_manager = self.draft_kv_cache_manager if draft_manager is not None: if not draft_manager.is_draft: @@ -1162,23 +977,12 @@ def _validate_v2_compatibility(self) -> None: ) self._v2_validated = True - # The framework drives all request-lifecycle hooks. TriAttention resolves - # calibration on request init, evicts periodically at generation-step end, - # and removes per-request state at finish. It scores from offline - # calibration, not from live queries or attention scores, so it needs no - # per-layer attention hook: the whole eviction runs once per period in - # on_generation_step_end, which loops the layers and reads each layer's keys - # straight from the KV pool. - def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Compact after native KV-cache updates have finalized this iteration. - The compression manager is ordered after KVCacheManagerV2, so capacity - already reflects the written token and any rewind. The overlap scheduler - may already have enqueued the next forward; CUDA stream ordering keeps - compaction after that reader. The resize happens only after compaction; - it detaches the compacted tail without blocking the host, while V2's - per-slot finish events prevent early page reuse. + Runs after KVCacheManagerV2 (capacity reflects the written token and + any rewind); CUDA stream ordering keeps compaction after any already + enqueued overlap forward. """ with nvtx_range_debug("triattention.generation_step_end", color="blue"): self._periodic_evict(scheduled_batch) @@ -1191,12 +995,7 @@ def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwarg def _inflight_generation_growth( self, scheduled_batch: "ScheduledRequests", request_id: int ) -> int: - """Return exact newer target allocation width under overlap scheduling. - - The width is the fixed-linear constant ``1 + reserved draft`` for - members of the prepared batch (the reserve bounds every step's - actual draft; the configured-tail guard enforces the bound). - """ + """Return the in-flight allocation width (1 + reserved draft) under overlap.""" prepared = self._prepared_generation_batch if prepared is None or scheduled_batch is prepared: return 0 @@ -1216,8 +1015,7 @@ def _periodic_evict( self, scheduled_batch: "ScheduledRequests", ) -> None: - """Count confirmed tokens; every ``beta`` tokens score the cache - and physically evict to the pinned prompt plus top-B decode tokens.""" + """Every ``beta`` confirmed tokens, evict to the pinned prompt + top-B set.""" gen_requests = scheduled_batch.generation_requests if not gen_requests: return @@ -1247,17 +1045,12 @@ def _periodic_evict( prepared: List[Dict[str, object]] = [] protected_tail_capacity = self._configured_protected_tail_capacity() - # Resolve every active target cache before changing cadence state (the - # captured cache objects thread all the way to resize -- V2 map - # lookups happen once per step), building the due cohort's - # per-request eviction metadata in the same pass -- - # ``_evict_requests`` trusts it as-is. + # The resolved cache objects thread all the way to resize; the due + # cohort's metadata is built in the same pass. with nvtx_range("triattention.metadata", color="cyan"): for request, request_id, kv_cache in resolved_requests: - # Step accounting and the beta cadence gate come first: a - # non-due request costs two dict ops and two int ops. The - # capacity/tail math and the consistency raises run in the - # due branch, at the point their values are consumed. + # Cadence gate first; capacity/tail math and the consistency + # raises run in the due branch. request_state = self._request_states[request_id] previous_step = request_state["generation_steps"] step = previous_step + 1 + int(request.py_num_accepted_draft_tokens) @@ -1265,11 +1058,8 @@ def _periodic_evict( if previous_step // self.beta >= step // self.beta: continue raw_capacity = int(kv_cache.capacity) - # One-engine speculative decoding keeps a fixed reserve E. - # Under overlap, B(n) is allocated/enqueued before finalizing - # B(n-1), so its exact scheduler growth Q is also opaque. Both - # spans are contiguous after the stable target prefix and move - # byte-for-byte. + # Speculative reserve + in-flight overlap growth: contiguous + # after the stable prefix, moved byte-for-byte. protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( scheduled_batch, request_id ) @@ -1310,9 +1100,7 @@ def _periodic_evict( "kv_cache": kv_cache, "draft_kv_cache": draft_kv_cache, "seq_len": int(seq_len), - # Restore the uncompressed confirmed logical position - # from the physical prefix and cumulative eviction - # count. + # Uncompressed logical position (prefix + evicted). "round_start": int(seq_len + request_state["evicted_tokens"]), "prompt_len": min(int(request.py_prompt_len), int(seq_len)), "expected_keep_count": expected_keep_count, @@ -1320,18 +1108,10 @@ def _periodic_evict( } ) - # Compact all affected dense and kernel-masked SWA layers, then release - # the unreachable tail directly through V2's public resize primitive. - # Prompt lengths and tails are per-request metadata, so the whole due - # cohort runs as one batched round (the buffers hold max_batch_size - # requests, which bounds any generation batch). if not prepared: return num_layers = self._num_layers_from_manager() - # Ungated NVTX with the due count in the message, so any nsys capture - # shows how many requests each eviction round carries. This path runs - # outside CUDA-graph capture, so the dynamic message is safe; the cost - # is one host-side f-string per eviction round. + # Ungated NVTX: the due count in the message shows each round's size. with nvtx_range( f"triattention.evict_request_group reqs={len(prepared)}", color="purple", @@ -1340,13 +1120,7 @@ def _periodic_evict( self._resize_compacted_requests(capacity_targets, protected_tails) def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: - """Release each compacted tail through the caches resolved this hook. - - ``capacity_targets`` threads the ``(request_id, kv_cache, - draft_kv_cache, keep_count)`` tuples resolved by ``_periodic_evict`` - within this same synchronous hook, so no V2 map re-lookup happens - here; only the cheap is_active guard remains. - """ + """Release each compacted tail through the resolved cache objects.""" if not capacity_targets: return with nvtx_range("triattention.resize", color="red"): @@ -1367,9 +1141,8 @@ def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: f"to {resized_capacity} tokens" ) if self.draft_kv_cache_manager is not None: - # The draft cache was compacted with the same kept token - # set, so it shrinks to the same retained length plus its - # own protected tail. + # Same kept set: the draft shrinks to the same retained + # length plus its own tail. draft_protected_tail = self._draft_protected_tail_capacity() for rid, _, draft_kv_cache, target_capacity in capacity_targets: if not draft_kv_cache.is_active: @@ -1382,13 +1155,8 @@ def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: ) def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: - """Return the largest cache length for which selection is an identity. - - With a decode-only budget, pinned prompt tokens do not consume ``budget``. - Selection therefore keeps every token until the cache exceeds - ``prompt_len + budget``. The constructor guarantees the decode-only - budget (``pin_prefill=True``, ``count_prompt_tokens=False``). - """ + """Return the largest cache length for which selection is an identity + (decode-only budget: everything is kept up to ``prompt_len + budget``).""" prompt_len = min(int(request.py_prompt_len), seq_len) return prompt_len + self.budget @@ -1428,11 +1196,8 @@ def _configured_protected_tail_capacity(self) -> int: return self._protected_tail_cache def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: - """Drop this request's per-request length and eviction state.""" + """Drop this request's eviction state; the buffers stay resident.""" self._request_states.pop(request.py_request_id, None) - # The buffers stay resident across idle periods: their memory is a - # deliberate one-time cost and rebuilding it per burst would reintroduce - # allocation on the decode hot path. # ================================================================== # # Helpers (eviction / scoring / V2 cache access / calibration) # @@ -1485,9 +1250,8 @@ def _attention_layer_partition( ) -> Tuple[List[int], List[int], Optional[int]]: """Return dense layers, kernel-masked SWA layers, and the SWA window. - TriAttention initialization has already rejected real V2 windowed - lifecycles. A sliding layer found here is therefore stored at full length - and applies its window only in the attention kernel. + SWA layers here are stored at full length; the window applies only in + the attention kernel. """ cached = self._attention_layer_partition_cache if cached is not None: @@ -1553,10 +1317,8 @@ def _attention_layer_partition( def _runtime_kv_layout(self, num_layers: int) -> Dict[str, object]: """Return stable V2 pool views and layer groups for eviction. - KVCacheManagerV2 keeps GPU virtual addresses and layer geometry stable, - while opt-in pool rebalance can change the page dimension. Cache all - layer views, then query the live page count for one representative per - physical pool before reuse (fail-closed rebalance check). + Cached; re-checks live pool page counts before reuse (pool rebalance + is rejected fail-closed). """ cached = self._runtime_kv_layout_cache manager = self.kv_cache_manager @@ -1608,9 +1370,8 @@ def _build_runtime_kv_layout( ) -> Dict[str, object]: """Build the manager-lifetime layer and pool views one eviction reads. - ``dense_storage_groups`` restricts the compaction groups to the dense - layers (target cache); None groups every layer (draft cache, which has - no SWA partition). ``what`` prefixes error messages ("" or "draft "). + ``dense_storage_groups=None`` groups every layer (draft cache); + ``what`` prefixes error messages. """ num_layers = len(global_layers) maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] @@ -1655,12 +1416,7 @@ def _build_runtime_kv_layout( ) def _draft_runtime_kv_layout(self) -> Dict[str, object]: - """Return stable draft V2 pool views, mirroring ``_runtime_kv_layout``. - - The draft cache is compacted with the target's kept token set, so its - layout has no scoring role: every draft layer is dense and there is no - SWA partition. - """ + """Return stable draft V2 pool views, mirroring ``_runtime_kv_layout``.""" manager = self.draft_kv_cache_manager if manager is None: raise RuntimeError("TriAttention has no draft KV cache manager to lay out") @@ -1734,12 +1490,7 @@ def _buffers_for( ) -> SimpleNamespace: """Return the eviction buffers, building them once at first use. - The request capacity follows the executor's max batch size (memory - scales linearly with it) and the decode-width capacity follows the - eviction bound (compaction keeps the scored decode region near - ``budget`` plus one period of growth), so one set of buffers serves - every round. It is rebuilt only when the pool views change or a - round outgrows it. + Rebuilt only when the pool views change or a round outgrows them. """ if not prepared: raise ValueError("TriAttention eviction requires at least one request") @@ -1779,25 +1530,15 @@ def _buffers_for( needed_width, self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), ) - # Bucket the score scratch by what cohorts actually present instead of - # pinning it to max_seq_len: with pinned prompts the post-compaction - # length is bounded by prompt + budget + slack, so one power-of-two - # bucket serves the steady state, and a cohort that outgrows it simply - # rebuilds the buffers through the capacity check above. A - # max_seq_len floor would make the scratch unindexable in 32 bits and - # tens of GiB at large batch for work that never scores past ~1K - # tokens per request. + # Power-of-two bucket sized by the presented cohorts, NOT max_seq_len + # (a max_seq_len floor would break 32-bit indexing at large batch). seq_capacity = max(int(needed_page_tokens), 1024) seq_capacity = 1 << (seq_capacity - 1).bit_length() seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) - # The CuTe score kernel stores full compute tiles into a scratch - # strided by this bucket capacity, so the capacity must be - # tile-aligned (the geometry gate rejects anything else). Rounding up - # costs at most one tile of scratch per segment. + # The bucket capacity must be tile-aligned (mis-tiled buckets stripe + # the score scratch silently). score_tile_tokens = max(64, int(mgr.tokens_per_block)) seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens - # Mis-tiled buckets stripe the score scratch silently; the builder - # guarantees alignment right here. assert seq_capacity % score_tile_tokens == 0 page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) @@ -1827,7 +1568,7 @@ def _buffers_for( first_pool = layout["layer_pools"][layout["dense_layers"][0]] if self._offsets is None: - # Upstream pruning_utils.build_geometric_offsets: [1, 2, 4, ... <= max]. + # Upstream geometric offsets [1, 2, 4, ... <= max]. self._offsets = torch.tensor( [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], device=first_pool.device, @@ -1883,11 +1624,8 @@ def _move_offsets_for( prepared: Sequence[Dict[str, object]], capacity: int, ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - """Build this round's per-family move offsets, padded to the capacity. - - Rows past the cohort repeat the final offset, so padded requests move - nothing in the pack kernel and the C++ compact launches. - """ + """Build this round's per-family move offsets, padded to the capacity + (padded rows repeat the final offset and move nothing).""" def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets = [0] @@ -1945,13 +1683,10 @@ def _evict_requests( prepared: List[Dict[str, object]], num_layers: int, ) -> List[Tuple[int, int]]: - """Score and compact a prepared cohort, returning ``(request_id, capacity)`` targets. + """Score and compact a prepared cohort, returning the resize targets. - ``prepared`` carries the per-request eviction metadata resolved by - ``_periodic_evict`` (every entry is due and evictable). Only - full-attention layers participate in scoring. For kernel-masked SWA - layers, the latest model window is rebased to the tail of the common - compacted prefix before the request-wide capacity is reduced. + Only full-attention layers are scored; kernel-masked SWA layers keep + their latest window rebased into the compacted prefix. """ with nvtx_range_debug("triattention.resolve_layout", color="blue"): layout = self._runtime_kv_layout(num_layers) @@ -1992,10 +1727,8 @@ def _evict_requests( raise RuntimeError("TriAttention attempted an identity compaction") request_state = self._request_states[item["request_id"]] request_state["evicted_tokens"] += evicted - # Publish the cumulative count on the request: this is the - # manager's only channel to the runtime. The model engine - # reads it back where it builds num_cached_tokens_per_seq, - # so the kernels see the compacted KV length next step. + # The manager's only channel to the runtime: the engine subtracts + # it where it builds num_cached_tokens_per_seq. item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] capacity_targets.append( (item["request_id"], item["kv_cache"], item["draft_kv_cache"], keep_count) 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 index 6db07534b07a..35854d1789d1 100644 --- 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 @@ -128,10 +128,8 @@ def __init__( # cos/sin/mlr coefficient planes per frequency. self.k_coeff = 3 * num_freqs self.tokens_per_block = tokens_per_block - # One 128-token compute tile either matches a page exactly - # (128-token pages: one TMA box per phase) or spans several pages - # (32-token pages: four page fragments per phase, one TMA box each - # into the same transaction barrier). + # One compute tile = one page (128-token pages) or four page + # fragments (32-token pages), one TMA box each. self.box_tokens = min(CTA_M, tokens_per_block) self.fragments_per_phase = CTA_M // self.box_tokens self.pages_per_tile = self.fragments_per_phase @@ -233,12 +231,7 @@ def _stage_raw_band_copies( ): """One raw-K band's fragment copies into an acquired pipeline stage. - The CALLER owns the pipeline handshake (producer_acquire before, - state.advance after): the double-buffered prefetch shares ONE - acquire/advance across its dynamic destination arms, so the - handshake cannot live here. ``band`` is the trace-time real/imag - plane index; ``page_fragments`` is only read for the multi-fragment - specialization (None otherwise). + The caller owns the pipeline handshake (acquire before, advance after). """ 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): @@ -305,10 +298,8 @@ def __call__( tcgen05.CtaGroup.ONE, self.mma_tiler[:2], ) - # Split raw-K transport: one real and one imaginary stage per - # raw-page buffer, each a CTA_M x num_freqs bf16 tile - # (raw_tma_copy_bytes per phase); the descriptor encoder picks the - # SW64/SW128 swizzle from the num_freqs row width. + # One real + one imaginary CTA_M x num_freqs bf16 stage 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, N, self.num_freqs), @@ -511,11 +502,7 @@ def kernel( valid_seq_len = valid_seq_lens[req_id] page_off = seg_page_off[segment] out_base = seg_out_offset[segment] - # Per-request score window start (the request's pinned prompt - # length), loaded like the other per-segment metadata: each CTA - # owns one segment, so its whole schedule derives from one start. - # Scratch writes stay absolute; only the scoring/stats domain and - # the first scored page move per request. + # Per-request score window start; scratch writes stay absolute. score_start = cutlass.Int32(token_starts[req_id]) smem = utils.SmemAllocator() @@ -555,11 +542,8 @@ def kernel( (self.raw_tma_feature_extent, self.box_tokens), coord=(None, None, None), ) - # One smem view and TMA partition per page fragment of the - # 128-token tile, for each of the four stage slices. Fragment f - # lands box_tokens rows deeper in the same stage; the offset is a - # whole multiple of the swizzle period, so the descriptor swizzle - # stays phase-aligned. + # One smem view/TMA partition per page fragment; fragment offsets + # are whole swizzle periods. raw_tma_shared_partition_real = [] raw_tma_shared_partition_imag = [] raw_tma_shared_partition_real_next = [] @@ -675,10 +659,7 @@ def kernel( ) if cutlass.const_expr(self.fragments_per_phase > 1): for fragment in cutlass.range_constexpr(1, self.pages_per_tile): - # The tail tile may not reach this fragment's - # page; clamp to the first fragment (those - # scores lie past the valid width and are - # masked downstream) so the TMA never + # Clamp the tail fragment's page so the TMA never # dereferences an unstaged block entry. fragment_page_id = producer_prefetched_page_id_lane0 if tile_start_token + fragment * self.box_tokens < valid_seq_len: @@ -735,10 +716,8 @@ def kernel( num_threads=EPILOGUE_THREADS, ) cute.arch.mbarrier_init_fence() - # Build the per-(head, frequency) score coefficients: the cos/sin - # bands rotated by the mean future phase (scale*(qr*C - qi*S), - # scale*(qr*S + qi*C)) split into bf16 value+residual pairs, and the - # MLR magnitude coefficient split into fp16 pairs. + # Per-(head, frequency) score coefficients, split into bf16/fp16 + # value+residual pairs. for weight_round in cutlass.range_constexpr(N * self.k_coeff // THREADS): linear_index = tidx + weight_round * THREADS qg = linear_index // self.k_coeff @@ -746,11 +725,8 @@ def kernel( coefficient_kind = feature // self.num_freqs frequency = feature % self.num_freqs mean_offset = req_id * self.num_freqs + frequency - # GQA groups below the minimum MMA tile N=8 ride padded - # columns: they read the group's first head (any valid - # address) and force zero coefficients, so the padded score - # columns come out zero and land in scratch rows the union - # finalizer never reads. + # Padded GQA columns read the group's first head and force + # zero coefficients. qg_read = qg if cutlass.const_expr(self.group_size < N): if qg_read >= self.group_size: @@ -879,10 +855,7 @@ def kernel( 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 - # (CTA_M x num_freqs bf16, raw_tma_copy_bytes). - # Every producer-warp lane participates in the - # PipelineTmaAsync barrier election. + # 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, @@ -1178,13 +1151,8 @@ def kernel( # phase after all of its asynchronous consumers finish. raw_tma_pipeline.consumer_release(raw_tma_consumer_state) raw_tma_consumer_state.advance() - # Every term multiplying sum_seq must stay 64-bit: with - # request*layer segments of seq_len columns the score-plane - # stride alone can exceed 2^31. The scratch head axis is - # padded to the MMA tile N=8 per KV head (group-4 columns - # 4..7 land in padded score planes holding zero scores). The - # host audit in triattention.init_eviction_buffers bounds the - # 32-bit head-axis fold (kv_head * N <= 7 planes) here. + # Every term multiplying sum_seq must stay 64-bit (the plane + # stride can exceed 2^31; the host audit bounds the head fold). output_offset = cutlass.Int64(kv_head * N) * sum_seq + out_base + tile_start_token page_output = cute.make_tensor( output.iterator + output_offset, @@ -1217,10 +1185,7 @@ def kernel( tTR_rAcc, ) tTR_rC.store(tTR_rAcc.load().to(self.c_dtype)) - # The window start is a per-request runtime value, so - # the tile-interior fast path is a dynamic predicate; - # only the straddling first tile takes the per-token - # branch. + # 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.seq_len ): @@ -1300,10 +1265,8 @@ def kernel( sStats[stats_scratch_base + 1] = stats_square_sum cute.arch.barrier() if warp_idx == 0: - # Padded head columns (GQA group below the MMA tile N=8) - # carry zero scores; only the real heads' statistics are - # merged and written, in the compact row layout the union - # finalizer reads (row = segment * num_q_heads + q_head). + # Only real heads' statistics merge into the compact row + # layout (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) @@ -1577,11 +1540,8 @@ def __init__( variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] if max_requests > 1: variants.append((max_requests, 2)) - # ONE compile ritual: per-head runners compile the score-only entry, - # the union runner ONLY its fused stats+union pipeline (plus the - # normalize finalizer below). Same cache keys and compile order as - # the former per-variant blocks; write_partial_stats=False is the - # kernel's default, so passing it explicitly is a no-op. + # Per-head runners compile the score-only entry; the union runner + # only its fused stats+union pipeline. kernel_kwargs = dict( num_layers=num_layers, seq_len=seq_len, @@ -1676,10 +1636,8 @@ def __init__( num_layers=num_layers, seq_len=seq_len, num_q_heads=num_q_heads, - # The score scratch pads each KV head's group - # of score planes to the MMA tile N=8; the - # finalizer maps real head rows onto those - # padded planes (identity for GQA group 8). + # The finalizer maps real head rows onto the + # N=8-padded score planes. num_kv_heads=num_kv_heads, page_shards=page_shards, tokens_per_lane=tokens_per_lane, 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 index 7423ac70b205..da6cae393729 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -50,13 +50,8 @@ def _select_normalize_union_config( width: int, sm_count: int, ) -> tuple[int, int, int]: - """Return tokens per lane, token subtiles, and row-cluster CTAs. - - Heuristic: prefer the small token tile (more CTAs, cluster-widened per - row) while the whole grid still fits the residency bound below; past it, - fall back to the large tile with no clustering so each CTA carries more - tokens instead of oversubscribing the SMs. - """ + """Return tokens per lane, token subtiles, and row-cluster CTAs + (small clustered tile while the grid fits residency, else the large tile).""" 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 @@ -125,10 +120,8 @@ def _ld_cluster_f32(mapped_addr): def _gmem_lane_tile(iterator, flat_index, tokens_per_lane, assumed_align): """One lane's gmem tile of ``tokens_per_lane`` fp32 values. - Folds the (possibly 64-bit) flat index into the pointer BEFORE any - element access, so nothing routes through the DSL's 32-bit dynamic - coordinate; ``assumed_align`` is the only per-site difference between - the vectorized and fallback arms. + Folds the 64-bit flat index into the pointer BEFORE element access, so + nothing routes through the DSL's 32-bit dynamic coordinate. """ return cute.make_tensor( cute.make_ptr( @@ -207,10 +200,8 @@ def __call__( stream=stream, ) else: - # 1D-flattened cluster grid: the X extent is divisible by the - # cluster size and cluster peers are consecutive CTAs, so the - # kernel's (request, token-tile, rank) decode must keep exactly - # this factor order. + # Cluster peers are consecutive CTAs: the kernel's (request, + # token-tile, rank) decode must keep this factor order. kernel.launch( grid=( request_count * self.num_token_tiles * self.row_cluster_ctas, @@ -236,12 +227,8 @@ def _reduce_and_store_union_rows( lane_idx: cutlass.Int32, from_cluster_peers: cutlass.Constexpr, ): - """Final peer reduce + union-row store, shared by both arms. - - The peer-reduction source is the ONLY difference between the - single-CTA and cluster arms (other warps' smem rows vs DSM cluster - peers), selected at trace time by ``from_cluster_peers``. - """ + """Final peer reduce + union-row store, shared by both arms + (the peer source is the only difference, 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): @@ -269,13 +256,8 @@ def _reduce_and_store_union_rows( ) reduced_values[token_slot] = union_value subtile_first_token = first_token + token_subtile * self.subtile_token_tile - # The output row covers this request's own window, - # [0, valid - start); the straddling subtile falls back to the - # per-token stores. union_scores spans request_count * width - # < 2^31 by the host score-plane audit - # (triattention.init_eviction_buffers), so the i32 union_index - # cannot wrap here -- unlike the scratch, whose offsets fold - # through Int64. + # Straddling subtiles fall back to per-token stores; union_scores + # stays < 2^31 by the host audit, so the i32 index cannot wrap. if cutlass.const_expr(self.width % self.tokens_per_lane == 0) and cutlass.dynamic_expr( subtile_first_token + self.tokens_per_lane <= valid_width ): @@ -472,15 +454,10 @@ def kernel( cute.coalesce(score_value_tiles[token_subtile]), ) else: - # Fold the 64-bit flat index into the pointer BEFORE the - # per-element loads, exactly like the vectorized branch: - # the score scratch exceeds 2^31 elements at large - # request counts, and indexing ``scores`` with the flat - # Int64 goes through the DSL's 32-bit dynamic coordinate, - # which wraps. This branch runs whenever a request's - # window start is not lane-aligned (any pinned prompt - # length not divisible by ``tokens_per_lane``), so real - # serve cohorts hit it on every eviction round. + # Fold the 64-bit index into the pointer BEFORE the loads: + # the scratch exceeds 2^31 elements and the DSL's 32-bit + # dynamic coordinate wraps. Unaligned window starts hit + # this arm on every real serve round. score_tail = _gmem_lane_tile( scores.iterator, score_index, self.tokens_per_lane, 4 ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 1d25857f1b2e..1ed1b7f3958e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -1,22 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""GPU kernels for the TriAttention KV-eviction pipeline. - -Scoring runs EXCLUSIVELY through the SM100 CuTe-DSL fused score pack -(``triattention_cute_score_fused.py``): mean aggregation, BF16 KV pools, -head size 64/128, 32/128-token pages, GQA group 4 or 8, per-request score -window starts. Any geometry outside that contract raises loudly at -workspace construction. The per-head modes use the pack's score-only entry -plus the row-stats kernel; union eviction runs the fused -score+stats+union pipeline. One-time buffer staging and runner compilation -live in ``triattention.init_eviction_buffers``; this module keeps the -Triton kernels, their launch helpers, and the mean-phase table builders. - -House rules honored throughout: - * fp32 math (loads up-cast to fp32, fp32 accumulators, fp32 score output). - * int64 for every flat buffer offset that can exceed 2^31. - * mask ragged valid-width tails (and frequency tails) in every load and store. - * the kernels are vendored in this module (no lazy-load hub). +"""Triton kernels, launch helpers, and mean-phase table builders for TriAttention. + +Scoring itself runs through the SM100 CuTe pack in +``triattention_cute_score_fused.py``. House rules: fp32 math, int64 for any +flat offset that can exceed 2^31, mask every ragged tail, kernels vendored here. """ from __future__ import annotations @@ -32,15 +20,11 @@ # --------------------------------------------------------------------------- # -# Positions past this row count are no longer exactly representable in fp32, -# so a larger table would silently degrade every downstream phase. +# Positions past this row count are not exactly representable in fp32. _MEAN_PHASE_MAX_ROWS = 1 << 24 -# Score z-normalization epsilon, shared with the CuTe union pipeline -# (triattention_cute_selection.py imports it): both kernels bake the same -# literal, keeping the eviction modes' normalization consistent. Plain float -# (the CuTe DSL traces it directly); the Triton stats kernel receives it as -# an explicit constexpr parameter. +# Score z-normalization epsilon, shared with the CuTe union pipeline. Plain +# float: the CuTe DSL traces it; the Triton kernel takes it as a constexpr arg. STD_EPSILON = 1e-6 @@ -61,19 +45,13 @@ def _gather_mean_phase_kernel( F_BLOCK: tl.constexpr, HAS_SWA: tl.constexpr, ): - """Copy each request's precomputed phase-table row into the fixed buffers. - - The same launch derives the request's valid decode width (valid length - minus window start), and with SWA layers also rebases the round's SWA - landing positions (window start plus the init-frozen delta), so the - round needs no separate subtraction or rebase launches. - """ + """Copy each request's phase-table row; derive valid widths and, with + SWA layers, the rebased SWA landing positions in the same launch.""" request = tl.program_id(0) frequency = tl.arange(0, F_BLOCK) frequency_mask = frequency < NUM_FREQS table_row = tl.load(round_starts + request).to(tl.int64) - # Clamp stale or padded round starts into the table instead of faulting; - # staged cohorts are host-validated, so live rows are never clamped. + # Clamp stale or padded round starts instead of faulting. table_row = tl.minimum(tl.maximum(table_row, 0), table_rows - 1) source_offset = table_row * NUM_FREQS + frequency output_offset = request * NUM_FREQS + frequency @@ -90,15 +68,9 @@ def _gather_mean_phase_kernel( def build_mean_phase_table( offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int ) -> Dict[str, object]: - """Build the plain-dict mean-phase table shared by every workspace. - - Row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` over the - calibration offsets for every frequency, so refreshing a round's - ``mean_cos``/``mean_sin`` is one pure-gather launch over the staged round - starts. The dict is shared BY REFERENCE between the manager and its - workspace, so ``grow_mean_phase_table`` reaches both. Grow the table while - the round starts are still host integers; the gather kernel clamps stale - rows instead of faulting. + """Build the plain-dict mean-phase table shared (by reference) by every workspace. + + Row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` per frequency. """ phase: Dict[str, object] = { "offsets": offsets.contiguous(), @@ -127,8 +99,7 @@ def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: positions = torch.arange(target, device=omega.device, dtype=torch.float32) cos_table = torch.zeros((target, omega.numel()), dtype=torch.float32, device=omega.device) sin_table = torch.zeros_like(cos_table) - # Accumulate offset-by-offset in fp32 (fixed summation order keeps the - # table bit-stable across rebuilds). + # Fixed summation order keeps the table bit-stable across rebuilds. for offset in phase["offset_values"]: angle = torch.outer(positions + offset, omega) cos_table += torch.cos(angle) @@ -152,13 +123,8 @@ def gather_mean_phases( swa_destination_bases: Optional[torch.Tensor] = None, rebase_delta: int = 0, ) -> None: - """Refresh the fixed mean buffers and valid widths from staged metadata. - - With SWA layers the same launch rebases ``swa_destination_bases`` - (window start + ``rebase_delta``); without them the HAS_SWA branch is - compiled away and the pointer argument is None. Writes in place because - the compiled CuTe score launch captured the destination buffers' device - pointers. CUDA-only; eviction never runs under CUDA graph capture. + """Refresh the mean buffers, valid widths, and (with SWA) the rebased + landing positions, in place: the compiled score launch captured the pointers. """ num_freqs = phase["omega"].numel() _gather_mean_phase_kernel[(request_count,)]( @@ -241,10 +207,8 @@ def _score_per_head_reduce_kernel( ): """Reduce query-head score rows into one selector row per KV-head domain. - per_layer: row (layer, kv_head) = max over the KV head's query group. - Otherwise: row kv_head = mean over layers of that per-layer group max. - Optionally z-normalizes each query-head row with the precomputed - mean/inv-std before reducing. + per_layer: row (layer, kv_head) = max over the KV head's query group; + otherwise row kv_head = mean over layers of that per-layer group max. """ request = tl.program_id(0) selection_row = tl.program_id(1) @@ -318,8 +282,7 @@ def prepare_per_head_scores( num_kv_heads = int(num_kv_heads) _, num_layers, num_query_heads, width = scores.shape selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - # 256 lanes / 4 warps: one program spans the row in a few static loop - # trips without starving occupancy; matches the settle/pack shape. + # 256 lanes / 4 warps, matching the settle/pack shape. stats_block = 256 rows = num_layers * num_query_heads if normalize_scores: @@ -362,10 +325,7 @@ def prepare_per_head_scores( # --------------------------------------------------------------------------- # -# Launch shape of the settle/pack kernel: tokens per program along the -# width/move axis, and its warp count. One pair for every launch site (the -# fused settle, the draft pack, and the standalone test packs) so a retune -# cannot diverge silently. +# Settle/pack launch shape, shared by every launch site. SETTLE_PACK_BLOCK = 256 SETTLE_PACK_NUM_WARPS = 4 @@ -398,24 +358,12 @@ def _settle_ties_and_pack_compaction_sources_kernel( ): """Settle one selection row's ties, then pack its compaction move sources. - One program per (request, selection row). The settle half recovers the - threshold from the provisional top-k, counts the strictly greater - scores, then emits the kept ordinals in increasing order - (lowest-index-wins ties), rebased by the row's pinned prompt length. - The pack half then writes the move sources for the packed rows this - selection row feeds: the kept ordinals it just wrote, the request's - protected tail, plus the SWA rows (latest window). Union selection has - one row per request feeding every KV head's packed row, so that single - program writes all of them. ``HAS_SETTLE=False`` compiles the settle - half away, packing pre-settled ordinals read from ``output_indices`` -- - the draft co-compaction flow, whose keep set is the target's and needs - no settling. - - The increasing-ordinal emission and the tail placement are LOAD-BEARING - for the consumer: sparseKvCacheCompactOp.cpp's forward tiled in-place - copy requires increasing source ordinals with - ``destinationBases[request] + move <= source[move]`` (the SWA window - inequality is the SWA-family instance of the same invariant). + One program per (request, selection row): settle emits the kept ordinals + (lowest-index-wins ties, prompt-rebased), pack writes the move sources + (kept ordinals + protected tail + SWA rows). ``HAS_SETTLE=False`` packs + pre-settled ordinals from ``output_indices`` (draft co-compaction). + Emission order is load-bearing: the C++ compact requires increasing + ordinals with ``destination_bases[request] + move <= source[move]``. """ request = tl.program_id(0) selection_domain = tl.program_id(1) @@ -424,9 +372,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( if HAS_SETTLE: row_scores = scores + row * WIDTH row_selected = provisional_indices + row * KEEP_COUNT - # Scores are decode-relative; this row's pinned prompt length rebases - # the emitted ordinals to absolute positions (per row, so one launch - # may mix prompt lengths). + # Rebases the decode-relative ordinals to absolute positions. prompt_len = tl.load(prompt_offsets + row) threshold = float("inf") @@ -438,11 +384,8 @@ def _settle_ties_and_pack_compaction_sources_kernel( mask=selected_mask, other=0, ) - # Rows shorter than KEEP_COUNT arrive padded with -1 sentinels - # from the top-k's short-row path (zero-width padded rows are all - # sentinels). Mask those lanes out of the gather so no lane - # dereferences ``row_scores - 1`` and a sentinel never joins the - # threshold; rows without sentinels load exactly as before. + # Short rows arrive padded with -1 sentinels from the top-k; + # masked out so no lane dereferences ``row_scores - 1``. selected_valid = selected_mask & (token_index >= 0) selected_score = tl.load( row_scores + token_index, @@ -490,9 +433,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( ties_seen += tl.sum(tied_i32) if HAS_SETTLE: - # The emission above scatters through other lanes of this program; - # make those global stores visible to every lane before the pack - # half reads the row back. + # Make the settled row's scattered stores visible before the pack half. tl.debug_barrier() dense_begin = tl.load(dense_offsets + request) dense_end = tl.load(dense_offsets + request + 1) @@ -511,8 +452,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( ) dense_source = tl.where(move < KEEP_COUNT, selected, valid_len + move - KEEP_COUNT) if UNION: - # The one union row per request feeds every KV head's packed - # row with the same move sources. + # The one union row per request feeds every KV head's packed row. for head in tl.static_range(0, NUM_KV_HEADS): tl.store( dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, @@ -536,9 +476,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( else: swa_mask = move < swa_count if PER_LAYER: - # Per-layer selection has one dense domain per (layer, - # head). SWA uses one shared source row per head, so - # only the first layer's domains write it. + # SWA has one shared row per head; first layer's domains write it. swa_mask = swa_mask & (selection_domain < NUM_KV_HEADS) head = selection_domain % NUM_KV_HEADS swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move From 6d8047e0f4999c9ee5442fb6f4dab5a9aaa6b212 Mon Sep 17 00:00:00 2001 From: tianruih Date: Wed, 22 Jul 2026 23:52:20 -0700 Subject: [PATCH 108/178] [None][chore] One-line the branch-added comments outside the compression module Signed-off-by: tianruih --- tensorrt_llm/_torch/modules/attention.py | 7 +------ tensorrt_llm/_torch/pyexecutor/_util.py | 4 +--- tensorrt_llm/_torch/pyexecutor/model_engine.py | 5 +---- 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 906db2b77f0c..72865589ada1 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -639,12 +639,7 @@ def __init__( if (config.kv_cache_compression_config is not None and config.kv_cache_compression_config. kv_cache_compression_mode.is_eviction_method()): - # The configured method physically evicts cached tokens, so the - # KV length no longer equals the logical sequence length. The - # fused path derives each new token's rotary position from the KV - # length inside the kernel; the unfused path consumes the engine's - # logical position_ids, so surviving keys retain their original - # phases. + # Fused RoPE derives positions from the KV length, which eviction shortens; stay unfused. logger.warning_once( "disable rope_fusion for KV-cache compression " f"({config.kv_cache_compression_config.algorithm}): " diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 56800e6125be..3a57ce241e42 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2140,9 +2140,7 @@ def validate_kv_cache_compression_with_spec( if spec_config.draft_len_schedule is not None: raise ValueError("TriAttention does not yet support dynamic " "speculative draft lengths") - # Rejection sampling (and MTP's relaxed thinking acceptance, an - # MTP-only field) makes the per-step accepted-token count stochastic; - # compression eviction is only validated with greedy acceptance. + # Compression eviction is only validated with greedy acceptance. if spec_config.use_rejection_sampling or getattr( spec_config, "use_relaxed_acceptance_for_thinking", False): raise ValueError("TriAttention does not support speculative " diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5719f92fc955..e01c594e4c9a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4935,10 +4935,7 @@ def previous_seq_slots_device(): and not multimodal_params_list and not lora_params and attn_metadata.padded_num_tokens is None and self._get_position_id_offset() == 0 - # A KV-cache compression manager shrinks the physical - # cache mid-generation, so positions and cached-token - # counts stop advancing in lockstep; keep the full - # prepare path, which rebuilds both every step. + # KV compression shrinks the cache mid-generation; take the full prepare path. and not getattr(kv_cache_manager, "kv_compression_manages_history", False)): self._steady_gen_positions_pinned[:_n_gen].copy_( From 2caf43cf6a59033b5be76b5355774c3ef17712a6 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 00:01:49 -0700 Subject: [PATCH 109/178] [None][perf] Cache launch-path from_dlpack wraps per persistent buffer (knife 18 item 19) Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) 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 index 35854d1789d1..de5682ae5079 100644 --- 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 @@ -1517,6 +1517,9 @@ def __init__( _to_cute(seg_out_offset), _to_cute(token_starts, assumed_align=4), ) + self._cute_partial_stats = _to_cute(self.partial_stats) + # Launch-path from_dlpack wraps, cached per persistent buffer identity. + self._cute_launch_cache: dict[str, tuple[torch.Tensor, cute.Tensor]] = {} static_geometry = ( max_requests, num_layers, @@ -1659,6 +1662,15 @@ def __init__( compiled_configs[config_key] = compiled_selection self._compiled_normalize_union[request_count] = compiled_selection + def _cute_cached(self, key: str, tensor: torch.Tensor) -> cute.Tensor: + """One from_dlpack per persistent buffer; rewrap only on identity change.""" + cached = self._cute_launch_cache.get(key) + if cached is not None and cached[0] is tensor: + return cached[1] + wrapped = _to_cute(tensor.view(-1)) + self._cute_launch_cache[key] = (tensor, wrapped) + return wrapped + def launch( self, request_count: int, @@ -1669,8 +1681,8 @@ def launch( stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) self._compiled[request_count]( *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), + self._cute_cached("mean_cos", mean_cos), + self._cute_cached("mean_sin", mean_sin), *self._cute_tail, request_count, stream, @@ -1687,16 +1699,16 @@ def launch_union_fusion( stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) self._compiled_stats[request_count]( *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), + self._cute_cached("mean_cos", mean_cos), + self._cute_cached("mean_sin", mean_sin), *self._cute_tail, request_count, stream, ) self._compiled_normalize_union[request_count]( - _to_cute(self.partial_stats), + self._cute_partial_stats, *self._cute_selection_prefix, - _to_cute(union_scores.view(-1)), + self._cute_cached("union_scores", union_scores), request_count, stream, ) From 16f34c9e9dae8ea5268117cf62f09413f606687a Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 00:06:13 -0700 Subject: [PATCH 110/178] [None][perf] Stage cohort metadata with in-place numpy row writes (knife 18 item 22) Signed-off-by: tianruih --- .../triattention/triattention.py | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 22612a2426d6..bd75e527f0b1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -279,6 +279,8 @@ def init_eviction_buffers( bufs.request_metadata_host = torch.empty( (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) + # numpy view over the pinned rows: per-round staging writes lists in place. + bufs.request_metadata_host_np = bufs.request_metadata_host.numpy() bufs._bulk_copy_idx_src = torch.arange( max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) @@ -563,9 +565,32 @@ def stage_eviction_cohort( stream = torch.cuda.current_stream(bufs.device) if bufs.page_tables_active: raise RuntimeError("previous page-table cohort is still active") - request_metadata = torch.as_tensor((round_starts, seq_lens, token_starts), dtype=torch.int32) - # Must grow while the round starts are still host integers. - grow_mean_phase_table(bufs.phase, int(max(round_starts)) + 1) + # int32 gate first (before buffers or device work): the in-place numpy + # writes below can wrap silently, and round starts and move offsets are + # the two proven overflow families. + max_round_start = max(round_starts) + rows = ( + (0, round_starts), + (1, seq_lens), + (2, token_starts), + (3, dense_move_offsets), + (4, swa_move_offsets), + (5, draft_move_offsets), + ) + for row, values in rows: + if values is not None and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF: + raise ValueError(f"staged metadata row {row} exceeds the int32 range") + # The previous cohort's metadata H2D must complete before the pinned + # rows are rewritten (same guard _stage_block_offsets applies later). + if bufs.copy_pending and not bufs.copy_done.query(): + bufs.copy_done.synchronize() + host_table = bufs.request_metadata_host_np + for row, values in rows: + if values is not None: + host_table[row, : len(values)] = values + # Zero lengths keep the score kernel and selection inert for padded rows. + host_table[:3, request_count:] = 0 + grow_mean_phase_table(bufs.phase, int(max_round_start) + 1) _stage_block_offsets( bufs, manager, @@ -585,18 +610,6 @@ def stage_eviction_cohort( bufs.draft_block_offsets_device, bufs.draft_copy_block_count, ) - bufs.request_metadata_host[:3, :request_count].copy_(request_metadata) - # Zero lengths keep the score kernel and selection inert for padded rows. - bufs.request_metadata_host[:3, request_count:].zero_() - for row, family_offsets in ( - (3, dense_move_offsets), - (4, swa_move_offsets), - (5, draft_move_offsets), - ): - if family_offsets is not None: - bufs.request_metadata_host[row, : len(family_offsets)].copy_( - torch.as_tensor(family_offsets, dtype=torch.int32) - ) try: bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) finally: From 71bbedc7175f2735c9998f779385452ba99635d3 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 00:10:33 -0700 Subject: [PATCH 111/178] [None][refactor] Inline the mean-phase table wrappers into their call sites Signed-off-by: tianruih --- .../triattention/triattention.py | 46 ++++++++++----- .../triattention/triattention_kernels.py | 57 +------------------ 2 files changed, 34 insertions(+), 69 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index bd75e527f0b1..de740f391d46 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple import torch +import triton from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState @@ -42,9 +43,8 @@ from .triattention_kernels import ( SETTLE_PACK_BLOCK, SETTLE_PACK_NUM_WARPS, + _gather_mean_phase_kernel, _settle_ties_and_pack_compaction_sources_kernel, - build_mean_phase_table, - gather_mean_phases, grow_mean_phase_table, prepare_per_head_scores, ) @@ -299,8 +299,18 @@ def init_eviction_buffers( bufs.mean_sin = torch.empty_like(bufs.mean_cos) # ``phase=None`` is the one documented test seam (private table built here). if phase is None: - phase = build_mean_phase_table(offsets, omega, initial_rows=seq_len) + phase = { + "offsets": offsets.contiguous(), + "omega": omega.contiguous(), + "offset_values": offsets.tolist(), + "cos": None, + "sin": None, + "rows": 0, + } + grow_mean_phase_table(phase, max(int(seq_len), 1)) bufs.phase = phase + bufs.phase_num_freqs = int(phase["omega"].numel()) + bufs.phase_f_block = triton.next_power_of_2(bufs.phase_num_freqs) # ---- score state: one fused group across all dense layers -------------- p0 = layer_pools[dense_layers[0]] @@ -675,17 +685,23 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: union = bufs.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): # In-place refresh: the compiled score launches captured these pointers. - gather_mean_phases( - bufs.phase, + # The phase table tensors are read at launch time (growth replaces them). + _gather_mean_phase_kernel[(request_count,)]( bufs.round_starts_device, + bufs.phase["cos"], + bufs.phase["sin"], bufs.mean_cos, bufs.mean_sin, bufs.valid_seq_lens_device, bufs.token_starts_device, bufs.valid_widths, - request_count, - swa_destination_bases=bufs.swa_destination_bases, - rebase_delta=bufs.swa_rebase_delta, + bufs.swa_destination_bases, + bufs.phase["rows"], + bufs.swa_rebase_delta, + NUM_FREQS=bufs.phase_num_freqs, + F_BLOCK=bufs.phase_f_block, + HAS_SWA=bufs.swa_destination_bases is not None, + num_warps=1, ) if union: bufs.runner.launch_union_fusion( @@ -1588,13 +1604,17 @@ def _buffers_for( dtype=torch.float32, ) if self._phase is None: - self._phase = build_mean_phase_table( - self._offsets, - self.calibration["omega"] + self._phase = { + "offsets": self._offsets.contiguous(), + "omega": self.calibration["omega"] .to(device=first_pool.device, dtype=torch.float32) .contiguous(), - initial_rows=seq_capacity, - ) + "offset_values": self._offsets.tolist(), + "cos": None, + "sin": None, + "rows": 0, + } + grow_mean_phase_table(self._phase, max(int(seq_capacity), 1)) q_real, q_imag, mlr_coef = self._local_score_calibration( layout["num_layers"], layout["global_layers"] ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 1ed1b7f3958e..e57b85ccf285 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -9,7 +9,7 @@ from __future__ import annotations -from typing import Dict, Optional +from typing import Dict import torch import triton @@ -65,25 +65,6 @@ def _gather_mean_phase_kernel( tl.store(swa_destination_bases + request, token_start + rebase_delta) -def build_mean_phase_table( - offsets: torch.Tensor, omega: torch.Tensor, initial_rows: int -) -> Dict[str, object]: - """Build the plain-dict mean-phase table shared (by reference) by every workspace. - - Row ``p`` holds ``mean_o(trig((p + offset_o) * omega_f))`` per frequency. - """ - phase: Dict[str, object] = { - "offsets": offsets.contiguous(), - "omega": omega.contiguous(), - "offset_values": offsets.tolist(), - "cos": None, - "sin": None, - "rows": 0, - } - grow_mean_phase_table(phase, max(int(initial_rows), 1)) - return phase - - def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" rows = int(rows) @@ -110,42 +91,6 @@ def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: phase["rows"] = target -def gather_mean_phases( - phase: Dict[str, object], - round_starts: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - valid_seq_lens: torch.Tensor, - token_starts: torch.Tensor, - valid_widths: torch.Tensor, - request_count: int, - *, - swa_destination_bases: Optional[torch.Tensor] = None, - rebase_delta: int = 0, -) -> None: - """Refresh the mean buffers, valid widths, and (with SWA) the rebased - landing positions, in place: the compiled score launch captured the pointers. - """ - num_freqs = phase["omega"].numel() - _gather_mean_phase_kernel[(request_count,)]( - round_starts, - phase["cos"], - phase["sin"], - mean_cos, - mean_sin, - valid_seq_lens, - token_starts, - valid_widths, - swa_destination_bases, - phase["rows"], - 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. # # --------------------------------------------------------------------------- # From 3a7c0192b4993680513ecf9a4a76b291a4d3f5f6 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 00:29:22 -0700 Subject: [PATCH 112/178] [None][chore] Compress docstrings to one or two lines Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 6 +- .../triattention/triattention.py | 162 ++++-------------- .../triattention_cute_score_fused.py | 139 ++++----------- .../triattention_cute_selection.py | 51 ++---- .../triattention/triattention_kernels.py | 46 ++--- 5 files changed, 95 insertions(+), 309 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index d753aacf0091..6d741879ddac 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -31,8 +31,6 @@ def _make_move_indices( request_count: int, device: torch.device, ) -> torch.Tensor: - """Packed source-index buffer sized for the widest per-request moves - (the move offsets are caller-owned rows).""" return torch.empty( (*index_prefix, moves_per_request * request_count), dtype=torch.int32, device=device ) @@ -44,9 +42,7 @@ def _compact_groups( device: torch.device, per_layer_slots: Optional[Dict[int, int]] = None, ) -> Tuple[Dict[str, object], ...]: - """Batch layers into one ``sparse_kv_cache_compact_layers`` launch per - uniform V2 pool. ``per_layer_slots`` maps layers to selection rows - (per-layer eviction only).""" + """Batch layers into one ``sparse_kv_cache_compact_layers`` launch per uniform V2 pool.""" grouped = OrderedDict() for layer, pool, page_table in entries: key = ( diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index de740f391d46..a0c93374b032 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -68,7 +68,6 @@ def _protected_tail_capacity(manager: KVCacheManagerV2, what: str) -> int: - """The V2 tail (extra KV + draft reserve + 1) moved with every compaction.""" capacity = int(manager.num_extra_kv_tokens) + int(manager._kv_reserve_draft_tokens) + 1 if capacity <= 0: raise RuntimeError(f"{what}KVCacheManagerV2 exposes an invalid protected-tail capacity") @@ -84,10 +83,6 @@ def _allocate_page_table_plane( max_requests: int, device: torch.device, ) -> Tuple[Dict[int, int], int, torch.Tensor, torch.Tensor]: - """Allocate one staged block-offset plane (host pinned + device). - - The ``("pool", id)`` page-table keys are the snapshot slot numbering. - """ representative_slots = { representative: int(key[1]) for representative, key in zip(page_representatives, page_table_keys) @@ -208,11 +203,8 @@ def init_eviction_buffers( draft_page_table_token_capacity: Optional[int] = None, draft_protected_tail_capacity: int = 0, ) -> SimpleNamespace: - """Build the one namespace of buffers, compiled launches, and compaction data. - - Runs once per geometry, outside CUDA graph capture. The compiled kernels - capture raw pool addresses, so the scored pools must stay alive and stay put. - """ + """Build the one namespace of buffers, compiled launches, and compaction data + (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put).""" from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS from .triattention_cute_score_fused import TriAttentionCuteScoreRunner @@ -273,9 +265,7 @@ def init_eviction_buffers( device, ) - # ---- per-round metadata table: one host-to-device copy per round ------- - # Rows: logical position, valid length, prompt length, then one - # move-offsets row per family (offsets rows need the +1 column). + # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- bufs.request_metadata_host = torch.empty( (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) @@ -320,8 +310,7 @@ def init_eviction_buffers( bufs.num_kv_heads = int(num_kv_heads) bufs.num_freqs = int(num_freqs) bufs.tokens_per_block = int(tokens_per_block) - # Segments index calibration by absolute layer id on device, where it - # cannot be range-checked: validate the extent here. + # Device-side calibration indexing cannot be range-checked: validate the layer extent here. num_calibrated_layers = q_real.numel() // (bufs.num_q_heads * bufs.num_freqs) if min(dense_layers) < 0 or max(dense_layers) >= num_calibrated_layers: raise ValueError("scored layer index exceeds the calibrated layer extent") @@ -528,11 +517,7 @@ def _stage_block_offsets( destination: torch.Tensor, copy_block_count: int, ) -> None: - """Copy one request group's V2 block offsets before live compaction. - - Gathers an immutable pinned snapshot of the beam-0 K block offsets before - the asynchronous device copy (resize later mutates the live host table). - """ + """Gather the pinned snapshot before the async device copy: resize mutates the live host table.""" if bufs.copy_pending and not bufs.copy_done.query(): bufs.copy_done.synchronize() manager.index_mapper.gather_k_block_offsets( @@ -566,18 +551,12 @@ def stage_eviction_cohort( swa_move_offsets: Optional[List[int]] = None, draft_move_offsets: Optional[List[int]] = None, ) -> None: - """Copy one eviction cohort into the reusable device buffers. - - ``token_starts`` carries each request's pinned prompt length (per-request - decode window start), so one cohort may mix prompt lengths. - """ + """Copy one eviction cohort into the reusable device buffers.""" request_count = len(request_ids) stream = torch.cuda.current_stream(bufs.device) if bufs.page_tables_active: raise RuntimeError("previous page-table cohort is still active") - # int32 gate first (before buffers or device work): the in-place numpy - # writes below can wrap silently, and round starts and move offsets are - # the two proven overflow families. + # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. max_round_start = max(round_starts) rows = ( (0, round_starts), @@ -590,8 +569,7 @@ def stage_eviction_cohort( for row, values in rows: if values is not None and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF: raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # The previous cohort's metadata H2D must complete before the pinned - # rows are rewritten (same guard _stage_block_offsets applies later). + # The previous cohort's metadata H2D must complete before the pinned rows are rewritten. if bufs.copy_pending and not bufs.copy_done.query(): bufs.copy_done.synchronize() host_table = bufs.request_metadata_host_np @@ -645,11 +623,7 @@ def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cud def settle_top_tokens(bufs: SimpleNamespace) -> None: - """Pick the top-k, settle ties to sorted ordinals, and pack the move sources. - - The settle kernel resolves the selector's arbitrary tie-breaks with - lowest-index-wins and rebases each row by its prompt offset. - """ + """Pick the top-k, settle ties to sorted ordinals, and pack the move sources.""" # The trailing 1 is next_n: decode scores one query token per request. torch.ops.trtllm.cute_dsl_indexer_topk_decode( bufs.selection_scores_rows, @@ -676,16 +650,12 @@ def settle_top_tokens(bufs: SimpleNamespace) -> None: def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: - """Fire one staged eviction round: score, select, settle, compact. - - Every launch covers the full request capacity; padded rows carry zero - lengths and stay inert. - """ + """Fire one staged eviction round: score, select, settle, compact + (every launch covers the full request capacity; padded rows carry zero lengths and stay inert).""" request_count = bufs.max_requests union = bufs.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): # In-place refresh: the compiled score launches captured these pointers. - # The phase table tensors are read at launch time (growth replaces them). _gather_mean_phase_kernel[(request_count,)]( bufs.round_starts_device, bufs.phase["cos"], @@ -711,9 +681,7 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: bufs.combined[:request_count, :columns].copy_(bufs.union_rows[:request_count, :columns]) else: bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) - # Gather each decode window from the head-major scratch into the - # [request, layer, head, token] layout the reduce kernels read; - # columns past a request's valid width are masked downstream. + # Gather each decode window into the [request, layer, head, token] layout of the reduces. group_size = bufs.num_q_heads // bufs.num_kv_heads num_segments = request_count * bufs.num_layers pad = bufs.padded_head_columns @@ -767,12 +735,7 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: class TriAttention(BaseKVCacheCompressionManager): - """Periodic physical KV eviction driven by trigonometric importance scoring. - - Scores full-attention layers every ``beta`` confirmed tokens and evicts - below the keep set; kernel-masked SWA layers keep their latest window. - Every layer ends with the same request-wide cached length. - """ + """Periodic physical KV eviction driven by trigonometric importance scoring.""" adjusts_generation_kv_length = True @@ -814,8 +777,7 @@ def __init__( "TriAttention physical KV reclaim requires pin_prefill=True and " "count_prompt_tokens=False so finalized prompt KV is preserved" ) - # Calibration is the official TriAttention .pt, converted on the - # first request; TRT-LLM does not compute calibration. + # Calibration is the official TriAttention .pt; TRT-LLM does not compute calibration. self.model_path = model_path if self.model_path is None: raise ValueError( @@ -837,8 +799,7 @@ def __init__( # Per-request {generation_steps, evicted_tokens}. self._request_states: Dict[int, Dict[str, object]] = {} - # In-flight overlap batch reference; membership id-set and the growth - # constant (1 + reserved draft width) resolve lazily. + # In-flight overlap batch reference; membership and growth resolve lazily. self._prepared_generation_batch: Optional[object] = None self._prepared_generation_ids: Optional[set] = None self._generation_growth: Optional[int] = None @@ -870,7 +831,6 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: self._ensure_calibrated() def _validate_request_capacity(self, request: "LlmRequest") -> None: - """Require enough target page-table capacity to reach first eviction.""" manager = self.kv_cache_manager speculative_overshoot = int(manager.max_draft_len) first_eviction_decode_length = ( @@ -923,11 +883,9 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: ) def _draft_protected_tail_capacity(self) -> int: - """Return the draft tail moved and re-reserved by every co-compression.""" return _protected_tail_capacity(self.draft_kv_cache_manager, "draft ") def _ensure_calibrated(self) -> None: - """Resolve calibration once for the first request.""" if self._calibrated: return self.calibration = self._resolve_calibration() @@ -944,7 +902,6 @@ def _ensure_calibrated(self) -> None: self._calibrated = True def _validate_v2_compatibility(self) -> None: - """Reject runtime modes outside the V2 physical-compaction contract (memoized).""" if self._v2_validated: return manager = self.kv_cache_manager @@ -963,8 +920,7 @@ def _validate_v2_compatibility(self) -> None: raise ValueError("TriAttention requires beam-width-one decoding") if manager.enable_swa_scratch_reuse: raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") - # Speculative feature gates run in the factory; the draft cache itself - # is validated here. + # Speculative feature gates run in the factory; the draft cache itself is validated here. draft_manager = self.draft_kv_cache_manager if draft_manager is not None: if not draft_manager.is_draft: @@ -1007,12 +963,8 @@ def _validate_v2_compatibility(self) -> None: self._v2_validated = True def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Compact after native KV-cache updates have finalized this iteration. - - Runs after KVCacheManagerV2 (capacity reflects the written token and - any rewind); CUDA stream ordering keeps compaction after any already - enqueued overlap forward. - """ + """Compact after native KV-cache updates have finalized this iteration + (must run after KVCacheManagerV2 so capacity reflects the written token and any rewind).""" with nvtx_range_debug("triattention.generation_step_end", color="blue"): self._periodic_evict(scheduled_batch) @@ -1024,7 +976,6 @@ def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwarg def _inflight_generation_growth( self, scheduled_batch: "ScheduledRequests", request_id: int ) -> int: - """Return the in-flight allocation width (1 + reserved draft) under overlap.""" prepared = self._prepared_generation_batch if prepared is None or scheduled_batch is prepared: return 0 @@ -1044,7 +995,6 @@ def _periodic_evict( self, scheduled_batch: "ScheduledRequests", ) -> None: - """Every ``beta`` confirmed tokens, evict to the pinned prompt + top-B set.""" gen_requests = scheduled_batch.generation_requests if not gen_requests: return @@ -1074,12 +1024,10 @@ def _periodic_evict( prepared: List[Dict[str, object]] = [] protected_tail_capacity = self._configured_protected_tail_capacity() - # The resolved cache objects thread all the way to resize; the due - # cohort's metadata is built in the same pass. + # The resolved cache objects thread all the way to resize. with nvtx_range("triattention.metadata", color="cyan"): for request, request_id, kv_cache in resolved_requests: - # Cadence gate first; capacity/tail math and the consistency - # raises run in the due branch. + # Cadence gate first; capacity math and consistency raises run in the due branch. request_state = self._request_states[request_id] previous_step = request_state["generation_steps"] step = previous_step + 1 + int(request.py_num_accepted_draft_tokens) @@ -1087,8 +1035,7 @@ def _periodic_evict( if previous_step // self.beta >= step // self.beta: continue raw_capacity = int(kv_cache.capacity) - # Speculative reserve + in-flight overlap growth: contiguous - # after the stable prefix, moved byte-for-byte. + # Speculative reserve + in-flight overlap growth: contiguous tail moved byte-for-byte. protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( scheduled_batch, request_id ) @@ -1149,7 +1096,6 @@ def _periodic_evict( self._resize_compacted_requests(capacity_targets, protected_tails) def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: - """Release each compacted tail through the resolved cache objects.""" if not capacity_targets: return with nvtx_range("triattention.resize", color="red"): @@ -1170,8 +1116,7 @@ def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: f"to {resized_capacity} tokens" ) if self.draft_kv_cache_manager is not None: - # Same kept set: the draft shrinks to the same retained - # length plus its own tail. + # Same kept set: the draft shrinks to the same retained length plus its own tail. draft_protected_tail = self._draft_protected_tail_capacity() for rid, _, draft_kv_cache, target_capacity in capacity_targets: if not draft_kv_cache.is_active: @@ -1184,8 +1129,7 @@ def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: ) def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: - """Return the largest cache length for which selection is an identity - (decode-only budget: everything is kept up to ``prompt_len + budget``).""" + """Return the largest cache length for which selection is an identity.""" prompt_len = min(int(request.py_prompt_len), seq_len) return prompt_len + self.budget @@ -1194,7 +1138,6 @@ def _local_score_calibration( num_layers: int, global_layers: List[int], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Return calibration tensors indexed in this PP rank's local layer order.""" if global_layers and max(global_layers) >= self._triattn_q_real.shape[0]: raise ValueError( f"TriAttention calibration has {self._triattn_q_real.shape[0]} layers, " @@ -1219,7 +1162,6 @@ def _local_score_calibration( ) def _configured_protected_tail_capacity(self) -> int: - """Return the largest target tail reserved by the native V2 lifecycle.""" if self._protected_tail_cache is None: self._protected_tail_cache = _protected_tail_capacity(self.kv_cache_manager, "") return self._protected_tail_cache @@ -1228,12 +1170,9 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: """Drop this request's eviction state; the buffers stay resident.""" self._request_states.pop(request.py_request_id, None) - # ================================================================== # - # Helpers (eviction / scoring / V2 cache access / calibration) # - # ================================================================== # + # ---- helpers (eviction / scoring / V2 cache access / calibration) ---- def _local_to_global_layers(self, num_layers: int) -> List[int]: - """Return V2's global layer id for every local TriAttention layer slot.""" cached = self._local_to_global_layers_cache if cached is not None: if len(cached) != num_layers: @@ -1253,7 +1192,6 @@ def _local_to_global_layers(self, num_layers: int) -> List[int]: @staticmethod def _has_sliding_window_signal(config: Dict[str, object]) -> bool: - """Return whether config metadata hints at sliding attention.""" use_sliding_window = config.get("use_sliding_window") if isinstance(use_sliding_window, bool): return use_sliding_window @@ -1277,11 +1215,7 @@ def _has_sliding_window_signal(config: Dict[str, object]) -> bool: def _attention_layer_partition( self, num_layers: int ) -> Tuple[List[int], List[int], Optional[int]]: - """Return dense layers, kernel-masked SWA layers, and the SWA window. - - SWA layers here are stored at full length; the window applies only in - the attention kernel. - """ + """SWA layers here are stored at full length; the window applies only in the kernel.""" cached = self._attention_layer_partition_cache if cached is not None: return cached @@ -1344,11 +1278,6 @@ def _attention_layer_partition( return result def _runtime_kv_layout(self, num_layers: int) -> Dict[str, object]: - """Return stable V2 pool views and layer groups for eviction. - - Cached; re-checks live pool page counts before reuse (pool rebalance - is rejected fail-closed). - """ cached = self._runtime_kv_layout_cache manager = self.kv_cache_manager if cached is not None: @@ -1397,11 +1326,6 @@ def _build_runtime_kv_layout( dense_storage_groups: Optional[Dict[object, List[int]]], what: str, ) -> Dict[str, object]: - """Build the manager-lifetime layer and pool views one eviction reads. - - ``dense_storage_groups=None`` groups every layer (draft cache); - ``what`` prefixes error messages. - """ num_layers = len(global_layers) maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] if any(pool is None for pool in maybe_layer_pools): @@ -1445,7 +1369,6 @@ def _build_runtime_kv_layout( ) def _draft_runtime_kv_layout(self) -> Dict[str, object]: - """Return stable draft V2 pool views, mirroring ``_runtime_kv_layout``.""" manager = self.draft_kv_cache_manager if manager is None: raise RuntimeError("TriAttention has no draft KV cache manager to lay out") @@ -1486,7 +1409,6 @@ def _pool_page_counts( global_layers: Sequence[int], pool_representatives: Sequence[int], ) -> Tuple[int, ...]: - """Read the only pool-view dimension that V2 rebalance can change.""" return tuple( int( manager.impl.get_page_index_upper_bound( @@ -1500,7 +1422,6 @@ def _pool_page_counts( @staticmethod def _pool_view_fingerprint(pools: List[torch.Tensor]) -> Tuple[tuple, ...]: - """Identify the V2 pool properties consumed by score and compact kernels.""" return tuple( ( pool.data_ptr(), @@ -1517,10 +1438,6 @@ def _buffers_for( layout: Dict[str, object], prepared: Sequence[Dict[str, object]], ) -> SimpleNamespace: - """Return the eviction buffers, building them once at first use. - - Rebuilt only when the pool views change or a round outgrows them. - """ if not prepared: raise ValueError("TriAttention eviction requires at least one request") needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) @@ -1559,13 +1476,11 @@ def _buffers_for( needed_width, self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), ) - # Power-of-two bucket sized by the presented cohorts, NOT max_seq_len - # (a max_seq_len floor would break 32-bit indexing at large batch). + # Bucket sized by the presented cohorts, NOT max_seq_len (a floor there breaks 32-bit indexing). seq_capacity = max(int(needed_page_tokens), 1024) seq_capacity = 1 << (seq_capacity - 1).bit_length() seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) - # The bucket capacity must be tile-aligned (mis-tiled buckets stripe - # the score scratch silently). + # The bucket capacity must be tile-aligned (mis-tiling stripes the score scratch silently). score_tile_tokens = max(64, int(mgr.tokens_per_block)) seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens assert seq_capacity % score_tile_tokens == 0 @@ -1657,9 +1572,6 @@ def _move_offsets_for( prepared: Sequence[Dict[str, object]], capacity: int, ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - """Build this round's per-family move offsets, padded to the capacity - (padded rows repeat the final offset and move nothing).""" - def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets = [0] for moves in moves_per_request: @@ -1684,7 +1596,6 @@ def _page_table_pool_keys( global_layers: List[int], manager: Optional[KVCacheManagerV2] = None, ) -> List[object]: - """Return stable V2-pool keys for the representative layers.""" if manager is None: manager = self.kv_cache_manager layer_offsets = manager.layer_offsets @@ -1702,7 +1613,6 @@ def _dense_layer_pool_groups( dense_layers: List[int], global_layers: List[int], ) -> Dict[object, List[int]]: - """Group layers that use the same V2 page table.""" groups: Dict[object, List[int]] = {} for layer, pool_key in zip( dense_layers, @@ -1716,16 +1626,10 @@ def _evict_requests( prepared: List[Dict[str, object]], num_layers: int, ) -> List[Tuple[int, int]]: - """Score and compact a prepared cohort, returning the resize targets. - - Only full-attention layers are scored; kernel-masked SWA layers keep - their latest window rebased into the compacted prefix. - """ with nvtx_range_debug("triattention.resolve_layout", color="blue"): layout = self._runtime_kv_layout(num_layers) with nvtx_range_debug("triattention.staging_lookup", color="blue"): - # Retained spans always cover the model window: construction - # rejects budget < window, and the pinned prompt only adds. + # Retained spans always cover the model window (construction rejects budget < window). bufs = self._buffers_for(layout, prepared) with nvtx_range_debug("triattention.page_table_stage", color="orange"): dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( @@ -1760,8 +1664,7 @@ def _evict_requests( raise RuntimeError("TriAttention attempted an identity compaction") request_state = self._request_states[item["request_id"]] request_state["evicted_tokens"] += evicted - # The manager's only channel to the runtime: the engine subtracts - # it where it builds num_cached_tokens_per_seq. + # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] capacity_targets.append( (item["request_id"], item["kv_cache"], item["draft_kv_cache"], keep_count) @@ -1771,9 +1674,7 @@ def _evict_requests( def _num_layers_from_manager(self) -> int: return len(self.kv_cache_manager.pp_layers) - # ------------------------------------------------------------------ # - # Helpers: calibration loading # - # ------------------------------------------------------------------ # + # ---- helpers: calibration loading ---- def _resolve_calibration(self) -> Dict[str, torch.Tensor]: """Load the user-supplied calibration .pt and return our runtime schema. @@ -1876,8 +1777,7 @@ def analytic_inv_freq(): return (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() if rope_type == "default": - # The original RoPE formula (transformers>=5.5 computes it per-model - # and no longer keys "default" in ROPE_INIT_FUNCTIONS). + # transformers>=5.5 no longer keys "default" in ROPE_INIT_FUNCTIONS: use the formula. omega, scale_sq = analytic_inv_freq(), 1.0 else: try: 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 index de5682ae5079..5422ca8db00e 100644 --- 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 @@ -1,13 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""SM100 CuTe-DSL scorer for the TriAttention mean-score path. - -This is the ONLY score implementation: the per-head modes launch its score-only -entry and union eviction launches its fused score+stats+union pipeline. It -uses split real/imag TMA loads, BF16 and FP16 compensated UMMA, sqrt FTZ, -and producer-only page-ID lookahead. Geometry outside the exact contract -raises loudly at kernel construction; there is no fallback path. -""" +"""SM100 CuTe-DSL scorer for the TriAttention mean-score path (score-only and fused +score+stats+union entries); geometry outside the exact contract raises at construction.""" from __future__ import annotations @@ -28,12 +22,7 @@ @dsl_user_op def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Float32: - """Emit the approximate FTZ square root as inline PTX. - - ``cute.math.sqrt``'s fast-sqrt keyword spelling varies across CuTe DSL - releases; the inline-asm form is release-independent and matches the - ``sqrt.approx.ftz.f32`` the fused score path has always used. - """ + """Inline-PTX sqrt.approx.ftz.f32 (cute.math.sqrt's fast-sqrt kwarg varies across releases).""" return cutlass.Float32( llvm.inline_asm( T.f32(), @@ -50,24 +39,19 @@ def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Fl CTA_M = 128 -# Minimum tcgen05 MMA tile N: GQA groups below 8 ride zero-padded head -# columns (see the weight-builder loop and the partial-stats epilogue). +# Minimum tcgen05 MMA tile N; GQA groups below 8 ride zero-padded head columns. N = 8 THREADS = 256 EPILOGUE_THREADS = 128 RAW_PAGE_BUFFERS = 2 -# partial_stats is a flat [stats_row, page_shard, {count, mean, m2}] record -# array (stats_row = segment * num_q_heads + q_head); the selection finalizer -# imports these field constants. +# 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: N first-tile score origins, then one (sum, square-sum) -# pair per (epilogue warp, padded head column). +# Stats smem scratch: N score origins + one (sum, square-sum) pair per (warp, head column). STATS_ORIGIN_SLOTS = N STATS_SCRATCH_ELEMENTS = STATS_ORIGIN_SLOTS + (EPILOGUE_THREADS // 32) * N * 2 -# HND K pools interleave K/V planes: staged block-offset entries encode -# physical_page * K_PLANES_PER_POOL_PAGE + plane. +# Staged block-offset entries encode physical_page * K_PLANES_PER_POOL_PAGE + plane. K_PLANES_PER_POOL_PAGE = 2 RAW_K_VECTOR_ELEMENTS = 8 @@ -108,9 +92,7 @@ def __init__( 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: - # The schedule assumes one page never spans multiple compute - # tiles (the retired page_half loop generalized this; both - # supported page sizes make it a single iteration). + # 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") @@ -128,8 +110,7 @@ def __init__( # cos/sin/mlr coefficient planes per frequency. self.k_coeff = 3 * num_freqs self.tokens_per_block = tokens_per_block - # One compute tile = one page (128-token pages) or four page - # fragments (32-token pages), one TMA box each. + # 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.pages_per_tile = self.fragments_per_phase @@ -139,9 +120,7 @@ def __init__( # Producer staging constants baked into the generated code. self.prefetch_depth = 4 self.raw_tma_feature_extent = num_freqs - # Barrier transaction bytes for one phase: the full 128-token tile - # of one coefficient plane, regardless of how many page fragments - # deliver it. + # 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 self.accumulator_pipeline_stages = 1 @@ -229,10 +208,7 @@ def _stage_raw_band_copies( shared_partition, stage_args, ): - """One raw-K band's fragment copies into an acquired pipeline stage. - - The caller owns the pipeline handshake (acquire before, advance after). - """ + """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 @@ -298,8 +274,7 @@ def __call__( tcgen05.CtaGroup.ONE, self.mma_tiler[:2], ) - # One real + one imaginary CTA_M x num_freqs bf16 stage per - # raw-page buffer; swizzle follows the num_freqs row width. + # 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, N, self.num_freqs), @@ -363,8 +338,7 @@ def __call__( 1, ) acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) - # One accumulator slot; the explicit slot mode keeps the shared - # producer/consumer slicing protocol and folds away in codegen. + # 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) @@ -419,9 +393,7 @@ class SharedStorage: ] self.shared_storage = SharedStorage - # 64-bit: at large request counts this product exceeds 2^31 (the - # score scratch spans request*layer segments of seq_len columns), so - # the score-plane stride must reach the kernel as Int64. + # The score-plane stride product can exceed 2^31, so it must reach the kernel as Int64. sum_seq = cutlass.Int64(request_count * self.num_layers * self.seq_len) num_ctas = request_count * self.num_layers * self.num_kv_heads * self.page_shards self.kernel( @@ -534,16 +506,13 @@ def kernel( cpasync_raw_k_imag = cpasync_raw_k_0[(None, None, None, 1)] 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)] - # Each stage slice retains the swizzled smem pointer flags. Reuse - # only the feature-first outer mapping for the corresponding TMA - # destination so the swizzle is not applied twice. + # 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 page fragment; fragment offsets - # are whole swizzle periods. + # 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 = [] @@ -608,8 +577,7 @@ def kernel( (raw_tma_descriptors.iterator + layer_id * TMA_DESCRIPTOR_QWORDS).align(128), cute.AddressSpace.generic, ) - # Trace-time invariants of every raw-band stage copy (see - # _stage_raw_band_copies); bound once, unpacked in the helper. + # 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, @@ -636,9 +604,7 @@ def kernel( physical_fragments_arg = None prefetched_fragments_arg = None if cutlass.const_expr(self.fragments_per_phase > 1): - # Per-fragment page-id registers for multi-page compute tiles. - # Slot 0 is unused: fragment 0 keeps the scalar broadcast - # registers of the validated single-fragment schedule. + # Per-fragment page-id registers; slot 0 unused (fragment 0 uses the scalar registers). producer_prefetched_page_ids_lane0 = cute.make_rmem_tensor( (self.pages_per_tile,), cutlass.Int32 ) @@ -651,16 +617,14 @@ def kernel( if cutlass.dynamic_expr(shard_has_page): if warp_idx == self.producer_warp_id: if lane_idx == 0: - # The staged K-plane entries encode physical_page * - # kv_factor (2); decode to the pool page index here. + # Staged entries encode physical_page * kv_factor; decode to the pool page. producer_prefetched_page_id_lane0 = ( cutlass.Int32(page_ids[page_off + tile_index * self.pages_per_tile]) // K_PLANES_PER_POOL_PAGE ) if cutlass.const_expr(self.fragments_per_phase > 1): for fragment in cutlass.range_constexpr(1, self.pages_per_tile): - # Clamp the tail fragment's page so the TMA never - # dereferences an unstaged block entry. + # 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 < valid_seq_len: fragment_page_id = ( @@ -716,8 +680,7 @@ def kernel( num_threads=EPILOGUE_THREADS, ) cute.arch.mbarrier_init_fence() - # Per-(head, frequency) score coefficients, split into bf16/fp16 - # value+residual pairs. + # Per-(head, frequency) score coefficients, split into bf16/fp16 value+residual pairs. for weight_round in cutlass.range_constexpr(N * self.k_coeff // THREADS): linear_index = tidx + weight_round * THREADS qg = linear_index // self.k_coeff @@ -725,8 +688,7 @@ def kernel( 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. + # Padded GQA columns read the group's first head and force zero coefficients. qg_read = qg if cutlass.const_expr(self.group_size < N): if qg_read >= self.group_size: @@ -903,9 +865,7 @@ def kernel( ) if cutlass.const_expr(self.fragments_per_phase > 1): for fragment in cutlass.range_constexpr(1, self.pages_per_tile): - # Same tail-tile clamp as the initial - # prefetch: fall back to the first - # fragment's page. + # 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 @@ -965,9 +925,7 @@ def kernel( 0, ) if cutlass.dynamic_expr(prefetch_next_raw): - # ONE acquire/advance per band is SHARED across the - # dynamic destination arms (current vs next buffer); - # only the smem destination differs per arm. + # 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): @@ -1013,12 +971,10 @@ def kernel( raw_stage_args, ) raw_tma_producer_state.advance() - # Each of the 32 lanes stages one frequency per pass; 64- - # frequency heads take two passes. + # 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 from the - # raw-K shared buffer before consuming any of them. + # Stage prefetch_depth independent token loads before consuming any of them. for token_base in cutlass.range( 0, CTA_M // (THREADS // 32), @@ -1109,9 +1065,7 @@ def kernel( tCrRawBf16B1[(None, None, imag_b_block, 0)], tCtAcc, ) - # Compensated FP16 magnitude accumulation: - # |K|*coeff = A0*B0 + A0*B1 + A1*B0 (the A1*B1 term is - # below fp32 accumulation resolution and dropped). + # Compensated FP16 magnitude: |K|*coeff = A0*B0 + A0*B1 + A1*B0 (A1*B1 dropped). magnitude_lo_tiled_mma.set( tcgen05.Field.ACCUMULATE, True, @@ -1146,13 +1100,10 @@ def kernel( 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): - # The alternate raw-page buffer was filled while this - # page's UMMA completed. Release only the current imag - # phase after all of its asynchronous consumers finish. + # 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 sum_seq must stay 64-bit (the plane - # stride can exceed 2^31; the host audit bounds the head fold). + # Every term multiplying sum_seq must stay 64-bit; the plane stride can exceed 2^31. output_offset = cutlass.Int64(kv_head * N) * sum_seq + out_base + tile_start_token page_output = cute.make_tensor( output.iterator + output_offset, @@ -1265,8 +1216,7 @@ def kernel( sStats[stats_scratch_base + 1] = stats_square_sum cute.arch.barrier() if warp_idx == 0: - # Only real heads' statistics merge into the compact row - # layout (row = segment * num_q_heads + q_head). + # 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) @@ -1318,7 +1268,6 @@ def _encode_tma_descriptors( num_freqs: int, tokens_per_block: int, ) -> torch.Tensor: - """Encode one immutable feature-first TensorMap per layer index.""" anchor = layer_pools[layer_indices[0]] active_layers = set(layer_indices) uint32 = cuda.cuuint32_t @@ -1353,8 +1302,7 @@ def _encode_tma_descriptors( 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 - # (the whole page for the validated 128-token geometry). + # 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, @@ -1365,9 +1313,7 @@ def _encode_tma_descriptors( [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 smem layout the TMA lands in; the - # sm100 helpers pick it from the inner-row byte count (one - # coefficient plane: num_freqs bf16 elements). + # 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 @@ -1444,16 +1390,13 @@ def __init__( ) -> None: self.max_requests = int(max_requests) self.num_layers = int(num_layers) - # The score window start is a per-request runtime input - # (``token_starts``), so the widest window — the whole bucket — - # sizes every start-dependent buffer. + # The widest score window (the whole bucket) sizes every start-dependent buffer. self.width = int(seq_len) self.num_q_heads = int(num_q_heads) self.num_kv_heads = int(num_kv_heads) self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) self.enable_partial_stats = bool(enable_partial_stats) - # One [stats_row, page_shard, {count, mean, m2}] record array (see the - # STATS_FIELDS constants above). + # One [stats_row, page_shard, {count, mean, m2}] record array. partial_stats_elements = ( max_requests * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * STATS_FIELDS if self.enable_partial_stats @@ -1486,9 +1429,7 @@ def __init__( layer_pools[layer_indices[0]], self.descriptors, ) - # valid_seq_lens/token_starts are row views into the staged metadata - # table (byte offset 4*(max_requests+1)*row): only 4-byte aligned, - # and only ever read as per-CTA scalars. + # valid_seq_lens/token_starts are only 4-byte-aligned row views, read as per-CTA scalars. prefix_aligns = (16, 16, 16, 16, 4, 16, 4, 16, 16, 16) self._cute_prefix = tuple( _to_cute(tensor, assumed_align=align) @@ -1543,8 +1484,7 @@ def __init__( variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] if max_requests > 1: variants.append((max_requests, 2)) - # Per-head runners compile the score-only entry; the union runner - # only its fused stats+union pipeline. + # Per-head runners compile the score-only entry; the union runner the fused pipeline. kernel_kwargs = dict( num_layers=num_layers, seq_len=seq_len, @@ -1596,9 +1536,7 @@ def __init__( small = compiled_entries.get(1) large = compiled_entries.get(max_requests) for request_count in range(1, max_requests + 1): - # Shard-pick heuristic: give small cohorts the extra page - # shard while the 2-shard grid stays under two waves - # (2 * sm_count CTAs); larger cohorts already fill the GPU. + # Give small cohorts the extra shard while the 2-shard grid stays under two waves. two_shard_ctas = request_count * num_layers * num_kv_heads * 2 use_extra_score_shard = two_shard_ctas < 2 * self.sm_count compiled_entries[request_count] = small if use_extra_score_shard else large @@ -1639,8 +1577,7 @@ def __init__( num_layers=num_layers, seq_len=seq_len, num_q_heads=num_q_heads, - # The finalizer maps real head rows onto the - # N=8-padded score planes. + # The finalizer maps real head rows onto N=8-padded planes. num_kv_heads=num_kv_heads, page_shards=page_shards, tokens_per_lane=tokens_per_lane, 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 index da6cae393729..8583bac57150 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -28,9 +28,7 @@ from .triattention_cute_score_fused import STATS_FIELDS as _STATS_FIELDS from .triattention_cute_score_fused import STATS_M2, STATS_MEAN -# Single-sourced constants: the fused score file owns the padded-head tile -# (N) and the partial-stats record layout it writes; the Triton kernels -# module owns the z-normalization epsilon shared by both selection paths. +# Single-sourced constants: the score file owns N and the stats layout; Triton owns the epsilon. from .triattention_cute_score_fused import N as _PADDED_HEAD_COLUMNS from .triattention_kernels import STD_EPSILON as _STD_EPSILON @@ -50,8 +48,7 @@ def _select_normalize_union_config( width: int, sm_count: int, ) -> tuple[int, int, int]: - """Return tokens per lane, token subtiles, and row-cluster CTAs - (small clustered tile while the grid fits residency, else the large tile).""" + """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 @@ -118,11 +115,7 @@ def _ld_cluster_f32(mapped_addr): def _gmem_lane_tile(iterator, flat_index, tokens_per_lane, assumed_align): - """One lane's gmem tile of ``tokens_per_lane`` fp32 values. - - Folds the 64-bit flat index into the pointer BEFORE element access, so - nothing routes through the DSL's 32-bit dynamic coordinate. - """ + """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, @@ -149,16 +142,12 @@ def __init__( token_subtiles: int, row_cluster_ctas: int, ) -> None: - # The score scratch pads each KV head's group of score planes up to - # the MMA tile: real head row ``q_head`` lives in score plane - # ``kv * 8 + qg``. The partial-stats rows are always compact. + # 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.seq_len = seq_len - # The score window start is per-request runtime metadata - # (``token_starts``); the widest window — the whole bucket — - # sizes the output rows and the token-tile grid. + # The widest score window (the whole bucket) sizes the output rows and token-tile grid. self.width = seq_len self.num_q_heads = num_q_heads self.num_rows = num_layers * num_q_heads @@ -200,8 +189,7 @@ def __call__( stream=stream, ) else: - # Cluster peers are consecutive CTAs: the kernel's (request, - # token-tile, rank) decode must keep this factor order. + # 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, @@ -227,15 +215,13 @@ def _reduce_and_store_union_rows( lane_idx: cutlass.Int32, from_cluster_peers: cutlass.Constexpr, ): - """Final peer reduce + union-row store, shared by both arms - (the peer source is the only difference, picked at trace time).""" + """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 object the smem - # tensor was built with (no hand-derived strides). + # 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 ) @@ -256,8 +242,7 @@ def _reduce_and_store_union_rows( ) reduced_values[token_slot] = union_value subtile_first_token = first_token + token_subtile * self.subtile_token_tile - # Straddling subtiles fall back to per-token stores; union_scores - # stays < 2^31 by the host audit, so the i32 index cannot wrap. + # Straddling subtiles store per token; union_scores stays < 2^31 so i32 cannot wrap. if cutlass.const_expr(self.width % self.tokens_per_lane == 0) and cutlass.dynamic_expr( subtile_first_token + self.tokens_per_lane <= valid_width ): @@ -305,8 +290,7 @@ def kernel( 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 - # Per-request score window start: the normalization domain and the - # union output row both cover [0, valid - start) for this request. + # The normalization domain and the union output row both cover [0, valid - start). score_start = cutlass.Int32(token_starts[request_idx]) valid_width = valid_seq_lens[request_idx] - score_start warp_max_ptr = cute.arch.alloc_smem( @@ -388,8 +372,7 @@ def kernel( 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 (possibly padded) score plane; - # the padded planes carry zero scores and are never visited. + # 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 @@ -433,9 +416,7 @@ def kernel( + score_start + subtile_first_token ) - # The vectorized load also needs the runtime start aligned to - # the lane width (subtile_first_token and the segment stride - # are aligned whenever seq_len is). + # The vectorized load needs the runtime start aligned to the lane width. if cutlass.const_expr( self.seq_len % self.tokens_per_lane == 0 ) and cutlass.dynamic_expr( @@ -454,10 +435,7 @@ def kernel( cute.coalesce(score_value_tiles[token_subtile]), ) else: - # Fold the 64-bit index into the pointer BEFORE the loads: - # the scratch exceeds 2^31 elements and the DSL's 32-bit - # dynamic coordinate wraps. Unaligned window starts hit - # this arm on every real serve round. + # 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 ) @@ -528,8 +506,7 @@ def kernel( False, ) else: - # Warp 0 reduces its CTA's row partition, then CTA 0 combines the - # cluster's partial maxima through distributed shared memory. + # 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): diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index e57b85ccf285..bbed90a0d3e9 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -1,11 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Triton kernels, launch helpers, and mean-phase table builders for TriAttention. - -Scoring itself runs through the SM100 CuTe pack in -``triattention_cute_score_fused.py``. House rules: fp32 math, int64 for any -flat offset that can exceed 2^31, mask every ragged tail, kernels vendored here. -""" +"""Triton kernels, launch helpers, and mean-phase table builders for TriAttention +(fp32 math; int64 past-2^31 flat offsets; masked ragged tails; scoring in the CuTe pack).""" from __future__ import annotations @@ -15,16 +11,13 @@ import triton import triton.language as tl -# --------------------------------------------------------------------------- # -# Mean-phase table: RoPE-style position table of mean trig phases. # -# --------------------------------------------------------------------------- # +# ---- Mean-phase table: RoPE-style position table of mean trig phases ---- # Positions past this row count are not exactly representable in fp32. _MEAN_PHASE_MAX_ROWS = 1 << 24 -# Score z-normalization epsilon, shared with the CuTe union pipeline. Plain -# float: the CuTe DSL traces it; the Triton kernel takes it as a constexpr arg. +# Score z-normalization epsilon; must stay a plain float (the CuTe DSL traces it). STD_EPSILON = 1e-6 @@ -45,8 +38,7 @@ def _gather_mean_phase_kernel( F_BLOCK: tl.constexpr, HAS_SWA: tl.constexpr, ): - """Copy each request's phase-table row; derive valid widths and, with - SWA layers, the rebased SWA landing positions in the same launch.""" + """Copy each request's phase-table row; derive valid widths and SWA landing bases.""" request = tl.program_id(0) frequency = tl.arange(0, F_BLOCK) frequency_mask = frequency < NUM_FREQS @@ -91,9 +83,7 @@ def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: phase["rows"] = target -# --------------------------------------------------------------------------- # -# Selection: combine scores per mode, then finalize the top-k set. # -# --------------------------------------------------------------------------- # +# ---- Selection: combine scores per mode, then finalize the top-k set ---- @triton.jit @@ -150,11 +140,7 @@ def _score_per_head_reduce_kernel( NORMALIZE: tl.constexpr, BLOCK: tl.constexpr, ): - """Reduce query-head score rows into one selector row per KV-head domain. - - per_layer: row (layer, kv_head) = max over the KV head's query group; - otherwise row kv_head = mean over layers of that per-layer group max. - """ + """Reduce query-head score rows into one selector row per KV-head domain.""" request = tl.program_id(0) selection_row = tl.program_id(1) token_block = tl.program_id(2) @@ -265,9 +251,7 @@ def prepare_per_head_scores( ) -# --------------------------------------------------------------------------- # -# Compaction: pack the kept ordinals into per-request move indices. # -# --------------------------------------------------------------------------- # +# ---- Compaction: pack the kept ordinals into per-request move indices ---- # Settle/pack launch shape, shared by every launch site. @@ -301,15 +285,8 @@ def _settle_ties_and_pack_compaction_sources_kernel( HAS_SETTLE: tl.constexpr, BLOCK: tl.constexpr, ): - """Settle one selection row's ties, then pack its compaction move sources. - - One program per (request, selection row): settle emits the kept ordinals - (lowest-index-wins ties, prompt-rebased), pack writes the move sources - (kept ordinals + protected tail + SWA rows). ``HAS_SETTLE=False`` packs - pre-settled ordinals from ``output_indices`` (draft co-compaction). - Emission order is load-bearing: the C++ compact requires increasing - ordinals with ``destination_bases[request] + move <= source[move]``. - """ + """Settle one selection row's ties and pack its move sources (increasing + kept ordinals; C++ in-place copy contract).""" request = tl.program_id(0) selection_domain = tl.program_id(1) row = request * SELECTION_ROWS + selection_domain @@ -329,8 +306,7 @@ def _settle_ties_and_pack_compaction_sources_kernel( mask=selected_mask, other=0, ) - # Short rows arrive padded with -1 sentinels from the top-k; - # masked out so no lane dereferences ``row_scores - 1``. + # 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, From a321e3343ab02fa1a862881e9bfd73f65867b845 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 01:07:50 -0700 Subject: [PATCH 113/178] [None][refactor] Collapse init parameter transport; compaction becomes an algorithm-neutral mover (knife 25a) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 216 +++++----- .../triattention/triattention.py | 408 ++++++++---------- .../_torch/kv_cache_compression/conftest.py | 228 ++++++---- .../test_triattention_draft_cocompaction.py | 30 +- .../test_triattention_pipeline.py | 27 +- .../test_triattention_selection_compaction.py | 93 ++-- 6 files changed, 512 insertions(+), 490 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 6d741879ddac..1fa61a6c29f3 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -13,10 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Batched physical KV-cache compaction for eviction-based compression. +"""Batched physical KV-cache compaction: an algorithm-neutral mover. -``init_compaction_buffers`` allocates plain-dict launch data once per -geometry; the eviction driver fires the bundle's fields directly each round. +``init_compaction_buffers`` agrees on the decision buffers (move-source +indices; offsets ride the caller's staged rows) once per geometry and retains +one launch contract. The caller materializes its keep decision into those +buffers each round, then ``compact`` fires the native target and draft +launches. This module knows cache-family geometry and the decision format +only; the contract's launch tuples are private to it. """ from collections import OrderedDict @@ -81,52 +85,64 @@ def _compact_groups( return tuple(result) +def _launch_tuples( + groups: Tuple[Dict[str, object], ...], + source: torch.Tensor, + offsets: torch.Tensor, + destination_bases: torch.Tensor, +) -> Tuple[tuple, ...]: + return tuple( + ( + group["pools"], + group["pool_pointers"], + group["page_table"], + source, + offsets, + destination_bases, + group["source_layer_indices"], + ) + for group in groups + ) + + def init_compaction_buffers( *, - union: bool, - per_layer: bool, - layer_pools: List[torch.Tensor], - dense_layers: List[int], - swa_layers: List[int], - layer_group_representative: Dict[int, int], - valid_sequence_lengths: torch.Tensor, - kv_block_offsets: torch.Tensor, - page_table_slots: Dict[int, int], - request_count: int, - prompt_offsets: torch.Tensor, - decode_keep_count: int, - swa_window: Optional[int], - layer_pool_keys: List[object], - protected_tail_capacity: int = 0, - draft_layer_pools: Optional[List[torch.Tensor]] = None, - draft_layers: Optional[List[int]] = None, - draft_layer_group_representative: Optional[Dict[int, int]] = None, - draft_layer_pool_keys: Optional[List[object]] = None, - draft_protected_tail_capacity: Optional[int] = None, - draft_kv_block_offsets: Optional[torch.Tensor] = None, - draft_page_table_slots: Optional[Dict[int, int]] = None, - dense_move_offsets: torch.Tensor, - swa_move_offsets: Optional[torch.Tensor] = None, - draft_move_offsets: Optional[torch.Tensor] = None, + target: Dict[str, object], + capacities: Dict[str, int], + draft: Optional[Dict[str, object]] = None, ) -> Dict[str, object]: - """Allocate per-geometry compaction launch data for the driver's round. + """Agree on the decision buffers and retain one launch contract per geometry. - Selection rows: union = one per request; per_layer = one per (layer, KV - head); else one per KV head. Move sources must be increasing kept ordinals - with destination_bases[request] + move <= source[move] (C++ in-place copy - contract). Returns the launch bundle the round function fires directly. + Move sources must be increasing kept ordinals with + destination_bases[request] + move <= source[move] (C++ in-place copy + contract). ``target`` carries the resolved dense/SWA grouping inputs from + the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows); + ``draft`` is one all-or-none resolved branch; ``capacities`` the + request/keep/tail capacity numbers. The returned contract exposes the + agreed move-source buffers and geometry constants; its launch tuples are + private to :func:`compact`. """ + layer_pools = target["layer_pools"] + dense_layers = tuple(int(layer) for layer in target["dense_layers"]) + swa_layers = tuple(int(layer) for layer in target["swa_layers"]) + layer_pool_keys = tuple(target["layer_pool_keys"]) + kv_block_offsets = target["kv_block_offsets"] + page_table_slots = target["page_table_slots"] + layer_group_representative = target["layer_group_representative"] + prompt_offsets = target["prompt_offsets"] + dense_move_offsets = target["dense_move_offsets"] + swa_move_offsets = target["swa_move_offsets"] + swa_window = target["swa_window"] + per_layer_sources = bool(target["per_layer_sources"]) + device = layer_pools[dense_layers[0]].device - request_count = int(request_count) - decode_keep_count = int(decode_keep_count) - protected_tail_capacity = int(protected_tail_capacity) - dense_layers = tuple(int(layer) for layer in dense_layers) - swa_layers = tuple(int(layer) for layer in swa_layers) - layer_pool_keys = tuple(layer_pool_keys) + request_count = int(capacities["request_capacity"]) + decode_keep_count = int(capacities["decode_keep_count"]) + protected_tail_capacity = int(capacities["protected_tail_capacity"]) # Pool shape [pages, K/V, heads, tokens, dim]. num_kv_heads = int(layer_pools[dense_layers[0]].shape[2]) - dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer else (num_kv_heads,) + dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer_sources else (num_kv_heads,) dense_move_indices = _make_move_indices( dense_index_prefix, decode_keep_count + protected_tail_capacity, @@ -169,54 +185,40 @@ def init_compaction_buffers( for layer in swa_layers ] - dense_slots = {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer else None - # Without SWA layers the SWA pointer args are None (HAS_SWA=False). + dense_slots = ( + {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer_sources else None + ) has_swa = swa_move_indices is not None swa_total = int(swa_move_indices.shape[-1]) if has_swa else 0 # Widest per-request move count any staged offsets may express. move_capacity = decode_keep_count + protected_tail_capacity if has_swa: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) - settle_pack_tensors = ( - valid_sequence_lengths, - dense_move_offsets, - dense_move_indices, - swa_move_offsets if has_swa else None, - swa_move_indices if has_swa else None, - ) - settle_pack_shape = dict( - DENSE_TOTAL=int(dense_move_indices.shape[-1]), - SWA_TOTAL=swa_total, - MOVE_CAPACITY=move_capacity, - NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window, - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=has_swa, - ) - families = [ - dict( - name="dense", - groups=_compact_groups(dense_entries, layer_pool_keys, device, dense_slots), - source=dense_move_indices, - offsets=dense_move_offsets, - destination_bases=prompt_offsets, + + target_launches = list( + _launch_tuples( + _compact_groups(dense_entries, layer_pool_keys, device, dense_slots), + dense_move_indices, + dense_move_offsets, + prompt_offsets, ) - ] + ) if swa_layers: - families.append( - dict( - name="swa", - groups=_compact_groups(swa_entries, layer_pool_keys, device), - source=swa_move_indices, - offsets=swa_move_offsets, - destination_bases=swa_destination_bases, + target_launches.extend( + _launch_tuples( + _compact_groups(swa_entries, layer_pool_keys, device), + swa_move_indices, + swa_move_offsets, + swa_destination_bases, ) ) - if draft_layers: - draft_tail = int(draft_protected_tail_capacity or 0) - draft_layers = tuple(int(layer) for layer in draft_layers) + draft_launches: Tuple[tuple, ...] = () + draft_move_indices = None + if draft is not None: + draft_layer_pools = draft["layer_pools"] + draft_layers = tuple(int(layer) for layer in draft["layers"]) + draft_tail = int(draft["protected_tail_capacity"]) # Own launch groups: the draft may use a different KV-head count. draft_num_kv_heads = int(draft_layer_pools[draft_layers[0]].shape[2]) draft_move_indices = _make_move_indices( @@ -229,40 +231,52 @@ def init_compaction_buffers( ( layer, draft_layer_pools[layer], - draft_kv_block_offsets[ - draft_page_table_slots[draft_layer_group_representative[layer]], + draft["kv_block_offsets"][ + draft["page_table_slots"][draft["layer_group_representative"][layer]], :request_count, 0, ], ) for layer in draft_layers ] - # Geometry constants for the draft's own pack launch. - draft_pack = dict( - indices=draft_move_indices, - offsets=draft_move_offsets, - dense_total=int(draft_move_indices.shape[-1]), - move_capacity=decode_keep_count + draft_tail, - num_kv_heads=draft_num_kv_heads, - ) - families.append( - dict( - name="draft", - groups=_compact_groups(draft_entries, tuple(draft_layer_pool_keys), device), - source=draft_move_indices, - offsets=draft_move_offsets, - destination_bases=prompt_offsets, - ) + draft_launches = _launch_tuples( + _compact_groups(draft_entries, tuple(draft["layer_pool_keys"]), device), + draft_move_indices, + draft["move_offsets"], + prompt_offsets, ) - else: - draft_pack = None return dict( - families=families, - settle_pack_tensors=settle_pack_tensors, - settle_pack_shape=settle_pack_shape, - draft_pack=draft_pack, + # Agreed decision buffers and geometry constants (the public interface). + dense_move_indices=dense_move_indices, + swa_move_indices=swa_move_indices, + draft_move_indices=draft_move_indices, + dense_total=int(dense_move_indices.shape[-1]), + swa_total=swa_total, + move_capacity=move_capacity, + num_kv_heads=num_kv_heads, + swa_window=swa_window, + has_swa=has_swa, swa_destination_bases=swa_destination_bases, # Per-round SWA destination rebase delta. swa_rebase_delta=decode_keep_count - swa_window, + # Completion event: compact() records it after the last native launch. + consume_done=torch.cuda.Event(), + # Private launch tuples: only compact() interprets these. + target_launches=tuple(target_launches), + draft_launches=draft_launches, + has_draft=draft is not None, ) + + +def compact(compaction: Dict[str, object], request_count: int) -> None: + """Fire the native target compacts, then the draft compacts, and record completion. + + Pure mover: the caller has already materialized its keep decision into the + agreed move-source buffers for the active ``request_count`` cohort. + """ + for launch in compaction["target_launches"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) + for launch in compaction["draft_launches"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) + compaction["consume_done"].record(torch.cuda.current_stream()) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index a0c93374b032..f1e527e0ddcb 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -39,7 +39,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig -from ..compaction import init_compaction_buffers +from ..compaction import compact, init_compaction_buffers from .triattention_kernels import ( SETTLE_PACK_BLOCK, SETTLE_PACK_NUM_WARPS, @@ -96,129 +96,52 @@ def _allocate_page_table_plane( return representative_slots, copy_block_count, host, dev -def attach_compaction_bundle(bufs: SimpleNamespace, compaction: Dict[str, object]) -> None: - """Flatten one compaction bundle into the prebound launch fields the round fires.""" - bufs.compaction_families = compaction["families"] - bufs.settle_pack_tensors = compaction["settle_pack_tensors"] - bufs.settle_pack_shape = compaction["settle_pack_shape"] - bufs.draft_pack = compaction["draft_pack"] - bufs.swa_destination_bases = compaction["swa_destination_bases"] - bufs.swa_rebase_delta = compaction["swa_rebase_delta"] - - def group_args(family): - return tuple( - ( - group["pools"], - group["pool_pointers"], - group["page_table"], - family["source"], - family["offsets"], - family["destination_bases"], - group["source_layer_indices"], - ) - for group in family["groups"] - ) - - target_args: List[tuple] = [] - draft_args: Tuple[tuple, ...] = () - for family in compaction["families"]: - if family["name"] == "draft": - # Kept separate: the draft pack launch must precede the draft moves. - draft_args = group_args(family) - else: - target_args.extend(group_args(family)) - bufs.compact_launch_args = tuple(target_args) - bufs.draft_compact_launch_args = draft_args - draft_pack = compaction["draft_pack"] - if draft_pack is None: - bufs.draft_pack_args = None - bufs.draft_pack_kwargs = None - return - # Pack-only launch (HAS_SETTLE=False): settle-side pointers are None. - bufs.draft_pack_args = ( - None, - None, - None, - None, - bufs.keep, - bufs.valid_seq_lens_device, - draft_pack["offsets"], - draft_pack["indices"], - None, - None, - ) - bufs.draft_pack_kwargs = dict( - WIDTH=bufs.keep_count, - KEEP_COUNT=bufs.keep_count, - SELECTION_ROWS=1, - DENSE_TOTAL=draft_pack["dense_total"], - SWA_TOTAL=0, - MOVE_CAPACITY=draft_pack["move_capacity"], - NUM_KV_HEADS=draft_pack["num_kv_heads"], - SWA_WINDOW=0, - UNION=True, - PER_LAYER=False, - HAS_SWA=False, - HAS_SETTLE=False, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, - ) - - def init_eviction_buffers( *, eviction_mode: str, - layer_pools: List[torch.Tensor], - dense_groups: List[List[int]], - dense_layers: List[int], - swa_layers: Sequence[int] = (), - swa_window: Optional[int] = None, - layer_group_representative: Optional[Dict[int, int]] = None, - layer_pool_keys: Optional[List[object]] = None, - page_representatives: List[int], - max_requests: int, - seq_len: int, - num_q_heads: int, - num_freqs: int, - keep_count: int, - q_real: torch.Tensor, - q_imag: torch.Tensor, - mlr_coef: torch.Tensor, - freq_scale_sq: torch.Tensor, - offsets: torch.Tensor, - omega: torch.Tensor, - phase: Optional[Dict[str, object]] = None, - page_table_keys: List[object], - num_page_table_slots: int, - decode_width: int, - page_table_token_capacity: int, - protected_tail_capacity: int = 0, - draft_layer_pools: Optional[List[torch.Tensor]] = None, - draft_layers: Optional[List[int]] = None, - draft_layer_group_representative: Optional[Dict[int, int]] = None, - draft_layer_pool_keys: Optional[List[object]] = None, - draft_page_representatives: Optional[List[int]] = None, - draft_page_table_keys: Optional[List[object]] = None, - draft_num_page_table_slots: Optional[int] = None, - draft_page_table_token_capacity: Optional[int] = None, - draft_protected_tail_capacity: int = 0, + layout: Dict[str, object], + calibration: Dict[str, torch.Tensor], + phase: Dict[str, object], + capacities: Dict[str, int], + draft: Optional[Dict[str, object]] = None, ) -> SimpleNamespace: """Build the one namespace of buffers, compiled launches, and compaction data - (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put).""" + (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put). + + ``layout`` is the runtime KV layout dict, passed whole; ``calibration`` + carries the local q_real/q_imag/mlr_coef [L, H, F] slices and + freq_scale_sq; ``draft`` is one all-or-none resolved branch (its layout + dict plus tail/page-table capacities); ``capacities`` the capacity numbers. + """ from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS from .triattention_cute_score_fused import TriAttentionCuteScoreRunner + layer_pools = layout["layer_pools"] + dense_layers = list(layout["dense_layers"]) + swa_layers = list(layout["swa_layers"]) + swa_window = layout["swa_window"] + layer_group_representative = layout["layer_group_representative"] + layer_pool_keys = list(layout["layer_pool_keys"]) + dense_groups = list(layout["storage_groups"].values()) + page_representatives = [group[0] for group in dense_groups] + page_representatives.extend(layer for layer in swa_layers if layer not in page_representatives) + page_table_keys = [layer_pool_keys[layer] for layer in page_representatives] + num_page_table_slots = int(layout["manager"].num_pools) + device = layer_pools[page_representatives[0]].device - max_requests = int(max_requests) - seq_len = int(seq_len) - page_table_token_capacity = int(page_table_token_capacity) - decode_width = int(decode_width) - keep_count = int(keep_count) - - q_real, q_imag, mlr_coef, freq_scale_sq, offsets, omega = ( - tensor.to(device=device, dtype=torch.float32).contiguous() - for tensor in (q_real, q_imag, mlr_coef, freq_scale_sq, offsets, omega) + max_requests = int(capacities["max_requests"]) + seq_len = int(capacities["bucket_seq_len"]) + page_table_token_capacity = int(capacities["page_table_token_capacity"]) + decode_width = int(capacities["decode_width"]) + keep_count = int(capacities["keep_count"]) + protected_tail_capacity = int(capacities["protected_tail_capacity"]) + + q_real, q_imag, mlr_coef, freq_scale_sq = ( + calibration[key].to(device=device, dtype=torch.float32).contiguous() + for key in ("q_real", "q_imag", "mlr_coef", "freq_scale_sq") ) + num_q_heads = int(q_real.shape[1]) + num_freqs = int(q_real.shape[2]) bufs = SimpleNamespace() bufs.eviction_mode = eviction_mode @@ -249,18 +172,20 @@ def init_eviction_buffers( bufs._draft_bulk_offsets_src = None bufs.draft_copy_block_count = 0 draft_page_slots: Dict[int, int] = {} - if draft_layer_pools is not None: + if draft is not None: + draft_layout = draft["layout"] + draft_representatives = list(draft_layout["pool_representatives"]) ( draft_page_slots, bufs.draft_copy_block_count, bufs._draft_bulk_offsets_src, bufs.draft_block_offsets_device, ) = _allocate_page_table_plane( - draft_layer_pools, - draft_page_representatives, - draft_page_table_keys, - int(draft_num_page_table_slots), - int(draft_page_table_token_capacity), + draft_layout["layer_pools"], + draft_representatives, + [draft_layout["layer_pool_keys"][layer] for layer in draft_representatives], + int(draft_layout["manager"].num_pools), + int(draft["page_table_token_capacity"]), max_requests, device, ) @@ -287,17 +212,6 @@ def init_eviction_buffers( draft_move_offsets_row = bufs.request_metadata_device[5] bufs.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) bufs.mean_sin = torch.empty_like(bufs.mean_cos) - # ``phase=None`` is the one documented test seam (private table built here). - if phase is None: - phase = { - "offsets": offsets.contiguous(), - "omega": omega.contiguous(), - "offset_values": offsets.tolist(), - "cos": None, - "sin": None, - "rows": 0, - } - grow_mean_phase_table(phase, max(int(seq_len), 1)) bufs.phase = phase bufs.phase_num_freqs = int(phase["omega"].numel()) bufs.phase_f_block = triton.next_power_of_2(bufs.phase_num_freqs) @@ -461,42 +375,110 @@ def init_eviction_buffers( bufs.keep_rows = bufs.keep.view(-1, keep_count) bufs.top_indices_i32.zero_() - # ---- compaction launch data + settle/pack fusion ------------------------ + # ---- compaction contract + decision-materialization prebinds ------------ bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) - draft_kwargs = {} - if draft_layers: - draft_kwargs = dict( - draft_layer_pools=draft_layer_pools, - draft_layers=list(draft_layers), - draft_layer_group_representative=draft_layer_group_representative, - draft_layer_pool_keys=draft_layer_pool_keys, - draft_protected_tail_capacity=int(draft_protected_tail_capacity), - draft_kv_block_offsets=bufs.draft_block_offsets_device, - draft_page_table_slots=draft_page_slots, - draft_move_offsets=draft_move_offsets_row, + per_layer = eviction_mode == "per_layer_perhead" + draft_contract = None + if draft is not None: + draft_layout = draft["layout"] + draft_contract = dict( + layer_pools=draft_layout["layer_pools"], + layers=list(draft_layout["dense_layers"]), + layer_group_representative=draft_layout["layer_group_representative"], + layer_pool_keys=list(draft_layout["layer_pool_keys"]), + kv_block_offsets=bufs.draft_block_offsets_device, + page_table_slots=draft_page_slots, + move_offsets=draft_move_offsets_row, + protected_tail_capacity=int(draft["protected_tail_capacity"]), ) - compaction = init_compaction_buffers( - union=union, - per_layer=eviction_mode == "per_layer_perhead", - layer_pools=layer_pools, - dense_layers=list(dense_layers), - swa_layers=list(swa_layers), - layer_group_representative=layer_group_representative, - valid_sequence_lengths=bufs.valid_seq_lens_device, - kv_block_offsets=bufs.block_offsets_device, - page_table_slots=bufs.representative_slots, - request_count=max_requests, - prompt_offsets=bufs.token_starts_device, - decode_keep_count=keep_count, - swa_window=swa_window, - layer_pool_keys=list(layer_pool_keys), - protected_tail_capacity=int(protected_tail_capacity), - # Per-round tails: the move offsets ride the staged metadata rows. - dense_move_offsets=dense_move_offsets_row, - swa_move_offsets=swa_move_offsets_row, - **draft_kwargs, + contract = init_compaction_buffers( + target=dict( + layer_pools=layer_pools, + dense_layers=list(dense_layers), + swa_layers=list(swa_layers), + swa_window=swa_window, + layer_group_representative=layer_group_representative, + layer_pool_keys=list(layer_pool_keys), + kv_block_offsets=bufs.block_offsets_device, + page_table_slots=bufs.representative_slots, + prompt_offsets=bufs.token_starts_device, + # Per-round tails: the move offsets ride the staged metadata rows. + dense_move_offsets=dense_move_offsets_row, + swa_move_offsets=swa_move_offsets_row, + per_layer_sources=per_layer, + ), + capacities=dict( + request_capacity=max_requests, + decode_keep_count=keep_count, + protected_tail_capacity=int(protected_tail_capacity), + ), + draft=draft_contract, + ) + bufs.compaction = contract + bufs.swa_destination_bases = contract["swa_destination_bases"] + bufs.swa_rebase_delta = contract["swa_rebase_delta"] + # The decision side: settle+pack launch args against the agreed buffers. + bufs.settle_args = ( + bufs.selection_scores_rows, + bufs.selection_row_lengths, + bufs.row_prompt_offsets, + bufs.provisional_rows, + bufs.keep_rows, + bufs.valid_seq_lens_device, + dense_move_offsets_row, + contract["dense_move_indices"], + swa_move_offsets_row if contract["has_swa"] else None, + contract["swa_move_indices"], + ) + bufs.settle_kwargs = dict( + WIDTH=decode_width, + KEEP_COUNT=keep_count, + SELECTION_ROWS=bufs.selection_rows_per_request, + DENSE_TOTAL=contract["dense_total"], + SWA_TOTAL=contract["swa_total"], + MOVE_CAPACITY=contract["move_capacity"], + NUM_KV_HEADS=contract["num_kv_heads"], + SWA_WINDOW=contract["swa_window"], + UNION=union, + PER_LAYER=per_layer, + HAS_SWA=contract["has_swa"], + HAS_SETTLE=True, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, ) - attach_compaction_bundle(bufs, compaction) + bufs.draft_pack_launch = None + if draft is not None: + draft_move_indices = contract["draft_move_indices"] + # Pack-only decision broadcast (HAS_SETTLE=False): settle pointers are None. + broadcast_pack_args = ( + None, + None, + None, + None, + bufs.keep, + bufs.valid_seq_lens_device, + draft_move_offsets_row, + draft_move_indices, + None, + None, + ) + broadcast_pack_kwargs = dict( + WIDTH=keep_count, + KEEP_COUNT=keep_count, + SELECTION_ROWS=1, + DENSE_TOTAL=int(draft_move_indices.shape[-1]), + SWA_TOTAL=0, + MOVE_CAPACITY=int(draft_move_indices.shape[-1]) // max_requests, + NUM_KV_HEADS=int(draft_move_indices.shape[0]), + SWA_WINDOW=0, + UNION=True, + PER_LAYER=False, + HAS_SWA=False, + HAS_SETTLE=False, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, + ) + bufs.draft_pack_launch = (broadcast_pack_args, broadcast_pack_kwargs) # ---- round-ordering events ---------------------------------------------- bufs.copy_done = torch.cuda.Event() @@ -623,7 +605,8 @@ def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cud def settle_top_tokens(bufs: SimpleNamespace) -> None: - """Pick the top-k, settle ties to sorted ordinals, and pack the move sources.""" + """Pick the top-k, settle ties, and materialize the move-source decision + (target settle+pack, then the draft broadcast pack).""" # The trailing 1 is next_n: decode scores one query token per request. torch.ops.trtllm.cute_dsl_indexer_topk_decode( bufs.selection_scores_rows, @@ -633,20 +616,13 @@ def settle_top_tokens(bufs: SimpleNamespace) -> None: 1, ) _settle_ties_and_pack_compaction_sources_kernel[bufs.settle_grid]( - bufs.selection_scores_rows, - bufs.selection_row_lengths, - bufs.row_prompt_offsets, - bufs.provisional_rows, - bufs.keep_rows, - *bufs.settle_pack_tensors, - WIDTH=bufs.decode_width, - KEEP_COUNT=bufs.keep_count, - SELECTION_ROWS=bufs.selection_rows_per_request, - **bufs.settle_pack_shape, - HAS_SETTLE=True, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, + *bufs.settle_args, **bufs.settle_kwargs ) + if bufs.draft_pack_launch is not None: + broadcast_pack_args, broadcast_pack_kwargs = bufs.draft_pack_launch + _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( + *broadcast_pack_args, **broadcast_pack_kwargs + ) def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: @@ -723,15 +699,7 @@ def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: ) settle_top_tokens(bufs) with nvtx_range("triattention.compact", color="purple"): - # Prebound calls; the draft pack launch precedes the draft moves. - for args in bufs.compact_launch_args: - torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) - if bufs.draft_pack_args is not None: - _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( - *bufs.draft_pack_args, **bufs.draft_pack_kwargs - ) - for args in bufs.draft_compact_launch_args: - torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) + compact(bufs.compaction, request_count) class TriAttention(BaseKVCacheCompressionManager): @@ -1486,28 +1454,13 @@ def _buffers_for( assert seq_capacity % score_tile_tokens == 0 page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) - dense_groups = list(layout["storage_groups"].values()) - representatives = [group[0] for group in dense_groups] - representatives.extend( - layer for layer in layout["swa_layers"] if layer not in representatives - ) - draft_kwargs = {} + draft = None if self.draft_kv_cache_manager is not None: - draft_layout = self._draft_runtime_kv_layout() draft_tail_capacity = self._draft_protected_tail_capacity() - draft_representatives = list(draft_layout["pool_representatives"]) - draft_kwargs = dict( - draft_layer_pools=draft_layout["layer_pools"], - draft_layers=draft_layout["dense_layers"], - draft_layer_group_representative=draft_layout["layer_group_representative"], - draft_layer_pool_keys=list(draft_layout["layer_pool_keys"]), - draft_page_representatives=draft_representatives, - draft_page_table_keys=[ - draft_layout["layer_pool_keys"][layer] for layer in draft_representatives - ], - draft_num_page_table_slots=self.draft_kv_cache_manager.num_pools, - draft_page_table_token_capacity=seq_capacity + draft_tail_capacity, - draft_protected_tail_capacity=draft_tail_capacity, + draft = dict( + layout=self._draft_runtime_kv_layout(), + protected_tail_capacity=draft_tail_capacity, + page_table_token_capacity=seq_capacity + draft_tail_capacity, ) first_pool = layout["layer_pools"][layout["dense_layers"][0]] @@ -1535,32 +1488,23 @@ def _buffers_for( ) bufs = init_eviction_buffers( eviction_mode=self.eviction_mode, - layer_pools=layout["layer_pools"], - dense_groups=dense_groups, - dense_layers=layout["dense_layers"], - swa_layers=layout["swa_layers"], - swa_window=layout["swa_window"], - layer_group_representative=layout["layer_group_representative"], - layer_pool_keys=list(layout["layer_pool_keys"]), - page_representatives=representatives, - max_requests=request_capacity, - seq_len=seq_capacity, - num_q_heads=int(self._H), - num_freqs=int(self._F), - keep_count=self.budget, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=self._freq_scale_sq, - offsets=self._offsets, - omega=self.calibration["omega"], + layout=layout, + calibration=dict( + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=self._freq_scale_sq, + ), phase=self._phase, - page_table_keys=self._page_table_pool_keys(representatives, layout["global_layers"]), - num_page_table_slots=layout["manager"].num_pools, - decode_width=decode_width, - page_table_token_capacity=page_table_token_capacity, - protected_tail_capacity=tail_capacity, - **draft_kwargs, + capacities=dict( + max_requests=request_capacity, + bucket_seq_len=seq_capacity, + decode_width=decode_width, + page_table_token_capacity=page_table_token_capacity, + keep_count=self.budget, + protected_tail_capacity=tail_capacity, + ), + draft=draft, ) self._buffers = bufs self._buffers_fingerprint = fingerprint diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 850ebe93b1e8..76f9d1f99354 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -37,14 +37,6 @@ def encode_block_offsets(page_ids: torch.Tensor) -> torch.Tensor: return encoded -def compaction_family(compaction, name): - """Return one cache family dict ("dense", "swa", or "draft") or None.""" - for family in compaction["families"]: - if family["name"] == name: - return family - return None - - def _write_move_offsets(compaction, offsets, moves_per_request): cumulative = [0] for count in moves_per_request: @@ -55,25 +47,23 @@ def _write_move_offsets(compaction, offsets, moves_per_request): def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): - """Load per-request protected tails into the staged move offsets.""" + """Load per-request protected tails into the caller-owned move offsets.""" if len(tail_lengths) > compaction["request_count"]: raise ValueError("the cohort exceeds the compaction request capacity") if any(tail < 0 or tail > compaction["protected_tail_capacity"] for tail in tail_lengths): raise ValueError("a protected tail exceeds the configured capacity") _write_move_offsets( compaction, - compaction_family(compaction, "dense")["offsets"], + compaction["dense_move_offsets"], [compaction["decode_keep_count"] + int(tail) for tail in tail_lengths], ) - swa_family = compaction_family(compaction, "swa") - if swa_family is not None: + if compaction["has_swa"]: _write_move_offsets( compaction, - swa_family["offsets"], + compaction["swa_move_offsets"], [compaction["swa_window"] + int(tail) for tail in tail_lengths], ) - draft_family = compaction_family(compaction, "draft") - if draft_family is not None: + if compaction["has_draft"]: if draft_tail_lengths is None: draft_tail_lengths = [0] * len(tail_lengths) if len(draft_tail_lengths) != len(tail_lengths): @@ -85,7 +75,7 @@ def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): raise ValueError("a draft protected tail exceeds the configured capacity") _write_move_offsets( compaction, - draft_family["offsets"], + compaction["draft_move_offsets"], [compaction["decode_keep_count"] + int(tail) for tail in draft_tail_lengths], ) @@ -124,8 +114,9 @@ def make_ramp_pools( def build_compaction(**overrides): """``init_compaction_buffers`` with the suite's 2-layer defaults: - translates ``eviction_mode``, allocates the caller-owned offset rows - (capacity cumsum), and mirrors geometry inputs onto the bundle.""" + translates ``eviction_mode`` into ``per_layer_sources``, allocates the + caller-owned move-offset rows (capacity cumsum), and mirrors the + decision inputs onto the returned contract.""" from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers args = dict( @@ -148,6 +139,7 @@ def build_compaction(**overrides): keep_count = args["decode_keep_count"] tail = int(args.get("protected_tail_capacity", 0)) draft_tail = int(args.get("draft_protected_tail_capacity") or 0) + has_draft = bool(args.get("draft_layers")) device = args["layer_pools"][args["dense_layers"][0]].device swa_window = int(args["swa_window"] or 0) if args["swa_layers"] else 0 @@ -155,23 +147,62 @@ def capacity_offsets(count): return torch.arange(0, (request_count + 1) * count, count, dtype=torch.int32, device=device) args.setdefault("dense_move_offsets", capacity_offsets(keep_count + tail)) - if args["swa_layers"]: - args.setdefault("swa_move_offsets", capacity_offsets(swa_window + tail)) - if args.get("draft_layers"): + args.setdefault( + "swa_move_offsets", capacity_offsets(swa_window + tail) if args["swa_layers"] else None + ) + if has_draft: args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) num_kv_heads = int(args["layer_pools"][args["dense_layers"][0]].shape[2]) - compaction = init_compaction_buffers(union=union, per_layer=per_layer, **args) + draft = None + if has_draft: + draft = dict( + layer_pools=args["draft_layer_pools"], + layers=args["draft_layers"], + layer_group_representative=args["draft_layer_group_representative"], + layer_pool_keys=args["draft_layer_pool_keys"], + kv_block_offsets=args["draft_kv_block_offsets"], + page_table_slots=args["draft_page_table_slots"], + move_offsets=args["draft_move_offsets"], + protected_tail_capacity=draft_tail, + ) + compaction = init_compaction_buffers( + target=dict( + layer_pools=args["layer_pools"], + dense_layers=args["dense_layers"], + swa_layers=args["swa_layers"], + swa_window=args["swa_window"], + layer_group_representative=args["layer_group_representative"], + layer_pool_keys=args["layer_pool_keys"], + kv_block_offsets=args["kv_block_offsets"], + page_table_slots=args["page_table_slots"], + prompt_offsets=args["prompt_offsets"], + dense_move_offsets=args["dense_move_offsets"], + swa_move_offsets=args["swa_move_offsets"], + per_layer_sources=per_layer, + ), + capacities=dict( + request_capacity=request_count, + decode_keep_count=keep_count, + protected_tail_capacity=tail, + ), + draft=draft, + ) # Test-side mirror of the construction inputs: production reads only the - # launch fields; the standalone helpers here need the geometry back. + # contract's public keys; the standalone helpers here need the decision + # inputs and the caller-owned move-offset rows back. compaction.update( + union=union, + per_layer=per_layer, kept_token_ordinals=kept, valid_sequence_lengths=args["valid_sequence_lengths"], prompt_offsets=args["prompt_offsets"], request_count=request_count, decode_keep_count=keep_count, protected_tail_capacity=tail, - draft_protected_tail_capacity=draft_tail if args.get("draft_layers") else 0, - swa_window=swa_window, + draft_protected_tail_capacity=draft_tail if has_draft else 0, + dense_move_offsets=args["dense_move_offsets"], + swa_move_offsets=args["swa_move_offsets"], + draft_move_offsets=args["draft_move_offsets"] if has_draft else None, selection_rows=( 1 if union @@ -182,26 +213,27 @@ def capacity_offsets(count): def launch_family_pack(compaction, name): - """Standalone HAS_SETTLE=False pack for one family so the C++ moves - read initialized indices (settle-side pointers are compiled away).""" + """Standalone HAS_SETTLE=False pack for one family ("dense" or "draft") + so the C++ moves read initialized indices (settle-side pointers are + compiled away); the dense launch also packs the SWA rows when present.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( SETTLE_PACK_BLOCK, SETTLE_PACK_NUM_WARPS, _settle_ties_and_pack_compaction_sources_kernel, ) - family = compaction_family(compaction, name) kept = compaction["kept_token_ordinals"] valid = compaction["valid_sequence_lengths"] keep_count = compaction["decode_keep_count"] if name == "draft": - draft = compaction["draft_pack"] + indices = compaction["draft_move_indices"] + offsets = compaction["draft_move_offsets"] rows = 1 shape = dict( - DENSE_TOTAL=draft["dense_total"], + DENSE_TOTAL=int(indices.shape[-1]), SWA_TOTAL=0, - MOVE_CAPACITY=draft["move_capacity"], - NUM_KV_HEADS=draft["num_kv_heads"], + MOVE_CAPACITY=int(indices.shape[-1]) // compaction["request_count"], + NUM_KV_HEADS=int(indices.shape[0]), SWA_WINDOW=0, UNION=True, PER_LAYER=False, @@ -209,11 +241,21 @@ def launch_family_pack(compaction, name): ) swa_offsets, swa_indices = None, None else: + indices = compaction["dense_move_indices"] + offsets = compaction["dense_move_offsets"] rows = compaction["selection_rows"] - shape = compaction["settle_pack_shape"] - swa_family = compaction_family(compaction, "swa") - swa_offsets = swa_family["offsets"] if swa_family else None - swa_indices = swa_family["source"] if swa_family else None + shape = dict( + DENSE_TOTAL=compaction["dense_total"], + SWA_TOTAL=compaction["swa_total"], + MOVE_CAPACITY=compaction["move_capacity"], + NUM_KV_HEADS=compaction["num_kv_heads"], + SWA_WINDOW=compaction["swa_window"], + UNION=compaction["union"], + PER_LAYER=compaction["per_layer"], + HAS_SWA=compaction["has_swa"], + ) + swa_offsets = compaction["swa_move_offsets"] if compaction["has_swa"] else None + swa_indices = compaction["swa_move_indices"] _settle_ties_and_pack_compaction_sources_kernel[(compaction["request_count"], rows)]( None, None, @@ -221,8 +263,8 @@ def launch_family_pack(compaction, name): None, kept, valid, - family["offsets"], - family["source"], + offsets, + indices, swa_offsets, swa_indices, WIDTH=keep_count, @@ -236,26 +278,22 @@ def launch_family_pack(compaction, name): def run_compaction(compaction, pack=("dense", "draft")): - """Replica of the round's compact stage: packs, SWA rebase, C++ moves.""" + """Replica of the round's decision-plus-move stage in production order: + target pack, draft broadcast pack, SWA destination rebase, then + ``compact`` fires the native target and draft moves.""" + from tensorrt_llm._torch.kv_cache_compression.compaction import compact + + if "dense" in pack: + launch_family_pack(compaction, "dense") + if "draft" in pack and compaction["has_draft"]: + launch_family_pack(compaction, "draft") if compaction["swa_destination_bases"] is not None: torch.add( compaction["prompt_offsets"], compaction["swa_rebase_delta"], out=compaction["swa_destination_bases"], ) - for family in compaction["families"]: - if family["name"] in pack and family["name"] != "swa": - launch_family_pack(compaction, family["name"]) - for group in family["groups"]: - torch.ops.trtllm.sparse_kv_cache_compact_layers( - group["pools"], - group["pool_pointers"], - group["page_table"], - family["source"], - family["offsets"], - family["destination_bases"], - group["source_layer_indices"], - ) + compact(compaction, compaction["request_count"]) def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None): @@ -466,6 +504,25 @@ def torch_tri_score_oracle( return scores +def make_phase_table(offsets, omega, initial_rows): + """Build the mean-phase table dict exactly like the product's inlined + form and grow it to cover positions ``[0, initial_rows)``.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + grow_mean_phase_table, + ) + + phase = { + "offsets": offsets.contiguous(), + "omega": omega.to(dtype=torch.float32).contiguous(), + "offset_values": offsets.tolist(), + "cos": None, + "sin": None, + "rows": 0, + } + grow_mean_phase_table(phase, max(int(initial_rows), 1)) + return phase + + def make_cute_buffers( *, eviction_mode, @@ -480,41 +537,62 @@ def make_cute_buffers( omega, offsets, decode_width=None, + keep_count=1, + page_table_token_capacity=None, + protected_tail_capacity=0, + storage_groups=None, + layer_pool_keys=None, ): - """Real eviction buffers over one shared page-table slot; split - reference legs use ``eviction_mode="per_head"`` over the same pools.""" + """Real eviction buffers over the one-shared-slot default layout; split + reference legs use ``eviction_mode="per_head"`` over the same pools. + ``storage_groups``/``layer_pool_keys`` override the page-table grouping + (keys are ``(name, slot)`` tuples; the slot is ``key[1]``).""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( init_eviction_buffers, ) num_layers = len(layer_pools) + assert int(q_real.shape[1]) == num_q_heads # The constructor takes every capacity explicitly (no test-only # None-derive path); the widest window defaults keep old call sites. if decode_width is None: decode_width = seq_len - return init_eviction_buffers( - eviction_mode=eviction_mode, + if storage_groups is None: + storage_groups = {("pool", 0): list(range(num_layers))} + if layer_pool_keys is None: + layer_pool_keys = [("pool", 0)] * num_layers + layout = dict( + manager=SimpleNamespace(num_pools=max(key[1] for key in layer_pool_keys) + 1), layer_pools=layer_pools, - dense_groups=[list(range(num_layers))], dense_layers=list(range(num_layers)), - page_representatives=[0], - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=num_q_heads, - num_freqs=int(q_real.shape[-1]), - keep_count=1, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=freq_scale_sq, - offsets=offsets, - omega=omega, - decode_width=decode_width, - page_table_keys=[("pool", 0)], - num_page_table_slots=1, - page_table_token_capacity=seq_len, - layer_group_representative={layer: 0 for layer in range(num_layers)}, - layer_pool_keys=[("pool", 0)] * num_layers, + swa_layers=[], + swa_window=None, + storage_groups=storage_groups, + layer_group_representative={ + layer: layers[0] for layers in storage_groups.values() for layer in layers + }, + layer_pool_keys=layer_pool_keys, + ) + return init_eviction_buffers( + eviction_mode=eviction_mode, + layout=layout, + calibration=dict( + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + ), + phase=make_phase_table(offsets, omega, seq_len), + capacities=dict( + max_requests=max_requests, + bucket_seq_len=seq_len, + decode_width=decode_width, + page_table_token_capacity=( + seq_len if page_table_token_capacity is None else page_table_token_capacity + ), + keep_count=keep_count, + protected_tail_capacity=protected_tail_capacity, + ), ) 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 index 40bc6f0b8751..e001064bf297 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -13,7 +13,6 @@ import pytest import torch from conftest import build_compaction as _build_compaction -from conftest import compaction_family as _compaction_family from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_buffer_stubs as _make_buffer_stubs from conftest import make_fake_v2 as _make_fake_v2 @@ -160,11 +159,11 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect expected_moves.append(draft_source.to(torch.int32)) expected_offsets.append(expected_offsets[-1] + int(draft_source.numel())) - # Packed indices must match the same broadcast-plus-tail oracle. - draft_family = _compaction_family(built.compaction, "draft") + # Packed indices must match the same broadcast-plus-tail oracle: the + # test-owned draft offsets row and the contract's draft move sources. expected_row = torch.cat(expected_moves) - assert draft_family["offsets"].cpu().tolist() == expected_offsets - draft_indices = draft_family["source"] + assert built.compaction["draft_move_offsets"].cpu().tolist() == expected_offsets + draft_indices = built.compaction["draft_move_indices"] # Capacity-sized buffer; this round's moves pack at the front. capacity_total = built.request_count * ( int(built.keep.shape[1]) + max(built.draft_protected_tails) @@ -370,17 +369,20 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): # bucket follows what the cohort actually presents (power-of-two, # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. assert resources is buffers - assert prepare.call_args.kwargs["eviction_mode"] == "union" - assert prepare.call_args.kwargs["max_requests"] == 8 - assert prepare.call_args.kwargs["decode_width"] == 4 + 2 * 128 - assert prepare.call_args.kwargs["seq_len"] == 1024 - assert prepare.call_args.kwargs["page_table_token_capacity"] == 1024 + 1 - assert prepare.call_args.kwargs["draft_page_table_token_capacity"] == 1024 + 1 + kwargs = prepare.call_args.kwargs + assert kwargs["eviction_mode"] == "union" + assert kwargs["capacities"]["max_requests"] == 8 + assert kwargs["capacities"]["decode_width"] == 4 + 2 * 128 + assert kwargs["capacities"]["bucket_seq_len"] == 1024 + assert kwargs["capacities"]["page_table_token_capacity"] == 1024 + 1 + assert kwargs["draft"]["page_table_token_capacity"] == 1024 + 1 + assert kwargs["draft"]["layout"] is manager._draft_runtime_kv_layout.return_value # Migrated from the pipeline buffer-kwargs test: the budget, the # shared phase-table dict, and the pool keys thread through unchanged. - assert prepare.call_args.kwargs["keep_count"] == manager.budget - assert prepare.call_args.kwargs["phase"] is manager._phase - assert prepare.call_args.kwargs["layer_pool_keys"] == list(layout["layer_pool_keys"]) + assert kwargs["capacities"]["keep_count"] == manager.budget + assert kwargs["phase"] is manager._phase + assert kwargs["layout"] is layout + assert list(kwargs["layout"]["layer_pool_keys"]) == list(layout["layer_pool_keys"]) # A second round with unchanged pools reuses the resident buffers # (and with them the compaction launch data they carry). diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index b2fb3a633a6f..5111a3d13ebb 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -27,6 +27,7 @@ import torch from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_bare_staging as _make_bare_staging +from conftest import make_cute_buffers as _make_cute_buffers 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 @@ -814,28 +815,22 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun round_starts = round_device[:request_count].tolist() seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) - bufs = module.init_eviction_buffers( + bufs = _make_cute_buffers( eviction_mode="per_head", layer_pools=pools, - dense_groups=[[layer] for layer in layer_order], - dense_layers=layer_order, - page_representatives=layer_order, max_requests=max_requests, seq_len=seq_len, num_q_heads=num_q_heads, - num_freqs=num_freqs, - keep_count=4, q_real=q_real, q_imag=q_imag, mlr_coef=mlr, freq_scale_sq=freq, - offsets=offsets, omega=omega, + offsets=offsets, decode_width=seq_len - prompt_len, - page_table_keys=[("pool", layer) for layer in layer_order], - num_page_table_slots=num_layers, - page_table_token_capacity=seq_len, - layer_group_representative={layer: layer for layer in layer_order}, + keep_count=4, + # One storage group and page-table slot per layer (distinct pools). + storage_groups={("pool", layer): [layer] for layer in layer_order}, layer_pool_keys=[("pool", layer) for layer in layer_order], ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) @@ -849,11 +844,12 @@ def stage_round(): ) score_sentinel = -12345.0 - # Empty prebound launch args: the compact stage is a no-op. - bufs.compact_launch_args = () stage_round() bufs.score_output.fill_(score_sentinel) - module.run_eviction_round(bufs, normalize_scores=False) + # The compact stage is stubbed to a no-op: this test owns the score + # buffers only, never a staged move decision. + with mock.patch.object(module, "compact"): + module.run_eviction_round(bufs, normalize_scores=False) fixed = bufs.score_output.clone() assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] @@ -892,7 +888,8 @@ def stage_round(): stage_round() bufs.score_output.fill_(score_sentinel) bufs.valid_widths.fill_(-1) - module.run_eviction_round(bufs, normalize_scores=False) + with mock.patch.object(module, "compact"): + module.run_eviction_round(bufs, normalize_scores=False) second_launch = bufs.score_output.clone() assert torch.equal(bufs.valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index f5d615a709a0..9fbe1391aa96 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -8,12 +8,15 @@ import torch from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_cute_buffers as _make_cute_buffers from conftest import make_ramp_pools as _make_ramp_pools from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + SETTLE_PACK_BLOCK, + SETTLE_PACK_NUM_WARPS, prepare_per_head_scores, ) @@ -103,27 +106,41 @@ def _make_selection_buffers( bufs.keep_rows = bufs.keep.view(-1, keep_count) bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) # The settle launch always packs now; zero per-request move counts mask - # every pack store off, so these buffers stay selection-only. + # every pack store off, so these buffers stay selection-only. The launch + # args mirror the product's ``bufs.settle_args``/``bufs.settle_kwargs`` + # order and keys exactly (swa pointers are None without SWA layers). zero_offsets = torch.zeros(max_requests + 1, dtype=torch.int32, device=device) zero_lengths = torch.zeros(max_requests, dtype=torch.int32, device=device) zero_indices = torch.zeros(1, dtype=torch.int32, device=device) - bufs.settle_pack_tensors = ( + bufs.settle_args = ( + bufs.selection_scores_rows, + bufs.selection_row_lengths, + bufs.row_prompt_offsets, + bufs.provisional_rows, + bufs.keep_rows, zero_lengths, zero_offsets, zero_indices, - zero_offsets, - zero_indices, + None, + None, ) - bufs.settle_pack_shape = dict( + bufs.settle_kwargs = dict( + WIDTH=width, + KEEP_COUNT=keep_count, + SELECTION_ROWS=bufs.selection_rows_per_request, DENSE_TOTAL=0, SWA_TOTAL=0, MOVE_CAPACITY=keep_count, NUM_KV_HEADS=1, SWA_WINDOW=0, - UNION=False, - PER_LAYER=False, + UNION=eviction_mode == "union", + PER_LAYER=eviction_mode == "per_layer_perhead", HAS_SWA=False, + HAS_SETTLE=True, + BLOCK=SETTLE_PACK_BLOCK, + num_warps=SETTLE_PACK_NUM_WARPS, ) + bufs.draft_pack_launch = None return bufs @@ -449,8 +466,6 @@ 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") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - attach_compaction_bundle, - init_eviction_buffers, run_eviction_round, ) @@ -468,9 +483,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): tokens_per_block = 32 head_dim = 64 num_freqs = head_dim // 2 - dense_layers = [0, 1, 2] dense_groups = [[0, 2], [1]] - layer_group_representative = {0: 0, 1: 1, 2: 0} page_tables = ( torch.tensor([[1, 0]], dtype=torch.int32, device=device), torch.tensor([[0, 1]], dtype=torch.int32, device=device), @@ -498,55 +511,38 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): mlr_coef[:, :, 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - bufs = init_eviction_buffers( + bufs = _make_cute_buffers( eviction_mode="per_layer_perhead", layer_pools=pools, - dense_groups=dense_groups, - dense_layers=dense_layers, - page_representatives=[0, 1], max_requests=1, seq_len=bucket_capacity, num_q_heads=num_q_heads, - num_freqs=num_freqs, - keep_count=keep_count, q_real=q_real, q_imag=q_imag, mlr_coef=mlr_coef, freq_scale_sq=freq_scale_sq, - offsets=torch.zeros(1, dtype=torch.float32, device=device), omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), + offsets=torch.zeros(1, dtype=torch.float32, device=device), decode_width=bucket_capacity, - page_table_keys=[("pool", 0), ("pool", 1)], - num_page_table_slots=2, - page_table_token_capacity=bucket_capacity, - layer_group_representative=layer_group_representative, + keep_count=keep_count, + storage_groups={ + ("pool", 0): dense_groups[0], + ("pool", 1): dense_groups[1], + }, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], ) + assert bufs.compaction["has_swa"] is False bufs.block_offsets_device.zero_() bufs.block_offsets_device[..., :2].copy_(_encode_block_offsets(torch.stack(page_tables))) bufs.round_starts_device.fill_(0) bufs.valid_seq_lens_device.fill_(seq_len) bufs.token_starts_device.fill_(0) - - # Replace the constructor-built bundle wholesale; the settle launch - # packs this bundle's construction-time move offsets. - compaction = _build_compaction( - eviction_mode="per_layer_perhead", - layer_pools=pools, - dense_layers=dense_layers, - layer_group_representative=layer_group_representative, - layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], - kept_token_ordinals=bufs.keep[:1], - valid_sequence_lengths=bufs.valid_seq_lens_device[:1], - kv_block_offsets=bufs.block_offsets_device, - page_table_slots=bufs.representative_slots, - request_count=1, - prompt_offsets=torch.zeros(1, dtype=torch.int32, device=device), - decode_keep_count=keep_count, - protected_tail_capacity=0, + # Stage this round's move offsets into the buffers' OWN metadata row: + # the fused settle launch and the native moves consume the buffers' + # own contract (keep_count moves per request, no protected tail). + bufs.request_metadata_device[3, :2].copy_( + torch.tensor([0, keep_count], dtype=torch.int32, device=device) ) - _set_protected_tails(compaction, [0]) - attach_compaction_bundle(bufs, compaction) run_eviction_round(bufs, normalize_scores=False) assert torch.equal(bufs.keep, expected_keep) torch.cuda.synchronize(device) @@ -572,7 +568,6 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - init_eviction_buffers, mark_page_tables_consumed, run_eviction_round, stage_eviction_cohort, @@ -678,28 +673,20 @@ def expected_keep() -> torch.Tensor: mlr_coef[..., 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - bufs = init_eviction_buffers( + bufs = _make_cute_buffers( eviction_mode="union", layer_pools=[pool], - dense_groups=[[0]], - dense_layers=[0], - layer_group_representative={0: 0}, - layer_pool_keys=[("pool", 0)], - page_representatives=[0], max_requests=1, seq_len=seq_len, num_q_heads=num_q_heads, - num_freqs=num_freqs, - keep_count=keep_count, q_real=q_real, q_imag=q_imag, mlr_coef=mlr_coef, freq_scale_sq=freq_scale_sq, - offsets=torch.zeros(1, dtype=torch.float32, device=device), omega=torch.zeros(num_freqs, dtype=torch.float32, device=device), - page_table_keys=[("pool", 0)], - num_page_table_slots=1, + offsets=torch.zeros(1, dtype=torch.float32, device=device), decode_width=seq_len - prompt_len, + keep_count=keep_count, page_table_token_capacity=seq_len + protected_tail, protected_tail_capacity=protected_tail, ) From b423dbc27a2e1aaf73359d5300934f7793e1d001 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 01:08:43 -0700 Subject: [PATCH 114/178] [None][perf] Size the score-only raw-K smem by its single-buffer stage count (knife 18 item 1) Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 66 +++++++++++-------- 1 file changed, 37 insertions(+), 29 deletions(-) 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 index 5422ca8db00e..251598a3207f 100644 --- 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 @@ -123,6 +123,9 @@ def __init__( # 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 self.physical_threads = THREADS @@ -279,7 +282,7 @@ def __call__( raw_bf16_tiled_mma, (CTA_M, N, self.num_freqs), cutlass.BFloat16, - 2 * RAW_PAGE_BUFFERS, + 2 * self.raw_page_buffers, ) raw_tma_smem_layout = cute.make_composed_layout( raw_bf16_direct_a_smem_layout.inner, @@ -343,8 +346,8 @@ def __call__( 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) - # Two double-buffered raw-K stages (real+imag halves per page buffer). - raw_k_elements = CTA_M * 2 * self.num_freqs * RAW_PAGE_BUFFERS + # 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) @@ -504,8 +507,10 @@ def kernel( ) cpasync_raw_k_real = cpasync_raw_k_0[(None, None, None, 0)] cpasync_raw_k_imag = cpasync_raw_k_0[(None, None, None, 1)] - 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)] + 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, @@ -528,14 +533,15 @@ def kernel( cpasync_raw_k_imag.iterator + fragment_offset, raw_tma_smem_layout.outer, ) - 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, - ) + 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, @@ -550,24 +556,26 @@ def kernel( cute.group_modes(fragment_imag, 0, 2), cute.group_modes(raw_tma_source_tiles, 0, 2), ) - 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), - ) + 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) - raw_tma_shared_partition_real_next.append(partition_real_next) - raw_tma_shared_partition_imag_next.append(partition_imag_next) + 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, From 9cb7af78216bc60195bbf942258ffd4faac74851 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 01:28:27 -0700 Subject: [PATCH 115/178] [None][refactor] Fission the fused settle+pack kernel: decision settle stays, mover pack moves to compaction (knife 25b) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 187 ++++++++++++++-- .../triattention/triattention.py | 76 ++----- .../triattention/triattention_kernels.py | 203 ++++++------------ .../_torch/kv_cache_compression/conftest.py | 102 ++------- .../test_triattention_fused_settle_pack.py | 202 +++++++++-------- .../test_triattention_selection_compaction.py | 44 ++-- 6 files changed, 395 insertions(+), 419 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 1fa61a6c29f3..d29027f34bee 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,18 +15,103 @@ """Batched physical KV-cache compaction: an algorithm-neutral mover. -``init_compaction_buffers`` agrees on the decision buffers (move-source -indices; offsets ride the caller's staged rows) once per geometry and retains +``init_compaction_buffers`` agrees on the decision rows (kept ordinals; +move offsets ride the caller's staged rows) once per geometry and retains one launch contract. The caller materializes its keep decision into those -buffers each round, then ``compact`` fires the native target and draft -launches. This module knows cache-family geometry and the decision format -only; the contract's launch tuples are private to it. +rows each round, then ``compact`` packs them into per-family move sources +and fires the native target and draft launches. This module knows +cache-family geometry and the decision format only; the contract's launch +tuples are private to it. """ from collections import OrderedDict from typing import Dict, List, Optional, Tuple import torch +import triton +import triton.language as tl + +# Pack launch shape, shared by every family's pack launch. +PACK_BLOCK = 256 +PACK_NUM_WARPS = 4 + + +@triton.jit +def _pack_move_sources_kernel( + kept_ordinal_rows, + valid_seq_lens, + dense_offsets, + dense_indices, + swa_offsets, + swa_indices, + KEEP_COUNT: tl.constexpr, + DECISION_ROWS: tl.constexpr, + DENSE_TOTAL: tl.constexpr, + SWA_TOTAL: tl.constexpr, + MOVE_CAPACITY: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + SWA_WINDOW: tl.constexpr, + BROADCAST: tl.constexpr, + PER_LAYER: tl.constexpr, + HAS_SWA: tl.constexpr, + BLOCK: tl.constexpr, +): + """Pack one decision row into one family's move sources (increasing + kept ordinals; C++ in-place copy contract): dense rows forward the row + content verbatim for the first KEEP_COUNT moves, then append the + protected tail; SWA rows write latest-window ordinals once per KV head.""" + request = tl.program_id(0) + decision_row = tl.program_id(1) + row = request * DECISION_ROWS + decision_row + kept_row = kept_ordinal_rows + row * KEEP_COUNT + dense_begin = tl.load(dense_offsets + request) + dense_end = tl.load(dense_offsets + request + 1) + dense_count = dense_end - dense_begin + valid_len = tl.load(valid_seq_lens + request) + if HAS_SWA: + swa_begin = tl.load(swa_offsets + request) + swa_end = tl.load(swa_offsets + request + 1) + swa_count = swa_end - swa_begin + for move_start in tl.static_range(0, MOVE_CAPACITY, BLOCK): + move = move_start + tl.arange(0, BLOCK) + kept = tl.load( + kept_row + move, + mask=move < KEEP_COUNT, + other=0, + ) + dense_source = tl.where(move < KEEP_COUNT, kept, valid_len + move - KEEP_COUNT) + if BROADCAST: + # The one decision row per request feeds every KV head's packed row. + for head in tl.static_range(0, NUM_KV_HEADS): + tl.store( + dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, + dense_source, + mask=move < dense_count, + ) + else: + dense_output = decision_row.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move + tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) + if HAS_SWA: + swa_source = valid_len - SWA_WINDOW + move + if BROADCAST: + for head in tl.static_range(0, NUM_KV_HEADS): + tl.store( + swa_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, + swa_source, + mask=move < swa_count, + ) + else: + swa_mask = move < swa_count + if PER_LAYER: + # SWA has one shared row per head; the first layer's decision rows write it. + swa_mask = swa_mask & (decision_row < NUM_KV_HEADS) + head = decision_row % NUM_KV_HEADS + swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move + tl.store( + swa_indices + swa_output, + swa_source, + mask=swa_mask, + ) def _make_move_indices( @@ -111,16 +196,22 @@ def init_compaction_buffers( capacities: Dict[str, int], draft: Optional[Dict[str, object]] = None, ) -> Dict[str, object]: - """Agree on the decision buffers and retain one launch contract per geometry. + """Agree on the decision rows and retain one launch contract per geometry. Move sources must be increasing kept ordinals with destination_bases[request] + move <= source[move] (C++ in-place copy contract). ``target`` carries the resolved dense/SWA grouping inputs from - the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows); - ``draft`` is one all-or-none resolved branch; ``capacities`` the - request/keep/tail capacity numbers. The returned contract exposes the - agreed move-source buffers and geometry constants; its launch tuples are - private to :func:`compact`. + the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows) + plus the decision inputs :func:`compact` packs each round: + ``kept_ordinal_rows`` (``request_capacity * decision_rows`` rows of + ``decode_keep_count`` int32 kept ordinals, forwarded verbatim), the + per-request ``decision_rows`` count (1 = one shared row broadcast over + every KV head), and the staged per-request ``valid_seq_lens`` the + protected tail rides after. ``draft`` is one all-or-none resolved branch + (its dense-only moves broadcast the one shared decision row over the + draft's own KV heads); ``capacities`` the request/keep/tail capacity + numbers. The returned contract exposes the agreed move-source buffers and + geometry constants; its launch tuples are private to :func:`compact`. """ layer_pools = target["layer_pools"] dense_layers = tuple(int(layer) for layer in target["dense_layers"]) @@ -134,6 +225,9 @@ def init_compaction_buffers( swa_move_offsets = target["swa_move_offsets"] swa_window = target["swa_window"] per_layer_sources = bool(target["per_layer_sources"]) + kept_ordinal_rows = target["kept_ordinal_rows"] + decision_rows = int(target["decision_rows"]) + valid_seq_lens = target["valid_seq_lens"] device = layer_pools[dense_layers[0]].device request_count = int(capacities["request_capacity"]) @@ -213,9 +307,41 @@ def init_compaction_buffers( ) ) + target_pack_launch = ( + decision_rows, + ( + kept_ordinal_rows, + valid_seq_lens, + dense_move_offsets, + dense_move_indices, + swa_move_offsets, + swa_move_indices, + ), + dict( + KEEP_COUNT=decode_keep_count, + DECISION_ROWS=decision_rows, + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=swa_total, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + SWA_WINDOW=swa_window, + BROADCAST=decision_rows == 1, + PER_LAYER=per_layer_sources, + HAS_SWA=has_swa, + BLOCK=PACK_BLOCK, + num_warps=PACK_NUM_WARPS, + ), + ) + draft_launches: Tuple[tuple, ...] = () draft_move_indices = None + draft_pack_launch = None if draft is not None: + if decision_rows != 1: + raise ValueError( + "draft packing broadcasts one shared decision row per request; " + f"got {decision_rows} decision rows" + ) draft_layer_pools = draft["layer_pools"] draft_layers = tuple(int(layer) for layer in draft["layers"]) draft_tail = int(draft["protected_tail_capacity"]) @@ -245,6 +371,31 @@ def init_compaction_buffers( draft["move_offsets"], prompt_offsets, ) + draft_pack_launch = ( + 1, + ( + kept_ordinal_rows, + valid_seq_lens, + draft["move_offsets"], + draft_move_indices, + None, + None, + ), + dict( + KEEP_COUNT=decode_keep_count, + DECISION_ROWS=1, + DENSE_TOTAL=int(draft_move_indices.shape[-1]), + SWA_TOTAL=0, + MOVE_CAPACITY=int(draft_move_indices.shape[-1]) // request_count, + NUM_KV_HEADS=draft_num_kv_heads, + SWA_WINDOW=0, + BROADCAST=True, + PER_LAYER=False, + HAS_SWA=False, + BLOCK=PACK_BLOCK, + num_warps=PACK_NUM_WARPS, + ), + ) return dict( # Agreed decision buffers and geometry constants (the public interface). @@ -265,18 +416,26 @@ def init_compaction_buffers( # Private launch tuples: only compact() interprets these. target_launches=tuple(target_launches), draft_launches=draft_launches, + target_pack_launch=target_pack_launch, + draft_pack_launch=draft_pack_launch, has_draft=draft is not None, ) def compact(compaction: Dict[str, object], request_count: int) -> None: - """Fire the native target compacts, then the draft compacts, and record completion. + """Pack each family's move sources, fire its native compacts, and record completion. - Pure mover: the caller has already materialized its keep decision into the - agreed move-source buffers for the active ``request_count`` cohort. + Pure mover: the caller has already materialized its kept ordinals into the + agreed decision rows for the active ``request_count`` cohort. """ + rows, pack_args, pack_kwargs = compaction["target_pack_launch"] + _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) for launch in compaction["target_launches"]: torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) + draft_pack_launch = compaction["draft_pack_launch"] + if draft_pack_launch is not None: + rows, pack_args, pack_kwargs = draft_pack_launch + _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) for launch in compaction["draft_launches"]: torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) compaction["consume_done"].record(torch.cuda.current_stream()) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index f1e527e0ddcb..28f00ec1c227 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -41,10 +41,10 @@ from ..compaction import compact, init_compaction_buffers from .triattention_kernels import ( - SETTLE_PACK_BLOCK, - SETTLE_PACK_NUM_WARPS, + SETTLE_BLOCK, + SETTLE_NUM_WARPS, _gather_mean_phase_kernel, - _settle_ties_and_pack_compaction_sources_kernel, + _settle_ties_kernel, grow_mean_phase_table, prepare_per_head_scores, ) @@ -406,6 +406,10 @@ def init_eviction_buffers( dense_move_offsets=dense_move_offsets_row, swa_move_offsets=swa_move_offsets_row, per_layer_sources=per_layer, + # The decision rows the contract packs into move sources. + kept_ordinal_rows=bufs.keep_rows, + decision_rows=bufs.selection_rows_per_request, + valid_seq_lens=bufs.valid_seq_lens_device, ), capacities=dict( request_capacity=max_requests, @@ -417,68 +421,21 @@ def init_eviction_buffers( bufs.compaction = contract bufs.swa_destination_bases = contract["swa_destination_bases"] bufs.swa_rebase_delta = contract["swa_rebase_delta"] - # The decision side: settle+pack launch args against the agreed buffers. + # The decision side: the settle launch materializes the kept-ordinal rows. bufs.settle_args = ( bufs.selection_scores_rows, bufs.selection_row_lengths, bufs.row_prompt_offsets, bufs.provisional_rows, bufs.keep_rows, - bufs.valid_seq_lens_device, - dense_move_offsets_row, - contract["dense_move_indices"], - swa_move_offsets_row if contract["has_swa"] else None, - contract["swa_move_indices"], ) bufs.settle_kwargs = dict( WIDTH=decode_width, KEEP_COUNT=keep_count, SELECTION_ROWS=bufs.selection_rows_per_request, - DENSE_TOTAL=contract["dense_total"], - SWA_TOTAL=contract["swa_total"], - MOVE_CAPACITY=contract["move_capacity"], - NUM_KV_HEADS=contract["num_kv_heads"], - SWA_WINDOW=contract["swa_window"], - UNION=union, - PER_LAYER=per_layer, - HAS_SWA=contract["has_swa"], - HAS_SETTLE=True, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, + BLOCK=SETTLE_BLOCK, + num_warps=SETTLE_NUM_WARPS, ) - bufs.draft_pack_launch = None - if draft is not None: - draft_move_indices = contract["draft_move_indices"] - # Pack-only decision broadcast (HAS_SETTLE=False): settle pointers are None. - broadcast_pack_args = ( - None, - None, - None, - None, - bufs.keep, - bufs.valid_seq_lens_device, - draft_move_offsets_row, - draft_move_indices, - None, - None, - ) - broadcast_pack_kwargs = dict( - WIDTH=keep_count, - KEEP_COUNT=keep_count, - SELECTION_ROWS=1, - DENSE_TOTAL=int(draft_move_indices.shape[-1]), - SWA_TOTAL=0, - MOVE_CAPACITY=int(draft_move_indices.shape[-1]) // max_requests, - NUM_KV_HEADS=int(draft_move_indices.shape[0]), - SWA_WINDOW=0, - UNION=True, - PER_LAYER=False, - HAS_SWA=False, - HAS_SETTLE=False, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, - ) - bufs.draft_pack_launch = (broadcast_pack_args, broadcast_pack_kwargs) # ---- round-ordering events ---------------------------------------------- bufs.copy_done = torch.cuda.Event() @@ -605,8 +562,8 @@ def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cud def settle_top_tokens(bufs: SimpleNamespace) -> None: - """Pick the top-k, settle ties, and materialize the move-source decision - (target settle+pack, then the draft broadcast pack).""" + """Pick the top-k and settle ties into the kept-ordinal decision rows + (the compaction contract packs them into move sources).""" # The trailing 1 is next_n: decode scores one query token per request. torch.ops.trtllm.cute_dsl_indexer_topk_decode( bufs.selection_scores_rows, @@ -615,14 +572,7 @@ def settle_top_tokens(bufs: SimpleNamespace) -> None: bufs.keep_count, 1, ) - _settle_ties_and_pack_compaction_sources_kernel[bufs.settle_grid]( - *bufs.settle_args, **bufs.settle_kwargs - ) - if bufs.draft_pack_launch is not None: - broadcast_pack_args, broadcast_pack_kwargs = bufs.draft_pack_launch - _settle_ties_and_pack_compaction_sources_kernel[(bufs.max_requests, 1)]( - *broadcast_pack_args, **broadcast_pack_kwargs - ) + _settle_ties_kernel[bufs.settle_grid](*bufs.settle_args, **bufs.settle_kwargs) def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index bbed90a0d3e9..a4f8fcbc5e40 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -213,7 +213,7 @@ def prepare_per_head_scores( num_kv_heads = int(num_kv_heads) _, num_layers, num_query_heads, width = scores.shape selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - # 256 lanes / 4 warps, matching the settle/pack shape. + # 256 lanes / 4 warps, matching the settle shape. stats_block = 256 rows = num_layers * num_query_heads if normalize_scores: @@ -251,158 +251,91 @@ def prepare_per_head_scores( ) -# ---- Compaction: pack the kept ordinals into per-request move indices ---- +# ---- Selection finalize: settle threshold ties into the kept-ordinal rows ---- -# Settle/pack launch shape, shared by every launch site. -SETTLE_PACK_BLOCK = 256 -SETTLE_PACK_NUM_WARPS = 4 +# Settle launch shape, shared by every launch site. +SETTLE_BLOCK = 256 +SETTLE_NUM_WARPS = 4 @triton.jit -def _settle_ties_and_pack_compaction_sources_kernel( +def _settle_ties_kernel( scores, seq_lens, prompt_offsets, provisional_indices, output_indices, - valid_seq_lens, - dense_offsets, - dense_indices, - swa_offsets, - swa_indices, WIDTH: tl.constexpr, KEEP_COUNT: tl.constexpr, SELECTION_ROWS: tl.constexpr, - DENSE_TOTAL: tl.constexpr, - SWA_TOTAL: tl.constexpr, - MOVE_CAPACITY: tl.constexpr, - NUM_KV_HEADS: tl.constexpr, - SWA_WINDOW: tl.constexpr, - UNION: tl.constexpr, - PER_LAYER: tl.constexpr, - HAS_SWA: tl.constexpr, - HAS_SETTLE: tl.constexpr, BLOCK: tl.constexpr, ): - """Settle one selection row's ties and pack its move sources (increasing - kept ordinals; C++ in-place copy contract).""" + """Settle one selection row's ties into its kept-ordinal output row + (threshold recovery with sentinel-skip, strictly-greater count, + lowest-index tie quota, ascending prompt-rebased emission; entries past + the emitted count keep their previous value).""" request = tl.program_id(0) selection_domain = tl.program_id(1) row = request * SELECTION_ROWS + selection_domain row_output = output_indices + row * KEEP_COUNT - if HAS_SETTLE: - row_scores = scores + row * WIDTH - row_selected = provisional_indices + row * KEEP_COUNT - # Rebases the decode-relative ordinals to absolute positions. - prompt_len = tl.load(prompt_offsets + row) - - 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)) - - seq_len = tl.load(seq_lens + 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 < seq_len) - 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 < seq_len) - 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_len, - mask=selected, - ) - output_count += tl.sum(selected_i32) - ties_seen += tl.sum(tied_i32) - - if HAS_SETTLE: - # Make the settled row's scattered stores visible before the pack half. - tl.debug_barrier() - dense_begin = tl.load(dense_offsets + request) - dense_end = tl.load(dense_offsets + request + 1) - dense_count = dense_end - dense_begin - valid_len = tl.load(valid_seq_lens + request) - if HAS_SWA: - swa_begin = tl.load(swa_offsets + request) - swa_end = tl.load(swa_offsets + request + 1) - swa_count = swa_end - swa_begin - for move_start in tl.static_range(0, MOVE_CAPACITY, BLOCK): - move = move_start + tl.arange(0, BLOCK) - selected = tl.load( - row_output + move, - mask=move < KEEP_COUNT, + row_scores = scores + row * WIDTH + row_selected = provisional_indices + row * KEEP_COUNT + # Rebases the decode-relative ordinals to absolute positions. + prompt_len = tl.load(prompt_offsets + row) + + 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, ) - dense_source = tl.where(move < KEEP_COUNT, selected, valid_len + move - KEEP_COUNT) - if UNION: - # The one union row per request feeds every KV head's packed row. - for head in tl.static_range(0, NUM_KV_HEADS): - tl.store( - dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, - dense_source, - mask=move < dense_count, - ) - else: - dense_output = ( - selection_domain.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move - ) - tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) - if HAS_SWA: - swa_source = valid_len - SWA_WINDOW + move - if UNION: - for head in tl.static_range(0, NUM_KV_HEADS): - tl.store( - swa_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, - swa_source, - mask=move < swa_count, - ) - else: - swa_mask = move < swa_count - if PER_LAYER: - # SWA has one shared row per head; first layer's domains write it. - swa_mask = swa_mask & (selection_domain < NUM_KV_HEADS) - head = selection_domain % NUM_KV_HEADS - swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move - tl.store( - swa_indices + swa_output, - swa_source, - mask=swa_mask, - ) + # 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)) + + seq_len = tl.load(seq_lens + 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 < seq_len) + 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 < seq_len) + 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_len, + mask=selected, + ) + output_count += tl.sum(selected_i32) + ties_seen += tl.sum(tied_i32) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 76f9d1f99354..266ae7bac1ba 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -114,9 +114,9 @@ def make_ramp_pools( def build_compaction(**overrides): """``init_compaction_buffers`` with the suite's 2-layer defaults: - translates ``eviction_mode`` into ``per_layer_sources``, allocates the - caller-owned move-offset rows (capacity cumsum), and mirrors the - decision inputs onto the returned contract.""" + translates ``eviction_mode`` into ``per_layer_sources``/``decision_rows``, + allocates the caller-owned move-offset rows (capacity cumsum), and hands + the test's pre-settled ``kept_token_ordinals`` in as the decision rows.""" from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers args = dict( @@ -153,6 +153,9 @@ def capacity_offsets(count): if has_draft: args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) num_kv_heads = int(args["layer_pools"][args["dense_layers"][0]].shape[2]) + selection_rows = ( + 1 if union else (len(args["dense_layers"]) * num_kv_heads if per_layer else num_kv_heads) + ) draft = None if has_draft: draft = dict( @@ -179,6 +182,9 @@ def capacity_offsets(count): dense_move_offsets=args["dense_move_offsets"], swa_move_offsets=args["swa_move_offsets"], per_layer_sources=per_layer, + kept_ordinal_rows=kept.reshape(-1, keep_count), + decision_rows=selection_rows, + valid_seq_lens=args["valid_sequence_lengths"], ), capacities=dict( request_capacity=request_count, @@ -188,13 +194,9 @@ def capacity_offsets(count): draft=draft, ) # Test-side mirror of the construction inputs: production reads only the - # contract's public keys; the standalone helpers here need the decision - # inputs and the caller-owned move-offset rows back. + # contract's public keys; the standalone helpers here need the + # caller-owned move-offset rows back. compaction.update( - union=union, - per_layer=per_layer, - kept_token_ordinals=kept, - valid_sequence_lengths=args["valid_sequence_lengths"], prompt_offsets=args["prompt_offsets"], request_count=request_count, decode_keep_count=keep_count, @@ -203,90 +205,16 @@ def capacity_offsets(count): dense_move_offsets=args["dense_move_offsets"], swa_move_offsets=args["swa_move_offsets"], draft_move_offsets=args["draft_move_offsets"] if has_draft else None, - selection_rows=( - 1 - if union - else (len(args["dense_layers"]) * num_kv_heads if per_layer else num_kv_heads) - ), ) return compaction -def launch_family_pack(compaction, name): - """Standalone HAS_SETTLE=False pack for one family ("dense" or "draft") - so the C++ moves read initialized indices (settle-side pointers are - compiled away); the dense launch also packs the SWA rows when present.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_PACK_BLOCK, - SETTLE_PACK_NUM_WARPS, - _settle_ties_and_pack_compaction_sources_kernel, - ) - - kept = compaction["kept_token_ordinals"] - valid = compaction["valid_sequence_lengths"] - keep_count = compaction["decode_keep_count"] - if name == "draft": - indices = compaction["draft_move_indices"] - offsets = compaction["draft_move_offsets"] - rows = 1 - shape = dict( - DENSE_TOTAL=int(indices.shape[-1]), - SWA_TOTAL=0, - MOVE_CAPACITY=int(indices.shape[-1]) // compaction["request_count"], - NUM_KV_HEADS=int(indices.shape[0]), - SWA_WINDOW=0, - UNION=True, - PER_LAYER=False, - HAS_SWA=False, - ) - swa_offsets, swa_indices = None, None - else: - indices = compaction["dense_move_indices"] - offsets = compaction["dense_move_offsets"] - rows = compaction["selection_rows"] - shape = dict( - DENSE_TOTAL=compaction["dense_total"], - SWA_TOTAL=compaction["swa_total"], - MOVE_CAPACITY=compaction["move_capacity"], - NUM_KV_HEADS=compaction["num_kv_heads"], - SWA_WINDOW=compaction["swa_window"], - UNION=compaction["union"], - PER_LAYER=compaction["per_layer"], - HAS_SWA=compaction["has_swa"], - ) - swa_offsets = compaction["swa_move_offsets"] if compaction["has_swa"] else None - swa_indices = compaction["swa_move_indices"] - _settle_ties_and_pack_compaction_sources_kernel[(compaction["request_count"], rows)]( - None, - None, - None, - None, - kept, - valid, - offsets, - indices, - swa_offsets, - swa_indices, - WIDTH=keep_count, - KEEP_COUNT=keep_count, - SELECTION_ROWS=rows, - **shape, - HAS_SETTLE=False, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, - ) - - -def run_compaction(compaction, pack=("dense", "draft")): - """Replica of the round's decision-plus-move stage in production order: - target pack, draft broadcast pack, SWA destination rebase, then - ``compact`` fires the native target and draft moves.""" +def run_compaction(compaction): + """Replica of the round's move stage in production order: SWA + destination rebase, then ``compact`` packs the decision rows into move + sources and fires the native target and draft moves.""" from tensorrt_llm._torch.kv_cache_compression.compaction import compact - if "dense" in pack: - launch_family_pack(compaction, "dense") - if "draft" in pack and compaction["has_draft"]: - launch_family_pack(compaction, "draft") if compaction["swa_destination_bases"] is not None: torch.add( compaction["prompt_offsets"], 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 index cb1acd5f402d..585dd12076d3 100644 --- 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 @@ -1,24 +1,37 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Fused settle-and-pack vs a pure-torch integer oracle (threshold -recovery with sentinel-skip, strictly-greater count, lowest-index tie -quota, ascending prompt-rebased emission, dense/SWA packing). All outputs -are integers, so comparisons are ``torch.equal`` including stale regions.""" +"""The settle and pack kernels vs pure-torch integer oracles: +``_settle_ties_kernel`` (threshold recovery with sentinel-skip, +strictly-greater count, lowest-index tie quota, ascending prompt-rebased +emission) and ``_pack_move_sources_kernel`` (dense/SWA packing, fed +pre-settled rows). All outputs are integers, so comparisons are +``torch.equal`` including stale regions.""" import pytest import torch +from tensorrt_llm._torch.kv_cache_compression.compaction import PACK_BLOCK as _PACK_BLOCK +from tensorrt_llm._torch.kv_cache_compression.compaction import PACK_NUM_WARPS as _PACK_NUM_WARPS +from tensorrt_llm._torch.kv_cache_compression.compaction import _pack_move_sources_kernel from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_PACK_BLOCK as _BLOCK, + SETTLE_BLOCK as _SETTLE_BLOCK, ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_PACK_NUM_WARPS as _NUM_WARPS, + SETTLE_NUM_WARPS as _SETTLE_NUM_WARPS, ) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _settle_ties_and_pack_compaction_sources_kernel, + _settle_ties_kernel, ) +# Both kernels share the settle geometry parameter grid. +_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 in place: threshold = min over non-sentinel provisional @@ -99,18 +112,77 @@ def _staged_offsets(counts, device): return torch.tensor(offsets, dtype=torch.int32, device=device) +def _make_settle_inputs(rows_total, width, keep_count, seed, device): + """One seeded settle problem: heavily tied scores, ragged rows, a + top-k stand-in, and per-row prompt rebases.""" + 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, + ) + row_prompt_offsets = torch.tensor( + [3 * (row % 3) for row in range(rows_total)], 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, row_prompt_offsets, provisional + + +@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "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, row_prompt_offsets, provisional = _make_settle_inputs( + rows_total, width, keep_count, seed, device + ) + # 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_kernel[(request_count, selection_rows)]( + scores, + row_lengths, + row_prompt_offsets, + provisional, + output_actual, + WIDTH=width, + KEEP_COUNT=keep_count, + SELECTION_ROWS=selection_rows, + BLOCK=_SETTLE_BLOCK, + num_warps=_SETTLE_NUM_WARPS, + ) + torch.cuda.synchronize(device) + + assert torch.equal(output_actual, output_reference), f"kept ordinals differ (seed {seed})" + + @pytest.mark.parametrize("has_swa", [False, True]) @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) -@pytest.mark.parametrize( - "width,keep_count", - [ - # 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 test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, keep_count): +@pytest.mark.parametrize("width,keep_count", _WIDTH_KEEP_CASES) +def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, width, keep_count): device = torch.device("cuda", torch.cuda.current_device()) request_count, num_layers, num_kv_heads = 3, 2, 2 union = eviction_mode == "union" @@ -135,48 +207,29 @@ def test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, k swa_offsets = _staged_offsets(swa_counts, device) for seed in range(5): - 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, + scores, row_lengths, row_prompt_offsets, provisional = _make_settle_inputs( + rows_total, width, keep_count, seed, device ) - row_prompt_offsets = torch.tensor( - [3 * (row % 3) for row in range(rows_total)], dtype=torch.int32, device=device + # Pre-settled decision rows straight from the settle oracle: short + # rows keep stale garbage past their emitted count, which the pack + # must forward verbatim. + settled = torch.randint( + -(2**30), 2**30, (rows_total, keep_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() + _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, settled, keep_count) # 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 - ) dense_stale = torch.randint( -(2**30), 2**30, (packed_rows, dense_total), dtype=torch.int32, device=device ) swa_stale = torch.randint( -(2**30), 2**30, (num_kv_heads, swa_total), dtype=torch.int32, device=device ) - - output_reference = output_stale.clone() dense_reference = dense_stale.clone() swa_reference = swa_stale.clone() - _settle_oracle( - scores, row_lengths, row_prompt_offsets, provisional, output_reference, keep_count - ) _pack_oracle( - output_reference, + settled, valid_seq_lens, dense_offsets, dense_reference, @@ -191,41 +244,32 @@ def test_fused_settle_pack_matches_torch_oracle(eviction_mode, has_swa, width, k has_swa=has_swa, ) - output_fused = output_stale.clone() - dense_fused = dense_stale.clone() - swa_fused = swa_stale.clone() - swa_fused_arg = swa_fused if has_swa else dense_fused - _settle_ties_and_pack_compaction_sources_kernel[(request_count, selection_rows)]( - scores, - row_lengths, - row_prompt_offsets, - provisional, - output_fused, + dense_actual = dense_stale.clone() + swa_actual = swa_stale.clone() + _pack_move_sources_kernel[(request_count, selection_rows)]( + settled, valid_seq_lens, dense_offsets, - dense_fused, - swa_offsets if has_swa else dense_offsets, - swa_fused_arg, - WIDTH=width, + dense_actual, + swa_offsets if has_swa else None, + swa_actual if has_swa else None, KEEP_COUNT=keep_count, - SELECTION_ROWS=selection_rows, + DECISION_ROWS=selection_rows, DENSE_TOTAL=dense_total, SWA_TOTAL=swa_total if has_swa else 0, MOVE_CAPACITY=move_capacity, NUM_KV_HEADS=num_kv_heads, SWA_WINDOW=swa_window if has_swa else 0, - UNION=union, + BROADCAST=union, PER_LAYER=per_layer, HAS_SWA=has_swa, - HAS_SETTLE=True, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, + BLOCK=_PACK_BLOCK, + num_warps=_PACK_NUM_WARPS, ) torch.cuda.synchronize(device) - assert torch.equal(output_fused, output_reference), f"kept ordinals differ (seed {seed})" - assert torch.equal(dense_fused, dense_reference), f"dense moves differ (seed {seed})" - assert torch.equal(swa_fused, swa_reference), f"SWA moves differ (seed {seed})" + assert torch.equal(dense_actual, dense_reference), f"dense moves differ (seed {seed})" + assert torch.equal(swa_actual, swa_reference), f"SWA moves differ (seed {seed})" def test_settle_handles_topk_sentinel_padding(): @@ -261,35 +305,17 @@ def test_settle_handles_topk_sentinel_padding(): stale = 0x5EED output = torch.full((rows_total, keep_count), stale, dtype=torch.int32, device=device) - # The pack half always runs now; zero per-request move counts mask every - # pack store off, so the settle assertions below stay byte-exact. - dense_offsets = torch.zeros(rows_total + 1, dtype=torch.int32, device=device) - dense_indices = torch.zeros(1, dtype=torch.int32, device=device) - _settle_ties_and_pack_compaction_sources_kernel[(rows_total, 1)]( + _settle_ties_kernel[(rows_total, 1)]( scores, row_lengths, row_prompt_offsets, provisional, output, - row_lengths, - dense_offsets, - dense_indices, - dense_offsets, - dense_indices, WIDTH=width, KEEP_COUNT=keep_count, SELECTION_ROWS=1, - DENSE_TOTAL=0, - SWA_TOTAL=0, - MOVE_CAPACITY=keep_count, - NUM_KV_HEADS=1, - SWA_WINDOW=0, - UNION=False, - PER_LAYER=False, - HAS_SWA=False, - HAS_SETTLE=True, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, + BLOCK=_SETTLE_BLOCK, + num_warps=_SETTLE_NUM_WARPS, ) torch.cuda.synchronize(device) 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 index 9fbe1391aa96..1f56c83223e5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -15,8 +15,8 @@ from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_PACK_BLOCK, - SETTLE_PACK_NUM_WARPS, + SETTLE_BLOCK, + SETTLE_NUM_WARPS, prepare_per_head_scores, ) @@ -48,7 +48,7 @@ def _make_selection_buffers( num_kv_heads=1, ): """Selection-only buffers for the mode, without CuTe score state or - compaction (the settle's pack half is masked off).""" + compaction (the settle launch writes only the kept-ordinal rows).""" bufs = SimpleNamespace( eviction_mode=eviction_mode, device=device, @@ -105,42 +105,22 @@ def _make_selection_buffers( bufs.provisional_rows = bufs.top_indices_i32.view(-1, keep_count) bufs.keep_rows = bufs.keep.view(-1, keep_count) bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) - # The settle launch always packs now; zero per-request move counts mask - # every pack store off, so these buffers stay selection-only. The launch - # args mirror the product's ``bufs.settle_args``/``bufs.settle_kwargs`` - # order and keys exactly (swa pointers are None without SWA layers). - zero_offsets = torch.zeros(max_requests + 1, dtype=torch.int32, device=device) - zero_lengths = torch.zeros(max_requests, dtype=torch.int32, device=device) - zero_indices = torch.zeros(1, dtype=torch.int32, device=device) + # The launch args mirror the product's ``bufs.settle_args``/ + # ``bufs.settle_kwargs`` order and keys exactly. bufs.settle_args = ( bufs.selection_scores_rows, bufs.selection_row_lengths, bufs.row_prompt_offsets, bufs.provisional_rows, bufs.keep_rows, - zero_lengths, - zero_offsets, - zero_indices, - None, - None, ) bufs.settle_kwargs = dict( WIDTH=width, KEEP_COUNT=keep_count, SELECTION_ROWS=bufs.selection_rows_per_request, - DENSE_TOTAL=0, - SWA_TOTAL=0, - MOVE_CAPACITY=keep_count, - NUM_KV_HEADS=1, - SWA_WINDOW=0, - UNION=eviction_mode == "union", - PER_LAYER=eviction_mode == "per_layer_perhead", - HAS_SWA=False, - HAS_SETTLE=True, - BLOCK=SETTLE_PACK_BLOCK, - num_warps=SETTLE_PACK_NUM_WARPS, + BLOCK=SETTLE_BLOCK, + num_warps=SETTLE_NUM_WARPS, ) - bufs.draft_pack_launch = None return bufs @@ -418,9 +398,9 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) protected_tail_capacity=max(protected_tails), ) _set_protected_tails(compaction, protected_tails) - # Production packs these buffers inside the fused settle launch; with - # pre-settled ordinals the standalone pack in run_compaction is its - # exact analog. + # Production settles the kept ordinals into the contract's decision + # rows; with pre-settled ordinals the pack launch inside compact() is + # its exact analog. _run_compaction(compaction) torch.cuda.synchronize(device) @@ -538,8 +518,8 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): bufs.valid_seq_lens_device.fill_(seq_len) bufs.token_starts_device.fill_(0) # Stage this round's move offsets into the buffers' OWN metadata row: - # the fused settle launch and the native moves consume the buffers' - # own contract (keep_count moves per request, no protected tail). + # the pack launch and the native moves consume the buffers' own + # contract (keep_count moves per request, no protected tail). bufs.request_metadata_device[3, :2].copy_( torch.tensor([0, keep_count], dtype=torch.int32, device=device) ) From b83f80cf42fdbad6306783f5d20206696c1b5715 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 01:30:56 -0700 Subject: [PATCH 116/178] [None][refactor] Remove the single-value pin_prefill and count_prompt_tokens knobs (user-approved) Signed-off-by: tianruih --- examples/triattention/README.md | 1 - .../triattention/triattention.py | 12 ++---------- tensorrt_llm/llmapi/llm_args.py | 11 ----------- tensorrt_llm/usage/llm_args_golden_manifest.json | 14 -------------- 4 files changed, 2 insertions(+), 36 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index 0389f38332ad..8fc067e7a29b 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -121,6 +121,5 @@ trtllm-eval --model --config config.yaml longbench_v2 --max_outp * `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 requires `True` (the fused union pipeline always z-normalizes; construction rejects `False`). -* **`pin_prefill`** (bool, default=True): Always preserve the prompt (prefill) tokens; only decode tokens compete for the budget (upstream behaviour). * **`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/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 28f00ec1c227..af736ab41141 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -667,8 +667,6 @@ def __init__( calibration_path: Optional[str] = None, eviction_mode: str = "union", normalize_scores: bool = True, - pin_prefill: bool = True, - count_prompt_tokens: bool = False, ): super().__init__(kv_cache_manager, draft_kv_cache_manager) self.budget = budget @@ -687,14 +685,8 @@ def __init__( "TriAttention union eviction requires normalize_scores=True: " "the fused union pipeline always z-normalizes score rows" ) - self.pin_prefill = bool(pin_prefill) - # False (default): the budget counts decode tokens only. - self.count_prompt_tokens = bool(count_prompt_tokens) - if not self.pin_prefill or self.count_prompt_tokens: - raise ValueError( - "TriAttention physical KV reclaim requires pin_prefill=True and " - "count_prompt_tokens=False so finalized prompt KV is preserved" - ) + # Hard-coded semantics: the prompt is always pinned and the budget + # counts decode tokens only (physical KV reclaim requires both). # Calibration is the official TriAttention .pt; TRT-LLM does not compute calibration. self.model_path = model_path if self.model_path is None: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index cf2ba9ec9b75..896f4aa6fdef 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3452,10 +3452,6 @@ class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): 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.") - pin_prefill: bool = Field( - default=True, - description="Always preserve the prompt (prefill) tokens; only decode " - "tokens compete for the budget (upstream behaviour).") budget: int = Field( default=2048, gt=0, @@ -3479,11 +3475,6 @@ class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): "(produced by github.com/WeianMao/triattention). TRT-LLM does not " "compute calibration; it converts this file to the runtime schema at " "load.") - count_prompt_tokens: bool = Field( - default=False, - description="If False (default), the KV budget counts only DECODE tokens " - "(the pinned prompt is kept on top). Physical capacity reclaim currently " - "requires False.") @model_validator(mode="after") def _require_calibration_inputs(self): @@ -3505,8 +3496,6 @@ def to_manager_kwargs(self) -> dict: "calibration_path": self.calibration_path, "eviction_mode": self.eviction_mode, "normalize_scores": self.normalize_scores, - "pin_prefill": self.pin_prefill, - "count_prompt_tokens": self.count_prompt_tokens, } diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 047aaa3bbbb5..1c3f6acb0c5d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -567,13 +567,6 @@ "kind": "value", "path": "kv_cache_compression_config.budget" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_cache_compression_config.count_prompt_tokens" - }, { "allowed_values": [ "union", @@ -592,13 +585,6 @@ "kind": "value", "path": "kv_cache_compression_config.normalize_scores" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "kv_cache_compression_config.pin_prefill" - }, { "allowed_values": [], "annotation": "", From a611ebd3d093beb1d5abb68daed1c5ba4d85f1a8 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 02:01:18 -0700 Subject: [PATCH 117/178] [None][refactor] One round owner: merge stage, run, and consume; delete fingerprints; one prepared cohort (knife 25c) Signed-off-by: tianruih --- .../triattention/triattention.py | 507 ++++++++---------- .../_torch/kv_cache_compression/conftest.py | 48 +- .../test_triattention_draft_cocompaction.py | 162 ++++-- .../test_triattention_pipeline.py | 159 +++--- .../test_triattention_selection_compaction.py | 64 ++- 5 files changed, 496 insertions(+), 444 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index af736ab41141..627f2990b24c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -443,7 +443,6 @@ def init_eviction_buffers( bufs.bulk_copy_done = torch.cuda.Event() bufs.bulk_consume_done = torch.cuda.Event() bufs.copy_pending = False - bufs.page_tables_active = False return bufs @@ -478,87 +477,32 @@ def _stage_block_offsets( current_stream.wait_event(bufs.bulk_copy_done) -def stage_eviction_cohort( +def _cohort_move_offsets( bufs: SimpleNamespace, - manager: KVCacheManagerV2, - request_ids: List[int], - round_starts: List[int], - token_starts: List[int], - seq_lens: List[int], - draft_manager: Optional[KVCacheManagerV2] = None, - dense_move_offsets: Optional[List[int]] = None, - swa_move_offsets: Optional[List[int]] = None, - draft_move_offsets: Optional[List[int]] = None, -) -> None: - """Copy one eviction cohort into the reusable device buffers.""" - request_count = len(request_ids) - stream = torch.cuda.current_stream(bufs.device) - if bufs.page_tables_active: - raise RuntimeError("previous page-table cohort is still active") - # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. - max_round_start = max(round_starts) - rows = ( - (0, round_starts), - (1, seq_lens), - (2, token_starts), - (3, dense_move_offsets), - (4, swa_move_offsets), - (5, draft_move_offsets), - ) - for row, values in rows: - if values is not None and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF: - raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # The previous cohort's metadata H2D must complete before the pinned rows are rewritten. - if bufs.copy_pending and not bufs.copy_done.query(): - bufs.copy_done.synchronize() - host_table = bufs.request_metadata_host_np - for row, values in rows: - if values is not None: - host_table[row, : len(values)] = values - # Zero lengths keep the score kernel and selection inert for padded rows. - host_table[:3, request_count:] = 0 - grow_mean_phase_table(bufs.phase, int(max_round_start) + 1) - _stage_block_offsets( - bufs, - manager, - request_ids, - stream, - bufs._bulk_offsets_src, - bufs.block_offsets_device, - bufs.copy_block_count, - ) + prepared: Sequence[Dict[str, object]], + draft_manager: Optional[KVCacheManagerV2], +) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: + """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep + set plus protected tail per request; rows past the cohort repeat the final + offset and contribute no moves).""" + + def padded_offsets(moves_per_request: List[int]) -> List[int]: + offsets = [0] + for moves in moves_per_request: + offsets.append(offsets[-1] + moves) + offsets.extend(offsets[-1:] * (bufs.max_requests - len(moves_per_request))) + return offsets + + tails = [int(item["protected_tail"]) for item in prepared] + dense = padded_offsets([bufs.keep_count + tail for tail in tails]) + swa = None + if bufs.compaction["has_swa"]: + swa = padded_offsets([int(bufs.compaction["swa_window"]) + tail for tail in tails]) + draft = None if draft_manager is not None: - _stage_block_offsets( - bufs, - draft_manager, - request_ids, - stream, - bufs._draft_bulk_offsets_src, - bufs.draft_block_offsets_device, - bufs.draft_copy_block_count, - ) - try: - bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) - finally: - # Guards the pinned metadata until the asynchronous copies complete. - bufs.copy_done.record(stream) - bufs.copy_pending = True - bufs.page_tables_active = True - # Per-head modes re-expand the prompt lengths into their row-major view. - if bufs.row_prompt_offsets is not bufs.prompt_offsets: - bufs.row_prompt_offsets.view(bufs.max_requests, bufs.selection_rows_per_request).copy_( - bufs.prompt_offsets.unsqueeze(1).expand(-1, bufs.selection_rows_per_request) - ) - - -def mark_page_tables_consumed(bufs: SimpleNamespace, *manager_streams: torch.cuda.Stream) -> None: - """Order V2 page-table reuse and resize after this cohort's compact.""" - if not bufs.page_tables_active: - raise RuntimeError("TriAttention page tables were not staged") - bufs.bulk_consume_done.record(torch.cuda.current_stream(bufs.device)) - for manager_stream in manager_streams: - manager_stream.wait_event(bufs.bulk_consume_done) - bufs.page_tables_active = False + draft_tail = _protected_tail_capacity(draft_manager, "draft ") + draft = padded_offsets([bufs.keep_count + draft_tail] * len(prepared)) + return dense, swa, draft def settle_top_tokens(bufs: SimpleNamespace) -> None: @@ -575,81 +519,169 @@ def settle_top_tokens(bufs: SimpleNamespace) -> None: _settle_ties_kernel[bufs.settle_grid](*bufs.settle_args, **bufs.settle_kwargs) -def run_eviction_round(bufs: SimpleNamespace, normalize_scores: bool) -> None: - """Fire one staged eviction round: score, select, settle, compact - (every launch covers the full request capacity; padded rows carry zero lengths and stay inert).""" - request_count = bufs.max_requests - union = bufs.eviction_mode == "union" - with nvtx_range("triattention.score", color="blue"): - # In-place refresh: the compiled score launches captured these pointers. - _gather_mean_phase_kernel[(request_count,)]( - bufs.round_starts_device, - bufs.phase["cos"], - bufs.phase["sin"], - bufs.mean_cos, - bufs.mean_sin, - bufs.valid_seq_lens_device, - bufs.token_starts_device, - bufs.valid_widths, - bufs.swa_destination_bases, - bufs.phase["rows"], - bufs.swa_rebase_delta, - NUM_FREQS=bufs.phase_num_freqs, - F_BLOCK=bufs.phase_f_block, - HAS_SWA=bufs.swa_destination_bases is not None, - num_warps=1, +def execute_eviction_round( + bufs: SimpleNamespace, + manager: KVCacheManagerV2, + draft_manager: Optional[KVCacheManagerV2], + prepared: Sequence[Dict[str, object]], + *, + normalize_scores: bool, +) -> None: + """Run one eviction round over the prepared cohort: stage the page-table + snapshots and round metadata, then score, select, settle, and compact, and + finally order the manager streams after this cohort's compact (every launch + covers the full request capacity; padded rows carry zero lengths and stay + inert).""" + with nvtx_range_debug("triattention.page_table_stage", color="orange"): + request_ids = [item["request_id"] for item in prepared] + round_starts = [item["round_start"] for item in prepared] + token_starts = [item["prompt_len"] for item in prepared] + seq_lens = [item["seq_len"] for item in prepared] + dense_move_offsets, swa_move_offsets, draft_move_offsets = _cohort_move_offsets( + bufs, prepared, draft_manager ) - if union: - bufs.runner.launch_union_fusion( - request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_rows[:request_count] - ) - columns = min(bufs.union_rows.shape[1], bufs.combined.shape[1]) - bufs.combined[:request_count, :columns].copy_(bufs.union_rows[:request_count, :columns]) - else: - bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) - # Gather each decode window into the [request, layer, head, token] layout of the reduces. - group_size = bufs.num_q_heads // bufs.num_kv_heads - num_segments = request_count * bufs.num_layers - pad = bufs.padded_head_columns - source = ( - bufs.cute_scratch[: bufs.num_kv_heads * pad * num_segments * bufs.bucket_seq_len] - .view(bufs.num_kv_heads, pad, request_count, bufs.num_layers, bufs.bucket_seq_len)[ - :, :group_size - ] - .permute(2, 3, 0, 1, 4) - ) - columns = ( - bufs.token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) - + bufs.gather_columns - ) - columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + stream = torch.cuda.current_stream(bufs.device) + # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. + max_round_start = max(round_starts) + rows = ( + (0, round_starts), + (1, seq_lens), + (2, token_starts), + (3, dense_move_offsets), + (4, swa_move_offsets), + (5, draft_move_offsets), + ) + for row, values in rows: + if values is not None and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF: + raise ValueError(f"staged metadata row {row} exceeds the int32 range") + # The previous cohort's metadata H2D must complete before the pinned rows are rewritten. + if bufs.copy_pending and not bufs.copy_done.query(): + bufs.copy_done.synchronize() + host_table = bufs.request_metadata_host_np + for row, values in rows: + if values is not None: + host_table[row, : len(values)] = values + # Zero lengths keep the score kernel and selection inert for padded rows. + host_table[:3, len(prepared) :] = 0 + grow_mean_phase_table(bufs.phase, int(max_round_start) + 1) + _stage_block_offsets( + bufs, + manager, + request_ids, + stream, + bufs._bulk_offsets_src, + bufs.block_offsets_device, + bufs.copy_block_count, + ) + if draft_manager is not None: + _stage_block_offsets( + bufs, + draft_manager, + request_ids, + stream, + bufs._draft_bulk_offsets_src, + bufs.draft_block_offsets_device, + bufs.draft_copy_block_count, ) - torch.gather( - source, - 4, - columns, - out=bufs.score_output[:request_count].view( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width - ), + try: + bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) + finally: + # Guards the pinned metadata until the asynchronous copies complete. + bufs.copy_done.record(stream) + bufs.copy_pending = True + # Per-head modes re-expand the prompt lengths into their row-major view. + if bufs.row_prompt_offsets is not bufs.prompt_offsets: + bufs.row_prompt_offsets.view(bufs.max_requests, bufs.selection_rows_per_request).copy_( + bufs.prompt_offsets.unsqueeze(1).expand(-1, bufs.selection_rows_per_request) ) - with nvtx_range("triattention.select", color="yellow"): - if not union: - prepare_per_head_scores( - bufs.score_output[:request_count], + + request_count = bufs.max_requests + union = bufs.eviction_mode == "union" + try: + with nvtx_range("triattention.score", color="blue"): + # In-place refresh: the compiled score launches captured these pointers. + _gather_mean_phase_kernel[(request_count,)]( + bufs.round_starts_device, + bufs.phase["cos"], + bufs.phase["sin"], + bufs.mean_cos, + bufs.mean_sin, + bufs.valid_seq_lens_device, + bufs.token_starts_device, bufs.valid_widths, - bufs.row_mean, - bufs.row_std, - bufs.selection_scores, - bufs.row_seq_lens, - request_count, - num_kv_heads=bufs.num_kv_heads, - per_layer=bufs.eviction_mode == "per_layer_perhead", - normalize_scores=normalize_scores, + bufs.swa_destination_bases, + bufs.phase["rows"], + bufs.swa_rebase_delta, + NUM_FREQS=bufs.phase_num_freqs, + F_BLOCK=bufs.phase_f_block, + HAS_SWA=bufs.swa_destination_bases is not None, + num_warps=1, ) - settle_top_tokens(bufs) - with nvtx_range("triattention.compact", color="purple"): - compact(bufs.compaction, request_count) + if union: + bufs.runner.launch_union_fusion( + request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_rows[:request_count] + ) + columns = min(bufs.union_rows.shape[1], bufs.combined.shape[1]) + bufs.combined[:request_count, :columns].copy_( + bufs.union_rows[:request_count, :columns] + ) + else: + bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) + # Gather each decode window into the [request, layer, head, token] layout of the reduces. + group_size = bufs.num_q_heads // bufs.num_kv_heads + num_segments = request_count * bufs.num_layers + pad = bufs.padded_head_columns + source = ( + bufs.cute_scratch[ + : bufs.num_kv_heads * pad * num_segments * bufs.bucket_seq_len + ] + .view( + bufs.num_kv_heads, pad, request_count, bufs.num_layers, bufs.bucket_seq_len + )[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + columns = ( + bufs.token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + + bufs.gather_columns + ) + columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( + request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + ) + torch.gather( + source, + 4, + columns, + out=bufs.score_output[:request_count].view( + request_count, + bufs.num_layers, + bufs.num_kv_heads, + group_size, + bufs.decode_width, + ), + ) + with nvtx_range("triattention.select", color="yellow"): + if not union: + prepare_per_head_scores( + bufs.score_output[:request_count], + bufs.valid_widths, + bufs.row_mean, + bufs.row_std, + bufs.selection_scores, + bufs.row_seq_lens, + request_count, + num_kv_heads=bufs.num_kv_heads, + per_layer=bufs.eviction_mode == "per_layer_perhead", + normalize_scores=normalize_scores, + ) + settle_top_tokens(bufs) + with nvtx_range("triattention.compact", color="purple"): + compact(bufs.compaction, request_count) + finally: + # Order V2 page-table reuse and resize after this cohort's compact. + bufs.bulk_consume_done.record(torch.cuda.current_stream(bufs.device)) + manager._stream.wait_event(bufs.bulk_consume_done) + if draft_manager is not None: + draft_manager._stream.wait_event(bufs.bulk_consume_done) class TriAttention(BaseKVCacheCompressionManager): @@ -709,16 +741,17 @@ def __init__( # Per-request {generation_steps, evicted_tokens}. self._request_states: Dict[int, Dict[str, object]] = {} - # In-flight overlap batch reference; membership and growth resolve lazily. + # In-flight overlap batch reference; membership resolves lazily. self._prepared_generation_batch: Optional[object] = None self._prepared_generation_ids: Optional[set] = None - self._generation_growth: Optional[int] = None - # Memoized manager invariants. - self._v2_validated = False - self._protected_tail_cache: Optional[int] = None + # Manager-lifetime capability gates: everything read there is fixed at + # construction, so validation runs once here. + self._validate_v2_compatibility() + # Manager-lifetime constants (V2 fixes both inputs at construction). + self._protected_tail_capacity = _protected_tail_capacity(kv_cache_manager, "") + self._generation_growth = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) # Built once at the first eviction, reused for the manager's lifetime. self._buffers: Optional[SimpleNamespace] = None - self._buffers_fingerprint: Optional[tuple] = None self._local_to_global_layers_cache: Optional[List[int]] = None self._attention_layer_partition_cache: Optional[ Tuple[List[int], List[int], Optional[int]] @@ -727,10 +760,9 @@ def __init__( self._draft_runtime_kv_layout_cache: Optional[Dict[str, object]] = None def on_request_init(self, request: "LlmRequest", **kwargs) -> None: - """Validate once and resolve the official calibration for the first request.""" + """Track the request and resolve the official calibration on first use.""" request_id = request.py_request_id if request_id not in self._request_states: - self._validate_v2_compatibility() self._validate_request_capacity(request) num_layers = self._num_layers_from_manager() self._attention_layer_partition(num_layers) @@ -748,7 +780,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: ) * self.beta + speculative_overshoot decode_capacity = min(int(request.py_max_new_tokens), first_eviction_decode_length) confirmed_capacity = int(request.py_prompt_len) + decode_capacity - protected_tail_capacity = self._configured_protected_tail_capacity() + protected_tail_capacity = self._protected_tail_capacity required_capacity = confirmed_capacity + protected_tail_capacity pool_confirmed_capacity = manager.get_num_available_tokens( token_num_upper_bound=confirmed_capacity, @@ -812,8 +844,6 @@ def _ensure_calibrated(self) -> None: self._calibrated = True def _validate_v2_compatibility(self) -> None: - if self._v2_validated: - return manager = self.kv_cache_manager if not isinstance(manager, KVCacheManagerV2): raise ValueError("TriAttention physical eviction requires KVCacheManagerV2") @@ -870,7 +900,6 @@ def _validate_v2_compatibility(self) -> None: "TriAttention requires full-attention V2 lifecycles; native SWA, " "VSWA, and SSM pools are not supported" ) - self._v2_validated = True def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Compact after native KV-cache updates have finalized this iteration @@ -895,11 +924,7 @@ def _inflight_generation_growth( self._prepared_generation_ids = member_ids if request_id not in member_ids: return 0 - growth = self._generation_growth - if growth is None: - growth = 1 + int(self.kv_cache_manager._kv_reserve_draft_tokens) - self._generation_growth = growth - return growth + return self._generation_growth def _periodic_evict( self, @@ -930,9 +955,8 @@ def _periodic_evict( resolved_requests.append((request, request_id, kv_cache)) if not resolved_requests: return - protected_tails: Dict[int, int] = {} prepared: List[Dict[str, object]] = [] - protected_tail_capacity = self._configured_protected_tail_capacity() + protected_tail_capacity = self._protected_tail_capacity # The resolved cache objects thread all the way to resize. with nvtx_range("triattention.metadata", color="cyan"): @@ -978,7 +1002,6 @@ def _periodic_evict( f"suspended draft KV cache; request {request_id} must " "be resumed before the final update hook" ) - protected_tails[request_id] = protected_tail prepared.append( { "request": request, @@ -1002,40 +1025,42 @@ def _periodic_evict( f"triattention.evict_request_group reqs={len(prepared)}", color="purple", ): - capacity_targets = self._evict_requests(prepared, num_layers) - self._resize_compacted_requests(capacity_targets, protected_tails) + compacted = self._evict_requests(prepared, num_layers) + self._resize_compacted_requests(compacted) - def _resize_compacted_requests(self, capacity_targets, protected_tails) -> None: - if not capacity_targets: + def _resize_compacted_requests(self, prepared) -> None: + if not prepared: return with nvtx_range("triattention.resize", color="red"): with nvtx_range_debug("triattention.v2_resize", color="red"): - for rid, kv_cache, _, target_capacity in capacity_targets: + for item in prepared: + kv_cache = item["kv_cache"] if not kv_cache.is_active: continue + target_capacity = item["expected_keep_count"] if target_capacity > kv_cache.capacity: raise RuntimeError( - f"Request {rid} compacted capacity {target_capacity} exceeds " - f"current capacity {kv_cache.capacity}" + f"Request {item['request_id']} compacted capacity " + f"{target_capacity} exceeds current capacity {kv_cache.capacity}" ) - protected_tail = protected_tails[rid] - resized_capacity = target_capacity + protected_tail + resized_capacity = target_capacity + item["protected_tail"] if not kv_cache.resize(resized_capacity, None): raise RuntimeError( - f"Failed to resize compacted KV cache for request {rid} " - f"to {resized_capacity} tokens" + f"Failed to resize compacted KV cache for request " + f"{item['request_id']} to {resized_capacity} tokens" ) if self.draft_kv_cache_manager is not None: # Same kept set: the draft shrinks to the same retained length plus its own tail. draft_protected_tail = self._draft_protected_tail_capacity() - for rid, _, draft_kv_cache, target_capacity in capacity_targets: + for item in prepared: + draft_kv_cache = item["draft_kv_cache"] if not draft_kv_cache.is_active: continue - draft_capacity = target_capacity + draft_protected_tail + draft_capacity = item["expected_keep_count"] + draft_protected_tail if not draft_kv_cache.resize(draft_capacity, None): raise RuntimeError( - "Failed to resize co-compressed draft KV cache " - f"for request {rid} to {draft_capacity} tokens" + "Failed to resize co-compressed draft KV cache for " + f"request {item['request_id']} to {draft_capacity} tokens" ) def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: @@ -1071,11 +1096,6 @@ def _local_score_calibration( self._triattn_mlr_coef.index_select(0, layer_ids), ) - def _configured_protected_tail_capacity(self) -> int: - if self._protected_tail_cache is None: - self._protected_tail_cache = _protected_tail_capacity(self.kv_cache_manager, "") - return self._protected_tail_cache - def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: """Drop this request's eviction state; the buffers stay resident.""" self._request_states.pop(request.py_request_id, None) @@ -1273,9 +1293,6 @@ def _build_runtime_kv_layout( pool_page_counts=tuple( int(layer_pools[layer].shape[0]) for layer in pool_representatives ), - pool_view_fingerprint=self._pool_view_fingerprint( - [layer_pools[layer] for layer in pool_representatives] - ), ) def _draft_runtime_kv_layout(self) -> Dict[str, object]: @@ -1330,19 +1347,6 @@ def _pool_page_counts( for layer in pool_representatives ) - @staticmethod - def _pool_view_fingerprint(pools: List[torch.Tensor]) -> Tuple[tuple, ...]: - return tuple( - ( - pool.data_ptr(), - tuple(int(value) for value in pool.shape), - tuple(int(value) for value in pool.stride()), - pool.dtype, - pool.device, - ) - for pool in pools - ) - def _buffers_for( self, layout: Dict[str, object], @@ -1353,34 +1357,23 @@ def _buffers_for( needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) needed_requests = len(prepared) - draft_fingerprint = None if self.draft_kv_cache_manager is not None: - draft_layout = self._draft_runtime_kv_layout() - draft_fingerprint = ( - draft_layout["pool_page_counts"], - draft_layout["pool_view_fingerprint"], - ) - fingerprint = ( - self.eviction_mode, - self.budget, - tuple(layout["dense_layers"]), - layout["pool_view_fingerprint"], - draft_fingerprint, - ) + # The cached layout lookup enforces draft V2 pool page-count + # stability every round, exactly like the target's lookup. + self._draft_runtime_kv_layout() bufs = self._buffers if bufs is not None: if ( - self._buffers_fingerprint == fingerprint - and needed_width <= bufs.decode_width + needed_width <= bufs.decode_width and needed_page_tokens <= bufs.page_table_token_capacity and needed_requests <= bufs.max_requests ): return bufs - # Pools changed or this round outgrew the buffers: rebuild. + # This round outgrew the buffers: rebuild. self._buffers = None mgr = self.kv_cache_manager - tail_capacity = self._configured_protected_tail_capacity() + tail_capacity = self._protected_tail_capacity request_capacity = max(needed_requests, int(mgr.max_batch_size)) decode_width = max( needed_width, @@ -1449,33 +1442,8 @@ def _buffers_for( draft=draft, ) self._buffers = bufs - self._buffers_fingerprint = fingerprint return bufs - def _move_offsets_for( - self, - layout: Dict[str, object], - prepared: Sequence[Dict[str, object]], - capacity: int, - ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - def padded_offsets(moves_per_request: List[int]) -> List[int]: - offsets = [0] - for moves in moves_per_request: - offsets.append(offsets[-1] + moves) - offsets.extend(offsets[-1:] * (capacity - len(moves_per_request))) - return offsets - - tails = [item["protected_tail"] for item in prepared] - dense = padded_offsets([self.budget + tail for tail in tails]) - swa = None - if layout["swa_layers"] and layout["swa_window"]: - swa = padded_offsets([int(layout["swa_window"]) + tail for tail in tails]) - draft = None - if self.draft_kv_cache_manager is not None: - draft_tail = self._draft_protected_tail_capacity() - draft = padded_offsets([self.budget + draft_tail] * len(prepared)) - return dense, swa, draft - def _page_table_pool_keys( self, representatives: List[int], @@ -1511,51 +1479,28 @@ def _evict_requests( self, prepared: List[Dict[str, object]], num_layers: int, - ) -> List[Tuple[int, int]]: + ) -> List[Dict[str, object]]: with nvtx_range_debug("triattention.resolve_layout", color="blue"): layout = self._runtime_kv_layout(num_layers) with nvtx_range_debug("triattention.staging_lookup", color="blue"): # Retained spans always cover the model window (construction rejects budget < window). bufs = self._buffers_for(layout, prepared) - with nvtx_range_debug("triattention.page_table_stage", color="orange"): - dense_offsets, swa_offsets, draft_offsets = self._move_offsets_for( - layout, prepared, bufs.max_requests - ) - stage_eviction_cohort( - bufs, - self.kv_cache_manager, - [item["request_id"] for item in prepared], - [item["round_start"] for item in prepared], - [item["prompt_len"] for item in prepared], - [item["seq_len"] for item in prepared], - draft_manager=self.draft_kv_cache_manager, - dense_move_offsets=dense_offsets, - swa_move_offsets=swa_offsets, - draft_move_offsets=draft_offsets, - ) - - try: - run_eviction_round(bufs, self.normalize_scores) - finally: - consumer_streams = [self.kv_cache_manager._stream] - if self.draft_kv_cache_manager is not None: - consumer_streams.append(self.draft_kv_cache_manager._stream) - mark_page_tables_consumed(bufs, *consumer_streams) - - capacity_targets = [] + execute_eviction_round( + bufs, + self.kv_cache_manager, + self.draft_kv_cache_manager, + prepared, + normalize_scores=self.normalize_scores, + ) for item in prepared: - keep_count = item["expected_keep_count"] - evicted = item["seq_len"] - keep_count + evicted = item["seq_len"] - item["expected_keep_count"] if evicted <= 0: raise RuntimeError("TriAttention attempted an identity compaction") request_state = self._request_states[item["request_id"]] request_state["evicted_tokens"] += evicted # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] - capacity_targets.append( - (item["request_id"], item["kv_cache"], item["draft_kv_cache"], keep_count) - ) - return capacity_targets + return prepared def _num_layers_from_manager(self) -> int: return len(self.kv_cache_manager.pp_layers) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 266ae7bac1ba..3efa4634c213 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -232,9 +232,10 @@ def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None staging.copy_block_count = copy_block_count if page_count is not None: staging.page_count = page_count + staging.keep_count = 4 + staging.compaction = {"has_swa": False} staging.bulk_copy_done = torch.cuda.Event() staging.bulk_consume_done = torch.cuda.Event() - staging.page_tables_active = False staging.copy_done = torch.cuda.Event() staging.copy_pending = False staging._bulk_offsets_src = torch.empty( @@ -283,7 +284,6 @@ def make_buffer_stubs(manager, *, decode_width=260): storage_groups={0: [0, 1]}, layer_group_representative={0: 0, 1: 0}, layer_pool_keys=(("pool", 0), ("pool", 0)), - pool_view_fingerprint=(("fixed",),), ) buffers = SimpleNamespace( decode_width=decode_width, @@ -337,6 +337,32 @@ def make_triattention(**overrides): return TriAttention(make_fake_v2(), **options) +def make_prepared_item( + request=None, + *, + request_id=0, + seq_len, + round_start=None, + prompt_len=0, + expected_keep_count=0, + protected_tail=0, + kv_cache=None, + draft_kv_cache=None, +): + """One prepared-cohort item shaped exactly like ``_periodic_evict`` builds.""" + return { + "request": request, + "request_id": request_id, + "kv_cache": kv_cache, + "draft_kv_cache": draft_kv_cache, + "seq_len": int(seq_len), + "round_start": int(seq_len if round_start is None else round_start), + "prompt_len": int(prompt_len), + "expected_keep_count": int(expected_keep_count), + "protected_tail": int(protected_tail), + } + + def make_request(request_id, **overrides): """Build the explicit request fields consumed by TriAttention.""" from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState @@ -357,24 +383,16 @@ def make_request(request_id, **overrides): @contextmanager def mocked_eviction_internals(manager): - """Run the real ``_evict_requests`` body around mocked GPU launches.""" + """Run the real ``_evict_requests`` body around a mocked round executor.""" from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module buffers = SimpleNamespace(max_requests=8) - layout = dict(swa_layers=[], swa_window=None) with ( - mock.patch.object(manager, "_runtime_kv_layout", return_value=layout), + mock.patch.object(manager, "_runtime_kv_layout", return_value={}), mock.patch.object(manager, "_buffers_for", return_value=buffers), - mock.patch.object(module, "stage_eviction_cohort") as stage, - mock.patch.object(module, "run_eviction_round") as run_round, - mock.patch.object(module, "mark_page_tables_consumed") as consumed, + mock.patch.object(module, "execute_eviction_round") as execute, ): - yield SimpleNamespace( - buffers=buffers, - stage=stage, - run_round=run_round, - consumed=consumed, - ) + yield SimpleNamespace(buffers=buffers, execute=execute) def torch_tri_score_oracle( @@ -546,7 +564,7 @@ def launch_split_scores( bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin ): """The production score-only leg plus the decode-window gather - (``run_eviction_round``'s per-head sequence, parameterized by count).""" + (``execute_eviction_round``'s per-head sequence, parameterized by count).""" stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) assert request_count in bufs.runner._compiled bufs.runner.launch(request_count, mean_cos, mean_sin) 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 index e001064bf297..b0005d16e517 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -5,7 +5,7 @@ the draft's own KV heads, the draft's tail appends as ordinals, both land at ``destination_base = prompt_len``. Covers the physical moves, packed indices, stream ordering, admission gates, the published compressed-token -invariant, and buffer rebuild/invalidation.""" +invariant, and buffer reuse/rebuild.""" from types import SimpleNamespace from unittest import mock @@ -16,6 +16,7 @@ from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_buffer_stubs as _make_buffer_stubs from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_prepared_item as _make_prepared_item from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_request as _make_request from conftest import make_triattention as _make_triattention @@ -23,10 +24,7 @@ from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - TriAttention, - mark_page_tables_consumed, -) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention def _fresh_request_state(): @@ -174,29 +172,81 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect assert torch.equal(draft_indices[head, : expected_offsets[-1]], expected_row) -def test_mark_page_tables_consumed_orders_both_manager_streams(): +def test_execute_eviction_round_orders_both_manager_streams(): + """The round executor snapshots both page-table planes, then records one + completion event and BOTH cache-manager streams wait on it -- even when + the round body fails -- so neither manager can free or reallocate pages + this cohort is still reading.""" + from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module + event = mock.Mock() + host = torch.zeros(6, 9, dtype=torch.int32) + prompt_offsets = object() buffers = SimpleNamespace( device=torch.device("cuda", torch.cuda.current_device()), - page_tables_active=True, + max_requests=8, + keep_count=4, + eviction_mode="union", + compaction={"has_swa": False}, + copy_pending=False, + copy_done=mock.Mock(), bulk_consume_done=event, + request_metadata_host=host, + request_metadata_host_np=host.numpy(), + request_metadata_device=torch.zeros_like(host), + phase={"cos": None, "sin": None, "rows": 8}, + phase_num_freqs=1, + phase_f_block=1, + prompt_offsets=prompt_offsets, + row_prompt_offsets=prompt_offsets, + round_starts_device=None, + valid_seq_lens_device=None, + token_starts_device=None, + valid_widths=None, + mean_cos=None, + mean_sin=None, + swa_destination_bases=None, + swa_rebase_delta=0, + copy_block_count=0, + _bulk_offsets_src=None, + block_offsets_device=None, + draft_copy_block_count=0, + _draft_bulk_offsets_src=None, + draft_block_offsets_device=None, ) target_stream = mock.Mock() draft_stream = mock.Mock() + manager = SimpleNamespace(_stream=target_stream) + draft_manager = SimpleNamespace( + _stream=draft_stream, num_extra_kv_tokens=0, _kv_reserve_draft_tokens=0 + ) compute_stream = SimpleNamespace() + prepared = [_make_prepared_item(request_id=7, seq_len=8)] + + class Boom(RuntimeError): + pass + + score_kernel = mock.MagicMock() + score_kernel.__getitem__.return_value.side_effect = Boom + with ( + mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream), + mock.patch.object(module, "grow_mean_phase_table"), + mock.patch.object(module, "_stage_block_offsets") as stage, + mock.patch.object(module, "_gather_mean_phase_kernel", score_kernel), + mock.patch.object(module, "compact") as compact, + ): + with pytest.raises(Boom): + module.execute_eviction_round( + buffers, manager, draft_manager, prepared, normalize_scores=True + ) - with mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream): - mark_page_tables_consumed(buffers, target_stream, draft_stream) - - # One event records the compact launches; BOTH cache managers wait on it, - # so neither can free or reallocate pages this cohort is still reading. + # Both page-table planes were snapshotted before the round body fired. + assert stage.call_count == 2 + compact.assert_not_called() + # One event records the round; BOTH cache managers wait on it. event.record.assert_called_once_with(compute_stream) target_stream.wait_event.assert_called_once_with(event) draft_stream.wait_event.assert_called_once_with(event) - assert buffers.page_tables_active is False - - with pytest.raises(RuntimeError, match="not staged"): - mark_page_tables_consumed(buffers, target_stream, draft_stream) @pytest.mark.parametrize( @@ -234,16 +284,25 @@ def test_draft_admission_gates_raise(gate, match): draft_manager, ) return - manager = TriAttention( - _make_fake_v2(), - budget=8, - model_path="/models/test", - eviction_mode="per_head" if gate == "union_only_per_head" else "union", - draft_kv_cache_manager=None if gate == "target_kv_factor" else draft_manager, - ) + + def construct(): + return TriAttention( + _make_fake_v2(), + budget=8, + model_path="/models/test", + eviction_mode="per_head" if gate == "union_only_per_head" else "union", + draft_kv_cache_manager=None if gate == "target_kv_factor" else draft_manager, + ) + + if gate in ("union_only_per_head", "full_attention_draft"): + # Manager-lifetime capability gates run once, at construction. + with pytest.raises(ValueError, match=match): + construct() + return + manager = construct() + # Flipping kv_factor after construction exercises TriAttention's own + # capability gate directly. if gate == "draft_kv_factor": - # Flipping kv_factor after construction exercises TriAttention's own - # runtime gate. draft_manager.kv_factor = 1 if gate == "target_kv_factor": manager.kv_cache_manager.kv_factor = 1 @@ -301,10 +360,9 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert confirmed == 2 + 4 # The staged logical position restores the uncompressed # length: physical confirmed plus everything evicted so far - # (stage_eviction_cohort args: bufs, manager, ids, round_starts, - # prompt lengths, seq lens, page-table lens). - round_starts = internals.stage.call_args.args[3] - assert round_starts[0] == uncompressed + # (the prepared item's round_start). + prepared = internals.execute.call_args.args[3] + assert prepared[0]["round_start"] == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. assert request.py_num_compressed_tokens == uncompressed - confirmed @@ -313,15 +371,19 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert eviction_rounds == 3 assert previous_published == 12 - # Each round the draft cache shrinks with the target and both manager - # streams are ordered after the compact launches. + # Each round the draft cache shrinks with the target, and the one + # executor call carries both managers, whose streams it orders after the + # compact launches. assert draft_cache.resize.call_args_list == [mock.call(7, None)] * eviction_rounds - assert internals.consumed.call_args_list == ( - [mock.call(internals.buffers, target._stream, draft_manager._stream)] * eviction_rounds - ) + assert len(internals.execute.call_args_list) == eviction_rounds + for call in internals.execute.call_args_list: + assert call.args[0] is internals.buffers + assert call.args[1] is target + assert call.args[2] is draft_manager + assert call.kwargs == {"normalize_scores": True} -def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): +def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module manager = _make_triattention(budget=4) @@ -343,19 +405,10 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): pool_representatives=(), layer_pool_keys=(), pool_page_counts=(4,), - pool_view_fingerprint=(), ) ) prepared = [ - { - "request": _make_request(7), - "request_id": 7, - "seq_len": 8, - "round_start": 8, - "prompt_len": 0, - "expected_keep_count": 4, - "protected_tail": 0, - } + _make_prepared_item(_make_request(7), request_id=7, seq_len=8, expected_keep_count=4) ] with mock.patch.object( @@ -384,21 +437,28 @@ def test_pool_change_rebuilds_buffers_and_drops_cached_compaction(): assert kwargs["layout"] is layout assert list(kwargs["layout"]["layer_pool_keys"]) == list(layout["layer_pool_keys"]) - # A second round with unchanged pools reuses the resident buffers + # A second round within the resident capacities reuses the buffers # (and with them the compaction launch data they carry). assert manager._buffers_for(layout, prepared) is resources assert prepare.call_count == 1 - # A pool change invalidates the whole buffer namespace, compaction - # included. - layout["pool_view_fingerprint"] = (("moved",),) + # A cohort that outgrows the resident capacities rebuilds the whole + # buffer namespace, compaction included. + grown = [ + _make_prepared_item( + _make_request(7), + request_id=7, + seq_len=8 + buffers.decode_width, + expected_keep_count=4, + ) + ] rebuilt_buffers = SimpleNamespace( - decode_width=buffers.decode_width, + decode_width=buffers.decode_width + 8, page_table_token_capacity=buffers.page_table_token_capacity, max_requests=buffers.max_requests, ) prepare.return_value = rebuilt_buffers - rebuilt = manager._buffers_for(layout, prepared) + rebuilt = manager._buffers_for(layout, grown) assert rebuilt is not resources assert rebuilt is rebuilt_buffers assert prepare.call_count == 2 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 5111a3d13ebb..7b33488e1634 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -25,10 +25,10 @@ import pytest import torch -from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_bare_staging as _make_bare_staging from conftest import make_cute_buffers as _make_cute_buffers from conftest import make_fake_v2 as _make_fake_v2 +from conftest import make_prepared_item as _make_prepared_item from conftest import make_request as _make_request from conftest import make_staging_manager as _make_staging_manager from conftest import make_triattention as _make_triattention @@ -64,29 +64,6 @@ def _set_request_state(manager, request_id, *, generation_steps=0, evicted_token return state -def _prepared_eviction( - request, - *, - seq_len, - expected_keep_count, - protected_tail=0, - request_id=0, - round_start=None, - prompt_len=0, -): - return { - "request": request, - "request_id": request_id, - "kv_cache": None, - "draft_kv_cache": None, - "seq_len": seq_len, - "round_start": int(seq_len if round_start is None else round_start), - "prompt_len": prompt_len, - "expected_keep_count": expected_keep_count, - "protected_tail": protected_tail, - } - - @pytest.fixture def flat_calibration_pt(tmp_path): """Build a minimal valid calibration ``.pt`` in our flat runtime schema.""" @@ -180,7 +157,6 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): # These are local layer slots. Layer 2 shares layer 0's pool. pool_representatives=(0, 1), pool_page_counts=(4, 8), - pool_view_fingerprint=(), ) triattention._runtime_kv_layout_cache = cached @@ -350,7 +326,7 @@ def test_identity_compaction_is_rejected_instead_of_published(self): with _mocked_eviction_internals(manager) as internals: manager._periodic_evict(SimpleNamespace(generation_requests=[request])) - internals.run_round.assert_not_called() + internals.execute.assert_not_called() assert request.py_num_compressed_tokens == 0 assert state["evicted_tokens"] == 0 cache.resize.assert_not_called() @@ -360,7 +336,7 @@ def test_identity_compaction_is_rejected_instead_of_published(self): with pytest.raises(RuntimeError, match="identity compaction"): manager._evict_requests( [ - _prepared_eviction( + _make_prepared_item( request, request_id=7, seq_len=6, @@ -407,14 +383,19 @@ def test_unregistered_generation_request_is_rejected(self): manager._periodic_evict(SimpleNamespace(generation_requests=[request])) @staticmethod - def _make_due_decode_request(seq_len): + 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]) - mgr = _make_triattention() + 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 + mgr = TriAttention(fake_v2, budget=8, model_path="/models/test") mgr._calibrated = True cache = SimpleNamespace( capacity=seq_len, @@ -427,8 +408,8 @@ def _make_due_decode_request(seq_len): kv_cache_map={7: cache}, pp_layers=[0, 1], _stream=mock.Mock(), - num_extra_kv_tokens=0, - _kv_reserve_draft_tokens=0, + num_extra_kv_tokens=num_extra_kv_tokens, + _kv_reserve_draft_tokens=kv_reserve_draft_tokens, ) mgr._L = 2 mgr._request_states = {} @@ -458,13 +439,15 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): current_growth = 4 tail = reserve + current_growth retained = 1024 + 4096 - mgr, request, batch = self._make_due_decode_request(seq_len=confirmed) + # 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.kv_cache_manager.num_extra_kv_tokens = reserve - # Growth constant = 1 + _kv_reserve_draft_tokens for batch members. - mgr.kv_cache_manager._kv_reserve_draft_tokens = current_growth - 1 mgr._prepared_generation_batch = SimpleNamespace(generation_requests=[request]) draft_manager = _make_fake_v2(is_draft=True) draft_cache = SimpleNamespace(is_active=True, resize=mock.Mock(return_value=True)) @@ -472,8 +455,9 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): draft_manager._stream = mock.Mock() mgr.draft_kv_cache_manager = draft_manager - def compact(*_args, **_kwargs): - return [(7, cache, draft_cache, retained)] + def compact(prepared, _num_layers): + # Publish and resize consume the same prepared cohort. + return prepared with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: mgr._periodic_evict(batch) @@ -594,11 +578,11 @@ def test_union_rejects_unnormalized_scores(self): with pytest.raises(ValueError, match="normalize_scores=True"): _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) - def test_stage_rejects_int32_overflowing_round_starts(self): + def test_execute_rejects_int32_overflowing_round_starts(self): # Round starts past the int32 metadata range fail loudly (in the host # metadata build) before any GPU work is enqueued. from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - stage_eviction_cohort, + execute_eviction_round, ) device = torch.device("cuda", torch.cuda.current_device()) @@ -607,9 +591,10 @@ def test_stage_rejects_int32_overflowing_round_starts(self): manager = _make_staging_manager( torch.zeros(1, 2, 2, 12, dtype=torch.int32), gather, torch.cuda.Stream(device=device) ) + prepared = [_make_prepared_item(request_id=7, seq_len=64, round_start=2**31)] with pytest.raises((RuntimeError, OverflowError, ValueError)): - stage_eviction_cohort(staging, manager, [7], [2**31], [0], [64]) + execute_eviction_round(staging, manager, None, prepared, normalize_scores=True) assert gather.call_count == 0 def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): @@ -617,14 +602,14 @@ def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): waits for the previous cohort's consumers. Both persistent V2 host inputs (the block-offset table and the - index-mapper slot assignment) may mutate as soon as ``stage`` returns; + index-mapper slot assignment) may mutate as soon as staging returns; the staged device tables must reflect the values at staging time. A subsequent bulk copy must also wait until the previous round's - consumers (recorded by ``mark_page_tables_consumed``) are done. + consumers (ordered by ``execute_eviction_round``'s completion event) + are done. """ from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _stage_block_offsets, - mark_page_tables_consumed, ) device = torch.device("cuda", torch.cuda.current_device()) @@ -698,16 +683,18 @@ def stage_once(): 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 delayed consumer read (snapshot) queued before - # ``mark_page_tables_consumed`` must complete before the next bulk + # Round 3: a delayed consumer read (snapshot) queued before the + # round's completion ordering must complete before the next bulk # copy overwrites the device tables. selected_slot[0] = 0 manager_stream.synchronize() snapshot = torch.empty_like(staging.block_offsets_device) torch.cuda._sleep(20_000_000) snapshot.copy_(staging.block_offsets_device) - staging.page_tables_active = True - mark_page_tables_consumed(staging, manager_stream) + # ``execute_eviction_round``'s completion ordering: one event records + # the consumers and the manager stream waits on it. + staging.bulk_consume_done.record(torch.cuda.current_stream(device)) + manager_stream.wait_event(staging.bulk_consume_done) stage_once() current_stream.synchronize() @@ -730,7 +717,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): with _mocked_eviction_internals(manager) as internals: manager._evict_requests( [ - _prepared_eviction( + _make_prepared_item( first, request_id=7, seq_len=8, @@ -738,7 +725,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): expected_keep_count=7, protected_tail=2, ), - _prepared_eviction( + _make_prepared_item( second, request_id=8, seq_len=10, @@ -750,19 +737,32 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): 2, ) - # One batched staging call carries the whole cohort; budget=4 moves - # are keep + tail = [6, 7], padded rows repeat the final offset. - args = internals.stage.call_args + # One round-executor call carries the whole cohort (with the target + # and draft managers it orders after the compact launches). + args = internals.execute.call_args assert args.args[0] is internals.buffers assert args.args[1] is manager.kv_cache_manager - assert args.args[2:6] == ([7, 8], [8, 10], [3, 5], [8, 10]) - assert args.kwargs["draft_manager"] is None - assert args.kwargs["dense_move_offsets"] == [0, 6, 13, 13, 13, 13, 13, 13, 13] - assert args.kwargs["swa_move_offsets"] is None - assert args.kwargs["draft_move_offsets"] is None - internals.consumed.assert_called_once_with( - internals.buffers, manager.kv_cache_manager._stream + assert args.args[2] is None + prepared = args.args[3] + assert [item["request_id"] for item in prepared] == [7, 8] + assert [item["round_start"] for item in prepared] == [8, 10] + assert [item["prompt_len"] for item in prepared] == [3, 5] + assert [item["seq_len"] for item in prepared] == [8, 10] + assert args.kwargs == {"normalize_scores": True} + + # The derived move offsets stage keep + tail moves per request + # (budget=4 -> [6, 7]); padded rows repeat the final offset. + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _cohort_move_offsets, + ) + + offsets_buffers = SimpleNamespace( + max_requests=8, keep_count=4, compaction={"has_swa": False} ) + dense, swa, draft = _cohort_move_offsets(offsets_buffers, prepared, None) + assert dense == [0, 6, 13, 13, 13, 13, 13, 13, 13] + assert swa is None + assert draft is None @requires_sm100 @pytest.mark.parametrize("request_count", [1, 8]) @@ -835,21 +835,41 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) - def stage_round(): - bufs.round_starts_device.copy_(round_device) - bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens) - bufs.token_starts_device.fill_(prompt_len) - bufs.block_offsets_device[:, :, :, :page_count].copy_( - _encode_block_offsets(page_ids_3d) + # Rounds stage through the production executor: the gather double + # writes each layer's K page ids and the bulk copy encodes K/V rows. + def gather_k_block_offsets(host_table, source, request_ids, num_blocks): + assert request_ids == list(range(request_count)) + source[..., 0, :].zero_() + source[:, :request_count, 0, :page_count].copy_( + page_ids_3d[:, :request_count].to(torch.int32).cpu() ) + manager = _make_staging_manager( + torch.zeros(num_layers, max_requests, 2, 8, dtype=torch.int32), + gather_k_block_offsets, + torch.cuda.Stream(device=device), + num_slots=num_layers, + ) + + def prepared_cohort(): + return [ + _make_prepared_item( + request_id=request, + seq_len=int(valid_seq_lens[request]), + round_start=int(round_device[request]), + prompt_len=prompt_len, + ) + for request in range(request_count) + ] + score_sentinel = -12345.0 - stage_round() bufs.score_output.fill_(score_sentinel) # The compact stage is stubbed to a no-op: this test owns the score # buffers only, never a staged move decision. with mock.patch.object(module, "compact"): - module.run_eviction_round(bufs, normalize_scores=False) + module.execute_eviction_round( + bufs, manager, None, prepared_cohort(), normalize_scores=False + ) fixed = bufs.score_output.clone() assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] @@ -885,11 +905,12 @@ def stage_round(): ) ) expected_second_widths = valid_seq_lens - prompt_len - stage_round() bufs.score_output.fill_(score_sentinel) bufs.valid_widths.fill_(-1) with mock.patch.object(module, "compact"): - module.run_eviction_round(bufs, normalize_scores=False) + module.execute_eviction_round( + bufs, manager, None, prepared_cohort(), normalize_scores=False + ) second_launch = bufs.score_output.clone() assert torch.equal(bufs.valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index 1f56c83223e5..41d0cd3ed824 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -9,7 +9,9 @@ from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_cute_buffers as _make_cute_buffers +from conftest import make_prepared_item as _make_prepared_item from conftest import make_ramp_pools as _make_ramp_pools +from conftest import make_staging_manager as _make_staging_manager from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails @@ -446,7 +448,7 @@ 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") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - run_eviction_round, + execute_eviction_round, ) device = torch.device("cuda", torch.cuda.current_device()) @@ -512,18 +514,25 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], ) assert bufs.compaction["has_swa"] is False - bufs.block_offsets_device.zero_() - bufs.block_offsets_device[..., :2].copy_(_encode_block_offsets(torch.stack(page_tables))) - bufs.round_starts_device.fill_(0) - bufs.valid_seq_lens_device.fill_(seq_len) - bufs.token_starts_device.fill_(0) - # Stage this round's move offsets into the buffers' OWN metadata row: - # the pack launch and the native moves consume the buffers' own - # contract (keep_count moves per request, no protected tail). - bufs.request_metadata_device[3, :2].copy_( - torch.tensor([0, keep_count], dtype=torch.int32, device=device) + + # 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, ) - run_eviction_round(bufs, normalize_scores=False) + prepared = [_make_prepared_item(request_id=7, seq_len=seq_len, round_start=0)] + execute_eviction_round(bufs, manager, None, prepared, normalize_scores=False) assert torch.equal(bufs.keep, expected_keep) torch.cuda.synchronize(device) @@ -548,9 +557,7 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): import tensorrt_llm import tensorrt_llm.bindings from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - mark_page_tables_consumed, - run_eviction_round, - stage_eviction_cohort, + execute_eviction_round, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -673,20 +680,21 @@ def expected_keep() -> torch.Tensor: def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) - stage_eviction_cohort( - bufs, - manager, - [request_id], - [0], - [prompt_len], - [seq_len], - dense_move_offsets=[0, keep_count + protected_tail], - ) - # THE union path (fused pipeline). Z-normalization is monotonic - # per row, so the raw-score keep set is unchanged. - run_eviction_round(bufs, normalize_scores=True) + prepared = [ + _make_prepared_item( + request_id=request_id, + seq_len=seq_len, + round_start=0, + prompt_len=prompt_len, + protected_tail=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. + execute_eviction_round(bufs, manager, None, prepared, normalize_scores=True) selected = bufs.keep[0].clone().to(torch.long) - mark_page_tables_consumed(bufs, manager._stream) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) after = snapshot(compacted_capacity) From 21add9ccd7d2be25f76816e49eedfee2449fa7de Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 02:02:33 -0700 Subject: [PATCH 118/178] [None][perf] Persist the per-head gather index; refresh only its base in place (knife 18 item 20) Signed-off-by: tianruih --- .../triattention/triattention.py | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 627f2990b24c..369009d969a0 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -265,6 +265,22 @@ def init_eviction_buffers( ) # Compile the mode's SM100 CuTe entries; no other score path, no fallback. union = eviction_mode == "union" + # Persistent gather index (per-head modes): per round only the + # token-start base is re-added in place; the expanded view is fixed. + bufs.gather_index_base = None + bufs.gather_index = None + if not union: + num_kv_heads_early = int(layer_pools[dense_layers[0]].shape[2]) + bufs.gather_index_base = torch.empty( + (max_requests, 1, 1, 1, decode_width), dtype=torch.int64, device=device + ) + bufs.gather_index = bufs.gather_index_base.expand( + max_requests, + len(dense_layers), + num_kv_heads_early, + num_q_heads // num_kv_heads_early, + decode_width, + ) bufs.union_rows = None if union: # Bucket-wide rows; consumers mask by the per-request widths. @@ -640,13 +656,13 @@ def execute_eviction_round( )[:, :group_size] .permute(2, 3, 0, 1, 4) ) - columns = ( - bufs.token_starts_device[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) - + bufs.gather_columns - ) - columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + torch.add( + bufs.token_starts_device[:request_count].view(-1, 1, 1, 1, 1), + bufs.gather_columns, + out=bufs.gather_index_base[:request_count], ) + bufs.gather_index_base[:request_count].clamp_(max=bufs.bucket_seq_len - 1) + columns = bufs.gather_index[:request_count] torch.gather( source, 4, From 4de514a329eb32536d6c740090506dd57b114548 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 02:09:46 -0700 Subject: [PATCH 119/178] [None][perf] Settle reads per-request prompt offsets; drop the per-round row expansion (knife 18 item 21) Signed-off-by: tianruih --- .../triattention/triattention.py | 12 +--------- .../triattention/triattention_kernels.py | 5 ++-- .../test_triattention_draft_cocompaction.py | 1 - .../test_triattention_fused_settle_pack.py | 24 +++++++++++-------- .../test_triattention_selection_compaction.py | 8 +------ 5 files changed, 19 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 369009d969a0..d63fe830f8ad 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -327,7 +327,6 @@ def init_eviction_buffers( bufs.prompt_offsets = bufs.token_starts_device if union: bufs.selection_rows_per_request = 1 - bufs.row_prompt_offsets = bufs.prompt_offsets bufs.combined = torch.empty( (max_requests, decode_width), dtype=torch.float32, device=device ) @@ -359,9 +358,6 @@ def init_eviction_buffers( f"scores {score_rect}, selection {selection_rect}" ) bufs.selection_rows_per_request = selection_rows - bufs.row_prompt_offsets = torch.zeros( - (max_requests * selection_rows,), dtype=torch.int32, device=device - ) # [request, layer, head, token] layout read by the reduce kernels. bufs.score_output = torch.empty( max_requests, @@ -441,7 +437,7 @@ def init_eviction_buffers( bufs.settle_args = ( bufs.selection_scores_rows, bufs.selection_row_lengths, - bufs.row_prompt_offsets, + bufs.prompt_offsets, bufs.provisional_rows, bufs.keep_rows, ) @@ -605,12 +601,6 @@ def execute_eviction_round( # Guards the pinned metadata until the asynchronous copies complete. bufs.copy_done.record(stream) bufs.copy_pending = True - # Per-head modes re-expand the prompt lengths into their row-major view. - if bufs.row_prompt_offsets is not bufs.prompt_offsets: - bufs.row_prompt_offsets.view(bufs.max_requests, bufs.selection_rows_per_request).copy_( - bufs.prompt_offsets.unsqueeze(1).expand(-1, bufs.selection_rows_per_request) - ) - request_count = bufs.max_requests union = bufs.eviction_mode == "union" try: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index a4f8fcbc5e40..b139950bac45 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -281,8 +281,9 @@ def _settle_ties_kernel( row_output = output_indices + row * KEEP_COUNT row_scores = scores + row * WIDTH row_selected = provisional_indices + row * KEEP_COUNT - # Rebases the decode-relative ordinals to absolute positions. - prompt_len = tl.load(prompt_offsets + row) + # Rebases the decode-relative ordinals to absolute positions (per request: + # every selection row of a request shares its pinned prompt length). + prompt_len = tl.load(prompt_offsets + request) threshold = float("inf") for start in tl.static_range(0, KEEP_COUNT, BLOCK): 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 index b0005d16e517..1d0e06f89709 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -198,7 +198,6 @@ def test_execute_eviction_round_orders_both_manager_streams(): phase_num_freqs=1, phase_f_block=1, prompt_offsets=prompt_offsets, - row_prompt_offsets=prompt_offsets, round_starts_device=None, valid_seq_lens_device=None, token_starts_device=None, 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 index 585dd12076d3..07d74000030a 100644 --- 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 @@ -112,9 +112,10 @@ def _staged_offsets(counts, device): return torch.tensor(offsets, dtype=torch.int32, device=device) -def _make_settle_inputs(rows_total, width, keep_count, seed, device): +def _make_settle_inputs(request_count, selection_rows, width, keep_count, seed, device): """One seeded settle problem: heavily tied scores, ragged rows, a - top-k stand-in, and per-row prompt rebases.""" + top-k stand-in, and per-request prompt rebases.""" + 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( @@ -127,8 +128,9 @@ def _make_settle_inputs(rows_total, width, keep_count, seed, device): dtype=torch.int32, device=device, ) - row_prompt_offsets = torch.tensor( - [3 * (row % 3) 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. @@ -136,7 +138,7 @@ def _make_settle_inputs(rows_total, width, keep_count, seed, device): 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, row_prompt_offsets, provisional + return scores, row_lengths, prompt_offsets, provisional @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) @@ -148,9 +150,10 @@ def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): rows_total = request_count * selection_rows for seed in range(5): - scores, row_lengths, row_prompt_offsets, provisional = _make_settle_inputs( - rows_total, width, keep_count, seed, device + 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( @@ -165,7 +168,7 @@ def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): _settle_ties_kernel[(request_count, selection_rows)]( scores, row_lengths, - row_prompt_offsets, + prompt_offsets, provisional, output_actual, WIDTH=width, @@ -207,9 +210,10 @@ def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, w swa_offsets = _staged_offsets(swa_counts, device) for seed in range(5): - scores, row_lengths, row_prompt_offsets, provisional = _make_settle_inputs( - rows_total, width, keep_count, seed, device + 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) # Pre-settled decision rows straight from the settle oracle: short # rows keep stale garbage past their emitted count, which the pack # must forward verbatim. 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 index 41d0cd3ed824..43c9928fa1a9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -66,7 +66,6 @@ def _make_selection_buffers( bufs.prompt_offsets = torch.zeros(max_requests, dtype=torch.int32, device=device) if eviction_mode == "union": bufs.selection_rows_per_request = 1 - bufs.row_prompt_offsets = bufs.prompt_offsets bufs.combined = torch.empty((max_requests, width), dtype=torch.float32, device=device) # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. @@ -81,9 +80,6 @@ def _make_selection_buffers( else: selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads bufs.selection_rows_per_request = selection_rows - bufs.row_prompt_offsets = torch.zeros( - max_requests * selection_rows, dtype=torch.int32, device=device - ) bufs.row_mean = torch.empty( max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device ) @@ -112,7 +108,7 @@ def _make_selection_buffers( bufs.settle_args = ( bufs.selection_scores_rows, bufs.selection_row_lengths, - bufs.row_prompt_offsets, + bufs.prompt_offsets, bufs.provisional_rows, bufs.keep_rows, ) @@ -258,11 +254,9 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid max_requests=request_count, ) bufs.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) - # The union row-major view aliases the per-request buffer. bufs.prompt_offsets[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) - assert bufs.row_prompt_offsets is bufs.prompt_offsets bufs.combined.copy_(scores.amax(dim=1)) settle_top_tokens(bufs) actual = bufs.keep.cpu() From 5da79899dda0f4a2a505cd610045580ffb76fdfb Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 02:12:21 -0700 Subject: [PATCH 120/178] [None][perf] Launch eviction rounds at the active cohort size (knife 18 item 17) Signed-off-by: tianruih --- .../triattention/triattention.py | 18 ++++++++++-------- .../test_triattention_selection_compaction.py | 5 ++--- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d63fe830f8ad..ea98a424635e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -388,7 +388,6 @@ def init_eviction_buffers( bufs.top_indices_i32.zero_() # ---- compaction contract + decision-materialization prebinds ------------ - bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) per_layer = eviction_mode == "per_layer_perhead" draft_contract = None if draft is not None: @@ -517,18 +516,21 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: return dense, swa, draft -def settle_top_tokens(bufs: SimpleNamespace) -> None: +def settle_top_tokens(bufs: SimpleNamespace, request_count: int) -> None: """Pick the top-k and settle ties into the kept-ordinal decision rows (the compaction contract packs them into move sources).""" + rows = request_count * bufs.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( - bufs.selection_scores_rows, - bufs.selection_row_lengths, - bufs.provisional_rows, + bufs.selection_scores_rows[:rows], + bufs.selection_row_lengths[:rows], + bufs.provisional_rows[:rows], bufs.keep_count, 1, ) - _settle_ties_kernel[bufs.settle_grid](*bufs.settle_args, **bufs.settle_kwargs) + _settle_ties_kernel[(request_count, bufs.selection_rows_per_request)]( + *bufs.settle_args, **bufs.settle_kwargs + ) def execute_eviction_round( @@ -601,7 +603,7 @@ def execute_eviction_round( # Guards the pinned metadata until the asynchronous copies complete. bufs.copy_done.record(stream) bufs.copy_pending = True - request_count = bufs.max_requests + request_count = len(prepared) union = bufs.eviction_mode == "union" try: with nvtx_range("triattention.score", color="blue"): @@ -679,7 +681,7 @@ def execute_eviction_round( per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - settle_top_tokens(bufs) + settle_top_tokens(bufs, request_count) with nvtx_range("triattention.compact", color="purple"): compact(bufs.compaction, request_count) finally: 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 index 43c9928fa1a9..05d92039961c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -102,7 +102,6 @@ def _make_selection_buffers( bufs.selection_row_lengths = bufs.row_seq_lens.view(-1) bufs.provisional_rows = bufs.top_indices_i32.view(-1, keep_count) bufs.keep_rows = bufs.keep.view(-1, keep_count) - bufs.settle_grid = (max_requests, bufs.selection_rows_per_request) # The launch args mirror the product's ``bufs.settle_args``/ # ``bufs.settle_kwargs`` order and keys exactly. bufs.settle_args = ( @@ -136,7 +135,7 @@ def _select_per_head(bufs, scores, *, normalize_scores): per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - settle_top_tokens(bufs) + settle_top_tokens(bufs, bufs.max_requests) def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: @@ -258,7 +257,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) bufs.combined.copy_(scores.amax(dim=1)) - settle_top_tokens(bufs) + settle_top_tokens(bufs, bufs.max_requests) actual = bufs.keep.cpu() combined = scores.amax(dim=1).cpu() From 470acdee91aa816ad1d6af5ac0fa5d2a03b4c64a Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 02:24:31 -0700 Subject: [PATCH 121/178] [None][chore] Return one named plane bundle; no anonymous tuple unpacks (knife 25d) Signed-off-by: tianruih --- .../triattention/triattention.py | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index ea98a424635e..d73332750974 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -82,7 +82,7 @@ def _allocate_page_table_plane( token_capacity: int, max_requests: int, device: torch.device, -) -> Tuple[Dict[int, int], int, torch.Tensor, torch.Tensor]: +) -> Dict[str, object]: representative_slots = { representative: int(key[1]) for representative, key in zip(page_representatives, page_table_keys) @@ -93,7 +93,7 @@ def _allocate_page_table_plane( plane_shape = (num_page_table_slots, max_requests, 2, copy_block_count) host = torch.empty(plane_shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) dev = torch.empty(plane_shape, dtype=torch.int32, device=device) - return representative_slots, copy_block_count, host, dev + return dict(slots=representative_slots, copy_block_count=copy_block_count, host=host, dev=dev) def init_eviction_buffers( @@ -153,12 +153,7 @@ def init_eviction_buffers( bufs.page_table_token_capacity = page_table_token_capacity # ---- staged page-table planes (target, plus the co-compressed draft) --- - ( - bufs.representative_slots, - bufs.copy_block_count, - bufs._bulk_offsets_src, - bufs.block_offsets_device, - ) = _allocate_page_table_plane( + plane = _allocate_page_table_plane( layer_pools, page_representatives, page_table_keys, @@ -167,6 +162,10 @@ def init_eviction_buffers( max_requests, device, ) + bufs.representative_slots = plane["slots"] + bufs.copy_block_count = plane["copy_block_count"] + bufs._bulk_offsets_src = plane["host"] + bufs.block_offsets_device = plane["dev"] # The draft is never scored: these offsets feed only the draft compacts. bufs.draft_block_offsets_device = None bufs._draft_bulk_offsets_src = None @@ -175,12 +174,7 @@ def init_eviction_buffers( if draft is not None: draft_layout = draft["layout"] draft_representatives = list(draft_layout["pool_representatives"]) - ( - draft_page_slots, - bufs.draft_copy_block_count, - bufs._draft_bulk_offsets_src, - bufs.draft_block_offsets_device, - ) = _allocate_page_table_plane( + draft_plane = _allocate_page_table_plane( draft_layout["layer_pools"], draft_representatives, [draft_layout["layer_pool_keys"][layer] for layer in draft_representatives], @@ -189,6 +183,10 @@ def init_eviction_buffers( max_requests, device, ) + draft_page_slots = draft_plane["slots"] + bufs.draft_copy_block_count = draft_plane["copy_block_count"] + bufs._draft_bulk_offsets_src = draft_plane["host"] + bufs.draft_block_offsets_device = draft_plane["dev"] # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- bufs.request_metadata_host = torch.empty( From da7065ea3f6d88d6c0e17e1c422711940bf4c8a4 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 03:14:08 -0700 Subject: [PATCH 122/178] [None][chore] Parameter-hygiene sweep: cluster, rename, and close orphans (knife 26) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 150 ++++++++---------- .../triattention/triattention.py | 133 +++++++--------- .../triattention_cute_score_fused.py | 40 ++--- .../triattention/triattention_kernels.py | 92 +++++------ .../_torch/kv_cache_compression/conftest.py | 12 +- .../test_triattention_cute_union_fusion.py | 5 +- .../test_triattention_draft_cocompaction.py | 12 +- .../test_triattention_fused_settle_pack.py | 22 +-- .../test_triattention_pipeline.py | 30 ++-- .../test_triattention_selection_compaction.py | 42 ++--- 10 files changed, 216 insertions(+), 322 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index d29027f34bee..e79debc2c8b5 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -31,46 +31,42 @@ import triton import triton.language as tl -# Pack launch shape, shared by every family's pack launch. -PACK_BLOCK = 256 -PACK_NUM_WARPS = 4 - @triton.jit def _pack_move_sources_kernel( kept_ordinal_rows, valid_seq_lens, - dense_offsets, - dense_indices, - swa_offsets, - swa_indices, + dense_move_offsets, + dense_move_indices, + swa_move_offsets, + swa_move_indices, KEEP_COUNT: tl.constexpr, DECISION_ROWS: tl.constexpr, - DENSE_TOTAL: tl.constexpr, - SWA_TOTAL: tl.constexpr, MOVE_CAPACITY: tl.constexpr, NUM_KV_HEADS: tl.constexpr, - SWA_WINDOW: tl.constexpr, - BROADCAST: tl.constexpr, PER_LAYER: tl.constexpr, - HAS_SWA: tl.constexpr, - BLOCK: tl.constexpr, + DENSE_TOTAL: tl.constexpr, + SWA_TOTAL: tl.constexpr, + SWA_WINDOW: tl.constexpr, + BLOCK: tl.constexpr = 256, ): """Pack one decision row into one family's move sources (increasing kept ordinals; C++ in-place copy contract): dense rows forward the row content verbatim for the first KEEP_COUNT moves, then append the protected tail; SWA rows write latest-window ordinals once per KV head.""" + BROADCAST: tl.constexpr = DECISION_ROWS == 1 + HAS_SWA: tl.constexpr = SWA_TOTAL > 0 request = tl.program_id(0) decision_row = tl.program_id(1) row = request * DECISION_ROWS + decision_row kept_row = kept_ordinal_rows + row * KEEP_COUNT - dense_begin = tl.load(dense_offsets + request) - dense_end = tl.load(dense_offsets + request + 1) + dense_begin = tl.load(dense_move_offsets + request) + dense_end = tl.load(dense_move_offsets + request + 1) dense_count = dense_end - dense_begin valid_len = tl.load(valid_seq_lens + request) if HAS_SWA: - swa_begin = tl.load(swa_offsets + request) - swa_end = tl.load(swa_offsets + request + 1) + swa_begin = tl.load(swa_move_offsets + request) + swa_end = tl.load(swa_move_offsets + request + 1) swa_count = swa_end - swa_begin for move_start in tl.static_range(0, MOVE_CAPACITY, BLOCK): move = move_start + tl.arange(0, BLOCK) @@ -84,19 +80,19 @@ def _pack_move_sources_kernel( # The one decision row per request feeds every KV head's packed row. for head in tl.static_range(0, NUM_KV_HEADS): tl.store( - dense_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, + dense_move_indices + head * DENSE_TOTAL + dense_begin.to(tl.int64) + move, dense_source, mask=move < dense_count, ) else: dense_output = decision_row.to(tl.int64) * DENSE_TOTAL + dense_begin.to(tl.int64) + move - tl.store(dense_indices + dense_output, dense_source, mask=move < dense_count) + tl.store(dense_move_indices + dense_output, dense_source, mask=move < dense_count) if HAS_SWA: swa_source = valid_len - SWA_WINDOW + move if BROADCAST: for head in tl.static_range(0, NUM_KV_HEADS): tl.store( - swa_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, + swa_move_indices + head * SWA_TOTAL + swa_begin.to(tl.int64) + move, swa_source, mask=move < swa_count, ) @@ -108,7 +104,7 @@ def _pack_move_sources_kernel( head = decision_row % NUM_KV_HEADS swa_output = head.to(tl.int64) * SWA_TOTAL + swa_begin.to(tl.int64) + move tl.store( - swa_indices + swa_output, + swa_move_indices + swa_output, swa_source, mask=swa_mask, ) @@ -117,21 +113,21 @@ def _pack_move_sources_kernel( def _make_move_indices( index_prefix: Tuple[int, ...], moves_per_request: int, - request_count: int, + max_requests: int, device: torch.device, ) -> torch.Tensor: return torch.empty( - (*index_prefix, moves_per_request * request_count), dtype=torch.int32, device=device + (*index_prefix, moves_per_request * max_requests), dtype=torch.int32, device=device ) def _compact_groups( entries: List[Tuple[int, torch.Tensor, torch.Tensor]], pool_keys: Tuple[object, ...], - device: torch.device, per_layer_slots: Optional[Dict[int, int]] = None, ) -> Tuple[Dict[str, object], ...]: """Batch layers into one ``sparse_kv_cache_compact_layers`` launch per uniform V2 pool.""" + device = entries[0][1].device grouped = OrderedDict() for layer, pool, page_table in entries: key = ( @@ -172,8 +168,8 @@ def _compact_groups( def _launch_tuples( groups: Tuple[Dict[str, object], ...], - source: torch.Tensor, - offsets: torch.Tensor, + move_indices: torch.Tensor, + move_offsets: torch.Tensor, destination_bases: torch.Tensor, ) -> Tuple[tuple, ...]: return tuple( @@ -181,8 +177,8 @@ def _launch_tuples( group["pools"], group["pool_pointers"], group["page_table"], - source, - offsets, + move_indices, + move_offsets, destination_bases, group["source_layer_indices"], ) @@ -203,15 +199,15 @@ def init_compaction_buffers( contract). ``target`` carries the resolved dense/SWA grouping inputs from the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows) plus the decision inputs :func:`compact` packs each round: - ``kept_ordinal_rows`` (``request_capacity * decision_rows`` rows of - ``decode_keep_count`` int32 kept ordinals, forwarded verbatim), the + ``kept_ordinal_rows`` (``max_requests * decision_rows`` rows of + ``keep_count`` int32 kept ordinals, forwarded verbatim), the per-request ``decision_rows`` count (1 = one shared row broadcast over every KV head), and the staged per-request ``valid_seq_lens`` the protected tail rides after. ``draft`` is one all-or-none resolved branch (its dense-only moves broadcast the one shared decision row over the draft's own KV heads); ``capacities`` the request/keep/tail capacity - numbers. The returned contract exposes the agreed move-source buffers and - geometry constants; its launch tuples are private to :func:`compact`. + numbers. The returned contract exposes the SWA/draft geometry the caller + stages against; its launch tuples are private to :func:`compact`. """ layer_pools = target["layer_pools"] dense_layers = tuple(int(layer) for layer in target["dense_layers"]) @@ -220,7 +216,7 @@ def init_compaction_buffers( kv_block_offsets = target["kv_block_offsets"] page_table_slots = target["page_table_slots"] layer_group_representative = target["layer_group_representative"] - prompt_offsets = target["prompt_offsets"] + token_starts = target["token_starts"] dense_move_offsets = target["dense_move_offsets"] swa_move_offsets = target["swa_move_offsets"] swa_window = target["swa_window"] @@ -230,8 +226,8 @@ def init_compaction_buffers( valid_seq_lens = target["valid_seq_lens"] device = layer_pools[dense_layers[0]].device - request_count = int(capacities["request_capacity"]) - decode_keep_count = int(capacities["decode_keep_count"]) + max_requests = int(capacities["max_requests"]) + keep_count = int(capacities["keep_count"]) protected_tail_capacity = int(capacities["protected_tail_capacity"]) # Pool shape [pages, K/V, heads, tokens, dim]. @@ -239,17 +235,15 @@ def init_compaction_buffers( dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer_sources else (num_kv_heads,) dense_move_indices = _make_move_indices( dense_index_prefix, - decode_keep_count + protected_tail_capacity, - request_count, + keep_count + protected_tail_capacity, + max_requests, device, ) dense_entries = [ ( layer, layer_pools[layer], - kv_block_offsets[ - page_table_slots[layer_group_representative[layer]], :request_count, 0 - ], + kv_block_offsets[page_table_slots[layer_group_representative[layer]], :max_requests, 0], ) for layer in dense_layers ] @@ -262,11 +256,11 @@ def init_compaction_buffers( swa_window = 0 else: swa_window = int(swa_window) - swa_destination_bases = torch.empty_like(prompt_offsets) + swa_destination_bases = torch.empty_like(token_starts) swa_move_indices = _make_move_indices( (num_kv_heads,), swa_window + protected_tail_capacity, - request_count, + max_requests, device, ) # SWA layers are staged as their own page-table representatives. @@ -274,7 +268,7 @@ def init_compaction_buffers( ( layer, layer_pools[layer], - kv_block_offsets[page_table_slots[layer], :request_count, 0], + kv_block_offsets[page_table_slots[layer], :max_requests, 0], ) for layer in swa_layers ] @@ -285,22 +279,22 @@ def init_compaction_buffers( has_swa = swa_move_indices is not None swa_total = int(swa_move_indices.shape[-1]) if has_swa else 0 # Widest per-request move count any staged offsets may express. - move_capacity = decode_keep_count + protected_tail_capacity + move_capacity = keep_count + protected_tail_capacity if has_swa: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) target_launches = list( _launch_tuples( - _compact_groups(dense_entries, layer_pool_keys, device, dense_slots), + _compact_groups(dense_entries, layer_pool_keys, dense_slots), dense_move_indices, dense_move_offsets, - prompt_offsets, + token_starts, ) ) if swa_layers: target_launches.extend( _launch_tuples( - _compact_groups(swa_entries, layer_pool_keys, device), + _compact_groups(swa_entries, layer_pool_keys), swa_move_indices, swa_move_offsets, swa_destination_bases, @@ -318,18 +312,14 @@ def init_compaction_buffers( swa_move_indices, ), dict( - KEEP_COUNT=decode_keep_count, + KEEP_COUNT=keep_count, DECISION_ROWS=decision_rows, - DENSE_TOTAL=int(dense_move_indices.shape[-1]), - SWA_TOTAL=swa_total, MOVE_CAPACITY=move_capacity, NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window, - BROADCAST=decision_rows == 1, PER_LAYER=per_layer_sources, - HAS_SWA=has_swa, - BLOCK=PACK_BLOCK, - num_warps=PACK_NUM_WARPS, + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=swa_total, + SWA_WINDOW=swa_window, ), ) @@ -343,14 +333,14 @@ def init_compaction_buffers( f"got {decision_rows} decision rows" ) draft_layer_pools = draft["layer_pools"] - draft_layers = tuple(int(layer) for layer in draft["layers"]) + draft_dense_layers = tuple(int(layer) for layer in draft["dense_layers"]) draft_tail = int(draft["protected_tail_capacity"]) # Own launch groups: the draft may use a different KV-head count. - draft_num_kv_heads = int(draft_layer_pools[draft_layers[0]].shape[2]) + draft_num_kv_heads = int(draft_layer_pools[draft_dense_layers[0]].shape[2]) draft_move_indices = _make_move_indices( (draft_num_kv_heads,), - decode_keep_count + draft_tail, - request_count, + keep_count + draft_tail, + max_requests, device, ) draft_entries = [ @@ -359,58 +349,49 @@ def init_compaction_buffers( draft_layer_pools[layer], draft["kv_block_offsets"][ draft["page_table_slots"][draft["layer_group_representative"][layer]], - :request_count, + :max_requests, 0, ], ) - for layer in draft_layers + for layer in draft_dense_layers ] draft_launches = _launch_tuples( - _compact_groups(draft_entries, tuple(draft["layer_pool_keys"]), device), + _compact_groups(draft_entries, tuple(draft["layer_pool_keys"])), draft_move_indices, - draft["move_offsets"], - prompt_offsets, + draft["dense_move_offsets"], + token_starts, ) draft_pack_launch = ( 1, ( kept_ordinal_rows, valid_seq_lens, - draft["move_offsets"], + draft["dense_move_offsets"], draft_move_indices, None, None, ), dict( - KEEP_COUNT=decode_keep_count, + KEEP_COUNT=keep_count, DECISION_ROWS=1, + MOVE_CAPACITY=int(draft_move_indices.shape[-1]) // max_requests, + NUM_KV_HEADS=draft_num_kv_heads, + PER_LAYER=False, DENSE_TOTAL=int(draft_move_indices.shape[-1]), SWA_TOTAL=0, - MOVE_CAPACITY=int(draft_move_indices.shape[-1]) // request_count, - NUM_KV_HEADS=draft_num_kv_heads, SWA_WINDOW=0, - BROADCAST=True, - PER_LAYER=False, - HAS_SWA=False, - BLOCK=PACK_BLOCK, - num_warps=PACK_NUM_WARPS, ), ) return dict( - # Agreed decision buffers and geometry constants (the public interface). - dense_move_indices=dense_move_indices, - swa_move_indices=swa_move_indices, - draft_move_indices=draft_move_indices, - dense_total=int(dense_move_indices.shape[-1]), - swa_total=swa_total, - move_capacity=move_capacity, - num_kv_heads=num_kv_heads, - swa_window=swa_window, + # SWA/draft geometry the caller stages against (the public interface). has_swa=has_swa, + swa_window=swa_window, swa_destination_bases=swa_destination_bases, # Per-round SWA destination rebase delta. - swa_rebase_delta=decode_keep_count - swa_window, + swa_rebase_delta=keep_count - swa_window, + has_draft=draft is not None, + draft_move_indices=draft_move_indices, # Completion event: compact() records it after the last native launch. consume_done=torch.cuda.Event(), # Private launch tuples: only compact() interprets these. @@ -418,7 +399,6 @@ def init_compaction_buffers( draft_launches=draft_launches, target_pack_launch=target_pack_launch, draft_pack_launch=draft_pack_launch, - has_draft=draft is not None, ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d73332750974..2a340de9a91b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -41,8 +41,6 @@ from ..compaction import compact, init_compaction_buffers from .triattention_kernels import ( - SETTLE_BLOCK, - SETTLE_NUM_WARPS, _gather_mean_phase_kernel, _settle_ties_kernel, grow_mean_phase_table, @@ -81,12 +79,12 @@ def _allocate_page_table_plane( num_page_table_slots: int, token_capacity: int, max_requests: int, - device: torch.device, ) -> Dict[str, object]: representative_slots = { representative: int(key[1]) for representative, key in zip(page_representatives, page_table_keys) } + device = layer_pools[page_representatives[0]].device tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block copy_block_count = (page_count + 3) // 4 * 4 @@ -160,15 +158,14 @@ def init_eviction_buffers( num_page_table_slots, page_table_token_capacity, max_requests, - device, ) bufs.representative_slots = plane["slots"] bufs.copy_block_count = plane["copy_block_count"] - bufs._bulk_offsets_src = plane["host"] + bufs._block_offsets_host = plane["host"] bufs.block_offsets_device = plane["dev"] # The draft is never scored: these offsets feed only the draft compacts. bufs.draft_block_offsets_device = None - bufs._draft_bulk_offsets_src = None + bufs._draft_block_offsets_host = None bufs.draft_copy_block_count = 0 draft_page_slots: Dict[int, int] = {} if draft is not None: @@ -181,11 +178,10 @@ def init_eviction_buffers( int(draft_layout["manager"].num_pools), int(draft["page_table_token_capacity"]), max_requests, - device, ) draft_page_slots = draft_plane["slots"] bufs.draft_copy_block_count = draft_plane["copy_block_count"] - bufs._draft_bulk_offsets_src = draft_plane["host"] + bufs._draft_block_offsets_host = draft_plane["host"] bufs.draft_block_offsets_device = draft_plane["dev"] # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- @@ -228,15 +224,15 @@ def init_eviction_buffers( raise ValueError("scored layer index exceeds the calibrated layer extent") _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} page_table_slots = [bufs.representative_slots[_rep_of[layer]] for layer in dense_layers] - bufs.seg_req = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( + seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( bufs.num_layers ) - seg_layer = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( + seg_layer_id = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( max_requests ) block_offsets = bufs.block_offsets_device slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) - req_idx = bufs.seg_req.to(torch.int64) + req_idx = seg_req_id.to(torch.int64) slot_idx = slots_t.repeat(max_requests) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) @@ -279,25 +275,22 @@ def init_eviction_buffers( num_q_heads // num_kv_heads_early, decode_width, ) - bufs.union_rows = None + bufs.union_scores = None if union: # Bucket-wide rows; consumers mask by the per-request widths. - bufs.union_rows = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) + bufs.union_scores = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) try: bufs.runner = TriAttentionCuteScoreRunner( layer_pools=list(layer_pools), layer_indices=[int(layer) for layer in dense_layers], max_requests=max_requests, - num_layers=bufs.num_layers, seq_len=seq_len, num_q_heads=bufs.num_q_heads, - num_kv_heads=bufs.num_kv_heads, num_freqs=bufs.num_freqs, - tokens_per_block=bufs.tokens_per_block, page_ids=block_offsets.view(-1), seg_page_off=seg_page_off, - seg_req_id=bufs.seg_req, - seg_layer_id=seg_layer, + seg_req_id=seg_req_id, + seg_layer_id=seg_layer_id, # Pointer capture of the staged metadata rows. valid_seq_lens=bufs.valid_seq_lens_device, seg_out_offset=seg_out_offset, @@ -322,13 +315,12 @@ def init_eviction_buffers( # ---- selection buffers -------------------------------------------------- bufs.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) - bufs.prompt_offsets = bufs.token_starts_device if union: bufs.selection_rows_per_request = 1 bufs.combined = torch.empty( (max_requests, decode_width), dtype=torch.float32, device=device ) - bufs.final_indices = torch.empty( + bufs.provisional_indices = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=device ) # Kept decode ordinals only (prompt-length independent rows). @@ -336,10 +328,10 @@ def init_eviction_buffers( # Row-major views consumed by the top-k settle launch. bufs.selection_scores_rows = bufs.combined bufs.selection_row_lengths = bufs.valid_widths - bufs.provisional_rows = bufs.final_indices - bufs.keep_rows = bufs.keep + bufs.provisional_rows = bufs.provisional_indices + bufs.kept_ordinal_rows = bufs.keep # Padded rows still need in-range ordinals for the finalizer's gather. - bufs.final_indices.zero_() + bufs.provisional_indices.zero_() bufs.score_output = None else: selection_rows = ( @@ -367,23 +359,23 @@ def init_eviction_buffers( ) score_shape = (max_requests, bufs.num_layers, bufs.num_q_heads, 1) bufs.row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) - bufs.row_std = torch.empty_like(bufs.row_mean) + bufs.row_inv_std = torch.empty_like(bufs.row_mean) bufs.selection_scores = torch.empty( (max_requests, selection_rows, decode_width), dtype=torch.float32, device=device ) - bufs.row_seq_lens = torch.full( + bufs.selection_seq_lens = torch.full( (max_requests, selection_rows), decode_width, dtype=torch.int32, device=device ) selection_shape = (max_requests, selection_rows, keep_count) - bufs.top_indices_i32 = torch.empty(selection_shape, dtype=torch.int32, device=device) + bufs.provisional_indices = torch.empty(selection_shape, dtype=torch.int32, device=device) bufs.keep = torch.empty(selection_shape, dtype=torch.int32, device=device) bufs.selection_scores_rows = bufs.selection_scores.view( max_requests * selection_rows, decode_width ) - bufs.selection_row_lengths = bufs.row_seq_lens.view(-1) - bufs.provisional_rows = bufs.top_indices_i32.view(-1, keep_count) - bufs.keep_rows = bufs.keep.view(-1, keep_count) - bufs.top_indices_i32.zero_() + bufs.selection_row_lengths = bufs.selection_seq_lens.view(-1) + bufs.provisional_rows = bufs.provisional_indices.view(-1, keep_count) + bufs.kept_ordinal_rows = bufs.keep.view(-1, keep_count) + bufs.provisional_indices.zero_() # ---- compaction contract + decision-materialization prebinds ------------ per_layer = eviction_mode == "per_layer_perhead" @@ -392,12 +384,12 @@ def init_eviction_buffers( draft_layout = draft["layout"] draft_contract = dict( layer_pools=draft_layout["layer_pools"], - layers=list(draft_layout["dense_layers"]), + dense_layers=list(draft_layout["dense_layers"]), layer_group_representative=draft_layout["layer_group_representative"], layer_pool_keys=list(draft_layout["layer_pool_keys"]), kv_block_offsets=bufs.draft_block_offsets_device, page_table_slots=draft_page_slots, - move_offsets=draft_move_offsets_row, + dense_move_offsets=draft_move_offsets_row, protected_tail_capacity=int(draft["protected_tail_capacity"]), ) contract = init_compaction_buffers( @@ -410,19 +402,19 @@ def init_eviction_buffers( layer_pool_keys=list(layer_pool_keys), kv_block_offsets=bufs.block_offsets_device, page_table_slots=bufs.representative_slots, - prompt_offsets=bufs.token_starts_device, + token_starts=bufs.token_starts_device, # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, swa_move_offsets=swa_move_offsets_row, per_layer_sources=per_layer, # The decision rows the contract packs into move sources. - kept_ordinal_rows=bufs.keep_rows, + kept_ordinal_rows=bufs.kept_ordinal_rows, decision_rows=bufs.selection_rows_per_request, valid_seq_lens=bufs.valid_seq_lens_device, ), capacities=dict( - request_capacity=max_requests, - decode_keep_count=keep_count, + max_requests=max_requests, + keep_count=keep_count, protected_tail_capacity=int(protected_tail_capacity), ), draft=draft_contract, @@ -434,16 +426,14 @@ def init_eviction_buffers( bufs.settle_args = ( bufs.selection_scores_rows, bufs.selection_row_lengths, - bufs.prompt_offsets, + bufs.token_starts_device, bufs.provisional_rows, - bufs.keep_rows, + bufs.kept_ordinal_rows, ) bufs.settle_kwargs = dict( WIDTH=decode_width, KEEP_COUNT=keep_count, SELECTION_ROWS=bufs.selection_rows_per_request, - BLOCK=SETTLE_BLOCK, - num_warps=SETTLE_NUM_WARPS, ) # ---- round-ordering events ---------------------------------------------- @@ -459,9 +449,8 @@ def _stage_block_offsets( bufs: SimpleNamespace, manager: KVCacheManagerV2, request_ids: List[int], - current_stream: torch.cuda.Stream, - source: torch.Tensor, - destination: torch.Tensor, + host_block_offsets: torch.Tensor, + device_block_offsets: torch.Tensor, copy_block_count: int, ) -> None: """Gather the pinned snapshot before the async device copy: resize mutates the live host table.""" @@ -469,21 +458,21 @@ def _stage_block_offsets( bufs.copy_done.synchronize() manager.index_mapper.gather_k_block_offsets( manager.host_kv_cache_block_offsets, - source, + host_block_offsets, request_ids, copy_block_count, ) manager._stream.wait_event(bufs.copy_done) copy_batch_block_offsets_to_device( - source, - destination, + host_block_offsets, + device_block_offsets, bufs._bulk_copy_idx_src[: len(request_ids)], manager.index_scales, manager.kv_offset, manager._stream.cuda_stream, ) bufs.bulk_copy_done.record(manager._stream) - current_stream.wait_event(bufs.bulk_copy_done) + torch.cuda.current_stream(bufs.device).wait_event(bufs.bulk_copy_done) def _cohort_move_offsets( @@ -534,8 +523,8 @@ def settle_top_tokens(bufs: SimpleNamespace, request_count: int) -> None: def execute_eviction_round( bufs: SimpleNamespace, manager: KVCacheManagerV2, - draft_manager: Optional[KVCacheManagerV2], prepared: Sequence[Dict[str, object]], + draft_manager: Optional[KVCacheManagerV2] = None, *, normalize_scores: bool, ) -> None: @@ -580,8 +569,7 @@ def execute_eviction_round( bufs, manager, request_ids, - stream, - bufs._bulk_offsets_src, + bufs._block_offsets_host, bufs.block_offsets_device, bufs.copy_block_count, ) @@ -590,8 +578,7 @@ def execute_eviction_round( bufs, draft_manager, request_ids, - stream, - bufs._draft_bulk_offsets_src, + bufs._draft_block_offsets_host, bufs.draft_block_offsets_device, bufs.draft_copy_block_count, ) @@ -610,13 +597,13 @@ def execute_eviction_round( bufs.round_starts_device, bufs.phase["cos"], bufs.phase["sin"], - bufs.mean_cos, - bufs.mean_sin, + bufs.phase["rows"], bufs.valid_seq_lens_device, bufs.token_starts_device, + bufs.mean_cos, + bufs.mean_sin, bufs.valid_widths, bufs.swa_destination_bases, - bufs.phase["rows"], bufs.swa_rebase_delta, NUM_FREQS=bufs.phase_num_freqs, F_BLOCK=bufs.phase_f_block, @@ -625,11 +612,11 @@ def execute_eviction_round( ) if union: bufs.runner.launch_union_fusion( - request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_rows[:request_count] + request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_scores[:request_count] ) - columns = min(bufs.union_rows.shape[1], bufs.combined.shape[1]) + columns = min(bufs.union_scores.shape[1], bufs.combined.shape[1]) bufs.combined[:request_count, :columns].copy_( - bufs.union_rows[:request_count, :columns] + bufs.union_scores[:request_count, :columns] ) else: bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) @@ -671,11 +658,9 @@ def execute_eviction_round( bufs.score_output[:request_count], bufs.valid_widths, bufs.row_mean, - bufs.row_std, + bufs.row_inv_std, bufs.selection_scores, - bufs.row_seq_lens, - request_count, - num_kv_heads=bufs.num_kv_heads, + bufs.selection_seq_lens, per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) @@ -1025,13 +1010,12 @@ def _periodic_evict( if not prepared: return - num_layers = self._num_layers_from_manager() # Ungated NVTX: the due count in the message shows each round's size. with nvtx_range( f"triattention.evict_request_group reqs={len(prepared)}", color="purple", ): - compacted = self._evict_requests(prepared, num_layers) + compacted = self._evict_requests(prepared) self._resize_compacted_requests(compacted) def _resize_compacted_requests(self, prepared) -> None: @@ -1076,9 +1060,9 @@ def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: def _local_score_calibration( self, - num_layers: int, global_layers: List[int], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + num_layers = len(global_layers) if global_layers and max(global_layers) >= self._triattn_q_real.shape[0]: raise ValueError( f"TriAttention calibration has {self._triattn_q_real.shape[0]} layers, " @@ -1243,9 +1227,9 @@ def _runtime_kv_layout(self, num_layers: int) -> Dict[str, object]: manager, global_layers, dense_layers=dense_layers, + dense_storage_groups=self._dense_layer_pool_groups(dense_layers, global_layers), swa_layers=swa_layers, swa_window=swa_window, - dense_storage_groups=self._dense_layer_pool_groups(dense_layers, global_layers), what="", ) self._runtime_kv_layout_cache = layout @@ -1257,9 +1241,9 @@ def _build_runtime_kv_layout( global_layers: List[int], *, dense_layers: List[int], + dense_storage_groups: Optional[Dict[object, List[int]]], swa_layers: List[int], swa_window: Optional[int], - dense_storage_groups: Optional[Dict[object, List[int]]], what: str, ) -> Dict[str, object]: num_layers = len(global_layers) @@ -1328,9 +1312,9 @@ def _draft_runtime_kv_layout(self) -> Dict[str, object]: manager, global_layers, dense_layers=list(range(len(global_layers))), + dense_storage_groups=None, swa_layers=[], swa_window=None, - dense_storage_groups=None, what="draft ", ) self._draft_runtime_kv_layout_cache = layout @@ -1424,9 +1408,7 @@ def _buffers_for( "rows": 0, } grow_mean_phase_table(self._phase, max(int(seq_capacity), 1)) - q_real, q_imag, mlr_coef = self._local_score_calibration( - layout["num_layers"], layout["global_layers"] - ) + q_real, q_imag, mlr_coef = self._local_score_calibration(layout["global_layers"]) bufs = init_eviction_buffers( eviction_mode=self.eviction_mode, layout=layout, @@ -1452,7 +1434,7 @@ def _buffers_for( def _page_table_pool_keys( self, - representatives: List[int], + local_layers: List[int], global_layers: List[int], manager: Optional[KVCacheManagerV2] = None, ) -> List[object]: @@ -1463,7 +1445,7 @@ def _page_table_pool_keys( try: return [ ("pool", int(layer_to_pool[layer_offsets[global_layers[layer]]])) - for layer in representatives + for layer in local_layers ] except (IndexError, KeyError, TypeError, ValueError) as exc: raise RuntimeError("KVCacheManagerV2 exposes an invalid layer-to-pool mapping") from exc @@ -1484,18 +1466,17 @@ def _dense_layer_pool_groups( def _evict_requests( self, prepared: List[Dict[str, object]], - num_layers: int, ) -> List[Dict[str, object]]: with nvtx_range_debug("triattention.resolve_layout", color="blue"): - layout = self._runtime_kv_layout(num_layers) + layout = self._runtime_kv_layout(self._num_layers_from_manager()) with nvtx_range_debug("triattention.staging_lookup", color="blue"): # Retained spans always cover the model window (construction rejects budget < window). bufs = self._buffers_for(layout, prepared) execute_eviction_round( bufs, self.kv_cache_manager, - self.draft_kv_cache_manager, prepared, + self.draft_kv_cache_manager, normalize_scores=self.normalize_scores, ) for item in prepared: 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 index 251598a3207f..9de16a53ac7e 100644 --- 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 @@ -73,18 +73,14 @@ def __init__( num_layers: int, seq_len: int, num_q_heads: int, - num_kv_heads: int, num_freqs: int, - tokens_per_block: int, pool_shape: tuple[int, int, int, int, int], pool_strides: tuple[int, int, int, int, int], - pool_dtype: type[cutlass.Numeric], page_shards: int, write_partial_stats: bool = False, ) -> None: """Build the single validated production specialization.""" - if pool_dtype is not cutlass.BFloat16: - raise ValueError("TriAttention CuTe score requires BF16 K pages") + 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)" @@ -130,12 +126,7 @@ def __init__( self.producer_warp_id = 0 self.physical_threads = THREADS - self.num_physical_pages, _, pool_kv_heads, pool_tokens, pool_dim = pool_shape - if ( - pool_kv_heads != num_kv_heads - or pool_tokens != tokens_per_block - or pool_dim != 2 * num_freqs - ): + 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_slot, self.s_dim = pool_strides if self.s_slot != 2 * num_freqs or self.s_dim != 1: @@ -1374,12 +1365,9 @@ def __init__( layer_pools: list[torch.Tensor], layer_indices: list[int], max_requests: int, - num_layers: int, seq_len: int, num_q_heads: int, - num_kv_heads: int, num_freqs: int, - tokens_per_block: int, page_ids: torch.Tensor, seg_page_off: torch.Tensor, seg_req_id: torch.Tensor, @@ -1396,12 +1384,17 @@ def __init__( output: torch.Tensor, enable_partial_stats: bool = False, ) -> None: + # Pool shape [pages, K/V, heads, tokens, dim] of the anchor scored layer. + anchor_pool = layer_pools[layer_indices[0]] + num_layers = len(layer_indices) + num_kv_heads = int(anchor_pool.shape[2]) + tokens_per_block = int(anchor_pool.shape[3]) self.max_requests = int(max_requests) - self.num_layers = int(num_layers) + self.num_layers = num_layers # The widest score window (the whole bucket) sizes every start-dependent buffer. self.width = int(seq_len) self.num_q_heads = int(num_q_heads) - self.num_kv_heads = int(num_kv_heads) + self.num_kv_heads = num_kv_heads self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) self.enable_partial_stats = bool(enable_partial_stats) # One [stats_row, page_shard, {count, mean, m2}] record array. @@ -1434,7 +1427,7 @@ def __init__( freq_scale_sq, output, self.partial_stats, - layer_pools[layer_indices[0]], + anchor_pool, self.descriptors, ) # valid_seq_lens/token_starts are only 4-byte-aligned row views, read as per-CTA scalars. @@ -1447,7 +1440,7 @@ def __init__( _to_cute(freq_scale_sq), _to_cute(output), _to_cute(self.partial_stats), - _to_cute(layer_pools[layer_indices[0]]), + _to_cute(anchor_pool), _to_cute(self.descriptors, assumed_align=128), ) self._compiled: dict[int, object] = {} @@ -1477,8 +1470,8 @@ def __init__( num_kv_heads, num_freqs, tokens_per_block, - tuple(int(value) for value in layer_pools[layer_indices[0]].shape), - tuple(int(value) for value in layer_pools[layer_indices[0]].stride()), + tuple(int(value) for value in anchor_pool.shape), + tuple(int(value) for value in anchor_pool.stride()), ) tensor_specs = tuple( _tensor_spec(tensor) @@ -1497,12 +1490,9 @@ def __init__( num_layers=num_layers, seq_len=seq_len, num_q_heads=num_q_heads, - num_kv_heads=num_kv_heads, num_freqs=num_freqs, - tokens_per_block=tokens_per_block, - pool_shape=tuple(int(value) for value in layer_pools[layer_indices[0]].shape), - pool_strides=tuple(int(value) for value in layer_pools[layer_indices[0]].stride()), - pool_dtype=cutlass.BFloat16, + pool_shape=tuple(int(value) for value in anchor_pool.shape), + pool_strides=tuple(int(value) for value in anchor_pool.stride()), ) if self.enable_partial_stats: variant_key = "triattention_cute_score_stats" diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index b139950bac45..e1feee15af60 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -24,16 +24,16 @@ @triton.jit def _gather_mean_phase_kernel( round_starts, - table_cos, - table_sin, - mean_cos, - mean_sin, + phase_cos, + phase_sin, + phase_rows, valid_seq_lens, token_starts, + mean_cos, + mean_sin, valid_widths, swa_destination_bases, - table_rows, - rebase_delta, + swa_rebase_delta, NUM_FREQS: tl.constexpr, F_BLOCK: tl.constexpr, HAS_SWA: tl.constexpr, @@ -44,17 +44,17 @@ def _gather_mean_phase_kernel( frequency_mask = frequency < NUM_FREQS table_row = tl.load(round_starts + request).to(tl.int64) # Clamp stale or padded round starts instead of faulting. - table_row = tl.minimum(tl.maximum(table_row, 0), table_rows - 1) + table_row = tl.minimum(tl.maximum(table_row, 0), phase_rows - 1) source_offset = table_row * NUM_FREQS + frequency output_offset = request * NUM_FREQS + frequency - row_cos = tl.load(table_cos + source_offset, mask=frequency_mask, other=0.0) - row_sin = tl.load(table_sin + source_offset, mask=frequency_mask, other=0.0) + 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) token_start = tl.load(token_starts + request) tl.store(valid_widths + request, tl.load(valid_seq_lens + request) - token_start) if HAS_SWA: - tl.store(swa_destination_bases + request, token_start + rebase_delta) + tl.store(swa_destination_bases + request, token_start + swa_rebase_delta) def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: @@ -94,8 +94,7 @@ def _score_row_stats_kernel( row_inv_std, ROWS: tl.constexpr, WIDTH: tl.constexpr, - BLOCK: tl.constexpr, - EPSILON: tl.constexpr, + BLOCK: tl.constexpr = 256, ): """Compute one valid-prefix mean and inverse standard deviation per score row.""" flat_row = tl.program_id(0) @@ -119,7 +118,7 @@ def _score_row_stats_kernel( square_sum += tl.sum(centered * centered, axis=0) std = tl.sqrt(square_sum / valid_width) tl.store(row_mean + flat_row, mean) - tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, EPSILON)) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, STD_EPSILON)) @triton.jit @@ -131,16 +130,16 @@ def _score_per_head_reduce_kernel( selection_scores, selection_seq_lens, NUM_LAYERS: tl.constexpr, - NUM_QUERY_HEADS: tl.constexpr, + NUM_Q_HEADS: tl.constexpr, NUM_KV_HEADS: tl.constexpr, - QUERY_GROUP_SIZE: tl.constexpr, - SELECTION_ROWS: tl.constexpr, WIDTH: tl.constexpr, PER_LAYER: tl.constexpr, NORMALIZE: tl.constexpr, - BLOCK: tl.constexpr, + BLOCK: tl.constexpr = 256, ): """Reduce query-head score rows into one selector row per KV-head domain.""" + 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) @@ -160,7 +159,7 @@ def _score_per_head_reduce_kernel( 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_QUERY_HEADS + query_head + flat_row = (request * NUM_LAYERS + layer) * NUM_Q_HEADS + query_head value = tl.load( scores + flat_row * WIDTH + token, mask=valid_token, @@ -177,7 +176,7 @@ def _score_per_head_reduce_kernel( 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_QUERY_HEADS + query_head + flat_row = (request * NUM_LAYERS + layer) * NUM_Q_HEADS + query_head value = tl.load( scores + flat_row * WIDTH + token, mask=valid_token, @@ -202,20 +201,15 @@ def prepare_per_head_scores( row_inv_std: torch.Tensor, selection_scores: torch.Tensor, selection_seq_lens: torch.Tensor, - request_count: int, *, - num_kv_heads: int, per_layer: bool, normalize_scores: bool, ) -> None: """Normalize and reduce score rows for either per-head eviction mode.""" - request_count = int(request_count) - num_kv_heads = int(num_kv_heads) - _, num_layers, num_query_heads, width = scores.shape - selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - # 256 lanes / 4 warps, matching the settle shape. - stats_block = 256 - rows = num_layers * num_query_heads + request_count, num_layers, num_q_heads, width = scores.shape + selection_rows = int(selection_scores.shape[1]) + num_kv_heads = selection_rows // num_layers if per_layer else selection_rows + rows = num_layers * num_q_heads if normalize_scores: _score_row_stats_kernel[(request_count * rows,)]( scores, @@ -224,14 +218,9 @@ def prepare_per_head_scores( row_inv_std, ROWS=rows, WIDTH=width, - BLOCK=stats_block, - EPSILON=STD_EPSILON, - num_warps=4, ) - reduction_block = 256 - _score_per_head_reduce_kernel[ - (request_count, selection_rows, triton.cdiv(width, reduction_block)) - ]( + # 256-token tiles match the reduce kernel's BLOCK default. + _score_per_head_reduce_kernel[(request_count, selection_rows, triton.cdiv(width, 256))]( scores, valid_widths, row_mean, @@ -239,37 +228,28 @@ def prepare_per_head_scores( selection_scores, selection_seq_lens, NUM_LAYERS=num_layers, - NUM_QUERY_HEADS=num_query_heads, + NUM_Q_HEADS=num_q_heads, NUM_KV_HEADS=num_kv_heads, - QUERY_GROUP_SIZE=num_query_heads // num_kv_heads, - SELECTION_ROWS=selection_rows, WIDTH=width, PER_LAYER=per_layer, NORMALIZE=normalize_scores, - BLOCK=reduction_block, - num_warps=4, ) # ---- Selection finalize: settle threshold ties into the kept-ordinal rows ---- -# Settle launch shape, shared by every launch site. -SETTLE_BLOCK = 256 -SETTLE_NUM_WARPS = 4 - - @triton.jit def _settle_ties_kernel( - scores, - seq_lens, - prompt_offsets, - provisional_indices, - output_indices, + selection_scores_rows, + selection_row_lengths, + token_starts, + provisional_rows, + kept_ordinal_rows, WIDTH: tl.constexpr, KEEP_COUNT: tl.constexpr, SELECTION_ROWS: tl.constexpr, - BLOCK: tl.constexpr, + BLOCK: tl.constexpr = 256, ): """Settle one selection row's ties into its kept-ordinal output row (threshold recovery with sentinel-skip, strictly-greater count, @@ -278,12 +258,12 @@ def _settle_ties_kernel( request = tl.program_id(0) selection_domain = tl.program_id(1) row = request * SELECTION_ROWS + selection_domain - row_output = output_indices + row * KEEP_COUNT - row_scores = scores + row * WIDTH - row_selected = provisional_indices + row * KEEP_COUNT + 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_len = tl.load(prompt_offsets + request) + prompt_len = tl.load(token_starts + request) threshold = float("inf") for start in tl.static_range(0, KEEP_COUNT, BLOCK): @@ -303,7 +283,7 @@ def _settle_ties_kernel( ).to(tl.float32) threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) - seq_len = tl.load(seq_lens + row) + seq_len = 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) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 3efa4634c213..db729bab296b 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -160,12 +160,12 @@ def capacity_offsets(count): if has_draft: draft = dict( layer_pools=args["draft_layer_pools"], - layers=args["draft_layers"], + dense_layers=args["draft_layers"], layer_group_representative=args["draft_layer_group_representative"], layer_pool_keys=args["draft_layer_pool_keys"], kv_block_offsets=args["draft_kv_block_offsets"], page_table_slots=args["draft_page_table_slots"], - move_offsets=args["draft_move_offsets"], + dense_move_offsets=args["draft_move_offsets"], protected_tail_capacity=draft_tail, ) compaction = init_compaction_buffers( @@ -178,7 +178,7 @@ def capacity_offsets(count): layer_pool_keys=args["layer_pool_keys"], kv_block_offsets=args["kv_block_offsets"], page_table_slots=args["page_table_slots"], - prompt_offsets=args["prompt_offsets"], + token_starts=args["prompt_offsets"], dense_move_offsets=args["dense_move_offsets"], swa_move_offsets=args["swa_move_offsets"], per_layer_sources=per_layer, @@ -187,8 +187,8 @@ def capacity_offsets(count): valid_seq_lens=args["valid_sequence_lengths"], ), capacities=dict( - request_capacity=request_count, - decode_keep_count=keep_count, + max_requests=request_count, + keep_count=keep_count, protected_tail_capacity=tail, ), draft=draft, @@ -238,7 +238,7 @@ def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None staging.bulk_consume_done = torch.cuda.Event() staging.copy_done = torch.cuda.Event() staging.copy_pending = False - staging._bulk_offsets_src = torch.empty( + staging._block_offsets_host = torch.empty( 1, max_requests, 2, copy_block_count, dtype=torch.int32, device="cpu", pin_memory=True ) staging._bulk_copy_idx_src = torch.arange( 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 index 64f5b8912a11..0581bab90a0a 100644 --- 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 @@ -184,7 +184,7 @@ def test_union_fusion_matches_split_pipeline( def test_union_fusion_frequency_count_guard_raises() -> None: """16 frequencies (head size 32) sit outside the fused kernel contract and are rejected at kernel construction.""" - cutlass = pytest.importorskip("cutlass") + pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_cute_score_fused import ( # noqa: E501 _TriAttentionScoreKernel, @@ -195,12 +195,9 @@ def test_union_fusion_frequency_count_guard_raises() -> None: num_layers=1, seq_len=256, num_q_heads=8, - num_kv_heads=1, num_freqs=16, - tokens_per_block=128, pool_shape=(2, 2, 1, 128, 32), pool_strides=(8192, 4096, 4096, 32, 1), - pool_dtype=cutlass.BFloat16, page_shards=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 index 1d0e06f89709..ad29d76f83ef 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -181,7 +181,6 @@ def test_execute_eviction_round_orders_both_manager_streams(): event = mock.Mock() host = torch.zeros(6, 9, dtype=torch.int32) - prompt_offsets = object() buffers = SimpleNamespace( device=torch.device("cuda", torch.cuda.current_device()), max_requests=8, @@ -197,7 +196,6 @@ def test_execute_eviction_round_orders_both_manager_streams(): phase={"cos": None, "sin": None, "rows": 8}, phase_num_freqs=1, phase_f_block=1, - prompt_offsets=prompt_offsets, round_starts_device=None, valid_seq_lens_device=None, token_starts_device=None, @@ -207,10 +205,10 @@ def test_execute_eviction_round_orders_both_manager_streams(): swa_destination_bases=None, swa_rebase_delta=0, copy_block_count=0, - _bulk_offsets_src=None, + _block_offsets_host=None, block_offsets_device=None, draft_copy_block_count=0, - _draft_bulk_offsets_src=None, + _draft_block_offsets_host=None, draft_block_offsets_device=None, ) target_stream = mock.Mock() @@ -236,7 +234,7 @@ class Boom(RuntimeError): ): with pytest.raises(Boom): module.execute_eviction_round( - buffers, manager, draft_manager, prepared, normalize_scores=True + buffers, manager, prepared, draft_manager, normalize_scores=True ) # Both page-table planes were snapshotted before the round body fired. @@ -360,7 +358,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): # The staged logical position restores the uncompressed # length: physical confirmed plus everything evicted so far # (the prepared item's round_start). - prepared = internals.execute.call_args.args[3] + prepared = internals.execute.call_args.args[2] assert prepared[0]["round_start"] == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. @@ -378,7 +376,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): for call in internals.execute.call_args_list: assert call.args[0] is internals.buffers assert call.args[1] is target - assert call.args[2] is draft_manager + assert call.args[3] is draft_manager assert call.kwargs == {"normalize_scores": True} 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 index 07d74000030a..3422841644c2 100644 --- 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 @@ -11,15 +11,7 @@ import pytest import torch -from tensorrt_llm._torch.kv_cache_compression.compaction import PACK_BLOCK as _PACK_BLOCK -from tensorrt_llm._torch.kv_cache_compression.compaction import PACK_NUM_WARPS as _PACK_NUM_WARPS from tensorrt_llm._torch.kv_cache_compression.compaction import _pack_move_sources_kernel -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_BLOCK as _SETTLE_BLOCK, -) -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_NUM_WARPS as _SETTLE_NUM_WARPS, -) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _settle_ties_kernel, ) @@ -174,8 +166,6 @@ def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): WIDTH=width, KEEP_COUNT=keep_count, SELECTION_ROWS=selection_rows, - BLOCK=_SETTLE_BLOCK, - num_warps=_SETTLE_NUM_WARPS, ) torch.cuda.synchronize(device) @@ -259,16 +249,12 @@ def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, w swa_actual if has_swa else None, KEEP_COUNT=keep_count, DECISION_ROWS=selection_rows, - DENSE_TOTAL=dense_total, - SWA_TOTAL=swa_total if has_swa else 0, MOVE_CAPACITY=move_capacity, NUM_KV_HEADS=num_kv_heads, - SWA_WINDOW=swa_window if has_swa else 0, - BROADCAST=union, PER_LAYER=per_layer, - HAS_SWA=has_swa, - BLOCK=_PACK_BLOCK, - num_warps=_PACK_NUM_WARPS, + DENSE_TOTAL=dense_total, + SWA_TOTAL=swa_total if has_swa else 0, + SWA_WINDOW=swa_window if has_swa else 0, ) torch.cuda.synchronize(device) @@ -318,8 +304,6 @@ def test_settle_handles_topk_sentinel_padding(): WIDTH=width, KEEP_COUNT=keep_count, SELECTION_ROWS=1, - BLOCK=_SETTLE_BLOCK, - num_warps=_SETTLE_NUM_WARPS, ) torch.cuda.synchronize(device) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 7b33488e1634..7602866e09cf 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -343,8 +343,7 @@ def test_identity_compaction_is_rejected_instead_of_published(self): expected_keep_count=6, prompt_len=2, ) - ], - 2, + ] ) assert request.py_num_compressed_tokens == 0 @@ -455,7 +454,7 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): draft_manager._stream = mock.Mock() mgr.draft_kv_cache_manager = draft_manager - def compact(prepared, _num_layers): + def compact(prepared): # Publish and resize consume the same prepared cohort. return prepared @@ -476,8 +475,7 @@ def compact(prepared, _num_layers): "expected_keep_count": retained, "protected_tail": tail, } - ], - 2, + ] ) cache.resize.assert_called_once_with(retained + tail, None) draft_cache.resize.assert_called_once_with(retained + 1, None) @@ -594,7 +592,7 @@ def test_execute_rejects_int32_overflowing_round_starts(self): prepared = [_make_prepared_item(request_id=7, seq_len=64, round_start=2**31)] with pytest.raises((RuntimeError, OverflowError, ValueError)): - execute_eviction_round(staging, manager, None, prepared, normalize_scores=True) + execute_eviction_round(staging, manager, prepared, normalize_scores=True) assert gather.call_count == 0 def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): @@ -645,8 +643,7 @@ def stage_once(): staging, manager, [7], - current_stream, - staging._bulk_offsets_src, + staging._block_offsets_host, staging.block_offsets_device, staging.copy_block_count, ) @@ -662,7 +659,7 @@ def stage_once(): side_effect=AssertionError("page-table staging used torch.index_select"), ): stage_once() - assert staging._bulk_offsets_src.shape == (1, 1, 2, 8) + 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() @@ -733,8 +730,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): expected_keep_count=9, protected_tail=3, ), - ], - 2, + ] ) # One round-executor call carries the whole cohort (with the target @@ -742,8 +738,8 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): args = internals.execute.call_args assert args.args[0] is internals.buffers assert args.args[1] is manager.kv_cache_manager - assert args.args[2] is None - prepared = args.args[3] + assert args.args[3] is None + prepared = args.args[2] assert [item["request_id"] for item in prepared] == [7, 8] assert [item["round_start"] for item in prepared] == [8, 10] assert [item["prompt_len"] for item in prepared] == [3, 5] @@ -867,9 +863,7 @@ def prepared_cohort(): # The compact stage is stubbed to a no-op: this test owns the score # buffers only, never a staged move decision. with mock.patch.object(module, "compact"): - module.execute_eviction_round( - bufs, manager, None, prepared_cohort(), normalize_scores=False - ) + module.execute_eviction_round(bufs, manager, prepared_cohort(), normalize_scores=False) fixed = bufs.score_output.clone() assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] @@ -908,9 +902,7 @@ def prepared_cohort(): bufs.score_output.fill_(score_sentinel) bufs.valid_widths.fill_(-1) with mock.patch.object(module, "compact"): - module.execute_eviction_round( - bufs, manager, None, prepared_cohort(), normalize_scores=False - ) + module.execute_eviction_round(bufs, manager, prepared_cohort(), normalize_scores=False) second_launch = bufs.score_output.clone() assert torch.equal(bufs.valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index 05d92039961c..ad6d8692e51c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -17,8 +17,6 @@ from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - SETTLE_BLOCK, - SETTLE_NUM_WARPS, prepare_per_head_scores, ) @@ -63,34 +61,34 @@ def _make_selection_buffers( stream=None, ) bufs.valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) - bufs.prompt_offsets = torch.zeros(max_requests, dtype=torch.int32, device=device) + bufs.token_starts_device = torch.zeros(max_requests, dtype=torch.int32, device=device) if eviction_mode == "union": bufs.selection_rows_per_request = 1 bufs.combined = torch.empty((max_requests, width), dtype=torch.float32, device=device) # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. - bufs.final_indices = torch.zeros( + bufs.provisional_indices = torch.zeros( (max_requests, keep_count), dtype=torch.int32, device=device ) bufs.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) bufs.selection_scores_rows = bufs.combined bufs.selection_row_lengths = bufs.valid_widths - bufs.provisional_rows = bufs.final_indices - bufs.keep_rows = bufs.keep + bufs.provisional_rows = bufs.provisional_indices + bufs.kept_ordinal_rows = bufs.keep else: selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads bufs.selection_rows_per_request = selection_rows bufs.row_mean = torch.empty( max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device ) - bufs.row_std = torch.empty_like(bufs.row_mean) + bufs.row_inv_std = torch.empty_like(bufs.row_mean) bufs.selection_scores = torch.empty( (max_requests, selection_rows, width), dtype=torch.float32, device=device ) - bufs.row_seq_lens = torch.full( + bufs.selection_seq_lens = torch.full( (max_requests, selection_rows), width, dtype=torch.int32, device=device ) - bufs.top_indices_i32 = torch.zeros( + bufs.provisional_indices = torch.zeros( (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device ) bufs.keep = torch.empty( @@ -99,24 +97,22 @@ def _make_selection_buffers( bufs.selection_scores_rows = bufs.selection_scores.view( max_requests * selection_rows, width ) - bufs.selection_row_lengths = bufs.row_seq_lens.view(-1) - bufs.provisional_rows = bufs.top_indices_i32.view(-1, keep_count) - bufs.keep_rows = bufs.keep.view(-1, keep_count) + bufs.selection_row_lengths = bufs.selection_seq_lens.view(-1) + bufs.provisional_rows = bufs.provisional_indices.view(-1, keep_count) + bufs.kept_ordinal_rows = bufs.keep.view(-1, keep_count) # The launch args mirror the product's ``bufs.settle_args``/ # ``bufs.settle_kwargs`` order and keys exactly. bufs.settle_args = ( bufs.selection_scores_rows, bufs.selection_row_lengths, - bufs.prompt_offsets, + bufs.token_starts_device, bufs.provisional_rows, - bufs.keep_rows, + bufs.kept_ordinal_rows, ) bufs.settle_kwargs = dict( WIDTH=width, KEEP_COUNT=keep_count, SELECTION_ROWS=bufs.selection_rows_per_request, - BLOCK=SETTLE_BLOCK, - num_warps=SETTLE_NUM_WARPS, ) return bufs @@ -127,11 +123,9 @@ def _select_per_head(bufs, scores, *, normalize_scores): scores, bufs.valid_widths, bufs.row_mean, - bufs.row_std, + bufs.row_inv_std, bufs.selection_scores, - bufs.row_seq_lens, - bufs.max_requests, - num_kv_heads=bufs.num_kv_heads, + bufs.selection_seq_lens, per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) @@ -253,7 +247,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid max_requests=request_count, ) bufs.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) - bufs.prompt_offsets[:request_count].copy_( + bufs.token_starts_device[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) bufs.combined.copy_(scores.amax(dim=1)) @@ -303,8 +297,6 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no row_inv_std, selection_scores, selection_seq_lens, - request_count, - num_kv_heads=kv_heads, per_layer=per_layer, normalize_scores=normalize_scores, ) @@ -525,7 +517,7 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): num_slots=2, ) prepared = [_make_prepared_item(request_id=7, seq_len=seq_len, round_start=0)] - execute_eviction_round(bufs, manager, None, prepared, normalize_scores=False) + execute_eviction_round(bufs, manager, prepared, normalize_scores=False) assert torch.equal(bufs.keep, expected_keep) torch.cuda.synchronize(device) @@ -686,7 +678,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: # the derived move offsets stage keep_count + protected_tail # moves. Z-normalization is monotonic per row, so the raw-score # keep set is unchanged. - execute_eviction_round(bufs, manager, None, prepared, normalize_scores=True) + execute_eviction_round(bufs, manager, prepared, normalize_scores=True) selected = bufs.keep[0].clone().to(torch.long) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) From d604fa1567b59e419753f154accb3e301746c80c Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 03:49:45 -0700 Subject: [PATCH 123/178] [None][refactor] P0 closure: dead fields, duplicate guards, per-round transport, canonical names (knife 27) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 9 +- .../triattention/triattention.py | 530 ++++++++---------- .../triattention/triattention_kernels.py | 17 +- .../_torch/kv_cache_compression/conftest.py | 38 +- .../test_triattention_cute_score.py | 6 +- .../test_triattention_cute_union_fusion.py | 2 +- .../test_triattention_draft_cocompaction.py | 26 +- .../test_triattention_pipeline.py | 86 ++- .../test_triattention_selection_compaction.py | 90 +-- 9 files changed, 357 insertions(+), 447 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index e79debc2c8b5..40754eed42b8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -390,10 +390,7 @@ def init_compaction_buffers( swa_destination_bases=swa_destination_bases, # Per-round SWA destination rebase delta. swa_rebase_delta=keep_count - swa_window, - has_draft=draft is not None, draft_move_indices=draft_move_indices, - # Completion event: compact() records it after the last native launch. - consume_done=torch.cuda.Event(), # Private launch tuples: only compact() interprets these. target_launches=tuple(target_launches), draft_launches=draft_launches, @@ -403,10 +400,11 @@ def init_compaction_buffers( def compact(compaction: Dict[str, object], request_count: int) -> None: - """Pack each family's move sources, fire its native compacts, and record completion. + """Pack each family's move sources and fire its native compacts. Pure mover: the caller has already materialized its kept ordinals into the - agreed decision rows for the active ``request_count`` cohort. + agreed decision rows for the active ``request_count`` cohort, and the + caller owns the completion ordering of the whole round. """ rows, pack_args, pack_kwargs = compaction["target_pack_launch"] _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) @@ -418,4 +416,3 @@ def compact(compaction: Dict[str, object], request_count: int) -> None: _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) for launch in compaction["draft_launches"]: torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) - compaction["consume_done"].record(torch.cuda.current_stream()) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 2a340de9a91b..dd0265cf29ae 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -65,33 +65,23 @@ _OFFSET_MAX_LENGTH = 65536 -def _protected_tail_capacity(manager: KVCacheManagerV2, what: str) -> int: - capacity = int(manager.num_extra_kv_tokens) + int(manager._kv_reserve_draft_tokens) + 1 - if capacity <= 0: - raise RuntimeError(f"{what}KVCacheManagerV2 exposes an invalid protected-tail capacity") - return capacity - - -def _allocate_page_table_plane( - layer_pools: List[torch.Tensor], - page_representatives: List[int], - page_table_keys: List[object], - num_page_table_slots: int, - token_capacity: int, +def _allocate_block_offset_staging( + anchor_pool: torch.Tensor, + *, + num_pools: int, max_requests: int, -) -> Dict[str, object]: - representative_slots = { - representative: int(key[1]) - for representative, key in zip(page_representatives, page_table_keys) - } - device = layer_pools[page_representatives[0]].device - tokens_per_block = int(layer_pools[page_representatives[0]].shape[3]) + token_capacity: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """One pinned host snapshot + persistent device table pair in the native + V2 ``[pool, request, K/V, block]`` layout (block width 4-aligned for the + ``PackedInt`` copy ABI); the device follows the anchor KV pool.""" + tokens_per_block = int(anchor_pool.shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block - copy_block_count = (page_count + 3) // 4 * 4 - plane_shape = (num_page_table_slots, max_requests, 2, copy_block_count) - host = torch.empty(plane_shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) - dev = torch.empty(plane_shape, dtype=torch.int32, device=device) - return dict(slots=representative_slots, copy_block_count=copy_block_count, host=host, dev=dev) + staged_blocks_per_seq = (page_count + 3) // 4 * 4 + shape = (num_pools, max_requests, 2, staged_blocks_per_seq) + host = torch.empty(shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) + device_table = torch.empty(shape, dtype=torch.int32, device=anchor_pool.device) + return host, device_table def init_eviction_buffers( @@ -102,6 +92,7 @@ def init_eviction_buffers( phase: Dict[str, object], capacities: Dict[str, int], draft: Optional[Dict[str, object]] = None, + normalize_scores: bool = True, ) -> SimpleNamespace: """Build the one namespace of buffers, compiled launches, and compaction data (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put). @@ -110,6 +101,7 @@ def init_eviction_buffers( carries the local q_real/q_imag/mlr_coef [L, H, F] slices and freq_scale_sq; ``draft`` is one all-or-none resolved branch (its layout dict plus tail/page-table capacities); ``capacities`` the capacity numbers. + Static round policy (``normalize_scores``, mode) binds here, once. """ from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS from .triattention_cute_score_fused import TriAttentionCuteScoreRunner @@ -143,46 +135,51 @@ def init_eviction_buffers( bufs = SimpleNamespace() bufs.eviction_mode = eviction_mode - bufs.device = device + bufs.normalize_scores = bool(normalize_scores) bufs.max_requests = max_requests bufs.bucket_seq_len = seq_len bufs.decode_width = decode_width bufs.keep_count = keep_count bufs.page_table_token_capacity = page_table_token_capacity - # ---- staged page-table planes (target, plus the co-compressed draft) --- - plane = _allocate_page_table_plane( - layer_pools, - page_representatives, - page_table_keys, - num_page_table_slots, - page_table_token_capacity, - max_requests, + # ---- block-offset staging (target, plus the co-compressed draft) ------- + # Representative layer -> V2 pool id, resolved once here (init-local). + pool_id_by_representative_layer = { + representative: int(key[1]) + for representative, key in zip(page_representatives, page_table_keys) + } + bufs.block_offsets_host, bufs.block_offsets_device = _allocate_block_offset_staging( + layer_pools[page_representatives[0]], + num_pools=num_page_table_slots, + max_requests=max_requests, + token_capacity=page_table_token_capacity, ) - bufs.representative_slots = plane["slots"] - bufs.copy_block_count = plane["copy_block_count"] - bufs._block_offsets_host = plane["host"] - bufs.block_offsets_device = plane["dev"] # The draft is never scored: these offsets feed only the draft compacts. bufs.draft_block_offsets_device = None - bufs._draft_block_offsets_host = None - bufs.draft_copy_block_count = 0 + bufs.draft_block_offsets_host = None + bufs.draft_protected_tail_capacity = None draft_page_slots: Dict[int, int] = {} if draft is not None: draft_layout = draft["layout"] draft_representatives = list(draft_layout["pool_representatives"]) - draft_plane = _allocate_page_table_plane( - draft_layout["layer_pools"], - draft_representatives, - [draft_layout["layer_pool_keys"][layer] for layer in draft_representatives], - int(draft_layout["manager"].num_pools), - int(draft["page_table_token_capacity"]), - max_requests, + draft_page_slots = { + representative: int(draft_layout["layer_pool_keys"][representative][1]) + for representative in draft_representatives + } + draft_anchor_pool = draft_layout["layer_pools"][draft_representatives[0]] + # Construction-boundary invariant: the round shares one stream/event + # contract, so the draft pools must live on the target device. + if draft_anchor_pool.device != device: + raise RuntimeError("TriAttention draft KV pools must share the target KV pool device") + bufs.draft_block_offsets_host, bufs.draft_block_offsets_device = ( + _allocate_block_offset_staging( + draft_anchor_pool, + num_pools=int(draft_layout["manager"].num_pools), + max_requests=max_requests, + token_capacity=int(draft["page_table_token_capacity"]), + ) ) - draft_page_slots = draft_plane["slots"] - bufs.draft_copy_block_count = draft_plane["copy_block_count"] - bufs._draft_block_offsets_host = draft_plane["host"] - bufs.draft_block_offsets_device = draft_plane["dev"] + bufs.draft_protected_tail_capacity = int(draft["protected_tail_capacity"]) # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- bufs.request_metadata_host = torch.empty( @@ -190,7 +187,7 @@ def init_eviction_buffers( ) # numpy view over the pinned rows: per-round staging writes lists in place. bufs.request_metadata_host_np = bufs.request_metadata_host.numpy() - bufs._bulk_copy_idx_src = torch.arange( + bufs.identity_copy_indices_host = torch.arange( max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) # Zero-filled: an unstaged cohort must gather the phase table's row 0. @@ -218,12 +215,8 @@ def init_eviction_buffers( bufs.num_kv_heads = int(num_kv_heads) bufs.num_freqs = int(num_freqs) bufs.tokens_per_block = int(tokens_per_block) - # Device-side calibration indexing cannot be range-checked: validate the layer extent here. - num_calibrated_layers = q_real.numel() // (bufs.num_q_heads * bufs.num_freqs) - if min(dense_layers) < 0 or max(dense_layers) >= num_calibrated_layers: - raise ValueError("scored layer index exceeds the calibrated layer extent") _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} - page_table_slots = [bufs.representative_slots[_rep_of[layer]] for layer in dense_layers] + dense_layer_slots = [pool_id_by_representative_layer[_rep_of[layer]] for layer in dense_layers] seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( bufs.num_layers ) @@ -231,7 +224,7 @@ def init_eviction_buffers( max_requests ) block_offsets = bufs.block_offsets_device - slots_t = torch.tensor(page_table_slots, dtype=torch.int64, device=device) + slots_t = torch.tensor(dense_layer_slots, dtype=torch.int64, device=device) req_idx = seg_req_id.to(torch.int64) slot_idx = slots_t.repeat(max_requests) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) @@ -245,7 +238,7 @@ def init_eviction_buffers( ) # Persistent buffers: the compiled kernels capture their device pointers. bufs.padded_head_columns = PADDED_HEAD_COLUMNS - bufs.cute_scratch = torch.empty( + bufs.score_scratch = torch.empty( bufs.num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * seq_len, dtype=torch.float32, device=device, @@ -279,59 +272,54 @@ def init_eviction_buffers( if union: # Bucket-wide rows; consumers mask by the per-request widths. bufs.union_scores = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) - try: - bufs.runner = TriAttentionCuteScoreRunner( - layer_pools=list(layer_pools), - layer_indices=[int(layer) for layer in dense_layers], - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=bufs.num_q_heads, - num_freqs=bufs.num_freqs, - page_ids=block_offsets.view(-1), - seg_page_off=seg_page_off, - seg_req_id=seg_req_id, - seg_layer_id=seg_layer_id, - # Pointer capture of the staged metadata rows. - valid_seq_lens=bufs.valid_seq_lens_device, - seg_out_offset=seg_out_offset, - token_starts=bufs.token_starts_device, - q_real=q_real.view(-1), - q_imag=q_imag.view(-1), - mlr_coef=mlr_coef.view(-1), - mean_cos=bufs.mean_cos, - mean_sin=bufs.mean_sin, - freq_scale_sq=freq_scale_sq, - output=bufs.cute_scratch, - enable_partial_stats=union, - ) - except (ImportError, RuntimeError, TypeError, ValueError, AssertionError) as error: - raise RuntimeError( - "TriAttention CuTe score setup failed and no other score path exists" - ) from error + # THE score path (no fallback): construction failures raise the runner's + # own dtype/shape/TMA error. + bufs.runner = TriAttentionCuteScoreRunner( + layer_pools=list(layer_pools), + layer_indices=[int(layer) for layer in dense_layers], + max_requests=max_requests, + seq_len=seq_len, + num_q_heads=bufs.num_q_heads, + num_freqs=bufs.num_freqs, + page_ids=block_offsets.view(-1), + seg_page_off=seg_page_off, + seg_req_id=seg_req_id, + seg_layer_id=seg_layer_id, + # Pointer capture of the staged metadata rows. + valid_seq_lens=bufs.valid_seq_lens_device, + seg_out_offset=seg_out_offset, + token_starts=bufs.token_starts_device, + q_real=q_real.view(-1), + q_imag=q_imag.view(-1), + mlr_coef=mlr_coef.view(-1), + mean_cos=bufs.mean_cos, + mean_sin=bufs.mean_sin, + freq_scale_sq=freq_scale_sq, + output=bufs.score_scratch, + enable_partial_stats=union, + ) logger.info( f"TriAttention CuTe score enabled: {bufs.num_q_heads}q/{bufs.num_kv_heads}kv heads, " f"{bufs.num_freqs} freqs, {bufs.tokens_per_block}-token pages" ) - # ---- selection buffers -------------------------------------------------- + # ---- selection buffers (canonical row-major, one name per storage) ----- bufs.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) if union: bufs.selection_rows_per_request = 1 - bufs.combined = torch.empty( + bufs.selection_scores_rows = torch.empty( (max_requests, decode_width), dtype=torch.float32, device=device ) - bufs.provisional_indices = torch.empty( + # One selection row per request: its length IS the staged valid width. + bufs.selection_row_lengths = bufs.valid_widths + # Padded rows still need in-range ordinals for the finalizer's gather. + bufs.provisional_rows = torch.zeros( (max_requests, keep_count), dtype=torch.int32, device=device ) # Kept decode ordinals only (prompt-length independent rows). - bufs.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) - # Row-major views consumed by the top-k settle launch. - bufs.selection_scores_rows = bufs.combined - bufs.selection_row_lengths = bufs.valid_widths - bufs.provisional_rows = bufs.provisional_indices - bufs.kept_ordinal_rows = bufs.keep - # Padded rows still need in-range ordinals for the finalizer's gather. - bufs.provisional_indices.zero_() + bufs.kept_ordinal_rows = torch.empty( + (max_requests, keep_count), dtype=torch.int32, device=device + ) bufs.score_output = None else: selection_rows = ( @@ -360,22 +348,18 @@ def init_eviction_buffers( score_shape = (max_requests, bufs.num_layers, bufs.num_q_heads, 1) bufs.row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) bufs.row_inv_std = torch.empty_like(bufs.row_mean) - bufs.selection_scores = torch.empty( - (max_requests, selection_rows, decode_width), dtype=torch.float32, device=device + bufs.selection_scores_rows = torch.empty( + (max_requests * selection_rows, decode_width), dtype=torch.float32, device=device + ) + bufs.selection_row_lengths = torch.full( + (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device ) - bufs.selection_seq_lens = torch.full( - (max_requests, selection_rows), decode_width, dtype=torch.int32, device=device + bufs.provisional_rows = torch.zeros( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - selection_shape = (max_requests, selection_rows, keep_count) - bufs.provisional_indices = torch.empty(selection_shape, dtype=torch.int32, device=device) - bufs.keep = torch.empty(selection_shape, dtype=torch.int32, device=device) - bufs.selection_scores_rows = bufs.selection_scores.view( - max_requests * selection_rows, decode_width + bufs.kept_ordinal_rows = torch.empty( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - bufs.selection_row_lengths = bufs.selection_seq_lens.view(-1) - bufs.provisional_rows = bufs.provisional_indices.view(-1, keep_count) - bufs.kept_ordinal_rows = bufs.keep.view(-1, keep_count) - bufs.provisional_indices.zero_() # ---- compaction contract + decision-materialization prebinds ------------ per_layer = eviction_mode == "per_layer_perhead" @@ -401,7 +385,7 @@ def init_eviction_buffers( layer_group_representative=layer_group_representative, layer_pool_keys=list(layer_pool_keys), kv_block_offsets=bufs.block_offsets_device, - page_table_slots=bufs.representative_slots, + page_table_slots=pool_id_by_representative_layer, token_starts=bufs.token_starts_device, # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, @@ -419,7 +403,7 @@ def init_eviction_buffers( ), draft=draft_contract, ) - bufs.compaction = contract + bufs.compaction_plan = contract bufs.swa_destination_bases = contract["swa_destination_bases"] bufs.swa_rebase_delta = contract["swa_rebase_delta"] # The decision side: the settle launch materializes the kept-ordinal rows. @@ -437,10 +421,13 @@ def init_eviction_buffers( ) # ---- round-ordering events ---------------------------------------------- - bufs.copy_done = torch.cuda.Event() - bufs.copy_done.record(torch.cuda.current_stream(device)) - bufs.bulk_copy_done = torch.cuda.Event() - bufs.bulk_consume_done = torch.cuda.Event() + # Host staging (pinned metadata + snapshots) reuse fence. + bufs.staging_reuse_event = torch.cuda.Event() + bufs.staging_reuse_event.record(torch.cuda.current_stream(device)) + # Manager-stream H2D of the block-offset tables has completed. + bufs.block_offsets_ready_event = torch.cuda.Event() + # This cohort's compact is done: manager may resize/reuse pages. + bufs.compaction_done_event = torch.cuda.Event() bufs.copy_pending = False return bufs @@ -451,34 +438,33 @@ def _stage_block_offsets( request_ids: List[int], host_block_offsets: torch.Tensor, device_block_offsets: torch.Tensor, - copy_block_count: int, ) -> None: - """Gather the pinned snapshot before the async device copy: resize mutates the live host table.""" - if bufs.copy_pending and not bufs.copy_done.query(): - bufs.copy_done.synchronize() + """Gather the pinned snapshot before the async device copy: resize mutates + the live host table. The round owner has already fenced host-staging reuse.""" manager.index_mapper.gather_k_block_offsets( manager.host_kv_cache_block_offsets, host_block_offsets, request_ids, - copy_block_count, + host_block_offsets.shape[-1], ) - manager._stream.wait_event(bufs.copy_done) + manager._stream.wait_event(bufs.staging_reuse_event) copy_batch_block_offsets_to_device( host_block_offsets, device_block_offsets, - bufs._bulk_copy_idx_src[: len(request_ids)], + bufs.identity_copy_indices_host[: len(request_ids)], manager.index_scales, manager.kv_offset, manager._stream.cuda_stream, ) - bufs.bulk_copy_done.record(manager._stream) - torch.cuda.current_stream(bufs.device).wait_event(bufs.bulk_copy_done) + bufs.block_offsets_ready_event.record(manager._stream) + torch.cuda.current_stream(device_block_offsets.device).wait_event( + bufs.block_offsets_ready_event + ) def _cohort_move_offsets( bufs: SimpleNamespace, prepared: Sequence[Dict[str, object]], - draft_manager: Optional[KVCacheManagerV2], ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep set plus protected tail per request; rows past the cohort repeat the final @@ -494,12 +480,13 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: tails = [int(item["protected_tail"]) for item in prepared] dense = padded_offsets([bufs.keep_count + tail for tail in tails]) swa = None - if bufs.compaction["has_swa"]: - swa = padded_offsets([int(bufs.compaction["swa_window"]) + tail for tail in tails]) + if bufs.compaction_plan["has_swa"]: + swa = padded_offsets([int(bufs.compaction_plan["swa_window"]) + tail for tail in tails]) draft = None - if draft_manager is not None: - draft_tail = _protected_tail_capacity(draft_manager, "draft ") - draft = padded_offsets([bufs.keep_count + draft_tail] * len(prepared)) + if bufs.draft_protected_tail_capacity is not None: + draft = padded_offsets( + [bufs.keep_count + bufs.draft_protected_tail_capacity] * len(prepared) + ) return dense, swa, draft @@ -525,8 +512,6 @@ def execute_eviction_round( manager: KVCacheManagerV2, prepared: Sequence[Dict[str, object]], draft_manager: Optional[KVCacheManagerV2] = None, - *, - normalize_scores: bool, ) -> None: """Run one eviction round over the prepared cohort: stage the page-table snapshots and round metadata, then score, select, settle, and compact, and @@ -539,9 +524,9 @@ def execute_eviction_round( token_starts = [item["prompt_len"] for item in prepared] seq_lens = [item["seq_len"] for item in prepared] dense_move_offsets, swa_move_offsets, draft_move_offsets = _cohort_move_offsets( - bufs, prepared, draft_manager + bufs, prepared ) - stream = torch.cuda.current_stream(bufs.device) + stream = torch.cuda.current_stream(bufs.block_offsets_device.device) # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. max_round_start = max(round_starts) rows = ( @@ -555,9 +540,11 @@ def execute_eviction_round( for row, values in rows: if values is not None and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF: raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # The previous cohort's metadata H2D must complete before the pinned rows are rewritten. - if bufs.copy_pending and not bufs.copy_done.query(): - bufs.copy_done.synchronize() + # The one host-staging reuse fence: the previous cohort's async copies + # must complete before the pinned metadata rows AND the pinned + # target/draft block-offset snapshots are rewritten. + if bufs.copy_pending and not bufs.staging_reuse_event.query(): + bufs.staging_reuse_event.synchronize() host_table = bufs.request_metadata_host_np for row, values in rows: if values is not None: @@ -569,24 +556,22 @@ def execute_eviction_round( bufs, manager, request_ids, - bufs._block_offsets_host, + bufs.block_offsets_host, bufs.block_offsets_device, - bufs.copy_block_count, ) if draft_manager is not None: _stage_block_offsets( bufs, draft_manager, request_ids, - bufs._draft_block_offsets_host, + bufs.draft_block_offsets_host, bufs.draft_block_offsets_device, - bufs.draft_copy_block_count, ) try: bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) finally: - # Guards the pinned metadata until the asynchronous copies complete. - bufs.copy_done.record(stream) + # Guards the pinned staging until the asynchronous copies complete. + bufs.staging_reuse_event.record(stream) bufs.copy_pending = True request_count = len(prepared) union = bufs.eviction_mode == "union" @@ -614,8 +599,8 @@ def execute_eviction_round( bufs.runner.launch_union_fusion( request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_scores[:request_count] ) - columns = min(bufs.union_scores.shape[1], bufs.combined.shape[1]) - bufs.combined[:request_count, :columns].copy_( + columns = min(bufs.union_scores.shape[1], bufs.selection_scores_rows.shape[1]) + bufs.selection_scores_rows[:request_count, :columns].copy_( bufs.union_scores[:request_count, :columns] ) else: @@ -625,7 +610,7 @@ def execute_eviction_round( num_segments = request_count * bufs.num_layers pad = bufs.padded_head_columns source = ( - bufs.cute_scratch[ + bufs.score_scratch[ : bufs.num_kv_heads * pad * num_segments * bufs.bucket_seq_len ] .view( @@ -659,20 +644,20 @@ def execute_eviction_round( bufs.valid_widths, bufs.row_mean, bufs.row_inv_std, - bufs.selection_scores, - bufs.selection_seq_lens, + bufs.selection_scores_rows, + bufs.selection_row_lengths, per_layer=bufs.eviction_mode == "per_layer_perhead", - normalize_scores=normalize_scores, + normalize_scores=bufs.normalize_scores, ) settle_top_tokens(bufs, request_count) with nvtx_range("triattention.compact", color="purple"): - compact(bufs.compaction, request_count) + compact(bufs.compaction_plan, request_count) finally: # Order V2 page-table reuse and resize after this cohort's compact. - bufs.bulk_consume_done.record(torch.cuda.current_stream(bufs.device)) - manager._stream.wait_event(bufs.bulk_consume_done) + bufs.compaction_done_event.record(stream) + manager._stream.wait_event(bufs.compaction_done_event) if draft_manager is not None: - draft_manager._stream.wait_event(bufs.bulk_consume_done) + draft_manager._stream.wait_event(bufs.compaction_done_event) class TriAttention(BaseKVCacheCompressionManager): @@ -692,16 +677,11 @@ def __init__( normalize_scores: bool = True, ): super().__init__(kv_cache_manager, draft_kv_cache_manager) + # budget/beta positivity and the eviction_mode literal are validated at + # the config boundary (TriAttentionKvCacheCompressionConfig). self.budget = budget self.beta = beta - if self.budget <= 0 or self.beta <= 0: - raise ValueError("TriAttention budget and beta must both be positive") self.eviction_mode = eviction_mode - if self.eviction_mode not in ("union", "per_head", "per_layer_perhead"): - raise ValueError( - f"Unknown eviction_mode {self.eviction_mode!r}; expected one of " - "'union', 'per_head', 'per_layer_perhead'" - ) self.normalize_scores = bool(normalize_scores) if self.eviction_mode == "union" and not self.normalize_scores: raise ValueError( @@ -710,23 +690,14 @@ def __init__( ) # Hard-coded semantics: the prompt is always pinned and the budget # counts decode tokens only (physical KV reclaim requires both). - # Calibration is the official TriAttention .pt; TRT-LLM does not compute calibration. + # Calibration is the official TriAttention .pt; TRT-LLM does not + # compute calibration. The config boundary requires both paths. self.model_path = model_path - if self.model_path is None: - raise ValueError( - "TriAttention requires model_path so kernel-masked " - "sliding-attention layers can be classified safely" - ) self.calibration_path = calibration_path self.calibration: Optional[Dict[str, torch.Tensor]] = None self._calibrated = False - # Calibration-derived dims + stats, filled in on_request_init. - self._H: Optional[int] = None - self._F: Optional[int] = None self._freq_scale_sq: Optional[torch.Tensor] = None - # Geometric integration offsets, built lazily on first eviction. - self._offsets: Optional[torch.Tensor] = None # Mean-phase table dict, shared by reference with every buffer namespace. self._phase: Optional[Dict[str, object]] = None @@ -738,8 +709,20 @@ def __init__( # Manager-lifetime capability gates: everything read there is fixed at # construction, so validation runs once here. self._validate_v2_compatibility() - # Manager-lifetime constants (V2 fixes both inputs at construction). - self._protected_tail_capacity = _protected_tail_capacity(kv_cache_manager, "") + # Manager-lifetime constants (V2 fixes every input at construction): + # protected tails are num_extra + reserved draft width + 1 sampled token. + self._protected_tail_capacity = ( + int(kv_cache_manager.num_extra_kv_tokens) + + int(kv_cache_manager._kv_reserve_draft_tokens) + + 1 + ) + self._draft_protected_tail_capacity: Optional[int] = None + 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 + ) self._generation_growth = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) # Built once at the first eviction, reused for the manager's lifetime. self._buffers: Optional[SimpleNamespace] = None @@ -755,8 +738,7 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: request_id = request.py_request_id if request_id not in self._request_states: self._validate_request_capacity(request) - num_layers = self._num_layers_from_manager() - self._attention_layer_partition(num_layers) + self._attention_layer_partition() self._request_states[request_id] = { "generation_steps": 0, "evicted_tokens": 0, @@ -793,7 +775,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: draft_manager = self.draft_kv_cache_manager if draft_manager is None: return - draft_protected_tail = self._draft_protected_tail_capacity() + draft_protected_tail = self._draft_protected_tail_capacity draft_required_capacity = confirmed_capacity + draft_protected_tail draft_pool_capacity = draft_manager.get_num_available_tokens( token_num_upper_bound=confirmed_capacity, @@ -815,15 +797,10 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: f"page table covers {draft_table_capacity} tokens" ) - def _draft_protected_tail_capacity(self) -> int: - return _protected_tail_capacity(self.draft_kv_cache_manager, "draft ") - def _ensure_calibrated(self) -> None: if self._calibrated: return self.calibration = self._resolve_calibration() - self._H = int(self.calibration["E_q"].shape[1]) - self._F = int(self.calibration["E_q"].shape[2]) self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) # Pre-split query stats + MLR coefficient, shapes [L, H, F]. _Eq = self.calibration["E_q"] @@ -835,9 +812,9 @@ def _ensure_calibrated(self) -> None: self._calibrated = True def _validate_v2_compatibility(self) -> None: + # The base manager already enforces KVCacheManagerV2 target/draft types, + # and V2 construction guarantees beam width one. manager = self.kv_cache_manager - if not isinstance(manager, KVCacheManagerV2): - raise ValueError("TriAttention physical eviction requires KVCacheManagerV2") if manager.kv_factor != 2: raise ValueError( "TriAttention requires a standard key/value KV cache; " @@ -847,11 +824,10 @@ def _validate_v2_compatibility(self) -> None: raise ValueError("TriAttention does not support attention DP") if manager.is_disagg: raise ValueError("TriAttention does not support disaggregated serving") - if manager.max_beam_width != 1: - raise ValueError("TriAttention requires beam-width-one decoding") if manager.enable_swa_scratch_reuse: raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") - # Speculative feature gates run in the factory; the draft cache itself is validated here. + # Speculative feature gates run in the factory; the draft cache itself + # is validated here (V2 already forces scratch reuse off for drafts). draft_manager = self.draft_kv_cache_manager if draft_manager is not None: if not draft_manager.is_draft: @@ -865,10 +841,6 @@ def _validate_v2_compatibility(self) -> None: "the target, so the draft cache must be a standard " "key/value cache" ) - if draft_manager.enable_swa_scratch_reuse: - raise RuntimeError( - "TriAttention does not support V2 SWA scratch page-table remapping" - ) if self.eviction_mode != "union": raise ValueError( "TriAttention draft KV co-compression supports only " @@ -934,20 +906,13 @@ def _periodic_evict( if kv_cache is None: continue if not kv_cache.is_active: - raise RuntimeError( - "TriAttention cannot finalize a suspended target KV cache; " - f"request {request_id} must be resumed before " - "the final update hook" - ) - if request_id not in self._request_states: - raise RuntimeError( - f"request {request_id} reached generation without on_request_init" - ) + # Overlap scheduling may suspend a cache mid-flight; defer this + # request (pre-launch) instead of failing the whole batch. + continue resolved_requests.append((request, request_id, kv_cache)) if not resolved_requests: return prepared: List[Dict[str, object]] = [] - protected_tail_capacity = self._protected_tail_capacity # The resolved cache objects thread all the way to resize. with nvtx_range("triattention.metadata", color="cyan"): @@ -965,17 +930,6 @@ def _periodic_evict( scheduled_batch, request_id ) seq_len = raw_capacity - protected_tail - if seq_len < 0 or protected_tail < 0: - raise RuntimeError( - f"Request {request_id} has an inconsistent protected target tail: " - f"confirmed={seq_len}, capacity={raw_capacity}, " - f"protected_tail={protected_tail}" - ) - if protected_tail > protected_tail_capacity: - raise RuntimeError( - f"Request {request_id} protected tail {protected_tail} exceeds " - f"configured capacity {protected_tail_capacity}" - ) if seq_len < kv_cache.history_length: raise RuntimeError( f"Request {request_id} KV length {seq_len} is below finalized " @@ -987,12 +941,15 @@ def _periodic_evict( draft_kv_cache = None if self.draft_kv_cache_manager is not None: draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) - if draft_kv_cache is None or not draft_kv_cache.is_active: + if draft_kv_cache is None: + # A missing draft cache is a wiring/lifecycle bug. raise RuntimeError( - "TriAttention cannot co-compress a missing or " - f"suspended draft KV cache; request {request_id} must " - "be resumed before the final update hook" + "TriAttention cannot co-compress a missing draft KV " + f"cache for request {request_id}" ) + if not draft_kv_cache.is_active: + # Target and draft defer together (pre-launch). + continue prepared.append( { "request": request, @@ -1026,14 +983,15 @@ def _resize_compacted_requests(self, prepared) -> None: for item in prepared: kv_cache = item["kv_cache"] if not kv_cache.is_active: - continue - target_capacity = item["expected_keep_count"] - if target_capacity > kv_cache.capacity: + # Bytes already moved: skipping the ledger resize here + # would leave silent corruption. The compact-to-resize + # window is owned by this hook; a suspension inside it + # breaks the lifecycle contract. raise RuntimeError( - f"Request {item['request_id']} compacted capacity " - f"{target_capacity} exceeds current capacity {kv_cache.capacity}" + f"Request {item['request_id']} target KV cache was " + "suspended between compact and resize" ) - resized_capacity = target_capacity + item["protected_tail"] + resized_capacity = item["expected_keep_count"] + item["protected_tail"] if not kv_cache.resize(resized_capacity, None): raise RuntimeError( f"Failed to resize compacted KV cache for request " @@ -1041,11 +999,14 @@ def _resize_compacted_requests(self, prepared) -> None: ) if self.draft_kv_cache_manager is not None: # Same kept set: the draft shrinks to the same retained length plus its own tail. - draft_protected_tail = self._draft_protected_tail_capacity() + draft_protected_tail = self._draft_protected_tail_capacity for item in prepared: draft_kv_cache = item["draft_kv_cache"] if not draft_kv_cache.is_active: - continue + raise RuntimeError( + f"Request {item['request_id']} draft KV cache was " + "suspended between compact and resize" + ) draft_capacity = item["expected_keep_count"] + draft_protected_tail if not draft_kv_cache.resize(draft_capacity, None): raise RuntimeError( @@ -1092,23 +1053,12 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: # ---- helpers (eviction / scoring / V2 cache access / calibration) ---- - def _local_to_global_layers(self, num_layers: int) -> List[int]: + def _local_to_global_layers(self) -> List[int]: cached = self._local_to_global_layers_cache - if cached is not None: - if len(cached) != num_layers: - raise ValueError( - f"TriAttention layer count changed from {len(cached)} to {num_layers}" - ) - return cached - - global_layers = [int(layer) for layer in self.kv_cache_manager.pp_layers] - if len(global_layers) != num_layers: - raise ValueError( - f"KVCacheManagerV2 exposes {len(global_layers)} PP layers, " - f"but TriAttention received {num_layers} local layers" - ) - self._local_to_global_layers_cache = global_layers - return global_layers + if cached is None: + cached = [int(layer) for layer in self.kv_cache_manager.pp_layers] + self._local_to_global_layers_cache = cached + return cached @staticmethod def _has_sliding_window_signal(config: Dict[str, object]) -> bool: @@ -1132,17 +1082,15 @@ def _has_sliding_window_signal(config: Dict[str, object]) -> bool: return True return False - def _attention_layer_partition( - self, num_layers: int - ) -> Tuple[List[int], List[int], Optional[int]]: + def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int]]: """SWA layers here are stored at full length; the window applies only in the kernel.""" cached = self._attention_layer_partition_cache if cached is not None: return cached model_path = self.model_path - if model_path is None: - raise ValueError("TriAttention requires model_path") + global_layers = self._local_to_global_layers() + num_layers = len(global_layers) try: from transformers import AutoConfig @@ -1165,7 +1113,6 @@ def _attention_layer_partition( result = (list(range(num_layers)), [], None) self._attention_layer_partition_cache = result return result - global_layers = self._local_to_global_layers(num_layers) if global_layers and max(global_layers) >= len(layer_types): raise ValueError( f"Model config has {len(layer_types)} layer_types entries, " @@ -1197,16 +1144,13 @@ def _attention_layer_partition( self._attention_layer_partition_cache = result return result - def _runtime_kv_layout(self, num_layers: int) -> Dict[str, object]: - cached = self._runtime_kv_layout_cache + def _runtime_kv_layout(self) -> Dict[str, object]: + # The manager identity and layer count are manager-lifetime owner + # contracts; only the pool page counts are polled (stale-pointer + # safety until V2 exposes a layout epoch). manager = self.kv_cache_manager + cached = self._runtime_kv_layout_cache if cached is not None: - if cached["num_layers"] != num_layers: - raise ValueError( - f"TriAttention layer count changed from {cached['num_layers']} to {num_layers}" - ) - if cached["manager"] is not manager: - raise RuntimeError("TriAttention target KV cache manager changed at runtime") current_page_counts = self._pool_page_counts( manager, cached["global_layers"], @@ -1219,8 +1163,8 @@ def _runtime_kv_layout(self, num_layers: int) -> Dict[str, object]: ) return cached - global_layers = self._local_to_global_layers(num_layers) - dense_layers, swa_layers, swa_window = self._attention_layer_partition(num_layers) + global_layers = self._local_to_global_layers() + dense_layers, swa_layers, swa_window = self._attention_layer_partition() if not dense_layers: raise ValueError("TriAttention requires at least one full-attention layer") layout = self._build_runtime_kv_layout( @@ -1270,7 +1214,6 @@ def _build_runtime_kv_layout( pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) return dict( manager=manager, - num_layers=num_layers, global_layers=global_layers, layer_pools=layer_pools, dense_layers=dense_layers, @@ -1286,13 +1229,10 @@ def _build_runtime_kv_layout( ) def _draft_runtime_kv_layout(self) -> Dict[str, object]: + # Production callers gate on ``draft_kv_cache_manager is not None``. manager = self.draft_kv_cache_manager - if manager is None: - raise RuntimeError("TriAttention has no draft KV cache manager to lay out") cached = self._draft_runtime_kv_layout_cache if cached is not None: - if cached["manager"] is not manager: - raise RuntimeError("TriAttention draft KV cache manager changed at runtime") current_page_counts = self._pool_page_counts( manager, cached["global_layers"], @@ -1342,8 +1282,7 @@ def _buffers_for( layout: Dict[str, object], prepared: Sequence[Dict[str, object]], ) -> SimpleNamespace: - if not prepared: - raise ValueError("TriAttention eviction requires at least one request") + # Empty cohorts never reach here: _periodic_evict no-ops pre-launch. needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) needed_requests = len(prepared) @@ -1373,15 +1312,15 @@ def _buffers_for( seq_capacity = max(int(needed_page_tokens), 1024) seq_capacity = 1 << (seq_capacity - 1).bit_length() seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) - # The bucket capacity must be tile-aligned (mis-tiling stripes the score scratch silently). + # The bucket capacity must be tile-aligned (mis-tiling stripes the + # score scratch silently); the ceiling division constructs that fact. score_tile_tokens = max(64, int(mgr.tokens_per_block)) seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens - assert seq_capacity % score_tile_tokens == 0 page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) draft = None if self.draft_kv_cache_manager is not None: - draft_tail_capacity = self._draft_protected_tail_capacity() + draft_tail_capacity = self._draft_protected_tail_capacity draft = dict( layout=self._draft_runtime_kv_layout(), protected_tail_capacity=draft_tail_capacity, @@ -1389,20 +1328,14 @@ def _buffers_for( ) first_pool = layout["layer_pools"][layout["dense_layers"][0]] - if self._offsets is None: - # Upstream geometric offsets [1, 2, 4, ... <= max]. - self._offsets = torch.tensor( - [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], - device=first_pool.device, - dtype=torch.float32, - ) if self._phase is None: + # Upstream geometric offsets [1, 2, 4, ... <= max]: the table + # builder consumes them as host floats only (no device copy). self._phase = { - "offsets": self._offsets.contiguous(), "omega": self.calibration["omega"] .to(device=first_pool.device, dtype=torch.float32) .contiguous(), - "offset_values": self._offsets.tolist(), + "offset_values": [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], "cos": None, "sin": None, "rows": 0, @@ -1428,6 +1361,7 @@ def _buffers_for( protected_tail_capacity=tail_capacity, ), draft=draft, + normalize_scores=self.normalize_scores, ) self._buffers = bufs return bufs @@ -1442,13 +1376,11 @@ def _page_table_pool_keys( manager = self.kv_cache_manager layer_offsets = manager.layer_offsets layer_to_pool = manager.layer_to_pool_mapping_dict - try: - return [ - ("pool", int(layer_to_pool[layer_offsets[global_layers[layer]]])) - for layer in local_layers - ] - except (IndexError, KeyError, TypeError, ValueError) as exc: - raise RuntimeError("KVCacheManagerV2 exposes an invalid layer-to-pool mapping") from exc + # V2 owns the mapping; its own lookup errors are the precise ones. + return [ + ("pool", int(layer_to_pool[layer_offsets[global_layers[layer]]])) + for layer in local_layers + ] def _dense_layer_pool_groups( self, @@ -1468,7 +1400,7 @@ def _evict_requests( prepared: List[Dict[str, object]], ) -> List[Dict[str, object]]: with nvtx_range_debug("triattention.resolve_layout", color="blue"): - layout = self._runtime_kv_layout(self._num_layers_from_manager()) + layout = self._runtime_kv_layout() with nvtx_range_debug("triattention.staging_lookup", color="blue"): # Retained spans always cover the model window (construction rejects budget < window). bufs = self._buffers_for(layout, prepared) @@ -1477,21 +1409,16 @@ def _evict_requests( self.kv_cache_manager, prepared, self.draft_kv_cache_manager, - normalize_scores=self.normalize_scores, ) for item in prepared: + # Identity cohorts were filtered pre-launch (_periodic_evict). evicted = item["seq_len"] - item["expected_keep_count"] - if evicted <= 0: - raise RuntimeError("TriAttention attempted an identity compaction") request_state = self._request_states[item["request_id"]] request_state["evicted_tokens"] += evicted # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] return prepared - def _num_layers_from_manager(self) -> int: - return len(self.kv_cache_manager.pp_layers) - # ---- helpers: calibration loading ---- def _resolve_calibration(self) -> Dict[str, torch.Tensor]: @@ -1503,17 +1430,11 @@ def _resolve_calibration(self) -> Dict[str, torch.Tensor]: (``{metadata, stats{"layerLL_headHH": {q_mean_real, q_mean_imag, q_abs_mean}}}``) and our already-converted flat layout are accepted -- the official one is converted here. Calibration resolves lazily on the - first request (``on_request_init``), not at manager construction.""" - if self.calibration_path is None: - raise ValueError( - "TriAttention requires `calibration_path`: a calibration .pt from " - "the official tool (github.com/WeianMao/triattention). TRT-LLM does " - "not compute calibration -- see examples/ for the Qwen3-8B file and " - "the official calibration instructions." - ) + first request (``on_request_init``), not at manager construction, and + stays on CPU: runtime construction moves it to the pool device once.""" raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): - return {k: (v.to("cuda") if torch.is_tensor(v) else v) for k, v in raw.items()} + return raw if isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): return self._convert_official_calibration(raw) got = sorted(raw.keys()) if isinstance(raw, dict) else type(raw).__name__ @@ -1550,11 +1471,12 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: E_q[layer, h] = torch.complex(s["q_mean_real"].float(), s["q_mean_imag"].float()) E_q_norm[layer, h] = s["q_abs_mean"].float() omega, freq_scale_sq = self._rope_tables(freq_count) + # CPU schema: runtime construction moves tensors to the pool device once. calib = { - "E_q": E_q.to("cuda"), - "E_q_norm": E_q_norm.to("cuda"), - "omega": omega.to("cuda"), - "freq_scale_sq": freq_scale_sq.to("cuda"), + "E_q": E_q, + "E_q_norm": E_q_norm, + "omega": omega, + "freq_scale_sq": freq_scale_sq, } logger.info( f"TriAttention: converted official calibration {self.calibration_path}" diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index e1feee15af60..1db25cdfc6c3 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -199,15 +199,20 @@ def prepare_per_head_scores( valid_widths: torch.Tensor, row_mean: torch.Tensor, row_inv_std: torch.Tensor, - selection_scores: torch.Tensor, - selection_seq_lens: torch.Tensor, + selection_scores_rows: torch.Tensor, + selection_row_lengths: torch.Tensor, *, per_layer: bool, normalize_scores: bool, ) -> None: - """Normalize and reduce score rows for either per-head eviction mode.""" + """Normalize and reduce score rows for either per-head eviction mode. + + ``selection_scores_rows``/``selection_row_lengths`` are the canonical + row-major selection buffers over the full request capacity + (``capacity * selection_rows`` rows); ``valid_widths`` also spans the + capacity, so the per-request row count derives from the shapes.""" request_count, num_layers, num_q_heads, width = scores.shape - selection_rows = int(selection_scores.shape[1]) + selection_rows = int(selection_scores_rows.shape[0]) // int(valid_widths.shape[0]) num_kv_heads = selection_rows // num_layers if per_layer else selection_rows rows = num_layers * num_q_heads if normalize_scores: @@ -225,8 +230,8 @@ def prepare_per_head_scores( valid_widths, row_mean, row_inv_std, - selection_scores, - selection_seq_lens, + selection_scores_rows, + selection_row_lengths, NUM_LAYERS=num_layers, NUM_Q_HEADS=num_q_heads, NUM_KV_HEADS=num_kv_heads, diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index db729bab296b..949e307cbe62 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -63,7 +63,7 @@ def set_protected_tails(compaction, tail_lengths, draft_tail_lengths=None): compaction["swa_move_offsets"], [compaction["swa_window"] + int(tail) for tail in tail_lengths], ) - if compaction["has_draft"]: + if compaction["draft_move_offsets"] is not None: if draft_tail_lengths is None: draft_tail_lengths = [0] * len(tail_lengths) if len(draft_tail_lengths) != len(tail_lengths): @@ -224,34 +224,31 @@ def run_compaction(compaction): compact(compaction, compaction["request_count"]) -def make_bare_staging(device, *, max_requests, copy_block_count, page_count=None): +def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): """A bare buffer namespace for the bulk page-table copy tests.""" staging = SimpleNamespace() - staging.device = device staging.max_requests = max_requests - staging.copy_block_count = copy_block_count - if page_count is not None: - staging.page_count = page_count staging.keep_count = 4 - staging.compaction = {"has_swa": False} - staging.bulk_copy_done = torch.cuda.Event() - staging.bulk_consume_done = torch.cuda.Event() - staging.copy_done = torch.cuda.Event() + staging.compaction_plan = {"has_swa": False} + staging.draft_protected_tail_capacity = None + staging.block_offsets_ready_event = torch.cuda.Event() + staging.compaction_done_event = torch.cuda.Event() + staging.staging_reuse_event = torch.cuda.Event() staging.copy_pending = False - staging._block_offsets_host = torch.empty( - 1, max_requests, 2, copy_block_count, dtype=torch.int32, device="cpu", pin_memory=True + staging.block_offsets_host = torch.empty( + 1, max_requests, 2, staged_blocks_per_seq, dtype=torch.int32, device="cpu", pin_memory=True ) - staging._bulk_copy_idx_src = torch.arange( + 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, copy_block_count, dtype=torch.int32, device=device + 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_page_tables_bulk``/``stage`` consume.""" + """The manager surface ``_stage_block_offsets`` consumes.""" return SimpleNamespace( host_kv_cache_block_offsets=host_table, kv_factor=2, @@ -264,10 +261,7 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): def make_buffer_stubs(manager, *, decode_width=260): """Stub the calibration/layout surfaces around ``_buffers_for``.""" - manager._H = 2 - manager._F = 2 manager._freq_scale_sq = torch.ones(2) - manager._offsets = torch.ones(2) manager._phase = {"rows": 8} manager.calibration = {"omega": torch.ones(2)} manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) @@ -275,7 +269,6 @@ def make_buffer_stubs(manager, *, decode_width=260): pool = torch.empty(8, 2, 1, 4, 4) layout = dict( manager=SimpleNamespace(num_pools=1), - num_layers=2, global_layers=[0, 1], layer_pools=[pool, pool], dense_layers=[0, 1], @@ -458,7 +451,6 @@ def make_phase_table(offsets, omega, initial_rows): ) phase = { - "offsets": offsets.contiguous(), "omega": omega.to(dtype=torch.float32).contiguous(), "offset_values": offsets.tolist(), "cos": None, @@ -488,6 +480,7 @@ def make_cute_buffers( protected_tail_capacity=0, storage_groups=None, layer_pool_keys=None, + normalize_scores=True, ): """Real eviction buffers over the one-shared-slot default layout; split reference legs use ``eviction_mode="per_head"`` over the same pools. @@ -539,6 +532,7 @@ def make_cute_buffers( keep_count=keep_count, protected_tail_capacity=protected_tail_capacity, ), + normalize_scores=normalize_scores, ) @@ -571,7 +565,7 @@ def launch_split_scores( num_segments = request_count * bufs.num_layers group_size = bufs.num_q_heads // bufs.num_kv_heads source = ( - bufs.cute_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] + bufs.score_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ :, :group_size ] @@ -587,7 +581,7 @@ def launch_split_scores( (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), float("nan"), dtype=torch.float32, - device=bufs.device, + device=bufs.score_scratch.device, ) torch.gather( source, 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 index 391a705c6fbc..1e9a129db573 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -141,7 +141,7 @@ def test_cute_kernel_matches_torch_oracle(case): mean_sin, oracle_inputs, ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) - device = bufs.device + device = bufs.score_scratch.device oracle = _torch_tri_score_oracle( pools, @@ -204,7 +204,9 @@ def test_unsupported_geometry_raises_at_buffer_construction(): for _ in range(num_layers) ] calib = torch.randn(num_layers, 2, num_freqs, device=device) - with pytest.raises(RuntimeError, match="no other score path exists"): + # No rewrap: the runner'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, 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 index 0581bab90a0a..ba671184af71 100644 --- 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 @@ -284,7 +284,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: decode_width=decode_window, ) bufs = _make_cute_buffers(eviction_mode="union", max_requests=max_requests, **common) - assert (bufs.cute_scratch.numel() > 2**31) == (max_requests == 64) + assert (bufs.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_bufs = _make_cute_buffers(eviction_mode="per_head", max_requests=request_count, **common) 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 index ad29d76f83ef..3be09125c566 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -182,14 +182,14 @@ def test_execute_eviction_round_orders_both_manager_streams(): event = mock.Mock() host = torch.zeros(6, 9, dtype=torch.int32) buffers = SimpleNamespace( - device=torch.device("cuda", torch.cuda.current_device()), max_requests=8, keep_count=4, eviction_mode="union", - compaction={"has_swa": False}, + compaction_plan={"has_swa": False}, + draft_protected_tail_capacity=1, copy_pending=False, - copy_done=mock.Mock(), - bulk_consume_done=event, + staging_reuse_event=mock.Mock(), + compaction_done_event=event, request_metadata_host=host, request_metadata_host_np=host.numpy(), request_metadata_device=torch.zeros_like(host), @@ -204,11 +204,9 @@ def test_execute_eviction_round_orders_both_manager_streams(): mean_sin=None, swa_destination_bases=None, swa_rebase_delta=0, - copy_block_count=0, - _block_offsets_host=None, - block_offsets_device=None, - draft_copy_block_count=0, - _draft_block_offsets_host=None, + block_offsets_host=None, + block_offsets_device=torch.zeros(1, dtype=torch.int32), + draft_block_offsets_host=None, draft_block_offsets_device=None, ) target_stream = mock.Mock() @@ -233,9 +231,7 @@ class Boom(RuntimeError): mock.patch.object(module, "compact") as compact, ): with pytest.raises(Boom): - module.execute_eviction_round( - buffers, manager, prepared, draft_manager, normalize_scores=True - ) + module.execute_eviction_round(buffers, manager, prepared, draft_manager) # Both page-table planes were snapshotted before the round body fired. assert stage.call_count == 2 @@ -327,6 +323,8 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): 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) manager._request_states[7] = _fresh_request_state() @@ -377,7 +375,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert call.args[0] is internals.buffers assert call.args[1] is target assert call.args[3] is draft_manager - assert call.kwargs == {"normalize_scores": True} + assert call.kwargs == {} def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): @@ -394,6 +392,8 @@ def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): draft_manager.num_pools = 1 draft_manager.host_kv_cache_block_offsets = torch.zeros(1, 8, 2, 4, dtype=torch.int32) manager.draft_kv_cache_manager = draft_manager + # Injected post-construction: mirror the ctor-cached manager-lifetime tail. + manager._draft_protected_tail_capacity = 1 manager._draft_runtime_kv_layout = mock.Mock( return_value=dict( layer_pools=[], diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 7602866e09cf..e8d16cdd4b58 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -145,7 +145,6 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): triattention.kv_cache_manager = manager cached = dict( manager=manager, - num_layers=3, global_layers=[10, 11, 12], layer_pools=[torch.empty(4), torch.empty(8), torch.empty(4)], dense_layers=[0, 1, 2], @@ -160,7 +159,7 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): ) triattention._runtime_kv_layout_cache = cached - assert triattention._runtime_kv_layout(3) is cached + assert triattention._runtime_kv_layout() is cached manager.get_buffers.assert_not_called() assert page_count_query.call_args_list == [ mock.call(100, Role.KEY), @@ -168,7 +167,7 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): ] with pytest.raises(RuntimeError, match="pool layout changed"): - triattention._runtime_kv_layout(3) + triattention._runtime_kv_layout() manager.get_buffers.assert_not_called() assert page_count_query.call_args_list == [ mock.call(100, Role.KEY), @@ -311,7 +310,9 @@ class TestCompressedTokenPublication: # The monotone publication contract is covered end to end in # test_triattention_draft_cocompaction.py. - def test_identity_compaction_is_rejected_instead_of_published(self): + def test_identity_selection_is_filtered_before_launch(self): + # Identity cohorts (seq_len == prompt + budget) are the pre-launch + # owner no-op: nothing launches and nothing is published. manager = _make_triattention(budget=4) manager.kv_cache_manager._stream = mock.Mock() request = _make_request(7, py_prompt_len=2) @@ -331,22 +332,6 @@ def test_identity_compaction_is_rejected_instead_of_published(self): assert state["evicted_tokens"] == 0 cache.resize.assert_not_called() - # A keep-contract violation fails loudly, never publishes. - with _mocked_eviction_internals(manager): - with pytest.raises(RuntimeError, match="identity compaction"): - manager._evict_requests( - [ - _make_prepared_item( - request, - request_id=7, - seq_len=6, - expected_keep_count=6, - prompt_len=2, - ) - ] - ) - assert request.py_num_compressed_tokens == 0 - class TestEvictionLifecycle: def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): @@ -370,17 +355,6 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): periodic_evict.assert_called_once_with(batch) - def test_unregistered_generation_request_is_rejected(self): - # Missing on_request_init = framework ordering bug: fail loudly. - manager = _make_triattention() - manager._calibrated = True - cache = SimpleNamespace(capacity=8, history_length=0, is_active=True) - manager.kv_cache_manager.kv_cache_map = {7: cache} - request = _make_request(7, py_prompt_len=2) - - with pytest.raises(RuntimeError, match="without on_request_init"): - manager._periodic_evict(SimpleNamespace(generation_requests=[request])) - @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 @@ -410,14 +384,16 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft num_extra_kv_tokens=num_extra_kv_tokens, _kv_reserve_draft_tokens=kv_reserve_draft_tokens, ) - mgr._L = 2 mgr._request_states = {} _set_request_state(mgr, 7, generation_steps=127) mgr.beta = 128 mgr.budget = 4096 return mgr, request, batch - def test_suspended_cache_rejects_batch_before_cadence_mutation(self): + def test_suspended_cache_defers_that_request_pre_launch(self): + # A suspended cache is a legal overlap-scheduler transient: that + # request defers (pre-launch, no cadence mutation) while the rest of + # the cohort proceeds. manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) second_request = _make_request(8, py_prompt_len=1024) manager.kv_cache_manager.kv_cache_map[8] = SimpleNamespace(is_active=False) @@ -425,10 +401,13 @@ def test_suspended_cache_rejects_batch_before_cadence_mutation(self): second_state = _set_request_state(manager, 8, generation_steps=127) batch = SimpleNamespace(generation_requests=[first_request, second_request]) - with pytest.raises(RuntimeError, match="request 8 must be resumed"): + with mock.patch.object(manager, "_evict_requests", side_effect=lambda p: p) as evict: manager._periodic_evict(batch) - assert first_state["generation_steps"] == 127 + # Only the active request launched; the suspended one deferred whole. + prepared = evict.call_args.args[0] + assert [item["request_id"] for item in prepared] == [7] + assert first_state["generation_steps"] == 128 assert second_state["generation_steps"] == 127 @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) @@ -453,6 +432,8 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): 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 def compact(prepared): # Publish and resize consume the same prepared cohort. @@ -584,7 +565,7 @@ def test_execute_rejects_int32_overflowing_round_starts(self): ) device = torch.device("cuda", torch.cuda.current_device()) - staging = _make_bare_staging(device, max_requests=1, copy_block_count=8) + staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) gather = mock.Mock() manager = _make_staging_manager( torch.zeros(1, 2, 2, 12, dtype=torch.int32), gather, torch.cuda.Stream(device=device) @@ -592,7 +573,7 @@ def test_execute_rejects_int32_overflowing_round_starts(self): prepared = [_make_prepared_item(request_id=7, seq_len=64, round_start=2**31)] with pytest.raises((RuntimeError, OverflowError, ValueError)): - execute_eviction_round(staging, manager, prepared, normalize_scores=True) + execute_eviction_round(staging, manager, prepared) assert gather.call_count == 0 def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): @@ -633,8 +614,8 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): ) gather = mock.Mock(side_effect=gather_k_block_offsets) - staging = _make_bare_staging(device, max_requests=1, copy_block_count=8, page_count=5) - staging.copy_done.record(current_stream) + staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) + staging.staging_reuse_event.record(current_stream) manager = _make_staging_manager(host_table, gather, manager_stream) def stage_once(): @@ -643,9 +624,8 @@ def stage_once(): staging, manager, [7], - staging._block_offsets_host, + staging.block_offsets_host, staging.block_offsets_device, - staging.copy_block_count, ) # Round 1: mutate the host table and the slot assignment right after @@ -659,7 +639,7 @@ def stage_once(): side_effect=AssertionError("page-table staging used torch.index_select"), ): stage_once() - assert staging._block_offsets_host.shape == (1, 1, 2, 8) + 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() @@ -690,8 +670,8 @@ def stage_once(): snapshot.copy_(staging.block_offsets_device) # ``execute_eviction_round``'s completion ordering: one event records # the consumers and the manager stream waits on it. - staging.bulk_consume_done.record(torch.cuda.current_stream(device)) - manager_stream.wait_event(staging.bulk_consume_done) + staging.compaction_done_event.record(torch.cuda.current_stream(device)) + manager_stream.wait_event(staging.compaction_done_event) stage_once() current_stream.synchronize() @@ -744,7 +724,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): assert [item["round_start"] for item in prepared] == [8, 10] assert [item["prompt_len"] for item in prepared] == [3, 5] assert [item["seq_len"] for item in prepared] == [8, 10] - assert args.kwargs == {"normalize_scores": True} + assert args.kwargs == {} # The derived move offsets stage keep + tail moves per request # (budget=4 -> [6, 7]); padded rows repeat the final offset. @@ -753,9 +733,12 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): ) offsets_buffers = SimpleNamespace( - max_requests=8, keep_count=4, compaction={"has_swa": False} + max_requests=8, + keep_count=4, + compaction_plan={"has_swa": False}, + draft_protected_tail_capacity=None, ) - dense, swa, draft = _cohort_move_offsets(offsets_buffers, prepared, None) + dense, swa, draft = _cohort_move_offsets(offsets_buffers, prepared) assert dense == [0, 6, 13, 13, 13, 13, 13, 13, 13] assert swa is None assert draft is None @@ -828,6 +811,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun # One storage group and page-table slot per layer (distinct pools). storage_groups={("pool", layer): [layer] for layer in layer_order}, layer_pool_keys=[("pool", layer) for layer in layer_order], + normalize_scores=False, ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) @@ -863,7 +847,7 @@ def prepared_cohort(): # The compact stage is stubbed to a no-op: this test owns the score # buffers only, never a staged move decision. with mock.patch.object(module, "compact"): - module.execute_eviction_round(bufs, manager, prepared_cohort(), normalize_scores=False) + module.execute_eviction_round(bufs, manager, prepared_cohort()) fixed = bufs.score_output.clone() assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] @@ -902,7 +886,7 @@ def prepared_cohort(): bufs.score_output.fill_(score_sentinel) bufs.valid_widths.fill_(-1) with mock.patch.object(module, "compact"): - module.execute_eviction_round(bufs, manager, prepared_cohort(), normalize_scores=False) + module.execute_eviction_round(bufs, manager, prepared_cohort()) second_launch = bufs.score_output.clone() assert torch.equal(bufs.valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) @@ -929,9 +913,9 @@ def test_layer_partition_uses_local_config_and_validates_window(self, budget, fi if not fits_window: # The decode budget must cover the kernel-masked SWA window. with pytest.raises(ValueError, match="budget=127"): - mgr._attention_layer_partition(4) + mgr._attention_layer_partition() return - dense, sliding, window = mgr._attention_layer_partition(4) + dense, sliding, window = mgr._attention_layer_partition() load.assert_called_once_with( "/models/gpt-oss", trust_remote_code=True, local_files_only=True 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 index ad6d8692e51c..1f1a8b846c9a 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -48,33 +48,33 @@ def _make_selection_buffers( num_kv_heads=1, ): """Selection-only buffers for the mode, without CuTe score state or - compaction (the settle launch writes only the kept-ordinal rows).""" + compaction (the settle launch writes only the kept-ordinal rows). + Mirrors the product's canonical row-major selection allocation.""" bufs = SimpleNamespace( eviction_mode=eviction_mode, - device=device, max_requests=max_requests, decode_width=width, keep_count=keep_count, num_layers=num_layers, num_q_heads=num_query_heads, num_kv_heads=num_kv_heads, - stream=None, ) bufs.valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) bufs.token_starts_device = torch.zeros(max_requests, dtype=torch.int32, device=device) if eviction_mode == "union": bufs.selection_rows_per_request = 1 - bufs.combined = torch.empty((max_requests, width), dtype=torch.float32, device=device) + bufs.selection_scores_rows = torch.empty( + (max_requests, width), dtype=torch.float32, device=device + ) + bufs.selection_row_lengths = bufs.valid_widths # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. - bufs.provisional_indices = torch.zeros( + bufs.provisional_rows = torch.zeros( + (max_requests, keep_count), dtype=torch.int32, device=device + ) + bufs.kept_ordinal_rows = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=device ) - bufs.keep = torch.empty((max_requests, keep_count), dtype=torch.int32, device=device) - bufs.selection_scores_rows = bufs.combined - bufs.selection_row_lengths = bufs.valid_widths - bufs.provisional_rows = bufs.provisional_indices - bufs.kept_ordinal_rows = bufs.keep else: selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads bufs.selection_rows_per_request = selection_rows @@ -82,24 +82,18 @@ def _make_selection_buffers( max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device ) bufs.row_inv_std = torch.empty_like(bufs.row_mean) - bufs.selection_scores = torch.empty( - (max_requests, selection_rows, width), dtype=torch.float32, device=device - ) - bufs.selection_seq_lens = torch.full( - (max_requests, selection_rows), width, dtype=torch.int32, device=device + bufs.selection_scores_rows = torch.empty( + (max_requests * selection_rows, width), dtype=torch.float32, device=device ) - bufs.provisional_indices = torch.zeros( - (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device + bufs.selection_row_lengths = torch.full( + (max_requests * selection_rows,), width, dtype=torch.int32, device=device ) - bufs.keep = torch.empty( - (max_requests, selection_rows, keep_count), dtype=torch.int32, device=device + bufs.provisional_rows = torch.zeros( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - bufs.selection_scores_rows = bufs.selection_scores.view( - max_requests * selection_rows, width + bufs.kept_ordinal_rows = torch.empty( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - bufs.selection_row_lengths = bufs.selection_seq_lens.view(-1) - bufs.provisional_rows = bufs.provisional_indices.view(-1, keep_count) - bufs.kept_ordinal_rows = bufs.keep.view(-1, keep_count) # The launch args mirror the product's ``bufs.settle_args``/ # ``bufs.settle_kwargs`` order and keys exactly. bufs.settle_args = ( @@ -124,8 +118,8 @@ def _select_per_head(bufs, scores, *, normalize_scores): bufs.valid_widths, bufs.row_mean, bufs.row_inv_std, - bufs.selection_scores, - bufs.selection_seq_lens, + bufs.selection_scores_rows, + bufs.selection_row_lengths, per_layer=bufs.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) @@ -211,10 +205,11 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( ) bufs.valid_widths.copy_(valid_widths.to(device)) scores = scores_cpu.to(device) + keep_shape = (request_count, bufs.selection_rows_per_request, keep_count) _select_per_head(bufs, scores, normalize_scores=normalize_scores) - first = bufs.keep.cpu() + first = bufs.kept_ordinal_rows.view(keep_shape).cpu() _select_per_head(bufs, scores, normalize_scores=normalize_scores) - second = bufs.keep.cpu() + second = bufs.kept_ordinal_rows.view(keep_shape).cpu() stream.synchronize() assert torch.equal(first, expected) @@ -250,9 +245,9 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid bufs.token_starts_device[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) - bufs.combined.copy_(scores.amax(dim=1)) + bufs.selection_scores_rows.copy_(scores.amax(dim=1)) settle_top_tokens(bufs, bufs.max_requests) - actual = bufs.keep.cpu() + actual = bufs.kept_ordinal_rows.cpu() combined = scores.amax(dim=1).cpu() for request, valid_width in enumerate(valid_widths): @@ -283,11 +278,12 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no ) row_inv_std = torch.empty_like(row_mean) selection_rows = layers * kv_heads if per_layer else kv_heads - selection_scores = torch.empty( - request_count, selection_rows, width, dtype=torch.float32, device=device + # 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_seq_lens = torch.empty( - request_count, selection_rows, dtype=torch.int32, device=device + selection_row_lengths = torch.empty( + request_count * selection_rows, dtype=torch.int32, device=device ) prepare_per_head_scores( @@ -295,15 +291,16 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no valid_widths, row_mean, row_inv_std, - selection_scores, - selection_seq_lens, + selection_scores_rows, + selection_row_lengths, 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_seq_lens.cpu(), + selection_row_lengths.view(request_count, selection_rows).cpu(), valid_widths.cpu().view(request_count, 1).expand(-1, selection_rows), ) query_group_size = query_heads // kv_heads @@ -497,8 +494,17 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): ("pool", 1): dense_groups[1], }, layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], + normalize_scores=False, ) - assert bufs.compaction["has_swa"] is False + assert bufs.compaction_plan["has_swa"] is False + # 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 bufs.block_offsets_host.shape == bufs.block_offsets_device.shape + assert bufs.block_offsets_host.shape[:3] == (2, 1, 2) + assert bufs.block_offsets_host.shape[-1] % 4 == 0 + assert bufs.block_offsets_host.dtype == bufs.block_offsets_device.dtype == torch.int32 + assert bufs.block_offsets_host.is_contiguous() and bufs.block_offsets_device.is_contiguous() + assert bufs.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; @@ -517,8 +523,8 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): num_slots=2, ) prepared = [_make_prepared_item(request_id=7, seq_len=seq_len, round_start=0)] - execute_eviction_round(bufs, manager, prepared, normalize_scores=False) - assert torch.equal(bufs.keep, expected_keep) + execute_eviction_round(bufs, manager, prepared) + assert torch.equal(bufs.kept_ordinal_rows.view_as(expected_keep), expected_keep) torch.cuda.synchronize(device) for before_pool, after_pool, table, layer in zip( @@ -678,8 +684,8 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: # the derived move offsets stage keep_count + protected_tail # moves. Z-normalization is monotonic per row, so the raw-score # keep set is unchanged. - execute_eviction_round(bufs, manager, prepared, normalize_scores=True) - selected = bufs.keep[0].clone().to(torch.long) + execute_eviction_round(bufs, manager, prepared) + selected = bufs.kept_ordinal_rows[0].clone().to(torch.long) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) after = snapshot(compacted_capacity) From 82c4893aee89023233c1f0e681f82c757cc326b7 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 04:57:37 -0700 Subject: [PATCH 124/178] [None][fix] Bind the settle-stats epsilon as a constexpr default; Triton rejects plain-global capture Signed-off-by: tianruih --- .../triattention/triattention_kernels.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 1db25cdfc6c3..ec19091f617b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -95,6 +95,9 @@ def _score_row_stats_kernel( ROWS: 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 one valid-prefix mean and inverse standard deviation per score row.""" flat_row = tl.program_id(0) @@ -118,7 +121,7 @@ def _score_row_stats_kernel( square_sum += tl.sum(centered * centered, axis=0) std = tl.sqrt(square_sum / valid_width) tl.store(row_mean + flat_row, mean) - tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, STD_EPSILON)) + tl.store(row_inv_std + flat_row, 1.0 / tl.maximum(std, EPSILON)) @triton.jit From df979c8fa1fd9465ec79160108965429198d8485 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 05:17:53 -0700 Subject: [PATCH 125/178] [None][refactor] P1 owner consolidation: canonical pool ids, opaque compaction plans, ctor-bound runner buffers (knife 28) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 294 +++++++++--------- .../triattention/triattention.py | 112 +++---- .../triattention_cute_score_fused.py | 75 ++--- .../_torch/kv_cache_compression/conftest.py | 75 +++-- .../test_triattention_cute_union_fusion.py | 9 +- .../test_triattention_draft_cocompaction.py | 30 +- .../test_triattention_pipeline.py | 8 +- .../test_triattention_selection_compaction.py | 14 +- 8 files changed, 294 insertions(+), 323 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 40754eed42b8..4a1ed0727c82 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -16,12 +16,12 @@ """Batched physical KV-cache compaction: an algorithm-neutral mover. ``init_compaction_buffers`` agrees on the decision rows (kept ordinals; -move offsets ride the caller's staged rows) once per geometry and retains -one launch contract. The caller materializes its keep decision into those -rows each round, then ``compact`` packs them into per-family move sources -and fires the native target and draft launches. This module knows -cache-family geometry and the decision format only; the contract's launch -tuples are private to it. +move offsets ride the caller's staged rows) once per geometry and returns +opaque launch plans. The caller materializes its keep decision into the +agreed rows each round, then ``compact`` loops the plans: each packs its +family's move sources and fires its native launches. This module knows +cache-family geometry and the decision format only; the plans' internals +are private to it. """ from collections import OrderedDict @@ -121,17 +121,23 @@ def _make_move_indices( ) -def _compact_groups( +def _make_compact_launches( entries: List[Tuple[int, torch.Tensor, torch.Tensor]], - pool_keys: Tuple[object, ...], + layer_pool_ids: Tuple[int, ...], + *, + move_indices: torch.Tensor, + move_offsets: torch.Tensor, + destination_bases: torch.Tensor, per_layer_slots: Optional[Dict[int, int]] = None, -) -> Tuple[Dict[str, object], ...]: - """Batch layers into one ``sparse_kv_cache_compact_layers`` launch per uniform V2 pool.""" +) -> Tuple[tuple, ...]: + """Batch layers into one prepacked ``sparse_kv_cache_compact_layers`` + argument tuple per uniform V2 pool (grouped by canonical pool id plus + pool/page-table geometry; native ABI order).""" device = entries[0][1].device grouped = OrderedDict() for layer, pool, page_table in entries: key = ( - pool_keys[layer], + layer_pool_ids[layer], str(pool.dtype), str(pool.device), tuple(int(value) for value in pool.shape[1:]), @@ -139,10 +145,10 @@ def _compact_groups( ) grouped.setdefault(key, []).append((layer, pool, page_table)) - result = [] + launches = [] for group_entries in grouped.values(): layers = tuple(entry[0] for entry in group_entries) - pools = tuple(entry[1] for entry in group_entries) + pools = list(entry[1] for entry in group_entries) page_tables = tuple(entry[2] for entry in group_entries) source_layer_indices = None if per_layer_slots is not None: @@ -151,38 +157,67 @@ def _compact_groups( dtype=torch.int32, device=device, ) - result.append( - dict( - pools=list(pools), - page_table=page_tables[0], - pool_pointers=torch.tensor( + launches.append( + ( + pools, + torch.tensor( [pool.data_ptr() for pool in pools], dtype=torch.int64, device=device, ), - source_layer_indices=source_layer_indices, + page_tables[0], + move_indices, + move_offsets, + destination_bases, + source_layer_indices, ) ) - return tuple(result) + return tuple(launches) -def _launch_tuples( - groups: Tuple[Dict[str, object], ...], - move_indices: torch.Tensor, - move_offsets: torch.Tensor, - destination_bases: torch.Tensor, -) -> Tuple[tuple, ...]: - return tuple( +def _make_pack_launch( + *, + # Decision inputs, packed each round. + kept_ordinal_rows: torch.Tensor, + valid_seq_lens: torch.Tensor, + decision_rows: int, + per_layer_sources: bool, + # Dense move family. + dense_move_offsets: torch.Tensor, + dense_move_indices: torch.Tensor, + # SWA move family (all-or-none). + swa_move_offsets: Optional[torch.Tensor] = None, + swa_move_indices: Optional[torch.Tensor] = None, + swa_window: int = 0, + # Launch geometry. + keep_count: int, + move_capacity: int, + num_kv_heads: int, +) -> Tuple[int, tuple, dict]: + """Assemble one family's prepacked pack launch ``(grid_rows, args, + kwargs)``: stable tensor references plus constexpr dispatch (the frozen + ``_pack_move_sources_kernel`` ABI), built once per geometry so the round + path constructs no Python containers.""" + return ( + decision_rows, ( - group["pools"], - group["pool_pointers"], - group["page_table"], - move_indices, - move_offsets, - destination_bases, - group["source_layer_indices"], - ) - for group in groups + kept_ordinal_rows, + valid_seq_lens, + dense_move_offsets, + dense_move_indices, + swa_move_offsets, + swa_move_indices, + ), + dict( + KEEP_COUNT=keep_count, + DECISION_ROWS=decision_rows, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + PER_LAYER=per_layer_sources, + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, + SWA_WINDOW=swa_window, + ), ) @@ -191,14 +226,16 @@ def init_compaction_buffers( target: Dict[str, object], capacities: Dict[str, int], draft: Optional[Dict[str, object]] = None, -) -> Dict[str, object]: - """Agree on the decision rows and retain one launch contract per geometry. +) -> Tuple[Tuple[Tuple[int, tuple, dict], Tuple[tuple, ...]], ...]: + """Agree on the decision rows and return opaque launch plans per geometry. Move sources must be increasing kept ordinals with destination_bases[request] + move <= source[move] (C++ in-place copy contract). ``target`` carries the resolved dense/SWA grouping inputs from - the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows) - plus the decision inputs :func:`compact` packs each round: + the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows; + ``layer_pool_ids`` the canonical layer -> V2 pool id tuple that indexes + ``kv_block_offsets``; ``swa_destination_bases`` the caller-owned SWA + rebase row) plus the decision inputs :func:`compact` packs each round: ``kept_ordinal_rows`` (``max_requests * decision_rows`` rows of ``keep_count`` int32 kept ordinals, forwarded verbatim), the per-request ``decision_rows`` count (1 = one shared row broadcast over @@ -206,15 +243,17 @@ def init_compaction_buffers( protected tail rides after. ``draft`` is one all-or-none resolved branch (its dense-only moves broadcast the one shared decision row over the draft's own KV heads); ``capacities`` the request/keep/tail capacity - numbers. The returned contract exposes the SWA/draft geometry the caller - stages against; its launch tuples are private to :func:`compact`. + numbers. + + Returns one ``(pack_launch, native_launches)`` plan per compacted cache + family (target, then the optional draft); the plans hold their own + move-index buffers and only :func:`compact` interprets them. """ layer_pools = target["layer_pools"] dense_layers = tuple(int(layer) for layer in target["dense_layers"]) swa_layers = tuple(int(layer) for layer in target["swa_layers"]) - layer_pool_keys = tuple(target["layer_pool_keys"]) + layer_pool_ids = tuple(int(pool_id) for pool_id in target["layer_pool_ids"]) kv_block_offsets = target["kv_block_offsets"] - page_table_slots = target["page_table_slots"] layer_group_representative = target["layer_group_representative"] token_starts = target["token_starts"] dense_move_offsets = target["dense_move_offsets"] @@ -243,89 +282,78 @@ def init_compaction_buffers( ( layer, layer_pools[layer], - kv_block_offsets[page_table_slots[layer_group_representative[layer]], :max_requests, 0], + kv_block_offsets[layer_pool_ids[layer_group_representative[layer]], :max_requests, 0], ) for layer in dense_layers ] - swa_destination_bases = None swa_move_indices = None - swa_entries = [] if not swa_layers: swa_move_offsets = None swa_window = 0 else: swa_window = int(swa_window) - swa_destination_bases = torch.empty_like(token_starts) swa_move_indices = _make_move_indices( (num_kv_heads,), swa_window + protected_tail_capacity, max_requests, device, ) - # SWA layers are staged as their own page-table representatives. - swa_entries = [ - ( - layer, - layer_pools[layer], - kv_block_offsets[page_table_slots[layer], :max_requests, 0], - ) - for layer in swa_layers - ] dense_slots = ( {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer_sources else None ) - has_swa = swa_move_indices is not None - swa_total = int(swa_move_indices.shape[-1]) if has_swa else 0 # Widest per-request move count any staged offsets may express. move_capacity = keep_count + protected_tail_capacity - if has_swa: + if swa_move_indices is not None: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) target_launches = list( - _launch_tuples( - _compact_groups(dense_entries, layer_pool_keys, dense_slots), - dense_move_indices, - dense_move_offsets, - token_starts, + _make_compact_launches( + dense_entries, + layer_pool_ids, + move_indices=dense_move_indices, + move_offsets=dense_move_offsets, + destination_bases=token_starts, + per_layer_slots=dense_slots, ) ) if swa_layers: + # SWA layers stage against their own page-table slots. + swa_entries = [ + ( + layer, + layer_pools[layer], + kv_block_offsets[layer_pool_ids[layer], :max_requests, 0], + ) + for layer in swa_layers + ] target_launches.extend( - _launch_tuples( - _compact_groups(swa_entries, layer_pool_keys), - swa_move_indices, - swa_move_offsets, - swa_destination_bases, + _make_compact_launches( + swa_entries, + layer_pool_ids, + move_indices=swa_move_indices, + move_offsets=swa_move_offsets, + destination_bases=target["swa_destination_bases"], ) ) - target_pack_launch = ( - decision_rows, - ( - kept_ordinal_rows, - valid_seq_lens, - dense_move_offsets, - dense_move_indices, - swa_move_offsets, - swa_move_indices, - ), - dict( - KEEP_COUNT=keep_count, - DECISION_ROWS=decision_rows, - MOVE_CAPACITY=move_capacity, - NUM_KV_HEADS=num_kv_heads, - PER_LAYER=per_layer_sources, - DENSE_TOTAL=int(dense_move_indices.shape[-1]), - SWA_TOTAL=swa_total, - SWA_WINDOW=swa_window, - ), + target_pack_launch = _make_pack_launch( + kept_ordinal_rows=kept_ordinal_rows, + valid_seq_lens=valid_seq_lens, + decision_rows=decision_rows, + per_layer_sources=per_layer_sources, + dense_move_offsets=dense_move_offsets, + dense_move_indices=dense_move_indices, + swa_move_offsets=swa_move_offsets, + swa_move_indices=swa_move_indices, + swa_window=swa_window, + keep_count=keep_count, + move_capacity=move_capacity, + num_kv_heads=num_kv_heads, ) - draft_launches: Tuple[tuple, ...] = () - draft_move_indices = None - draft_pack_launch = None + plans = [(target_pack_launch, tuple(target_launches))] if draft is not None: if decision_rows != 1: raise ValueError( @@ -334,6 +362,7 @@ def init_compaction_buffers( ) draft_layer_pools = draft["layer_pools"] draft_dense_layers = tuple(int(layer) for layer in draft["dense_layers"]) + draft_layer_pool_ids = tuple(int(pool_id) for pool_id in draft["layer_pool_ids"]) draft_tail = int(draft["protected_tail_capacity"]) # Own launch groups: the draft may use a different KV-head count. draft_num_kv_heads = int(draft_layer_pools[draft_dense_layers[0]].shape[2]) @@ -348,71 +377,46 @@ def init_compaction_buffers( layer, draft_layer_pools[layer], draft["kv_block_offsets"][ - draft["page_table_slots"][draft["layer_group_representative"][layer]], + draft_layer_pool_ids[draft["layer_group_representative"][layer]], :max_requests, 0, ], ) for layer in draft_dense_layers ] - draft_launches = _launch_tuples( - _compact_groups(draft_entries, tuple(draft["layer_pool_keys"])), - draft_move_indices, - draft["dense_move_offsets"], - token_starts, + draft_launches = _make_compact_launches( + draft_entries, + draft_layer_pool_ids, + move_indices=draft_move_indices, + move_offsets=draft["dense_move_offsets"], + destination_bases=token_starts, ) - draft_pack_launch = ( - 1, - ( - kept_ordinal_rows, - valid_seq_lens, - draft["dense_move_offsets"], - draft_move_indices, - None, - None, - ), - dict( - KEEP_COUNT=keep_count, - DECISION_ROWS=1, - MOVE_CAPACITY=int(draft_move_indices.shape[-1]) // max_requests, - NUM_KV_HEADS=draft_num_kv_heads, - PER_LAYER=False, - DENSE_TOTAL=int(draft_move_indices.shape[-1]), - SWA_TOTAL=0, - SWA_WINDOW=0, - ), + draft_pack_launch = _make_pack_launch( + kept_ordinal_rows=kept_ordinal_rows, + valid_seq_lens=valid_seq_lens, + decision_rows=1, + per_layer_sources=False, + dense_move_offsets=draft["dense_move_offsets"], + dense_move_indices=draft_move_indices, + keep_count=keep_count, + move_capacity=keep_count + draft_tail, + num_kv_heads=draft_num_kv_heads, ) - - return dict( - # SWA/draft geometry the caller stages against (the public interface). - has_swa=has_swa, - swa_window=swa_window, - swa_destination_bases=swa_destination_bases, - # Per-round SWA destination rebase delta. - swa_rebase_delta=keep_count - swa_window, - draft_move_indices=draft_move_indices, - # Private launch tuples: only compact() interprets these. - target_launches=tuple(target_launches), - draft_launches=draft_launches, - target_pack_launch=target_pack_launch, - draft_pack_launch=draft_pack_launch, - ) + plans.append((draft_pack_launch, draft_launches)) + return tuple(plans) -def compact(compaction: Dict[str, object], request_count: int) -> None: - """Pack each family's move sources and fire its native compacts. +def compact( + plans: Tuple[Tuple[Tuple[int, tuple, dict], Tuple[tuple, ...]], ...], + request_count: int, +) -> None: + """Pack each plan's move sources and fire its native compacts, in plan order. Pure mover: the caller has already materialized its kept ordinals into the agreed decision rows for the active ``request_count`` cohort, and the caller owns the completion ordering of the whole round. """ - rows, pack_args, pack_kwargs = compaction["target_pack_launch"] - _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) - for launch in compaction["target_launches"]: - torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) - draft_pack_launch = compaction["draft_pack_launch"] - if draft_pack_launch is not None: - rows, pack_args, pack_kwargs = draft_pack_launch + for (rows, pack_args, pack_kwargs), native_launches in plans: _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) - for launch in compaction["draft_launches"]: - torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) + for launch in native_launches: + torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index dd0265cf29ae..0b68286a4764 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -111,11 +111,11 @@ def init_eviction_buffers( swa_layers = list(layout["swa_layers"]) swa_window = layout["swa_window"] layer_group_representative = layout["layer_group_representative"] - layer_pool_keys = list(layout["layer_pool_keys"]) + # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. + layer_pool_ids = tuple(layout["layer_pool_ids"]) dense_groups = list(layout["storage_groups"].values()) page_representatives = [group[0] for group in dense_groups] page_representatives.extend(layer for layer in swa_layers if layer not in page_representatives) - page_table_keys = [layer_pool_keys[layer] for layer in page_representatives] num_page_table_slots = int(layout["manager"].num_pools) device = layer_pools[page_representatives[0]].device @@ -143,11 +143,6 @@ def init_eviction_buffers( bufs.page_table_token_capacity = page_table_token_capacity # ---- block-offset staging (target, plus the co-compressed draft) ------- - # Representative layer -> V2 pool id, resolved once here (init-local). - pool_id_by_representative_layer = { - representative: int(key[1]) - for representative, key in zip(page_representatives, page_table_keys) - } bufs.block_offsets_host, bufs.block_offsets_device = _allocate_block_offset_staging( layer_pools[page_representatives[0]], num_pools=num_page_table_slots, @@ -158,14 +153,9 @@ def init_eviction_buffers( bufs.draft_block_offsets_device = None bufs.draft_block_offsets_host = None bufs.draft_protected_tail_capacity = None - draft_page_slots: Dict[int, int] = {} if draft is not None: draft_layout = draft["layout"] draft_representatives = list(draft_layout["pool_representatives"]) - draft_page_slots = { - representative: int(draft_layout["layer_pool_keys"][representative][1]) - for representative in draft_representatives - } draft_anchor_pool = draft_layout["layer_pools"][draft_representatives[0]] # Construction-boundary invariant: the round shares one stream/event # contract, so the draft pools must live on the target device. @@ -201,6 +191,11 @@ def init_eviction_buffers( dense_move_offsets_row = bufs.request_metadata_device[3] swa_move_offsets_row = bufs.request_metadata_device[4] draft_move_offsets_row = bufs.request_metadata_device[5] + # SWA staging geometry, bound once (the compaction plans stay opaque): + # the phase gather rebases each request's SWA destination base in place. + bufs.swa_window = int(swa_window) if swa_layers else None + bufs.swa_destination_bases = torch.empty_like(bufs.token_starts_device) if swa_layers else None + bufs.swa_rebase_delta = keep_count - bufs.swa_window if swa_layers else 0 bufs.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) bufs.mean_sin = torch.empty_like(bufs.mean_cos) bufs.phase = phase @@ -216,7 +211,7 @@ def init_eviction_buffers( bufs.num_freqs = int(num_freqs) bufs.tokens_per_block = int(tokens_per_block) _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} - dense_layer_slots = [pool_id_by_representative_layer[_rep_of[layer]] for layer in dense_layers] + dense_layer_slots = [layer_pool_ids[_rep_of[layer]] for layer in dense_layers] seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( bufs.num_layers ) @@ -296,6 +291,7 @@ def init_eviction_buffers( mean_sin=bufs.mean_sin, freq_scale_sq=freq_scale_sq, output=bufs.score_scratch, + union_scores=bufs.union_scores, enable_partial_stats=union, ) logger.info( @@ -361,7 +357,7 @@ def init_eviction_buffers( (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - # ---- compaction contract + decision-materialization prebinds ------------ + # ---- compaction plans + decision-materialization prebinds --------------- per_layer = eviction_mode == "per_layer_perhead" draft_contract = None if draft is not None: @@ -370,28 +366,28 @@ def init_eviction_buffers( layer_pools=draft_layout["layer_pools"], dense_layers=list(draft_layout["dense_layers"]), layer_group_representative=draft_layout["layer_group_representative"], - layer_pool_keys=list(draft_layout["layer_pool_keys"]), + layer_pool_ids=tuple(draft_layout["layer_pool_ids"]), kv_block_offsets=bufs.draft_block_offsets_device, - page_table_slots=draft_page_slots, dense_move_offsets=draft_move_offsets_row, protected_tail_capacity=int(draft["protected_tail_capacity"]), ) - contract = init_compaction_buffers( + # Opaque launch plans: only compact() interprets them. + bufs.compaction_plan = init_compaction_buffers( target=dict( layer_pools=layer_pools, dense_layers=list(dense_layers), swa_layers=list(swa_layers), swa_window=swa_window, layer_group_representative=layer_group_representative, - layer_pool_keys=list(layer_pool_keys), + layer_pool_ids=layer_pool_ids, kv_block_offsets=bufs.block_offsets_device, - page_table_slots=pool_id_by_representative_layer, token_starts=bufs.token_starts_device, + swa_destination_bases=bufs.swa_destination_bases, # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, swa_move_offsets=swa_move_offsets_row, per_layer_sources=per_layer, - # The decision rows the contract packs into move sources. + # The decision rows the plans pack into move sources. kept_ordinal_rows=bufs.kept_ordinal_rows, decision_rows=bufs.selection_rows_per_request, valid_seq_lens=bufs.valid_seq_lens_device, @@ -403,9 +399,6 @@ def init_eviction_buffers( ), draft=draft_contract, ) - bufs.compaction_plan = contract - bufs.swa_destination_bases = contract["swa_destination_bases"] - bufs.swa_rebase_delta = contract["swa_rebase_delta"] # The decision side: the settle launch materializes the kept-ordinal rows. bufs.settle_args = ( bufs.selection_scores_rows, @@ -480,8 +473,8 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: tails = [int(item["protected_tail"]) for item in prepared] dense = padded_offsets([bufs.keep_count + tail for tail in tails]) swa = None - if bufs.compaction_plan["has_swa"]: - swa = padded_offsets([int(bufs.compaction_plan["swa_window"]) + tail for tail in tails]) + if bufs.swa_window is not None: + swa = padded_offsets([bufs.swa_window + tail for tail in tails]) draft = None if bufs.draft_protected_tail_capacity is not None: draft = padded_offsets( @@ -596,15 +589,15 @@ def execute_eviction_round( num_warps=1, ) if union: - bufs.runner.launch_union_fusion( - request_count, bufs.mean_cos, bufs.mean_sin, bufs.union_scores[:request_count] - ) + # The runner's ctor bound the mean/union buffers; the launch + # takes only the active cohort size. + bufs.runner.launch_union_fusion(request_count) columns = min(bufs.union_scores.shape[1], bufs.selection_scores_rows.shape[1]) bufs.selection_scores_rows[:request_count, :columns].copy_( bufs.union_scores[:request_count, :columns] ) else: - bufs.runner.launch(request_count, bufs.mean_cos, bufs.mean_sin) + bufs.runner.launch(request_count) # Gather each decode window into the [request, layer, head, token] layout of the reduces. group_size = bufs.num_q_heads // bufs.num_kv_heads num_segments = request_count * bufs.num_layers @@ -1171,7 +1164,6 @@ def _runtime_kv_layout(self) -> Dict[str, object]: manager, global_layers, dense_layers=dense_layers, - dense_storage_groups=self._dense_layer_pool_groups(dense_layers, global_layers), swa_layers=swa_layers, swa_window=swa_window, what="", @@ -1185,12 +1177,10 @@ def _build_runtime_kv_layout( global_layers: List[int], *, dense_layers: List[int], - dense_storage_groups: Optional[Dict[object, List[int]]], swa_layers: List[int], swa_window: Optional[int], what: str, ) -> Dict[str, object]: - num_layers = len(global_layers) maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] if any(pool is None for pool in maybe_layer_pools): missing = [ @@ -1198,16 +1188,16 @@ def _build_runtime_kv_layout( ] raise RuntimeError(f"Missing {what}KV pools for attention layers {missing}") layer_pools = [pool for pool in maybe_layer_pools if pool is not None] - all_layers = list(range(num_layers)) - layer_pool_keys = tuple( - self._page_table_pool_keys(all_layers, global_layers, manager=manager) - ) - all_storage_groups: Dict[object, List[int]] = {} - for layer, pool_key in zip(all_layers, layer_pool_keys): - all_storage_groups.setdefault(pool_key, []).append(layer) - storage_groups = ( - dense_storage_groups if dense_storage_groups is not None else all_storage_groups - ) + # Canonical pool IDs, resolved once; every grouping derives from them. + layer_pool_ids = self._page_table_pool_ids(manager, global_layers) + all_storage_groups: Dict[int, List[int]] = {} + for layer, pool_id in enumerate(layer_pool_ids): + all_storage_groups.setdefault(pool_id, []).append(layer) + # Scored/compacted groups cover the dense layers only; SWA layers + # stage and compact as their own representatives. + storage_groups: Dict[int, List[int]] = {} + for layer in dense_layers: + storage_groups.setdefault(layer_pool_ids[layer], []).append(layer) layer_group_representative = { layer: layers[0] for layers in storage_groups.values() for layer in layers } @@ -1221,7 +1211,7 @@ def _build_runtime_kv_layout( swa_window=swa_window, storage_groups=storage_groups, layer_group_representative=layer_group_representative, - layer_pool_keys=layer_pool_keys, + layer_pool_ids=layer_pool_ids, pool_representatives=pool_representatives, pool_page_counts=tuple( int(layer_pools[layer].shape[0]) for layer in pool_representatives @@ -1252,7 +1242,6 @@ def _draft_runtime_kv_layout(self) -> Dict[str, object]: manager, global_layers, dense_layers=list(range(len(global_layers))), - dense_storage_groups=None, swa_layers=[], swa_window=None, what="draft ", @@ -1366,34 +1355,19 @@ def _buffers_for( self._buffers = bufs return bufs - def _page_table_pool_keys( - self, - local_layers: List[int], + @staticmethod + def _page_table_pool_ids( + manager: KVCacheManagerV2, global_layers: List[int], - manager: Optional[KVCacheManagerV2] = None, - ) -> List[object]: - if manager is None: - manager = self.kv_cache_manager + ) -> Tuple[int, ...]: + """Canonical local layer -> V2 pool id tuple (the staged plane slots). + + V2 owns the mapping; its own lookup errors are the precise ones.""" layer_offsets = manager.layer_offsets layer_to_pool = manager.layer_to_pool_mapping_dict - # V2 owns the mapping; its own lookup errors are the precise ones. - return [ - ("pool", int(layer_to_pool[layer_offsets[global_layers[layer]]])) - for layer in local_layers - ] - - def _dense_layer_pool_groups( - self, - dense_layers: List[int], - global_layers: List[int], - ) -> Dict[object, List[int]]: - groups: Dict[object, List[int]] = {} - for layer, pool_key in zip( - dense_layers, - self._page_table_pool_keys(dense_layers, global_layers), - ): - groups.setdefault(pool_key, []).append(layer) - return groups + return tuple( + int(layer_to_pool[layer_offsets[global_layer]]) for global_layer in global_layers + ) def _evict_requests( self, 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 index 9de16a53ac7e..02b311cf38a8 100644 --- 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 @@ -1382,6 +1382,7 @@ def __init__( mean_sin: torch.Tensor, freq_scale_sq: torch.Tensor, output: torch.Tensor, + union_scores: torch.Tensor | None = None, enable_partial_stats: bool = False, ) -> None: # Pool shape [pages, K/V, heads, tokens, dim] of the anchor scored layer. @@ -1395,8 +1396,13 @@ def __init__( self.width = int(seq_len) self.num_q_heads = int(num_q_heads) self.num_kv_heads = num_kv_heads + self.device = output.device self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) self.enable_partial_stats = bool(enable_partial_stats) + if self.enable_partial_stats and union_scores is None: + raise ValueError( + "TriAttention union fusion requires the persistent union_scores output" + ) # One [stats_row, page_shard, {count, mean, m2}] record array. partial_stats_elements = ( max_requests * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * STATS_FIELDS @@ -1443,16 +1449,18 @@ def __init__( _to_cute(anchor_pool), _to_cute(self.descriptors, assumed_align=128), ) + # Ctor-bound persistent launch operands (one from_dlpack wrap each): + # the round path refreshes their contents in place and launches with + # the active cohort size only. + self._cute_mean_cos = _to_cute(mean_cos.view(-1)) + self._cute_mean_sin = _to_cute(mean_sin.view(-1)) + self._cute_union_scores = ( + _to_cute(union_scores.view(-1)) if self.enable_partial_stats else None + ) self._compiled: dict[int, object] = {} self._compiled_stats: dict[int, object] = {} self._compiled_normalize_union: dict[int, object] = {} self._page_shards: dict[int, int] = {} - compile_output_rows = max_requests if self.enable_partial_stats else 1 - self._normalize_union_compile_output = torch.empty( - (compile_output_rows, self.width), - dtype=torch.float32, - device=output.device, - ) self._cute_selection_prefix = ( _to_cute(output), _to_cute(valid_seq_lens, assumed_align=4), @@ -1460,8 +1468,6 @@ def __init__( _to_cute(token_starts, assumed_align=4), ) self._cute_partial_stats = _to_cute(self.partial_stats) - # Launch-path from_dlpack wraps, cached per persistent buffer identity. - self._cute_launch_cache: dict[str, tuple[torch.Tensor, cute.Tensor]] = {} static_geometry = ( max_requests, num_layers, @@ -1520,8 +1526,8 @@ def __init__( compiled = cute.compile( kernel, *self._cute_prefix, - _to_cute(mean_cos.view(-1)), - _to_cute(mean_sin.view(-1)), + self._cute_mean_cos, + self._cute_mean_sin, *self._cute_tail, cutlass.Int32(1), stream, @@ -1564,7 +1570,7 @@ def __init__( static_geometry, tensor_specs, config_key, - _tensor_spec(self._normalize_union_compile_output), + _tensor_spec(union_scores), _tensor_spec(self.partial_stats), ) with _COMPILE_LOCK: @@ -1589,7 +1595,7 @@ def __init__( kernel, _to_cute(self.partial_stats), *self._cute_selection_prefix, - _to_cute(self._normalize_union_compile_output.view(-1)), + self._cute_union_scores, cutlass.Int32(1), stream, ) @@ -1597,45 +1603,28 @@ def __init__( compiled_configs[config_key] = compiled_selection self._compiled_normalize_union[request_count] = compiled_selection - def _cute_cached(self, key: str, tensor: torch.Tensor) -> cute.Tensor: - """One from_dlpack per persistent buffer; rewrap only on identity change.""" - cached = self._cute_launch_cache.get(key) - if cached is not None and cached[0] is tensor: - return cached[1] - wrapped = _to_cute(tensor.view(-1)) - self._cute_launch_cache[key] = (tensor, wrapped) - return wrapped - - def launch( - self, - request_count: int, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - ) -> None: - """Launch the CuTe score kernel on the current PyTorch stream.""" - stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) + def launch(self, request_count: int) -> None: + """Launch the CuTe score kernel on the current PyTorch stream over the + active cohort (every tensor operand is ctor-bound).""" + stream = cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream) self._compiled[request_count]( *self._cute_prefix, - self._cute_cached("mean_cos", mean_cos), - self._cute_cached("mean_sin", mean_sin), + self._cute_mean_cos, + self._cute_mean_sin, *self._cute_tail, request_count, stream, ) - def launch_union_fusion( - self, - request_count: int, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - union_scores: torch.Tensor, - ) -> None: - """Launch score plus stats followed by normalized union reduction.""" - stream = cuda.CUstream(torch.cuda.current_stream(mean_cos.device).cuda_stream) + def launch_union_fusion(self, request_count: int) -> None: + """Launch score plus stats followed by normalized union reduction over + the active cohort (every tensor operand is ctor-bound; the union rows + land in the ctor-bound ``union_scores``).""" + stream = cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream) self._compiled_stats[request_count]( *self._cute_prefix, - self._cute_cached("mean_cos", mean_cos), - self._cute_cached("mean_sin", mean_sin), + self._cute_mean_cos, + self._cute_mean_sin, *self._cute_tail, request_count, stream, @@ -1643,7 +1632,7 @@ def launch_union_fusion( self._compiled_normalize_union[request_count]( self._cute_partial_stats, *self._cute_selection_prefix, - self._cute_cached("union_scores", union_scores), + self._cute_union_scores, request_count, stream, ) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 949e307cbe62..bf4b20a4927f 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -115,8 +115,10 @@ def make_ramp_pools( def build_compaction(**overrides): """``init_compaction_buffers`` with the suite's 2-layer defaults: translates ``eviction_mode`` into ``per_layer_sources``/``decision_rows``, - allocates the caller-owned move-offset rows (capacity cumsum), and hands - the test's pre-settled ``kept_token_ordinals`` in as the decision rows.""" + allocates the caller-owned move-offset rows (capacity cumsum) and SWA + destination bases, and hands the test's pre-settled + ``kept_token_ordinals`` in as the decision rows. Returns the opaque + ``plans`` plus a test-side mirror of the caller-owned inputs.""" from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers args = dict( @@ -124,8 +126,7 @@ def build_compaction(**overrides): dense_layers=[0, 1], swa_layers=[], layer_group_representative={0: 0, 1: 1}, - layer_pool_keys=[("dense", 0), ("dense", 0)], - page_table_slots={0: 0, 1: 0}, + layer_pool_ids=[0, 0], request_count=2, decode_keep_count=4, swa_window=None, @@ -140,16 +141,16 @@ def build_compaction(**overrides): tail = int(args.get("protected_tail_capacity", 0)) draft_tail = int(args.get("draft_protected_tail_capacity") or 0) has_draft = bool(args.get("draft_layers")) + has_swa = bool(args["swa_layers"]) device = args["layer_pools"][args["dense_layers"][0]].device - swa_window = int(args["swa_window"] or 0) if args["swa_layers"] else 0 + swa_window = int(args["swa_window"] or 0) if has_swa else 0 + swa_destination_bases = torch.empty_like(args["prompt_offsets"]) if has_swa else None def capacity_offsets(count): return torch.arange(0, (request_count + 1) * count, count, dtype=torch.int32, device=device) args.setdefault("dense_move_offsets", capacity_offsets(keep_count + tail)) - args.setdefault( - "swa_move_offsets", capacity_offsets(swa_window + tail) if args["swa_layers"] else None - ) + args.setdefault("swa_move_offsets", capacity_offsets(swa_window + tail) if has_swa else None) if has_draft: args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) num_kv_heads = int(args["layer_pools"][args["dense_layers"][0]].shape[2]) @@ -162,23 +163,22 @@ def capacity_offsets(count): layer_pools=args["draft_layer_pools"], dense_layers=args["draft_layers"], layer_group_representative=args["draft_layer_group_representative"], - layer_pool_keys=args["draft_layer_pool_keys"], + layer_pool_ids=args["draft_layer_pool_ids"], kv_block_offsets=args["draft_kv_block_offsets"], - page_table_slots=args["draft_page_table_slots"], dense_move_offsets=args["draft_move_offsets"], protected_tail_capacity=draft_tail, ) - compaction = init_compaction_buffers( + plans = init_compaction_buffers( target=dict( layer_pools=args["layer_pools"], dense_layers=args["dense_layers"], swa_layers=args["swa_layers"], swa_window=args["swa_window"], layer_group_representative=args["layer_group_representative"], - layer_pool_keys=args["layer_pool_keys"], + layer_pool_ids=args["layer_pool_ids"], kv_block_offsets=args["kv_block_offsets"], - page_table_slots=args["page_table_slots"], token_starts=args["prompt_offsets"], + swa_destination_bases=swa_destination_bases, dense_move_offsets=args["dense_move_offsets"], swa_move_offsets=args["swa_move_offsets"], per_layer_sources=per_layer, @@ -193,10 +193,11 @@ def capacity_offsets(count): ), draft=draft, ) - # Test-side mirror of the construction inputs: production reads only the - # contract's public keys; the standalone helpers here need the - # caller-owned move-offset rows back. - compaction.update( + # Opaque plans plus a test-side mirror of the caller-owned construction + # inputs (production binds the same values on its buffer namespace); the + # standalone helpers here need the move-offset rows and SWA staging back. + return dict( + plans=plans, prompt_offsets=args["prompt_offsets"], request_count=request_count, decode_keep_count=keep_count, @@ -205,14 +206,17 @@ def capacity_offsets(count): dense_move_offsets=args["dense_move_offsets"], swa_move_offsets=args["swa_move_offsets"], draft_move_offsets=args["draft_move_offsets"] if has_draft else None, + has_swa=has_swa, + swa_window=swa_window, + swa_destination_bases=swa_destination_bases, + swa_rebase_delta=keep_count - swa_window, ) - return compaction def run_compaction(compaction): """Replica of the round's move stage in production order: SWA - destination rebase, then ``compact`` packs the decision rows into move - sources and fires the native target and draft moves.""" + destination rebase, then ``compact`` loops the opaque plans (each packs + its decision rows into move sources and fires its native moves).""" from tensorrt_llm._torch.kv_cache_compression.compaction import compact if compaction["swa_destination_bases"] is not None: @@ -221,7 +225,7 @@ def run_compaction(compaction): compaction["swa_rebase_delta"], out=compaction["swa_destination_bases"], ) - compact(compaction, compaction["request_count"]) + compact(compaction["plans"], compaction["request_count"]) def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): @@ -229,7 +233,7 @@ def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): staging = SimpleNamespace() staging.max_requests = max_requests staging.keep_count = 4 - staging.compaction_plan = {"has_swa": False} + staging.swa_window = None staging.draft_protected_tail_capacity = None staging.block_offsets_ready_event = torch.cuda.Event() staging.compaction_done_event = torch.cuda.Event() @@ -265,7 +269,6 @@ def make_buffer_stubs(manager, *, decode_width=260): manager._phase = {"rows": 8} manager.calibration = {"omega": torch.ones(2)} manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) - manager._page_table_pool_keys = mock.Mock(return_value=[("pool", 0)]) pool = torch.empty(8, 2, 1, 4, 4) layout = dict( manager=SimpleNamespace(num_pools=1), @@ -276,7 +279,7 @@ def make_buffer_stubs(manager, *, decode_width=260): swa_window=None, storage_groups={0: [0, 1]}, layer_group_representative={0: 0, 1: 0}, - layer_pool_keys=(("pool", 0), ("pool", 0)), + layer_pool_ids=(0, 0), ) buffers = SimpleNamespace( decode_width=decode_width, @@ -479,13 +482,13 @@ def make_cute_buffers( page_table_token_capacity=None, protected_tail_capacity=0, storage_groups=None, - layer_pool_keys=None, + layer_pool_ids=None, normalize_scores=True, ): """Real eviction buffers over the one-shared-slot default layout; split reference legs use ``eviction_mode="per_head"`` over the same pools. - ``storage_groups``/``layer_pool_keys`` override the page-table grouping - (keys are ``(name, slot)`` tuples; the slot is ``key[1]``).""" + ``storage_groups``/``layer_pool_ids`` override the page-table grouping + (``layer_pool_ids`` is the canonical per-layer V2 pool id list).""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( init_eviction_buffers, ) @@ -497,11 +500,11 @@ def make_cute_buffers( if decode_width is None: decode_width = seq_len if storage_groups is None: - storage_groups = {("pool", 0): list(range(num_layers))} - if layer_pool_keys is None: - layer_pool_keys = [("pool", 0)] * num_layers + storage_groups = {0: list(range(num_layers))} + if layer_pool_ids is None: + layer_pool_ids = [0] * num_layers layout = dict( - manager=SimpleNamespace(num_pools=max(key[1] for key in layer_pool_keys) + 1), + manager=SimpleNamespace(num_pools=max(layer_pool_ids) + 1), layer_pools=layer_pools, dense_layers=list(range(num_layers)), swa_layers=[], @@ -510,7 +513,7 @@ def make_cute_buffers( layer_group_representative={ layer: layers[0] for layers in storage_groups.values() for layer in layers }, - layer_pool_keys=layer_pool_keys, + layer_pool_ids=layer_pool_ids, ) return init_eviction_buffers( eviction_mode=eviction_mode, @@ -558,10 +561,14 @@ def launch_split_scores( bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin ): """The production score-only leg plus the decode-window gather - (``execute_eviction_round``'s per-head sequence, parameterized by count).""" + (``execute_eviction_round``'s per-head sequence, parameterized by count). + Test mean phases load into the runner's ctor-bound buffers, exactly like + the production in-place gather refresh.""" stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) + bufs.mean_cos[:request_count].copy_(mean_cos[:request_count]) + bufs.mean_sin[:request_count].copy_(mean_sin[:request_count]) assert request_count in bufs.runner._compiled - bufs.runner.launch(request_count, mean_cos, mean_sin) + bufs.runner.launch(request_count) num_segments = request_count * bufs.num_layers group_size = bufs.num_q_heads // bufs.num_kv_heads source = ( 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 index ba671184af71..dc4d49d3a006 100644 --- 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 @@ -22,13 +22,18 @@ def _launch_union_fusion( bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out ): - """The fused score+stats+normalized-union pipeline (THE union path).""" + """The fused score+stats+normalized-union pipeline (THE union path). + Test mean phases load into the runner's ctor-bound buffers; the fused + rows land in the ctor-bound ``bufs.union_scores`` and copy out.""" _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) + bufs.mean_cos[:request_count].copy_(mean_cos[:request_count]) + bufs.mean_sin[:request_count].copy_(mean_sin[:request_count]) assert ( request_count in bufs.runner._compiled_stats and request_count in bufs.runner._compiled_normalize_union ) - bufs.runner.launch_union_fusion(request_count, mean_cos, mean_sin, union_out[:request_count]) + bufs.runner.launch_union_fusion(request_count) + union_out[:request_count].copy_(bufs.union_scores[:request_count]) def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: 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 index 3be09125c566..b5cb48f2be51 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -60,7 +60,7 @@ def _launched_draft_compaction(draft_protected_tails): compaction = _build_compaction( layer_pools=target_pools, - layer_pool_keys=[("pool", 0), ("pool", 0)], + layer_pool_ids=[0, 0], kept_token_ordinals=keep.to(torch.int32), valid_sequence_lengths=torch.tensor(valid_seq_lens, dtype=torch.int32, device=device), kv_block_offsets=_encode_block_offsets(target_tables), @@ -69,10 +69,9 @@ def _launched_draft_compaction(draft_protected_tails): draft_layer_pools=[draft_pool], draft_layers=[0], draft_layer_group_representative={0: 0}, - draft_layer_pool_keys=[("draft_pool", 0)], + draft_layer_pool_ids=[0], draft_protected_tail_capacity=max(draft_protected_tails), draft_kv_block_offsets=_encode_block_offsets(draft_tables), - draft_page_table_slots={0: 0}, ) _set_protected_tails(compaction, target_protected_tails, draft_protected_tails) _run_compaction(compaction) @@ -103,7 +102,6 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect prompt_len = built.prompt_len expected_offsets = [0] - expected_moves = [] for request in range(built.request_count): valid = built.valid_seq_lens[request] # Target dense layers compact the union keep set plus the target tail. @@ -154,22 +152,13 @@ def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protect before[:, head].index_select(1, draft_source), ) - expected_moves.append(draft_source.to(torch.int32)) expected_offsets.append(expected_offsets[-1] + int(draft_source.numel())) - # Packed indices must match the same broadcast-plus-tail oracle: the - # test-owned draft offsets row and the contract's draft move sources. - expected_row = torch.cat(expected_moves) + # The test-owned draft move-offset row must match the broadcast-plus-tail + # oracle; the packed move sources themselves are covered byte-exactly by + # the pool assertions above (the ramp payload makes every wrong move land + # on different bytes) and by the pack-kernel oracle suite. assert built.compaction["draft_move_offsets"].cpu().tolist() == expected_offsets - draft_indices = built.compaction["draft_move_indices"] - # Capacity-sized buffer; this round's moves pack at the front. - capacity_total = built.request_count * ( - int(built.keep.shape[1]) + max(built.draft_protected_tails) - ) - assert draft_indices.shape == (int(built.draft_pool.shape[2]), capacity_total) - for head in range(int(draft_indices.shape[0])): - # Union mode broadcasts one keep set over every draft KV head. - assert torch.equal(draft_indices[head, : expected_offsets[-1]], expected_row) def test_execute_eviction_round_orders_both_manager_streams(): @@ -185,7 +174,8 @@ def test_execute_eviction_round_orders_both_manager_streams(): max_requests=8, keep_count=4, eviction_mode="union", - compaction_plan={"has_swa": False}, + swa_window=None, + compaction_plan=(), draft_protected_tail_capacity=1, copy_pending=False, staging_reuse_event=mock.Mock(), @@ -400,7 +390,7 @@ def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): dense_layers=[], layer_group_representative={}, pool_representatives=(), - layer_pool_keys=(), + layer_pool_ids=(), pool_page_counts=(4,), ) ) @@ -432,7 +422,7 @@ def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): assert kwargs["capacities"]["keep_count"] == manager.budget assert kwargs["phase"] is manager._phase assert kwargs["layout"] is layout - assert list(kwargs["layout"]["layer_pool_keys"]) == list(layout["layer_pool_keys"]) + assert list(kwargs["layout"]["layer_pool_ids"]) == list(layout["layer_pool_ids"]) # A second round within the resident capacities reuses the buffers # (and with them the compaction launch data they carry). diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index e8d16cdd4b58..9a291e2d59ee 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -152,7 +152,7 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): swa_window=None, storage_groups={0: [0, 2], 1: [1]}, layer_group_representative={0: 0, 1: 1, 2: 0}, - layer_pool_keys=(0, 1, 0), + layer_pool_ids=(0, 1, 0), # These are local layer slots. Layer 2 shares layer 0's pool. pool_representatives=(0, 1), pool_page_counts=(4, 8), @@ -735,7 +735,7 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): offsets_buffers = SimpleNamespace( max_requests=8, keep_count=4, - compaction_plan={"has_swa": False}, + swa_window=None, draft_protected_tail_capacity=None, ) dense, swa, draft = _cohort_move_offsets(offsets_buffers, prepared) @@ -809,8 +809,8 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self, request_coun decode_width=seq_len - prompt_len, keep_count=4, # One storage group and page-table slot per layer (distinct pools). - storage_groups={("pool", layer): [layer] for layer in layer_order}, - layer_pool_keys=[("pool", layer) for layer in layer_order], + storage_groups={layer: [layer] for layer in layer_order}, + layer_pool_ids=list(layer_order), normalize_scores=False, ) valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) 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 index 1f1a8b846c9a..c4497cc66c05 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -490,13 +490,15 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): decode_width=bucket_capacity, keep_count=keep_count, storage_groups={ - ("pool", 0): dense_groups[0], - ("pool", 1): dense_groups[1], + 0: dense_groups[0], + 1: dense_groups[1], }, - layer_pool_keys=[("pool", 0), ("pool", 1), ("pool", 0)], + layer_pool_ids=[0, 1, 0], normalize_scores=False, ) - assert bufs.compaction_plan["has_swa"] is False + # No SWA in this layout: no window, no rebase row for the phase gather. + assert bufs.swa_window is None + assert bufs.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 bufs.block_offsets_host.shape == bufs.block_offsets_device.shape @@ -771,11 +773,11 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): dense_layers=[0], swa_layers=[1], layer_group_representative={0: 0}, - layer_pool_keys=[("dense", 0), ("swa", 0)], + # Dense layer 0 stages in plane 0, the SWA layer in its own plane 1. + layer_pool_ids=[0, 1], kept_token_ordinals=keep.to(torch.int32), valid_sequence_lengths=valid_seq_lens, kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), - page_table_slots={0: 0, 1: 1}, prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), swa_window=2, protected_tail_capacity=max(protected_tails), From d2c119f8ac62ec7245de6ea116d80ef8694b75a1 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 05:39:10 -0700 Subject: [PATCH 126/178] [None][refactor] Move the mean-phase table builder into its owner module Signed-off-by: tianruih --- .../triattention/triattention.py | 30 +++++++++++++++++- .../triattention/triattention_kernels.py | 31 +------------------ .../_torch/kv_cache_compression/conftest.py | 2 +- 3 files changed, 31 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 0b68286a4764..fa6e23382a50 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -43,7 +43,6 @@ from .triattention_kernels import ( _gather_mean_phase_kernel, _settle_ties_kernel, - grow_mean_phase_table, prepare_per_head_scores, ) @@ -84,6 +83,35 @@ def _allocate_block_offset_staging( return host, device_table +_MEAN_PHASE_MAX_ROWS = 1 << 24 + + +def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: + """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" + rows = int(rows) + if rows <= phase["rows"]: + return + if rows > _MEAN_PHASE_MAX_ROWS: + raise ValueError(f"a {rows}-row mean-phase table exceeds the exact-FP32 position range") + target = 1 + while target < rows: + target *= 2 + target = min(max(target, 2 * phase["rows"]), _MEAN_PHASE_MAX_ROWS) + omega = phase["omega"] + positions = torch.arange(target, device=omega.device, dtype=torch.float32) + cos_table = torch.zeros((target, omega.numel()), dtype=torch.float32, device=omega.device) + sin_table = torch.zeros_like(cos_table) + # Fixed summation order keeps the table bit-stable across rebuilds. + for offset in phase["offset_values"]: + angle = torch.outer(positions + offset, omega) + cos_table += torch.cos(angle) + sin_table += torch.sin(angle) + scale = 1.0 / len(phase["offset_values"]) + phase["cos"] = cos_table.mul_(scale) + phase["sin"] = sin_table.mul_(scale) + phase["rows"] = target + + def init_eviction_buffers( *, eviction_mode: str, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index ec19091f617b..6b9bd397e1dd 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -5,17 +5,14 @@ from __future__ import annotations -from typing import Dict - import torch import triton import triton.language as tl -# ---- Mean-phase table: RoPE-style position table of mean trig phases ---- +# ---- Mean-phase gather: per-request phase-row fetch + width derivation ---- # Positions past this row count are not exactly representable in fp32. -_MEAN_PHASE_MAX_ROWS = 1 << 24 # Score z-normalization epsilon; must stay a plain float (the CuTe DSL traces it). STD_EPSILON = 1e-6 @@ -57,32 +54,6 @@ def _gather_mean_phase_kernel( tl.store(swa_destination_bases + request, token_start + swa_rebase_delta) -def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: - """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" - rows = int(rows) - if rows <= phase["rows"]: - return - if rows > _MEAN_PHASE_MAX_ROWS: - raise ValueError(f"a {rows}-row mean-phase table exceeds the exact-FP32 position range") - target = 1 - while target < rows: - target *= 2 - target = min(max(target, 2 * phase["rows"]), _MEAN_PHASE_MAX_ROWS) - omega = phase["omega"] - positions = torch.arange(target, device=omega.device, dtype=torch.float32) - cos_table = torch.zeros((target, omega.numel()), dtype=torch.float32, device=omega.device) - sin_table = torch.zeros_like(cos_table) - # Fixed summation order keeps the table bit-stable across rebuilds. - for offset in phase["offset_values"]: - angle = torch.outer(positions + offset, omega) - cos_table += torch.cos(angle) - sin_table += torch.sin(angle) - scale = 1.0 / len(phase["offset_values"]) - phase["cos"] = cos_table.mul_(scale) - phase["sin"] = sin_table.mul_(scale) - phase["rows"] = target - - # ---- Selection: combine scores per mode, then finalize the top-k set ---- diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index bf4b20a4927f..9e5370bb5ff0 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -449,7 +449,7 @@ def torch_tri_score_oracle( def make_phase_table(offsets, omega, initial_rows): """Build the mean-phase table dict exactly like the product's inlined form and grow it to cover positions ``[0, initial_rows)``.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( grow_mean_phase_table, ) From f9615bb6199c222835119f52ce0a63a5b7d20b00 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 06:10:56 -0700 Subject: [PATCH 127/178] [None][test] Test tranche 4: post-restructure orphan leg, matrix squeeze, thop family convergence (knife 31) Signed-off-by: tianruih --- .../test_triattention_cute_union_fusion.py | 8 +- .../test_triattention_draft_cocompaction.py | 7 +- .../test_triattention_fused_settle_pack.py | 23 +++-- .../test_triattention_pipeline.py | 68 ++++----------- .../serial/test_sparse_kv_cache_compact.py | 86 ++++--------------- 5 files changed, 58 insertions(+), 134 deletions(-) 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 index dc4d49d3a006..c692b7c425c4 100644 --- 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 @@ -72,11 +72,9 @@ def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tenso # only the real heads' rows, and the union finalizer maps head rows # onto the padded score planes. (128, 32, 4, 0, None), - # Mixed-prompt cohorts: each request scores its own window (one - # start mid-tile, one page-aligned) — the case the fused pipeline - # previously declined. - (32, 32, 8, [37, 128], [250, 198]), - (32, 64, 4, [37, 128], None), + # Mixed-prompt cohort (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]), ], ) 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 index b5cb48f2be51..7fea8b84fb92 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -95,9 +95,10 @@ def _launched_draft_compaction(draft_protected_tails): ) -@pytest.mark.parametrize("draft_protected_tails", [[1, 1], [1, 2]]) -def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(draft_protected_tails): - built = _launched_draft_compaction(draft_protected_tails=draft_protected_tails) +def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(): + # Ragged draft tails [1, 2] against target tails [2, 1]: one request's + # draft tail below and one above its target, subsuming the uniform row. + built = _launched_draft_compaction(draft_protected_tails=[1, 2]) device = built.device prompt_len = built.prompt_len 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 index 3422841644c2..ccdcff65a479 100644 --- 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 @@ -16,7 +16,8 @@ _settle_ties_kernel, ) -# Both kernels share the settle geometry parameter grid. +# 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), @@ -24,6 +25,18 @@ (350, 300), ] +# Pack rows: one per distinct kernel path combo (BROADCAST = union; PER_LAYER +# only branches inside the non-broadcast SWA store, so it is compiled out +# without SWA); keep 5 vs 300 flips the MOVE_CAPACITY block trip count on +# both BROADCAST sides. +_PACK_PATH_ROWS = [ + ("union", False, 21, 5), + ("union", True, 350, 300), + ("per_head", False, 350, 300), + ("per_head", True, 21, 5), + ("per_layer_perhead", True, 350, 300), +] + def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, keep_count): """Settle in place: threshold = min over non-sentinel provisional @@ -133,7 +146,9 @@ def _make_settle_inputs(request_count, selection_rows, width, keep_count, seed, return scores, row_lengths, prompt_offsets, provisional -@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) +# 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()) @@ -172,9 +187,7 @@ def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): assert torch.equal(output_actual, output_reference), f"kept ordinals differ (seed {seed})" -@pytest.mark.parametrize("has_swa", [False, True]) -@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) -@pytest.mark.parametrize("width,keep_count", _WIDTH_KEEP_CASES) +@pytest.mark.parametrize("eviction_mode,has_swa,width,keep_count", _PACK_PATH_ROWS) def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, width, keep_count): device = torch.device("cuda", torch.cuda.current_device()) request_count, num_layers, num_kv_heads = 3, 2, 2 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 9a291e2d59ee..5bc69ce72652 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -410,7 +410,9 @@ def test_suspended_cache_defers_that_request_pre_launch(self): assert first_state["generation_steps"] == 128 assert second_state["generation_steps"] == 127 - @pytest.mark.parametrize("accepted", [0, 1, 2, 3]) + # ``accepted`` enters the prepared item linearly; the zero and maximal + # boundary rows pin the whole family. + @pytest.mark.parametrize("accepted", [0, 3]) def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): confirmed = 1024 + 4096 + 1 + accepted reserve = 2 @@ -679,55 +681,11 @@ def stage_once(): 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] - def test_staged_page_tables_bypass_per_request_cuda_materialization(self): - manager = _make_triattention(budget=4) - manager.kv_cache_manager.num_extra_kv_tokens = 3 - manager.kv_cache_manager._stream = mock.Mock() - manager.kv_cache_manager.get_batch_cache_indices = mock.Mock( - side_effect=AssertionError("eviction staged page tables per request") - ) - first = _make_request(7, py_prompt_len=3) - second = _make_request(8, py_prompt_len=5) - _set_request_state(manager, 7) - _set_request_state(manager, 8) - - with _mocked_eviction_internals(manager) as internals: - manager._evict_requests( - [ - _make_prepared_item( - first, - request_id=7, - seq_len=8, - prompt_len=3, - expected_keep_count=7, - protected_tail=2, - ), - _make_prepared_item( - second, - request_id=8, - seq_len=10, - prompt_len=5, - expected_keep_count=9, - protected_tail=3, - ), - ] - ) - - # One round-executor call carries the whole cohort (with the target - # and draft managers it orders after the compact launches). - args = internals.execute.call_args - assert args.args[0] is internals.buffers - assert args.args[1] is manager.kv_cache_manager - assert args.args[3] is None - prepared = args.args[2] - assert [item["request_id"] for item in prepared] == [7, 8] - assert [item["round_start"] for item in prepared] == [8, 10] - assert [item["prompt_len"] for item in prepared] == [3, 5] - assert [item["seq_len"] for item in prepared] == [8, 10] - assert args.kwargs == {} - + def test_cohort_move_offsets_stage_keep_plus_tail_and_pad_rows(self): # The derived move offsets stage keep + tail moves per request - # (budget=4 -> [6, 7]); padded rows repeat the final offset. + # (keep_count=4 -> [6, 7]); padded rows past the cohort repeat the + # final offset and contribute no moves. (The executor-call contract + # itself is pinned by the overlap-tail and draft publication tests.) from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _cohort_move_offsets, ) @@ -738,20 +696,26 @@ def test_staged_page_tables_bypass_per_request_cuda_materialization(self): swa_window=None, draft_protected_tail_capacity=None, ) + prepared = [ + _make_prepared_item(request_id=7, seq_len=8, protected_tail=2), + _make_prepared_item(request_id=8, seq_len=10, protected_tail=3), + ] dense, swa, draft = _cohort_move_offsets(offsets_buffers, prepared) assert dense == [0, 6, 13, 13, 13, 13, 13, 13, 13] assert swa is None assert draft is None @requires_sm100 - @pytest.mark.parametrize("request_count", [1, 8]) - def test_fused_score_spans_distinct_storages_and_block_tables(self, request_count): + def test_fused_score_spans_distinct_storages_and_block_tables(self): """ONE launch over layers in DISTINCT storages with DISTINCT block tables (the production V2 shape), checked against the Torch oracle, - then relaunched after a round-start advance and a table rebind.""" + then relaunched after a round-start advance and a table rebind. + (Single-request launches are pinned by the CuTe score oracle's + request-count loop; the distinct-storages property is layer-axis.)""" pytest.importorskip("cutlass") from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module + request_count = 8 device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(20260707 + request_count) max_requests = request_count diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 81bb0aa72d3e..1e065b751613 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -48,9 +48,8 @@ def _make_pools( dtype: torch.dtype, head_dim: int, page_index_scale: int = _PAGE_INDEX_DIVISOR, - pages_per_seq: int = 3, - sequential_pages: bool = False, ) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: + pages_per_seq = 3 num_pages = _BATCH_SIZE * pages_per_seq * page_index_scale // _PAGE_INDEX_DIVISOR shape = ( num_pages, @@ -65,15 +64,9 @@ def _make_pools( for layer in range(num_layers) ] pools = [pool.cuda() for pool in pools_cpu] - if sequential_pages: - raw_page_table = torch.arange( - _BATCH_SIZE * pages_per_seq, dtype=torch.int32, device="cuda" - ).reshape(_BATCH_SIZE, pages_per_seq) - else: - assert pages_per_seq == 3 - raw_pages = [[4, 1, 5], [2, 0, 3]] - assert set(raw_pages[0]).isdisjoint(raw_pages[1]) - raw_page_table = torch.tensor(raw_pages, dtype=torch.int32, device="cuda") + raw_pages = [[4, 1, 5], [2, 0, 3]] + assert set(raw_pages[0]).isdisjoint(raw_pages[1]) + raw_page_table = torch.tensor(raw_pages, dtype=torch.int32, device="cuda") page_table = _encode_k_block_offsets(raw_page_table, page_index_scale) page_tables = [page_table] * num_layers assert page_tables[0].stride(0) == 2 * page_tables[0].shape[1] @@ -181,78 +174,33 @@ def _compact( _SMALL_ROW = [2, 5, 8, 3, 7, 10] -# One byte-equality case per launch shape: head_dim/destination/page-scale -# sweep, per-request bases, 3-D per-layer routing, multi-tile. +# The production-shaped fast-geometry matrix below is the byte-exact anchor +# (both head dims, both page sizes, per-request destination bases, 3-D +# per-layer routing, multi-tile ragged moves). These two rows keep the +# op-level contracts it does not pin: the destination-base-0 (prompt 0) +# boundary and the fixed //2 K-offset decode against a scale-4 encoder. _LAYER_CASES = [ - pytest.param( - dict(head_dim=head_dim, dest=dest, scale=scale), - id=f"bf16_h{head_dim}_dest{dest}_scale{scale}", - ) - for head_dim in (64, 128) - for dest, scale in ((0, 2), (2, 2), (2, 4)) -] + [ - pytest.param( - dict(head_dim=64, num_layers=2, row=[3, 5, 8, 6, 7, 10], dest=[2, 5]), - id="per_request_destination_bases", - ), - pytest.param( - dict( - head_dim=64, - num_layers=2, - dest=2, - indices3d=[ - [[2, 5, 8, 3, 7, 10], [3, 6, 9, 2, 5, 8]], - [[3, 7, 10, 2, 6, 9], [2, 5, 9, 3, 6, 10]], - [[4, 7, 9, 3, 6, 8], [3, 5, 8, 4, 7, 10]], - ], - layer_indices=[2, 0], - ), - id="per_layer_source", - ), - pytest.param( - dict( - head_dim=64, - num_layers=2, - dest=2, - pages_per_seq=24, - sequential_pages=True, - offsets=(0, 40, 75), - row=list(range(40, 80)) + list(range(36, 71)), - ), - id="multiple_tiles", - ), + pytest.param(dict(head_dim=64, dest=0, scale=2), id="bf16_h64_dest0_scale2"), + pytest.param(dict(head_dim=64, dest=2, scale=4), id="bf16_h64_dest2_scale4"), ] @pytest.mark.parametrize("case", _LAYER_CASES) def test_sparse_kv_cache_compact_layers(case): - pools_cpu, pools, page_tables = _make_pools( - case.get("num_layers", 3), - torch.bfloat16, - case["head_dim"], - case.get("scale", _PAGE_INDEX_DIVISOR), - pages_per_seq=case.get("pages_per_seq", 3), - sequential_pages=case.get("sequential_pages", False), - ) + pools_cpu, pools, page_tables = _make_pools(3, torch.bfloat16, case["head_dim"], case["scale"]) page_tables_cpu = [page_table.cpu() for page_table in page_tables] - source_offsets = torch.tensor(case.get("offsets", (0, 3, 6)), dtype=torch.int32) - if "indices3d" in case: - source_indices = torch.tensor(case["indices3d"], dtype=torch.int32) - source_layer_indices = torch.tensor(case["layer_indices"], dtype=torch.int32) - else: - source_row = torch.tensor(case.get("row", _SMALL_ROW), dtype=torch.int32) - source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() - source_layer_indices = None - destination_base = case.get("dest", 0) + source_offsets = torch.tensor((0, 3, 6), dtype=torch.int32) + source_row = torch.tensor(_SMALL_ROW, dtype=torch.int32) + source_indices = source_row.view(1, -1).expand(_NUM_KV_HEADS, -1).contiguous() + destination_base = case["dest"] expected = _reference_compact( pools_cpu, page_tables_cpu, source_indices, source_offsets, destination_base, - source_layer_indices, ) - arguments = _device_arguments(pools, source_indices, source_offsets, source_layer_indices) + arguments = _device_arguments(pools, source_indices, source_offsets) _compact(pools, page_tables, arguments, destination_base) torch.cuda.synchronize() From 53268b371d9dcc98b876f5e7f09a06500f253a91 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 06:58:53 -0700 Subject: [PATCH 128/178] [None][refactor] Constructor takes the whole config; drop the Base prefix (knife 32) Signed-off-by: tianruih --- .../triattention/triattention.py | 24 ++++----- tensorrt_llm/_torch/pyexecutor/_util.py | 6 +-- .../_torch/pyexecutor/resource_manager.py | 2 +- tensorrt_llm/llmapi/llm_args.py | 53 ++++++------------- .../test_kv_cache_compression_manager.py | 42 +++++++-------- .../_torch/kv_cache_compression/conftest.py | 18 +++++-- .../test_triattention_draft_cocompaction.py | 8 +-- .../test_triattention_pipeline.py | 14 ++--- 8 files changed, 77 insertions(+), 90 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index fa6e23382a50..52cdd7e6f243 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -31,7 +31,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.resource_manager import BaseKVCacheCompressionManager +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheCompressionManager from tensorrt_llm._utils import nvtx_range, nvtx_range_debug, prefer_pinned from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( copy_batch_block_offsets_to_device, @@ -49,6 +49,7 @@ if TYPE_CHECKING: from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests + from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig # Required keys for the calibration ``.pt`` consumed by TriAttention. @@ -681,7 +682,7 @@ def execute_eviction_round( draft_manager._stream.wait_event(bufs.compaction_done_event) -class TriAttention(BaseKVCacheCompressionManager): +class TriAttention(KVCacheCompressionManager): """Periodic physical KV eviction driven by trigonometric importance scoring.""" adjusts_generation_kv_length = True @@ -689,21 +690,16 @@ class TriAttention(BaseKVCacheCompressionManager): def __init__( self, kv_cache_manager: KVCacheManagerV2, - budget: int, + config: "TriAttentionKvCacheCompressionConfig", draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, - beta: int = 128, - model_path: Optional[str] = None, - calibration_path: Optional[str] = None, - eviction_mode: str = "union", - normalize_scores: bool = True, ): super().__init__(kv_cache_manager, draft_kv_cache_manager) # budget/beta positivity and the eviction_mode literal are validated at # the config boundary (TriAttentionKvCacheCompressionConfig). - self.budget = budget - self.beta = beta - self.eviction_mode = eviction_mode - self.normalize_scores = bool(normalize_scores) + self.budget = config.budget + self.beta = config.beta + self.eviction_mode = config.eviction_mode + self.normalize_scores = bool(config.normalize_scores) if self.eviction_mode == "union" and not self.normalize_scores: raise ValueError( "TriAttention union eviction requires normalize_scores=True: " @@ -713,8 +709,8 @@ def __init__( # counts decode tokens only (physical KV reclaim requires both). # Calibration is the official TriAttention .pt; TRT-LLM does not # compute calibration. The config boundary requires both paths. - self.model_path = model_path - self.calibration_path = calibration_path + self.model_path = config.model_path + self.calibration_path = config.calibration_path self.calibration: Optional[Dict[str, torch.Tensor]] = None self._calibrated = False self._freq_scale_sq: Optional[torch.Tensor] = None diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 3a57ce241e42..74ba92dba01d 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -62,7 +62,7 @@ use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor -from .resource_manager import (BaseKVCacheCompressionManager, KVCacheManager, +from .resource_manager import (KVCacheCompressionManager, KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType) from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler, @@ -2156,7 +2156,7 @@ def create_kv_cache_compression_manager( config: KvCacheCompressionConfig, kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, -) -> Optional[BaseKVCacheCompressionManager]: +) -> Optional[KVCacheCompressionManager]: """Build the KV-cache compression manager for ``config.algorithm``, or return None if no algorithm matches. @@ -2171,8 +2171,8 @@ def create_kv_cache_compression_manager( return TriAttention( kv_cache_manager, + config=config, draft_kv_cache_manager=draft_kv_cache_manager, - **config.to_manager_kwargs(), ) logger.warning( diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 157e2d7d720a..513007418e69 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2412,7 +2412,7 @@ def _free_blocks(self, block_list: list): # --------------------------------------------------------------------- # -class BaseKVCacheCompressionManager(BaseResourceManager): +class KVCacheCompressionManager(BaseResourceManager): """Framework-level base class for all KV-cache compression managers. Inherits :class:`BaseResourceManager` so PyExecutor's main loop diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 896f4aa6fdef..0d6bf83401d9 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3398,7 +3398,7 @@ def supports_backend(self, backend: str) -> bool: ] -class BaseKvCacheCompressionConfig(StrictBaseModel): +class KvCacheCompressionConfig(StrictBaseModel): """Config for KV-cache compression: a compression manager runs a KV-reduction algorithm (e.g. periodic token eviction) alongside KVCacheManagerV2. @@ -3422,21 +3422,11 @@ def kv_cache_compression_mode(self): return KvCacheCompressionMode.from_string(self.algorithm) -class TriAttentionKvCacheCompressionConfig(BaseKvCacheCompressionConfig): - """KV-cache compression config for TriAttention. - - TriAttention periodically evicts cached tokens during generation, guided by - an offline-calibrated trigonometric importance score - (github.com/WeianMao/triattention). It runs on the KV-cache compression - framework with the standard ``KVCacheManagerV2``, whose ``update_resources`` - returns eviction-freed blocks to the pool for the capacity gain. TRT-LLM does - not compute calibration: supply the official tool's ``.pt`` via ``calibration_path`` and - it is converted to the runtime schema at load. TriAttention is a pure - compression method: it has no sparse-attention config and no attention - backend of its own -- decode runs the model's standard attention over the - compacted cache, and the manager publishes each request's evicted count on - ``LlmRequest.py_num_compressed_tokens`` for the engine to subtract. - """ +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.""" algorithm: Literal["triattention"] = "triattention" eviction_mode: Literal["union", "per_head", "per_layer_perhead"] = Field( default="union", @@ -3487,23 +3477,6 @@ def _require_calibration_inputs(self): "compute one.") return self - def to_manager_kwargs(self) -> dict: - """Constructor kwargs for the TriAttention manager.""" - return { - "budget": self.budget, - "beta": self.beta, - "model_path": self.model_path, - "calibration_path": self.calibration_path, - "eviction_mode": self.eviction_mode, - "normalize_scores": self.normalize_scores, - } - - -KvCacheCompressionConfig: TypeAlias = Annotated[ - Union[TriAttentionKvCacheCompressionConfig], - Field(discriminator="algorithm"), -] - @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) class AgentTreeConfig(StrictBaseModel, PybindMirror): @@ -4254,11 +4227,15 @@ 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( - default=None, - description="KV-cache compression config; None disables compression.", - status="prototype") + # KV is stored, not the attention computation). Dispatch is by the + # ``algorithm`` tag; grow this into a discriminated union when a second + # algorithm lands. + kv_cache_compression_config: Optional[ + TriAttentionKvCacheCompressionConfig] = Field( + default=None, + description= + "KV-cache compression config; None disables compression.", + status="prototype") # Speculative decoding parameters speculative_config: Optional[SpeculativeConfig] = Field( 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 4e1c70f758dd..7c11357c3210 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -2,11 +2,11 @@ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Unit tests for the KV-cache compression manager framework -(``BaseKVCacheCompressionManager`` in ``resource_manager.py``) — the +(``KVCacheCompressionManager`` in ``resource_manager.py``) — the ``BaseResourceManager``-based single-manager design. Covers: -- :class:`BaseKVCacheCompressionManager` contract: the four lifecycle hooks +- :class:`KVCacheCompressionManager` contract: the four lifecycle hooks default to no-op, zero resource counts, and it inherits :class:`BaseResourceManager` (so PyExecutor auto-drives it once registered). - The resource-manager API -> lifecycle-hook translation, gated on PyExecutor's @@ -31,8 +31,8 @@ from tensorrt_llm._torch.pyexecutor import _util as util_mod from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager from tensorrt_llm._torch.pyexecutor.resource_manager import ( - BaseKVCacheCompressionManager, BaseResourceManager, + KVCacheCompressionManager, ResourceManager, ResourceManagerType, ) @@ -55,7 +55,7 @@ def _record(self, hook_name: str): self._record_list.append(f"{self._name}:{hook_name}") -class _MockCompressionManager(_RecordingMixin, BaseKVCacheCompressionManager): +class _MockCompressionManager(_RecordingMixin, KVCacheCompressionManager): """Mock manager that records the four lifecycle hooks.""" def on_request_init(self, request): @@ -71,7 +71,7 @@ def on_request_finish(self, request): self._record("on_request_finish") -class _LengthAdjustingCompressionManager(BaseKVCacheCompressionManager): +class _LengthAdjustingCompressionManager(KVCacheCompressionManager): adjusts_generation_kv_length: ClassVar[bool] = True @@ -108,17 +108,17 @@ def _batch(context=(), generation=(), last_chunk=()): # ---------------------------------------------------------------------- # -# 1. BaseKVCacheCompressionManager contract # +# 1. KVCacheCompressionManager contract # # ---------------------------------------------------------------------- # class TestBaseABC: def test_inherits_base_resource_manager(self): # So PyExecutor's main loop auto-invokes prepare/update/free_resources. - assert issubclass(BaseKVCacheCompressionManager, BaseResourceManager) + assert issubclass(KVCacheCompressionManager, BaseResourceManager) def test_four_hooks_default_noop(self, fake_kv_cache_manager): - m = BaseKVCacheCompressionManager(fake_kv_cache_manager) + m = KVCacheCompressionManager(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,12 +128,12 @@ 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 = BaseKVCacheCompressionManager(fake_kv_cache_manager) + m = KVCacheCompressionManager(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 = BaseKVCacheCompressionManager(fake_kv_cache_manager) + m = KVCacheCompressionManager(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 @@ -155,9 +155,9 @@ def test_length_adjustment_marks_target_and_draft_v2(self): def test_rejects_non_v2_ownership(self): with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - BaseKVCacheCompressionManager(MagicMock()) + KVCacheCompressionManager(MagicMock()) with pytest.raises(TypeError, match="requires KVCacheManagerV2"): - BaseKVCacheCompressionManager(_v2_manager(is_draft=False), MagicMock()) + KVCacheCompressionManager(_v2_manager(is_draft=False), MagicMock()) def test_request_field_defaults_to_zero(self): """LlmRequest carries the compression count (the manager's only @@ -296,20 +296,20 @@ def test_factory_accepts_independent_draft_manager(self): 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 BaseKvCacheCompressionConfig + from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig - config = BaseKvCacheCompressionConfig(algorithm="offload") + config = KvCacheCompressionConfig(algorithm="offload") assert config.kv_cache_compression_mode.is_eviction_method() is False - m = BaseKVCacheCompressionManager(_v2_manager(is_draft=False)) + m = KVCacheCompressionManager(_v2_manager(is_draft=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 from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode - from tensorrt_llm.llmapi.llm_args import BaseKvCacheCompressionConfig + from tensorrt_llm.llmapi.llm_args import KvCacheCompressionConfig # Non-evicting methods pass with any speculative mode; no exception. - config = BaseKvCacheCompressionConfig(algorithm="offload") + config = KvCacheCompressionConfig(algorithm="offload") 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) @@ -326,7 +326,7 @@ def test_names_importable_from_canonical_modules(self): # Base class stays in resource_manager (it IS a resource manager); the # factory lives in _util next to _create_kv_cache_manager. - assert hasattr(resource_manager, "BaseKVCacheCompressionManager") + assert hasattr(resource_manager, "KVCacheCompressionManager") assert hasattr(_util, "create_kv_cache_compression_manager") def test_names_not_in_sparse_module(self): @@ -334,7 +334,7 @@ def test_names_not_in_sparse_module(self): # sparse-attention backend); the sparse package no longer exports it. from tensorrt_llm._torch.attention_backend import sparse - assert not hasattr(sparse, "BaseKVCacheCompressionManager") + assert not hasattr(sparse, "KVCacheCompressionManager") assert not hasattr(sparse, "create_kv_cache_compression_manager") @@ -354,7 +354,7 @@ def _mgr(self, enable_block_reuse): def test_raises_when_reuse_on(self): with pytest.raises(ValueError, match="block reuse"): - BaseKVCacheCompressionManager(self._mgr(enable_block_reuse=True)) + KVCacheCompressionManager(self._mgr(enable_block_reuse=True)) def test_ok_when_reuse_off(self): - BaseKVCacheCompressionManager(self._mgr(enable_block_reuse=False)) # no raise + KVCacheCompressionManager(self._mgr(enable_block_reuse=False)) # no raise diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 9e5370bb5ff0..eb963fd2e5ce 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -324,13 +324,25 @@ def make_fake_v2(enable_block_reuse=False, *, is_draft=False): return fake_v2 +def make_tri_config(**overrides): + """A real TriAttentionKvCacheCompressionConfig with test calibration inputs + (the config validator requires both ``model_path`` and ``calibration_path``).""" + from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig + + options = { + "budget": 8, + "model_path": "/models/test", + "calibration_path": "/calib/test.pt", + } + options.update(overrides) + return TriAttentionKvCacheCompressionConfig(**options) + + def make_triattention(**overrides): """Construct a fully initialized manager for method-level unit tests.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention - options = {"budget": 8, "model_path": "/models/test"} - options.update(overrides) - return TriAttention(make_fake_v2(), **options) + return TriAttention(make_fake_v2(), make_tri_config(**overrides)) def make_prepared_item( 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 index 7fea8b84fb92..e5777b5dcc51 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -19,6 +19,7 @@ from conftest import make_prepared_item as _make_prepared_item from conftest import make_ramp_pools as _make_ramp_pools 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 conftest import run_compaction as _run_compaction @@ -272,9 +273,10 @@ def test_draft_admission_gates_raise(gate, match): def construct(): return TriAttention( _make_fake_v2(), - budget=8, - model_path="/models/test", - eviction_mode="per_head" if gate == "union_only_per_head" else "union", + _make_tri_config( + budget=8, + eviction_mode="per_head" if gate == "union_only_per_head" else "union", + ), draft_kv_cache_manager=None if gate == "target_kv_factor" else draft_manager, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 5bc69ce72652..d6c3dd8916f5 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -31,6 +31,7 @@ from conftest import make_prepared_item as _make_prepared_item 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 from conftest import torch_tri_score_oracle as _torch_tri_score_oracle @@ -103,8 +104,8 @@ def test_llm_args_dispatch_and_validation(self): assert tri_args.kv_cache_compression_config.budget == 2048 assert tri_args.kv_cache_compression_config.beta == 128 - # The union dispatches on the algorithm tag, so an unknown algorithm - # fails config validation instead of falling back to a base config. + # Dispatch is by the algorithm tag, so an unknown algorithm fails + # config validation instead of falling back to a base config. with pytest.raises(ValidationError): TorchLlmArgs( model="dummy", @@ -182,7 +183,7 @@ def test_request_init_and_finish_lifecycle(self): manager = _make_fake_v2() manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 - triattention = TriAttention(manager, budget=8, model_path="/models/test") + triattention = TriAttention(manager, _make_tri_config(budget=8)) triattention._attention_layer_partition_cache = ([], [], None) triattention._calibrated = True @@ -368,7 +369,7 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft 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 - mgr = TriAttention(fake_v2, budget=8, model_path="/models/test") + mgr = TriAttention(fake_v2, _make_tri_config(budget=8)) mgr._calibrated = True cache = SimpleNamespace( capacity=seq_len, @@ -500,8 +501,7 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): draft_manager.max_seq_len = 8192 manager = TriAttention( _make_fake_v2(), - budget=8, - model_path="/models/test", + _make_tri_config(budget=8), draft_kv_cache_manager=draft_manager, ) @@ -536,7 +536,7 @@ def test_prepare_snapshots_fixed_linear_generation_growth( manager.kv_cache_map = { 7: SimpleNamespace(capacity=106, is_active=True), } - triattention = TriAttention(manager, budget=8, model_path="/models/test") + triattention = TriAttention(manager, _make_tri_config(budget=8)) batch = SimpleNamespace( context_requests=[], generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], From 368fa8c6a894009aad0a5159cbc49e134981e035 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 07:35:36 -0700 Subject: [PATCH 129/178] [None][refactor] Dissolve the runner shell and the bufs bag; TriAttention owns its state (knives 33/34/30) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 133 +- .../triattention/triattention.py | 1400 ++++++++++------- .../triattention_cute_score_fused.py | 282 ---- .../_torch/kv_cache_compression/conftest.py | 169 +- .../test_triattention_cute_score.py | 16 +- .../test_triattention_cute_union_fusion.py | 76 +- .../test_triattention_draft_cocompaction.py | 131 +- .../test_triattention_pipeline.py | 90 +- .../test_triattention_selection_compaction.py | 150 +- 9 files changed, 1191 insertions(+), 1256 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 4a1ed0727c82..5be22bf03969 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -129,10 +129,10 @@ def _make_compact_launches( move_offsets: torch.Tensor, destination_bases: torch.Tensor, per_layer_slots: Optional[Dict[int, int]] = None, -) -> Tuple[tuple, ...]: - """Batch layers into one prepacked ``sparse_kv_cache_compact_layers`` - argument tuple per uniform V2 pool (grouped by canonical pool id plus - pool/page-table geometry; native ABI order).""" +) -> Tuple[Dict[str, object], ...]: + """Batch layers into one launch record per uniform V2 pool (grouped by + canonical pool id plus pool/page-table geometry), carrying the native + ``sparse_kv_cache_compact_layers`` operands as plain named fields.""" device = entries[0][1].device grouped = OrderedDict() for layer, pool, page_table in entries: @@ -158,75 +158,29 @@ def _make_compact_launches( device=device, ) launches.append( - ( - pools, - torch.tensor( + dict( + pools=pools, + pool_pointers=torch.tensor( [pool.data_ptr() for pool in pools], dtype=torch.int64, device=device, ), - page_tables[0], - move_indices, - move_offsets, - destination_bases, - source_layer_indices, + page_table=page_tables[0], + move_indices=move_indices, + move_offsets=move_offsets, + destination_bases=destination_bases, + source_layer_indices=source_layer_indices, ) ) return tuple(launches) -def _make_pack_launch( - *, - # Decision inputs, packed each round. - kept_ordinal_rows: torch.Tensor, - valid_seq_lens: torch.Tensor, - decision_rows: int, - per_layer_sources: bool, - # Dense move family. - dense_move_offsets: torch.Tensor, - dense_move_indices: torch.Tensor, - # SWA move family (all-or-none). - swa_move_offsets: Optional[torch.Tensor] = None, - swa_move_indices: Optional[torch.Tensor] = None, - swa_window: int = 0, - # Launch geometry. - keep_count: int, - move_capacity: int, - num_kv_heads: int, -) -> Tuple[int, tuple, dict]: - """Assemble one family's prepacked pack launch ``(grid_rows, args, - kwargs)``: stable tensor references plus constexpr dispatch (the frozen - ``_pack_move_sources_kernel`` ABI), built once per geometry so the round - path constructs no Python containers.""" - return ( - decision_rows, - ( - kept_ordinal_rows, - valid_seq_lens, - dense_move_offsets, - dense_move_indices, - swa_move_offsets, - swa_move_indices, - ), - dict( - KEEP_COUNT=keep_count, - DECISION_ROWS=decision_rows, - MOVE_CAPACITY=move_capacity, - NUM_KV_HEADS=num_kv_heads, - PER_LAYER=per_layer_sources, - DENSE_TOTAL=int(dense_move_indices.shape[-1]), - SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, - SWA_WINDOW=swa_window, - ), - ) - - def init_compaction_buffers( *, target: Dict[str, object], capacities: Dict[str, int], draft: Optional[Dict[str, object]] = None, -) -> Tuple[Tuple[Tuple[int, tuple, dict], Tuple[tuple, ...]], ...]: +) -> Tuple[Dict[str, object], ...]: """Agree on the decision rows and return opaque launch plans per geometry. Move sources must be increasing kept ordinals with @@ -245,9 +199,10 @@ def init_compaction_buffers( draft's own KV heads); ``capacities`` the request/keep/tail capacity numbers. - Returns one ``(pack_launch, native_launches)`` plan per compacted cache - family (target, then the optional draft); the plans hold their own - move-index buffers and only :func:`compact` interprets them. + Returns one plan per compacted cache family (target, then the optional + draft): plain contract fields -- decision inputs, move-index buffers, + pack geometry ints, and the grouped native launch records -- that only + :func:`compact` interprets. """ layer_pools = target["layer_pools"] dense_layers = tuple(int(layer) for layer in target["dense_layers"]) @@ -338,7 +293,7 @@ def init_compaction_buffers( ) ) - target_pack_launch = _make_pack_launch( + target_plan = dict( kept_ordinal_rows=kept_ordinal_rows, valid_seq_lens=valid_seq_lens, decision_rows=decision_rows, @@ -351,9 +306,12 @@ def init_compaction_buffers( keep_count=keep_count, move_capacity=move_capacity, num_kv_heads=num_kv_heads, + dense_total=int(dense_move_indices.shape[-1]), + swa_total=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, + launches=tuple(target_launches), ) - plans = [(target_pack_launch, tuple(target_launches))] + plans = [target_plan] if draft is not None: if decision_rows != 1: raise ValueError( @@ -391,32 +349,63 @@ def init_compaction_buffers( move_offsets=draft["dense_move_offsets"], destination_bases=token_starts, ) - draft_pack_launch = _make_pack_launch( + draft_plan = dict( kept_ordinal_rows=kept_ordinal_rows, valid_seq_lens=valid_seq_lens, decision_rows=1, per_layer_sources=False, dense_move_offsets=draft["dense_move_offsets"], dense_move_indices=draft_move_indices, + swa_move_offsets=None, + swa_move_indices=None, + swa_window=0, keep_count=keep_count, move_capacity=keep_count + draft_tail, num_kv_heads=draft_num_kv_heads, + dense_total=int(draft_move_indices.shape[-1]), + swa_total=0, + launches=draft_launches, ) - plans.append((draft_pack_launch, draft_launches)) + plans.append(draft_plan) return tuple(plans) def compact( - plans: Tuple[Tuple[Tuple[int, tuple, dict], Tuple[tuple, ...]], ...], + plans: Tuple[Dict[str, object], ...], request_count: int, ) -> None: """Pack each plan's move sources and fire its native compacts, in plan order. Pure mover: the caller has already materialized its kept ordinals into the agreed decision rows for the active ``request_count`` cohort, and the - caller owns the completion ordering of the whole round. + caller owns the completion ordering of the whole round. Every launch + argument comes straight off the plan's init-built contract fields; the + round constructs no containers and runs no torch ops of its own. """ - for (rows, pack_args, pack_kwargs), native_launches in plans: - _pack_move_sources_kernel[(request_count, rows)](*pack_args, **pack_kwargs) - for launch in native_launches: - torch.ops.trtllm.sparse_kv_cache_compact_layers(*launch) + for plan in plans: + _pack_move_sources_kernel[(request_count, plan["decision_rows"])]( + plan["kept_ordinal_rows"], + plan["valid_seq_lens"], + plan["dense_move_offsets"], + plan["dense_move_indices"], + plan["swa_move_offsets"], + plan["swa_move_indices"], + KEEP_COUNT=plan["keep_count"], + DECISION_ROWS=plan["decision_rows"], + MOVE_CAPACITY=plan["move_capacity"], + NUM_KV_HEADS=plan["num_kv_heads"], + PER_LAYER=plan["per_layer_sources"], + DENSE_TOTAL=plan["dense_total"], + SWA_TOTAL=plan["swa_total"], + SWA_WINDOW=plan["swa_window"], + ) + for launch in plan["launches"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers( + launch["pools"], + launch["pool_pointers"], + launch["page_table"], + launch["move_indices"], + launch["move_offsets"], + launch["destination_bases"], + launch["source_layer_indices"], + ) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 52cdd7e6f243..6d4a3eb83ac1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -23,9 +23,9 @@ official tool (github.com/WeianMao/triattention) and is converted at load. """ -from types import SimpleNamespace from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple +import cuda.bindings.driver as cuda_driver import torch import triton @@ -113,575 +113,6 @@ def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: phase["rows"] = target -def init_eviction_buffers( - *, - eviction_mode: str, - layout: Dict[str, object], - calibration: Dict[str, torch.Tensor], - phase: Dict[str, object], - capacities: Dict[str, int], - draft: Optional[Dict[str, object]] = None, - normalize_scores: bool = True, -) -> SimpleNamespace: - """Build the one namespace of buffers, compiled launches, and compaction data - (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put). - - ``layout`` is the runtime KV layout dict, passed whole; ``calibration`` - carries the local q_real/q_imag/mlr_coef [L, H, F] slices and - freq_scale_sq; ``draft`` is one all-or-none resolved branch (its layout - dict plus tail/page-table capacities); ``capacities`` the capacity numbers. - Static round policy (``normalize_scores``, mode) binds here, once. - """ - from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS - from .triattention_cute_score_fused import TriAttentionCuteScoreRunner - - layer_pools = layout["layer_pools"] - dense_layers = list(layout["dense_layers"]) - swa_layers = list(layout["swa_layers"]) - swa_window = layout["swa_window"] - layer_group_representative = layout["layer_group_representative"] - # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. - layer_pool_ids = tuple(layout["layer_pool_ids"]) - dense_groups = list(layout["storage_groups"].values()) - page_representatives = [group[0] for group in dense_groups] - page_representatives.extend(layer for layer in swa_layers if layer not in page_representatives) - num_page_table_slots = int(layout["manager"].num_pools) - - device = layer_pools[page_representatives[0]].device - max_requests = int(capacities["max_requests"]) - seq_len = int(capacities["bucket_seq_len"]) - page_table_token_capacity = int(capacities["page_table_token_capacity"]) - decode_width = int(capacities["decode_width"]) - keep_count = int(capacities["keep_count"]) - protected_tail_capacity = int(capacities["protected_tail_capacity"]) - - q_real, q_imag, mlr_coef, freq_scale_sq = ( - calibration[key].to(device=device, dtype=torch.float32).contiguous() - for key in ("q_real", "q_imag", "mlr_coef", "freq_scale_sq") - ) - num_q_heads = int(q_real.shape[1]) - num_freqs = int(q_real.shape[2]) - - bufs = SimpleNamespace() - bufs.eviction_mode = eviction_mode - bufs.normalize_scores = bool(normalize_scores) - bufs.max_requests = max_requests - bufs.bucket_seq_len = seq_len - bufs.decode_width = decode_width - bufs.keep_count = keep_count - bufs.page_table_token_capacity = page_table_token_capacity - - # ---- block-offset staging (target, plus the co-compressed draft) ------- - bufs.block_offsets_host, bufs.block_offsets_device = _allocate_block_offset_staging( - layer_pools[page_representatives[0]], - num_pools=num_page_table_slots, - max_requests=max_requests, - token_capacity=page_table_token_capacity, - ) - # The draft is never scored: these offsets feed only the draft compacts. - bufs.draft_block_offsets_device = None - bufs.draft_block_offsets_host = None - bufs.draft_protected_tail_capacity = None - if draft is not None: - draft_layout = draft["layout"] - draft_representatives = list(draft_layout["pool_representatives"]) - draft_anchor_pool = draft_layout["layer_pools"][draft_representatives[0]] - # Construction-boundary invariant: the round shares one stream/event - # contract, so the draft pools must live on the target device. - if draft_anchor_pool.device != device: - raise RuntimeError("TriAttention draft KV pools must share the target KV pool device") - bufs.draft_block_offsets_host, bufs.draft_block_offsets_device = ( - _allocate_block_offset_staging( - draft_anchor_pool, - num_pools=int(draft_layout["manager"].num_pools), - max_requests=max_requests, - token_capacity=int(draft["page_table_token_capacity"]), - ) - ) - bufs.draft_protected_tail_capacity = int(draft["protected_tail_capacity"]) - - # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- - bufs.request_metadata_host = torch.empty( - (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() - ) - # numpy view over the pinned rows: per-round staging writes lists in place. - bufs.request_metadata_host_np = bufs.request_metadata_host.numpy() - bufs.identity_copy_indices_host = torch.arange( - max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() - ) - # Zero-filled: an unstaged cohort must gather the phase table's row 0. - bufs.request_metadata_device = torch.zeros( - (6, max_requests + 1), dtype=torch.int32, device=device - ) - bufs.round_starts_device = bufs.request_metadata_device[0, :max_requests] - bufs.valid_seq_lens_device = bufs.request_metadata_device[1, :max_requests] - # Per-request pinned prompt lengths (per-request decode window starts). - bufs.token_starts_device = bufs.request_metadata_device[2, :max_requests] - dense_move_offsets_row = bufs.request_metadata_device[3] - swa_move_offsets_row = bufs.request_metadata_device[4] - draft_move_offsets_row = bufs.request_metadata_device[5] - # SWA staging geometry, bound once (the compaction plans stay opaque): - # the phase gather rebases each request's SWA destination base in place. - bufs.swa_window = int(swa_window) if swa_layers else None - bufs.swa_destination_bases = torch.empty_like(bufs.token_starts_device) if swa_layers else None - bufs.swa_rebase_delta = keep_count - bufs.swa_window if swa_layers else 0 - bufs.mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) - bufs.mean_sin = torch.empty_like(bufs.mean_cos) - bufs.phase = phase - bufs.phase_num_freqs = int(phase["omega"].numel()) - bufs.phase_f_block = triton.next_power_of_2(bufs.phase_num_freqs) - - # ---- score state: one fused group across all dense layers -------------- - p0 = layer_pools[dense_layers[0]] - _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape - bufs.num_layers = len(dense_layers) - bufs.num_q_heads = int(num_q_heads) - bufs.num_kv_heads = int(num_kv_heads) - bufs.num_freqs = int(num_freqs) - bufs.tokens_per_block = int(tokens_per_block) - _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} - dense_layer_slots = [layer_pool_ids[_rep_of[layer]] for layer in dense_layers] - seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( - bufs.num_layers - ) - seg_layer_id = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( - max_requests - ) - block_offsets = bufs.block_offsets_device - slots_t = torch.tensor(dense_layer_slots, dtype=torch.int64, device=device) - req_idx = seg_req_id.to(torch.int64) - slot_idx = slots_t.repeat(max_requests) - seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - - max_segments = max_requests * bufs.num_layers - # The score plane must stay 32-bit indexable (wraparound = silent wild read). - if (PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len >= 2**31: - raise ValueError( - "score bucket overflows the 32-bit score plane: " - f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len}" - ) - # Persistent buffers: the compiled kernels capture their device pointers. - bufs.padded_head_columns = PADDED_HEAD_COLUMNS - bufs.score_scratch = torch.empty( - bufs.num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * seq_len, - dtype=torch.float32, - device=device, - ) - # int32 is safe here: covered by the 2^31 score-plane audit above. - seg_out_offset = (torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len).to( - torch.int32 - ) - bufs.gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( - 1, 1, 1, 1, -1 - ) - # Compile the mode's SM100 CuTe entries; no other score path, no fallback. - union = eviction_mode == "union" - # Persistent gather index (per-head modes): per round only the - # token-start base is re-added in place; the expanded view is fixed. - bufs.gather_index_base = None - bufs.gather_index = None - if not union: - num_kv_heads_early = int(layer_pools[dense_layers[0]].shape[2]) - bufs.gather_index_base = torch.empty( - (max_requests, 1, 1, 1, decode_width), dtype=torch.int64, device=device - ) - bufs.gather_index = bufs.gather_index_base.expand( - max_requests, - len(dense_layers), - num_kv_heads_early, - num_q_heads // num_kv_heads_early, - decode_width, - ) - bufs.union_scores = None - if union: - # Bucket-wide rows; consumers mask by the per-request widths. - bufs.union_scores = torch.empty((max_requests, seq_len), dtype=torch.float32, device=device) - # THE score path (no fallback): construction failures raise the runner's - # own dtype/shape/TMA error. - bufs.runner = TriAttentionCuteScoreRunner( - layer_pools=list(layer_pools), - layer_indices=[int(layer) for layer in dense_layers], - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=bufs.num_q_heads, - num_freqs=bufs.num_freqs, - page_ids=block_offsets.view(-1), - seg_page_off=seg_page_off, - seg_req_id=seg_req_id, - seg_layer_id=seg_layer_id, - # Pointer capture of the staged metadata rows. - valid_seq_lens=bufs.valid_seq_lens_device, - seg_out_offset=seg_out_offset, - token_starts=bufs.token_starts_device, - q_real=q_real.view(-1), - q_imag=q_imag.view(-1), - mlr_coef=mlr_coef.view(-1), - mean_cos=bufs.mean_cos, - mean_sin=bufs.mean_sin, - freq_scale_sq=freq_scale_sq, - output=bufs.score_scratch, - union_scores=bufs.union_scores, - enable_partial_stats=union, - ) - logger.info( - f"TriAttention CuTe score enabled: {bufs.num_q_heads}q/{bufs.num_kv_heads}kv heads, " - f"{bufs.num_freqs} freqs, {bufs.tokens_per_block}-token pages" - ) - - # ---- selection buffers (canonical row-major, one name per storage) ----- - bufs.valid_widths = torch.full((max_requests,), decode_width, dtype=torch.int32, device=device) - if union: - bufs.selection_rows_per_request = 1 - bufs.selection_scores_rows = torch.empty( - (max_requests, decode_width), dtype=torch.float32, device=device - ) - # One selection row per request: its length IS the staged valid width. - bufs.selection_row_lengths = bufs.valid_widths - # Padded rows still need in-range ordinals for the finalizer's gather. - bufs.provisional_rows = torch.zeros( - (max_requests, keep_count), dtype=torch.int32, device=device - ) - # Kept decode ordinals only (prompt-length independent rows). - bufs.kept_ordinal_rows = torch.empty( - (max_requests, keep_count), dtype=torch.int32, device=device - ) - bufs.score_output = None - else: - selection_rows = ( - bufs.num_kv_heads - if eviction_mode == "per_head" - else bufs.num_layers * bufs.num_kv_heads - ) - # Both rectangles must stay 32-bit indexable (wraparound = wild reads). - score_rect = max_requests * bufs.num_layers * bufs.num_q_heads * decode_width - selection_rect = max_requests * selection_rows * max(decode_width, keep_count) - if max(score_rect, selection_rect) >= 2**31: - raise ValueError( - f"per-head score rectangles overflow 32-bit indexing: " - f"scores {score_rect}, selection {selection_rect}" - ) - bufs.selection_rows_per_request = selection_rows - # [request, layer, head, token] layout read by the reduce kernels. - bufs.score_output = torch.empty( - max_requests, - bufs.num_layers, - bufs.num_q_heads, - decode_width, - dtype=torch.float32, - device=device, - ) - score_shape = (max_requests, bufs.num_layers, bufs.num_q_heads, 1) - bufs.row_mean = torch.empty(score_shape, dtype=torch.float32, device=device) - bufs.row_inv_std = torch.empty_like(bufs.row_mean) - bufs.selection_scores_rows = torch.empty( - (max_requests * selection_rows, decode_width), dtype=torch.float32, device=device - ) - bufs.selection_row_lengths = torch.full( - (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device - ) - bufs.provisional_rows = torch.zeros( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) - bufs.kept_ordinal_rows = torch.empty( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) - - # ---- compaction plans + decision-materialization prebinds --------------- - per_layer = eviction_mode == "per_layer_perhead" - draft_contract = None - if draft is not None: - draft_layout = draft["layout"] - draft_contract = dict( - layer_pools=draft_layout["layer_pools"], - dense_layers=list(draft_layout["dense_layers"]), - layer_group_representative=draft_layout["layer_group_representative"], - layer_pool_ids=tuple(draft_layout["layer_pool_ids"]), - kv_block_offsets=bufs.draft_block_offsets_device, - dense_move_offsets=draft_move_offsets_row, - protected_tail_capacity=int(draft["protected_tail_capacity"]), - ) - # Opaque launch plans: only compact() interprets them. - bufs.compaction_plan = init_compaction_buffers( - target=dict( - layer_pools=layer_pools, - dense_layers=list(dense_layers), - swa_layers=list(swa_layers), - swa_window=swa_window, - layer_group_representative=layer_group_representative, - layer_pool_ids=layer_pool_ids, - kv_block_offsets=bufs.block_offsets_device, - token_starts=bufs.token_starts_device, - swa_destination_bases=bufs.swa_destination_bases, - # Per-round tails: the move offsets ride the staged metadata rows. - dense_move_offsets=dense_move_offsets_row, - swa_move_offsets=swa_move_offsets_row, - per_layer_sources=per_layer, - # The decision rows the plans pack into move sources. - kept_ordinal_rows=bufs.kept_ordinal_rows, - decision_rows=bufs.selection_rows_per_request, - valid_seq_lens=bufs.valid_seq_lens_device, - ), - capacities=dict( - max_requests=max_requests, - keep_count=keep_count, - protected_tail_capacity=int(protected_tail_capacity), - ), - draft=draft_contract, - ) - # The decision side: the settle launch materializes the kept-ordinal rows. - bufs.settle_args = ( - bufs.selection_scores_rows, - bufs.selection_row_lengths, - bufs.token_starts_device, - bufs.provisional_rows, - bufs.kept_ordinal_rows, - ) - bufs.settle_kwargs = dict( - WIDTH=decode_width, - KEEP_COUNT=keep_count, - SELECTION_ROWS=bufs.selection_rows_per_request, - ) - - # ---- round-ordering events ---------------------------------------------- - # Host staging (pinned metadata + snapshots) reuse fence. - bufs.staging_reuse_event = torch.cuda.Event() - bufs.staging_reuse_event.record(torch.cuda.current_stream(device)) - # Manager-stream H2D of the block-offset tables has completed. - bufs.block_offsets_ready_event = torch.cuda.Event() - # This cohort's compact is done: manager may resize/reuse pages. - bufs.compaction_done_event = torch.cuda.Event() - bufs.copy_pending = False - return bufs - - -def _stage_block_offsets( - bufs: SimpleNamespace, - manager: KVCacheManagerV2, - request_ids: List[int], - host_block_offsets: torch.Tensor, - device_block_offsets: torch.Tensor, -) -> None: - """Gather the pinned snapshot before the async device copy: resize mutates - the live host table. The round owner has already fenced host-staging reuse.""" - manager.index_mapper.gather_k_block_offsets( - manager.host_kv_cache_block_offsets, - host_block_offsets, - request_ids, - host_block_offsets.shape[-1], - ) - manager._stream.wait_event(bufs.staging_reuse_event) - copy_batch_block_offsets_to_device( - host_block_offsets, - device_block_offsets, - bufs.identity_copy_indices_host[: len(request_ids)], - manager.index_scales, - manager.kv_offset, - manager._stream.cuda_stream, - ) - bufs.block_offsets_ready_event.record(manager._stream) - torch.cuda.current_stream(device_block_offsets.device).wait_event( - bufs.block_offsets_ready_event - ) - - -def _cohort_move_offsets( - bufs: SimpleNamespace, - prepared: Sequence[Dict[str, object]], -) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep - set plus protected tail per request; rows past the cohort repeat the final - offset and contribute no moves).""" - - def padded_offsets(moves_per_request: List[int]) -> List[int]: - offsets = [0] - for moves in moves_per_request: - offsets.append(offsets[-1] + moves) - offsets.extend(offsets[-1:] * (bufs.max_requests - len(moves_per_request))) - return offsets - - tails = [int(item["protected_tail"]) for item in prepared] - dense = padded_offsets([bufs.keep_count + tail for tail in tails]) - swa = None - if bufs.swa_window is not None: - swa = padded_offsets([bufs.swa_window + tail for tail in tails]) - draft = None - if bufs.draft_protected_tail_capacity is not None: - draft = padded_offsets( - [bufs.keep_count + bufs.draft_protected_tail_capacity] * len(prepared) - ) - return dense, swa, draft - - -def settle_top_tokens(bufs: SimpleNamespace, request_count: int) -> None: - """Pick the top-k and settle ties into the kept-ordinal decision rows - (the compaction contract packs them into move sources).""" - rows = request_count * bufs.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( - bufs.selection_scores_rows[:rows], - bufs.selection_row_lengths[:rows], - bufs.provisional_rows[:rows], - bufs.keep_count, - 1, - ) - _settle_ties_kernel[(request_count, bufs.selection_rows_per_request)]( - *bufs.settle_args, **bufs.settle_kwargs - ) - - -def execute_eviction_round( - bufs: SimpleNamespace, - manager: KVCacheManagerV2, - prepared: Sequence[Dict[str, object]], - draft_manager: Optional[KVCacheManagerV2] = None, -) -> None: - """Run one eviction round over the prepared cohort: stage the page-table - snapshots and round metadata, then score, select, settle, and compact, and - finally order the manager streams after this cohort's compact (every launch - covers the full request capacity; padded rows carry zero lengths and stay - inert).""" - with nvtx_range_debug("triattention.page_table_stage", color="orange"): - request_ids = [item["request_id"] for item in prepared] - round_starts = [item["round_start"] for item in prepared] - token_starts = [item["prompt_len"] for item in prepared] - seq_lens = [item["seq_len"] for item in prepared] - dense_move_offsets, swa_move_offsets, draft_move_offsets = _cohort_move_offsets( - bufs, prepared - ) - stream = torch.cuda.current_stream(bufs.block_offsets_device.device) - # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. - max_round_start = max(round_starts) - rows = ( - (0, round_starts), - (1, seq_lens), - (2, token_starts), - (3, dense_move_offsets), - (4, swa_move_offsets), - (5, draft_move_offsets), - ) - for row, values in rows: - if values is not None and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF: - raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # The one host-staging reuse fence: the previous cohort's async copies - # must complete before the pinned metadata rows AND the pinned - # target/draft block-offset snapshots are rewritten. - if bufs.copy_pending and not bufs.staging_reuse_event.query(): - bufs.staging_reuse_event.synchronize() - host_table = bufs.request_metadata_host_np - for row, values in rows: - if values is not None: - host_table[row, : len(values)] = values - # Zero lengths keep the score kernel and selection inert for padded rows. - host_table[:3, len(prepared) :] = 0 - grow_mean_phase_table(bufs.phase, int(max_round_start) + 1) - _stage_block_offsets( - bufs, - manager, - request_ids, - bufs.block_offsets_host, - bufs.block_offsets_device, - ) - if draft_manager is not None: - _stage_block_offsets( - bufs, - draft_manager, - request_ids, - bufs.draft_block_offsets_host, - bufs.draft_block_offsets_device, - ) - try: - bufs.request_metadata_device.copy_(bufs.request_metadata_host, non_blocking=True) - finally: - # Guards the pinned staging until the asynchronous copies complete. - bufs.staging_reuse_event.record(stream) - bufs.copy_pending = True - request_count = len(prepared) - union = bufs.eviction_mode == "union" - try: - with nvtx_range("triattention.score", color="blue"): - # In-place refresh: the compiled score launches captured these pointers. - _gather_mean_phase_kernel[(request_count,)]( - bufs.round_starts_device, - bufs.phase["cos"], - bufs.phase["sin"], - bufs.phase["rows"], - bufs.valid_seq_lens_device, - bufs.token_starts_device, - bufs.mean_cos, - bufs.mean_sin, - bufs.valid_widths, - bufs.swa_destination_bases, - bufs.swa_rebase_delta, - NUM_FREQS=bufs.phase_num_freqs, - F_BLOCK=bufs.phase_f_block, - HAS_SWA=bufs.swa_destination_bases is not None, - num_warps=1, - ) - if union: - # The runner's ctor bound the mean/union buffers; the launch - # takes only the active cohort size. - bufs.runner.launch_union_fusion(request_count) - columns = min(bufs.union_scores.shape[1], bufs.selection_scores_rows.shape[1]) - bufs.selection_scores_rows[:request_count, :columns].copy_( - bufs.union_scores[:request_count, :columns] - ) - else: - bufs.runner.launch(request_count) - # Gather each decode window into the [request, layer, head, token] layout of the reduces. - group_size = bufs.num_q_heads // bufs.num_kv_heads - num_segments = request_count * bufs.num_layers - pad = bufs.padded_head_columns - source = ( - bufs.score_scratch[ - : bufs.num_kv_heads * pad * num_segments * bufs.bucket_seq_len - ] - .view( - bufs.num_kv_heads, pad, request_count, bufs.num_layers, bufs.bucket_seq_len - )[:, :group_size] - .permute(2, 3, 0, 1, 4) - ) - torch.add( - bufs.token_starts_device[:request_count].view(-1, 1, 1, 1, 1), - bufs.gather_columns, - out=bufs.gather_index_base[:request_count], - ) - bufs.gather_index_base[:request_count].clamp_(max=bufs.bucket_seq_len - 1) - columns = bufs.gather_index[:request_count] - torch.gather( - source, - 4, - columns, - out=bufs.score_output[:request_count].view( - request_count, - bufs.num_layers, - bufs.num_kv_heads, - group_size, - bufs.decode_width, - ), - ) - with nvtx_range("triattention.select", color="yellow"): - if not union: - prepare_per_head_scores( - bufs.score_output[:request_count], - bufs.valid_widths, - bufs.row_mean, - bufs.row_inv_std, - bufs.selection_scores_rows, - bufs.selection_row_lengths, - per_layer=bufs.eviction_mode == "per_layer_perhead", - normalize_scores=bufs.normalize_scores, - ) - settle_top_tokens(bufs, request_count) - with nvtx_range("triattention.compact", color="purple"): - compact(bufs.compaction_plan, request_count) - finally: - # Order V2 page-table reuse and resize after this cohort's compact. - bufs.compaction_done_event.record(stream) - manager._stream.wait_event(bufs.compaction_done_event) - if draft_manager is not None: - draft_manager._stream.wait_event(bufs.compaction_done_event) - - class TriAttention(KVCacheCompressionManager): """Periodic physical KV eviction driven by trigonometric importance scoring.""" @@ -715,7 +146,7 @@ def __init__( self._calibrated = False self._freq_scale_sq: Optional[torch.Tensor] = None - # Mean-phase table dict, shared by reference with every buffer namespace. + # Mean-phase table dict; buffer builds bind its device tables in place. self._phase: Optional[Dict[str, object]] = None # Per-request {generation_steps, evicted_tokens}. @@ -741,8 +172,9 @@ def __init__( + 1 ) self._generation_growth = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) - # Built once at the first eviction, reused for the manager's lifetime. - self._buffers: Optional[SimpleNamespace] = None + # Buffers build once at the first eviction as plain attributes on this + # manager and stay resident for the manager's lifetime. + self._buffers_built = False self._local_to_global_layers_cache: Optional[List[int]] = None self._attention_layer_partition_cache: Optional[ Tuple[List[int], List[int], Optional[int]] @@ -1290,11 +722,11 @@ def _pool_page_counts( for layer in pool_representatives ) - def _buffers_for( + def _ensure_buffers( self, layout: Dict[str, object], prepared: Sequence[Dict[str, object]], - ) -> SimpleNamespace: + ) -> None: # Empty cohorts never reach here: _periodic_evict no-ops pre-launch. needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) @@ -1303,16 +735,15 @@ def _buffers_for( # The cached layout lookup enforces draft V2 pool page-count # stability every round, exactly like the target's lookup. self._draft_runtime_kv_layout() - bufs = self._buffers - if bufs is not None: + if self._buffers_built: if ( - needed_width <= bufs.decode_width - and needed_page_tokens <= bufs.page_table_token_capacity - and needed_requests <= bufs.max_requests + needed_width <= self._decode_width + and needed_page_tokens <= self._page_table_token_capacity + and needed_requests <= self._max_requests ): - return bufs + return # This round outgrew the buffers: rebuild. - self._buffers = None + self._buffers_built = False mgr = self.kv_cache_manager tail_capacity = self._protected_tail_capacity @@ -1355,8 +786,7 @@ def _buffers_for( } grow_mean_phase_table(self._phase, max(int(seq_capacity), 1)) q_real, q_imag, mlr_coef = self._local_score_calibration(layout["global_layers"]) - bufs = init_eviction_buffers( - eviction_mode=self.eviction_mode, + self._build_buffers( layout=layout, calibration=dict( q_real=q_real, @@ -1364,7 +794,6 @@ def _buffers_for( mlr_coef=mlr_coef, freq_scale_sq=self._freq_scale_sq, ), - phase=self._phase, capacities=dict( max_requests=request_capacity, bucket_seq_len=seq_capacity, @@ -1374,10 +803,536 @@ def _buffers_for( protected_tail_capacity=tail_capacity, ), draft=draft, - normalize_scores=self.normalize_scores, ) - self._buffers = bufs - return bufs + self._buffers_built = True + + def _build_buffers( + self, + *, + layout: Dict[str, object], + calibration: Dict[str, torch.Tensor], + capacities: Dict[str, int], + draft: Optional[Dict[str, object]] = None, + ) -> None: + """Build the round's buffers, compiled score launches, and compaction data + in place as plain attributes on this manager + (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put). + + ``layout`` is the runtime KV layout dict, passed whole; ``calibration`` + carries the local q_real/q_imag/mlr_coef [L, H, F] slices and + freq_scale_sq; ``draft`` is one all-or-none resolved branch (its layout + dict plus tail/page-table capacities); ``capacities`` the capacity + numbers. Static round policy (mode, ``normalize_scores``, the shared + ``self._phase`` table) reads off the manager itself. + """ + import cutlass + import cutlass.cute as cute + + from .triattention_cute_score_fused import ( + _COMPILE_LOCK, + _COMPILED_KERNELS, + SMALL_WORKLOAD_PAGE_SHARDS, + STATS_FIELDS, + _encode_tma_descriptors, + _tensor_spec, + _to_cute, + _TriAttentionScoreKernel, + ) + from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS + + layer_pools = layout["layer_pools"] + dense_layers = list(layout["dense_layers"]) + swa_layers = list(layout["swa_layers"]) + swa_window = layout["swa_window"] + layer_group_representative = layout["layer_group_representative"] + # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. + layer_pool_ids = tuple(layout["layer_pool_ids"]) + dense_groups = list(layout["storage_groups"].values()) + page_representatives = [group[0] for group in dense_groups] + page_representatives.extend( + layer for layer in swa_layers if layer not in page_representatives + ) + num_page_table_slots = int(layout["manager"].num_pools) + + device = layer_pools[page_representatives[0]].device + max_requests = int(capacities["max_requests"]) + seq_len = int(capacities["bucket_seq_len"]) + page_table_token_capacity = int(capacities["page_table_token_capacity"]) + decode_width = int(capacities["decode_width"]) + keep_count = int(capacities["keep_count"]) + protected_tail_capacity = int(capacities["protected_tail_capacity"]) + + q_real, q_imag, mlr_coef, freq_scale_sq = ( + calibration[key].to(device=device, dtype=torch.float32).contiguous() + for key in ("q_real", "q_imag", "mlr_coef", "freq_scale_sq") + ) + num_q_heads = int(q_real.shape[1]) + num_freqs = int(q_real.shape[2]) + + self._max_requests = max_requests + self._bucket_seq_len = seq_len + self._decode_width = decode_width + self._keep_count = keep_count + self._page_table_token_capacity = page_table_token_capacity + + # ---- block-offset staging (target, plus the co-compressed draft) ------- + self._block_offsets_host, self._block_offsets_device = _allocate_block_offset_staging( + layer_pools[page_representatives[0]], + num_pools=num_page_table_slots, + max_requests=max_requests, + token_capacity=page_table_token_capacity, + ) + # The draft is never scored: these offsets feed only the draft compacts. + self._draft_block_offsets_device = None + self._draft_block_offsets_host = None + if draft is not None: + draft_layout = draft["layout"] + draft_representatives = list(draft_layout["pool_representatives"]) + draft_anchor_pool = draft_layout["layer_pools"][draft_representatives[0]] + # Construction-boundary invariant: the round shares one stream/event + # contract, so the draft pools must live on the target device. + if draft_anchor_pool.device != device: + raise RuntimeError( + "TriAttention draft KV pools must share the target KV pool device" + ) + self._draft_block_offsets_host, self._draft_block_offsets_device = ( + _allocate_block_offset_staging( + draft_anchor_pool, + num_pools=int(draft_layout["manager"].num_pools), + max_requests=max_requests, + token_capacity=int(draft["page_table_token_capacity"]), + ) + ) + + # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- + self._request_metadata_host = torch.empty( + (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + ) + # numpy view over the pinned rows: per-round staging writes lists in place. + self._request_metadata_host_np = self._request_metadata_host.numpy() + self._identity_copy_indices_host = torch.arange( + max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + ) + # Zero-filled: an unstaged cohort must gather the phase table's row 0. + self._request_metadata_device = torch.zeros( + (6, max_requests + 1), dtype=torch.int32, device=device + ) + self._round_starts_device = self._request_metadata_device[0, :max_requests] + self._valid_seq_lens_device = self._request_metadata_device[1, :max_requests] + # Per-request pinned prompt lengths (per-request decode window starts). + self._token_starts_device = self._request_metadata_device[2, :max_requests] + dense_move_offsets_row = self._request_metadata_device[3] + swa_move_offsets_row = self._request_metadata_device[4] + draft_move_offsets_row = self._request_metadata_device[5] + # SWA staging geometry, bound once (the compaction plans stay opaque): + # the phase gather rebases each request's SWA destination base in place. + self._swa_window = int(swa_window) if swa_layers else None + self._swa_destination_bases = ( + torch.empty_like(self._token_starts_device) if swa_layers else None + ) + self._swa_rebase_delta = keep_count - self._swa_window if swa_layers else 0 + self._mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) + self._mean_sin = torch.empty_like(self._mean_cos) + self._phase_num_freqs = int(self._phase["omega"].numel()) + self._phase_f_block = triton.next_power_of_2(self._phase_num_freqs) + + # ---- score state: one fused group across all dense layers -------------- + p0 = layer_pools[dense_layers[0]] + _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape + self._num_layers = len(dense_layers) + self._num_q_heads = int(num_q_heads) + self._num_kv_heads = int(num_kv_heads) + self._num_freqs = int(num_freqs) + self._tokens_per_block = int(tokens_per_block) + _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} + dense_layer_slots = [layer_pool_ids[_rep_of[layer]] for layer in dense_layers] + seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( + self._num_layers + ) + seg_layer_id = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( + max_requests + ) + block_offsets = self._block_offsets_device + slots_t = torch.tensor(dense_layer_slots, dtype=torch.int64, device=device) + req_idx = seg_req_id.to(torch.int64) + slot_idx = slots_t.repeat(max_requests) + seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) + + max_segments = max_requests * self._num_layers + # The score plane must stay 32-bit indexable (wraparound = silent wild read). + if (PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len >= 2**31: + raise ValueError( + "score bucket overflows the 32-bit score plane: " + f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len}" + ) + # Persistent buffers: the compiled kernels capture their device pointers. + self._padded_head_columns = PADDED_HEAD_COLUMNS + self._score_scratch = torch.empty( + self._num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * seq_len, + dtype=torch.float32, + device=device, + ) + # int32 is safe here: covered by the 2^31 score-plane audit above. + seg_out_offset = ( + torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len + ).to(torch.int32) + self._gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( + 1, 1, 1, 1, -1 + ) + # Compile the mode's SM100 CuTe entries; no other score path, no fallback. + union = self.eviction_mode == "union" + # Persistent gather index (per-head modes): per round only the + # token-start base is re-added in place; the expanded view is fixed. + self._gather_index_base = None + self._gather_index = None + if not union: + num_kv_heads_early = int(layer_pools[dense_layers[0]].shape[2]) + self._gather_index_base = torch.empty( + (max_requests, 1, 1, 1, decode_width), dtype=torch.int64, device=device + ) + self._gather_index = self._gather_index_base.expand( + max_requests, + len(dense_layers), + num_kv_heads_early, + num_q_heads // num_kv_heads_early, + decode_width, + ) + self._union_scores = None + if union: + # Bucket-wide rows; consumers mask by the per-request widths. + self._union_scores = torch.empty( + (max_requests, seq_len), dtype=torch.float32, device=device + ) + + # ---- THE score path (no fallback): per request-count/page-shard + # compiled variants over ctor-bound cute tensor handles; construction + # failures raise the kernels' own dtype/shape/TMA errors. + anchor_pool = p0 + sm_count = int(torch.cuda.get_device_properties(device).multi_processor_count) + # One [stats_row, page_shard, {count, mean, m2}] record array. + partial_stats_elements = ( + max_requests + * self._num_layers + * num_q_heads + * SMALL_WORKLOAD_PAGE_SHARDS + * STATS_FIELDS + if union + else 1 + ) + self._partial_stats = torch.empty( + partial_stats_elements, + dtype=torch.float32, + device=device, + ) + self._tma_descriptors = _encode_tma_descriptors( + list(layer_pools), + [int(layer) for layer in dense_layers], + int(num_freqs), + int(tokens_per_block), + ) + torch_prefix = ( + block_offsets.view(-1), + seg_page_off, + seg_req_id, + seg_layer_id, + # Pointer capture of the staged metadata rows. + self._valid_seq_lens_device, + seg_out_offset, + self._token_starts_device, + q_real.view(-1), + q_imag.view(-1), + mlr_coef.view(-1), + ) + torch_tail = ( + freq_scale_sq, + self._score_scratch, + self._partial_stats, + anchor_pool, + self._tma_descriptors, + ) + # Keep-alive twin of the cute handles below: the compiled launches + # capture these raw device pointers for the buffers' lifetime. + self._score_torch_operands = (*torch_prefix, *torch_tail) + # valid_seq_lens/token_starts are only 4-byte-aligned row views, read as per-CTA scalars. + prefix_aligns = (16, 16, 16, 16, 4, 16, 4, 16, 16, 16) + self._cute_score_prefix = tuple( + _to_cute(tensor, assumed_align=align) + for tensor, align in zip(torch_prefix, prefix_aligns) + ) + self._cute_score_tail = ( + _to_cute(freq_scale_sq), + _to_cute(self._score_scratch), + _to_cute(self._partial_stats), + _to_cute(anchor_pool), + _to_cute(self._tma_descriptors, assumed_align=128), + ) + # Ctor-bound persistent launch operands (one from_dlpack wrap each): + # the round path refreshes their contents in place and launches with + # the active cohort size only. + self._cute_mean_cos = _to_cute(self._mean_cos.view(-1)) + self._cute_mean_sin = _to_cute(self._mean_sin.view(-1)) + self._cute_union_scores = _to_cute(self._union_scores.view(-1)) if union else None + self._compiled_score: Dict[int, object] = {} + self._compiled_score_stats: Dict[int, object] = {} + self._compiled_normalize_union: Dict[int, object] = {} + page_shards_by_count: Dict[int, int] = {} + self._cute_selection_prefix = ( + _to_cute(self._score_scratch), + _to_cute(self._valid_seq_lens_device, assumed_align=4), + _to_cute(seg_out_offset), + _to_cute(self._token_starts_device, assumed_align=4), + ) + self._cute_partial_stats = _to_cute(self._partial_stats) + static_geometry = ( + max_requests, + self._num_layers, + seq_len, + num_q_heads, + self._num_kv_heads, + num_freqs, + self._tokens_per_block, + tuple(int(value) for value in anchor_pool.shape), + tuple(int(value) for value in anchor_pool.stride()), + ) + tensor_specs = tuple( + _tensor_spec(tensor) + for tensor in ( + *torch_prefix, + self._mean_cos.view(-1), + self._mean_sin.view(-1), + *torch_tail, + ) + ) + variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] + if max_requests > 1: + variants.append((max_requests, 2)) + # Per-head modes compile the score-only entry; union the fused pipeline. + kernel_kwargs = dict( + num_layers=self._num_layers, + seq_len=seq_len, + 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()), + ) + if union: + variant_key = "triattention_cute_score_stats" + compiled_entries = self._compiled_score_stats + else: + variant_key = "triattention_cute_score" + compiled_entries = self._compiled_score + for request_count, page_shards in variants: + cache_key = ( + variant_key, + static_geometry, + tensor_specs, + request_count, + page_shards, + ) + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + kernel = _TriAttentionScoreKernel( + **kernel_kwargs, + page_shards=page_shards, + write_partial_stats=union, + ) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + compiled = cute.compile( + kernel, + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled + compiled_entries[request_count] = compiled + page_shards_by_count[request_count] = page_shards + + if max_requests > 1: + small = compiled_entries.get(1) + large = compiled_entries.get(max_requests) + for request_count in range(1, max_requests + 1): + # Give small cohorts the extra shard while the 2-shard grid stays under two waves. + two_shard_ctas = request_count * self._num_layers * self._num_kv_heads * 2 + use_extra_score_shard = two_shard_ctas < 2 * sm_count + compiled_entries[request_count] = small if use_extra_score_shard else large + page_shards_by_count[request_count] = ( + SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 + ) + + if union: + from .triattention_cute_selection import ( + _select_normalize_union_config, + _TriAttentionNormalizeUnionKernel, + ) + + compiled_configs: Dict[Tuple[int, int, int, int], object] = {} + for request_count in range(1, max_requests + 1): + page_shards = page_shards_by_count[request_count] + config = _select_normalize_union_config( + request_count, + seq_len, + sm_count, + ) + config_key = (page_shards, *config) + compiled_selection = compiled_configs.get(config_key) + if compiled_selection is None: + cache_key = ( + "triattention_cute_normalize_union", + static_geometry, + tensor_specs, + config_key, + _tensor_spec(self._union_scores), + _tensor_spec(self._partial_stats), + ) + with _COMPILE_LOCK: + compiled_selection = _COMPILED_KERNELS.get(cache_key) + if compiled_selection is None: + tokens_per_lane, token_subtiles, row_cluster_ctas = config + kernel = _TriAttentionNormalizeUnionKernel( + num_layers=self._num_layers, + seq_len=seq_len, + num_q_heads=num_q_heads, + # The finalizer maps real head rows onto N=8-padded planes. + num_kv_heads=self._num_kv_heads, + page_shards=page_shards, + tokens_per_lane=tokens_per_lane, + token_subtiles=token_subtiles, + row_cluster_ctas=row_cluster_ctas, + ) + stream = cuda_driver.CUstream( + torch.cuda.current_stream(device).cuda_stream + ) + compiled_selection = cute.compile( + kernel, + _to_cute(self._partial_stats), + *self._cute_selection_prefix, + self._cute_union_scores, + cutlass.Int32(1), + stream, + ) + _COMPILED_KERNELS[cache_key] = compiled_selection + compiled_configs[config_key] = compiled_selection + self._compiled_normalize_union[request_count] = compiled_selection + logger.info( + f"TriAttention CuTe score enabled: {self._num_q_heads}q/{self._num_kv_heads}kv heads, " + f"{self._num_freqs} freqs, {self._tokens_per_block}-token pages" + ) + + # ---- selection buffers (canonical row-major, one name per storage) ----- + self._valid_widths = torch.full( + (max_requests,), decode_width, dtype=torch.int32, device=device + ) + if union: + self._selection_rows_per_request = 1 + self._selection_scores_rows = torch.empty( + (max_requests, decode_width), dtype=torch.float32, device=device + ) + # One selection row per request: its length IS the staged valid width. + self._selection_row_lengths = self._valid_widths + # Padded rows still need in-range ordinals for the finalizer's gather. + self._provisional_rows = torch.zeros( + (max_requests, keep_count), dtype=torch.int32, device=device + ) + # Kept decode ordinals only (prompt-length independent rows). + self._kept_ordinal_rows = torch.empty( + (max_requests, keep_count), dtype=torch.int32, device=device + ) + self._score_output = None + else: + selection_rows = ( + self._num_kv_heads + if self.eviction_mode == "per_head" + else self._num_layers * self._num_kv_heads + ) + # Both rectangles must stay 32-bit indexable (wraparound = wild reads). + score_rect = max_requests * self._num_layers * self._num_q_heads * decode_width + selection_rect = max_requests * selection_rows * max(decode_width, keep_count) + if max(score_rect, selection_rect) >= 2**31: + raise ValueError( + f"per-head score rectangles overflow 32-bit indexing: " + f"scores {score_rect}, selection {selection_rect}" + ) + self._selection_rows_per_request = selection_rows + # [request, layer, head, token] layout read by the reduce kernels. + self._score_output = torch.empty( + max_requests, + self._num_layers, + self._num_q_heads, + decode_width, + dtype=torch.float32, + device=device, + ) + score_shape = (max_requests, 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( + (max_requests * selection_rows, decode_width), dtype=torch.float32, device=device + ) + self._selection_row_lengths = torch.full( + (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device + ) + self._provisional_rows = torch.zeros( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + ) + self._kept_ordinal_rows = torch.empty( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + ) + + # ---- compaction plans (opaque: only compact() interprets them) --------- + per_layer = self.eviction_mode == "per_layer_perhead" + draft_contract = None + if draft is not None: + draft_layout = draft["layout"] + draft_contract = dict( + layer_pools=draft_layout["layer_pools"], + dense_layers=list(draft_layout["dense_layers"]), + layer_group_representative=draft_layout["layer_group_representative"], + layer_pool_ids=tuple(draft_layout["layer_pool_ids"]), + kv_block_offsets=self._draft_block_offsets_device, + dense_move_offsets=draft_move_offsets_row, + protected_tail_capacity=int(draft["protected_tail_capacity"]), + ) + self.compaction_plan = init_compaction_buffers( + target=dict( + layer_pools=layer_pools, + dense_layers=list(dense_layers), + swa_layers=list(swa_layers), + swa_window=swa_window, + layer_group_representative=layer_group_representative, + layer_pool_ids=layer_pool_ids, + kv_block_offsets=self._block_offsets_device, + token_starts=self._token_starts_device, + swa_destination_bases=self._swa_destination_bases, + # Per-round tails: the move offsets ride the staged metadata rows. + dense_move_offsets=dense_move_offsets_row, + swa_move_offsets=swa_move_offsets_row, + per_layer_sources=per_layer, + # The decision rows the plans pack into move sources. + kept_ordinal_rows=self._kept_ordinal_rows, + decision_rows=self._selection_rows_per_request, + valid_seq_lens=self._valid_seq_lens_device, + ), + capacities=dict( + max_requests=max_requests, + keep_count=keep_count, + protected_tail_capacity=int(protected_tail_capacity), + ), + draft=draft_contract, + ) + + # ---- round-ordering events ---------------------------------------------- + # Host staging (pinned metadata + snapshots) reuse fence. + self._staging_reuse_event = torch.cuda.Event() + self._staging_reuse_event.record(torch.cuda.current_stream(device)) + # Manager-stream H2D of the block-offset tables has completed. + self._block_offsets_ready_event = torch.cuda.Event() + # This cohort's compact is done: manager may resize/reuse pages. + self._compaction_done_event = torch.cuda.Event() + self._copy_pending = False @staticmethod def _page_table_pool_ids( @@ -1401,13 +1356,8 @@ def _evict_requests( layout = self._runtime_kv_layout() with nvtx_range_debug("triattention.staging_lookup", color="blue"): # Retained spans always cover the model window (construction rejects budget < window). - bufs = self._buffers_for(layout, prepared) - execute_eviction_round( - bufs, - self.kv_cache_manager, - prepared, - self.draft_kv_cache_manager, - ) + self._ensure_buffers(layout, prepared) + self._execute_eviction_round(prepared) for item in prepared: # Identity cohorts were filtered pre-launch (_periodic_evict). evicted = item["seq_len"] - item["expected_keep_count"] @@ -1417,6 +1367,266 @@ def _evict_requests( item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] return prepared + def _stage_block_offsets( + self, + manager: KVCacheManagerV2, + request_ids: List[int], + host_block_offsets: torch.Tensor, + device_block_offsets: torch.Tensor, + ) -> None: + """Gather the pinned snapshot before the async device copy: resize mutates + the live host table. The round owner has already fenced host-staging reuse.""" + manager.index_mapper.gather_k_block_offsets( + manager.host_kv_cache_block_offsets, + host_block_offsets, + request_ids, + host_block_offsets.shape[-1], + ) + manager._stream.wait_event(self._staging_reuse_event) + 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, + manager._stream.cuda_stream, + ) + self._block_offsets_ready_event.record(manager._stream) + torch.cuda.current_stream(device_block_offsets.device).wait_event( + self._block_offsets_ready_event + ) + + def _cohort_move_offsets( + self, + prepared: Sequence[Dict[str, object]], + ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: + """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep + set plus protected tail per request; rows past the cohort repeat the final + offset and contribute no moves).""" + + def padded_offsets(moves_per_request: List[int]) -> List[int]: + offsets = [0] + for moves in moves_per_request: + offsets.append(offsets[-1] + moves) + offsets.extend(offsets[-1:] * (self._max_requests - len(moves_per_request))) + return offsets + + tails = [int(item["protected_tail"]) for item in prepared] + dense = padded_offsets([self._keep_count + tail for tail in tails]) + swa = None + if self._swa_window is not None: + swa = padded_offsets([self._swa_window + tail for tail in tails]) + draft = None + if self._draft_protected_tail_capacity is not None: + draft = padded_offsets( + [self._keep_count + self._draft_protected_tail_capacity] * len(prepared) + ) + return dense, swa, draft + + def _settle_top_tokens(self, request_count: int) -> None: + """Pick the top-k and settle ties into the kept-ordinal decision rows + (the compaction contract packs them into move sources).""" + 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._keep_count, + 1, + ) + _settle_ties_kernel[(request_count, self._selection_rows_per_request)]( + self._selection_scores_rows, + self._selection_row_lengths, + self._token_starts_device, + self._provisional_rows, + self._kept_ordinal_rows, + WIDTH=self._decode_width, + KEEP_COUNT=self._keep_count, + SELECTION_ROWS=self._selection_rows_per_request, + ) + + def _execute_eviction_round( + self, + prepared: Sequence[Dict[str, object]], + ) -> None: + """Run one eviction round over the prepared cohort: stage the page-table + snapshots and round metadata, then score, select, settle, and compact, and + finally order the manager streams after this cohort's compact (every launch + covers the full request capacity; padded rows carry zero lengths and stay + inert).""" + manager = self.kv_cache_manager + draft_manager = self.draft_kv_cache_manager + with nvtx_range_debug("triattention.page_table_stage", color="orange"): + request_ids = [item["request_id"] for item in prepared] + round_starts = [item["round_start"] for item in prepared] + token_starts = [item["prompt_len"] for item in prepared] + seq_lens = [item["seq_len"] for item in prepared] + dense_move_offsets, swa_move_offsets, draft_move_offsets = self._cohort_move_offsets( + prepared + ) + stream = torch.cuda.current_stream(self._block_offsets_device.device) + # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. + max_round_start = max(round_starts) + rows = ( + (0, round_starts), + (1, seq_lens), + (2, token_starts), + (3, dense_move_offsets), + (4, swa_move_offsets), + (5, draft_move_offsets), + ) + for row, values in rows: + if ( + values is not None + and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF + ): + raise ValueError(f"staged metadata row {row} exceeds the int32 range") + # The one host-staging reuse fence: the previous cohort's async copies + # must complete before the pinned metadata rows AND the pinned + # target/draft block-offset snapshots are rewritten. + if self._copy_pending and not self._staging_reuse_event.query(): + self._staging_reuse_event.synchronize() + host_table = self._request_metadata_host_np + for row, values in rows: + if values is not None: + host_table[row, : len(values)] = values + # Zero lengths keep the score kernel and selection inert for padded rows. + host_table[:3, len(prepared) :] = 0 + grow_mean_phase_table(self._phase, int(max_round_start) + 1) + self._stage_block_offsets( + manager, + request_ids, + self._block_offsets_host, + self._block_offsets_device, + ) + if draft_manager is not None: + self._stage_block_offsets( + draft_manager, + request_ids, + self._draft_block_offsets_host, + self._draft_block_offsets_device, + ) + try: + self._request_metadata_device.copy_(self._request_metadata_host, non_blocking=True) + finally: + # Guards the pinned staging until the asynchronous copies complete. + self._staging_reuse_event.record(stream) + self._copy_pending = True + request_count = len(prepared) + union = self.eviction_mode == "union" + try: + with nvtx_range("triattention.score", color="blue"): + # In-place refresh: the compiled score launches captured these pointers. + _gather_mean_phase_kernel[(request_count,)]( + self._round_starts_device, + self._phase["cos"], + self._phase["sin"], + self._phase["rows"], + self._valid_seq_lens_device, + self._token_starts_device, + self._mean_cos, + self._mean_sin, + self._valid_widths, + self._swa_destination_bases, + self._swa_rebase_delta, + NUM_FREQS=self._phase_num_freqs, + F_BLOCK=self._phase_f_block, + HAS_SWA=self._swa_destination_bases is not None, + num_warps=1, + ) + if union: + # THE union path: the fused score+stats entry, then the + # normalized union reduction, straight off the ctor-bound + # cute handles (only the active cohort size varies). + cu_stream = cuda_driver.CUstream(stream.cuda_stream) + self._compiled_score_stats[request_count]( + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + request_count, + cu_stream, + ) + self._compiled_normalize_union[request_count]( + self._cute_partial_stats, + *self._cute_selection_prefix, + self._cute_union_scores, + request_count, + cu_stream, + ) + columns = min(self._union_scores.shape[1], self._selection_scores_rows.shape[1]) + self._selection_scores_rows[:request_count, :columns].copy_( + self._union_scores[:request_count, :columns] + ) + else: + cu_stream = cuda_driver.CUstream(stream.cuda_stream) + self._compiled_score[request_count]( + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + request_count, + cu_stream, + ) + # Gather each decode window into the [request, layer, head, token] layout of the reduces. + group_size = self._num_q_heads // self._num_kv_heads + num_segments = request_count * self._num_layers + pad = self._padded_head_columns + source = ( + self._score_scratch[ + : self._num_kv_heads * pad * num_segments * self._bucket_seq_len + ] + .view( + self._num_kv_heads, + pad, + request_count, + self._num_layers, + self._bucket_seq_len, + )[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + torch.add( + self._token_starts_device[:request_count].view(-1, 1, 1, 1, 1), + self._gather_columns, + out=self._gather_index_base[:request_count], + ) + self._gather_index_base[:request_count].clamp_(max=self._bucket_seq_len - 1) + columns = self._gather_index[:request_count] + torch.gather( + source, + 4, + columns, + out=self._score_output[:request_count].view( + request_count, + self._num_layers, + self._num_kv_heads, + group_size, + self._decode_width, + ), + ) + with nvtx_range("triattention.select", color="yellow"): + if not union: + prepare_per_head_scores( + self._score_output[:request_count], + self._valid_widths, + self._row_mean, + self._row_inv_std, + self._selection_scores_rows, + self._selection_row_lengths, + per_layer=self.eviction_mode == "per_layer_perhead", + normalize_scores=self.normalize_scores, + ) + self._settle_top_tokens(request_count) + with nvtx_range("triattention.compact", color="purple"): + compact(self.compaction_plan, request_count) + finally: + # Order V2 page-table reuse and resize after this cohort's compact. + self._compaction_done_event.record(stream) + manager._stream.wait_event(self._compaction_done_event) + if draft_manager is not None: + draft_manager._stream.wait_event(self._compaction_done_event) + # ---- helpers: calibration loading ---- def _resolve_calibration(self) -> Dict[str, torch.Tensor]: 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 index 02b311cf38a8..687701e49bb2 100644 --- 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 @@ -1354,285 +1354,3 @@ def _tensor_spec(tensor: torch.Tensor) -> tuple: def _to_cute(tensor: torch.Tensor, *, assumed_align: int = 16) -> cute.Tensor: return from_dlpack(tensor, assumed_align=assumed_align) - - -class TriAttentionCuteScoreRunner: - """Compile and launch the exact SM100 mean-score specialization.""" - - def __init__( - self, - *, - layer_pools: list[torch.Tensor], - layer_indices: list[int], - max_requests: int, - seq_len: int, - num_q_heads: int, - num_freqs: int, - page_ids: torch.Tensor, - seg_page_off: torch.Tensor, - seg_req_id: torch.Tensor, - seg_layer_id: torch.Tensor, - valid_seq_lens: torch.Tensor, - seg_out_offset: torch.Tensor, - token_starts: torch.Tensor, - q_real: torch.Tensor, - q_imag: torch.Tensor, - mlr_coef: torch.Tensor, - mean_cos: torch.Tensor, - mean_sin: torch.Tensor, - freq_scale_sq: torch.Tensor, - output: torch.Tensor, - union_scores: torch.Tensor | None = None, - enable_partial_stats: bool = False, - ) -> None: - # Pool shape [pages, K/V, heads, tokens, dim] of the anchor scored layer. - anchor_pool = layer_pools[layer_indices[0]] - num_layers = len(layer_indices) - num_kv_heads = int(anchor_pool.shape[2]) - tokens_per_block = int(anchor_pool.shape[3]) - self.max_requests = int(max_requests) - self.num_layers = num_layers - # The widest score window (the whole bucket) sizes every start-dependent buffer. - self.width = int(seq_len) - self.num_q_heads = int(num_q_heads) - self.num_kv_heads = num_kv_heads - self.device = output.device - self.sm_count = int(torch.cuda.get_device_properties(output.device).multi_processor_count) - self.enable_partial_stats = bool(enable_partial_stats) - if self.enable_partial_stats and union_scores is None: - raise ValueError( - "TriAttention union fusion requires the persistent union_scores output" - ) - # One [stats_row, page_shard, {count, mean, m2}] record array. - partial_stats_elements = ( - max_requests * num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS * STATS_FIELDS - if self.enable_partial_stats - else 1 - ) - self.partial_stats = torch.empty( - partial_stats_elements, - dtype=torch.float32, - device=output.device, - ) - self.descriptors = _encode_tma_descriptors( - layer_pools, layer_indices, int(num_freqs), int(tokens_per_block) - ) - self._torch_prefix = ( - page_ids, - seg_page_off, - seg_req_id, - seg_layer_id, - valid_seq_lens, - seg_out_offset, - token_starts, - q_real, - q_imag, - mlr_coef, - ) - self._torch_tail = ( - freq_scale_sq, - output, - self.partial_stats, - anchor_pool, - self.descriptors, - ) - # valid_seq_lens/token_starts are only 4-byte-aligned row views, read as per-CTA scalars. - prefix_aligns = (16, 16, 16, 16, 4, 16, 4, 16, 16, 16) - self._cute_prefix = tuple( - _to_cute(tensor, assumed_align=align) - for tensor, align in zip(self._torch_prefix, prefix_aligns) - ) - self._cute_tail = ( - _to_cute(freq_scale_sq), - _to_cute(output), - _to_cute(self.partial_stats), - _to_cute(anchor_pool), - _to_cute(self.descriptors, assumed_align=128), - ) - # Ctor-bound persistent launch operands (one from_dlpack wrap each): - # the round path refreshes their contents in place and launches with - # the active cohort size only. - self._cute_mean_cos = _to_cute(mean_cos.view(-1)) - self._cute_mean_sin = _to_cute(mean_sin.view(-1)) - self._cute_union_scores = ( - _to_cute(union_scores.view(-1)) if self.enable_partial_stats else None - ) - self._compiled: dict[int, object] = {} - self._compiled_stats: dict[int, object] = {} - self._compiled_normalize_union: dict[int, object] = {} - self._page_shards: dict[int, int] = {} - self._cute_selection_prefix = ( - _to_cute(output), - _to_cute(valid_seq_lens, assumed_align=4), - _to_cute(seg_out_offset), - _to_cute(token_starts, assumed_align=4), - ) - self._cute_partial_stats = _to_cute(self.partial_stats) - static_geometry = ( - max_requests, - num_layers, - seq_len, - 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()), - ) - tensor_specs = tuple( - _tensor_spec(tensor) - for tensor in ( - *self._torch_prefix, - mean_cos.view(-1), - mean_sin.view(-1), - *self._torch_tail, - ) - ) - variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] - if max_requests > 1: - variants.append((max_requests, 2)) - # Per-head runners compile the score-only entry; the union runner the fused pipeline. - kernel_kwargs = dict( - num_layers=num_layers, - seq_len=seq_len, - 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()), - ) - if self.enable_partial_stats: - variant_key = "triattention_cute_score_stats" - compiled_entries = self._compiled_stats - else: - variant_key = "triattention_cute_score" - compiled_entries = self._compiled - for request_count, page_shards in variants: - cache_key = ( - variant_key, - static_geometry, - tensor_specs, - request_count, - page_shards, - ) - with _COMPILE_LOCK: - compiled = _COMPILED_KERNELS.get(cache_key) - if compiled is None: - kernel = _TriAttentionScoreKernel( - **kernel_kwargs, - page_shards=page_shards, - write_partial_stats=self.enable_partial_stats, - ) - stream = cuda.CUstream(torch.cuda.current_stream(output.device).cuda_stream) - compiled = cute.compile( - kernel, - *self._cute_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_tail, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled - compiled_entries[request_count] = compiled - self._page_shards[request_count] = page_shards - - if max_requests > 1: - small = compiled_entries.get(1) - large = compiled_entries.get(max_requests) - for request_count in range(1, max_requests + 1): - # Give small cohorts the extra shard while the 2-shard grid stays under two waves. - two_shard_ctas = request_count * num_layers * num_kv_heads * 2 - use_extra_score_shard = two_shard_ctas < 2 * self.sm_count - compiled_entries[request_count] = small if use_extra_score_shard else large - self._page_shards[request_count] = ( - SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 - ) - - if self.enable_partial_stats: - from .triattention_cute_selection import ( - _select_normalize_union_config, - _TriAttentionNormalizeUnionKernel, - ) - - compiled_configs: dict[tuple[int, int, int, int], object] = {} - for request_count in range(1, max_requests + 1): - page_shards = self._page_shards[request_count] - config = _select_normalize_union_config( - request_count, - self.width, - self.sm_count, - ) - config_key = (page_shards, *config) - compiled_selection = compiled_configs.get(config_key) - if compiled_selection is None: - cache_key = ( - "triattention_cute_normalize_union", - static_geometry, - tensor_specs, - config_key, - _tensor_spec(union_scores), - _tensor_spec(self.partial_stats), - ) - with _COMPILE_LOCK: - compiled_selection = _COMPILED_KERNELS.get(cache_key) - if compiled_selection is None: - tokens_per_lane, token_subtiles, row_cluster_ctas = config - kernel = _TriAttentionNormalizeUnionKernel( - num_layers=num_layers, - seq_len=seq_len, - num_q_heads=num_q_heads, - # The finalizer maps real head rows onto N=8-padded planes. - 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, - ) - stream = cuda.CUstream( - torch.cuda.current_stream(output.device).cuda_stream - ) - compiled_selection = cute.compile( - kernel, - _to_cute(self.partial_stats), - *self._cute_selection_prefix, - self._cute_union_scores, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled_selection - compiled_configs[config_key] = compiled_selection - self._compiled_normalize_union[request_count] = compiled_selection - - def launch(self, request_count: int) -> None: - """Launch the CuTe score kernel on the current PyTorch stream over the - active cohort (every tensor operand is ctor-bound).""" - stream = cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream) - self._compiled[request_count]( - *self._cute_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_tail, - request_count, - stream, - ) - - def launch_union_fusion(self, request_count: int) -> None: - """Launch score plus stats followed by normalized union reduction over - the active cohort (every tensor operand is ctor-bound; the union rows - land in the ctor-bound ``union_scores``).""" - stream = cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream) - self._compiled_stats[request_count]( - *self._cute_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_tail, - request_count, - stream, - ) - self._compiled_normalize_union[request_count]( - self._cute_partial_stats, - *self._cute_selection_prefix, - self._cute_union_scores, - request_count, - stream, - ) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index eb963fd2e5ce..8de107cbfd87 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -194,7 +194,7 @@ def capacity_offsets(count): draft=draft, ) # Opaque plans plus a test-side mirror of the caller-owned construction - # inputs (production binds the same values on its buffer namespace); the + # inputs (production binds the same values as manager attributes); the # standalone helpers here need the move-offset rows and SWA staging back. return dict( plans=plans, @@ -229,23 +229,28 @@ def run_compaction(compaction): def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): - """A bare buffer namespace for the bulk page-table copy tests.""" - staging = SimpleNamespace() - staging.max_requests = max_requests - staging.keep_count = 4 - staging.swa_window = None - staging.draft_protected_tail_capacity = None - staging.block_offsets_ready_event = torch.cuda.Event() - staging.compaction_done_event = torch.cuda.Event() - staging.staging_reuse_event = torch.cuda.Event() - staging.copy_pending = False - staging.block_offsets_host = torch.empty( + """A bare manager carrying only the staging attributes, for the bulk + page-table copy tests (mirrors the product's ``_build_buffers`` names).""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention + + staging = TriAttention.__new__(TriAttention) + staging.kv_cache_manager = None + staging.draft_kv_cache_manager = None + staging._max_requests = max_requests + staging._keep_count = 4 + staging._swa_window = None + staging._draft_protected_tail_capacity = None + staging._block_offsets_ready_event = torch.cuda.Event() + staging._compaction_done_event = torch.cuda.Event() + staging._staging_reuse_event = torch.cuda.Event() + staging._copy_pending = False + 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( + staging._identity_copy_indices_host = torch.arange( max_requests, dtype=torch.int32, device="cpu", pin_memory=True ) - staging.block_offsets_device = torch.empty( + staging._block_offsets_device = torch.empty( 1, max_requests, 2, staged_blocks_per_seq, dtype=torch.int32, device=device ) return staging @@ -264,7 +269,10 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): def make_buffer_stubs(manager, *, decode_width=260): - """Stub the calibration/layout surfaces around ``_buffers_for``.""" + """Stub the calibration/layout surfaces around ``_ensure_buffers``. + + Returns the layout dict plus the manager attributes a stubbed + ``_build_buffers`` should set (production sets them in place).""" manager._freq_scale_sq = torch.ones(2) manager._phase = {"rows": 8} manager.calibration = {"omega": torch.ones(2)} @@ -281,14 +289,14 @@ def make_buffer_stubs(manager, *, decode_width=260): layer_group_representative={0: 0, 1: 0}, layer_pool_ids=(0, 0), ) - buffers = SimpleNamespace( - decode_width=decode_width, - page_table_token_capacity=65537, - max_requests=8, - token_starts_device=torch.zeros(8, dtype=torch.int32), - valid_widths=torch.empty(8, dtype=torch.int32), + built_attributes = dict( + _decode_width=decode_width, + _page_table_token_capacity=65537, + _max_requests=8, + _token_starts_device=torch.zeros(8, dtype=torch.int32), + _valid_widths=torch.empty(8, dtype=torch.int32), ) - return layout, buffers + return layout, built_attributes def make_fake_v2(enable_block_reuse=False, *, is_draft=False): @@ -392,15 +400,12 @@ def make_request(request_id, **overrides): @contextmanager def mocked_eviction_internals(manager): """Run the real ``_evict_requests`` body around a mocked round executor.""" - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - - buffers = SimpleNamespace(max_requests=8) with ( mock.patch.object(manager, "_runtime_kv_layout", return_value={}), - mock.patch.object(manager, "_buffers_for", return_value=buffers), - mock.patch.object(module, "execute_eviction_round") as execute, + mock.patch.object(manager, "_ensure_buffers"), + mock.patch.object(manager, "_execute_eviction_round") as execute, ): - yield SimpleNamespace(buffers=buffers, execute=execute) + yield SimpleNamespace(execute=execute) def torch_tri_score_oracle( @@ -497,18 +502,17 @@ def make_cute_buffers( layer_pool_ids=None, normalize_scores=True, ): - """Real eviction buffers over the one-shared-slot default layout; split - reference legs use ``eviction_mode="per_head"`` over the same pools. - ``storage_groups``/``layer_pool_ids`` override the page-table grouping - (``layer_pool_ids`` is the canonical per-layer V2 pool id list).""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - init_eviction_buffers, - ) + """A bare manager with real eviction buffers built over the one-shared-slot + default layout; split reference legs use ``eviction_mode="per_head"`` over + the same pools. ``storage_groups``/``layer_pool_ids`` override the + page-table grouping (``layer_pool_ids`` is the canonical per-layer V2 pool + id list).""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention num_layers = len(layer_pools) assert int(q_real.shape[1]) == num_q_heads - # The constructor takes every capacity explicitly (no test-only - # None-derive path); the widest window defaults keep old call sites. + # The build takes every capacity explicitly (no test-only None-derive + # path); the widest window defaults keep old call sites. if decode_width is None: decode_width = seq_len if storage_groups is None: @@ -527,8 +531,14 @@ def make_cute_buffers( }, layer_pool_ids=layer_pool_ids, ) - return init_eviction_buffers( - eviction_mode=eviction_mode, + manager = TriAttention.__new__(TriAttention) + manager.kv_cache_manager = None + manager.draft_kv_cache_manager = None + manager._draft_protected_tail_capacity = None + manager.eviction_mode = eviction_mode + manager.normalize_scores = normalize_scores + manager._phase = make_phase_table(offsets, omega, seq_len) + manager._build_buffers( layout=layout, calibration=dict( q_real=q_real, @@ -536,7 +546,6 @@ def make_cute_buffers( mlr_coef=mlr_coef, freq_scale_sq=freq_scale_sq, ), - phase=make_phase_table(offsets, omega, seq_len), capacities=dict( max_requests=max_requests, bucket_seq_len=seq_len, @@ -547,67 +556,87 @@ def make_cute_buffers( keep_count=keep_count, protected_tail_capacity=protected_tail_capacity, ), - normalize_scores=normalize_scores, ) + return manager -def write_block_offsets(bufs, encoded): +def write_block_offsets(manager, encoded): """Load a test page table into the staged block-offset plane.""" - bufs.block_offsets_device.zero_() - bufs.block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) + manager._block_offsets_device.zero_() + manager._block_offsets_device[:, : encoded.shape[1], :, : encoded.shape[-1]].copy_(encoded) -def stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts): +def stage_score_metadata(manager, request_count, valid_seq_lens, valid_widths, token_starts): """Stage the per-round score metadata exactly like production (the - compiled runner reads the staged rows via pointer capture).""" + compiled score launches read the staged rows via pointer capture).""" torch.sub( valid_seq_lens[:request_count], token_starts[:request_count], out=valid_widths[:request_count], ) - bufs.valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) - bufs.token_starts_device[:request_count].copy_(token_starts[:request_count]) + manager._valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) + manager._token_starts_device[:request_count].copy_(token_starts[:request_count]) def launch_split_scores( - bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin + manager, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin ): - """The production score-only leg plus the decode-window gather - (``execute_eviction_round``'s per-head sequence, parameterized by count). - Test mean phases load into the runner's ctor-bound buffers, exactly like - the production in-place gather refresh.""" - stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) - bufs.mean_cos[:request_count].copy_(mean_cos[:request_count]) - bufs.mean_sin[:request_count].copy_(mean_sin[:request_count]) - assert request_count in bufs.runner._compiled - bufs.runner.launch(request_count) - num_segments = request_count * bufs.num_layers - group_size = bufs.num_q_heads // bufs.num_kv_heads + """The production score-only leg plus the decode-window gather (the round + executor's per-head sequence, parameterized by count). Test mean phases + load into the build-bound buffers, exactly like the production in-place + gather refresh; the compiled entry fires directly, like the round does.""" + import cuda.bindings.driver as cuda_driver + + stage_score_metadata(manager, request_count, valid_seq_lens, valid_widths, token_starts) + manager._mean_cos[:request_count].copy_(mean_cos[:request_count]) + manager._mean_sin[:request_count].copy_(mean_sin[:request_count]) + assert request_count in manager._compiled_score + stream = cuda_driver.CUstream( + torch.cuda.current_stream(manager._score_scratch.device).cuda_stream + ) + manager._compiled_score[request_count]( + *manager._cute_score_prefix, + manager._cute_mean_cos, + manager._cute_mean_sin, + *manager._cute_score_tail, + request_count, + stream, + ) + num_segments = request_count * manager._num_layers + group_size = manager._num_q_heads // manager._num_kv_heads source = ( - bufs.score_scratch[: bufs.num_kv_heads * 8 * num_segments * bufs.bucket_seq_len] - .view(bufs.num_kv_heads, 8, request_count, bufs.num_layers, bufs.bucket_seq_len)[ - :, :group_size - ] + manager._score_scratch[: manager._num_kv_heads * 8 * num_segments * manager._bucket_seq_len] + .view( + manager._num_kv_heads, 8, request_count, manager._num_layers, manager._bucket_seq_len + )[:, :group_size] .permute(2, 3, 0, 1, 4) ) columns = ( - token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + bufs.gather_columns + token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + manager._gather_columns ) - columns = columns.clamp_(max=bufs.bucket_seq_len - 1).expand( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + columns = columns.clamp_(max=manager._bucket_seq_len - 1).expand( + request_count, + manager._num_layers, + manager._num_kv_heads, + group_size, + manager._decode_width, ) output = torch.full( - (request_count, bufs.num_layers, bufs.num_q_heads, bufs.decode_width), + (request_count, manager._num_layers, manager._num_q_heads, manager._decode_width), float("nan"), dtype=torch.float32, - device=bufs.score_scratch.device, + device=manager._score_scratch.device, ) torch.gather( source, 4, columns, out=output.view( - request_count, bufs.num_layers, bufs.num_kv_heads, group_size, bufs.decode_width + request_count, + manager._num_layers, + manager._num_kv_heads, + group_size, + manager._decode_width, ), ) return output 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 index 1e9a129db573..6ffd40a22e17 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -58,7 +58,7 @@ def _build_case( 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 - bufs = _make_cute_buffers( + tri = _make_cute_buffers( eviction_mode="per_head", layer_pools=pools, max_requests=max_requests, @@ -72,7 +72,7 @@ def _build_case( offsets=offsets_t, decode_width=capacity - prompt_len, ) - _write_block_offsets(bufs, _encode_block_offsets(page_ids)) + _write_block_offsets(tri, _encode_block_offsets(page_ids)) round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) # Mid-page/mid-tile tails; 58 leaves a fully-invalid trailing fragment. @@ -92,7 +92,7 @@ def _build_case( offsets=offsets_t, ) return ( - bufs, + tri, pools, token_starts, valid_seq_lens, @@ -132,7 +132,7 @@ def test_cute_kernel_matches_torch_oracle(case): max_requests = case["max_requests"] num_layers = case["num_layers"] ( - bufs, + tri, pools, token_starts, valid_seq_lens, @@ -141,7 +141,7 @@ def test_cute_kernel_matches_torch_oracle(case): mean_sin, oracle_inputs, ) = _build_case(prompt_len=prompt_len, seed=20260719, **case) - device = bufs.score_scratch.device + device = tri._score_scratch.device oracle = _torch_tri_score_oracle( pools, @@ -158,11 +158,11 @@ def test_cute_kernel_matches_torch_oracle(case): ) # Every count up to capacity is served, nothing beyond. - assert max_requests + 1 not in bufs.runner._compiled + assert max_requests + 1 not in tri._compiled_score for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) scores = _launch_split_scores( - bufs, + tri, request_count, valid_seq_lens, valid_widths, @@ -174,7 +174,7 @@ def test_cute_kernel_matches_torch_oracle(case): request_count, num_layers, case["num_q_heads"], - bufs.decode_width, + tri._decode_width, ) # The score leg owns the per-request decode widths the selection # reduce kernels consume. 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 index c692b7c425c4..315b331c315c 100644 --- 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 @@ -19,21 +19,39 @@ ) -def _launch_union_fusion( - bufs, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out +def _run_fused_union( + tri, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out ): - """The fused score+stats+normalized-union pipeline (THE union path). - Test mean phases load into the runner's ctor-bound buffers; the fused - rows land in the ctor-bound ``bufs.union_scores`` and copy out.""" - _stage_score_metadata(bufs, request_count, valid_seq_lens, valid_widths, token_starts) - bufs.mean_cos[:request_count].copy_(mean_cos[:request_count]) - bufs.mean_sin[:request_count].copy_(mean_sin[:request_count]) + """The fused score+stats+normalized-union pipeline (THE union path), fired + directly off the compiled entries exactly like the round. Test mean phases + load into the build-bound buffers; the fused rows land in the build-bound + ``tri._union_scores`` and copy out.""" + import cuda.bindings.driver as cuda_driver + + _stage_score_metadata(tri, request_count, valid_seq_lens, valid_widths, token_starts) + tri._mean_cos[:request_count].copy_(mean_cos[:request_count]) + tri._mean_sin[:request_count].copy_(mean_sin[:request_count]) assert ( - request_count in bufs.runner._compiled_stats - and request_count in bufs.runner._compiled_normalize_union + request_count in tri._compiled_score_stats + and request_count in tri._compiled_normalize_union + ) + stream = cuda_driver.CUstream(torch.cuda.current_stream(tri._score_scratch.device).cuda_stream) + tri._compiled_score_stats[request_count]( + *tri._cute_score_prefix, + tri._cute_mean_cos, + tri._cute_mean_sin, + *tri._cute_score_tail, + request_count, + stream, + ) + tri._compiled_normalize_union[request_count]( + tri._cute_partial_stats, + *tri._cute_selection_prefix, + tri._cute_union_scores, + request_count, + stream, ) - bufs.runner.launch_union_fusion(request_count) - union_out[:request_count].copy_(bufs.union_scores[:request_count]) + union_out[:request_count].copy_(tri._union_scores[:request_count]) def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: @@ -123,16 +141,16 @@ def test_union_fusion_matches_split_pipeline( omega=omega, offsets=offsets, ) - bufs = _make_cute_buffers(eviction_mode="union", **common) + tri = _make_cute_buffers(eviction_mode="union", **common) # The split reference leg runs on its own score-only buffers. - ref_bufs = _make_cute_buffers(eviction_mode="per_head", **common) + 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(bufs, encoded) - _write_block_offsets(ref_bufs, encoded) + _write_block_offsets(tri, encoded) + _write_block_offsets(ref_tri, encoded) if valid_lens is None: valid_lens = [seq_len, seq_len] valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) @@ -146,7 +164,7 @@ def test_union_fusion_matches_split_pipeline( split_widths = torch.empty(request_count, dtype=torch.int32, device=device) token_starts = torch.tensor(score_starts, dtype=torch.int32, device=device) per_head = _launch_split_scores( - ref_bufs, + ref_tri, request_count, valid_seq_lens, split_widths, @@ -162,8 +180,8 @@ def test_union_fusion_matches_split_pipeline( fused_out = torch.full( (request_count, seq_len), float("nan"), dtype=torch.float32, device=device ) - _launch_union_fusion( - bufs, + _run_fused_union( + tri, request_count, valid_seq_lens, fused_widths, @@ -286,16 +304,16 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: offsets=offsets, decode_width=decode_window, ) - bufs = _make_cute_buffers(eviction_mode="union", max_requests=max_requests, **common) - assert (bufs.score_scratch.numel() > 2**31) == (max_requests == 64) + 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_bufs = _make_cute_buffers(eviction_mode="per_head", max_requests=request_count, **common) + 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 (bufs, ref_bufs): - 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 + 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 valid_seq_lens = torch.zeros(max_requests, dtype=torch.int32, device=device) token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) @@ -306,7 +324,7 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: # the pure-torch union oracle. split_widths = torch.zeros(max_requests, dtype=torch.int32, device=device) per_head = _launch_split_scores( - ref_bufs, + ref_tri, request_count, valid_seq_lens, split_widths, @@ -323,8 +341,8 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: fused_out = torch.full( (max_requests, seq_len), float("nan"), dtype=torch.float32, device=device ) - _launch_union_fusion( - bufs, + _run_fused_union( + tri, max_requests, valid_seq_lens, fused_widths, 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 index e5777b5dcc51..ff84e4067004 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -172,41 +172,42 @@ def test_execute_eviction_round_orders_both_manager_streams(): event = mock.Mock() host = torch.zeros(6, 9, dtype=torch.int32) - buffers = SimpleNamespace( - max_requests=8, - keep_count=4, - eviction_mode="union", - swa_window=None, - compaction_plan=(), - draft_protected_tail_capacity=1, - copy_pending=False, - staging_reuse_event=mock.Mock(), - compaction_done_event=event, - request_metadata_host=host, - request_metadata_host_np=host.numpy(), - request_metadata_device=torch.zeros_like(host), - phase={"cos": None, "sin": None, "rows": 8}, - phase_num_freqs=1, - phase_f_block=1, - round_starts_device=None, - valid_seq_lens_device=None, - token_starts_device=None, - valid_widths=None, - mean_cos=None, - mean_sin=None, - swa_destination_bases=None, - swa_rebase_delta=0, - block_offsets_host=None, - block_offsets_device=torch.zeros(1, dtype=torch.int32), - draft_block_offsets_host=None, - draft_block_offsets_device=None, - ) + tri = TriAttention.__new__(TriAttention) + tri._max_requests = 8 + tri._keep_count = 4 + tri.eviction_mode = "union" + tri._swa_window = None + tri.compaction_plan = () + tri._draft_protected_tail_capacity = 1 + tri._copy_pending = False + tri._staging_reuse_event = mock.Mock() + tri._compaction_done_event = event + tri._request_metadata_host = host + tri._request_metadata_host_np = host.numpy() + tri._request_metadata_device = torch.zeros_like(host) + tri._phase = {"cos": None, "sin": None, "rows": 8} + tri._phase_num_freqs = 1 + tri._phase_f_block = 1 + tri._round_starts_device = None + tri._valid_seq_lens_device = None + tri._token_starts_device = None + tri._valid_widths = None + tri._mean_cos = None + tri._mean_sin = None + tri._swa_destination_bases = None + tri._swa_rebase_delta = 0 + 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 target_stream = mock.Mock() draft_stream = mock.Mock() manager = SimpleNamespace(_stream=target_stream) draft_manager = SimpleNamespace( _stream=draft_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 = SimpleNamespace() prepared = [_make_prepared_item(request_id=7, seq_len=8)] @@ -218,12 +219,12 @@ class Boom(RuntimeError): with ( mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream), mock.patch.object(module, "grow_mean_phase_table"), - mock.patch.object(module, "_stage_block_offsets") as stage, + mock.patch.object(tri, "_stage_block_offsets") as stage, mock.patch.object(module, "_gather_mean_phase_kernel", score_kernel), mock.patch.object(module, "compact") as compact, ): with pytest.raises(Boom): - module.execute_eviction_round(buffers, manager, prepared, draft_manager) + tri._execute_eviction_round(prepared) # Both page-table planes were snapshotted before the round body fired. assert stage.call_count == 2 @@ -349,7 +350,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): # The staged logical position restores the uncompressed # length: physical confirmed plus everything evicted so far # (the prepared item's round_start). - prepared = internals.execute.call_args.args[2] + prepared = internals.execute.call_args.args[0] assert prepared[0]["round_start"] == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. @@ -360,22 +361,20 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert eviction_rounds == 3 assert previous_published == 12 # Each round the draft cache shrinks with the target, and the one - # executor call carries both managers, whose streams it orders after the - # compact launches. + # executor call runs on the manager itself, which carries both cache + # managers whose streams it orders after the compact launches. 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 call.args[0] is internals.buffers - assert call.args[1] is target - assert call.args[3] is draft_manager + assert len(call.args) == 1 assert call.kwargs == {} def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - manager = _make_triattention(budget=4) - layout, buffers = _make_buffer_stubs(manager) + layout, built_attributes = _make_buffer_stubs(manager) # The one-time host block-offset table shape gate reads the real manager # tables (int32 [pools, slots, K/V, blocks]). manager.kv_cache_manager.host_kv_cache_block_offsets = torch.zeros( @@ -401,55 +400,59 @@ def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): _make_prepared_item(_make_request(7), request_id=7, seq_len=8, expected_keep_count=4) ] - with mock.patch.object( - module, - "init_eviction_buffers", - return_value=buffers, - ) as prepare: - resources = manager._buffers_for(layout, prepared) + def apply_built(**kwargs): + for name, value in built_attributes.items(): + setattr(manager, name, value) + + phase = manager._phase + with mock.patch.object(manager, "_build_buffers", side_effect=apply_built) as prepare: + manager._ensure_buffers(layout, prepared) # Request capacity follows the executor limits, while the score # bucket follows what the cohort actually presents (power-of-two, # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. - assert resources is buffers + # The stubbed build's capacities became the resident manager state. + assert manager._buffers_built + assert manager._decode_width == built_attributes["_decode_width"] kwargs = prepare.call_args.kwargs - assert kwargs["eviction_mode"] == "union" + # The mode and the shared phase-table dict live on the manager itself + # and thread through unchanged (no longer build arguments). + assert manager.eviction_mode == "union" + assert manager._phase is phase assert kwargs["capacities"]["max_requests"] == 8 assert kwargs["capacities"]["decode_width"] == 4 + 2 * 128 assert kwargs["capacities"]["bucket_seq_len"] == 1024 assert kwargs["capacities"]["page_table_token_capacity"] == 1024 + 1 assert kwargs["draft"]["page_table_token_capacity"] == 1024 + 1 assert kwargs["draft"]["layout"] is manager._draft_runtime_kv_layout.return_value - # Migrated from the pipeline buffer-kwargs test: the budget, the - # shared phase-table dict, and the pool keys thread through unchanged. + # Migrated from the pipeline buffer-kwargs test: the budget and the + # pool keys thread through unchanged. assert kwargs["capacities"]["keep_count"] == manager.budget - assert kwargs["phase"] is manager._phase assert kwargs["layout"] is layout assert list(kwargs["layout"]["layer_pool_ids"]) == list(layout["layer_pool_ids"]) # A second round within the resident capacities reuses the buffers # (and with them the compaction launch data they carry). - assert manager._buffers_for(layout, prepared) is resources + manager._ensure_buffers(layout, prepared) assert prepare.call_count == 1 # A cohort that outgrows the resident capacities rebuilds the whole - # buffer namespace, compaction included. + # buffer state, compaction included. grown = [ _make_prepared_item( _make_request(7), request_id=7, - seq_len=8 + buffers.decode_width, + seq_len=8 + built_attributes["_decode_width"], expected_keep_count=4, ) ] - rebuilt_buffers = SimpleNamespace( - decode_width=buffers.decode_width + 8, - page_table_token_capacity=buffers.page_table_token_capacity, - max_requests=buffers.max_requests, - ) - prepare.return_value = rebuilt_buffers - rebuilt = manager._buffers_for(layout, grown) - assert rebuilt is not resources - assert rebuilt is rebuilt_buffers + + def apply_rebuilt(**kwargs): + apply_built() + manager._decode_width = built_attributes["_decode_width"] + 8 + + prepare.side_effect = apply_rebuilt + manager._ensure_buffers(layout, grown) assert prepare.call_count == 2 - assert manager._buffers is rebuilt_buffers + assert manager._buffers_built + assert manager._decode_width == built_attributes["_decode_width"] + 8 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index d6c3dd8916f5..5eed75eb5060 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -195,14 +195,15 @@ def test_request_init_and_finish_lifecycle(self): assert set(triattention._request_states) == {11, 12} buffers = object() - triattention._buffers = buffers + triattention._buffers_built = True + triattention._score_scratch = buffers batch = SimpleNamespace() triattention._prepared_generation_batch = batch triattention.on_request_finish(_make_request(11)) triattention.on_request_finish(_make_request(12)) assert triattention._request_states == {} assert triattention._prepared_generation_batch is batch - assert triattention._buffers is buffers + assert triattention._buffers_built and triattention._score_scratch is buffers def test_resolve_accepts_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() @@ -562,20 +563,17 @@ def test_union_rejects_unnormalized_scores(self): def test_execute_rejects_int32_overflowing_round_starts(self): # Round starts past the int32 metadata range fail loudly (in the host # metadata build) before any GPU work is enqueued. - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - execute_eviction_round, - ) - device = torch.device("cuda", torch.cuda.current_device()) staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) gather = mock.Mock() manager = _make_staging_manager( torch.zeros(1, 2, 2, 12, dtype=torch.int32), gather, torch.cuda.Stream(device=device) ) + staging.kv_cache_manager = manager prepared = [_make_prepared_item(request_id=7, seq_len=64, round_start=2**31)] with pytest.raises((RuntimeError, OverflowError, ValueError)): - execute_eviction_round(staging, manager, prepared) + staging._execute_eviction_round(prepared) assert gather.call_count == 0 def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): @@ -586,13 +584,8 @@ def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): index-mapper slot assignment) may mutate as soon as staging returns; the staged device tables must reflect the values at staging time. A subsequent bulk copy must also wait until the previous round's - consumers (ordered by ``execute_eviction_round``'s completion event) - are done. + consumers (ordered by the round executor's completion event) are done. """ - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _stage_block_offsets, - ) - device = torch.device("cuda", torch.cuda.current_device()) current_stream = torch.cuda.current_stream(device) manager_stream = torch.cuda.Stream(device=device) @@ -617,17 +610,16 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): gather = mock.Mock(side_effect=gather_k_block_offsets) staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) - staging.staging_reuse_event.record(current_stream) + staging._staging_reuse_event.record(current_stream) manager = _make_staging_manager(host_table, gather, manager_stream) def stage_once(): # Raises on any staging failure; success returns None. - _stage_block_offsets( - staging, + staging._stage_block_offsets( manager, [7], - staging.block_offsets_host, - staging.block_offsets_device, + staging._block_offsets_host, + staging._block_offsets_device, ) # Round 1: mutate the host table and the slot assignment right after @@ -641,13 +633,13 @@ def stage_once(): side_effect=AssertionError("page-table staging used torch.index_select"), ): stage_once() - assert staging.block_offsets_host.shape == (1, 1, 2, 8) + 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] + 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 cohort. host_table[0, 0, 0, :5] = torch.tensor([18, 19, 20, 21, 22], dtype=torch.int32) @@ -659,48 +651,43 @@ def stage_once(): 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] + 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 delayed consumer read (snapshot) queued before the # round's completion ordering must complete before the next bulk # copy overwrites the device tables. selected_slot[0] = 0 manager_stream.synchronize() - snapshot = torch.empty_like(staging.block_offsets_device) + snapshot = torch.empty_like(staging._block_offsets_device) torch.cuda._sleep(20_000_000) - snapshot.copy_(staging.block_offsets_device) - # ``execute_eviction_round``'s completion ordering: one event records - # the consumers and the manager stream waits on it. - staging.compaction_done_event.record(torch.cuda.current_stream(device)) - manager_stream.wait_event(staging.compaction_done_event) + snapshot.copy_(staging._block_offsets_device) + # The round executor's completion ordering: one event records the + # consumers and the manager stream waits on it. + staging._compaction_done_event.record(torch.cuda.current_stream(device)) + manager_stream.wait_event(staging._compaction_done_event) 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] + assert staging._block_offsets_device[0, 0, 0, :5].tolist() == [46, 48, 50, 52, 54] def test_cohort_move_offsets_stage_keep_plus_tail_and_pad_rows(self): # The derived move offsets stage keep + tail moves per request # (keep_count=4 -> [6, 7]); padded rows past the cohort repeat the # final offset and contribute no moves. (The executor-call contract # itself is pinned by the overlap-tail and draft publication tests.) - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _cohort_move_offsets, - ) - - offsets_buffers = SimpleNamespace( - max_requests=8, - keep_count=4, - swa_window=None, - draft_protected_tail_capacity=None, - ) + offsets_manager = TriAttention.__new__(TriAttention) + offsets_manager._max_requests = 8 + offsets_manager._keep_count = 4 + offsets_manager._swa_window = None + offsets_manager._draft_protected_tail_capacity = None prepared = [ _make_prepared_item(request_id=7, seq_len=8, protected_tail=2), _make_prepared_item(request_id=8, seq_len=10, protected_tail=3), ] - dense, swa, draft = _cohort_move_offsets(offsets_buffers, prepared) + dense, swa, draft = offsets_manager._cohort_move_offsets(prepared) assert dense == [0, 6, 13, 13, 13, 13, 13, 13, 13] assert swa is None assert draft is None @@ -758,7 +745,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self): round_starts = round_device[:request_count].tolist() seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) - bufs = _make_cute_buffers( + tri = _make_cute_buffers( eviction_mode="per_head", layer_pools=pools, max_requests=max_requests, @@ -794,6 +781,7 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): torch.cuda.Stream(device=device), num_slots=num_layers, ) + tri.kv_cache_manager = manager def prepared_cohort(): return [ @@ -807,13 +795,13 @@ def prepared_cohort(): ] score_sentinel = -12345.0 - bufs.score_output.fill_(score_sentinel) + tri._score_output.fill_(score_sentinel) # The compact stage is stubbed to a no-op: this test owns the score # buffers only, never a staged move decision. with mock.patch.object(module, "compact"): - module.execute_eviction_round(bufs, manager, prepared_cohort()) - fixed = bufs.score_output.clone() - assert bufs.valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + tri._execute_eviction_round(prepared_cohort()) + fixed = tri._score_output.clone() + assert tri._valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] oracle = _torch_tri_score_oracle( pools, @@ -847,12 +835,12 @@ def prepared_cohort(): ) ) expected_second_widths = valid_seq_lens - prompt_len - bufs.score_output.fill_(score_sentinel) - bufs.valid_widths.fill_(-1) + tri._score_output.fill_(score_sentinel) + tri._valid_widths.fill_(-1) with mock.patch.object(module, "compact"): - module.execute_eviction_round(bufs, manager, prepared_cohort()) - second_launch = bufs.score_output.clone() - assert torch.equal(bufs.valid_widths, expected_second_widths) + tri._execute_eviction_round(prepared_cohort()) + second_launch = tri._score_output.clone() + assert torch.equal(tri._valid_widths, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index c4497cc66c05..7912e7a2a550 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -2,8 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 -from types import SimpleNamespace - import pytest import torch from conftest import build_compaction as _build_compaction @@ -15,7 +13,7 @@ from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import settle_top_tokens +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( prepare_per_head_scores, ) @@ -47,83 +45,69 @@ def _make_selection_buffers( num_query_heads=1, num_kv_heads=1, ): - """Selection-only buffers for the mode, without CuTe score state or - compaction (the settle launch writes only the kept-ordinal rows). - Mirrors the product's canonical row-major selection allocation.""" - bufs = SimpleNamespace( - eviction_mode=eviction_mode, - max_requests=max_requests, - decode_width=width, - keep_count=keep_count, - num_layers=num_layers, - num_q_heads=num_query_heads, - num_kv_heads=num_kv_heads, - ) - bufs.valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) - bufs.token_starts_device = torch.zeros(max_requests, dtype=torch.int32, device=device) + """A bare manager with selection-only attributes for the mode, without + CuTe score state or compaction (the settle launch writes only the + kept-ordinal rows). Mirrors the product's canonical row-major selection + allocation; ``_settle_top_tokens`` reads exactly these attributes.""" + tri = TriAttention.__new__(TriAttention) + tri.eviction_mode = eviction_mode + tri._max_requests = max_requests + tri._decode_width = width + tri._keep_count = keep_count + tri._num_layers = num_layers + tri._num_q_heads = num_query_heads + tri._num_kv_heads = num_kv_heads + tri._valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) + tri._token_starts_device = torch.zeros(max_requests, dtype=torch.int32, device=device) if eviction_mode == "union": - bufs.selection_rows_per_request = 1 - bufs.selection_scores_rows = torch.empty( + tri._selection_rows_per_request = 1 + tri._selection_scores_rows = torch.empty( (max_requests, width), dtype=torch.float32, device=device ) - bufs.selection_row_lengths = bufs.valid_widths + tri._selection_row_lengths = tri._valid_widths # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. - bufs.provisional_rows = torch.zeros( + tri._provisional_rows = torch.zeros( (max_requests, keep_count), dtype=torch.int32, device=device ) - bufs.kept_ordinal_rows = torch.empty( + tri._kept_ordinal_rows = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=device ) else: selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads - bufs.selection_rows_per_request = selection_rows - bufs.row_mean = torch.empty( + tri._selection_rows_per_request = selection_rows + tri._row_mean = torch.empty( max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device ) - bufs.row_inv_std = torch.empty_like(bufs.row_mean) - bufs.selection_scores_rows = torch.empty( + tri._row_inv_std = torch.empty_like(tri._row_mean) + tri._selection_scores_rows = torch.empty( (max_requests * selection_rows, width), dtype=torch.float32, device=device ) - bufs.selection_row_lengths = torch.full( + tri._selection_row_lengths = torch.full( (max_requests * selection_rows,), width, dtype=torch.int32, device=device ) - bufs.provisional_rows = torch.zeros( + tri._provisional_rows = torch.zeros( (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - bufs.kept_ordinal_rows = torch.empty( + tri._kept_ordinal_rows = torch.empty( (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device ) - # The launch args mirror the product's ``bufs.settle_args``/ - # ``bufs.settle_kwargs`` order and keys exactly. - bufs.settle_args = ( - bufs.selection_scores_rows, - bufs.selection_row_lengths, - bufs.token_starts_device, - bufs.provisional_rows, - bufs.kept_ordinal_rows, - ) - bufs.settle_kwargs = dict( - WIDTH=width, - KEEP_COUNT=keep_count, - SELECTION_ROWS=bufs.selection_rows_per_request, - ) - return bufs + return tri -def _select_per_head(bufs, scores, *, normalize_scores): +def _select_per_head(tri, scores, *, normalize_scores): """The per-head selection flow: reduce kernels, then top-k settle.""" prepare_per_head_scores( scores, - bufs.valid_widths, - bufs.row_mean, - bufs.row_inv_std, - bufs.selection_scores_rows, - bufs.selection_row_lengths, - per_layer=bufs.eviction_mode == "per_layer_perhead", + tri._valid_widths, + tri._row_mean, + tri._row_inv_std, + tri._selection_scores_rows, + tri._selection_row_lengths, + per_layer=tri.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - settle_top_tokens(bufs, bufs.max_requests) + tri._settle_top_tokens(tri._max_requests) def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: @@ -193,7 +177,7 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( device = torch.device("cuda", torch.cuda.current_device()) stream = torch.cuda.Stream(device=device) with torch.cuda.stream(stream): - bufs = _make_selection_buffers( + tri = _make_selection_buffers( eviction_mode=eviction_mode, width=width, keep_count=keep_count, @@ -203,13 +187,13 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( num_query_heads=query_heads, num_kv_heads=kv_heads, ) - bufs.valid_widths.copy_(valid_widths.to(device)) + tri._valid_widths.copy_(valid_widths.to(device)) scores = scores_cpu.to(device) - keep_shape = (request_count, bufs.selection_rows_per_request, keep_count) - _select_per_head(bufs, scores, normalize_scores=normalize_scores) - first = bufs.kept_ordinal_rows.view(keep_shape).cpu() - _select_per_head(bufs, scores, normalize_scores=normalize_scores) - second = bufs.kept_ordinal_rows.view(keep_shape).cpu() + 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) @@ -234,20 +218,20 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid device=device, ).to(torch.float32) valid_widths = (width, width - 32) - bufs = _make_selection_buffers( + tri = _make_selection_buffers( eviction_mode="union", width=width, keep_count=keep_count, device=device, max_requests=request_count, ) - bufs.valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) - bufs.token_starts_device[:request_count].copy_( + tri._valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) + tri._token_starts_device[:request_count].copy_( torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) - bufs.selection_scores_rows.copy_(scores.amax(dim=1)) - settle_top_tokens(bufs, bufs.max_requests) - actual = bufs.kept_ordinal_rows.cpu() + tri._selection_scores_rows.copy_(scores.amax(dim=1)) + tri._settle_top_tokens(tri._max_requests) + actual = tri._kept_ordinal_rows.cpu() combined = scores.amax(dim=1).cpu() for request, valid_width in enumerate(valid_widths): @@ -429,9 +413,6 @@ def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode) 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") - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - execute_eviction_round, - ) device = torch.device("cuda", torch.cuda.current_device()) num_layers = 3 @@ -475,7 +456,7 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): mlr_coef[:, :, 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - bufs = _make_cute_buffers( + tri = _make_cute_buffers( eviction_mode="per_layer_perhead", layer_pools=pools, max_requests=1, @@ -497,16 +478,16 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): normalize_scores=False, ) # No SWA in this layout: no window, no rebase row for the phase gather. - assert bufs.swa_window is None - assert bufs.swa_destination_bases is None + 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 bufs.block_offsets_host.shape == bufs.block_offsets_device.shape - assert bufs.block_offsets_host.shape[:3] == (2, 1, 2) - assert bufs.block_offsets_host.shape[-1] % 4 == 0 - assert bufs.block_offsets_host.dtype == bufs.block_offsets_device.dtype == torch.int32 - assert bufs.block_offsets_host.is_contiguous() and bufs.block_offsets_device.is_contiguous() - assert bufs.block_offsets_host.is_pinned() + 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; @@ -525,8 +506,9 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): num_slots=2, ) prepared = [_make_prepared_item(request_id=7, seq_len=seq_len, round_start=0)] - execute_eviction_round(bufs, manager, prepared) - assert torch.equal(bufs.kept_ordinal_rows.view_as(expected_keep), expected_keep) + tri.kv_cache_manager = manager + tri._execute_eviction_round(prepared) + 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( @@ -549,9 +531,6 @@ def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): pytest.importorskip("cutlass") import tensorrt_llm import tensorrt_llm.bindings - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - execute_eviction_round, - ) 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 @@ -653,7 +632,7 @@ def expected_keep() -> torch.Tensor: mlr_coef[..., 0] = 1 freq_scale_sq = torch.zeros(num_freqs, dtype=torch.float32, device=device) freq_scale_sq[0] = 1 - bufs = _make_cute_buffers( + tri = _make_cute_buffers( eviction_mode="union", layer_pools=[pool], max_requests=1, @@ -670,6 +649,7 @@ def expected_keep() -> torch.Tensor: page_table_token_capacity=seq_len + protected_tail, protected_tail_capacity=protected_tail, ) + tri.kv_cache_manager = manager def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) @@ -686,8 +666,8 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: # the derived move offsets stage keep_count + protected_tail # moves. Z-normalization is monotonic per row, so the raw-score # keep set is unchanged. - execute_eviction_round(bufs, manager, prepared) - selected = bufs.kept_ordinal_rows[0].clone().to(torch.long) + tri._execute_eviction_round(prepared) + selected = tri._kept_ordinal_rows[0].clone().to(torch.long) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) after = snapshot(compacted_capacity) From e7819dd635aa17cf174417bf9303e1110058e79f Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 07:40:00 -0700 Subject: [PATCH 130/178] [None][chore] Reflow the compression config docstring for the legacy lint Signed-off-by: tianruih --- tensorrt_llm/llmapi/llm_args.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 0d6bf83401d9..ec52667438f7 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3423,10 +3423,12 @@ def kv_cache_compression_mode(self): 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.""" + """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. + """ algorithm: Literal["triattention"] = "triattention" eviction_mode: Literal["union", "per_head", "per_layer_perhead"] = Field( default="union", From e04902cca99c7e9668fe78c1dd4ba75a773ba2ef Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 08:15:58 -0700 Subject: [PATCH 131/178] [None][test] Final backstop: drop the config tautology and reflection block; reword the runner-era comment (knife 31b) Signed-off-by: tianruih --- .../test_triattention_cute_score.py | 4 +-- .../test_triattention_pipeline.py | 35 ------------------- 2 files changed, 2 insertions(+), 37 deletions(-) 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 index 6ffd40a22e17..b87391776347 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -204,8 +204,8 @@ def test_unsupported_geometry_raises_at_buffer_construction(): for _ in range(num_layers) ] calib = torch.randn(num_layers, 2, num_freqs, device=device) - # No rewrap: the runner's own contract error surfaces directly (the - # fp32 pools trip the BF16 gate first, at TMA descriptor encoding). + # 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", diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 5eed75eb5060..729b88319b6d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -35,7 +35,6 @@ from conftest import make_triattention as _make_triattention from conftest import mocked_eviction_internals as _mocked_eviction_internals from conftest import torch_tri_score_oracle as _torch_tri_score_oracle -from pydantic import ValidationError # TriAttention lives in the kv_cache_compression package. It exposes only the # compression manager -- no attention classes or KV-cache-manager subclass. @@ -86,34 +85,6 @@ def _make_hf_config(**values): class TestConfigAndFactory: - def test_llm_args_dispatch_and_validation(self): - from tensorrt_llm.llmapi.llm_args import TorchLlmArgs - - tri_args = TorchLlmArgs( - model="dummy", - kv_cache_compression_config={ - "algorithm": "triattention", - "model_path": "/models/test", - "calibration_path": "/calib/test.pt", - }, - ) - assert isinstance( - tri_args.kv_cache_compression_config, - TriAttentionKvCacheCompressionConfig, - ) - assert tri_args.kv_cache_compression_config.budget == 2048 - assert tri_args.kv_cache_compression_config.beta == 128 - - # Dispatch is by the algorithm tag, so an unknown algorithm fails - # config validation instead of falling back to a base config. - with pytest.raises(ValidationError): - TorchLlmArgs( - model="dummy", - kv_cache_compression_config={"algorithm": "future_method"}, - ) - with pytest.raises(ValidationError): - TriAttentionKvCacheCompressionConfig(eviction_mode="made_up_mode") - def test_factory_returns_triattention_and_propagates_config_fields(self): # Calibration is deferred to the first request, so construction needs # no calibration file or CUDA. @@ -337,12 +308,6 @@ def test_identity_selection_is_filtered_before_launch(self): class TestEvictionLifecycle: def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): - # Hooks only, never the base template methods. - assert "prepare_resources" not in TriAttention.__dict__ - assert "update_resources" not in TriAttention.__dict__ - assert "on_generation_step_begin" in TriAttention.__dict__ - assert "on_generation_step_end" in TriAttention.__dict__ - manager = _make_triattention() batch = SimpleNamespace( context_requests=[], From 9a99caf2a454b5d338795cc6dac4eaecf6707ad4 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 08:24:42 -0700 Subject: [PATCH 132/178] [None][refactor] Fold the launch-record grouping into init as one families loop (knife 35) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 196 +++++++++--------- 1 file changed, 101 insertions(+), 95 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 5be22bf03969..b46964b88ce1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -121,60 +121,6 @@ def _make_move_indices( ) -def _make_compact_launches( - entries: List[Tuple[int, torch.Tensor, torch.Tensor]], - layer_pool_ids: Tuple[int, ...], - *, - move_indices: torch.Tensor, - move_offsets: torch.Tensor, - destination_bases: torch.Tensor, - per_layer_slots: Optional[Dict[int, int]] = None, -) -> Tuple[Dict[str, object], ...]: - """Batch layers into one launch record per uniform V2 pool (grouped by - canonical pool id plus pool/page-table geometry), carrying the native - ``sparse_kv_cache_compact_layers`` operands as plain named fields.""" - device = entries[0][1].device - grouped = OrderedDict() - for layer, pool, page_table in entries: - key = ( - layer_pool_ids[layer], - str(pool.dtype), - str(pool.device), - tuple(int(value) for value in pool.shape[1:]), - tuple(int(value) for value in page_table.shape), - ) - grouped.setdefault(key, []).append((layer, pool, page_table)) - - launches = [] - for group_entries in grouped.values(): - layers = tuple(entry[0] for entry in group_entries) - pools = list(entry[1] for entry in group_entries) - page_tables = tuple(entry[2] for entry in group_entries) - source_layer_indices = None - if per_layer_slots is not None: - source_layer_indices = torch.tensor( - [per_layer_slots[layer] for layer in layers], - dtype=torch.int32, - device=device, - ) - launches.append( - dict( - pools=pools, - pool_pointers=torch.tensor( - [pool.data_ptr() for pool in pools], - dtype=torch.int64, - device=device, - ), - page_table=page_tables[0], - move_indices=move_indices, - move_offsets=move_offsets, - destination_bases=destination_bases, - source_layer_indices=source_layer_indices, - ) - ) - return tuple(launches) - - def init_compaction_buffers( *, target: Dict[str, object], @@ -201,8 +147,8 @@ def init_compaction_buffers( Returns one plan per compacted cache family (target, then the optional draft): plain contract fields -- decision inputs, move-index buffers, - pack geometry ints, and the grouped native launch records -- that only - :func:`compact` interprets. + pack geometry ints, and per-pool batched native-operand records built + here at init time -- that only :func:`compact` reads at launch. """ layer_pools = target["layer_pools"] dense_layers = tuple(int(layer) for layer in target["dense_layers"]) @@ -263,16 +209,20 @@ def init_compaction_buffers( if swa_move_indices is not None: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) - target_launches = list( - _make_compact_launches( + # Families: (entries, pool ids, move indices, move offsets, destination + # bases, per-layer slots, is_draft). One grouping loop below batches each + # family into per-pool native-operand records (all init-time construction). + families = [ + ( dense_entries, layer_pool_ids, - move_indices=dense_move_indices, - move_offsets=dense_move_offsets, - destination_bases=token_starts, - per_layer_slots=dense_slots, - ) - ) + dense_move_indices, + dense_move_offsets, + token_starts, + dense_slots, + False, + ), + ] if swa_layers: # SWA layers stage against their own page-table slots. swa_entries = [ @@ -283,35 +233,18 @@ def init_compaction_buffers( ) for layer in swa_layers ] - target_launches.extend( - _make_compact_launches( + families.append( + ( swa_entries, layer_pool_ids, - move_indices=swa_move_indices, - move_offsets=swa_move_offsets, - destination_bases=target["swa_destination_bases"], + swa_move_indices, + swa_move_offsets, + target["swa_destination_bases"], + None, + False, ) ) - - target_plan = dict( - kept_ordinal_rows=kept_ordinal_rows, - valid_seq_lens=valid_seq_lens, - decision_rows=decision_rows, - per_layer_sources=per_layer_sources, - dense_move_offsets=dense_move_offsets, - dense_move_indices=dense_move_indices, - swa_move_offsets=swa_move_offsets, - swa_move_indices=swa_move_indices, - swa_window=swa_window, - keep_count=keep_count, - move_capacity=move_capacity, - num_kv_heads=num_kv_heads, - dense_total=int(dense_move_indices.shape[-1]), - swa_total=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, - launches=tuple(target_launches), - ) - - plans = [target_plan] + draft_move_indices = None if draft is not None: if decision_rows != 1: raise ValueError( @@ -342,13 +275,86 @@ def init_compaction_buffers( ) for layer in draft_dense_layers ] - draft_launches = _make_compact_launches( - draft_entries, - draft_layer_pool_ids, - move_indices=draft_move_indices, - move_offsets=draft["dense_move_offsets"], - destination_bases=token_starts, + families.append( + ( + draft_entries, + draft_layer_pool_ids, + draft_move_indices, + draft["dense_move_offsets"], + token_starts, + None, + True, + ) ) + + target_launches: List[Dict[str, object]] = [] + draft_launch_records: List[Dict[str, object]] = [] + for ( + entries, + pool_ids, + move_indices, + move_offsets, + destination_bases, + slots, + is_draft, + ) in families: + grouped = OrderedDict() + for layer, pool, page_table in entries: + key = ( + pool_ids[layer], + str(pool.dtype), + str(pool.device), + tuple(int(value) for value in pool.shape[1:]), + tuple(int(value) for value in page_table.shape), + ) + grouped.setdefault(key, []).append((layer, pool, page_table)) + for group_entries in grouped.values(): + layers = tuple(entry[0] for entry in group_entries) + pools = list(entry[1] for entry in group_entries) + page_tables = tuple(entry[2] for entry in group_entries) + source_layer_indices = None + if slots is not None: + source_layer_indices = torch.tensor( + [slots[layer] for layer in layers], + dtype=torch.int32, + device=device, + ) + record = dict( + pools=pools, + pool_pointers=torch.tensor( + [pool.data_ptr() for pool in pools], + dtype=torch.int64, + device=device, + ), + page_table=page_tables[0], + move_indices=move_indices, + move_offsets=move_offsets, + destination_bases=destination_bases, + source_layer_indices=source_layer_indices, + ) + (draft_launch_records if is_draft else target_launches).append(record) + draft_launches = tuple(draft_launch_records) + + target_plan = dict( + kept_ordinal_rows=kept_ordinal_rows, + valid_seq_lens=valid_seq_lens, + decision_rows=decision_rows, + per_layer_sources=per_layer_sources, + dense_move_offsets=dense_move_offsets, + dense_move_indices=dense_move_indices, + swa_move_offsets=swa_move_offsets, + swa_move_indices=swa_move_indices, + swa_window=swa_window, + keep_count=keep_count, + move_capacity=move_capacity, + num_kv_heads=num_kv_heads, + dense_total=int(dense_move_indices.shape[-1]), + swa_total=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, + launches=tuple(target_launches), + ) + + plans = [target_plan] + if draft is not None: draft_plan = dict( kept_ordinal_rows=kept_ordinal_rows, valid_seq_lens=valid_seq_lens, From 702954a4b997dcd29f2b6f771dd7ab566304974d Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 09:17:10 -0700 Subject: [PATCH 133/178] [None][refactor] Delete the construction-time V2 compatibility validator (knife 37) Capability walls fall to point-of-use owners: kv_factor is enforced by the native compact op's TORCH_CHECKs, non-V2 managers by the base manager ctor, and the remaining surface is documented in the TriAttention README. The draft admission checks move to the executor call-site gate (validate_kv_cache_compression_with_spec), which already owns the draft/spec domain. Signed-off-by: tianruih --- examples/triattention/README.md | 1 + .../triattention/triattention.py | 57 ---------------- tensorrt_llm/_torch/pyexecutor/_util.py | 17 +++++ .../_torch/kv_cache_compression/conftest.py | 3 - .../test_triattention_draft_cocompaction.py | 68 ++++++------------- .../test_triattention_pipeline.py | 1 - 6 files changed, 38 insertions(+), 109 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index 8fc067e7a29b..afc75c48fd7d 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -27,6 +27,7 @@ TriAttention is integrated into TensorRT LLM as a KV-cache compression manager o 1. TriAttention requires `enable_block_reuse=False` in the KV-cache configuration — the eviction physically rewrites stored keys, which is incompatible with block reuse. The construction step rejects a cache manager that has block reuse enabled. 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. Requires full-attention KVCacheManagerV2 lifecycles; attention-DP, disaggregated serving, native SWA/VSWA/SSM pools, and MLA caches are unsupported. ## Calibration diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 6d4a3eb83ac1..de25942140d0 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -37,7 +37,6 @@ copy_batch_block_offsets_to_device, ) from tensorrt_llm.logger import logger -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig from ..compaction import compact, init_compaction_buffers from .triattention_kernels import ( @@ -154,9 +153,6 @@ def __init__( # In-flight overlap batch reference; membership resolves lazily. self._prepared_generation_batch: Optional[object] = None self._prepared_generation_ids: Optional[set] = None - # Manager-lifetime capability gates: everything read there is fixed at - # construction, so validation runs once here. - self._validate_v2_compatibility() # Manager-lifetime constants (V2 fixes every input at construction): # protected tails are num_extra + reserved draft width + 1 sampled token. self._protected_tail_capacity = ( @@ -260,59 +256,6 @@ def _ensure_calibrated(self) -> None: ).contiguous() self._calibrated = True - def _validate_v2_compatibility(self) -> None: - # The base manager already enforces KVCacheManagerV2 target/draft types, - # and V2 construction guarantees beam width one. - manager = self.kv_cache_manager - if manager.kv_factor != 2: - raise ValueError( - "TriAttention requires a standard key/value KV cache; " - "MLA/SELFKONLY caches are not supported" - ) - if manager.mapping.enable_attention_dp: - raise ValueError("TriAttention does not support attention DP") - if manager.is_disagg: - raise ValueError("TriAttention does not support disaggregated serving") - if manager.enable_swa_scratch_reuse: - raise RuntimeError("TriAttention does not support V2 SWA scratch page-table remapping") - # Speculative feature gates run in the factory; the draft cache itself - # is validated here (V2 already forces scratch reuse off for drafts). - draft_manager = self.draft_kv_cache_manager - if draft_manager is not None: - if not draft_manager.is_draft: - raise ValueError( - "TriAttention speculative compatibility requires the actual " - "separate draft KV cache manager" - ) - if draft_manager.kv_factor != 2: - raise ValueError( - "TriAttention compresses the draft KV cache together with " - "the target, so the draft cache must be a standard " - "key/value cache" - ) - if self.eviction_mode != "union": - raise ValueError( - "TriAttention draft KV co-compression supports only " - "eviction_mode='union'; per-head keep sets are not defined " - "for draft layers, which are never scored" - ) - if any(window is not None for window in draft_manager.max_attention_window_vec) or any( - not isinstance(layer, AttentionLayerConfig) or layer.sliding_window_size is not None - for layer in draft_manager.kv_cache_manager_py_config.layers - ): - raise ValueError( - "TriAttention draft KV co-compression requires full-attention " - "draft V2 lifecycles" - ) - if any(window is not None for window in manager.max_attention_window_vec) or any( - not isinstance(layer, AttentionLayerConfig) or layer.sliding_window_size is not None - for layer in manager.kv_cache_manager_py_config.layers - ): - raise ValueError( - "TriAttention requires full-attention V2 lifecycles; native SWA, " - "VSWA, and SSM pools are not supported" - ) - def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Compact after native KV-cache updates have finalized this iteration (must run after KVCacheManagerV2 so capacity reflects the written token and any rewind).""" diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 74ba92dba01d..38cfeb38f594 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -39,6 +39,7 @@ get_default_trtllm_modules_to_hf_modules) from tensorrt_llm.lora_manager import load_torch_lora from tensorrt_llm.mapping import CpType, Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig from ..attention_backend import get_sparse_attn_kv_cache_manager from ..hostfunc import set_low_latency_dispatch @@ -2150,6 +2151,22 @@ def validate_kv_cache_compression_with_spec( "TriAttention speculative compatibility requires a separate " "draft KV cache; shared target/draft pools cannot be " "compacted safely") + if not draft_kv_cache_manager.is_draft: + raise ValueError( + "TriAttention speculative compatibility requires the actual " + "separate draft KV cache manager") + if config.eviction_mode != "union": + raise ValueError( + "TriAttention draft KV co-compression supports only " + "eviction_mode='union'; draft layers are never scored") + if any(window is not None for window in + draft_kv_cache_manager.max_attention_window_vec) or any( + not isinstance(layer, AttentionLayerConfig) + or layer.sliding_window_size is not None for layer in + draft_kv_cache_manager.kv_cache_manager_py_config.layers): + raise ValueError( + "TriAttention draft KV co-compression requires " + "full-attention draft V2 lifecycles") def create_kv_cache_compression_manager( diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 8de107cbfd87..d7ac86fc8818 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -305,12 +305,9 @@ def make_fake_v2(enable_block_reuse=False, *, is_draft=False): fake_v2 = KVCacheManagerV2.__new__(KVCacheManagerV2) fake_v2.enable_block_reuse = enable_block_reuse - fake_v2.enable_swa_scratch_reuse = False fake_v2.is_draft = is_draft fake_v2.kv_compression_manages_history = False fake_v2.kv_factor = 2 - fake_v2.mapping = SimpleNamespace(enable_attention_dp=False) - fake_v2.is_disagg = False fake_v2.max_beam_width = 1 fake_v2.max_batch_size = 8 fake_v2.num_extra_kv_tokens = 0 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 index ff84e4067004..253be33246d9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -238,64 +238,36 @@ class Boom(RuntimeError): @pytest.mark.parametrize( "gate,match", [ - # One representative per guard family (the per-mode/per-config - # variants raise through the same checks). + # One representative per call-site guard family (the per-mode/per-config + # variants raise through the same checks). kv_factor geometry needs no + # admission gate: the native compact op TORCH_CHECKs every pool's K/V + # plane count at the first compact. + ("callsite_dflash", "standard paged cache compacted together"), ("union_only_per_head", "union"), - ("draft_kv_factor", "standard key/value cache"), - # Same check family on the TARGET cache (MLA SELFKONLY, kv_factor 1). - ("target_kv_factor", "standard key/value KV cache"), ("full_attention_draft", "full-attention draft"), - ("callsite_dflash", "standard paged cache compacted together"), ], ) def test_draft_admission_gates_raise(gate, match): + # Draft/spec admission is owned by the executor call-site gate: rejected + # before any compression manager exists. + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec + from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig, MTPDecodingConfig + draft_manager = _make_fake_v2(is_draft=True) if gate == "full_attention_draft": draft_manager.max_attention_window_vec = [128] - if gate.startswith("callsite_"): - # Call-site speculative gate: rejected before any manager exists. - from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec - from tensorrt_llm.llmapi.llm_args import ( - DFlashDecodingConfig, - TriAttentionKvCacheCompressionConfig, - ) - - spec_config = DFlashDecodingConfig(max_draft_len=3) - with pytest.raises(ValueError, match=match): - validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", budget=8 - ), - spec_config, - draft_manager, - ) - return - - def construct(): - return TriAttention( - _make_fake_v2(), - _make_tri_config( - budget=8, - eviction_mode="per_head" if gate == "union_only_per_head" else "union", - ), - draft_kv_cache_manager=None if gate == "target_kv_factor" else draft_manager, - ) - - if gate in ("union_only_per_head", "full_attention_draft"): - # Manager-lifetime capability gates run once, at construction. - with pytest.raises(ValueError, match=match): - construct() - return - manager = construct() - # Flipping kv_factor after construction exercises TriAttention's own - # capability gate directly. - if gate == "draft_kv_factor": - draft_manager.kv_factor = 1 - if gate == "target_kv_factor": - manager.kv_cache_manager.kv_factor = 1 + spec_config = ( + DFlashDecodingConfig(max_draft_len=3) + if gate == "callsite_dflash" + else 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): - manager._validate_v2_compatibility() + validate_kv_cache_compression_with_spec(config, spec_config, draft_manager) def test_compressed_count_is_monotone_and_tracks_confirmed_length(): diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 729b88319b6d..6832c2e3a327 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -471,7 +471,6 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): draft_kv_cache_manager=draft_manager, ) - manager._validate_v2_compatibility() assert manager.kv_cache_manager.kv_compression_manages_history is True assert draft_manager.kv_compression_manages_history is True from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec From 216ffb14b70dc3912f1919dbfa9085cb3f155daf Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 09:56:30 -0700 Subject: [PATCH 134/178] [None][refactor] Resolve manager-lifetime layer facts at construction and merge the layout resolvers (knife 38) Global layers and the attention layer partition are manager-lifetime facts: resolve them once in the constructor and delete the three manual caches. The target and draft runtime KV layout twins collapse into one cached resolver with the pool page-count poll written once. In the buffer build, each cute handle owns its operand's DLPack capsule, so the keep-alive twin tuple dies; alignments move next to the operands they describe; the compile loops share one get-or-compile helper with verbatim cache keys and compile order. Signed-off-by: tianruih --- .../triattention/triattention.py | 252 ++++++++---------- .../_torch/kv_cache_compression/conftest.py | 27 +- .../test_triattention_draft_cocompaction.py | 8 +- .../test_triattention_pipeline.py | 13 +- 4 files changed, 139 insertions(+), 161 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index de25942140d0..280b2d2e748c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -171,19 +171,21 @@ def __init__( # Buffers build once at the first eviction as plain attributes on this # manager and stay resident for the manager's lifetime. self._buffers_built = False - self._local_to_global_layers_cache: Optional[List[int]] = None - self._attention_layer_partition_cache: Optional[ - Tuple[List[int], List[int], Optional[int]] - ] = None - self._runtime_kv_layout_cache: Optional[Dict[str, object]] = None - self._draft_runtime_kv_layout_cache: Optional[Dict[str, object]] = None + # 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._layer_partition = self._attention_layer_partition() + # Target/draft runtime KV layouts, cached by the one resolver. + self._kv_layout_caches: Dict[bool, Optional[Dict[str, object]]] = { + False: None, + True: None, + } def on_request_init(self, request: "LlmRequest", **kwargs) -> None: """Track the request and resolve the official calibration on first use.""" request_id = request.py_request_id if request_id not in self._request_states: self._validate_request_capacity(request) - self._attention_layer_partition() self._request_states[request_id] = { "generation_steps": 0, "evicted_tokens": 0, @@ -445,43 +447,25 @@ def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: # ---- helpers (eviction / scoring / V2 cache access / calibration) ---- - def _local_to_global_layers(self) -> List[int]: - cached = self._local_to_global_layers_cache - if cached is None: - cached = [int(layer) for layer in self.kv_cache_manager.pp_layers] - self._local_to_global_layers_cache = cached - return cached - @staticmethod def _has_sliding_window_signal(config: Dict[str, object]) -> bool: use_sliding_window = config.get("use_sliding_window") if isinstance(use_sliding_window, bool): return use_sliding_window - for field in ( - "sliding_window", - "sliding_window_size", - "sliding_window_pattern", - "max_window_layers", - ): - value = config.get(field) - if isinstance(value, bool): - if value: - return True - elif isinstance(value, (int, float)): - if value > 0: - return True - elif value: - return True - return False + return any( + config.get(field) + for field in ( + "sliding_window", + "sliding_window_size", + "sliding_window_pattern", + "max_window_layers", + ) + ) def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int]]: """SWA layers here are stored at full length; the window applies only in the kernel.""" - cached = self._attention_layer_partition_cache - if cached is not None: - return cached - model_path = self.model_path - global_layers = self._local_to_global_layers() + global_layers = self._global_layers num_layers = len(global_layers) try: @@ -490,7 +474,7 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int config = AutoConfig.from_pretrained( model_path, trust_remote_code=True, local_files_only=True ) - except Exception as exc: + except (OSError, ValueError) as exc: raise ValueError( f"TriAttention could not load the local model config from {model_path!r}" ) from exc @@ -502,9 +486,7 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int "Model config exposes sliding-window metadata but no layer_types; " "TriAttention cannot classify kernel-masked SWA layers safely" ) - result = (list(range(num_layers)), [], None) - self._attention_layer_partition_cache = result - return result + 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, " @@ -532,16 +514,15 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int f"the kernel-masked SWA window size {raw_window}" ) window_size = raw_window - result = (dense_layers, swa_layers, window_size) - self._attention_layer_partition_cache = result - return result + return (dense_layers, swa_layers, window_size) - def _runtime_kv_layout(self) -> Dict[str, object]: + def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: # The manager identity and layer count are manager-lifetime owner # contracts; only the pool page counts are polled (stale-pointer - # safety until V2 exposes a layout epoch). - manager = self.kv_cache_manager - cached = self._runtime_kv_layout_cache + # safety until V2 exposes a layout epoch). Production draft callers + # gate on ``draft_kv_cache_manager is not None``. + manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager + cached = self._kv_layout_caches[draft] if cached is not None: current_page_counts = self._pool_page_counts( manager, @@ -550,24 +531,33 @@ def _runtime_kv_layout(self) -> Dict[str, object]: ) if current_page_counts != cached["pool_page_counts"]: raise RuntimeError( - "TriAttention V2 pool layout changed after the layout was built; " - "KV pool rebalance is not supported" + f"TriAttention {'draft ' if draft else ''}V2 pool layout changed " + "after the layout was built; KV pool rebalance is not supported" ) return cached - global_layers = self._local_to_global_layers() - dense_layers, swa_layers, swa_window = self._attention_layer_partition() - if not dense_layers: - raise ValueError("TriAttention requires at least one full-attention layer") + if draft: + global_layers = [int(layer) for layer in manager.pp_layers] + if not global_layers: + raise RuntimeError("TriAttention draft KV cache manager exposes no 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, swa_layers, swa_window = self._layer_partition + if not dense_layers: + raise ValueError("TriAttention requires at least one full-attention layer") layout = self._build_runtime_kv_layout( manager, global_layers, dense_layers=dense_layers, swa_layers=swa_layers, swa_window=swa_window, - what="", + what="draft " if draft else "", ) - self._runtime_kv_layout_cache = layout + self._kv_layout_caches[draft] = layout return layout def _build_runtime_kv_layout( @@ -617,37 +607,6 @@ def _build_runtime_kv_layout( ), ) - def _draft_runtime_kv_layout(self) -> Dict[str, object]: - # Production callers gate on ``draft_kv_cache_manager is not None``. - manager = self.draft_kv_cache_manager - cached = self._draft_runtime_kv_layout_cache - if cached is not None: - current_page_counts = self._pool_page_counts( - manager, - cached["global_layers"], - cached["pool_representatives"], - ) - if current_page_counts != cached["pool_page_counts"]: - raise RuntimeError( - "TriAttention draft V2 pool layout changed after the layout " - "was built; KV pool rebalance is not supported" - ) - return cached - - global_layers = [int(layer) for layer in manager.pp_layers] - if not global_layers: - raise RuntimeError("TriAttention draft KV cache manager exposes no layers") - layout = self._build_runtime_kv_layout( - manager, - global_layers, - dense_layers=list(range(len(global_layers))), - swa_layers=[], - swa_window=None, - what="draft ", - ) - self._draft_runtime_kv_layout_cache = layout - return layout - @staticmethod def _pool_page_counts( manager: KVCacheManagerV2, @@ -677,7 +636,7 @@ def _ensure_buffers( if self.draft_kv_cache_manager is not None: # The cached layout lookup enforces draft V2 pool page-count # stability every round, exactly like the target's lookup. - self._draft_runtime_kv_layout() + self._runtime_kv_layout(draft=True) if self._buffers_built: if ( needed_width <= self._decode_width @@ -709,7 +668,7 @@ def _ensure_buffers( if self.draft_kv_cache_manager is not None: draft_tail_capacity = self._draft_protected_tail_capacity draft = dict( - layout=self._draft_runtime_kv_layout(), + layout=self._runtime_kv_layout(draft=True), protected_tail_capacity=draft_tail_capacity, page_table_token_capacity=seq_capacity + draft_tail_capacity, ) @@ -929,15 +888,14 @@ def _build_buffers( self._gather_index_base = None self._gather_index = None if not union: - num_kv_heads_early = int(layer_pools[dense_layers[0]].shape[2]) self._gather_index_base = torch.empty( (max_requests, 1, 1, 1, decode_width), dtype=torch.int64, device=device ) self._gather_index = self._gather_index_base.expand( max_requests, len(dense_layers), - num_kv_heads_early, - num_q_heads // num_kv_heads_early, + self._num_kv_heads, + num_q_heads // self._num_kv_heads, decode_width, ) self._union_scores = None @@ -947,9 +905,7 @@ def _build_buffers( (max_requests, seq_len), dtype=torch.float32, device=device ) - # ---- THE score path (no fallback): per request-count/page-shard - # compiled variants over ctor-bound cute tensor handles; construction - # failures raise the kernels' own dtype/shape/TMA errors. + # ---- THE score path (no fallback): compiled per request-count/page-shard ---- anchor_pool = p0 sm_count = int(torch.cuda.get_device_properties(device).multi_processor_count) # One [stats_row, page_shard, {count, mean, m2}] record array. @@ -973,19 +929,21 @@ def _build_buffers( int(num_freqs), int(tokens_per_block), ) - torch_prefix = ( - block_offsets.view(-1), - seg_page_off, - seg_req_id, - seg_layer_id, - # Pointer capture of the staged metadata rows. - self._valid_seq_lens_device, - seg_out_offset, - self._token_starts_device, - q_real.view(-1), - q_imag.view(-1), - mlr_coef.view(-1), + # Alignment sits with the operand it describes; valid_seq_lens and + # token_starts are only 4-byte-aligned row views, read as per-CTA scalars. + prefix_operands = ( + (block_offsets.view(-1), 16), + (seg_page_off, 16), + (seg_req_id, 16), + (seg_layer_id, 16), + (self._valid_seq_lens_device, 4), + (seg_out_offset, 16), + (self._token_starts_device, 4), + (q_real.view(-1), 16), + (q_imag.view(-1), 16), + (mlr_coef.view(-1), 16), ) + torch_prefix = tuple(tensor for tensor, _ in prefix_operands) torch_tail = ( freq_scale_sq, self._score_scratch, @@ -993,14 +951,10 @@ def _build_buffers( anchor_pool, self._tma_descriptors, ) - # Keep-alive twin of the cute handles below: the compiled launches - # capture these raw device pointers for the buffers' lifetime. - self._score_torch_operands = (*torch_prefix, *torch_tail) - # valid_seq_lens/token_starts are only 4-byte-aligned row views, read as per-CTA scalars. - prefix_aligns = (16, 16, 16, 16, 4, 16, 4, 16, 16, 16) + # No keep-alive twin: each cute handle below owns its operand's DLPack + # capsule, which retains the underlying torch storage. self._cute_score_prefix = tuple( - _to_cute(tensor, assumed_align=align) - for tensor, align in zip(torch_prefix, prefix_aligns) + _to_cute(tensor, assumed_align=align) for tensor, align in prefix_operands ) self._cute_score_tail = ( _to_cute(freq_scale_sq), @@ -1058,6 +1012,16 @@ def _build_buffers( pool_shape=tuple(int(value) for value in anchor_pool.shape), pool_strides=tuple(int(value) for value in anchor_pool.stride()), ) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + + def _compiled_kernel(cache_key, build): + with _COMPILE_LOCK: + compiled = _COMPILED_KERNELS.get(cache_key) + if compiled is None: + compiled = build() + _COMPILED_KERNELS[cache_key] = compiled + return compiled + if union: variant_key = "triattention_cute_score_stats" compiled_entries = self._compiled_score_stats @@ -1072,26 +1036,22 @@ def _build_buffers( request_count, page_shards, ) - with _COMPILE_LOCK: - compiled = _COMPILED_KERNELS.get(cache_key) - if compiled is None: - kernel = _TriAttentionScoreKernel( + compiled_entries[request_count] = _compiled_kernel( + cache_key, + lambda page_shards=page_shards: cute.compile( + _TriAttentionScoreKernel( **kernel_kwargs, page_shards=page_shards, write_partial_stats=union, - ) - stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) - compiled = cute.compile( - kernel, - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled - compiled_entries[request_count] = compiled + ), + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + cutlass.Int32(1), + stream, + ), + ) page_shards_by_count[request_count] = page_shards if max_requests > 1: @@ -1131,11 +1091,14 @@ def _build_buffers( _tensor_spec(self._union_scores), _tensor_spec(self._partial_stats), ) - with _COMPILE_LOCK: - compiled_selection = _COMPILED_KERNELS.get(cache_key) - if compiled_selection is None: - tokens_per_lane, token_subtiles, row_cluster_ctas = config - kernel = _TriAttentionNormalizeUnionKernel( + tokens_per_lane, token_subtiles, row_cluster_ctas = config + compiled_selection = _compiled_kernel( + cache_key, + lambda page_shards=page_shards, + tokens_per_lane=tokens_per_lane, + token_subtiles=token_subtiles, + row_cluster_ctas=row_cluster_ctas: cute.compile( + _TriAttentionNormalizeUnionKernel( num_layers=self._num_layers, seq_len=seq_len, num_q_heads=num_q_heads, @@ -1145,19 +1108,14 @@ def _build_buffers( tokens_per_lane=tokens_per_lane, token_subtiles=token_subtiles, row_cluster_ctas=row_cluster_ctas, - ) - stream = cuda_driver.CUstream( - torch.cuda.current_stream(device).cuda_stream - ) - compiled_selection = cute.compile( - kernel, - _to_cute(self._partial_stats), - *self._cute_selection_prefix, - self._cute_union_scores, - cutlass.Int32(1), - stream, - ) - _COMPILED_KERNELS[cache_key] = compiled_selection + ), + self._cute_partial_stats, + *self._cute_selection_prefix, + self._cute_union_scores, + cutlass.Int32(1), + stream, + ), + ) compiled_configs[config_key] = compiled_selection self._compiled_normalize_union[request_count] = compiled_selection logger.info( diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index d7ac86fc8818..d60117dfe0e4 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -13,8 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +import os +import tempfile from contextlib import contextmanager from types import SimpleNamespace +from typing import Optional from unittest import mock import torch @@ -329,6 +333,27 @@ def make_fake_v2(enable_block_reuse=False, *, is_draft=False): return fake_v2 +_TEST_MODEL_DIR: Optional[str] = None + + +def make_test_model_dir() -> str: + """A real on-disk dense model config: construction-time layer partition + resolves through the production AutoConfig path, no mocks.""" + 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_tri_config(**overrides): """A real TriAttentionKvCacheCompressionConfig with test calibration inputs (the config validator requires both ``model_path`` and ``calibration_path``).""" @@ -336,7 +361,7 @@ def make_tri_config(**overrides): options = { "budget": 8, - "model_path": "/models/test", + "model_path": make_test_model_dir(), "calibration_path": "/calib/test.pt", } options.update(overrides) 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 index 253be33246d9..50616a525c9f 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -273,7 +273,7 @@ def test_draft_admission_gates_raise(gate, match): def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager = _make_triattention(budget=4, beta=4) manager._calibrated = True - manager._attention_layer_partition_cache = ([0, 1], [], None) + manager._layer_partition = ([0, 1], [], None) target = manager.kv_cache_manager target._stream = mock.Mock() target.pp_layers = [0, 1] @@ -358,7 +358,9 @@ def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): manager.draft_kv_cache_manager = draft_manager # Injected post-construction: mirror the ctor-cached manager-lifetime tail. manager._draft_protected_tail_capacity = 1 - manager._draft_runtime_kv_layout = mock.Mock( + # The one resolver serves both sides; only its draft arm runs here (the + # target layout arrives as the explicit _ensure_buffers argument). + manager._runtime_kv_layout = mock.Mock( return_value=dict( layer_pools=[], dense_layers=[], @@ -396,7 +398,7 @@ def apply_built(**kwargs): assert kwargs["capacities"]["bucket_seq_len"] == 1024 assert kwargs["capacities"]["page_table_token_capacity"] == 1024 + 1 assert kwargs["draft"]["page_table_token_capacity"] == 1024 + 1 - assert kwargs["draft"]["layout"] is manager._draft_runtime_kv_layout.return_value + assert kwargs["draft"]["layout"] is manager._runtime_kv_layout.return_value # Migrated from the pipeline buffer-kwargs test: the budget and the # pool keys thread through unchanged. assert kwargs["capacities"]["keep_count"] == manager.budget diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 6832c2e3a327..97f06d274a60 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -89,13 +89,7 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): # Calibration is deferred to the first request, so construction needs # no calibration file or CUDA. fake_v2 = _make_fake_v2(enable_block_reuse=False) - cfg = TriAttentionKvCacheCompressionConfig( - budget=32, - beta=16, - eviction_mode="per_head", - model_path="/models/test", - calibration_path="/calib/test.pt", - ) + cfg = _make_tri_config(budget=32, beta=16, eviction_mode="per_head") mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) assert mgr.budget == 32 @@ -129,7 +123,7 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): pool_representatives=(0, 1), pool_page_counts=(4, 8), ) - triattention._runtime_kv_layout_cache = cached + triattention._kv_layout_caches[False] = cached assert triattention._runtime_kv_layout() is cached manager.get_buffers.assert_not_called() @@ -155,7 +149,6 @@ def test_request_init_and_finish_lifecycle(self): manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 triattention = TriAttention(manager, _make_tri_config(budget=8)) - triattention._attention_layer_partition_cache = ([], [], None) triattention._calibrated = True triattention.on_request_init(_make_request(11)) @@ -814,7 +807,7 @@ def test_layer_partition_uses_local_config_and_validates_window(self, budget, fi mgr = _make_triattention() mgr.model_path = "/models/gpt-oss" mgr.budget = budget - mgr.kv_cache_manager = SimpleNamespace(pp_layers=[0, 1, 2, 3]) + mgr._global_layers = [0, 1, 2, 3] config = _make_hf_config( layer_types=[ "sliding_attention", From fe11a3d23aac47cac199816b4374cf49d85c7fed Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 10:19:18 -0700 Subject: [PATCH 135/178] [None][chore] Trim regrown narration comments to the banner bar (knife 39a) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 36 +++---------------- .../triattention/triattention.py | 20 +++-------- .../triattention/triattention_kernels.py | 10 +----- 3 files changed, 10 insertions(+), 56 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index b46964b88ce1..0346a9799488 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -127,29 +127,9 @@ def init_compaction_buffers( capacities: Dict[str, int], draft: Optional[Dict[str, object]] = None, ) -> Tuple[Dict[str, object], ...]: - """Agree on the decision rows and return opaque launch plans per geometry. - - Move sources must be increasing kept ordinals with - destination_bases[request] + move <= source[move] (C++ in-place copy - contract). ``target`` carries the resolved dense/SWA grouping inputs from - the runtime layout (``per_layer_sources`` selects 3-D per-layer move rows; - ``layer_pool_ids`` the canonical layer -> V2 pool id tuple that indexes - ``kv_block_offsets``; ``swa_destination_bases`` the caller-owned SWA - rebase row) plus the decision inputs :func:`compact` packs each round: - ``kept_ordinal_rows`` (``max_requests * decision_rows`` rows of - ``keep_count`` int32 kept ordinals, forwarded verbatim), the - per-request ``decision_rows`` count (1 = one shared row broadcast over - every KV head), and the staged per-request ``valid_seq_lens`` the - protected tail rides after. ``draft`` is one all-or-none resolved branch - (its dense-only moves broadcast the one shared decision row over the - draft's own KV heads); ``capacities`` the request/keep/tail capacity - numbers. - - Returns one plan per compacted cache family (target, then the optional - draft): plain contract fields -- decision inputs, move-index buffers, - pack geometry ints, and per-pool batched native-operand records built - here at init time -- that only :func:`compact` reads at launch. - """ + """Agree on the decision rows and return opaque launch plans per geometry: + one plan per compacted cache family (target, then the optional draft), + with init-built contract fields that only :func:`compact` reads at launch.""" layer_pools = target["layer_pools"] dense_layers = tuple(int(layer) for layer in target["dense_layers"]) swa_layers = tuple(int(layer) for layer in target["swa_layers"]) @@ -380,14 +360,8 @@ def compact( plans: Tuple[Dict[str, object], ...], request_count: int, ) -> None: - """Pack each plan's move sources and fire its native compacts, in plan order. - - Pure mover: the caller has already materialized its kept ordinals into the - agreed decision rows for the active ``request_count`` cohort, and the - caller owns the completion ordering of the whole round. Every launch - argument comes straight off the plan's init-built contract fields; the - round constructs no containers and runs no torch ops of its own. - """ + """Pack each plan's move sources and fire its native compacts, in plan order + (pure mover: the caller owns the decision rows and the round's completion ordering).""" for plan in plans: _pack_move_sources_kernel[(request_count, plan["decision_rows"])]( plan["kept_ordinal_rows"], diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 280b2d2e748c..131d056f94c0 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -135,10 +135,8 @@ def __init__( "TriAttention union eviction requires normalize_scores=True: " "the fused union pipeline always z-normalizes score rows" ) - # Hard-coded semantics: the prompt is always pinned and the budget - # counts decode tokens only (physical KV reclaim requires both). - # Calibration is the official TriAttention .pt; TRT-LLM does not - # compute calibration. The config boundary requires both paths. + # Prompt always pinned; budget counts decode tokens only. Calibration + # is the official TriAttention .pt (TRT-LLM never computes it). self.model_path = config.model_path self.calibration_path = config.calibration_path self.calibration: Optional[Dict[str, torch.Tensor]] = None @@ -719,13 +717,6 @@ def _build_buffers( """Build the round's buffers, compiled score launches, and compaction data in place as plain attributes on this manager (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put). - - ``layout`` is the runtime KV layout dict, passed whole; ``calibration`` - carries the local q_real/q_imag/mlr_coef [L, H, F] slices and - freq_scale_sq; ``draft`` is one all-or-none resolved branch (its layout - dict plus tail/page-table capacities); ``capacities`` the capacity - numbers. Static round policy (mode, ``normalize_scores``, the shared - ``self._phase`` table) reads off the manager itself. """ import cutlass import cutlass.cute as cute @@ -1351,11 +1342,8 @@ def _execute_eviction_round( self, prepared: Sequence[Dict[str, object]], ) -> None: - """Run one eviction round over the prepared cohort: stage the page-table - snapshots and round metadata, then score, select, settle, and compact, and - finally order the manager streams after this cohort's compact (every launch - covers the full request capacity; padded rows carry zero lengths and stay - inert).""" + """Run one eviction round over the prepared cohort (every launch covers + the full request capacity; padded rows carry zero lengths and stay inert).""" manager = self.kv_cache_manager draft_manager = self.draft_kv_cache_manager with nvtx_range_debug("triattention.page_table_stage", color="orange"): diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 6b9bd397e1dd..9526e0f84e82 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -11,9 +11,6 @@ # ---- Mean-phase gather: per-request phase-row fetch + width derivation ---- - -# Positions past this row count are not exactly representable in fp32. - # Score z-normalization epsilon; must stay a plain float (the CuTe DSL traces it). STD_EPSILON = 1e-6 @@ -179,12 +176,7 @@ def prepare_per_head_scores( per_layer: bool, normalize_scores: bool, ) -> None: - """Normalize and reduce score rows for either per-head eviction mode. - - ``selection_scores_rows``/``selection_row_lengths`` are the canonical - row-major selection buffers over the full request capacity - (``capacity * selection_rows`` rows); ``valid_widths`` also spans the - capacity, so the per-request row count derives from the shapes.""" + """Normalize and reduce score rows for either per-head eviction mode.""" request_count, num_layers, num_q_heads, width = scores.shape selection_rows = int(selection_scores_rows.shape[0]) // int(valid_widths.shape[0]) num_kv_heads = selection_rows // num_layers if per_layer else selection_rows From 2c5ef3722ba84ad4b4fc3417dbb3e236ba1f1ae0 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 10:22:22 -0700 Subject: [PATCH 136/178] [None][refactor] One source of truth for calibration/copy state; merge the resize twins (knife 39b) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/interface.py | 5 +- .../triattention/triattention.py | 66 +++++++------------ .../_torch/kv_cache_compression/conftest.py | 1 - .../test_triattention_draft_cocompaction.py | 5 +- .../test_triattention_pipeline.py | 8 +-- 5 files changed, 31 insertions(+), 54 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index 06890572f0fe..8f22dd7f9b31 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 from enum import IntEnum, auto -from typing import Optional class KvCacheCompressionMode(IntEnum): @@ -21,9 +20,7 @@ def is_eviction_method(self): return self == KvCacheCompressionMode.TRIATTENTION @staticmethod - def from_string(name: Optional[str]) -> "KvCacheCompressionMode": - if name is None: - return KvCacheCompressionMode.NONE + def from_string(name: str) -> "KvCacheCompressionMode": try: return KvCacheCompressionMode[name.upper()] except KeyError: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 131d056f94c0..db1cda45326d 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -140,7 +140,6 @@ def __init__( self.model_path = config.model_path self.calibration_path = config.calibration_path self.calibration: Optional[Dict[str, torch.Tensor]] = None - self._calibrated = False self._freq_scale_sq: Optional[torch.Tensor] = None # Mean-phase table dict; buffer builds bind its device tables in place. @@ -243,7 +242,7 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: ) def _ensure_calibrated(self) -> None: - if self._calibrated: + if self.calibration is not None: return self.calibration = self._resolve_calibration() self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) @@ -254,7 +253,6 @@ def _ensure_calibrated(self) -> None: self._triattn_mlr_coef = ( self.calibration["E_q_norm"].to(torch.float32) - _Eq.abs().to(torch.float32) ).contiguous() - self._calibrated = True def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Compact after native KV-cache updates have finalized this iteration @@ -332,13 +330,8 @@ def _periodic_evict( continue draft_kv_cache = None if self.draft_kv_cache_manager is not None: - draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map.get(request_id) - if draft_kv_cache is None: - # A missing draft cache is a wiring/lifecycle bug. - raise RuntimeError( - "TriAttention cannot co-compress a missing draft KV " - f"cache for request {request_id}" - ) + # A missing draft cache is a wiring bug: the dict's KeyError is the report. + draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] if not draft_kv_cache.is_active: # Target and draft defer together (pre-launch). continue @@ -372,38 +365,30 @@ def _resize_compacted_requests(self, prepared) -> None: return with nvtx_range("triattention.resize", color="red"): with nvtx_range_debug("triattention.v2_resize", color="red"): - for item in prepared: - kv_cache = item["kv_cache"] - if not kv_cache.is_active: - # Bytes already moved: skipping the ledger resize here - # would leave silent corruption. The compact-to-resize - # window is owned by this hook; a suspension inside it - # breaks the lifecycle contract. - raise RuntimeError( - f"Request {item['request_id']} target KV cache was " - "suspended between compact and resize" - ) - resized_capacity = item["expected_keep_count"] + item["protected_tail"] - if not kv_cache.resize(resized_capacity, None): - raise RuntimeError( - f"Failed to resize compacted KV cache for request " - f"{item['request_id']} to {resized_capacity} tokens" - ) + families = [("target", "kv_cache", None)] if self.draft_kv_cache_manager is not None: - # Same kept set: the draft shrinks to the same retained length plus its own tail. - draft_protected_tail = self._draft_protected_tail_capacity + # Same kept set: the draft shrinks to the same retained + # length plus its own fixed tail. + families.append( + ("draft", "draft_kv_cache", self._draft_protected_tail_capacity) + ) + for label, cache_key, fixed_tail in families: for item in prepared: - draft_kv_cache = item["draft_kv_cache"] - if not draft_kv_cache.is_active: + kv_cache = item[cache_key] + if not kv_cache.is_active: + # Bytes already moved: skipping the ledger resize would + # leave silent corruption; the compact-to-resize window + # is owned by this hook. raise RuntimeError( - f"Request {item['request_id']} draft KV cache was " + f"Request {item['request_id']} {label} KV cache was " "suspended between compact and resize" ) - draft_capacity = item["expected_keep_count"] + draft_protected_tail - if not draft_kv_cache.resize(draft_capacity, None): + tail = item["protected_tail"] if fixed_tail is None else fixed_tail + resized_capacity = item["expected_keep_count"] + tail + if not kv_cache.resize(resized_capacity, None): raise RuntimeError( - "Failed to resize co-compressed draft KV cache for " - f"request {item['request_id']} to {draft_capacity} tokens" + f"Failed to resize compacted {label} KV cache for " + f"request {item['request_id']} to {resized_capacity} tokens" ) def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: @@ -1188,7 +1173,7 @@ def _compiled_kernel(cache_key, build): dense_move_offsets=draft_move_offsets_row, protected_tail_capacity=int(draft["protected_tail_capacity"]), ) - self.compaction_plan = init_compaction_buffers( + self._compaction_plan = init_compaction_buffers( target=dict( layer_pools=layer_pools, dense_layers=list(dense_layers), @@ -1224,7 +1209,6 @@ def _compiled_kernel(cache_key, build): self._block_offsets_ready_event = torch.cuda.Event() # This cohort's compact is done: manager may resize/reuse pages. self._compaction_done_event = torch.cuda.Event() - self._copy_pending = False @staticmethod def _page_table_pool_ids( @@ -1374,8 +1358,7 @@ def _execute_eviction_round( # The one host-staging reuse fence: the previous cohort's async copies # must complete before the pinned metadata rows AND the pinned # target/draft block-offset snapshots are rewritten. - if self._copy_pending and not self._staging_reuse_event.query(): - self._staging_reuse_event.synchronize() + self._staging_reuse_event.synchronize() host_table = self._request_metadata_host_np for row, values in rows: if values is not None: @@ -1401,7 +1384,6 @@ def _execute_eviction_round( finally: # Guards the pinned staging until the asynchronous copies complete. self._staging_reuse_event.record(stream) - self._copy_pending = True request_count = len(prepared) union = self.eviction_mode == "union" try: @@ -1508,7 +1490,7 @@ def _execute_eviction_round( ) self._settle_top_tokens(request_count) with nvtx_range("triattention.compact", color="purple"): - compact(self.compaction_plan, request_count) + compact(self._compaction_plan, request_count) finally: # Order V2 page-table reuse and resize after this cohort's compact. self._compaction_done_event.record(stream) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index d60117dfe0e4..b3e68410d0b3 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -247,7 +247,6 @@ def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): staging._block_offsets_ready_event = torch.cuda.Event() staging._compaction_done_event = torch.cuda.Event() staging._staging_reuse_event = torch.cuda.Event() - staging._copy_pending = False staging._block_offsets_host = torch.empty( 1, max_requests, 2, staged_blocks_per_seq, dtype=torch.int32, device="cpu", pin_memory=True ) 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 index 50616a525c9f..d626f46a9dc6 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -177,9 +177,8 @@ def test_execute_eviction_round_orders_both_manager_streams(): tri._keep_count = 4 tri.eviction_mode = "union" tri._swa_window = None - tri.compaction_plan = () + tri._compaction_plan = () tri._draft_protected_tail_capacity = 1 - tri._copy_pending = False tri._staging_reuse_event = mock.Mock() tri._compaction_done_event = event tri._request_metadata_host = host @@ -272,7 +271,7 @@ def test_draft_admission_gates_raise(gate, match): def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager = _make_triattention(budget=4, beta=4) - manager._calibrated = True + manager.calibration = {} manager._layer_partition = ([0, 1], [], None) target = manager.kv_cache_manager target._stream = mock.Mock() diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 97f06d274a60..996a22eb011f 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -149,7 +149,7 @@ def test_request_init_and_finish_lifecycle(self): manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 triattention = TriAttention(manager, _make_tri_config(budget=8)) - triattention._calibrated = True + triattention.calibration = {} triattention.on_request_init(_make_request(11)) triattention.on_request_init(_make_request(12)) @@ -287,7 +287,7 @@ def test_identity_selection_is_filtered_before_launch(self): capacity=6, history_length=0, is_active=True, resize=mock.Mock(return_value=True) ) manager.kv_cache_manager.kv_cache_map = {7: cache} - manager._calibrated = True + manager.calibration = {} state = _set_request_state(manager, 7, generation_steps=127) with _mocked_eviction_internals(manager) as internals: @@ -329,7 +329,7 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft fake_v2.num_extra_kv_tokens = num_extra_kv_tokens fake_v2._kv_reserve_draft_tokens = kv_reserve_draft_tokens mgr = TriAttention(fake_v2, _make_tri_config(budget=8)) - mgr._calibrated = True + mgr.calibration = {} cache = SimpleNamespace( capacity=seq_len, history_length=1024, @@ -428,7 +428,7 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): # (capacity minus the protected tail), never the logical length. physical_confirmed = 6100 manager = _make_triattention(beta=128) - manager._calibrated = True + manager.calibration = {} _set_request_state(manager, 7, generation_steps=127, evicted_tokens=100) cache = SimpleNamespace( capacity=physical_confirmed, From f5e9c6249b080934d72c77633e0798bf1cd50613 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 10:24:55 -0700 Subject: [PATCH 137/178] [None][refactor] Buffer build reads the layout's own maps; direct build arguments (knife 39c) Signed-off-by: tianruih --- .../triattention/triattention.py | 94 +++++++++---------- .../_torch/kv_cache_compression/conftest.py | 26 +++-- .../test_triattention_draft_cocompaction.py | 10 +- 3 files changed, 59 insertions(+), 71 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index db1cda45326d..c2c3a9bb1904 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -560,8 +560,13 @@ def _build_runtime_kv_layout( ] raise RuntimeError(f"Missing {what}KV pools for attention layers {missing}") layer_pools = [pool for pool in maybe_layer_pools if pool is not None] - # Canonical pool IDs, resolved once; every grouping derives from them. - layer_pool_ids = self._page_table_pool_ids(manager, global_layers) + # Canonical pool IDs, resolved once; every grouping derives from them + # (V2 owns the mapping; its own 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 + ) all_storage_groups: Dict[int, List[int]] = {} for layer, pool_id in enumerate(layer_pool_ids): all_storage_groups.setdefault(pool_id, []).append(layer) @@ -673,20 +678,16 @@ def _ensure_buffers( q_real, q_imag, mlr_coef = self._local_score_calibration(layout["global_layers"]) self._build_buffers( layout=layout, - calibration=dict( - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=self._freq_scale_sq, - ), - capacities=dict( - max_requests=request_capacity, - bucket_seq_len=seq_capacity, - decode_width=decode_width, - page_table_token_capacity=page_table_token_capacity, - keep_count=self.budget, - protected_tail_capacity=tail_capacity, - ), + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=self._freq_scale_sq, + max_requests=request_capacity, + bucket_seq_len=seq_capacity, + decode_width=decode_width, + page_table_token_capacity=page_table_token_capacity, + keep_count=self.budget, + protected_tail_capacity=tail_capacity, draft=draft, ) self._buffers_built = True @@ -695,8 +696,16 @@ def _build_buffers( self, *, layout: Dict[str, object], - calibration: Dict[str, torch.Tensor], - capacities: Dict[str, int], + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + freq_scale_sq: torch.Tensor, + max_requests: int, + bucket_seq_len: int, + decode_width: int, + page_table_token_capacity: int, + keep_count: int, + protected_tail_capacity: int, draft: Optional[Dict[str, object]] = None, ) -> None: """Build the round's buffers, compiled score launches, and compaction data @@ -725,24 +734,21 @@ def _build_buffers( layer_group_representative = layout["layer_group_representative"] # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. layer_pool_ids = tuple(layout["layer_pool_ids"]) - dense_groups = list(layout["storage_groups"].values()) - page_representatives = [group[0] for group in dense_groups] - page_representatives.extend( - layer for layer in swa_layers if layer not in page_representatives - ) num_page_table_slots = int(layout["manager"].num_pools) - device = layer_pools[page_representatives[0]].device - max_requests = int(capacities["max_requests"]) - seq_len = int(capacities["bucket_seq_len"]) - page_table_token_capacity = int(capacities["page_table_token_capacity"]) - decode_width = int(capacities["decode_width"]) - keep_count = int(capacities["keep_count"]) - protected_tail_capacity = int(capacities["protected_tail_capacity"]) + # The first dense layer anchors device and staging geometry. + p0 = layer_pools[dense_layers[0]] + device = p0.device + max_requests = int(max_requests) + seq_len = int(bucket_seq_len) + page_table_token_capacity = int(page_table_token_capacity) + decode_width = int(decode_width) + keep_count = int(keep_count) + protected_tail_capacity = int(protected_tail_capacity) q_real, q_imag, mlr_coef, freq_scale_sq = ( - calibration[key].to(device=device, dtype=torch.float32).contiguous() - for key in ("q_real", "q_imag", "mlr_coef", "freq_scale_sq") + tensor.to(device=device, dtype=torch.float32).contiguous() + for tensor in (q_real, q_imag, mlr_coef, freq_scale_sq) ) num_q_heads = int(q_real.shape[1]) num_freqs = int(q_real.shape[2]) @@ -755,7 +761,7 @@ def _build_buffers( # ---- block-offset staging (target, plus the co-compressed draft) ------- self._block_offsets_host, self._block_offsets_device = _allocate_block_offset_staging( - layer_pools[page_representatives[0]], + p0, num_pools=num_page_table_slots, max_requests=max_requests, token_capacity=page_table_token_capacity, @@ -815,15 +821,15 @@ def _build_buffers( self._phase_f_block = triton.next_power_of_2(self._phase_num_freqs) # ---- score state: one fused group across all dense layers -------------- - p0 = layer_pools[dense_layers[0]] - _, kv_factor, num_kv_heads, tokens_per_block, head_dim = p0.shape + _, _, num_kv_heads, tokens_per_block, _ = p0.shape self._num_layers = len(dense_layers) self._num_q_heads = int(num_q_heads) self._num_kv_heads = int(num_kv_heads) self._num_freqs = int(num_freqs) self._tokens_per_block = int(tokens_per_block) - _rep_of = {layer: layers[0] for layers in dense_groups for layer in layers} - dense_layer_slots = [layer_pool_ids[_rep_of[layer]] for layer in dense_layers] + dense_layer_slots = [ + layer_pool_ids[layer_group_representative[layer]] for layer in dense_layers + ] seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( self._num_layers ) @@ -1210,20 +1216,6 @@ def _compiled_kernel(cache_key, build): # This cohort's compact is done: manager may resize/reuse pages. self._compaction_done_event = torch.cuda.Event() - @staticmethod - def _page_table_pool_ids( - manager: KVCacheManagerV2, - global_layers: List[int], - ) -> Tuple[int, ...]: - """Canonical local layer -> V2 pool id tuple (the staged plane slots). - - V2 owns the mapping; its own lookup errors are the precise ones.""" - layer_offsets = manager.layer_offsets - layer_to_pool = manager.layer_to_pool_mapping_dict - return tuple( - int(layer_to_pool[layer_offsets[global_layer]]) for global_layer in global_layers - ) - def _evict_requests( self, prepared: List[Dict[str, object]], diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index b3e68410d0b3..9fd02c0a520a 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -561,22 +561,18 @@ def make_cute_buffers( manager._phase = make_phase_table(offsets, omega, seq_len) manager._build_buffers( layout=layout, - calibration=dict( - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=freq_scale_sq, - ), - capacities=dict( - max_requests=max_requests, - bucket_seq_len=seq_len, - decode_width=decode_width, - page_table_token_capacity=( - seq_len if page_table_token_capacity is None else page_table_token_capacity - ), - keep_count=keep_count, - protected_tail_capacity=protected_tail_capacity, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=freq_scale_sq, + max_requests=max_requests, + bucket_seq_len=seq_len, + decode_width=decode_width, + page_table_token_capacity=( + seq_len if page_table_token_capacity is None else page_table_token_capacity ), + keep_count=keep_count, + protected_tail_capacity=protected_tail_capacity, ) return manager 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 index d626f46a9dc6..e1350bff6203 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -392,15 +392,15 @@ def apply_built(**kwargs): # and thread through unchanged (no longer build arguments). assert manager.eviction_mode == "union" assert manager._phase is phase - assert kwargs["capacities"]["max_requests"] == 8 - assert kwargs["capacities"]["decode_width"] == 4 + 2 * 128 - assert kwargs["capacities"]["bucket_seq_len"] == 1024 - assert kwargs["capacities"]["page_table_token_capacity"] == 1024 + 1 + assert kwargs["max_requests"] == 8 + assert kwargs["decode_width"] == 4 + 2 * 128 + assert kwargs["bucket_seq_len"] == 1024 + assert kwargs["page_table_token_capacity"] == 1024 + 1 assert kwargs["draft"]["page_table_token_capacity"] == 1024 + 1 assert kwargs["draft"]["layout"] is manager._runtime_kv_layout.return_value # Migrated from the pipeline buffer-kwargs test: the budget and the # pool keys thread through unchanged. - assert kwargs["capacities"]["keep_count"] == manager.budget + assert kwargs["keep_count"] == manager.budget assert kwargs["layout"] is layout assert list(kwargs["layout"]["layer_pool_ids"]) == list(layout["layer_pool_ids"]) From c5efed4dd5d624b763c99a7ce5bbd3f2c8d0d488 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 10:25:35 -0700 Subject: [PATCH 138/178] [None][refactor] Compaction grouping drops its anachronisms; one owner for the union-only draft rule (knife 39d) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 0346a9799488..76f673a22213 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -24,7 +24,6 @@ are private to it. """ -from collections import OrderedDict from typing import Dict, List, Optional, Tuple import torch @@ -226,11 +225,8 @@ def init_compaction_buffers( ) draft_move_indices = None if draft is not None: - if decision_rows != 1: - raise ValueError( - "draft packing broadcasts one shared decision row per request; " - f"got {decision_rows} decision rows" - ) + # Union-only drafts (one shared decision row) are enforced by the + # executor call-site gate; the plan builder just packs. draft_layer_pools = draft["layer_pools"] draft_dense_layers = tuple(int(layer) for layer in draft["dense_layers"]) draft_layer_pool_ids = tuple(int(pool_id) for pool_id in draft["layer_pool_ids"]) @@ -278,7 +274,7 @@ def init_compaction_buffers( slots, is_draft, ) in families: - grouped = OrderedDict() + grouped = {} for layer, pool, page_table in entries: key = ( pool_ids[layer], @@ -291,7 +287,6 @@ def init_compaction_buffers( for group_entries in grouped.values(): layers = tuple(entry[0] for entry in group_entries) pools = list(entry[1] for entry in group_entries) - page_tables = tuple(entry[2] for entry in group_entries) source_layer_indices = None if slots is not None: source_layer_indices = torch.tensor( @@ -306,7 +301,7 @@ def init_compaction_buffers( dtype=torch.int64, device=device, ), - page_table=page_tables[0], + page_table=group_entries[0][2], move_indices=move_indices, move_offsets=move_offsets, destination_bases=destination_bases, From 7539fcdb697d6c0b2dd8baa0e5786a294287ba2f Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 10:26:52 -0700 Subject: [PATCH 139/178] [None][refactor] Drop the dead epilogue parameter and host alias twins (knife 39e) Signed-off-by: tianruih --- .../triattention_cute_score_fused.py | 63 +++++++++---------- 1 file changed, 31 insertions(+), 32 deletions(-) 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 index 687701e49bb2..d6fa3dc3b7ca 100644 --- 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 @@ -109,9 +109,7 @@ def __init__( # 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.pages_per_tile = self.fragments_per_phase - self.tile_tokens = CTA_M - self.max_tiles = (seq_len + self.tile_tokens - 1) // self.tile_tokens + self.max_tiles = (seq_len + CTA_M - 1) // CTA_M # Producer staging constants baked into the generated code. self.prefetch_depth = 4 @@ -124,7 +122,6 @@ def __init__( self.raw_page_buffers = RAW_PAGE_BUFFERS if write_partial_stats else 1 self.accumulator_pipeline_stages = 1 self.producer_warp_id = 0 - self.physical_threads = THREADS if pool_dim != 2 * num_freqs: raise ValueError("K pool shape does not match the CuTe score specialization") @@ -140,7 +137,6 @@ def epilog_tmem_copy_and_partition( accumulator: cute.Tensor, output: cute.Tensor, epilogue_tile: cute.Tile, - use_2cta_instrs: bool, ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: copy_atom = sm100_utils.get_tmem_load_op( self.cta_tile_shape_mnk, @@ -148,7 +144,7 @@ def epilog_tmem_copy_and_partition( self.c_dtype, self.acc_dtype, epilogue_tile, - use_2cta_instrs, + False, ) accumulator_epilogue = cute.flat_divide( accumulator[((None, None), 0, 0)], @@ -420,7 +416,7 @@ class SharedStorage: magnitude_fp16_b_smem_layout, ).launch( grid=(num_ctas, 1, 1), - block=(self.physical_threads, 1, 1), + block=(THREADS, 1, 1), stream=stream, ) @@ -587,8 +583,8 @@ def kernel( swizzle=raw_bf16_b_smem_layout.inner, ) - tile_index = score_start // self.tile_tokens + page_shard - tile_start_token = tile_index * self.tile_tokens + 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): @@ -605,10 +601,14 @@ def kernel( 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.pages_per_tile,), cutlass.Int32 + (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_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) - prefetched_page_fragments = cute.make_rmem_tensor((self.pages_per_tile,), cutlass.Int32) physical_fragments_arg = physical_page_fragments prefetched_fragments_arg = prefetched_page_fragments shard_has_page = valid_seq_len > score_start and tile_start_token < valid_seq_len @@ -618,18 +618,20 @@ def kernel( if lane_idx == 0: # Staged entries encode physical_page * kv_factor; decode to the pool page. producer_prefetched_page_id_lane0 = ( - cutlass.Int32(page_ids[page_off + tile_index * self.pages_per_tile]) + cutlass.Int32(page_ids[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.pages_per_tile): + 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 < valid_seq_len: fragment_page_id = ( cutlass.Int32( page_ids[ - page_off + tile_index * self.pages_per_tile + fragment + page_off + + tile_index * self.fragments_per_phase + + fragment ] ) // K_PLANES_PER_POOL_PAGE @@ -765,7 +767,7 @@ def kernel( 0, ) if cutlass.const_expr(self.fragments_per_phase > 1): - for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + 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, @@ -804,7 +806,7 @@ def kernel( 0, ) if cutlass.const_expr(self.fragments_per_phase > 1): - for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + 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, @@ -848,7 +850,7 @@ def kernel( 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 + self.tile_tokens * self.page_shards + next_tile_start_token = tile_start_token + CTA_M * self.page_shards next_pages_processed = tiles_processed + 1 if ( next_tile_start_token < valid_seq_len @@ -857,13 +859,14 @@ def kernel( next_page_id_lane0 = ( cutlass.Int32( page_ids[ - page_off + (tile_index + self.page_shards) * self.pages_per_tile + 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.pages_per_tile): + 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 ( @@ -875,7 +878,7 @@ def kernel( page_ids[ page_off + (tile_index + self.page_shards) - * self.pages_per_tile + * self.fragments_per_phase + fragment ] ) @@ -907,7 +910,7 @@ def kernel( 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 + self.tile_tokens * self.page_shards + 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 < valid_seq_len @@ -918,7 +921,7 @@ def kernel( 0, ) if cutlass.const_expr(self.fragments_per_phase > 1): - for fragment in cutlass.range_constexpr(1, self.pages_per_tile): + 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, @@ -1119,7 +1122,7 @@ def kernel( tCgC = thr_mma.partition_C(gC_mnl) epilogue_tidx = tidx % EPILOGUE_THREADS tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epilogue_tidx, tCtAcc, tCgC, self.epi_tile, False + epilogue_tidx, tCtAcc, tCgC, self.epi_tile ) simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile @@ -1195,7 +1198,7 @@ def kernel( stats_square_sums_m128[stats_head] + stats_delta * stats_delta ) tile_index += self.page_shards - tile_start_token += self.tile_tokens * 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) @@ -1223,17 +1226,13 @@ def kernel( stats_scratch_base = STATS_ORIGIN_SLOTS + (stats_warp * N + 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 * self.tile_tokens + 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 - self.tile_tokens * self.page_shards - ) - stats_invalid_tail = ( - stats_last_tile_start_token + self.tile_tokens - valid_seq_len - ) + stats_last_tile_start_token = tile_start_token - CTA_M * self.page_shards + stats_invalid_tail = stats_last_tile_start_token + CTA_M - valid_seq_len if cutlass.dynamic_expr(stats_invalid_tail > 0): stats_count_i32 = stats_count_i32 - stats_invalid_tail stats_count = cutlass.Float32(stats_count_i32) From 4877392f4cc8643d8895d4b0850361dabfc2c3e9 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 10:27:54 -0700 Subject: [PATCH 140/178] [None][test] Trim the draft-contract acceptance test to its live assertions (knife 39f) Signed-off-by: tianruih --- .../test_triattention_pipeline.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 996a22eb011f..9cadcc612225 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -44,7 +44,6 @@ # in pyexecutor._util (next to _create_kv_cache_manager), matching #15106. from tensorrt_llm._torch.pyexecutor._util import create_kv_cache_compression_manager from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role -from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig # The SM100 CuTe kernel is the only score path, so every test that actually # launches scores (or builds the real staging buffers, whose constructor @@ -454,25 +453,21 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): cache.resize.assert_not_called() def test_one_model_draft_co_compression_contract_is_accepted(self): - # Draft physical length tracks the target's: smaller draft - # max_seq_len is accepted and both managers are marked. + # Construction accepts the separate draft manager, and the executor + # call-site gate accepts the one-model MTP roundtrip (base-ctor + # marking is asserted in the executor manager tests). draft_manager = _make_fake_v2(is_draft=True) - draft_manager.max_seq_len = 8192 - manager = TriAttention( + TriAttention( _make_fake_v2(), _make_tri_config(budget=8), draft_kv_cache_manager=draft_manager, ) - assert manager.kv_cache_manager.kv_compression_manages_history is True - assert draft_manager.kv_compression_manages_history is True from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig validate_kv_cache_compression_with_spec( - TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt", budget=8 - ), + _make_tri_config(budget=8), MTPDecodingConfig(max_draft_len=1), draft_manager, ) From aa9e0e82c2b97cb39bfdc91146dd85469fb742b3 Mon Sep 17 00:00:00 2001 From: tianruih Date: Thu, 23 Jul 2026 20:28:35 -0700 Subject: [PATCH 141/178] [None][chore] Trim comments to plain functional descriptions (knife 40) Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 21 ++--- .../triattention/triattention.py | 84 +++++-------------- .../triattention/triattention_kernels.py | 6 +- 3 files changed, 29 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 76f673a22213..54f2acd60fac 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,13 +15,9 @@ """Batched physical KV-cache compaction: an algorithm-neutral mover. -``init_compaction_buffers`` agrees on the decision rows (kept ordinals; -move offsets ride the caller's staged rows) once per geometry and returns -opaque launch plans. The caller materializes its keep decision into the -agreed rows each round, then ``compact`` loops the plans: each packs its -family's move sources and fires its native launches. This module knows -cache-family geometry and the decision format only; the plans' internals -are private to it. +``init_compaction_buffers`` returns opaque launch plans once per geometry; +each round the caller writes its keep decision into the agreed rows and +``compact`` packs every family's move sources and fires its native launches. """ from typing import Dict, List, Optional, Tuple @@ -49,10 +45,9 @@ def _pack_move_sources_kernel( SWA_WINDOW: tl.constexpr, BLOCK: tl.constexpr = 256, ): - """Pack one decision row into one family's move sources (increasing - kept ordinals; C++ in-place copy contract): dense rows forward the row - content verbatim for the first KEEP_COUNT moves, then append the - protected tail; SWA rows write latest-window ordinals once per KV head.""" + """Pack one decision row into one family's move sources: dense rows emit the + kept tokens then the protected tail; SWA rows emit the latest window + (ascending order: the native copy moves in place).""" BROADCAST: tl.constexpr = DECISION_ROWS == 1 HAS_SWA: tl.constexpr = SWA_TOTAL > 0 request = tl.program_id(0) @@ -188,9 +183,7 @@ def init_compaction_buffers( if swa_move_indices is not None: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) - # Families: (entries, pool ids, move indices, move offsets, destination - # bases, per-layer slots, is_draft). One grouping loop below batches each - # family into per-pool native-operand records (all init-time construction). + # One move-descriptor tuple per cache family (dense / SWA / draft). families = [ ( dense_entries, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index c2c3a9bb1904..bb60905b4dae 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -124,8 +124,6 @@ def __init__( draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, ): super().__init__(kv_cache_manager, draft_kv_cache_manager) - # budget/beta positivity and the eviction_mode literal are validated at - # the config boundary (TriAttentionKvCacheCompressionConfig). self.budget = config.budget self.beta = config.beta self.eviction_mode = config.eviction_mode @@ -135,8 +133,7 @@ def __init__( "TriAttention union eviction requires normalize_scores=True: " "the fused union pipeline always z-normalizes score rows" ) - # Prompt always pinned; budget counts decode tokens only. Calibration - # is the official TriAttention .pt (TRT-LLM never computes it). + # Prompt always pinned; budget counts decode tokens only. self.model_path = config.model_path self.calibration_path = config.calibration_path self.calibration: Optional[Dict[str, torch.Tensor]] = None @@ -145,13 +142,12 @@ def __init__( # Mean-phase table dict; buffer builds bind its device tables in place. self._phase: Optional[Dict[str, object]] = None - # Per-request {generation_steps, evicted_tokens}. + # Per-request eviction progress. self._request_states: Dict[int, Dict[str, object]] = {} # In-flight overlap batch reference; membership resolves lazily. self._prepared_generation_batch: Optional[object] = None self._prepared_generation_ids: Optional[set] = None - # Manager-lifetime constants (V2 fixes every input at construction): - # protected tails are num_extra + reserved draft width + 1 sampled token. + # Manager-lifetime constants. self._protected_tail_capacity = ( int(kv_cache_manager.num_extra_kv_tokens) + int(kv_cache_manager._kv_reserve_draft_tokens) @@ -342,7 +338,7 @@ def _periodic_evict( "kv_cache": kv_cache, "draft_kv_cache": draft_kv_cache, "seq_len": int(seq_len), - # Uncompressed logical position (prefix + evicted). + # Uncompressed logical position. "round_start": int(seq_len + request_state["evicted_tokens"]), "prompt_len": min(int(request.py_prompt_len), int(seq_len)), "expected_keep_count": expected_keep_count, @@ -376,9 +372,7 @@ def _resize_compacted_requests(self, prepared) -> None: for item in prepared: kv_cache = item[cache_key] if not kv_cache.is_active: - # Bytes already moved: skipping the ledger resize would - # leave silent corruption; the compact-to-resize window - # is owned by this hook. + # Bytes already moved: the compact-to-resize window is owned by this hook. raise RuntimeError( f"Request {item['request_id']} {label} KV cache was " "suspended between compact and resize" @@ -500,10 +494,7 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int return (dense_layers, swa_layers, window_size) def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: - # The manager identity and layer count are manager-lifetime owner - # contracts; only the pool page counts are polled (stale-pointer - # safety until V2 exposes a layout epoch). Production draft callers - # gate on ``draft_kv_cache_manager is not None``. + # Only pool page counts are polled; manager identity and layer count are manager-lifetime contracts. manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager cached = self._kv_layout_caches[draft] if cached is not None: @@ -663,8 +654,7 @@ def _ensure_buffers( first_pool = layout["layer_pools"][layout["dense_layers"][0]] if self._phase is None: - # Upstream geometric offsets [1, 2, 4, ... <= max]: the table - # builder consumes them as host floats only (no device copy). + # Host-only width offsets for the table builder (no device copy). self._phase = { "omega": self.calibration["omega"] .to(device=first_pool.device, dtype=torch.float32) @@ -708,10 +698,8 @@ def _build_buffers( protected_tail_capacity: int, draft: Optional[Dict[str, object]] = None, ) -> None: - """Build the round's buffers, compiled score launches, and compaction data - in place as plain attributes on this manager - (the compiled kernels capture raw pool addresses: scored pools must stay alive and stay put). - """ + """Build the round's buffers, compiled launches, and compaction data as attributes + (compiled kernels capture raw pool addresses: pools must stay alive and stay put).""" import cutlass import cutlass.cute as cute @@ -803,7 +791,7 @@ def _build_buffers( ) self._round_starts_device = self._request_metadata_device[0, :max_requests] self._valid_seq_lens_device = self._request_metadata_device[1, :max_requests] - # Per-request pinned prompt lengths (per-request decode window starts). + # Pinned per-request decode-window starts. self._token_starts_device = self._request_metadata_device[2, :max_requests] dense_move_offsets_row = self._request_metadata_device[3] swa_move_offsets_row = self._request_metadata_device[4] @@ -863,10 +851,8 @@ def _build_buffers( self._gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( 1, 1, 1, 1, -1 ) - # Compile the mode's SM100 CuTe entries; no other score path, no fallback. union = self.eviction_mode == "union" - # Persistent gather index (per-head modes): per round only the - # token-start base is re-added in place; the expanded view is fixed. + # Per-head gather index: per round only the token-start base is re-added in place. self._gather_index_base = None self._gather_index = None if not union: @@ -887,10 +873,10 @@ def _build_buffers( (max_requests, seq_len), dtype=torch.float32, device=device ) - # ---- THE score path (no fallback): compiled per request-count/page-shard ---- + # ---- score path: compiled per request-count/page-shard ---- anchor_pool = p0 sm_count = int(torch.cuda.get_device_properties(device).multi_processor_count) - # One [stats_row, page_shard, {count, mean, m2}] record array. + # Per-shard partial score statistics. partial_stats_elements = ( max_requests * self._num_layers @@ -945,9 +931,7 @@ def _build_buffers( _to_cute(anchor_pool), _to_cute(self._tma_descriptors, assumed_align=128), ) - # Ctor-bound persistent launch operands (one from_dlpack wrap each): - # the round path refreshes their contents in place and launches with - # the active cohort size only. + # Ctor-bound persistent launch operands: refreshed in place each round. self._cute_mean_cos = _to_cute(self._mean_cos.view(-1)) self._cute_mean_sin = _to_cute(self._mean_sin.view(-1)) self._cute_union_scores = _to_cute(self._union_scores.view(-1)) if union else None @@ -1120,7 +1104,7 @@ def _compiled_kernel(cache_key, build): self._provisional_rows = torch.zeros( (max_requests, keep_count), dtype=torch.int32, device=device ) - # Kept decode ordinals only (prompt-length independent rows). + # Kept decode ordinals. self._kept_ordinal_rows = torch.empty( (max_requests, keep_count), dtype=torch.int32, device=device ) @@ -1347,9 +1331,7 @@ def _execute_eviction_round( and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF ): raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # The one host-staging reuse fence: the previous cohort's async copies - # must complete before the pinned metadata rows AND the pinned - # target/draft block-offset snapshots are rewritten. + # Host-staging reuse fence: prior cohort's async copies must finish before the pinned rows are rewritten. self._staging_reuse_event.synchronize() host_table = self._request_metadata_host_np for row, values in rows: @@ -1399,9 +1381,7 @@ def _execute_eviction_round( num_warps=1, ) if union: - # THE union path: the fused score+stats entry, then the - # normalized union reduction, straight off the ctor-bound - # cute handles (only the active cohort size varies). + # Union path: fused score+stats, then the normalized union reduction. cu_stream = cuda_driver.CUstream(stream.cuda_stream) self._compiled_score_stats[request_count]( *self._cute_score_prefix, @@ -1493,16 +1473,7 @@ def _execute_eviction_round( # ---- helpers: calibration loading ---- def _resolve_calibration(self) -> Dict[str, torch.Tensor]: - """Load the user-supplied calibration .pt and return our runtime schema. - - TriAttention does NOT compute calibration -- the user calibrates with the - official tool (github.com/WeianMao/triattention) and passes that file via - ``calibration_path``; we only run inference. Both the official R-KV layout - (``{metadata, stats{"layerLL_headHH": {q_mean_real, q_mean_imag, - q_abs_mean}}}``) and our already-converted flat layout are accepted -- the - official one is converted here. Calibration resolves lazily on the - first request (``on_request_init``), not at manager construction, and - stays on CPU: runtime construction moves it to the pool device once.""" + """Load the calibration file, converting the official layout if needed.""" raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): return raw @@ -1516,13 +1487,7 @@ def _resolve_calibration(self) -> Dict[str, torch.Tensor]: ) def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: - """Convert the official per-(layer, head) stats to our flat runtime schema. - - ``E_q[l,h] = q_mean_real + i*q_mean_imag`` and ``E_q_norm[l,h] = - q_abs_mean`` are the same statistic, just restacked into ``[L, H, F]``. - ``omega`` / ``freq_scale_sq`` are not in the official file (its runtime - recomputes them from the model rotary), so we derive them from the model - config -- model-intrinsic and corpus-independent.""" + """Convert the official calibration format to the runtime schema.""" stats = raw["stats"] meta = raw.get("metadata", {}) if "sampled_heads" in meta: @@ -1542,7 +1507,6 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: E_q[layer, h] = torch.complex(s["q_mean_real"].float(), s["q_mean_imag"].float()) E_q_norm[layer, h] = s["q_abs_mean"].float() omega, freq_scale_sq = self._rope_tables(freq_count) - # CPU schema: runtime construction moves tensors to the pool device once. calib = { "E_q": E_q, "E_q_norm": E_q_norm, @@ -1556,15 +1520,7 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: return calib def _rope_tables(self, freq_count: int): - """RoPE ``omega`` (inv_freq) + ``freq_scale_sq`` (squared position-0 - amplitude) from the model config -- model-intrinsic, corpus-independent - (the official file does not store them). Reads both config generations: - transformers>=5.5 ``rope_parameters`` (rope_theta folded inside) and the - legacy top-level ``rope_scaling``/``rope_theta``. Plain RoPE uses the - standard formula with the resolved theta (attention_factor 1); scaled - variants (yarn, llama3, ...) go through transformers' rope-init so their - attention_factor is honored. The analytic fallback survives ONLY for - ImportError (rope-init module absent); every other failure raises.""" + """Derive the RoPE frequency tables from the model config.""" import transformers from transformers import AutoConfig diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 9526e0f84e82..2bd4fa15bdb2 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -222,10 +222,8 @@ def _settle_ties_kernel( SELECTION_ROWS: tl.constexpr, BLOCK: tl.constexpr = 256, ): - """Settle one selection row's ties into its kept-ordinal output row - (threshold recovery with sentinel-skip, strictly-greater count, - lowest-index tie quota, ascending prompt-rebased emission; entries past - the emitted count keep their previous value).""" + """Settle one selection row's score ties into its final kept-token row + (deterministic lowest-index tie break, ascending output).""" request = tl.program_id(0) selection_domain = tl.program_id(1) row = request * SELECTION_ROWS + selection_domain From be747b628ab8a37caddf807c9ac43909db9e6bf9 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 00:48:41 -0700 Subject: [PATCH 142/178] [None][chore] Drop the block-offset gather check wall and dead debug declaration Signed-off-by: tianruih --- .../batch_manager/kvCacheManagerV2Utils.cpp | 26 +++---------------- .../kernels/unfusedAttentionKernels.h | 17 +++--------- .../thop/sparseKvCacheCompactOp.cpp | 17 +++--------- .../test_disagg_index_mapper_early_release.py | 9 ++----- 4 files changed, 12 insertions(+), 57 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp index 1f735fa112e1..e7e0946a2bf3 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cpp @@ -19,8 +19,8 @@ #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/memoryUtils.h" #include -#include #include +#include #include #include #include @@ -221,36 +221,16 @@ at::Tensor IndexMapper::getCopyIndex( void IndexMapper::gatherKBlockOffsets(at::Tensor const& source, at::Tensor destination, std::vector const& requestIds, SizeType32 numBlocks) { - TLLM_CHECK_WITH_INFO(source.device().is_cpu() && destination.device().is_cpu(), - "Block-offset gather requires CPU tensors"); - TLLM_CHECK_WITH_INFO(source.scalar_type() == at::kInt && destination.scalar_type() == at::kInt, - "Block-offset gather requires int32 tensors"); - TLLM_CHECK_WITH_INFO(source.dim() == 4 && destination.dim() == 4 && source.size(2) == 2 - && destination.size(2) == 2, - "Block-offset gather requires [pool, sequence, K/V, block] tensors"); - TLLM_CHECK_WITH_INFO(source.is_contiguous() && destination.is_contiguous(), - "Block-offset gather requires contiguous tensors"); - TLLM_CHECK_WITH_INFO(source.storage().data_ptr().get() != destination.storage().data_ptr().get(), - "Block-offset gather requires distinct source and destination storage"); - TLLM_CHECK_WITH_INFO(source.size(0) == destination.size(0), "Block-offset gather pool counts must match"); - TLLM_CHECK_WITH_INFO(!requestIds.empty(), "Block-offset gather requires at least one request"); - TLLM_CHECK_WITH_INFO(numBlocks > 0 && numBlocks <= source.size(3) && numBlocks <= destination.size(3), - "Block-offset gather block count exceeds the source or destination capacity"); - TLLM_CHECK_WITH_INFO(static_cast(requestIds.size()) <= destination.size(1), - "Block-offset gather request count exceeds the destination capacity"); - - auto const sourceRows = source.size(1); std::vector sourceRowsByRequest; sourceRowsByRequest.reserve(requestIds.size()); for (auto const requestId : requestIds) { - auto const sourceRow = static_cast(getIndex(requestId)) * maxBeamWidth_; - TLLM_CHECK_WITH_INFO(sourceRow < sourceRows, "IndexMapper slot exceeds the source tensor capacity"); - sourceRowsByRequest.push_back(sourceRow); + 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); diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 06d58bdbe9cf..baa4ede44217 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -413,15 +413,9 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams template void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams params, cudaStream_t stream); -//! Adapt a uniform group of KVCacheManagerV2 layer pools to the batched -//! compaction kernels (double-buffered cp.async pipelined bf16 copies, one -//! CTA per layer/KV-head/request). Device pointer arrays allow one layered -//! launch. Within each request and head, source ordinals must increase -//! strictly and satisfy destinationBases[request] + move <= source[move]; -//! the per-request bases let one launch cover a cohort with mixed -//! pinned-prompt lengths. sourceHeadStride is the head-row stride of -//! sparseKvIndices: the index buffers may be wider than one round's total -//! move count. +//! Compact a uniform group of KVCacheManagerV2 layer pools in one batched launch +//! (per request and head, moves are ascending and never overtake their sources: +//! the copy runs in place). template void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const* pageTable, int32_t numLayers, int64_t pageTableRequestStride, int32_t const* sparseKvIndices, int32_t const* sourceLayerIndices, @@ -429,11 +423,6 @@ void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const int32_t const* destinationBases, int32_t batchSize, int32_t numKvHeads, int32_t tokensPerBlock, int32_t headDim, cudaStream_t stream); -// Debug function to test basic parameter access -template -void invokeDebugSparseKvCacheParams( - QKVPreprocessingParams params, int* debug_output, cudaStream_t stream); - template void invokeKvCachePostprocessing(QKVPreprocessingParams params, cudaStream_t stream) { diff --git a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp index b63845156c9d..b0446b0ef574 100644 --- a/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp +++ b/cpp/tensorrt_llm/thop/sparseKvCacheCompactOp.cpp @@ -32,15 +32,9 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -//! Adapt one uniform group of KVCacheManagerV2 HND layer pools to the -//! batched compaction kernels: dedicated double-buffered cp.async pipelined -//! bf16 kernels, one CTA per (layer, KV head, request), addressed through a -//! flat V2 K-plane block-offset table shared by all layers. Per-request -//! destinationBases replace the former arbitrary destination tensor because -//! every compaction move targets one contiguous interval per request. -//! Within each request and KV head, the caller must supply increasing -//! source ordinals with destinationBases[request] + move <= source[move], -//! which makes the forward tiled in-place copy safe. +//! Compact one uniform group of KVCacheManagerV2 HND layer pools in one batched +//! launch (per request and KV head, moves are ascending and never overtake their +//! sources: the copy runs in place). void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor const& poolPointers, th::Tensor const& pageTable, th::Tensor const& sourceIndices, th::Tensor const& sourceOffsets, th::Tensor const& destinationBases, std::optional const& sourceLayerIndices) @@ -115,10 +109,7 @@ void sparseKvCacheCompactLayers(std::vector const& pools, th::Tensor sourceLayerPtr = layerIndices.data_ptr(); } - // source_offsets carve each request's move range out of source_indices; - // the values live on device and the kernel trusts them. The index buffer - // may be wider than one round's total move count, so the head-row - // stride passed below comes from the tensor shape, not from the offsets. + // source_offsets carve each request's move range; device-resident, the kernel trusts them. TORCH_CHECK(sourceOffsets.is_cuda() && sourceOffsets.get_device() == device && sourceOffsets.scalar_type() == th::kInt32 && sourceOffsets.is_contiguous() && sourceOffsets.dim() == 1 && sourceOffsets.size(0) == batchSize + 1, 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 fdf4ecfb75b1..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 @@ -217,7 +217,7 @@ def test_gather_k_block_offsets_uses_request_order_and_beam_zero(self): assert torch.count_nonzero(destination[:, :, 1] != -1) == 0 assert torch.count_nonzero(destination[:, 2, 0] != -1) == 0 - def test_gather_k_block_offsets_rejects_invalid_request_or_block_count(self): + def test_gather_k_block_offsets_rejects_unknown_request(self): from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( IndexMapper, ) @@ -225,16 +225,11 @@ def test_gather_k_block_offsets_rejects_invalid_request_or_block_count(self): 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.zeros((1, 2, 2, 3), dtype=torch.int32) + destination = torch.full((1, 2, 2, 3), -1, dtype=torch.int32) - destination.fill_(-1) 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 - with pytest.raises(Exception, match="block count"): - index_mapper.gather_k_block_offsets(source, destination, [11], 4) - with pytest.raises(Exception, match="distinct source and destination storage"): - index_mapper.gather_k_block_offsets(source, source, [11], 4) 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 ( From ae0e190f7feedf9839a205223e30d5dab0635b59 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 00:48:42 -0700 Subject: [PATCH 143/178] [None][refactor] Order TriAttention by lifecycle and slim the compaction plan API Signed-off-by: tianruih --- examples/triattention/README.md | 7 +- .../_torch/kv_cache_compression/compaction.py | 321 ++-- .../triattention/triattention.py | 1617 ++++++++--------- tensorrt_llm/_torch/pyexecutor/_util.py | 18 +- .../_torch/kv_cache_compression/conftest.py | 38 +- .../test_triattention_draft_cocompaction.py | 2 +- .../test_triattention_pipeline.py | 17 +- 7 files changed, 961 insertions(+), 1059 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index afc75c48fd7d..5c2adc7e25c7 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -13,12 +13,11 @@ TriAttention runs entirely in the generation phase and reuses the standard dense 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 it and converts it to its runtime schema at the first request. 2. **Periodic eviction (Stage during generation).** Every `beta` confirmed generation tokens, once a sequence is over budget, TriAttention scores the whole cache, selects `budget` tokens to keep (the prompt tokens are preserved on top of the budget), and physically compacts the KV cache down to the kept 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`. The scoring and compaction kernels are implemented in **Triton**. +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 -* GPU Compute Capability >= 9.0 (Hopper or newer) -* FP16 / BF16 +* GPU Compute Capability >= 10.0 (Blackwell or newer) * Paged KV Cache (`KVCacheManagerV2`) * Tensor Parallel * PyTorch backend @@ -121,6 +120,6 @@ trtllm-eval --model --config config.yaml longbench_v2 --max_outp * `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 requires `True` (the fused union pipeline always z-normalizes; construction rejects `False`). +* **`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/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 54f2acd60fac..192a9fe8aa05 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,18 +15,35 @@ """Batched physical KV-cache compaction: an algorithm-neutral mover. -``init_compaction_buffers`` returns opaque launch plans once per geometry; +``build_compaction_plans`` returns opaque launch plans once per geometry; each round the caller writes its keep decision into the agreed rows and -``compact`` packs every family's move sources and fires its native launches. +``compact`` packs every plan's move sources and fires its native launches. """ -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Tuple, TypedDict import torch import triton import triton.language as tl +class CompactionLaunch(TypedDict): + pools: List[torch.Tensor] + pool_pointers: torch.Tensor + page_table: torch.Tensor + move_indices: torch.Tensor + move_offsets: torch.Tensor + destination_bases: torch.Tensor + source_layer_indices: Optional[torch.Tensor] + + +class CompactionPlan(TypedDict): + decision_rows: int + pack_args: Tuple[Optional[torch.Tensor], ...] + pack_constexprs: Dict[str, object] + launches: Tuple[CompactionLaunch, ...] + + @triton.jit def _pack_move_sources_kernel( kept_ordinal_rows, @@ -104,54 +121,55 @@ def _pack_move_sources_kernel( ) -def _make_move_indices( - index_prefix: Tuple[int, ...], - moves_per_request: int, - max_requests: int, - device: torch.device, -) -> torch.Tensor: - return torch.empty( - (*index_prefix, moves_per_request * max_requests), dtype=torch.int32, device=device - ) - - -def init_compaction_buffers( +def build_compaction_plans( *, target: Dict[str, object], - capacities: Dict[str, int], draft: Optional[Dict[str, object]] = None, -) -> Tuple[Dict[str, object], ...]: - """Agree on the decision rows and return opaque launch plans per geometry: - one plan per compacted cache family (target, then the optional draft), - with init-built contract fields that only :func:`compact` reads at launch.""" - layer_pools = target["layer_pools"] - dense_layers = tuple(int(layer) for layer in target["dense_layers"]) - swa_layers = tuple(int(layer) for layer in target["swa_layers"]) - layer_pool_ids = tuple(int(pool_id) for pool_id in target["layer_pool_ids"]) - kv_block_offsets = target["kv_block_offsets"] - layer_group_representative = target["layer_group_representative"] - token_starts = target["token_starts"] - dense_move_offsets = target["dense_move_offsets"] - swa_move_offsets = target["swa_move_offsets"] - swa_window = target["swa_window"] - per_layer_sources = bool(target["per_layer_sources"]) - kept_ordinal_rows = target["kept_ordinal_rows"] - decision_rows = int(target["decision_rows"]) - valid_seq_lens = target["valid_seq_lens"] +) -> Tuple[CompactionPlan, ...]: + """One plan per compacted cache family (target, then the optional draft); + only :func:`compact` reads the plans' contents at launch.""" + plans = [_build_plan(target)] + if draft is not None: + # The draft compacts under the target's decision rows and token starts. + plans.append( + _build_plan( + dict( + draft, + token_starts=target["token_starts"], + kept_ordinal_rows=target["kept_ordinal_rows"], + valid_seq_lens=target["valid_seq_lens"], + ) + ) + ) + return tuple(plans) + - device = layer_pools[dense_layers[0]].device - max_requests = int(capacities["max_requests"]) - keep_count = int(capacities["keep_count"]) - protected_tail_capacity = int(capacities["protected_tail_capacity"]) +def _build_plan(cache: Dict[str, object]) -> CompactionPlan: + layer_pools = cache["layer_pools"] + dense_layers = tuple(int(layer) for layer in cache["dense_layers"]) + swa_layers = tuple(int(layer) for layer in cache.get("swa_layers", ())) + layer_pool_ids = tuple(int(pool_id) for pool_id in cache["layer_pool_ids"]) + kv_block_offsets = cache["kv_block_offsets"] + layer_group_representative = cache["layer_group_representative"] + token_starts = cache["token_starts"] + dense_move_offsets = cache["dense_move_offsets"] + per_layer_sources = bool(cache.get("per_layer_sources", False)) + kept_ordinal_rows = cache["kept_ordinal_rows"] + valid_seq_lens = cache["valid_seq_lens"] + protected_tail_capacity = int(cache["protected_tail_capacity"]) + first_pool = layer_pools[dense_layers[0]] + device = first_pool.device + max_requests = int(valid_seq_lens.shape[0]) + keep_count = int(kept_ordinal_rows.shape[1]) + decision_rows = int(kept_ordinal_rows.shape[0]) // max_requests # Pool shape [pages, K/V, heads, tokens, dim]. - num_kv_heads = int(layer_pools[dense_layers[0]].shape[2]) + num_kv_heads = int(first_pool.shape[2]) dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer_sources else (num_kv_heads,) - dense_move_indices = _make_move_indices( - dense_index_prefix, - keep_count + protected_tail_capacity, - max_requests, - device, + dense_move_indices = torch.empty( + (*dense_index_prefix, (keep_count + protected_tail_capacity) * max_requests), + dtype=torch.int32, + device=device, ) dense_entries = [ ( @@ -161,41 +179,25 @@ def init_compaction_buffers( ) for layer in dense_layers ] - - swa_move_indices = None - if not swa_layers: - swa_move_offsets = None - swa_window = 0 - else: - swa_window = int(swa_window) - swa_move_indices = _make_move_indices( - (num_kv_heads,), - swa_window + protected_tail_capacity, - max_requests, - device, - ) - dense_slots = ( {layer: slot for slot, layer in enumerate(dense_layers)} if per_layer_sources else None ) - # Widest per-request move count any staged offsets may express. - move_capacity = keep_count + protected_tail_capacity - if swa_move_indices is not None: - move_capacity = max(move_capacity, swa_window + protected_tail_capacity) - # One move-descriptor tuple per cache family (dense / SWA / draft). - families = [ - ( - dense_entries, - layer_pool_ids, - dense_move_indices, - dense_move_offsets, - token_starts, - dense_slots, - False, - ), + swa_move_indices = None + swa_move_offsets = None + swa_window = 0 + # One move group per family axis (dense / SWA): the layers + the tensors driving their moves. + move_groups = [ + (dense_entries, dense_move_indices, dense_move_offsets, token_starts, dense_slots), ] if swa_layers: + swa_move_offsets = cache["swa_move_offsets"] + swa_window = int(cache["swa_window"]) + swa_move_indices = torch.empty( + (num_kv_heads, (swa_window + protected_tail_capacity) * max_requests), + dtype=torch.int32, + device=device, + ) # SWA layers stage against their own page-table slots. swa_entries = [ ( @@ -205,72 +207,21 @@ def init_compaction_buffers( ) for layer in swa_layers ] - families.append( - ( - swa_entries, - layer_pool_ids, - swa_move_indices, - swa_move_offsets, - target["swa_destination_bases"], - None, - False, - ) - ) - draft_move_indices = None - if draft is not None: - # Union-only drafts (one shared decision row) are enforced by the - # executor call-site gate; the plan builder just packs. - draft_layer_pools = draft["layer_pools"] - draft_dense_layers = tuple(int(layer) for layer in draft["dense_layers"]) - draft_layer_pool_ids = tuple(int(pool_id) for pool_id in draft["layer_pool_ids"]) - draft_tail = int(draft["protected_tail_capacity"]) - # Own launch groups: the draft may use a different KV-head count. - draft_num_kv_heads = int(draft_layer_pools[draft_dense_layers[0]].shape[2]) - draft_move_indices = _make_move_indices( - (draft_num_kv_heads,), - keep_count + draft_tail, - max_requests, - device, - ) - draft_entries = [ - ( - layer, - draft_layer_pools[layer], - draft["kv_block_offsets"][ - draft_layer_pool_ids[draft["layer_group_representative"][layer]], - :max_requests, - 0, - ], - ) - for layer in draft_dense_layers - ] - families.append( - ( - draft_entries, - draft_layer_pool_ids, - draft_move_indices, - draft["dense_move_offsets"], - token_starts, - None, - True, - ) + move_groups.append( + (swa_entries, swa_move_indices, swa_move_offsets, cache["swa_destination_bases"], None) ) - target_launches: List[Dict[str, object]] = [] - draft_launch_records: List[Dict[str, object]] = [] - for ( - entries, - pool_ids, - move_indices, - move_offsets, - destination_bases, - slots, - is_draft, - ) in families: + # Widest per-request move count any staged offsets may express. + move_capacity = keep_count + protected_tail_capacity + if swa_layers: + move_capacity = max(move_capacity, swa_window + protected_tail_capacity) + + launches: List[CompactionLaunch] = [] + for entries, move_indices, move_offsets, destination_bases, slots in move_groups: grouped = {} for layer, pool, page_table in entries: key = ( - pool_ids[layer], + layer_pool_ids[layer], str(pool.dtype), str(pool.device), tuple(int(value) for value in pool.shape[1:]), @@ -287,85 +238,55 @@ def init_compaction_buffers( dtype=torch.int32, device=device, ) - record = dict( - pools=pools, - pool_pointers=torch.tensor( - [pool.data_ptr() for pool in pools], - dtype=torch.int64, - device=device, - ), - page_table=group_entries[0][2], - move_indices=move_indices, - move_offsets=move_offsets, - destination_bases=destination_bases, - source_layer_indices=source_layer_indices, + launches.append( + CompactionLaunch( + pools=pools, + pool_pointers=torch.tensor( + [pool.data_ptr() for pool in pools], + dtype=torch.int64, + device=device, + ), + page_table=group_entries[0][2], + move_indices=move_indices, + move_offsets=move_offsets, + destination_bases=destination_bases, + source_layer_indices=source_layer_indices, + ) ) - (draft_launch_records if is_draft else target_launches).append(record) - draft_launches = tuple(draft_launch_records) - target_plan = dict( - kept_ordinal_rows=kept_ordinal_rows, - valid_seq_lens=valid_seq_lens, + return CompactionPlan( decision_rows=decision_rows, - per_layer_sources=per_layer_sources, - dense_move_offsets=dense_move_offsets, - dense_move_indices=dense_move_indices, - swa_move_offsets=swa_move_offsets, - swa_move_indices=swa_move_indices, - swa_window=swa_window, - keep_count=keep_count, - move_capacity=move_capacity, - num_kv_heads=num_kv_heads, - dense_total=int(dense_move_indices.shape[-1]), - swa_total=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, - launches=tuple(target_launches), + pack_args=( + kept_ordinal_rows, + valid_seq_lens, + dense_move_offsets, + dense_move_indices, + swa_move_offsets, + swa_move_indices, + ), + pack_constexprs=dict( + KEEP_COUNT=keep_count, + DECISION_ROWS=decision_rows, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + PER_LAYER=per_layer_sources, + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, + SWA_WINDOW=swa_window, + ), + launches=tuple(launches), ) - plans = [target_plan] - if draft is not None: - draft_plan = dict( - kept_ordinal_rows=kept_ordinal_rows, - valid_seq_lens=valid_seq_lens, - decision_rows=1, - per_layer_sources=False, - dense_move_offsets=draft["dense_move_offsets"], - dense_move_indices=draft_move_indices, - swa_move_offsets=None, - swa_move_indices=None, - swa_window=0, - keep_count=keep_count, - move_capacity=keep_count + draft_tail, - num_kv_heads=draft_num_kv_heads, - dense_total=int(draft_move_indices.shape[-1]), - swa_total=0, - launches=draft_launches, - ) - plans.append(draft_plan) - return tuple(plans) - def compact( - plans: Tuple[Dict[str, object], ...], + plans: Tuple[CompactionPlan, ...], request_count: int, ) -> None: """Pack each plan's move sources and fire its native compacts, in plan order (pure mover: the caller owns the decision rows and the round's completion ordering).""" for plan in plans: _pack_move_sources_kernel[(request_count, plan["decision_rows"])]( - plan["kept_ordinal_rows"], - plan["valid_seq_lens"], - plan["dense_move_offsets"], - plan["dense_move_indices"], - plan["swa_move_offsets"], - plan["swa_move_indices"], - KEEP_COUNT=plan["keep_count"], - DECISION_ROWS=plan["decision_rows"], - MOVE_CAPACITY=plan["move_capacity"], - NUM_KV_HEADS=plan["num_kv_heads"], - PER_LAYER=plan["per_layer_sources"], - DENSE_TOTAL=plan["dense_total"], - SWA_TOTAL=plan["swa_total"], - SWA_WINDOW=plan["swa_window"], + *plan["pack_args"], **plan["pack_constexprs"] ) for launch in plan["launches"]: torch.ops.trtllm.sparse_kv_cache_compact_layers( diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index bb60905b4dae..17acf4a78374 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -38,7 +38,7 @@ ) from tensorrt_llm.logger import logger -from ..compaction import compact, init_compaction_buffers +from ..compaction import build_compaction_plans, compact from .triattention_kernels import ( _gather_mean_phase_kernel, _settle_ties_kernel, @@ -117,10 +117,12 @@ class TriAttention(KVCacheCompressionManager): adjusts_generation_kv_length = True + # ---- construction ---- + def __init__( self, - kv_cache_manager: KVCacheManagerV2, config: "TriAttentionKvCacheCompressionConfig", + kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, ): super().__init__(kv_cache_manager, draft_kv_cache_manager) @@ -129,15 +131,15 @@ def __init__( self.eviction_mode = config.eviction_mode self.normalize_scores = bool(config.normalize_scores) if self.eviction_mode == "union" and not self.normalize_scores: - raise ValueError( - "TriAttention union eviction requires normalize_scores=True: " - "the fused union pipeline always z-normalizes score rows" + logger.warning( + "TriAttention union eviction always z-normalizes scores; " + "forcing normalize_scores=True" ) + self.normalize_scores = True # Prompt always pinned; budget counts decode tokens only. self.model_path = config.model_path self.calibration_path = config.calibration_path - self.calibration: Optional[Dict[str, torch.Tensor]] = None - self._freq_scale_sq: Optional[torch.Tensor] = None + self._load_calibration() # Mean-phase table dict; buffer builds bind its device tables in place. self._phase: Optional[Dict[str, object]] = None @@ -174,72 +176,76 @@ def __init__( True: None, } - def on_request_init(self, request: "LlmRequest", **kwargs) -> None: - """Track the request and resolve the official calibration on first use.""" - request_id = request.py_request_id - if request_id not in self._request_states: - self._validate_request_capacity(request) - self._request_states[request_id] = { - "generation_steps": 0, - "evicted_tokens": 0, - } - self._ensure_calibrated() + def _attention_layer_partition(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) - def _validate_request_capacity(self, request: "LlmRequest") -> None: - manager = self.kv_cache_manager - speculative_overshoot = int(manager.max_draft_len) - first_eviction_decode_length = ( - self.budget // self.beta + 1 - ) * self.beta + speculative_overshoot - decode_capacity = min(int(request.py_max_new_tokens), first_eviction_decode_length) - confirmed_capacity = int(request.py_prompt_len) + decode_capacity - protected_tail_capacity = self._protected_tail_capacity - required_capacity = confirmed_capacity + protected_tail_capacity - pool_confirmed_capacity = manager.get_num_available_tokens( - token_num_upper_bound=confirmed_capacity, - max_num_draft_tokens=int(manager._kv_reserve_draft_tokens) + 1, - ) - table_capacity = manager.max_blocks_per_seq * manager.tokens_per_block - if confirmed_capacity > pool_confirmed_capacity or required_capacity > table_capacity: - raise ValueError( - "TriAttention target KV capacity is too small to reach its first " - f"eviction: request requires {required_capacity} tokens " - f"(prompt={request.py_prompt_len}, budget={self.budget}, " - f"beta={self.beta}, decode before eviction or completion=" - f"{decode_capacity}, speculative overshoot=" - f"{speculative_overshoot}, protected tail=" - f"{protected_tail_capacity}), " - f"but the V2 pool covers {pool_confirmed_capacity + protected_tail_capacity} " - f"tokens and its page table covers {table_capacity} tokens" + try: + from transformers import AutoConfig + + config = AutoConfig.from_pretrained( + model_path, trust_remote_code=True, local_files_only=True ) - draft_manager = self.draft_kv_cache_manager - if draft_manager is None: - return - draft_protected_tail = self._draft_protected_tail_capacity - draft_required_capacity = confirmed_capacity + draft_protected_tail - draft_pool_capacity = draft_manager.get_num_available_tokens( - token_num_upper_bound=confirmed_capacity, - max_num_draft_tokens=int(draft_manager._kv_reserve_draft_tokens) + 1, - ) - draft_table_capacity = draft_manager.max_blocks_per_seq * draft_manager.tokens_per_block - if ( - confirmed_capacity > draft_pool_capacity - or draft_required_capacity > draft_table_capacity - ): + except (OSError, ValueError) as exc: + raise ValueError( + f"TriAttention could not load the local model config from {model_path!r}" + ) from exc + config_values = config.get_text_config().to_dict() + layer_types = config_values.get("layer_types") + if not layer_types: + if self._has_sliding_window_signal(config_values): + 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( - "TriAttention draft KV capacity is too small to reach the first " - f"co-compression: request requires {draft_required_capacity} " - f"tokens (prompt={request.py_prompt_len}, budget={self.budget}, " - f"beta={self.beta}, decode before eviction or completion=" - f"{decode_capacity}, draft protected tail={draft_protected_tail}), " - f"but the draft V2 pool covers " - f"{draft_pool_capacity + draft_protected_tail} tokens and its " - f"page table covers {draft_table_capacity} tokens" + f"Model config has {len(layer_types)} layer_types entries, " + f"but this PP rank references global layer {max(global_layers)}" ) - def _ensure_calibrated(self) -> None: - if self.calibration is not None: - return + 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] + 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) + + @staticmethod + def _has_sliding_window_signal(config: Dict[str, object]) -> bool: + use_sliding_window = config.get("use_sliding_window") + if isinstance(use_sliding_window, bool): + return use_sliding_window + return any( + config.get(field) + for field in ( + "sliding_window", + "sliding_window_size", + "sliding_window_pattern", + "max_window_layers", + ) + ) + + def _load_calibration(self) -> None: self.calibration = self._resolve_calibration() self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) # Pre-split query stats + MLR coefficient, shapes [L, H, F]. @@ -250,30 +256,168 @@ def _ensure_calibrated(self) -> None: self.calibration["E_q_norm"].to(torch.float32) - _Eq.abs().to(torch.float32) ).contiguous() + def _resolve_calibration(self) -> Dict[str, torch.Tensor]: + """Load the calibration file, converting the official layout if needed.""" + raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) + if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): + return raw + if isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): + return self._convert_official_calibration(raw) + got = sorted(raw.keys()) 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}." + ) + + def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: + """Convert the official calibration format to the runtime schema.""" + stats = raw["stats"] + meta = raw.get("metadata", {}) + if "sampled_heads" in meta: + heads = [(int(a), int(b)) for a, b in meta["sampled_heads"]] + else: + heads = [ + (int(k[len("layer") : k.index("_head")]), int(k[k.index("_head") + len("_head") :])) + for k in stats + ] + num_layers = max(layer for layer, _ in heads) + 1 + num_heads = max(h for _, h 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, h in heads: + s = stats[f"layer{layer:02d}_head{h:02d}"] + E_q[layer, h] = torch.complex(s["q_mean_real"].float(), s["q_mean_imag"].float()) + E_q_norm[layer, h] = s["q_abs_mean"].float() + omega, freq_scale_sq = self._rope_tables(freq_count) + calib = { + "E_q": E_q, + "E_q_norm": E_q_norm, + "omega": omega, + "freq_scale_sq": freq_scale_sq, + } + logger.info( + f"TriAttention: converted official calibration {self.calibration_path}" + f" -> E_q[L={num_layers}, H={num_heads}, F={freq_count}]" + ) + return calib + + def _rope_tables(self, freq_count: int): + """Derive the RoPE frequency tables from the model config.""" + import transformers + from transformers import AutoConfig + + cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() + config_values = cfg.to_dict() + head_dim = freq_count * 2 + rope_params = ( + config_values.get("rope_parameters") or config_values.get("rope_scaling") or {} + ) + if rope_params and all(isinstance(v, dict) for v in rope_params.values()): + raise ValueError( + f"TriAttention: layer-type-keyed rope_parameters are not supported for " + f"calibration conversion (model {self.model_path}); got {rope_params!r}." + ) + rope_type = rope_params.get("rope_type") or rope_params.get("type") or "default" + theta_seen = rope_params.get("rope_theta", config_values.get("rope_theta")) + base = float(theta_seen) if theta_seen is not None else 10000.0 + + def analytic_inv_freq(): + idx = torch.arange(0, head_dim, 2, dtype=torch.float32) + return (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() + + if rope_type == "default": + # transformers>=5.5 no longer keys "default" in ROPE_INIT_FUNCTIONS: use the formula. + omega, scale_sq = analytic_inv_freq(), 1.0 + else: + try: + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + except ImportError: + logger.warning( + f"TriAttention: transformers rope-init unavailable; using the analytic " + f"inv_freq with theta={base} for {self.model_path} and IGNORING " + f"rope_type={rope_type!r} scaling corrections." + ) + return analytic_inv_freq(), torch.ones(freq_count, dtype=torch.float32) + if rope_type not in ROPE_INIT_FUNCTIONS: + raise ValueError( + f"TriAttention: unknown rope_type {rope_type!r} for {self.model_path} " + f"(transformers {transformers.__version__} provides " + f"{sorted(ROPE_INIT_FUNCTIONS)}); rope config seen: {rope_params!r}." + ) + try: + inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") + except Exception as exc: + raise ValueError( + f"TriAttention: rope-init {rope_type!r} failed for {self.model_path}; " + f"rope config seen: {rope_params!r}." + ) from exc + omega = inv_freq.to(torch.float32)[:freq_count].clone() + scale_sq = float(attention_factor) ** 2 + return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) + + # ---- framework hooks (call order) ---- + + def on_request_init(self, request: "LlmRequest", **kwargs) -> None: + """Register the request for eviction tracking.""" + request_id = request.py_request_id + if request_id not in self._request_states: + self._validate_request_capacity(request) + self._request_states[request_id] = { + "generation_steps": 0, + "evicted_tokens": 0, + } + + def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: + """Snapshot the prepared batch; mutation remains in final update.""" + self._prepared_generation_batch = scheduled_batch + self._prepared_generation_ids = None + def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Compact after native KV-cache updates have finalized this iteration (must run after KVCacheManagerV2 so capacity reflects the written token and any rewind).""" with nvtx_range_debug("triattention.generation_step_end", color="blue"): self._periodic_evict(scheduled_batch) - def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Snapshot the prepared batch; mutation remains in final update.""" - self._prepared_generation_batch = scheduled_batch - self._prepared_generation_ids = None + def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: + """Drop this request's eviction state; the buffers stay resident.""" + self._request_states.pop(request.py_request_id, None) - def _inflight_generation_growth( - self, scheduled_batch: "ScheduledRequests", request_id: int - ) -> int: - prepared = self._prepared_generation_batch - if prepared is None or scheduled_batch is prepared: - return 0 - member_ids = self._prepared_generation_ids - if member_ids is None: - member_ids = {request.py_request_id for request in prepared.generation_requests} - self._prepared_generation_ids = member_ids - if request_id not in member_ids: - return 0 - return self._generation_growth + # ---- request capacity ---- + + def _validate_request_capacity(self, request: "LlmRequest") -> None: + """Reject a request whose pre-first-eviction peak cannot fit (only + TriAttention can compute it: the framework's dense guards are off).""" + speculative_overshoot = int(self.kv_cache_manager.max_draft_len) + first_eviction_decode_length = ( + self.budget // self.beta + 1 + ) * self.beta + speculative_overshoot + decode_capacity = min(int(request.py_max_new_tokens), first_eviction_decode_length) + confirmed_capacity = int(request.py_prompt_len) + decode_capacity + checked = [(self.kv_cache_manager, self._protected_tail_capacity, "target")] + if self.draft_kv_cache_manager is not None: + checked.append( + (self.draft_kv_cache_manager, self._draft_protected_tail_capacity, "draft") + ) + for manager, protected_tail, label in checked: + required_capacity = confirmed_capacity + protected_tail + pool_capacity = manager.get_num_available_tokens( + token_num_upper_bound=confirmed_capacity, + max_num_draft_tokens=int(manager._kv_reserve_draft_tokens) + 1, + ) + table_capacity = manager.max_blocks_per_seq * manager.tokens_per_block + if confirmed_capacity > pool_capacity or required_capacity > table_capacity: + raise ValueError( + f"TriAttention {label} KV capacity is too small to reach the first " + f"eviction: request requires {required_capacity} tokens " + f"(prompt={request.py_prompt_len}, budget={self.budget}, " + f"beta={self.beta}, protected tail={protected_tail}), but the " + f"V2 pool covers {pool_capacity + protected_tail} tokens and " + f"its page table covers {table_capacity} tokens" + ) + + # ---- eviction round ---- def _periodic_evict( self, @@ -356,344 +500,417 @@ def _periodic_evict( compacted = self._evict_requests(prepared) self._resize_compacted_requests(compacted) - def _resize_compacted_requests(self, prepared) -> None: - if not prepared: - return - with nvtx_range("triattention.resize", color="red"): - with nvtx_range_debug("triattention.v2_resize", color="red"): - families = [("target", "kv_cache", None)] - if self.draft_kv_cache_manager is not None: - # Same kept set: the draft shrinks to the same retained - # length plus its own fixed tail. - families.append( - ("draft", "draft_kv_cache", self._draft_protected_tail_capacity) - ) - for label, cache_key, fixed_tail in families: - for item in prepared: - kv_cache = item[cache_key] - if not kv_cache.is_active: - # Bytes already moved: the compact-to-resize window is owned by this hook. - raise RuntimeError( - f"Request {item['request_id']} {label} KV cache was " - "suspended between compact and resize" - ) - tail = item["protected_tail"] if fixed_tail is None else fixed_tail - resized_capacity = item["expected_keep_count"] + tail - if not kv_cache.resize(resized_capacity, None): - raise RuntimeError( - f"Failed to resize compacted {label} KV cache for " - f"request {item['request_id']} to {resized_capacity} tokens" - ) + def _inflight_generation_growth( + self, scheduled_batch: "ScheduledRequests", request_id: int + ) -> int: + prepared = self._prepared_generation_batch + if prepared is None or scheduled_batch is prepared: + return 0 + member_ids = self._prepared_generation_ids + if member_ids is None: + member_ids = {request.py_request_id for request in prepared.generation_requests} + self._prepared_generation_ids = member_ids + if request_id not in member_ids: + return 0 + return self._generation_growth def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: """Return the largest cache length for which selection is an identity.""" prompt_len = min(int(request.py_prompt_len), seq_len) return prompt_len + self.budget - def _local_score_calibration( + def _evict_requests( self, - global_layers: List[int], - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - num_layers = len(global_layers) - if global_layers and max(global_layers) >= self._triattn_q_real.shape[0]: - raise ValueError( - f"TriAttention calibration has {self._triattn_q_real.shape[0]} layers, " - f"but this PP rank references global layer {max(global_layers)}" - ) - if global_layers == list(range(global_layers[0], global_layers[0] + num_layers)): - layer_slice = slice(global_layers[0], global_layers[0] + num_layers) - return ( - self._triattn_q_real[layer_slice], - self._triattn_q_imag[layer_slice], - self._triattn_mlr_coef[layer_slice], - ) - layer_ids = torch.as_tensor( - global_layers, - device=self._triattn_q_real.device, - dtype=torch.long, - ) - return ( - self._triattn_q_real.index_select(0, layer_ids), - self._triattn_q_imag.index_select(0, layer_ids), - self._triattn_mlr_coef.index_select(0, layer_ids), - ) - - def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: - """Drop this request's eviction state; the buffers stay resident.""" - self._request_states.pop(request.py_request_id, None) - - # ---- helpers (eviction / scoring / V2 cache access / calibration) ---- - - @staticmethod - def _has_sliding_window_signal(config: Dict[str, object]) -> bool: - use_sliding_window = config.get("use_sliding_window") - if isinstance(use_sliding_window, bool): - return use_sliding_window - return any( - config.get(field) - for field in ( - "sliding_window", - "sliding_window_size", - "sliding_window_pattern", - "max_window_layers", - ) - ) - - def _attention_layer_partition(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) - - try: - from transformers import AutoConfig + prepared: List[Dict[str, object]], + ) -> List[Dict[str, object]]: + with nvtx_range_debug("triattention.resolve_layout", color="blue"): + layout = self._runtime_kv_layout() + with nvtx_range_debug("triattention.staging_lookup", color="blue"): + # Retained spans always cover the model window (construction rejects budget < window). + self._ensure_buffers(layout, prepared) + self._execute_eviction_round(prepared) + for item in prepared: + # Identity cohorts were filtered pre-launch (_periodic_evict). + evicted = item["seq_len"] - item["expected_keep_count"] + request_state = self._request_states[item["request_id"]] + request_state["evicted_tokens"] += evicted + # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). + item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] + return prepared - config = AutoConfig.from_pretrained( - model_path, trust_remote_code=True, local_files_only=True + def _execute_eviction_round( + self, + prepared: Sequence[Dict[str, object]], + ) -> None: + """Run one eviction round over the prepared cohort (every launch covers + the full request capacity; padded rows carry zero lengths and stay inert).""" + manager = self.kv_cache_manager + draft_manager = self.draft_kv_cache_manager + with nvtx_range_debug("triattention.page_table_stage", color="orange"): + request_ids = [item["request_id"] for item in prepared] + round_starts = [item["round_start"] for item in prepared] + token_starts = [item["prompt_len"] for item in prepared] + seq_lens = [item["seq_len"] for item in prepared] + dense_move_offsets, swa_move_offsets, draft_move_offsets = self._cohort_move_offsets( + prepared ) - except (OSError, ValueError) as exc: - raise ValueError( - f"TriAttention could not load the local model config from {model_path!r}" - ) from exc - config_values = config.get_text_config().to_dict() - layer_types = config_values.get("layer_types") - if not layer_types: - if self._has_sliding_window_signal(config_values): - 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)}" + stream = torch.cuda.current_stream(self._block_offsets_device.device) + # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. + max_round_start = max(round_starts) + rows = ( + (0, round_starts), + (1, seq_lens), + (2, token_starts), + (3, dense_move_offsets), + (4, swa_move_offsets), + (5, draft_move_offsets), ) - - 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] - 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 _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: - # Only pool page counts are polled; manager identity and layer count are manager-lifetime contracts. - manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager - cached = self._kv_layout_caches[draft] - if cached is not None: - current_page_counts = self._pool_page_counts( + for row, values in rows: + if ( + values is not None + and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF + ): + raise ValueError(f"staged metadata row {row} exceeds the int32 range") + # Host-staging reuse fence: prior cohort's async copies must finish before the pinned rows are rewritten. + self._staging_reuse_event.synchronize() + host_table = self._request_metadata_host_np + for row, values in rows: + if values is not None: + host_table[row, : len(values)] = values + # Zero lengths keep the score kernel and selection inert for padded rows. + host_table[:3, len(prepared) :] = 0 + grow_mean_phase_table(self._phase, int(max_round_start) + 1) + self._stage_block_offsets( manager, - cached["global_layers"], - cached["pool_representatives"], + request_ids, + self._block_offsets_host, + self._block_offsets_device, ) - if current_page_counts != cached["pool_page_counts"]: - raise RuntimeError( - f"TriAttention {'draft ' if draft else ''}V2 pool layout changed " - "after the layout was built; KV pool rebalance is not supported" + if draft_manager is not None: + self._stage_block_offsets( + draft_manager, + request_ids, + self._draft_block_offsets_host, + self._draft_block_offsets_device, ) - return cached - - if draft: - global_layers = [int(layer) for layer in manager.pp_layers] - if not global_layers: - raise RuntimeError("TriAttention draft KV cache manager exposes no 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, swa_layers, swa_window = self._layer_partition - if not dense_layers: - raise ValueError("TriAttention requires at least one full-attention layer") - layout = self._build_runtime_kv_layout( - manager, - global_layers, - dense_layers=dense_layers, - swa_layers=swa_layers, - swa_window=swa_window, - what="draft " if draft else "", - ) - self._kv_layout_caches[draft] = layout - return layout - - def _build_runtime_kv_layout( - self, - manager: KVCacheManagerV2, - global_layers: List[int], - *, - dense_layers: List[int], - swa_layers: List[int], - swa_window: Optional[int], - what: str, - ) -> Dict[str, object]: - maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] - if any(pool is None for pool in maybe_layer_pools): - missing = [ - layer for layer, pool in zip(global_layers, maybe_layer_pools) if pool is None - ] - raise RuntimeError(f"Missing {what}KV pools for attention layers {missing}") - layer_pools = [pool for pool in maybe_layer_pools if pool is not None] - # Canonical pool IDs, resolved once; every grouping derives from them - # (V2 owns the mapping; its own 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 - ) - all_storage_groups: Dict[int, List[int]] = {} - for layer, pool_id in enumerate(layer_pool_ids): - all_storage_groups.setdefault(pool_id, []).append(layer) - # Scored/compacted groups cover the dense layers only; SWA layers - # stage and compact as their own representatives. - storage_groups: Dict[int, List[int]] = {} - for layer in dense_layers: - storage_groups.setdefault(layer_pool_ids[layer], []).append(layer) - layer_group_representative = { - layer: layers[0] for layers in storage_groups.values() for layer in layers - } - pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) - return dict( - manager=manager, - global_layers=global_layers, - layer_pools=layer_pools, - dense_layers=dense_layers, - swa_layers=swa_layers, - swa_window=swa_window, - storage_groups=storage_groups, - layer_group_representative=layer_group_representative, - layer_pool_ids=layer_pool_ids, - pool_representatives=pool_representatives, - pool_page_counts=tuple( - int(layer_pools[layer].shape[0]) for layer in pool_representatives - ), - ) - - @staticmethod - def _pool_page_counts( - manager: KVCacheManagerV2, - global_layers: Sequence[int], - pool_representatives: Sequence[int], - ) -> Tuple[int, ...]: - return tuple( - int( - manager.impl.get_page_index_upper_bound( - manager.layer_offsets[global_layers[layer]], - Role.KEY, + try: + self._request_metadata_device.copy_(self._request_metadata_host, non_blocking=True) + finally: + # Guards the pinned staging until the asynchronous copies complete. + self._staging_reuse_event.record(stream) + request_count = len(prepared) + union = self.eviction_mode == "union" + try: + with nvtx_range("triattention.score", color="blue"): + # In-place refresh: the compiled score launches captured these pointers. + _gather_mean_phase_kernel[(request_count,)]( + self._round_starts_device, + self._phase["cos"], + self._phase["sin"], + self._phase["rows"], + self._valid_seq_lens_device, + self._token_starts_device, + self._mean_cos, + self._mean_sin, + self._valid_widths, + self._swa_destination_bases, + self._swa_rebase_delta, + NUM_FREQS=self._phase_num_freqs, + F_BLOCK=self._phase_f_block, + HAS_SWA=self._swa_destination_bases is not None, + num_warps=1, ) - ) - // int(manager.kv_factor) - for layer in pool_representatives - ) + if union: + # Union path: fused score+stats, then the normalized union reduction. + cu_stream = cuda_driver.CUstream(stream.cuda_stream) + self._compiled_score_stats[request_count]( + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + request_count, + cu_stream, + ) + self._compiled_normalize_union[request_count]( + self._cute_partial_stats, + *self._cute_selection_prefix, + self._cute_union_scores, + request_count, + cu_stream, + ) + columns = min(self._union_scores.shape[1], self._selection_scores_rows.shape[1]) + self._selection_scores_rows[:request_count, :columns].copy_( + self._union_scores[:request_count, :columns] + ) + else: + cu_stream = cuda_driver.CUstream(stream.cuda_stream) + self._compiled_score[request_count]( + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + request_count, + cu_stream, + ) + # Gather each decode window into the [request, layer, head, token] layout of the reduces. + group_size = self._num_q_heads // self._num_kv_heads + num_segments = request_count * self._num_layers + pad = self._padded_head_columns + source = ( + self._score_scratch[ + : self._num_kv_heads * pad * num_segments * self._bucket_seq_len + ] + .view( + self._num_kv_heads, + pad, + request_count, + self._num_layers, + self._bucket_seq_len, + )[:, :group_size] + .permute(2, 3, 0, 1, 4) + ) + torch.add( + self._token_starts_device[:request_count].view(-1, 1, 1, 1, 1), + self._gather_columns, + out=self._gather_index_base[:request_count], + ) + self._gather_index_base[:request_count].clamp_(max=self._bucket_seq_len - 1) + columns = self._gather_index[:request_count] + torch.gather( + source, + 4, + columns, + out=self._score_output[:request_count].view( + request_count, + self._num_layers, + self._num_kv_heads, + group_size, + self._decode_width, + ), + ) + with nvtx_range("triattention.select", color="yellow"): + if not union: + prepare_per_head_scores( + self._score_output[:request_count], + self._valid_widths, + self._row_mean, + self._row_inv_std, + self._selection_scores_rows, + self._selection_row_lengths, + per_layer=self.eviction_mode == "per_layer_perhead", + normalize_scores=self.normalize_scores, + ) + self._settle_top_tokens(request_count) + with nvtx_range("triattention.compact", color="purple"): + compact(self._compaction_plans, request_count) + finally: + # Order V2 page-table reuse and resize after this cohort's compact. + self._compaction_done_event.record(stream) + manager._stream.wait_event(self._compaction_done_event) + if draft_manager is not None: + draft_manager._stream.wait_event(self._compaction_done_event) - def _ensure_buffers( + def _cohort_move_offsets( self, - layout: Dict[str, object], prepared: Sequence[Dict[str, object]], - ) -> None: - # Empty cohorts never reach here: _periodic_evict no-ops pre-launch. - needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) - needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) - needed_requests = len(prepared) - if self.draft_kv_cache_manager is not None: - # The cached layout lookup enforces draft V2 pool page-count - # stability every round, exactly like the target's lookup. - self._runtime_kv_layout(draft=True) - if self._buffers_built: - if ( - needed_width <= self._decode_width - and needed_page_tokens <= self._page_table_token_capacity - and needed_requests <= self._max_requests - ): - return - # This round outgrew the buffers: rebuild. - self._buffers_built = False + ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: + """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep + set plus protected tail per request; rows past the cohort repeat the final + offset and contribute no moves).""" - mgr = self.kv_cache_manager - tail_capacity = self._protected_tail_capacity - request_capacity = max(needed_requests, int(mgr.max_batch_size)) - decode_width = max( - needed_width, - self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), - ) - # Bucket sized by the presented cohorts, NOT max_seq_len (a floor there breaks 32-bit indexing). - seq_capacity = max(int(needed_page_tokens), 1024) - seq_capacity = 1 << (seq_capacity - 1).bit_length() - seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) - # The bucket capacity must be tile-aligned (mis-tiling stripes the - # score scratch silently); the ceiling division constructs that fact. - score_tile_tokens = max(64, int(mgr.tokens_per_block)) - seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens - page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) + def padded_offsets(moves_per_request: List[int]) -> List[int]: + offsets = [0] + for moves in moves_per_request: + offsets.append(offsets[-1] + moves) + offsets.extend(offsets[-1:] * (self._max_requests - len(moves_per_request))) + return offsets + tails = [int(item["protected_tail"]) for item in prepared] + dense = padded_offsets([self._keep_count + tail for tail in tails]) + swa = None + if self._swa_window is not None: + swa = padded_offsets([self._swa_window + tail for tail in tails]) draft = None - if self.draft_kv_cache_manager is not None: - draft_tail_capacity = self._draft_protected_tail_capacity - draft = dict( - layout=self._runtime_kv_layout(draft=True), - protected_tail_capacity=draft_tail_capacity, - page_table_token_capacity=seq_capacity + draft_tail_capacity, + if self._draft_protected_tail_capacity is not None: + draft = padded_offsets( + [self._keep_count + self._draft_protected_tail_capacity] * len(prepared) ) + return dense, swa, draft - first_pool = layout["layer_pools"][layout["dense_layers"][0]] - if self._phase is None: - # Host-only width offsets for the table builder (no device copy). - self._phase = { - "omega": self.calibration["omega"] - .to(device=first_pool.device, dtype=torch.float32) - .contiguous(), - "offset_values": [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], - "cos": None, - "sin": None, - "rows": 0, - } - grow_mean_phase_table(self._phase, max(int(seq_capacity), 1)) - q_real, q_imag, mlr_coef = self._local_score_calibration(layout["global_layers"]) - self._build_buffers( - layout=layout, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=self._freq_scale_sq, - max_requests=request_capacity, - bucket_seq_len=seq_capacity, - decode_width=decode_width, - page_table_token_capacity=page_table_token_capacity, - keep_count=self.budget, - protected_tail_capacity=tail_capacity, - draft=draft, - ) - self._buffers_built = True - - def _build_buffers( + def _stage_block_offsets( self, - *, - layout: Dict[str, object], - q_real: torch.Tensor, - q_imag: torch.Tensor, - mlr_coef: torch.Tensor, - freq_scale_sq: torch.Tensor, - max_requests: int, - bucket_seq_len: int, - decode_width: int, - page_table_token_capacity: int, + manager: KVCacheManagerV2, + request_ids: List[int], + host_block_offsets: torch.Tensor, + device_block_offsets: torch.Tensor, + ) -> None: + """Gather the pinned snapshot before the async device copy: resize mutates + the live host table. The round owner has already fenced host-staging reuse.""" + manager.index_mapper.gather_k_block_offsets( + manager.host_kv_cache_block_offsets, + host_block_offsets, + request_ids, + host_block_offsets.shape[-1], + ) + manager._stream.wait_event(self._staging_reuse_event) + 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, + manager._stream.cuda_stream, + ) + self._block_offsets_ready_event.record(manager._stream) + torch.cuda.current_stream(device_block_offsets.device).wait_event( + self._block_offsets_ready_event + ) + + def _settle_top_tokens(self, request_count: int) -> None: + """Pick the top-k and settle ties into the kept-ordinal decision rows + (the compaction contract packs them into move sources).""" + 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._keep_count, + 1, + ) + _settle_ties_kernel[(request_count, self._selection_rows_per_request)]( + self._selection_scores_rows, + self._selection_row_lengths, + self._token_starts_device, + self._provisional_rows, + self._kept_ordinal_rows, + WIDTH=self._decode_width, + KEEP_COUNT=self._keep_count, + SELECTION_ROWS=self._selection_rows_per_request, + ) + + def _resize_compacted_requests(self, prepared) -> None: + if not prepared: + return + with nvtx_range("triattention.resize", color="red"): + with nvtx_range_debug("triattention.v2_resize", color="red"): + families = [("target", "kv_cache", None)] + if self.draft_kv_cache_manager is not None: + # Same kept set: the draft shrinks to the same retained + # length plus its own fixed tail. + families.append( + ("draft", "draft_kv_cache", self._draft_protected_tail_capacity) + ) + for label, cache_key, fixed_tail in families: + for item in prepared: + kv_cache = item[cache_key] + if not kv_cache.is_active: + # Bytes already moved: the compact-to-resize window is owned by this hook. + raise RuntimeError( + f"Request {item['request_id']} {label} KV cache was " + "suspended between compact and resize" + ) + tail = item["protected_tail"] if fixed_tail is None else fixed_tail + resized_capacity = item["expected_keep_count"] + tail + if not kv_cache.resize(resized_capacity, None): + raise RuntimeError( + f"Failed to resize compacted {label} KV cache for " + f"request {item['request_id']} to {resized_capacity} tokens" + ) + + # ---- buffers + layout ---- + + def _ensure_buffers( + self, + layout: Dict[str, object], + prepared: Sequence[Dict[str, object]], + ) -> None: + # Empty cohorts never reach here: _periodic_evict no-ops pre-launch. + needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) + needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) + needed_requests = len(prepared) + if self.draft_kv_cache_manager is not None: + # The cached layout lookup enforces draft V2 pool page-count + # stability every round, exactly like the target's lookup. + self._runtime_kv_layout(draft=True) + if self._buffers_built: + if ( + needed_width <= self._decode_width + and needed_page_tokens <= self._page_table_token_capacity + and needed_requests <= self._max_requests + ): + return + # This round outgrew the buffers: rebuild. + self._buffers_built = False + + mgr = self.kv_cache_manager + tail_capacity = self._protected_tail_capacity + request_capacity = max(needed_requests, int(mgr.max_batch_size)) + decode_width = max( + needed_width, + self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), + ) + # Bucket sized by the presented cohorts, NOT max_seq_len (a floor there breaks 32-bit indexing). + seq_capacity = max(int(needed_page_tokens), 1024) + seq_capacity = 1 << (seq_capacity - 1).bit_length() + seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) + # The bucket capacity must be tile-aligned (mis-tiling stripes the + # score scratch silently); the ceiling division constructs that fact. + score_tile_tokens = max(64, int(mgr.tokens_per_block)) + seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens + page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) + + draft = None + if self.draft_kv_cache_manager is not None: + draft_tail_capacity = self._draft_protected_tail_capacity + draft = dict( + layout=self._runtime_kv_layout(draft=True), + protected_tail_capacity=draft_tail_capacity, + page_table_token_capacity=seq_capacity + draft_tail_capacity, + ) + + first_pool = layout["layer_pools"][layout["dense_layers"][0]] + if self._phase is None: + # Host-only width offsets for the table builder (no device copy). + self._phase = { + "omega": self.calibration["omega"] + .to(device=first_pool.device, dtype=torch.float32) + .contiguous(), + "offset_values": [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], + "cos": None, + "sin": None, + "rows": 0, + } + grow_mean_phase_table(self._phase, max(int(seq_capacity), 1)) + q_real, q_imag, mlr_coef = self._local_score_calibration(layout["global_layers"]) + self._build_buffers( + layout=layout, + q_real=q_real, + q_imag=q_imag, + mlr_coef=mlr_coef, + freq_scale_sq=self._freq_scale_sq, + max_requests=request_capacity, + bucket_seq_len=seq_capacity, + decode_width=decode_width, + page_table_token_capacity=page_table_token_capacity, + keep_count=self.budget, + protected_tail_capacity=tail_capacity, + draft=draft, + ) + self._buffers_built = True + + def _build_buffers( + self, + *, + layout: Dict[str, object], + q_real: torch.Tensor, + q_imag: torch.Tensor, + mlr_coef: torch.Tensor, + freq_scale_sq: torch.Tensor, + max_requests: int, + bucket_seq_len: int, + decode_width: int, + page_table_token_capacity: int, keep_count: int, protected_tail_capacity: int, draft: Optional[Dict[str, object]] = None, @@ -1142,433 +1359,193 @@ def _compiled_kernel(cache_key, build): self._selection_row_lengths = torch.full( (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device ) - self._provisional_rows = torch.zeros( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) - self._kept_ordinal_rows = torch.empty( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) - - # ---- compaction plans (opaque: only compact() interprets them) --------- - per_layer = self.eviction_mode == "per_layer_perhead" - draft_contract = None - if draft is not None: - draft_layout = draft["layout"] - draft_contract = dict( - layer_pools=draft_layout["layer_pools"], - dense_layers=list(draft_layout["dense_layers"]), - layer_group_representative=draft_layout["layer_group_representative"], - layer_pool_ids=tuple(draft_layout["layer_pool_ids"]), - kv_block_offsets=self._draft_block_offsets_device, - dense_move_offsets=draft_move_offsets_row, - protected_tail_capacity=int(draft["protected_tail_capacity"]), - ) - self._compaction_plan = init_compaction_buffers( - target=dict( - layer_pools=layer_pools, - dense_layers=list(dense_layers), - swa_layers=list(swa_layers), - swa_window=swa_window, - layer_group_representative=layer_group_representative, - layer_pool_ids=layer_pool_ids, - kv_block_offsets=self._block_offsets_device, - token_starts=self._token_starts_device, - swa_destination_bases=self._swa_destination_bases, - # Per-round tails: the move offsets ride the staged metadata rows. - dense_move_offsets=dense_move_offsets_row, - swa_move_offsets=swa_move_offsets_row, - per_layer_sources=per_layer, - # The decision rows the plans pack into move sources. - kept_ordinal_rows=self._kept_ordinal_rows, - decision_rows=self._selection_rows_per_request, - valid_seq_lens=self._valid_seq_lens_device, - ), - capacities=dict( - max_requests=max_requests, - keep_count=keep_count, - protected_tail_capacity=int(protected_tail_capacity), - ), - draft=draft_contract, - ) - - # ---- round-ordering events ---------------------------------------------- - # Host staging (pinned metadata + snapshots) reuse fence. - self._staging_reuse_event = torch.cuda.Event() - self._staging_reuse_event.record(torch.cuda.current_stream(device)) - # Manager-stream H2D of the block-offset tables has completed. - self._block_offsets_ready_event = torch.cuda.Event() - # This cohort's compact is done: manager may resize/reuse pages. - self._compaction_done_event = torch.cuda.Event() - - def _evict_requests( - self, - prepared: List[Dict[str, object]], - ) -> List[Dict[str, object]]: - with nvtx_range_debug("triattention.resolve_layout", color="blue"): - layout = self._runtime_kv_layout() - with nvtx_range_debug("triattention.staging_lookup", color="blue"): - # Retained spans always cover the model window (construction rejects budget < window). - self._ensure_buffers(layout, prepared) - self._execute_eviction_round(prepared) - for item in prepared: - # Identity cohorts were filtered pre-launch (_periodic_evict). - evicted = item["seq_len"] - item["expected_keep_count"] - request_state = self._request_states[item["request_id"]] - request_state["evicted_tokens"] += evicted - # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). - item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] - return prepared - - def _stage_block_offsets( - self, - manager: KVCacheManagerV2, - request_ids: List[int], - host_block_offsets: torch.Tensor, - device_block_offsets: torch.Tensor, - ) -> None: - """Gather the pinned snapshot before the async device copy: resize mutates - the live host table. The round owner has already fenced host-staging reuse.""" - manager.index_mapper.gather_k_block_offsets( - manager.host_kv_cache_block_offsets, - host_block_offsets, - request_ids, - host_block_offsets.shape[-1], - ) - manager._stream.wait_event(self._staging_reuse_event) - 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, - manager._stream.cuda_stream, - ) - self._block_offsets_ready_event.record(manager._stream) - torch.cuda.current_stream(device_block_offsets.device).wait_event( - self._block_offsets_ready_event - ) - - def _cohort_move_offsets( - self, - prepared: Sequence[Dict[str, object]], - ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep - set plus protected tail per request; rows past the cohort repeat the final - offset and contribute no moves).""" - - def padded_offsets(moves_per_request: List[int]) -> List[int]: - offsets = [0] - for moves in moves_per_request: - offsets.append(offsets[-1] + moves) - offsets.extend(offsets[-1:] * (self._max_requests - len(moves_per_request))) - return offsets - - tails = [int(item["protected_tail"]) for item in prepared] - dense = padded_offsets([self._keep_count + tail for tail in tails]) - swa = None - if self._swa_window is not None: - swa = padded_offsets([self._swa_window + tail for tail in tails]) - draft = None - if self._draft_protected_tail_capacity is not None: - draft = padded_offsets( - [self._keep_count + self._draft_protected_tail_capacity] * len(prepared) - ) - return dense, swa, draft - - def _settle_top_tokens(self, request_count: int) -> None: - """Pick the top-k and settle ties into the kept-ordinal decision rows - (the compaction contract packs them into move sources).""" - 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._keep_count, - 1, - ) - _settle_ties_kernel[(request_count, self._selection_rows_per_request)]( - self._selection_scores_rows, - self._selection_row_lengths, - self._token_starts_device, - self._provisional_rows, - self._kept_ordinal_rows, - WIDTH=self._decode_width, - KEEP_COUNT=self._keep_count, - SELECTION_ROWS=self._selection_rows_per_request, - ) - - def _execute_eviction_round( - self, - prepared: Sequence[Dict[str, object]], - ) -> None: - """Run one eviction round over the prepared cohort (every launch covers - the full request capacity; padded rows carry zero lengths and stay inert).""" - manager = self.kv_cache_manager - draft_manager = self.draft_kv_cache_manager - with nvtx_range_debug("triattention.page_table_stage", color="orange"): - request_ids = [item["request_id"] for item in prepared] - round_starts = [item["round_start"] for item in prepared] - token_starts = [item["prompt_len"] for item in prepared] - seq_lens = [item["seq_len"] for item in prepared] - dense_move_offsets, swa_move_offsets, draft_move_offsets = self._cohort_move_offsets( - prepared - ) - stream = torch.cuda.current_stream(self._block_offsets_device.device) - # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. - max_round_start = max(round_starts) - rows = ( - (0, round_starts), - (1, seq_lens), - (2, token_starts), - (3, dense_move_offsets), - (4, swa_move_offsets), - (5, draft_move_offsets), - ) - for row, values in rows: - if ( - values is not None - and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF - ): - raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # Host-staging reuse fence: prior cohort's async copies must finish before the pinned rows are rewritten. - self._staging_reuse_event.synchronize() - host_table = self._request_metadata_host_np - for row, values in rows: - if values is not None: - host_table[row, : len(values)] = values - # Zero lengths keep the score kernel and selection inert for padded rows. - host_table[:3, len(prepared) :] = 0 - grow_mean_phase_table(self._phase, int(max_round_start) + 1) - self._stage_block_offsets( - manager, - request_ids, - self._block_offsets_host, - self._block_offsets_device, - ) - if draft_manager is not None: - self._stage_block_offsets( - draft_manager, - request_ids, - self._draft_block_offsets_host, - self._draft_block_offsets_device, - ) - try: - self._request_metadata_device.copy_(self._request_metadata_host, non_blocking=True) - finally: - # Guards the pinned staging until the asynchronous copies complete. - self._staging_reuse_event.record(stream) - request_count = len(prepared) - union = self.eviction_mode == "union" - try: - with nvtx_range("triattention.score", color="blue"): - # In-place refresh: the compiled score launches captured these pointers. - _gather_mean_phase_kernel[(request_count,)]( - self._round_starts_device, - self._phase["cos"], - self._phase["sin"], - self._phase["rows"], - self._valid_seq_lens_device, - self._token_starts_device, - self._mean_cos, - self._mean_sin, - self._valid_widths, - self._swa_destination_bases, - self._swa_rebase_delta, - NUM_FREQS=self._phase_num_freqs, - F_BLOCK=self._phase_f_block, - HAS_SWA=self._swa_destination_bases is not None, - num_warps=1, - ) - if union: - # Union path: fused score+stats, then the normalized union reduction. - cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score_stats[request_count]( - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, - request_count, - cu_stream, - ) - self._compiled_normalize_union[request_count]( - self._cute_partial_stats, - *self._cute_selection_prefix, - self._cute_union_scores, - request_count, - cu_stream, - ) - columns = min(self._union_scores.shape[1], self._selection_scores_rows.shape[1]) - self._selection_scores_rows[:request_count, :columns].copy_( - self._union_scores[:request_count, :columns] - ) - else: - cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score[request_count]( - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, - request_count, - cu_stream, - ) - # Gather each decode window into the [request, layer, head, token] layout of the reduces. - group_size = self._num_q_heads // self._num_kv_heads - num_segments = request_count * self._num_layers - pad = self._padded_head_columns - source = ( - self._score_scratch[ - : self._num_kv_heads * pad * num_segments * self._bucket_seq_len - ] - .view( - self._num_kv_heads, - pad, - request_count, - self._num_layers, - self._bucket_seq_len, - )[:, :group_size] - .permute(2, 3, 0, 1, 4) - ) - torch.add( - self._token_starts_device[:request_count].view(-1, 1, 1, 1, 1), - self._gather_columns, - out=self._gather_index_base[:request_count], - ) - self._gather_index_base[:request_count].clamp_(max=self._bucket_seq_len - 1) - columns = self._gather_index[:request_count] - torch.gather( - source, - 4, - columns, - out=self._score_output[:request_count].view( - request_count, - self._num_layers, - self._num_kv_heads, - group_size, - self._decode_width, - ), - ) - with nvtx_range("triattention.select", color="yellow"): - if not union: - prepare_per_head_scores( - self._score_output[:request_count], - self._valid_widths, - self._row_mean, - self._row_inv_std, - self._selection_scores_rows, - self._selection_row_lengths, - per_layer=self.eviction_mode == "per_layer_perhead", - normalize_scores=self.normalize_scores, - ) - self._settle_top_tokens(request_count) - with nvtx_range("triattention.compact", color="purple"): - compact(self._compaction_plan, request_count) - finally: - # Order V2 page-table reuse and resize after this cohort's compact. - self._compaction_done_event.record(stream) - manager._stream.wait_event(self._compaction_done_event) - if draft_manager is not None: - draft_manager._stream.wait_event(self._compaction_done_event) - - # ---- helpers: calibration loading ---- - - def _resolve_calibration(self) -> Dict[str, torch.Tensor]: - """Load the calibration file, converting the official layout if needed.""" - raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) - if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): - return raw - if isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): - return self._convert_official_calibration(raw) - got = sorted(raw.keys()) 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._provisional_rows = torch.zeros( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + ) + self._kept_ordinal_rows = torch.empty( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + ) - def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: - """Convert the official calibration format to the runtime schema.""" - stats = raw["stats"] - meta = raw.get("metadata", {}) - if "sampled_heads" in meta: - heads = [(int(a), int(b)) for a, b in meta["sampled_heads"]] - else: - heads = [ - (int(k[len("layer") : k.index("_head")]), int(k[k.index("_head") + len("_head") :])) - for k in stats - ] - num_layers = max(layer for layer, _ in heads) + 1 - num_heads = max(h for _, h 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, h in heads: - s = stats[f"layer{layer:02d}_head{h:02d}"] - E_q[layer, h] = torch.complex(s["q_mean_real"].float(), s["q_mean_imag"].float()) - E_q_norm[layer, h] = s["q_abs_mean"].float() - omega, freq_scale_sq = self._rope_tables(freq_count) - calib = { - "E_q": E_q, - "E_q_norm": E_q_norm, - "omega": omega, - "freq_scale_sq": freq_scale_sq, - } - logger.info( - f"TriAttention: converted official calibration {self.calibration_path}" - f" -> E_q[L={num_layers}, H={num_heads}, F={freq_count}]" + # ---- compaction plans (opaque: only compact() interprets them) --------- + per_layer = self.eviction_mode == "per_layer_perhead" + draft_contract = None + if draft is not None: + draft_layout = draft["layout"] + draft_contract = dict( + layer_pools=draft_layout["layer_pools"], + dense_layers=list(draft_layout["dense_layers"]), + layer_group_representative=draft_layout["layer_group_representative"], + layer_pool_ids=tuple(draft_layout["layer_pool_ids"]), + kv_block_offsets=self._draft_block_offsets_device, + dense_move_offsets=draft_move_offsets_row, + protected_tail_capacity=int(draft["protected_tail_capacity"]), + ) + self._compaction_plans = build_compaction_plans( + target=dict( + layer_pools=layer_pools, + dense_layers=list(dense_layers), + swa_layers=list(swa_layers), + swa_window=swa_window, + layer_group_representative=layer_group_representative, + layer_pool_ids=layer_pool_ids, + kv_block_offsets=self._block_offsets_device, + token_starts=self._token_starts_device, + swa_destination_bases=self._swa_destination_bases, + # Per-round tails: the move offsets ride the staged metadata rows. + dense_move_offsets=dense_move_offsets_row, + swa_move_offsets=swa_move_offsets_row, + per_layer_sources=per_layer, + # The decision rows the plans pack into move sources. + kept_ordinal_rows=self._kept_ordinal_rows, + valid_seq_lens=self._valid_seq_lens_device, + protected_tail_capacity=int(protected_tail_capacity), + ), + draft=draft_contract, ) - return calib - def _rope_tables(self, freq_count: int): - """Derive the RoPE frequency tables from the model config.""" - import transformers - from transformers import AutoConfig + # ---- round-ordering events ---------------------------------------------- + # Host staging (pinned metadata + snapshots) reuse fence. + self._staging_reuse_event = torch.cuda.Event() + self._staging_reuse_event.record(torch.cuda.current_stream(device)) + # Manager-stream H2D of the block-offset tables has completed. + self._block_offsets_ready_event = torch.cuda.Event() + # This cohort's compact is done: manager may resize/reuse pages. + self._compaction_done_event = torch.cuda.Event() - cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() - config_values = cfg.to_dict() - head_dim = freq_count * 2 - rope_params = ( - config_values.get("rope_parameters") or config_values.get("rope_scaling") or {} - ) - if rope_params and all(isinstance(v, dict) for v in rope_params.values()): + def _local_score_calibration( + self, + global_layers: List[int], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + num_layers = len(global_layers) + if global_layers and max(global_layers) >= self._triattn_q_real.shape[0]: raise ValueError( - f"TriAttention: layer-type-keyed rope_parameters are not supported for " - f"calibration conversion (model {self.model_path}); got {rope_params!r}." + f"TriAttention calibration has {self._triattn_q_real.shape[0]} layers, " + f"but this PP rank references global layer {max(global_layers)}" ) - rope_type = rope_params.get("rope_type") or rope_params.get("type") or "default" - theta_seen = rope_params.get("rope_theta", config_values.get("rope_theta")) - base = float(theta_seen) if theta_seen is not None else 10000.0 + if global_layers == list(range(global_layers[0], global_layers[0] + num_layers)): + layer_slice = slice(global_layers[0], global_layers[0] + num_layers) + return ( + self._triattn_q_real[layer_slice], + self._triattn_q_imag[layer_slice], + self._triattn_mlr_coef[layer_slice], + ) + layer_ids = torch.as_tensor( + global_layers, + device=self._triattn_q_real.device, + dtype=torch.long, + ) + return ( + self._triattn_q_real.index_select(0, layer_ids), + self._triattn_q_imag.index_select(0, layer_ids), + self._triattn_mlr_coef.index_select(0, layer_ids), + ) - def analytic_inv_freq(): - idx = torch.arange(0, head_dim, 2, dtype=torch.float32) - return (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() + def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: + # Only pool page counts are polled; manager identity and layer count are manager-lifetime contracts. + manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager + cached = self._kv_layout_caches[draft] + if cached is not None: + current_page_counts = self._pool_page_counts( + manager, + cached["global_layers"], + cached["pool_representatives"], + ) + if current_page_counts != cached["pool_page_counts"]: + raise RuntimeError( + f"TriAttention {'draft ' if draft else ''}V2 pool layout changed " + "after the layout was built; KV pool rebalance is not supported" + ) + return cached - if rope_type == "default": - # transformers>=5.5 no longer keys "default" in ROPE_INIT_FUNCTIONS: use the formula. - omega, scale_sq = analytic_inv_freq(), 1.0 + if draft: + global_layers = [int(layer) for layer in manager.pp_layers] + if not global_layers: + raise RuntimeError("TriAttention draft KV cache manager exposes no 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: - try: - from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS - except ImportError: - logger.warning( - f"TriAttention: transformers rope-init unavailable; using the analytic " - f"inv_freq with theta={base} for {self.model_path} and IGNORING " - f"rope_type={rope_type!r} scaling corrections." - ) - return analytic_inv_freq(), torch.ones(freq_count, dtype=torch.float32) - if rope_type not in ROPE_INIT_FUNCTIONS: - raise ValueError( - f"TriAttention: unknown rope_type {rope_type!r} for {self.model_path} " - f"(transformers {transformers.__version__} provides " - f"{sorted(ROPE_INIT_FUNCTIONS)}); rope config seen: {rope_params!r}." + global_layers = self._global_layers + dense_layers, swa_layers, swa_window = self._layer_partition + if not dense_layers: + raise ValueError("TriAttention requires at least one full-attention layer") + layout = self._build_runtime_kv_layout( + manager, + global_layers, + dense_layers=dense_layers, + swa_layers=swa_layers, + swa_window=swa_window, + what="draft " if draft else "", + ) + self._kv_layout_caches[draft] = layout + return layout + + def _build_runtime_kv_layout( + self, + manager: KVCacheManagerV2, + global_layers: List[int], + *, + dense_layers: List[int], + swa_layers: List[int], + swa_window: Optional[int], + what: str, + ) -> Dict[str, object]: + maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] + if any(pool is None for pool in maybe_layer_pools): + missing = [ + layer for layer, pool in zip(global_layers, maybe_layer_pools) if pool is None + ] + raise RuntimeError(f"Missing {what}KV pools for attention layers {missing}") + layer_pools = [pool for pool in maybe_layer_pools if pool is not None] + # Canonical pool IDs, resolved once; every grouping derives from them + # (V2 owns the mapping; its own 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 + ) + all_storage_groups: Dict[int, List[int]] = {} + for layer, pool_id in enumerate(layer_pool_ids): + all_storage_groups.setdefault(pool_id, []).append(layer) + # Scored/compacted groups cover the dense layers only; SWA layers + # stage and compact as their own representatives. + storage_groups: Dict[int, List[int]] = {} + for layer in dense_layers: + storage_groups.setdefault(layer_pool_ids[layer], []).append(layer) + layer_group_representative = { + layer: layers[0] for layers in storage_groups.values() for layer in layers + } + pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) + return dict( + manager=manager, + global_layers=global_layers, + layer_pools=layer_pools, + dense_layers=dense_layers, + swa_layers=swa_layers, + swa_window=swa_window, + storage_groups=storage_groups, + layer_group_representative=layer_group_representative, + layer_pool_ids=layer_pool_ids, + pool_representatives=pool_representatives, + pool_page_counts=tuple( + int(layer_pools[layer].shape[0]) for layer in pool_representatives + ), + ) + + @staticmethod + def _pool_page_counts( + manager: KVCacheManagerV2, + global_layers: Sequence[int], + pool_representatives: Sequence[int], + ) -> Tuple[int, ...]: + return tuple( + int( + manager.impl.get_page_index_upper_bound( + manager.layer_offsets[global_layers[layer]], + Role.KEY, ) - try: - inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") - except Exception as exc: - raise ValueError( - f"TriAttention: rope-init {rope_type!r} failed for {self.model_path}; " - f"rope config seen: {rope_params!r}." - ) from exc - omega = inv_freq.to(torch.float32)[:freq_count].clone() - scale_sq = float(attention_factor) ** 2 - return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) + ) + // int(manager.kv_factor) + for layer in pool_representatives + ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 38cfeb38f594..b3339524c655 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2159,14 +2159,14 @@ def validate_kv_cache_compression_with_spec( raise ValueError( "TriAttention draft KV co-compression supports only " "eviction_mode='union'; draft layers are never scored") - if any(window is not None for window in - draft_kv_cache_manager.max_attention_window_vec) or any( - not isinstance(layer, AttentionLayerConfig) - or layer.sliding_window_size is not None for layer in - draft_kv_cache_manager.kv_cache_manager_py_config.layers): - raise ValueError( - "TriAttention draft KV co-compression requires " - "full-attention draft V2 lifecycles") + if any(window is not None + for window in draft_kv_cache_manager.max_attention_window_vec + ) or any(not isinstance(layer, AttentionLayerConfig) + or layer.sliding_window_size is not None + for layer in draft_kv_cache_manager. + kv_cache_manager_py_config.layers): + raise ValueError("TriAttention draft KV co-compression requires " + "full-attention draft V2 lifecycles") def create_kv_cache_compression_manager( @@ -2187,8 +2187,8 @@ def create_kv_cache_compression_manager( TriAttention return TriAttention( + config, kv_cache_manager, - config=config, draft_kv_cache_manager=draft_kv_cache_manager, ) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 9fd02c0a520a..b25da55924c4 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -117,13 +117,12 @@ def make_ramp_pools( def build_compaction(**overrides): - """``init_compaction_buffers`` with the suite's 2-layer defaults: - translates ``eviction_mode`` into ``per_layer_sources``/``decision_rows``, + """``build_compaction_plans`` with the suite's 2-layer defaults: allocates the caller-owned move-offset rows (capacity cumsum) and SWA destination bases, and hands the test's pre-settled ``kept_token_ordinals`` in as the decision rows. Returns the opaque ``plans`` plus a test-side mirror of the caller-owned inputs.""" - from tensorrt_llm._torch.kv_cache_compression.compaction import init_compaction_buffers + from tensorrt_llm._torch.kv_cache_compression.compaction import build_compaction_plans args = dict( eviction_mode="union", @@ -137,7 +136,6 @@ def build_compaction(**overrides): ) args.update(overrides) mode = args.pop("eviction_mode") - union = mode == "union" per_layer = mode == "per_layer_perhead" kept = args.pop("kept_token_ordinals") request_count = args["request_count"] @@ -157,10 +155,6 @@ def capacity_offsets(count): args.setdefault("swa_move_offsets", capacity_offsets(swa_window + tail) if has_swa else None) if has_draft: args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) - num_kv_heads = int(args["layer_pools"][args["dense_layers"][0]].shape[2]) - selection_rows = ( - 1 if union else (len(args["dense_layers"]) * num_kv_heads if per_layer else num_kv_heads) - ) draft = None if has_draft: draft = dict( @@ -172,7 +166,7 @@ def capacity_offsets(count): dense_move_offsets=args["draft_move_offsets"], protected_tail_capacity=draft_tail, ) - plans = init_compaction_buffers( + plans = build_compaction_plans( target=dict( layer_pools=args["layer_pools"], dense_layers=args["dense_layers"], @@ -187,12 +181,7 @@ def capacity_offsets(count): swa_move_offsets=args["swa_move_offsets"], per_layer_sources=per_layer, kept_ordinal_rows=kept.reshape(-1, keep_count), - decision_rows=selection_rows, valid_seq_lens=args["valid_sequence_lengths"], - ), - capacities=dict( - max_requests=request_count, - keep_count=keep_count, protected_tail_capacity=tail, ), draft=draft, @@ -353,6 +342,23 @@ def make_test_model_dir() -> str: 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): """A real TriAttentionKvCacheCompressionConfig with test calibration inputs (the config validator requires both ``model_path`` and ``calibration_path``).""" @@ -361,7 +367,7 @@ def make_tri_config(**overrides): options = { "budget": 8, "model_path": make_test_model_dir(), - "calibration_path": "/calib/test.pt", + "calibration_path": make_test_calibration_pt(), } options.update(overrides) return TriAttentionKvCacheCompressionConfig(**options) @@ -371,7 +377,7 @@ def make_triattention(**overrides): """Construct a fully initialized manager for method-level unit tests.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention - return TriAttention(make_fake_v2(), make_tri_config(**overrides)) + return TriAttention(make_tri_config(**overrides), make_fake_v2()) def make_prepared_item( 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 index e1350bff6203..dbd568af259b 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -177,7 +177,7 @@ def test_execute_eviction_round_orders_both_manager_streams(): tri._keep_count = 4 tri.eviction_mode = "union" tri._swa_window = None - tri._compaction_plan = () + tri._compaction_plans = () tri._draft_protected_tail_capacity = 1 tri._staging_reuse_event = mock.Mock() tri._compaction_done_event = event diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 9cadcc612225..84b2a9e07c95 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -147,7 +147,7 @@ def test_request_init_and_finish_lifecycle(self): manager = _make_fake_v2() manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 - triattention = TriAttention(manager, _make_tri_config(budget=8)) + triattention = TriAttention(_make_tri_config(budget=8), manager) triattention.calibration = {} triattention.on_request_init(_make_request(11)) @@ -327,7 +327,7 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft 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 - mgr = TriAttention(fake_v2, _make_tri_config(budget=8)) + mgr = TriAttention(_make_tri_config(budget=8), fake_v2) mgr.calibration = {} cache = SimpleNamespace( capacity=seq_len, @@ -458,8 +458,8 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): # marking is asserted in the executor manager tests). draft_manager = _make_fake_v2(is_draft=True) TriAttention( - _make_fake_v2(), _make_tri_config(budget=8), + _make_fake_v2(), draft_kv_cache_manager=draft_manager, ) @@ -489,7 +489,7 @@ def test_prepare_snapshots_fixed_linear_generation_growth( manager.kv_cache_map = { 7: SimpleNamespace(capacity=106, is_active=True), } - triattention = TriAttention(manager, _make_tri_config(budget=8)) + triattention = TriAttention(_make_tri_config(budget=8), manager) batch = SimpleNamespace( context_requests=[], generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], @@ -506,11 +506,10 @@ def test_prepare_snapshots_fixed_linear_generation_growth( class TestFixedScoreMetadata: - def test_union_rejects_unnormalized_scores(self): - # The fused pipeline (THE union path) always z-normalizes, so - # normalize_scores=False is rejected loudly at construction. - with pytest.raises(ValueError, match="normalize_scores=True"): - _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) + 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_execute_rejects_int32_overflowing_round_starts(self): # Round starts past the int32 metadata range fail loudly (in the host From 9b676a6bac32ba204767be64fd51b67249e6b0bf Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 01:14:56 -0700 Subject: [PATCH 144/178] [None][refactor] Give compaction a single-cache typed plan builder Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 68 ++++++++---------- .../triattention/triattention.py | 58 +++++++--------- .../_torch/kv_cache_compression/conftest.py | 69 ++++++++++--------- 3 files changed, 89 insertions(+), 106 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 192a9fe8aa05..c430ffad7102 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,7 +15,7 @@ """Batched physical KV-cache compaction: an algorithm-neutral mover. -``build_compaction_plans`` returns opaque launch plans once per geometry; +``build_compaction_plan`` returns one opaque launch plan per compacted cache; each round the caller writes its keep decision into the agreed rows and ``compact`` packs every plan's move sources and fires its native launches. """ @@ -121,42 +121,29 @@ def _pack_move_sources_kernel( ) -def build_compaction_plans( +def build_compaction_plan( + layout: Dict[str, object], *, - target: Dict[str, object], - draft: Optional[Dict[str, object]] = None, -) -> Tuple[CompactionPlan, ...]: - """One plan per compacted cache family (target, then the optional draft); - only :func:`compact` reads the plans' contents at launch.""" - plans = [_build_plan(target)] - if draft is not None: - # The draft compacts under the target's decision rows and token starts. - plans.append( - _build_plan( - dict( - draft, - token_starts=target["token_starts"], - kept_ordinal_rows=target["kept_ordinal_rows"], - valid_seq_lens=target["valid_seq_lens"], - ) - ) - ) - return tuple(plans) - - -def _build_plan(cache: Dict[str, object]) -> CompactionPlan: - layer_pools = cache["layer_pools"] - dense_layers = tuple(int(layer) for layer in cache["dense_layers"]) - swa_layers = tuple(int(layer) for layer in cache.get("swa_layers", ())) - layer_pool_ids = tuple(int(pool_id) for pool_id in cache["layer_pool_ids"]) - kv_block_offsets = cache["kv_block_offsets"] - layer_group_representative = cache["layer_group_representative"] - token_starts = cache["token_starts"] - dense_move_offsets = cache["dense_move_offsets"] - per_layer_sources = bool(cache.get("per_layer_sources", False)) - kept_ordinal_rows = cache["kept_ordinal_rows"] - valid_seq_lens = cache["valid_seq_lens"] - protected_tail_capacity = int(cache["protected_tail_capacity"]) + block_offsets: torch.Tensor, + kept_ordinals: torch.Tensor, + source_lengths: torch.Tensor, + dense_destination_bases: torch.Tensor, + dense_move_offsets: torch.Tensor, + protected_tail_capacity: int, + swa_move_offsets: Optional[torch.Tensor] = None, + swa_destination_bases: Optional[torch.Tensor] = None, +) -> CompactionPlan: + """One opaque launch plan for one compacted cache; only :func:`compact` + reads its contents at launch.""" + layer_pools = layout["layer_pools"] + dense_layers = tuple(int(layer) for layer in layout["dense_layers"]) + swa_layers = tuple(int(layer) for layer in layout["swa_layers"]) + layer_pool_ids = tuple(int(pool_id) for pool_id in layout["layer_pool_ids"]) + kv_block_offsets = block_offsets + kept_ordinal_rows = kept_ordinals + valid_seq_lens = source_lengths + token_starts = dense_destination_bases + protected_tail_capacity = int(protected_tail_capacity) first_pool = layer_pools[dense_layers[0]] device = first_pool.device @@ -165,6 +152,7 @@ def _build_plan(cache: Dict[str, object]) -> CompactionPlan: decision_rows = int(kept_ordinal_rows.shape[0]) // max_requests # Pool shape [pages, K/V, heads, tokens, dim]. num_kv_heads = int(first_pool.shape[2]) + per_layer_sources = len(dense_layers) > 1 and decision_rows == len(dense_layers) * num_kv_heads dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer_sources else (num_kv_heads,) dense_move_indices = torch.empty( (*dense_index_prefix, (keep_count + protected_tail_capacity) * max_requests), @@ -175,7 +163,7 @@ def _build_plan(cache: Dict[str, object]) -> CompactionPlan: ( layer, layer_pools[layer], - kv_block_offsets[layer_pool_ids[layer_group_representative[layer]], :max_requests, 0], + kv_block_offsets[layer_pool_ids[layer], :max_requests, 0], ) for layer in dense_layers ] @@ -184,15 +172,13 @@ def _build_plan(cache: Dict[str, object]) -> CompactionPlan: ) swa_move_indices = None - swa_move_offsets = None swa_window = 0 # One move group per family axis (dense / SWA): the layers + the tensors driving their moves. move_groups = [ (dense_entries, dense_move_indices, dense_move_offsets, token_starts, dense_slots), ] if swa_layers: - swa_move_offsets = cache["swa_move_offsets"] - swa_window = int(cache["swa_window"]) + swa_window = int(layout["swa_window"]) swa_move_indices = torch.empty( (num_kv_heads, (swa_window + protected_tail_capacity) * max_requests), dtype=torch.int32, @@ -208,7 +194,7 @@ def _build_plan(cache: Dict[str, object]) -> CompactionPlan: for layer in swa_layers ] move_groups.append( - (swa_entries, swa_move_indices, swa_move_offsets, cache["swa_destination_bases"], None) + (swa_entries, swa_move_indices, swa_move_offsets, swa_destination_bases, None) ) # Widest per-request move count any staged offsets may express. diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 17acf4a78374..3bf903b6bb1d 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -38,7 +38,7 @@ ) from tensorrt_llm.logger import logger -from ..compaction import build_compaction_plans, compact +from ..compaction import build_compaction_plan, compact from .triattention_kernels import ( _gather_mean_phase_kernel, _settle_ties_kernel, @@ -1367,41 +1367,33 @@ def _compiled_kernel(cache_key, build): ) # ---- compaction plans (opaque: only compact() interprets them) --------- - per_layer = self.eviction_mode == "per_layer_perhead" - draft_contract = None - if draft is not None: - draft_layout = draft["layout"] - draft_contract = dict( - layer_pools=draft_layout["layer_pools"], - dense_layers=list(draft_layout["dense_layers"]), - layer_group_representative=draft_layout["layer_group_representative"], - layer_pool_ids=tuple(draft_layout["layer_pool_ids"]), - kv_block_offsets=self._draft_block_offsets_device, - dense_move_offsets=draft_move_offsets_row, - protected_tail_capacity=int(draft["protected_tail_capacity"]), - ) - self._compaction_plans = build_compaction_plans( - target=dict( - layer_pools=layer_pools, - dense_layers=list(dense_layers), - swa_layers=list(swa_layers), - swa_window=swa_window, - layer_group_representative=layer_group_representative, - layer_pool_ids=layer_pool_ids, - kv_block_offsets=self._block_offsets_device, - token_starts=self._token_starts_device, - swa_destination_bases=self._swa_destination_bases, + plans = [ + build_compaction_plan( + layout, + block_offsets=self._block_offsets_device, + kept_ordinals=self._kept_ordinal_rows, + source_lengths=self._valid_seq_lens_device, + dense_destination_bases=self._token_starts_device, # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, - swa_move_offsets=swa_move_offsets_row, - per_layer_sources=per_layer, - # The decision rows the plans pack into move sources. - kept_ordinal_rows=self._kept_ordinal_rows, - valid_seq_lens=self._valid_seq_lens_device, protected_tail_capacity=int(protected_tail_capacity), - ), - draft=draft_contract, - ) + swa_move_offsets=swa_move_offsets_row, + swa_destination_bases=self._swa_destination_bases, + ) + ] + if draft is not None: + plans.append( + build_compaction_plan( + draft["layout"], + block_offsets=self._draft_block_offsets_device, + kept_ordinals=self._kept_ordinal_rows, + source_lengths=self._valid_seq_lens_device, + dense_destination_bases=self._token_starts_device, + dense_move_offsets=draft_move_offsets_row, + protected_tail_capacity=int(draft["protected_tail_capacity"]), + ) + ) + self._compaction_plans = tuple(plans) # ---- round-ordering events ---------------------------------------------- # Host staging (pinned metadata + snapshots) reuse fence. diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index b25da55924c4..a1ac14a03c3b 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -122,7 +122,7 @@ def build_compaction(**overrides): destination bases, and hands the test's pre-settled ``kept_token_ordinals`` in as the decision rows. Returns the opaque ``plans`` plus a test-side mirror of the caller-owned inputs.""" - from tensorrt_llm._torch.kv_cache_compression.compaction import build_compaction_plans + from tensorrt_llm._torch.kv_cache_compression.compaction import build_compaction_plan args = dict( eviction_mode="union", @@ -135,8 +135,7 @@ def build_compaction(**overrides): swa_window=None, ) args.update(overrides) - mode = args.pop("eviction_mode") - per_layer = mode == "per_layer_perhead" + args.pop("eviction_mode") kept = args.pop("kept_token_ordinals") request_count = args["request_count"] keep_count = args["decode_keep_count"] @@ -155,37 +154,43 @@ def capacity_offsets(count): args.setdefault("swa_move_offsets", capacity_offsets(swa_window + tail) if has_swa else None) if has_draft: args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) - draft = None - if has_draft: - draft = dict( - layer_pools=args["draft_layer_pools"], - dense_layers=args["draft_layers"], - layer_group_representative=args["draft_layer_group_representative"], - layer_pool_ids=args["draft_layer_pool_ids"], - kv_block_offsets=args["draft_kv_block_offsets"], - dense_move_offsets=args["draft_move_offsets"], - protected_tail_capacity=draft_tail, - ) - plans = build_compaction_plans( - target=dict( - layer_pools=args["layer_pools"], - dense_layers=args["dense_layers"], - swa_layers=args["swa_layers"], - swa_window=args["swa_window"], - layer_group_representative=args["layer_group_representative"], - layer_pool_ids=args["layer_pool_ids"], - kv_block_offsets=args["kv_block_offsets"], - token_starts=args["prompt_offsets"], - swa_destination_bases=swa_destination_bases, + plan_list = [ + build_compaction_plan( + dict( + layer_pools=args["layer_pools"], + dense_layers=args["dense_layers"], + swa_layers=args["swa_layers"], + swa_window=args["swa_window"], + layer_pool_ids=args["layer_pool_ids"], + ), + block_offsets=args["kv_block_offsets"], + kept_ordinals=kept.reshape(-1, keep_count), + source_lengths=args["valid_sequence_lengths"], + dense_destination_bases=args["prompt_offsets"], dense_move_offsets=args["dense_move_offsets"], - swa_move_offsets=args["swa_move_offsets"], - per_layer_sources=per_layer, - kept_ordinal_rows=kept.reshape(-1, keep_count), - valid_seq_lens=args["valid_sequence_lengths"], protected_tail_capacity=tail, - ), - draft=draft, - ) + swa_move_offsets=args["swa_move_offsets"], + swa_destination_bases=swa_destination_bases, + ) + ] + if has_draft: + plan_list.append( + build_compaction_plan( + dict( + layer_pools=args["draft_layer_pools"], + dense_layers=args["draft_layers"], + swa_layers=[], + layer_pool_ids=args["draft_layer_pool_ids"], + ), + block_offsets=args["draft_kv_block_offsets"], + kept_ordinals=kept.reshape(-1, keep_count), + source_lengths=args["valid_sequence_lengths"], + dense_destination_bases=args["prompt_offsets"], + dense_move_offsets=args["draft_move_offsets"], + protected_tail_capacity=draft_tail, + ) + ) + plans = tuple(plan_list) # Opaque plans plus a test-side mirror of the caller-owned construction # inputs (production binds the same values as manager attributes); the # standalone helpers here need the move-offset rows and SWA staging back. From 7b28954fbae234f885184ed449dd809b30c46054 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 01:26:28 -0700 Subject: [PATCH 145/178] [None][refactor] Flatten compaction launches into CompactionParams Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 112 ++++++++---------- .../triattention/triattention.py | 14 +-- .../_torch/kv_cache_compression/conftest.py | 22 ++-- .../test_triattention_draft_cocompaction.py | 2 +- 4 files changed, 66 insertions(+), 84 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index c430ffad7102..51d12b4a1e51 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -15,9 +15,9 @@ """Batched physical KV-cache compaction: an algorithm-neutral mover. -``build_compaction_plan`` returns one opaque launch plan per compacted cache; +``build_compaction_params`` pre-binds one cache's launch parameters; each round the caller writes its keep decision into the agreed rows and -``compact`` packs every plan's move sources and fires its native launches. +``compact`` packs every cache's move sources and fires its native launches. """ from typing import Dict, List, Optional, Tuple, TypedDict @@ -27,21 +27,11 @@ import triton.language as tl -class CompactionLaunch(TypedDict): - pools: List[torch.Tensor] - pool_pointers: torch.Tensor - page_table: torch.Tensor - move_indices: torch.Tensor - move_offsets: torch.Tensor - destination_bases: torch.Tensor - source_layer_indices: Optional[torch.Tensor] - - -class CompactionPlan(TypedDict): +class CompactionParams(TypedDict): decision_rows: int pack_args: Tuple[Optional[torch.Tensor], ...] pack_constexprs: Dict[str, object] - launches: Tuple[CompactionLaunch, ...] + compact_args: List[Tuple[object, ...]] @triton.jit @@ -121,7 +111,7 @@ def _pack_move_sources_kernel( ) -def build_compaction_plan( +def build_compaction_params( layout: Dict[str, object], *, block_offsets: torch.Tensor, @@ -132,9 +122,8 @@ def build_compaction_plan( protected_tail_capacity: int, swa_move_offsets: Optional[torch.Tensor] = None, swa_destination_bases: Optional[torch.Tensor] = None, -) -> CompactionPlan: - """One opaque launch plan for one compacted cache; only :func:`compact` - reads its contents at launch.""" +) -> CompactionParams: + """Pre-bind one compacted cache's launch parameters; only :func:`compact` reads them.""" layer_pools = layout["layer_pools"] dense_layers = tuple(int(layer) for layer in layout["dense_layers"]) swa_layers = tuple(int(layer) for layer in layout["swa_layers"]) @@ -202,7 +191,29 @@ def build_compaction_plan( if swa_layers: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) - launches: List[CompactionLaunch] = [] + params = CompactionParams( + decision_rows=decision_rows, + pack_args=( + kept_ordinal_rows, + valid_seq_lens, + dense_move_offsets, + dense_move_indices, + swa_move_offsets, + swa_move_indices, + ), + pack_constexprs=dict( + KEEP_COUNT=keep_count, + DECISION_ROWS=decision_rows, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + PER_LAYER=per_layer_sources, + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, + SWA_WINDOW=swa_window, + ), + # One positional-args tuple per native compact call, in its op signature order. + compact_args=[], + ) for entries, move_indices, move_offsets, destination_bases, slots in move_groups: grouped = {} for layer, pool, page_table in entries: @@ -224,63 +235,34 @@ def build_compaction_plan( dtype=torch.int32, device=device, ) - launches.append( - CompactionLaunch( - pools=pools, - pool_pointers=torch.tensor( + params["compact_args"].append( + ( + pools, + torch.tensor( [pool.data_ptr() for pool in pools], dtype=torch.int64, device=device, ), - page_table=group_entries[0][2], - move_indices=move_indices, - move_offsets=move_offsets, - destination_bases=destination_bases, - source_layer_indices=source_layer_indices, + group_entries[0][2], + move_indices, + move_offsets, + destination_bases, + source_layer_indices, ) ) - return CompactionPlan( - decision_rows=decision_rows, - pack_args=( - kept_ordinal_rows, - valid_seq_lens, - dense_move_offsets, - dense_move_indices, - swa_move_offsets, - swa_move_indices, - ), - pack_constexprs=dict( - KEEP_COUNT=keep_count, - DECISION_ROWS=decision_rows, - MOVE_CAPACITY=move_capacity, - NUM_KV_HEADS=num_kv_heads, - PER_LAYER=per_layer_sources, - DENSE_TOTAL=int(dense_move_indices.shape[-1]), - SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, - SWA_WINDOW=swa_window, - ), - launches=tuple(launches), - ) + return params def compact( - plans: Tuple[CompactionPlan, ...], + params: Tuple[CompactionParams, ...], request_count: int, ) -> None: - """Pack each plan's move sources and fire its native compacts, in plan order + """Pack each cache's move sources and fire its native compacts, in order (pure mover: the caller owns the decision rows and the round's completion ordering).""" - for plan in plans: - _pack_move_sources_kernel[(request_count, plan["decision_rows"])]( - *plan["pack_args"], **plan["pack_constexprs"] + for cache_params in params: + _pack_move_sources_kernel[(request_count, cache_params["decision_rows"])]( + *cache_params["pack_args"], **cache_params["pack_constexprs"] ) - for launch in plan["launches"]: - torch.ops.trtllm.sparse_kv_cache_compact_layers( - launch["pools"], - launch["pool_pointers"], - launch["page_table"], - launch["move_indices"], - launch["move_offsets"], - launch["destination_bases"], - launch["source_layer_indices"], - ) + for args in cache_params["compact_args"]: + torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 3bf903b6bb1d..8e627fb46de6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -38,7 +38,7 @@ ) from tensorrt_llm.logger import logger -from ..compaction import build_compaction_plan, compact +from ..compaction import build_compaction_params, compact from .triattention_kernels import ( _gather_mean_phase_kernel, _settle_ties_kernel, @@ -702,7 +702,7 @@ def _execute_eviction_round( ) self._settle_top_tokens(request_count) with nvtx_range("triattention.compact", color="purple"): - compact(self._compaction_plans, request_count) + compact(self._compaction_params, request_count) finally: # Order V2 page-table reuse and resize after this cohort's compact. self._compaction_done_event.record(stream) @@ -1367,8 +1367,8 @@ def _compiled_kernel(cache_key, build): ) # ---- compaction plans (opaque: only compact() interprets them) --------- - plans = [ - build_compaction_plan( + compaction_params = [ + build_compaction_params( layout, block_offsets=self._block_offsets_device, kept_ordinals=self._kept_ordinal_rows, @@ -1382,8 +1382,8 @@ def _compiled_kernel(cache_key, build): ) ] if draft is not None: - plans.append( - build_compaction_plan( + compaction_params.append( + build_compaction_params( draft["layout"], block_offsets=self._draft_block_offsets_device, kept_ordinals=self._kept_ordinal_rows, @@ -1393,7 +1393,7 @@ def _compiled_kernel(cache_key, build): protected_tail_capacity=int(draft["protected_tail_capacity"]), ) ) - self._compaction_plans = tuple(plans) + self._compaction_params = tuple(compaction_params) # ---- round-ordering events ---------------------------------------------- # Host staging (pinned metadata + snapshots) reuse fence. diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index a1ac14a03c3b..9a82ed188c6d 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -117,12 +117,12 @@ def make_ramp_pools( def build_compaction(**overrides): - """``build_compaction_plans`` with the suite's 2-layer defaults: + """``build_compaction_params`` with the suite's 2-layer defaults: allocates the caller-owned move-offset rows (capacity cumsum) and SWA destination bases, and hands the test's pre-settled ``kept_token_ordinals`` in as the decision rows. Returns the opaque - ``plans`` plus a test-side mirror of the caller-owned inputs.""" - from tensorrt_llm._torch.kv_cache_compression.compaction import build_compaction_plan + ``params`` plus a test-side mirror of the caller-owned inputs.""" + from tensorrt_llm._torch.kv_cache_compression.compaction import build_compaction_params args = dict( eviction_mode="union", @@ -154,8 +154,8 @@ def capacity_offsets(count): args.setdefault("swa_move_offsets", capacity_offsets(swa_window + tail) if has_swa else None) if has_draft: args.setdefault("draft_move_offsets", capacity_offsets(keep_count + draft_tail)) - plan_list = [ - build_compaction_plan( + params_list = [ + build_compaction_params( dict( layer_pools=args["layer_pools"], dense_layers=args["dense_layers"], @@ -174,8 +174,8 @@ def capacity_offsets(count): ) ] if has_draft: - plan_list.append( - build_compaction_plan( + params_list.append( + build_compaction_params( dict( layer_pools=args["draft_layer_pools"], dense_layers=args["draft_layers"], @@ -190,12 +190,12 @@ def capacity_offsets(count): protected_tail_capacity=draft_tail, ) ) - plans = tuple(plan_list) + params = tuple(params_list) # Opaque plans plus a test-side mirror of the caller-owned construction # inputs (production binds the same values as manager attributes); the # standalone helpers here need the move-offset rows and SWA staging back. return dict( - plans=plans, + params=params, prompt_offsets=args["prompt_offsets"], request_count=request_count, decode_keep_count=keep_count, @@ -213,7 +213,7 @@ def capacity_offsets(count): def run_compaction(compaction): """Replica of the round's move stage in production order: SWA - destination rebase, then ``compact`` loops the opaque plans (each packs + destination rebase, then ``compact`` loops the opaque params (each packs its decision rows into move sources and fires its native moves).""" from tensorrt_llm._torch.kv_cache_compression.compaction import compact @@ -223,7 +223,7 @@ def run_compaction(compaction): compaction["swa_rebase_delta"], out=compaction["swa_destination_bases"], ) - compact(compaction["plans"], compaction["request_count"]) + compact(compaction["params"], compaction["request_count"]) def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): 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 index dbd568af259b..d3e673629dca 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -177,7 +177,7 @@ def test_execute_eviction_round_orders_both_manager_streams(): tri._keep_count = 4 tri.eviction_mode = "union" tri._swa_window = None - tri._compaction_plans = () + tri._compaction_params = () tri._draft_protected_tail_capacity = 1 tri._staging_reuse_event = mock.Mock() tri._compaction_done_event = event From a2614bfb7d39313f04642ec0c9fa44503226dd7e Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 01:33:26 -0700 Subject: [PATCH 146/178] [None][refactor] Make CompactionParams a frozen dataclass Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 51d12b4a1e51..7ad808523ffd 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -20,18 +20,20 @@ ``compact`` packs every cache's move sources and fires its native launches. """ -from typing import Dict, List, Optional, Tuple, TypedDict +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple import torch import triton import triton.language as tl -class CompactionParams(TypedDict): +@dataclass(kw_only=True, frozen=True) +class CompactionParams: decision_rows: int pack_args: Tuple[Optional[torch.Tensor], ...] pack_constexprs: Dict[str, object] - compact_args: List[Tuple[object, ...]] + compact_args: List[Tuple[object, ...]] = field(default_factory=list) @triton.jit @@ -235,7 +237,7 @@ def build_compaction_params( dtype=torch.int32, device=device, ) - params["compact_args"].append( + params.compact_args.append( ( pools, torch.tensor( @@ -261,8 +263,8 @@ def compact( """Pack each cache's move sources and fire its native compacts, in order (pure mover: the caller owns the decision rows and the round's completion ordering).""" for cache_params in params: - _pack_move_sources_kernel[(request_count, cache_params["decision_rows"])]( - *cache_params["pack_args"], **cache_params["pack_constexprs"] + _pack_move_sources_kernel[(request_count, cache_params.decision_rows)]( + *cache_params.pack_args, **cache_params.pack_constexprs ) - for args in cache_params["compact_args"]: + for args in cache_params.compact_args: torch.ops.trtllm.sparse_kv_cache_compact_layers(*args) From ff8bb73d7f6a35f3877ff02e6a037f67e65f108b Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 01:38:43 -0700 Subject: [PATCH 147/178] [None][refactor] Fill CompactionParams progressively Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 54 +++++++++---------- 1 file changed, 26 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 7ad808523ffd..579c051f1b9c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -28,11 +28,11 @@ import triton.language as tl -@dataclass(kw_only=True, frozen=True) +@dataclass class CompactionParams: - decision_rows: int - pack_args: Tuple[Optional[torch.Tensor], ...] - pack_constexprs: Dict[str, object] + decision_rows: int = 0 + pack_args: Tuple[Optional[torch.Tensor], ...] = () + pack_constexprs: Dict[str, object] = field(default_factory=dict) compact_args: List[Tuple[object, ...]] = field(default_factory=list) @@ -136,14 +136,17 @@ def build_compaction_params( token_starts = dense_destination_bases protected_tail_capacity = int(protected_tail_capacity) + params = CompactionParams() first_pool = layer_pools[dense_layers[0]] device = first_pool.device max_requests = int(valid_seq_lens.shape[0]) keep_count = int(kept_ordinal_rows.shape[1]) - decision_rows = int(kept_ordinal_rows.shape[0]) // max_requests + params.decision_rows = int(kept_ordinal_rows.shape[0]) // max_requests # Pool shape [pages, K/V, heads, tokens, dim]. num_kv_heads = int(first_pool.shape[2]) - per_layer_sources = len(dense_layers) > 1 and decision_rows == len(dense_layers) * num_kv_heads + per_layer_sources = ( + len(dense_layers) > 1 and params.decision_rows == len(dense_layers) * num_kv_heads + ) dense_index_prefix = (len(dense_layers), num_kv_heads) if per_layer_sources else (num_kv_heads,) dense_move_indices = torch.empty( (*dense_index_prefix, (keep_count + protected_tail_capacity) * max_requests), @@ -193,28 +196,23 @@ def build_compaction_params( if swa_layers: move_capacity = max(move_capacity, swa_window + protected_tail_capacity) - params = CompactionParams( - decision_rows=decision_rows, - pack_args=( - kept_ordinal_rows, - valid_seq_lens, - dense_move_offsets, - dense_move_indices, - swa_move_offsets, - swa_move_indices, - ), - pack_constexprs=dict( - KEEP_COUNT=keep_count, - DECISION_ROWS=decision_rows, - MOVE_CAPACITY=move_capacity, - NUM_KV_HEADS=num_kv_heads, - PER_LAYER=per_layer_sources, - DENSE_TOTAL=int(dense_move_indices.shape[-1]), - SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, - SWA_WINDOW=swa_window, - ), - # One positional-args tuple per native compact call, in its op signature order. - compact_args=[], + params.pack_args = ( + kept_ordinal_rows, + valid_seq_lens, + dense_move_offsets, + dense_move_indices, + swa_move_offsets, + swa_move_indices, + ) + params.pack_constexprs = dict( + KEEP_COUNT=keep_count, + DECISION_ROWS=params.decision_rows, + MOVE_CAPACITY=move_capacity, + NUM_KV_HEADS=num_kv_heads, + PER_LAYER=per_layer_sources, + DENSE_TOTAL=int(dense_move_indices.shape[-1]), + SWA_TOTAL=int(swa_move_indices.shape[-1]) if swa_move_indices is not None else 0, + SWA_WINDOW=swa_window, ) for entries, move_indices, move_offsets, destination_bases, slots in move_groups: grouped = {} From 18c9695c57a37d219cad9aea7f0ea268596caa18 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 02:10:59 -0700 Subject: [PATCH 148/178] [None][chore] Order compaction params below the pack kernel Signed-off-by: tianruih --- .../_torch/kv_cache_compression/compaction.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 579c051f1b9c..d1bea53de529 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -28,14 +28,6 @@ import triton.language as tl -@dataclass -class CompactionParams: - decision_rows: int = 0 - pack_args: Tuple[Optional[torch.Tensor], ...] = () - pack_constexprs: Dict[str, object] = field(default_factory=dict) - compact_args: List[Tuple[object, ...]] = field(default_factory=list) - - @triton.jit def _pack_move_sources_kernel( kept_ordinal_rows, @@ -113,6 +105,14 @@ def _pack_move_sources_kernel( ) +@dataclass +class CompactionParams: + decision_rows: int = 0 + pack_args: Tuple[Optional[torch.Tensor], ...] = () + pack_constexprs: Dict[str, object] = field(default_factory=dict) + compact_args: List[Tuple[object, ...]] = field(default_factory=list) + + def build_compaction_params( layout: Dict[str, object], *, From cd50a4e413fd4b14b95356dc7043fa333c259b8f Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 02:32:25 -0700 Subject: [PATCH 149/178] [None][chore] Simplify the compaction kernel comments Signed-off-by: tianruih --- .../unfusedAttentionKernels_2_template.h | 98 +++++-------------- 1 file changed, 23 insertions(+), 75 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h index ceab9ac55c42..22e17b438d58 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h @@ -1758,40 +1758,17 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams #ifdef ENABLE_BF16 -// Pipelined bf16 compaction kernels, ported from Fanrong Li's optimized -// compact kernels (snapshot 2026-07-19). The port keeps the double-buffered -// cp.async pipeline intact and adapts only the addressing to this -// repository's KVCacheManagerV2 ABI: -// (a) the per-layer page-table pointer array became one flat int32 V2 -// K-plane block-offset table shared by all layers (entries encode -// 2 * page + plane with plane == 0, so >> 1 recovers the page), strided -// per request; -// (b) the host-scalar destination base became per-request device bases, read -// once per CTA (one launch covers a cohort with mixed prompt lengths); -// (c) the head-row stride of the move-source indices is an explicit -// parameter instead of being derived from sourceOffsets[batchSize] on -// device: the move buffers are allocation-wide, so a device-derived -// stride would silently read the wrong plane for every KV head above -// head 0. -// The original kernels were written for 128-token pages and Dh = 64; the port -// additionally parameterizes the page and head-vector math so 32-token pages -// and Dh = 128 (the production geometry here) take the same pipeline. +// Pipelined bf16 compaction kernels: double-buffered cp.async copies that +// compact paged KV pools in place through the V2 block-offset table. namespace compact_detail { -// Vendored cp.async wrappers, equivalent to the ones in -// cpp/kernels/xqa/ldgsts.cuh. That header cannot be included from this -// widely-included template header because it drags in xqa's cuda_hint.cuh / -// barriers.cuh, whose macros and helpers would leak into every translation -// unit that includes this file. +// Vendored cp.async wrappers (xqa's ldgsts.cuh cannot be included here). template __device__ __forceinline__ void copyAsync(void* dst, void const* src, uint32_t srcSize = size) { static_assert(size == 16, "only the 16B cp.async variant is vendored"); - // srcSize == 0 turns the copy into a shared-memory zero fill; predicated - // lanes use it so ragged tiles never touch global memory. Nulling src is - // the same workaround as the xqa original, which observed speculative - // global reads without it. + // srcSize == 0 zero-fills shared memory instead of reading global. if (srcSize == 0) { src = nullptr; @@ -1835,11 +1812,8 @@ struct SparseKvCacheCompactBf16Params size_t bytesPerPage; }; -//! Double-buffered cp.async pipeline: while the current 32-token tile drains -//! from shared memory into its destination pages, the next tile's K/V vectors -//! are already streaming global -> shared into the other buffer. One CTA per -//! (layer, KV head, request); threadIdx.x walks the 16B vectors of one head, -//! threadIdx.y walks the tokens of a tile. +//! Double-buffered cp.async pipeline: the next tile streams in while the +//! current tile drains. One CTA per (layer, KV head, request). template __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) * kSparseKvCompactTokensPerTile) void sparseKvCacheCompactV2Bf16PipelineKernel(SparseKvCacheCompactBf16Params @@ -1850,8 +1824,6 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) // 128-token pages are the geometry the kernel was written for; 32-token // pages cover the supported production configuration (one tile == one page). static_assert(TokensPerBlock == 32 || TokensPerBlock == 128); - // 16B vectors per head: Dh64 -> 8 lanes (block 8x32 = 256 threads), - // Dh128 -> 16 lanes (block 16x32 = 512 threads). constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); constexpr int32_t kTokensPerTile = kSparseKvCompactTokensPerTile; constexpr int32_t kVectorsPerTile = kTokensPerTile * kVectorsPerHead; @@ -1869,22 +1841,13 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) return; } - // Layer resolution rule shared with the packed move-source layout: - // without an explicit map, launch layer i reads source plane i (the flat - // layout passes sourceLayerStride == 0, which collapses the term). + // Without an explicit layer map, launch layer i reads source plane i. int32_t const sourceLayer = params.sourceLayerIndices == nullptr ? layerIdx : params.sourceLayerIndices[layerIdx]; - // ABI adaptation (c) -- see the port note above the compact_detail - // namespace: head rows are strided by the allocation width - // of the move buffers; this request's range within a plane starts at - // moveBegin. int64_t const sourceMoveBase = static_cast(sourceLayer) * params.sourceLayerStride + static_cast(kvHeadIdx) * params.sourceHeadStride + moveBegin; - // ABI adaptation (b) -- see the port note above the compact_detail namespace: per-request landing position. int32_t const destinationBase = params.destinationBases[batchIdx]; auto* const pool = reinterpret_cast(static_cast(params.poolPointers[layerIdx])); - // ABI adaptation (a) -- see the port note above the compact_detail namespace: flat V2 K-plane block-offset table; each lookup - // below decodes an entry to a page with >> 1. TokensPerBlock is a - // compile-time power of two, so / and % lower to shifts and masks. + // Block-offset entries decode to a page with >> 1. int32_t const* const pageTable = params.pageTable + static_cast(batchIdx) * params.pageTableRequestStride; extern __shared__ uint4 sharedVectors[]; @@ -1945,10 +1908,8 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) compact_detail::copyAsync(&nextSharedV[sharedVector], nextSourceVVector, nextSourceBytes); compact_detail::commitGroup(); - // The compaction contract provides strictly increasing sources per request/head and - // dst(i) = destinationBase + i <= src(i). For current i and future j, i < j implies - // dst(i) < destinationBase + j <= src(j), so current stores cannot alias future prefetch sources. - // The current tile itself completed its wait and CTA barrier before reaching this store phase. + // In-place safety: sources are strictly increasing with dst(i) <= src(i), + // so current stores never alias future prefetch sources. int32_t const destinationToken = destinationBase + currentRequestMove; if (currentValid && currentSourceToken != destinationToken) { @@ -1963,8 +1924,7 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) destinationV[localVector] = currentSharedV[sharedVector]; } - // commitGroup only closes the group; waitGroup<0> completes this thread's next-tile copies. The CTA - // barrier then makes the ping-pong buffer visible to all threads before it becomes current. + // waitGroup<0> completes the next tile's copies before the buffer swap. compact_detail::waitGroup<0>(); __syncthreads(); currentRequestMove = nextRequestMove; @@ -1989,26 +1949,18 @@ __global__ __launch_bounds__(HeadDim * sizeof(T) / sizeof(uint4) } } -// A destination-page-staging variant (needs a host-side proof that every -// destination base is tile-aligned) is parked on branch -// tr-parked-destination-page-kernel pending compaction.py alignment-flag plumbing. template void launchSparseKvCacheCompactV2Bf16Pipeline(SparseKvCacheCompactBf16Params const& params, cudaStream_t stream) { constexpr int32_t kVectorsPerHead = HeadDim * sizeof(T) / sizeof(uint4); dim3 const block(kVectorsPerHead, kSparseKvCompactTokensPerTile); dim3 const grid(params.numLayers, params.numKvHeads, params.batchSize); - // Two ping-pong buffers x (K tile + V tile) of 32 tokens x kVectorsPerHead - // 16B vectors: - // Dh64: 4 * 32 * 8 * 16 B = 16 KiB - // Dh128: 4 * 32 * 16 * 16 B = 32 KiB - // Both fit the 48 KiB per-CTA dynamic shared memory default, so no - // cudaFuncSetAttribute opt-in is required. + // Two ping-pong buffers of (K + V) tiles; both geometries fit the 48 KiB + // per-CTA dynamic shared memory default. size_t const sharedBytes = 4 * kSparseKvCompactTokensPerTile * kVectorsPerHead * sizeof(uint4); sparseKvCacheCompactV2Bf16PipelineKernel<<>>(params); } - #endif // ENABLE_BF16 template @@ -2022,7 +1974,6 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( int const batch_idx = blockIdx.z; int const kv_head_idx = blockIdx.y; - // Head-row stride of the packed indices. int const total_num_sparse_kv_tokens = params.sparse_kv_offsets[params.batch_size]; int const sparse_start_idx = params.sparse_kv_offsets[batch_idx]; @@ -2044,6 +1995,7 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( { int const global_sparse_idx = sparse_start_idx + sparse_token_offset; int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; + int const src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; void* src_k_ptr = params.kv_cache_buffer.getKBlockPtr(batch_idx, src_token_idx); @@ -2057,6 +2009,7 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); auto const src_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx(src_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_k_block_ptr[src_k_vec_idx]; v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx] = src_v_block_ptr[src_v_vec_idx]; } @@ -2067,6 +2020,7 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( { int const global_sparse_idx = sparse_start_idx + sparse_token_offset; int const sparse_idx_offset = kv_head_idx * total_num_sparse_kv_tokens + global_sparse_idx; + int const src_token_idx = params.sparse_kv_indices[sparse_idx_offset]; int const dst_token_idx = sparse_token_offset; @@ -2079,10 +2033,10 @@ __global__ __launch_bounds__(BLOCK_SIZE) void updateSparseKvCacheAfterFmha( for (int head_vec_idx = threadIdx.x; head_vec_idx < VECS_PER_HEAD; head_vec_idx += vecs_per_block) { - auto const dst_k_vec_idx = params.kv_cache_buffer.getKVLocalIdx( - dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); - auto const dst_v_vec_idx = params.kv_cache_buffer.getKVLocalIdx( - dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + auto const dst_k_vec_idx + = params.kv_cache_buffer.getKVLocalIdx(dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); + auto const dst_v_vec_idx + = params.kv_cache_buffer.getKVLocalIdx(dst_token_idx, kv_head_idx, VECS_PER_HEAD, head_vec_idx); dst_k_block_ptr[dst_k_vec_idx] = k_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; dst_v_block_ptr[dst_v_vec_idx] = v_smem[threadIdx.y * VECS_PER_HEAD + head_vec_idx]; } @@ -2106,8 +2060,7 @@ void kernelSparseDispatchHeadSize(QKVPreprocessingParams param // grid.x is always 1 to avoid data races dim3 grid(1, params.kv_head_num, params.batch_size); - updateSparseKvCacheAfterFmha - <<>>(params); + updateSparseKvCacheAfterFmha<<>>(params); } template @@ -2142,11 +2095,6 @@ void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const #ifdef ENABLE_BF16 if constexpr (std::is_same_v) { - // The pipelined kernels are the only shipped path: they won the A/B - // comparison against the retired register-staging kernel everywhere - // (verified 2026-07-20: 1.47x at batch 1, 1.09-1.30x at batch 32, - // 1.09-1.13x at batch 256, byte-identical outputs). Unsupported pool - // dtypes and geometries fail the check below instead of falling back. if ((headDim == 64 || headDim == 128) && (tokensPerBlock == 32 || tokensPerBlock == 128)) { SparseKvCacheCompactBf16Params fastParams{}; @@ -2206,8 +2154,8 @@ void invokeSparseKvCacheCompactLayers(int64_t const* poolPointers, int32_t const QKVPreprocessingParams params, cudaStream_t stream); \ //////////////////////////////////////////////////////////////////////////////////////////////////// -#define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_LAYERS(T) \ - template void invokeSparseKvCacheCompactLayers(int64_t const*, int32_t const*, int32_t, int64_t, \ +#define INSTANTIATE_SPARSE_KV_CACHE_COMPACT_LAYERS(T) \ + template void invokeSparseKvCacheCompactLayers(int64_t const*, int32_t const*, int32_t, int64_t, \ int32_t const*, int32_t const*, int64_t, int64_t, int32_t const*, int32_t const*, int32_t, int32_t, int32_t, \ int32_t, cudaStream_t); From 1bad4ac28ee3dc38412cdd4bf0aa67ad41901555 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 09:35:05 -0700 Subject: [PATCH 150/178] [None][chore] Sync review remediation from compaction PR branch Signed-off-by: tianruih --- tensorrt_llm/_torch/kv_cache_compression/compaction.py | 1 + tensorrt_llm/_torch/pyexecutor/model_engine.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index d1bea53de529..5d7ef1765262 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -260,6 +260,7 @@ def compact( ) -> None: """Pack each cache's move sources and fire its native compacts, in order (pure mover: the caller owns the decision rows and the round's completion ordering).""" + # One launch per (cache, pool group); each launch covers every layer in the group. for cache_params in params: _pack_move_sources_kernel[(request_count, cache_params.decision_rows)]( *cache_params.pack_args, **cache_params.pack_constexprs diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 33309c76fef7..3795c479bf0d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4941,7 +4941,6 @@ def previous_seq_slots_device(): and not multimodal_params_list and not lora_params and attn_metadata.padded_num_tokens is None and self._get_position_id_offset() == 0 - # KV compression shrinks the cache mid-generation; take the full prepare path. and not getattr(kv_cache_manager, "kv_compression_manages_history", False)): self._steady_gen_positions_pinned[:_n_gen].copy_( From 597ac68b1f735e32dc752c4657ebd90b9fb3bf31 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 10:21:19 -0700 Subject: [PATCH 151/178] [None][fix] Close three eviction-runtime lifetime gaps: score-capacity reuse gate, staged block width clamp, event epoch reuse Signed-off-by: tianruih --- .../triattention/triattention.py | 44 +++++++++---- .../_torch/kv_cache_compression/conftest.py | 15 ++++- .../test_triattention_draft_cocompaction.py | 64 +++++++++++++++++++ 3 files changed, 110 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 8e627fb46de6..8266bcb2666e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -70,13 +70,16 @@ def _allocate_block_offset_staging( num_pools: int, max_requests: int, token_capacity: int, + max_source_blocks: int, ) -> Tuple[torch.Tensor, torch.Tensor]: """One pinned host snapshot + persistent device table pair in the native V2 ``[pool, request, K/V, block]`` layout (block width 4-aligned for the - ``PackedInt`` copy ABI); the device follows the anchor KV pool.""" + ``PackedInt`` copy ABI); the device follows the anchor KV pool. The staged + width is clamped to the manager's live source-table width: static bucket + slack never holds valid tokens and the native gather copies the full width.""" tokens_per_block = int(anchor_pool.shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block - staged_blocks_per_seq = (page_count + 3) // 4 * 4 + staged_blocks_per_seq = min((page_count + 3) // 4 * 4, int(max_source_blocks)) shape = (num_pools, max_requests, 2, staged_blocks_per_seq) host = torch.empty(shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) device_table = torch.empty(shape, dtype=torch.int32, device=anchor_pool.device) @@ -163,9 +166,14 @@ def __init__( + 1 ) self._generation_growth = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) - # Buffers build once at the first eviction as plain attributes on this - # manager and stay resident for the manager's lifetime. + # Lazy resident eviction runtime: built at the first eviction, reused + # across rounds, and replaced as a whole when a capacity axis grows. self._buffers_built = False + # Round-ordering events: device-lifetime, created at the first build + # and reused across capacity rebuilds. + self._staging_reuse_event: Optional[torch.cuda.Event] = None + self._block_offsets_ready_event: Optional[torch.cuda.Event] = None + self._compaction_done_event: Optional[torch.cuda.Event] = None # 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] @@ -827,6 +835,7 @@ def _ensure_buffers( ) -> None: # Empty cohorts never reach here: _periodic_evict no-ops pre-launch. needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) + needed_score_tokens = max(item["seq_len"] for item in prepared) needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) needed_requests = len(prepared) if self.draft_kv_cache_manager is not None: @@ -836,11 +845,15 @@ def _ensure_buffers( if self._buffers_built: if ( needed_width <= self._decode_width + and needed_score_tokens <= self._bucket_seq_len and needed_page_tokens <= self._page_table_token_capacity and needed_requests <= self._max_requests ): return - # This round outgrew the buffers: rebuild. + # This round outgrew the buffers: wait out the prior round (its + # completion event orders after every use of the old epoch), then rebuild. + if self._compaction_done_event is not None: + self._compaction_done_event.synchronize() self._buffers_built = False mgr = self.kv_cache_manager @@ -970,6 +983,7 @@ def _build_buffers( num_pools=num_page_table_slots, max_requests=max_requests, token_capacity=page_table_token_capacity, + max_source_blocks=int(layout["manager"].host_kv_cache_block_offsets.shape[-1]), ) # The draft is never scored: these offsets feed only the draft compacts. self._draft_block_offsets_device = None @@ -990,6 +1004,9 @@ def _build_buffers( num_pools=int(draft_layout["manager"].num_pools), max_requests=max_requests, token_capacity=int(draft["page_table_token_capacity"]), + max_source_blocks=int( + draft_layout["manager"].host_kv_cache_block_offsets.shape[-1] + ), ) ) @@ -1396,13 +1413,16 @@ def _compiled_kernel(cache_key, build): self._compaction_params = tuple(compaction_params) # ---- round-ordering events ---------------------------------------------- - # Host staging (pinned metadata + snapshots) reuse fence. - self._staging_reuse_event = torch.cuda.Event() - self._staging_reuse_event.record(torch.cuda.current_stream(device)) - # Manager-stream H2D of the block-offset tables has completed. - self._block_offsets_ready_event = torch.cuda.Event() - # This cohort's compact is done: manager may resize/reuse pages. - self._compaction_done_event = torch.cuda.Event() + # Device-lifetime: created once and reused across capacity rebuilds + # (they carry no pointer state; replacing them orphans in-flight ordering). + if self._staging_reuse_event is None: + # Host staging (pinned metadata + snapshots) reuse fence. + self._staging_reuse_event = torch.cuda.Event() + self._staging_reuse_event.record(torch.cuda.current_stream(device)) + # Manager-stream H2D of the block-offset tables has completed. + self._block_offsets_ready_event = torch.cuda.Event() + # This cohort's compact is done: manager may resize/reuse pages. + self._compaction_done_event = torch.cuda.Event() def _local_score_calibration( self, diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 9a82ed188c6d..0a27ebf26f23 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -288,6 +288,7 @@ def make_buffer_stubs(manager, *, decode_width=260): ) built_attributes = dict( _decode_width=decode_width, + _bucket_seq_len=1024, _page_table_token_capacity=65537, _max_requests=8, _token_starts_device=torch.zeros(8, dtype=torch.int32), @@ -551,8 +552,17 @@ def make_cute_buffers( storage_groups = {0: list(range(num_layers))} if layer_pool_ids is None: layer_pool_ids = [0] * num_layers + # A live-manager source table exactly wide enough for the requested + # capacity (the staged width clamps to it). + requested_tokens = seq_len if page_table_token_capacity is None else page_table_token_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( - manager=SimpleNamespace(num_pools=max(layer_pool_ids) + 1), + manager=SimpleNamespace( + num_pools=max(layer_pool_ids) + 1, + host_kv_cache_block_offsets=torch.empty(1, 1, 2, source_blocks, dtype=torch.int32), + ), layer_pools=layer_pools, dense_layers=list(range(num_layers)), swa_layers=[], @@ -569,6 +579,9 @@ def make_cute_buffers( manager._draft_protected_tail_capacity = None manager.eviction_mode = eviction_mode manager.normalize_scores = normalize_scores + manager._staging_reuse_event = None + manager._block_offsets_ready_event = None + manager._compaction_done_event = None manager._phase = make_phase_table(offsets, omega, seq_len) manager._build_buffers( layout=layout, 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 index d3e673629dca..7f792b784f58 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -429,3 +429,67 @@ def apply_rebuilt(**kwargs): assert prepare.call_count == 2 assert manager._buffers_built assert manager._decode_width == built_attributes["_decode_width"] + 8 + + +def test_source_growth_beyond_score_bucket_rebuilds_buffers(): + """A later cohort can grow max(source_length) past the compiled score + bucket while decode width, page tokens, and request count all still fit; + the reuse gate must rebuild instead of scoring past the static geometry.""" + manager = _make_triattention(budget=4) + layout, built_attributes = _make_buffer_stubs(manager) + prepared = [ + _make_prepared_item( + _make_request(7), + request_id=7, + seq_len=1024, + prompt_len=1020, + expected_keep_count=4, + ) + ] + + def apply_built(**kwargs): + for name, value in built_attributes.items(): + setattr(manager, name, value) + + with mock.patch.object(manager, "_build_buffers", side_effect=apply_built) as prepare: + manager._ensure_buffers(layout, prepared) + assert prepare.call_count == 1 + # One more source token, same decode width and request count. + grown = [ + _make_prepared_item( + _make_request(7), + request_id=7, + seq_len=1025, + prompt_len=1021, + expected_keep_count=4, + ) + ] + manager._ensure_buffers(layout, grown) + assert prepare.call_count == 2 + + +def test_staged_block_width_clamps_to_manager_source_width(): + """Score tile rounding can request more page-table blocks than the live V2 + source table holds; the staged width must clamp to the manager width so the + native gather never reads past the K plane (tpb=32, max_seq_len=96, tail=1).""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _allocate_block_offset_staging, + ) + + anchor_pool = torch.empty(1, 2, 1, 32, 4) + host, device_table = _allocate_block_offset_staging( + anchor_pool, + num_pools=1, + max_requests=2, + token_capacity=129, + max_source_blocks=4, + ) + assert host.shape[-1] == 4 and device_table.shape[-1] == 4 + host, device_table = _allocate_block_offset_staging( + anchor_pool, + num_pools=1, + max_requests=2, + token_capacity=129, + max_source_blocks=64, + ) + assert host.shape[-1] == 8 and device_table.shape[-1] == 8 From a626ae0c3dde4c4aa9bed7474a54b3205e838b43 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 20:56:09 -0700 Subject: [PATCH 152/178] [None][refactor] Straighten the eviction control plane into one transaction owner Signed-off-by: tianruih --- .../triattention/triattention.py | 363 ++++++++---------- .../_torch/kv_cache_compression/conftest.py | 78 ++-- .../test_triattention_draft_cocompaction.py | 108 +++--- .../test_triattention_pipeline.py | 93 +++-- .../test_triattention_selection_compaction.py | 23 +- 5 files changed, 301 insertions(+), 364 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 8266bcb2666e..d2fe0637e6fc 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -32,6 +32,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheCompressionManager +from tensorrt_llm._torch.utils import next_positive_power_of_2 from tensorrt_llm._utils import nvtx_range, nvtx_range_debug, prefer_pinned from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( copy_batch_block_offsets_to_device, @@ -150,8 +151,8 @@ def __init__( # Per-request eviction progress. self._request_states: Dict[int, Dict[str, object]] = {} # In-flight overlap batch reference; membership resolves lazily. - self._prepared_generation_batch: Optional[object] = None - self._prepared_generation_ids: Optional[set] = None + self._inflight_scheduled_batch: Optional[object] = None + self._inflight_generation_request_ids: Optional[set] = None # Manager-lifetime constants. self._protected_tail_capacity = ( int(kv_cache_manager.num_extra_kv_tokens) @@ -378,15 +379,15 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: } def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Snapshot the prepared batch; mutation remains in final update.""" - self._prepared_generation_batch = scheduled_batch - self._prepared_generation_ids = None + """Snapshot the in-flight batch; mutation remains in final update.""" + self._inflight_scheduled_batch = scheduled_batch + self._inflight_generation_request_ids = None def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Compact after native KV-cache updates have finalized this iteration (must run after KVCacheManagerV2 so capacity reflects the written token and any rewind).""" with nvtx_range_debug("triattention.generation_step_end", color="blue"): - self._periodic_evict(scheduled_batch) + self._evict_due_requests(scheduled_batch) def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: """Drop this request's eviction state; the buffers stay resident.""" @@ -427,140 +428,121 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: # ---- eviction round ---- - def _periodic_evict( + def _evict_due_requests( self, scheduled_batch: "ScheduledRequests", ) -> None: - gen_requests = scheduled_batch.generation_requests - if not gen_requests: - return - mgr = self.kv_cache_manager - resolved_requests = [] - for request in gen_requests: - if request.is_dummy or request.state in _SKIP_REQUEST_STATES: - continue - request_id = request.py_request_id - kv_cache = mgr.kv_cache_map.get(request_id) - if kv_cache is None: - continue - if not kv_cache.is_active: - # Overlap scheduling may suspend a cache mid-flight; defer this - # request (pre-launch) instead of failing the whole batch. - continue - resolved_requests.append((request, request_id, kv_cache)) - if not resolved_requests: - return - prepared: List[Dict[str, object]] = [] - - # The resolved cache objects thread all the way to resize. + """Owner of the full eviction transaction: admission, cadence, launch, + publication, and cache resize.""" + manager = self.kv_cache_manager + eviction_inputs: List[Dict[str, object]] = [] with nvtx_range("triattention.metadata", color="cyan"): - for request, request_id, kv_cache in resolved_requests: + for request in scheduled_batch.generation_requests: + if request.is_dummy or request.state in _SKIP_REQUEST_STATES: + 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 # Cadence gate first; capacity math and consistency raises run in the due branch. - request_state = self._request_states[request_id] - previous_step = request_state["generation_steps"] + state = self._request_states[request_id] + previous_step = state["generation_steps"] step = previous_step + 1 + int(request.py_num_accepted_draft_tokens) - request_state["generation_steps"] = step + state["generation_steps"] = step if previous_step // self.beta >= step // self.beta: continue - raw_capacity = int(kv_cache.capacity) # Speculative reserve + in-flight overlap growth: contiguous tail moved byte-for-byte. - protected_tail = int(mgr.num_extra_kv_tokens) + self._inflight_generation_growth( - scheduled_batch, request_id - ) - seq_len = raw_capacity - protected_tail - if seq_len < kv_cache.history_length: + target_tail_length = int( + manager.num_extra_kv_tokens + ) + self._inflight_generation_growth(scheduled_batch, request_id) + source_length = int(target_cache.capacity) - target_tail_length + if source_length < target_cache.history_length: raise RuntimeError( - f"Request {request_id} KV length {seq_len} is below finalized " - f"history {kv_cache.history_length}" + f"Request {request_id} KV length {source_length} is below " + f"finalized history {target_cache.history_length}" ) - expected_keep_count = self._minimum_evictable_length(request, seq_len) - if seq_len <= expected_keep_count: + prompt_length = min(int(request.py_prompt_len), source_length) + if source_length <= prompt_length + self.budget: + # Selection would be an identity: nothing to evict yet. continue - draft_kv_cache = None + draft_cache = None if self.draft_kv_cache_manager is not None: # A missing draft cache is a wiring bug: the dict's KeyError is the report. - draft_kv_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] - if not draft_kv_cache.is_active: + draft_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] + if not draft_cache.is_active: # Target and draft defer together (pre-launch). continue - prepared.append( + eviction_inputs.append( { "request": request, - "request_id": request_id, - "kv_cache": kv_cache, - "draft_kv_cache": draft_kv_cache, - "seq_len": int(seq_len), + "target_cache": target_cache, + "draft_cache": draft_cache, + "source_length": source_length, # Uncompressed logical position. - "round_start": int(seq_len + request_state["evicted_tokens"]), - "prompt_len": min(int(request.py_prompt_len), int(seq_len)), - "expected_keep_count": expected_keep_count, - "protected_tail": protected_tail, + "logical_source_length": source_length + state["evicted_tokens"], + "prompt_length": prompt_length, + "target_tail_length": target_tail_length, } ) - - if not prepared: + if not eviction_inputs: return + + with nvtx_range_debug("triattention.resolve_layout", color="blue"): + target_layout = self._runtime_kv_layout() + draft_layout = ( + self._runtime_kv_layout(draft=True) + if self.draft_kv_cache_manager is not None + else None + ) + with nvtx_range_debug("triattention.staging_lookup", color="blue"): + # Retained spans always cover the model window (construction rejects budget < window). + self._ensure_eviction_runtime(target_layout, draft_layout, eviction_inputs) # Ungated NVTX: the due count in the message shows each round's size. with nvtx_range( - f"triattention.evict_request_group reqs={len(prepared)}", + f"triattention.evict_request_group reqs={len(eviction_inputs)}", color="purple", ): - compacted = self._evict_requests(prepared) - self._resize_compacted_requests(compacted) + self._execute_eviction_round(eviction_inputs) + for item in eviction_inputs: + request = item["request"] + evicted = item["source_length"] - item["prompt_length"] - self.budget + state = self._request_states[request.py_request_id] + state["evicted_tokens"] += evicted + # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). + request.py_num_compressed_tokens = state["evicted_tokens"] + self._resize_compacted_caches(eviction_inputs) def _inflight_generation_growth( self, scheduled_batch: "ScheduledRequests", request_id: int ) -> int: - prepared = self._prepared_generation_batch - if prepared is None or scheduled_batch is prepared: + inflight = self._inflight_scheduled_batch + if inflight is None or scheduled_batch is inflight: return 0 - member_ids = self._prepared_generation_ids + member_ids = self._inflight_generation_request_ids if member_ids is None: - member_ids = {request.py_request_id for request in prepared.generation_requests} - self._prepared_generation_ids = member_ids + member_ids = {request.py_request_id for request in inflight.generation_requests} + self._inflight_generation_request_ids = member_ids if request_id not in member_ids: return 0 return self._generation_growth - def _minimum_evictable_length(self, request: "LlmRequest", seq_len: int) -> int: - """Return the largest cache length for which selection is an identity.""" - prompt_len = min(int(request.py_prompt_len), seq_len) - return prompt_len + self.budget - - def _evict_requests( - self, - prepared: List[Dict[str, object]], - ) -> List[Dict[str, object]]: - with nvtx_range_debug("triattention.resolve_layout", color="blue"): - layout = self._runtime_kv_layout() - with nvtx_range_debug("triattention.staging_lookup", color="blue"): - # Retained spans always cover the model window (construction rejects budget < window). - self._ensure_buffers(layout, prepared) - self._execute_eviction_round(prepared) - for item in prepared: - # Identity cohorts were filtered pre-launch (_periodic_evict). - evicted = item["seq_len"] - item["expected_keep_count"] - request_state = self._request_states[item["request_id"]] - request_state["evicted_tokens"] += evicted - # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). - item["request"].py_num_compressed_tokens = request_state["evicted_tokens"] - return prepared - def _execute_eviction_round( self, - prepared: Sequence[Dict[str, object]], + eviction_inputs: Sequence[Dict[str, object]], ) -> None: - """Run one eviction round over the prepared cohort (every launch covers + """Run one eviction round over the due cohort (every launch covers the full request capacity; padded rows carry zero lengths and stay inert).""" manager = self.kv_cache_manager draft_manager = self.draft_kv_cache_manager with nvtx_range_debug("triattention.page_table_stage", color="orange"): - request_ids = [item["request_id"] for item in prepared] - round_starts = [item["round_start"] for item in prepared] - token_starts = [item["prompt_len"] for item in prepared] - seq_lens = [item["seq_len"] for item in prepared] - dense_move_offsets, swa_move_offsets, draft_move_offsets = self._cohort_move_offsets( - prepared + request_ids = [item["request"].py_request_id for item in eviction_inputs] + round_starts = [item["logical_source_length"] for item in eviction_inputs] + token_starts = [item["prompt_length"] for item in eviction_inputs] + seq_lens = [item["source_length"] for item in eviction_inputs] + dense_move_offsets, swa_move_offsets, draft_move_offsets = ( + self._compute_compaction_move_offsets(eviction_inputs) ) stream = torch.cuda.current_stream(self._block_offsets_device.device) # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. @@ -586,7 +568,7 @@ def _execute_eviction_round( if values is not None: host_table[row, : len(values)] = values # Zero lengths keep the score kernel and selection inert for padded rows. - host_table[:3, len(prepared) :] = 0 + host_table[:3, len(eviction_inputs) :] = 0 grow_mean_phase_table(self._phase, int(max_round_start) + 1) self._stage_block_offsets( manager, @@ -606,7 +588,7 @@ def _execute_eviction_round( finally: # Guards the pinned staging until the asynchronous copies complete. self._staging_reuse_event.record(stream) - request_count = len(prepared) + request_count = len(eviction_inputs) union = self.eviction_mode == "union" try: with nvtx_range("triattention.score", color="blue"): @@ -718,11 +700,11 @@ def _execute_eviction_round( if draft_manager is not None: draft_manager._stream.wait_event(self._compaction_done_event) - def _cohort_move_offsets( + def _compute_compaction_move_offsets( self, - prepared: Sequence[Dict[str, object]], + eviction_inputs: Sequence[Dict[str, object]], ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - """Cumulative dense/SWA/draft move offsets for one prepared cohort (keep + """Cumulative dense/SWA/draft move offsets for one due cohort (keep set plus protected tail per request; rows past the cohort repeat the final offset and contribute no moves).""" @@ -733,7 +715,7 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets.extend(offsets[-1:] * (self._max_requests - len(moves_per_request))) return offsets - tails = [int(item["protected_tail"]) for item in prepared] + tails = [int(item["target_tail_length"]) for item in eviction_inputs] dense = padded_offsets([self._keep_count + tail for tail in tails]) swa = None if self._swa_window is not None: @@ -741,7 +723,7 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: draft = None if self._draft_protected_tail_capacity is not None: draft = padded_offsets( - [self._keep_count + self._draft_protected_tail_capacity] * len(prepared) + [self._keep_count + self._draft_protected_tail_capacity] * len(eviction_inputs) ) return dense, swa, draft @@ -797,92 +779,81 @@ def _settle_top_tokens(self, request_count: int) -> None: SELECTION_ROWS=self._selection_rows_per_request, ) - def _resize_compacted_requests(self, prepared) -> None: - if not prepared: - return + def _resize_compacted_caches(self, eviction_inputs) -> None: with nvtx_range("triattention.resize", color="red"): with nvtx_range_debug("triattention.v2_resize", color="red"): - families = [("target", "kv_cache", None)] + families = [("target", "target_cache", None)] if self.draft_kv_cache_manager is not None: # Same kept set: the draft shrinks to the same retained # length plus its own fixed tail. - families.append( - ("draft", "draft_kv_cache", self._draft_protected_tail_capacity) - ) + families.append(("draft", "draft_cache", self._draft_protected_tail_capacity)) for label, cache_key, fixed_tail in families: - for item in prepared: - kv_cache = item[cache_key] - if not kv_cache.is_active: + for item in eviction_inputs: + cache = item[cache_key] + request_id = item["request"].py_request_id + if not cache.is_active: # Bytes already moved: the compact-to-resize window is owned by this hook. raise RuntimeError( - f"Request {item['request_id']} {label} KV cache was " + f"Request {request_id} {label} KV cache was " "suspended between compact and resize" ) - tail = item["protected_tail"] if fixed_tail is None else fixed_tail - resized_capacity = item["expected_keep_count"] + tail - if not kv_cache.resize(resized_capacity, None): + tail = item["target_tail_length"] if fixed_tail is None else fixed_tail + resized_capacity = item["prompt_length"] + self.budget + tail + if not cache.resize(resized_capacity, None): raise RuntimeError( f"Failed to resize compacted {label} KV cache for " - f"request {item['request_id']} to {resized_capacity} tokens" + f"request {request_id} to {resized_capacity} tokens" ) # ---- buffers + layout ---- - def _ensure_buffers( + def _ensure_eviction_runtime( self, - layout: Dict[str, object], - prepared: Sequence[Dict[str, object]], + target_layout: Dict[str, object], + draft_layout: Optional[Dict[str, object]], + eviction_inputs: Sequence[Dict[str, object]], ) -> None: - # Empty cohorts never reach here: _periodic_evict no-ops pre-launch. - needed_width = max(item["seq_len"] - item["prompt_len"] for item in prepared) - needed_score_tokens = max(item["seq_len"] for item in prepared) - needed_page_tokens = max(item["seq_len"] + item["protected_tail"] for item in prepared) - needed_requests = len(prepared) - if self.draft_kv_cache_manager is not None: - # The cached layout lookup enforces draft V2 pool page-count - # stability every round, exactly like the target's lookup. - self._runtime_kv_layout(draft=True) + """Per-round reuse gate over the three capacity axes; first round or + growth replaces the resident runtime as a whole.""" + # Empty cohorts never reach here: _evict_due_requests no-ops pre-launch. + needed_width = max( + item["source_length"] - item["prompt_length"] for item in eviction_inputs + ) + needed_score_tokens = max(item["source_length"] for item in eviction_inputs) + needed_page_tokens = max( + item["source_length"] + item["target_tail_length"] for item in eviction_inputs + ) + needed_requests = len(eviction_inputs) if self._buffers_built: if ( needed_width <= self._decode_width and needed_score_tokens <= self._bucket_seq_len - and needed_page_tokens <= self._page_table_token_capacity and needed_requests <= self._max_requests ): return - # This round outgrew the buffers: wait out the prior round (its + # This round outgrew the runtime: wait out the prior round (its # completion event orders after every use of the old epoch), then rebuild. if self._compaction_done_event is not None: self._compaction_done_event.synchronize() self._buffers_built = False mgr = self.kv_cache_manager - tail_capacity = self._protected_tail_capacity request_capacity = max(needed_requests, int(mgr.max_batch_size)) - decode_width = max( + selection_width_capacity = max( needed_width, self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), ) # Bucket sized by the presented cohorts, NOT max_seq_len (a floor there breaks 32-bit indexing). - seq_capacity = max(int(needed_page_tokens), 1024) - seq_capacity = 1 << (seq_capacity - 1).bit_length() - seq_capacity = min(seq_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens))) + score_token_capacity = next_positive_power_of_2(max(int(needed_page_tokens), 1024)) + score_token_capacity = min( + score_token_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens)) + ) # The bucket capacity must be tile-aligned (mis-tiling stripes the # score scratch silently); the ceiling division constructs that fact. score_tile_tokens = max(64, int(mgr.tokens_per_block)) - seq_capacity = -(-seq_capacity // score_tile_tokens) * score_tile_tokens - page_table_token_capacity = max(needed_page_tokens, seq_capacity + tail_capacity) - - draft = None - if self.draft_kv_cache_manager is not None: - draft_tail_capacity = self._draft_protected_tail_capacity - draft = dict( - layout=self._runtime_kv_layout(draft=True), - protected_tail_capacity=draft_tail_capacity, - page_table_token_capacity=seq_capacity + draft_tail_capacity, - ) + score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens - first_pool = layout["layer_pools"][layout["dense_layers"][0]] + first_pool = target_layout["layer_pools"][target_layout["dense_layers"][0]] if self._phase is None: # Host-only width offsets for the table builder (no device copy). self._phase = { @@ -894,41 +865,26 @@ def _ensure_buffers( "sin": None, "rows": 0, } - grow_mean_phase_table(self._phase, max(int(seq_capacity), 1)) - q_real, q_imag, mlr_coef = self._local_score_calibration(layout["global_layers"]) - self._build_buffers( - layout=layout, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=self._freq_scale_sq, - max_requests=request_capacity, - bucket_seq_len=seq_capacity, - decode_width=decode_width, - page_table_token_capacity=page_table_token_capacity, - keep_count=self.budget, - protected_tail_capacity=tail_capacity, - draft=draft, + grow_mean_phase_table(self._phase, max(int(score_token_capacity), 1)) + self._rebuild_eviction_runtime( + target_layout, + draft_layout, + request_capacity=request_capacity, + score_token_capacity=score_token_capacity, + selection_width_capacity=selection_width_capacity, ) self._buffers_built = True - def _build_buffers( + def _rebuild_eviction_runtime( self, + target_layout: Dict[str, object], + draft_layout: Optional[Dict[str, object]], *, - layout: Dict[str, object], - q_real: torch.Tensor, - q_imag: torch.Tensor, - mlr_coef: torch.Tensor, - freq_scale_sq: torch.Tensor, - max_requests: int, - bucket_seq_len: int, - decode_width: int, - page_table_token_capacity: int, - keep_count: int, - protected_tail_capacity: int, - draft: Optional[Dict[str, object]] = None, + request_capacity: int, + score_token_capacity: int, + selection_width_capacity: int, ) -> None: - """Build the round's buffers, compiled launches, and compaction data as attributes + """Build the resident eviction runtime for one capacity epoch as attributes (compiled kernels capture raw pool addresses: pools must stay alive and stay put).""" import cutlass import cutlass.cute as cute @@ -945,28 +901,29 @@ def _build_buffers( ) from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS - layer_pools = layout["layer_pools"] - dense_layers = list(layout["dense_layers"]) - swa_layers = list(layout["swa_layers"]) - swa_window = layout["swa_window"] - layer_group_representative = layout["layer_group_representative"] + layer_pools = target_layout["layer_pools"] + dense_layers = list(target_layout["dense_layers"]) + swa_layers = list(target_layout["swa_layers"]) + swa_window = target_layout["swa_window"] + layer_group_representative = target_layout["layer_group_representative"] # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. - layer_pool_ids = tuple(layout["layer_pool_ids"]) - num_page_table_slots = int(layout["manager"].num_pools) + layer_pool_ids = tuple(target_layout["layer_pool_ids"]) + num_page_table_slots = int(target_layout["manager"].num_pools) # The first dense layer anchors device and staging geometry. p0 = layer_pools[dense_layers[0]] device = p0.device - max_requests = int(max_requests) - seq_len = int(bucket_seq_len) - page_table_token_capacity = int(page_table_token_capacity) - decode_width = int(decode_width) - keep_count = int(keep_count) - protected_tail_capacity = int(protected_tail_capacity) - + max_requests = int(request_capacity) + seq_len = int(score_token_capacity) + decode_width = int(selection_width_capacity) + keep_count = int(self.budget) + protected_tail_capacity = int(self._protected_tail_capacity) + page_table_token_capacity = seq_len + protected_tail_capacity + + q_real, q_imag, mlr_coef = self._local_score_calibration(target_layout["global_layers"]) q_real, q_imag, mlr_coef, freq_scale_sq = ( tensor.to(device=device, dtype=torch.float32).contiguous() - for tensor in (q_real, q_imag, mlr_coef, freq_scale_sq) + for tensor in (q_real, q_imag, mlr_coef, self._freq_scale_sq) ) num_q_heads = int(q_real.shape[1]) num_freqs = int(q_real.shape[2]) @@ -975,7 +932,6 @@ def _build_buffers( self._bucket_seq_len = seq_len self._decode_width = decode_width self._keep_count = keep_count - self._page_table_token_capacity = page_table_token_capacity # ---- block-offset staging (target, plus the co-compressed draft) ------- self._block_offsets_host, self._block_offsets_device = _allocate_block_offset_staging( @@ -983,13 +939,12 @@ def _build_buffers( num_pools=num_page_table_slots, max_requests=max_requests, token_capacity=page_table_token_capacity, - max_source_blocks=int(layout["manager"].host_kv_cache_block_offsets.shape[-1]), + max_source_blocks=int(target_layout["manager"].host_kv_cache_block_offsets.shape[-1]), ) # The draft is never scored: these offsets feed only the draft compacts. self._draft_block_offsets_device = None self._draft_block_offsets_host = None - if draft is not None: - draft_layout = draft["layout"] + if draft_layout is not None: draft_representatives = list(draft_layout["pool_representatives"]) draft_anchor_pool = draft_layout["layer_pools"][draft_representatives[0]] # Construction-boundary invariant: the round shares one stream/event @@ -1003,7 +958,7 @@ def _build_buffers( draft_anchor_pool, num_pools=int(draft_layout["manager"].num_pools), max_requests=max_requests, - token_capacity=int(draft["page_table_token_capacity"]), + token_capacity=seq_len + int(self._draft_protected_tail_capacity), max_source_blocks=int( draft_layout["manager"].host_kv_cache_block_offsets.shape[-1] ), @@ -1386,28 +1341,28 @@ def _compiled_kernel(cache_key, build): # ---- compaction plans (opaque: only compact() interprets them) --------- compaction_params = [ build_compaction_params( - layout, + target_layout, block_offsets=self._block_offsets_device, kept_ordinals=self._kept_ordinal_rows, source_lengths=self._valid_seq_lens_device, dense_destination_bases=self._token_starts_device, # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, - protected_tail_capacity=int(protected_tail_capacity), + protected_tail_capacity=protected_tail_capacity, swa_move_offsets=swa_move_offsets_row, swa_destination_bases=self._swa_destination_bases, ) ] - if draft is not None: + if draft_layout is not None: compaction_params.append( build_compaction_params( - draft["layout"], + draft_layout, block_offsets=self._draft_block_offsets_device, kept_ordinals=self._kept_ordinal_rows, source_lengths=self._valid_seq_lens_device, dense_destination_bases=self._token_starts_device, dense_move_offsets=draft_move_offsets_row, - protected_tail_capacity=int(draft["protected_tail_capacity"]), + protected_tail_capacity=int(self._draft_protected_tail_capacity), ) ) self._compaction_params = tuple(compaction_params) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 0a27ebf26f23..c05b52241c61 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -228,7 +228,7 @@ def run_compaction(compaction): def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): """A bare manager carrying only the staging attributes, for the bulk - page-table copy tests (mirrors the product's ``_build_buffers`` names).""" + page-table copy tests (mirrors the product's ``_rebuild_eviction_runtime`` names).""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention staging = TriAttention.__new__(TriAttention) @@ -266,10 +266,10 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): def make_buffer_stubs(manager, *, decode_width=260): - """Stub the calibration/layout surfaces around ``_ensure_buffers``. + """Stub the calibration/layout surfaces around ``_ensure_eviction_runtime``. Returns the layout dict plus the manager attributes a stubbed - ``_build_buffers`` should set (production sets them in place).""" + ``_rebuild_eviction_runtime`` should set (production sets them in place).""" manager._freq_scale_sq = torch.ones(2) manager._phase = {"rows": 8} manager.calibration = {"omega": torch.ones(2)} @@ -289,7 +289,6 @@ def make_buffer_stubs(manager, *, decode_width=260): built_attributes = dict( _decode_width=decode_width, _bucket_seq_len=1024, - _page_table_token_capacity=65537, _max_requests=8, _token_starts_device=torch.zeros(8, dtype=torch.int32), _valid_widths=torch.empty(8, dtype=torch.int32), @@ -386,29 +385,30 @@ def make_triattention(**overrides): return TriAttention(make_tri_config(**overrides), make_fake_v2()) -def make_prepared_item( +def make_eviction_input( request=None, *, request_id=0, - seq_len, - round_start=None, - prompt_len=0, - expected_keep_count=0, - protected_tail=0, - kv_cache=None, - draft_kv_cache=None, + source_length, + logical_source_length=None, + prompt_length=0, + target_tail_length=0, + target_cache=None, + draft_cache=None, ): - """One prepared-cohort item shaped exactly like ``_periodic_evict`` builds.""" + """One due-cohort item shaped exactly like ``_evict_due_requests`` builds.""" + if request is None: + request = SimpleNamespace(py_request_id=request_id, py_num_compressed_tokens=0) return { "request": request, - "request_id": request_id, - "kv_cache": kv_cache, - "draft_kv_cache": draft_kv_cache, - "seq_len": int(seq_len), - "round_start": int(seq_len if round_start is None else round_start), - "prompt_len": int(prompt_len), - "expected_keep_count": int(expected_keep_count), - "protected_tail": int(protected_tail), + "target_cache": target_cache, + "draft_cache": draft_cache, + "source_length": int(source_length), + "logical_source_length": int( + source_length if logical_source_length is None else logical_source_length + ), + "prompt_length": int(prompt_length), + "target_tail_length": int(target_tail_length), } @@ -432,10 +432,10 @@ def make_request(request_id, **overrides): @contextmanager def mocked_eviction_internals(manager): - """Run the real ``_evict_requests`` body around a mocked round executor.""" + """Run the real ``_evict_due_requests`` transaction around a mocked round executor.""" with ( mock.patch.object(manager, "_runtime_kv_layout", return_value={}), - mock.patch.object(manager, "_ensure_buffers"), + mock.patch.object(manager, "_ensure_eviction_runtime"), mock.patch.object(manager, "_execute_eviction_round") as execute, ): yield SimpleNamespace(execute=execute) @@ -529,7 +529,6 @@ def make_cute_buffers( offsets, decode_width=None, keep_count=1, - page_table_token_capacity=None, protected_tail_capacity=0, storage_groups=None, layer_pool_ids=None, @@ -554,7 +553,7 @@ def make_cute_buffers( layer_pool_ids = [0] * num_layers # A live-manager source table exactly wide enough for the requested # capacity (the staged width clamps to it). - requested_tokens = seq_len if page_table_token_capacity is None else page_table_token_capacity + 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 @@ -563,6 +562,7 @@ def make_cute_buffers( num_pools=max(layer_pool_ids) + 1, host_kv_cache_block_offsets=torch.empty(1, 1, 2, source_blocks, dtype=torch.int32), ), + global_layers=list(range(num_layers)), layer_pools=layer_pools, dense_layers=list(range(num_layers)), swa_layers=[], @@ -583,20 +583,20 @@ def make_cute_buffers( manager._block_offsets_ready_event = None manager._compaction_done_event = None manager._phase = make_phase_table(offsets, omega, seq_len) - manager._build_buffers( - layout=layout, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr_coef, - freq_scale_sq=freq_scale_sq, - max_requests=max_requests, - bucket_seq_len=seq_len, - decode_width=decode_width, - page_table_token_capacity=( - seq_len if page_table_token_capacity is None else page_table_token_capacity - ), - keep_count=keep_count, - protected_tail_capacity=protected_tail_capacity, + # Real owner state consumed by the cold builder (production sets these at + # construction / calibration load). + manager.budget = keep_count + manager._protected_tail_capacity = protected_tail_capacity + manager._freq_scale_sq = freq_scale_sq + manager._triattn_q_real = q_real + manager._triattn_q_imag = q_imag + manager._triattn_mlr_coef = mlr_coef + manager._rebuild_eviction_runtime( + layout, + None, + request_capacity=max_requests, + score_token_capacity=seq_len, + selection_width_capacity=decode_width, ) return manager 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 index 7f792b784f58..822a43064a8c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -15,8 +15,8 @@ from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_buffer_stubs as _make_buffer_stubs +from conftest import make_eviction_input as _make_eviction_input from conftest import make_fake_v2 as _make_fake_v2 -from conftest import make_prepared_item as _make_prepared_item from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_request as _make_request from conftest import make_tri_config as _make_tri_config @@ -208,7 +208,7 @@ def test_execute_eviction_round_orders_both_manager_streams(): tri.kv_cache_manager = manager tri.draft_kv_cache_manager = draft_manager compute_stream = SimpleNamespace() - prepared = [_make_prepared_item(request_id=7, seq_len=8)] + eviction_inputs = [_make_eviction_input(request_id=7, source_length=8)] class Boom(RuntimeError): pass @@ -223,7 +223,7 @@ class Boom(RuntimeError): mock.patch.object(module, "compact") as compact, ): with pytest.raises(Boom): - tri._execute_eviction_round(prepared) + tri._execute_eviction_round(eviction_inputs) # Both page-table planes were snapshotted before the round body fired. assert stage.call_count == 2 @@ -308,7 +308,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): confirmed += 2 cache.capacity = confirmed - manager._periodic_evict(batch) + manager._evict_due_requests(batch) state = manager._request_states[7] if state["evicted_tokens"] > previous_evicted: @@ -320,9 +320,9 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert confirmed == 2 + 4 # The staged logical position restores the uncompressed # length: physical confirmed plus everything evicted so far - # (the prepared item's round_start). + # (the eviction input's logical_source_length). prepared = internals.execute.call_args.args[0] - assert prepared[0]["round_start"] == uncompressed + assert prepared[0]["logical_source_length"] == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. assert request.py_num_compressed_tokens == uncompressed - confirmed @@ -346,40 +346,28 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): manager = _make_triattention(budget=4) layout, built_attributes = _make_buffer_stubs(manager) - # The one-time host block-offset table shape gate reads the real manager - # tables (int32 [pools, slots, K/V, blocks]). - manager.kv_cache_manager.host_kv_cache_block_offsets = torch.zeros( - 1, 8, 2, 4, dtype=torch.int32 + draft_layout = dict( + layer_pools=[], + dense_layers=[], + layer_group_representative={}, + pool_representatives=(), + layer_pool_ids=(), + pool_page_counts=(4,), ) - draft_manager = _make_fake_v2(is_draft=True) - draft_manager.num_pools = 1 - draft_manager.host_kv_cache_block_offsets = torch.zeros(1, 8, 2, 4, dtype=torch.int32) - manager.draft_kv_cache_manager = draft_manager + manager.draft_kv_cache_manager = _make_fake_v2(is_draft=True) # Injected post-construction: mirror the ctor-cached manager-lifetime tail. manager._draft_protected_tail_capacity = 1 - # The one resolver serves both sides; only its draft arm runs here (the - # target layout arrives as the explicit _ensure_buffers argument). - manager._runtime_kv_layout = mock.Mock( - return_value=dict( - layer_pools=[], - dense_layers=[], - layer_group_representative={}, - pool_representatives=(), - layer_pool_ids=(), - pool_page_counts=(4,), - ) - ) - prepared = [ - _make_prepared_item(_make_request(7), request_id=7, seq_len=8, expected_keep_count=4) - ] + eviction_inputs = [_make_eviction_input(_make_request(7), request_id=7, source_length=8)] - def apply_built(**kwargs): + def apply_built(*args, **kwargs): for name, value in built_attributes.items(): setattr(manager, name, value) phase = manager._phase - with mock.patch.object(manager, "_build_buffers", side_effect=apply_built) as prepare: - manager._ensure_buffers(layout, prepared) + with mock.patch.object( + manager, "_rebuild_eviction_runtime", side_effect=apply_built + ) as prepare: + manager._ensure_eviction_runtime(layout, draft_layout, eviction_inputs) # Request capacity follows the executor limits, while the score # bucket follows what the cohort actually presents (power-of-two, @@ -387,45 +375,39 @@ def apply_built(**kwargs): # The stubbed build's capacities became the resident manager state. assert manager._buffers_built assert manager._decode_width == built_attributes["_decode_width"] - kwargs = prepare.call_args.kwargs # The mode and the shared phase-table dict live on the manager itself # and thread through unchanged (no longer build arguments). assert manager.eviction_mode == "union" assert manager._phase is phase - assert kwargs["max_requests"] == 8 - assert kwargs["decode_width"] == 4 + 2 * 128 - assert kwargs["bucket_seq_len"] == 1024 - assert kwargs["page_table_token_capacity"] == 1024 + 1 - assert kwargs["draft"]["page_table_token_capacity"] == 1024 + 1 - assert kwargs["draft"]["layout"] is manager._runtime_kv_layout.return_value - # Migrated from the pipeline buffer-kwargs test: the budget and the - # pool keys thread through unchanged. - assert kwargs["keep_count"] == manager.budget - assert kwargs["layout"] is layout - assert list(kwargs["layout"]["layer_pool_ids"]) == list(layout["layer_pool_ids"]) + args = prepare.call_args.args + kwargs = prepare.call_args.kwargs + assert args[0] is layout + assert args[1] is draft_layout + assert kwargs["request_capacity"] == 8 + assert kwargs["selection_width_capacity"] == 4 + 2 * 128 + assert kwargs["score_token_capacity"] == 1024 # A second round within the resident capacities reuses the buffers # (and with them the compaction launch data they carry). - manager._ensure_buffers(layout, prepared) + manager._ensure_eviction_runtime(layout, draft_layout, eviction_inputs) assert prepare.call_count == 1 # A cohort that outgrows the resident capacities rebuilds the whole # buffer state, compaction included. grown = [ - _make_prepared_item( + _make_eviction_input( _make_request(7), request_id=7, - seq_len=8 + built_attributes["_decode_width"], - expected_keep_count=4, + source_length=8 + built_attributes["_decode_width"], ) ] - def apply_rebuilt(**kwargs): + def apply_rebuilt(*args, **kwargs): apply_built() manager._decode_width = built_attributes["_decode_width"] + 8 prepare.side_effect = apply_rebuilt - manager._ensure_buffers(layout, grown) + manager._ensure_eviction_runtime(layout, draft_layout, grown) assert prepare.call_count == 2 assert manager._buffers_built assert manager._decode_width == built_attributes["_decode_width"] + 8 @@ -437,34 +419,34 @@ def test_source_growth_beyond_score_bucket_rebuilds_buffers(): the reuse gate must rebuild instead of scoring past the static geometry.""" manager = _make_triattention(budget=4) layout, built_attributes = _make_buffer_stubs(manager) - prepared = [ - _make_prepared_item( + eviction_inputs = [ + _make_eviction_input( _make_request(7), request_id=7, - seq_len=1024, - prompt_len=1020, - expected_keep_count=4, + source_length=1024, + prompt_length=1020, ) ] - def apply_built(**kwargs): + def apply_built(*args, **kwargs): for name, value in built_attributes.items(): setattr(manager, name, value) - with mock.patch.object(manager, "_build_buffers", side_effect=apply_built) as prepare: - manager._ensure_buffers(layout, prepared) + with mock.patch.object( + manager, "_rebuild_eviction_runtime", side_effect=apply_built + ) as prepare: + manager._ensure_eviction_runtime(layout, None, eviction_inputs) assert prepare.call_count == 1 # One more source token, same decode width and request count. grown = [ - _make_prepared_item( + _make_eviction_input( _make_request(7), request_id=7, - seq_len=1025, - prompt_len=1021, - expected_keep_count=4, + source_length=1025, + prompt_length=1021, ) ] - manager._ensure_buffers(layout, grown) + manager._ensure_eviction_runtime(layout, None, grown) assert prepare.call_count == 2 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 84b2a9e07c95..a8192a5db696 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -27,8 +27,8 @@ import torch from conftest import make_bare_staging as _make_bare_staging from conftest import make_cute_buffers as _make_cute_buffers +from conftest import make_eviction_input as _make_eviction_input from conftest import make_fake_v2 as _make_fake_v2 -from conftest import make_prepared_item as _make_prepared_item 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 @@ -161,11 +161,11 @@ def test_request_init_and_finish_lifecycle(self): triattention._buffers_built = True triattention._score_scratch = buffers batch = SimpleNamespace() - triattention._prepared_generation_batch = batch + triattention._inflight_scheduled_batch = batch triattention.on_request_finish(_make_request(11)) triattention.on_request_finish(_make_request(12)) assert triattention._request_states == {} - assert triattention._prepared_generation_batch is batch + assert triattention._inflight_scheduled_batch is batch assert triattention._buffers_built and triattention._score_scratch is buffers def test_resolve_accepts_flat_pt(self, flat_calibration_pt): @@ -290,7 +290,7 @@ def test_identity_selection_is_filtered_before_launch(self): state = _set_request_state(manager, 7, generation_steps=127) with _mocked_eviction_internals(manager) as internals: - manager._periodic_evict(SimpleNamespace(generation_requests=[request])) + manager._evict_due_requests(SimpleNamespace(generation_requests=[request])) internals.execute.assert_not_called() assert request.py_num_compressed_tokens == 0 @@ -306,13 +306,13 @@ def test_prepare_does_not_evict_and_update_runs_final_hook_once(self): context_requests_last_chunk=[], generation_requests=[_make_request(7)], ) - with mock.patch.object(manager, "_periodic_evict") as periodic_evict: + with mock.patch.object(manager, "_evict_due_requests") as evict_due: manager.prepare_resources(batch) - periodic_evict.assert_not_called() + evict_due.assert_not_called() manager.update_resources(batch) - periodic_evict.assert_called_once_with(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): @@ -360,12 +360,12 @@ def test_suspended_cache_defers_that_request_pre_launch(self): second_state = _set_request_state(manager, 8, generation_steps=127) batch = SimpleNamespace(generation_requests=[first_request, second_request]) - with mock.patch.object(manager, "_evict_requests", side_effect=lambda p: p) as evict: - manager._periodic_evict(batch) + with _mocked_eviction_internals(manager) as internals: + manager._evict_due_requests(batch) # Only the active request launched; the suspended one deferred whole. - prepared = evict.call_args.args[0] - assert [item["request_id"] for item in prepared] == [7] + eviction_inputs = internals.execute.call_args.args[0] + assert [item["request"].py_request_id for item in eviction_inputs] == [7] assert first_state["generation_steps"] == 128 assert second_state["generation_steps"] == 127 @@ -387,7 +387,7 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): request.py_num_accepted_draft_tokens = accepted cache = mgr.kv_cache_manager.kv_cache_map[7] cache.capacity = confirmed + tail - mgr._prepared_generation_batch = SimpleNamespace(generation_requests=[request]) + mgr._inflight_scheduled_batch = 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} @@ -396,35 +396,30 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): # Injected post-construction: mirror the ctor-cached manager-lifetime tail. mgr._draft_protected_tail_capacity = 1 - def compact(prepared): - # Publish and resize consume the same prepared cohort. - return prepared - - with mock.patch.object(mgr, "_evict_requests", side_effect=compact) as evict: - mgr._periodic_evict(batch) + with _mocked_eviction_internals(mgr) as internals: + mgr._evict_due_requests(batch) - # Tail excluded from seq_len; keep target = prompt + budget. - evict.assert_called_once_with( + # Tail excluded from the source length; keep target = prompt + budget. + internals.execute.assert_called_once_with( [ { "request": request, - "request_id": 7, - "kv_cache": cache, - "draft_kv_cache": draft_cache, - "seq_len": confirmed, - "round_start": confirmed, - "prompt_len": 1024, - "expected_keep_count": retained, - "protected_tail": tail, + "target_cache": cache, + "draft_cache": draft_cache, + "source_length": confirmed, + "logical_source_length": confirmed, + "prompt_length": 1024, + "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 seq_len must come from the physical capacity ledger - # (capacity minus the protected tail), never the logical length. + # The due-branch source length must come from the physical capacity + # ledger (capacity minus the protected tail), never the logical length. physical_confirmed = 6100 manager = _make_triattention(beta=128) manager.calibration = {} @@ -445,12 +440,14 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): py_draft_tokens=[1, 2, 3, 4], ) - with mock.patch.object(manager, "_evict_requests", return_value=[]) as evict: - manager._periodic_evict(SimpleNamespace(generation_requests=[request])) + with _mocked_eviction_internals(manager) as internals: + manager._evict_due_requests(SimpleNamespace(generation_requests=[request])) - prepared = evict.call_args.args[0] - assert prepared[0]["seq_len"] == physical_confirmed - cache.resize.assert_not_called() + eviction_inputs = internals.execute.call_args.args[0] + assert eviction_inputs[0]["source_length"] == physical_confirmed + # The logical position restores everything already evicted. + assert eviction_inputs[0]["logical_source_length"] == physical_confirmed + 100 + cache.resize.assert_called_once_with(1024 + manager.budget, None) def test_one_model_draft_co_compression_contract_is_accepted(self): # Construction accepts the separate draft manager, and the executor @@ -497,7 +494,7 @@ def test_prepare_snapshots_fixed_linear_generation_growth( triattention.prepare_resources(batch) - assert triattention._prepared_generation_batch is batch + assert triattention._inflight_scheduled_batch is batch # Members of the prepared batch grow by the cached constant; others # by zero. The prepared batch itself is the identity early-out. assert triattention._inflight_generation_growth(SimpleNamespace(), 7) == expected_growth @@ -521,10 +518,12 @@ def test_execute_rejects_int32_overflowing_round_starts(self): torch.zeros(1, 2, 2, 12, dtype=torch.int32), gather, torch.cuda.Stream(device=device) ) staging.kv_cache_manager = manager - prepared = [_make_prepared_item(request_id=7, seq_len=64, round_start=2**31)] + eviction_inputs = [ + _make_eviction_input(request_id=7, source_length=64, logical_source_length=2**31) + ] with pytest.raises((RuntimeError, OverflowError, ValueError)): - staging._execute_eviction_round(prepared) + staging._execute_eviction_round(eviction_inputs) assert gather.call_count == 0 def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): @@ -624,7 +623,7 @@ def stage_once(): 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] - def test_cohort_move_offsets_stage_keep_plus_tail_and_pad_rows(self): + def test_compaction_move_offsets_stage_keep_plus_tail_and_pad_rows(self): # The derived move offsets stage keep + tail moves per request # (keep_count=4 -> [6, 7]); padded rows past the cohort repeat the # final offset and contribute no moves. (The executor-call contract @@ -634,11 +633,11 @@ def test_cohort_move_offsets_stage_keep_plus_tail_and_pad_rows(self): offsets_manager._keep_count = 4 offsets_manager._swa_window = None offsets_manager._draft_protected_tail_capacity = None - prepared = [ - _make_prepared_item(request_id=7, seq_len=8, protected_tail=2), - _make_prepared_item(request_id=8, seq_len=10, protected_tail=3), + eviction_inputs = [ + _make_eviction_input(request_id=7, source_length=8, target_tail_length=2), + _make_eviction_input(request_id=8, source_length=10, target_tail_length=3), ] - dense, swa, draft = offsets_manager._cohort_move_offsets(prepared) + dense, swa, draft = offsets_manager._compute_compaction_move_offsets(eviction_inputs) assert dense == [0, 6, 13, 13, 13, 13, 13, 13, 13] assert swa is None assert draft is None @@ -736,11 +735,11 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): def prepared_cohort(): return [ - _make_prepared_item( + _make_eviction_input( request_id=request, - seq_len=int(valid_seq_lens[request]), - round_start=int(round_device[request]), - prompt_len=prompt_len, + source_length=int(valid_seq_lens[request]), + logical_source_length=int(round_device[request]), + prompt_length=prompt_len, ) for request in range(request_count) ] 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 index 7912e7a2a550..36f55b2c75ba 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -7,7 +7,7 @@ from conftest import build_compaction as _build_compaction from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_cute_buffers as _make_cute_buffers -from conftest import make_prepared_item as _make_prepared_item +from conftest import make_eviction_input as _make_eviction_input from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_staging_manager as _make_staging_manager from conftest import run_compaction as _run_compaction @@ -505,9 +505,11 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): torch.cuda.Stream(device=device), num_slots=2, ) - prepared = [_make_prepared_item(request_id=7, seq_len=seq_len, round_start=0)] + eviction_inputs = [ + _make_eviction_input(request_id=7, source_length=seq_len, logical_source_length=0) + ] tri.kv_cache_manager = manager - tri._execute_eviction_round(prepared) + tri._execute_eviction_round(eviction_inputs) assert torch.equal(tri._kept_ordinal_rows.view_as(expected_keep), expected_keep) torch.cuda.synchronize(device) @@ -646,27 +648,26 @@ def expected_keep() -> torch.Tensor: offsets=torch.zeros(1, dtype=torch.float32, device=device), decode_width=seq_len - prompt_len, keep_count=keep_count, - page_table_token_capacity=seq_len + protected_tail, protected_tail_capacity=protected_tail, ) tri.kv_cache_manager = manager def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) - prepared = [ - _make_prepared_item( + eviction_inputs = [ + _make_eviction_input( request_id=request_id, - seq_len=seq_len, - round_start=0, - prompt_len=prompt_len, - protected_tail=protected_tail, + source_length=seq_len, + logical_source_length=0, + prompt_length=prompt_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(prepared) + tri._execute_eviction_round(eviction_inputs) selected = tri._kept_ordinal_rows[0].clone().to(torch.long) torch.cuda.synchronize(device) assert cache.resize(compacted_capacity, None) From c87378c2b74f0baac07bc1b34231fa7b6578df59 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 21:05:51 -0700 Subject: [PATCH 153/178] [None][chore] Fail fast on rope and model-config loading; drop the import fallback and version probes Signed-off-by: tianruih --- .../triattention/triattention.py | 58 +++++-------------- 1 file changed, 13 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d2fe0637e6fc..17c53782efb6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -191,16 +191,9 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int global_layers = self._global_layers num_layers = len(global_layers) - try: - from transformers import AutoConfig + from transformers import AutoConfig - config = AutoConfig.from_pretrained( - model_path, trust_remote_code=True, local_files_only=True - ) - except (OSError, ValueError) as exc: - raise ValueError( - f"TriAttention could not load the local model config from {model_path!r}" - ) from exc + 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: @@ -314,54 +307,29 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: def _rope_tables(self, freq_count: int): """Derive the RoPE frequency tables from the model config.""" - import transformers from transformers import AutoConfig + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() config_values = cfg.to_dict() - head_dim = freq_count * 2 rope_params = ( config_values.get("rope_parameters") or config_values.get("rope_scaling") or {} ) - if rope_params and all(isinstance(v, dict) for v in rope_params.values()): + if rope_params and all(isinstance(value, dict) for value in rope_params.values()): raise ValueError( - f"TriAttention: layer-type-keyed rope_parameters are not supported for " - f"calibration conversion (model {self.model_path}); got {rope_params!r}." + f"TriAttention does not support per-layer-type rope parameters ({self.model_path})" ) rope_type = rope_params.get("rope_type") or rope_params.get("type") or "default" - theta_seen = rope_params.get("rope_theta", config_values.get("rope_theta")) - base = float(theta_seen) if theta_seen is not None else 10000.0 - - def analytic_inv_freq(): - idx = torch.arange(0, head_dim, 2, dtype=torch.float32) - return (1.0 / (base ** (idx / head_dim)))[:freq_count].clone() - if rope_type == "default": - # transformers>=5.5 no longer keys "default" in ROPE_INIT_FUNCTIONS: use the formula. - omega, scale_sq = analytic_inv_freq(), 1.0 + # "default" has no ROPE_INIT_FUNCTIONS entry; the analytic formula is its definition. + head_dim = freq_count * 2 + theta = rope_params.get("rope_theta", config_values.get("rope_theta")) + base = float(theta) if theta is not None else 10000.0 + positions = torch.arange(0, head_dim, 2, dtype=torch.float32) + omega = (1.0 / (base ** (positions / head_dim)))[:freq_count].clone() + scale_sq = 1.0 else: - try: - from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS - except ImportError: - logger.warning( - f"TriAttention: transformers rope-init unavailable; using the analytic " - f"inv_freq with theta={base} for {self.model_path} and IGNORING " - f"rope_type={rope_type!r} scaling corrections." - ) - return analytic_inv_freq(), torch.ones(freq_count, dtype=torch.float32) - if rope_type not in ROPE_INIT_FUNCTIONS: - raise ValueError( - f"TriAttention: unknown rope_type {rope_type!r} for {self.model_path} " - f"(transformers {transformers.__version__} provides " - f"{sorted(ROPE_INIT_FUNCTIONS)}); rope config seen: {rope_params!r}." - ) - try: - inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") - except Exception as exc: - raise ValueError( - f"TriAttention: rope-init {rope_type!r} failed for {self.model_path}; " - f"rope config seen: {rope_params!r}." - ) from exc + inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") omega = inv_freq.to(torch.float32)[:freq_count].clone() scale_sq = float(attention_factor) ** 2 return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) From bccdc8d6820b1b59d7b045bd53dbd2170b20806c Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 21:14:43 -0700 Subject: [PATCH 154/178] [None][refactor] Close the resident-state names: one compiled-score map, derived pool slots, capacity-named attributes Signed-off-by: tianruih --- .../triattention/triattention.py | 120 ++++++++---------- .../_torch/kv_cache_compression/conftest.py | 69 +++++----- .../test_triattention_cute_score.py | 4 +- .../test_triattention_cute_union_fusion.py | 4 +- .../test_triattention_draft_cocompaction.py | 22 ++-- .../test_triattention_pipeline.py | 14 +- .../test_triattention_selection_compaction.py | 30 ++--- 7 files changed, 118 insertions(+), 145 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 17c53782efb6..71bfc0607466 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -193,7 +193,9 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int from transformers import AutoConfig - config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, local_files_only=True) + 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: @@ -562,15 +564,15 @@ def _execute_eviction_round( with nvtx_range("triattention.score", color="blue"): # In-place refresh: the compiled score launches captured these pointers. _gather_mean_phase_kernel[(request_count,)]( - self._round_starts_device, + self._logical_source_lengths_device, self._phase["cos"], self._phase["sin"], self._phase["rows"], - self._valid_seq_lens_device, - self._token_starts_device, + self._source_lengths_device, + self._prompt_lengths_device, self._mean_cos, self._mean_sin, - self._valid_widths, + self._decode_lengths_device, self._swa_destination_bases, self._swa_rebase_delta, NUM_FREQS=self._phase_num_freqs, @@ -581,7 +583,7 @@ def _execute_eviction_round( if union: # Union path: fused score+stats, then the normalized union reduction. cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score_stats[request_count]( + self._compiled_score_by_request_count[request_count]( *self._cute_score_prefix, self._cute_mean_cos, self._cute_mean_sin, @@ -602,7 +604,7 @@ def _execute_eviction_round( ) else: cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score[request_count]( + self._compiled_score_by_request_count[request_count]( *self._cute_score_prefix, self._cute_mean_cos, self._cute_mean_sin, @@ -616,23 +618,25 @@ def _execute_eviction_round( pad = self._padded_head_columns source = ( self._score_scratch[ - : self._num_kv_heads * pad * num_segments * self._bucket_seq_len + : self._num_kv_heads * pad * num_segments * self._score_token_capacity ] .view( self._num_kv_heads, pad, request_count, self._num_layers, - self._bucket_seq_len, + self._score_token_capacity, )[:, :group_size] .permute(2, 3, 0, 1, 4) ) torch.add( - self._token_starts_device[:request_count].view(-1, 1, 1, 1, 1), + self._prompt_lengths_device[:request_count].view(-1, 1, 1, 1, 1), self._gather_columns, out=self._gather_index_base[:request_count], ) - self._gather_index_base[:request_count].clamp_(max=self._bucket_seq_len - 1) + self._gather_index_base[:request_count].clamp_( + max=self._score_token_capacity - 1 + ) columns = self._gather_index[:request_count] torch.gather( source, @@ -643,14 +647,14 @@ def _execute_eviction_round( self._num_layers, self._num_kv_heads, group_size, - self._decode_width, + self._selection_width_capacity, ), ) with nvtx_range("triattention.select", color="yellow"): if not union: prepare_per_head_scores( self._score_output[:request_count], - self._valid_widths, + self._decode_lengths_device, self._row_mean, self._row_inv_std, self._selection_scores_rows, @@ -680,7 +684,7 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets = [0] for moves in moves_per_request: offsets.append(offsets[-1] + moves) - offsets.extend(offsets[-1:] * (self._max_requests - len(moves_per_request))) + offsets.extend(offsets[-1:] * (self._request_capacity - len(moves_per_request))) return offsets tails = [int(item["target_tail_length"]) for item in eviction_inputs] @@ -710,7 +714,6 @@ def _stage_block_offsets( request_ids, host_block_offsets.shape[-1], ) - manager._stream.wait_event(self._staging_reuse_event) copy_batch_block_offsets_to_device( host_block_offsets, device_block_offsets, @@ -739,10 +742,10 @@ def _settle_top_tokens(self, request_count: int) -> None: _settle_ties_kernel[(request_count, self._selection_rows_per_request)]( self._selection_scores_rows, self._selection_row_lengths, - self._token_starts_device, + self._prompt_lengths_device, self._provisional_rows, self._kept_ordinal_rows, - WIDTH=self._decode_width, + WIDTH=self._selection_width_capacity, KEEP_COUNT=self._keep_count, SELECTION_ROWS=self._selection_rows_per_request, ) @@ -794,9 +797,9 @@ def _ensure_eviction_runtime( needed_requests = len(eviction_inputs) if self._buffers_built: if ( - needed_width <= self._decode_width - and needed_score_tokens <= self._bucket_seq_len - and needed_requests <= self._max_requests + needed_width <= self._selection_width_capacity + and needed_score_tokens <= self._score_token_capacity + and needed_requests <= self._request_capacity ): return # This round outgrew the runtime: wait out the prior round (its @@ -873,10 +876,9 @@ def _rebuild_eviction_runtime( dense_layers = list(target_layout["dense_layers"]) swa_layers = list(target_layout["swa_layers"]) swa_window = target_layout["swa_window"] - layer_group_representative = target_layout["layer_group_representative"] # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. layer_pool_ids = tuple(target_layout["layer_pool_ids"]) - num_page_table_slots = int(target_layout["manager"].num_pools) + num_page_table_slots = int(self.kv_cache_manager.num_pools) # The first dense layer anchors device and staging geometry. p0 = layer_pools[dense_layers[0]] @@ -896,9 +898,9 @@ def _rebuild_eviction_runtime( num_q_heads = int(q_real.shape[1]) num_freqs = int(q_real.shape[2]) - self._max_requests = max_requests - self._bucket_seq_len = seq_len - self._decode_width = decode_width + self._request_capacity = max_requests + self._score_token_capacity = seq_len + self._selection_width_capacity = decode_width self._keep_count = keep_count # ---- block-offset staging (target, plus the co-compressed draft) ------- @@ -907,7 +909,7 @@ def _rebuild_eviction_runtime( num_pools=num_page_table_slots, max_requests=max_requests, token_capacity=page_table_token_capacity, - max_source_blocks=int(target_layout["manager"].host_kv_cache_block_offsets.shape[-1]), + max_source_blocks=int(self.kv_cache_manager.host_kv_cache_block_offsets.shape[-1]), ) # The draft is never scored: these offsets feed only the draft compacts. self._draft_block_offsets_device = None @@ -924,11 +926,11 @@ def _rebuild_eviction_runtime( self._draft_block_offsets_host, self._draft_block_offsets_device = ( _allocate_block_offset_staging( draft_anchor_pool, - num_pools=int(draft_layout["manager"].num_pools), + num_pools=int(self.draft_kv_cache_manager.num_pools), max_requests=max_requests, token_capacity=seq_len + int(self._draft_protected_tail_capacity), max_source_blocks=int( - draft_layout["manager"].host_kv_cache_block_offsets.shape[-1] + self.draft_kv_cache_manager.host_kv_cache_block_offsets.shape[-1] ), ) ) @@ -946,10 +948,10 @@ def _rebuild_eviction_runtime( self._request_metadata_device = torch.zeros( (6, max_requests + 1), dtype=torch.int32, device=device ) - self._round_starts_device = self._request_metadata_device[0, :max_requests] - self._valid_seq_lens_device = self._request_metadata_device[1, :max_requests] + self._logical_source_lengths_device = self._request_metadata_device[0, :max_requests] + self._source_lengths_device = self._request_metadata_device[1, :max_requests] # Pinned per-request decode-window starts. - self._token_starts_device = self._request_metadata_device[2, :max_requests] + self._prompt_lengths_device = self._request_metadata_device[2, :max_requests] dense_move_offsets_row = self._request_metadata_device[3] swa_move_offsets_row = self._request_metadata_device[4] draft_move_offsets_row = self._request_metadata_device[5] @@ -957,7 +959,7 @@ def _rebuild_eviction_runtime( # the phase gather rebases each request's SWA destination base in place. self._swa_window = int(swa_window) if swa_layers else None self._swa_destination_bases = ( - torch.empty_like(self._token_starts_device) if swa_layers else None + torch.empty_like(self._prompt_lengths_device) if swa_layers else None ) self._swa_rebase_delta = keep_count - self._swa_window if swa_layers else 0 self._mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) @@ -970,11 +972,7 @@ def _rebuild_eviction_runtime( self._num_layers = len(dense_layers) self._num_q_heads = int(num_q_heads) self._num_kv_heads = int(num_kv_heads) - self._num_freqs = int(num_freqs) - self._tokens_per_block = int(tokens_per_block) - dense_layer_slots = [ - layer_pool_ids[layer_group_representative[layer]] for layer in dense_layers - ] + dense_layer_slots = [layer_pool_ids[layer] for layer in dense_layers] seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( self._num_layers ) @@ -1061,9 +1059,9 @@ def _rebuild_eviction_runtime( (seg_page_off, 16), (seg_req_id, 16), (seg_layer_id, 16), - (self._valid_seq_lens_device, 4), + (self._source_lengths_device, 4), (seg_out_offset, 16), - (self._token_starts_device, 4), + (self._prompt_lengths_device, 4), (q_real.view(-1), 16), (q_imag.view(-1), 16), (mlr_coef.view(-1), 16), @@ -1088,19 +1086,18 @@ def _rebuild_eviction_runtime( _to_cute(anchor_pool), _to_cute(self._tma_descriptors, assumed_align=128), ) - # Ctor-bound persistent launch operands: refreshed in place each round. + # Build-bound persistent launch operands: refreshed in place each round. self._cute_mean_cos = _to_cute(self._mean_cos.view(-1)) self._cute_mean_sin = _to_cute(self._mean_sin.view(-1)) self._cute_union_scores = _to_cute(self._union_scores.view(-1)) if union else None - self._compiled_score: Dict[int, object] = {} - self._compiled_score_stats: Dict[int, object] = {} + self._compiled_score_by_request_count: Dict[int, object] = {} self._compiled_normalize_union: Dict[int, object] = {} page_shards_by_count: Dict[int, int] = {} self._cute_selection_prefix = ( _to_cute(self._score_scratch), - _to_cute(self._valid_seq_lens_device, assumed_align=4), + _to_cute(self._source_lengths_device, assumed_align=4), _to_cute(seg_out_offset), - _to_cute(self._token_starts_device, assumed_align=4), + _to_cute(self._prompt_lengths_device, assumed_align=4), ) self._cute_partial_stats = _to_cute(self._partial_stats) static_geometry = ( @@ -1110,7 +1107,7 @@ def _rebuild_eviction_runtime( num_q_heads, self._num_kv_heads, num_freqs, - self._tokens_per_block, + int(tokens_per_block), tuple(int(value) for value in anchor_pool.shape), tuple(int(value) for value in anchor_pool.stride()), ) @@ -1145,12 +1142,8 @@ def _compiled_kernel(cache_key, build): _COMPILED_KERNELS[cache_key] = compiled return compiled - if union: - variant_key = "triattention_cute_score_stats" - compiled_entries = self._compiled_score_stats - else: - variant_key = "triattention_cute_score" - compiled_entries = self._compiled_score + variant_key = "triattention_cute_score_stats" if union else "triattention_cute_score" + compiled_entries = self._compiled_score_by_request_count for request_count, page_shards in variants: cache_key = ( variant_key, @@ -1243,11 +1236,11 @@ def _compiled_kernel(cache_key, build): self._compiled_normalize_union[request_count] = compiled_selection logger.info( f"TriAttention CuTe score enabled: {self._num_q_heads}q/{self._num_kv_heads}kv heads, " - f"{self._num_freqs} freqs, {self._tokens_per_block}-token pages" + f"{num_freqs} freqs, {int(tokens_per_block)}-token pages" ) # ---- selection buffers (canonical row-major, one name per storage) ----- - self._valid_widths = torch.full( + self._decode_lengths_device = torch.full( (max_requests,), decode_width, dtype=torch.int32, device=device ) if union: @@ -1256,7 +1249,7 @@ def _compiled_kernel(cache_key, build): (max_requests, decode_width), dtype=torch.float32, device=device ) # One selection row per request: its length IS the staged valid width. - self._selection_row_lengths = self._valid_widths + self._selection_row_lengths = self._decode_lengths_device # Padded rows still need in-range ordinals for the finalizer's gather. self._provisional_rows = torch.zeros( (max_requests, keep_count), dtype=torch.int32, device=device @@ -1312,8 +1305,8 @@ def _compiled_kernel(cache_key, build): target_layout, block_offsets=self._block_offsets_device, kept_ordinals=self._kept_ordinal_rows, - source_lengths=self._valid_seq_lens_device, - dense_destination_bases=self._token_starts_device, + source_lengths=self._source_lengths_device, + dense_destination_bases=self._prompt_lengths_device, # Per-round tails: the move offsets ride the staged metadata rows. dense_move_offsets=dense_move_offsets_row, protected_tail_capacity=protected_tail_capacity, @@ -1327,8 +1320,8 @@ def _compiled_kernel(cache_key, build): draft_layout, block_offsets=self._draft_block_offsets_device, kept_ordinals=self._kept_ordinal_rows, - source_lengths=self._valid_seq_lens_device, - dense_destination_bases=self._token_starts_device, + source_lengths=self._source_lengths_device, + dense_destination_bases=self._prompt_lengths_device, dense_move_offsets=draft_move_offsets_row, protected_tail_capacity=int(self._draft_protected_tail_capacity), ) @@ -1443,24 +1436,13 @@ def _build_runtime_kv_layout( all_storage_groups: Dict[int, List[int]] = {} for layer, pool_id in enumerate(layer_pool_ids): all_storage_groups.setdefault(pool_id, []).append(layer) - # Scored/compacted groups cover the dense layers only; SWA layers - # stage and compact as their own representatives. - storage_groups: Dict[int, List[int]] = {} - for layer in dense_layers: - storage_groups.setdefault(layer_pool_ids[layer], []).append(layer) - layer_group_representative = { - layer: layers[0] for layers in storage_groups.values() for layer in layers - } pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) return dict( - manager=manager, global_layers=global_layers, layer_pools=layer_pools, dense_layers=dense_layers, swa_layers=swa_layers, swa_window=swa_window, - storage_groups=storage_groups, - layer_group_representative=layer_group_representative, layer_pool_ids=layer_pool_ids, pool_representatives=pool_representatives, pool_page_counts=tuple( diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index c05b52241c61..635cc2dbeca2 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -128,7 +128,6 @@ def build_compaction(**overrides): eviction_mode="union", dense_layers=[0, 1], swa_layers=[], - layer_group_representative={0: 0, 1: 1}, layer_pool_ids=[0, 0], request_count=2, decode_keep_count=4, @@ -234,7 +233,7 @@ def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): staging = TriAttention.__new__(TriAttention) staging.kv_cache_manager = None staging.draft_kv_cache_manager = None - staging._max_requests = max_requests + staging._request_capacity = max_requests staging._keep_count = 4 staging._swa_window = None staging._draft_protected_tail_capacity = None @@ -276,22 +275,19 @@ def make_buffer_stubs(manager, *, decode_width=260): manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) pool = torch.empty(8, 2, 1, 4, 4) layout = dict( - manager=SimpleNamespace(num_pools=1), global_layers=[0, 1], layer_pools=[pool, pool], dense_layers=[0, 1], swa_layers=[], swa_window=None, - storage_groups={0: [0, 1]}, - layer_group_representative={0: 0, 1: 0}, layer_pool_ids=(0, 0), ) built_attributes = dict( - _decode_width=decode_width, - _bucket_seq_len=1024, - _max_requests=8, - _token_starts_device=torch.zeros(8, dtype=torch.int32), - _valid_widths=torch.empty(8, dtype=torch.int32), + _selection_width_capacity=decode_width, + _score_token_capacity=1024, + _request_capacity=8, + _prompt_lengths_device=torch.zeros(8, dtype=torch.int32), + _decode_lengths_device=torch.empty(8, dtype=torch.int32), ) return layout, built_attributes @@ -530,15 +526,13 @@ def make_cute_buffers( decode_width=None, keep_count=1, protected_tail_capacity=0, - storage_groups=None, layer_pool_ids=None, normalize_scores=True, ): """A bare manager with real eviction buffers built over the one-shared-slot default layout; split reference legs use ``eviction_mode="per_head"`` over - the same pools. ``storage_groups``/``layer_pool_ids`` override the - page-table grouping (``layer_pool_ids`` is the canonical per-layer V2 pool - id list).""" + the same pools. ``layer_pool_ids`` is the canonical per-layer V2 pool id + list and drives the page-table grouping.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention num_layers = len(layer_pools) @@ -547,8 +541,6 @@ def make_cute_buffers( # path); the widest window defaults keep old call sites. if decode_width is None: decode_width = seq_len - if storage_groups is None: - storage_groups = {0: list(range(num_layers))} if layer_pool_ids is None: layer_pool_ids = [0] * num_layers # A live-manager source table exactly wide enough for the requested @@ -558,23 +550,19 @@ def make_cute_buffers( source_blocks = -(-int(requested_tokens) // tokens_per_block) source_blocks = (source_blocks + 3) // 4 * 4 layout = dict( - manager=SimpleNamespace( - num_pools=max(layer_pool_ids) + 1, - host_kv_cache_block_offsets=torch.empty(1, 1, 2, source_blocks, dtype=torch.int32), - ), global_layers=list(range(num_layers)), layer_pools=layer_pools, dense_layers=list(range(num_layers)), swa_layers=[], swa_window=None, - storage_groups=storage_groups, - layer_group_representative={ - layer: layers[0] for layers in storage_groups.values() for layer in layers - }, layer_pool_ids=layer_pool_ids, ) manager = TriAttention.__new__(TriAttention) - manager.kv_cache_manager = None + # The cold builder reads staging geometry from the owning manager. + manager.kv_cache_manager = SimpleNamespace( + num_pools=max(layer_pool_ids) + 1, + host_kv_cache_block_offsets=torch.empty(1, 1, 2, source_blocks, dtype=torch.int32), + ) manager.draft_kv_cache_manager = None manager._draft_protected_tail_capacity = None manager.eviction_mode = eviction_mode @@ -615,8 +603,8 @@ def stage_score_metadata(manager, request_count, valid_seq_lens, valid_widths, t token_starts[:request_count], out=valid_widths[:request_count], ) - manager._valid_seq_lens_device[:request_count].copy_(valid_seq_lens[:request_count]) - manager._token_starts_device[:request_count].copy_(token_starts[:request_count]) + manager._source_lengths_device[:request_count].copy_(valid_seq_lens[:request_count]) + manager._prompt_lengths_device[:request_count].copy_(token_starts[:request_count]) def launch_split_scores( @@ -631,11 +619,11 @@ def launch_split_scores( stage_score_metadata(manager, request_count, valid_seq_lens, valid_widths, token_starts) manager._mean_cos[:request_count].copy_(mean_cos[:request_count]) manager._mean_sin[:request_count].copy_(mean_sin[:request_count]) - assert request_count in manager._compiled_score + assert request_count in manager._compiled_score_by_request_count stream = cuda_driver.CUstream( torch.cuda.current_stream(manager._score_scratch.device).cuda_stream ) - manager._compiled_score[request_count]( + manager._compiled_score_by_request_count[request_count]( *manager._cute_score_prefix, manager._cute_mean_cos, manager._cute_mean_sin, @@ -646,24 +634,35 @@ def launch_split_scores( num_segments = request_count * manager._num_layers group_size = manager._num_q_heads // manager._num_kv_heads source = ( - manager._score_scratch[: manager._num_kv_heads * 8 * num_segments * manager._bucket_seq_len] + manager._score_scratch[ + : manager._num_kv_heads * 8 * num_segments * manager._score_token_capacity + ] .view( - manager._num_kv_heads, 8, request_count, manager._num_layers, manager._bucket_seq_len + manager._num_kv_heads, + 8, + request_count, + manager._num_layers, + manager._score_token_capacity, )[:, :group_size] .permute(2, 3, 0, 1, 4) ) columns = ( token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + manager._gather_columns ) - columns = columns.clamp_(max=manager._bucket_seq_len - 1).expand( + columns = columns.clamp_(max=manager._score_token_capacity - 1).expand( request_count, manager._num_layers, manager._num_kv_heads, group_size, - manager._decode_width, + manager._selection_width_capacity, ) output = torch.full( - (request_count, manager._num_layers, manager._num_q_heads, manager._decode_width), + ( + request_count, + manager._num_layers, + manager._num_q_heads, + manager._selection_width_capacity, + ), float("nan"), dtype=torch.float32, device=manager._score_scratch.device, @@ -677,7 +676,7 @@ def launch_split_scores( manager._num_layers, manager._num_kv_heads, group_size, - manager._decode_width, + manager._selection_width_capacity, ), ) return output 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 index b87391776347..8799e0869150 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -158,7 +158,7 @@ def test_cute_kernel_matches_torch_oracle(case): ) # Every count up to capacity is served, nothing beyond. - assert max_requests + 1 not in tri._compiled_score + assert max_requests + 1 not in tri._compiled_score_by_request_count for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) scores = _launch_split_scores( @@ -174,7 +174,7 @@ def test_cute_kernel_matches_torch_oracle(case): request_count, num_layers, case["num_q_heads"], - tri._decode_width, + tri._selection_width_capacity, ) # The score leg owns the per-request decode widths the selection # reduce kernels consume. 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 index 315b331c315c..b1cd23798242 100644 --- 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 @@ -32,11 +32,11 @@ def _run_fused_union( tri._mean_cos[:request_count].copy_(mean_cos[:request_count]) tri._mean_sin[:request_count].copy_(mean_sin[:request_count]) assert ( - request_count in tri._compiled_score_stats + request_count in tri._compiled_score_by_request_count and request_count in tri._compiled_normalize_union ) stream = cuda_driver.CUstream(torch.cuda.current_stream(tri._score_scratch.device).cuda_stream) - tri._compiled_score_stats[request_count]( + tri._compiled_score_by_request_count[request_count]( *tri._cute_score_prefix, tri._cute_mean_cos, tri._cute_mean_sin, 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 index 822a43064a8c..5f791d6238aa 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -69,7 +69,6 @@ def _launched_draft_compaction(draft_protected_tails): protected_tail_capacity=max(target_protected_tails), draft_layer_pools=[draft_pool], draft_layers=[0], - draft_layer_group_representative={0: 0}, draft_layer_pool_ids=[0], draft_protected_tail_capacity=max(draft_protected_tails), draft_kv_block_offsets=_encode_block_offsets(draft_tables), @@ -173,7 +172,7 @@ def test_execute_eviction_round_orders_both_manager_streams(): event = mock.Mock() host = torch.zeros(6, 9, dtype=torch.int32) tri = TriAttention.__new__(TriAttention) - tri._max_requests = 8 + tri._request_capacity = 8 tri._keep_count = 4 tri.eviction_mode = "union" tri._swa_window = None @@ -187,10 +186,10 @@ def test_execute_eviction_round_orders_both_manager_streams(): tri._phase = {"cos": None, "sin": None, "rows": 8} tri._phase_num_freqs = 1 tri._phase_f_block = 1 - tri._round_starts_device = None - tri._valid_seq_lens_device = None - tri._token_starts_device = None - tri._valid_widths = None + tri._logical_source_lengths_device = None + tri._source_lengths_device = None + tri._prompt_lengths_device = None + tri._decode_lengths_device = None tri._mean_cos = None tri._mean_sin = None tri._swa_destination_bases = None @@ -349,7 +348,6 @@ def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): draft_layout = dict( layer_pools=[], dense_layers=[], - layer_group_representative={}, pool_representatives=(), layer_pool_ids=(), pool_page_counts=(4,), @@ -374,7 +372,7 @@ def apply_built(*args, **kwargs): # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. # The stubbed build's capacities became the resident manager state. assert manager._buffers_built - assert manager._decode_width == built_attributes["_decode_width"] + assert manager._selection_width_capacity == built_attributes["_selection_width_capacity"] # The mode and the shared phase-table dict live on the manager itself # and thread through unchanged (no longer build arguments). assert manager.eviction_mode == "union" @@ -398,19 +396,21 @@ def apply_built(*args, **kwargs): _make_eviction_input( _make_request(7), request_id=7, - source_length=8 + built_attributes["_decode_width"], + source_length=8 + built_attributes["_selection_width_capacity"], ) ] def apply_rebuilt(*args, **kwargs): apply_built() - manager._decode_width = built_attributes["_decode_width"] + 8 + manager._selection_width_capacity = built_attributes["_selection_width_capacity"] + 8 prepare.side_effect = apply_rebuilt manager._ensure_eviction_runtime(layout, draft_layout, grown) assert prepare.call_count == 2 assert manager._buffers_built - assert manager._decode_width == built_attributes["_decode_width"] + 8 + assert ( + manager._selection_width_capacity == built_attributes["_selection_width_capacity"] + 8 + ) def test_source_growth_beyond_score_bucket_rebuilds_buffers(): diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index a8192a5db696..f0c3f1746fa6 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -109,14 +109,11 @@ def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): triattention = _make_triattention() triattention.kv_cache_manager = manager cached = dict( - manager=manager, global_layers=[10, 11, 12], layer_pools=[torch.empty(4), torch.empty(8), torch.empty(4)], dense_layers=[0, 1, 2], swa_layers=[], swa_window=None, - storage_groups={0: [0, 2], 1: [1]}, - layer_group_representative={0: 0, 1: 1, 2: 0}, layer_pool_ids=(0, 1, 0), # These are local layer slots. Layer 2 shares layer 0's pool. pool_representatives=(0, 1), @@ -629,7 +626,7 @@ def test_compaction_move_offsets_stage_keep_plus_tail_and_pad_rows(self): # final offset and contribute no moves. (The executor-call contract # itself is pinned by the overlap-tail and draft publication tests.) offsets_manager = TriAttention.__new__(TriAttention) - offsets_manager._max_requests = 8 + offsets_manager._request_capacity = 8 offsets_manager._keep_count = 4 offsets_manager._swa_window = None offsets_manager._draft_protected_tail_capacity = None @@ -709,8 +706,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self): offsets=offsets, decode_width=seq_len - prompt_len, keep_count=4, - # One storage group and page-table slot per layer (distinct pools). - storage_groups={layer: [layer] for layer in layer_order}, + # One page-table slot per layer (distinct pools). layer_pool_ids=list(layer_order), normalize_scores=False, ) @@ -751,7 +747,7 @@ def prepared_cohort(): with mock.patch.object(module, "compact"): tri._execute_eviction_round(prepared_cohort()) fixed = tri._score_output.clone() - assert tri._valid_widths.tolist() == [seq_len - prompt_len for seq_len in seq_lens] + assert tri._decode_lengths_device.tolist() == [seq_len - prompt_len for seq_len in seq_lens] oracle = _torch_tri_score_oracle( pools, @@ -786,11 +782,11 @@ def prepared_cohort(): ) expected_second_widths = valid_seq_lens - prompt_len tri._score_output.fill_(score_sentinel) - tri._valid_widths.fill_(-1) + tri._decode_lengths_device.fill_(-1) with mock.patch.object(module, "compact"): tri._execute_eviction_round(prepared_cohort()) second_launch = tri._score_output.clone() - assert torch.equal(tri._valid_widths, expected_second_widths) + assert torch.equal(tri._decode_lengths_device, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index 36f55b2c75ba..9a0b7b892600 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -51,20 +51,22 @@ def _make_selection_buffers( allocation; ``_settle_top_tokens`` reads exactly these attributes.""" tri = TriAttention.__new__(TriAttention) tri.eviction_mode = eviction_mode - tri._max_requests = max_requests - tri._decode_width = width + tri._request_capacity = max_requests + tri._selection_width_capacity = width tri._keep_count = keep_count tri._num_layers = num_layers tri._num_q_heads = num_query_heads tri._num_kv_heads = num_kv_heads - tri._valid_widths = torch.full((max_requests,), width, dtype=torch.int32, device=device) - tri._token_starts_device = torch.zeros(max_requests, dtype=torch.int32, device=device) + tri._decode_lengths_device = torch.full( + (max_requests,), width, dtype=torch.int32, device=device + ) + tri._prompt_lengths_device = torch.zeros(max_requests, dtype=torch.int32, device=device) if eviction_mode == "union": tri._selection_rows_per_request = 1 tri._selection_scores_rows = torch.empty( (max_requests, width), dtype=torch.float32, device=device ) - tri._selection_row_lengths = tri._valid_widths + tri._selection_row_lengths = tri._decode_lengths_device # Padded rows carry zero valid width; their provisional TopK entries # must still be in-range ordinals for the finalizer's score gather. tri._provisional_rows = torch.zeros( @@ -99,7 +101,7 @@ def _select_per_head(tri, scores, *, normalize_scores): """The per-head selection flow: reduce kernels, then top-k settle.""" prepare_per_head_scores( scores, - tri._valid_widths, + tri._decode_lengths_device, tri._row_mean, tri._row_inv_std, tri._selection_scores_rows, @@ -107,7 +109,7 @@ def _select_per_head(tri, scores, *, normalize_scores): per_layer=tri.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - tri._settle_top_tokens(tri._max_requests) + tri._settle_top_tokens(tri._request_capacity) def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: @@ -187,7 +189,7 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( num_query_heads=query_heads, num_kv_heads=kv_heads, ) - tri._valid_widths.copy_(valid_widths.to(device)) + tri._decode_lengths_device.copy_(valid_widths.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) @@ -225,12 +227,12 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid device=device, max_requests=request_count, ) - tri._valid_widths.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) - tri._token_starts_device[:request_count].copy_( + tri._decode_lengths_device.copy_(torch.tensor(valid_widths, 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._settle_top_tokens(tri._max_requests) + tri._settle_top_tokens(tri._request_capacity) actual = tri._kept_ordinal_rows.cpu() combined = scores.amax(dim=1).cpu() @@ -428,7 +430,6 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): tokens_per_block = 32 head_dim = 64 num_freqs = head_dim // 2 - dense_groups = [[0, 2], [1]] page_tables = ( torch.tensor([[1, 0]], dtype=torch.int32, device=device), torch.tensor([[0, 1]], dtype=torch.int32, device=device), @@ -470,10 +471,6 @@ def test_per_layer_score_selection_and_compaction_preserve_dense_layer_order(): offsets=torch.zeros(1, dtype=torch.float32, device=device), decode_width=bucket_capacity, keep_count=keep_count, - storage_groups={ - 0: dense_groups[0], - 1: dense_groups[1], - }, layer_pool_ids=[0, 1, 0], normalize_scores=False, ) @@ -753,7 +750,6 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): layer_pools=pools, dense_layers=[0], swa_layers=[1], - layer_group_representative={0: 0}, # Dense layer 0 stages in plane 0, the SWA layer in its own plane 1. layer_pool_ids=[0, 1], kept_token_ordinals=keep.to(torch.int32), From e6d2cb74b83e135e3c255242e18bf7e1901f1125 Mon Sep 17 00:00:00 2001 From: tianruih Date: Fri, 24 Jul 2026 21:41:19 -0700 Subject: [PATCH 155/178] [None][perf] Reduce score windows straight from the scratch and write union rows in kernel Signed-off-by: tianruih --- .../triattention/triattention.py | 230 ++++++------------ .../triattention_cute_selection.py | 19 +- .../triattention/triattention_kernels.py | 95 ++++++-- .../_torch/kv_cache_compression/conftest.py | 26 +- .../test_triattention_cute_union_fusion.py | 9 +- .../test_triattention_pipeline.py | 41 +++- .../test_triattention_selection_compaction.py | 20 +- 7 files changed, 243 insertions(+), 197 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 71bfc0607466..d691112bbf5e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -580,85 +580,40 @@ def _execute_eviction_round( HAS_SWA=self._swa_destination_bases is not None, num_warps=1, ) + cu_stream = cuda_driver.CUstream(stream.cuda_stream) + self._compiled_score_by_request_count[request_count]( + *self._cute_score_prefix, + self._cute_mean_cos, + self._cute_mean_sin, + *self._cute_score_tail, + request_count, + cu_stream, + ) if union: - # Union path: fused score+stats, then the normalized union reduction. - cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score_by_request_count[request_count]( - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, - request_count, - cu_stream, - ) + # Normalized union reduction, written straight into the selection rows. self._compiled_normalize_union[request_count]( self._cute_partial_stats, *self._cute_selection_prefix, - self._cute_union_scores, - request_count, - cu_stream, - ) - columns = min(self._union_scores.shape[1], self._selection_scores_rows.shape[1]) - self._selection_scores_rows[:request_count, :columns].copy_( - self._union_scores[:request_count, :columns] - ) - else: - cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score_by_request_count[request_count]( - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, + self._cute_selection_scores_rows, request_count, cu_stream, ) - # Gather each decode window into the [request, layer, head, token] layout of the reduces. - group_size = self._num_q_heads // self._num_kv_heads - num_segments = request_count * self._num_layers - pad = self._padded_head_columns - source = ( - self._score_scratch[ - : self._num_kv_heads * pad * num_segments * self._score_token_capacity - ] - .view( - self._num_kv_heads, - pad, - request_count, - self._num_layers, - self._score_token_capacity, - )[:, :group_size] - .permute(2, 3, 0, 1, 4) - ) - torch.add( - self._prompt_lengths_device[:request_count].view(-1, 1, 1, 1, 1), - self._gather_columns, - out=self._gather_index_base[:request_count], - ) - self._gather_index_base[:request_count].clamp_( - max=self._score_token_capacity - 1 - ) - columns = self._gather_index[:request_count] - torch.gather( - source, - 4, - columns, - out=self._score_output[:request_count].view( - request_count, - self._num_layers, - self._num_kv_heads, - group_size, - self._selection_width_capacity, - ), - ) with nvtx_range("triattention.select", color="yellow"): if not union: + # Per-head reduces read each decode window straight out of the scratch. prepare_per_head_scores( - self._score_output[:request_count], + 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, + num_layers=self._num_layers, + num_q_heads=self._num_q_heads, + padded_head_columns=self._padded_head_columns, + score_token_capacity=self._score_token_capacity, per_layer=self.eviction_mode == "per_layer_perhead", normalize_scores=self.normalize_scores, ) @@ -1003,30 +958,7 @@ def _rebuild_eviction_runtime( seg_out_offset = ( torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len ).to(torch.int32) - self._gather_columns = torch.arange(decode_width, dtype=torch.int64, device=device).view( - 1, 1, 1, 1, -1 - ) union = self.eviction_mode == "union" - # Per-head gather index: per round only the token-start base is re-added in place. - self._gather_index_base = None - self._gather_index = None - if not union: - self._gather_index_base = torch.empty( - (max_requests, 1, 1, 1, decode_width), dtype=torch.int64, device=device - ) - self._gather_index = self._gather_index_base.expand( - max_requests, - len(dense_layers), - self._num_kv_heads, - num_q_heads // self._num_kv_heads, - decode_width, - ) - self._union_scores = None - if union: - # Bucket-wide rows; consumers mask by the per-request widths. - self._union_scores = torch.empty( - (max_requests, seq_len), dtype=torch.float32, device=device - ) # ---- score path: compiled per request-count/page-shard ---- anchor_pool = p0 @@ -1089,7 +1021,6 @@ def _rebuild_eviction_runtime( # Build-bound persistent launch operands: refreshed in place each round. self._cute_mean_cos = _to_cute(self._mean_cos.view(-1)) self._cute_mean_sin = _to_cute(self._mean_sin.view(-1)) - self._cute_union_scores = _to_cute(self._union_scores.view(-1)) if union else None self._compiled_score_by_request_count: Dict[int, object] = {} self._compiled_normalize_union: Dict[int, object] = {} page_shards_by_count: Dict[int, int] = {} @@ -1182,6 +1113,60 @@ def _compiled_kernel(cache_key, build): SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 ) + logger.info( + f"TriAttention CuTe score enabled: {self._num_q_heads}q/{self._num_kv_heads}kv heads, " + f"{num_freqs} freqs, {int(tokens_per_block)}-token pages" + ) + + # ---- selection buffers (canonical row-major, one name per storage) ----- + self._decode_lengths_device = torch.full( + (max_requests,), decode_width, dtype=torch.int32, device=device + ) + if union: + self._selection_rows_per_request = 1 + self._selection_scores_rows = torch.empty( + (max_requests, decode_width), dtype=torch.float32, device=device + ) + # One selection row per request: its length IS the staged valid width. + self._selection_row_lengths = self._decode_lengths_device + # Padded rows still need in-range ordinals for the finalizer's gather. + self._provisional_rows = torch.zeros( + (max_requests, keep_count), dtype=torch.int32, device=device + ) + # Kept decode ordinals. + self._kept_ordinal_rows = torch.empty( + (max_requests, keep_count), dtype=torch.int32, device=device + ) + self._cute_selection_scores_rows = _to_cute(self._selection_scores_rows.view(-1)) + else: + selection_rows = ( + self._num_kv_heads + if self.eviction_mode == "per_head" + else self._num_layers * self._num_kv_heads + ) + # The selection rectangle must stay 32-bit indexable (wraparound = wild reads). + selection_rect = max_requests * selection_rows * max(decode_width, keep_count) + if selection_rect >= 2**31: + raise ValueError( + f"per-head selection rectangle overflows 32-bit indexing: {selection_rect}" + ) + self._selection_rows_per_request = selection_rows + score_shape = (max_requests, 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( + (max_requests * selection_rows, decode_width), dtype=torch.float32, device=device + ) + self._selection_row_lengths = torch.full( + (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device + ) + self._provisional_rows = torch.zeros( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + ) + self._kept_ordinal_rows = torch.empty( + (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + ) + if union: from .triattention_cute_selection import ( _select_normalize_union_config, @@ -1204,7 +1189,7 @@ def _compiled_kernel(cache_key, build): static_geometry, tensor_specs, config_key, - _tensor_spec(self._union_scores), + _tensor_spec(self._selection_scores_rows), _tensor_spec(self._partial_stats), ) tokens_per_lane, token_subtiles, row_cluster_ctas = config @@ -1224,80 +1209,17 @@ def _compiled_kernel(cache_key, build): tokens_per_lane=tokens_per_lane, token_subtiles=token_subtiles, row_cluster_ctas=row_cluster_ctas, + output_row_stride=decode_width, ), self._cute_partial_stats, *self._cute_selection_prefix, - self._cute_union_scores, + self._cute_selection_scores_rows, cutlass.Int32(1), stream, ), ) compiled_configs[config_key] = compiled_selection self._compiled_normalize_union[request_count] = compiled_selection - logger.info( - f"TriAttention CuTe score enabled: {self._num_q_heads}q/{self._num_kv_heads}kv heads, " - f"{num_freqs} freqs, {int(tokens_per_block)}-token pages" - ) - - # ---- selection buffers (canonical row-major, one name per storage) ----- - self._decode_lengths_device = torch.full( - (max_requests,), decode_width, dtype=torch.int32, device=device - ) - if union: - self._selection_rows_per_request = 1 - self._selection_scores_rows = torch.empty( - (max_requests, decode_width), dtype=torch.float32, device=device - ) - # One selection row per request: its length IS the staged valid width. - self._selection_row_lengths = self._decode_lengths_device - # Padded rows still need in-range ordinals for the finalizer's gather. - self._provisional_rows = torch.zeros( - (max_requests, keep_count), dtype=torch.int32, device=device - ) - # Kept decode ordinals. - self._kept_ordinal_rows = torch.empty( - (max_requests, keep_count), dtype=torch.int32, device=device - ) - self._score_output = None - else: - selection_rows = ( - self._num_kv_heads - if self.eviction_mode == "per_head" - else self._num_layers * self._num_kv_heads - ) - # Both rectangles must stay 32-bit indexable (wraparound = wild reads). - score_rect = max_requests * self._num_layers * self._num_q_heads * decode_width - selection_rect = max_requests * selection_rows * max(decode_width, keep_count) - if max(score_rect, selection_rect) >= 2**31: - raise ValueError( - f"per-head score rectangles overflow 32-bit indexing: " - f"scores {score_rect}, selection {selection_rect}" - ) - self._selection_rows_per_request = selection_rows - # [request, layer, head, token] layout read by the reduce kernels. - self._score_output = torch.empty( - max_requests, - self._num_layers, - self._num_q_heads, - decode_width, - dtype=torch.float32, - device=device, - ) - score_shape = (max_requests, 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( - (max_requests * selection_rows, decode_width), dtype=torch.float32, device=device - ) - self._selection_row_lengths = torch.full( - (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device - ) - self._provisional_rows = torch.zeros( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) - self._kept_ordinal_rows = torch.empty( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) # ---- compaction plans (opaque: only compact() interprets them) --------- compaction_params = [ 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 index 8583bac57150..22887ef8bbc8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -141,14 +141,17 @@ def __init__( 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.seq_len = seq_len - # The widest score window (the whole bucket) sizes the output rows and token-tile grid. + # The widest score window (the whole bucket) sizes the token-tile grid; + # output rows are the TopK selection rows with their own stride. self.width = seq_len + 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 @@ -242,11 +245,11 @@ def _reduce_and_store_union_rows( ) reduced_values[token_slot] = union_value subtile_first_token = first_token + token_subtile * self.subtile_token_tile - # Straddling subtiles store per token; union_scores stays < 2^31 so i32 cannot wrap. - if cutlass.const_expr(self.width % self.tokens_per_lane == 0) and cutlass.dynamic_expr( - subtile_first_token + self.tokens_per_lane <= valid_width - ): - union_index = request_idx * self.width + subtile_first_token + # 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 <= valid_width): + union_index = request_idx * self.output_row_stride + subtile_first_token union_tile = _gmem_lane_tile( union_scores.iterator, union_index, @@ -262,7 +265,9 @@ def _reduce_and_store_union_rows( for token_slot in cutlass.range_constexpr(self.tokens_per_lane): token = subtile_first_token + token_slot if cutlass.dynamic_expr(token < valid_width): - union_scores[request_idx * self.width + token] = reduced_values[token_slot] + union_scores[request_idx * self.output_row_stride + token] = reduced_values[ + token_slot + ] @cute.kernel def kernel( diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 2bd4fa15bdb2..175f5465fd35 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -56,22 +56,41 @@ def _gather_mean_phase_kernel( @triton.jit def _score_row_stats_kernel( - scores, - valid_widths, + 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 one valid-prefix mean and inverse standard deviation per score row.""" + """Compute one valid-window mean and inverse standard deviation per score row, + reading each row's decode window straight out of the score scratch.""" + QUERY_GROUP_SIZE: tl.constexpr = NUM_Q_HEADS // NUM_KV_HEADS flat_row = tl.program_id(0) request = flat_row // ROWS - valid_width = tl.load(valid_widths + request) - score_row = scores + flat_row * WIDTH + 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 + valid_width = 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): @@ -94,28 +113,34 @@ def _score_row_stats_kernel( @triton.jit def _score_per_head_reduce_kernel( - scores, - valid_widths, + score_scratch, + decode_lengths, + prompt_lengths, row_mean, row_inv_std, selection_scores, selection_seq_lens, + 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 query-head score rows into one selector row per KV-head domain.""" + """Reduce each KV-head domain's decode window straight out of the score + scratch into one 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) - valid_width = tl.load(valid_widths + request) + valid_width = tl.load(decode_lengths + request) + prompt_start = tl.load(prompt_lengths + request) valid_token = token < valid_width if token_block == 0: @@ -131,8 +156,12 @@ def _score_per_head_reduce_kernel( 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( - scores + flat_row * WIDTH + token, + 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) @@ -148,8 +177,12 @@ def _score_per_head_reduce_kernel( 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( - scores + flat_row * WIDTH + token, + 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) @@ -166,50 +199,66 @@ def _score_per_head_reduce_kernel( def prepare_per_head_scores( - scores: torch.Tensor, - valid_widths: torch.Tensor, + 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, + num_layers: int, + num_q_heads: int, + padded_head_columns: int, + score_token_capacity: int, per_layer: bool, normalize_scores: bool, ) -> None: - """Normalize and reduce score rows for either per-head eviction mode.""" - request_count, num_layers, num_q_heads, width = scores.shape - selection_rows = int(selection_scores_rows.shape[0]) // int(valid_widths.shape[0]) + """Normalize and reduce the scratch's decode windows for either per-head + eviction mode.""" + width = int(selection_scores_rows.shape[1]) + selection_rows = int(selection_scores_rows.shape[0]) // int(decode_lengths.shape[0]) num_kv_heads = selection_rows // num_layers if per_layer else selection_rows 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,)]( - scores, - valid_widths, + 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=width, ) # 256-token tiles match the reduce kernel's BLOCK default. _score_per_head_reduce_kernel[(request_count, selection_rows, triton.cdiv(width, 256))]( - scores, - valid_widths, + 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=width, PER_LAYER=per_layer, NORMALIZE=normalize_scores, ) -# ---- Selection finalize: settle threshold ties into the kept-ordinal rows ---- - - @triton.jit def _settle_ties_kernel( selection_scores_rows, diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 635cc2dbeca2..968556e06a41 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -595,6 +595,24 @@ def write_block_offsets(manager, encoded): 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 a [request, layer, q_head, token] rectangle into the fused + scorer's scratch layout (prompt starts at zero, bucket = rectangle width).""" + 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, valid_seq_lens, valid_widths, token_starts): """Stage the per-round score metadata exactly like production (the compiled score launches read the staged rows via pointer capture).""" @@ -646,9 +664,11 @@ def launch_split_scores( )[:, :group_size] .permute(2, 3, 0, 1, 4) ) - columns = ( - token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + manager._gather_columns - ) + columns = token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + torch.arange( + manager._selection_width_capacity, + dtype=torch.int64, + device=manager._score_scratch.device, + ).view(1, 1, 1, 1, -1) columns = columns.clamp_(max=manager._score_token_capacity - 1).expand( request_count, manager._num_layers, 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 index b1cd23798242..ecbbf235e9ab 100644 --- 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 @@ -24,8 +24,8 @@ def _run_fused_union( ): """The fused score+stats+normalized-union pipeline (THE union path), fired directly off the compiled entries exactly like the round. Test mean phases - load into the build-bound buffers; the fused rows land in the build-bound - ``tri._union_scores`` and copy out.""" + load into the build-bound buffers; the fused rows land straight in the + build-bound ``tri._selection_scores_rows``.""" import cuda.bindings.driver as cuda_driver _stage_score_metadata(tri, request_count, valid_seq_lens, valid_widths, token_starts) @@ -47,11 +47,12 @@ def _run_fused_union( tri._compiled_normalize_union[request_count]( tri._cute_partial_stats, *tri._cute_selection_prefix, - tri._cute_union_scores, + tri._cute_selection_scores_rows, request_count, stream, ) - union_out[:request_count].copy_(tri._union_scores[: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, valid_widths: torch.Tensor) -> torch.Tensor: diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index f0c3f1746fa6..f31b4e949c1f 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -740,13 +740,46 @@ def prepared_cohort(): for request in range(request_count) ] + def score_rectangle(): + # Test-side extraction of the decode-window rectangle from the + # scratch (production reduces the scratch in-kernel). + group = tri._num_q_heads // tri._num_kv_heads + segments = request_count * tri._num_layers + source = ( + tri._score_scratch[: tri._num_kv_heads * 8 * segments * tri._score_token_capacity] + .view( + tri._num_kv_heads, + 8, + request_count, + tri._num_layers, + tri._score_token_capacity, + )[:, :group] + .permute(2, 3, 0, 1, 4) + .reshape( + request_count, + tri._num_layers, + tri._num_q_heads, + tri._score_token_capacity, + ) + ) + columns = prompt_len + torch.arange( + tri._selection_width_capacity, dtype=torch.int64, device=device + ).view(1, 1, 1, -1) + columns = columns.clamp_(max=tri._score_token_capacity - 1).expand( + request_count, + tri._num_layers, + tri._num_q_heads, + tri._selection_width_capacity, + ) + return torch.gather(source, 3, columns) + score_sentinel = -12345.0 - tri._score_output.fill_(score_sentinel) + tri._score_scratch.fill_(score_sentinel) # The compact stage is stubbed to a no-op: this test owns the score # buffers only, never a staged move decision. with mock.patch.object(module, "compact"): tri._execute_eviction_round(prepared_cohort()) - fixed = tri._score_output.clone() + fixed = score_rectangle() assert tri._decode_lengths_device.tolist() == [seq_len - prompt_len for seq_len in seq_lens] oracle = _torch_tri_score_oracle( @@ -781,11 +814,11 @@ def prepared_cohort(): ) ) expected_second_widths = valid_seq_lens - prompt_len - tri._score_output.fill_(score_sentinel) + tri._score_scratch.fill_(score_sentinel) tri._decode_lengths_device.fill_(-1) with mock.patch.object(module, "compact"): tri._execute_eviction_round(prepared_cohort()) - second_launch = tri._score_output.clone() + second_launch = score_rectangle() assert torch.equal(tri._decode_lengths_device, expected_second_widths) assert not torch.equal(second_launch, fixed) 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 index 9a0b7b892600..af60affecc8e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -10,6 +10,7 @@ from conftest import make_eviction_input as _make_eviction_input from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_staging_manager as _make_staging_manager +from conftest import rect_to_score_scratch as _rect_to_score_scratch from conftest import run_compaction as _run_compaction from conftest import set_protected_tails as _set_protected_tails @@ -99,13 +100,21 @@ def _make_selection_buffers( def _select_per_head(tri, scores, *, normalize_scores): """The per-head selection flow: reduce kernels, then top-k settle.""" + request_count, num_layers, num_q_heads, width = scores.shape + score_scratch, prompt_lengths = _rect_to_score_scratch(scores, tri._num_kv_heads) prepare_per_head_scores( - 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, + num_layers=num_layers, + num_q_heads=num_q_heads, + padded_head_columns=8, + score_token_capacity=width, per_layer=tri.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) @@ -272,13 +281,20 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no request_count * selection_rows, dtype=torch.int32, device=device ) + score_scratch, prompt_lengths = _rect_to_score_scratch(scores, kv_heads) prepare_per_head_scores( - scores, + score_scratch, valid_widths, + prompt_lengths, row_mean, row_inv_std, selection_scores_rows, selection_row_lengths, + request_count=request_count, + num_layers=layers, + num_q_heads=query_heads, + padded_head_columns=8, + score_token_capacity=width, per_layer=per_layer, normalize_scores=normalize_scores, ) From 1a32167a523cd84809997492a02681fa81a9ec5c Mon Sep 17 00:00:00 2001 From: tianruih Date: Sat, 25 Jul 2026 09:17:27 -0700 Subject: [PATCH 156/178] [None][fix] Slice the score calibration to this rank's attention heads under TP Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 7 +++++++ tests/unittest/_torch/kv_cache_compression/conftest.py | 1 + 2 files changed, 8 insertions(+) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index d691112bbf5e..2e191f165070 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -846,6 +846,13 @@ def _rebuild_eviction_runtime( page_table_token_capacity = seq_len + protected_tail_capacity q_real, q_imag, mlr_coef = self._local_score_calibration(target_layout["global_layers"]) + # TP splits attention heads contiguously per rank; the calibration is global. + 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, freq_scale_sq = ( tensor.to(device=device, dtype=torch.float32).contiguous() for tensor in (q_real, q_imag, mlr_coef, self._freq_scale_sq) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 968556e06a41..27466ca3b324 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -562,6 +562,7 @@ def make_cute_buffers( manager.kv_cache_manager = SimpleNamespace( num_pools=max(layer_pool_ids) + 1, 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 = None From 5b532d9ff0cd54ece229f140ec2bb56b4a26276d Mon Sep 17 00:00:00 2001 From: tianruih Date: Sat, 25 Jul 2026 09:54:13 -0700 Subject: [PATCH 157/178] [None][fix] Restore the exact global union under TP: gather rank unions and max-fold Signed-off-by: tianruih --- .../triattention/triattention.py | 24 +++++++++++++++++ .../triattention/triattention_kernels.py | 26 +++++++++++++++++++ .../test_triattention_selection_compaction.py | 25 ++++++++++++++++++ 3 files changed, 75 insertions(+) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 2e191f165070..7fa3191b7b98 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -29,6 +29,7 @@ import torch import triton +from tensorrt_llm._torch.distributed import allgather from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheCompressionManager @@ -41,6 +42,7 @@ from ..compaction import build_compaction_params, compact from .triattention_kernels import ( + _fold_union_ranks_kernel, _gather_mean_phase_kernel, _settle_ties_kernel, prepare_per_head_scores, @@ -598,6 +600,22 @@ def _execute_eviction_round( request_count, cu_stream, ) + if self._union_tp_mapping is not None: + # Max-fold the rank-local unions into the global union (exact: + # max is order-free), so every rank keeps the same ordinals. + gathered = allgather( + self._selection_scores_rows[:request_count], + self._union_tp_mapping, + dim=0, + ) + width = int(self._selection_scores_rows.shape[1]) + _fold_union_ranks_kernel[(request_count, triton.cdiv(width, 1024))]( + gathered, + self._selection_scores_rows, + request_count, + TP_SIZE=self._union_tp_size, + WIDTH=width, + ) with nvtx_range("triattention.select", color="yellow"): if not union: # Per-head reduces read each decode window straight out of the scratch. @@ -853,6 +871,12 @@ def _rebuild_eviction_runtime( 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] + # Union reduces over ALL heads: rank-local rows are max-folded across the + # TP group each round so the kept set matches the single-rank algorithm. + self._union_tp_mapping = ( + mapping if (self.eviction_mode == "union" and tp_size > 1) else None + ) + self._union_tp_size = tp_size q_real, q_imag, mlr_coef, freq_scale_sq = ( tensor.to(device=device, dtype=torch.float32).contiguous() for tensor in (q_real, q_imag, mlr_coef, self._freq_scale_sq) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 175f5465fd35..93ebfc81dcc9 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -259,6 +259,32 @@ def prepare_per_head_scores( ) +@triton.jit +def _fold_union_ranks_kernel( + gathered_rows, + selection_rows, + request_count, + TP_SIZE: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr = 1024, +): + """Fold the TP-gathered rank-local union rows into the global union row + (elementwise max over the rank blocks).""" + 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_rows + request * WIDTH + token, folded, mask=mask) + + @triton.jit def _settle_ties_kernel( selection_scores_rows, 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 index af60affecc8e..faf7749f2ecb 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -815,3 +815,28 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): swa_after.index_select(2, swa_destination), swa_before.index_select(2, swa_source), ) + + +def test_fold_union_ranks_matches_max_oracle(): + """The TP union fold is an exact elementwise max over the gathered rank blocks.""" + import triton + + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( + _fold_union_ranks_kernel, + ) + + 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_kernel[(request_count, triton.cdiv(width, 1024))]( + gathered, + folded, + request_count, + TP_SIZE=tp_size, + WIDTH=width, + ) + expected = gathered.view(tp_size, request_count, width).amax(dim=0) + torch.cuda.synchronize(device) + assert torch.equal(folded, expected) From faba603e93606f103f247f0da0077d200297274c Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 26 Jul 2026 02:52:06 -0700 Subject: [PATCH 158/178] [None][chore] Align TriAttention kernel-boundary names to the module standard and slim dead guards Census batch (knife 47 A/B, 39 items): rename all pre-standard kernel parameters and locals to the source_length/prompt_length/decode_length families, collapse capacity aliases (seq_len/width/max_requests/ decode_width), replace the per-round eviction-input dict with a NamedTuple, hoist per-round constants to rebuild time, and delete constructor-guaranteed clamps, defaults, and single-use wrappers. Behaviour-preserving: kvcc 72 + executor 13 green, digest triplet token-identical vs the remerge bank. Signed-off-by: tianruih --- .../triattention/triattention.py | 408 +++++++++--------- .../triattention_cute_score_fused.py | 231 +++++----- .../triattention_cute_selection.py | 70 ++- .../triattention/triattention_kernels.py | 73 ++-- .../_torch/kv_cache_compression/conftest.py | 52 +-- .../test_triattention_cute_score.py | 30 +- .../test_triattention_cute_union_fusion.py | 69 +-- .../test_triattention_draft_cocompaction.py | 7 +- .../test_triattention_fused_settle_pack.py | 12 +- .../test_triattention_pipeline.py | 53 +-- .../test_triattention_selection_compaction.py | 71 +-- 11 files changed, 527 insertions(+), 549 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 7fa3191b7b98..b077ddae93c4 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -23,7 +23,7 @@ official tool (github.com/WeianMao/triattention) and is converted at load. """ -from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple import cuda.bindings.driver as cuda_driver import torch @@ -64,14 +64,27 @@ ) # Upper bound of the geometric integration offset ladder [1, 2, 4, ...]. -_OFFSET_MAX_LENGTH = 65536 +_MAX_INTEGRATION_OFFSET = 65536 + + +class _EvictionInput(NamedTuple): + """One due request's eviction operands for a single round.""" + + request: "LlmRequest" + target_cache: object + draft_cache: Optional[object] + state: Dict[str, object] + source_length: int + logical_source_length: int + prompt_length: int + target_tail_length: int def _allocate_block_offset_staging( anchor_pool: torch.Tensor, *, num_pools: int, - max_requests: int, + request_capacity: int, token_capacity: int, max_source_blocks: int, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -83,7 +96,7 @@ def _allocate_block_offset_staging( tokens_per_block = int(anchor_pool.shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block staged_blocks_per_seq = min((page_count + 3) // 4 * 4, int(max_source_blocks)) - shape = (num_pools, max_requests, 2, staged_blocks_per_seq) + shape = (num_pools, request_capacity, 2, staged_blocks_per_seq) host = torch.empty(shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) device_table = torch.empty(shape, dtype=torch.int32, device=anchor_pool.device) return host, device_table @@ -156,6 +169,7 @@ def __init__( self._inflight_scheduled_batch: Optional[object] = None self._inflight_generation_request_ids: Optional[set] = None # Manager-lifetime constants. + self._num_extra_kv_tokens = int(kv_cache_manager.num_extra_kv_tokens) self._protected_tail_capacity = ( int(kv_cache_manager.num_extra_kv_tokens) + int(kv_cache_manager._kv_reserve_draft_tokens) @@ -201,7 +215,21 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int config_values = config.get_text_config().to_dict() layer_types = config_values.get("layer_types") if not layer_types: - if self._has_sliding_window_signal(config_values): + 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" @@ -236,30 +264,16 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int window_size = raw_window return (dense_layers, swa_layers, window_size) - @staticmethod - def _has_sliding_window_signal(config: Dict[str, object]) -> bool: - use_sliding_window = config.get("use_sliding_window") - if isinstance(use_sliding_window, bool): - return use_sliding_window - return any( - config.get(field) - for field in ( - "sliding_window", - "sliding_window_size", - "sliding_window_pattern", - "max_window_layers", - ) - ) - def _load_calibration(self) -> None: - self.calibration = self._resolve_calibration() - self._freq_scale_sq = self.calibration["freq_scale_sq"].to(dtype=torch.float32) + calibration = self._resolve_calibration() + self._freq_scale_sq = calibration["freq_scale_sq"].to(dtype=torch.float32) + self._omega = calibration["omega"] # Pre-split query stats + MLR coefficient, shapes [L, H, F]. - _Eq = self.calibration["E_q"] - self._triattn_q_real = _Eq.real.to(torch.float32).contiguous() - self._triattn_q_imag = _Eq.imag.to(torch.float32).contiguous() - self._triattn_mlr_coef = ( - self.calibration["E_q_norm"].to(torch.float32) - _Eq.abs().to(torch.float32) + e_q = calibration["E_q"] + 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 = ( + calibration["E_q_norm"].to(torch.float32) - e_q.abs().to(torch.float32) ).contiguous() def _resolve_calibration(self) -> Dict[str, torch.Tensor]: @@ -279,7 +293,7 @@ def _resolve_calibration(self) -> Dict[str, torch.Tensor]: def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: """Convert the official calibration format to the runtime schema.""" stats = raw["stats"] - meta = raw.get("metadata", {}) + meta = raw["metadata"] if "sampled_heads" in meta: heads = [(int(a), int(b)) for a, b in meta["sampled_heads"]] else: @@ -314,8 +328,10 @@ def _rope_tables(self, freq_count: int): from transformers import AutoConfig from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS - cfg = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True).get_text_config() - config_values = cfg.to_dict() + config = AutoConfig.from_pretrained( + self.model_path, trust_remote_code=True + ).get_text_config() + config_values = config.to_dict() rope_params = ( config_values.get("rope_parameters") or config_values.get("rope_scaling") or {} ) @@ -333,7 +349,7 @@ def _rope_tables(self, freq_count: int): omega = (1.0 / (base ** (positions / head_dim)))[:freq_count].clone() scale_sq = 1.0 else: - inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](cfg, device="cpu") + inv_freq, attention_factor = ROPE_INIT_FUNCTIONS[rope_type](config, device="cpu") omega = inv_freq.to(torch.float32)[:freq_count].clone() scale_sq = float(attention_factor) ** 2 return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) @@ -342,13 +358,11 @@ def _rope_tables(self, freq_count: int): def on_request_init(self, request: "LlmRequest", **kwargs) -> None: """Register the request for eviction tracking.""" - request_id = request.py_request_id - if request_id not in self._request_states: - self._validate_request_capacity(request) - self._request_states[request_id] = { - "generation_steps": 0, - "evicted_tokens": 0, - } + self._validate_request_capacity(request) + self._request_states[request.py_request_id] = { + "generation_steps": 0, + "evicted_tokens": 0, + } def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Snapshot the in-flight batch; mutation remains in final update.""" @@ -407,7 +421,7 @@ def _evict_due_requests( """Owner of the full eviction transaction: admission, cadence, launch, publication, and cache resize.""" manager = self.kv_cache_manager - eviction_inputs: List[Dict[str, object]] = [] + eviction_inputs: List[_EvictionInput] = [] with nvtx_range("triattention.metadata", color="cyan"): for request in scheduled_batch.generation_requests: if request.is_dummy or request.state in _SKIP_REQUEST_STATES: @@ -426,16 +440,16 @@ def _evict_due_requests( if previous_step // self.beta >= step // self.beta: continue # Speculative reserve + in-flight overlap growth: contiguous tail moved byte-for-byte. - target_tail_length = int( - manager.num_extra_kv_tokens - ) + self._inflight_generation_growth(scheduled_batch, request_id) + target_tail_length = self._num_extra_kv_tokens + ( + self._inflight_generation_growth(scheduled_batch, request_id) + ) 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 = min(int(request.py_prompt_len), source_length) + prompt_length = int(request.py_prompt_len) if source_length <= prompt_length + self.budget: # Selection would be an identity: nothing to evict yet. continue @@ -447,16 +461,17 @@ def _evict_due_requests( # Target and draft defer together (pre-launch). continue eviction_inputs.append( - { - "request": request, - "target_cache": target_cache, - "draft_cache": draft_cache, - "source_length": source_length, + _EvictionInput( + request=request, + target_cache=target_cache, + draft_cache=draft_cache, + state=state, + source_length=source_length, # Uncompressed logical position. - "logical_source_length": source_length + state["evicted_tokens"], - "prompt_length": prompt_length, - "target_tail_length": target_tail_length, - } + logical_source_length=source_length + state["evicted_tokens"], + prompt_length=prompt_length, + target_tail_length=target_tail_length, + ) ) if not eviction_inputs: return @@ -478,12 +493,10 @@ def _evict_due_requests( ): self._execute_eviction_round(eviction_inputs) for item in eviction_inputs: - request = item["request"] - evicted = item["source_length"] - item["prompt_length"] - self.budget - state = self._request_states[request.py_request_id] - state["evicted_tokens"] += evicted + evicted = item.source_length - item.prompt_length - self.budget + item.state["evicted_tokens"] += evicted # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). - request.py_num_compressed_tokens = state["evicted_tokens"] + item.request.py_num_compressed_tokens = item.state["evicted_tokens"] self._resize_compacted_caches(eviction_inputs) def _inflight_generation_growth( @@ -502,27 +515,27 @@ def _inflight_generation_growth( def _execute_eviction_round( self, - eviction_inputs: Sequence[Dict[str, object]], + eviction_inputs: Sequence[_EvictionInput], ) -> None: """Run one eviction round over the due cohort (every launch covers the full request capacity; padded rows carry zero lengths and stay inert).""" manager = self.kv_cache_manager draft_manager = self.draft_kv_cache_manager with nvtx_range_debug("triattention.page_table_stage", color="orange"): - request_ids = [item["request"].py_request_id for item in eviction_inputs] - round_starts = [item["logical_source_length"] for item in eviction_inputs] - token_starts = [item["prompt_length"] for item in eviction_inputs] - seq_lens = [item["source_length"] for item in eviction_inputs] + request_ids = [item.request.py_request_id for item in eviction_inputs] + logical_source_lengths = [item.logical_source_length for item in eviction_inputs] + prompt_lengths = [item.prompt_length for item in eviction_inputs] + source_lengths = [item.source_length for item in eviction_inputs] dense_move_offsets, swa_move_offsets, draft_move_offsets = ( self._compute_compaction_move_offsets(eviction_inputs) ) stream = torch.cuda.current_stream(self._block_offsets_device.device) # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. - max_round_start = max(round_starts) + max_logical_source_length = max(logical_source_lengths) rows = ( - (0, round_starts), - (1, seq_lens), - (2, token_starts), + (0, logical_source_lengths), + (1, source_lengths), + (2, prompt_lengths), (3, dense_move_offsets), (4, swa_move_offsets), (5, draft_move_offsets), @@ -541,7 +554,7 @@ def _execute_eviction_round( host_table[row, : len(values)] = values # Zero lengths keep the score kernel and selection inert for padded rows. host_table[:3, len(eviction_inputs) :] = 0 - grow_mean_phase_table(self._phase, int(max_round_start) + 1) + grow_mean_phase_table(self._phase, int(max_logical_source_length) + 1) self._stage_block_offsets( manager, request_ids, @@ -569,7 +582,6 @@ def _execute_eviction_round( self._logical_source_lengths_device, self._phase["cos"], self._phase["sin"], - self._phase["rows"], self._source_lengths_device, self._prompt_lengths_device, self._mean_cos, @@ -593,7 +605,7 @@ def _execute_eviction_round( ) if union: # Normalized union reduction, written straight into the selection rows. - self._compiled_normalize_union[request_count]( + self._compiled_normalize_union_by_request_count[request_count]( self._cute_partial_stats, *self._cute_selection_prefix, self._cute_selection_scores_rows, @@ -608,13 +620,12 @@ def _execute_eviction_round( self._union_tp_mapping, dim=0, ) - width = int(self._selection_scores_rows.shape[1]) - _fold_union_ranks_kernel[(request_count, triton.cdiv(width, 1024))]( + _fold_union_ranks_kernel[(request_count, self._fold_width_blocks)]( gathered, self._selection_scores_rows, request_count, TP_SIZE=self._union_tp_size, - WIDTH=width, + WIDTH=self._selection_width_capacity, ) with nvtx_range("triattention.select", color="yellow"): if not union: @@ -630,8 +641,10 @@ def _execute_eviction_round( request_count=request_count, num_layers=self._num_layers, num_q_heads=self._num_q_heads, + num_kv_heads=self._num_kv_heads, padded_head_columns=self._padded_head_columns, score_token_capacity=self._score_token_capacity, + selection_width=self._selection_width_capacity, per_layer=self.eviction_mode == "per_layer_perhead", normalize_scores=self.normalize_scores, ) @@ -647,7 +660,7 @@ def _execute_eviction_round( def _compute_compaction_move_offsets( self, - eviction_inputs: Sequence[Dict[str, object]], + eviction_inputs: Sequence[_EvictionInput], ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: """Cumulative dense/SWA/draft move offsets for one due cohort (keep set plus protected tail per request; rows past the cohort repeat the final @@ -660,7 +673,7 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets.extend(offsets[-1:] * (self._request_capacity - len(moves_per_request))) return offsets - tails = [int(item["target_tail_length"]) for item in eviction_inputs] + tails = [int(item.target_tail_length) for item in eviction_inputs] dense = padded_offsets([self._keep_count + tail for tail in tails]) swa = None if self._swa_window is not None: @@ -733,16 +746,10 @@ def _resize_compacted_caches(self, eviction_inputs) -> None: families.append(("draft", "draft_cache", self._draft_protected_tail_capacity)) for label, cache_key, fixed_tail in families: for item in eviction_inputs: - cache = item[cache_key] - request_id = item["request"].py_request_id - if not cache.is_active: - # Bytes already moved: the compact-to-resize window is owned by this hook. - raise RuntimeError( - f"Request {request_id} {label} KV cache was " - "suspended between compact and resize" - ) - tail = item["target_tail_length"] if fixed_tail is None else fixed_tail - resized_capacity = item["prompt_length"] + self.budget + tail + cache = getattr(item, cache_key) + request_id = item.request.py_request_id + tail = item.target_tail_length if fixed_tail is None else fixed_tail + resized_capacity = item.prompt_length + self.budget + tail if not cache.resize(resized_capacity, None): raise RuntimeError( f"Failed to resize compacted {label} KV cache for " @@ -755,18 +762,13 @@ def _ensure_eviction_runtime( self, target_layout: Dict[str, object], draft_layout: Optional[Dict[str, object]], - eviction_inputs: Sequence[Dict[str, object]], + eviction_inputs: Sequence[_EvictionInput], ) -> None: """Per-round reuse gate over the three capacity axes; first round or growth replaces the resident runtime as a whole.""" # Empty cohorts never reach here: _evict_due_requests no-ops pre-launch. - needed_width = max( - item["source_length"] - item["prompt_length"] for item in eviction_inputs - ) - needed_score_tokens = max(item["source_length"] for item in eviction_inputs) - needed_page_tokens = max( - item["source_length"] + item["target_tail_length"] for item in eviction_inputs - ) + needed_width = max(item.source_length - item.prompt_length for item in eviction_inputs) + needed_score_tokens = max(item.source_length for item in eviction_inputs) needed_requests = len(eviction_inputs) if self._buffers_built: if ( @@ -781,35 +783,38 @@ def _ensure_eviction_runtime( self._compaction_done_event.synchronize() self._buffers_built = False - mgr = self.kv_cache_manager - request_capacity = max(needed_requests, int(mgr.max_batch_size)) + needed_page_tokens = max( + item.source_length + item.target_tail_length for item in eviction_inputs + ) + manager = self.kv_cache_manager + request_capacity = max(needed_requests, int(manager.max_batch_size)) selection_width_capacity = max( needed_width, - self.budget + 2 * self.beta + int(mgr.max_total_draft_tokens or 0), + self.budget + 2 * self.beta + int(manager.max_total_draft_tokens or 0), ) # Bucket sized by the presented cohorts, NOT max_seq_len (a floor there breaks 32-bit indexing). score_token_capacity = next_positive_power_of_2(max(int(needed_page_tokens), 1024)) score_token_capacity = min( - score_token_capacity, max(int(mgr.max_seq_len), int(needed_page_tokens)) + score_token_capacity, max(int(manager.max_seq_len), int(needed_page_tokens)) ) # The bucket capacity must be tile-aligned (mis-tiling stripes the # score scratch silently); the ceiling division constructs that fact. - score_tile_tokens = max(64, int(mgr.tokens_per_block)) + score_tile_tokens = max(64, int(manager.tokens_per_block)) score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens first_pool = target_layout["layer_pools"][target_layout["dense_layers"][0]] if self._phase is None: # Host-only width offsets for the table builder (no device copy). self._phase = { - "omega": self.calibration["omega"] - .to(device=first_pool.device, dtype=torch.float32) - .contiguous(), - "offset_values": [float(1 << i) for i in range(_OFFSET_MAX_LENGTH.bit_length())], + "omega": self._omega.to(device=first_pool.device, dtype=torch.float32).contiguous(), + "offset_values": [ + float(1 << i) for i in range(_MAX_INTEGRATION_OFFSET.bit_length()) + ], "cos": None, "sin": None, "rows": 0, } - grow_mean_phase_table(self._phase, max(int(score_token_capacity), 1)) + grow_mean_phase_table(self._phase, int(score_token_capacity)) self._rebuild_eviction_runtime( target_layout, draft_layout, @@ -836,6 +841,7 @@ def _rebuild_eviction_runtime( from .triattention_cute_score_fused import ( _COMPILE_LOCK, _COMPILED_KERNELS, + PADDED_HEAD_COLUMNS, SMALL_WORKLOAD_PAGE_SHARDS, STATS_FIELDS, _encode_tma_descriptors, @@ -843,7 +849,6 @@ def _rebuild_eviction_runtime( _to_cute, _TriAttentionScoreKernel, ) - from .triattention_cute_score_fused import N as PADDED_HEAD_COLUMNS layer_pools = target_layout["layer_pools"] dense_layers = list(target_layout["dense_layers"]) @@ -854,14 +859,14 @@ def _rebuild_eviction_runtime( num_page_table_slots = int(self.kv_cache_manager.num_pools) # The first dense layer anchors device and staging geometry. - p0 = layer_pools[dense_layers[0]] - device = p0.device - max_requests = int(request_capacity) - seq_len = int(score_token_capacity) - decode_width = int(selection_width_capacity) + anchor_pool = layer_pools[dense_layers[0]] + device = anchor_pool.device + request_capacity = int(request_capacity) + score_token_capacity = int(score_token_capacity) + selection_width_capacity = int(selection_width_capacity) keep_count = int(self.budget) protected_tail_capacity = int(self._protected_tail_capacity) - page_table_token_capacity = seq_len + protected_tail_capacity + page_table_token_capacity = score_token_capacity + protected_tail_capacity q_real, q_imag, mlr_coef = self._local_score_calibration(target_layout["global_layers"]) # TP splits attention heads contiguously per rank; the calibration is global. @@ -884,16 +889,17 @@ def _rebuild_eviction_runtime( num_q_heads = int(q_real.shape[1]) num_freqs = int(q_real.shape[2]) - self._request_capacity = max_requests - self._score_token_capacity = seq_len - self._selection_width_capacity = decode_width + self._request_capacity = request_capacity + self._score_token_capacity = score_token_capacity + self._selection_width_capacity = selection_width_capacity + self._fold_width_blocks = triton.cdiv(selection_width_capacity, 1024) self._keep_count = keep_count # ---- block-offset staging (target, plus the co-compressed draft) ------- self._block_offsets_host, self._block_offsets_device = _allocate_block_offset_staging( - p0, + anchor_pool, num_pools=num_page_table_slots, - max_requests=max_requests, + request_capacity=request_capacity, token_capacity=page_table_token_capacity, max_source_blocks=int(self.kv_cache_manager.host_kv_cache_block_offsets.shape[-1]), ) @@ -913,8 +919,8 @@ def _rebuild_eviction_runtime( _allocate_block_offset_staging( draft_anchor_pool, num_pools=int(self.draft_kv_cache_manager.num_pools), - max_requests=max_requests, - token_capacity=seq_len + int(self._draft_protected_tail_capacity), + request_capacity=request_capacity, + token_capacity=score_token_capacity + int(self._draft_protected_tail_capacity), max_source_blocks=int( self.draft_kv_cache_manager.host_kv_cache_block_offsets.shape[-1] ), @@ -923,21 +929,21 @@ def _rebuild_eviction_runtime( # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- self._request_metadata_host = torch.empty( - (6, max_requests + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + (6, request_capacity + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) # numpy view over the pinned rows: per-round staging writes lists in place. self._request_metadata_host_np = self._request_metadata_host.numpy() self._identity_copy_indices_host = torch.arange( - max_requests, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + request_capacity, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() ) # Zero-filled: an unstaged cohort must gather the phase table's row 0. self._request_metadata_device = torch.zeros( - (6, max_requests + 1), dtype=torch.int32, device=device + (6, request_capacity + 1), dtype=torch.int32, device=device ) - self._logical_source_lengths_device = self._request_metadata_device[0, :max_requests] - self._source_lengths_device = self._request_metadata_device[1, :max_requests] + self._logical_source_lengths_device = self._request_metadata_device[0, :request_capacity] + self._source_lengths_device = self._request_metadata_device[1, :request_capacity] # Pinned per-request decode-window starts. - self._prompt_lengths_device = self._request_metadata_device[2, :max_requests] + self._prompt_lengths_device = self._request_metadata_device[2, :request_capacity] dense_move_offsets_row = self._request_metadata_device[3] swa_move_offsets_row = self._request_metadata_device[4] draft_move_offsets_row = self._request_metadata_device[5] @@ -948,55 +954,56 @@ def _rebuild_eviction_runtime( torch.empty_like(self._prompt_lengths_device) if swa_layers else None ) self._swa_rebase_delta = keep_count - self._swa_window if swa_layers else 0 - self._mean_cos = torch.empty((max_requests, num_freqs), dtype=torch.float32, device=device) + self._mean_cos = torch.empty( + (request_capacity, num_freqs), dtype=torch.float32, device=device + ) self._mean_sin = torch.empty_like(self._mean_cos) self._phase_num_freqs = int(self._phase["omega"].numel()) self._phase_f_block = triton.next_power_of_2(self._phase_num_freqs) # ---- score state: one fused group across all dense layers -------------- - _, _, num_kv_heads, tokens_per_block, _ = p0.shape + _, _, num_kv_heads, tokens_per_block, _ = anchor_pool.shape self._num_layers = len(dense_layers) self._num_q_heads = int(num_q_heads) self._num_kv_heads = int(num_kv_heads) dense_layer_slots = [layer_pool_ids[layer] for layer in dense_layers] - seg_req_id = torch.arange(max_requests, dtype=torch.int32, device=device).repeat_interleave( - self._num_layers - ) + seg_req_id = torch.arange( + request_capacity, dtype=torch.int32, device=device + ).repeat_interleave(self._num_layers) seg_layer_id = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( - max_requests + request_capacity ) block_offsets = self._block_offsets_device slots_t = torch.tensor(dense_layer_slots, dtype=torch.int64, device=device) req_idx = seg_req_id.to(torch.int64) - slot_idx = slots_t.repeat(max_requests) + slot_idx = slots_t.repeat(request_capacity) seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - max_segments = max_requests * self._num_layers + max_segments = request_capacity * self._num_layers # The score plane must stay 32-bit indexable (wraparound = silent wild read). - if (PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len >= 2**31: + if (PADDED_HEAD_COLUMNS - 1) * max_segments * score_token_capacity >= 2**31: raise ValueError( "score bucket overflows the 32-bit score plane: " - f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * seq_len}" + f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * score_token_capacity}" ) # Persistent buffers: the compiled kernels capture their device pointers. self._padded_head_columns = PADDED_HEAD_COLUMNS self._score_scratch = torch.empty( - self._num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * seq_len, + self._num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * score_token_capacity, dtype=torch.float32, device=device, ) # int32 is safe here: covered by the 2^31 score-plane audit above. seg_out_offset = ( - torch.arange(max_segments, dtype=torch.int64, device=device) * seq_len + torch.arange(max_segments, dtype=torch.int64, device=device) * score_token_capacity ).to(torch.int32) union = self.eviction_mode == "union" # ---- score path: compiled per request-count/page-shard ---- - anchor_pool = p0 sm_count = int(torch.cuda.get_device_properties(device).multi_processor_count) # Per-shard partial score statistics. partial_stats_elements = ( - max_requests + request_capacity * self._num_layers * num_q_heads * SMALL_WORKLOAD_PAGE_SHARDS @@ -1015,8 +1022,8 @@ def _rebuild_eviction_runtime( int(num_freqs), int(tokens_per_block), ) - # Alignment sits with the operand it describes; valid_seq_lens and - # token_starts are only 4-byte-aligned row views, read as per-CTA scalars. + # Alignment sits with the operand it describes; source_lengths and + # prompt_lengths are only 4-byte-aligned row views, read as per-CTA scalars. prefix_operands = ( (block_offsets.view(-1), 16), (seg_page_off, 16), @@ -1053,8 +1060,8 @@ def _rebuild_eviction_runtime( self._cute_mean_cos = _to_cute(self._mean_cos.view(-1)) self._cute_mean_sin = _to_cute(self._mean_sin.view(-1)) self._compiled_score_by_request_count: Dict[int, object] = {} - self._compiled_normalize_union: Dict[int, object] = {} - page_shards_by_count: Dict[int, int] = {} + self._compiled_normalize_union_by_request_count: Dict[int, object] = {} + page_shards_by_request_count: Dict[int, int] = {} self._cute_selection_prefix = ( _to_cute(self._score_scratch), _to_cute(self._source_lengths_device, assumed_align=4), @@ -1063,9 +1070,9 @@ def _rebuild_eviction_runtime( ) self._cute_partial_stats = _to_cute(self._partial_stats) static_geometry = ( - max_requests, + request_capacity, self._num_layers, - seq_len, + score_token_capacity, num_q_heads, self._num_kv_heads, num_freqs, @@ -1083,12 +1090,12 @@ def _rebuild_eviction_runtime( ) ) variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] - if max_requests > 1: - variants.append((max_requests, 2)) + if request_capacity > 1: + variants.append((request_capacity, 2)) # Per-head modes compile the score-only entry; union the fused pipeline. kernel_kwargs = dict( num_layers=self._num_layers, - seq_len=seq_len, + 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), @@ -1130,17 +1137,17 @@ def _compiled_kernel(cache_key, build): stream, ), ) - page_shards_by_count[request_count] = page_shards + page_shards_by_request_count[request_count] = page_shards - if max_requests > 1: - small = compiled_entries.get(1) - large = compiled_entries.get(max_requests) - for request_count in range(1, max_requests + 1): + if request_capacity > 1: + small = compiled_entries[1] + large = compiled_entries[request_capacity] + for request_count in range(1, request_capacity + 1): # Give small cohorts the extra shard while the 2-shard grid stays under two waves. two_shard_ctas = request_count * self._num_layers * self._num_kv_heads * 2 use_extra_score_shard = two_shard_ctas < 2 * sm_count compiled_entries[request_count] = small if use_extra_score_shard else large - page_shards_by_count[request_count] = ( + page_shards_by_request_count[request_count] = ( SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 ) @@ -1151,22 +1158,22 @@ def _compiled_kernel(cache_key, build): # ---- selection buffers (canonical row-major, one name per storage) ----- self._decode_lengths_device = torch.full( - (max_requests,), decode_width, dtype=torch.int32, device=device + (request_capacity,), selection_width_capacity, dtype=torch.int32, device=device ) if union: self._selection_rows_per_request = 1 self._selection_scores_rows = torch.empty( - (max_requests, decode_width), dtype=torch.float32, device=device + (request_capacity, selection_width_capacity), dtype=torch.float32, device=device ) # One selection row per request: its length IS the staged valid width. self._selection_row_lengths = self._decode_lengths_device # Padded rows still need in-range ordinals for the finalizer's gather. self._provisional_rows = torch.zeros( - (max_requests, keep_count), dtype=torch.int32, device=device + (request_capacity, keep_count), dtype=torch.int32, device=device ) # Kept decode ordinals. self._kept_ordinal_rows = torch.empty( - (max_requests, keep_count), dtype=torch.int32, device=device + (request_capacity, keep_count), dtype=torch.int32, device=device ) self._cute_selection_scores_rows = _to_cute(self._selection_scores_rows.view(-1)) else: @@ -1176,26 +1183,33 @@ def _compiled_kernel(cache_key, build): else self._num_layers * self._num_kv_heads ) # The selection rectangle must stay 32-bit indexable (wraparound = wild reads). - selection_rect = max_requests * selection_rows * max(decode_width, keep_count) + selection_rect = ( + request_capacity * selection_rows * max(selection_width_capacity, keep_count) + ) if selection_rect >= 2**31: raise ValueError( f"per-head selection rectangle overflows 32-bit indexing: {selection_rect}" ) self._selection_rows_per_request = selection_rows - score_shape = (max_requests, self._num_layers, self._num_q_heads, 1) + 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( - (max_requests * selection_rows, decode_width), dtype=torch.float32, device=device + (request_capacity * selection_rows, selection_width_capacity), + dtype=torch.float32, + device=device, ) self._selection_row_lengths = torch.full( - (max_requests * selection_rows,), decode_width, dtype=torch.int32, device=device + (request_capacity * selection_rows,), + selection_width_capacity, + dtype=torch.int32, + device=device, ) self._provisional_rows = torch.zeros( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + (request_capacity * selection_rows, keep_count), dtype=torch.int32, device=device ) self._kept_ordinal_rows = torch.empty( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device + (request_capacity * selection_rows, keep_count), dtype=torch.int32, device=device ) if union: @@ -1205,11 +1219,11 @@ def _compiled_kernel(cache_key, build): ) compiled_configs: Dict[Tuple[int, int, int, int], object] = {} - for request_count in range(1, max_requests + 1): - page_shards = page_shards_by_count[request_count] + 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, - seq_len, + score_token_capacity, sm_count, ) config_key = (page_shards, *config) @@ -1232,7 +1246,7 @@ def _compiled_kernel(cache_key, build): row_cluster_ctas=row_cluster_ctas: cute.compile( _TriAttentionNormalizeUnionKernel( num_layers=self._num_layers, - seq_len=seq_len, + score_token_capacity=score_token_capacity, num_q_heads=num_q_heads, # The finalizer maps real head rows onto N=8-padded planes. num_kv_heads=self._num_kv_heads, @@ -1240,7 +1254,7 @@ def _compiled_kernel(cache_key, build): tokens_per_lane=tokens_per_lane, token_subtiles=token_subtiles, row_cluster_ctas=row_cluster_ctas, - output_row_stride=decode_width, + output_row_stride=selection_width_capacity, ), self._cute_partial_stats, *self._cute_selection_prefix, @@ -1250,7 +1264,7 @@ def _compiled_kernel(cache_key, build): ), ) compiled_configs[config_key] = compiled_selection - self._compiled_normalize_union[request_count] = compiled_selection + self._compiled_normalize_union_by_request_count[request_count] = compiled_selection # ---- compaction plans (opaque: only compact() interprets them) --------- compaction_params = [ @@ -1298,27 +1312,27 @@ def _local_score_calibration( global_layers: List[int], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: num_layers = len(global_layers) - if global_layers and max(global_layers) >= self._triattn_q_real.shape[0]: + if global_layers and max(global_layers) >= self._calibration_q_real.shape[0]: raise ValueError( - f"TriAttention calibration has {self._triattn_q_real.shape[0]} layers, " + f"TriAttention calibration has {self._calibration_q_real.shape[0]} layers, " f"but this PP rank references global layer {max(global_layers)}" ) if global_layers == list(range(global_layers[0], global_layers[0] + num_layers)): layer_slice = slice(global_layers[0], global_layers[0] + num_layers) return ( - self._triattn_q_real[layer_slice], - self._triattn_q_imag[layer_slice], - self._triattn_mlr_coef[layer_slice], + self._calibration_q_real[layer_slice], + self._calibration_q_imag[layer_slice], + self._calibration_mlr_coef[layer_slice], ) layer_ids = torch.as_tensor( global_layers, - device=self._triattn_q_real.device, + device=self._calibration_q_real.device, dtype=torch.long, ) return ( - self._triattn_q_real.index_select(0, layer_ids), - self._triattn_q_imag.index_select(0, layer_ids), - self._triattn_mlr_coef.index_select(0, layer_ids), + self._calibration_q_real.index_select(0, layer_ids), + self._calibration_q_imag.index_select(0, layer_ids), + self._calibration_mlr_coef.index_select(0, layer_ids), ) def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: @@ -1326,10 +1340,15 @@ def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager cached = self._kv_layout_caches[draft] if cached is not None: - current_page_counts = self._pool_page_counts( - manager, - cached["global_layers"], - cached["pool_representatives"], + current_page_counts = tuple( + int( + manager.impl.get_page_index_upper_bound( + manager.layer_offsets[cached["global_layers"][layer]], + Role.KEY, + ) + ) + // int(manager.kv_factor) + for layer in cached["pool_representatives"] ) if current_page_counts != cached["pool_page_counts"]: raise RuntimeError( @@ -1357,7 +1376,7 @@ def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: dense_layers=dense_layers, swa_layers=swa_layers, swa_window=swa_window, - what="draft " if draft else "", + label="draft " if draft else "", ) self._kv_layout_caches[draft] = layout return layout @@ -1370,15 +1389,9 @@ def _build_runtime_kv_layout( dense_layers: List[int], swa_layers: List[int], swa_window: Optional[int], - what: str, + label: str, ) -> Dict[str, object]: - maybe_layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] - if any(pool is None for pool in maybe_layer_pools): - missing = [ - layer for layer, pool in zip(global_layers, maybe_layer_pools) if pool is None - ] - raise RuntimeError(f"Missing {what}KV pools for attention layers {missing}") - layer_pools = [pool for pool in maybe_layer_pools if pool is not None] + layer_pools = [manager.get_buffers(layer, kv_layout="HND") for layer in global_layers] # Canonical pool IDs, resolved once; every grouping derives from them # (V2 owns the mapping; its own lookup errors are the precise ones). layer_offsets = manager.layer_offsets @@ -1402,20 +1415,3 @@ def _build_runtime_kv_layout( int(layer_pools[layer].shape[0]) for layer in pool_representatives ), ) - - @staticmethod - def _pool_page_counts( - manager: KVCacheManagerV2, - global_layers: Sequence[int], - pool_representatives: Sequence[int], - ) -> Tuple[int, ...]: - return tuple( - int( - manager.impl.get_page_index_upper_bound( - manager.layer_offsets[global_layers[layer]], - Role.KEY, - ) - ) - // int(manager.kv_factor) - for layer in pool_representatives - ) 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 index d6fa3dc3b7ca..ca3b5c6107be 100644 --- 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 @@ -39,8 +39,9 @@ def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Fl CTA_M = 128 -# Minimum tcgen05 MMA tile N; GQA groups below 8 ride zero-padded head columns. -N = 8 +# 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 @@ -48,9 +49,9 @@ def _sqrt_approx_ftz(value: cutlass.Float32, *, loc=None, ip=None) -> cutlass.Fl STATS_FIELDS = 3 STATS_MEAN = 1 STATS_M2 = 2 -# Stats smem scratch: N score origins + one (sum, square-sum) pair per (warp, head column). -STATS_ORIGIN_SLOTS = N -STATS_SCRATCH_ELEMENTS = STATS_ORIGIN_SLOTS + (EPILOGUE_THREADS // 32) * N * 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 @@ -71,7 +72,7 @@ def __init__( self, *, num_layers: int, - seq_len: int, + score_token_capacity: int, num_q_heads: int, num_freqs: int, pool_shape: tuple[int, int, int, int, int], @@ -95,7 +96,7 @@ def __init__( if page_shards not in _SUPPORTED_PAGE_SHARDS: raise ValueError("TriAttention CuTe score has unsupported page shards") - self.seq_len = seq_len + 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 @@ -109,7 +110,7 @@ def __init__( # 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 = (seq_len + CTA_M - 1) // CTA_M + self.max_tiles = (score_token_capacity + CTA_M - 1) // CTA_M # Producer staging constants baked into the generated code. self.prefetch_depth = 4 @@ -125,68 +126,12 @@ def __init__( 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_slot, self.s_dim = pool_strides - if self.s_slot != 2 * num_freqs or self.s_dim != 1: + 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") - def epilog_tmem_copy_and_partition( - self, - tidx: cutlass.Int32, - accumulator: cute.Tensor, - output: cute.Tensor, - epilogue_tile: cute.Tile, - ) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: - copy_atom = sm100_utils.get_tmem_load_op( - self.cta_tile_shape_mnk, - self.c_layout, - self.c_dtype, - self.acc_dtype, - epilogue_tile, - False, - ) - accumulator_epilogue = cute.flat_divide( - accumulator[((None, None), 0, 0)], - epilogue_tile, - ) - tiled_copy = tcgen05.make_tmem_copy( - copy_atom, - accumulator_epilogue[(None, None, 0, 0)], - ) - thread_copy = tiled_copy.get_slice(tidx) - thread_accumulator = thread_copy.partition_S(accumulator_epilogue) - output_epilogue = cute.flat_divide( - output[((None, None), 0, 0, None, None, None)], - epilogue_tile, - ) - thread_output = thread_copy.partition_D(output_epilogue) - register_accumulator = cute.make_rmem_tensor( - thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.acc_dtype, - ) - return tiled_copy, thread_accumulator, register_accumulator - - def epilog_gmem_copy_and_partition( - self, - tidx: cutlass.Int32, - tiled_copy: cute.TiledCopy, - output: cute.Tensor, - epilogue_tile: cute.Tile, - ) -> tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: - output_epilogue = cute.flat_divide( - output[((None, None), 0, 0, None, None, None)], - epilogue_tile, - ) - thread_copy = tiled_copy.get_slice(tidx) - thread_output = thread_copy.partition_D(output_epilogue) - register_output = cute.make_rmem_tensor( - thread_output[(None, None, None, 0, 0, 0, 0, 0)].shape, - self.c_dtype, - ) - copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) - return copy_atom, register_output, thread_output - @cute.jit def _stage_raw_band_copies( self, @@ -222,13 +167,13 @@ def _stage_raw_band_copies( @cute.jit def __call__( self, - page_ids: cute.Tensor, + block_offset_entries: cute.Tensor, seg_page_off: cute.Tensor, seg_req_id: cute.Tensor, seg_layer_id: cute.Tensor, - valid_seq_lens: cute.Tensor, + source_lengths: cute.Tensor, seg_out_offset: cute.Tensor, - token_starts: cute.Tensor, + prompt_lengths: cute.Tensor, q_real: cute.Tensor, q_imag: cute.Tensor, mlr_coef: cute.Tensor, @@ -237,16 +182,16 @@ def __call__( freq_scale_sq: cute.Tensor, output: cute.Tensor, partial_stats: cute.Tensor, - pool_template: 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, N, self.k_coeff) + self.mma_tiler = (CTA_M, PADDED_HEAD_COLUMNS, self.k_coeff) self.cta_tile_shape_mnk = self.mma_tiler - self.epi_tile = (CTA_M, N) + self.epi_tile = (CTA_M, PADDED_HEAD_COLUMNS) tiled_mma = sm100_utils.make_trivial_tiled_mma( cutlass.Float32, @@ -267,7 +212,7 @@ def __call__( # 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, N, self.num_freqs), + (CTA_M, PADDED_HEAD_COLUMNS, self.num_freqs), cutlass.BFloat16, 2 * self.raw_page_buffers, ) @@ -287,12 +232,12 @@ def __call__( ), stride=( self.s_dim, - self.s_slot, + self.s_token, (self.s_kv_head, self.s_page), ), ) raw_tma_source = cute.make_tensor( - pool_template.iterator, + anchor_pool.iterator, raw_tma_source_layout, ) raw_tma_atom, raw_tma_tensor = cpasync.make_tiled_tma_atom( @@ -303,7 +248,7 @@ def __call__( ) raw_bf16_b_smem_layout = sm100_utils.make_smem_layout_b( raw_bf16_tiled_mma, - (CTA_M, N, 2 * self.num_freqs), + (CTA_M, PADDED_HEAD_COLUMNS, 2 * self.num_freqs), cutlass.BFloat16, 1, ) @@ -317,13 +262,13 @@ def __call__( ) magnitude_fp16_a_smem_layout = sm100_utils.make_smem_layout_a( magnitude_lo_tiled_mma, - (CTA_M, N, self.num_freqs), + (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, N, self.num_freqs), + (CTA_M, PADDED_HEAD_COLUMNS, self.num_freqs), cutlass.Float16, 1, ) @@ -384,7 +329,7 @@ class SharedStorage: self.shared_storage = SharedStorage # The score-plane stride product can exceed 2^31, so it must reach the kernel as Int64. - sum_seq = cutlass.Int64(request_count * self.num_layers * self.seq_len) + 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, @@ -393,13 +338,13 @@ class SharedStorage: raw_tma_atom, raw_tma_tensor, raw_tma_descriptors, - page_ids, + block_offset_entries, seg_page_off, seg_req_id, seg_layer_id, - valid_seq_lens, + source_lengths, seg_out_offset, - token_starts, + prompt_lengths, q_real, q_imag, mlr_coef, @@ -408,7 +353,7 @@ class SharedStorage: freq_scale_sq, output, partial_stats, - sum_seq, + segment_tokens, raw_bf16_direct_a_smem_layout, raw_tma_smem_layout, raw_bf16_b_smem_layout, @@ -429,13 +374,13 @@ def kernel( raw_tma_atom: cute.CopyAtom, raw_tma_source: cute.Tensor, raw_tma_descriptors: cute.Tensor, - page_ids: cute.Tensor, + block_offset_entries: cute.Tensor, seg_page_off: cute.Tensor, seg_req_id: cute.Tensor, seg_layer_id: cute.Tensor, - valid_seq_lens: cute.Tensor, + source_lengths: cute.Tensor, seg_out_offset: cute.Tensor, - token_starts: cute.Tensor, + prompt_lengths: cute.Tensor, q_real: cute.Tensor, q_imag: cute.Tensor, mlr_coef: cute.Tensor, @@ -444,7 +389,7 @@ def kernel( freq_scale_sq: cute.Tensor, output: cute.Tensor, partial_stats: cute.Tensor, - sum_seq: cutlass.Int64, + 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, @@ -461,11 +406,11 @@ def kernel( kv_head = task % self.num_kv_heads req_id = seg_req_id[segment] layer_id = seg_layer_id[segment] - valid_seq_len = valid_seq_lens[req_id] + 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(token_starts[req_id]) + score_start = cutlass.Int32(prompt_lengths[req_id]) smem = utils.SmemAllocator() storage = smem.allocate(self.shared_storage) @@ -588,11 +533,11 @@ def kernel( 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((N,), cutlass.Float32) - stats_origins_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) - stats_sums_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) - stats_square_sums_m128 = cute.make_rmem_tensor((N,), cutlass.Float32) - for stats_head in cutlass.range_constexpr(N): + 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) @@ -611,24 +556,26 @@ def kernel( ) physical_fragments_arg = physical_page_fragments prefetched_fragments_arg = prefetched_page_fragments - shard_has_page = valid_seq_len > score_start and tile_start_token < valid_seq_len - empty_shard = valid_seq_len <= score_start or tile_start_token >= valid_seq_len + 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(page_ids[page_off + tile_index * self.fragments_per_phase]) + 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 < valid_seq_len: + if tile_start_token + fragment * self.box_tokens < source_length: fragment_page_id = ( cutlass.Int32( - page_ids[ + block_offset_entries[ page_off + tile_index * self.fragments_per_phase + fragment @@ -682,7 +629,7 @@ def kernel( ) cute.arch.mbarrier_init_fence() # Per-(head, frequency) score coefficients, split into bf16/fp16 value+residual pairs. - for weight_round in cutlass.range_constexpr(N * self.k_coeff // THREADS): + 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 @@ -691,7 +638,7 @@ def kernel( 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 < N): + 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 @@ -708,7 +655,7 @@ def kernel( value = scale * (qr * msin + qi * mcos) else: value = scale * cutlass.Float32(mlr_coef[calib_offset]) - if cutlass.const_expr(self.group_size < N): + 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 @@ -795,8 +742,8 @@ def kernel( ) raw_tma_producer_state.advance() while ( - valid_seq_len > score_start - and tile_start_token < valid_seq_len + source_length > score_start + and tile_start_token < source_length and tiles_processed < self.max_tiles ): physical_page = cutlass.Int32(0) @@ -853,12 +800,12 @@ def kernel( next_tile_start_token = tile_start_token + CTA_M * self.page_shards next_pages_processed = tiles_processed + 1 if ( - next_tile_start_token < valid_seq_len + next_tile_start_token < source_length and next_pages_processed < self.max_tiles ): next_page_id_lane0 = ( cutlass.Int32( - page_ids[ + block_offset_entries[ page_off + (tile_index + self.page_shards) * self.fragments_per_phase ] @@ -871,11 +818,11 @@ def kernel( next_fragment_page_id = next_page_id_lane0 if ( next_tile_start_token + fragment * self.box_tokens - < valid_seq_len + < source_length ): next_fragment_page_id = ( cutlass.Int32( - page_ids[ + block_offset_entries[ page_off + (tile_index + self.page_shards) * self.fragments_per_phase @@ -913,7 +860,7 @@ def kernel( 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 < valid_seq_len + next_tile_start_token < source_length and next_pages_processed < self.max_tiles ) prefetched_physical_page = cute.arch.shuffle_sync( @@ -1105,28 +1052,53 @@ def kernel( # 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 sum_seq must stay 64-bit; the plane stride can exceed 2^31. - output_offset = cutlass.Int64(kv_head * N) * sum_seq + out_base + tile_start_token + # 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, N, 1), + (CTA_M, PADDED_HEAD_COLUMNS, 1), stride=( 1, - sum_seq, - N * sum_seq, + 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 - tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition( - epilogue_tidx, tCtAcc, tCgC, self.epi_tile + 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, ) - simt_atom, tTR_rC, tTR_gC = self.epilog_gmem_copy_and_partition( - epilogue_tidx, tiled_copy_t2r, tCgC, self.epi_tile + 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)) @@ -1140,7 +1112,8 @@ def kernel( 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.seq_len + tile_start_token >= score_start + and tile_start_token + CTA_M <= self.score_token_capacity ): cute.copy( simt_atom, @@ -1150,7 +1123,7 @@ def kernel( else: output_token = tile_start_token + epilogue_tidx if cutlass.dynamic_expr( - output_token >= score_start and output_token < self.seq_len + output_token >= score_start and output_token < self.score_token_capacity ): cute.copy( simt_atom, @@ -1182,14 +1155,14 @@ def kernel( raw_tma_consumer_state.advance() if cutlass.const_expr(self.write_partial_stats): if tiles_processed == 0: - for stats_head in cutlass.range_constexpr(N): + 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 < valid_seq_len + stats_token >= score_start and stats_token < source_length ): - for stats_head in cutlass.range_constexpr(N): + for stats_head in cutlass.range_constexpr(PADDED_HEAD_COLUMNS): stats_delta = ( stats_page_scores_m128[stats_head] - stats_origins_m128[stats_head] ) @@ -1204,7 +1177,7 @@ def kernel( 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(N): + 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): @@ -1213,7 +1186,9 @@ def kernel( stats_square_sum, stats_offset ) if lane_idx == 0 and warp_idx < EPILOGUE_THREADS // 32: - stats_scratch_base = STATS_ORIGIN_SLOTS + (warp_idx * N + stats_head) * 2 + 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() @@ -1223,7 +1198,9 @@ def kernel( 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 * N + lane_idx) * 2 + 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 @@ -1232,7 +1209,7 @@ def kernel( 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 - valid_seq_len + 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) 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 index 22887ef8bbc8..0942551ecb32 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -25,11 +25,9 @@ from cutlass.cute.typing import Pointer as CutePointer from cutlass.cutlass_dsl import T, dsl_user_op -from .triattention_cute_score_fused import STATS_FIELDS as _STATS_FIELDS -from .triattention_cute_score_fused import STATS_M2, STATS_MEAN - # Single-sourced constants: the score file owns N and the stats layout; Triton owns the epsilon. -from .triattention_cute_score_fused import N as _PADDED_HEAD_COLUMNS +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 @@ -82,11 +80,6 @@ def _mapa_shared_cluster( ) -@cute.jit -def _mapa_cluster(smem_ptr, peer_rank): - return _mapa_shared_cluster(smem_ptr, peer_rank) - - @dsl_user_op def _ld_shared_cluster_f32( mapped_addr: CuteInt32, @@ -109,11 +102,6 @@ def _ld_shared_cluster_f32( ) -@cute.jit -def _ld_cluster_f32(mapped_addr): - return _ld_shared_cluster_f32(mapped_addr) - - 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( @@ -134,7 +122,7 @@ def __init__( self, *, num_layers: int, - seq_len: int, + score_token_capacity: int, num_q_heads: int, num_kv_heads: int, page_shards: int, @@ -145,12 +133,9 @@ def __init__( ) -> 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.score_head_pad = PADDED_HEAD_COLUMNS - self.score_group_size self.num_layers = num_layers - self.seq_len = seq_len - # The widest score window (the whole bucket) sizes the token-tile grid; - # output rows are the TopK selection rows with their own stride. - self.width = seq_len + 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 @@ -162,16 +147,18 @@ def __init__( self.reduce_threads = _REDUCE_THREADS self.reduce_warps = _REDUCE_WARPS self.row_cluster_ctas = row_cluster_ctas - self.num_token_tiles = (self.width + self.token_tile - 1) // self.token_tile + # 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, - valid_seq_lens: cute.Tensor, + source_lengths: cute.Tensor, seg_out_offset: cute.Tensor, - token_starts: cute.Tensor, + prompt_lengths: cute.Tensor, union_scores: cute.Tensor, request_count: cutlass.Int32, stream: cuda.CUstream, @@ -179,9 +166,9 @@ def __call__( kernel = self.kernel( partial_stats, scores, - valid_seq_lens, + source_lengths, seg_out_offset, - token_starts, + prompt_lengths, union_scores, request_count, ) @@ -213,7 +200,7 @@ def _reduce_and_store_union_rows( warp_max_ptr, score_copy_atom, request_idx: cutlass.Int32, - valid_width: cutlass.Int32, + decode_length: cutlass.Int32, first_token: cutlass.Int32, lane_idx: cutlass.Int32, from_cluster_peers: cutlass.Constexpr, @@ -229,10 +216,10 @@ def _reduce_and_store_union_rows( (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_cluster(warp_max_ptr, cutlass.Int32(peer_rank)) + remote_addr = _mapa_shared_cluster(warp_max_ptr, cutlass.Int32(peer_rank)) union_value = cute.arch.fmax( union_value, - _ld_cluster_f32( + _ld_shared_cluster_f32( remote_addr + shared_offset * (cutlass.Float32.width // 8) ), ) @@ -248,7 +235,7 @@ def _reduce_and_store_union_rows( # 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 <= valid_width): + ) 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, @@ -264,7 +251,7 @@ def _reduce_and_store_union_rows( else: for token_slot in cutlass.range_constexpr(self.tokens_per_lane): token = subtile_first_token + token_slot - if cutlass.dynamic_expr(token < valid_width): + if cutlass.dynamic_expr(token < decode_length): union_scores[request_idx * self.output_row_stride + token] = reduced_values[ token_slot ] @@ -274,9 +261,9 @@ def kernel( self, partial_stats: cute.Tensor, scores: cute.Tensor, - valid_seq_lens: cute.Tensor, + source_lengths: cute.Tensor, seg_out_offset: cute.Tensor, - token_starts: cute.Tensor, + prompt_lengths: cute.Tensor, union_scores: cute.Tensor, request_count: cutlass.Int32, ): @@ -296,8 +283,8 @@ def kernel( 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(token_starts[request_idx]) - valid_width = valid_seq_lens[request_idx] - score_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, @@ -416,17 +403,20 @@ def kernel( 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.seq_len + 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.seq_len % self.tokens_per_lane == 0 + 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 <= valid_width + and subtile_first_token + self.tokens_per_lane <= decode_length ): score_tile = _gmem_lane_tile( scores.iterator, @@ -446,7 +436,7 @@ def kernel( ) for token_slot in cutlass.range_constexpr(self.tokens_per_lane): token = subtile_first_token + token_slot - if cutlass.dynamic_expr(token < valid_width): + 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( @@ -505,7 +495,7 @@ def kernel( warp_max_ptr, score_copy_atom, request_idx, - valid_width, + decode_length, first_token, lane_idx, False, @@ -533,7 +523,7 @@ def kernel( warp_max_ptr, score_copy_atom, request_idx, - valid_width, + decode_length, first_token, lane_idx, True, diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 93ebfc81dcc9..51c82ce7fa8e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Triton kernels, launch helpers, and mean-phase table builders for TriAttention +"""Triton kernels and launch helpers for TriAttention (fp32 math; int64 past-2^31 flat offsets; masked ragged tails; scoring in the CuTe pack).""" from __future__ import annotations @@ -17,38 +17,35 @@ @triton.jit def _gather_mean_phase_kernel( - round_starts, + logical_source_lengths, phase_cos, phase_sin, - phase_rows, - valid_seq_lens, - token_starts, + source_lengths, + prompt_lengths, mean_cos, mean_sin, - valid_widths, + 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 valid widths and SWA landing bases.""" + """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(round_starts + request).to(tl.int64) - # Clamp stale or padded round starts instead of faulting. - table_row = tl.minimum(tl.maximum(table_row, 0), phase_rows - 1) + 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) - token_start = tl.load(token_starts + request) - tl.store(valid_widths + request, tl.load(valid_seq_lens + request) - token_start) + 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, token_start + swa_rebase_delta) + tl.store(swa_destination_bases + request, prompt_length + swa_rebase_delta) # ---- Selection: combine scores per mode, then finalize the top-k set ---- @@ -84,7 +81,7 @@ def _score_row_stats_kernel( 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 - valid_width = tl.load(decode_lengths + request) + decode_length = tl.load(decode_lengths + request) prompt_start = tl.load(prompt_lengths + request) score_row = ( score_scratch @@ -95,18 +92,18 @@ def _score_row_stats_kernel( score_sum = 0.0 for start in tl.static_range(0, WIDTH, BLOCK): token = start + lane - valid = token < valid_width + 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 / valid_width + mean = score_sum / decode_length square_sum = 0.0 for start in tl.static_range(0, WIDTH, BLOCK): token = start + lane - valid = token < valid_width + 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 / valid_width) + 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)) @@ -119,7 +116,7 @@ def _score_per_head_reduce_kernel( row_mean, row_inv_std, selection_scores, - selection_seq_lens, + selection_row_lengths, segment_tokens, NUM_LAYERS: tl.constexpr, NUM_Q_HEADS: tl.constexpr, @@ -139,14 +136,14 @@ def _score_per_head_reduce_kernel( selection_row = tl.program_id(1) token_block = tl.program_id(2) token = token_block * BLOCK + tl.arange(0, BLOCK) - valid_width = tl.load(decode_lengths + request) + decode_length = tl.load(decode_lengths + request) prompt_start = tl.load(prompt_lengths + request) - valid_token = token < valid_width + valid_token = token < decode_length if token_block == 0: tl.store( - selection_seq_lens + request * SELECTION_ROWS + selection_row, - valid_width, + selection_row_lengths + request * SELECTION_ROWS + selection_row, + decode_length, ) kv_head = selection_row % NUM_KV_HEADS @@ -210,16 +207,16 @@ def prepare_per_head_scores( request_count: int, num_layers: int, num_q_heads: int, + num_kv_heads: int, padded_head_columns: int, score_token_capacity: int, + selection_width: int, per_layer: bool, normalize_scores: bool, ) -> None: """Normalize and reduce the scratch's decode windows for either per-head eviction mode.""" - width = int(selection_scores_rows.shape[1]) - selection_rows = int(selection_scores_rows.shape[0]) // int(decode_lengths.shape[0]) - num_kv_heads = selection_rows // num_layers if per_layer else selection_rows + selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads rows = num_layers * num_q_heads segment_tokens = request_count * num_layers * score_token_capacity if normalize_scores: @@ -236,10 +233,12 @@ def prepare_per_head_scores( NUM_KV_HEADS=num_kv_heads, PADDED_COLUMNS=padded_head_columns, BUCKET=score_token_capacity, - WIDTH=width, + WIDTH=selection_width, ) # 256-token tiles match the reduce kernel's BLOCK default. - _score_per_head_reduce_kernel[(request_count, selection_rows, triton.cdiv(width, 256))]( + _score_per_head_reduce_kernel[ + (request_count, selection_rows, triton.cdiv(selection_width, 256)) + ]( score_scratch, decode_lengths, prompt_lengths, @@ -253,7 +252,7 @@ def prepare_per_head_scores( NUM_KV_HEADS=num_kv_heads, PADDED_COLUMNS=padded_head_columns, BUCKET=score_token_capacity, - WIDTH=width, + WIDTH=selection_width, PER_LAYER=per_layer, NORMALIZE=normalize_scores, ) @@ -262,7 +261,7 @@ def prepare_per_head_scores( @triton.jit def _fold_union_ranks_kernel( gathered_rows, - selection_rows, + selection_scores_rows, request_count, TP_SIZE: tl.constexpr, WIDTH: tl.constexpr, @@ -282,14 +281,14 @@ def _fold_union_ranks_kernel( other=-float("inf"), ) folded = tl.maximum(folded, value) - tl.store(selection_rows + request * WIDTH + token, folded, mask=mask) + tl.store(selection_scores_rows + request * WIDTH + token, folded, mask=mask) @triton.jit def _settle_ties_kernel( selection_scores_rows, selection_row_lengths, - token_starts, + prompt_lengths, provisional_rows, kept_ordinal_rows, WIDTH: tl.constexpr, @@ -307,7 +306,7 @@ def _settle_ties_kernel( 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_len = tl.load(token_starts + request) + prompt_length = tl.load(prompt_lengths + request) threshold = float("inf") for start in tl.static_range(0, KEEP_COUNT, BLOCK): @@ -327,11 +326,11 @@ def _settle_ties_kernel( ).to(tl.float32) threshold = tl.minimum(threshold, tl.min(selected_score, axis=0)) - seq_len = tl.load(selection_row_lengths + row) + 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 < seq_len) + valid = (token_index < WIDTH) & (token_index < row_length) score = tl.load( row_scores + token_index, mask=valid, @@ -344,7 +343,7 @@ def _settle_ties_kernel( 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 < seq_len) + valid = (token_index < WIDTH) & (token_index < row_length) score = tl.load( row_scores + token_index, mask=valid, @@ -359,7 +358,7 @@ def _settle_ties_kernel( write_offset = output_count + tl.cumsum(selected_i32, axis=0) - selected_i32 tl.store( row_output + write_offset, - token_index + prompt_len, + token_index + prompt_length, mask=selected, ) output_count += tl.sum(selected_i32) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 27466ca3b324..e5651ad28081 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -271,7 +271,7 @@ def make_buffer_stubs(manager, *, decode_width=260): ``_rebuild_eviction_runtime`` should set (production sets them in place).""" manager._freq_scale_sq = torch.ones(2) manager._phase = {"rows": 8} - manager.calibration = {"omega": torch.ones(2)} + manager._omega = torch.ones(2) manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) pool = torch.empty(8, 2, 1, 4, 4) layout = dict( @@ -391,21 +391,25 @@ def make_eviction_input( target_tail_length=0, target_cache=None, draft_cache=None, + state=None, ): """One due-cohort item shaped exactly like ``_evict_due_requests`` builds.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import _EvictionInput + if request is None: request = SimpleNamespace(py_request_id=request_id, py_num_compressed_tokens=0) - return { - "request": request, - "target_cache": target_cache, - "draft_cache": draft_cache, - "source_length": int(source_length), - "logical_source_length": int( + return _EvictionInput( + request=request, + target_cache=target_cache, + draft_cache=draft_cache, + state={"generation_steps": 0, "evicted_tokens": 0} if state is None else state, + source_length=int(source_length), + logical_source_length=int( source_length if logical_source_length is None else logical_source_length ), - "prompt_length": int(prompt_length), - "target_tail_length": int(target_tail_length), - } + prompt_length=int(prompt_length), + target_tail_length=int(target_tail_length), + ) def make_request(request_id, **overrides): @@ -441,7 +445,7 @@ def torch_tri_score_oracle( layer_pools, page_ids, seq_lens, - round_starts, + logical_source_lengths, q_real, q_imag, mlr_coef, @@ -455,7 +459,7 @@ def torch_tri_score_oracle( scores = [] num_q_heads = int(q_real.shape[1]) for request, seq_len in enumerate(seq_lens): - phase = (round_starts[request] + offsets[:, None]) * omega[None, :] + 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: @@ -577,9 +581,9 @@ def make_cute_buffers( manager.budget = keep_count manager._protected_tail_capacity = protected_tail_capacity manager._freq_scale_sq = freq_scale_sq - manager._triattn_q_real = q_real - manager._triattn_q_imag = q_imag - manager._triattn_mlr_coef = mlr_coef + manager._calibration_q_real = q_real + manager._calibration_q_imag = q_imag + manager._calibration_mlr_coef = mlr_coef manager._rebuild_eviction_runtime( layout, None, @@ -614,20 +618,20 @@ def rect_to_score_scratch(scores, num_kv_heads, padded_head_columns=8): return scratch, prompt_lengths -def stage_score_metadata(manager, request_count, valid_seq_lens, valid_widths, token_starts): +def stage_score_metadata(manager, request_count, source_lengths, decode_lengths, prompt_lengths): """Stage the per-round score metadata exactly like production (the compiled score launches read the staged rows via pointer capture).""" torch.sub( - valid_seq_lens[:request_count], - token_starts[:request_count], - out=valid_widths[:request_count], + source_lengths[:request_count], + prompt_lengths[:request_count], + out=decode_lengths[:request_count], ) - manager._source_lengths_device[:request_count].copy_(valid_seq_lens[:request_count]) - manager._prompt_lengths_device[:request_count].copy_(token_starts[: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, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin + manager, request_count, source_lengths, decode_lengths, prompt_lengths, mean_cos, mean_sin ): """The production score-only leg plus the decode-window gather (the round executor's per-head sequence, parameterized by count). Test mean phases @@ -635,7 +639,7 @@ def launch_split_scores( gather refresh; the compiled entry fires directly, like the round does.""" import cuda.bindings.driver as cuda_driver - stage_score_metadata(manager, request_count, valid_seq_lens, valid_widths, token_starts) + 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]) assert request_count in manager._compiled_score_by_request_count @@ -665,7 +669,7 @@ def launch_split_scores( )[:, :group_size] .permute(2, 3, 0, 1, 4) ) - columns = token_starts[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + torch.arange( + columns = prompt_lengths[:request_count].to(torch.int64).view(-1, 1, 1, 1, 1) + torch.arange( manager._selection_width_capacity, dtype=torch.int64, device=manager._score_scratch.device, 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 index 8799e0869150..31c36ff7e556 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -73,13 +73,17 @@ def _build_case( decode_width=capacity - prompt_len, ) _write_block_offsets(tri, _encode_block_offsets(page_ids)) - round_starts = (torch.arange(max_requests, dtype=torch.int32, device=device) + 9).contiguous() - token_starts = torch.full((max_requests,), prompt_len, dtype=torch.int32, device=device) + 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)] - valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) - phase = (round_starts.float()[:, None, None] + offsets_t[None, :, None]) * omega[None, None, :] + 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( @@ -94,8 +98,8 @@ def _build_case( return ( tri, pools, - token_starts, - valid_seq_lens, + prompt_lengths, + source_lengths, seq_lens, mean_cos, mean_sin, @@ -134,8 +138,8 @@ def test_cute_kernel_matches_torch_oracle(case): ( tri, pools, - token_starts, - valid_seq_lens, + prompt_lengths, + source_lengths, seq_lens, mean_cos, mean_sin, @@ -160,13 +164,13 @@ def test_cute_kernel_matches_torch_oracle(case): # Every count up to capacity is served, nothing beyond. assert max_requests + 1 not in tri._compiled_score_by_request_count for request_count in dict.fromkeys((1, max_requests - 1, max_requests)): - valid_widths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) + decode_lengths = torch.full((max_requests,), -1, dtype=torch.int32, device=device) scores = _launch_split_scores( tri, request_count, - valid_seq_lens, - valid_widths, - token_starts, + source_lengths, + decode_lengths, + prompt_lengths, mean_cos, mean_sin, ) @@ -178,7 +182,7 @@ def test_cute_kernel_matches_torch_oracle(case): ) # The score leg owns the per-request decode widths the selection # reduce kernels consume. - assert valid_widths[:request_count].tolist() == [ + assert decode_lengths[:request_count].tolist() == [ seq_lens[request] - prompt_len for request in range(request_count) ] for request in range(request_count): 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 index ecbbf235e9ab..fe2b64d24c0f 100644 --- 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 @@ -20,7 +20,14 @@ def _run_fused_union( - tri, request_count, valid_seq_lens, valid_widths, token_starts, mean_cos, mean_sin, union_out + tri, + request_count, + source_lengths, + decode_lengths, + prompt_lengths, + mean_cos, + mean_sin, + union_out, ): """The fused score+stats+normalized-union pipeline (THE union path), fired directly off the compiled entries exactly like the round. Test mean phases @@ -28,12 +35,12 @@ def _run_fused_union( build-bound ``tri._selection_scores_rows``.""" import cuda.bindings.driver as cuda_driver - _stage_score_metadata(tri, request_count, valid_seq_lens, valid_widths, token_starts) + _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]) assert ( request_count in tri._compiled_score_by_request_count - and request_count in tri._compiled_normalize_union + and request_count in tri._compiled_normalize_union_by_request_count ) stream = cuda_driver.CUstream(torch.cuda.current_stream(tri._score_scratch.device).cuda_stream) tri._compiled_score_by_request_count[request_count]( @@ -44,7 +51,7 @@ def _run_fused_union( request_count, stream, ) - tri._compiled_normalize_union[request_count]( + tri._compiled_normalize_union_by_request_count[request_count]( tri._cute_partial_stats, *tri._cute_selection_prefix, tri._cute_selection_scores_rows, @@ -55,7 +62,9 @@ def _run_fused_union( union_out[:request_count, :columns].copy_(tri._selection_scores_rows[:request_count, :columns]) -def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tensor) -> torch.Tensor: +def _reference_union_scores( + scores_rows: torch.Tensor, decode_lengths: torch.Tensor +) -> torch.Tensor: """Union oracle: per-row mean/biased-std z-norm over the valid prefix (std clamped at 1e-6), union-max across rows, ``-inf`` past the width.""" request_count, _, width = scores_rows.shape @@ -63,13 +72,13 @@ def _reference_union_scores(scores_rows: torch.Tensor, valid_widths: torch.Tenso (request_count, width), float("-inf"), dtype=torch.float32, device=scores_rows.device ) for request in range(request_count): - valid_width = int(valid_widths[request]) - if valid_width <= 0: + decode_length = int(decode_lengths[request]) + if decode_length <= 0: continue - valid = scores_rows[request, :, :valid_width].to(torch.float32) + 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) / valid_width).sqrt() - combined[request, :valid_width] = ((valid - mean) / std.clamp_min(1e-6)).amax(dim=0) + 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 @@ -125,8 +134,8 @@ def test_union_fusion_matches_split_pipeline( 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) - round_starts = torch.tensor([float(seq_len), float(seq_len + 1)], device=device) - phase = (round_starts[:, None, None] + offsets[None, :, None]) * omega[None, None] + 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() @@ -154,7 +163,7 @@ def test_union_fusion_matches_split_pipeline( _write_block_offsets(ref_tri, encoded) if valid_lens is None: valid_lens = [seq_len, seq_len] - valid_seq_lens = torch.tensor(valid_lens, dtype=torch.int32, device=device) + 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 @@ -163,13 +172,13 @@ def test_union_fusion_matches_split_pipeline( # 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) - token_starts = torch.tensor(score_starts, 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, - valid_seq_lens, + source_lengths, split_widths, - token_starts, + prompt_lengths, mean_cos, mean_sin, ) @@ -184,9 +193,9 @@ def test_union_fusion_matches_split_pipeline( _run_fused_union( tri, request_count, - valid_seq_lens, + source_lengths, fused_widths, - token_starts, + prompt_lengths, mean_cos, mean_sin, fused_out, @@ -215,7 +224,7 @@ def test_union_fusion_frequency_count_guard_raises() -> None: with pytest.raises(ValueError, match="frequencies"): _TriAttentionScoreKernel( num_layers=1, - seq_len=256, + score_token_capacity=256, num_q_heads=8, num_freqs=16, pool_shape=(2, 2, 1, 128, 32), @@ -288,8 +297,10 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: 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) - round_starts = torch.arange(max_requests, dtype=torch.float32, device=device) + seq_len - phase = (round_starts[:, None, None] + offsets[None, :, None]) * omega[None, None] + 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() @@ -316,10 +327,10 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: 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 - valid_seq_lens = torch.zeros(max_requests, dtype=torch.int32, device=device) - token_starts = torch.zeros(max_requests, dtype=torch.int32, device=device) - valid_seq_lens[:request_count] = torch.tensor(valid_lens, dtype=torch.int32, device=device) - token_starts[:request_count] = torch.tensor(score_starts, dtype=torch.int32, device=device) + 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. @@ -327,9 +338,9 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: per_head = _launch_split_scores( ref_tri, request_count, - valid_seq_lens, + source_lengths, split_widths, - token_starts, + prompt_lengths, mean_cos, mean_sin, ) @@ -345,9 +356,9 @@ def test_union_fusion_giant_scratch_unaligned_start(max_requests: int) -> None: _run_fused_union( tri, max_requests, - valid_seq_lens, + source_lengths, fused_widths, - token_starts, + prompt_lengths, mean_cos, mean_sin, fused_out, 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 index 5f791d6238aa..d5f4866e449d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -270,7 +270,6 @@ def test_draft_admission_gates_raise(gate, match): def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager = _make_triattention(budget=4, beta=4) - manager.calibration = {} manager._layer_partition = ([0, 1], [], None) target = manager.kv_cache_manager target._stream = mock.Mock() @@ -321,7 +320,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): # length: physical confirmed plus everything evicted so far # (the eviction input's logical_source_length). prepared = internals.execute.call_args.args[0] - assert prepared[0]["logical_source_length"] == uncompressed + assert prepared[0].logical_source_length == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. assert request.py_num_compressed_tokens == uncompressed - confirmed @@ -462,7 +461,7 @@ def test_staged_block_width_clamps_to_manager_source_width(): host, device_table = _allocate_block_offset_staging( anchor_pool, num_pools=1, - max_requests=2, + request_capacity=2, token_capacity=129, max_source_blocks=4, ) @@ -470,7 +469,7 @@ def test_staged_block_width_clamps_to_manager_source_width(): host, device_table = _allocate_block_offset_staging( anchor_pool, num_pools=1, - max_requests=2, + request_capacity=2, token_capacity=129, max_source_blocks=64, ) 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 index ccdcff65a479..144691087af5 100644 --- 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 @@ -61,7 +61,7 @@ def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, def _pack_oracle( settled, - valid_seq_lens, + source_lengths, dense_offsets, dense_out, swa_offsets, @@ -78,12 +78,12 @@ def _pack_oracle( """Pack in place: dense rows forward settled content verbatim (stale included) then append the tail ``seq_len + move - keep_count``; SWA rows write latest-window ordinals once per KV head.""" - request_count = int(valid_seq_lens.shape[0]) + request_count = int(source_lengths.shape[0]) packed_rows = int(dense_out.shape[0]) dense_total = int(dense_out.shape[1]) swa_total = int(swa_out.shape[1]) for request in range(request_count): - seq_len = int(valid_seq_lens[request]) + seq_len = int(source_lengths[request]) dense_begin = int(dense_offsets[request]) dense_count = int(dense_offsets[request + 1]) - dense_begin for domain in range(packed_rows): @@ -208,7 +208,7 @@ def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, w move_capacity = max(move_capacity, swa_window + tail_capacity) dense_total = request_count * (keep_count + tail_capacity) swa_total = request_count * (swa_window + tail_capacity) - valid_seq_lens = torch.tensor([10, 8, 0], dtype=torch.int32, device=device) + source_lengths = torch.tensor([10, 8, 0], dtype=torch.int32, device=device) dense_offsets = _staged_offsets(dense_counts, device) swa_offsets = _staged_offsets(swa_counts, device) @@ -237,7 +237,7 @@ def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, w swa_reference = swa_stale.clone() _pack_oracle( settled, - valid_seq_lens, + source_lengths, dense_offsets, dense_reference, swa_offsets if has_swa else dense_offsets, @@ -255,7 +255,7 @@ def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, w swa_actual = swa_stale.clone() _pack_move_sources_kernel[(request_count, selection_rows)]( settled, - valid_seq_lens, + source_lengths, dense_offsets, dense_actual, swa_offsets if has_swa else None, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index f31b4e949c1f..1fc326f23ba4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -145,7 +145,6 @@ def test_request_init_and_finish_lifecycle(self): manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 triattention = TriAttention(_make_tri_config(budget=8), manager) - triattention.calibration = {} triattention.on_request_init(_make_request(11)) triattention.on_request_init(_make_request(12)) @@ -283,7 +282,6 @@ def test_identity_selection_is_filtered_before_launch(self): capacity=6, history_length=0, is_active=True, resize=mock.Mock(return_value=True) ) manager.kv_cache_manager.kv_cache_map = {7: cache} - manager.calibration = {} state = _set_request_state(manager, 7, generation_steps=127) with _mocked_eviction_internals(manager) as internals: @@ -325,7 +323,6 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft fake_v2.num_extra_kv_tokens = num_extra_kv_tokens fake_v2._kv_reserve_draft_tokens = kv_reserve_draft_tokens mgr = TriAttention(_make_tri_config(budget=8), fake_v2) - mgr.calibration = {} cache = SimpleNamespace( capacity=seq_len, history_length=1024, @@ -362,7 +359,7 @@ def test_suspended_cache_defers_that_request_pre_launch(self): # Only the active request launched; the suspended one deferred whole. eviction_inputs = internals.execute.call_args.args[0] - assert [item["request"].py_request_id for item in eviction_inputs] == [7] + assert [item.request.py_request_id for item in eviction_inputs] == [7] assert first_state["generation_steps"] == 128 assert second_state["generation_steps"] == 127 @@ -397,19 +394,15 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): mgr._evict_due_requests(batch) # Tail excluded from the source length; keep target = prompt + budget. - internals.execute.assert_called_once_with( - [ - { - "request": request, - "target_cache": cache, - "draft_cache": draft_cache, - "source_length": confirmed, - "logical_source_length": confirmed, - "prompt_length": 1024, - "target_tail_length": tail, - } - ] - ) + 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.logical_source_length == confirmed + assert launched.prompt_length == 1024 + 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) @@ -419,7 +412,6 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): # ledger (capacity minus the protected tail), never the logical length. physical_confirmed = 6100 manager = _make_triattention(beta=128) - manager.calibration = {} _set_request_state(manager, 7, generation_steps=127, evicted_tokens=100) cache = SimpleNamespace( capacity=physical_confirmed, @@ -429,7 +421,6 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): ) manager.kv_cache_manager.kv_cache_map = {7: cache} manager.kv_cache_manager.pp_layers = [0, 1] - manager.kv_cache_manager.num_extra_kv_tokens = 0 request = _make_request( 7, py_prompt_len=1024, @@ -441,9 +432,9 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): manager._evict_due_requests(SimpleNamespace(generation_requests=[request])) eviction_inputs = internals.execute.call_args.args[0] - assert eviction_inputs[0]["source_length"] == physical_confirmed + assert eviction_inputs[0].source_length == physical_confirmed # The logical position restores everything already evicted. - assert eviction_inputs[0]["logical_source_length"] == physical_confirmed + 100 + assert eviction_inputs[0].logical_source_length == physical_confirmed + 100 cache.resize.assert_called_once_with(1024 + manager.budget, None) def test_one_model_draft_co_compression_contract_is_accepted(self): @@ -505,7 +496,7 @@ def test_union_forces_normalized_scores(self): triattention = _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) assert triattention.normalize_scores is True - def test_execute_rejects_int32_overflowing_round_starts(self): + def test_execute_rejects_int32_overflowing_logical_source_lengths(self): # Round starts past the int32 metadata range fail loudly (in the host # metadata build) before any GPU work is enqueued. device = torch.device("cuda", torch.cuda.current_device()) @@ -689,7 +680,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self): omega = torch.rand(num_freqs, device=device) * 0.05 offsets = torch.tensor([1.0, 2.0, 4.0], device=device) round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 - round_starts = round_device[:request_count].tolist() + logical_source_lengths = round_device[:request_count].tolist() seq_lens = [seq_len - request % 2 for request in range(request_count)] layer_order = list(range(num_layers)) tri = _make_cute_buffers( @@ -710,7 +701,7 @@ def test_fused_score_spans_distinct_storages_and_block_tables(self): layer_pool_ids=list(layer_order), normalize_scores=False, ) - valid_seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) + source_lengths = torch.tensor(seq_lens, dtype=torch.int32, device=device) # Rounds stage through the production executor: the gather double # writes each layer's K page ids and the bulk copy encodes K/V rows. @@ -733,7 +724,7 @@ def prepared_cohort(): return [ _make_eviction_input( request_id=request, - source_length=int(valid_seq_lens[request]), + source_length=int(source_lengths[request]), logical_source_length=int(round_device[request]), prompt_length=prompt_len, ) @@ -786,7 +777,7 @@ def score_rectangle(): pools, {layer: page_ids_3d[layer, :request_count] for layer in layer_order}, seq_lens, - round_starts, + logical_source_lengths, q_real, q_imag, mlr, @@ -797,23 +788,23 @@ def score_rectangle(): ) for request in range(request_count): for layer_slot, layer in enumerate(layer_order): - valid_width = seq_lens[request] - prompt_len - segment = fixed[request, layer_slot, :, :valid_width] + decode_length = seq_lens[request] - prompt_len + segment = fixed[request, layer_slot, :, :decode_length] expected = oracle[request * num_layers + layer][ - :, prompt_len : prompt_len + valid_width + :, prompt_len : prompt_len + decode_length ] torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) round_device.add_(17) page_ids_3d = page_ids_3d.roll(1, dims=2) - valid_seq_lens.copy_( + source_lengths.copy_( torch.tensor( [seq_len - (request + 1) % 2 for request in range(request_count)], dtype=torch.int32, device=device, ) ) - expected_second_widths = valid_seq_lens - prompt_len + expected_second_widths = source_lengths - prompt_len tri._score_scratch.fill_(score_sentinel) tri._decode_lengths_device.fill_(-1) with mock.patch.object(module, "compact"): 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 index faf7749f2ecb..05f55b6909e1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -113,8 +113,10 @@ def _select_per_head(tri, scores, *, normalize_scores): request_count=request_count, num_layers=num_layers, num_q_heads=num_q_heads, + num_kv_heads=tri._num_kv_heads, padded_head_columns=8, score_token_capacity=width, + selection_width=tri._selection_width_capacity, per_layer=tri.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) @@ -129,7 +131,7 @@ def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor def _per_head_keep_oracle( scores: torch.Tensor, - valid_widths: torch.Tensor, + decode_lengths: torch.Tensor, keep_count: int, eviction_mode: str, normalize_scores: bool, @@ -139,23 +141,26 @@ def _per_head_keep_oracle( num_kv_heads = 2 rows = [] for request in range(request_count): - valid_width = int(valid_widths[request]) - valid = scores[request, ..., :valid_width].clone() + 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) / (valid_width**0.5) + 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, valid_width + 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, valid_width) + selection = grouped.reshape(num_layers * num_kv_heads, decode_length) rows.append( torch.stack( - [torch.sort(_stable_topk(row, valid_width, keep_count)).values for row in selection] + [ + torch.sort(_stable_topk(row, decode_length, keep_count)).values + for row in selection + ] ) ) return torch.stack(rows) @@ -179,10 +184,10 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( generator=generator, dtype=torch.int32, ).to(torch.float32) - valid_widths = torch.tensor([83, 91], dtype=torch.int32) + decode_lengths = torch.tensor([83, 91], dtype=torch.int32) expected = _per_head_keep_oracle( - scores_cpu, valid_widths, keep_count, eviction_mode, normalize_scores + scores_cpu, decode_lengths, keep_count, eviction_mode, normalize_scores ) device = torch.device("cuda", torch.cuda.current_device()) @@ -198,7 +203,7 @@ def test_per_head_selection_matches_torch_oracle_on_selector_stream( num_query_heads=query_heads, num_kv_heads=kv_heads, ) - tri._decode_lengths_device.copy_(valid_widths.to(device)) + 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) @@ -228,7 +233,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid dtype=torch.int32, device=device, ).to(torch.float32) - valid_widths = (width, width - 32) + decode_lengths = (width, width - 32) tri = _make_selection_buffers( eviction_mode="union", width=width, @@ -236,7 +241,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid device=device, max_requests=request_count, ) - tri._decode_lengths_device.copy_(torch.tensor(valid_widths, dtype=torch.int32, device=device)) + 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) ) @@ -245,9 +250,9 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid actual = tri._kept_ordinal_rows.cpu() combined = scores.amax(dim=1).cpu() - for request, valid_width in enumerate(valid_widths): + for request, decode_length in enumerate(decode_lengths): expected_decode = torch.sort( - _stable_topk(combined[request], valid_width, keep_count).to(torch.int32) + prompt_len + _stable_topk(combined[request], decode_length, keep_count).to(torch.int32) + prompt_len ).values assert torch.equal(actual[request], expected_decode) @@ -267,7 +272,7 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no dtype=torch.float32, device=device, ) - valid_widths = torch.tensor([83, 91], dtype=torch.int32, 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 ) @@ -284,7 +289,7 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no score_scratch, prompt_lengths = _rect_to_score_scratch(scores, kv_heads) prepare_per_head_scores( score_scratch, - valid_widths, + decode_lengths, prompt_lengths, row_mean, row_inv_std, @@ -293,8 +298,10 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no request_count=request_count, num_layers=layers, num_q_heads=query_heads, + num_kv_heads=kv_heads, padded_head_columns=8, score_token_capacity=width, + selection_width=width, per_layer=per_layer, normalize_scores=normalize_scores, ) @@ -303,26 +310,26 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no selection_scores = selection_scores_rows.view(request_count, selection_rows, width) assert torch.equal( selection_row_lengths.view(request_count, selection_rows).cpu(), - valid_widths.cpu().view(request_count, 1).expand(-1, selection_rows), + decode_lengths.cpu().view(request_count, 1).expand(-1, selection_rows), ) query_group_size = query_heads // kv_heads - for request, valid_width in enumerate(valid_widths.tolist()): - valid = scores[request, :, :, :valid_width] + 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 / valid_width**0.5).clamp_min(1e-6) + std = (std / decode_length**0.5).clamp_min(1e-6) valid = (valid - mean) / std - grouped = valid.view(layers, kv_heads, query_group_size, valid_width).amax(dim=2) + 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, valid_width) + expected = expected.reshape(selection_rows, decode_length) assert torch.allclose( - selection_scores[request, :, :valid_width], + selection_scores[request, :, :decode_length], expected, rtol=2e-5, atol=2e-5, ) - assert torch.isneginf(selection_scores[request, :, valid_width:]).all() + assert torch.isneginf(selection_scores[request, :, decode_length:]).all() @pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) @@ -760,7 +767,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): dtype=torch.int64, device=device, ) - valid_seq_lens = torch.tensor([64, 56], dtype=torch.int32, device=device) + source_lengths = torch.tensor([64, 56], dtype=torch.int32, device=device) protected_tails = [2, 1] compaction = _build_compaction( layer_pools=pools, @@ -769,7 +776,7 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): # Dense layer 0 stages in plane 0, the SWA layer in its own plane 1. layer_pool_ids=[0, 1], kept_token_ordinals=keep.to(torch.int32), - valid_sequence_lengths=valid_seq_lens, + valid_sequence_lengths=source_lengths, kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), swa_window=2, @@ -779,8 +786,8 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): _run_compaction(compaction) torch.cuda.synchronize(device) - for request, (valid_seq_len, tail_length) in enumerate( - zip(valid_seq_lens.tolist(), protected_tails) + for request, (source_length, tail_length) in enumerate( + zip(source_lengths.tolist(), protected_tails) ): dense_pages = dense_tables[request].to(torch.long) swa_pages = swa_tables[request].to(torch.long) @@ -789,8 +796,8 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): swa_before = initial_pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) swa_after = pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape_as(swa_before) tail = torch.arange( - valid_seq_len, - valid_seq_len + tail_length, + source_length, + source_length + tail_length, dtype=torch.int64, device=device, ) @@ -799,8 +806,8 @@ def test_eager_compaction_rebases_masked_swa_window_and_tail(): 2, 2 + dense_source.numel(), dtype=torch.int64, device=device ) swa_source = torch.arange( - valid_seq_len - 2, - valid_seq_len + tail_length, + source_length - 2, + source_length + tail_length, dtype=torch.int64, device=device, ) From 0d68666fa370a542d5c860c4ff51f1639137a81c Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 26 Jul 2026 08:41:06 -0700 Subject: [PATCH 159/178] [None][chore] Drop the per-round pool-drift poll and legacy rope config fallbacks Knife 47 decision batch: V2 pools are allocated once at manager init, so the cached KV layout no longer re-polls page counts every eviction round; rope tables read rope_parameters directly (transformers >= 5.5 folds rope_theta and rope_type into it), removing the rope_scaling / legacy type-key / 10000.0 silent-default fallbacks. Missing config keys now fail loudly. kvcc 71 + executor 13 green, digest triplet token-identical vs the remerge bank. Signed-off-by: tianruih --- .../triattention/triattention.py | 35 +++------------ .../test_triattention_pipeline.py | 43 +------------------ 2 files changed, 8 insertions(+), 70 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index b077ddae93c4..e1a275d588b8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -30,7 +30,7 @@ import triton from tensorrt_llm._torch.distributed import allgather -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheCompressionManager from tensorrt_llm._torch.utils import next_positive_power_of_2 @@ -331,20 +331,17 @@ def _rope_tables(self, freq_count: int): config = AutoConfig.from_pretrained( self.model_path, trust_remote_code=True ).get_text_config() - config_values = config.to_dict() - rope_params = ( - config_values.get("rope_parameters") or config_values.get("rope_scaling") or {} - ) - if rope_params and all(isinstance(value, dict) for value in rope_params.values()): + # 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( f"TriAttention does not support per-layer-type rope parameters ({self.model_path})" ) - rope_type = rope_params.get("rope_type") or rope_params.get("type") or "default" + rope_type = rope_params["rope_type"] if rope_type == "default": # "default" has no ROPE_INIT_FUNCTIONS entry; the analytic formula is its definition. head_dim = freq_count * 2 - theta = rope_params.get("rope_theta", config_values.get("rope_theta")) - base = float(theta) if theta is not None else 10000.0 + 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() scale_sq = 1.0 @@ -1336,25 +1333,10 @@ def _local_score_calibration( ) def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: - # Only pool page counts are polled; manager identity and layer count are manager-lifetime contracts. + # V2 pools are allocated once at manager init; the layout is a manager-lifetime contract. manager = self.draft_kv_cache_manager if draft else self.kv_cache_manager cached = self._kv_layout_caches[draft] if cached is not None: - current_page_counts = tuple( - int( - manager.impl.get_page_index_upper_bound( - manager.layer_offsets[cached["global_layers"][layer]], - Role.KEY, - ) - ) - // int(manager.kv_factor) - for layer in cached["pool_representatives"] - ) - if current_page_counts != cached["pool_page_counts"]: - raise RuntimeError( - f"TriAttention {'draft ' if draft else ''}V2 pool layout changed " - "after the layout was built; KV pool rebalance is not supported" - ) return cached if draft: @@ -1411,7 +1393,4 @@ def _build_runtime_kv_layout( swa_window=swa_window, layer_pool_ids=layer_pool_ids, pool_representatives=pool_representatives, - pool_page_counts=tuple( - int(layer_pools[layer].shape[0]) for layer in pool_representatives - ), ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 1fc326f23ba4..9e3f97a5a1d4 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -43,7 +43,6 @@ # 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 -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role # The SM100 CuTe kernel is the only score path, so every test that actually # launches scores (or builds the real staging buffers, whose constructor @@ -98,46 +97,6 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): class TestTriAttentionClass: - def test_cached_layout_checks_page_counts_without_rebuilding_pool_views(self): - page_count_query = mock.Mock(side_effect=[8, 16, 8, 18]) - manager = SimpleNamespace( - get_buffers=mock.Mock(side_effect=AssertionError("pool view was rebuilt")), - impl=SimpleNamespace(get_page_index_upper_bound=page_count_query), - kv_factor=2, - layer_offsets={10: 100, 11: 101, 12: 102}, - ) - triattention = _make_triattention() - triattention.kv_cache_manager = manager - cached = dict( - global_layers=[10, 11, 12], - layer_pools=[torch.empty(4), torch.empty(8), torch.empty(4)], - dense_layers=[0, 1, 2], - swa_layers=[], - swa_window=None, - layer_pool_ids=(0, 1, 0), - # These are local layer slots. Layer 2 shares layer 0's pool. - pool_representatives=(0, 1), - pool_page_counts=(4, 8), - ) - triattention._kv_layout_caches[False] = cached - - assert triattention._runtime_kv_layout() is cached - manager.get_buffers.assert_not_called() - assert page_count_query.call_args_list == [ - mock.call(100, Role.KEY), - mock.call(101, Role.KEY), - ] - - with pytest.raises(RuntimeError, match="pool layout changed"): - triattention._runtime_kv_layout() - manager.get_buffers.assert_not_called() - assert page_count_query.call_args_list == [ - mock.call(100, Role.KEY), - mock.call(101, Role.KEY), - mock.call(100, Role.KEY), - mock.call(101, Role.KEY), - ] - def test_request_init_and_finish_lifecycle(self): # Init: speculative capacity accepted, manager marked, state tracked. # Finish: state cleared; buffers and the step's batch stay resident. @@ -191,7 +150,7 @@ def test_resolve_converts_official_layout(self, tmp_path): torch.save({"metadata": {"sampled_heads": sampled}, "stats": stats}, path) mgr = _make_triattention() mgr.calibration_path = str(path) - config = _make_hf_config(rope_theta=10000.0) + config = _make_hf_config(rope_parameters={"rope_type": "default", "rope_theta": 10000.0}) with mock.patch("transformers.AutoConfig.from_pretrained", return_value=config): converted = mgr._resolve_calibration() From 77017bcac7a641875e4fe6260c189eae23dc21b4 Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 26 Jul 2026 21:34:42 -0700 Subject: [PATCH 160/178] [None][doc] Note that the native compact op derives batch from tensor geometry Signed-off-by: tianruih --- tensorrt_llm/_torch/kv_cache_compression/compaction.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index 5d7ef1765262..f622687b7780 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -259,7 +259,9 @@ def compact( request_count: int, ) -> None: """Pack each cache's move sources and fire its native compacts, in order - (pure mover: the caller owns the decision rows and the round's completion ordering).""" + (pure mover: the caller owns the decision rows and the round's completion ordering). + ``request_count`` only sizes the Triton pack grid; the native op takes the batch + from the page-table view's row count and validates the companion tensors against it.""" # One launch per (cache, pool group); each launch covers every layer in the group. for cache_params in params: _pack_move_sources_kernel[(request_count, cache_params.decision_rows)]( From cb0101b974490913e829ff27c18a0fb93ce408ba Mon Sep 17 00:00:00 2001 From: tianruih Date: Sun, 26 Jul 2026 21:41:53 -0700 Subject: [PATCH 161/178] [None][fix] Validate the decision-row layout at the compaction API boundary Reject kept-ordinal row counts that do not divide the request count evenly and decision layouts other than broadcast, per-head, or per-layer-per-head, so a mismatched cache geometry fails loudly instead of packing uninitialized move indices. Signed-off-by: tianruih --- tensorrt_llm/_torch/kv_cache_compression/compaction.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/_torch/kv_cache_compression/compaction.py b/tensorrt_llm/_torch/kv_cache_compression/compaction.py index f622687b7780..876d80d010ec 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/compaction.py +++ b/tensorrt_llm/_torch/kv_cache_compression/compaction.py @@ -144,6 +144,13 @@ def build_compaction_params( params.decision_rows = int(kept_ordinal_rows.shape[0]) // max_requests # Pool shape [pages, K/V, heads, tokens, dim]. num_kv_heads = int(first_pool.shape[2]) + if params.decision_rows * max_requests != int(kept_ordinal_rows.shape[0]): + raise ValueError("kept_ordinal_rows rows must be a multiple of max_requests") + if params.decision_rows not in (1, num_kv_heads, len(dense_layers) * num_kv_heads): + raise ValueError( + f"unsupported decision layout: {params.decision_rows} rows for " + f"{num_kv_heads} heads x {len(dense_layers)} dense layers" + ) per_layer_sources = ( len(dense_layers) > 1 and params.decision_rows == len(dense_layers) * num_kv_heads ) From bb89ca106b1bc71e510d4e0902b7f2040835459a Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Mon, 27 Jul 2026 01:54:08 -0700 Subject: [PATCH 162/178] [None][test] Preserve compaction tests in TriAttention stack Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- .../kv_cache_compression/test_compaction.py | 324 ++++++++++++++++++ 1 file changed, 324 insertions(+) create mode 100644 tests/unittest/_torch/kv_cache_compression/test_compaction.py diff --git a/tests/unittest/_torch/kv_cache_compression/test_compaction.py b/tests/unittest/_torch/kv_cache_compression/test_compaction.py new file mode 100644 index 000000000000..1e7913016851 --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_compaction.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Physical KV-cache compaction: packed moves, protected tails, SWA windows, +and draft co-compaction, checked byte-exactly against torch oracles.""" + +from types import SimpleNamespace + +import pytest +import torch +from conftest import build_compaction as _build_compaction +from conftest import encode_block_offsets as _encode_block_offsets +from conftest import make_ramp_pools as _make_ramp_pools +from conftest import run_compaction as _run_compaction +from conftest import set_protected_tails as _set_protected_tails + +requires_sm100 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0), + reason="KV-cache compaction kernels require SM100", +) + + +def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: + """Gather one request's pages into [K/V, head, token, dim] order.""" + num_kv_heads = int(pool.shape[2]) + head_dim = int(pool.shape[4]) + return pool.index_select(0, pages).permute(1, 2, 0, 3, 4).reshape(2, num_kv_heads, -1, head_dim) + + +@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) +def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode): + # Supported bf16 geometry; kept ordinals span all three pages so moves + # cross page boundaries. + device = torch.device("cuda", torch.cuda.current_device()) + request_count = 2 + num_layers = 2 + num_kv_heads = 2 + # Mixed prompt lengths prove per-request destination rebasing. + prompt_lens = [2, 5] + decode_keep_count = 4 + seq_len = 80 + tokens_per_block = 32 + pages_per_request = 3 + head_dim = 64 + protected_tails = [2, 1] + page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + initial_pools = _make_ramp_pools(num_layers, device=device) + pools = [pool.clone() for pool in initial_pools] + + # Decode-only kept ordinals holding absolute positions. + union_decode = torch.tensor( + [[16, 32, 56, 72], [24, 40, 48, 64]], dtype=torch.int64, device=device + ) + if eviction_mode == "union": + keep = union_decode + selection_rows = 1 + else: + selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads + keep = torch.empty( + request_count, + selection_rows, + decode_keep_count, + dtype=torch.int64, + device=device, + ) + for request in range(request_count): + for row in range(selection_rows): + keep[request, row] = torch.tensor( + sorted( + { + prompt_lens[request] + ((request + row + offset * 2) % 8) * 8 + for offset in range(decode_keep_count) + } + ), + dtype=torch.int64, + device=device, + ) + + compaction = _build_compaction( + eviction_mode=eviction_mode, + layer_pools=pools, + kept_token_ordinals=keep.to(torch.int32), + valid_sequence_lengths=torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device), + kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), + prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), + protected_tail_capacity=max(protected_tails), + ) + _set_protected_tails(compaction, protected_tails) + # Production settles the kept ordinals into the contract's decision + # rows; with pre-settled ordinals the pack launch inside compact() is + # its exact analog. + _run_compaction(compaction) + torch.cuda.synchronize(device) + + for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): + for request in range(request_count): + prompt_len = prompt_lens[request] + pages = page_tables[request].to(torch.long) + before = ( + before_pool[pages] + .permute(1, 2, 0, 3, 4) + .reshape(2, num_kv_heads, pages_per_request * tokens_per_block, head_dim) + ) + after = after_pool[pages].permute(1, 2, 0, 3, 4).reshape_as(before) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + for head in range(num_kv_heads): + if eviction_mode == "union": + selected = keep[request] + elif eviction_mode == "per_head": + selected = keep[request, head] + else: + selected = keep[request, layer * num_kv_heads + head] + tail = torch.arange( + seq_len, + seq_len + protected_tails[request], + dtype=torch.int64, + device=device, + ) + source = torch.cat((selected, tail)) + destination = torch.arange( + prompt_len, + prompt_len + source.numel(), + dtype=torch.int64, + device=device, + ) + assert torch.equal( + after[:, head].index_select(1, destination), + before[:, head].index_select(1, source), + ) + + +@requires_sm100 +def test_eager_compaction_rebases_masked_swa_window_and_tail(): + # Supported bf16 geometry; dense and SWA moves stay page-crossing. + device = torch.device("cuda", torch.cuda.current_device()) + dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) + swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) + initial_pools = _make_ramp_pools(2, num_kv_heads=1, device=device) + pools = [pool.clone() for pool in initial_pools] + # Decode-only kept ordinals holding absolute positions past the prompt. + keep = torch.tensor( + [[16, 32, 40, 56], [16, 24, 40, 48]], + dtype=torch.int64, + device=device, + ) + valid_seq_lens = torch.tensor([64, 56], dtype=torch.int32, device=device) + protected_tails = [2, 1] + compaction = _build_compaction( + layer_pools=pools, + dense_layers=[0], + swa_layers=[1], + layer_group_representative={0: 0}, + # Dense layer 0 stages in plane 0, the SWA layer in its own plane 1. + layer_pool_ids=[0, 1], + kept_token_ordinals=keep.to(torch.int32), + valid_sequence_lengths=valid_seq_lens, + kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), + prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), + swa_window=2, + protected_tail_capacity=max(protected_tails), + ) + _set_protected_tails(compaction, protected_tails) + _run_compaction(compaction) + torch.cuda.synchronize(device) + + for request, (valid_seq_len, tail_length) in enumerate( + zip(valid_seq_lens.tolist(), protected_tails) + ): + dense_pages = dense_tables[request].to(torch.long) + swa_pages = swa_tables[request].to(torch.long) + dense_before = initial_pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) + dense_after = pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape_as(dense_before) + swa_before = initial_pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) + swa_after = pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape_as(swa_before) + tail = torch.arange( + valid_seq_len, + valid_seq_len + tail_length, + dtype=torch.int64, + device=device, + ) + dense_source = torch.cat((keep[request], tail)) + dense_destination = torch.arange( + 2, 2 + dense_source.numel(), dtype=torch.int64, device=device + ) + swa_source = torch.arange( + valid_seq_len - 2, + valid_seq_len + tail_length, + dtype=torch.int64, + device=device, + ) + swa_destination = torch.arange(4, 4 + swa_source.numel(), dtype=torch.int64, device=device) + assert torch.equal(dense_after[:, :, :2], dense_before[:, :, :2]) + assert torch.equal(swa_after[:, :, :2], swa_before[:, :, :2]) + assert torch.equal( + dense_after.index_select(2, dense_destination), + dense_before.index_select(2, dense_source), + ) + assert torch.equal( + swa_after.index_select(2, swa_destination), + swa_before.index_select(2, swa_source), + ) + + +def _launched_draft_compaction(draft_protected_tails): + """Target and draft pools with distinct head counts (supported bf16 + geometry, mod-251 ramp payload), compacted in one round.""" + device = torch.device("cuda", torch.cuda.current_device()) + request_count = 2 + prompt_len = 2 + target_protected_tails = [2, 1] + valid_seq_lens = [10, 9] + + target_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) + draft_tables = torch.tensor([[1, 0, 2], [5, 4, 3]], dtype=torch.int32, device=device) + target_pools = _make_ramp_pools(2, num_kv_heads=2, device=device) + draft_pool = _make_ramp_pools(1, num_kv_heads=4, base=149, device=device)[0] + assert target_pools[0].shape[2] != draft_pool.shape[2] + initial_target = [pool.clone() for pool in target_pools] + initial_draft = draft_pool.clone() + + keep = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) + + compaction = _build_compaction( + layer_pools=target_pools, + layer_pool_ids=[0, 0], + kept_token_ordinals=keep.to(torch.int32), + valid_sequence_lengths=torch.tensor(valid_seq_lens, dtype=torch.int32, device=device), + kv_block_offsets=_encode_block_offsets(target_tables), + prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), + protected_tail_capacity=max(target_protected_tails), + draft_layer_pools=[draft_pool], + draft_layers=[0], + draft_layer_group_representative={0: 0}, + draft_layer_pool_ids=[0], + draft_protected_tail_capacity=max(draft_protected_tails), + draft_kv_block_offsets=_encode_block_offsets(draft_tables), + ) + _set_protected_tails(compaction, target_protected_tails, draft_protected_tails) + _run_compaction(compaction) + torch.cuda.synchronize(device) + + return SimpleNamespace( + device=device, + request_count=request_count, + prompt_len=prompt_len, + keep=keep, + valid_seq_lens=valid_seq_lens, + target_protected_tails=target_protected_tails, + draft_protected_tails=draft_protected_tails, + target_tables=target_tables, + draft_tables=draft_tables, + target_pools=target_pools, + draft_pool=draft_pool, + initial_target=initial_target, + initial_draft=initial_draft, + compaction=compaction, + ) + + +def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(): + # Ragged draft tails [1, 2] against target tails [2, 1]: one request's + # draft tail below and one above its target, subsuming the uniform row. + built = _launched_draft_compaction(draft_protected_tails=[1, 2]) + device = built.device + prompt_len = built.prompt_len + + expected_offsets = [0] + for request in range(built.request_count): + valid = built.valid_seq_lens[request] + # Target dense layers compact the union keep set plus the target tail. + target_pages = built.target_tables[request].to(torch.long) + target_tail = torch.arange( + valid, + valid + built.target_protected_tails[request], + dtype=torch.int64, + device=device, + ) + target_source = torch.cat((built.keep[request], target_tail)) + target_destination = torch.arange( + prompt_len, + prompt_len + target_source.numel(), + dtype=torch.int64, + device=device, + ) + for before_pool, after_pool in zip(built.initial_target, built.target_pools): + before = _logical_view(before_pool, target_pages) + after = _logical_view(after_pool, target_pages) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + assert torch.equal( + after.index_select(2, target_destination), + before.index_select(2, target_source), + ) + + # Same kept ordinals through the draft's OWN table/heads/tail. + draft_pages = built.draft_tables[request].to(torch.long) + draft_tail = torch.arange( + valid, + valid + built.draft_protected_tails[request], + dtype=torch.int64, + device=device, + ) + draft_source = torch.cat((built.keep[request], draft_tail)) + draft_destination = torch.arange( + prompt_len, + prompt_len + draft_source.numel(), + dtype=torch.int64, + device=device, + ) + before = _logical_view(built.initial_draft, draft_pages) + after = _logical_view(built.draft_pool, draft_pages) + assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) + for head in range(int(built.draft_pool.shape[2])): + assert torch.equal( + after[:, head].index_select(1, draft_destination), + before[:, head].index_select(1, draft_source), + ) + + expected_offsets.append(expected_offsets[-1] + int(draft_source.numel())) + + # The test-owned draft move-offset row must match the broadcast-plus-tail + # oracle; the packed move sources themselves are covered byte-exactly by + # the pool assertions above (the ramp payload makes every wrong move land + # on different bytes) and by the pack-kernel oracle suite. + assert built.compaction["draft_move_offsets"].cpu().tolist() == expected_offsets From 395e9c840f35ba7f6daca941f479c4121d9d94ca Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Mon, 27 Jul 2026 02:13:54 -0700 Subject: [PATCH 163/178] [None][chore] Keep TriAttention cleanup inside its PR boundary Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h | 5 +++++ .../unfusedAttentionKernels_2_bf16_bf16.cu | 2 +- .../_torch/kv_cache_compression/interface.py | 13 ++++++++----- .../_torch/kv_cache_compression/conftest.py | 3 +++ .../thop/serial/test_sparse_kv_cache_compact.py | 2 +- 5 files changed, 18 insertions(+), 7 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index baa4ede44217..3a757c3263d9 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -413,6 +413,11 @@ void invokeUpdateCyclicKvCacheAfterFmha(QKVPreprocessingParams template void invokeUpdateSparseKvCacheAfterFmha(QKVPreprocessingParams params, cudaStream_t stream); +// Debug function to test basic parameter access +template +void invokeDebugSparseKvCacheParams( + QKVPreprocessingParams params, int* debug_output, cudaStream_t stream); + //! Compact a uniform group of KVCacheManagerV2 layer pools in one batched launch //! (per request and head, moves are ascending and never overtake their sources: //! the copy runs in place). diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu index 023702c4a72f..8c2d11804d38 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_bf16_bf16.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. * * Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index 8f22dd7f9b31..cd6f3c2483ef 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -2,25 +2,28 @@ # 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; feature gates - read the ``is_*`` trait predicates (algorithm dispatch itself matches the - config's ``algorithm`` string). + Configs map their ``algorithm`` string to a member here; callers read the + ``is_*`` predicates instead of comparing strings. """ TRIATTENTION = auto() NONE = auto() def is_eviction_method(self): - """Whether this method physically evicts cached tokens.""" + """Whether this method physically evicts cached tokens. Evicting + algorithms add their member and extend this predicate.""" return self == KvCacheCompressionMode.TRIATTENTION @staticmethod - def from_string(name: str) -> "KvCacheCompressionMode": + def from_string(name: Optional[str]) -> "KvCacheCompressionMode": + if name is None: + return KvCacheCompressionMode.NONE try: return KvCacheCompressionMode[name.upper()] except KeyError: diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index e5651ad28081..1cf42d70fc76 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Shared harness for the KV-cache compaction tests.""" + import json import os import tempfile @@ -128,6 +130,7 @@ def build_compaction(**overrides): eviction_mode="union", dense_layers=[0, 1], swa_layers=[], + layer_group_representative={0: 0, 1: 1}, layer_pool_ids=[0, 0], request_count=2, decode_keep_count=4, diff --git a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py index 1e065b751613..ca40c09667b9 100644 --- a/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py +++ b/tests/unittest/_torch/thop/serial/test_sparse_kv_cache_compact.py @@ -210,7 +210,7 @@ def test_sparse_kv_cache_compact_layers(case): def test_sparse_kv_cache_compact_layers_cuda_graph_replay(): - """Check operation-level capture safety, not a standalone TriAttention graph.""" + """Check operation-level capture safety inside an externally captured graph.""" pools_cpu, pools, page_tables = _make_pools(3, torch.bfloat16, 64) page_tables_cpu = [page_table.cpu() for page_table in page_tables] source_offsets = torch.tensor([0, 3, 6], dtype=torch.int32) From 36090b887d2d763d980170f808bf86c6b51330ef Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 27 Jul 2026 07:32:24 -0700 Subject: [PATCH 164/178] [None][chore] Simplify TriAttention eviction runtime Signed-off-by: tianruih --- .../_torch/kv_cache_compression/interface.py | 7 +- .../triattention/triattention.py | 1198 ++++++----------- .../triattention_cute_score_fused.py | 243 +++- .../triattention/triattention_kernels.py | 31 +- .../_torch/kv_cache_compression/conftest.py | 194 +-- .../test_rope_fusion_gate.py | 54 +- .../test_triattention_cute_score.py | 6 +- .../test_triattention_cute_union_fusion.py | 42 +- .../test_triattention_draft_cocompaction.py | 372 ++--- .../test_triattention_fused_settle_pack.py | 165 +-- .../test_triattention_pipeline.py | 390 ++---- .../test_triattention_selection_compaction.py | 246 +--- 12 files changed, 947 insertions(+), 2001 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/interface.py b/tensorrt_llm/_torch/kv_cache_compression/interface.py index cd6f3c2483ef..6e36001d7ce8 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -15,10 +15,9 @@ class KvCacheCompressionMode(IntEnum): TRIATTENTION = auto() NONE = auto() - def is_eviction_method(self): - """Whether this method physically evicts cached tokens. Evicting - algorithms add their member and extend this predicate.""" - return self == KvCacheCompressionMode.TRIATTENTION + def is_eviction_method(self) -> bool: + """Return whether this mode physically evicts cached tokens.""" + return self in (KvCacheCompressionMode.TRIATTENTION,) @staticmethod def from_string(name: Optional[str]) -> "KvCacheCompressionMode": diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index e1a275d588b8..2aacc7094e79 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -23,11 +23,13 @@ official tool (github.com/WeianMao/triattention) and is converted at load. """ +from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple -import cuda.bindings.driver as cuda_driver import torch import triton +from transformers import AutoConfig +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS from tensorrt_llm._torch.distributed import allgather from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 @@ -41,11 +43,12 @@ from tensorrt_llm.logger import logger 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_kernel, _gather_mean_phase_kernel, _settle_ties_kernel, - prepare_per_head_scores, + reduce_per_head_scores, ) if TYPE_CHECKING: @@ -67,13 +70,19 @@ _MAX_INTEGRATION_OFFSET = 65536 +@dataclass +class _RequestState: + confirmed_tokens: int = 0 + evicted_tokens: int = 0 + + class _EvictionInput(NamedTuple): """One due request's eviction operands for a single round.""" request: "LlmRequest" target_cache: object draft_cache: Optional[object] - state: Dict[str, object] + state: _RequestState source_length: int logical_source_length: int prompt_length: int @@ -88,15 +97,15 @@ def _allocate_block_offset_staging( token_capacity: int, max_source_blocks: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """One pinned host snapshot + persistent device table pair in the native - V2 ``[pool, request, K/V, block]`` layout (block width 4-aligned for the - ``PackedInt`` copy ABI); the device follows the anchor KV pool. The staged - width is clamped to the manager's live source-table width: static bucket - slack never holds valid tokens and the native gather copies the full width.""" + """Allocate a host snapshot and persistent device table. + + Both use the native V2 ``[pool, request, K/V, block]`` layout and grow only + when a newly admitted request raises the high-water mark. + """ tokens_per_block = int(anchor_pool.shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block - staged_blocks_per_seq = min((page_count + 3) // 4 * 4, int(max_source_blocks)) - shape = (num_pools, request_capacity, 2, staged_blocks_per_seq) + block_capacity = min((page_count + 3) // 4 * 4, max_source_blocks) + shape = (num_pools, request_capacity, 2, block_capacity) host = torch.empty(shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) device_table = torch.empty(shape, dtype=torch.int32, device=anchor_pool.device) return host, device_table @@ -105,30 +114,43 @@ def _allocate_block_offset_staging( _MEAN_PHASE_MAX_ROWS = 1 << 24 -def grow_mean_phase_table(phase: Dict[str, object], rows: int) -> None: - """Cover positions ``[0, rows)``, rebuilding the table if it must grow.""" - rows = int(rows) - if rows <= phase["rows"]: - return - if rows > _MEAN_PHASE_MAX_ROWS: - raise ValueError(f"a {rows}-row mean-phase table exceeds the exact-FP32 position range") - target = 1 - while target < rows: - target *= 2 - target = min(max(target, 2 * phase["rows"]), _MEAN_PHASE_MAX_ROWS) - omega = phase["omega"] - positions = torch.arange(target, device=omega.device, dtype=torch.float32) - cos_table = torch.zeros((target, omega.numel()), dtype=torch.float32, device=omega.device) - sin_table = torch.zeros_like(cos_table) - # Fixed summation order keeps the table bit-stable across rebuilds. - for offset in phase["offset_values"]: - angle = torch.outer(positions + offset, omega) - cos_table += torch.cos(angle) - sin_table += torch.sin(angle) - scale = 1.0 / len(phase["offset_values"]) - phase["cos"] = cos_table.mul_(scale) - phase["sin"] = sin_table.mul_(scale) - phase["rows"] = target +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._offsets = tuple(float(1 << i) for i in range(_MAX_INTEGRATION_OFFSET.bit_length())) + self.cos: Optional[torch.Tensor] = None + self.sin: Optional[torch.Tensor] = None + self.rows = 0 + self.num_freqs = int(self._omega.numel()) + self.frequency_block = triton.next_power_of_2(self.num_freqs) + + 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 self._offsets: + angle = torch.outer(positions + offset, self._omega) + cos_table += torch.cos(angle) + sin_table += torch.sin(angle) + scale = 1.0 / len(self._offsets) + self.cos = cos_table.mul_(scale) + self.sin = sin_table.mul_(scale) + self.rows = target class TriAttention(KVCacheCompressionManager): @@ -143,39 +165,30 @@ def __init__( config: "TriAttentionKvCacheCompressionConfig", kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, - ): + ) -> None: super().__init__(kv_cache_manager, draft_kv_cache_manager) self.budget = config.budget self.beta = config.beta self.eviction_mode = config.eviction_mode - self.normalize_scores = bool(config.normalize_scores) - if self.eviction_mode == "union" and not self.normalize_scores: - logger.warning( - "TriAttention union eviction always z-normalizes scores; " - "forcing normalize_scores=True" - ) - self.normalize_scores = True + 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() - # Mean-phase table dict; buffer builds bind its device tables in place. - self._phase: Optional[Dict[str, object]] = None - # Per-request eviction progress. - self._request_states: Dict[int, Dict[str, object]] = {} + self._request_states: Dict[int, _RequestState] = {} # In-flight overlap batch reference; membership resolves lazily. - self._inflight_scheduled_batch: Optional[object] = None - self._inflight_generation_request_ids: Optional[set] = None + self._inflight_scheduled_batch: Optional["ScheduledRequests"] = None + self._inflight_generation_request_ids: Optional[set[int]] = None # Manager-lifetime constants. self._num_extra_kv_tokens = int(kv_cache_manager.num_extra_kv_tokens) self._protected_tail_capacity = ( - int(kv_cache_manager.num_extra_kv_tokens) - + int(kv_cache_manager._kv_reserve_draft_tokens) - + 1 + self._num_extra_kv_tokens + int(kv_cache_manager._kv_reserve_draft_tokens) + 1 ) - self._draft_protected_tail_capacity: Optional[int] = None + 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) @@ -183,32 +196,34 @@ def __init__( + 1 ) self._generation_growth = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) - # Lazy resident eviction runtime: built at the first eviction, reused - # across rounds, and replaced as a whole when a capacity axis grows. - self._buffers_built = False - # Round-ordering events: device-lifetime, created at the first build - # and reused across capacity rebuilds. - self._staging_reuse_event: Optional[torch.cuda.Event] = None - self._block_offsets_ready_event: Optional[torch.cuda.Event] = None - self._compaction_done_event: Optional[torch.cuda.Event] = None + # Fixed by the manager/config; only absolute source length can grow + # when later requests arrive with longer prompts. + self._request_capacity = int(kv_cache_manager.max_batch_size) + self._selection_width_capacity = ( + self.budget + 2 * self.beta + int(kv_cache_manager.max_total_draft_tokens) + ) + 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._layer_partition = self._attention_layer_partition() - # Target/draft runtime KV layouts, cached by the one resolver. - self._kv_layout_caches: Dict[bool, Optional[Dict[str, object]]] = { - False: None, - True: None, - } - - def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int]]: + ( + 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) - from transformers import AutoConfig - config = AutoConfig.from_pretrained( model_path, trust_remote_code=True, local_files_only=True ) @@ -248,6 +263,8 @@ def _attention_layer_partition(self) -> Tuple[List[int], List[int], Optional[int ] swa_set = set(swa_layers) dense_layers = [layer for layer in range(num_layers) if layer not in swa_set] + 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") @@ -325,9 +342,6 @@ def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: def _rope_tables(self, freq_count: int): """Derive the RoPE frequency tables from the model config.""" - from transformers import AutoConfig - from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS - config = AutoConfig.from_pretrained( self.model_path, trust_remote_code=True ).get_text_config() @@ -356,10 +370,8 @@ def _rope_tables(self, freq_count: int): def on_request_init(self, request: "LlmRequest", **kwargs) -> None: """Register the request for eviction tracking.""" self._validate_request_capacity(request) - self._request_states[request.py_request_id] = { - "generation_steps": 0, - "evicted_tokens": 0, - } + self._reserve_eviction_capacity(request) + self._request_states[request.py_request_id] = _RequestState() def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: """Snapshot the in-flight batch; mutation remains in final update.""" @@ -367,25 +379,26 @@ def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwarg self._inflight_generation_request_ids = None def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Compact after native KV-cache updates have finalized this iteration - (must run after KVCacheManagerV2 so capacity reflects the written token and any rewind).""" + """Compact after native KV-cache updates finalize the iteration. + + KVCacheManagerV2 must run first so capacity includes the written token and any rewind. + """ with nvtx_range_debug("triattention.generation_step_end", color="blue"): self._evict_due_requests(scheduled_batch) def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: - """Drop this request's eviction state; the buffers stay resident.""" + """Drop this request's eviction state; persistent allocations stay resident.""" self._request_states.pop(request.py_request_id, None) # ---- request capacity ---- def _validate_request_capacity(self, request: "LlmRequest") -> None: - """Reject a request whose pre-first-eviction peak cannot fit (only - TriAttention can compute it: the framework's dense guards are off).""" + """Reject requests whose maximum pre-eviction peak cannot fit.""" speculative_overshoot = int(self.kv_cache_manager.max_draft_len) - first_eviction_decode_length = ( - self.budget // self.beta + 1 - ) * self.beta + speculative_overshoot - decode_capacity = min(int(request.py_max_new_tokens), first_eviction_decode_length) + decode_capacity = min( + int(request.py_max_new_tokens), + self.budget + self.beta + speculative_overshoot, + ) confirmed_capacity = int(request.py_prompt_len) + decode_capacity checked = [(self.kv_cache_manager, self._protected_tail_capacity, "target")] if self.draft_kv_cache_manager is not None: @@ -401,8 +414,8 @@ def _validate_request_capacity(self, request: "LlmRequest") -> None: table_capacity = manager.max_blocks_per_seq * manager.tokens_per_block if confirmed_capacity > pool_capacity or required_capacity > table_capacity: raise ValueError( - f"TriAttention {label} KV capacity is too small to reach the first " - f"eviction: request requires {required_capacity} tokens " + f"TriAttention {label} KV capacity is too small for eviction: " + f"request requires {required_capacity} tokens " f"(prompt={request.py_prompt_len}, budget={self.budget}, " f"beta={self.beta}, protected tail={protected_tail}), but the " f"V2 pool covers {pool_capacity + protected_tail} tokens and " @@ -415,8 +428,7 @@ def _evict_due_requests( self, scheduled_batch: "ScheduledRequests", ) -> None: - """Owner of the full eviction transaction: admission, cadence, launch, - publication, and cache resize.""" + """Collect due requests, execute one eviction round, publish, and resize.""" manager = self.kv_cache_manager eviction_inputs: List[_EvictionInput] = [] with nvtx_range("triattention.metadata", color="cyan"): @@ -429,12 +441,20 @@ def _evict_due_requests( # Overlap scheduling may suspend a cache mid-flight; defer # this request (pre-launch) instead of failing the batch. continue - # Cadence gate first; capacity math and consistency raises run in the due branch. + 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 + # Only requests with active target/draft caches advance the cadence ledger. state = self._request_states[request_id] - previous_step = state["generation_steps"] - step = previous_step + 1 + int(request.py_num_accepted_draft_tokens) - state["generation_steps"] = step - if previous_step // self.beta >= step // self.beta: + previous_confirmed_tokens = state.confirmed_tokens + confirmed_tokens = ( + previous_confirmed_tokens + 1 + int(request.py_num_accepted_draft_tokens) + ) + state.confirmed_tokens = confirmed_tokens + if previous_confirmed_tokens // self.beta >= confirmed_tokens // self.beta: continue # Speculative reserve + in-flight overlap growth: contiguous tail moved byte-for-byte. target_tail_length = self._num_extra_kv_tokens + ( @@ -450,13 +470,6 @@ def _evict_due_requests( if source_length <= prompt_length + self.budget: # Selection would be an identity: nothing to evict yet. continue - draft_cache = None - if self.draft_kv_cache_manager is not None: - # A missing draft cache is a wiring bug: the dict's KeyError is the report. - draft_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] - if not draft_cache.is_active: - # Target and draft defer together (pre-launch). - continue eviction_inputs.append( _EvictionInput( request=request, @@ -465,7 +478,7 @@ def _evict_due_requests( state=state, source_length=source_length, # Uncompressed logical position. - logical_source_length=source_length + state["evicted_tokens"], + logical_source_length=source_length + state.evicted_tokens, prompt_length=prompt_length, target_tail_length=target_tail_length, ) @@ -473,16 +486,6 @@ def _evict_due_requests( if not eviction_inputs: return - with nvtx_range_debug("triattention.resolve_layout", color="blue"): - target_layout = self._runtime_kv_layout() - draft_layout = ( - self._runtime_kv_layout(draft=True) - if self.draft_kv_cache_manager is not None - else None - ) - with nvtx_range_debug("triattention.staging_lookup", color="blue"): - # Retained spans always cover the model window (construction rejects budget < window). - self._ensure_eviction_runtime(target_layout, draft_layout, eviction_inputs) # Ungated NVTX: the due count in the message shows each round's size. with nvtx_range( f"triattention.evict_request_group reqs={len(eviction_inputs)}", @@ -491,9 +494,9 @@ def _evict_due_requests( self._execute_eviction_round(eviction_inputs) for item in eviction_inputs: evicted = item.source_length - item.prompt_length - self.budget - item.state["evicted_tokens"] += evicted + item.state.evicted_tokens += evicted # The manager's only channel to the runtime (feeds num_cached_tokens_per_seq). - item.request.py_num_compressed_tokens = item.state["evicted_tokens"] + item.request.py_num_compressed_tokens = item.state.evicted_tokens self._resize_compacted_caches(eviction_inputs) def _inflight_generation_growth( @@ -514,71 +517,66 @@ def _execute_eviction_round( self, eviction_inputs: Sequence[_EvictionInput], ) -> None: - """Run one eviction round over the due cohort (every launch covers - the full request capacity; padded rows carry zero lengths and stay inert).""" + """Score, select, and compact one due request group.""" manager = self.kv_cache_manager draft_manager = self.draft_kv_cache_manager - with nvtx_range_debug("triattention.page_table_stage", color="orange"): - request_ids = [item.request.py_request_id for item in eviction_inputs] - logical_source_lengths = [item.logical_source_length for item in eviction_inputs] - prompt_lengths = [item.prompt_length for item in eviction_inputs] - source_lengths = [item.source_length for item in eviction_inputs] - dense_move_offsets, swa_move_offsets, draft_move_offsets = ( - self._compute_compaction_move_offsets(eviction_inputs) - ) - stream = torch.cuda.current_stream(self._block_offsets_device.device) - # int32 gate before any buffer or device work: the in-place numpy writes below wrap silently. - max_logical_source_length = max(logical_source_lengths) - rows = ( - (0, logical_source_lengths), - (1, source_lengths), - (2, prompt_lengths), - (3, dense_move_offsets), - (4, swa_move_offsets), - (5, draft_move_offsets), - ) - for row, values in rows: - if ( - values is not None - and not -0x80000000 <= min(values) <= max(values) <= 0x7FFFFFFF - ): - raise ValueError(f"staged metadata row {row} exceeds the int32 range") - # Host-staging reuse fence: prior cohort's async copies must finish before the pinned rows are rewritten. - self._staging_reuse_event.synchronize() - host_table = self._request_metadata_host_np - for row, values in rows: - if values is not None: - host_table[row, : len(values)] = values - # Zero lengths keep the score kernel and selection inert for padded rows. - host_table[:3, len(eviction_inputs) :] = 0 - grow_mean_phase_table(self._phase, int(max_logical_source_length) + 1) - self._stage_block_offsets( - manager, - request_ids, - self._block_offsets_host, - self._block_offsets_device, - ) - if draft_manager is not None: - self._stage_block_offsets( - draft_manager, - request_ids, - self._draft_block_offsets_host, - self._draft_block_offsets_device, - ) - try: - self._request_metadata_device.copy_(self._request_metadata_host, non_blocking=True) - finally: - # Guards the pinned staging until the asynchronous copies complete. - self._staging_reuse_event.record(stream) - request_count = len(eviction_inputs) - union = self.eviction_mode == "union" + 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: + with nvtx_range_debug("triattention.page_table_stage", color="orange"): + request_ids = [item.request.py_request_id for item in eviction_inputs] + logical_source_lengths = [item.logical_source_length for item in eviction_inputs] + prompt_lengths = [item.prompt_length for item in eviction_inputs] + source_lengths = [item.source_length for item in eviction_inputs] + dense_move_offsets, swa_move_offsets, draft_move_offsets = ( + self._compute_compaction_move_offsets(eviction_inputs) + ) + 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_inputs) :] = 0 + try: + self._stage_block_offsets( + manager, + request_ids, + self._block_offsets_host, + self._block_offsets_device, + ) + if draft_manager is not None: + self._stage_block_offsets( + 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_inputs) + union = self.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): # In-place refresh: the compiled score launches captured these pointers. _gather_mean_phase_kernel[(request_count,)]( self._logical_source_lengths_device, - self._phase["cos"], - self._phase["sin"], + self._phase.cos, + self._phase.sin, self._source_lengths_device, self._prompt_lengths_device, self._mean_cos, @@ -586,48 +584,34 @@ def _execute_eviction_round( self._decode_lengths_device, self._swa_destination_bases, self._swa_rebase_delta, - NUM_FREQS=self._phase_num_freqs, - F_BLOCK=self._phase_f_block, + NUM_FREQS=self._phase.num_freqs, + F_BLOCK=self._phase.frequency_block, HAS_SWA=self._swa_destination_bases is not None, num_warps=1, ) - cu_stream = cuda_driver.CUstream(stream.cuda_stream) - self._compiled_score_by_request_count[request_count]( - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, - request_count, - cu_stream, - ) - if union: - # Normalized union reduction, written straight into the selection rows. - self._compiled_normalize_union_by_request_count[request_count]( - self._cute_partial_stats, - *self._cute_selection_prefix, - self._cute_selection_scores_rows, - request_count, - cu_stream, + 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, ) - if self._union_tp_mapping is not None: - # Max-fold the rank-local unions into the global union (exact: - # max is order-free), so every rank keeps the same ordinals. - gathered = allgather( - self._selection_scores_rows[:request_count], - self._union_tp_mapping, - dim=0, - ) - _fold_union_ranks_kernel[(request_count, self._fold_width_blocks)]( - gathered, - self._selection_scores_rows, + _fold_union_ranks_kernel[ + ( request_count, - TP_SIZE=self._union_tp_size, - WIDTH=self._selection_width_capacity, + triton.cdiv(self._selection_width_capacity, 1024), ) + ]( + gathered, + self._selection_scores_rows, + request_count, + TP_SIZE=int(self._union_tp_mapping.tp_size), + WIDTH=self._selection_width_capacity, + ) with nvtx_range("triattention.select", color="yellow"): if not union: - # Per-head reduces read each decode window straight out of the scratch. - prepare_per_head_scores( + reduce_per_head_scores( self._score_scratch, self._decode_lengths_device, self._prompt_lengths_device, @@ -636,32 +620,25 @@ def _execute_eviction_round( self._selection_scores_rows, self._selection_row_lengths, request_count=request_count, - num_layers=self._num_layers, - num_q_heads=self._num_q_heads, - num_kv_heads=self._num_kv_heads, - padded_head_columns=self._padded_head_columns, + padded_head_columns=PADDED_HEAD_COLUMNS, score_token_capacity=self._score_token_capacity, - selection_width=self._selection_width_capacity, per_layer=self.eviction_mode == "per_layer_perhead", normalize_scores=self.normalize_scores, ) - self._settle_top_tokens(request_count) + self._select_top_tokens(request_count) with nvtx_range("triattention.compact", color="purple"): compact(self._compaction_params, request_count) finally: - # Order V2 page-table reuse and resize after this cohort's compact. + # Target and draft V2 managers share this execution stream. self._compaction_done_event.record(stream) - manager._stream.wait_event(self._compaction_done_event) - if draft_manager is not None: - draft_manager._stream.wait_event(self._compaction_done_event) + if manager._stream != stream: + manager._stream.wait_event(self._compaction_done_event) def _compute_compaction_move_offsets( self, eviction_inputs: Sequence[_EvictionInput], ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: - """Cumulative dense/SWA/draft move offsets for one due cohort (keep - set plus protected tail per request; rows past the cohort repeat the final - offset and contribute no moves).""" + """Build padded cumulative dense, SWA, and draft move offsets.""" def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets = [0] @@ -671,14 +648,14 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: return offsets tails = [int(item.target_tail_length) for item in eviction_inputs] - dense = padded_offsets([self._keep_count + tail for tail in tails]) + dense = padded_offsets([self.budget + tail for tail in tails]) swa = None if self._swa_window is not None: swa = padded_offsets([self._swa_window + tail for tail in tails]) draft = None - if self._draft_protected_tail_capacity is not None: + if self.draft_kv_cache_manager is not None: draft = padded_offsets( - [self._keep_count + self._draft_protected_tail_capacity] * len(eviction_inputs) + [self.budget + self._draft_protected_tail_capacity] * len(eviction_inputs) ) return dense, swa, draft @@ -689,8 +666,7 @@ def _stage_block_offsets( host_block_offsets: torch.Tensor, device_block_offsets: torch.Tensor, ) -> None: - """Gather the pinned snapshot before the async device copy: resize mutates - the live host table. The round owner has already fenced host-staging reuse.""" + """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, @@ -703,23 +679,18 @@ def _stage_block_offsets( self._identity_copy_indices_host[: len(request_ids)], manager.index_scales, manager.kv_offset, - manager._stream.cuda_stream, - ) - self._block_offsets_ready_event.record(manager._stream) - torch.cuda.current_stream(device_block_offsets.device).wait_event( - self._block_offsets_ready_event + torch.cuda.current_stream(device_block_offsets.device).cuda_stream, ) - def _settle_top_tokens(self, request_count: int) -> None: - """Pick the top-k and settle ties into the kept-ordinal decision rows - (the compaction contract packs them into move sources).""" + def _select_top_tokens(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._keep_count, + self.budget, 1, ) _settle_ties_kernel[(request_count, self._selection_rows_per_request)]( @@ -729,585 +700,319 @@ def _settle_top_tokens(self, request_count: int) -> None: self._provisional_rows, self._kept_ordinal_rows, WIDTH=self._selection_width_capacity, - KEEP_COUNT=self._keep_count, + KEEP_COUNT=self.budget, SELECTION_ROWS=self._selection_rows_per_request, ) - def _resize_compacted_caches(self, eviction_inputs) -> None: + def _resize_compacted_caches(self, eviction_inputs: Sequence[_EvictionInput]) -> None: with nvtx_range("triattention.resize", color="red"): - with nvtx_range_debug("triattention.v2_resize", color="red"): - families = [("target", "target_cache", None)] - if self.draft_kv_cache_manager is not None: - # Same kept set: the draft shrinks to the same retained - # length plus its own fixed tail. - families.append(("draft", "draft_cache", self._draft_protected_tail_capacity)) - for label, cache_key, fixed_tail in families: - for item in eviction_inputs: - cache = getattr(item, cache_key) - request_id = item.request.py_request_id - tail = item.target_tail_length if fixed_tail is None else fixed_tail - resized_capacity = item.prompt_length + self.budget + tail - if not cache.resize(resized_capacity, None): - raise RuntimeError( - f"Failed to resize compacted {label} KV cache for " - f"request {request_id} to {resized_capacity} tokens" - ) - - # ---- buffers + layout ---- - - def _ensure_eviction_runtime( - self, - target_layout: Dict[str, object], - draft_layout: Optional[Dict[str, object]], - eviction_inputs: Sequence[_EvictionInput], - ) -> None: - """Per-round reuse gate over the three capacity axes; first round or - growth replaces the resident runtime as a whole.""" - # Empty cohorts never reach here: _evict_due_requests no-ops pre-launch. - needed_width = max(item.source_length - item.prompt_length for item in eviction_inputs) - needed_score_tokens = max(item.source_length for item in eviction_inputs) - needed_requests = len(eviction_inputs) - if self._buffers_built: - if ( - needed_width <= self._selection_width_capacity - and needed_score_tokens <= self._score_token_capacity - and needed_requests <= self._request_capacity - ): + for item in eviction_inputs: + resized_capacity = item.prompt_length + self.budget + item.target_tail_length + if not item.target_cache.resize(resized_capacity, None): + raise RuntimeError( + "Failed to resize compacted target KV cache for " + f"request {item.request.py_request_id} to " + f"{resized_capacity} tokens" + ) + if self.draft_kv_cache_manager is None: return - # This round outgrew the runtime: wait out the prior round (its - # completion event orders after every use of the old epoch), then rebuild. - if self._compaction_done_event is not None: - self._compaction_done_event.synchronize() - self._buffers_built = False - - needed_page_tokens = max( - item.source_length + item.target_tail_length for item in eviction_inputs - ) - manager = self.kv_cache_manager - request_capacity = max(needed_requests, int(manager.max_batch_size)) - selection_width_capacity = max( - needed_width, - self.budget + 2 * self.beta + int(manager.max_total_draft_tokens or 0), - ) - # Bucket sized by the presented cohorts, NOT max_seq_len (a floor there breaks 32-bit indexing). - score_token_capacity = next_positive_power_of_2(max(int(needed_page_tokens), 1024)) - score_token_capacity = min( - score_token_capacity, max(int(manager.max_seq_len), int(needed_page_tokens)) - ) - # The bucket capacity must be tile-aligned (mis-tiling stripes the - # score scratch silently); the ceiling division constructs that fact. - score_tile_tokens = max(64, int(manager.tokens_per_block)) - score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens + # Target and draft retain the same selected tokens. + for item in eviction_inputs: + resized_capacity = ( + item.prompt_length + self.budget + self._draft_protected_tail_capacity + ) + if not item.draft_cache.resize(resized_capacity, None): + raise RuntimeError( + "Failed to resize compacted draft KV cache for " + f"request {item.request.py_request_id} to " + f"{resized_capacity} tokens" + ) - first_pool = target_layout["layer_pools"][target_layout["dense_layers"][0]] - if self._phase is None: - # Host-only width offsets for the table builder (no device copy). - self._phase = { - "omega": self._omega.to(device=first_pool.device, dtype=torch.float32).contiguous(), - "offset_values": [ - float(1 << i) for i in range(_MAX_INTEGRATION_OFFSET.bit_length()) - ], - "cos": None, - "sin": None, - "rows": 0, - } - grow_mean_phase_table(self._phase, int(score_token_capacity)) - self._rebuild_eviction_runtime( - target_layout, - draft_layout, - request_capacity=request_capacity, - score_token_capacity=score_token_capacity, - selection_width_capacity=selection_width_capacity, - ) - self._buffers_built = True + # ---- persistent state + request capacity ---- - def _rebuild_eviction_runtime( - self, - target_layout: Dict[str, object], - draft_layout: Optional[Dict[str, object]], - *, - request_capacity: int, - score_token_capacity: int, - selection_width_capacity: int, - ) -> None: - """Build the resident eviction runtime for one capacity epoch as attributes - (compiled kernels capture raw pool addresses: pools must stay alive and stay put).""" - import cutlass - import cutlass.cute as cute - - from .triattention_cute_score_fused import ( - _COMPILE_LOCK, - _COMPILED_KERNELS, - PADDED_HEAD_COLUMNS, - SMALL_WORKLOAD_PAGE_SHARDS, - STATS_FIELDS, - _encode_tma_descriptors, - _tensor_spec, - _to_cute, - _TriAttentionScoreKernel, + 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 = list(target_layout["dense_layers"]) - swa_layers = list(target_layout["swa_layers"]) - swa_window = target_layout["swa_window"] - # Canonical layer -> V2 pool id tuple; it IS the staged plane slot map. - layer_pool_ids = tuple(target_layout["layer_pool_ids"]) - num_page_table_slots = int(self.kv_cache_manager.num_pools) - - # The first dense layer anchors device and staging geometry. + dense_layers = target_layout["dense_layers"] anchor_pool = layer_pools[dense_layers[0]] device = anchor_pool.device - request_capacity = int(request_capacity) - score_token_capacity = int(score_token_capacity) - selection_width_capacity = int(selection_width_capacity) - keep_count = int(self.budget) - protected_tail_capacity = int(self._protected_tail_capacity) - page_table_token_capacity = score_token_capacity + protected_tail_capacity - - q_real, q_imag, mlr_coef = self._local_score_calibration(target_layout["global_layers"]) - # TP splits attention heads contiguously per rank; the calibration is global. + _, _, 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 + + q_real, q_imag, mlr_coef = self._local_score_calibration() 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] - # Union reduces over ALL heads: rank-local rows are max-folded across the - # TP group each round so the kept set matches the single-rank algorithm. - self._union_tp_mapping = ( - mapping if (self.eviction_mode == "union" and tp_size > 1) else None - ) - self._union_tp_size = tp_size - q_real, q_imag, mlr_coef, freq_scale_sq = ( + 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._request_capacity = request_capacity - self._score_token_capacity = score_token_capacity - self._selection_width_capacity = selection_width_capacity - self._fold_width_blocks = triton.cdiv(selection_width_capacity, 1024) - self._keep_count = keep_count + 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) - # ---- block-offset staging (target, plus the co-compressed draft) ------- - self._block_offsets_host, self._block_offsets_device = _allocate_block_offset_staging( - anchor_pool, - num_pools=num_page_table_slots, - request_capacity=request_capacity, - token_capacity=page_table_token_capacity, - max_source_blocks=int(self.kv_cache_manager.host_kv_cache_block_offsets.shape[-1]), + 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" ) - # The draft is never scored: these offsets feed only the draft compacts. - self._draft_block_offsets_device = None - self._draft_block_offsets_host = None - if draft_layout is not None: - draft_representatives = list(draft_layout["pool_representatives"]) - draft_anchor_pool = draft_layout["layer_pools"][draft_representatives[0]] - # Construction-boundary invariant: the round shares one stream/event - # contract, so the draft pools must live on the target device. - if draft_anchor_pool.device != device: - raise RuntimeError( - "TriAttention draft KV pools must share the target KV pool device" - ) - self._draft_block_offsets_host, self._draft_block_offsets_device = ( - _allocate_block_offset_staging( - draft_anchor_pool, - num_pools=int(self.draft_kv_cache_manager.num_pools), - request_capacity=request_capacity, - token_capacity=score_token_capacity + int(self._draft_protected_tail_capacity), - max_source_blocks=int( - self.draft_kv_cache_manager.host_kv_cache_block_offsets.shape[-1] - ), - ) - ) - # ---- per-round metadata table: one H2D copy; move-offsets rows need the +1 column ---- + def _allocate_metadata_buffers( + self, + device: torch.device, + *, + num_freqs: int, + ) -> None: + """Allocate manager-lifetime host staging and device metadata.""" + row_count = 6 + request_capacity = self._request_capacity self._request_metadata_host = torch.empty( - (6, request_capacity + 1), dtype=torch.int32, device="cpu", pin_memory=prefer_pinned() + (row_count, request_capacity + 1), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), ) - # numpy view over the pinned rows: per-round staging writes lists in place. 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() + request_capacity, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), ) - # Zero-filled: an unstaged cohort must gather the phase table's row 0. self._request_metadata_device = torch.zeros( - (6, request_capacity + 1), dtype=torch.int32, device=device + (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] - # Pinned per-request decode-window starts. self._prompt_lengths_device = self._request_metadata_device[2, :request_capacity] - dense_move_offsets_row = self._request_metadata_device[3] - swa_move_offsets_row = self._request_metadata_device[4] - draft_move_offsets_row = self._request_metadata_device[5] - # SWA staging geometry, bound once (the compaction plans stay opaque): - # the phase gather rebases each request's SWA destination base in place. - self._swa_window = int(swa_window) if swa_layers else None + 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 swa_layers else None + 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._swa_rebase_delta = keep_count - self._swa_window if swa_layers 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) - self._phase_num_freqs = int(self._phase["omega"].numel()) - self._phase_f_block = triton.next_power_of_2(self._phase_num_freqs) - # ---- score state: one fused group across all dense layers -------------- - _, _, num_kv_heads, tokens_per_block, _ = anchor_pool.shape - self._num_layers = len(dense_layers) - self._num_q_heads = int(num_q_heads) - self._num_kv_heads = int(num_kv_heads) - dense_layer_slots = [layer_pool_ids[layer] for layer in dense_layers] - seg_req_id = torch.arange( - request_capacity, dtype=torch.int32, device=device - ).repeat_interleave(self._num_layers) - seg_layer_id = torch.tensor(list(dense_layers), dtype=torch.int32, device=device).repeat( - request_capacity - ) - block_offsets = self._block_offsets_device - slots_t = torch.tensor(dense_layer_slots, dtype=torch.int64, device=device) - req_idx = seg_req_id.to(torch.int64) - slot_idx = slots_t.repeat(request_capacity) - seg_page_off = slot_idx * block_offsets.stride(0) + req_idx * block_offsets.stride(1) - - max_segments = request_capacity * self._num_layers - # The score plane must stay 32-bit indexable (wraparound = silent wild read). - if (PADDED_HEAD_COLUMNS - 1) * max_segments * score_token_capacity >= 2**31: - raise ValueError( - "score bucket overflows the 32-bit score plane: " - f"{(PADDED_HEAD_COLUMNS - 1) * max_segments * score_token_capacity}" - ) - # Persistent buffers: the compiled kernels capture their device pointers. - self._padded_head_columns = PADDED_HEAD_COLUMNS - self._score_scratch = torch.empty( - self._num_kv_heads * PADDED_HEAD_COLUMNS * max_segments * score_token_capacity, - dtype=torch.float32, - device=device, - ) - # int32 is safe here: covered by the 2^31 score-plane audit above. - seg_out_offset = ( - torch.arange(max_segments, dtype=torch.int64, device=device) * score_token_capacity - ).to(torch.int32) + def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None: + """Allocate fixed TopK inputs and outputs for the configured mode.""" + request_capacity = self._request_capacity + selection_width = self._selection_width_capacity union = self.eviction_mode == "union" - - # ---- score path: compiled per request-count/page-shard ---- - sm_count = int(torch.cuda.get_device_properties(device).multi_processor_count) - # Per-shard partial score statistics. - partial_stats_elements = ( - request_capacity - * self._num_layers - * num_q_heads - * SMALL_WORKLOAD_PAGE_SHARDS - * STATS_FIELDS - if union - else 1 - ) - self._partial_stats = torch.empty( - partial_stats_elements, - dtype=torch.float32, - device=device, - ) - self._tma_descriptors = _encode_tma_descriptors( - list(layer_pools), - [int(layer) for layer in dense_layers], - int(num_freqs), - int(tokens_per_block), - ) - # Alignment sits with the operand it describes; source_lengths and - # prompt_lengths are only 4-byte-aligned row views, read as per-CTA scalars. - prefix_operands = ( - (block_offsets.view(-1), 16), - (seg_page_off, 16), - (seg_req_id, 16), - (seg_layer_id, 16), - (self._source_lengths_device, 4), - (seg_out_offset, 16), - (self._prompt_lengths_device, 4), - (q_real.view(-1), 16), - (q_imag.view(-1), 16), - (mlr_coef.view(-1), 16), - ) - torch_prefix = tuple(tensor for tensor, _ in prefix_operands) - torch_tail = ( - freq_scale_sq, - self._score_scratch, - self._partial_stats, - anchor_pool, - self._tma_descriptors, - ) - # No keep-alive twin: each cute handle below owns its operand's DLPack - # capsule, which retains the underlying torch storage. - self._cute_score_prefix = tuple( - _to_cute(tensor, assumed_align=align) for tensor, align in prefix_operands - ) - self._cute_score_tail = ( - _to_cute(freq_scale_sq), - _to_cute(self._score_scratch), - _to_cute(self._partial_stats), - _to_cute(anchor_pool), - _to_cute(self._tma_descriptors, assumed_align=128), - ) - # Build-bound persistent launch operands: refreshed in place each round. - self._cute_mean_cos = _to_cute(self._mean_cos.view(-1)) - self._cute_mean_sin = _to_cute(self._mean_sin.view(-1)) - self._compiled_score_by_request_count: Dict[int, object] = {} - self._compiled_normalize_union_by_request_count: Dict[int, object] = {} - page_shards_by_request_count: Dict[int, int] = {} - self._cute_selection_prefix = ( - _to_cute(self._score_scratch), - _to_cute(self._source_lengths_device, assumed_align=4), - _to_cute(seg_out_offset), - _to_cute(self._prompt_lengths_device, assumed_align=4), - ) - self._cute_partial_stats = _to_cute(self._partial_stats) - static_geometry = ( - request_capacity, - self._num_layers, - score_token_capacity, - num_q_heads, - self._num_kv_heads, - num_freqs, - int(tokens_per_block), - tuple(int(value) for value in anchor_pool.shape), - tuple(int(value) for value in anchor_pool.stride()), - ) - tensor_specs = tuple( - _tensor_spec(tensor) - for tensor in ( - *torch_prefix, - self._mean_cos.view(-1), - self._mean_sin.view(-1), - *torch_tail, - ) - ) - variants = [(1, SMALL_WORKLOAD_PAGE_SHARDS)] - if request_capacity > 1: - variants.append((request_capacity, 2)) - # Per-head modes compile the score-only entry; union the fused pipeline. - kernel_kwargs = dict( - num_layers=self._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()), - ) - stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) - - def _compiled_kernel(cache_key, build): - with _COMPILE_LOCK: - compiled = _COMPILED_KERNELS.get(cache_key) - if compiled is None: - compiled = build() - _COMPILED_KERNELS[cache_key] = compiled - return compiled - - variant_key = "triattention_cute_score_stats" if union else "triattention_cute_score" - compiled_entries = self._compiled_score_by_request_count - for request_count, page_shards in variants: - cache_key = ( - variant_key, - static_geometry, - tensor_specs, - request_count, - page_shards, - ) - compiled_entries[request_count] = _compiled_kernel( - cache_key, - lambda page_shards=page_shards: cute.compile( - _TriAttentionScoreKernel( - **kernel_kwargs, - page_shards=page_shards, - write_partial_stats=union, - ), - *self._cute_score_prefix, - self._cute_mean_cos, - self._cute_mean_sin, - *self._cute_score_tail, - cutlass.Int32(1), - stream, - ), + if union: + self._selection_rows_per_request = 1 + elif self.eviction_mode == "per_head": + self._selection_rows_per_request = self._num_kv_heads + else: + self._selection_rows_per_request = self._num_layers * self._num_kv_heads + selection_rows = request_capacity * self._selection_rows_per_request + selection_rect = selection_rows * max(selection_width, self.budget) + if union: + selection_rect = max( + selection_rect, + tp_size * request_capacity * selection_width, ) - page_shards_by_request_count[request_count] = page_shards - - if request_capacity > 1: - small = compiled_entries[1] - large = compiled_entries[request_capacity] - for request_count in range(1, request_capacity + 1): - # Give small cohorts the extra shard while the 2-shard grid stays under two waves. - two_shard_ctas = request_count * self._num_layers * self._num_kv_heads * 2 - use_extra_score_shard = two_shard_ctas < 2 * sm_count - compiled_entries[request_count] = small if use_extra_score_shard else large - page_shards_by_request_count[request_count] = ( - SMALL_WORKLOAD_PAGE_SHARDS if use_extra_score_shard else 2 - ) - - logger.info( - f"TriAttention CuTe score enabled: {self._num_q_heads}q/{self._num_kv_heads}kv heads, " - f"{num_freqs} freqs, {int(tokens_per_block)}-token pages" - ) + if selection_rect >= 2**31: + raise ValueError(f"selection rectangle overflows 32-bit indexing: {selection_rect}") - # ---- selection buffers (canonical row-major, one name per storage) ----- self._decode_lengths_device = torch.full( - (request_capacity,), selection_width_capacity, dtype=torch.int32, device=device + (request_capacity,), selection_width, dtype=torch.int32, device=device ) if union: - self._selection_rows_per_request = 1 self._selection_scores_rows = torch.empty( - (request_capacity, selection_width_capacity), dtype=torch.float32, device=device + (request_capacity, selection_width), + dtype=torch.float32, + device=device, ) - # One selection row per request: its length IS the staged valid width. self._selection_row_lengths = self._decode_lengths_device - # Padded rows still need in-range ordinals for the finalizer's gather. - self._provisional_rows = torch.zeros( - (request_capacity, keep_count), dtype=torch.int32, device=device - ) - # Kept decode ordinals. - self._kept_ordinal_rows = torch.empty( - (request_capacity, keep_count), dtype=torch.int32, device=device - ) - self._cute_selection_scores_rows = _to_cute(self._selection_scores_rows.view(-1)) else: - selection_rows = ( - self._num_kv_heads - if self.eviction_mode == "per_head" - else self._num_layers * self._num_kv_heads + score_shape = ( + request_capacity, + self._num_layers, + self._num_q_heads, + 1, ) - # The selection rectangle must stay 32-bit indexable (wraparound = wild reads). - selection_rect = ( - request_capacity * selection_rows * max(selection_width_capacity, keep_count) - ) - if selection_rect >= 2**31: - raise ValueError( - f"per-head selection rectangle overflows 32-bit indexing: {selection_rect}" - ) - self._selection_rows_per_request = selection_rows - 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( - (request_capacity * selection_rows, selection_width_capacity), + (selection_rows, selection_width), dtype=torch.float32, device=device, ) self._selection_row_lengths = torch.full( - (request_capacity * selection_rows,), - selection_width_capacity, + (selection_rows,), + selection_width, dtype=torch.int32, device=device, ) - self._provisional_rows = torch.zeros( - (request_capacity * selection_rows, keep_count), dtype=torch.int32, device=device - ) - self._kept_ordinal_rows = torch.empty( - (request_capacity * selection_rows, keep_count), dtype=torch.int32, device=device - ) - if union: - from .triattention_cute_selection import ( - _select_normalize_union_config, - _TriAttentionNormalizeUnionKernel, + 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 _reserve_eviction_capacity(self, request: "LlmRequest") -> None: + """Reserve all request-dependent runtime capacity at admission.""" + first_evict_step = (self.budget // self.beta + 1) * self.beta + if int(request.py_max_new_tokens) < first_evict_step: + return + self._phase.reserve(int(request.py_prompt_len) + int(request.py_max_new_tokens) + 1) + decode_tokens = min(int(request.py_max_new_tokens), self._selection_width_capacity) + required_source_tokens = int(request.py_prompt_len) + decode_tokens + if required_source_tokens <= self._score_token_capacity: + return + + # A newly admitted request may require larger score state while work + # from an older request is still in flight. + if self._launch_score is not None: + self._compaction_done_event.synchronize() + + manager = self.kv_cache_manager + # Bucket the largest admitted source rather than eagerly reserving + # max_seq_len; score scratch is the dominant allocation. + score_token_capacity = next_positive_power_of_2(max(required_source_tokens, 1024)) + score_token_capacity = min( + score_token_capacity, max(int(manager.max_seq_len), required_source_tokens) + ) + # The bucket capacity must be tile-aligned (mis-tiling stripes the + # score scratch silently); the ceiling division constructs that fact. + score_tile_tokens = max(64, int(manager.tokens_per_block)) + score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens + + self._build_eviction_capacity( + score_token_capacity=score_token_capacity, + ) + + def _build_eviction_capacity( + self, + *, + score_token_capacity: int, + ) -> None: + """Build and publish all state bound to one score-token capacity.""" + 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_staging( + anchor_pool, + num_pools=int(self.kv_cache_manager.num_pools), + request_capacity=request_capacity, + token_capacity=score_token_capacity + self._protected_tail_capacity, + max_source_blocks=int(self.kv_cache_manager.host_kv_cache_block_offsets.shape[-1]), + ) + 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_staging( + draft_anchor_pool, + num_pools=int(self.draft_kv_cache_manager.num_pools), + request_capacity=request_capacity, + token_capacity=(score_token_capacity + self._draft_protected_tail_capacity), + max_source_blocks=int( + self.draft_kv_cache_manager.host_kv_cache_block_offsets.shape[-1] + ), ) - compiled_configs: Dict[Tuple[int, int, int, int], object] = {} - 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) - compiled_selection = compiled_configs.get(config_key) - if compiled_selection is None: - cache_key = ( - "triattention_cute_normalize_union", - static_geometry, - tensor_specs, - config_key, - _tensor_spec(self._selection_scores_rows), - _tensor_spec(self._partial_stats), - ) - tokens_per_lane, token_subtiles, row_cluster_ctas = config - compiled_selection = _compiled_kernel( - cache_key, - lambda page_shards=page_shards, - tokens_per_lane=tokens_per_lane, - token_subtiles=token_subtiles, - row_cluster_ctas=row_cluster_ctas: cute.compile( - _TriAttentionNormalizeUnionKernel( - num_layers=self._num_layers, - score_token_capacity=score_token_capacity, - num_q_heads=num_q_heads, - # The finalizer maps real head rows onto N=8-padded planes. - num_kv_heads=self._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=selection_width_capacity, - ), - self._cute_partial_stats, - *self._cute_selection_prefix, - self._cute_selection_scores_rows, - cutlass.Int32(1), - stream, - ), - ) - compiled_configs[config_key] = compiled_selection - self._compiled_normalize_union_by_request_count[request_count] = compiled_selection + 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 plans (opaque: only compact() interprets them) --------- compaction_params = [ build_compaction_params( - target_layout, - block_offsets=self._block_offsets_device, + 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, - # Per-round tails: the move offsets ride the staged metadata rows. - dense_move_offsets=dense_move_offsets_row, - protected_tail_capacity=protected_tail_capacity, - swa_move_offsets=swa_move_offsets_row, + 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 draft_layout is not None: + if self._draft_layout is not None: compaction_params.append( build_compaction_params( - draft_layout, - block_offsets=self._draft_block_offsets_device, + 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=draft_move_offsets_row, - protected_tail_capacity=int(self._draft_protected_tail_capacity), + dense_move_offsets=self._draft_move_offsets_device, + protected_tail_capacity=self._draft_protected_tail_capacity, ) ) - self._compaction_params = tuple(compaction_params) - # ---- round-ordering events ---------------------------------------------- - # Device-lifetime: created once and reused across capacity rebuilds - # (they carry no pointer state; replacing them orphans in-flight ordering). - if self._staging_reuse_event is None: - # Host staging (pinned metadata + snapshots) reuse fence. - self._staging_reuse_event = torch.cuda.Event() - self._staging_reuse_event.record(torch.cuda.current_stream(device)) - # Manager-stream H2D of the block-offset tables has completed. - self._block_offsets_ready_event = torch.cuda.Event() - # This cohort's compact is done: manager may resize/reuse pages. - self._compaction_done_event = torch.cuda.Event() + # 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 _local_score_calibration( self, - global_layers: List[int], ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + global_layers = self._global_layers num_layers = len(global_layers) if global_layers and max(global_layers) >= self._calibration_q_real.shape[0]: raise ValueError( @@ -1332,65 +1037,32 @@ def _local_score_calibration( self._calibration_mlr_coef.index_select(0, layer_ids), ) - def _runtime_kv_layout(self, *, draft: bool = False) -> Dict[str, object]: - # V2 pools are allocated once at manager init; the layout is a manager-lifetime contract. + 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 - cached = self._kv_layout_caches[draft] - if cached is not None: - return cached if draft: global_layers = [int(layer) for layer in manager.pp_layers] - if not global_layers: - raise RuntimeError("TriAttention draft KV cache manager exposes no 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, swa_layers, swa_window = self._layer_partition - if not dense_layers: - raise ValueError("TriAttention requires at least one full-attention layer") - layout = self._build_runtime_kv_layout( - manager, - global_layers, - dense_layers=dense_layers, - swa_layers=swa_layers, - swa_window=swa_window, - label="draft " if draft else "", - ) - self._kv_layout_caches[draft] = layout - return layout - - def _build_runtime_kv_layout( - self, - manager: KVCacheManagerV2, - global_layers: List[int], - *, - dense_layers: List[int], - swa_layers: List[int], - swa_window: Optional[int], - label: str, - ) -> Dict[str, object]: + 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, resolved once; every grouping derives from them - # (V2 owns the mapping; its own lookup errors are the precise ones). + # 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 ) - all_storage_groups: Dict[int, List[int]] = {} - for layer, pool_id in enumerate(layer_pool_ids): - all_storage_groups.setdefault(pool_id, []).append(layer) - pool_representatives = tuple(layers[0] for layers in all_storage_groups.values()) return dict( - global_layers=global_layers, layer_pools=layer_pools, dense_layers=dense_layers, swa_layers=swa_layers, swa_window=swa_window, layer_pool_ids=layer_pool_ids, - pool_representatives=pool_representatives, ) 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 index ca3b5c6107be..90905663022f 100644 --- 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 @@ -1,11 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""SM100 CuTe-DSL scorer for the TriAttention mean-score path (score-only and fused -score+stats+union entries); geometry outside the exact contract raises at construction.""" +"""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 @@ -1330,3 +1330,242 @@ def _tensor_spec(tensor: torch.Tensor) -> tuple: 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_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 51c82ce7fa8e..11657229e1c5 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Triton kernels and launch helpers for TriAttention -(fp32 math; int64 past-2^31 flat offsets; masked ragged tails; scoring in the CuTe pack).""" +"""Triton reduction, TP-fold, and selection kernels for TriAttention.""" from __future__ import annotations @@ -71,8 +70,7 @@ def _score_row_stats_kernel( # at def time (STD_EPSILON itself must stay a plain float for the CuTe import). EPSILON: tl.constexpr = STD_EPSILON, ): - """Compute one valid-window mean and inverse standard deviation per score row, - reading each row's decode window straight out of the score scratch.""" + """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 @@ -128,8 +126,7 @@ def _score_per_head_reduce_kernel( NORMALIZE: tl.constexpr, BLOCK: tl.constexpr = 256, ): - """Reduce each KV-head domain's decode window straight out of the score - scratch into one selector row.""" + """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) @@ -195,7 +192,7 @@ def _score_per_head_reduce_kernel( tl.store(selection_scores + output, reduced, mask=token < WIDTH) -def prepare_per_head_scores( +def reduce_per_head_scores( score_scratch: torch.Tensor, decode_lengths: torch.Tensor, prompt_lengths: torch.Tensor, @@ -205,18 +202,18 @@ def prepare_per_head_scores( selection_row_lengths: torch.Tensor, *, request_count: int, - num_layers: int, - num_q_heads: int, - num_kv_heads: int, padded_head_columns: int, score_token_capacity: int, - selection_width: int, per_layer: bool, normalize_scores: bool, ) -> None: - """Normalize and reduce the scratch's decode windows for either per-head - eviction mode.""" - selection_rows = num_layers * num_kv_heads if per_layer else num_kv_heads + """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: @@ -267,8 +264,7 @@ def _fold_union_ranks_kernel( WIDTH: tl.constexpr, BLOCK: tl.constexpr = 1024, ): - """Fold the TP-gathered rank-local union rows into the global union row - (elementwise max over the rank blocks).""" + """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) @@ -296,8 +292,7 @@ def _settle_ties_kernel( SELECTION_ROWS: tl.constexpr, BLOCK: tl.constexpr = 256, ): - """Settle one selection row's score ties into its final kept-token row - (deterministic lowest-index tie break, ascending output).""" + """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 diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 1cf42d70fc76..5eeae90e03b7 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -97,8 +97,7 @@ def make_ramp_pools( base=0, device=None, ): - """bf16 pools with a shifted ``arange % 251`` ramp: every wrong move - lands on a different byte pattern (supported geometry defaults).""" + """Build bf16 ramp pools whose moved regions have distinct byte patterns.""" return [ ( ( @@ -119,11 +118,7 @@ def make_ramp_pools( def build_compaction(**overrides): - """``build_compaction_params`` with the suite's 2-layer defaults: - allocates the caller-owned move-offset rows (capacity cumsum) and SWA - destination bases, and hands the test's pre-settled - ``kept_token_ordinals`` in as the decision rows. Returns the opaque - ``params`` plus a test-side mirror of the caller-owned inputs.""" + """Build opaque compaction parameters and caller-owned test inputs.""" from tensorrt_llm._torch.kv_cache_compression.compaction import build_compaction_params args = dict( @@ -214,9 +209,7 @@ def capacity_offsets(count): def run_compaction(compaction): - """Replica of the round's move stage in production order: SWA - destination rebase, then ``compact`` loops the opaque params (each packs - its decision rows into move sources and fires its native moves).""" + """Run SWA rebasing and opaque compaction in production order.""" from tensorrt_llm._torch.kv_cache_compression.compaction import compact if compaction["swa_destination_bases"] is not None: @@ -229,18 +222,16 @@ def run_compaction(compaction): def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): - """A bare manager carrying only the staging attributes, for the bulk - page-table copy tests (mirrors the product's ``_rebuild_eviction_runtime`` names).""" + """A bare manager carrying only the page-table staging attributes.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention staging = TriAttention.__new__(TriAttention) staging.kv_cache_manager = None staging.draft_kv_cache_manager = None staging._request_capacity = max_requests - staging._keep_count = 4 + staging.budget = 4 staging._swa_window = None - staging._draft_protected_tail_capacity = None - staging._block_offsets_ready_event = torch.cuda.Event() + 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( @@ -267,34 +258,6 @@ def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): ) -def make_buffer_stubs(manager, *, decode_width=260): - """Stub the calibration/layout surfaces around ``_ensure_eviction_runtime``. - - Returns the layout dict plus the manager attributes a stubbed - ``_rebuild_eviction_runtime`` should set (production sets them in place).""" - manager._freq_scale_sq = torch.ones(2) - manager._phase = {"rows": 8} - manager._omega = torch.ones(2) - manager._local_score_calibration = mock.Mock(return_value=(torch.ones(2, 2, 2),) * 3) - pool = torch.empty(8, 2, 1, 4, 4) - layout = dict( - global_layers=[0, 1], - layer_pools=[pool, pool], - dense_layers=[0, 1], - swa_layers=[], - swa_window=None, - layer_pool_ids=(0, 0), - ) - built_attributes = dict( - _selection_width_capacity=decode_width, - _score_token_capacity=1024, - _request_capacity=8, - _prompt_lengths_device=torch.zeros(8, dtype=torch.int32), - _decode_lengths_device=torch.empty(8, dtype=torch.int32), - ) - return layout, built_attributes - - 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 @@ -329,8 +292,7 @@ def make_fake_v2(enable_block_reuse=False, *, is_draft=False): def make_test_model_dir() -> str: - """A real on-disk dense model config: construction-time layer partition - resolves through the production AutoConfig path, no mocks.""" + """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_") @@ -364,8 +326,7 @@ def make_test_calibration_pt() -> str: def make_tri_config(**overrides): - """A real TriAttentionKvCacheCompressionConfig with test calibration inputs - (the config validator requires both ``model_path`` and ``calibration_path``).""" + """Build a real TriAttention config with test calibration inputs.""" from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig options = { @@ -378,10 +339,11 @@ def make_tri_config(**overrides): def make_triattention(**overrides): - """Construct a fully initialized manager for method-level unit tests.""" + """Construct a manager while isolating GPU-owned persistent state.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention - return TriAttention(make_tri_config(**overrides), make_fake_v2()) + with mock.patch.object(TriAttention, "_initialize_eviction_state"): + return TriAttention(make_tri_config(**overrides), make_fake_v2()) def make_eviction_input( @@ -396,8 +358,11 @@ def make_eviction_input( draft_cache=None, state=None, ): - """One due-cohort item shaped exactly like ``_evict_due_requests`` builds.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import _EvictionInput + """One due-request item shaped exactly like ``_evict_due_requests`` builds.""" + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + _EvictionInput, + _RequestState, + ) if request is None: request = SimpleNamespace(py_request_id=request_id, py_num_compressed_tokens=0) @@ -405,7 +370,7 @@ def make_eviction_input( request=request, target_cache=target_cache, draft_cache=draft_cache, - state={"generation_steps": 0, "evicted_tokens": 0} if state is None else state, + state=_RequestState() if state is None else state, source_length=int(source_length), logical_source_length=int( source_length if logical_source_length is None else logical_source_length @@ -436,11 +401,7 @@ def make_request(request_id, **overrides): @contextmanager def mocked_eviction_internals(manager): """Run the real ``_evict_due_requests`` transaction around a mocked round executor.""" - with ( - mock.patch.object(manager, "_runtime_kv_layout", return_value={}), - mock.patch.object(manager, "_ensure_eviction_runtime"), - mock.patch.object(manager, "_execute_eviction_round") as execute, - ): + with mock.patch.object(manager, "_execute_eviction_round") as execute: yield SimpleNamespace(execute=execute) @@ -457,8 +418,7 @@ def torch_tri_score_oracle( offsets, layer_indices, ): - """Independent Torch oracle of the paged mean score (GQA mapping via - ``head // group_size`` plus the position-independent MLR term).""" + """Compute paged mean scores independently with Torch.""" scores = [] num_q_heads = int(q_real.shape[1]) for request, seq_len in enumerate(seq_lens): @@ -500,22 +460,18 @@ def torch_tri_score_oracle( def make_phase_table(offsets, omega, initial_rows): - """Build the mean-phase table dict exactly like the product's inlined - form and grow it to cover positions ``[0, initial_rows)``.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - grow_mean_phase_table, + """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, + frequency_block=1 << (num_freqs - 1).bit_length(), ) - phase = { - "omega": omega.to(dtype=torch.float32).contiguous(), - "offset_values": offsets.tolist(), - "cos": None, - "sin": None, - "rows": 0, - } - grow_mean_phase_table(phase, max(int(initial_rows), 1)) - return phase - def make_cute_buffers( *, @@ -536,28 +492,20 @@ def make_cute_buffers( layer_pool_ids=None, normalize_scores=True, ): - """A bare manager with real eviction buffers built over the one-shared-slot - default layout; split reference legs use ``eviction_mode="per_head"`` over - the same pools. ``layer_pool_ids`` is the canonical per-layer V2 pool id - list and drives the page-table grouping.""" + """Build a bare manager with a real score pipeline over test pools.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention num_layers = len(layer_pools) assert int(q_real.shape[1]) == num_q_heads - # The build takes every capacity explicitly (no test-only None-derive - # path); the widest window defaults keep old call sites. if decode_width is None: decode_width = seq_len if layer_pool_ids is None: layer_pool_ids = [0] * num_layers - # A live-manager source table exactly wide enough for the requested - # capacity (the staged width clamps to it). 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( - global_layers=list(range(num_layers)), layer_pools=layer_pools, dense_layers=list(range(num_layers)), swa_layers=[], @@ -565,35 +513,45 @@ def make_cute_buffers( layer_pool_ids=layer_pool_ids, ) manager = TriAttention.__new__(TriAttention) - # The cold builder reads staging geometry from the owning manager. manager.kv_cache_manager = SimpleNamespace( num_pools=max(layer_pool_ids) + 1, 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 = None + manager._draft_protected_tail_capacity = 0 manager.eviction_mode = eviction_mode manager.normalize_scores = normalize_scores - manager._staging_reuse_event = None - manager._block_offsets_ready_event = None - manager._compaction_done_event = None + manager._request_capacity = max_requests + manager._selection_width_capacity = decode_width manager._phase = make_phase_table(offsets, omega, seq_len) - # Real owner state consumed by the cold builder (production sets these at - # construction / calibration load). manager.budget = keep_count manager._protected_tail_capacity = protected_tail_capacity manager._freq_scale_sq = freq_scale_sq - manager._calibration_q_real = q_real - manager._calibration_q_imag = q_imag - manager._calibration_mlr_coef = mlr_coef - manager._rebuild_eviction_runtime( - layout, - None, - request_capacity=max_requests, - score_token_capacity=seq_len, - selection_width_capacity=decode_width, + 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_eviction_capacity(score_token_capacity=seq_len) return manager @@ -604,8 +562,7 @@ def write_block_offsets(manager, encoded): def rect_to_score_scratch(scores, num_kv_heads, padded_head_columns=8): - """Scatter a [request, layer, q_head, token] rectangle into the fused - scorer's scratch layout (prompt starts at zero, bucket = rectangle width).""" + """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( @@ -622,8 +579,7 @@ def rect_to_score_scratch(scores, num_kv_heads, padded_head_columns=8): def stage_score_metadata(manager, request_count, source_lengths, decode_lengths, prompt_lengths): - """Stage the per-round score metadata exactly like production (the - compiled score launches read the staged rows via pointer capture).""" + """Stage per-round score metadata exactly as production does.""" torch.sub( source_lengths[:request_count], prompt_lengths[:request_count], @@ -636,48 +592,32 @@ def stage_score_metadata(manager, request_count, source_lengths, decode_lengths, def launch_split_scores( manager, request_count, source_lengths, decode_lengths, prompt_lengths, mean_cos, mean_sin ): - """The production score-only leg plus the decode-window gather (the round - executor's per-head sequence, parameterized by count). Test mean phases - load into the build-bound buffers, exactly like the production in-place - gather refresh; the compiled entry fires directly, like the round does.""" - import cuda.bindings.driver as cuda_driver - + """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]) - assert request_count in manager._compiled_score_by_request_count - stream = cuda_driver.CUstream( - torch.cuda.current_stream(manager._score_scratch.device).cuda_stream - ) - manager._compiled_score_by_request_count[request_count]( - *manager._cute_score_prefix, - manager._cute_mean_cos, - manager._cute_mean_sin, - *manager._cute_score_tail, - request_count, - stream, - ) + 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 = ( - manager._score_scratch[ - : manager._num_kv_heads * 8 * num_segments * manager._score_token_capacity - ] + score_scratch[: manager._num_kv_heads * 8 * num_segments * score_token_capacity] .view( manager._num_kv_heads, 8, request_count, manager._num_layers, - manager._score_token_capacity, + 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=manager._score_scratch.device, + device=score_scratch.device, ).view(1, 1, 1, 1, -1) - columns = columns.clamp_(max=manager._score_token_capacity - 1).expand( + columns = columns.clamp_(max=score_token_capacity - 1).expand( request_count, manager._num_layers, manager._num_kv_heads, @@ -693,7 +633,7 @@ def launch_split_scores( ), float("nan"), dtype=torch.float32, - device=manager._score_scratch.device, + device=score_scratch.device, ) torch.gather( source, 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 index b68e2be63b51..718909f509de 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py +++ b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py @@ -39,55 +39,11 @@ def test_plain_attention_defaults_to_fused_rope() -> None: def test_kv_cache_compression_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 - - -def test_unfused_yarn_rope_is_applied_exactly_once() -> None: - """With rope_fusion=False the Python-side rotary module owns RoPE, so the - backend must receive no position-embedding params. yarn is not listed in - PositionEmbeddingType.is_rope(), which used to leak the params through and - made the C++ QKV preprocess rotate a second time (double RoPE).""" - from tensorrt_llm._torch.attention_backend.interface import ( - PositionalEmbeddingParams, - RopeParams, - ) - from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType - - yarn_params = PositionalEmbeddingParams( - type=PositionEmbeddingType.yarn, - rope=RopeParams( - dim=32, - theta=150000, - scale_type=RotaryScalingType.yarn, - scale=32.0, - max_positions=1024, - original_max_positions=256, - beta_fast=32, - beta_slow=1, - duplicate_data=False, - ), - is_neox=True, - ) - model_config = ModelConfig(kv_cache_compression_config=TriAttentionKvCacheCompressionConfig( - model_path="/models/test", calibration_path="/calib/test.pt")) - attn = Attention( - hidden_size=256, - num_attention_heads=8, - num_key_value_heads=8, - max_position_embeddings=1024, - bias=False, - pos_embd_params=yarn_params, - layer_idx=0, - dtype=torch.bfloat16, - config=model_config, + 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 - assert attn.rotary_emb is not None - # The TRTLLM backend keeps the type as an int; 0 means no position - # embedding was handed to the kernel side. - assert attn.attn.position_embedding_type == 0 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 index 31c36ff7e556..58c8a7afb4fb 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_cute_score.py @@ -3,7 +3,8 @@ """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.""" +pure-PyTorch oracle; the contract test pins the no-fallback loud raise. +""" import pytest import torch @@ -161,8 +162,7 @@ def test_cute_kernel_matches_torch_oracle(case): list(range(num_layers)), ) - # Every count up to capacity is served, nothing beyond. - assert max_requests + 1 not in tri._compiled_score_by_request_count + # 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( 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 index fe2b64d24c0f..b49b731085de 100644 --- 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 @@ -4,7 +4,8 @@ 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.""" +unchanged from the retired Triton reference copies. +""" import pytest import torch @@ -29,35 +30,11 @@ def _run_fused_union( mean_sin, union_out, ): - """The fused score+stats+normalized-union pipeline (THE union path), fired - directly off the compiled entries exactly like the round. Test mean phases - load into the build-bound buffers; the fused rows land straight in the - build-bound ``tri._selection_scores_rows``.""" - import cuda.bindings.driver as cuda_driver - + """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]) - assert ( - request_count in tri._compiled_score_by_request_count - and request_count in tri._compiled_normalize_union_by_request_count - ) - stream = cuda_driver.CUstream(torch.cuda.current_stream(tri._score_scratch.device).cuda_stream) - tri._compiled_score_by_request_count[request_count]( - *tri._cute_score_prefix, - tri._cute_mean_cos, - tri._cute_mean_sin, - *tri._cute_score_tail, - request_count, - stream, - ) - tri._compiled_normalize_union_by_request_count[request_count]( - tri._cute_partial_stats, - *tri._cute_selection_prefix, - tri._cute_selection_scores_rows, - request_count, - stream, - ) + 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]) @@ -65,8 +42,7 @@ def _run_fused_union( def _reference_union_scores( scores_rows: torch.Tensor, decode_lengths: torch.Tensor ) -> torch.Tensor: - """Union oracle: per-row mean/biased-std z-norm over the valid prefix - (std clamped at 1e-6), union-max across rows, ``-inf`` past the width.""" + """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 @@ -100,7 +76,7 @@ def _reference_union_scores( # only the real heads' rows, and the union finalizer maps head rows # onto the padded score planes. (128, 32, 4, 0, None), - # Mixed-prompt cohort (one start mid-tile, one page-aligned) — the + # 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]), @@ -113,8 +89,7 @@ def test_union_fusion_matches_split_pipeline( score_starts: "int | list", valid_lens: "list | None", ) -> None: - """Fused rows must reproduce split score->normalize->union rows; - ``score_starts`` is uniform or per-request (read at runtime).""" + """Check fused union rows against the split score-normalize-union path.""" pytest.importorskip("cutlass") torch.manual_seed(20260721) @@ -213,8 +188,7 @@ def test_union_fusion_matches_split_pipeline( @_SM100_ONLY def test_union_fusion_frequency_count_guard_raises() -> None: - """16 frequencies (head size 32) sit outside the fused kernel contract - and are rejected at kernel construction.""" + """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 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 index d5f4866e449d..5a2446f50789 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -1,236 +1,75 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Draft KV co-compression: the target's union keep set broadcasts over -the draft's own KV heads, the draft's tail appends as ordinals, both land -at ``destination_base = prompt_len``. Covers the physical moves, packed -indices, stream ordering, admission gates, the published compressed-token -invariant, and buffer reuse/rebuild.""" +"""TriAttention draft lifecycle, ordering, admission, and publication.""" from types import SimpleNamespace from unittest import mock import pytest import torch -from conftest import build_compaction as _build_compaction -from conftest import encode_block_offsets as _encode_block_offsets -from conftest import make_buffer_stubs as _make_buffer_stubs from conftest import make_eviction_input as _make_eviction_input from conftest import make_fake_v2 as _make_fake_v2 -from conftest import make_ramp_pools as _make_ramp_pools 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 conftest import run_compaction as _run_compaction -from conftest import set_protected_tails as _set_protected_tails - -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention - - -def _fresh_request_state(): - """One request's compression ledger, as the manager initializes it.""" - return {"generation_steps": 0, "evicted_tokens": 0} - - -def _logical_view(pool: torch.Tensor, pages: torch.Tensor) -> torch.Tensor: - """Gather one request's pages into [K/V, head, token, dim] order.""" - num_kv_heads = int(pool.shape[2]) - head_dim = int(pool.shape[4]) - return pool.index_select(0, pages).permute(1, 2, 0, 3, 4).reshape(2, num_kv_heads, -1, head_dim) - - -def _launched_draft_compaction(draft_protected_tails): - """Target and draft pools with distinct head counts (supported bf16 - geometry, mod-251 ramp payload), compacted in one round.""" - device = torch.device("cuda", torch.cuda.current_device()) - request_count = 2 - prompt_len = 2 - target_protected_tails = [2, 1] - valid_seq_lens = [10, 9] - - target_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - draft_tables = torch.tensor([[1, 0, 2], [5, 4, 3]], dtype=torch.int32, device=device) - target_pools = _make_ramp_pools(2, num_kv_heads=2, device=device) - draft_pool = _make_ramp_pools(1, num_kv_heads=4, base=149, device=device)[0] - assert target_pools[0].shape[2] != draft_pool.shape[2] - initial_target = [pool.clone() for pool in target_pools] - initial_draft = draft_pool.clone() - - keep = torch.tensor([[2, 4, 7, 9], [3, 5, 6, 8]], dtype=torch.int64, device=device) - - compaction = _build_compaction( - layer_pools=target_pools, - layer_pool_ids=[0, 0], - kept_token_ordinals=keep.to(torch.int32), - valid_sequence_lengths=torch.tensor(valid_seq_lens, dtype=torch.int32, device=device), - kv_block_offsets=_encode_block_offsets(target_tables), - prompt_offsets=torch.full((request_count,), prompt_len, dtype=torch.int32, device=device), - protected_tail_capacity=max(target_protected_tails), - draft_layer_pools=[draft_pool], - draft_layers=[0], - draft_layer_pool_ids=[0], - draft_protected_tail_capacity=max(draft_protected_tails), - draft_kv_block_offsets=_encode_block_offsets(draft_tables), - ) - _set_protected_tails(compaction, target_protected_tails, draft_protected_tails) - _run_compaction(compaction) - torch.cuda.synchronize(device) - - return SimpleNamespace( - device=device, - request_count=request_count, - prompt_len=prompt_len, - keep=keep, - valid_seq_lens=valid_seq_lens, - target_protected_tails=target_protected_tails, - draft_protected_tails=draft_protected_tails, - target_tables=target_tables, - draft_tables=draft_tables, - target_pools=target_pools, - draft_pool=draft_pool, - initial_target=initial_target, - initial_draft=initial_draft, - compaction=compaction, - ) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttention, + _RequestState, +) -def test_draft_moves_and_pack_match_keep_broadcast_and_tail_oracle(): - # Ragged draft tails [1, 2] against target tails [2, 1]: one request's - # draft tail below and one above its target, subsuming the uniform row. - built = _launched_draft_compaction(draft_protected_tails=[1, 2]) - device = built.device - prompt_len = built.prompt_len - - expected_offsets = [0] - for request in range(built.request_count): - valid = built.valid_seq_lens[request] - # Target dense layers compact the union keep set plus the target tail. - target_pages = built.target_tables[request].to(torch.long) - target_tail = torch.arange( - valid, - valid + built.target_protected_tails[request], - dtype=torch.int64, - device=device, - ) - target_source = torch.cat((built.keep[request], target_tail)) - target_destination = torch.arange( - prompt_len, - prompt_len + target_source.numel(), - dtype=torch.int64, - device=device, - ) - for before_pool, after_pool in zip(built.initial_target, built.target_pools): - before = _logical_view(before_pool, target_pages) - after = _logical_view(after_pool, target_pages) - assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) - assert torch.equal( - after.index_select(2, target_destination), - before.index_select(2, target_source), - ) - - # Same kept ordinals through the draft's OWN table/heads/tail. - draft_pages = built.draft_tables[request].to(torch.long) - draft_tail = torch.arange( - valid, - valid + built.draft_protected_tails[request], - dtype=torch.int64, - device=device, - ) - draft_source = torch.cat((built.keep[request], draft_tail)) - draft_destination = torch.arange( - prompt_len, - prompt_len + draft_source.numel(), - dtype=torch.int64, - device=device, - ) - before = _logical_view(built.initial_draft, draft_pages) - after = _logical_view(built.draft_pool, draft_pages) - assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) - for head in range(int(built.draft_pool.shape[2])): - assert torch.equal( - after[:, head].index_select(1, draft_destination), - before[:, head].index_select(1, draft_source), - ) - - expected_offsets.append(expected_offsets[-1] + int(draft_source.numel())) - - # The test-owned draft move-offset row must match the broadcast-plus-tail - # oracle; the packed move sources themselves are covered byte-exactly by - # the pool assertions above (the ramp payload makes every wrong move land - # on different bytes) and by the pack-kernel oracle suite. - assert built.compaction["draft_move_offsets"].cpu().tolist() == expected_offsets - - -def test_execute_eviction_round_orders_both_manager_streams(): - """The round executor snapshots both page-table planes, then records one - completion event and BOTH cache-manager streams wait on it -- even when - the round body fails -- so neither manager can free or reallocate pages - this cohort is still reading.""" - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module +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 = TriAttention.__new__(TriAttention) tri._request_capacity = 8 - tri._keep_count = 4 - tri.eviction_mode = "union" + tri.budget = 4 tri._swa_window = None - tri._compaction_params = () 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() - tri._request_metadata_device = torch.zeros_like(host) - tri._phase = {"cos": None, "sin": None, "rows": 8} - tri._phase_num_freqs = 1 - tri._phase_f_block = 1 - tri._logical_source_lengths_device = None - tri._source_lengths_device = None - tri._prompt_lengths_device = None - tri._decode_lengths_device = None - tri._mean_cos = None - tri._mean_sin = None - tri._swa_destination_bases = None - tri._swa_rebase_delta = 0 + 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 - target_stream = mock.Mock() - draft_stream = mock.Mock() - manager = SimpleNamespace(_stream=target_stream) + execution_stream = mock.Mock() + manager = SimpleNamespace(_stream=execution_stream) draft_manager = SimpleNamespace( - _stream=draft_stream, num_extra_kv_tokens=0, _kv_reserve_draft_tokens=0 + _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 = SimpleNamespace() + compute_stream = mock.Mock() eviction_inputs = [_make_eviction_input(request_id=7, source_length=8)] class Boom(RuntimeError): pass - score_kernel = mock.MagicMock() - score_kernel.__getitem__.return_value.side_effect = Boom + metadata_device.copy_.side_effect = Boom with ( - mock.patch.object(torch.cuda, "current_stream", return_value=compute_stream), - mock.patch.object(module, "grow_mean_phase_table"), + mock.patch.object( + torch.cuda, "current_stream", return_value=compute_stream + ) as current_stream, mock.patch.object(tri, "_stage_block_offsets") as stage, - mock.patch.object(module, "_gather_mean_phase_kernel", score_kernel), - mock.patch.object(module, "compact") as compact, ): with pytest.raises(Boom): tri._execute_eviction_round(eviction_inputs) # Both page-table planes were snapshotted before the round body fired. assert stage.call_count == 2 - compact.assert_not_called() - # One event records the round; BOTH cache managers wait on it. + # 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) - target_stream.wait_event.assert_called_once_with(event) - draft_stream.wait_event.assert_called_once_with(event) + execution_stream.wait_event.assert_called_once_with(event) @pytest.mark.parametrize( @@ -270,7 +109,6 @@ def test_draft_admission_gates_raise(gate, match): def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager = _make_triattention(budget=4, beta=4) - manager._layer_partition = ([0, 1], [], None) target = manager.kv_cache_manager target._stream = mock.Mock() target.pp_layers = [0, 1] @@ -290,7 +128,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager._draft_protected_tail_capacity = 1 request = _make_request(7, py_prompt_len=2, py_num_accepted_draft_tokens=1) - manager._request_states[7] = _fresh_request_state() + manager._request_states[7] = _RequestState() batch = SimpleNamespace(generation_requests=[request]) # Every step confirms one sampled token plus one accepted draft token. @@ -309,11 +147,11 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager._evict_due_requests(batch) state = manager._request_states[7] - if state["evicted_tokens"] > previous_evicted: + if state.evicted_tokens > previous_evicted: # An eviction round compacted the cache to prompt + budget. eviction_rounds += 1 - confirmed -= state["evicted_tokens"] - previous_evicted - previous_evicted = state["evicted_tokens"] + confirmed -= state.evicted_tokens - previous_evicted + previous_evicted = state.evicted_tokens cache.capacity = confirmed assert confirmed == 2 + 4 # The staged logical position restores the uncompressed @@ -330,8 +168,8 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): 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 manager itself, which carries both cache - # managers whose streams it orders after the compact launches. + # 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 @@ -341,118 +179,58 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): assert call.kwargs == {} -def test_cohort_growth_rebuilds_buffers_and_drops_cached_compaction(): - manager = _make_triattention(budget=4) - layout, built_attributes = _make_buffer_stubs(manager) - draft_layout = dict( - layer_pools=[], - dense_layers=[], - pool_representatives=(), - layer_pool_ids=(), - pool_page_counts=(4,), - ) - manager.draft_kv_cache_manager = _make_fake_v2(is_draft=True) - # Injected post-construction: mirror the ctor-cached manager-lifetime tail. - manager._draft_protected_tail_capacity = 1 - eviction_inputs = [_make_eviction_input(_make_request(7), request_id=7, source_length=8)] - - def apply_built(*args, **kwargs): - for name, value in built_attributes.items(): - setattr(manager, name, value) - - phase = manager._phase - with mock.patch.object( - manager, "_rebuild_eviction_runtime", side_effect=apply_built - ) as prepare: - manager._ensure_eviction_runtime(layout, draft_layout, eviction_inputs) - - # Request capacity follows the executor limits, while the score - # bucket follows what the cohort actually presents (power-of-two, - # 1024 floor) instead of pinning tens-of-GiB scratch to max_seq_len. - # The stubbed build's capacities became the resident manager state. - assert manager._buffers_built - assert manager._selection_width_capacity == built_attributes["_selection_width_capacity"] - # The mode and the shared phase-table dict live on the manager itself - # and thread through unchanged (no longer build arguments). - assert manager.eviction_mode == "union" - assert manager._phase is phase - args = prepare.call_args.args - kwargs = prepare.call_args.kwargs - assert args[0] is layout - assert args[1] is draft_layout - assert kwargs["request_capacity"] == 8 - assert kwargs["selection_width_capacity"] == 4 + 2 * 128 - assert kwargs["score_token_capacity"] == 1024 - - # A second round within the resident capacities reuses the buffers - # (and with them the compaction launch data they carry). - manager._ensure_eviction_runtime(layout, draft_layout, eviction_inputs) - assert prepare.call_count == 1 - - # A cohort that outgrows the resident capacities rebuilds the whole - # buffer state, compaction included. - grown = [ - _make_eviction_input( - _make_request(7), - request_id=7, - source_length=8 + built_attributes["_selection_width_capacity"], - ) - ] - - def apply_rebuilt(*args, **kwargs): - apply_built() - manager._selection_width_capacity = built_attributes["_selection_width_capacity"] + 8 - - prepare.side_effect = apply_rebuilt - manager._ensure_eviction_runtime(layout, draft_layout, grown) - assert prepare.call_count == 2 - assert manager._buffers_built - assert ( - manager._selection_width_capacity == built_attributes["_selection_width_capacity"] + 8 - ) - - -def test_source_growth_beyond_score_bucket_rebuilds_buffers(): - """A later cohort can grow max(source_length) past the compiled score - bucket while decode width, page tokens, and request count all still fit; - the reuse gate must rebuild instead of scoring past the static geometry.""" - manager = _make_triattention(budget=4) - layout, built_attributes = _make_buffer_stubs(manager) - eviction_inputs = [ - _make_eviction_input( - _make_request(7), - request_id=7, - source_length=1024, - prompt_length=1020, - ) +def test_request_admission_reserves_score_high_watermark(): + manager = _make_triattention(budget=128, beta=64) + 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_eviction_capacity = 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), ] - def apply_built(*args, **kwargs): - for name, value in built_attributes.items(): - setattr(manager, name, value) - - with mock.patch.object( - manager, "_rebuild_eviction_runtime", side_effect=apply_built - ) as prepare: - manager._ensure_eviction_runtime(layout, None, eviction_inputs) - assert prepare.call_count == 1 - # One more source token, same decode width and request count. - grown = [ - _make_eviction_input( - _make_request(7), - request_id=7, - source_length=1025, - prompt_length=1021, - ) - ] - manager._ensure_eviction_runtime(layout, None, grown) - assert prepare.call_count == 2 + for request in requests: + manager._reserve_eviction_capacity(request) + + assert manager._build_eviction_capacity.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_eviction_capacity = mock.Mock() + + manager._reserve_eviction_capacity(_make_request(1, py_prompt_len=1025, py_max_new_tokens=192)) + + manager._build_eviction_capacity.assert_called_once_with(score_token_capacity=1280) + manager._compaction_done_event.synchronize.assert_not_called() def test_staged_block_width_clamps_to_manager_source_width(): - """Score tile rounding can request more page-table blocks than the live V2 - source table holds; the staged width must clamp to the manager width so the - native gather never reads past the K plane (tpb=32, max_seq_len=96, tail=1).""" + """Clamp staged block width to the live V2 source-table width.""" from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( _allocate_block_offset_staging, ) 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 index 144691087af5..c5b6ec4c2c7e 100644 --- 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 @@ -1,17 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""The settle and pack kernels vs pure-torch integer oracles: -``_settle_ties_kernel`` (threshold recovery with sentinel-skip, -strictly-greater count, lowest-index tie quota, ascending prompt-rebased -emission) and ``_pack_move_sources_kernel`` (dense/SWA packing, fed -pre-settled rows). All outputs are integers, so comparisons are -``torch.equal`` including stale regions.""" +"""The TriAttention settle kernel vs a pure-torch integer oracle.""" import pytest import torch -from tensorrt_llm._torch.kv_cache_compression.compaction import _pack_move_sources_kernel from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( _settle_ties_kernel, ) @@ -25,23 +19,9 @@ (350, 300), ] -# Pack rows: one per distinct kernel path combo (BROADCAST = union; PER_LAYER -# only branches inside the non-broadcast SWA store, so it is compiled out -# without SWA); keep 5 vs 300 flips the MOVE_CAPACITY block trip count on -# both BROADCAST sides. -_PACK_PATH_ROWS = [ - ("union", False, 21, 5), - ("union", True, 350, 300), - ("per_head", False, 350, 300), - ("per_head", True, 21, 5), - ("per_layer_perhead", True, 350, 300), -] - def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, keep_count): - """Settle in place: threshold = min over non-sentinel provisional - lanes; keep strictly-greater, fill quota with lowest-index ties, emit - ascending rebased by prompt; entries past the emitted count stay.""" + """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] @@ -59,49 +39,6 @@ def _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, output, ) -def _pack_oracle( - settled, - source_lengths, - dense_offsets, - dense_out, - swa_offsets, - swa_out, - *, - selection_rows, - keep_count, - num_kv_heads, - swa_window, - union, - per_layer, - has_swa, -): - """Pack in place: dense rows forward settled content verbatim (stale - included) then append the tail ``seq_len + move - keep_count``; SWA - rows write latest-window ordinals once per KV head.""" - request_count = int(source_lengths.shape[0]) - packed_rows = int(dense_out.shape[0]) - dense_total = int(dense_out.shape[1]) - swa_total = int(swa_out.shape[1]) - for request in range(request_count): - seq_len = int(source_lengths[request]) - dense_begin = int(dense_offsets[request]) - dense_count = int(dense_offsets[request + 1]) - dense_begin - for domain in range(packed_rows): - selection_domain = 0 if union else domain - settled_row = settled[request * selection_rows + selection_domain] - for move in range(dense_count): - value = int(settled_row[move]) if move < keep_count else seq_len + move - keep_count - dense_out.view(-1)[domain * dense_total + dense_begin + move] = value - if has_swa and (domain < num_kv_heads if per_layer else True): - swa_begin = int(swa_offsets[request]) - swa_count = int(swa_offsets[request + 1]) - swa_begin - head = domain % num_kv_heads - for move in range(swa_count): - swa_out.view(-1)[head * swa_total + swa_begin + move] = ( - seq_len - swa_window + move - ) - - def _selection_rows_for(eviction_mode: str, num_layers: int, num_kv_heads: int) -> int: if eviction_mode == "union": return 1 @@ -110,16 +47,8 @@ def _selection_rows_for(eviction_mode: str, num_layers: int, num_kv_heads: int) return num_layers * num_kv_heads -def _staged_offsets(counts, device): - offsets = [0] - for count in counts: - offsets.append(offsets[-1] + count) - return torch.tensor(offsets, dtype=torch.int32, device=device) - - def _make_settle_inputs(request_count, selection_rows, width, keep_count, seed, device): - """One seeded settle problem: heavily tied scores, ragged rows, a - top-k stand-in, and per-request prompt rebases.""" + """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. @@ -187,94 +116,6 @@ def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): assert torch.equal(output_actual, output_reference), f"kept ordinals differ (seed {seed})" -@pytest.mark.parametrize("eviction_mode,has_swa,width,keep_count", _PACK_PATH_ROWS) -def test_pack_matches_torch_oracle_on_pre_settled_rows(eviction_mode, has_swa, width, keep_count): - device = torch.device("cuda", torch.cuda.current_device()) - request_count, num_layers, num_kv_heads = 3, 2, 2 - union = eviction_mode == "union" - per_layer = eviction_mode == "per_layer_perhead" - selection_rows = _selection_rows_for(eviction_mode, num_layers, num_kv_heads) - rows_total = request_count * selection_rows - packed_rows = num_layers * num_kv_heads if per_layer else num_kv_heads - - # Move geometry: the last request is a padded row that moves nothing. - protected_tails = [2, 0, 0] - tail_capacity = max(protected_tails) - dense_counts = [keep_count + protected_tails[0], keep_count + protected_tails[1], 0] - swa_window = 6 - swa_counts = [swa_window + protected_tails[0], swa_window + protected_tails[1], 0] - move_capacity = keep_count + tail_capacity - if has_swa: - move_capacity = max(move_capacity, swa_window + tail_capacity) - dense_total = request_count * (keep_count + tail_capacity) - swa_total = request_count * (swa_window + tail_capacity) - source_lengths = torch.tensor([10, 8, 0], dtype=torch.int32, device=device) - dense_offsets = _staged_offsets(dense_counts, device) - swa_offsets = _staged_offsets(swa_counts, device) - - 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) - # Pre-settled decision rows straight from the settle oracle: short - # rows keep stale garbage past their emitted count, which the pack - # must forward verbatim. - settled = torch.randint( - -(2**30), 2**30, (rows_total, keep_count), dtype=torch.int32, device=device - ) - _settle_oracle(scores, row_lengths, row_prompt_offsets, provisional, settled, keep_count) - - # Identical stale garbage on both sides so untouched regions must - # match too. - dense_stale = torch.randint( - -(2**30), 2**30, (packed_rows, dense_total), dtype=torch.int32, device=device - ) - swa_stale = torch.randint( - -(2**30), 2**30, (num_kv_heads, swa_total), dtype=torch.int32, device=device - ) - dense_reference = dense_stale.clone() - swa_reference = swa_stale.clone() - _pack_oracle( - settled, - source_lengths, - dense_offsets, - dense_reference, - swa_offsets if has_swa else dense_offsets, - swa_reference if has_swa else dense_reference, - selection_rows=selection_rows, - keep_count=keep_count, - num_kv_heads=num_kv_heads, - swa_window=swa_window if has_swa else 0, - union=union, - per_layer=per_layer, - has_swa=has_swa, - ) - - dense_actual = dense_stale.clone() - swa_actual = swa_stale.clone() - _pack_move_sources_kernel[(request_count, selection_rows)]( - settled, - source_lengths, - dense_offsets, - dense_actual, - swa_offsets if has_swa else None, - swa_actual if has_swa else None, - KEEP_COUNT=keep_count, - DECISION_ROWS=selection_rows, - MOVE_CAPACITY=move_capacity, - NUM_KV_HEADS=num_kv_heads, - PER_LAYER=per_layer, - DENSE_TOTAL=dense_total, - SWA_TOTAL=swa_total if has_swa else 0, - SWA_WINDOW=swa_window if has_swa else 0, - ) - torch.cuda.synchronize(device) - - assert torch.equal(dense_actual, dense_reference), f"dense moves differ (seed {seed})" - assert torch.equal(swa_actual, swa_reference), f"SWA moves differ (seed {seed})" - - def test_settle_handles_topk_sentinel_padding(): """Rows shorter than KEEP_COUNT arrive -1-padded and must settle inertly. diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 9e3f97a5a1d4..7e57faf5ab04 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -15,10 +15,11 @@ """Unit tests for the TriAttention compression-manager pipeline. -Config, construction, eviction lifecycle, page-table staging, and the fixed -score buffers; the manager publishes evicted counts via +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``.""" +``test_triattention_draft_cocompaction.py``. +""" from types import SimpleNamespace from unittest import mock @@ -26,38 +27,30 @@ import pytest import torch from conftest import make_bare_staging as _make_bare_staging -from conftest import make_cute_buffers as _make_cute_buffers -from conftest import make_eviction_input as _make_eviction_input 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 -from conftest import torch_tri_score_oracle as _torch_tri_score_oracle # 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 TriAttention +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttention, + _RequestState, +) # 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 -# The SM100 CuTe kernel is the only score path, so every test that actually -# launches scores (or builds the real staging buffers, whose constructor -# compiles the kernel) is 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 _set_request_state(manager, request_id, *, generation_steps=0, evicted_tokens=0): - state = { - "generation_steps": generation_steps, - "evicted_tokens": evicted_tokens, - } +def _set_request_state(manager, request_id, *, confirmed_tokens=0, evicted_tokens=0): + state = _RequestState( + confirmed_tokens=confirmed_tokens, + evicted_tokens=evicted_tokens, + ) manager._request_states[request_id] = state return state @@ -84,44 +77,66 @@ def _make_hf_config(**values): class TestConfigAndFactory: def test_factory_returns_triattention_and_propagates_config_fields(self): - # Calibration is deferred to the first request, so construction needs - # no calibration file or CUDA. + # The factory contract is independent of GPU-owned persistent buffers. fake_v2 = _make_fake_v2(enable_block_reuse=False) cfg = _make_tri_config(budget=32, beta=16, eviction_mode="per_head") - mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) + with mock.patch.object(TriAttention, "_initialize_eviction_state") as initialize: + mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) assert mgr.budget == 32 assert mgr.beta == 16 assert mgr.eviction_mode == "per_head" assert mgr.kv_cache_manager is fake_v2 + initialize.assert_called_once_with() class TestTriAttentionClass: def test_request_init_and_finish_lifecycle(self): - # Init: speculative capacity accepted, manager marked, state tracked. - # Finish: state cleared; buffers and the step's batch stay resident. + # Init reserves request-dependent capacity and tracks state. Finish + # clears only request state; persistent runtime objects stay resident. manager = _make_fake_v2() manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 - triattention = TriAttention(_make_tri_config(budget=8), manager) - - triattention.on_request_init(_make_request(11)) - triattention.on_request_init(_make_request(12)) + with mock.patch.object(TriAttention, "_initialize_eviction_state"): + triattention = TriAttention(_make_tri_config(budget=8), manager) + with ( + mock.patch.object(triattention, "_validate_request_capacity") as validate, + mock.patch.object(triattention, "_reserve_eviction_capacity") as reserve, + ): + request_11 = _make_request(11) + request_12 = _make_request(12) + triattention.on_request_init(request_11) + triattention.on_request_init(request_12) assert triattention.adjusts_generation_kv_length is True assert manager.kv_compression_manages_history assert set(triattention._request_states) == {11, 12} + assert validate.call_args_list == [mock.call(request_11), mock.call(request_12)] + assert reserve.call_args_list == [mock.call(request_11), mock.call(request_12)] - buffers = object() - triattention._buffers_built = True - triattention._score_scratch = buffers + phase = object() + triattention._phase = phase batch = SimpleNamespace() triattention._inflight_scheduled_batch = batch triattention.on_request_finish(_make_request(11)) triattention.on_request_finish(_make_request(12)) assert triattention._request_states == {} assert triattention._inflight_scheduled_batch is batch - assert triattention._buffers_built and triattention._score_scratch is buffers + assert triattention._phase is phase + + def test_capacity_guard_uses_maximum_steady_eviction_peak(self): + manager = _make_triattention(budget=10, beta=8) + manager.kv_cache_manager.get_num_available_tokens = mock.Mock(return_value=17) + + with pytest.raises(ValueError, match="requires 19 tokens"): + manager._validate_request_capacity( + _make_request(7, py_prompt_len=0, py_max_new_tokens=100) + ) + + manager.kv_cache_manager.get_num_available_tokens.assert_called_once_with( + token_num_upper_bound=18, + max_num_draft_tokens=1, + ) def test_resolve_accepts_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() @@ -231,7 +246,7 @@ class TestCompressedTokenPublication: # test_triattention_draft_cocompaction.py. def test_identity_selection_is_filtered_before_launch(self): - # Identity cohorts (seq_len == prompt + budget) are the pre-launch + # Identity requests (seq_len == prompt + budget) are the pre-launch # owner no-op: nothing launches and nothing is published. manager = _make_triattention(budget=4) manager.kv_cache_manager._stream = mock.Mock() @@ -241,14 +256,14 @@ def test_identity_selection_is_filtered_before_launch(self): capacity=6, history_length=0, is_active=True, resize=mock.Mock(return_value=True) ) manager.kv_cache_manager.kv_cache_map = {7: cache} - state = _set_request_state(manager, 7, generation_steps=127) + state = _set_request_state(manager, 7, confirmed_tokens=127) 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 - assert state["evicted_tokens"] == 0 + assert state.evicted_tokens == 0 cache.resize.assert_not_called() @@ -281,7 +296,8 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft 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 - mgr = TriAttention(_make_tri_config(budget=8), fake_v2) + with mock.patch.object(TriAttention, "_initialize_eviction_state"): + mgr = TriAttention(_make_tri_config(budget=8), fake_v2) cache = SimpleNamespace( capacity=seq_len, history_length=1024, @@ -297,7 +313,7 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft _kv_reserve_draft_tokens=kv_reserve_draft_tokens, ) mgr._request_states = {} - _set_request_state(mgr, 7, generation_steps=127) + _set_request_state(mgr, 7, confirmed_tokens=127) mgr.beta = 128 mgr.budget = 4096 return mgr, request, batch @@ -305,12 +321,12 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft def test_suspended_cache_defers_that_request_pre_launch(self): # A suspended cache is a legal overlap-scheduler transient: that # request defers (pre-launch, no cadence mutation) while the rest of - # the cohort proceeds. + # the request group proceeds. manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) second_request = _make_request(8, py_prompt_len=1024) manager.kv_cache_manager.kv_cache_map[8] = SimpleNamespace(is_active=False) first_state = manager._request_states[7] - second_state = _set_request_state(manager, 8, generation_steps=127) + second_state = _set_request_state(manager, 8, confirmed_tokens=127) batch = SimpleNamespace(generation_requests=[first_request, second_request]) with _mocked_eviction_internals(manager) as internals: @@ -319,8 +335,8 @@ def test_suspended_cache_defers_that_request_pre_launch(self): # Only the active request launched; the suspended one deferred whole. eviction_inputs = internals.execute.call_args.args[0] assert [item.request.py_request_id for item in eviction_inputs] == [7] - assert first_state["generation_steps"] == 128 - assert second_state["generation_steps"] == 127 + assert first_state.confirmed_tokens == 128 + assert second_state.confirmed_tokens == 127 # ``accepted`` enters the prepared item linearly; the zero and maximal # boundary rows pin the whole family. @@ -371,7 +387,7 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): # ledger (capacity minus the protected tail), never the logical length. physical_confirmed = 6100 manager = _make_triattention(beta=128) - _set_request_state(manager, 7, generation_steps=127, evicted_tokens=100) + _set_request_state(manager, 7, confirmed_tokens=127, evicted_tokens=100) cache = SimpleNamespace( capacity=physical_confirmed, history_length=1024, @@ -401,11 +417,12 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): # call-site gate accepts the one-model MTP roundtrip (base-ctor # marking is asserted in the executor manager tests). draft_manager = _make_fake_v2(is_draft=True) - TriAttention( - _make_tri_config(budget=8), - _make_fake_v2(), - draft_kv_cache_manager=draft_manager, - ) + with mock.patch.object(TriAttention, "_initialize_eviction_state"): + TriAttention( + _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_with_spec from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig @@ -433,7 +450,8 @@ def test_prepare_snapshots_fixed_linear_generation_growth( manager.kv_cache_map = { 7: SimpleNamespace(capacity=106, is_active=True), } - triattention = TriAttention(_make_tri_config(budget=8), manager) + with mock.patch.object(TriAttention, "_initialize_eviction_state"): + triattention = TriAttention(_make_tri_config(budget=8), manager) batch = SimpleNamespace( context_requests=[], generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], @@ -455,37 +473,10 @@ def test_union_forces_normalized_scores(self): triattention = _make_triattention(budget=4, eviction_mode="union", normalize_scores=False) assert triattention.normalize_scores is True - def test_execute_rejects_int32_overflowing_logical_source_lengths(self): - # Round starts past the int32 metadata range fail loudly (in the host - # metadata build) before any GPU work is enqueued. - device = torch.device("cuda", torch.cuda.current_device()) - staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) - gather = mock.Mock() - manager = _make_staging_manager( - torch.zeros(1, 2, 2, 12, dtype=torch.int32), gather, torch.cuda.Stream(device=device) - ) - staging.kv_cache_manager = manager - eviction_inputs = [ - _make_eviction_input(request_id=7, source_length=64, logical_source_length=2**31) - ] - - with pytest.raises((RuntimeError, OverflowError, ValueError)): - staging._execute_eviction_round(eviction_inputs) - assert gather.call_count == 0 - - def test_bulk_page_table_copy_snapshots_and_orders_consumers(self): - """The bulk copy stages immutable host snapshots, and the next copy - waits for the previous cohort's consumers. - - Both persistent V2 host inputs (the block-offset table and the - index-mapper slot assignment) may mutate as soon as staging returns; - the staged device tables must reflect the values at staging time. A - subsequent bulk copy must also wait until the previous round's - consumers (ordered by the round executor's completion event) are done. - """ + 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) - manager_stream = torch.cuda.Stream(device=device) host_table = torch.zeros( 1, 2, @@ -507,23 +498,20 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): gather = mock.Mock(side_effect=gather_k_block_offsets) staging = _make_bare_staging(device, max_requests=1, staged_blocks_per_seq=8) - staging._staging_reuse_event.record(current_stream) - manager = _make_staging_manager(host_table, gather, manager_stream) + manager = _make_staging_manager(host_table, gather, current_stream) def stage_once(): # Raises on any staging failure; success returns None. - staging._stage_block_offsets( - manager, - [7], - staging._block_offsets_host, - staging._block_offsets_device, - ) + with torch.cuda.stream(current_stream): + staging._stage_block_offsets( + manager, + [7], + staging._block_offsets_host, + staging._block_offsets_device, + ) # Round 1: mutate the host table and the slot assignment right after - # staging, with the manager stream artificially delayed. The staged - # result must still reflect the values at staging time. - with torch.cuda.stream(manager_stream): - torch.cuda._sleep(50_000_000) + # staging. The staged result must still reflect the gathered snapshot. with mock.patch.object( torch, "index_select", @@ -538,11 +526,9 @@ def stage_once(): 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 cohort. + # 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 - with torch.cuda.stream(manager_stream): - torch.cuda._sleep(50_000_000) stage_once() host_table[0, 0, 0, :5] = torch.tensor([23, 24, 25, 26, 27], dtype=torch.int32) selected_slot[0] = 1 @@ -551,18 +537,10 @@ def stage_once(): 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 delayed consumer read (snapshot) queued before the - # round's completion ordering must complete before the next bulk - # copy overwrites the device tables. + # Round 3: a consumer queued before restaging sees the prior table. selected_slot[0] = 0 - manager_stream.synchronize() snapshot = torch.empty_like(staging._block_offsets_device) - torch.cuda._sleep(20_000_000) snapshot.copy_(staging._block_offsets_device) - # The round executor's completion ordering: one event records the - # consumers and the manager stream waits on it. - staging._compaction_done_event.record(torch.cuda.current_stream(device)) - manager_stream.wait_event(staging._compaction_done_event) stage_once() current_stream.synchronize() @@ -570,212 +548,10 @@ def stage_once(): 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] - def test_compaction_move_offsets_stage_keep_plus_tail_and_pad_rows(self): - # The derived move offsets stage keep + tail moves per request - # (keep_count=4 -> [6, 7]); padded rows past the cohort repeat the - # final offset and contribute no moves. (The executor-call contract - # itself is pinned by the overlap-tail and draft publication tests.) - offsets_manager = TriAttention.__new__(TriAttention) - offsets_manager._request_capacity = 8 - offsets_manager._keep_count = 4 - offsets_manager._swa_window = None - offsets_manager._draft_protected_tail_capacity = None - eviction_inputs = [ - _make_eviction_input(request_id=7, source_length=8, target_tail_length=2), - _make_eviction_input(request_id=8, source_length=10, target_tail_length=3), - ] - dense, swa, draft = offsets_manager._compute_compaction_move_offsets(eviction_inputs) - assert dense == [0, 6, 13, 13, 13, 13, 13, 13, 13] - assert swa is None - assert draft is None - - @requires_sm100 - def test_fused_score_spans_distinct_storages_and_block_tables(self): - """ONE launch over layers in DISTINCT storages with DISTINCT block - tables (the production V2 shape), checked against the Torch oracle, - then relaunched after a round-start advance and a table rebind. - (Single-request launches are pinned by the CuTe score oracle's - request-count loop; the distinct-storages property is layer-axis.)""" - pytest.importorskip("cutlass") - from tensorrt_llm._torch.kv_cache_compression.triattention import triattention as module - - request_count = 8 - device = torch.device("cuda", torch.cuda.current_device()) - torch.manual_seed(20260707 + request_count) - max_requests = request_count - page_count = 2 - tokens_per_block = 32 - head_dim = 64 - num_freqs = head_dim // 2 - num_q_heads = 8 - seq_len = page_count * tokens_per_block - prompt_len = 1 - num_layers = 3 - pools = [ - ( - 0.125 - * torch.randn( - max_requests * page_count, 2, 1, tokens_per_block, head_dim, device=device - ) - ).to(torch.bfloat16) - for _ in range(num_layers) - ] - assert len({pool.untyped_storage().data_ptr() for pool in pools}) == num_layers - generator = torch.Generator(device="cpu").manual_seed(7 + request_count) - page_ids_3d = torch.stack( - [ - torch.randperm(max_requests * page_count, generator=generator)[ - : max_requests * page_count - ] - .view(max_requests, page_count) - .to(device=device, dtype=torch.int64) - for _ in range(num_layers) - ] - ).contiguous() - 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 = 0.125 * torch.randn(num_layers, num_q_heads, num_freqs, device=device) - freq = torch.rand(num_freqs, device=device) + 0.5 - omega = torch.rand(num_freqs, device=device) * 0.05 - offsets = torch.tensor([1.0, 2.0, 4.0], device=device) - round_device = torch.arange(max_requests, dtype=torch.int32, device=device) + 9 - logical_source_lengths = round_device[:request_count].tolist() - seq_lens = [seq_len - request % 2 for request in range(request_count)] - layer_order = list(range(num_layers)) - tri = _make_cute_buffers( - eviction_mode="per_head", - layer_pools=pools, - max_requests=max_requests, - seq_len=seq_len, - num_q_heads=num_q_heads, - q_real=q_real, - q_imag=q_imag, - mlr_coef=mlr, - freq_scale_sq=freq, - omega=omega, - offsets=offsets, - decode_width=seq_len - prompt_len, - keep_count=4, - # One page-table slot per layer (distinct pools). - layer_pool_ids=list(layer_order), - normalize_scores=False, - ) - source_lengths = torch.tensor(seq_lens, dtype=torch.int32, device=device) - - # Rounds stage through the production executor: the gather double - # writes each layer's K page ids and the bulk copy encodes K/V rows. - def gather_k_block_offsets(host_table, source, request_ids, num_blocks): - assert request_ids == list(range(request_count)) - source[..., 0, :].zero_() - source[:, :request_count, 0, :page_count].copy_( - page_ids_3d[:, :request_count].to(torch.int32).cpu() - ) - - manager = _make_staging_manager( - torch.zeros(num_layers, max_requests, 2, 8, dtype=torch.int32), - gather_k_block_offsets, - torch.cuda.Stream(device=device), - num_slots=num_layers, - ) - tri.kv_cache_manager = manager - - def prepared_cohort(): - return [ - _make_eviction_input( - request_id=request, - source_length=int(source_lengths[request]), - logical_source_length=int(round_device[request]), - prompt_length=prompt_len, - ) - for request in range(request_count) - ] - - def score_rectangle(): - # Test-side extraction of the decode-window rectangle from the - # scratch (production reduces the scratch in-kernel). - group = tri._num_q_heads // tri._num_kv_heads - segments = request_count * tri._num_layers - source = ( - tri._score_scratch[: tri._num_kv_heads * 8 * segments * tri._score_token_capacity] - .view( - tri._num_kv_heads, - 8, - request_count, - tri._num_layers, - tri._score_token_capacity, - )[:, :group] - .permute(2, 3, 0, 1, 4) - .reshape( - request_count, - tri._num_layers, - tri._num_q_heads, - tri._score_token_capacity, - ) - ) - columns = prompt_len + torch.arange( - tri._selection_width_capacity, dtype=torch.int64, device=device - ).view(1, 1, 1, -1) - columns = columns.clamp_(max=tri._score_token_capacity - 1).expand( - request_count, - tri._num_layers, - tri._num_q_heads, - tri._selection_width_capacity, - ) - return torch.gather(source, 3, columns) - - score_sentinel = -12345.0 - tri._score_scratch.fill_(score_sentinel) - # The compact stage is stubbed to a no-op: this test owns the score - # buffers only, never a staged move decision. - with mock.patch.object(module, "compact"): - tri._execute_eviction_round(prepared_cohort()) - fixed = score_rectangle() - assert tri._decode_lengths_device.tolist() == [seq_len - prompt_len for seq_len in seq_lens] - - oracle = _torch_tri_score_oracle( - pools, - {layer: page_ids_3d[layer, :request_count] for layer in layer_order}, - seq_lens, - logical_source_lengths, - q_real, - q_imag, - mlr, - freq, - omega, - offsets, - layer_order, - ) - for request in range(request_count): - for layer_slot, layer in enumerate(layer_order): - decode_length = seq_lens[request] - prompt_len - segment = fixed[request, layer_slot, :, :decode_length] - expected = oracle[request * num_layers + layer][ - :, prompt_len : prompt_len + decode_length - ] - torch.testing.assert_close(segment, expected, rtol=5e-3, atol=5e-3) - - round_device.add_(17) - page_ids_3d = page_ids_3d.roll(1, dims=2) - source_lengths.copy_( - torch.tensor( - [seq_len - (request + 1) % 2 for request in range(request_count)], - dtype=torch.int32, - device=device, - ) - ) - expected_second_widths = source_lengths - prompt_len - tri._score_scratch.fill_(score_sentinel) - tri._decode_lengths_device.fill_(-1) - with mock.patch.object(module, "compact"): - tri._execute_eviction_round(prepared_cohort()) - second_launch = score_rectangle() - assert torch.equal(tri._decode_lengths_device, expected_second_widths) - assert not torch.equal(second_launch, fixed) - class TestKernelMaskedSwa: @pytest.mark.parametrize("budget,fits_window", [(128, True), (127, False)]) - def test_layer_partition_uses_local_config_and_validates_window(self, budget, fits_window): + 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 @@ -794,9 +570,9 @@ def test_layer_partition_uses_local_config_and_validates_window(self, budget, fi if not fits_window: # The decode budget must cover the kernel-masked SWA window. with pytest.raises(ValueError, match="budget=127"): - mgr._attention_layer_partition() + mgr._resolve_attention_layers() return - dense, sliding, window = mgr._attention_layer_partition() + dense, sliding, window = mgr._resolve_attention_layers() load.assert_called_once_with( "/models/gpt-oss", trust_remote_code=True, local_files_only=True 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 index 05f55b6909e1..8b28bcdc36ac 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -4,19 +4,15 @@ import pytest import torch -from conftest import build_compaction as _build_compaction -from conftest import encode_block_offsets as _encode_block_offsets from conftest import make_cute_buffers as _make_cute_buffers from conftest import make_eviction_input as _make_eviction_input from conftest import make_ramp_pools as _make_ramp_pools from conftest import make_staging_manager as _make_staging_manager from conftest import rect_to_score_scratch as _rect_to_score_scratch -from conftest import run_compaction as _run_compaction -from conftest import set_protected_tails as _set_protected_tails from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - prepare_per_head_scores, + reduce_per_head_scores, ) @@ -46,63 +42,25 @@ def _make_selection_buffers( num_query_heads=1, num_kv_heads=1, ): - """A bare manager with selection-only attributes for the mode, without - CuTe score state or compaction (the settle launch writes only the - kept-ordinal rows). Mirrors the product's canonical row-major selection - allocation; ``_settle_top_tokens`` reads exactly these attributes.""" + """Allocate the production selection buffers without score or compaction.""" tri = TriAttention.__new__(TriAttention) tri.eviction_mode = eviction_mode tri._request_capacity = max_requests tri._selection_width_capacity = width - tri._keep_count = keep_count + tri.budget = keep_count tri._num_layers = num_layers tri._num_q_heads = num_query_heads tri._num_kv_heads = num_kv_heads - tri._decode_lengths_device = torch.full( - (max_requests,), width, dtype=torch.int32, device=device - ) tri._prompt_lengths_device = torch.zeros(max_requests, dtype=torch.int32, device=device) - if eviction_mode == "union": - tri._selection_rows_per_request = 1 - tri._selection_scores_rows = torch.empty( - (max_requests, width), dtype=torch.float32, device=device - ) - tri._selection_row_lengths = tri._decode_lengths_device - # Padded rows carry zero valid width; their provisional TopK entries - # must still be in-range ordinals for the finalizer's score gather. - tri._provisional_rows = torch.zeros( - (max_requests, keep_count), dtype=torch.int32, device=device - ) - tri._kept_ordinal_rows = torch.empty( - (max_requests, keep_count), dtype=torch.int32, device=device - ) - else: - selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads - tri._selection_rows_per_request = selection_rows - tri._row_mean = torch.empty( - max_requests, num_layers, num_query_heads, 1, dtype=torch.float32, device=device - ) - tri._row_inv_std = torch.empty_like(tri._row_mean) - tri._selection_scores_rows = torch.empty( - (max_requests * selection_rows, width), dtype=torch.float32, device=device - ) - tri._selection_row_lengths = torch.full( - (max_requests * selection_rows,), width, dtype=torch.int32, device=device - ) - tri._provisional_rows = torch.zeros( - (max_requests * selection_rows, keep_count), dtype=torch.int32, device=device - ) - tri._kept_ordinal_rows = torch.empty( - (max_requests * selection_rows, keep_count), 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, num_layers, num_q_heads, width = scores.shape + request_count, _, _, width = scores.shape score_scratch, prompt_lengths = _rect_to_score_scratch(scores, tri._num_kv_heads) - prepare_per_head_scores( + reduce_per_head_scores( score_scratch, tri._decode_lengths_device, prompt_lengths, @@ -111,16 +69,12 @@ def _select_per_head(tri, scores, *, normalize_scores): tri._selection_scores_rows, tri._selection_row_lengths, request_count=request_count, - num_layers=num_layers, - num_q_heads=num_q_heads, - num_kv_heads=tri._num_kv_heads, padded_head_columns=8, score_token_capacity=width, - selection_width=tri._selection_width_capacity, per_layer=tri.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - tri._settle_top_tokens(tri._request_capacity) + tri._select_top_tokens(tri._request_capacity) def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: @@ -246,7 +200,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) tri._selection_scores_rows.copy_(scores.amax(dim=1)) - tri._settle_top_tokens(tri._request_capacity) + tri._select_top_tokens(tri._request_capacity) actual = tri._kept_ordinal_rows.cpu() combined = scores.amax(dim=1).cpu() @@ -259,7 +213,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid @pytest.mark.parametrize("per_layer", [False, True]) @pytest.mark.parametrize("normalize_scores", [False, True]) -def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, normalize_scores): +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) @@ -287,7 +241,7 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no ) score_scratch, prompt_lengths = _rect_to_score_scratch(scores, kv_heads) - prepare_per_head_scores( + reduce_per_head_scores( score_scratch, decode_lengths, prompt_lengths, @@ -296,12 +250,8 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no selection_scores_rows, selection_row_lengths, request_count=request_count, - num_layers=layers, - num_q_heads=query_heads, - num_kv_heads=kv_heads, padded_head_columns=8, score_token_capacity=width, - selection_width=width, per_layer=per_layer, normalize_scores=normalize_scores, ) @@ -332,108 +282,6 @@ def test_fused_per_head_preparation_matches_ragged_torch_reference(per_layer, no assert torch.isneginf(selection_scores[request, :, decode_length:]).all() -@pytest.mark.parametrize("eviction_mode", ["union", "per_head", "per_layer_perhead"]) -def test_eager_compaction_preserves_exact_selected_bytes_and_tail(eviction_mode): - # Supported bf16 geometry; kept ordinals span all three pages so moves - # cross page boundaries. - device = torch.device("cuda", torch.cuda.current_device()) - request_count = 2 - num_layers = 2 - num_kv_heads = 2 - # Mixed prompt lengths prove per-request destination rebasing. - prompt_lens = [2, 5] - decode_keep_count = 4 - seq_len = 80 - tokens_per_block = 32 - pages_per_request = 3 - head_dim = 64 - protected_tails = [2, 1] - page_tables = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32, device=device) - initial_pools = _make_ramp_pools(num_layers, device=device) - pools = [pool.clone() for pool in initial_pools] - - # Decode-only kept ordinals holding absolute positions. - union_decode = torch.tensor( - [[16, 32, 56, 72], [24, 40, 48, 64]], dtype=torch.int64, device=device - ) - if eviction_mode == "union": - keep = union_decode - selection_rows = 1 - else: - selection_rows = num_kv_heads if eviction_mode == "per_head" else num_layers * num_kv_heads - keep = torch.empty( - request_count, - selection_rows, - decode_keep_count, - dtype=torch.int64, - device=device, - ) - for request in range(request_count): - for row in range(selection_rows): - keep[request, row] = torch.tensor( - sorted( - { - prompt_lens[request] + ((request + row + offset * 2) % 8) * 8 - for offset in range(decode_keep_count) - } - ), - dtype=torch.int64, - device=device, - ) - - compaction = _build_compaction( - eviction_mode=eviction_mode, - layer_pools=pools, - kept_token_ordinals=keep.to(torch.int32), - valid_sequence_lengths=torch.tensor([seq_len, seq_len], dtype=torch.int32, device=device), - kv_block_offsets=_encode_block_offsets(page_tables.unsqueeze(0)), - prompt_offsets=torch.tensor(prompt_lens, dtype=torch.int32, device=device), - protected_tail_capacity=max(protected_tails), - ) - _set_protected_tails(compaction, protected_tails) - # Production settles the kept ordinals into the contract's decision - # rows; with pre-settled ordinals the pack launch inside compact() is - # its exact analog. - _run_compaction(compaction) - torch.cuda.synchronize(device) - - for layer, (before_pool, after_pool) in enumerate(zip(initial_pools, pools)): - for request in range(request_count): - prompt_len = prompt_lens[request] - pages = page_tables[request].to(torch.long) - before = ( - before_pool[pages] - .permute(1, 2, 0, 3, 4) - .reshape(2, num_kv_heads, pages_per_request * tokens_per_block, head_dim) - ) - after = after_pool[pages].permute(1, 2, 0, 3, 4).reshape_as(before) - assert torch.equal(after[:, :, :prompt_len], before[:, :, :prompt_len]) - for head in range(num_kv_heads): - if eviction_mode == "union": - selected = keep[request] - elif eviction_mode == "per_head": - selected = keep[request, head] - else: - selected = keep[request, layer * num_kv_heads + head] - tail = torch.arange( - seq_len, - seq_len + protected_tails[request], - dtype=torch.int64, - device=device, - ) - source = torch.cat((selected, tail)) - destination = torch.arange( - prompt_len, - prompt_len + source.numel(), - dtype=torch.int64, - device=device, - ) - assert torch.equal( - after[:, head].index_select(1, destination), - before[:, head].index_select(1, source), - ) - - @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.""" @@ -547,9 +395,7 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): @requires_sm100 def test_union_two_rounds_preserve_bytes_tail_and_v2_page_reuse(): - """Two real eviction rounds through one live V2 cache: three pages - compact to two, releasing one physical page for reuse; expected keep - sets derive from a host-side score mirror.""" + """Preserve bytes, tails, and V2 page reuse across two real eviction rounds.""" pytest.importorskip("cutlass") import tensorrt_llm import tensorrt_llm.bindings @@ -754,76 +600,6 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: manager.shutdown() -def test_eager_compaction_rebases_masked_swa_window_and_tail(): - # Supported bf16 geometry; dense and SWA moves stay page-crossing. - device = torch.device("cuda", torch.cuda.current_device()) - dense_tables = torch.tensor([[2, 0, 1], [5, 3, 4]], dtype=torch.int32, device=device) - swa_tables = torch.tensor([[1, 2, 0], [4, 5, 3]], dtype=torch.int32, device=device) - initial_pools = _make_ramp_pools(2, num_kv_heads=1, device=device) - pools = [pool.clone() for pool in initial_pools] - # Decode-only kept ordinals holding absolute positions past the prompt. - keep = torch.tensor( - [[16, 32, 40, 56], [16, 24, 40, 48]], - dtype=torch.int64, - device=device, - ) - source_lengths = torch.tensor([64, 56], dtype=torch.int32, device=device) - protected_tails = [2, 1] - compaction = _build_compaction( - layer_pools=pools, - dense_layers=[0], - swa_layers=[1], - # Dense layer 0 stages in plane 0, the SWA layer in its own plane 1. - layer_pool_ids=[0, 1], - kept_token_ordinals=keep.to(torch.int32), - valid_sequence_lengths=source_lengths, - kv_block_offsets=_encode_block_offsets(torch.stack((dense_tables, swa_tables))), - prompt_offsets=torch.tensor([2, 2], dtype=torch.int32, device=device), - swa_window=2, - protected_tail_capacity=max(protected_tails), - ) - _set_protected_tails(compaction, protected_tails) - _run_compaction(compaction) - torch.cuda.synchronize(device) - - for request, (source_length, tail_length) in enumerate( - zip(source_lengths.tolist(), protected_tails) - ): - dense_pages = dense_tables[request].to(torch.long) - swa_pages = swa_tables[request].to(torch.long) - dense_before = initial_pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) - dense_after = pools[0][dense_pages].permute(1, 2, 0, 3, 4).reshape_as(dense_before) - swa_before = initial_pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape(2, 1, -1, 64) - swa_after = pools[1][swa_pages].permute(1, 2, 0, 3, 4).reshape_as(swa_before) - tail = torch.arange( - source_length, - source_length + tail_length, - dtype=torch.int64, - device=device, - ) - dense_source = torch.cat((keep[request], tail)) - dense_destination = torch.arange( - 2, 2 + dense_source.numel(), dtype=torch.int64, device=device - ) - swa_source = torch.arange( - source_length - 2, - source_length + tail_length, - dtype=torch.int64, - device=device, - ) - swa_destination = torch.arange(4, 4 + swa_source.numel(), dtype=torch.int64, device=device) - assert torch.equal(dense_after[:, :, :2], dense_before[:, :, :2]) - assert torch.equal(swa_after[:, :, :2], swa_before[:, :, :2]) - assert torch.equal( - dense_after.index_select(2, dense_destination), - dense_before.index_select(2, dense_source), - ) - assert torch.equal( - swa_after.index_select(2, swa_destination), - swa_before.index_select(2, swa_source), - ) - - def test_fold_union_ranks_matches_max_oracle(): """The TP union fold is an exact elementwise max over the gathered rank blocks.""" import triton From 1a0a26e7bdfcb930ef8e8a20664eafe155ea2206 Mon Sep 17 00:00:00 2001 From: tianruih Date: Mon, 27 Jul 2026 22:28:41 -0700 Subject: [PATCH 165/178] Simplify TriAttention eviction flow Signed-off-by: tianruih --- .../triattention/triattention.py | 180 ++++++++---------- tensorrt_llm/_torch/pyexecutor/_util.py | 71 +++---- .../_torch/pyexecutor/model_engine.py | 1 + tensorrt_llm/llmapi/llm_args.py | 21 +- .../test_kv_cache_compression_manager.py | 4 +- .../_torch/kv_cache_compression/conftest.py | 25 +-- .../test_triattention_draft_cocompaction.py | 89 ++++----- .../test_triattention_pipeline.py | 143 +++++--------- .../test_triattention_selection_compaction.py | 19 +- 9 files changed, 219 insertions(+), 334 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 2aacc7094e79..4035d3ee6cc1 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -23,7 +23,6 @@ official tool (github.com/WeianMao/triattention) and is converted at load. """ -from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Sequence, Tuple import torch @@ -60,42 +59,25 @@ # Required keys for the calibration ``.pt`` consumed by TriAttention. _REQUIRED_CALIBRATION_KEYS = frozenset({"E_q", "E_q_norm", "omega", "freq_scale_sq"}) -# Generation requests skipped by every eviction step. -_SKIP_REQUEST_STATES = ( - LlmRequestState.GENERATION_COMPLETE, - LlmRequestState.CONTEXT_INIT, -) - -# Upper bound of the geometric integration offset ladder [1, 2, 4, ...]. -_MAX_INTEGRATION_OFFSET = 65536 - +_MEAN_PHASE_OFFSETS = tuple(float(1 << exponent) for exponent in range(17)) -@dataclass -class _RequestState: - confirmed_tokens: int = 0 - evicted_tokens: int = 0 - -class _EvictionInput(NamedTuple): - """One due request's eviction operands for a single round.""" +class _EvictionRequest(NamedTuple): + """One due request and the cache state needed by its eviction round.""" request: "LlmRequest" target_cache: object draft_cache: Optional[object] - state: _RequestState source_length: int - logical_source_length: int - prompt_length: int target_tail_length: int def _allocate_block_offset_staging( + manager: KVCacheManagerV2, anchor_pool: torch.Tensor, *, - num_pools: int, request_capacity: int, token_capacity: int, - max_source_blocks: int, ) -> Tuple[torch.Tensor, torch.Tensor]: """Allocate a host snapshot and persistent device table. @@ -104,8 +86,9 @@ def _allocate_block_offset_staging( """ tokens_per_block = int(anchor_pool.shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block + max_source_blocks = int(manager.host_kv_cache_block_offsets.shape[-1]) block_capacity = min((page_count + 3) // 4 * 4, max_source_blocks) - shape = (num_pools, request_capacity, 2, block_capacity) + shape = (int(manager.num_pools), request_capacity, 2, block_capacity) host = torch.empty(shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) device_table = torch.empty(shape, dtype=torch.int32, device=anchor_pool.device) return host, device_table @@ -119,7 +102,6 @@ class _MeanPhaseTable: def __init__(self, omega: torch.Tensor, device: torch.device) -> None: self._omega = omega.to(device=device, dtype=torch.float32).contiguous() - self._offsets = tuple(float(1 << i) for i in range(_MAX_INTEGRATION_OFFSET.bit_length())) self.cos: Optional[torch.Tensor] = None self.sin: Optional[torch.Tensor] = None self.rows = 0 @@ -143,11 +125,11 @@ def reserve(self, rows: int) -> None: ) sin_table = torch.zeros_like(cos_table) # Fixed summation order keeps table rebuilds bit-stable. - for offset in self._offsets: + 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(self._offsets) + scale = 1.0 / len(_MEAN_PHASE_OFFSETS) self.cos = cos_table.mul_(scale) self.sin = sin_table.mul_(scale) self.rows = target @@ -178,11 +160,7 @@ def __init__( self.calibration_path = config.calibration_path self._load_calibration() - # Per-request eviction progress. - self._request_states: Dict[int, _RequestState] = {} - # In-flight overlap batch reference; membership resolves lazily. - self._inflight_scheduled_batch: Optional["ScheduledRequests"] = None - self._inflight_generation_request_ids: Optional[set[int]] = None + 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 = ( @@ -195,7 +173,9 @@ def __init__( + int(draft_kv_cache_manager._kv_reserve_draft_tokens) + 1 ) - self._generation_growth = 1 + int(kv_cache_manager._kv_reserve_draft_tokens) + # 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 by the manager/config; only absolute source length can grow # when later requests arrive with longer prompts. self._request_capacity = int(kv_cache_manager.max_batch_size) @@ -263,6 +243,8 @@ def _resolve_attention_layers(self) -> Tuple[List[int], List[int], Optional[int] ] 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 @@ -368,15 +350,13 @@ def _rope_tables(self, freq_count: int): # ---- framework hooks (call order) ---- def on_request_init(self, request: "LlmRequest", **kwargs) -> None: - """Register the request for eviction tracking.""" + """Validate and reserve request-dependent eviction capacity.""" self._validate_request_capacity(request) self._reserve_eviction_capacity(request) - self._request_states[request.py_request_id] = _RequestState() def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Snapshot the in-flight batch; mutation remains in final update.""" - self._inflight_scheduled_batch = scheduled_batch - self._inflight_generation_request_ids = None + """Remember the batch whose next-step KV capacity is already reserved.""" + 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. @@ -386,10 +366,6 @@ def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) with nvtx_range_debug("triattention.generation_step_end", color="blue"): self._evict_due_requests(scheduled_batch) - def on_request_finish(self, request: "LlmRequest", **kwargs) -> None: - """Drop this request's eviction state; persistent allocations stay resident.""" - self._request_states.pop(request.py_request_id, None) - # ---- request capacity ---- def _validate_request_capacity(self, request: "LlmRequest") -> None: @@ -430,10 +406,21 @@ def _evict_due_requests( ) -> None: """Collect due requests, execute one eviction round, publish, and resize.""" manager = self.kv_cache_manager - eviction_inputs: List[_EvictionInput] = [] + 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() + ) with nvtx_range("triattention.metadata", color="cyan"): for request in scheduled_batch.generation_requests: - if request.is_dummy or request.state in _SKIP_REQUEST_STATES: + 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) @@ -447,18 +434,8 @@ def _evict_due_requests( draft_cache = self.draft_kv_cache_manager.kv_cache_map[request_id] if not draft_cache.is_active: continue - # Only requests with active target/draft caches advance the cadence ledger. - state = self._request_states[request_id] - previous_confirmed_tokens = state.confirmed_tokens - confirmed_tokens = ( - previous_confirmed_tokens + 1 + int(request.py_num_accepted_draft_tokens) - ) - state.confirmed_tokens = confirmed_tokens - if previous_confirmed_tokens // self.beta >= confirmed_tokens // self.beta: - continue - # Speculative reserve + in-flight overlap growth: contiguous tail moved byte-for-byte. target_tail_length = self._num_extra_kv_tokens + ( - self._inflight_generation_growth(scheduled_batch, request_id) + 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: @@ -467,55 +444,47 @@ def _evict_due_requests( 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 - eviction_inputs.append( - _EvictionInput( + eviction_requests.append( + _EvictionRequest( request=request, target_cache=target_cache, draft_cache=draft_cache, - state=state, source_length=source_length, - # Uncompressed logical position. - logical_source_length=source_length + state.evicted_tokens, - prompt_length=prompt_length, target_tail_length=target_tail_length, ) ) - if not eviction_inputs: + if not eviction_requests: return # Ungated NVTX: the due count in the message shows each round's size. with nvtx_range( - f"triattention.evict_request_group reqs={len(eviction_inputs)}", + f"triattention.evict_request_group reqs={len(eviction_requests)}", color="purple", ): - self._execute_eviction_round(eviction_inputs) - for item in eviction_inputs: - evicted = item.source_length - item.prompt_length - self.budget - item.state.evicted_tokens += evicted + 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 = item.state.evicted_tokens - self._resize_compacted_caches(eviction_inputs) - - def _inflight_generation_growth( - self, scheduled_batch: "ScheduledRequests", request_id: int - ) -> int: - inflight = self._inflight_scheduled_batch - if inflight is None or scheduled_batch is inflight: - return 0 - member_ids = self._inflight_generation_request_ids - if member_ids is None: - member_ids = {request.py_request_id for request in inflight.generation_requests} - self._inflight_generation_request_ids = member_ids - if request_id not in member_ids: - return 0 - return self._generation_growth + item.request.py_num_compressed_tokens += evicted + self._resize_compacted_caches(eviction_requests) def _execute_eviction_round( self, - eviction_inputs: Sequence[_EvictionInput], + eviction_requests: Sequence[_EvictionRequest], ) -> None: """Score, select, and compact one due request group.""" manager = self.kv_cache_manager @@ -525,12 +494,15 @@ def _execute_eviction_round( # compression resource update, so the round can use caller current. try: with nvtx_range_debug("triattention.page_table_stage", color="orange"): - request_ids = [item.request.py_request_id for item in eviction_inputs] - logical_source_lengths = [item.logical_source_length for item in eviction_inputs] - prompt_lengths = [item.prompt_length for item in eviction_inputs] - source_lengths = [item.source_length for item in eviction_inputs] + 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_inputs) + self._compute_compaction_move_offsets(eviction_requests) ) metadata_rows = ( logical_source_lengths, @@ -548,7 +520,7 @@ def _execute_eviction_round( 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_inputs) :] = 0 + host_table[:3, len(eviction_requests) :] = 0 try: self._stage_block_offsets( manager, @@ -569,7 +541,7 @@ def _execute_eviction_round( finally: self._staging_reuse_event.record(stream) - request_count = len(eviction_inputs) + request_count = len(eviction_requests) union = self.eviction_mode == "union" with nvtx_range("triattention.score", color="blue"): # In-place refresh: the compiled score launches captured these pointers. @@ -636,7 +608,7 @@ def _execute_eviction_round( def _compute_compaction_move_offsets( self, - eviction_inputs: Sequence[_EvictionInput], + eviction_requests: Sequence[_EvictionRequest], ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: """Build padded cumulative dense, SWA, and draft move offsets.""" @@ -647,7 +619,7 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: offsets.extend(offsets[-1:] * (self._request_capacity - len(moves_per_request))) return offsets - tails = [int(item.target_tail_length) for item in eviction_inputs] + tails = [int(item.target_tail_length) for item in eviction_requests] dense = padded_offsets([self.budget + tail for tail in tails]) swa = None if self._swa_window is not None: @@ -655,7 +627,7 @@ def padded_offsets(moves_per_request: List[int]) -> List[int]: draft = None if self.draft_kv_cache_manager is not None: draft = padded_offsets( - [self.budget + self._draft_protected_tail_capacity] * len(eviction_inputs) + [self.budget + self._draft_protected_tail_capacity] * len(eviction_requests) ) return dense, swa, draft @@ -704,10 +676,12 @@ def _select_top_tokens(self, request_count: int) -> None: SELECTION_ROWS=self._selection_rows_per_request, ) - def _resize_compacted_caches(self, eviction_inputs: Sequence[_EvictionInput]) -> None: + def _resize_compacted_caches(self, eviction_requests: Sequence[_EvictionRequest]) -> None: with nvtx_range("triattention.resize", color="red"): - for item in eviction_inputs: - resized_capacity = item.prompt_length + self.budget + item.target_tail_length + for item in eviction_requests: + resized_capacity = ( + int(item.request.py_prompt_len) + self.budget + item.target_tail_length + ) if not item.target_cache.resize(resized_capacity, None): raise RuntimeError( "Failed to resize compacted target KV cache for " @@ -717,9 +691,11 @@ def _resize_compacted_caches(self, eviction_inputs: Sequence[_EvictionInput]) -> if self.draft_kv_cache_manager is None: return # Target and draft retain the same selected tokens. - for item in eviction_inputs: + for item in eviction_requests: resized_capacity = ( - item.prompt_length + self.budget + self._draft_protected_tail_capacity + int(item.request.py_prompt_len) + + self.budget + + self._draft_protected_tail_capacity ) if not item.draft_cache.resize(resized_capacity, None): raise RuntimeError( @@ -938,24 +914,20 @@ def _build_eviction_capacity( anchor_pool = self._target_layout["layer_pools"][dense_layer] request_capacity = self._request_capacity block_offsets_host, block_offsets_device = _allocate_block_offset_staging( + self.kv_cache_manager, anchor_pool, - num_pools=int(self.kv_cache_manager.num_pools), request_capacity=request_capacity, token_capacity=score_token_capacity + self._protected_tail_capacity, - max_source_blocks=int(self.kv_cache_manager.host_kv_cache_block_offsets.shape[-1]), ) 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_staging( + self.draft_kv_cache_manager, draft_anchor_pool, - num_pools=int(self.draft_kv_cache_manager.num_pools), request_capacity=request_capacity, token_capacity=(score_token_capacity + self._draft_protected_tail_capacity), - max_source_blocks=int( - self.draft_kv_cache_manager.host_kv_cache_block_offsets.shape[-1] - ), ) score_scratch, launch_score = build_score_pipeline( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index b3339524c655..4444dd0b4a83 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -39,7 +39,6 @@ get_default_trtllm_modules_to_hf_modules) from tensorrt_llm.lora_manager import load_torch_lora from tensorrt_llm.mapping import CpType, Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig from ..attention_backend import get_sparse_attn_kv_cache_manager from ..hostfunc import set_low_latency_dispatch @@ -2114,59 +2113,35 @@ def _create_kv_cache_manager( def validate_kv_cache_compression_with_spec( config: KvCacheCompressionConfig, spec_config: Optional[SpeculativeConfig], - draft_kv_cache_manager: Optional[KVCacheManagerV2], + draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, ) -> 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()): + if spec_config is None: + return + if not config.kv_cache_compression_mode.is_eviction_method(): 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 config.eviction_mode != "union": + raise ValueError( + "KV-cache compression with speculative decoding 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).") - if config.algorithm == "triattention": - if spec_config.max_draft_len is None: - raise ValueError( - "TriAttention speculative compatibility requires a resolved " - "max_draft_len") - if not spec_config.is_linear_tree: - raise ValueError( - "TriAttention speculative compatibility requires linear " - "drafting") - if spec_config.draft_len_schedule is not None: - raise ValueError("TriAttention does not yet support dynamic " - "speculative draft lengths") - # Compression eviction is only validated with greedy acceptance. - if spec_config.use_rejection_sampling or getattr( - spec_config, "use_relaxed_acceptance_for_thinking", False): - raise ValueError("TriAttention does not support speculative " - "rejection sampling or relaxed acceptance") - if draft_kv_cache_manager is None: - raise ValueError( - "TriAttention speculative compatibility requires a separate " - "draft KV cache; shared target/draft pools cannot be " - "compacted safely") - if not draft_kv_cache_manager.is_draft: - raise ValueError( - "TriAttention speculative compatibility requires the actual " - "separate draft KV cache manager") - if config.eviction_mode != "union": - raise ValueError( - "TriAttention draft KV co-compression supports only " - "eviction_mode='union'; draft layers are never scored") - if any(window is not None - for window in draft_kv_cache_manager.max_attention_window_vec - ) or any(not isinstance(layer, AttentionLayerConfig) - or layer.sliding_window_size is not None - for layer in draft_kv_cache_manager. - kv_cache_manager_py_config.layers): - raise ValueError("TriAttention draft KV co-compression requires " - "full-attention draft V2 lifecycles") + f"KV-cache compression does not support speculative decoding " + f"mode {mode.name}; use one-model MTP or EAGLE3") + if not spec_config.is_linear_tree: + raise ValueError( + "KV-cache compression requires linear speculative decoding") + if spec_config.draft_len_schedule is not None: + raise ValueError( + "KV-cache compression requires a fixed speculative draft length") + if draft_kv_cache_manager is None: + raise ValueError( + "KV-cache compression requires a separate draft KV cache") + if any(window is not None + for window in draft_kv_cache_manager.max_attention_window_vec): + raise ValueError( + "KV-cache compression requires full-attention draft KV") def create_kv_cache_compression_manager( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3795c479bf0d..08bb2d6c806a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4941,6 +4941,7 @@ def previous_seq_slots_device(): and not multimodal_params_list and not lora_params and attn_metadata.padded_num_tokens is None and self._get_position_id_offset() == 0 + # Compression may shrink KV history between decode steps. and not getattr(kv_cache_manager, "kv_compression_manages_history", False)): self._steady_gen_positions_pinned[:_n_gen].copy_( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index dc915ce9c0fb..80951a2f974b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3579,29 +3579,18 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): "`divide_length`): one speculative iteration may advance the counter " "by multiple accepted tokens; at most one eviction is coalesced per update." ) - model_path: Optional[str] = Field( - default=None, + 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. Required by TriAttention.") - calibration_path: Optional[str] = Field( - default=None, + "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.") - @model_validator(mode="after") - def _require_calibration_inputs(self): - # Both paths are consumed at manager construction; failing here surfaces - # the error at config-validation time instead of deep in executor setup. - if not self.model_path or not self.calibration_path: - raise ValueError( - "TriAttention requires both model_path and calibration_path; " - "TRT-LLM consumes an official calibration file and does not " - "compute one.") - return self - @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) class AgentTreeConfig(StrictBaseModel, PybindMirror): 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..6bf7511e46b9 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -311,8 +311,8 @@ def test_spec_gate_only_restricts_eviction_methods(self): # Non-evicting methods pass with any speculative mode; no exception. config = KvCacheCompressionConfig(algorithm="offload") 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) + validate_kv_cache_compression_with_spec(config, spec_config) + validate_kv_cache_compression_with_spec(config, None) # ---------------------------------------------------------------------- # diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 5eeae90e03b7..8ed4a1cc8d39 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -346,36 +346,29 @@ def make_triattention(**overrides): return TriAttention(make_tri_config(**overrides), make_fake_v2()) -def make_eviction_input( +def make_eviction_request( request=None, *, request_id=0, source_length, - logical_source_length=None, - prompt_length=0, target_tail_length=0, target_cache=None, draft_cache=None, - state=None, ): - """One due-request item shaped exactly like ``_evict_due_requests`` builds.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _EvictionInput, - _RequestState, - ) + """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_num_compressed_tokens=0) - return _EvictionInput( + 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, - state=_RequestState() if state is None else state, source_length=int(source_length), - logical_source_length=int( - source_length if logical_source_length is None else logical_source_length - ), - prompt_length=int(prompt_length), target_tail_length=int(target_tail_length), ) 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 index 5a2446f50789..63ff153625b1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -8,17 +8,14 @@ import pytest import torch -from conftest import make_eviction_input as _make_eviction_input +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 ( - TriAttention, - _RequestState, -) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention def test_execute_eviction_round_uses_current_stream_and_hands_back_to_target(): @@ -48,7 +45,7 @@ def test_execute_eviction_round_uses_current_stream_and_hands_back_to_target(): tri.kv_cache_manager = manager tri.draft_kv_cache_manager = draft_manager compute_stream = mock.Mock() - eviction_inputs = [_make_eviction_input(request_id=7, source_length=8)] + eviction_requests = [_make_eviction_request(request_id=7, source_length=8)] class Boom(RuntimeError): pass @@ -61,7 +58,7 @@ class Boom(RuntimeError): mock.patch.object(tri, "_stage_block_offsets") as stage, ): with pytest.raises(Boom): - tri._execute_eviction_round(eviction_inputs) + tri._execute_eviction_round(eviction_requests) # Both page-table planes were snapshotted before the round body fired. assert stage.call_count == 2 @@ -75,36 +72,47 @@ class Boom(RuntimeError): @pytest.mark.parametrize( "gate,match", [ - # One representative per call-site guard family (the per-mode/per-config - # variants raise through the same checks). kv_factor geometry needs no - # admission gate: the native compact op TORCH_CHECKs every pool's K/V - # plane count at the first compact. - ("callsite_dflash", "standard paged cache compacted together"), + ("callsite_dflash", "one-model MTP or EAGLE3"), ("union_only_per_head", "union"), - ("full_attention_draft", "full-attention draft"), + ("dynamic_tree", "linear speculative"), + ("dynamic_draft_length", "fixed speculative draft length"), + ("missing_draft", "separate draft"), + ("sliding_draft", "full-attention draft"), ], ) -def test_draft_admission_gates_raise(gate, match): - # Draft/spec admission is owned by the executor call-site gate: rejected - # before any compression manager exists. +def test_speculative_admission_gates_raise(gate, match): from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec from tensorrt_llm.llmapi.llm_args import DFlashDecodingConfig, MTPDecodingConfig - draft_manager = _make_fake_v2(is_draft=True) - if gate == "full_attention_draft": - draft_manager.max_attention_window_vec = [128] - spec_config = ( - DFlashDecodingConfig(max_draft_len=3) - if gate == "callsite_dflash" - else MTPDecodingConfig(max_draft_len=1) - ) + if gate == "callsite_dflash": + spec_config = DFlashDecodingConfig(max_draft_len=3) + elif gate == "dynamic_tree": + spec_config = MTPDecodingConfig( + max_draft_len=2, + use_dynamic_tree=True, + dynamic_tree_max_topK=2, + ) + elif gate == "dynamic_draft_length": + spec_config = MTPDecodingConfig( + max_draft_len=1, + draft_len_schedule={1: 1, 8: 0}, + ) + 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", ) + draft_manager = None if gate == "missing_draft" else _make_fake_v2(is_draft=True) + if gate == "sliding_draft": + draft_manager.max_attention_window_vec = [128] with pytest.raises(ValueError, match=match): - validate_kv_cache_compression_with_spec(config, spec_config, draft_manager) + validate_kv_cache_compression_with_spec( + config, + spec_config, + draft_manager, + ) def test_compressed_count_is_monotone_and_tracks_confirmed_length(): @@ -128,7 +136,6 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager._draft_protected_tail_capacity = 1 request = _make_request(7, py_prompt_len=2, py_num_accepted_draft_tokens=1) - manager._request_states[7] = _RequestState() batch = SimpleNamespace(generation_requests=[request]) # Every step confirms one sampled token plus one accepted draft token. @@ -136,7 +143,6 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): confirmed = uncompressed cache.capacity = confirmed previous_published = 0 - previous_evicted = 0 eviction_rounds = 0 with _mocked_eviction_internals(manager) as internals: for _ in range(6): @@ -146,24 +152,18 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): manager._evict_due_requests(batch) - state = manager._request_states[7] - if state.evicted_tokens > previous_evicted: + published = request.py_num_compressed_tokens + if published > previous_published: # An eviction round compacted the cache to prompt + budget. eviction_rounds += 1 - confirmed -= state.evicted_tokens - previous_evicted - previous_evicted = state.evicted_tokens + confirmed -= published - previous_published cache.capacity = confirmed assert confirmed == 2 + 4 - # The staged logical position restores the uncompressed - # length: physical confirmed plus everything evicted so far - # (the eviction input's logical_source_length). - prepared = internals.execute.call_args.args[0] - assert prepared[0].logical_source_length == uncompressed # The published count equals the uncompressed confirmed logical # length minus the physical confirmed length, and never decreases. - assert request.py_num_compressed_tokens == uncompressed - confirmed - assert request.py_num_compressed_tokens >= previous_published - previous_published = request.py_num_compressed_tokens + assert published == uncompressed - confirmed + assert published >= previous_published + previous_published = published assert eviction_rounds == 3 assert previous_published == 12 @@ -236,19 +236,22 @@ def test_staged_block_width_clamps_to_manager_source_width(): ) anchor_pool = torch.empty(1, 2, 1, 32, 4) + manager = SimpleNamespace( + num_pools=1, + host_kv_cache_block_offsets=torch.empty(1, 2, 2, 4, dtype=torch.int32), + ) host, device_table = _allocate_block_offset_staging( + manager, anchor_pool, - num_pools=1, request_capacity=2, token_capacity=129, - max_source_blocks=4, ) assert host.shape[-1] == 4 and device_table.shape[-1] == 4 + manager.host_kv_cache_block_offsets = torch.empty(1, 2, 2, 64, dtype=torch.int32) host, device_table = _allocate_block_offset_staging( + manager, anchor_pool, - num_pools=1, request_capacity=2, token_capacity=129, - max_source_blocks=64, ) assert host.shape[-1] == 8 and device_table.shape[-1] == 8 diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 7e57faf5ab04..e0d04a6c5b0b 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -36,25 +36,13 @@ # 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 ( - TriAttention, - _RequestState, -) +from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention # 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 -def _set_request_state(manager, request_id, *, confirmed_tokens=0, evicted_tokens=0): - state = _RequestState( - confirmed_tokens=confirmed_tokens, - evicted_tokens=evicted_tokens, - ) - manager._request_states[request_id] = state - return state - - @pytest.fixture def flat_calibration_pt(tmp_path): """Build a minimal valid calibration ``.pt`` in our flat runtime schema.""" @@ -91,9 +79,7 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): class TestTriAttentionClass: - def test_request_init_and_finish_lifecycle(self): - # Init reserves request-dependent capacity and tracks state. Finish - # clears only request state; persistent runtime objects stay resident. + def test_request_init_validates_and_reserves_capacity(self): manager = _make_fake_v2() manager.num_extra_kv_tokens = 4 manager._kv_reserve_draft_tokens = 4 @@ -110,20 +96,9 @@ def test_request_init_and_finish_lifecycle(self): assert triattention.adjusts_generation_kv_length is True assert manager.kv_compression_manages_history - assert set(triattention._request_states) == {11, 12} assert validate.call_args_list == [mock.call(request_11), mock.call(request_12)] assert reserve.call_args_list == [mock.call(request_11), mock.call(request_12)] - phase = object() - triattention._phase = phase - batch = SimpleNamespace() - triattention._inflight_scheduled_batch = batch - triattention.on_request_finish(_make_request(11)) - triattention.on_request_finish(_make_request(12)) - assert triattention._request_states == {} - assert triattention._inflight_scheduled_batch is batch - assert triattention._phase is phase - def test_capacity_guard_uses_maximum_steady_eviction_peak(self): manager = _make_triattention(budget=10, beta=8) manager.kv_cache_manager.get_num_available_tokens = mock.Mock(return_value=17) @@ -248,7 +223,7 @@ class TestCompressedTokenPublication: 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) + 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. @@ -256,14 +231,12 @@ def test_identity_selection_is_filtered_before_launch(self): capacity=6, history_length=0, is_active=True, resize=mock.Mock(return_value=True) ) manager.kv_cache_manager.kv_cache_map = {7: cache} - state = _set_request_state(manager, 7, confirmed_tokens=127) 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 - assert state.evicted_tokens == 0 cache.resize.assert_not_called() @@ -312,37 +285,45 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft num_extra_kv_tokens=num_extra_kv_tokens, _kv_reserve_draft_tokens=kv_reserve_draft_tokens, ) - mgr._request_states = {} - _set_request_state(mgr, 7, confirmed_tokens=127) mgr.beta = 128 mgr.budget = 4096 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 (pre-launch, no cadence mutation) while the rest of - # the request group proceeds. - manager, first_request, _ = self._make_due_decode_request(seq_len=1024 + 4096 + 1) + # 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) - manager.kv_cache_manager.kv_cache_map[8] = SimpleNamespace(is_active=False) - first_state = manager._request_states[7] - second_state = _set_request_state(manager, 8, confirmed_tokens=127) + 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 + manager._evict_due_requests(SimpleNamespace(generation_requests=[second_request])) - # Only the active request launched; the suspended one deferred whole. - eviction_inputs = internals.execute.call_args.args[0] - assert [item.request.py_request_id for item in eviction_inputs] == [7] - assert first_state.confirmed_tokens == 128 - assert second_state.confirmed_tokens == 127 + 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 == 128 + second_cache.resize.assert_called_once_with(1024 + 4096, None) - # ``accepted`` enters the prepared item linearly; the zero and maximal - # boundary rows pin the whole family. + # 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 + 1 + accepted + confirmed = 1024 + 4096 + 128 reserve = 2 current_growth = 4 tail = reserve + current_growth @@ -356,7 +337,7 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): request.py_num_accepted_draft_tokens = accepted cache = mgr.kv_cache_manager.kv_cache_map[7] cache.capacity = confirmed + tail - mgr._inflight_scheduled_batch = SimpleNamespace(generation_requests=[request]) + 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} @@ -375,8 +356,6 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): assert launched.target_cache is cache assert launched.draft_cache is draft_cache assert launched.source_length == confirmed - assert launched.logical_source_length == confirmed - assert launched.prompt_length == 1024 assert launched.target_tail_length == tail assert request.py_num_compressed_tokens == confirmed - retained cache.resize.assert_called_once_with(retained + tail, None) @@ -385,9 +364,8 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): 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. - physical_confirmed = 6100 + physical_confirmed = 6172 manager = _make_triattention(beta=128) - _set_request_state(manager, 7, confirmed_tokens=127, evicted_tokens=100) cache = SimpleNamespace( capacity=physical_confirmed, history_length=1024, @@ -399,6 +377,7 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): request = _make_request( 7, py_prompt_len=1024, + py_num_compressed_tokens=100, max_beam_num_tokens=999999, py_draft_tokens=[1, 2, 3, 4], ) @@ -406,16 +385,15 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): with _mocked_eviction_internals(manager) as internals: manager._evict_due_requests(SimpleNamespace(generation_requests=[request])) - eviction_inputs = internals.execute.call_args.args[0] - assert eviction_inputs[0].source_length == physical_confirmed - # The logical position restores everything already evicted. - assert eviction_inputs[0].logical_source_length == physical_confirmed + 100 + eviction_requests = internals.execute.call_args.args[0] + assert eviction_requests[0].source_length == physical_confirmed + assert request.py_num_compressed_tokens == ( + 100 + physical_confirmed - 1024 - manager.budget + ) cache.resize.assert_called_once_with(1024 + manager.budget, None) - def test_one_model_draft_co_compression_contract_is_accepted(self): - # Construction accepts the separate draft manager, and the executor - # call-site gate accepts the one-model MTP roundtrip (base-ctor - # marking is asserted in the executor manager tests). + @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(TriAttention, "_initialize_eviction_state"): TriAttention( @@ -425,47 +403,24 @@ def test_one_model_draft_co_compression_contract_is_accepted(self): ) from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec - from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + 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_with_spec( _make_tri_config(budget=8), - MTPDecodingConfig(max_draft_len=1), + spec_config, draft_manager, ) - @pytest.mark.parametrize( - "reserved_draft,expected_growth", - [ - (0, 1), - # The reserved draft width protects capacity regardless of any - # step's actual draft length: growth is the cached constant. - (6, 7), - ], - ) - def test_prepare_snapshots_fixed_linear_generation_growth( - self, reserved_draft, expected_growth - ): - manager = _make_fake_v2() - manager._kv_reserve_draft_tokens = reserved_draft - manager.kv_cache_map = { - 7: SimpleNamespace(capacity=106, is_active=True), - } - with mock.patch.object(TriAttention, "_initialize_eviction_state"): - triattention = TriAttention(_make_tri_config(budget=8), manager) - batch = SimpleNamespace( - context_requests=[], - generation_requests=[_make_request(7, py_draft_tokens=[1, 2, 3])], - ) - - triattention.prepare_resources(batch) - - assert triattention._inflight_scheduled_batch is batch - # Members of the prepared batch grow by the cached constant; others - # by zero. The prepared batch itself is the identity early-out. - assert triattention._inflight_generation_growth(SimpleNamespace(), 7) == expected_growth - assert triattention._inflight_generation_growth(SimpleNamespace(), 99) == 0 - assert triattention._inflight_generation_growth(batch, 7) == 0 - class TestFixedScoreMetadata: def test_union_forces_normalized_scores(self): 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 index 8b28bcdc36ac..ffa2ff11956a 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -5,8 +5,9 @@ import pytest import torch from conftest import make_cute_buffers as _make_cute_buffers -from conftest import make_eviction_input as _make_eviction_input +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 @@ -373,11 +374,9 @@ def gather_k_block_offsets(host_table, source, request_ids, num_blocks): torch.cuda.Stream(device=device), num_slots=2, ) - eviction_inputs = [ - _make_eviction_input(request_id=7, source_length=seq_len, logical_source_length=0) - ] + eviction_requests = [_make_eviction_request(request_id=7, source_length=seq_len)] tri.kv_cache_manager = manager - tri._execute_eviction_round(eviction_inputs) + tri._execute_eviction_round(eviction_requests) assert torch.equal(tri._kept_ordinal_rows.view_as(expected_keep), expected_keep) torch.cuda.synchronize(device) @@ -520,12 +519,10 @@ def expected_keep() -> torch.Tensor: def evict_once() -> tuple[torch.Tensor, torch.Tensor]: before = snapshot(seq_len + protected_tail) - eviction_inputs = [ - _make_eviction_input( - request_id=request_id, + eviction_requests = [ + _make_eviction_request( + request=_make_request(request_id, py_prompt_len=prompt_len), source_length=seq_len, - logical_source_length=0, - prompt_length=prompt_len, target_tail_length=protected_tail, ) ] @@ -533,7 +530,7 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: # 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_inputs) + 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) From 418c660c73f363cf5ff151bdbbfe14361f707831 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 28 Jul 2026 02:23:11 -0700 Subject: [PATCH 166/178] [None][refactor] simplify TriAttention eviction runtime Signed-off-by: tianruih --- .../triattention/triattention.py | 491 ++++++++---------- tensorrt_llm/_torch/modules/attention.py | 7 +- tensorrt_llm/_torch/pyexecutor/_util.py | 18 +- tensorrt_llm/llmapi/llm_args.py | 22 +- .../_torch/kv_cache_compression/conftest.py | 2 +- .../test_triattention_draft_cocompaction.py | 32 +- .../test_triattention_pipeline.py | 59 +-- tests/unittest/llmapi/test_llm_args.py | 25 + 8 files changed, 281 insertions(+), 375 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 4035d3ee6cc1..1d6fb03bd1f0 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -35,7 +35,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheCompressionManager from tensorrt_llm._torch.utils import next_positive_power_of_2 -from tensorrt_llm._utils import nvtx_range, nvtx_range_debug, prefer_pinned +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, ) @@ -79,11 +79,7 @@ def _allocate_block_offset_staging( request_capacity: int, token_capacity: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Allocate a host snapshot and persistent device table. - - Both use the native V2 ``[pool, request, K/V, block]`` layout and grow only - when a newly admitted request raises the high-water mark. - """ + """Allocate a V2 page-table snapshot for the scorer/compactor token span.""" tokens_per_block = int(anchor_pool.shape[3]) page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block max_source_blocks = int(manager.host_kv_cache_block_offsets.shape[-1]) @@ -176,12 +172,13 @@ def __init__( # 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 by the manager/config; only absolute source length can grow - # when later requests arrive with longer prompts. + # Fixed buffer geometry. These are TriAttention scratch dimensions, + # not KV capacities owned by KVCacheManagerV2. self._request_capacity = int(kv_cache_manager.max_batch_size) - self._selection_width_capacity = ( - self.budget + 2 * self.beta + int(kv_cache_manager.max_total_draft_tokens) - ) + 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. + self._selection_width_capacity = self.budget + self.beta + 2 * max_draft_tokens + 1 max_tail_capacity = max( self._protected_tail_capacity, self._draft_protected_tail_capacity, @@ -350,12 +347,35 @@ def _rope_tables(self, freq_count: int): # ---- framework hooks (call order) ---- def on_request_init(self, request: "LlmRequest", **kwargs) -> None: - """Validate and reserve request-dependent eviction capacity.""" - self._validate_request_capacity(request) - self._reserve_eviction_capacity(request) + """Cover this request's largest possible scorer span.""" + 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) + required_source_tokens = prompt_length + min( + max_decode_tokens, self._selection_width_capacity + ) + if required_source_tokens <= self._score_token_capacity: + return + + # CuTe launches capture score buffers; retire an older, smaller span first. + if self._launch_score is not None: + self._compaction_done_event.synchronize() + + score_token_capacity = next_positive_power_of_2(max(required_source_tokens, 1024)) + score_token_capacity = min(score_token_capacity, int(manager.max_seq_len)) + score_tile_tokens = max(64, int(manager.tokens_per_block)) + score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens + self._build_score_runtime(score_token_capacity=score_token_capacity) def on_generation_step_begin(self, scheduled_batch: "ScheduledRequests", **kwargs) -> None: - """Remember the batch whose next-step KV capacity is already reserved.""" + """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: @@ -363,40 +383,7 @@ def on_generation_step_end(self, scheduled_batch: "ScheduledRequests", **kwargs) KVCacheManagerV2 must run first so capacity includes the written token and any rewind. """ - with nvtx_range_debug("triattention.generation_step_end", color="blue"): - self._evict_due_requests(scheduled_batch) - - # ---- request capacity ---- - - def _validate_request_capacity(self, request: "LlmRequest") -> None: - """Reject requests whose maximum pre-eviction peak cannot fit.""" - speculative_overshoot = int(self.kv_cache_manager.max_draft_len) - decode_capacity = min( - int(request.py_max_new_tokens), - self.budget + self.beta + speculative_overshoot, - ) - confirmed_capacity = int(request.py_prompt_len) + decode_capacity - checked = [(self.kv_cache_manager, self._protected_tail_capacity, "target")] - if self.draft_kv_cache_manager is not None: - checked.append( - (self.draft_kv_cache_manager, self._draft_protected_tail_capacity, "draft") - ) - for manager, protected_tail, label in checked: - required_capacity = confirmed_capacity + protected_tail - pool_capacity = manager.get_num_available_tokens( - token_num_upper_bound=confirmed_capacity, - max_num_draft_tokens=int(manager._kv_reserve_draft_tokens) + 1, - ) - table_capacity = manager.max_blocks_per_seq * manager.tokens_per_block - if confirmed_capacity > pool_capacity or required_capacity > table_capacity: - raise ValueError( - f"TriAttention {label} KV capacity is too small for eviction: " - f"request requires {required_capacity} tokens " - f"(prompt={request.py_prompt_len}, budget={self.budget}, " - f"beta={self.beta}, protected tail={protected_tail}), but the " - f"V2 pool covers {pool_capacity + protected_tail} tokens and " - f"its page table covers {table_capacity} tokens" - ) + self._evict_due_requests(scheduled_batch) # ---- eviction round ---- @@ -415,67 +402,66 @@ def _evict_due_requests( if prepared_batch is not None and prepared_batch is not scheduled_batch else set() ) - with nvtx_range("triattention.metadata", color="cyan"): - 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. + 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 - 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 + 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}" ) - 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 - eviction_requests.append( - _EvictionRequest( - request=request, - target_cache=target_cache, - draft_cache=draft_cache, - source_length=source_length, - target_tail_length=target_tail_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 - # Ungated NVTX: the due count in the message shows each round's size. - with nvtx_range( - f"triattention.evict_request_group reqs={len(eviction_requests)}", - color="purple", - ): - self._execute_eviction_round(eviction_requests) + 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). @@ -493,113 +479,107 @@ def _execute_eviction_round( # PyExecutor already joins its execution stream before the final # compression resource update, so the round can use caller current. try: - with nvtx_range_debug("triattention.page_table_stage", color="orange"): - 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, + 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_offsets( + manager, + request_ids, + self._block_offsets_host, + self._block_offsets_device, ) - # 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: + if draft_manager is not None: self._stage_block_offsets( - manager, + draft_manager, request_ids, - self._block_offsets_host, - self._block_offsets_device, - ) - if draft_manager is not None: - self._stage_block_offsets( - 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 + self._draft_block_offsets_host, + self._draft_block_offsets_device, ) - finally: - self._staging_reuse_event.record(stream) + 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" - with nvtx_range("triattention.score", color="blue"): - # In-place refresh: the compiled score launches captured these pointers. - _gather_mean_phase_kernel[(request_count,)]( - 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, - self._swa_rebase_delta, - NUM_FREQS=self._phase.num_freqs, - F_BLOCK=self._phase.frequency_block, - HAS_SWA=self._swa_destination_bases is not None, - num_warps=1, + # In-place refresh: the compiled score launches captured these pointers. + _gather_mean_phase_kernel[(request_count,)]( + 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, + self._swa_rebase_delta, + NUM_FREQS=self._phase.num_freqs, + F_BLOCK=self._phase.frequency_block, + HAS_SWA=self._swa_destination_bases is not None, + num_warps=1, + ) + 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, ) - 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_kernel[ - ( - request_count, - triton.cdiv(self._selection_width_capacity, 1024), - ) - ]( - gathered, - self._selection_scores_rows, + _fold_union_ranks_kernel[ + ( request_count, - TP_SIZE=int(self._union_tp_mapping.tp_size), - WIDTH=self._selection_width_capacity, - ) - with nvtx_range("triattention.select", color="yellow"): - 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, + triton.cdiv(self._selection_width_capacity, 1024), ) - self._select_top_tokens(request_count) - with nvtx_range("triattention.compact", color="purple"): - compact(self._compaction_params, request_count) + ]( + gathered, + self._selection_scores_rows, + request_count, + TP_SIZE=int(self._union_tp_mapping.tp_size), + WIDTH=self._selection_width_capacity, + ) + 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_top_tokens(request_count) + compact(self._compaction_params, request_count) finally: # Target and draft V2 managers share this execution stream. self._compaction_done_event.record(stream) @@ -677,34 +657,31 @@ def _select_top_tokens(self, request_count: int) -> None: ) def _resize_compacted_caches(self, eviction_requests: Sequence[_EvictionRequest]) -> None: - with nvtx_range("triattention.resize", color="red"): - for item in eviction_requests: - resized_capacity = ( - int(item.request.py_prompt_len) + self.budget + item.target_tail_length + for item in eviction_requests: + resized_capacity = ( + int(item.request.py_prompt_len) + self.budget + item.target_tail_length + ) + if not item.target_cache.resize(resized_capacity, None): + raise RuntimeError( + "Failed to resize compacted target KV cache for " + f"request {item.request.py_request_id} to " + f"{resized_capacity} tokens" ) - if not item.target_cache.resize(resized_capacity, None): - raise RuntimeError( - "Failed to resize compacted target KV cache for " - f"request {item.request.py_request_id} to " - f"{resized_capacity} tokens" - ) - if self.draft_kv_cache_manager is None: - return - # Target and draft retain the same selected tokens. - for item in eviction_requests: - resized_capacity = ( - int(item.request.py_prompt_len) - + self.budget - + self._draft_protected_tail_capacity + if self.draft_kv_cache_manager is None: + return + # Target and draft retain the same selected tokens. + for item in eviction_requests: + resized_capacity = ( + int(item.request.py_prompt_len) + self.budget + self._draft_protected_tail_capacity + ) + if not item.draft_cache.resize(resized_capacity, None): + raise RuntimeError( + "Failed to resize compacted draft KV cache for " + f"request {item.request.py_request_id} to " + f"{resized_capacity} tokens" ) - if not item.draft_cache.resize(resized_capacity, None): - raise RuntimeError( - "Failed to resize compacted draft KV cache for " - f"request {item.request.py_request_id} to " - f"{resized_capacity} tokens" - ) - # ---- persistent state + request capacity ---- + # ---- persistent state + score runtime ---- def _initialize_eviction_state(self) -> None: """Create manager-lifetime state once.""" @@ -778,7 +755,7 @@ def _allocate_metadata_buffers( *, num_freqs: int, ) -> None: - """Allocate manager-lifetime host staging and device metadata.""" + """Allocate fixed manager-lifetime host staging and device metadata.""" row_count = 6 request_capacity = self._request_capacity self._request_metadata_host = torch.empty( @@ -816,18 +793,18 @@ def _allocate_metadata_buffers( self._mean_sin = torch.empty_like(self._mean_cos) def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> None: - """Allocate fixed TopK inputs and outputs for the configured mode.""" + """Allocate fixed manager-lifetime TopK inputs and outputs.""" request_capacity = self._request_capacity selection_width = self._selection_width_capacity union = self.eviction_mode == "union" - if union: - self._selection_rows_per_request = 1 - elif self.eviction_mode == "per_head": - self._selection_rows_per_request = self._num_kv_heads - else: - self._selection_rows_per_request = self._num_layers * self._num_kv_heads + 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 * max(selection_width, self.budget) + selection_rect = selection_rows * selection_width if union: selection_rect = max( selection_rect, @@ -872,44 +849,12 @@ def _allocate_selection_buffers(self, device: torch.device, *, tp_size: int) -> ) self._kept_ordinal_rows = torch.empty_like(self._provisional_rows) - def _reserve_eviction_capacity(self, request: "LlmRequest") -> None: - """Reserve all request-dependent runtime capacity at admission.""" - first_evict_step = (self.budget // self.beta + 1) * self.beta - if int(request.py_max_new_tokens) < first_evict_step: - return - self._phase.reserve(int(request.py_prompt_len) + int(request.py_max_new_tokens) + 1) - decode_tokens = min(int(request.py_max_new_tokens), self._selection_width_capacity) - required_source_tokens = int(request.py_prompt_len) + decode_tokens - if required_source_tokens <= self._score_token_capacity: - return - - # A newly admitted request may require larger score state while work - # from an older request is still in flight. - if self._launch_score is not None: - self._compaction_done_event.synchronize() - - manager = self.kv_cache_manager - # Bucket the largest admitted source rather than eagerly reserving - # max_seq_len; score scratch is the dominant allocation. - score_token_capacity = next_positive_power_of_2(max(required_source_tokens, 1024)) - score_token_capacity = min( - score_token_capacity, max(int(manager.max_seq_len), required_source_tokens) - ) - # The bucket capacity must be tile-aligned (mis-tiling stripes the - # score scratch silently); the ceiling division constructs that fact. - score_tile_tokens = max(64, int(manager.tokens_per_block)) - score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens - - self._build_eviction_capacity( - score_token_capacity=score_token_capacity, - ) - - def _build_eviction_capacity( + def _build_score_runtime( self, *, score_token_capacity: int, ) -> None: - """Build and publish all state bound to one score-token capacity.""" + """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 @@ -985,19 +930,11 @@ def _local_score_calibration( self, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: global_layers = self._global_layers - num_layers = len(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)}" ) - if global_layers == list(range(global_layers[0], global_layers[0] + num_layers)): - layer_slice = slice(global_layers[0], global_layers[0] + num_layers) - return ( - self._calibration_q_real[layer_slice], - self._calibration_q_imag[layer_slice], - self._calibration_mlr_coef[layer_slice], - ) layer_ids = torch.as_tensor( global_layers, device=self._calibration_q_real.device, diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 72865589ada1..2f7dd6c451d1 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -639,12 +639,9 @@ def __init__( if (config.kv_cache_compression_config is not None and config.kv_cache_compression_config. kv_cache_compression_mode.is_eviction_method()): - # Fused RoPE derives positions from the KV length, which eviction shortens; stay unfused. logger.warning_once( - "disable rope_fusion for KV-cache compression " - f"({config.kv_cache_compression_config.algorithm}): " - "rotary positions must come from logical position_ids, " - "not the compression-shortened KV length.", + "KV-cache eviction changes the physical cache length; " + "setting rope_fusion=False.", key="disable_rope_fusion_for_kv_cache_compression") self.rope_fusion = False diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 4444dd0b4a83..3f4d1b36bda0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2113,7 +2113,6 @@ def _create_kv_cache_manager( def validate_kv_cache_compression_with_spec( config: KvCacheCompressionConfig, spec_config: Optional[SpeculativeConfig], - draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, ) -> None: """Reject speculative setups the compression method cannot run with.""" if spec_config is None: @@ -2129,19 +2128,6 @@ def validate_kv_cache_compression_with_spec( raise ValueError( f"KV-cache compression does not support speculative decoding " f"mode {mode.name}; use one-model MTP or EAGLE3") - if not spec_config.is_linear_tree: - raise ValueError( - "KV-cache compression requires linear speculative decoding") - if spec_config.draft_len_schedule is not None: - raise ValueError( - "KV-cache compression requires a fixed speculative draft length") - if draft_kv_cache_manager is None: - raise ValueError( - "KV-cache compression requires a separate draft KV cache") - if any(window is not None - for window in draft_kv_cache_manager.max_attention_window_vec): - raise ValueError( - "KV-cache compression requires full-attention draft KV") def create_kv_cache_compression_manager( @@ -2158,6 +2144,7 @@ def create_kv_cache_compression_manager( ``validate_kv_cache_compression_with_spec``. """ if config.algorithm == "triattention": + # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import \ TriAttention @@ -2422,8 +2409,7 @@ def create_py_executor_instance( 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) + spec_config) compression_manager = create_kv_cache_compression_manager( kv_cache_compression_config, kv_cache_manager, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 80951a2f974b..80612b12bc50 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -58,9 +58,7 @@ CapacitySchedulerPolicy as _CapacitySchedulerPolicy, ContextChunkingPolicy as _ContextChunkingPolicy, DecodingConfig, - DecodingMode, DynamicBatchConfig as _DynamicBatchConfig, - EagleConfig as _EagleConfig, ExecutorConfig as _ExecutorConfig, ExtendedRuntimePerfKnobConfig as _ExtendedRuntimePerfKnobConfig, KvCacheConfig as _KvCacheConfig, @@ -3592,6 +3590,12 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): "load.") +KvCacheCompressionConfigType: TypeAlias = Annotated[ + Union[TriAttentionKvCacheCompressionConfig], + Field(discriminator="algorithm"), +] + + @PybindMirror.mirror_pybind_fields(_AgentTreeConfig) class AgentTreeConfig(StrictBaseModel, PybindMirror): """Configuration for agent tree scheduling. @@ -4341,15 +4345,11 @@ class BaseLlmArgs(StrictBaseModel): status="prototype") # KV cache compression config (separate from sparse attention: changes which - # KV is stored, not the attention computation). Dispatch is by the - # ``algorithm`` tag; grow this into a discriminated union when a second - # algorithm lands. - kv_cache_compression_config: Optional[ - TriAttentionKvCacheCompressionConfig] = Field( - default=None, - description= - "KV-cache compression config; None disables compression.", - status="prototype") + # 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") # Speculative decoding parameters speculative_config: Optional[SpeculativeConfig] = Field( diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 8ed4a1cc8d39..55ff2fb183ca 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -544,7 +544,7 @@ def make_cute_buffers( 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_eviction_capacity(score_token_capacity=seq_len) + manager._build_score_runtime(score_token_capacity=seq_len) return manager 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 index 63ff153625b1..27b93b179bc9 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -74,10 +74,6 @@ class Boom(RuntimeError): [ ("callsite_dflash", "one-model MTP or EAGLE3"), ("union_only_per_head", "union"), - ("dynamic_tree", "linear speculative"), - ("dynamic_draft_length", "fixed speculative draft length"), - ("missing_draft", "separate draft"), - ("sliding_draft", "full-attention draft"), ], ) def test_speculative_admission_gates_raise(gate, match): @@ -86,32 +82,17 @@ def test_speculative_admission_gates_raise(gate, match): if gate == "callsite_dflash": spec_config = DFlashDecodingConfig(max_draft_len=3) - elif gate == "dynamic_tree": - spec_config = MTPDecodingConfig( - max_draft_len=2, - use_dynamic_tree=True, - dynamic_tree_max_topK=2, - ) - elif gate == "dynamic_draft_length": - spec_config = MTPDecodingConfig( - max_draft_len=1, - draft_len_schedule={1: 1, 8: 0}, - ) 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", ) - draft_manager = None if gate == "missing_draft" else _make_fake_v2(is_draft=True) - if gate == "sliding_draft": - draft_manager.max_attention_window_vec = [128] with pytest.raises(ValueError, match=match): validate_kv_cache_compression_with_spec( config, spec_config, - draft_manager, ) @@ -181,6 +162,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): def test_request_admission_reserves_score_high_watermark(): manager = _make_triattention(budget=128, beta=64) + assert manager._selection_width_capacity == 193 manager._phase = mock.Mock() manager._selection_width_capacity = 260 manager._score_token_capacity = 0 @@ -192,7 +174,7 @@ def publish_score_state(*, score_token_capacity): manager._score_token_capacity = score_token_capacity manager._launch_score = object() - manager._build_eviction_capacity = mock.Mock(side_effect=publish_score_state) + 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), @@ -200,9 +182,9 @@ def publish_score_state(*, score_token_capacity): ] for request in requests: - manager._reserve_eviction_capacity(request) + manager.on_request_init(request) - assert manager._build_eviction_capacity.call_args_list == [ + assert manager._build_score_runtime.call_args_list == [ mock.call(score_token_capacity=1024), mock.call(score_token_capacity=2048), ] @@ -221,11 +203,11 @@ def test_request_admission_aligns_clamped_score_bucket_to_tile(): 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_eviction_capacity = mock.Mock() + manager._build_score_runtime = mock.Mock() - manager._reserve_eviction_capacity(_make_request(1, py_prompt_len=1025, py_max_new_tokens=192)) + manager.on_request_init(_make_request(1, py_prompt_len=850, py_max_new_tokens=200)) - manager._build_eviction_capacity.assert_called_once_with(score_token_capacity=1280) + manager._build_score_runtime.assert_called_once_with(score_token_capacity=1152) manager._compaction_done_event.synchronize.assert_not_called() diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index e0d04a6c5b0b..bfee69b6109e 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -79,40 +79,6 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): class TestTriAttentionClass: - def test_request_init_validates_and_reserves_capacity(self): - manager = _make_fake_v2() - manager.num_extra_kv_tokens = 4 - manager._kv_reserve_draft_tokens = 4 - with mock.patch.object(TriAttention, "_initialize_eviction_state"): - triattention = TriAttention(_make_tri_config(budget=8), manager) - with ( - mock.patch.object(triattention, "_validate_request_capacity") as validate, - mock.patch.object(triattention, "_reserve_eviction_capacity") as reserve, - ): - request_11 = _make_request(11) - request_12 = _make_request(12) - triattention.on_request_init(request_11) - triattention.on_request_init(request_12) - - assert triattention.adjusts_generation_kv_length is True - assert manager.kv_compression_manages_history - assert validate.call_args_list == [mock.call(request_11), mock.call(request_12)] - assert reserve.call_args_list == [mock.call(request_11), mock.call(request_12)] - - def test_capacity_guard_uses_maximum_steady_eviction_peak(self): - manager = _make_triattention(budget=10, beta=8) - manager.kv_cache_manager.get_num_available_tokens = mock.Mock(return_value=17) - - with pytest.raises(ValueError, match="requires 19 tokens"): - manager._validate_request_capacity( - _make_request(7, py_prompt_len=0, py_max_new_tokens=100) - ) - - manager.kv_cache_manager.get_num_available_tokens.assert_called_once_with( - token_num_upper_bound=18, - max_num_draft_tokens=1, - ) - def test_resolve_accepts_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() mgr.calibration_path = flat_calibration_pt @@ -287,6 +253,7 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft ) 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): @@ -312,13 +279,24 @@ def test_suspended_cache_defers_that_request_pre_launch(self): 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 == 128 + 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]) @@ -364,8 +342,10 @@ def test_overlap_tail_is_excluded_from_selection_and_compacted(self, accepted): 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. - physical_confirmed = 6172 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, @@ -377,8 +357,8 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): request = _make_request( 7, py_prompt_len=1024, - py_num_compressed_tokens=100, - max_beam_num_tokens=999999, + py_num_compressed_tokens=compressed_tokens, + max_beam_num_tokens=physical_confirmed + compressed_tokens + 1, py_draft_tokens=[1, 2, 3, 4], ) @@ -388,7 +368,7 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): eviction_requests = internals.execute.call_args.args[0] assert eviction_requests[0].source_length == physical_confirmed assert request.py_num_compressed_tokens == ( - 100 + physical_confirmed - 1024 - manager.budget + compressed_tokens + physical_confirmed - 1024 - manager.budget ) cache.resize.assert_called_once_with(1024 + manager.budget, None) @@ -418,7 +398,6 @@ def test_one_model_draft_co_compression_is_accepted(self, spec_mode): validate_kv_cache_compression_with_spec( _make_tri_config(budget=8), spec_config, - draft_manager, ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 558a57c70208..f7acf482ab4f 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2714,6 +2714,31 @@ 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 + + class TestSkipSoftmaxAttentionConfig: """Test LLM Skip Softmax Attention config behavior.""" From 1c309dd9307c648adc505a00d0ba8e50537afe24 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 28 Jul 2026 04:42:27 -0700 Subject: [PATCH 167/178] fix: align TriAttention selection buffers Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 9 ++++++++- .../test_triattention_draft_cocompaction.py | 2 +- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 1d6fb03bd1f0..365582969eeb 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -61,6 +61,9 @@ _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.""" @@ -178,7 +181,11 @@ def __init__( 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. - self._selection_width_capacity = self.budget + self.beta + 2 * max_draft_tokens + 1 + 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, 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 index 27b93b179bc9..3e76aca017a1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -162,7 +162,7 @@ def test_compressed_count_is_monotone_and_tracks_confirmed_length(): def test_request_admission_reserves_score_high_watermark(): manager = _make_triattention(budget=128, beta=64) - assert manager._selection_width_capacity == 193 + assert manager._selection_width_capacity == 256 manager._phase = mock.Mock() manager._selection_width_capacity = 260 manager._score_token_capacity = 0 From 3e3fc88f9cbd81344bb4bc3a9a36a5911230c9d0 Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 28 Jul 2026 05:39:44 -0700 Subject: [PATCH 168/178] style: use relative TriAttention imports Signed-off-by: tianruih --- .../triattention/triattention.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 365582969eeb..5f68311d704b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -30,17 +30,17 @@ from transformers import AutoConfig from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS -from tensorrt_llm._torch.distributed import allgather -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheCompressionManager -from tensorrt_llm._torch.utils import next_positive_power_of_2 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 ( @@ -51,10 +51,11 @@ ) if TYPE_CHECKING: - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest - from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests 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"}) From 66fe4e8c61fed37024c1a314e156e668795dd50d Mon Sep 17 00:00:00 2001 From: tianruih Date: Tue, 28 Jul 2026 05:45:05 -0700 Subject: [PATCH 169/178] style: make TriAttention imports consistently relative Signed-off-by: tianruih --- .../kv_cache_compression/triattention/triattention.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 5f68311d704b..594c7516a1e6 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -30,12 +30,11 @@ 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 ( +from ...._utils import prefer_pinned +from ....bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( copy_batch_block_offsets_to_device, ) -from tensorrt_llm.logger import logger - +from ....logger import logger from ...distributed import allgather from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from ...pyexecutor.llm_request import LlmRequestState @@ -51,8 +50,7 @@ ) if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig - + from ....llmapi.llm_args import TriAttentionKvCacheCompressionConfig from ...pyexecutor.llm_request import LlmRequest from ...pyexecutor.scheduler import ScheduledRequests From dadffc8298521d3aec67e988c42ad38e13e74293 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Tue, 28 Jul 2026 07:36:29 -0700 Subject: [PATCH 170/178] [None][refactor] Simplify TriAttention runtime flow Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- examples/triattention/README.md | 20 +- .../_torch/kv_cache_compression/interface.py | 2 +- .../triattention/triattention.py | 282 +++++++++--------- tensorrt_llm/_torch/pyexecutor/_util.py | 13 +- .../_torch/kv_cache_compression/conftest.py | 4 +- .../test_triattention_draft_cocompaction.py | 16 +- .../test_triattention_pipeline.py | 58 ++-- .../test_triattention_selection_compaction.py | 4 +- .../_torch/thop/parallel/test_indexer_topk.py | 127 -------- 9 files changed, 213 insertions(+), 313 deletions(-) diff --git a/examples/triattention/README.md b/examples/triattention/README.md index 5c2adc7e25c7..1c3b5c237855 100644 --- a/examples/triattention/README.md +++ b/examples/triattention/README.md @@ -1,6 +1,6 @@ # TriAttention KV-Cache Compression -This document describes enabling TriAttention KV-cache compression in TensorRT LLM. +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. @@ -10,23 +10,23 @@ For technical details see the paper [TriAttention](https://arxiv.org/abs/2604.04 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 it and converts it to its runtime schema at the first request. -2. **Periodic eviction (Stage during generation).** Every `beta` confirmed generation tokens, once a sequence is over budget, TriAttention scores the whole cache, selects `budget` tokens to keep (the prompt tokens are preserved on top of the budget), and physically compacts the KV cache down to the kept set. A speculative iteration may confirm multiple tokens; crossing multiple periods in one update is coalesced into one eviction. +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. +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 -* GPU Compute Capability >= 10.0 (Blackwell or newer) +* NVIDIA B200 (SM100; the current validated target) * Paged KV Cache (`KVCacheManagerV2`) -* Tensor Parallel * PyTorch backend **Notes:** 1. TriAttention requires `enable_block_reuse=False` in the KV-cache configuration — the eviction physically rewrites stored keys, which is incompatible with block reuse. The construction step rejects a cache manager that has block reuse enabled. 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. Requires full-attention KVCacheManagerV2 lifecycles; attention-DP, disaggregated serving, native SWA/VSWA/SSM pools, and MLA caches are unsupported. +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 @@ -50,7 +50,7 @@ python3 scripts/calibrate.py \ --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.) +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 @@ -82,7 +82,7 @@ llm = LLM( kv_cache_config=kv_config, ) -# 4. Generate +# 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) @@ -121,5 +121,5 @@ trtllm-eval --model --config config.yaml longbench_v2 --max_outp * `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. +* **`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 index 6e36001d7ce8..a1df7a9d5b93 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/interface.py +++ b/tensorrt_llm/_torch/kv_cache_compression/interface.py @@ -12,8 +12,8 @@ class KvCacheCompressionMode(IntEnum): ``is_*`` predicates instead of comparing strings. """ - TRIATTENTION = auto() NONE = auto() + TRIATTENTION = auto() def is_eviction_method(self) -> bool: """Return whether this mode physically evicts cached tokens.""" diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 594c7516a1e6..1e2880834195 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -74,22 +74,30 @@ class _EvictionRequest(NamedTuple): target_tail_length: int -def _allocate_block_offset_staging( +_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 a V2 page-table snapshot for the scorer/compactor token span.""" - tokens_per_block = int(anchor_pool.shape[3]) - page_count = (token_capacity + tokens_per_block - 1) // tokens_per_block - max_source_blocks = int(manager.host_kv_cache_block_offsets.shape[-1]) - block_capacity = min((page_count + 3) // 4 * 4, max_source_blocks) - shape = (int(manager.num_pools), request_capacity, 2, block_capacity) - host = torch.empty(shape, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned()) - device_table = torch.empty(shape, dtype=torch.int32, device=anchor_pool.device) - return host, device_table + """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 @@ -267,93 +275,93 @@ def _resolve_attention_layers(self) -> Tuple[List[int], List[int], Optional[int] return (dense_layers, swa_layers, window_size) def _load_calibration(self) -> None: - calibration = self._resolve_calibration() - self._freq_scale_sq = calibration["freq_scale_sq"].to(dtype=torch.float32) - self._omega = calibration["omega"] - # Pre-split query stats + MLR coefficient, shapes [L, H, F]. - e_q = calibration["E_q"] - 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 = ( - calibration["E_q_norm"].to(torch.float32) - e_q.abs().to(torch.float32) - ).contiguous() - - def _resolve_calibration(self) -> Dict[str, torch.Tensor]: - """Load the calibration file, converting the official layout if needed.""" raw = torch.load(self.calibration_path, map_location="cpu", weights_only=False) if isinstance(raw, dict) and _REQUIRED_CALIBRATION_KEYS <= set(raw): - return raw - if isinstance(raw, dict) and {"metadata", "stats"} <= set(raw): - return self._convert_official_calibration(raw) - got = sorted(raw.keys()) 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}." - ) - - def _convert_official_calibration(self, raw) -> Dict[str, torch.Tensor]: - """Convert the official calibration format to the runtime schema.""" - stats = raw["stats"] - meta = raw["metadata"] - if "sampled_heads" in meta: - heads = [(int(a), int(b)) for a, b in meta["sampled_heads"]] + 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: - heads = [ - (int(k[len("layer") : k.index("_head")]), int(k[k.index("_head") + len("_head") :])) - for k in stats - ] - num_layers = max(layer for layer, _ in heads) + 1 - num_heads = max(h for _, h 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, h in heads: - s = stats[f"layer{layer:02d}_head{h:02d}"] - E_q[layer, h] = torch.complex(s["q_mean_real"].float(), s["q_mean_imag"].float()) - E_q_norm[layer, h] = s["q_abs_mean"].float() - omega, freq_scale_sq = self._rope_tables(freq_count) - calib = { - "E_q": E_q, - "E_q_norm": E_q_norm, - "omega": omega, - "freq_scale_sq": freq_scale_sq, - } - logger.info( - f"TriAttention: converted official calibration {self.calibration_path}" - f" -> E_q[L={num_layers}, H={num_heads}, F={freq_count}]" - ) - return calib - - def _rope_tables(self, freq_count: int): - """Derive the RoPE frequency tables from the model config.""" - 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()): + got = sorted(raw) if isinstance(raw, dict) else type(raw).__name__ raise ValueError( - f"TriAttention does not support per-layer-type rope parameters ({self.model_path})" + f"Unrecognized calibration at {self.calibration_path}: expected the " + f"official {{metadata, stats}} layout or " + f"{sorted(_REQUIRED_CALIBRATION_KEYS)}; got {got}." ) - rope_type = rope_params["rope_type"] - if rope_type == "default": - # "default" has no ROPE_INIT_FUNCTIONS entry; the analytic formula is its definition. - 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() - 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() - scale_sq = float(attention_factor) ** 2 - return omega, torch.full((freq_count,), scale_sq, dtype=torch.float32) + + 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: - """Cover this request's largest possible scorer span.""" + """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( @@ -364,21 +372,21 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: if max_decode_tokens < first_evict_step: return self._phase.reserve(prompt_length + max_decode_tokens + 1) - required_source_tokens = prompt_length + min( + max_source_tokens = prompt_length + min( max_decode_tokens, self._selection_width_capacity ) - if required_source_tokens <= self._score_token_capacity: + if max_source_tokens <= self._score_token_capacity: return - # CuTe launches capture score buffers; retire an older, smaller span first. + # CuTe launches capture score buffers; retire the old capacity before replacing it. if self._launch_score is not None: self._compaction_done_event.synchronize() - score_token_capacity = next_positive_power_of_2(max(required_source_tokens, 1024)) - score_token_capacity = min(score_token_capacity, int(manager.max_seq_len)) - score_tile_tokens = max(64, int(manager.tokens_per_block)) - score_token_capacity = -(-score_token_capacity // score_tile_tokens) * score_tile_tokens - self._build_score_runtime(score_token_capacity=score_token_capacity) + 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.""" @@ -513,14 +521,14 @@ def _execute_eviction_round( # their unused request rows explicit no-ops. host_table[:3, len(eviction_requests) :] = 0 try: - self._stage_block_offsets( + 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_offsets( + self._stage_block_offset_snapshot( draft_manager, request_ids, self._draft_block_offsets_host, @@ -584,7 +592,7 @@ def _execute_eviction_round( per_layer=self.eviction_mode == "per_layer_perhead", normalize_scores=self.normalize_scores, ) - self._select_top_tokens(request_count) + self._select_kept_ordinals(request_count) compact(self._compaction_params, request_count) finally: # Target and draft V2 managers share this execution stream. @@ -598,26 +606,26 @@ def _compute_compaction_move_offsets( ) -> Tuple[List[int], Optional[List[int]], Optional[List[int]]]: """Build padded cumulative dense, SWA, and draft move offsets.""" - def padded_offsets(moves_per_request: List[int]) -> List[int]: + def cumulative_offsets(move_counts: List[int]) -> List[int]: offsets = [0] - for moves in moves_per_request: - offsets.append(offsets[-1] + moves) - offsets.extend(offsets[-1:] * (self._request_capacity - len(moves_per_request))) + 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 = padded_offsets([self.budget + tail for tail in tails]) - swa = None + dense_offsets = cumulative_offsets([self.budget + tail for tail in tails]) + swa_offsets = None if self._swa_window is not None: - swa = padded_offsets([self._swa_window + tail for tail in tails]) - draft = 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 = padded_offsets( + draft_offsets = cumulative_offsets( [self.budget + self._draft_protected_tail_capacity] * len(eviction_requests) ) - return dense, swa, draft + return dense_offsets, swa_offsets, draft_offsets - def _stage_block_offsets( + def _stage_block_offset_snapshot( self, manager: KVCacheManagerV2, request_ids: List[int], @@ -640,7 +648,7 @@ def _stage_block_offsets( torch.cuda.current_stream(device_block_offsets.device).cuda_stream, ) - def _select_top_tokens(self, request_count: int) -> None: + 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. @@ -664,27 +672,26 @@ def _select_top_tokens(self, request_count: int) -> None: def _resize_compacted_caches(self, eviction_requests: Sequence[_EvictionRequest]) -> None: for item in eviction_requests: - resized_capacity = ( + target_capacity = ( int(item.request.py_prompt_len) + self.budget + item.target_tail_length ) - if not item.target_cache.resize(resized_capacity, None): + 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"{resized_capacity} tokens" + f"{target_capacity} tokens" ) if self.draft_kv_cache_manager is None: return - # Target and draft retain the same selected tokens. for item in eviction_requests: - resized_capacity = ( + draft_capacity = ( int(item.request.py_prompt_len) + self.budget + self._draft_protected_tail_capacity ) - if not item.draft_cache.resize(resized_capacity, None): + 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"{resized_capacity} tokens" + f"{draft_capacity} tokens" ) # ---- persistent state + score runtime ---- @@ -709,7 +716,20 @@ def _initialize_eviction_state(self) -> None: self._draft_block_offsets_host = None self._draft_block_offsets_device = None - q_real, q_imag, mlr_coef = self._local_score_calibration() + 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: @@ -864,7 +884,7 @@ def _build_score_runtime( 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_staging( + block_offsets_host, block_offsets_device = _allocate_block_offset_snapshot( self.kv_cache_manager, anchor_pool, request_capacity=request_capacity, @@ -874,7 +894,7 @@ def _build_score_runtime( 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_staging( + draft_block_offsets_host, draft_block_offsets_device = _allocate_block_offset_snapshot( self.draft_kv_cache_manager, draft_anchor_pool, request_capacity=request_capacity, @@ -932,26 +952,6 @@ def _build_score_runtime( self._launch_score = launch_score self._compaction_params = tuple(compaction_params) - def _local_score_calibration( - self, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - 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, - ) - return ( - self._calibration_q_real.index_select(0, layer_ids), - self._calibration_q_imag.index_select(0, layer_ids), - self._calibration_mlr_coef.index_select(0, layer_ids), - ) - 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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index be5bd31a43eb..19158248deb0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2268,15 +2268,19 @@ def validate_kv_cache_compression_with_spec( return if not config.kv_cache_compression_mode.is_eviction_method(): return - if config.eviction_mode != "union": + if config.algorithm != "triattention": raise ValueError( - "KV-cache compression with speculative decoding requires " - "eviction_mode='union'") + f"KV-cache compression algorithm {config.algorithm!r} has no " + "speculative-decoding compatibility contract") 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 does not support speculative decoding " f"mode {mode.name}; use one-model MTP or EAGLE3") + if getattr(config, "eviction_mode", None) != "union": + raise ValueError( + "KV-cache compression with speculative decoding requires " + "eviction_mode='union'") def create_kv_cache_compression_manager( @@ -2294,8 +2298,7 @@ def create_kv_cache_compression_manager( """ if config.algorithm == "triattention": # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import \ - TriAttention + from ..kv_cache_compression.triattention.triattention import TriAttention return TriAttention( config, diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index dffcddaeb0c8..08125abff9af 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -1,7 +1,7 @@ # 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 @@ -242,7 +242,7 @@ def make_bare_staging(device, *, max_requests, staged_blocks_per_seq): def make_staging_manager(host_table, gather, manager_stream, *, num_slots=1): - """The manager surface ``_stage_block_offsets`` consumes.""" + """The manager surface ``_stage_block_offset_snapshot`` consumes.""" return SimpleNamespace( host_kv_cache_block_offsets=host_table, kv_factor=2, 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 index 3e76aca017a1..225e311bf3b3 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -55,7 +55,7 @@ class Boom(RuntimeError): mock.patch.object( torch.cuda, "current_stream", return_value=compute_stream ) as current_stream, - mock.patch.object(tri, "_stage_block_offsets") as stage, + mock.patch.object(tri, "_stage_block_offset_snapshot") as stage, ): with pytest.raises(Boom): tri._execute_eviction_round(eviction_requests) @@ -211,26 +211,26 @@ def test_request_admission_aligns_clamped_score_bucket_to_tile(): manager._compaction_done_event.synchronize.assert_not_called() -def test_staged_block_width_clamps_to_manager_source_width(): - """Clamp staged block width to the live V2 source-table width.""" +def test_block_offset_snapshot_width_is_aligned_and_capped(): from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( - _allocate_block_offset_staging, + _allocate_block_offset_snapshot, ) anchor_pool = torch.empty(1, 2, 1, 32, 4) manager = SimpleNamespace( num_pools=1, - host_kv_cache_block_offsets=torch.empty(1, 2, 2, 4, dtype=torch.int32), + tokens_per_block=32, + max_blocks_per_seq=4, ) - host, device_table = _allocate_block_offset_staging( + 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.host_kv_cache_block_offsets = torch.empty(1, 2, 2, 64, dtype=torch.int32) - host, device_table = _allocate_block_offset_staging( + manager.max_blocks_per_seq = 64 + host, device_table = _allocate_block_offset_snapshot( manager, anchor_pool, request_capacity=2, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index bfee69b6109e..46eeef777745 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -79,15 +79,19 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): class TestTriAttentionClass: - def test_resolve_accepts_flat_pt(self, flat_calibration_pt): + def test_loads_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() mgr.calibration_path = flat_calibration_pt mgr.model_path = None - loaded = mgr._resolve_calibration() - for key in ("E_q", "E_q_norm", "omega", "freq_scale_sq"): - assert key in loaded + mgr._load_calibration() - def test_resolve_converts_official_layout(self, tmp_path): + 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. @@ -109,23 +113,24 @@ def test_resolve_converts_official_layout(self, tmp_path): config = _make_hf_config(rope_parameters={"rope_type": "default", "rope_theta": 10000.0}) with mock.patch("transformers.AutoConfig.from_pretrained", return_value=config): - converted = mgr._resolve_calibration() + mgr._load_calibration() - assert set(converted) == {"E_q", "E_q_norm", "omega", "freq_scale_sq"} - assert converted["E_q"].shape == (num_layers, num_heads, freq_count) + 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( - converted["E_q"][1, 0].cpu(), - torch.complex(torch.full((freq_count,), 10.0), torch.full((freq_count,), 1.0)), + mgr._calibration_q_imag[1, 0].cpu(), torch.full((freq_count,), 1.0) ) torch.testing.assert_close( - converted["E_q_norm"][1, 1].cpu(), torch.full((freq_count,), 3.0) + mgr._calibration_mlr_coef[1, 1].cpu(), torch.full((freq_count,), -8.0) ) - assert converted["omega"].numel() == freq_count + assert mgr._omega.numel() == freq_count idx = torch.arange(0, 2 * freq_count, 2, dtype=torch.float32) torch.testing.assert_close( - converted["omega"].cpu(), 1.0 / (10000.0 ** (idx / (2 * freq_count))) + mgr._omega.cpu(), 1.0 / (10000.0 ** (idx / (2 * freq_count))) ) - assert torch.equal(converted["freq_scale_sq"].cpu(), torch.ones(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 @@ -167,15 +172,34 @@ def config_dir(name, body): ) 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 - omega, freq_scale_sq = mgr._rope_tables(freq_count) + 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 - omega_yarn, freq_scale_sq_yarn = mgr._rope_tables(freq_count) + 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)) @@ -437,7 +461,7 @@ def gather_k_block_offsets(source, destination, request_ids, num_blocks): def stage_once(): # Raises on any staging failure; success returns None. with torch.cuda.stream(current_stream): - staging._stage_block_offsets( + staging._stage_block_offset_snapshot( manager, [7], staging._block_offsets_host, 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 index ffa2ff11956a..1b847a676dc1 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -75,7 +75,7 @@ def _select_per_head(tri, scores, *, normalize_scores): per_layer=tri.eviction_mode == "per_layer_perhead", normalize_scores=normalize_scores, ) - tri._select_top_tokens(tri._request_capacity) + tri._select_kept_ordinals(tri._request_capacity) def _stable_topk(row: torch.Tensor, width: int, keep_count: int) -> torch.Tensor: @@ -201,7 +201,7 @@ def test_union_eager_cuda_resolves_heavy_ties_and_ragged_lengths(keep_count, wid torch.tensor([prompt_len] * request_count, dtype=torch.int32, device=device) ) tri._selection_scores_rows.copy_(scores.amax(dim=1)) - tri._select_top_tokens(tri._request_capacity) + tri._select_kept_ordinals(tri._request_capacity) actual = tri._kept_ordinal_rows.cpu() combined = scores.amax(dim=1).cpu() diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py index a26c1f24bd1b..3634336aa961 100644 --- a/tests/unittest/_torch/thop/parallel/test_indexer_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk.py @@ -1492,133 +1492,6 @@ def run_fn(logits, seq_lens): ) -def _large_k_indexer_case(top_k: int, generation: int): - """Build a deterministic, tie-heavy input for the direct indexer op.""" - batch_size = 2 - next_n = 2 - num_rows = batch_size * next_n - num_tokens = ((top_k + 521 + 31) // 32) * 32 - columns = torch.arange(num_tokens, dtype=torch.int64, device="cuda") - rows = torch.arange(num_rows, dtype=torch.int64, device="cuda")[:, None] - logits = ((columns[None, :] * 17 + rows * 29 + generation * 43) % 257).to(torch.float32) - seq_lens = torch.tensor( - [num_tokens - 3 - generation, num_tokens - 17 + generation], - dtype=torch.int32, - device="cuda", - ) - return logits.contiguous(), seq_lens, next_n - - -def _assert_large_k_indexer_result( - logits: torch.Tensor, - seq_lens: torch.Tensor, - indices: torch.Tensor, - top_k: int, - next_n: int, -): - """Compare indices with a tie-safe, independent torch.topk oracle.""" - num_rows = logits.shape[0] - row_offsets = torch.arange(num_rows, device="cuda") % next_n - row_seq_lens = seq_lens.repeat_interleave(next_n) - next_n + row_offsets + 1 - - assert indices.shape == (num_rows, top_k) - assert indices.dtype == torch.int32 - assert indices.is_contiguous() - assert torch.all(indices >= 0) - assert torch.all(indices < row_seq_lens[:, None]) - - sorted_indices = indices.sort(dim=1).values - assert torch.all(sorted_indices[:, 1:] > sorted_indices[:, :-1]) - - selected_values = logits.gather(1, indices.to(torch.int64)).sort(dim=1, descending=True).values - oracle_values = [] - for row in range(num_rows): - row_seq_len = int(row_seq_lens[row].item()) - oracle_values.append( - torch.topk(logits[row, :row_seq_len], top_k).values.sort(descending=True).values - ) - oracle_values = torch.stack(oracle_values) - torch.testing.assert_close(selected_values, oracle_values, rtol=0.0, atol=0.0) - - -@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") -@skip_pre_blackwell -@pytest.mark.parametrize("index_topk", [4096, 4097, 8192]) -def test_cute_dsl_indexer_topk_decode_large_k_cuda_graph(index_topk): - """Validate eager and captured direct-indexer execution above k=2048.""" - eager_logits, eager_seq_lens, next_n = _large_k_indexer_case(index_topk, 0) - num_rows = eager_logits.shape[0] - eager_indices = torch.empty(num_rows, index_topk, dtype=torch.int32, device="cuda") - eager_pointer = eager_indices.data_ptr() - torch.ops.trtllm.cute_dsl_indexer_topk_decode( - input_values=eager_logits, - seq_lens=eager_seq_lens, - output_indices=eager_indices, - top_k=index_topk, - next_n=next_n, - num_copy_bits=256, - ) - torch.cuda.synchronize() - assert eager_indices.data_ptr() == eager_pointer - _assert_large_k_indexer_result(eager_logits, eager_seq_lens, eager_indices, index_topk, next_n) - - graph_logits, graph_seq_lens, _ = _large_k_indexer_case(index_topk, 1) - graph_indices = torch.empty_like(eager_indices) - graph_pointers = ( - graph_logits.data_ptr(), - graph_seq_lens.data_ptr(), - graph_indices.data_ptr(), - ) - - warmup_stream = torch.cuda.Stream() - warmup_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(warmup_stream): - torch.ops.trtllm.cute_dsl_indexer_topk_decode( - input_values=graph_logits, - seq_lens=graph_seq_lens, - output_indices=graph_indices, - top_k=index_topk, - next_n=next_n, - num_copy_bits=256, - ) - torch.cuda.current_stream().wait_stream(warmup_stream) - torch.cuda.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - torch.ops.trtllm.cute_dsl_indexer_topk_decode( - input_values=graph_logits, - seq_lens=graph_seq_lens, - output_indices=graph_indices, - top_k=index_topk, - next_n=next_n, - num_copy_bits=256, - ) - torch.cuda.synchronize() - assert graph_pointers == ( - graph_logits.data_ptr(), - graph_seq_lens.data_ptr(), - graph_indices.data_ptr(), - ) - _assert_large_k_indexer_result(graph_logits, graph_seq_lens, graph_indices, index_topk, next_n) - - for generation in (2, 3): - replay_logits, replay_seq_lens, _ = _large_k_indexer_case(index_topk, generation) - graph_logits.copy_(replay_logits) - graph_seq_lens.copy_(replay_seq_lens) - graph_indices.fill_(-1) - graph.replay() - torch.cuda.synchronize() - assert graph_pointers == ( - graph_logits.data_ptr(), - graph_seq_lens.data_ptr(), - graph_indices.data_ptr(), - ) - _assert_large_k_indexer_result( - graph_logits, graph_seq_lens, graph_indices, index_topk, next_n - ) - - @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="CuTE DSL not available") @skip_pre_blackwell @pytest.mark.parametrize("batch_size", [1, 16, 256]) From e72af06681a752245c41ce5c85f73f1bf4fbdf38 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Tue, 28 Jul 2026 07:50:40 -0700 Subject: [PATCH 171/178] [None][chore] Format TriAttention sources Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- .../triattention/triattention.py | 21 +++++-------------- tensorrt_llm/_torch/pyexecutor/_util.py | 3 ++- .../test_triattention_pipeline.py | 4 +--- 3 files changed, 8 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 1e2880834195..67f2c4e8587c 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -94,9 +94,7 @@ def _allocate_block_offset_snapshot( 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 - ) + block_offsets_device = torch.empty(snapshot_shape, dtype=torch.int32, device=anchor_pool.device) return block_offsets_host, block_offsets_device @@ -285,10 +283,7 @@ def _load_calibration(self) -> None: stats = raw["stats"] metadata = raw["metadata"] if "sampled_heads" in metadata: - heads = [ - (int(layer), int(head)) - for layer, head in metadata["sampled_heads"] - ] + heads = [(int(layer), int(head)) for layer, head in metadata["sampled_heads"]] else: heads = [ ( @@ -329,14 +324,10 @@ def _load_calibration(self) -> None: 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" - ) + 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 - ) + 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}]" @@ -372,9 +363,7 @@ def on_request_init(self, request: "LlmRequest", **kwargs) -> None: 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 - ) + max_source_tokens = prompt_length + min(max_decode_tokens, self._selection_width_capacity) if max_source_tokens <= self._score_token_capacity: return diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 19158248deb0..497621494b45 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2298,7 +2298,8 @@ def create_kv_cache_compression_manager( """ if config.algorithm == "triattention": # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. - from ..kv_cache_compression.triattention.triattention import TriAttention + from ..kv_cache_compression.triattention.triattention import \ + TriAttention return TriAttention( config, diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 46eeef777745..c0515d2c8cec 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -127,9 +127,7 @@ def test_loads_official_layout(self, tmp_path): ) 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))) - ) + 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): From d3e379ea3f1f8f045c788cdc58d6bc699a744bb4 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Tue, 28 Jul 2026 07:52:13 -0700 Subject: [PATCH 172/178] [None][test] Align TriAttention fixtures with KV cache V2 Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- tests/unittest/_torch/kv_cache_compression/conftest.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 08125abff9af..a08cd2d017e6 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -503,6 +503,8 @@ def make_cute_buffers( manager = TriAttention.__new__(TriAttention) 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), ) From 5bf4f3a16bb37e5a2cf103da2b42c79a137d0016 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:39:08 -0700 Subject: [PATCH 173/178] [None][fix] Correct TriAttention union tail stores Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- .../triattention/triattention_cute_selection.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) 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 index 0942551ecb32..0e89740e20a5 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_cute_selection.py @@ -249,12 +249,17 @@ def _reduce_and_store_union_rows( 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_scores[request_idx * self.output_row_stride + token] = reduced_values[ - token_slot - ] + union_tile[token_slot] = reduced_values[token_slot] @cute.kernel def kernel( From ad35f62bb19bdce2747eeb84ec821dd7e6782240 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Thu, 30 Jul 2026 00:23:15 -0700 Subject: [PATCH 174/178] [None][refactor] Encapsulate TriAttention Triton kernel launches Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- .../triattention/triattention.py | 34 +++----- .../triattention/triattention_kernels.py | 77 +++++++++++++++++++ .../_torch/kv_cache_compression/conftest.py | 1 - .../test_triattention_fused_settle_pack.py | 18 ++--- .../test_triattention_selection_compaction.py | 10 +-- 5 files changed, 98 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 67f2c4e8587c..6db65a868d9a 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -43,10 +43,10 @@ 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_kernel, - _gather_mean_phase_kernel, - _settle_ties_kernel, + fold_union_ranks, + gather_mean_phase, reduce_per_head_scores, + settle_ties, ) if TYPE_CHECKING: @@ -110,7 +110,6 @@ def __init__(self, omega: torch.Tensor, device: torch.device) -> None: self.sin: Optional[torch.Tensor] = None self.rows = 0 self.num_freqs = int(self._omega.numel()) - self.frequency_block = triton.next_power_of_2(self.num_freqs) def reserve(self, rows: int) -> None: """Cover positions ``[0, rows)`` with a power-of-two table.""" @@ -530,7 +529,7 @@ def _execute_eviction_round( request_count = len(eviction_requests) union = self.eviction_mode == "union" # In-place refresh: the compiled score launches captured these pointers. - _gather_mean_phase_kernel[(request_count,)]( + gather_mean_phase( self._logical_source_lengths_device, self._phase.cos, self._phase.sin, @@ -540,11 +539,8 @@ def _execute_eviction_round( self._mean_sin, self._decode_lengths_device, self._swa_destination_bases, - self._swa_rebase_delta, - NUM_FREQS=self._phase.num_freqs, - F_BLOCK=self._phase.frequency_block, - HAS_SWA=self._swa_destination_bases is not None, - num_warps=1, + 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: @@ -554,17 +550,10 @@ def _execute_eviction_round( self._union_tp_mapping, dim=0, ) - _fold_union_ranks_kernel[ - ( - request_count, - triton.cdiv(self._selection_width_capacity, 1024), - ) - ]( + fold_union_ranks( gathered, self._selection_scores_rows, - request_count, - TP_SIZE=int(self._union_tp_mapping.tp_size), - WIDTH=self._selection_width_capacity, + request_count=request_count, ) if not union: reduce_per_head_scores( @@ -648,15 +637,14 @@ def _select_kept_ordinals(self, request_count: int) -> None: self.budget, 1, ) - _settle_ties_kernel[(request_count, self._selection_rows_per_request)]( + settle_ties( self._selection_scores_rows, self._selection_row_lengths, self._prompt_lengths_device, self._provisional_rows, self._kept_ordinal_rows, - WIDTH=self._selection_width_capacity, - KEEP_COUNT=self.budget, - SELECTION_ROWS=self._selection_rows_per_request, + request_count=request_count, + selection_rows_per_request=self._selection_rows_per_request, ) def _resize_compacted_caches(self, eviction_requests: Sequence[_EvictionRequest]) -> None: diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py index 11657229e1c5..8a923bc24d6e 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention_kernels.py @@ -47,6 +47,40 @@ def _gather_mean_phase_kernel( 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 ---- @@ -280,6 +314,26 @@ def _fold_union_ranks_kernel( 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, @@ -358,3 +412,26 @@ def _settle_ties_kernel( ) 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/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index a08cd2d017e6..22c21cee3216 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -457,7 +457,6 @@ def make_phase_table(offsets, omega, initial_rows): cos=torch.cos(angles).mean(dim=1).contiguous(), sin=torch.sin(angles).mean(dim=1).contiguous(), num_freqs=num_freqs, - frequency_block=1 << (num_freqs - 1).bit_length(), ) 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 index c5b6ec4c2c7e..4dc6f9921b79 100644 --- 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 @@ -6,9 +6,7 @@ import pytest import torch -from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _settle_ties_kernel, -) +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. @@ -101,15 +99,14 @@ def test_settle_matches_torch_oracle(eviction_mode, width, keep_count): ) output_actual = output_stale.clone() - _settle_ties_kernel[(request_count, selection_rows)]( + settle_ties( scores, row_lengths, prompt_offsets, provisional, output_actual, - WIDTH=width, - KEEP_COUNT=keep_count, - SELECTION_ROWS=selection_rows, + request_count=request_count, + selection_rows_per_request=selection_rows, ) torch.cuda.synchronize(device) @@ -149,15 +146,14 @@ def test_settle_handles_topk_sentinel_padding(): stale = 0x5EED output = torch.full((rows_total, keep_count), stale, dtype=torch.int32, device=device) - _settle_ties_kernel[(rows_total, 1)]( + settle_ties( scores, row_lengths, row_prompt_offsets, provisional, output, - WIDTH=width, - KEEP_COUNT=keep_count, - SELECTION_ROWS=1, + request_count=rows_total, + selection_rows_per_request=1, ) torch.cuda.synchronize(device) 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 index 1b847a676dc1..f11a15c308ae 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -599,10 +599,8 @@ def evict_once() -> tuple[torch.Tensor, torch.Tensor]: def test_fold_union_ranks_matches_max_oracle(): """The TP union fold is an exact elementwise max over the gathered rank blocks.""" - import triton - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention_kernels import ( - _fold_union_ranks_kernel, + fold_union_ranks, ) device = torch.device("cuda", torch.cuda.current_device()) @@ -610,12 +608,10 @@ def test_fold_union_ranks_matches_max_oracle(): 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_kernel[(request_count, triton.cdiv(width, 1024))]( + fold_union_ranks( gathered, folded, - request_count, - TP_SIZE=tp_size, - WIDTH=width, + request_count=request_count, ) expected = gathered.view(tp_size, request_count, width).amax(dim=0) torch.cuda.synchronize(device) From 333e7bf57849db996338f1cce1d81bb749540da8 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:53:23 -0700 Subject: [PATCH 175/178] [None][refactor] Centralize KV-cache compression capabilities Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- .../triattention.md} | 8 +- .../_torch/kv_cache_compression/interface.py | 29 ----- .../triattention/triattention.py | 14 +-- tensorrt_llm/_torch/modules/attention.py | 5 +- tensorrt_llm/_torch/pyexecutor/_util.py | 36 ++++--- .../_torch/pyexecutor/resource_manager.py | 23 ++-- tensorrt_llm/llmapi/llm_args.py | 25 +++-- .../test_kv_cache_compression_manager.py | 102 +++++++++++------- .../test_rope_fusion_gate.py | 30 ++++-- .../test_triattention_draft_cocompaction.py | 7 +- .../test_triattention_pipeline.py | 23 +++- tests/unittest/llmapi/test_llm_args.py | 4 + 12 files changed, 170 insertions(+), 136 deletions(-) rename examples/{triattention/README.md => kv_cache_compression/triattention.md} (94%) delete mode 100644 tensorrt_llm/_torch/kv_cache_compression/interface.py diff --git a/examples/triattention/README.md b/examples/kv_cache_compression/triattention.md similarity index 94% rename from examples/triattention/README.md rename to examples/kv_cache_compression/triattention.md index 1c3b5c237855..1a805d723c93 100644 --- a/examples/triattention/README.md +++ b/examples/kv_cache_compression/triattention.md @@ -22,7 +22,7 @@ TriAttention is integrated into TensorRT-LLM as a KV-cache compression manager o * PyTorch backend **Notes:** -1. TriAttention requires `enable_block_reuse=False` in the KV-cache configuration — the eviction physically rewrites stored keys, which is incompatible with block reuse. The construction step rejects a cache manager that has block reuse enabled. +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. @@ -72,8 +72,8 @@ compression_config = TriAttentionKvCacheCompressionConfig( model_path="", # used to derive the RoPE tables ) -# 2. TriAttention needs the V2 KV-cache manager and block reuse disabled. -kv_config = KvCacheConfig(enable_block_reuse=False, use_kv_cache_manager_v2=True) +# 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="", @@ -102,7 +102,7 @@ kv_cache_compression_config: calibration_path: /path/to/qwen3-8b-calibration.pt model_path: kv_cache_config: - enable_block_reuse: false + enable_block_reuse: true use_kv_cache_manager_v2: true ``` 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 a1df7a9d5b93..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() - TRIATTENTION = auto() - - def is_eviction_method(self) -> bool: - """Return whether this mode physically evicts cached tokens.""" - return self in (KvCacheCompressionMode.TRIATTENTION,) - - @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/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 6db65a868d9a..8eac8104fb10 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -30,11 +30,12 @@ from transformers import AutoConfig from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS -from ...._utils import prefer_pinned -from ....bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( +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 ....logger import logger +from tensorrt_llm.logger import logger + from ...distributed import allgather from ...pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from ...pyexecutor.llm_request import LlmRequestState @@ -50,7 +51,8 @@ ) if TYPE_CHECKING: - from ....llmapi.llm_args import TriAttentionKvCacheCompressionConfig + from tensorrt_llm.llmapi.llm_args import TriAttentionKvCacheCompressionConfig + from ...pyexecutor.llm_request import LlmRequest from ...pyexecutor.scheduler import ScheduledRequests @@ -141,8 +143,6 @@ def reserve(self, rows: int) -> None: class TriAttention(KVCacheCompressionManager): """Periodic physical KV eviction driven by trigonometric importance scoring.""" - adjusts_generation_kv_length = True - # ---- construction ---- def __init__( @@ -151,7 +151,7 @@ def __init__( kv_cache_manager: KVCacheManagerV2, draft_kv_cache_manager: Optional[KVCacheManagerV2] = None, ) -> None: - super().__init__(kv_cache_manager, draft_kv_cache_manager) + super().__init__(config, kv_cache_manager, draft_kv_cache_manager) self.budget = config.budget self.beta = config.beta self.eviction_mode = config.eviction_mode diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py index 2f7dd6c451d1..53ab4b86bd44 100644 --- a/tensorrt_llm/_torch/modules/attention.py +++ b/tensorrt_llm/_torch/modules/attention.py @@ -636,9 +636,8 @@ 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. - kv_cache_compression_mode.is_eviction_method()): + 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.", diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 497621494b45..5cc90b6537ee 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2259,28 +2259,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], ) -> None: - """Reject speculative setups the compression method cannot run with.""" + """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 - if not config.kv_cache_compression_mode.is_eviction_method(): - return - if config.algorithm != "triattention": + if not config.supports_speculative_decoding(): raise ValueError( - f"KV-cache compression algorithm {config.algorithm!r} has no " - "speculative-decoding compatibility contract") + 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 does not support speculative decoding " f"mode {mode.name}; use one-model MTP or EAGLE3") - if getattr(config, "eviction_mode", None) != "union": - raise ValueError( - "KV-cache compression with speculative decoding requires " - "eviction_mode='union'") def create_kv_cache_compression_manager( @@ -2293,8 +2294,7 @@ 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. 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": # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. @@ -2561,8 +2561,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) compression_manager = create_kv_cache_compression_manager( kv_cache_compression_config, kv_cache_manager, @@ -3032,6 +3030,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/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 121569840695..8178cfdd846d 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3528,19 +3528,21 @@ 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): @@ -3550,6 +3552,9 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): 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", @@ -3589,6 +3594,12 @@ class TriAttentionKvCacheCompressionConfig(KvCacheCompressionConfig): "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], 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 6bf7511e46b9..53217dbba93b 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,34 @@ 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_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) - validate_kv_cache_compression_with_spec(config, 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 +358,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/test_rope_fusion_gate.py b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py index 718909f509de..144358fce6e7 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py +++ b/tests/unittest/_torch/kv_cache_compression/test_rope_fusion_gate.py @@ -1,21 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""KV-cache compression forces the unfused-RoPE path. - -Compression physically evicts cached tokens, so the KV length stops matching -the logical sequence length. The fused path derives each new token's rotary -position from the KV length inside the attention kernel; the unfused path -consumes the engine's logical ``position_ids``. With compression enabled the -attention module must therefore keep RoPE unfused so rotary positions stay -logical (original absolute positions, matching the official TriAttention -implementations) while the shortened KV length only bounds attention extent. +"""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 TriAttentionKvCacheCompressionConfig +from tensorrt_llm.llmapi.llm_args import ( + KvCacheCompressionConfig, + TriAttentionKvCacheCompressionConfig, +) def _make_attention(model_config: ModelConfig) -> Attention: @@ -38,7 +37,16 @@ def test_plain_attention_defaults_to_fused_rope() -> None: assert attn.rope_fusion is True -def test_kv_cache_compression_forces_unfused_rope() -> None: +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" 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 index 225e311bf3b3..029dd2ad6a06 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -73,11 +73,11 @@ class Boom(RuntimeError): "gate,match", [ ("callsite_dflash", "one-model MTP or EAGLE3"), - ("union_only_per_head", "union"), + ("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_with_spec + 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": @@ -90,8 +90,9 @@ def test_speculative_admission_gates_raise(gate, match): ) with pytest.raises(ValueError, match=match): - validate_kv_cache_compression_with_spec( + validate_kv_cache_compression_compatibility( config, + SimpleNamespace(enable_block_reuse=False), spec_config, ) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index c0515d2c8cec..52fdf709d1ac 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -40,7 +40,10 @@ # 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 +from tensorrt_llm._torch.pyexecutor._util import ( + create_kv_cache_compression_manager, + validate_kv_cache_compression_compatibility, +) @pytest.fixture @@ -64,10 +67,15 @@ def _make_hf_config(**values): class TestConfigAndFactory: - def test_factory_returns_triattention_and_propagates_config_fields(self): + 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=False) + 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.object(TriAttention, "_initialize_eviction_state") as initialize: mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) @@ -75,6 +83,10 @@ def test_factory_returns_triattention_and_propagates_config_fields(self): 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() @@ -404,7 +416,7 @@ def test_one_model_draft_co_compression_is_accepted(self, spec_mode): draft_kv_cache_manager=draft_manager, ) - from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_with_spec + from tensorrt_llm._torch.pyexecutor._util import validate_kv_cache_compression_compatibility from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig, MTPDecodingConfig spec_config = ( @@ -417,8 +429,9 @@ def test_one_model_draft_co_compression_is_accepted(self, spec_mode): ) ) - validate_kv_cache_compression_with_spec( + validate_kv_cache_compression_compatibility( _make_tri_config(budget=8), + SimpleNamespace(enable_block_reuse=False), spec_config, ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 0f7567ab3fc7..5560d873b7df 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -2969,6 +2969,10 @@ def test_kv_cache_compression_config_dispatches_by_algorithm(): 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: From 39d5078e03542c0e784fbcd2f33c8564313627b0 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:04:02 -0700 Subject: [PATCH 176/178] [None][fix] Gate TriAttention on SM100-family GPUs Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 7 +++++-- .../_torch/executor/test_kv_cache_compression_manager.py | 9 +++++++++ .../kv_cache_compression/test_triattention_pipeline.py | 8 +++++++- 3 files changed, 21 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f2e44c3e8d3f..7e3d3aa470d2 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 @@ -2333,6 +2333,9 @@ def create_kv_cache_compression_manager( 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 \ TriAttention 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 53217dbba93b..4f628d6fbb94 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py +++ b/tests/unittest/_torch/executor/test_kv_cache_compression_manager.py @@ -304,6 +304,15 @@ def test_factory_accepts_independent_draft_manager(self): is None ) + 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") target = _v2_manager(is_draft=False) diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index 52fdf709d1ac..ecc87da1390d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -76,7 +76,13 @@ def test_factory_allows_block_reuse_and_propagates_config_fields(self): SimpleNamespace(enable_block_reuse=True), None, ) - with mock.patch.object(TriAttention, "_initialize_eviction_state") as initialize: + with ( + mock.patch( + "tensorrt_llm._torch.pyexecutor._util.is_sm_100f", + return_value=True, + ), + mock.patch.object(TriAttention, "_initialize_eviction_state") as initialize, + ): mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) assert isinstance(mgr, TriAttention) assert mgr.budget == 32 From 8f111025b810c74b359e423799c8c1f8b96dad7e Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Thu, 30 Jul 2026 22:09:50 -0700 Subject: [PATCH 177/178] [None][chore] Apply TriAttention factory formatting Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7e3d3aa470d2..f87428a9641f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2335,7 +2335,8 @@ def create_kv_cache_compression_manager( if config.algorithm == "triattention": if not is_sm_100f(): raise RuntimeError( - "TriAttention requires an SM100-family device (SM100 or SM103).") + "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 \ TriAttention From a27a55c7c4096bb37340ab3d187a755356928208 Mon Sep 17 00:00:00 2001 From: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> Date: Mon, 3 Aug 2026 00:51:50 -0700 Subject: [PATCH 178/178] [None][refactor] Rename TriAttention compression manager Signed-off-by: Tianrui 'Hudayday' Hu <32944717+Hudayday@users.noreply.github.com> --- .../triattention/triattention.py | 4 ++-- tensorrt_llm/_torch/pyexecutor/_util.py | 4 ++-- .../_torch/kv_cache_compression/conftest.py | 20 ++++++++++++------- .../test_triattention_draft_cocompaction.py | 6 ++++-- .../test_triattention_pipeline.py | 20 +++++++++++-------- .../test_triattention_selection_compaction.py | 6 ++++-- 6 files changed, 37 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py index 8eac8104fb10..cc43f71441e9 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py +++ b/tensorrt_llm/_torch/kv_cache_compression/triattention/triattention.py @@ -140,8 +140,8 @@ def reserve(self, rows: int) -> None: self.rows = target -class TriAttention(KVCacheCompressionManager): - """Periodic physical KV eviction driven by trigonometric importance scoring.""" +class TriAttentionCompressionManager(KVCacheCompressionManager): + """KV-cache compression manager for periodic TriAttention eviction.""" # ---- construction ---- diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 379c0ed2c60d..e08bb27938f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2336,9 +2336,9 @@ def create_kv_cache_compression_manager( ) # TriAttention imports CuTe/CUTLASS; keep normal executor startup lazy. from ..kv_cache_compression.triattention.triattention import \ - TriAttention + TriAttentionCompressionManager - return TriAttention( + return TriAttentionCompressionManager( config, kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, diff --git a/tests/unittest/_torch/kv_cache_compression/conftest.py b/tests/unittest/_torch/kv_cache_compression/conftest.py index 22c21cee3216..b815b2270729 100644 --- a/tests/unittest/_torch/kv_cache_compression/conftest.py +++ b/tests/unittest/_torch/kv_cache_compression/conftest.py @@ -218,9 +218,11 @@ def run_compaction(compaction): 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 TriAttention + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, + ) - staging = TriAttention.__new__(TriAttention) + staging = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) staging.kv_cache_manager = None staging.draft_kv_cache_manager = None staging._request_capacity = max_requests @@ -335,10 +337,12 @@ def make_tri_config(**overrides): def make_triattention(**overrides): """Construct a manager while isolating GPU-owned persistent state.""" - from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import TriAttention + from tensorrt_llm._torch.kv_cache_compression.triattention.triattention import ( + TriAttentionCompressionManager, + ) - with mock.patch.object(TriAttention, "_initialize_eviction_state"): - return TriAttention(make_tri_config(**overrides), make_fake_v2()) + with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): + return TriAttentionCompressionManager(make_tri_config(**overrides), make_fake_v2()) def make_eviction_request( @@ -480,7 +484,9 @@ def make_cute_buffers( 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 TriAttention + 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 @@ -499,7 +505,7 @@ def make_cute_buffers( swa_window=None, layer_pool_ids=layer_pool_ids, ) - manager = TriAttention.__new__(TriAttention) + manager = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) manager.kv_cache_manager = SimpleNamespace( num_pools=max(layer_pool_ids) + 1, tokens_per_block=tokens_per_block, 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 index 029dd2ad6a06..6824e5d1ce83 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_draft_cocompaction.py @@ -15,14 +15,16 @@ 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 TriAttention +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 = TriAttention.__new__(TriAttention) + tri = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) tri._request_capacity = 8 tri.budget = 4 tri._swa_window = None diff --git a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py index ecc87da1390d..2b45d3beb52d 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_pipeline.py @@ -36,7 +36,9 @@ # 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 TriAttention +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. @@ -81,10 +83,12 @@ def test_factory_allows_block_reuse_and_propagates_config_fields(self): "tensorrt_llm._torch.pyexecutor._util.is_sm_100f", return_value=True, ), - mock.patch.object(TriAttention, "_initialize_eviction_state") as initialize, + mock.patch.object( + TriAttentionCompressionManager, "_initialize_eviction_state" + ) as initialize, ): mgr = create_kv_cache_compression_manager(cfg, kv_cache_manager=fake_v2) - assert isinstance(mgr, TriAttention) + assert isinstance(mgr, TriAttentionCompressionManager) assert mgr.budget == 32 assert mgr.beta == 16 assert mgr.eviction_mode == "per_head" @@ -96,7 +100,7 @@ def test_factory_allows_block_reuse_and_propagates_config_fields(self): initialize.assert_called_once_with() -class TestTriAttentionClass: +class TestTriAttentionCompressionManager: def test_loads_flat_pt(self, flat_calibration_pt): mgr = _make_triattention() mgr.calibration_path = flat_calibration_pt @@ -275,8 +279,8 @@ def _make_due_decode_request(seq_len, *, num_extra_kv_tokens=0, kv_reserve_draft 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(TriAttention, "_initialize_eviction_state"): - mgr = TriAttention(_make_tri_config(budget=8), fake_v2) + 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, @@ -415,8 +419,8 @@ def test_confirmed_length_comes_from_capacity_ledger_not_logical_length(self): @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(TriAttention, "_initialize_eviction_state"): - TriAttention( + with mock.patch.object(TriAttentionCompressionManager, "_initialize_eviction_state"): + TriAttentionCompressionManager( _make_tri_config(budget=8), _make_fake_v2(), draft_kv_cache_manager=draft_manager, 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 index f11a15c308ae..a17a1b4c6731 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py +++ b/tests/unittest/_torch/kv_cache_compression/test_triattention_selection_compaction.py @@ -11,7 +11,9 @@ 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 TriAttention +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, ) @@ -44,7 +46,7 @@ def _make_selection_buffers( num_kv_heads=1, ): """Allocate the production selection buffers without score or compaction.""" - tri = TriAttention.__new__(TriAttention) + tri = TriAttentionCompressionManager.__new__(TriAttentionCompressionManager) tri.eviction_mode = eviction_mode tri._request_capacity = max_requests tri._selection_width_capacity = width