diff --git a/docs/source/developer-guide/sparse-attention-development-guide.md b/docs/source/developer-guide/sparse-attention-development-guide.md index 18c189cd5555..bc81970a7b60 100644 --- a/docs/source/developer-guide/sparse-attention-development-guide.md +++ b/docs/source/developer-guide/sparse-attention-development-guide.md @@ -146,7 +146,7 @@ Dynamic generation-phase KV eviction is tracked as future work. | Public config / lowering | `CSA2SparseAttentionConfig(algorithm="csa2")` → `CSA2Params`, `CSA2MetadataParams`; checkpoint-owned `CSA2Layout` | | Registry | TRTLLM sparse backend → `CSA2TrtllmAttention`; cache manager → `CSA2CacheManager` (V2) | | Module | Typed MLA adapter in `sparse/csa2/module.py`; compression, shared KV/index source routing and per-phase prediction | -| Metadata | `CSA2TrtllmMetadata`; SWA slots/visibility, owner page tables, compressor inputs, graph-stable index descriptors and query tiles | +| Metadata | `CSA2Metadata` extension; SWA slots/visibility, owner page tables, compressor inputs, graph-stable index descriptors and query tiles | | Compute | SM100-family trtllm-gen; SM90 FlashMLA BF16 and SM120/121 FlashInfer BF16 helpers | CSA2 GLOBAL pages couple 288-byte main and 68-byte index records under one diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/modeling.py index cd23782662cb..825aadaa5472 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/deepseek_v3/r1_0528_nvfp4__sm_10x__default/modeling.py @@ -83,6 +83,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig @@ -131,7 +132,7 @@ def _build_step_args(md: TrtllmAttentionMetadata) -> dict: num_ctx_tokens=md.num_ctx_tokens, trtllm_gen_jit_warmup=md.trtllm_gen_jit_warmup, max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + force_prepare_spec_dec_tree_mask=md.ext(SpecDecodingMetadata).force_prepare_tree_mask, ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py index 2eeadb915157..25a5618e32bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/modeling.py @@ -22,14 +22,16 @@ import torch from transformers import PretrainedConfig +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, AttentionMetadata, ) +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata from tensorrt_llm._torch.attention.backends.sparse.dsa.module import forward_generation_sparse_attn from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSABackendForwardArgs -from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareMetadata +from tensorrt_llm._torch.attention.backends.sparse.params import MTPIndexShareExtension from tensorrt_llm._torch.distributed import AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.model_config import ModelConfig @@ -43,7 +45,6 @@ to_trtllm_gen_act_type, ) from tensorrt_llm._torch.speculative import MTPEagleWorker, SpecMetadata -from tensorrt_llm._torch.speculative.eagle3 import _reset_mtp_index_share from tensorrt_llm._torch.speculative.interface import ( _FORCE_ACCEPT_RNG_COUNTER_STRIDE, _FORCE_ACCEPT_RNG_POOL_SIZE, @@ -627,6 +628,7 @@ def draft_first_step( among the verify rows; and the selection the later steps attend to, as KV pool rows, or None where they attend to every position (a short context) or no step follows. """ + dsa = attn_metadata.ext(DSAMetadata) draft = self._draft layer = draft.layer rows = hidden_states.shape[0] @@ -651,7 +653,7 @@ def draft_first_step( ) attn_out = self._attention(layer, position_ids, hidden_states, attn_metadata, True) pool_rows = None - if draft_len > 1 and not attn_metadata.skip_indexer_for_gen_reqs: + if draft_len > 1 and not dsa.skip_indexer_for_gen_reqs: pool_rows = self._stash_selection(attn_metadata, num_accepted, draft_len + 1) hidden_states, residual = self._post_attention_hop(layer, attn_out, residual) gather_ids, hidden_states, residual = glm_gather_last_accepted_rows( @@ -675,11 +677,12 @@ def draft_next_step( indexer, to every position up to their own, the dense selection that metadata lists (refreshed for this step by the loop). Returns the MTP layer's output rows. """ + dsa = attn_metadata.ext(DSAMetadata) draft = self._draft layer = draft.layer hidden_states, residual = eh_proj_output - if attn_metadata.skip_indexer_for_gen_reqs: - positions = attn_metadata.topk_indices_buffer[: hidden_states.shape[0]] + if dsa.skip_indexer_for_gen_reqs: + positions = dsa.topk_indices_buffer[: hidden_states.shape[0]] attn_out = self._attention( layer, None, hidden_states, attn_metadata, True, positions=positions ) @@ -854,6 +857,7 @@ def _stash_selection(self, attn_metadata, num_accepted: torch.Tensor, next_n: in stream after draft step 0's attention: each request's top-k of its last accepted row (the MTP index share), as KV pool rows of the MTP layer. The first later step joins it (``selection_ready``).""" + dsa = attn_metadata.ext(DSAMetadata) draft = self._draft manager = attn_metadata.kv_cache_manager requests = num_accepted.shape[0] @@ -872,7 +876,7 @@ def _stash_selection(self, attn_metadata, num_accepted: torch.Tensor, next_n: in ) pool_rows = convert_req_index_to_global( draft.request_rows[:requests], - attn_metadata.block_table[:requests], + dsa.block_table[:requests], stashed, tokens_per_block, self.index_topk, @@ -938,6 +942,7 @@ def _shard_indexer_keys(self, attn_metadata) -> bool: Applies to one generation request whose shard fits the top-k kernel; returns whether it did, and otherwise every indexer layer scores the whole context with the built-in indexer.""" + dsa = attn_metadata.ext(DSAMetadata) manager = attn_metadata.kv_cache_manager if self._hist_workspace is None or attn_metadata.num_seqs != 1: return False @@ -975,8 +980,8 @@ def _shard_indexer_keys(self, attn_metadata) -> bool: ) buffers = self._hist_buffers hist_shard_meta( - attn_metadata.gen_indexer_kv_lens_cuda_runtime, - attn_metadata.indexer_k_cache_block_offsets[0], + dsa.gen_indexer_kv_lens_cuda_runtime, + dsa.indexer_k_cache_block_offsets[0], buffers["shard_ctx"], buffers["shard_block_table"], self.mapping.tp_size, @@ -1004,6 +1009,7 @@ def _attention( On an MTP draft step after the first (``sharded``) the rows attend to a given selection and the indexer does not run: ``pool_rows``, KV pool rows, or ``positions``, sequence positions the flash decode maps to pool rows itself.""" + dsa = attn_metadata.ext(DSAMetadata) mla = layer.self_attn mqa = mla.mqa indexer = mqa.indexer @@ -1138,7 +1144,7 @@ def kv_half_then_q_b(): output, seq_len_q, self.index_topk, - attn_metadata.block_table[: attn_metadata.num_seqs], + dsa.block_table[: attn_metadata.num_seqs], manager.tokens_per_block, page_index_scale, layer_offset, @@ -1190,6 +1196,7 @@ def kv_half_then_q_b(): def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metadata, rows, mla): """The indexer of a layer that runs one, on this rank's key shard: writes the exact top-2048 cache positions of every row into the shared selection buffer.""" + dsa = attn_metadata.ext(DSAMetadata) buffers = self._hist_buffers manager = attn_metadata.kv_cache_manager rope = indexer.rotary_emb @@ -1217,8 +1224,8 @@ def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metada buffers["k_fp8"][:rows], buffers["k_scale"][:rows], k_cache, - attn_metadata.slot_mapping_fp8[:rows], - attn_metadata.slot_mapping_scale[:rows], + dsa.slot_mapping_fp8[:rows], + dsa.slot_mapping_scale[:rows], rows, True, True, @@ -1260,7 +1267,7 @@ def _select_rows(self, indexer, kv_a, latent_end, q_a, position_ids, attn_metada ) hist_topk( logits, - attn_metadata.gen_indexer_kv_lens_cuda_runtime, + dsa.gen_indexer_kv_lens_cuda_runtime, buffers["topk"][:rows], self._hist_workspace.comm_buffer, self._hist_workspace.state, @@ -1773,6 +1780,7 @@ def _draft_on_target( dense selection is refreshed from the new lengths. The host side of draft step 0's switch to one row per request comes before that. """ + dsa = attn_metadata.ext(DSAMetadata) draft_len = spec_metadata.runtime_draft_len # The accepted tokens' rows: draft step 0's input ids (prepare_1st_drafter_inputs). accepted_tokens = inputs["input_ids"].view(batch_size, draft_len + 1) @@ -1781,12 +1789,12 @@ def _draft_on_target( draft_tokens = accepted_tokens.new_empty((batch_size, draft_len)) positions = inputs["position_ids"] eh_proj_output = gather_ids = pool_rows = next_new_tokens = None - uses_mtp_index_share = isinstance(attn_metadata, MTPIndexShareMetadata) + index_share = attn_metadata.find_ext(MTPIndexShareExtension) with contextlib.ExitStack() as draft_scope: - if uses_mtp_index_share: - attn_metadata.set_in_mtp_draft_loop(True) - attn_metadata.set_mtp_num_accepted(num_accepted_tokens) - draft_scope.callback(_reset_mtp_index_share, attn_metadata) + if index_share is not None: + index_share.set_in_mtp_draft_loop(True) + index_share.set_mtp_num_accepted(num_accepted_tokens) + draft_scope.callback(index_share.reset_mtp_index_share) attn_metadata = draft_scope.enter_context( self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) ) @@ -1795,8 +1803,8 @@ def _draft_on_target( seq_lens = attn_metadata.seq_lens_cuda[:batch_size] kv_lens = attn_metadata.kv_lens_cuda[:batch_size] for i in range(draft_len): - if uses_mtp_index_share: - attn_metadata.set_skip_topk(i > 0) + if index_share is not None: + index_share.set_skip_topk(i > 0) if i == 0: hidden_states, gather_ids, pool_rows = decode.draft_first_step( attn_metadata, @@ -1843,12 +1851,10 @@ def _draft_on_target( kv_lens, rewind_accepted, ) - if not last and attn_metadata.skip_indexer_for_gen_reqs: - attn_metadata.prepare_dense_topk_indices( - attn_metadata.kv_lens_cuda, device=True - ) + if not last and dsa.skip_indexer_for_gen_reqs: + dsa.prepare_dense_topk_indices(attn_metadata.kv_lens_cuda, device=True) if i == 0: - attn_metadata.use_spec_decoding = False + attn_metadata.ext(SpecDecodingMetadata).use_spec_decoding = False self._drafted_rows = (draft_tokens, next_new_tokens) return draft_tokens diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/modeling.py index 7b443d2f4760..547e3fd8bde7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/gpt_oss/gpt_oss_120b__sm_10x__default/modeling.py @@ -75,6 +75,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.mxfp8_quantize import ( mxfp8_quantize, ) +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig @@ -153,6 +155,12 @@ _MOE_MAX_T = 8192 +def _flash_mla_field(md: TrtllmAttentionMetadata, name: str): + """A FlashMLA metadata field, or None when FlashMLA is not attached.""" + flash_mla = md.find_ext(FlashMlaMetadata) + return getattr(flash_mla, name) if flash_mla is not None else None + + def _in_chunks(experts: Callable, x: torch.Tensor, L: "_LayerWeights") -> torch.Tensor: """`experts(x, L)` over at most `_MOE_MAX_T` rows at a time, re-joined in order.""" parts = [experts(c, L) for c in torch.split(x, _MOE_MAX_T, 0)] @@ -529,6 +537,7 @@ def _check_step_contract(self, md) -> None: def step_args(self, md: TrtllmAttentionMetadata) -> dict: """Project the prepared metadata onto thop_attention's explicit batch state, once per forward.""" + spec = md.find_ext(SpecDecodingMetadata) return dict( sequence_length=md.kv_lens_cuda_runtime, host_past_key_value_lengths=md.kv_lens_runtime, @@ -553,21 +562,21 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: block_ids_per_seq=md.block_ids_per_seq, max_context_q_len_override=md.max_context_q_len_override, is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + is_spec_decoding_enabled=spec.is_spec_decoding_enabled, + use_spec_decoding=spec.use_spec_decoding, + is_spec_dec_tree=spec.is_spec_dec_tree, + spec_decoding_generation_lengths=spec.generation_lengths, + spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp, + spec_decoding_packed_mask=spec.packed_mask, + spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset, + spec_decoding_bl_tree_mask=spec.bl_tree_mask, + spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv, num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, + flash_mla_tile_scheduler_metadata=_flash_mla_field(md, "tile_scheduler_metadata"), + flash_mla_num_splits=_flash_mla_field(md, "num_splits"), max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask, ) def forward( @@ -766,6 +775,7 @@ def _check_step_contract(self, md) -> None: def step_args(self, md: TrtllmAttentionMetadata) -> dict: """Project the prepared metadata onto thop_attention's explicit batch state, once per forward.""" + spec = md.find_ext(SpecDecodingMetadata) return dict( sequence_length=md.kv_lens_cuda_runtime, host_past_key_value_lengths=md.kv_lens_runtime, @@ -790,21 +800,21 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: block_ids_per_seq=md.block_ids_per_seq, max_context_q_len_override=md.max_context_q_len_override, is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + is_spec_decoding_enabled=spec.is_spec_decoding_enabled, + use_spec_decoding=spec.use_spec_decoding, + is_spec_dec_tree=spec.is_spec_dec_tree, + spec_decoding_generation_lengths=spec.generation_lengths, + spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp, + spec_decoding_packed_mask=spec.packed_mask, + spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset, + spec_decoding_bl_tree_mask=spec.bl_tree_mask, + spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv, num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, + flash_mla_tile_scheduler_metadata=_flash_mla_field(md, "tile_scheduler_metadata"), + flash_mla_num_splits=_flash_mla_field(md, "num_splits"), max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask, ) def forward( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index 29ac78174102..3e6a1e6d04bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -129,7 +129,10 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.graph.eager_on_graph import ( eager_on_graph, ) +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata, RecurrentMetadata +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.qsa.metadata import QSAMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM @@ -234,19 +237,19 @@ # bypass routing. _SM = (10, 3) -# QSA backend metadata buffers rebuilt by prepare(). Their absence means the engine -# did not select the QSA metadata and the index plane would never be written. +# QSA metadata buffers rebuilt by prepare(). Their absence means the engine did not +# attach the QSA metadata and the index plane would never be written. _QSA_FIELDS = ( - "qsa_block_table", - "qsa_req_idx_per_token", - "qsa_logical_positions", - "qsa_sequence_lengths", - "qsa_visible_blocks", - "qsa_topk_indices", - "qsa_topk_row_starts", + "block_table", + "req_idx_per_token", + "logical_positions", + "sequence_lengths", + "visible_blocks", + "topk_indices", + "topk_row_starts", ) -# Fields read off `attn_metadata.mamba_metadata`. Absence is silent: a smoke run +# Fields read off `RecurrentMetadata.mamba_metadata`. Absence is silent: a smoke run # passes with no recurrent state at all. _MAMBA_FIELDS = ( "state_indices", @@ -268,6 +271,12 @@ _CUTE_DSL_MOE_MIN_ROWS = 4096 # capacity: fixed: a dispatch threshold, not a buffer size +def _flash_mla_field(md: TrtllmAttentionMetadata, name: str): + """A FlashMLA metadata field, or None when FlashMLA is not attached.""" + flash_mla = md.find_ext(FlashMlaMetadata) + return getattr(flash_mla, name) if flash_mla is not None else None + + def _decode_projection(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: """A graph-captured decode step's bf16 projection, always a PDL launch. @@ -1748,9 +1757,10 @@ def _bind_cache_manager(self, md, position_ids) -> None: ) self._binds += 1 - missing = [name for name in _QSA_FIELDS if not hasattr(md, name)] + qsa = md.find_ext(QSAMetadata) + missing = [name for name in _QSA_FIELDS if not hasattr(qsa, name)] assert not missing, ( - f"metadata fields missing: {missing}; the engine did not select QSA " + f"QSA metadata fields missing: {missing}; the engine did not attach QSA " "metadata for this model" ) @@ -1800,8 +1810,8 @@ def _bind_cache_manager(self, md, position_ids) -> None: # `qsa_topk_select` checks only that the buffer is at least `index_topk` wide, # then addresses it as flat `index_topk`-wide rows ignoring the real stride, so # a width mismatch silently mis-writes every row but the first. - assert md.qsa_topk_indices.shape[1] == self.block_topk, ( - f"the engine sized the QSA selection buffer {md.qsa_topk_indices.shape[1]} " + assert qsa.topk_indices.shape[1] == self.block_topk, ( + f"the engine sized the QSA selection buffer {qsa.topk_indices.shape[1]} " f"wide and this target selects {self.block_topk} blocks per row" ) position_cache = cm.get_qsa_position_buffer() @@ -1813,9 +1823,10 @@ def _bind_cache_manager(self, md, position_ids) -> None: # A missed hybrid predicate would still build and pass a smoke run with 36 # layers carrying no state. - mm = md.mamba_metadata - assert mm is not None and mm is not False, ( - "attn_metadata.mamba_metadata is absent: the runtime allocated no " + recurrent = md.find_ext(RecurrentMetadata) + mm = recurrent.mamba_metadata if recurrent is not None else None + assert mm is not None, ( + "the recurrent metadata is absent: the runtime allocated no " f"recurrent-state pool, so the {len(self.gdn_layers)} linear-attention " "layers would run with no state" ) @@ -1847,7 +1858,7 @@ def _bind_cache_manager(self, md, position_ids) -> None: # With speculation on, the Blackwell trtllm-gen spec-dec mask machinery is off # for a linear tree at sm_103: every `spec_decoding_*` tensor stays None and # `predicted_tokens_per_seq` alone produces the within-block mask. - assert not md.is_spec_decoding_enabled, ( + assert not md.ext(SpecDecodingMetadata).is_spec_decoding_enabled, ( "the attention metadata reports is_spec_decoding_enabled; on sm_103 a " "linear-tree MTP forces that machinery off" ) @@ -1960,7 +1971,9 @@ def _bind_draft_cache_manager(self, md) -> None: "this target resolved IndexShare on but the engine's QSA metadata did not" ) slots = cm.qsa_shared_topk_slots - check_shared_selection(md.qsa_shared_topk_indices, slots, layer, self.block_topk) + check_shared_selection( + md.ext(QSAMetadata).shared_topk_indices, slots, layer, self.block_topk + ) share_slot = slots[layer] self._draft_bound = _DraftBound(cm, index_cache, position_cache, kv, share_slot) @@ -2680,6 +2693,7 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: """Project the prepared metadata onto thop_attention's batch state, once per forward. Tensors are engine-owned persistent buffers refreshed in place; Python ints are per-capture constants.""" + spec = md.find_ext(SpecDecodingMetadata) return dict( sequence_length=md.kv_lens_cuda_runtime, host_past_key_value_lengths=md.kv_lens_runtime, @@ -2706,21 +2720,21 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: block_ids_per_seq=md.block_ids_per_seq, max_context_q_len_override=md.max_context_q_len_override, is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + is_spec_decoding_enabled=spec.is_spec_decoding_enabled, + use_spec_decoding=spec.use_spec_decoding, + is_spec_dec_tree=spec.is_spec_dec_tree, + spec_decoding_generation_lengths=spec.generation_lengths, + spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp, + spec_decoding_packed_mask=spec.packed_mask, + spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset, + spec_decoding_bl_tree_mask=spec.bl_tree_mask, + spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv, num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, + flash_mla_tile_scheduler_metadata=_flash_mla_field(md, "tile_scheduler_metadata"), + flash_mla_num_splits=_flash_mla_field(md, "num_splits"), max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask, ) # --------------------------------------------------------------- blocks --- @@ -3678,7 +3692,7 @@ def prepare( md = attn_metadata self._check_step_contract(md) attn_args = self.step_args(md) - mm = md.mamba_metadata + mm = md.ext(RecurrentMetadata).mamba_metadata # **The activation rows are the batch's rows, or a prefill CUDA graph bucket's.** A # breakable prefill graph pads the step's rows up to `padded_num_tokens` (zero ids, @@ -3802,14 +3816,15 @@ def prepare( # The `[T, 3]` transpose of the `[3, 1, T]` position buffer, which the QSA position # plane stores. coords = torch.transpose(torch.reshape(position_ids, [_QSA_POSITION_AXES, -1]), 0, 1) + qsa_md = md.ext(QSAMetadata) qsa = ( - md.qsa_req_idx_per_token[:num_tokens], - md.qsa_logical_positions[:num_tokens], - md.qsa_sequence_lengths[:num_tokens], - md.qsa_visible_blocks[:num_tokens], - md.qsa_topk_indices[:num_tokens], - md.qsa_topk_row_starts[:num_tokens], - md.qsa_block_table, + qsa_md.req_idx_per_token[:num_tokens], + qsa_md.logical_positions[:num_tokens], + qsa_md.sequence_lengths[:num_tokens], + qsa_md.visible_blocks[:num_tokens], + qsa_md.topk_indices[:num_tokens], + qsa_md.topk_row_starts[:num_tokens], + qsa_md.block_table, md.tokens_per_block, coords[:num_tokens], sparse, @@ -3832,9 +3847,9 @@ def prepare( num_gens=num_decodes, rows_per_gen=gen_p, slots=si_d, - request_indices=md.qsa_req_idx_per_token[num_prefill_tokens:num_tokens], - logical_positions=md.qsa_logical_positions[num_prefill_tokens:num_tokens], - block_table=md.qsa_block_table, + request_indices=qsa_md.req_idx_per_token[num_prefill_tokens:num_tokens], + logical_positions=qsa_md.logical_positions[num_prefill_tokens:num_tokens], + block_table=qsa_md.block_table, tokens_per_block=md.tokens_per_block, ) @@ -4007,25 +4022,26 @@ def _capture_selection(self, md, topk_indices, visible_blocks) -> None: (`num_accepted - 1` past its first). Capturing here rather than at the first draft decode step keeps the indexer at one call per draft loop for a batch with context rows too.""" + qsa = md.find_ext(QSAMetadata) batch = md.num_seqs - assert md.qsa_mtp_num_accepted is not None, ( + assert qsa.mtp_num_accepted is not None, ( "IndexShare needs the worker's accepted counts to pick each request's capture " "row; the draft loop did not supply them" ) seq_lens = md.seq_lens_cuda[:batch].to(torch.int64) starts = torch.cumsum(seq_lens, 0) - seq_lens - accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) + accepted = qsa.mtp_num_accepted[:batch].to(torch.int64) offsets = torch.minimum(torch.clamp(accepted - 1, min=0), seq_lens - 1) offsets[: md.num_contexts] = seq_lens[: md.num_contexts] - 1 rows = starts + offsets slot = self.core._draft_bound.share_slot - md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) - md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( + qsa.shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) + qsa.shared_topk_visible_blocks[slot, :batch].copy_( torch.index_select(visible_blocks, 0, rows) ) # Host-side readiness, recorded in program order: under capture this runs while the # graph is recorded, so a graph holds the reuse only after the capture it reads. - md.qsa_shared_topk_captured_slots.add(slot) + qsa.shared_topk_captured_slots.add(slot) def mtp_forward( self, @@ -4047,6 +4063,7 @@ def mtp_forward( "core before the worker, so this cannot be reached first" ) md = attn_metadata + qsa = md.find_ext(QSAMetadata) rows = hidden_states.shape[0] assert input_ids.shape[0] == rows, ( f"the draft step's {input_ids.shape[0]} token ids and {rows} " @@ -4098,19 +4115,17 @@ def mtp_forward( ) # The draft layer's QSA planes and the page table that addresses them. With a - # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the - # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`); + # separate draft manager the engine's QSA metadata swaps `block_table` for the + # table it decodes from the draft manager's offsets (`QSAMetadata.draft_swap`); # a table still decoding the trunk's offsets would read and write another pool's # pages. draft = core._draft_bound if draft is None or cm is not draft.cm: core._bind_draft_cache_manager(md) draft = core._draft_bound - assert cm is core._bound.cm or md.qsa_block_table is getattr( - md, "qsa_draft_block_table", None - ), ( + assert cm is core._bound.cm or qsa.block_table is qsa.draft_block_table, ( "the draft step is bound to a separate draft KV cache manager but " - "qsa_block_table is not that manager's table: the engine's QSA metadata did " + "the QSA block_table is not that manager's table: the engine's QSA metadata did " "not swap its page table for the draft forward" ) # The `qsa_*` row mapping is sized from the metadata's token count, which the @@ -4128,8 +4143,8 @@ def mtp_forward( # IndexShare: the only draft step with context rows is the loop's first, its # draft-extend step. It always runs the indexer and, with sharing on, captures the # selection the draft decode steps reuse. - share = core.mtp_index_share and md.qsa_in_mtp_draft_loop - assert not (share and md.qsa_indexer_skip_topk), ( + share = core.mtp_index_share and qsa.in_mtp_draft_loop + assert not (share and qsa.indexer_skip_topk), ( "the worker asked a draft step with context rows to reuse a selection; only " "the draft decode steps after it reuse" ) @@ -4213,9 +4228,9 @@ def mtp_forward( # then the compression of every group this step completes. Maintained on every # draft step at both ends of the threshold, so the draft plane is complete when a # request crosses it. - req_idx = md.qsa_req_idx_per_token[:rows] - logical = md.qsa_logical_positions[:rows] - block_table = md.qsa_block_table + req_idx = qsa.req_idx_per_token[:rows] + logical = qsa.logical_positions[:rows] + block_table = qsa.block_table tokens_per_block = md.tokens_per_block index_cache = draft.index_cache qk_index = cublas_mm(mixed, M.idx) @@ -4293,12 +4308,12 @@ def mtp_forward( only_visible_blocks=True, context_rows=True, ) - topk_indices = md.qsa_topk_indices[:rows] - visible_blocks = md.qsa_visible_blocks[:rows] + topk_indices = qsa.topk_indices[:rows] + visible_blocks = qsa.visible_blocks[:rows] _select( core, scores, - md.qsa_topk_row_starts[:rows], + qsa.topk_row_starts[:rows], visible_blocks, topk_indices, _live_score_columns(core, md, batch), @@ -4308,7 +4323,7 @@ def mtp_forward( selected = qsa_expand_block_indices( topk_indices, logical, - md.qsa_sequence_lengths[:rows], + qsa.sequence_lengths[:rows], compress_ratio=core.compress_ratio, token_topk=core.token_topk, ) @@ -4380,6 +4395,7 @@ def _check_step_contract(self, md) -> None: def step_args(self, md: TrtllmAttentionMetadata) -> dict: """See PrefillTarget.step_args. Context counts are stated, not read: routing guarantees them, and a decode-only CUDA graph would freeze whatever it read.""" + spec = md.find_ext(SpecDecodingMetadata) return dict( sequence_length=md.kv_lens_cuda_runtime, host_past_key_value_lengths=md.kv_lens_runtime, @@ -4406,21 +4422,21 @@ def step_args(self, md: TrtllmAttentionMetadata) -> dict: block_ids_per_seq=md.block_ids_per_seq, max_context_q_len_override=md.max_context_q_len_override, is_cross=md.is_cross, - is_spec_decoding_enabled=md.is_spec_decoding_enabled, - use_spec_decoding=md.use_spec_decoding, - is_spec_dec_tree=md.is_spec_dec_tree, - spec_decoding_generation_lengths=md.spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=md.spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=md.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=md.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=md.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=md.max_total_draft_tokens, - spec_bl_tree_first_sparse_mask_offset_kv=md.spec_bl_tree_first_sparse_mask_offset_kv, + is_spec_decoding_enabled=spec.is_spec_decoding_enabled, + use_spec_decoding=spec.use_spec_decoding, + is_spec_dec_tree=spec.is_spec_dec_tree, + spec_decoding_generation_lengths=spec.generation_lengths, + spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp, + spec_decoding_packed_mask=spec.packed_mask, + spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset, + spec_decoding_bl_tree_mask=spec.bl_tree_mask, + spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens, + spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv, num_sparse_topk=md.num_sparse_topk, - flash_mla_tile_scheduler_metadata=md.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=md.flash_mla_num_splits, + flash_mla_tile_scheduler_metadata=_flash_mla_field(md, "tile_scheduler_metadata"), + flash_mla_num_splits=_flash_mla_field(md, "num_splits"), max_num_sequences=md.max_num_sequences, - force_prepare_spec_dec_tree_mask=md.force_prepare_spec_dec_tree_mask, + force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask, ) # --------------------------------------------------------------- blocks --- @@ -5225,7 +5241,7 @@ def forward( self._check_step_contract(md) step = self.step_args(md) spec_metadata = kwargs.get("spec_metadata") - mm = md.mamba_metadata + mm = md.ext(RecurrentMetadata).mamba_metadata # The bf16 block projections of a graph-captured step run as PDL launches # (`_decode_projection`); an eager step keeps cuBLAS. gemm = _decode_projection if md.is_cuda_graph else cublas_mm @@ -5321,14 +5337,15 @@ def forward( # per-capture constant matching the engine's short/long decode graph families. sparse = batch > 0 and int(md.kv_lens_runtime[:batch].max()) > core.dense_threshold + qsa_md = md.ext(QSAMetadata) qsa = ( - md.qsa_req_idx_per_token[:num_tokens], - md.qsa_logical_positions[:num_tokens], - md.qsa_sequence_lengths[:num_tokens], - md.qsa_visible_blocks[:num_tokens], - md.qsa_topk_indices[:num_tokens], - md.qsa_topk_row_starts[:num_tokens], - md.qsa_block_table, + qsa_md.req_idx_per_token[:num_tokens], + qsa_md.logical_positions[:num_tokens], + qsa_md.sequence_lengths[:num_tokens], + qsa_md.visible_blocks[:num_tokens], + qsa_md.topk_indices[:num_tokens], + qsa_md.topk_row_starts[:num_tokens], + qsa_md.block_table, md.tokens_per_block, pos3, coords, @@ -5348,9 +5365,9 @@ def forward( num_gens=batch, rows_per_gen=gen_p, slots=si_all, - request_indices=md.qsa_req_idx_per_token[:num_tokens], - logical_positions=md.qsa_logical_positions[:num_tokens], - block_table=md.qsa_block_table, + request_indices=qsa_md.req_idx_per_token[:num_tokens], + logical_positions=qsa_md.logical_positions[:num_tokens], + block_table=qsa_md.block_table, tokens_per_block=md.tokens_per_block, ) @@ -5470,20 +5487,21 @@ def _capture_selection(self, md, topk_indices, visible_blocks, batch, p) -> None row with that row's complete-group count, in the draft layer's row of the shared buffers. Every request carries `p` rows here, and the worker continues the draft loop from row `num_accepted - 1` of each (the golden token counts as accepted).""" - assert md.qsa_mtp_num_accepted is not None, ( + qsa = md.find_ext(QSAMetadata) + assert qsa.mtp_num_accepted is not None, ( "IndexShare needs the worker's accepted counts to pick each request's capture " "row; the draft loop did not supply them" ) - accepted = md.qsa_mtp_num_accepted[:batch].to(torch.int64) + accepted = qsa.mtp_num_accepted[:batch].to(torch.int64) rows = torch.arange(batch, dtype=torch.int64, device=topk_indices.device) * p rows = rows + torch.clamp(accepted - 1, min=0, max=p - 1) slot = self.core._draft_bound.share_slot - md.qsa_shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) - md.qsa_shared_topk_visible_blocks[slot, :batch].copy_( + qsa.shared_topk_indices[slot, :batch].copy_(torch.index_select(topk_indices, 0, rows)) + qsa.shared_topk_visible_blocks[slot, :batch].copy_( torch.index_select(visible_blocks, 0, rows) ) # Host-side readiness; see PrefillTarget._capture_selection. - md.qsa_shared_topk_captured_slots.add(slot) + qsa.shared_topk_captured_slots.add(slot) def mtp_forward( self, @@ -5506,6 +5524,7 @@ def mtp_forward( "core before the worker, so this cannot be reached first" ) md = attn_metadata + qsa = md.find_ext(QSAMetadata) rows = hidden_states.shape[0] assert input_ids.shape[0] == rows, ( f"the draft step's {input_ids.shape[0]} token ids and {rows} " @@ -5552,19 +5571,17 @@ def mtp_forward( ) # The draft layer's QSA planes and the page table that addresses them. With a - # separate draft manager the engine's QSA metadata swaps `qsa_block_table` for the - # table it decodes from the draft manager's offsets (`prepare_for_draft_forward`); + # separate draft manager the engine's QSA metadata swaps `block_table` for the + # table it decodes from the draft manager's offsets (`QSAMetadata.draft_swap`); # a table still decoding the trunk's offsets would read and write another pool's # pages. draft = core._draft_bound if draft is None or cm is not draft.cm: core._bind_draft_cache_manager(md) draft = core._draft_bound - assert cm is core._bound.cm or md.qsa_block_table is getattr( - md, "qsa_draft_block_table", None - ), ( + assert cm is core._bound.cm or qsa.block_table is qsa.draft_block_table, ( "the draft step is bound to a separate draft KV cache manager but " - "qsa_block_table is not that manager's table: the engine's QSA metadata did " + "the QSA block_table is not that manager's table: the engine's QSA metadata did " "not swap its page table for the draft forward" ) # The `qsa_*` row mapping is sized from the metadata's token count, which the @@ -5581,10 +5598,10 @@ def mtp_forward( # IndexShare: the loop's draft-extend step (`p` rows per request) runs the indexer # and, with sharing on, captures each request's selection; every draft decode step # after it (one row per request) reuses that capture instead of running the indexer. - share = core.mtp_index_share and md.qsa_in_mtp_draft_loop - reuse = share and md.qsa_indexer_skip_topk + share = core.mtp_index_share and qsa.in_mtp_draft_loop + reuse = share and qsa.indexer_skip_topk assert not share or reuse == (p == 1), ( - f"the worker's reuse flag ({md.qsa_indexer_skip_topk}) disagrees with the " + f"the worker's reuse flag ({qsa.indexer_skip_topk}) disagrees with the " f"step's shape ({p} rows per request): only the draft decode steps after the " "draft-extend step reuse its selection" ) @@ -5663,9 +5680,9 @@ def mtp_forward( mrope_section1=core.mrope_section1, mrope_section2=core.mrope_section2, ) - req_idx = md.qsa_req_idx_per_token[:rows] - logical = md.qsa_logical_positions[:rows] - block_table = md.qsa_block_table + req_idx = qsa.req_idx_per_token[:rows] + logical = qsa.logical_positions[:rows] + block_table = qsa.block_table tokens_per_block = md.tokens_per_block index_cache = draft.index_cache # The QSA index branch, as in the trunk's `_attention_block` (see the prefill @@ -5737,20 +5754,20 @@ def mtp_forward( if sparse: kv = draft.kv _store_kv(core, kv, k_dst, v_dst, req_idx, logical, block_table, tokens_per_block) - topk_indices = md.qsa_topk_indices[:rows] - visible_blocks = md.qsa_visible_blocks[:rows] + topk_indices = qsa.topk_indices[:rows] + visible_blocks = qsa.visible_blocks[:rows] if reuse: # The capture this loop's draft-extend step made, with the groups completed # since appended; the expansion then adds the incomplete tail, so the # positions drafted since the capture stay visible. slot = draft.share_slot - assert slot in md.qsa_shared_topk_captured_slots, ( + assert slot in qsa.shared_topk_captured_slots, ( "a draft decode step would reuse a selection that no draft-extend step " "of this loop captured" ) qsa_mtp_reuse_topk( - shared_indices=md.qsa_shared_topk_indices[slot], - shared_visible_blocks=md.qsa_shared_topk_visible_blocks[slot], + shared_indices=qsa.shared_topk_indices[slot], + shared_visible_blocks=qsa.shared_topk_visible_blocks[slot], request_indices=req_idx, visible_blocks=visible_blocks, output=topk_indices, @@ -5767,13 +5784,13 @@ def mtp_forward( only_visible_blocks=True, context_rows=False, ) - _select(core, scores, md.qsa_topk_row_starts[:rows], visible_blocks, topk_indices) + _select(core, scores, qsa.topk_row_starts[:rows], visible_blocks, topk_indices) if share: self._capture_selection(md, topk_indices, visible_blocks, gen_seqs, p) selected = qsa_expand_block_indices( topk_indices, logical, - md.qsa_sequence_lengths[:rows], + qsa.sequence_lengths[:rows], compress_ratio=core.compress_ratio, token_topk=core.token_topk, ) @@ -6072,8 +6089,8 @@ class MTPLayer: lengths stay within `dense_seq_len_threshold`, and selection plus sparse GQA above it. With a separate draft KV manager the layer's K/V, index and position pages live in that manager's pools: the worker swaps `attn_metadata.kv_cache_manager` and the block - offsets for the draft loop, `QSAAttentionMetadata.prepare_for_draft_forward` swaps - `qsa_block_table` to the table it decodes from the draft manager's offsets, and the + offsets for the draft loop, `QSAMetadata.draft_swap` swaps + `block_table` to the table it decodes from the draft manager's offsets, and the core binds the draft manager's planes once per manager (`_DraftBound`). The planes need no rollback: each engine step's first draft call rewrites every position from the newly committed length on. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py index fff162545bc7..f72f140dc070 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/ops/attention/qsa_paged_index_scores.py @@ -124,7 +124,7 @@ class QsaPagedIndexScores(OpWrapper): note = """ Metadata consumed: none -- every input is an argument. What the target - passes: `block_table` is `QSAAttentionMetadata.qsa_block_table` (raw + passes: `block_table` is `QSAMetadata.block_table` (raw pool-page ids, int32, -1 past each sequence); `query_positions` / `request_indices` are `qsa_logical_positions` / `qsa_req_idx_per_token`; `index_cache` is the layer's bf16 index-K plane (`get_index_k_buffer`). diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 17cc33280f11..ed9d1800800c 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -217,7 +217,7 @@ registration alone. | Algorithm | Sparse backend | Metadata | Cache manager | |---|---|---|---| -| `csa2` | `CSA2TrtllmAttention`, with hardware-selected Flash helpers | `CSA2TrtllmMetadata` | `CSA2CacheManager` (V2) | +| `csa2` | `CSA2TrtllmAttention`, with hardware-selected Flash helpers | `TrtllmAttentionMetadata` + `CSA2Metadata` | `CSA2CacheManager` (V2) | CSA2 implements DeepSeek-V4.1 compressed sparse MLA. Its module adapter owns compression and shared KV/index sources; per-phase query tiles feed selected @@ -282,7 +282,8 @@ All backends implement the `AttentionBackend` interface. The core contract is: - `forward(q, k, v, metadata, forward_args=..., **kwargs)` -- `Metadata` subtype +- `Metadata` core family and `metadata_extensions`; build metadata with + `Backend.create_metadata(...)` so the backend's extensions are attached - `AttentionForwardArgs` for per-forward optional arguments such as masks, output buffers, scales, RoPE/mRoPE inputs, MLA buffers, and sparse inputs - coarse capability hooks: @@ -352,14 +353,35 @@ metadata and cache behavior. All backend metadata types inherit from `AttentionMetadata`. The base contract includes sequence-length and request-level state, KV-cache manager and parameters, runtime feature flags, and optional CUDA-graph buffer management. -Sparse metadata subtypes consume only `SparseMetadataParams`, not -backend-owned `SparseParams`. + +Everything a feature, a sparse algorithm or a model adds on top of a core +family is a `MetadataExtension` (`backends/extensions.py`) attached to the +metadata, not a metadata subclass. Callers reach one with `metadata.ext(T)` +(raises when absent) or `metadata.find_ext(T)` (returns None). The metadata +drives every extension through fixed hooks: `allocate()` on creation and on +every CUDA-graph copy, `prepare_scope()` / `pre_derive_lengths()` / +`override_kv_lens()` / `prepare()` / `post_derive()` around the core's +`prepare()`, `on_kv_lens_changed()` after device-side length updates, +`draft_swap()` / `draft_restore()` around a draft forward, and the +attention-DP and workspace hooks. Extensions come from three places: + +- the core family's own configuration (`_default_extension_types()`), such as + `HelixMetadata`, `DraftKVMetadata`, `Fp4MlaState`, `SpecDecodingMetadata` + and `RecurrentMetadata`; +- the backend's `metadata_extensions`, such as a sparse algorithm's + `SparseAlgorithmExtension` (at most one per metadata); and +- the caller's `extension_types`, such as `FlashMlaMetadata`. + +Sparse algorithm extensions consume only `SparseMetadataParams`, not +backend-owned `SparseParams`. Algorithms that share one MTP draft-loop +selection subclass `MTPIndexShareExtension`. **`TrtllmAttentionMetadata`** is the main metadata family. It adds paged-KV -block information, TRTLLM runtime state, chunked-prefill/speculative-decode/Helix -state, and MLA-specific state. If a source attention needs paged KV, chunked -prefill, FlashMLA, speculative decoding, or Helix-aware execution, the fit -question is mostly a `TrtllmAttentionMetadata` fit question. +block information, TRTLLM runtime state, and MLA-specific state; its features +(speculative decoding, Helix, FlashMLA, FP4 MLA, separate draft KV cache) are +extensions. If a source attention needs paged KV, chunked prefill, FlashMLA, +speculative decoding, or Helix-aware execution, the fit question is mostly a +`TrtllmAttentionMetadata` fit question. **`VanillaAttentionMetadata`** is lighter — base metadata plus simple cache-index information. Use it when the `Attention` module boundary fits but @@ -368,9 +390,9 @@ the fused TRTLLM path is too restrictive. **`FlashInferAttentionMetadata`** adds a planning-oriented contract with workspace, page-table KV metadata, and prefill/decode wrapper state. -**Sparse metadata** families extend the base backend metadata with -sparse-specific runtime state (indexer buffers, routing state, side-cache -state). +**Sparse metadata** is a `SparseAlgorithmExtension` on a core family that +holds sparse-specific runtime state (indexer buffers, routing state, +side-cache state). `SparseRuntimeParams` is the backend-to-`AttentionOp` carrier only on the `TrtllmAttention` path. Its fields are algorithm-specific, not a generic @@ -438,7 +460,7 @@ reuse is enabled (Encoder SWA is `PRIVATE` when reuse is disabled): | Decoder SWA | Per layer; same row format | `PRIVATE` for eligible bounded replay; ordinary lifecycle for full prefill | | Ratio-two compressor state | Source owners; FP32 partial-group values/scores | Managed with the owner's page lifecycle | -`CSA2TrtllmMetadata.prepare()` publishes owner tables, compression inputs, +Preparing `CSA2Metadata` publishes owner tables, compression inputs, visibility and SWA read/write slots. Live accepted-length updates refresh captured buffers in place. Graph clones share reserved index arenas and query tiles; eager shape eviction must preserve every graph-owned allocation. diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index ad315b78084f..031536acb2e4 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -8,6 +8,8 @@ import torch from torch import nn +from tensorrt_llm._torch.attention.backends.features.spec_decoding import \ + SpecDecodingMetadata from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -1161,13 +1163,11 @@ def forward( # For dynamic tree spec decoding with Python RoPE, adjust position_ids # to use tree offsets (same as C++ kernel: past_seq_len + offset). - if (not self.rope_fusion - and getattr(attn_metadata, 'is_spec_dec_dynamic_tree', False) - and getattr(attn_metadata, 'use_spec_decoding', False) - and getattr(attn_metadata, 'spec_decoding_position_offsets', - None) is not None - and attn_metadata.spec_decoding_position_offsets.dim() == - 1 # 1D layout ⇒ dynamic tree + spec = attn_metadata.find_ext(SpecDecodingMetadata) + if (not self.rope_fusion and spec is not None + and spec.is_spec_dec_dynamic_tree and spec.use_spec_decoding + and spec.position_offsets is not None + and spec.position_offsets.dim() == 1 # 1D layout ⇒ dynamic tree and position_ids is not None): position_ids = self._adjust_position_ids_for_spec_dec( position_ids, attn_metadata) @@ -1255,8 +1255,9 @@ def _adjust_position_ids_for_spec_dec(self, position_ids, attn_metadata): base_pos = attn_metadata.kv_lens_cuda[num_contexts:num_contexts + num_gens] - gen_len total = num_gens * gen_len - offsets = attn_metadata.spec_decoding_position_offsets[:total].view( - num_gens, gen_len) + offsets = attn_metadata.ext( + SpecDecodingMetadata).position_offsets[:total].view( + num_gens, gen_len) start = attn_metadata.num_ctx_tokens end = start + total adjusted = (base_pos.unsqueeze(1) + offsets).reshape(-1) diff --git a/tensorrt_llm/_torch/attention/backends/extensions.py b/tensorrt_llm/_torch/attention/backends/extensions.py new file mode 100644 index 000000000000..bd2c8c286b0e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/extensions.py @@ -0,0 +1,331 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Feature state attached to attention metadata. + +An ``AttentionMetadata`` holds the batch layout every attention backend needs: +request counts, sequence and KV lengths, the paged KV cache and the workspace. +Everything a feature, a sparse algorithm or a model adds on top lives in a +``MetadataExtension`` attached to it. The metadata drives each extension +through the hooks below at fixed points of its lifecycle, and callers reach an +extension's state with ``metadata.ext(T)`` or ``metadata.find_ext(T)``. + +Per-step order, all driven by the metadata: + +1. ``prepare_scope()`` wraps the whole prepare. +2. ``pre_derive_lengths()`` runs before the core derives KV lengths. +3. ``override_kv_lens(kv_lens)`` may rewrite the host KV lengths the core + just derived, before they are published to the device. +4. ``prepare()`` runs after the core, in dependency order. +5. ``post_derive()`` runs after every extension prepared. + +Outside ``prepare()``: ``allocate()`` on creation and on every CUDA graph or +cross-attention copy, ``begin_forward()`` before the model forward, +``on_kv_lens_changed()`` after device-side length changes, +``on_spec_dec_param()`` after the speculative-decoding parameters change, +``invalidate_length_derived_buffers()`` whenever the core invalidates its own +length prefixes, and ``draft_swap()`` / ``draft_restore()`` around a +draft-model forward. ``token_major_gen_view()`` and ``mla_scheduler_buffers()`` +let an extension replace the core's generation-row presentation. +""" + +from __future__ import annotations + +import contextlib +import copy +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + ContextManager, + Dict, + FrozenSet, + Iterable, + List, + Optional, + Sequence, + Tuple, + Type, + TypeVar, +) + +if TYPE_CHECKING: + import torch + + from ...memory_buffer_utils import Buffers + from ...pyexecutor.llm_request import LlmRequest + from ...speculative.spec_tree_manager import SpecTreeManager + from .interface import AttentionMetadata + from .sparse.dsa.metadata import TokenMajorGenView + +E = TypeVar("E", bound="MetadataExtension") + + +class MetadataExtension: + """State that one feature, sparse algorithm or model attaches to a metadata. + + Subclasses keep their state as attributes and override only the hooks they + need; every hook defaults to doing nothing. ``self.metadata`` is the + metadata the extension is attached to; read the core through it and never + assign the core's private fields. + """ + + #: Extensions that must be attached, and prepared, before this one. They + #: are attached automatically when this one is. + requires: ClassVar[Tuple[Type["MetadataExtension"], ...]] = () + #: Tensor attributes produced and consumed inside a captured graph; the + #: strict CUDA-graph buffer check skips them. + graph_temporary_attrs: ClassVar[FrozenSet[str]] = frozenset() + + def __init__(self, metadata: "AttentionMetadata") -> None: + self.metadata = metadata + + def clone_for(self, metadata: "AttentionMetadata") -> "MetadataExtension": + """Copy this extension for a CUDA graph or cross-attention copy. + + The default shallow copy shares every attribute with this extension; + ``allocate()`` then runs on the copy, so graph-stable buffers allocated + there are rebound per copy. Override to give a copy fresh per-copy + state that ``allocate()`` does not rebuild. + """ + clone = copy.copy(self) + clone.metadata = metadata + return clone + + def allocate(self, buffers: Optional["Buffers"]) -> None: + """Allocate persistent buffers; runs again on every metadata copy.""" + + def prepare_scope(self) -> ContextManager[None]: + """A context the whole ``prepare()`` runs in.""" + return contextlib.nullcontext() + + def pre_derive_lengths(self) -> None: + """Adjust the per-request inputs before the core derives KV lengths.""" + + def override_kv_lens(self, kv_lens: "torch.Tensor") -> None: + """Rewrite the host KV lengths in place before they are published. + + ``kv_lens`` holds the per-sequence KV lengths, without the extra KV + tokens, for the whole batch. + """ + + def prepare(self) -> None: + """Prepare this extension's per-step state after the core.""" + + def post_derive(self) -> None: + """Adjust the core's published views after every extension prepared.""" + + def begin_forward(self) -> None: + """Reset per-forward scratch state before the model forward.""" + + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: + """React to device-side KV-length changes. + + ``spec_dec_step`` is True when a speculative-decoding sub-step changed + the lengths (``update_for_spec_dec``) and False for an overlap or + draft-correction update (``on_update_kv_lens``). + """ + + def on_spec_dec_param( + self, + max_draft_len: int, + max_total_draft_tokens: int, + spec_tree_manager: Optional["SpecTreeManager"], + ) -> None: + """React to new speculative-decoding parameters. + + Runs after the core stored its own spec-dec fields, with the arguments + of ``update_spec_dec_param`` the core does not keep. + """ + + def invalidate_length_derived_buffers(self) -> None: + """Mark buffers derived from this step's lengths stale. + + Called whenever the core invalidates its own layer-invariant length + prefixes: at the start of every prepare and on each KV-length change. + """ + + def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: + """A one-row-per-query-token generation presentation, or None. + + The core presents the first non-None view to the MLA RoPE and + sparse-MLA generation ops. + """ + return None + + def mla_scheduler_buffers( + self, num_heads: int + ) -> Optional[Tuple["torch.Tensor", "torch.Tensor"]]: + """Cumulative MLA generation Q rows / KV lengths replacing the core's, or None.""" + return None + + def on_spec_dec_restored(self) -> None: + """Rebind views after ``restore_from_spec_dec`` restored the core tensors.""" + + def needs_full_block_table_width(self) -> bool: + """Whether every block-table column must be staged each step. + + True when device-side updates may advance KV lengths past the host + snapshot, so the core cannot cap the staged width at the batch maximum. + """ + return False + + def draft_swap(self) -> Any: + """Point this extension at the draft model's state; return what to restore.""" + return None + + def draft_restore(self, saved: Any) -> None: + """Undo ``draft_swap``.""" + + def swap_attrs_for_draft(self, draft_values: Dict[str, Any]) -> Dict[str, Any]: + """Rebind attributes to their draft values; return the originals. + + The originals are also kept in ``draft_replay_swapped_attrs`` until + ``restore_swapped_attrs``, so the strict CUDA-graph buffer check + validates the target tensors while a draft forward is bound. + """ + swapped = self.__dict__.setdefault("draft_replay_swapped_attrs", {}) + saved = {} + for name, value in draft_values.items(): + saved[name] = getattr(self, name, None) + swapped.setdefault(name, saved[name]) + setattr(self, name, value) + return saved + + def restore_swapped_attrs(self, saved: Dict[str, Any]) -> None: + """Undo ``swap_attrs_for_draft``.""" + for name, value in saved.items(): + setattr(self, name, value) + self.__dict__["draft_replay_swapped_attrs"] = {} + + def adp_token_counts(self) -> Optional[List[int]]: + """Local counts of the extra stages this extension adds to the attention-DP exchange.""" + return None + + def set_adp_token_counts(self, counts: List[List[int]], input_counts: List[int]) -> None: + """Receive every rank's counts for the stages ``adp_token_counts`` added. + + ``input_counts`` holds every rank's input-stage count. + """ + + def workspace_bytes(self) -> int: + """Extra attention workspace this extension needs.""" + return 0 + + def kv_cache_pools(self) -> Optional[Tuple["torch.Tensor", "torch.Tensor", int]]: + """The KV pools the attention kernels read instead of the manager's, or None. + + Returns host pool pointers, host pool mapping and tokens per block. + """ + return None + + def on_runtime_views_bound(self) -> None: + """React to the core rebinding its per-forward ``*_runtime`` views.""" + + def bind_context_requests( + self, requests: List["LlmRequest"], *, prefill_cuda_graphs: bool + ) -> None: + """Receive the scheduled context requests before ``prepare()``. + + ``prefill_cuda_graphs`` tells whether context forwards may run in CUDA + graphs. + """ + + def defer_decode_outputs(self) -> Optional[ContextManager[None]]: + """A scope for preparing a decode CUDA graph's metadata, or None. + + Inside it, outputs that the graph's own preprocessing derives after + ``on_update_kv_lens`` may be left unprepared. + """ + return None + + +def resolve_extension_order( + types: Iterable[Type[MetadataExtension]], +) -> List[Type[MetadataExtension]]: + """Order extension types so every requirement precedes its dependents. + + Requirements that are not listed are added. Listing order is otherwise + preserved. + """ + ordered: List[Type[MetadataExtension]] = [] + visiting: set = set() + + def visit(cls: Type[MetadataExtension]) -> None: + if cls in ordered: + return + if cls in visiting: + raise ValueError(f"Metadata extension requirements form a cycle at {cls.__name__}.") + visiting.add(cls) + for required in cls.requires: + visit(required) + visiting.discard(cls) + ordered.append(cls) + + for cls in types: + visit(cls) + return ordered + + +class ExtensionRegistry: + """The extensions attached to one metadata, in dependency order.""" + + def __init__(self, owner: Optional["AttentionMetadata"] = None) -> None: + self.owner = owner + self._by_type: Dict[Type[MetadataExtension], MetadataExtension] = {} + + def __iter__(self): + return iter(self._by_type.values()) + + def __len__(self) -> int: + return len(self._by_type) + + def __contains__(self, cls: Type[MetadataExtension]) -> bool: + return self.find(cls) is not None + + def types(self) -> List[Type[MetadataExtension]]: + return list(self._by_type) + + def find(self, cls: Type[E]) -> Optional[E]: + found = self._by_type.get(cls) + if found is not None: + return found + for extension in self._by_type.values(): + if isinstance(extension, cls): + return extension + return None + + def attach( + self, metadata: "AttentionMetadata", types: Sequence[Type[MetadataExtension]] + ) -> List[MetadataExtension]: + """Attach the given types and their requirements; return the new ones.""" + added = [] + for cls in resolve_extension_order(types): + if cls not in self._by_type: + extension = cls(metadata) + self._by_type[cls] = extension + added.append(extension) + return added + + def replace(self, extension: MetadataExtension) -> None: + """Put an already-built extension in place of any of its type.""" + self._by_type[type(extension)] = extension + + def remove(self, cls: Type[MetadataExtension]) -> None: + """Detach the extension of type ``cls``, if any.""" + found = self.find(cls) + if found is not None: + del self._by_type[type(found)] + + def bound_to(self, metadata: "AttentionMetadata") -> "ExtensionRegistry": + """This registry, owned by ``metadata``. + + Returns ``self`` when ``metadata`` already owns it; otherwise a new + registry holding a ``clone_for`` copy of each extension, so a shallow + copy of the metadata never shares the registry or its extensions with + its source. + """ + if self.owner is metadata: + return self + clone = ExtensionRegistry(metadata) + clone._by_type = {cls: e.clone_for(metadata) for cls, e in self._by_type.items()} + return clone diff --git a/tensorrt_llm/_torch/attention/backends/features/__init__.py b/tensorrt_llm/_torch/attention/backends/features/__init__.py new file mode 100644 index 000000000000..21861ec76bc3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/features/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Feature extensions of the attention metadata; see ``..extensions``.""" + +from .draft_kv import DraftKVMetadata +from .flash_mla import FlashMlaMetadata +from .helix import HelixMetadata, find_helix +from .recurrent import RecurrentMetadata + +__all__ = [ + "DraftKVMetadata", + "FlashMlaMetadata", + "HelixMetadata", + "RecurrentMetadata", + "find_helix", +] diff --git a/tensorrt_llm/_torch/attention/backends/features/draft_kv.py b/tensorrt_llm/_torch/attention/backends/features/draft_kv.py new file mode 100644 index 000000000000..078f5f5392a9 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/features/draft_kv.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Separate draft KV cache state of the attention metadata.""" + +from typing import Optional + +import torch + +from ..extensions import MetadataExtension + + +class DraftKVMetadata(MetadataExtension): + """Block offsets of the draft model's KV cache. + + One-model speculative decoding with a separate draft KV cache layout keeps + the draft cache's block offsets here, refreshed with the target's every + step, so a draft forward can swap them in without staging. The swap itself + is ``prepare_attn_metadata_for_draft_replay`` in the speculative package. + """ + + def __init__(self, metadata) -> None: + super().__init__(metadata) + self.kv_cache_block_offsets: Optional[torch.Tensor] = None + + def allocate(self, buffers) -> None: + md = self.metadata + manager = md.draft_kv_cache_manager + num_pools = getattr(manager, "num_attention_op_pools", manager.num_pools) + self.kv_cache_block_offsets = md.get_empty( + buffers, + [num_pools, md.max_num_sequences, 2, manager.max_blocks_per_seq], + cache_name="draft_kv_cache_block_offsets", + dtype=torch.int32, + capture_graph=md.is_cuda_graph, + ) + + def needs_full_block_table_width(self) -> bool: + # Draft sub-steps advance kv_lens_cuda on device past the host snapshot. + return self.metadata.draft_kv_cache_manager is not None + + def stage_block_offsets(self, max_blocks: Optional[int]) -> None: + """Copy this step's draft block offsets, as wide as the target's.""" + md = self.metadata + if md.draft_kv_cache_manager is None: + return + md.draft_kv_cache_manager.copy_batch_block_offsets( + self.kv_cache_block_offsets, + md.request_ids, + md.beam_width, + md.num_contexts, + md.num_seqs, + max_blocks=max_blocks, + ) diff --git a/tensorrt_llm/_torch/attention/backends/features/flash_mla.py b/tensorrt_llm/_torch/attention/backends/features/flash_mla.py new file mode 100644 index 000000000000..3314f14c7b6d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/features/flash_mla.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FlashMLA state of the attention metadata.""" + +from typing import Optional + +import torch + +from tensorrt_llm._utils import maybe_pin_memory +from tensorrt_llm.bindings.internal import thop + +from ..extensions import MetadataExtension + + +class FlashMlaMetadata(MetadataExtension): + """Block IDs and tile-scheduler metadata of the FlashMLA generation kernels. + + Provides the core's ``block_ids_per_seq`` (generation rows, compacted) and + ``kv_block_ids_per_seq`` (every row), plus draft-cache twins when a separate + draft KV cache exists. The tile-scheduler metadata is computed once per + forward by the first generation layer and reused until the KV lengths or + block IDs change. + """ + + def __init__(self, metadata) -> None: + super().__init__(metadata) + self.draft_block_ids_per_seq: Optional[torch.Tensor] = None + self.draft_kv_block_ids_per_seq: Optional[torch.Tensor] = None + self.tile_scheduler_metadata: Optional[torch.Tensor] = None + self.num_splits: Optional[torch.Tensor] = None + self.metadata_valid = False + + def allocate(self, buffers) -> None: + md = self.metadata + manager = md.kv_cache_manager + if manager is None: + return + capture_graph = md.is_cuda_graph + + def block_ids(owner, name): + return md.get_empty( + buffers, + [owner.max_batch_size, owner.max_blocks_per_seq], + cache_name=name, + dtype=torch.int32, + capture_graph=capture_graph, + ) + + md.block_ids_per_seq = block_ids(manager, "block_ids_per_seq") + md.kv_block_ids_per_seq = block_ids(manager, "kv_block_ids_per_seq") + if md.draft_kv_cache_manager is not None: + self.draft_block_ids_per_seq = block_ids( + md.draft_kv_cache_manager, "draft_block_ids_per_seq" + ) + self.draft_kv_block_ids_per_seq = block_ids( + md.draft_kv_cache_manager, "draft_kv_block_ids_per_seq" + ) + # Pre-allocated so their GPU addresses are stable across CUDA graph captures. + sm_count = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + self.tile_scheduler_metadata = md.get_empty( + buffers, + [sm_count * 8], # TileSchedulerMetaDataSize = 8 + cache_name="flash_mla_tile_scheduler_metadata", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.num_splits = md.get_empty( + buffers, + [manager.max_batch_size + 1], + cache_name="flash_mla_num_splits", + dtype=torch.int32, + capture_graph=capture_graph, + ) + + def prepare(self) -> None: + self.prepare_block_ids() + + def prepare_block_ids(self) -> None: + """Stage this batch's block IDs from the current KV cache manager.""" + md = self.metadata + self.metadata_valid = False + # Request-specific fills and H2D copies must happen before replay, not + # become fixed operations in the captured graph. + if torch.cuda.is_current_stream_capturing(): + return + + block_ids_per_seq = maybe_pin_memory( + md.kv_cache_manager.get_block_ids_per_seq(md.request_ids) + ) + num_blocks = block_ids_per_seq.shape[1] + md.kv_block_ids_per_seq.fill_(0) + md.kv_block_ids_per_seq[: md.num_seqs, :num_blocks].copy_( + block_ids_per_seq, non_blocking=True + ) + md.block_ids_per_seq.fill_(0) + md.block_ids_per_seq[: md.num_generations, :num_blocks].copy_( + block_ids_per_seq[md.num_contexts :], non_blocking=True + ) + + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: + self.metadata_valid = False + + def compute_scheduler_metadata(self, num_heads: int, kv_lora_rank: int) -> None: + """Compute the tile-scheduler metadata once per forward.""" + md = self.metadata + if self.metadata_valid: + return + num_generations = md.num_generations + if num_generations > 0: + generation_slice = slice(md.num_contexts, md.num_contexts + num_generations) + s_q = int(md.seq_lens[md.num_contexts].item()) + thop.compute_flash_mla_metadata( + md.kv_lens_cuda_runtime[generation_slice], + self.tile_scheduler_metadata, + self.num_splits[: num_generations + 1], + num_generations, + s_q, + num_heads, + 1, + kv_lora_rank, + ) + self.metadata_valid = True + + def draft_swap(self): + md = self.metadata + if self.draft_block_ids_per_seq is None or self.draft_kv_block_ids_per_seq is None: + raise RuntimeError( + "FlashMLA separate draft KV cache requires dedicated draft block-ID buffers" + ) + md.swap_for_draft("block_ids_per_seq", self.draft_block_ids_per_seq) + md.swap_for_draft("kv_block_ids_per_seq", self.draft_kv_block_ids_per_seq) + self.prepare_block_ids() + return None + + def draft_restore(self, saved) -> None: + # Target and draft block-ID buffers are independent. Restoring only + # needs to invalidate the scheduler metadata; refreshing the unchanged + # target buffers would repeat request-specific H2D work. + self.metadata_valid = False diff --git a/tensorrt_llm/_torch/attention/backends/features/helix.py b/tensorrt_llm/_torch/attention/backends/features/helix.py new file mode 100644 index 000000000000..815806af4368 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/features/helix.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Helix context-parallel state of the attention metadata.""" + +from typing import List, Optional + +import torch + +from tensorrt_llm._utils import prefer_pinned + +from ....utils import helix_local_len_tensor +from ..extensions import MetadataExtension +from ..fp4_mla.metadata import populate_helix_spec_buffers + + +def find_helix(metadata) -> Optional["HelixMetadata"]: + """The helix extension of ``metadata``, or None. + + Accepts the duck-typed metadata views the FP4 MLA helpers take, which may + carry no extensions at all. + """ + find_ext = getattr(metadata, "find_ext", None) + return find_ext(HelixMetadata) if find_ext is not None else None + + +class HelixMetadata(MetadataExtension): + """Per-token global positions and KV ownership under helix parallelism. + + Each helix rank holds a subset of a sequence's tokens. Only the active rank + appends the KV of a new token and attends to it; inactive ranks attend to + the previously cached tokens only. + + The per-sequence and per-token buffers are GENERATION-relative: the + packing loops in the executor append only for extend and plain generation + rows, so ``update`` writes exactly ``[0, num_generations)``. + """ + + def __init__(self, metadata) -> None: + super().__init__(metadata) + # Global position id of each token of each sequence. + self.position_offsets: Optional[torch.Tensor] = None + self.position_offsets_cpu: Optional[torch.Tensor] = None + # Whether this rank is inactive for each generation sequence. + self.is_inactive_rank: Optional[torch.Tensor] = None + self.is_inactive_rank_cpu: Optional[torch.Tensor] = None + # Per-token state for speculative verify groups: a 1 + draft_len group + # may straddle a ledger-page boundary onto two CP ranks, so the + # per-sequence flag above is not expressive enough. See + # recompute_spec_buffers for the derivation. + # local_slots[t]: rank-local KV write slot of generation token t on + # this rank, or -1 when another rank owns its position (consumed + # by the mla_rope_generation append kernel). + # kv_bounds[t]: number of rank-local KV entries token t may attend + # to, i.e. local_len(pos_t + 1) (consumed by the CuTe DSL MLA + # decode mask and the helix stats identity). + self.local_slots: Optional[torch.Tensor] = None + self.kv_bounds: Optional[torch.Tensor] = None + # Host per-sequence count of this step's new tokens owned by this rank + # (spec path; the single-token path derives it from the boolean flag). + self.owned_new_tokens_cpu: Optional[torch.Tensor] = None + self.spec_tokens_valid = False + + def allocate(self, buffers) -> None: + md = self.metadata + capture_graph = md.is_cuda_graph + self.position_offsets = md.get_empty( + buffers, + (md.max_num_tokens,), + cache_name="helix_position_offsets", + dtype=torch.int, + capture_graph=capture_graph, + ) + self.position_offsets_cpu = torch.empty_like( + self.position_offsets, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.is_inactive_rank = md.get_empty( + buffers, + (md.max_num_sequences,), + cache_name="helix_is_inactive_rank", + dtype=torch.bool, + capture_graph=capture_graph, + ) + self.is_inactive_rank_cpu = torch.empty_like( + self.is_inactive_rank, + device="cpu", + pin_memory=prefer_pinned(), + ) + self.local_slots = md.get_empty( + buffers, + (md.max_num_tokens,), + cache_name="helix_local_slots", + dtype=torch.int, + capture_graph=capture_graph, + ) + self.kv_bounds = md.get_empty( + buffers, + (md.max_num_tokens,), + cache_name="helix_kv_bounds", + dtype=torch.int, + capture_graph=capture_graph, + ) + self.owned_new_tokens_cpu = torch.zeros( + (md.max_num_sequences,), + device="cpu", + dtype=torch.int, + pin_memory=prefer_pinned(), + ) + self.spec_tokens_valid = False + + def update( + self, + position_offsets: List[int], + is_inactive_rank: List[bool], + owned_new_tokens: Optional[List[int]] = None, + ) -> None: + """Copy this step's helix parameters into the static buffers. + + Args: + position_offsets: Global position of each token, shape (num_tokens,). + is_inactive_rank: Whether this rank is inactive for each + generation sequence, shape (num_generations,). + owned_new_tokens: Per-sequence count of this step's new tokens + owned by this rank (speculative verify groups; one group may + straddle a page boundary onto two ranks). None on the + single-token path, where the boolean flag carries it. + """ + if position_offsets is not None: + num_tokens = len(position_offsets) + self.position_offsets_cpu[:num_tokens].copy_( + torch.tensor(position_offsets, dtype=torch.int) + ) + self.position_offsets[:num_tokens].copy_( + self.position_offsets_cpu[:num_tokens], non_blocking=True + ) + + if is_inactive_rank is not None: + batch_size = len(is_inactive_rank) + self.is_inactive_rank_cpu[:batch_size].copy_( + torch.tensor(is_inactive_rank, dtype=torch.bool) + ) + self.is_inactive_rank[:batch_size].copy_( + self.is_inactive_rank_cpu[:batch_size], non_blocking=True + ) + + self.spec_tokens_valid = False + if owned_new_tokens is not None: + batch_size = len(owned_new_tokens) + self.owned_new_tokens_cpu[:batch_size].copy_( + torch.tensor(owned_new_tokens, dtype=torch.int) + ) + self.spec_tokens_valid = True + + def override_kv_lens(self, kv_lens: torch.Tensor) -> None: + """Count only the new tokens this rank appends for generation rows. + + Context rows are not part of a verify group and are not packed into + the helix buffers; they append every one of their tokens. Generation + rows pair the generation-relative helix buffers with the generation + slice of the batch-indexed lengths. + """ + md = self.metadata + assert md.kv_cache_params.use_cache, "helix requires a KV cache" + num_gen = md.num_generations + gen = slice(md.num_contexts, md.num_seqs) + cached = kv_lens[gen] - md.seq_lens_kv[gen] + if self.spec_tokens_valid: + # Speculative verify groups: ownership of this step's new tokens is + # a per-sequence COUNT, not a boolean. Provisional host values; + # recompute_spec_buffers overrides the device copy after the + # overlap correction. + kv_lens[gen] = cached + self.owned_new_tokens_cpu[:num_gen] + else: + kv_lens[gen] = torch.where(self.is_inactive_rank_cpu[:num_gen], cached, kv_lens[gen]) + + def local_len_vec(self, global_lens: torch.Tensor) -> torch.Tensor: + """Vectorized rank-local prefix length for helix round-robin pages. + + For each global sequence length g, returns the number of the first g + tokens whose ledger page lives on this CP rank (page b -> rank + b % cp_size). The rule and its scalar twin live in ``_torch.utils``; + KVCacheManagerV2._helix_local_len and the host packing in the executor + use the same definition. + """ + md = self.metadata + return helix_local_len_tensor( + global_lens, + md.kv_cache_manager.tokens_per_block, + md.mapping.cp_size, + md.mapping.cp_rank, + ) + + def recompute_spec_buffers(self, num_gen_tokens: int, tokens_per_gen_seq: int) -> None: + """Derive per-token helix buffers from (corrected) global positions. + + Called after the overlap-scheduler device correction has been applied + to position_offsets, so every derived quantity reflects the real + committed length even though the host packed provisional values. + Static shapes only; safe under CUDA graph capture. + + ``tokens_per_gen_seq`` is the uniform verify-group width. Non-uniform + groups are rejected rather than silently mis-sliced: without the + overlap scheduler the extend loop packs a per-request + ``1 + get_draft_token_length(request)`` and a request entering with no + draft tokens is packed as a single-token generation row instead, so a + batch can arrive whose total happens to divide but whose rows do not + line up. + + Two index bases meet here, and they are not the same: + * position_offsets / local_slots / kv_bounds are GENERATION-RELATIVE. + * kv_lens_cuda is BATCH-indexed, hence the num_contexts offset on + the write below. + """ + md = self.metadata + assert num_gen_tokens % tokens_per_gen_seq == 0, ( + f"helix spec expects uniform verify groups: {num_gen_tokens} gen " + f"tokens not divisible by group size {tokens_per_gen_seq}" + ) + num_gen_seqs = num_gen_tokens // tokens_per_gen_seq + # Divisibility alone does not imply uniformity: a batch of mixed group + # widths can still divide and would then write the wrong number of + # kv_lens_cuda rows with values taken from the wrong tokens. + assert num_gen_seqs == md.num_generations, ( + f"helix spec expects uniform verify groups: {num_gen_tokens} gen " + f"tokens over {md.num_generations} generation rows do not all " + f"have width {tokens_per_gen_seq}" + ) + # One launch derives the write slots, the attention bounds and the + # per-sequence rank-local kv length (the bound of the sequence's last + # token: committed + owned in-flight tokens) from the positions; it + # replaces a chain of elementwise ops on the step's critical path. + # Static shapes and allocation-free, so it is safe under CUDA graph + # capture. + populate_helix_spec_buffers( + self.position_offsets[:num_gen_tokens], + self.local_slots[:num_gen_tokens], + self.kv_bounds[:num_gen_tokens], + md.kv_lens_cuda[md.num_contexts : md.num_contexts + num_gen_seqs], + num_tokens=num_gen_tokens, + tokens_per_seq=tokens_per_gen_seq, + tokens_per_block=md.kv_cache_manager.tokens_per_block, + cp_size=md.mapping.cp_size, + cp_rank=md.mapping.cp_rank, + ) diff --git a/tensorrt_llm/_torch/attention/backends/features/recurrent.py b/tensorrt_llm/_torch/attention/backends/features/recurrent.py new file mode 100644 index 000000000000..586f352b5005 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/features/recurrent.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Recurrent-layer state of the attention metadata of a hybrid model.""" + +from typing import TYPE_CHECKING, Optional, Type + +from ..extensions import MetadataExtension + +if TYPE_CHECKING: + from ....modules.mamba.mamba2_metadata import Mamba2Metadata + + +class RecurrentMetadata(MetadataExtension): + """Per-step metadata of the Mamba and linear-attention layers. + + Attached when the KV cache manager also holds recurrent state. The model + may choose a ``Mamba2Metadata`` subclass through ``mamba_metadata_cls``; + the instance is built on the first ``prepare()``. + """ + + def __init__(self, metadata) -> None: + super().__init__(metadata) + self.chunk_size = 128 + self.mamba_metadata_cls: Optional[Type["Mamba2Metadata"]] = None + self.mamba_metadata: Optional["Mamba2Metadata"] = None + + def prepare(self) -> None: + md = self.metadata + if self.mamba_metadata is None: + metadata_cls = self.mamba_metadata_cls + if metadata_cls is None: + from ....modules.mamba.mamba2_metadata import Mamba2Metadata + + metadata_cls = Mamba2Metadata + self.mamba_metadata = metadata_cls( + md.max_num_requests, self.chunk_size, max_num_tokens=md.max_num_tokens + ) + self.mamba_metadata.prepare(md) diff --git a/tensorrt_llm/_torch/attention/backends/features/spec_decoding.py b/tensorrt_llm/_torch/attention/backends/features/spec_decoding.py new file mode 100644 index 000000000000..bf5beb5e3ffd --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/features/spec_decoding.py @@ -0,0 +1,320 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding (multi-query generation) state of the TRTLLM metadata.""" + +import functools +import math +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._utils import get_sm_version + +from ..extensions import MetadataExtension + +if TYPE_CHECKING: + from ....speculative.interface import SpecMetadata + from ....speculative.spec_tree_manager import SpecTreeManager + + +def is_trtllm_gen_sm(sm: int) -> bool: + """Whether the TRTLLM attention op runs the trtllm-gen kernels on ``sm``.""" + return not (sm < 100 or sm in [120, 121]) + + +@functools.cache +def generate_spec_decoding_position_offsets(max_num_requests: int, draft_len: int) -> torch.Tensor: + width = draft_len + 1 + row = torch.arange(width, dtype=torch.int, device="cuda") + return row.unsqueeze(0).expand(max_num_requests, -1).contiguous() + + +@functools.cache +def generate_spec_decoding_packed_mask(max_num_requests: int, draft_len: int) -> torch.Tensor: + width = draft_len + 1 + num_blocks = math.ceil(width / 32) + mask = torch.zeros([max_num_requests, width, num_blocks], dtype=torch.int, device="cuda") + remaining = width + for blk in range(num_blocks): + if remaining <= 0: + break + n = min(32, remaining) + vals = (torch.pow(2, torch.arange(n) + 1) - 1).int() + mask[:, blk * 32 : blk * 32 + n, blk] = vals + remaining -= 32 + return mask + + +def set_use_spec_decoding(metadata, use: bool) -> None: + """Toggle spec-dec mode for this step on metadata whose backend has one.""" + spec = metadata.find_ext(SpecDecodingMetadata) + if spec is not None: + spec.use_spec_decoding = use + + +def find_spec_decoding(metadata) -> Optional["SpecDecodingMetadata"]: + """The spec-decoding extension of ``metadata``, or None. + + Accepts the duck-typed metadata views that may carry no extensions. + """ + find_ext = getattr(metadata, "find_ext", None) + return find_ext(SpecDecodingMetadata) if find_ext is not None else None + + +class SpecDecodingMetadata(MetadataExtension): + """Spec-dec mode flags and the draft-token buffers the attention op reads. + + The buffers stay None until ``update`` enables spec-dec mode; CUDA graph + copies share them with the metadata they were copied from. + """ + + def __init__(self, metadata) -> None: + super().__init__(metadata) + # Spec-dec mode (multi-query mode) of the TRTLLM XQA kernels, supported + # for the entire runtime. + self.is_spec_decoding_enabled = False + # Whether the attention layer runs in spec-dec mode at this step/layer. + self.use_spec_decoding = False + # Whether the spec-dec tree is a tree or a chain (linear tree). + self.is_spec_dec_tree = False + # A tree that never changes needs its mask computed only once. + self.is_spec_dec_dynamic_tree = False + self.force_prepare_tree_mask = False + self.max_total_draft_tokens: Optional[int] = None + self.position_offsets: Optional[torch.Tensor] = None + # The C++ attention op requires a 2-D position_offsets tensor and + # reads sizes()[1] as the generation length / packed-mask row stride. + self.position_offsets_cpp: Optional[torch.Tensor] = None + # Compact Hopper C++ row stride for 1D dynamic-tree offsets. + self.position_offsets_stride = 0 + self.packed_mask: Optional[torch.Tensor] = None + self.generation_lengths: Optional[torch.Tensor] = None + self.bl_tree_mask_offset: Optional[torch.Tensor] = None + self.bl_tree_mask: Optional[torch.Tensor] = None + self.bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor] = None + + @property + def position_offsets_for_cpp(self) -> Optional[torch.Tensor]: + """``position_offsets`` reshaped to the 2D layout the C++ kernel expects.""" + offsets = self.position_offsets + if offsets is not None and offsets.dim() == 1: + if self.position_offsets_cpp is not None and not is_trtllm_gen_sm(get_sm_version()): + return self.position_offsets_cpp + return offsets.view(self.metadata.max_num_requests, -1) + return offsets + + def update_position_offsets_for_cpp(self, query_len: int) -> None: + """Refresh the C++ view of spec-dec position offsets.""" + offsets = self.position_offsets + max_num_requests = self.metadata.max_num_requests + if offsets is None or offsets.dim() != 1: + self.position_offsets_cpp = offsets + self.position_offsets_stride = 0 + return + + if max_num_requests > 0 and query_len > 0: + self.position_offsets_stride = query_len + total = max_num_requests * query_len + self.position_offsets_cpp = offsets[:total].view(max_num_requests, query_len) + else: + self.position_offsets_cpp = offsets + self.position_offsets_stride = 0 + + def prepare_for_blackwell(self) -> None: + """Prepare the Blackwell parameters of the generation-phase spec-dec kernels. + + Uses persistent buffers (only allocated if None) for CUDA graph + compatibility. + """ + md = self.metadata + if self.bl_tree_mask_offset is None: + self.bl_tree_mask_offset = torch.zeros( + [md.max_num_requests], + dtype=torch.int64, + device="cuda", + ) + + if self.bl_tree_first_sparse_mask_offset_kv is None: + self.bl_tree_first_sparse_mask_offset_kv = torch.zeros( + [md.max_num_requests], + dtype=torch.int32, + device="cuda", + ) + + if self.bl_tree_mask is None: + seqLenQ = self.packed_mask.shape[1] if self.packed_mask is not None else 1 + mask_size = torch.ops.trtllm.spec_dec_tree_custom_mask_capacity( + md.max_num_requests, seqLenQ, md.max_num_heads_per_kv + ) + self.bl_tree_mask = torch.zeros( + mask_size, + dtype=torch.uint32, + device="cuda", + ) + + def update_blackwell_first_sparse_mask_offset(self) -> None: + """Fill gen slots [0:ng) with (kv_lens_cuda - seq_lens); in-place for CUDA graph.""" + md = self.metadata + nc = md.num_contexts + n = md.num_seqs + ng = n - nc + if ng > 0: + gen_offset = (md.kv_lens_cuda[nc:n] - md.seq_lens_cuda[nc:n]).to(torch.int32) + self.bl_tree_first_sparse_mask_offset_kv[:ng].copy_(gen_offset) + + def fill_generation_lengths(self, runtime_draft_len: int) -> None: + self.generation_lengths[: self.metadata.max_num_requests].fill_(runtime_draft_len + 1) + + def update( + self, + batch_size, + is_spec_decoding_enabled, + is_spec_dec_tree, + is_spec_dec_dynamic_tree, + max_draft_len, + max_total_draft_tokens, + spec_metadata: Optional["SpecMetadata"] = None, + spec_tree_manager: Optional["SpecTreeManager"] = None, + num_contexts: int = 0, + ) -> None: + """ + Update the spec-dec parameters for the TRTLLM attention layer. + Args: + batch_size: int, the number of requests in the batch. + is_spec_decoding_enabled: bool, whether the attention need to be spec_decoding mode, + which is determined by attention_need_spec_dec_mode() function. + is_spec_dec_tree: bool, whether the spec-dec mode is a tree, i.e., dynamic tree. + For linear-tree, it is always False. + is_spec_dec_dynamic_tree: bool, whether using dynamic tree. + max_draft_len: int, the number of the draft layers. + max_total_draft_tokens: int, the number of all nodes in the tree (except the root). + spec_metadata: Optional['SpecMetadata'] = None, the metadata of the spec-dec. + spec_tree_manager: Optional['SpecTreeManager'] = None, the spec_tree_manager for draft token tree. + num_contexts: int = 0, the number of context (prefill) requests in the + batch. The slot-storage layout is ``[ctx | gen]``; dynamic-tree + metadata only describes the gen rows, so we must source from + ``[num_contexts:batch_size]`` rather than ``[:batch_size]``. + """ + md = self.metadata + trtllm_gen = is_trtllm_gen_sm(get_sm_version()) + + # Blackwell trtllm-gen spec-dec is enabled only for dynamic-tree masks. + self.is_spec_decoding_enabled = is_spec_decoding_enabled and ( + not trtllm_gen or is_spec_dec_dynamic_tree + ) + + # use_spec_decoding is default to true by default, change in runtime by layers / requests + self.use_spec_decoding = self.is_spec_decoding_enabled + self.is_spec_dec_tree = is_spec_dec_tree + self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree + # A hybrid model's first executed attention layer can have a nonzero + # cache-local index because recurrent layers precede it. Do not rely + # on the C++ ``layer_idx == 0`` fallback to rebuild the target mask: + # the dynamic draft loop clears that mask before the next target step. + self.force_prepare_tree_mask = is_spec_dec_dynamic_tree + # Forward the tree length to FMHA kernel selection. + self.max_total_draft_tokens = max_total_draft_tokens + + if not self.is_spec_decoding_enabled: + return + + # These buffers are accessed more like removing input padding, + # rather than using max_total_draft_tokens + 1 as the offset between different requests. + if is_spec_dec_tree and self.position_offsets is None: + # Dynamic tree: use _internal_buf_dim which may be larger + # than max_total_draft_tokens+1 to accommodate K*max_draft_len. + # 1D layout for flexible view() in the drafting loop. + buf_dim = spec_tree_manager._internal_buf_dim + self.position_offsets = torch.empty( + (md.max_num_requests * buf_dim,), + dtype=torch.int, + device="cuda", + ) + if is_spec_dec_tree and self.packed_mask is None: + buf_dim = spec_tree_manager._internal_buf_dim + # Zero-init: dynamic-tree dst has inner dim + # ceil(buf_dim/32) but only ceil((max_total_draft_tokens+1)/32) + # is written each step. Unwritten cols would otherwise feed the + # C++ XQA kernel as live mask bits. + self.packed_mask = torch.zeros( + [md.max_num_requests, buf_dim, math.ceil(buf_dim / 32)], + dtype=torch.int, + device="cuda", + ) + + if self.generation_lengths is None: + self.generation_lengths = torch.empty( + [md.max_num_requests], + dtype=torch.int, + device="cuda", + ) + + if trtllm_gen: + self.prepare_for_blackwell() + else: + self.bl_tree_mask_offset = None + self.bl_tree_mask = None + self.bl_tree_first_sparse_mask_offset_kv = None + + cpp_query_len = 0 + + # Case 1: dynamic tree — copy per-request params from spec_tree_manager. + if self.is_spec_dec_dynamic_tree: + assert spec_tree_manager is not None, "spec_tree_manager is required for dynamic tree" + n_dt = spec_tree_manager.max_total_draft_tokens + 1 + buf_dim = spec_tree_manager._internal_buf_dim + + assert buf_dim >= n_dt, ( + f"Dynamic-tree copy requires buf_dim >= n_dt; got buf_dim={buf_dim}, n_dt={n_dt}." + ) + + slot_storage = spec_tree_manager.slot_storage + + num_gens = batch_size - num_contexts + if num_gens > 0: + slot_ids = slot_storage.all_ids_buf[num_contexts:batch_size] + + pos_src = torch.index_select(slot_storage.position_offsets, 0, slot_ids)[:, :n_dt] + pos_dst = self.position_offsets[: num_gens * n_dt].view(num_gens, n_dt) + pos_dst.copy_(pos_src, non_blocking=True) + + actual_mask_width = math.ceil(n_dt / 32) + mask_src = torch.index_select(slot_storage.packed_mask, 0, slot_ids)[ + :, :, :actual_mask_width + ] + if trtllm_gen: + # Blackwell reads the padded 3D mask layout. + self.packed_mask[:num_gens, :n_dt, :actual_mask_width].copy_( + mask_src, non_blocking=True + ) + else: + # Hopper XQA reads a compact flat prefix. + total = num_gens * n_dt * actual_mask_width + self.packed_mask.view(-1)[:total].copy_(mask_src.reshape(-1), non_blocking=True) + + self.generation_lengths[:batch_size].fill_(n_dt) + cpp_query_len = n_dt + + # Case 2: linear tree + else: + # Currently dynamic draft length is only supported for linear tree + # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length. + # So we create cache for position offsets and packed mask for each draft length to avoid reallocation. + assert max_draft_len == max_total_draft_tokens, ( + "max_draft_len should be equal to max_total_draft_tokens for linear tree" + ) + # For algos other than PARD, this equals runtime_draft_len (K); for PARD it's 2K-1. + runtime_draft_token_buffer_width = ( + spec_metadata.runtime_tokens_per_gen_step - 1 + if spec_metadata is not None + else max_draft_len + ) + self.fill_generation_lengths(runtime_draft_len=runtime_draft_token_buffer_width) + self.position_offsets = generate_spec_decoding_position_offsets( + md.max_num_requests, runtime_draft_token_buffer_width + ) + self.packed_mask = generate_spec_decoding_packed_mask( + md.max_num_requests, runtime_draft_token_buffer_width + ) + + self.update_position_offsets_for_cpp(cpp_query_len) diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index c7df494c58d4..3d9e07d1678e 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/flashinfer.py @@ -1078,8 +1078,8 @@ def update_for_spec_dec(self) -> None: self._prepare_full_draft_page_table() self._update_draft_kv_lengths() - def __post_init__(self) -> None: - super().__post_init__() + def _post_init_core(self) -> None: + super()._post_init_core() self._post_init_with_buffers(self.cuda_graph_buffers) def _post_init_with_buffers(self, buffers) -> None: @@ -1629,7 +1629,7 @@ def _refresh_fa2_cuda_graph_plans(self) -> None: wrappers.is_planned = False self._plan_with_params(plan_params) - def prepare(self) -> None: + def _prepare_core(self) -> None: def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: """Host int32 staging tensor for async H2D copies; pinned when @@ -1641,7 +1641,7 @@ def _to_int32_tensor(arr: np.ndarray) -> torch.Tensor: return maybe_pin_memory( torch.from_numpy(arr.astype(np.int32, copy=False))) - super().prepare() + super()._prepare_core() extra_attrs = get_model_extra_attrs() if extra_attrs is None: get_global_attrs().attention_metadata = weakref.ref(self) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index cf0f56d136ee..0ad8d3315e2d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -30,6 +30,8 @@ from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger +from ..features import HelixMetadata +from ..features.spec_decoding import find_spec_decoding from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha @@ -456,17 +458,19 @@ def _is_supported_with_reason( # [0, K - (seq_len_q - 1) + t)). Tree / dynamic-tree spec-decode carries # an explicit packed mask the kernel cannot express, and an explicit # CUSTOM mask is likewise unsupported -> fall back to TRTLLM for those. + spec = find_spec_decoding(meta) if ( fwd.attention_mask == CustomAttentionMask.CUSTOM or fwd.attention_mask_data is not None - or getattr(meta, "is_spec_dec_tree", False) - or getattr(meta, "is_spec_dec_dynamic_tree", False) + or (spec is not None and (spec.is_spec_dec_tree or spec.is_spec_dec_dynamic_tree)) ): return False, "CuTe DSL MLA FMHA does not support custom/tree speculative masks." seq_len_q = q.shape[0] // meta.num_generations batch_size = meta.num_generations - if meta.helix_position_offsets is not None: - if seq_len_q != 1 and not meta._helix_spec_tokens_valid: + helix = meta.find_ext(HelixMetadata) + helix_active = helix is not None and helix.position_offsets is not None + if helix_active: + if seq_len_q != 1 and not helix.spec_tokens_valid: # Multi-token decode under helix needs the per-token bound / # write-slot buffers of the speculative verify-group path. return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix." @@ -495,11 +499,9 @@ def _is_supported_with_reason( # _PERF_MIN_BATCH_FP8, which stays a pure measured-win table. # Single-token H=96 is already a measured table entry and still goes # through the gate. Helix with seq_len_q > 1 implies - # _helix_spec_tokens_valid -- the helix block above returns False + # spec_tokens_valid -- the helix block above returns False # otherwise. - helix_h96_verify_group = ( - attn.num_heads == 96 and seq_len_q > 1 and meta.helix_position_offsets is not None - ) + helix_h96_verify_group = attn.num_heads == 96 and seq_len_q > 1 and helix_active # The same failure without Helix: an FP8-KV multi-token verify group # (e.g. DSpark under attention DP, where every rank keeps all 96 heads) # is not in the FP8 perf table, and on SM100/SM103 the next library is @@ -512,7 +514,7 @@ def _is_supported_with_reason( fp8_wide_head_verify_group = ( 64 < attn.num_heads < 128 and seq_len_q > 1 - and meta.helix_position_offsets is None + and not helix_active and self._get_kernel_dtype(attn, q) == torch.float8_e4m3fn and get_sm_version() in (100, 103) ) @@ -571,6 +573,7 @@ def _run_mla_decode( ) -> None: attn = params.attn meta = params.meta + helix = meta.find_ext(HelixMetadata) # dtype / batch / MLA-dim validity is already enforced before dispatch # (``is_available`` + ``_is_supported_with_reason`` + the kernel's @@ -702,8 +705,12 @@ def _run_mla_decode( # Per-token rank-local bounds, filled by # recompute_helix_spec_buffers. None everywhere else. ( - meta.helix_kv_bounds[:num_tokens] - if (meta.helix_position_offsets is not None and meta._helix_spec_tokens_valid) + helix.kv_bounds[:num_tokens] + if ( + helix is not None + and helix.position_offsets is not None + and helix.spec_tokens_valid + ) else None ), block_split_kvs, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 0831ace484f8..dce8251a0031 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -26,6 +26,8 @@ from tensorrt_llm.bindings.internal import thop from tensorrt_llm.logger import logger +from ..features import FlashMlaMetadata, HelixMetadata +from ..features.spec_decoding import SpecDecodingMetadata from .interface import Fmha, FmhaPhase if TYPE_CHECKING: @@ -59,6 +61,11 @@ # literal value. _THOP_LITERALS: dict = {} +# Stand-ins for absent extensions: every field reads as None or off. +_NO_HELIX = HelixMetadata(None) +_NO_FLASH_MLA = FlashMlaMetadata(None) +_NO_SPEC = SpecDecodingMetadata(None) + class FallbackFmha(Fmha): """Fallback FMHA implementation using the fused TRT-LLM thop attention op.""" @@ -270,9 +277,11 @@ def _is_supported( # gate both assume the new KV entries are the trailing slots of one # rank's kv_len. Reject rather than run it silently wrong; being last # in the library list, this makes dispatch raise. + helix = metadata.find_ext(HelixMetadata) if ( - metadata.helix_position_offsets is not None - and metadata._helix_spec_tokens_valid + helix is not None + and helix.position_offsets is not None + and helix.spec_tokens_valid and metadata.num_generations > 0 # Count generation tokens only: in a mixed batch ``q`` also holds # the context tokens, which would otherwise trip this on a batch @@ -297,6 +306,10 @@ def forward( forward_args: AttentionForwardArgs, ) -> None: attn = self.attn + # Extensions this call reads; an absent one reads as all-None. + helix = metadata.find_ext(HelixMetadata) or _NO_HELIX + flash_mla = metadata.find_ext(FlashMlaMetadata) or _NO_FLASH_MLA + spec = metadata.find_ext(SpecDecodingMetadata) or _NO_SPEC if forward_args.sage_attn_smooth_v: logger.warning_once( "The attention op quantizes V without the V smoothing this SageAttention " @@ -333,22 +346,22 @@ def forward( max_num_sequences=metadata.max_num_sequences, beam_width=metadata.effective_beam_width, use_paged_context_fmha=metadata.use_paged_context_fmha, - helix_position_offsets=metadata.helix_position_offsets, - helix_is_inactive_rank=metadata.helix_is_inactive_rank, - is_spec_decoding_enabled=metadata.is_spec_decoding_enabled, - use_spec_decoding=metadata.use_spec_decoding, - is_spec_dec_tree=metadata.is_spec_dec_tree, - spec_decoding_generation_lengths=metadata.spec_decoding_generation_lengths, - spec_decoding_position_offsets_for_cpp=metadata.spec_decoding_position_offsets_for_cpp, - spec_decoding_packed_mask=metadata.spec_decoding_packed_mask, - spec_decoding_bl_tree_mask_offset=metadata.spec_decoding_bl_tree_mask_offset, - spec_decoding_bl_tree_mask=metadata.spec_decoding_bl_tree_mask, - spec_decoding_target_max_draft_tokens=metadata.max_total_draft_tokens, - force_prepare_spec_dec_tree_mask=metadata.force_prepare_spec_dec_tree_mask, - spec_bl_tree_first_sparse_mask_offset_kv=metadata.spec_bl_tree_first_sparse_mask_offset_kv, + helix_position_offsets=helix.position_offsets, + helix_is_inactive_rank=helix.is_inactive_rank, + is_spec_decoding_enabled=spec.is_spec_decoding_enabled, + use_spec_decoding=spec.use_spec_decoding, + is_spec_dec_tree=spec.is_spec_dec_tree, + spec_decoding_generation_lengths=spec.generation_lengths, + spec_decoding_position_offsets_for_cpp=spec.position_offsets_for_cpp, + spec_decoding_packed_mask=spec.packed_mask, + spec_decoding_bl_tree_mask_offset=spec.bl_tree_mask_offset, + spec_decoding_bl_tree_mask=spec.bl_tree_mask, + spec_decoding_target_max_draft_tokens=spec.max_total_draft_tokens, + force_prepare_spec_dec_tree_mask=spec.force_prepare_tree_mask, + spec_bl_tree_first_sparse_mask_offset_kv=spec.bl_tree_first_sparse_mask_offset_kv, num_sparse_topk=metadata.num_sparse_topk, - flash_mla_tile_scheduler_metadata=metadata.flash_mla_tile_scheduler_metadata, - flash_mla_num_splits=metadata.flash_mla_num_splits, + flash_mla_tile_scheduler_metadata=flash_mla.tile_scheduler_metadata, + flash_mla_num_splits=flash_mla.num_splits, num_contexts=metadata.num_contexts, num_ctx_tokens=metadata.num_ctx_tokens, max_context_length=metadata.max_context_length, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index e465eda7d093..5cc7a95ec728 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -61,6 +61,8 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.mode import QuantMode +from ..features import HelixMetadata +from ..features.spec_decoding import SpecDecodingMetadata from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha from .utils import ( @@ -662,7 +664,8 @@ def _is_supported_with_reason( or fwd.sage_attn_num_elts_per_blk_v > 0 ): return False, "sage attention." - if meta.helix_position_offsets is not None: + helix = meta.find_ext(HelixMetadata) + if helix is not None and helix.position_offsets is not None: return False, "helix parallelism." sparse_kv_indices = fwd.sparse_runtime_params.sparse_kv_indices sparse_attn_indices = fwd.sparse_runtime_params.sparse_attn_indices @@ -677,7 +680,8 @@ def _is_supported_with_reason( return False, "skip-softmax attention." if fwd.relative_attention_bias is not None: return False, "relative attention bias." - if meta.use_spec_decoding and meta.is_spec_dec_tree: + spec_dec = meta.ext(SpecDecodingMetadata) + if spec_dec.use_spec_decoding and spec_dec.is_spec_dec_tree: return False, "spec-dec tree/custom masks." if is_mla_enable and has_context_phase: return False, "MLA context attention." @@ -719,7 +723,7 @@ def _is_supported_with_reason( return False, "cross attention with NVFP4 KV cache." if is_mla_enable: return False, "cross attention with MLA." - if meta.is_spec_decoding_enabled or meta.use_spec_decoding: + if spec_dec.is_spec_decoding_enabled or spec_dec.use_spec_decoding: return False, "cross attention with speculative decoding." if fwd.update_kv_cache and fwd.cross_kv is None: return False, "cross attention with a missing cross_kv when update_kv_cache=True." diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py index 28f2e9996cb7..f64aa3b15ce0 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fp4_mla.py @@ -38,6 +38,8 @@ ) from tensorrt_llm.bindings import DataType +from ..features import find_helix +from ..fp4_mla.state import find_fp4_mla from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha @@ -109,7 +111,7 @@ def _is_supported( ) ): raise NotImplementedError("FP4 MLA does not support sparse attention.") - state = metadata.fp4_mla_state + state = find_fp4_mla(metadata) if state is None or state.hp_pool is None: raise RuntimeError("FP4 MLA requires prepared high-precision KV state.") if state.v_scale_pool is None: @@ -124,7 +126,8 @@ def _is_supported( raise RuntimeError("FP4 MLA generation expects a fused query input.") # Helix combines the per-rank partial outputs with these stats, so a # Helix generation without them would combine uninitialized memory. - if getattr(metadata, "helix_position_offsets", None) is not None: + helix = find_helix(metadata) + if helix is not None and helix.position_offsets is not None: softmax_stats = forward_args.softmax_stats_tensor expected_shape = (q.shape[0], self.attn.num_heads, 2) if softmax_stats is None: @@ -189,7 +192,7 @@ def run_mla_context(self, params: FmhaParams) -> None: raise RuntimeError("FP4 MLA context requires an output buffer.") if forward_args.latent_cache is None: raise RuntimeError("FP4 MLA context requires latent_cache.") - if metadata.fp4_mla_state.positions is None: + if find_fp4_mla(metadata).positions is None: raise RuntimeError("FP4 MLA context requires token positions.") if metadata.num_contexts <= 0: raise RuntimeError("FP4 MLA context requires context requests.") @@ -336,6 +339,7 @@ def forward_context_partition( def run_mla_generation(self, params: FmhaParams) -> None: attn = params.attn metadata = params.meta + fp4_mla_state = find_fp4_mla(metadata) forward_args = params.fwd q = params.query_input output = params.output @@ -353,18 +357,18 @@ def run_mla_generation(self, params: FmhaParams) -> None: qk_rope_head_dim = attn.qk_rope_head_dim or 0 fused_head_dim = kv_lora_rank + qk_rope_head_dim - if not bool(getattr(metadata.fp4_mla_state, "generation_cache_scattered", False)): + if not bool(getattr(fp4_mla_state, "generation_cache_scattered", False)): raise RuntimeError( "FP4 MLA generation requires fused RoPE/Q quantization and " "cache/HP-pool update before attention." ) - metadata.fp4_mla_state.generation_cache_scattered = False + fp4_mla_state.generation_cache_scattered = False query = q.view(q.shape[0], attn.num_heads, fused_head_dim) output_view = output.view(q.shape[0], attn.num_heads, kv_lora_rank) sm_scale = 1.0 / (attn.q_scaling * ((attn.qk_nope_head_dim or 0) + qk_rope_head_dim) ** 0.5) - prequantized_q = getattr(metadata.fp4_mla_state, "prequantized_q", None) - prequantized_q_sf = getattr(metadata.fp4_mla_state, "prequantized_q_sf", None) - q_batch_capacity = getattr(metadata.fp4_mla_state, "q_batch_capacity", None) + prequantized_q = getattr(fp4_mla_state, "prequantized_q", None) + prequantized_q_sf = getattr(fp4_mla_state, "prequantized_q_sf", None) + q_batch_capacity = getattr(fp4_mla_state, "q_batch_capacity", None) try: run_fp4_mla_attention_decode( metadata, @@ -381,10 +385,10 @@ def run_mla_generation(self, params: FmhaParams) -> None: softmax_stats_tensor=params.fwd.softmax_stats_tensor, ) finally: - metadata.fp4_mla_state.prequantized_q = None - metadata.fp4_mla_state.prequantized_q_sf = None - metadata.fp4_mla_state.q_batch_capacity = None - metadata.fp4_mla_state.prequantized_q_rows_per_seq = None + fp4_mla_state.prequantized_q = None + fp4_mla_state.prequantized_q_sf = None + fp4_mla_state.q_batch_capacity = None + fp4_mla_state.prequantized_q_rows_per_seq = None __all__ = ["Fp4MlaFmha"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index 5f86c100f8c7..5895a50a940f 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -35,6 +35,7 @@ from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.logger import logger +from ..features.spec_decoding import SpecDecodingMetadata from .combined import CombinedFmha from .interface import Fmha, FmhaPhase from .phased import PhasedFmha @@ -387,7 +388,7 @@ def _make_cache_key( generation_batch_size=generation_batch_size, generation_seq_len_q=generation_seq_len_q, attention_mask_type=attention_mask_type, - use_spec_decoding=metadata.use_spec_decoding, + use_spec_decoding=metadata.ext(SpecDecodingMetadata).use_spec_decoding, has_block_sparse_inputs=block_sparse_inputs is not None, output_dtype=output_dtype, output_sf_dtype=output_sf_dtype, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py index c6c6af4d08cd..c5a4f936da20 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_decode.py @@ -136,7 +136,10 @@ def run_generation(self, params: FmhaParams) -> None: from ..sparse.minimax_m3.kernels.msa_utils import write_msa_phase_kv metadata = params.meta - span = metadata.msa_decode_span + from ..sparse.minimax_m3.msa_backend import MiniMaxM3MsaMetadata + + msa = metadata.ext(MiniMaxM3MsaMetadata) + span = msa.msa_decode_span # Both kernels below map a query token to its request through the span, # while the phase params carry PhasedFmha's own derivation of the same # boundary. Checking them against each other rejects a step whose @@ -159,8 +162,8 @@ def run_generation(self, params: FmhaParams) -> None: ) row_first = params.seq_offset row_last = row_first + metadata.num_generations - block_table = metadata.msa_block_table[row_first:row_last] - seq_lens = metadata.msa_seq_lens_cuda[row_first:row_last] + block_table = msa.msa_block_table[row_first:row_last] + seq_lens = msa.msa_seq_lens_cuda[row_first:row_last] kv_block_indexes = params.fwd.sparse_runtime_params.sparse_attn_indices if kv_block_indexes is not None: @@ -238,12 +241,15 @@ def _run_dense( attn = params.attn metadata = params.meta + from ..sparse.minimax_m3.msa_backend import MiniMaxM3MsaMetadata + + msa = metadata.ext(MiniMaxM3MsaMetadata) head_dim = attn.head_dim num_tokens = params.num_tokens row_first = params.seq_offset # The sub-page block table prepare() staged, if it could; the kernel # expands its own when the factor does not match this layer's. - staged_table, staged_factor = metadata.msa_subpage_rows( + staged_table, staged_factor = msa.msa_subpage_rows( row_first, row_first + metadata.num_generations ) if self._dense_layout is None: @@ -261,7 +267,7 @@ def _run_dense( sm_scale=(head_dim**-0.5) / float(attn.q_scaling), output=params.output.view(num_tokens, attn.num_heads, head_dim), decode_query_len=params.input_seq_length, - max_seq_len=int(metadata.msa_max_kv_len), + max_seq_len=int(msa.msa_max_kv_len), max_num_requests=int(metadata.max_num_requests), staged_subpage_table=staged_table, staged_subpages_per_slot=staged_factor, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py index 120feaafeeee..4893b63df034 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_prefill.py @@ -181,6 +181,9 @@ def run_msa_nvfp4_sparse_gqa( require_msa_module, ) + from ..sparse.minimax_m3.msa_backend import MiniMaxM3MsaMetadata + + msa = metadata.ext(MiniMaxM3MsaMetadata) fmha_sm100 = require_msa_module() sparse = getattr(fmha_sm100, "sparse", None) if sparse is None: @@ -207,27 +210,27 @@ def run_msa_nvfp4_sparse_gqa( f"got {tuple(scale_buffers.shape)} for K {tuple(k_paged.shape)}" ) - batch = int(getattr(metadata, "_msa_live_batch", 0)) + batch = int(getattr(msa, "_msa_live_batch", 0)) if batch <= 0: raise RuntimeError("MiniMax-M3 NVFP4 sparse attention metadata was not prepared") - max_q_len = int(metadata._msa_max_q_len) - max_kv_len = int(metadata._msa_max_kv_len_all) - total_k = int(metadata._msa_total_k) - total_k_rows = int(metadata._msa_total_k_rows) + max_q_len = int(msa._msa_max_q_len) + max_kv_len = int(msa._msa_max_kv_len_all) + total_k = int(msa._msa_total_k) + total_k_rows = int(msa._msa_total_k_rows) if num_rows is not None: if not 0 < num_rows <= batch: raise ValueError(f"MiniMax-M3 NVFP4 row limit ({num_rows}) must lie in (0, {batch}].") batch = num_rows # Staged by prepare() over the context rows, which are exactly the rows # a narrowed call covers; see _stage_step_fields. - max_q_len, max_kv_len, total_k, total_k_rows = metadata._msa_context_prefix_bounds + max_q_len, max_kv_len, total_k, total_k_rows = msa._msa_context_prefix_bounds if total_k_rows <= 0: raise RuntimeError( "MiniMax-M3 NVFP4 attention was narrowed to a context prefix of " f"{num_rows} row(s), but prepare() staged no bounds for one." ) - cu_q = metadata.msa_cu_q_lens[: batch + 1] - cu_kv = metadata.msa_cu_kv_lens[: batch + 1] + cu_q = msa.msa_cu_q_lens[: batch + 1] + cu_kv = msa.msa_cu_kv_lens[: batch + 1] q2k = kv_block_indexes.permute(1, 0, 2).contiguous() topk = int(q2k.shape[-1]) page_size = int(k_paged.shape[2]) @@ -266,8 +269,8 @@ def run_msa_nvfp4_sparse_gqa( softmax_scale=sm_scale, partial_dtype=torch.bfloat16, return_softmax_lse=False, - page_table=metadata.msa_block_table[:batch], - seqused_k=metadata.msa_seq_lens_cuda[:batch], + page_table=msa.msa_block_table[:batch], + seqused_k=msa.msa_seq_lens_cuda[:batch], schedule=schedule, ) out.copy_(result) @@ -297,7 +300,9 @@ def run_msa_prefill_gqa( inline plan would read. """ from ..sparse.minimax_m3.kernels.msa_utils import msa_paged_kv + from ..sparse.minimax_m3.msa_backend import MiniMaxM3MsaMetadata + msa = metadata.ext(MiniMaxM3MsaMetadata) head_dim = attn.head_dim num_tokens = int(q.shape[0]) if num_tokens == 0: @@ -352,11 +357,11 @@ def rows_of(lens: Optional[torch.Tensor]) -> Optional[torch.Tensor]: k_paged, v_paged, kv_block_indexes, - kv_indices=metadata.msa_kv_indices, + kv_indices=msa.msa_kv_indices, sm_scale=sm_scale, - qo_lens_cpu=rows_of(metadata.msa_qo_lens_cpu), - kv_lens_cpu=rows_of(metadata.msa_kv_lens_cpu), - qo_offset_cpu=rows_of(metadata.msa_qo_offset_cpu), + qo_lens_cpu=rows_of(msa.msa_qo_lens_cpu), + kv_lens_cpu=rows_of(msa.msa_kv_lens_cpu), + qo_offset_cpu=rows_of(msa.msa_qo_offset_cpu), causal=True, head_dim=head_dim, plan=plan, @@ -405,6 +410,9 @@ def run_context(self, params: FmhaParams) -> None: from ..sparse.minimax_m3.kernels.msa_utils import write_msa_phase_kv metadata = params.meta + from ..sparse.minimax_m3.msa_backend import MiniMaxM3MsaMetadata + + msa = metadata.ext(MiniMaxM3MsaMetadata) write_msa_phase_kv( params.attn, params.key_input, @@ -428,11 +436,7 @@ def run_context(self, params: FmhaParams) -> None: metadata, params.output, kv_block_indexes=kv_block_indexes, - plan=( - metadata.msa_prefill_gqa_plan - if is_sparse_layer - else metadata.msa_prefill_dense_plan - ), + plan=(msa.msa_prefill_gqa_plan if is_sparse_layer else msa.msa_prefill_dense_plan), row_first=params.seq_offset, num_rows=metadata.num_contexts, kv_scale_quant_orig=params.fwd.kv_scale_quant_orig, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/phased.py b/tensorrt_llm/_torch/attention/backends/fmha/phased.py index f53aac9569a7..b2b133458df3 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/phased.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/phased.py @@ -27,6 +27,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from ..features.spec_decoding import SpecDecodingMetadata from .interface import Fmha if TYPE_CHECKING: @@ -172,6 +173,7 @@ def forward( metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, ) -> None: + spec = metadata.find_ext(SpecDecodingMetadata) attn = self.attn output = forward_args.output if output is None: @@ -288,9 +290,9 @@ def forward( predicted_tokens_per_seq = attn.predicted_tokens_per_seq spec_gen_lengths = None spec_pos_offsets = None - if metadata.is_spec_decoding_enabled and predicted_tokens_per_seq > 1: - spec_gen_lengths = metadata.spec_decoding_generation_lengths - position_offsets_for_cpp = metadata.spec_decoding_position_offsets_for_cpp + if spec.is_spec_decoding_enabled and predicted_tokens_per_seq > 1: + spec_gen_lengths = spec.generation_lengths + position_offsets_for_cpp = spec.position_offsets_for_cpp if position_offsets_for_cpp is not None and position_offsets_for_cpp.dim() == 1: position_offsets_for_cpp = position_offsets_for_cpp.view( metadata.max_num_requests, -1 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 2bf1554d3753..ce2791e04aa3 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -38,6 +38,8 @@ from tensorrt_llm.math_utils import ceil_div, pad_up from tensorrt_llm.quantization.mode import QuantMode +from ..features import HelixMetadata +from ..features.spec_decoding import SpecDecodingMetadata from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha from .utils import get_kv_page_offset, get_multi_processor_count_for_device @@ -107,13 +109,14 @@ def get_paged_kv_policy_unsupported_reason( metadata: "TrtllmAttentionMetadata", ) -> Optional[str]: """Return why the request's decoding policy is outside the fixed page-table envelope.""" + spec = metadata.find_ext(SpecDecodingMetadata) if metadata.beam_width != 1: return "beam search is not supported." if ( - metadata.is_spec_decoding_enabled - or metadata.use_spec_decoding - or metadata.is_spec_dec_tree - or metadata.is_spec_dec_dynamic_tree + spec.is_spec_decoding_enabled + or spec.use_spec_decoding + or spec.is_spec_dec_tree + or spec.is_spec_dec_dynamic_tree ): return "speculative decoding is not supported by the initial adapter." position_embedding_type = int(attn.position_embedding_type) @@ -139,7 +142,8 @@ def get_attention_feature_unsupported_reason( forward_args: "AttentionForwardArgs", ) -> Optional[str]: """Return which optional attention feature the fused-kernel adapters do not implement.""" - if metadata.helix_position_offsets is not None: + helix = metadata.find_ext(HelixMetadata) + if helix is not None and helix.position_offsets is not None: return "Helix parallelism is not supported." if forward_args.relative_attention_bias is not None: return "relative attention bias is not supported." diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 2cb531e21c3d..03455e557d48 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -63,6 +63,8 @@ from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_sm_version from tensorrt_llm.logger import logger +from ..features import HelixMetadata +from ..features.spec_decoding import SpecDecodingMetadata from .interface import Fmha if TYPE_CHECKING: @@ -72,15 +74,19 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( DeepseekV4CacheManager, ) - from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( - DeepseekV4TrtllmAttentionMetadata, - ) from tensorrt_llm._torch.attention.backends.trtllm import ( TrtllmAttention, TrtllmAttentionMetadata, ) +def _deepseek_v4_metadata(): + """The DeepSeek-V4 extension class, imported on first use to avoid an import cycle.""" + from ..sparse.deepseek_v4.metadata import DeepseekV4Metadata + + return DeepseekV4Metadata + + _MIN_CUTLASS_DSL_VERSION = Version("4.7.0") # Same compiler-target gate as fmha/prims_ts.py: the kernels share one CuTe DSL # task-scheduling stack, so a toolchain the dense libraries refuse is refused here too. @@ -255,13 +261,17 @@ def _unsupported_reason( return "a DeepSeek-V4 KV cache manager is required." if metadata.beam_width != 1: return "beam search is not supported." - if getattr(metadata, "token_positions_cuda", None) is None: + dsv4 = metadata.find_ext(_deepseek_v4_metadata()) + if dsv4 is None: + return "DeepSeek-V4 metadata is required." + if dsv4.token_positions_cuda is None: return "token positions are required for the output RoPE." - if getattr(metadata, "sparse_mla_swa_lens", None) is None: + if getattr(dsv4, "sparse_mla_swa_lens", None) is None: return "per-step live lengths are required (DeepSeek-V4 metadata)." if fwd.attention_mask_data is not None or fwd.relative_attention_bias is not None: return "custom masks and relative bias are not supported." - if metadata.helix_position_offsets is not None: + helix = metadata.find_ext(HelixMetadata) + if helix is not None and helix.position_offsets is not None: return "Helix parallelism is not supported." rows = int(q.shape[0]) if rows == 0: @@ -342,11 +352,11 @@ def _is_fp8_cache(metadata: "TrtllmAttentionMetadata") -> bool: return metadata.kv_cache_manager.dtype == DataType.FP8 - def _extra_capacity(self, metadata: "DeepseekV4TrtllmAttentionMetadata") -> int: + def _extra_capacity(self, metadata: "TrtllmAttentionMetadata") -> int: ratio = int(self.attn.compress_ratio) if ratio <= 1: return 0 - return int(metadata.max_compressed_indices[ratio]) + return int(metadata.ext(_deepseek_v4_metadata()).max_compressed_indices[ratio]) def _pool_view( self, @@ -605,7 +615,7 @@ def _warm_kernel_variants( # fails here instead of silently shrinking the warmed set and moving # the MTP row counts' compiles to the first serving step. requests = int(metadata.max_num_requests or 0) - draft = int(metadata.max_total_draft_tokens or 0) + draft = int(metadata.ext(SpecDecodingMetadata).max_total_draft_tokens or 0) tokens = int(metadata.max_num_tokens or 0) max_rows = requests * (1 + draft) if tokens: @@ -830,15 +840,16 @@ def forward( # come from the metadata too, derived once per step: the token's # position within the window, and the compressed rows selected beyond # the window. - positions = metadata.token_positions_cuda[start:end] - lengths = metadata.sparse_mla_swa_lens[start:end] + dsv4 = metadata.ext(_deepseek_v4_metadata()) + positions = dsv4.token_positions_cuda[start:end] + lengths = dsv4.sparse_mla_swa_lens[start:end] swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) extra_pool = None run_kwargs = {} if extra: extra_pool = self._pool_view(manager, DeepseekV4AttentionType.COMPRESS, ratio) run_kwargs["extra_indices"] = indices[:, window:] - run_kwargs["extra_lengths"] = metadata.sparse_mla_extra_lens[ratio][start:end] + run_kwargs["extra_lengths"] = dsv4.sparse_mla_extra_lens[ratio][start:end] if fused: # TRT-LLM's fused epilogue buffers are group-major [G, T, C] with diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py index fa8879296835..cc88ea8f3898 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_gather.py @@ -15,6 +15,7 @@ _validate_fp4_mla_kv_storage_shape, ) from .metadata import _materialize_fp4_mla_device_page_table_for_forward +from .state import find_fp4_mla def load_fp4_mla_chunked_kv_cache( @@ -30,6 +31,7 @@ def load_fp4_mla_chunked_kv_cache( qk_rope_head_dim: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Gather one cached-prefix partition from dense FP4 MLA V2 storage.""" + fp4_mla_state = find_fp4_mla(metadata) if out_dtype not in (torch.float16, torch.bfloat16, torch.float32): raise TypeError(f"FP4 MLA chunk gather does not support output dtype {out_dtype}.") if num_ctx_cached_tokens < 0: @@ -87,11 +89,11 @@ def load_fp4_mla_chunked_kv_cache( if num_ctx_cached_tokens == 0: return compressed_kv, k_pe - if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): + if not bool(getattr(fp4_mla_state, "device_page_table", False)): raise RuntimeError("FP4 MLA chunk gather requires fixed-stride device page metadata.") _materialize_fp4_mla_device_page_table_for_forward(metadata) - page_table_stride = int(metadata.fp4_mla_state.page_table_stride) - page_ids = metadata.fp4_mla_state._paged_kv_indices + page_table_stride = int(fp4_mla_state.page_table_stride) + page_ids = fp4_mla_state._paged_kv_indices if ( page_table_stride <= 0 or not isinstance(page_ids, torch.Tensor) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py index 83635d410b71..589f4c957f1c 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/cache_update.py @@ -56,6 +56,7 @@ _materialize_fp4_mla_device_page_table_for_forward, ) from .packed_plan import _fp4_mla_helix_plan_bounds, _fp4_mla_packed_plan_geometry +from .state import find_fp4_mla from .v_cache import ( _get_cutedsl_persistent_v_packed_cache, _get_fp4_mla_v_packed_pool_base, @@ -161,7 +162,11 @@ def _prepare_fp4_mla_q_buffers( or getattr(metadata, "max_num_requests", num_queries) or num_queries ) - max_query_width = 1 + int(getattr(metadata, "max_total_draft_tokens", None) or 0) + from ..features.spec_decoding import find_spec_decoding + + max_query_width = 1 + int( + getattr(find_spec_decoding(metadata), "max_total_draft_tokens", None) or 0 + ) capacity = min( max_num_tokens, max_num_sequences * max_query_width, @@ -272,6 +277,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( v_packed_base: Optional[torch.Tensor] = None, v_page_offset: int = 0, ) -> bool: + fp4_mla_state = find_fp4_mla(metadata) num_contexts = metadata.num_contexts if num_contexts > 0: prompt_lens_cpu = metadata.prompt_lens_cpu_runtime[:num_contexts] @@ -319,7 +325,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( q_context_view = latent_cache q_head_blocks = 0 - hp_pool = getattr(metadata.fp4_mla_state, "hp_pool", None) + hp_pool = getattr(fp4_mla_state, "hp_pool", None) if not isinstance(hp_pool, torch.Tensor): raise TypeError("FP4 MLA high-precision KV pool must be a tensor.") if hp_pool.device != latent_cache.device: @@ -338,7 +344,7 @@ def _scatter_fp4_mla_kv_cache_2d_context( ) if hp_pool.stride(-1) != 1: raise ValueError("FP4 MLA high-precision KV pool must be contiguous in head_dim.") - hp_page_ids = metadata.fp4_mla_state.hp_page_indices + hp_page_ids = fp4_mla_state.hp_page_indices if not isinstance(hp_page_ids, torch.Tensor): raise RuntimeError("FP4 MLA context cache update requires HP page metadata.") store_hp_tail = num_contexts > 0 @@ -367,13 +373,13 @@ def _scatter_fp4_mla_kv_cache_2d_context( rotary_cos_sin_ptr, hp_pool, hp_page_ids, - metadata.fp4_mla_state.batch_indices, - metadata.fp4_mla_state.positions, - metadata.fp4_mla_state.paged_kv_indices, - metadata.fp4_mla_state.paged_kv_indptr, - metadata.fp4_mla_state.paged_kv_indices.shape[0], - metadata.fp4_mla_state.paged_kv_indptr.shape[0], - metadata.fp4_mla_state.batch_indices.shape[0], + fp4_mla_state.batch_indices, + fp4_mla_state.positions, + fp4_mla_state.paged_kv_indices, + fp4_mla_state.paged_kv_indptr, + fp4_mla_state.paged_kv_indices.shape[0], + fp4_mla_state.paged_kv_indptr.shape[0], + fp4_mla_state.batch_indices.shape[0], v_sf.shape[1], v_sf.shape[0], num_contexts, @@ -448,6 +454,7 @@ def _scatter_fp4_mla_kv_cache_2d_generation( helix_position_offsets: Optional[torch.Tensor], helix_is_inactive_rank: Optional[torch.Tensor], ) -> Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + fp4_mla_state = find_fp4_mla(metadata) num_contexts = metadata.num_contexts num_seqs = metadata.num_seqs num_gen = num_seqs - num_contexts @@ -488,7 +495,7 @@ def _scatter_fp4_mla_kv_cache_2d_generation( ) _materialize_fp4_mla_device_page_table_for_forward(metadata, kv_lens_gen) - pool = getattr(metadata.fp4_mla_state, "hp_pool", None) + pool = getattr(fp4_mla_state, "hp_pool", None) if pool is None: raise RuntimeError("FP4 MLA 2D generation scatter requires the HP KV pool.") try: @@ -507,7 +514,10 @@ def _scatter_fp4_mla_kv_cache_2d_generation( max_gen_len = num_tokens // num_gen use_helix = helix_position_offsets is not None or helix_is_inactive_rank is not None - use_helix_local_slots = use_helix and bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + from ..features import find_helix + + helix = find_helix(metadata) + use_helix_local_slots = use_helix and helix is not None and helix.spec_tokens_valid if use_helix: if helix_position_offsets is None or helix_is_inactive_rank is None: raise RuntimeError( @@ -540,7 +550,7 @@ def _scatter_fp4_mla_kv_cache_2d_generation( "same-device bool tensor covering every generation sequence." ) if use_helix_local_slots: - helix_local_slots = getattr(metadata, "helix_local_slots", None) + helix_local_slots = helix.local_slots if ( not isinstance(helix_local_slots, torch.Tensor) or helix_local_slots.dtype != torch.int32 @@ -682,7 +692,7 @@ def launch_generation_update( helix_is_inactive_rank, page_ids, hp_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, + find_fp4_mla(metadata).paged_kv_indptr_decode, page_ids_len, hp_page_ids.numel(), indptr_len, @@ -743,7 +753,7 @@ def launch_generation_update( launch_generation_update( launch_grid, page_ids_len=page_ids.shape[0], - indptr_len=metadata.fp4_mla_state.paged_kv_indptr_decode.shape[0], + indptr_len=fp4_mla_state.paged_kv_indptr_decode.shape[0], max_gen_tiles_variant=max(max_gen_tiles, 1), q_prefix_block_dim_variant=q_prefix_block_dim, q_prefix_blocks_variant=q_prefix_blocks, @@ -804,12 +814,13 @@ def scatter_fp4_mla_kv_cache( backend-ready residual FP4 Q. The return value reports whether the current phase updated the HP pool. """ + fp4_mla_state = find_fp4_mla(metadata) if phase == "generation": - metadata.fp4_mla_state.generation_cache_scattered = False - metadata.fp4_mla_state.prequantized_q = None - metadata.fp4_mla_state.prequantized_q_sf = None - metadata.fp4_mla_state.q_batch_capacity = None - metadata.fp4_mla_state.prequantized_q_rows_per_seq = None + fp4_mla_state.generation_cache_scattered = False + fp4_mla_state.prequantized_q = None + fp4_mla_state.prequantized_q_sf = None + fp4_mla_state.q_batch_capacity = None + fp4_mla_state.prequantized_q_rows_per_seq = None if latent_cache.numel() == 0: raise ValueError("FP4 MLA cache scatter requires at least one latent token.") if q_context is not None and not latent_cache.is_contiguous(): @@ -825,8 +836,8 @@ def scatter_fp4_mla_kv_cache( raise ValueError( f"FP4 MLA KV head_dim must be divisible by {FP4_BLOCK_SIZE}, got {head_dim}." ) - indices_len = metadata.fp4_mla_state.batch_indices.shape[0] - positions_len = metadata.fp4_mla_state.positions.shape[0] + indices_len = fp4_mla_state.batch_indices.shape[0] + positions_len = fp4_mla_state.positions.shape[0] if token_offset + num_tokens > indices_len or token_offset + num_tokens > positions_len: raise RuntimeError( f"FP4 MLA scatter would read batch_indices[{token_offset}:" @@ -852,7 +863,7 @@ def scatter_fp4_mla_kv_cache( if phase not in ("context", "generation"): raise ValueError("FP4 MLA scatter requires phase='context' or 'generation'.") - if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: + if getattr(fp4_mla_state, "v_scale_pool", None) is None: raise RuntimeError("FP4 MLA scatter requires the auxiliary V scale pool.") if metadata.page_size % FP4_BLOCK_SIZE != 0: raise ValueError( @@ -918,10 +929,10 @@ def scatter_fp4_mla_kv_cache( q_quant_input.device, rows_per_seq=q_rows_per_seq, ) - metadata.fp4_mla_state.prequantized_q = q_fp4_out - metadata.fp4_mla_state.prequantized_q_sf = q_sf_out - metadata.fp4_mla_state.q_batch_capacity = q_batch_capacity - metadata.fp4_mla_state.prequantized_q_rows_per_seq = q_rows_per_seq + fp4_mla_state.prequantized_q = q_fp4_out + fp4_mla_state.prequantized_q_sf = q_sf_out + fp4_mla_state.q_batch_capacity = q_batch_capacity + fp4_mla_state.prequantized_q_rows_per_seq = q_rows_per_seq hp_pool_updated = True cutedsl_backend = _fp4_mla_attention_backend() == _FP4_MLA_CUTEDSL_BACKEND fused_v_transpose = cutedsl_backend and _fp4_mla_cutedsl_fused_v_transpose_enabled() @@ -1016,7 +1027,7 @@ def scatter_fp4_mla_kv_cache( v_packed_base=v_packed_base, v_page_offset=v_page_offset, ) - v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices + v_pack_page_ids = fp4_mla_state.paged_kv_indices else: generation_state = _scatter_fp4_mla_kv_cache_2d_generation( metadata, @@ -1058,20 +1069,18 @@ def scatter_fp4_mla_kv_cache( if cutedsl_backend and not fused_v_transpose and not direct_v_packed_write: if phase == "context": num_contexts = metadata.num_contexts - cutedsl_v_pack_page_ids = metadata.fp4_mla_state.paged_kv_indices[ - : metadata.fp4_mla_state.num_context_blocks + cutedsl_v_pack_page_ids = fp4_mla_state.paged_kv_indices[ + : fp4_mla_state.num_context_blocks ] - cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr[: num_contexts + 1] + cutedsl_repack_page_indptr = fp4_mla_state.paged_kv_indptr[: num_contexts + 1] cutedsl_repack_kv_lens = metadata.kv_lens_cuda_runtime[:num_contexts] cutedsl_repack_generation_lens = metadata.prompt_lens_cuda_runtime[:num_contexts] - cutedsl_repack_max_touched_pages = int( - metadata.fp4_mla_state.context_repack_max_touched_pages - ) + cutedsl_repack_max_touched_pages = int(fp4_mla_state.context_repack_max_touched_pages) elif generation_state is not None: kv_lens_gen, gen_lens_gen, generation_page_ids = generation_state num_gen = kv_lens_gen.numel() cutedsl_v_pack_page_ids = generation_page_ids - cutedsl_repack_page_indptr = metadata.fp4_mla_state.paged_kv_indptr_decode + cutedsl_repack_page_indptr = fp4_mla_state.paged_kv_indptr_decode cutedsl_repack_kv_lens = kv_lens_gen cutedsl_repack_generation_lens = gen_lens_gen cutedsl_repack_max_touched_pages = _ceil_div( @@ -1105,5 +1114,5 @@ def scatter_fp4_mla_kv_cache( num_valid_pages=v_pack_num_valid, ) if phase == "generation": - metadata.fp4_mla_state.generation_cache_scattered = hp_pool_updated + fp4_mla_state.generation_cache_scattered = hp_pool_updated return hp_pool_updated diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py index 372a442026fa..82f6b26ec9e9 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/decode.py @@ -55,6 +55,7 @@ _launch_fp4_mla_packed_plan, _PackedPlanGeometry, ) +from .state import find_fp4_mla from .v_cache import ( _get_cutedsl_persistent_v_packed_cache, _get_fp4_mla_v_packed_pool_base, @@ -166,6 +167,7 @@ def _run_triton_attention_decode( ``fp4_mla_triton.py``. Threads through the constexpr assume flags, TMA descriptors, occupancy/num-warps launch meta, and pipelined PV loop. """ + fp4_mla_state = find_fp4_mla(metadata) from .fp4_mla_triton import ( _fp4_mla_attention_group_reduce_stats_kernel as _attn_group_reduce_stats_kernel, ) @@ -312,7 +314,7 @@ def _tma_alloc(size: int, alignment: int, stream): global_scale, q_global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -377,8 +379,8 @@ def _tma_alloc(size: int, alignment: int, stream): # reduce during capture so we never allocate mid-capture. The single-level # reduce is numerically identical (it just launches fewer CTAs). if num_reduce_groups > 1 and torch.cuda.is_current_stream_capturing(): - gmax = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_max_buf") - gsum = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_group_sum_buf") + gmax = fp4_mla_state.workspaces.get("_fp4_mla_attention_group_max_buf") + gsum = fp4_mla_state.workspaces.get("_fp4_mla_attention_group_sum_buf") groups_ready = ( gmax is not None and gsum is not None @@ -458,7 +460,7 @@ def _tma_alloc(size: int, alignment: int, stream): max_scores, denom, page_max, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], max_scores.stride(0), @@ -528,7 +530,7 @@ def _tma_alloc(size: int, alignment: int, stream): # Fall back to the unsplit PV (numerically identical) during capture unless a # warmup forward already sized that workspace, so capture never allocates. if page_split > 1 and torch.cuda.is_current_stream_capturing(): - pbuf = metadata.fp4_mla_state.workspaces.get("_fp4_mla_attention_pv_partial_buf") + pbuf = fp4_mla_state.workspaces.get("_fp4_mla_attention_pv_partial_buf") partial_ready = ( pbuf is not None and pbuf.shape[0] >= num_queries @@ -558,7 +560,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -607,7 +609,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -681,7 +683,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -722,7 +724,7 @@ def _tma_alloc(size: int, alignment: int, stream): v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state.paged_kv_indptr_decode, kv_lens, src_page_ids.shape[0], kv_cache.shape[0], @@ -854,7 +856,7 @@ def _run_cutedsl_packed_decode( plan = plan_geometry.plan # The first FP4 MLA kernel of the step wrote the plan fused with the # generation lengths; a stale key means this call plans on other bounds. - if metadata.fp4_mla_state.packed_plan_key != plan_geometry.key: + if find_fp4_mla(metadata).packed_plan_key != plan_geometry.key: _launch_fp4_mla_packed_plan(plan_geometry, kv_lens, metadata) partial = _ensure_workspace_tensor( metadata, @@ -937,6 +939,7 @@ def run_fp4_mla_attention_decode( KV workspace is materialized on this path. Callers must supply the packed Q and scales produced by the fused generation cache update. """ + fp4_mla_state = find_fp4_mla(metadata) head_dim = kv_lora_rank + qk_rope_head_dim if qk_rope_head_dim != FP4_MLA_K_RESIDUAL_DIM: raise ValueError( @@ -973,7 +976,10 @@ def run_fp4_mla_attention_decode( raise ValueError("FP4 MLA attention output batch dimensions do not match.") backend = _fp4_mla_attention_backend() - helix_spec_tokens_valid = bool(getattr(metadata, "_helix_spec_tokens_valid", False)) + from ..features import find_helix + + helix = find_helix(metadata) + helix_spec_tokens_valid = helix is not None and helix.spec_tokens_valid helix_kv_bounds = None if softmax_stats_tensor is not None: if backend != _FP4_MLA_CUTEDSL_BACKEND: @@ -996,7 +1002,7 @@ def run_fp4_mla_attention_decode( f"stats with shape {expected_stats_shape}." ) if helix_spec_tokens_valid: - helix_kv_bounds = getattr(metadata, "helix_kv_bounds", None) + helix_kv_bounds = helix.kv_bounds if ( not isinstance(helix_kv_bounds, torch.Tensor) or helix_kv_bounds.dtype != torch.int32 @@ -1010,7 +1016,7 @@ def run_fp4_mla_attention_decode( "same-device int32 tensor covering every query token." ) helix_kv_bounds = helix_kv_bounds[:num_queries] - if getattr(metadata.fp4_mla_state, "v_scale_pool", None) is None: + if getattr(fp4_mla_state, "v_scale_pool", None) is None: raise RuntimeError( "FP4 MLA attention decode requires the auxiliary V scale pool to be allocated." ) @@ -1027,9 +1033,7 @@ def run_fp4_mla_attention_decode( capacity = int(q_batch_capacity) # The fused generation update records the Q layout it produced; metadata # that never went through it (contract tests) carries the legacy layout. - recorded_rows_per_seq = getattr( - getattr(metadata, "fp4_mla_state", None), "prequantized_q_rows_per_seq", None - ) + recorded_rows_per_seq = getattr(fp4_mla_state, "prequantized_q_rows_per_seq", None) packed_rows_per_seq = int(recorded_rows_per_seq or 0) if recorded_rows_per_seq is not None and packed_rows_per_seq != _fp4_mla_packed_rows_per_seq( num_heads, query_len_per_seq @@ -1089,7 +1093,7 @@ def run_fp4_mla_attention_decode( ) sf_cache = sf_cache.view(torch.float8_e4m3fn) - v_sf_pool = metadata.fp4_mla_state.v_scale_pool + v_sf_pool = fp4_mla_state.v_scale_pool v_sf = get_fp4_mla_v_scale_pool_view(metadata, v_head_dim=kv_lora_rank)[local_layer].view( torch.float8_e4m3fn ) @@ -1280,7 +1284,7 @@ def run_fp4_mla_attention_decode( global_scale, q_global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], + fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], kv_lens, output, num_heads=num_heads, @@ -1311,7 +1315,7 @@ def run_fp4_mla_attention_decode( core_v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], + fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], kv_lens, output, splits=select_split_count(max_pages * metadata.page_size), @@ -1363,7 +1367,7 @@ def run_fp4_mla_attention_decode( core_v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], + fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], kv_lens, kernel_output, splits=splits, @@ -1384,7 +1388,7 @@ def run_fp4_mla_attention_decode( core_v_sf, global_scale, src_page_ids, - metadata.fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], + fp4_mla_state.paged_kv_indptr_decode[: num_gen_seqs + 1], kv_lens, kernel_output, max_kv_len=max_pages * metadata.page_size, diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py index 482f4c9cd7ce..8f72a9c7be2a 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_context.py @@ -25,6 +25,7 @@ from tensorrt_llm.quantization.mode import QuantMode from .fp4_mla_kernels import _fp8_mla_context_block_table_kernel +from .state import Fp4MlaState, find_fp4_mla if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.trtllm import ( @@ -338,7 +339,7 @@ def _build_fp8_mla_context_metadata( if kv_lens_cpu is None: kv_lens_cpu = meta.prompt_lens_cpu_runtime[: meta.num_contexts] fp8_meta = copy.copy(meta) - fp8_meta.fp4_mla_state = None + fp8_meta.remove_ext(Fp4MlaState) fp8_meta.kv_cache_manager = scratch.cache_manager_view fp8_meta.kv_cache_block_offsets = scratch.block_offsets fp8_meta.block_ids_per_seq = scratch.block_ids_per_seq @@ -354,7 +355,12 @@ def _build_fp8_mla_context_metadata( # Scratch lengths intentionally start from zero. Preserve the actual # absolute positions for Q/K RoPE through the native kernel's explicit # per-token position-offset input. - fp8_meta.helix_position_offsets = meta.fp4_mla_state.positions[: meta.num_ctx_tokens] + from ..features import HelixMetadata + + helix = meta.find_ext(HelixMetadata) + helix = helix.clone_for(fp8_meta) if helix is not None else HelixMetadata(fp8_meta) + helix.position_offsets = find_fp4_mla(meta).positions[: meta.num_ctx_tokens] + fp8_meta.replace_ext(helix) return fp8_meta @@ -363,11 +369,12 @@ def _get_fp8_mla_context_metadata( scratch: _Fp8MlaContextScratch, ) -> "TrtllmAttentionMetadata": """Prepare and reuse one direct metadata view for all layers in a step.""" - state = meta.fp4_mla_state.fp8_context_state + fp4_mla_state = find_fp4_mla(meta) + state = fp4_mla_state.fp8_context_state if state is not None and state[0] is scratch: return state[1] scratch.prepare(meta) fp8_meta = _build_fp8_mla_context_metadata(meta, scratch) - meta.fp4_mla_state.fp8_context_state = (scratch, fp8_meta) + fp4_mla_state.fp8_context_state = (scratch, fp8_meta) return fp8_meta diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py index e4bdf6e2c739..d8ef69317a39 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/layout.py @@ -17,6 +17,7 @@ HP_BLOCK_SIZE, _ceil_div, ) +from .state import find_fp4_mla # FP4 MLA scale-layout helpers @@ -69,7 +70,7 @@ def get_fp4_mla_v_scale_pool_view( v_head_dim: int, ) -> torch.Tensor: """View the auxiliary MLA V-scale pool in Triton's block-scaled layout.""" - pool = getattr(metadata.fp4_mla_state, "v_scale_pool", None) + pool = getattr(find_fp4_mla(metadata), "v_scale_pool", None) if pool is None: raise RuntimeError("FP4 MLA V scale pool is not allocated.") @@ -100,7 +101,7 @@ def get_fp4_mla_v_scale_pool_view( def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata.fp4_mla_state, "kv_global_scale", None) + global_scale = getattr(find_fp4_mla(metadata), "kv_global_scale", None) if ( not isinstance(global_scale, torch.Tensor) or global_scale.device != device @@ -112,7 +113,7 @@ def _get_fp4_mla_global_scale(metadata: Any, device: torch.device) -> torch.Tens def _get_fp4_mla_q_global_scale(metadata: Any, device: torch.device) -> torch.Tensor: - global_scale = getattr(metadata.fp4_mla_state, "q_global_scale", None) + global_scale = getattr(find_fp4_mla(metadata), "q_global_scale", None) if ( not isinstance(global_scale, torch.Tensor) or global_scale.device != device @@ -235,7 +236,7 @@ def _ensure_workspace_tensor( dtype: torch.dtype, device: torch.device, ) -> torch.Tensor: - workspaces = metadata.fp4_mla_state.workspaces + workspaces = find_fp4_mla(metadata).workspaces tensor = workspaces.get(attr_name) needs_alloc = ( tensor is None diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py index bf1db8f08ba5..94cd6f6f519d 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/metadata.py @@ -10,6 +10,7 @@ from .config import _FP4_MLA_PAGE_TABLE_TILE_SIZE, FP4_MLA_TOKENS_PER_BLOCK, _ceil_div from .packed_plan import _launch_fp4_mla_packed_plan, _PackedPlanGeometry +from .state import find_fp4_mla @triton.jit @@ -488,10 +489,13 @@ def configure_fp4_mla_device_page_table( table on the GPU. The materialization kernel decodes V2 page indices and refreshes rows from the final device KV lengths before cache update. """ - metadata.fp4_mla_state.device_page_table = False - metadata.fp4_mla_state.device_page_table_valid = False - metadata.fp4_mla_state.page_table_stride = 0 - metadata.fp4_mla_state.context_repack_max_touched_pages = 1 + from ..features.spec_decoding import find_spec_decoding + + fp4_mla_state = find_fp4_mla(metadata) + fp4_mla_state.device_page_table = False + fp4_mla_state.device_page_table_valid = False + fp4_mla_state.page_table_stride = 0 + fp4_mla_state.context_repack_max_touched_pages = 1 kv_cache_manager = getattr(metadata, "kv_cache_manager", None) num_contexts = int(getattr(metadata, "num_contexts", 0)) @@ -501,9 +505,9 @@ def configure_fp4_mla_device_page_table( num_context_tokens = int(getattr(metadata, "num_ctx_tokens", 0)) num_generation_tokens = num_tokens - num_context_tokens block_offsets = getattr(metadata, "kv_cache_block_offsets", None) - page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) - paged_kv_indptr = getattr(metadata.fp4_mla_state, "_paged_kv_indptr", None) - paged_kv_indptr_decode = getattr(metadata.fp4_mla_state, "paged_kv_indptr_decode", None) + page_ids = getattr(fp4_mla_state, "_paged_kv_indices", None) + paged_kv_indptr = getattr(fp4_mla_state, "_paged_kv_indptr", None) + paged_kv_indptr_decode = getattr(fp4_mla_state, "paged_kv_indptr_decode", None) max_page_capacity = int(getattr(kv_cache_manager, "max_blocks_per_seq", 0) or 0) page_spec = _fp4_mla_page_table_spec(kv_cache_manager) page_index_scale = int(page_spec.cache_page_index_scale) @@ -523,7 +527,7 @@ def configure_fp4_mla_device_page_table( and kv_cache_manager is not None and has_valid_generation and int(getattr(metadata, "beam_width", 1)) == 1 - and not bool(getattr(metadata, "is_spec_dec_tree", False)) + and not getattr(find_spec_decoding(metadata), "is_spec_dec_tree", False) and not bool(getattr(metadata, "locality_domain_enabled", False)) and int(getattr(kv_cache_manager, "tokens_per_block", 0) or 0) == FP4_MLA_TOKENS_PER_BLOCK and max_page_capacity > 0 @@ -535,7 +539,7 @@ def configure_fp4_mla_device_page_table( if not supported: return False - hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) + hp_page_ids = getattr(fp4_mla_state, "hp_page_indices", None) max_pool_id = max(page_spec.cache_pool_id, page_spec.hp_pool_id) if ( not isinstance(hp_page_ids, torch.Tensor) @@ -544,10 +548,10 @@ def configure_fp4_mla_device_page_table( or block_offsets.shape[0] <= max_pool_id ): return False - metadata.fp4_mla_state.cache_pool_id = int(page_spec.cache_pool_id) - metadata.fp4_mla_state.cache_page_index_scale = int(page_spec.cache_page_index_scale) - metadata.fp4_mla_state.hp_pool_id = int(page_spec.hp_pool_id) - metadata.fp4_mla_state.hp_page_index_scale = int(page_spec.hp_page_index_scale) + fp4_mla_state.cache_pool_id = int(page_spec.cache_pool_id) + fp4_mla_state.cache_page_index_scale = int(page_spec.cache_page_index_scale) + fp4_mla_state.hp_pool_id = int(page_spec.hp_pool_id) + fp4_mla_state.hp_page_index_scale = int(page_spec.hp_page_index_scale) max_pages = max_page_capacity host_kv_lens_available = ( @@ -579,7 +583,7 @@ def configure_fp4_mla_device_page_table( max_context_len, FP4_MLA_TOKENS_PER_BLOCK, ) - metadata.fp4_mla_state.context_repack_max_touched_pages = min( + fp4_mla_state.context_repack_max_touched_pages = min( max_pages, triton.next_power_of_2(max(1, max_context_pages)), ) @@ -604,15 +608,15 @@ def configure_fp4_mla_device_page_table( ) if not buffers_cover_table: return False - if metadata.fp4_mla_state.hp_page_indices.numel() < required_page_ids: + if fp4_mla_state.hp_page_indices.numel() < required_page_ids: return False - metadata.fp4_mla_state.device_page_table = True - metadata.fp4_mla_state.num_sequences = num_sequences - metadata.fp4_mla_state.page_table_stride = max_pages - metadata.fp4_mla_state.num_blocks = None - metadata.fp4_mla_state.num_context_blocks = num_contexts * max_pages - metadata.fp4_mla_state.num_generation_blocks = num_generation_sequences * max_pages + fp4_mla_state.device_page_table = True + fp4_mla_state.num_sequences = num_sequences + fp4_mla_state.page_table_stride = max_pages + fp4_mla_state.num_blocks = None + fp4_mla_state.num_context_blocks = num_contexts * max_pages + fp4_mla_state.num_generation_blocks = num_generation_sequences * max_pages return True @@ -622,15 +626,16 @@ def materialize_fp4_mla_device_page_table( generation_kv_lens: Optional[torch.Tensor] = None, ) -> None: """Refresh the fixed-stride context and generation page table once per forward.""" - if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): + fp4_mla_state = find_fp4_mla(metadata) + if not bool(getattr(fp4_mla_state, "device_page_table", False)): raise RuntimeError("FP4 MLA requires fixed-stride device page metadata.") - if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): + if bool(getattr(fp4_mla_state, "device_page_table_valid", False)): return num_contexts = int(metadata.num_contexts) num_sequences = int(metadata.num_seqs) num_generation_sequences = num_sequences - num_contexts - max_pages = int(metadata.fp4_mla_state.page_table_stride) + max_pages = int(fp4_mla_state.page_table_stride) if num_sequences <= 0 or max_pages <= 0: raise RuntimeError( "FP4 MLA device page metadata requires positive sequence and page capacities." @@ -662,15 +667,15 @@ def materialize_fp4_mla_device_page_table( f"{num_generation_sequences} entries." ) - cache_pool_id = int(getattr(metadata.fp4_mla_state, "cache_pool_id", 0)) + cache_pool_id = int(getattr(fp4_mla_state, "cache_pool_id", 0)) block_offsets = metadata.kv_cache_block_offsets[ cache_pool_id, :num_sequences, 0, :max_pages, ] - page_ids = metadata.fp4_mla_state._paged_kv_indices[: num_sequences * max_pages] - page_index_scale = int(metadata.fp4_mla_state.cache_page_index_scale) + page_ids = fp4_mla_state._paged_kv_indices[: num_sequences * max_pages] + page_index_scale = int(fp4_mla_state.cache_page_index_scale) if page_index_scale <= 0: raise RuntimeError("FP4 MLA device page metadata requires a positive page-index scale.") grid = ( @@ -679,8 +684,8 @@ def materialize_fp4_mla_device_page_table( ) _fp4_mla_materialize_page_table_kernel[grid]( page_ids, - metadata.fp4_mla_state._paged_kv_indptr, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state._paged_kv_indptr, + fp4_mla_state.paged_kv_indptr_decode, block_offsets, kv_lens, generation_kv_lens, @@ -693,11 +698,11 @@ def materialize_fp4_mla_device_page_table( PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, num_warps=4, ) - hp_page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) + hp_page_ids = getattr(fp4_mla_state, "hp_page_indices", None) if not isinstance(hp_page_ids, torch.Tensor): raise RuntimeError("FP4 MLA requires an HP page-table output tensor.") - hp_pool_id = int(metadata.fp4_mla_state.hp_pool_id) - hp_page_index_scale = int(metadata.fp4_mla_state.hp_page_index_scale) + hp_pool_id = int(fp4_mla_state.hp_pool_id) + hp_page_index_scale = int(fp4_mla_state.hp_page_index_scale) hp_block_offsets = metadata.kv_cache_block_offsets[ hp_pool_id, :num_sequences, @@ -706,8 +711,8 @@ def materialize_fp4_mla_device_page_table( ] _fp4_mla_materialize_page_table_kernel[grid]( hp_page_ids, - metadata.fp4_mla_state._paged_kv_indptr, - metadata.fp4_mla_state.paged_kv_indptr_decode, + fp4_mla_state._paged_kv_indptr, + fp4_mla_state.paged_kv_indptr_decode, hp_block_offsets, kv_lens, generation_kv_lens, @@ -720,7 +725,7 @@ def materialize_fp4_mla_device_page_table( PAGE_TILE_SIZE=_FP4_MLA_PAGE_TABLE_TILE_SIZE, num_warps=4, ) - metadata.fp4_mla_state.device_page_table_valid = True + fp4_mla_state.device_page_table_valid = True def _host_int_list_during_forward(value: Any, start: int, end: int) -> Optional[list[int]]: @@ -737,7 +742,7 @@ def _fp4_mla_generation_num_blocks_device(metadata: Any) -> torch.Tensor: slots are never consumed. """ num_gen = metadata.num_seqs - metadata.num_contexts - return metadata.fp4_mla_state.paged_kv_indptr_decode[num_gen : num_gen + 1] + return find_fp4_mla(metadata).paged_kv_indptr_decode[num_gen : num_gen + 1] def _fp4_mla_uniform_generation_lengths( @@ -752,6 +757,7 @@ def _fp4_mla_uniform_generation_lengths( step's packed split plan (one CuTe DSL kernel instead of two launches), and a stale plan is rewritten even when the lengths are already current. """ + fp4_mla_state = find_fp4_mla(metadata) if num_gen <= 0 or num_gen_tokens % num_gen != 0: raise RuntimeError("FP4 MLA generation requires a non-empty uniform request batch.") if plan_geometry is not None and plan_geometry.num_requests != num_gen: @@ -761,8 +767,8 @@ def _fp4_mla_uniform_generation_lengths( num_seqs = metadata.num_seqs kv_lens_gen = metadata.kv_lens_cuda_runtime[num_contexts:num_seqs] prompt_lens_gen = metadata.prompt_lens_cuda_runtime[num_contexts:num_seqs] - corrected_kv_lens = getattr(metadata.fp4_mla_state, "generation_kv_lens", None) - generation_lens = getattr(metadata.fp4_mla_state, "generation_append_lens", None) + corrected_kv_lens = getattr(fp4_mla_state, "generation_kv_lens", None) + generation_lens = getattr(fp4_mla_state, "generation_append_lens", None) tensors = ( kv_lens_gen, prompt_lens_gen, @@ -780,18 +786,18 @@ def _fp4_mla_uniform_generation_lengths( record_for_capture = bool( getattr(metadata, "is_cuda_graph", False) and torch.cuda.is_current_stream_capturing() - and not getattr(metadata.fp4_mla_state, "generation_lengths_capture_recorded", False) + and not getattr(fp4_mla_state, "generation_lengths_capture_recorded", False) ) precomputed = ( not record_for_capture - and metadata.fp4_mla_state.generation_lengths_num_tokens == num_gen_tokens - and metadata.fp4_mla_state.generation_lengths_num_seqs == num_gen - and metadata.fp4_mla_state.generation_lengths_num_contexts == num_contexts + and fp4_mla_state.generation_lengths_num_tokens == num_gen_tokens + and fp4_mla_state.generation_lengths_num_seqs == num_gen + and fp4_mla_state.generation_lengths_num_contexts == num_contexts ) plan_pending = ( plan_geometry is not None and plan_geometry.key is not None - and metadata.fp4_mla_state.packed_plan_key != plan_geometry.key + and fp4_mla_state.packed_plan_key != plan_geometry.key ) if not precomputed: if plan_pending: @@ -815,11 +821,11 @@ def _fp4_mla_uniform_generation_lengths( num_gen_tokens=num_gen_tokens, num_gen=num_gen, ) - metadata.fp4_mla_state.generation_lengths_num_tokens = num_gen_tokens - metadata.fp4_mla_state.generation_lengths_num_seqs = num_gen - metadata.fp4_mla_state.generation_lengths_num_contexts = num_contexts + fp4_mla_state.generation_lengths_num_tokens = num_gen_tokens + fp4_mla_state.generation_lengths_num_seqs = num_gen + fp4_mla_state.generation_lengths_num_contexts = num_contexts if record_for_capture: - metadata.fp4_mla_state.generation_lengths_capture_recorded = True + fp4_mla_state.generation_lengths_capture_recorded = True elif plan_pending: _launch_fp4_mla_packed_plan(plan_geometry, corrected_kv_lens[:num_gen], metadata) return corrected_kv_lens[:num_gen], generation_lens[:num_gen] @@ -830,9 +836,10 @@ def _materialize_fp4_mla_device_page_table_for_forward( generation_kv_lens: Optional[torch.Tensor] = None, ) -> None: """Materialize all fixed-stride rows from final per-forward device lengths.""" - if not bool(getattr(metadata.fp4_mla_state, "device_page_table", False)): + fp4_mla_state = find_fp4_mla(metadata) + if not bool(getattr(fp4_mla_state, "device_page_table", False)): raise RuntimeError("FP4 MLA cache update requires fixed-stride device page metadata.") - if bool(getattr(metadata.fp4_mla_state, "device_page_table_valid", False)): + if bool(getattr(fp4_mla_state, "device_page_table_valid", False)): return num_contexts = int(metadata.num_contexts) @@ -856,12 +863,13 @@ def _materialize_fp4_mla_device_page_table_for_forward( def _max_generation_pages(metadata: Any) -> int: + fp4_mla_state = find_fp4_mla(metadata) num_gen = metadata.num_seqs - metadata.num_contexts if num_gen <= 0: return 0 - if not getattr(metadata.fp4_mla_state, "device_page_table", False): + if not getattr(fp4_mla_state, "device_page_table", False): raise RuntimeError("FP4 MLA generation requires fixed-stride device page metadata.") - max_pages = int(metadata.fp4_mla_state.page_table_stride) + max_pages = int(fp4_mla_state.page_table_stride) if max_pages <= 0: raise RuntimeError("FP4 MLA device page-table stride must be positive.") return max_pages @@ -876,7 +884,7 @@ def _fp4_mla_generation_page_ids(metadata: Any, num_gen_seqs: int) -> torch.Tens f"{num_gen_seqs} != {expected_num_gen}." ) max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata.fp4_mla_state, "_paged_kv_indices", None) + page_ids = getattr(find_fp4_mla(metadata), "_paged_kv_indices", None) start = metadata.num_contexts * max_pages end = start + num_gen_seqs * max_pages if ( @@ -899,7 +907,7 @@ def _fp4_mla_generation_hp_page_ids(metadata: Any, num_gen_seqs: int) -> torch.T f"{num_gen_seqs} != {expected_num_gen}." ) max_pages = _max_generation_pages(metadata) - page_ids = getattr(metadata.fp4_mla_state, "hp_page_indices", None) + page_ids = getattr(find_fp4_mla(metadata), "hp_page_indices", None) start = metadata.num_contexts * max_pages end = start + num_gen_seqs * max_pages if ( @@ -927,12 +935,13 @@ def _host_int_list(value: Any, start: int, end: int) -> Optional[list[int]]: def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> bool: + fp4_mla_state = find_fp4_mla(metadata) if getattr(metadata, "is_cuda_graph", False): return False start = metadata.num_contexts end = metadata.num_seqs - block_counts = _host_int_list(getattr(metadata.fp4_mla_state, "num_blocks", None), start, end) + block_counts = _host_int_list(getattr(fp4_mla_state, "num_blocks", None), start, end) if block_counts is not None and ( not block_counts or min(block_counts) != max_pages or max(block_counts) != max_pages ): @@ -948,12 +957,12 @@ def _infer_assume_full_pages(metadata: Any, max_pages: int, page_size: int) -> b tuple(block_counts) if block_counts is not None else None, kv_lens_cuda.data_ptr(), ) - cache = getattr(metadata.fp4_mla_state, "full_pages_cache", None) + cache = getattr(fp4_mla_state, "full_pages_cache", None) if cache is not None and cache[0] == cache_key: return bool(cache[1]) kv_lens = [int(item) for item in kv_lens_cuda[start:end].detach().cpu().tolist()] result = bool(kv_lens) and min(kv_lens) == max(kv_lens) == max_pages * page_size - setattr(metadata.fp4_mla_state, "full_pages_cache", (cache_key, result)) + setattr(fp4_mla_state, "full_pages_cache", (cache_key, result)) return result kv_cache_params = getattr(metadata, "kv_cache_params", None) diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py index adf3b723cfcb..3e0f984174f3 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py @@ -17,6 +17,7 @@ _fp4_mla_packed_rows_per_seq, ) from .layout import _ensure_workspace_tensor +from .state import find_fp4_mla _SM_COUNT_CACHE: dict[int, int] = {} @@ -103,13 +104,14 @@ def _fp4_mla_packed_plan_geometry( generation lengths (``_fp4_mla_uniform_generation_lengths``). ``key`` is None when every request keeps a single task and no plan is needed. """ + fp4_mla_state = find_fp4_mla(metadata) rows_per_seq = _fp4_mla_packed_rows_per_seq(num_heads, query_len) if rows_per_seq == 0 or num_requests <= 0: return None - if not metadata.fp4_mla_state.device_page_table: + if not fp4_mla_state.device_page_table: return None # The fixed-stride generation page-table capacity (see _max_generation_pages). - max_pages = int(metadata.fp4_mla_state.page_table_stride) + max_pages = int(fp4_mla_state.page_table_stride) if max_pages <= 0: return None max_kv_len = max_pages * metadata.page_size @@ -163,7 +165,7 @@ def _fp4_mla_packed_plan_geometry( dtype=torch.int32, device=device, ) - kv_lens = metadata.fp4_mla_state.generation_kv_lens + kv_lens = fp4_mla_state.generation_kv_lens key = ( num_requests, splits, @@ -193,9 +195,12 @@ def _fp4_mla_helix_plan_bounds( metadata: Any, num_tokens: int, *, use_helix: bool ) -> torch.Tensor | None: """Per-token Helix KV bounds the decode plans on, or None (validated at decode).""" - if not use_helix or not bool(getattr(metadata, "_helix_spec_tokens_valid", False)): + from ..features import find_helix + + helix = find_helix(metadata) if use_helix else None + if helix is None or not helix.spec_tokens_valid: return None - bounds = getattr(metadata, "helix_kv_bounds", None) + bounds = helix.kv_bounds if ( not isinstance(bounds, torch.Tensor) or bounds.dtype != torch.int32 @@ -228,4 +233,4 @@ def _launch_fp4_mla_packed_plan( helix_kv_bounds=geometry.helix_kv_bounds, generation_lengths=generation_lengths, ) - metadata.fp4_mla_state.packed_plan_key = geometry.key + find_fp4_mla(metadata).packed_plan_key = geometry.key diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py index 86de69be8607..9525629ffa68 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/state.py @@ -4,13 +4,14 @@ from __future__ import annotations -from dataclasses import dataclass, field -from typing import TYPE_CHECKING +import copy +from dataclasses import MISSING, dataclass, field, fields +from typing import TYPE_CHECKING, Any import torch +from ..extensions import MetadataExtension from .config import FP4_MLA_KV_GLOBAL_SCALE, FP4_MLA_Q_GLOBAL_SCALE, HP_BLOCK_SIZE -from .metadata import configure_fp4_mla_device_page_table, populate_fp4_mla_append_metadata if TYPE_CHECKING: from ....memory_buffer_utils import Buffers @@ -18,8 +19,18 @@ from .fp4_mla_context import _Fp8MlaContextScratch -@dataclass -class Fp4MlaState: +def find_fp4_mla(metadata: Any) -> Fp4MlaState | None: + """The FP4 MLA extension of ``metadata``, or None. + + Accepts the duck-typed metadata views the FP4 MLA helpers take, which may + carry no extensions at all. + """ + find_ext = getattr(metadata, "find_ext", None) + return find_ext(Fp4MlaState) if find_ext is not None else None + + +@dataclass(eq=False) +class Fp4MlaState(MetadataExtension): """State shared by RoPE/cache append and all FP4 FMHAs in one batch. Allocate through the metadata buffer allocator so eager forwards and each @@ -27,6 +38,7 @@ class Fp4MlaState: layer-local: MTP and metadata preparation update it before any FMHA runs. """ + metadata: TrtllmAttentionMetadata | None = None hp_pool: torch.Tensor | None = None hp_page_indices: torch.Tensor | None = None v_scale_pool: torch.Tensor | None = None @@ -80,8 +92,20 @@ def paged_kv_indptr(self) -> torch.Tensor: raise RuntimeError("FP4 MLA paged_kv_indptr is not allocated.") return self._paged_kv_indptr[: self.num_sequences + 1] - @classmethod - def create(cls, metadata: TrtllmAttentionMetadata, buffers: Buffers | None) -> Fp4MlaState: + def clone_for(self, metadata: TrtllmAttentionMetadata) -> Fp4MlaState: + clone = copy.copy(self) + clone.metadata = metadata + clone.workspaces = {} + clone.v_packed_cache_tags = {} + return clone + + def allocate(self, buffers: Buffers | None) -> None: + for f in fields(self): + if f.name == "metadata": + continue + default = f.default_factory() if f.default_factory is not MISSING else f.default + setattr(self, f.name, default) + metadata = self.metadata manager = metadata.kv_cache_manager hp_ring_size = manager.fp4_mla_hp_pool_size if hp_ring_size < HP_BLOCK_SIZE: @@ -117,7 +141,9 @@ def allocate( capture_graph=metadata.is_cuda_graph, ) - state = cls(hp_pool=hp_pool, v_scale_pool=manager.get_mla_v_scale_pool()) + state = self + state.hp_pool = hp_pool + state.v_scale_pool = manager.get_mla_v_scale_pool() tokens = (metadata.max_num_tokens,) sequences = (metadata.max_num_sequences,) indptr = (metadata.max_num_sequences + 1,) @@ -134,7 +160,6 @@ def allocate( state.q_global_scale.fill_(FP4_MLA_Q_GLOBAL_SCALE) state.kv_global_scale = allocate("kv_global_scale", (1,), torch.float32) state.kv_global_scale.fill_(FP4_MLA_KV_GLOBAL_SCALE) - return state def invalidate_generation_lengths(self) -> None: self.generation_lengths_num_tokens = -1 @@ -143,7 +168,17 @@ def invalidate_generation_lengths(self) -> None: self.generation_lengths_capture_recorded = False self.packed_plan_key = None - def prepare(self, metadata: TrtllmAttentionMetadata, kv_lens: torch.Tensor) -> None: + def prepare(self) -> None: + from .metadata import configure_fp4_mla_device_page_table + + metadata = self.metadata + num_seqs = metadata.num_seqs + # FP4 MLA needs the per-forward append lengths rather than the original + # prompt lengths. Context scratch-cache metadata and MTP generation + # both consume these runtime views. + metadata.prompt_lens_cuda_runtime = metadata.seq_lens_kv_cuda[:num_seqs] + metadata.prompt_lens_cpu_runtime = metadata.seq_lens_kv[:num_seqs] + kv_lens = metadata.kv_lens_runtime # All local layers reuse this FP8 view only within the current batch. self.fp8_context_state = None self.invalidate_generation_lengths() @@ -161,23 +196,23 @@ def prepare(self, metadata: TrtllmAttentionMetadata, kv_lens: torch.Tensor) -> N if not metadata.is_cuda_graph and metadata.num_tokens > 0 and metadata.num_contexts > 0: self.populate_append_metadata(metadata) - def on_update_kv_lens(self, metadata: TrtllmAttentionMetadata) -> None: + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: + metadata = self.metadata + if spec_dec_step: + num_seqs = metadata.num_seqs + metadata.prompt_lens_cuda_runtime = metadata.seq_lens_kv_cuda[:num_seqs] + if not torch.cuda.is_current_stream_capturing(): + metadata.prompt_lens_cpu_runtime = metadata.seq_lens_kv[:num_seqs] self.device_page_table_valid = False if metadata.num_tokens > 0: self.invalidate_generation_lengths() if metadata.num_contexts > 0: self.populate_append_metadata(metadata) - def update_for_spec_dec(self, metadata: TrtllmAttentionMetadata) -> None: - num_seqs = metadata.num_seqs - metadata.prompt_lens_cuda_runtime = metadata.seq_lens_kv_cuda[:num_seqs] - if not torch.cuda.is_current_stream_capturing(): - metadata.prompt_lens_cpu_runtime = metadata.seq_lens_kv[:num_seqs] - self.on_update_kv_lens(metadata) - - def restore_from_spec_dec(self, metadata: TrtllmAttentionMetadata) -> None: + def on_spec_dec_restored(self) -> None: # Rebind aliases to the restored stable tensors, never to the temporary # spec-dec clone whose storage may be released after graph capture. + metadata = self.metadata num_seqs = metadata.num_seqs metadata.kv_lens_cuda_runtime = metadata.kv_lens_cuda[:num_seqs] metadata.prompt_lens_cuda_runtime = metadata.seq_lens_kv_cuda[:num_seqs] @@ -185,6 +220,8 @@ def restore_from_spec_dec(self, metadata: TrtllmAttentionMetadata) -> None: metadata.prompt_lens_cpu_runtime = metadata.seq_lens_kv[:num_seqs] def populate_append_metadata(self, metadata: TrtllmAttentionMetadata) -> None: + from .metadata import populate_fp4_mla_append_metadata + num_seqs = metadata.num_contexts + metadata.num_generations if num_seqs == 0 or metadata.num_tokens == 0: return diff --git a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py index 5c92b130ea0b..32035bfd5230 100644 --- a/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py +++ b/tensorrt_llm/_torch/attention/backends/fp4_mla/v_cache.py @@ -26,6 +26,7 @@ ) from .fp4_mla_kernels import _fp4_mla_rebuild_v_scale_from_k_scale_kernel from .layout import _ensure_workspace_tensor, get_fp4_mla_v_scale_pool_size +from .state import find_fp4_mla def _shared_v_pack_storage_enabled() -> bool: @@ -255,7 +256,7 @@ def _set_triton_v_packed_cache_valid( if _shared_v_pack_storage_enabled() else _triton_v_packed_valid_attr(layer_idx) ) - metadata.fp4_mla_state.v_packed_cache_tags[valid_attr] = _triton_v_packed_cache_tag( + find_fp4_mla(metadata).v_packed_cache_tags[valid_attr] = _triton_v_packed_cache_tag( layer_idx, kv_cache, v_head_dim=v_head_dim, @@ -284,7 +285,7 @@ def _is_triton_v_packed_cache_valid( if _shared_v_pack_storage_enabled() else _triton_v_packed_valid_attr(layer_idx) ) - return metadata.fp4_mla_state.v_packed_cache_tags.get(valid_attr) == _triton_v_packed_cache_tag( + return find_fp4_mla(metadata).v_packed_cache_tags.get(valid_attr) == _triton_v_packed_cache_tag( layer_idx, kv_cache, v_head_dim=v_head_dim, @@ -322,7 +323,7 @@ def _get_triton_v_packed_cache( page_ids=page_ids, ): return None - v_packed = metadata.fp4_mla_state.workspaces.get(_triton_v_packed_attr(layer_idx)) + v_packed = find_fp4_mla(metadata).workspaces.get(_triton_v_packed_attr(layer_idx)) expected_shape = _v_packed_shape(kv_cache, v_head_dim, page_size, block_v) if ( v_packed is None diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 4b411b09edc8..61c36dabb65b 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -8,16 +8,16 @@ from dataclasses import dataclass, field from enum import Enum, IntEnum from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Literal, Optional, - Protocol, Tuple, Type, TypeVar, Union) + Protocol, Sequence, Tuple, Type, TypeVar, Union) import torch from typing_extensions import Self if TYPE_CHECKING: from ...model_config import ModelConfig - from ...modules.mamba.mamba2_metadata import Mamba2Metadata from ...pyexecutor.kv_cache.mixed_precision_staging import \ MixedPrecisionKVStaging + from ...pyexecutor.llm_request import LlmRequest from ...speculative.interface import SpecMetadata from ...speculative.spec_tree_manager import SpecTreeManager from ..kernels.sage_quant import SageQuantized @@ -29,6 +29,7 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig +from ...host_staging import copy_host_to_device from ...memory_buffer_utils import Buffers, get_memory_buffers from ...metadata import KVCacheParams from ...pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -36,9 +37,13 @@ from ...pyexecutor.resource_manager import KVCacheManager from ...pyexecutor.trace_log_utils import log_tensor_size from ...utils import get_model_extra_attrs +from .extensions import ExtensionRegistry, MetadataExtension from .sparse.params import (SparseBackendForwardArgs, SparseMetadataParams, SparseRuntimeParams) +ExtensionT = TypeVar("ExtensionT", bound=MetadataExtension) +_NO_EXTENSIONS = ExtensionRegistry() + try: # Transformers v5 from transformers.configuration_utils import ALLOWED_ATTENTION_LAYER_TYPES @@ -91,7 +96,6 @@ class AttentionMetadata: # HND: [max_num_pages, 2, num_kv_heads, page_size, head_dim] kv_layout: Literal["NHD", "HND"] = "HND" - enable_flash_mla: bool = False enable_context_mla_with_cached_kv: bool = False # Whether CUDA graph is enabled. is_cuda_graph: bool = field(default=False, repr=False) @@ -170,11 +174,24 @@ class AttentionMetadata: all_rank_num_tokens: Optional[List[int]] = None def get_adp_token_counts(self) -> List[int]: - """Local execution-stage counts carried by the input count exchange.""" - return [self.num_tokens] + """Local execution-stage counts carried by the input count exchange. + + The first stage is the input; each extension may append its own. + """ + counts = [self.num_tokens] + for extension in self._attached_extensions(): + counts.extend(extension.adp_token_counts() or ()) + return counts def set_adp_token_counts(self, counts: List[List[int]]) -> List[int]: """Save stage counts and return the input-stage counts for padding.""" + stage = 1 + for extension in self._attached_extensions(): + num_stages = len(extension.adp_token_counts() or ()) + if num_stages: + extension.set_adp_token_counts(counts[stage:stage + num_stages], + counts[0]) + stage += num_stages return counts[0] # These fields are set when changing seq_lens and _num_contexts to avoid computation @@ -184,10 +201,6 @@ def set_adp_token_counts(self, counts: List[List[int]]) -> List[int]: _num_ctx_tokens: int = field(init=False, default=0, repr=False) _num_tokens: int = field(init=False, default=0, repr=False) - mamba_metadata: Optional[Any] = None - mamba_chunk_size: int = 128 - mamba_metadata_cls: Optional[Type["Mamba2Metadata"]] = None - # The number of tokens in the padded sequence. padded_num_tokens: Optional[int] = None @@ -213,18 +226,150 @@ def set_adp_token_counts(self, counts: List[List[int]]) -> List[int]: mixed_precision_staging: Optional["MixedPrecisionKVStaging"] = field( init=False, default=None, repr=False, compare=False) + # Feature, sparse-algorithm and model state attached to this metadata; see + # extensions.py. The types are attached, with their requirements, when the + # metadata is created. + extension_types: Sequence[Type[MetadataExtension]] = () + _extensions: ExtensionRegistry = field(init=False, + default_factory=ExtensionRegistry, + repr=False, + compare=False) + # Pinned staging rings for stage_to_device, shared with metadata copies. + _host_staging: Dict[Any, Any] = field(init=False, + default_factory=dict, + repr=False, + compare=False) + def __post_init__(self) -> None: if self.is_cross: assert self.cross is None or self.cross is self, "Cross attention metadata should not have sub metadata" self.cross = self - return + else: + assert self.cross is None or type(self) is type( + self.cross + ), "Top level and cross attention sub metadata type mismatched" + self._post_init_core() + self._attach_extensions() + + def _post_init_core(self) -> None: + """Allocate the core's buffers; runs again on every metadata copy.""" + + def _default_extension_types(self) -> list: + """Extensions this core family attaches from its own configuration.""" + if isinstance(self.kv_cache_manager, BaseMambaCacheManager): + from .features.recurrent import RecurrentMetadata + return [RecurrentMetadata] + return [] + + def _attach_extensions(self) -> None: + # A copy made by copy.copy shares its source's registry; give it its own + # extensions before anything allocates into them. + self._extensions = self._attached_extensions().bound_to(self) + self._extensions.attach( + self, [*self._default_extension_types(), *self.extension_types]) + for extension in self._attached_extensions(): + extension.allocate(self.cuda_graph_buffers) + + def _attached_extensions(self) -> ExtensionRegistry: + # Metadata built without __init__ (as some tests do) has no registry. + return self.__dict__.get("_extensions", _NO_EXTENSIONS) + + def ext(self, cls: Type[ExtensionT]) -> ExtensionT: + """The attached extension of type ``cls``; raises when none is.""" + found = self._attached_extensions().find(cls) + if found is None: + raise LookupError( + f"{type(self).__name__} has no {cls.__name__} extension attached." + ) + return found + + def find_ext(self, cls: Type[ExtensionT]) -> Optional[ExtensionT]: + """The attached extension of type ``cls``, or None.""" + return self._attached_extensions().find(cls) - assert self.cross is None or type(self) is type( - self.cross - ), "Top level and cross attention sub metadata type mismatched" + def add_ext(self, cls: Type[ExtensionT]) -> ExtensionT: + """Attach ``cls`` and its requirements if absent; return the extension.""" + self._extensions = self._attached_extensions().bound_to(self) + for extension in self._extensions.attach(self, (cls, )): + extension.allocate(self.cuda_graph_buffers) + return self.ext(cls) + + def replace_ext(self, extension: MetadataExtension) -> None: + """Attach a built extension in place of any of its type. + + For derived views made by ``copy.copy``: binds this metadata to its own + registry first, so the source keeps its extensions. + """ + assert extension.metadata is self, "the extension must be built for this metadata" + self._extensions = self._attached_extensions().bound_to(self) + self._extensions.replace(extension) + + def remove_ext(self, cls: Type[MetadataExtension]) -> None: + """Detach the extension of type ``cls`` from a derived view. + + Like ``replace_ext``, binds this metadata to its own registry first. + """ + self._extensions = self._attached_extensions().bound_to(self) + self._extensions.remove(cls) + + @property + def extensions(self) -> List[MetadataExtension]: + """The attached extensions in dependency order.""" + return list(self._attached_extensions()) + + def get_buffer(self, shape: Sequence[int], dtype: torch.dtype, + name: str) -> torch.Tensor: + """A persistent buffer from this metadata's pool, graph-stable under capture.""" + return self.get_empty(self.cuda_graph_buffers, + list(shape), + dtype=dtype, + cache_name=name, + capture_graph=self.is_cuda_graph) + + def stage_to_device(self, key: Any, destination: torch.Tensor, + value: torch.Tensor) -> None: + """Copy a host tensor into ``destination`` through pinned staging.""" + copy_host_to_device(self._host_staging, key, destination, value) + + def begin_forward(self) -> None: + """Hook to be called before the model forward of each step.""" + for extension in self._attached_extensions(): + extension.begin_forward() + + def bind_context_requests(self, requests: List["LlmRequest"], *, + prefill_cuda_graphs: bool) -> None: + """Hand the scheduled context requests to the extensions before ``prepare()``.""" + for extension in self._attached_extensions(): + extension.bind_context_requests( + requests, prefill_cuda_graphs=prefill_cuda_graphs) + + def defer_decode_outputs( + self) -> Optional[contextlib.AbstractContextManager[None]]: + """A scope for preparing a decode CUDA graph's metadata, or None. + + None when no extension can leave outputs to the graph's own + preprocessing; otherwise the extensions' scopes entered together. + """ + scopes = [ + scope for scope in (extension.defer_decode_outputs() + for extension in self._attached_extensions()) + if scope is not None + ] + if not scopes: + return None + + @contextlib.contextmanager + def deferring(): + with contextlib.ExitStack() as stack: + for scope in scopes: + stack.enter_context(scope) + yield + + return deferring() def on_update_kv_lens(self): - pass + for extension in self._attached_extensions(): + extension.on_kv_lens_changed(False) def on_update(self): if (self._seq_lens is not None @@ -351,7 +496,50 @@ def prepare(self): """ Hook to be called before the forward step of the model. """ - self._prepare_mamba_metadata() + with contextlib.ExitStack() as scope: + for extension in self._attached_extensions(): + scope.enter_context(extension.prepare_scope()) + for extension in self._attached_extensions(): + extension.pre_derive_lengths() + self._prepare_core() + for extension in self._attached_extensions(): + extension.prepare() + for extension in self._attached_extensions(): + extension.post_derive() + + def _prepare_core(self) -> None: + """Derive the core's per-step state; ``prepare`` runs the extensions around it.""" + + def _override_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: + """Let the extensions rewrite freshly derived host KV lengths.""" + if len(self._attached_extensions()) == 0: + return kv_lens + if kv_lens is self._seq_lens or kv_lens is self._seq_lens_kv: + kv_lens = kv_lens.clone() + for extension in self._attached_extensions(): + extension.override_kv_lens(kv_lens) + return kv_lens + + def _needs_full_block_table_width(self) -> bool: + return any(extension.needs_full_block_table_width() + for extension in self._attached_extensions()) + + def prepare_for_draft_forward(self) -> dict | None: + """Point the extensions at the draft model's state before a draft forward.""" + saved = {} + for extension in self._attached_extensions(): + state = extension.draft_swap() + if state is not None: + saved[type(extension)] = state + return saved or None + + def restore_after_draft_forward(self, saved_state: dict | None) -> None: + """Undo prepare_for_draft_forward.""" + if not saved_state: + return + for extension in self._attached_extensions(): + if type(extension) in saved_state: + extension.draft_restore(saved_state[type(extension)]) def _has_mixed_precision_kv(self) -> bool: """Whether the KV cache manager holds mixed-precision layers.""" @@ -442,26 +630,6 @@ def prepare_encoder_only(self) -> None: """ self.prepare() - def _prepare_mamba_metadata(self): - if self.mamba_metadata is False: - return - - if self.mamba_metadata is None: - if isinstance(self.kv_cache_manager, BaseMambaCacheManager): - metadata_cls = self.mamba_metadata_cls - if metadata_cls is None: - from ...modules.mamba.mamba2_metadata import Mamba2Metadata - metadata_cls = Mamba2Metadata - self.mamba_metadata = metadata_cls( - self.max_num_requests, - self.mamba_chunk_size, - max_num_tokens=self.max_num_tokens) - else: - self.mamba_metadata = False - return - - self.mamba_metadata.prepare(self) - def create_cuda_graph_metadata(self, max_batch_size: int, sub_cross_metadata: bool = False, @@ -502,17 +670,6 @@ def create_cuda_graph_metadata(self, if self.is_cross: cuda_graph_metadata.seq_lens_kv = torch.zeros((max_batch_size, ), dtype=torch.int) - if self.enable_flash_mla: - if self.kv_cache_manager is not None: - cuda_graph_metadata.block_ids_per_seq = torch.zeros( - [ - self.kv_cache_manager.max_batch_size, - self.kv_cache_manager.max_blocks_per_seq - ], - dtype=torch.int32, - device='cuda', - ) - if not encode_only: # Decoder CUDA graphs are always generation-only (no context # requests). Encoder CUDA graphs are the opposite (all context); @@ -537,6 +694,8 @@ def restore_from_spec_dec(self) -> None: for f, v in self._saved_tensors.items(): setattr(self, f, v) self._saved_tensors.clear() + for extension in self._attached_extensions(): + extension.on_spec_dec_restored() @property def has_spec_dec_saved_state(self) -> bool: @@ -562,16 +721,9 @@ def update_spec_dec_param( Backends that consume gen-only slots (e.g. dynamic tree) must skip these rows to align with the XQA kernel's expected row layout. """ - - def update_helix_param( - self, - helix_position_offsets: List[int], - helix_is_inactive_rank: List[bool], - helix_owned_new_tokens: Optional[List[int]] = None, - ) -> None: - """ - Hook to be called when using helix parallelism. - """ + for extension in self._attached_extensions(): + extension.on_spec_dec_param(max_draft_len, max_total_draft_tokens, + spec_tree_manager) def create_cross_metadata( self, @@ -703,6 +855,8 @@ def update_for_spec_dec(self) -> None: """ Hook to be called during forward when using spec-dec one-model mode. """ + for extension in self._attached_extensions(): + extension.on_kv_lens_changed(True) @staticmethod def get_empty(buffers: Buffers, @@ -1135,6 +1289,16 @@ class AttentionBackend(Generic[TMetadata]): Base class for attention backends. """ Metadata: Type[TMetadata] = AttentionMetadata + # Extensions every metadata of this backend carries, e.g. a sparse + # algorithm's state; see extensions.py. + metadata_extensions: Tuple[Type[MetadataExtension], ...] = () + + @classmethod + def create_metadata(cls, **kwargs) -> TMetadata: + """Build this backend's metadata with its extensions attached.""" + extension_types = (*cls.metadata_extensions, + *kwargs.pop("extension_types", ())) + return cls.Metadata(extension_types=extension_types, **kwargs) def __init__( self, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/README.md b/tensorrt_llm/_torch/attention/backends/sparse/csa2/README.md index 901ec1e5d393..af27f88a72ca 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/README.md +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/README.md @@ -80,7 +80,7 @@ and delegates attention to the independent CSA2 module. ## Projections, compression and indexing `DeepseekV41Attention` owns projections, normalization, RoPE, compression and -grouped output projection. It consumes prepared `CSA2TrtllmMetadata`, retains +grouped output projection. It consumes the prepared `CSA2Metadata` extension, retains CSA2's omission of projected-Q head normalization, and derives index K from the compressed main latent before main-KV RoPE. Ratio-two compression reuses the native V4 non-overlap compressor with FP32 value/gate state and zero APE @@ -352,9 +352,10 @@ existing output contract; this does not add arbitrary hidden-state exports. ## Metadata, graphs and workspace -`CSA2TrtllmMetadata.prepare()` resolves real request IDs, cached lengths and V2 -page converters. Layer-specific SWA/visibility, owner page tables/write slots, -compression inputs, candidates and routing results live directly on metadata. +Preparing a metadata that carries `CSA2Metadata` resolves real request IDs, +cached lengths and V2 page converters. Layer-specific SWA/visibility, owner page +tables/write slots, compression inputs, candidates and routing results live on +the extension. `global_slot_tile()` resolves pages without a persistent token-by-context map. As in DeepSeek-V4, the manager converts every per-layer page table of a batch (SWA, GLOBAL and compressor state) in one CuTe DSL launch from a prepare-time diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/backend.py index 3b618112f5fb..5525ce428665 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/backend.py @@ -18,13 +18,13 @@ PredefinedAttentionMask, merge_attention_forward_args, ) -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo from .indexer import CSA2Indexer -from .metadata import CSA2TrtllmMetadata +from .metadata import CSA2Metadata from .params import CSA2BackendForwardArgs, CSA2Layout, CSA2Mode, CSA2Params, select_csa2_backend if TYPE_CHECKING: @@ -76,12 +76,13 @@ def _csa2_parent_quant_config(quant_config, native_fp8: bool): class CSA2TrtllmAttention(TrtllmAttention): """CSA2 sparse prediction with the inherited native trtllm-gen forward.""" - Metadata = CSA2TrtllmMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (CSA2Metadata,) compute_backend = "trtllm" @classmethod def runtime_workspace_bytes_per_token(cls, model_config, mapping) -> int: - return CSA2TrtllmMetadata.cache_gather_bytes_per_token(model_config) + return CSA2Metadata.cache_gather_bytes_per_token(model_config) @classmethod def runtime_workspace_is_chunked_prefill_bounded(cls, model_config) -> bool: @@ -115,7 +116,7 @@ def runtime_workspace_fixed_bytes( ).multi_processor_count def arena_bytes(queries: int, page_capacity: int) -> int: - return CSA2TrtllmMetadata.workspace_reservation_bytes( + return CSA2Metadata.workspace_reservation_bytes( max_batch_size, queries, page_capacity * 64, @@ -268,9 +269,10 @@ def create_output(self, q, *, is_quantize_output, **kwargs): def prepare_sparse_inputs( self, q: torch.Tensor, - metadata: CSA2TrtllmMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> CSA2BackendForwardArgs: + csa2 = metadata.ext(CSA2Metadata) if self.native_fp8 and ( forward_args.out_scale is not None or forward_args.out_scale_sf is not None @@ -279,7 +281,7 @@ def prepare_sparse_inputs( or (forward_args.output is not None and forward_args.output.dtype != torch.bfloat16) ): raise ValueError("CSA2 native FP8 attention requires unquantized BF16 output") - if metadata.swa_pool.dtype != self.staging_dtype: + if csa2.swa_pool.dtype != self.staging_dtype: raise ValueError("CSA2 tile staging dtype must match its attention backend") inputs = forward_args.sparse_backend_args if not isinstance(inputs, CSA2BackendForwardArgs): @@ -288,7 +290,7 @@ def prepare_sparse_inputs( if ( q.shape != (count, self.num_heads * _HEAD_DIM) or q.dtype != torch.bfloat16 - or metadata.num_query_heads != self.num_heads + or csa2.num_query_heads != self.num_heads ): raise ValueError("CSA2 Q geometry/dtype must match its BF16 metadata") if forward_args.attention_input_type not in ( @@ -299,6 +301,7 @@ def prepare_sparse_inputs( if inputs.state is not None: state = inputs.state source = state.metadata + source_csa2 = source.ext(CSA2Metadata) manager = source.kv_cache_manager layout = self.sparse_params.layout if layout is None: @@ -306,11 +309,11 @@ def prepare_sparse_inputs( layer = layout.layer(self.layer_idx) start, end = inputs.query_start, inputs.query_start + count if start == 0: - source.enter_layer(layer) + source_csa2.enter_layer(layer) global_rows = None if ( layer.mode == CSA2Mode.FULL - and self.layer_idx in source.csa2_precomputed_kv_layers + and self.layer_idx in source_csa2.csa2_precomputed_kv_layers ): if state.main_kv is not None or state.index_k is not None: raise ValueError("CED precomputed global KV must not be projected twice") @@ -321,7 +324,7 @@ def prepare_sparse_inputs( ) global_rows = ( layer.kv_source, - source.csa2_main_write_slots[layer.kv_source], + source_csa2.csa2_main_write_slots[layer.kv_source], state.main_kv, state.index_k, state.main_transform, @@ -338,7 +341,7 @@ def prepare_sparse_inputs( ) packed = manager.write_layer_rows( self.layer_idx, - source.csa2_swa_write_slots[self.layer_idx], + source_csa2.csa2_swa_write_slots[self.layer_idx], state.swa_kv, state.swa_transform, global_rows=global_rows, @@ -362,37 +365,37 @@ def prepare_sparse_inputs( # batch. FMHA query tiles only consume the published selections. if start == 0: self.indexer(state, 0, state.swa_kv.shape[0]) - logical = source.csa2_indices[self.layer_idx][start:end] + logical = source_csa2.csa2_indices[self.layer_idx][start:end] elif layer.mode == CSA2Mode.REUSE: - if layer.index_source not in source.csa2_indices: + if layer.index_source not in source_csa2.csa2_indices: raise ValueError("CSA2 index source did not run in this forward") - indices = source.csa2_indices[layer.index_source] + indices = source_csa2.csa2_indices[layer.index_source] if indices.shape[0] != state.swa_kv.shape[0]: raise ValueError("CSA2 Reuse rows must match their source's packed query order") logical = indices[start:end] slots = None mapping = None if logical is not None: - visible = source.csa2_visible_lengths[self.layer_idx][start:end] + visible = source_csa2.csa2_visible_lengths[self.layer_idx][start:end] if self._native_slot_mapping(logical): # The staging kernels resolve logical rows through the # owner's live page table themselves (one launch fewer). - owner = source.csa2_kv_sources[self.layer_idx] + owner = source_csa2.csa2_kv_sources[self.layer_idx] slots = logical mapping = ( - source.csa2_global_page_tables[owner], - source.csa2_global_page_sizes[owner], - source.csa2_global_max_positions[owner], - source.csa2_token_requests[start:end], + source_csa2.csa2_global_page_tables[owner], + source_csa2.csa2_global_page_sizes[owner], + source_csa2.csa2_global_max_positions[owner], + source_csa2.csa2_token_requests[start:end], visible, ) else: - slots = source.global_slot_tile( + slots = source_csa2.global_slot_tile( self.layer_idx, start, end, logical, visible_lengths=visible ) inputs = CSA2BackendForwardArgs( swa_pool=manager.get_swa_buffer(self.layer_idx), - swa_indices=source.csa2_swa_indices[self.layer_idx][start:end], + swa_indices=source_csa2.csa2_swa_indices[self.layer_idx][start:end], main_pool=None if logical is None else manager.get_main_buffer(layer.kv_source), topk_indices=slots, main_logical_indices=logical, @@ -410,7 +413,7 @@ def prepare_sparse_inputs( inputs.output_position_ids is not None or inputs.output_rotary_cos_sin is not None ) and not self.sparse_params.use_packed_sparse_attention: raise ValueError("CSA2 fused output RoPE requires packed sparse attention") - if metadata.swa_pool.device != q.device: + if csa2.swa_pool.device != q.device: raise ValueError("CSA2 Q and metadata must be on the same CUDA device") return inputs @@ -425,6 +428,7 @@ def _native_slot_mapping(self, logical: torch.Tensor) -> bool: ) def sparse_attn_predict(self, q, k, metadata, forward_args): + csa2 = metadata.ext(CSA2Metadata) if self.compute_backend != "trtllm" or self.sparse_params.use_packed_sparse_attention: raise ValueError("CSA2 Flash and packed attention require their explicit module helper") if ( @@ -443,7 +447,7 @@ def sparse_attn_predict(self, q, k, metadata, forward_args): # Repeating a forward on one carrier is fine: it hands back the very # buffers derived below. Any other scale describes a KV cache this path # never reads, so it could only contradict the staging pool. - owned = metadata.stage_kv_scales + owned = csa2.stage_kv_scales for scale in (forward_args.kv_scale_orig_quant, forward_args.kv_scale_quant_orig): if scale is not None and not any(scale is buffer for buffer in owned): raise ValueError("CSA2 FP8 staging derives its own Q/KV scales") @@ -465,7 +469,7 @@ def sparse_attn_predict(self, q, k, metadata, forward_args): # context does take range, bounded by Q's own amax and by the exponent the # main pool can actually need. Split-KV never runs here either way: every # MLA multi-CTA-KV cubin is a generation kernel. - stage_orig_quant, stage_quant_orig = metadata.derive_stage_kv_scales( + stage_orig_quant, stage_quant_orig = csa2.derive_stage_kv_scales( inputs, q=q, q_shared=context, @@ -478,7 +482,7 @@ def sparse_attn_predict(self, q, k, metadata, forward_args): # Q's pair is the staging pair read backwards: dequant_q = 1 / dequant_kv. forward_args.kv_scale_orig_quant = stage_quant_orig forward_args.kv_scale_quant_orig = stage_orig_quant - metadata.stage_selected( + csa2.stage_selected( inputs, kv_scale_orig_quant=stage_orig_quant, kv_scale_quant_orig=stage_quant_orig, @@ -497,32 +501,32 @@ def sparse_attn_predict(self, q, k, metadata, forward_args): forward_args.quant_q_buffer = quant_q.view( metadata.num_tokens, self.num_heads, _HEAD_DIM ) - forward_args.mla_bmm1_scale = metadata.mla_bmm1_scale - forward_args.mla_bmm2_scale = metadata.mla_bmm2_scale + forward_args.mla_bmm1_scale = csa2.mla_bmm1_scale + forward_args.mla_bmm2_scale = csa2.mla_bmm2_scale elif forward_args.attention_input_type == AttentionInputType.context_only: forward_args.quant_scale_qkv = forward_args.kv_scale_orig_quant else: - metadata.stage_selected(inputs, main_logical=inputs.main_logical_indices) + csa2.stage_selected(inputs, main_logical=inputs.main_logical_indices) sparse = forward_args.sparse_runtime_params - sparse.sparse_attn_kv_lens = metadata.prepared_lens + sparse.sparse_attn_kv_lens = csa2.prepared_lens sparse.aux_kv_cache_pool_ptr = ( - metadata.extra_pool.data_ptr() + csa2.extra_pool.data_ptr() if inputs.topk_indices is not None and metadata.num_sparse_topk > 128 else None ) - forward_args.fmha_scheduler_counter = metadata.prepared_counter + forward_args.fmha_scheduler_counter = csa2.prepared_counter forward_args.attention_window_size = metadata.num_sparse_topk # Context Q already includes RoPE and cache publication happened above. # A null latent cache skips the native context append/RoPE preprocessing. forward_args.latent_cache = ( None if forward_args.attention_input_type == AttentionInputType.context_only - else metadata.latent_placeholder + else csa2.latent_placeholder if self.native_fp8 - else metadata.swa_pool[:, 0] + else csa2.swa_pool[:, 0] ) forward_args.q_pe = q.view(metadata.num_tokens, self.num_heads, _HEAD_DIM)[..., 448:] - return metadata.prepared_indices, None + return csa2.prepared_indices, None def _supports_packed(self, q, inputs, forward_args) -> bool: """Validate the explicit packed attention input/output contract.""" @@ -585,26 +589,27 @@ def _validate_helper_output(q, forward_args) -> None: raise ValueError("CSA2 requires an attention sink") def prepare_selected(self, q, metadata, forward_args): + csa2 = metadata.ext(CSA2Metadata) self._validate_helper_output(q, forward_args) if self.sparse_params.use_packed_sparse_attention: raise ValueError("Packed CSA2 must use forward_packed without BF16 staging") inputs = self.prepare_sparse_inputs(q, metadata, forward_args) - metadata.stage_selected(inputs, main_logical=inputs.main_logical_indices) + csa2.stage_selected(inputs, main_logical=inputs.main_logical_indices) query = q.view(metadata.num_tokens, self.num_heads, _HEAD_DIM) - if metadata.shared_plan is not None: + if csa2.shared_plan is not None: return ( query, - metadata.shared_pool, - metadata.prepared_indices, + csa2.shared_pool, + csa2.prepared_indices, forward_args.attention_sinks, 1.0 / (self.q_scaling * math.sqrt(_HEAD_DIM)), ) - kv = metadata.swa_pool + kv = csa2.swa_pool if metadata.num_sparse_topk > _SWA_TILE: - kv = torch.cat((kv, metadata.extra_pool), dim=1) + kv = torch.cat((kv, csa2.extra_pool), dim=1) valid = ( torch.arange(metadata.num_sparse_topk, device=q.device)[None, :] - < metadata.prepared_lens[:, None] + < csa2.prepared_lens[:, None] ) return ( query, @@ -694,10 +699,11 @@ def forward_generation(self, q, metadata, forward_args): return self._forward_selected(q, metadata, forward_args) def _forward_selected(self, q, metadata, forward_args): + csa2 = metadata.ext(CSA2Metadata) if self.attention is None: raise ValueError("CSA2 helper requires its attention backend") inputs = self.attention.prepare_selected(q, metadata, forward_args) - run = self.run_shared if metadata.shared_plan is not None else self.run + run = self.run_shared if csa2.shared_plan is not None else self.run result = run(*inputs).flatten(1) if forward_args.output is not None: forward_args.output.copy_(result) @@ -872,10 +878,11 @@ def forward_generation(self, q, metadata, forward_args): return self._forward_selected(q, metadata, forward_args) def _forward_selected(self, q, metadata, forward_args): + csa2 = metadata.ext(CSA2Metadata) if self.attention is None: raise ValueError("CSA2 helper requires its attention backend") inputs = self.attention.prepare_selected(q, metadata, forward_args) - shared = metadata.shared_plan is not None + shared = csa2.shared_plan is not None query, values, selection, _, scale = inputs width = selection.shape[1] if shared else values.shape[1] key = (query.shape[0], query.shape[1], width, q.device, scale) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/ced.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/ced.py index d86635050f95..d9826ae96be7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/ced.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/ced.py @@ -11,14 +11,15 @@ from tensorrt_llm._torch.modules.mhc.hyper_connection import HCState, mHC from .....pyexecutor.ced_replay import EncoderReplay -from .metadata import CSA2TrtllmMetadata if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + from .....pyexecutor.llm_request import LlmRequest def prepare_ced_global_kv( - layer, hc_state: HCState, metadata: CSA2TrtllmMetadata + layer, hc_state: HCState, metadata: TrtllmAttentionMetadata ) -> torch.Tensor | None: """Write GLOBAL; a context-only source also returns its output before normalization.""" attention = layer.self_attn diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/decoder_replay.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/decoder_replay.py index 16c00b1a8a70..af6e4cdf8e02 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/decoder_replay.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/decoder_replay.py @@ -23,10 +23,12 @@ EncoderReplay, requires_full_decoder_prefill, ) +from .metadata import CSA2Metadata if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + from .....pyexecutor.llm_request import LlmRequest - from .metadata import CSA2TrtllmMetadata __all__ = [ "DecoderReplayPlan", @@ -85,7 +87,7 @@ def num_replay_tokens(self) -> int: def plan_decoder_replay( - metadata: "CSA2TrtllmMetadata", + metadata: TrtllmAttentionMetadata, window: int, requests: Optional[list[LlmRequest]] = None, *, @@ -126,7 +128,8 @@ def plan_decoder_replay( host_seq_lens = seq_lens[:num_requests].tolist() num_cached = list(metadata.kv_cache_params.num_cached_tokens_per_seq[:num_requests]) - context_ends = getattr(metadata, "decoder_context_ends", None) + csa2 = metadata.find_ext(CSA2Metadata) + context_ends = csa2.decoder_context_ends if csa2 is not None else None if context_ends is None and requests and all(hasattr(req, "prompt_len") for req in requests): context_ends = tuple(req.prompt_len for req in requests) replay_lens = list(host_seq_lens) @@ -208,7 +211,7 @@ def plan_decoder_replay( # moment they exist and the only moment they are about to be cleared. -def _swa_is_reused(metadata: "CSA2TrtllmMetadata") -> bool: +def _swa_is_reused(metadata: TrtllmAttentionMetadata) -> bool: """Whether a later request could read this pass's sliding-window KV. Read off the manager rather than the config so a manager that disabled the @@ -222,22 +225,23 @@ def _swa_is_reused(metadata: "CSA2TrtllmMetadata") -> bool: def enter_decoder_replay( - metadata: "CSA2TrtllmMetadata", + metadata: TrtllmAttentionMetadata, plan: DecoderReplayPlan, first_layer: Optional[int] = None, ) -> None: """Rebuild query views and carry encoder publications across the boundary.""" + csa2 = metadata.ext(CSA2Metadata) if plan.replay_all_rank_num_tokens is not None: metadata.all_rank_num_tokens = plan.replay_all_rank_num_tokens if not plan.updates_local_metadata: return - saved_indices = dict(metadata.csa2_indices) + saved_indices = dict(csa2.csa2_indices) saved_candidates = { - name: dict(getattr(metadata, name)) + name: dict(getattr(csa2, name)) for name in ("csa2_candidates", "csa2_candidate_blocks", "csa2_candidate_counts") } - saved_last_layer = metadata._csa2_last_layer - precomputed = set(getattr(metadata, "csa2_precomputed_kv_layers", ())) + saved_last_layer = csa2._csa2_last_layer + precomputed = set(getattr(csa2, "csa2_precomputed_kv_layers", ())) live_kv_lens = metadata.kv_lens_cuda[: metadata.num_seqs].clone() metadata.kv_cache_params.num_cached_tokens_per_seq[: len(plan.replay_num_cached)] = ( plan.replay_num_cached @@ -247,27 +251,27 @@ def enter_decoder_replay( if plan.num_replay_tokens == 0: # No local attention runs. Keep GLOBAL publication and live KV lengths, # but do not ask attention backends to prepare an empty query batch. - metadata.reset_routing() - metadata.csa2_precomputed_kv_layers = precomputed - metadata.csa2_replay_query_rows = plan.rows + csa2.reset_routing() + csa2.csa2_precomputed_kv_layers = precomputed + csa2.csa2_replay_query_rows = plan.rows return if first_layer is not None: - metadata.set_decoder_query_boundary(first_layer) + csa2.set_decoder_query_boundary(first_layer) metadata.prepare() metadata.kv_lens_cuda[: metadata.num_seqs].copy_(live_kv_lens) metadata.on_update_kv_lens() - metadata.csa2_precomputed_kv_layers = precomputed - metadata.csa2_replay_query_rows = plan.rows + csa2.csa2_precomputed_kv_layers = precomputed + csa2.csa2_replay_query_rows = plan.rows for layer, indices in saved_indices.items(): - metadata.csa2_indices[layer] = indices.index_select(0, plan.rows) + csa2.csa2_indices[layer] = indices.index_select(0, plan.rows) for name, published in saved_candidates.items(): - target = getattr(metadata, name) + target = getattr(csa2, name) for layer, value in published.items(): target[layer] = value.index_select(0, plan.rows) - metadata._csa2_last_layer = saved_last_layer + csa2._csa2_last_layer = saved_last_layer # Only context SWA is truncated. Generation rows and global visibility # retain their original logical positions and cache history. - starts = metadata._copy_csa2_tensor( + starts = csa2._copy_csa2_tensor( "decoder_starts", torch.tensor( plan.swa_floors, @@ -278,7 +282,7 @@ def enter_decoder_replay( _apply_decoder_swa_floor(metadata, starts) -def enter_remote_tail_decoder(metadata: "CSA2TrtllmMetadata", context_starts: List[int]) -> None: +def enter_remote_tail_decoder(metadata: TrtllmAttentionMetadata, context_starts: List[int]) -> None: """Hide pre-handoff SWA when a generation worker enters decoder layers. Unlike ordinary decoder replay, the query rows are already exactly the @@ -286,9 +290,10 @@ def enter_remote_tail_decoder(metadata: "CSA2TrtllmMetadata", context_starts: Li decoder layers must see only SWA written at or after each context's split. Generation rows retain their ordinary visibility. """ + csa2 = metadata.ext(CSA2Metadata) if len(context_starts) != metadata.num_contexts: raise ValueError("remote-tail starts must match the context rows") - starts = metadata._copy_csa2_tensor( + starts = csa2._copy_csa2_tensor( "remote_tail_starts", torch.tensor( list(context_starts) + [0] * metadata.num_generations, @@ -299,25 +304,26 @@ def enter_remote_tail_decoder(metadata: "CSA2TrtllmMetadata", context_starts: Li _apply_decoder_swa_floor(metadata, starts) -def _apply_decoder_swa_floor(metadata: "CSA2TrtllmMetadata", starts: torch.Tensor) -> None: +def _apply_decoder_swa_floor(metadata: TrtllmAttentionMetadata, starts: torch.Tensor) -> None: """Hide cached SWA positions below one logical floor per request from reads. Writes keep the prepared floors: query rows below the floor still store their SWA. The next layer resolves its slots against the raised floors. """ - replay = metadata.csa2_replay_start_positions - floors = metadata._get_csa2_buffer("swa_read_floors", tuple(replay.shape), torch.int64) - previous = getattr(metadata, "_csa2_swa_read_floors", None) + csa2 = metadata.ext(CSA2Metadata) + replay = csa2.csa2_replay_start_positions + floors = csa2._get_csa2_buffer("swa_read_floors", tuple(replay.shape), torch.int64) + previous = getattr(csa2, "_csa2_swa_read_floors", None) torch.maximum(replay if previous is None else previous, starts, out=floors) - metadata._csa2_swa_read_floors = floors + csa2._csa2_swa_read_floors = floors # Earlier layers of this forward may have resolved slots below the floors. - metadata._csa2_swa_resolved = False + csa2._csa2_swa_resolved = False # The steady fast path skips the full prepare that restores the floors. - metadata._csa2_steady_snapshot = None + csa2._csa2_steady_snapshot = None def exit_decoder_replay( - metadata: "CSA2TrtllmMetadata", + metadata: TrtllmAttentionMetadata, plan: DecoderReplayPlan, ) -> None: """Restore the encoder pass's shape. @@ -328,6 +334,7 @@ def exit_decoder_replay( values: every forward rebuilds them from these inputs in ``prepare()``, and copying them back would cost a second rebuild to no observable end. """ + csa2 = metadata.ext(CSA2Metadata) if plan.saved_all_rank_num_tokens is not None: metadata.all_rank_num_tokens = plan.saved_all_rank_num_tokens if not plan.updates_local_metadata: @@ -337,10 +344,10 @@ def exit_decoder_replay( ) metadata.prompt_lens = plan.saved_prompt_lens metadata.seq_lens = plan.saved_seq_lens - metadata.csa2_replay_query_rows = None + csa2.csa2_replay_query_rows = None # Publications are indexed in the compact query space and cannot survive # shape restoration. - metadata.reset_routing() + csa2.reset_routing() def gather_replayed_rows( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/indexer.py index 372d6c19cac5..a48fc4306a89 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/indexer.py @@ -29,6 +29,8 @@ except ImportError: # The bundled DeepGEMM predates the sparse indexer kernels _HAS_SPARSE_MQA_LOGITS = False +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + from ..dsa.indexer import _INDEXER_MQA_LOGITS_ELEM_BUDGET, Indexer, _split_prefill_queries from ..dsa.params import DSAParams from ..params import use_self_sampling_gvr @@ -39,7 +41,7 @@ mask_logits_, publish_candidates_, ) -from .metadata import CSA2TrtllmMetadata +from .metadata import CSA2Metadata from .params import CSA2ForwardState, CSA2Layout, CSA2Mode, CSA2Params from .quantization import ( INDEX_DATA_BYTES, @@ -540,7 +542,7 @@ def _select_mapped_logits( def _select_paged_logits( self, - metadata: CSA2TrtllmMetadata, + metadata: TrtllmAttentionMetadata, logits: torch.Tensor, candidates: Optional[torch.Tensor], output_indices: torch.Tensor, @@ -555,8 +557,9 @@ def _select_paged_logits( candidate columns instead. TopK runs over the masked scores, then one CuTe launch validates, sorts and pads the selection with ``-1``. """ - valid = metadata.csa2_indexer_valid_positions - visible = metadata.csa2_indexer_visible_lengths + csa2 = metadata.ext(CSA2Metadata) + valid = csa2.csa2_indexer_valid_positions + visible = csa2.csa2_indexer_visible_lengths width = logits.shape[1] if width == 0 or (candidates is not None and candidates.shape[1] == 0): return output_indices.fill_(-1) @@ -587,8 +590,8 @@ def _select_paged_logits( None, row_ends, is_prefill=False, - radix_aux_indices=metadata.csa2_indexer_radix_aux_indices, - radix_aux_logits=metadata.csa2_indexer_radix_aux_logits, + radix_aux_indices=csa2.csa2_indexer_radix_aux_indices, + radix_aux_logits=csa2.csa2_indexer_radix_aux_logits, gvr_prior_indices=gvr_prior_indices, ) if positions is None: @@ -924,7 +927,7 @@ def forward(self, state: CSA2ForwardState, query_start: int, count: int) -> torc def _paged_mqa_logits( self, - metadata: CSA2TrtllmMetadata, + metadata: TrtllmAttentionMetadata, q_data: torch.Tensor, weights: torch.Tensor, q_scale: torch.Tensor, @@ -932,6 +935,7 @@ def _paged_mqa_logits( ) -> torch.Tensor: # Existing native paged specializations support 32/64 heads. Zero # weighted padding extends smaller CSA2 index geometries exactly. + csa2 = metadata.ext(CSA2Metadata) use_dsl = self.options.use_cute_dsl_paged_mqa_logits and IS_CUTLASS_DSL_AVAILABLE padded_heads = 64 if use_dsl or self.n_heads > 32 else 32 if padded_heads != self.n_heads: @@ -942,22 +946,22 @@ def _paged_mqa_logits( return torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits( q_data.view(torch.uint8), q_scale, - metadata.csa2_indexer_k_cache, + csa2.csa2_indexer_k_cache, weights, - metadata.csa2_indexer_context_lengths.flatten(), - metadata.csa2_indexer_block_table, - metadata.csa2_indexer_scheduler_metadata, - metadata.csa2_indexer_max_seq_len, + csa2.csa2_indexer_context_lengths.flatten(), + csa2.csa2_indexer_block_table, + csa2.csa2_indexer_scheduler_metadata, + csa2.csa2_indexer_max_seq_len, **(emission_kwargs or {}), ) return self._call_paged_mqa_logits( q_data, - metadata.csa2_indexer_k_cache, + csa2.csa2_indexer_k_cache, weights, - metadata.csa2_indexer_context_lengths, - metadata.csa2_indexer_block_table, - metadata.csa2_indexer_scheduler_metadata, - metadata.csa2_indexer_max_seq_len, + csa2.csa2_indexer_context_lengths, + csa2.csa2_indexer_block_table, + csa2.csa2_indexer_scheduler_metadata, + csa2.csa2_indexer_max_seq_len, q_scale, ) @@ -1081,7 +1085,7 @@ def _sparse_candidate_logits( def sparse_attn_indexer( self, - metadata: CSA2TrtllmMetadata, + metadata: TrtllmAttentionMetadata, hidden_states: torch.Tensor, q_fp8: torch.Tensor, k_fp8: torch.Tensor | None, @@ -1095,34 +1099,35 @@ def sparse_attn_indexer( Inputs span the complete model batch. K is read from the owner cache; the base signature is retained for the backend's indexer contract. """ + csa2 = metadata.ext(CSA2Metadata) if is_generation is not None: raise ValueError("CSA2 prediction requires the complete model query batch") manager = metadata.kv_cache_manager count = hidden_states.shape[0] layer = self.layout.layer(self.layer_idx) - visible = metadata.csa2_visible_lengths[self.layer_idx] - max_positions = metadata.csa2_global_max_positions[layer.kv_source] + visible = csa2.csa2_visible_lengths[self.layer_idx] + max_positions = csa2.csa2_global_max_positions[layer.kv_source] logical = torch.empty((count, self.index_topk), dtype=torch.int32, device=q_fp8.device) - metadata.csa2_indices[self.layer_idx] = logical + csa2.csa2_indices[self.layer_idx] = logical prior = None if q_fp8.is_cuda and self.top_k.needs_gvr_prior: - metadata.register_indexer_reset( + csa2.register_indexer_reset( self.layer_idx, lambda: self.top_k.reset_gvr_emission_rows(slice(None)) ) - prior = metadata.prepare_indexer_prior(self.layer_idx, self.index_topk) + prior = csa2.prepare_indexer_prior(self.layer_idx, self.index_topk) candidates = None sparse_blocks = sparse_counts = None auxiliary = sharded = () score_hook = None sparse_logits = self.use_sparse_candidates and q_fp8.is_cuda if layer.candidate_source is not None and layer.candidate_source != self.layer_idx: - candidates = metadata.csa2_candidates.get(layer.candidate_source) + candidates = csa2.csa2_candidates.get(layer.candidate_source) if candidates is None or candidates.ndim != 2 or candidates.shape[0] != count: raise ValueError( "CSA2 candidate source must publish matching full-query rows first" ) - sparse_blocks = metadata.csa2_candidate_blocks.get(layer.candidate_source) - sparse_counts = metadata.csa2_candidate_counts.get(layer.candidate_source) + sparse_blocks = csa2.csa2_candidate_blocks.get(layer.candidate_source) + sparse_counts = csa2.csa2_candidate_counts.get(layer.candidate_source) sparse_logits = ( sparse_logits and sparse_blocks is not None @@ -1142,33 +1147,33 @@ def sparse_attn_indexer( published = torch.full( (count, candidate_width), -1, dtype=torch.int32, device=q_fp8.device ) - metadata.csa2_candidates[self.layer_idx] = published + csa2.csa2_candidates[self.layer_idx] = published auxiliary = (published,) sparse_block = self.sparse_block if sparse_block: # Rows no chunk scores keep block 0 and zero valid columns; gathered # across TP ranks with the candidates. - metadata.csa2_candidate_blocks[self.layer_idx] = torch.zeros( + csa2.csa2_candidate_blocks[self.layer_idx] = torch.zeros( (count, candidate_width // sparse_block), dtype=torch.int32, device=q_fp8.device ) - metadata.csa2_candidate_counts[self.layer_idx] = torch.zeros( + csa2.csa2_candidate_counts[self.layer_idx] = torch.zeros( (count, 1), dtype=torch.int32, device=q_fp8.device ) sharded = ( - metadata.csa2_candidate_blocks[self.layer_idx], - metadata.csa2_candidate_counts[self.layer_idx], + csa2.csa2_candidate_blocks[self.layer_idx], + csa2.csa2_candidate_counts[self.layer_idx], ) def score_hook(scores: torch.Tensor, start: int, end: int) -> None: self._publish_candidates( scores, visible[start:end], - metadata.csa2_candidates, + csa2.csa2_candidates, start, count, candidate_width, - metadata.csa2_candidate_blocks if sparse_block else None, - metadata.csa2_candidate_counts if sparse_block else None, + csa2.csa2_candidate_blocks if sparse_block else None, + csa2.csa2_candidate_counts if sparse_block else None, ) # Host prefix bounds are valid for eager execution only. A graph must @@ -1184,10 +1189,10 @@ def score_hook(scores: torch.Tensor, start: int, end: int) -> None: max( (prefix + end - begin) // layer.compress_ratio for prefix, (begin, end) in zip( - metadata.csa2_request_start_positions, metadata.csa2_request_query_ranges + csa2.csa2_request_start_positions, csa2.csa2_request_query_ranges ) ) - if metadata.csa2_request_query_ranges + if csa2.csa2_request_query_ranges else 0, ) ) @@ -1201,7 +1206,7 @@ def score_hook(scores: torch.Tensor, start: int, end: int) -> None: if candidates is not None else torch.arange(short_bound, device=q_fp8.device).expand(count, -1) ) - slots = metadata.global_slot_tile(self.layer_idx, 0, count, positions) + slots = csa2.global_slot_tile(self.layer_idx, 0, count, positions) positions = torch.where(slots >= 0, positions, -1) self._select_all_visible( positions, @@ -1210,12 +1215,12 @@ def score_hook(scores: torch.Tensor, start: int, end: int) -> None: logical, ) if prior is not None: - metadata.publish_indexer_prior(self.layer_idx, logical) + csa2.publish_indexer_prior(self.layer_idx, logical) return logical def candidate_tile(start: int, end: int) -> _ChunkInputs: positions = candidates[start:end] - slots = metadata.global_slot_tile(self.layer_idx, start, end, positions) + slots = csa2.global_slot_tile(self.layer_idx, start, end, positions) keys, scales = manager.gather_indexer_keys(layer.kv_source, slots.flatten()) width = positions.shape[1] starts = torch.arange(end - start, dtype=torch.int32, device=q_fp8.device) * width @@ -1232,9 +1237,7 @@ def shared_keys(start: int, end: int, width: int) -> _ChunkInputs: # A request's compressed prefix is gathered once and shared by all # query tiles. Cached and newly published keys have the same owner. positions = torch.arange(width, device=q_fp8.device) - slots = metadata.global_slot_tile(self.layer_idx, start, start + 1, positions[None, :])[ - 0 - ] + slots = csa2.global_slot_tile(self.layer_idx, start, start + 1, positions[None, :])[0] keys, scales = manager.gather_indexer_keys(layer.kv_source, slots) starts = torch.zeros(end - start, dtype=torch.int32, device=q_fp8.device) logical_positions = torch.where(slots >= 0, positions, -1).expand(end - start, -1) @@ -1248,8 +1251,8 @@ def shared_keys(start: int, end: int, width: int) -> _ChunkInputs: shared_columns=True, ) - ranges = metadata.csa2_request_query_ranges - num_contexts = metadata.csa2_num_context_requests + ranges = csa2.csa2_request_query_ranges + num_contexts = csa2.csa2_num_context_requests decode_start = ranges[num_contexts][0] if num_contexts < len(ranges) else count # Heuristic and GVR TopK address columns as key positions, so only # exact TopK may select over candidate-ordered columns. @@ -1269,11 +1272,11 @@ def shared_keys(start: int, end: int, width: int) -> _ChunkInputs: and 1 <= self.n_heads <= 64 and self.head_dim == 128 ) - paged = metadata.prepare_indexer(self.layer_idx) if native_decode else None + paged = csa2.prepare_indexer(self.layer_idx) if native_decode else None sparse_descriptors = None if paged is not None and candidates is not None: # Sparse paged logits read the native index pages in place - sparse_descriptors = metadata.prepare_sparse_indexer(self.layer_idx, self.sparse_block) + sparse_descriptors = csa2.prepare_sparse_indexer(self.layer_idx, self.sparse_block) if sparse_descriptors is None: paged = None chunks = [] @@ -1284,7 +1287,7 @@ def shared_keys(start: int, end: int, width: int) -> _ChunkInputs: # Gather the request's prefix once and score only each # query's candidate blocks in it. width = min( - max_positions, metadata.csa2_request_start_positions[request] + end - begin + max_positions, csa2.csa2_request_start_positions[request] + end - begin ) def load_shared(begin=begin, end=end, width=width): @@ -1324,7 +1327,7 @@ def load_shared(begin=begin, end=end, width=width): continue if not is_prefill and paged is not None: continue - chunk_size = metadata.indexer_max_chunk_size + chunk_size = csa2.indexer_max_chunk_size for start in range(begin, end, chunk_size): stop = min(start + chunk_size, end) # Decode graphs must continue seeing newly published keys as @@ -1334,7 +1337,7 @@ def load_shared(begin=begin, end=end, width=width): if not is_prefill or metadata.is_cuda_graph else min( max_positions, - (metadata.csa2_request_start_positions[request] + stop - begin) + (csa2.csa2_request_start_positions[request] + stop - begin) // layer.compress_ratio, ) ) @@ -1364,7 +1367,7 @@ def load_shared(begin=begin, end=end, width=width): q_scale, logical, metadata.mapping, - metadata.indexer_q_split_threshold, + csa2.indexer_q_split_threshold, score_hook, auxiliary, radix_indices, @@ -1381,7 +1384,7 @@ def load_shared(begin=begin, end=end, width=width): def hook(scores: torch.Tensor) -> None: score_hook(scores, decode_start, count) - decode_visible = metadata.csa2_indexer_visible_lengths + decode_visible = csa2.csa2_indexer_visible_lengths if candidates is not None: # Score only the published candidate blocks straight from the index pages scores = self._sparse_candidate_logits( @@ -1400,13 +1403,13 @@ def hook(scores: torch.Tensor) -> None: logical[decode_start:], None, is_prefill=False, - radix_aux_indices=metadata.csa2_indexer_radix_aux_indices, - radix_aux_logits=metadata.csa2_indexer_radix_aux_logits, + radix_aux_indices=csa2.csa2_indexer_radix_aux_indices, + radix_aux_logits=csa2.csa2_indexer_radix_aux_logits, gvr_prior_indices=prior, valid_counts=sparse_counts[decode_start:], ) if prior is not None: - metadata.publish_indexer_prior(self.layer_idx, logical) + csa2.publish_indexer_prior(self.layer_idx, logical) return logical decode_scale = ( q_scale[decode_start:] @@ -1418,13 +1421,13 @@ def hook(scores: torch.Tensor) -> None: if ( self.options.use_gvr_emission and prior is not None - and metadata.csa2_indexer_max_seq_len % 8 == 0 + and csa2.csa2_indexer_max_seq_len % 8 == 0 ): emission_kwargs = self.top_k.prepare_gvr_emission( decode_count, - metadata.csa2_indexer_max_seq_len, + csa2.csa2_indexer_max_seq_len, torch.cuda.get_device_properties(q_fp8.device).multi_processor_count, - metadata.csa2_indexer_prior_capacity[self.layer_idx], + csa2.csa2_indexer_prior_capacity[self.layer_idx], ) logits = self._paged_mqa_logits( metadata, @@ -1442,7 +1445,7 @@ def hook(scores: torch.Tensor) -> None: prior, ) if prior is not None: - metadata.publish_indexer_prior(self.layer_idx, logical) + csa2.publish_indexer_prior(self.layer_idx, logical) return logical diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/kernel.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/kernel.py index 111314e2594c..f3d5a61c287a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/kernel.py @@ -388,7 +388,7 @@ def kernel( class OwnerSlotsKernel: """GLOBAL write slots and compressed groups of every KV owner for live endpoints. - Mirrors ``CSA2TrtllmMetadata._refresh_owner_slots`` in one launch: grid + Mirrors ``CSA2Metadata._refresh_owner_slots`` in one launch: grid ``y`` selects the owner, whose buffers and geometry come from an int64 ``[owners, OWNER_DESC_WIDTH]`` descriptor table (see ``owner_slot_descriptor``). Per-request group counts and their prefix @@ -1385,7 +1385,7 @@ def refresh_owner_slots( requests: int, max_capacity: int, ) -> None: - """Device implementation of ``CSA2TrtllmMetadata._refresh_owner_slots`` for + """Device implementation of ``CSA2Metadata._refresh_owner_slots`` for every owner of the int64 ``[owners, OWNER_DESC_WIDTH]`` ``descriptors``.""" ns = _glue_kernels() _check(requests <= ns.OwnerSlotsKernel.MAX_REQUESTS, "CSA2 owner refresh request capacity") @@ -1426,7 +1426,7 @@ def fill_indexer_descriptors( visible: torch.Tensor, valid: torch.Tensor, ) -> None: - """Device implementation of ``CSA2TrtllmMetadata._fill_indexer_descriptors``.""" + """Device implementation of ``CSA2Metadata._fill_indexer_descriptors``.""" count = int(blocks.shape[0]) if count == 0: return @@ -1885,7 +1885,7 @@ def map_logical_slot( max_positions: Int64, visible: Int64, ) -> Int64: - """``CSA2TrtllmMetadata._map_global_slots`` for one entry (visibility applied).""" + """``CSA2Metadata._map_global_slots`` for one entry (visibility applied).""" slot = Int64(-1) if (table_rows > 0) & (table_cols > 0) & (logical >= 0): page = logical // page_size @@ -4110,7 +4110,7 @@ def stage_selected_rows( swa_slots, main_pool, main_slots, - metadata, + tile, scratch, *, compact=True, @@ -4126,14 +4126,14 @@ def stage_selected_rows( Main selections are either physical ``main_slots`` or ``main_logical`` positions resolved in-kernel through ``main_mapping`` (``page_table``, ``page_size``, ``max_positions``, ``requests``, ``visible``), exactly as - ``CSA2TrtllmMetadata.global_slot_tile`` with visibility. ``scratch`` holds + ``CSA2Metadata.global_slot_tile`` with visibility. ``scratch`` holds the group's ``tags``/``counts`` (and, in the mapped form, ``main_slots`` receives the physical rows); ``compact=False`` reuses them from an earlier layer of the same group in this forward. """ - count = metadata.num_tokens - capacity = metadata.num_sparse_topk - fp8 = metadata.swa_pool.dtype == torch.float8_e4m3fn + count = tile.metadata.num_tokens + capacity = tile.metadata.num_sparse_topk + fp8 = tile.swa_pool.dtype == torch.float8_e4m3fn device = swa_pool.device has_main = main_slots is not None or main_logical is not None if not has_main: @@ -4151,12 +4151,12 @@ def stage_selected_rows( tags, counts = scratch["tags"], scratch["counts"] _check(capacity <= tags.shape[1] and capacity <= 4096, "staging capacity") ns = _cache_kernels() - placeholder = metadata.prepared_counter + placeholder = tile.prepared_counter dequant = kv_scale_quant_orig if fp8 else placeholder quant = kv_scale_orig_quant if fp8 else placeholder q_dequant = q_scale_quant_orig if fp8 else placeholder - bmm1 = metadata.mla_bmm1_scale if fp8 else placeholder - bmm2 = metadata.mla_bmm2_scale if fp8 else placeholder + bmm1 = tile.mla_bmm1_scale if fp8 else placeholder + bmm2 = tile.mla_bmm2_scale if fp8 else placeholder logical_int64 = logical.dtype == torch.int64 table_ptr, table_rows, table_cols, table_stride, page_size, max_positions = table_args if compact: @@ -4190,7 +4190,7 @@ def stage_selected_rows( int(capacity), _stream(device), ) - extra_rows = int(metadata.extra_pool.shape[1]) + extra_rows = int(tile.extra_pool.shape[1]) _launch( ("stage_dequant", fp8, device.index), lambda: ns.StageDequantKernel(fp8), @@ -4200,12 +4200,12 @@ def stage_selected_rows( _ptr(slots, align=8), _ptr(tags), _ptr(counts), - _ptr(metadata.swa_pool, align=8), - _ptr(metadata.extra_pool, align=8), + _ptr(tile.swa_pool, align=8), + _ptr(tile.extra_pool, align=8), _ptr(quant), - _ptr(metadata.prepared_lens), - _ptr(metadata.prepared_indices), - _ptr(metadata.prepared_counter), + _ptr(tile.prepared_lens), + _ptr(tile.prepared_indices), + _ptr(tile.prepared_counter), _ptr(q_dequant), _ptr(dequant), _ptr(bmm1), @@ -4235,7 +4235,7 @@ def stage_shared_rows( main_pool, main_slots, main_logical, - metadata, + tile, plan, *, main_mapping=None, @@ -4251,10 +4251,10 @@ def stage_shared_rows( their ``main_logical`` positions, or logical positions resolved through ``main_mapping``. Outputs preserve selection order and duplicates. """ - count = metadata.num_tokens - capacity = metadata.num_sparse_topk - bank = metadata.shared_pool - selected = metadata._csa2_shared_selected + count = tile.metadata.num_tokens + capacity = tile.metadata.num_sparse_topk + bank = tile.shared_pool + selected = tile._csa2_shared_selected device = swa_pool.device _check( bank.shape == (1 + plan.swa_rows + plan.main_rows, 512) and bank.is_contiguous(), @@ -4297,7 +4297,7 @@ def stage_shared_rows( selected.zero_() bank[0].zero_() ns = _cache_kernels() - placeholder = metadata.prepared_counter + placeholder = tile.prepared_counter q_dequant = q_scale_quant_orig if fp8 else placeholder _launch( ("shared_compact", fp8, has_main, premapped, device.index), @@ -4313,13 +4313,13 @@ def stage_shared_rows( _ptr(visible, align=8), table_ptr, _ptr(selected), - _ptr(metadata.prepared_indices), - _ptr(metadata.prepared_lens), - _ptr(metadata.prepared_counter), + _ptr(tile.prepared_indices), + _ptr(tile.prepared_lens), + _ptr(tile.prepared_counter), _ptr(q_dequant), _ptr(kv_scale_quant_orig if fp8 else placeholder), - _ptr(metadata.mla_bmm1_scale if fp8 else placeholder), - _ptr(metadata.mla_bmm2_scale if fp8 else placeholder), + _ptr(tile.mla_bmm1_scale if fp8 else placeholder), + _ptr(tile.mla_bmm2_scale if fp8 else placeholder), int(count), swa_width, main_width, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/metadata.py index 6e2742c1edad..8fc81fafbe27 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/metadata.py @@ -12,9 +12,11 @@ import numpy as np import torch +from tensorrt_llm._torch.attention.backends.sparse.params import SparseAlgorithmExtension from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams +from ...features.spec_decoding import SpecDecodingMetadata from .params import CSA2BackendForwardArgs, CSA2Layer if TYPE_CHECKING: @@ -66,7 +68,7 @@ class _CoalescedUploads: _ALIGN = 16 - def __init__(self, metadata: CSA2TrtllmMetadata): + def __init__(self, metadata: CSA2Metadata): self._metadata = metadata self._pending = [] self._bytes = 0 @@ -115,11 +117,12 @@ def flush(self) -> None: self._bytes = 0 -class CSA2TrtllmMetadata(TrtllmAttentionMetadata): +class CSA2Metadata(SparseAlgorithmExtension): """Manager-backed request metadata with bounded native query staging. Normal prepare resolves manager pages for the packed model forward. - for_query_tile creates a separate compute-only metadata object. + for_query_tile creates a separate compute-only metadata object that + carries its own CSA2 extension. Eager context retains real query groups with virtual staged KV lengths. Generation and captured context use independent one-query generation rows. Source causality, request isolation and windows are encoded in selections; @@ -134,8 +137,10 @@ class CSA2TrtllmMetadata(TrtllmAttentionMetadata): _csa2_defer_decode_outputs: bool = False _csa2_deferred_decode_outputs: bool = False _csa2_graph_reserved: bool = False + # The manager's SWA publication token while ``prepare()`` runs. + _swa_publication: object | None = None - # Request metadata is held directly on this object. Main/index physical + # Request metadata is held directly on this extension. Main/index physical # tables and write slots are owner-keyed; SWA and visibility are layer-keyed. csa2_indices: dict[int, torch.Tensor] csa2_candidates: dict[int, torch.Tensor] @@ -159,53 +164,37 @@ class CSA2TrtllmMetadata(TrtllmAttentionMetadata): decoder_context_ends: tuple[int, ...] | None = None csa2_decoder_capture_lens: tuple[int, ...] | None = None - def get_adp_token_counts(self) -> list[int]: + def adp_token_counts(self) -> list[int]: plan = self.decoder_replay_plan - input_tokens = self.padded_num_tokens or self.num_tokens decoder_tokens = ( - plan.num_replay_tokens if plan and plan.replays_local_tokens else input_tokens + plan.num_replay_tokens + if plan and plan.replays_local_tokens + else self.metadata.num_tokens ) - return [input_tokens, decoder_tokens] + return [decoder_tokens] - def set_adp_token_counts(self, counts: list[list[int]]) -> list[int]: + def set_adp_token_counts(self, counts: list[list[int]], input_counts: list[int]) -> None: plan = self.decoder_replay_plan if plan is not None: - plan.saved_all_rank_num_tokens, plan.replay_all_rank_num_tokens = counts - if not plan.updates_local_metadata and counts[0] == counts[1]: + plan.saved_all_rank_num_tokens, plan.replay_all_rank_num_tokens = ( + input_counts, + counts[0], + ) + if not plan.updates_local_metadata and input_counts == counts[0]: self.decoder_replay_plan = None - return counts[0] - def __post_init__(self) -> None: - # Graph clones rerun post-init: eligibility belongs to their own next + def allocate(self, buffers) -> None: + # Graph clones re-allocate: eligibility belongs to their own next # successful preparation, not the metadata object they were copied from. self._csa2_deferred_decode_outputs = False self.decoder_replay_plan = None self.decoder_context_ends = None - super().__post_init__() - - @property - def tokens_per_block(self) -> int: - return ( - self.kv_cache_manager.tokens_per_block - if self.kv_cache_manager is not None - else _SWA_TILE - ) - - @property - def host_kv_cache_pool_pointers(self) -> torch.Tensor: - return ( - self.pool_pointers - if hasattr(self, "pool_pointers") - else super().host_kv_cache_pool_pointers - ) - @property - def host_kv_cache_pool_mapping(self) -> torch.Tensor: - return ( - self.pool_mapping - if hasattr(self, "pool_mapping") - else super().host_kv_cache_pool_mapping - ) + def kv_cache_pools(self) -> tuple[torch.Tensor, torch.Tensor, int] | None: + # A compute tile attends over its bounded staging pools, not the manager's. + if not hasattr(self, "pool_pointers"): + return None + return self.pool_pointers, self.pool_mapping, _SWA_TILE @property def stage_kv_scales(self) -> tuple[torch.Tensor, torch.Tensor]: @@ -226,6 +215,7 @@ def _allocate( staging_dtype: torch.dtype = torch.bfloat16, shared_rows: int | None = None, ) -> None: + md = self.metadata if staging_dtype not in (torch.bfloat16, torch.float8_e4m3fn): raise ValueError("CSA2 native staging requires BF16 or E4M3") self.staging_dtype = staging_dtype @@ -263,51 +253,51 @@ def _allocate( [[self.swa_pool.data_ptr(), 0]], dtype=torch.int64, device="cpu" ) self.pool_mapping = torch.zeros((1, 2), dtype=torch.int32, device="cpu") - self.num_sparse_topk = _SWA_TILE + extra_capacity - self.max_seq_len = self.num_sparse_topk - self.kv_cache_params = KVCacheParams(use_cache=True) - self.kv_cache_block_offsets = torch.zeros( - (1, capacity, 2, (self.num_sparse_topk + _SWA_TILE - 1) // _SWA_TILE), + md.num_sparse_topk = _SWA_TILE + extra_capacity + md.max_seq_len = md.num_sparse_topk + md.kv_cache_params = KVCacheParams(use_cache=True) + md.kv_cache_block_offsets = torch.zeros( + (1, capacity, 2, (md.num_sparse_topk + _SWA_TILE - 1) // _SWA_TILE), dtype=torch.int32, device=device, ) self.prepared_indices = torch.full( - (capacity, self.num_sparse_topk), -1, dtype=torch.int32, device=device + (capacity, md.num_sparse_topk), -1, dtype=torch.int32, device=device ) self.prepared_lens = torch.empty(capacity, dtype=torch.int32, device=device) self.prepared_counter = torch.zeros(1, dtype=torch.uint32, device=device) self.prepared_cu_q = torch.arange(capacity + 1, dtype=torch.int32, device=device) * heads self.prepared_cu_kv = ( - torch.arange(capacity + 1, dtype=torch.int32, device=device) * self.num_sparse_topk + torch.arange(capacity + 1, dtype=torch.int32, device=device) * md.num_sparse_topk ) self.query_lens_host = torch.ones(capacity, dtype=torch.int32, device="cpu") self.query_lens_device = torch.ones(capacity, dtype=torch.int32, device=device) - self.kv_lens.fill_(self.num_sparse_topk) - self.kv_lens_cuda.fill_(self.num_sparse_topk) - self.prompt_lens_cpu.fill_(1) - self.prompt_lens_cuda.fill_(1) - self.host_request_types.fill_(1) - self.host_total_kv_lens.zero_() + md.kv_lens.fill_(md.num_sparse_topk) + md.kv_lens_cuda.fill_(md.num_sparse_topk) + md.prompt_lens_cpu.fill_(1) + md.prompt_lens_cuda.fill_(1) + md.host_request_types.fill_(1) + md.host_total_kv_lens.zero_() # Each query count has independent native workspace storage. # Eager initialization and capture must use that same tensor. - self.cuda_graph_workspace = self.workspace + md.cuda_graph_workspace = md.workspace - @classmethod + @staticmethod def for_query_tile( - cls, q: torch.Tensor, extra_width: int, *, staging_dtype=torch.bfloat16, context_lengths: list[int] | None = None, shared_rows: int | None = None, - ) -> CSA2TrtllmMetadata: + ) -> TrtllmAttentionMetadata: """Prepare fixed geometry before TrtllmAttention.forward allocates output. q is BF16 [queries, heads, 512]. Extra capacity is a selected-main-row bound, not the persistent cache size. Callers retain one metadata object per geometry and warm it with the normal forward before graph capture. As with TrtllmAttentionMetadata, set is_cuda_graph for captured calls. + The returned metadata carries the tile's staging in its CSA2Metadata. """ count, heads, dim = q.shape if context_lengths and ( @@ -323,16 +313,18 @@ def for_query_tile( with torch.cuda.device(q.device): if torch.cuda.is_current_stream_capturing(): raise RuntimeError("Prepare and warm CSA2 metadata before CUDA Graph capture") - metadata = cls( + metadata = TrtllmAttentionMetadata( max_num_requests=len(context_lengths) if context_lengths else count, max_num_tokens=count, + extension_types=(CSA2Metadata,), ) - metadata._allocate( + tile = metadata.ext(CSA2Metadata) + tile._allocate( count, heads, (extra_width + 127) // 128 * 128, q.device, staging_dtype, shared_rows ) - metadata.num_query_heads = heads - metadata._seq_lens = metadata.query_lens_host - metadata._seq_lens_cuda = metadata.query_lens_device + tile.num_query_heads = heads + metadata._seq_lens = tile.query_lens_host + metadata._seq_lens_cuda = tile.query_lens_device metadata._num_contexts = metadata._num_ctx_tokens = 0 metadata._num_generations = metadata._num_tokens = count metadata._bind_runtime_views( @@ -343,23 +335,24 @@ def for_query_tile( host_request_types=metadata.host_request_types, ) metadata.host_total_kv_lens[1] = count * metadata.num_sparse_topk - metadata.cu_q_seqlens = metadata.prepared_cu_q - metadata.cu_kv_seqlens = metadata.prepared_cu_kv + metadata.cu_q_seqlens = tile.prepared_cu_q + metadata.cu_kv_seqlens = tile.prepared_cu_kv if context_lengths: - metadata._bind_context_tile(context_lengths) + tile._bind_context_tile(context_lengths) return metadata def _shared_context_plan(self, q, extra_width, layer_idx): # Graph frames are filtered by the caller. Without logical provenance, # keep the existing selected-row path. + md = self.metadata if ( layer_idx is None or _SWA_TILE + (extra_width + 127) // 128 * 128 > 4096 or not hasattr(self, "_csa2_main_domain_counts") - or q.shape[0] != self.num_ctx_tokens + or q.shape[0] != md.num_ctx_tokens ): return None - manager = self.kv_cache_manager + manager = md.kv_cache_manager requests = self.csa2_num_context_requests lengths = self.csa2_request_lengths[:requests] if not lengths or sum(lengths) != q.shape[0]: @@ -437,12 +430,13 @@ def get_query_tile_metadata( *, staging_dtype=torch.bfloat16, layer_idx: int | None = None, - ) -> CSA2TrtllmMetadata: + ) -> TrtllmAttentionMetadata: """Share bounded staging, retaining real request groups for context.""" + md = self.metadata context_lengths = [] - if query_start is not None and query_start < self.num_ctx_tokens: + if query_start is not None and query_start < md.num_ctx_tokens: end = query_start + q.shape[0] - if end > self.num_ctx_tokens: + if end > md.num_ctx_tokens: raise ValueError("CSA2 query tiles must not cross the context/generation boundary") for begin, stop in self.csa2_request_query_ranges: lo, hi = max(begin, query_start), min(stop, end) @@ -455,7 +449,7 @@ def get_query_tile_metadata( # path; choose it during graph warmup too, so its workspace is ready. with torch.cuda.device(q.device): capturing = torch.cuda.is_current_stream_capturing() - if context_lengths and (self.is_cuda_graph or capturing): + if context_lengths and (md.is_cuda_graph or capturing): context_lengths = [] shared_plan = ( self._shared_context_plan(q, extra_width, layer_idx) @@ -486,7 +480,7 @@ def get_query_tile_metadata( for previous, tile in tuple(self._csa2_query_tiles.items()): if ( previous != key - and not tile._csa2_graph_reserved + and not tile.ext(CSA2Metadata)._csa2_graph_reserved and previous[:2] == key[:2] and previous[3:6] == key[3:6] and previous[8] == key[8] @@ -509,22 +503,23 @@ def get_query_tile_metadata( if workspace is not None: metadata.workspace = metadata.cuda_graph_workspace = workspace self._csa2_query_tiles[key] = metadata - metadata.shared_plan = shared_plan - metadata._csa2_graph_reserved |= self.is_cuda_graph or capturing + tile = metadata.ext(CSA2Metadata) + tile.shared_plan = shared_plan + tile._csa2_graph_reserved |= md.is_cuda_graph or capturing metadata.is_cuda_graph = capturing if capturing: self._csa2_replay_capture_signature = getattr(self, "csa2_replay_signature", None) if context_lengths: - metadata._bind_context_tile(context_lengths) + tile._bind_context_tile(context_lengths) return metadata - def _bind_runtime_views(self, **kwargs) -> None: + def on_runtime_views_bound(self) -> None: # Generic prepare/rebinding can replace the geometry-only child views. self._csa2_context_geometry = None - super()._bind_runtime_views(**kwargs) def _bind_context_tile(self, query_lengths: list[int]) -> None: - geometry = (self.num_sparse_topk, tuple(query_lengths)) + md = self.metadata + geometry = (md.num_sparse_topk, tuple(query_lengths)) if getattr(self, "_csa2_context_geometry", None) == geometry: return # A failed partial update must not retain the previous successful key. @@ -533,27 +528,27 @@ def _bind_context_tile(self, query_lengths: list[int]) -> None: # context kernel still applies a causal upper bound in staged-column # coordinates: give even the first query the full sparse capacity, # otherwise short source prefixes would clip valid compressed keys. - kv_lengths = [self.num_sparse_topk + length - 1 for length in query_lengths] + kv_lengths = [md.num_sparse_topk + length - 1 for length in query_lengths] requests = len(query_lengths) query = torch.tensor(query_lengths, dtype=torch.int32, device="cpu") kv = torch.tensor(kv_lengths, dtype=torch.int32, device="cpu") self.query_lens_host[:requests].copy_(query) self._copy_host("context_query_lens", self.query_lens_device[:requests], query) - self._seq_lens = self.query_lens_host[:requests] - self._seq_lens_cuda = self.query_lens_device[:requests] - self._num_contexts = requests - self._num_ctx_tokens = self._num_tokens = sum(query_lengths) - self._num_generations = 0 - self.kv_lens[:requests].copy_(kv) - self._copy_host("context_kv_lens", self.kv_lens_cuda[:requests], kv) + md._seq_lens = self.query_lens_host[:requests] + md._seq_lens_cuda = self.query_lens_device[:requests] + md._num_contexts = requests + md._num_ctx_tokens = md._num_tokens = sum(query_lengths) + md._num_generations = 0 + md.kv_lens[:requests].copy_(kv) + self._copy_host("context_kv_lens", md.kv_lens_cuda[:requests], kv) # THOP interprets context_lengths as Q lengths, independently of # past/current KV lengths. Keep them consistent with unfolded cuQ. - self.prompt_lens_cpu[:requests].copy_(query) - self._copy_host("context_prompt_lens", self.prompt_lens_cuda[:requests], query) - self.host_request_types[:requests].zero_() - self.host_total_kv_lens.zero_() - self.host_total_kv_lens[0] = sum(kv_lengths) - self.max_seq_len = max(self.num_sparse_topk, max(kv_lengths)) + md.prompt_lens_cpu[:requests].copy_(query) + self._copy_host("context_prompt_lens", md.prompt_lens_cuda[:requests], query) + md.host_request_types[:requests].zero_() + md.host_total_kv_lens.zero_() + md.host_total_kv_lens[0] = sum(kv_lengths) + md.max_seq_len = max(md.num_sparse_topk, max(kv_lengths)) self._copy_host( "context_cu_q", self.prepared_cu_q[: requests + 1], @@ -564,14 +559,14 @@ def _bind_context_tile(self, query_lengths: list[int]) -> None: self.prepared_cu_kv[: requests + 1], torch.cat((kv.new_zeros(1), kv.cumsum(0).int())), ) - self.cu_q_seqlens = self.prepared_cu_q[: requests + 1] - self.cu_kv_seqlens = self.prepared_cu_kv[: requests + 1] - self._bind_runtime_views( - kv_lens_cuda=self.kv_lens_cuda[:requests], - kv_lens=self.kv_lens[:requests], - prompt_lens_cuda=self.prompt_lens_cuda[:requests], - prompt_lens_cpu=self.prompt_lens_cpu[:requests], - host_request_types=self.host_request_types[:requests], + md.cu_q_seqlens = self.prepared_cu_q[: requests + 1] + md.cu_kv_seqlens = self.prepared_cu_kv[: requests + 1] + md._bind_runtime_views( + kv_lens_cuda=md.kv_lens_cuda[:requests], + kv_lens=md.kv_lens[:requests], + prompt_lens_cuda=md.prompt_lens_cuda[:requests], + prompt_lens_cpu=md.prompt_lens_cpu[:requests], + host_request_types=md.host_request_types[:requests], ) self._csa2_context_geometry = geometry @@ -585,10 +580,15 @@ def _stage_scratch(self, source, group, count: int): compact. A capture records its own compaction so replays follow the live selections. """ + md, source_md = self.metadata, source.metadata groups = source.__dict__.setdefault("_csa2_stage_groups", {}) - key = (group, self.prepared_lens.shape[0], self.num_sparse_topk) + key = (group, self.prepared_lens.shape[0], md.num_sparse_topk) capturing = torch.cuda.is_current_stream_capturing() - frame = (id(source), id(source.kv_cache_manager), getattr(source, "_csa2_stage_epoch", 0)) + frame = ( + id(source), + id(source_md.kv_cache_manager), + getattr(source, "_csa2_stage_epoch", 0), + ) entry = groups.get(key) if entry is None: if capturing: @@ -602,18 +602,18 @@ def _stage_scratch(self, source, group, count: int): capacity, device = self.prepared_lens.shape[0], self.prepared_lens.device entry = groups[key] = { "tags": torch.empty( - (capacity, self.num_sparse_topk), dtype=torch.int32, device=device + (capacity, md.num_sparse_topk), dtype=torch.int32, device=device ), "counts": torch.empty(capacity, dtype=torch.int32, device=device), "main_slots": torch.empty( - (capacity, max(self.num_sparse_topk - _SWA_TILE, 1)), + (capacity, max(md.num_sparse_topk - _SWA_TILE, 1)), dtype=torch.int64, device=device, ), "serial": None, "graph_reserved": False, } - entry["graph_reserved"] |= self.is_cuda_graph or source.is_cuda_graph or capturing + entry["graph_reserved"] |= md.is_cuda_graph or source_md.is_cuda_graph or capturing serial = (*frame, count, capturing) compact = group is None or entry["serial"] != serial entry["serial"] = serial @@ -719,7 +719,8 @@ def stage_selected( """ from .quantization import _fused_gather_supported, gather_rows, row_bytes - count = self.num_tokens + md = self.metadata + count = md.num_tokens fp8 = self.swa_pool.dtype == torch.float8_e4m3fn if fp8: kv_scale_orig_quant = ( @@ -759,7 +760,7 @@ def stage_selected( if ( inputs.topk_indices.ndim != 2 or inputs.topk_indices.shape[0] != count - or inputs.topk_indices.shape[1] > self.num_sparse_topk - _SWA_TILE + or inputs.topk_indices.shape[1] > md.num_sparse_topk - _SWA_TILE ): raise ValueError("CSA2 selected main indices exceed metadata geometry") # Preserve gather_rows validation even for zero-width selections; an @@ -811,11 +812,11 @@ def stage_selected( and all(pool.is_cuda and pool.stride(1) == 1 for pool, _, _ in sources) and _fused_gather_supported(inputs.swa_pool.device.index) ) - if native and self.num_sparse_topk <= 4096: + if native and md.num_sparse_topk <= 4096: from .kernel import stage_selected_rows - source = inputs.state.metadata if inputs.state is not None else None - scratch, compact = self._stage_scratch(source or self, inputs.stage_group, count) + source = inputs.state.metadata.ext(CSA2Metadata) if inputs.state is not None else self + scratch, compact = self._stage_scratch(source, inputs.stage_group, count) main_slots = inputs.topk_indices if mapping is not None and main_slots is not None: main_slots = scratch["main_slots"][:count, : main_slots.shape[1]] @@ -891,7 +892,7 @@ def stage_selected( indices[:, :_SWA_TILE].copy_( torch.where(swa_positions < lengths[:, None], offsets * _SWA_TILE + swa_positions, -1) ) - extra_capacity = self.num_sparse_topk - _SWA_TILE + extra_capacity = md.num_sparse_topk - _SWA_TILE if extra_capacity: extra_positions = torch.arange(extra_capacity, device=swa.device)[None, :] indices[:, _SWA_TILE:].copy_( @@ -935,25 +936,31 @@ def set_swa_bounded_replay( ) self._csa2_pending_replay = (tuple(cached_prefix_lengths), decoder) + def bind_context_requests( + self, requests: list[LlmRequest], *, prefill_cuda_graphs: bool + ) -> None: + self.prepare_context_replay(requests, allow_final_window=not prefill_cuda_graphs) + def prepare_context_replay( self, requests: list[LlmRequest], *, allow_final_window: bool = True ) -> None: """Resolve recovery sources and read floors before preparing model inputs.""" from .....pyexecutor.ced_replay import EncoderCheckpoint, EncoderReplay + md = self.metadata # Graph selection may vary between chunks. Keep every chunk's Decoder # history when a later chunk can enter a graph that expects full inputs. self.decoder_context_ends = ( tuple(request.prompt_len for request in requests) if allow_final_window else () ) - prefixes = [None] * self.num_seqs - floors = [0] * self.num_seqs - starts = self.kv_cache_params.num_cached_tokens_per_seq - lengths = self.seq_lens.tolist() + prefixes = [None] * md.num_seqs + floors = [0] * md.num_seqs + starts = md.kv_cache_params.num_cached_tokens_per_seq + lengths = md.seq_lens.tolist() for row, request in enumerate(requests): plan = request.py_ced_replay if isinstance(plan, (EncoderReplay, EncoderCheckpoint)): - cache = self.kv_cache_manager.kv_cache_map.get(request.py_request_id) + cache = md.kv_cache_manager.kv_cache_map.get(request.py_request_id) if cache is None or plan.cache_identity != id(cache) or not cache.is_active: raise ValueError("Encoder recovery has lost its active Global claim") if isinstance(plan, EncoderReplay): @@ -978,7 +985,8 @@ def set_decoder_query_boundary(self, first_layer: int) -> None: are not dependencies of this query-only pass and must not be reserved again merely to build metadata for the longer decoder interval. """ - layout = self.kv_cache_manager.layout + md = self.metadata + layout = md.kv_cache_manager.layout if ( first_layer not in self.csa2_precomputed_kv_layers or first_layer != max(layout.kv_source_layer_ids) @@ -992,7 +1000,7 @@ def set_decoder_query_boundary(self, first_layer: int) -> None: ends = [ start + length for start, length in zip( - self.kv_cache_params.num_cached_tokens_per_seq, self.seq_lens.tolist() + md.kv_cache_params.num_cached_tokens_per_seq, md.seq_lens.tolist() ) ] self.set_source_batch([0] * len(ends), ends) @@ -1020,9 +1028,10 @@ def select_global_source( return hidden_states.index_select(0, indices) def _swa_replay_geometry(self, prefixes, decoder, starts, lengths) -> tuple: + md = self.metadata owner_shapes = [] - for owner in self.kv_cache_manager.layout.kv_source_layer_ids: - ratio = self.kv_cache_manager.layout.compress_ratios[owner] + for owner in md.kv_cache_manager.layout.kv_source_layer_ids: + ratio = md.kv_cache_manager.layout.compress_ratios[owner] source_lengths = tuple( length if prefix is None @@ -1035,7 +1044,7 @@ def _swa_replay_geometry(self, prefixes, decoder, starts, lengths) -> tuple: return ("decoder" if decoder else "encoder", tuple(lengths), tuple(owner_shapes)) @contextmanager - def defer_cuda_graph_decode_outputs(self): + def defer_decode_outputs(self): """Engine-only scope for a selected ordinary decode graph's preparation. The engine must run ``on_update_kv_lens`` before any model consumer, @@ -1049,12 +1058,14 @@ def defer_cuda_graph_decode_outputs(self): finally: self._csa2_defer_decode_outputs = previous - def prepare(self) -> None: + @contextmanager + def prepare_scope(self): from .cache_manager import CSA2CacheManager + md = self.metadata self._csa2_deferred_decode_outputs = False self._csa2_context_geometry = None - if self.is_cuda_graph and hasattr(self, "_csa2_replay_capture_signature"): + if md.is_cuda_graph and hasattr(self, "_csa2_replay_capture_signature"): pending = getattr(self, "_csa2_pending_replay", None) signature = ( None @@ -1062,34 +1073,41 @@ def prepare(self) -> None: else self._swa_replay_geometry( pending[0], pending[1], - self.kv_cache_params.num_cached_tokens_per_seq, - self.seq_lens.tolist(), + md.kv_cache_params.num_cached_tokens_per_seq, + md.seq_lens.tolist(), ) ) if signature != self._csa2_replay_capture_signature: raise ValueError( "CSA2 replay mode/source geometry changed; use fresh graph metadata and recapture" ) - with torch.cuda.device(self.kv_lens_cuda.device): + with torch.cuda.device(md.kv_lens_cuda.device): if torch.cuda.is_current_stream_capturing(): raise RuntimeError("Prepare CSA2 request metadata before CUDA Graph capture") self._csa2_ready_for_kv_update = False - manager = self.kv_cache_manager + manager = md.kv_cache_manager scope = ( manager._swa_publication_scope() if isinstance(manager, CSA2CacheManager) else nullcontext(None) ) with scope as publication: - super().prepare() - self.prepare_csa2(swa_publication=publication) + self._swa_publication = publication + try: + yield + finally: + self._swa_publication = None + + def prepare(self) -> None: + self.prepare_csa2(swa_publication=self._swa_publication) def _get_csa2_buffer(self, key, shape, dtype, *, like=None) -> torch.Tensor: # Preserve manager-local, shape-keyed addresses across graph replays. + md = self.metadata if not hasattr(self, "_csa2_buffers"): self._csa2_buffers = {} - cache_key = (key, tuple(shape), dtype, self.is_cuda_graph) - if not self.is_cuda_graph: + cache_key = (key, tuple(shape), dtype, md.is_cuda_graph) + if not md.is_cuda_graph: # Eager buffers retain only the latest shape per name. Track that # shape per name so a prepare with many keys does not rescan the # whole bank for every lookup. @@ -1105,9 +1123,9 @@ def _get_csa2_buffer(self, key, shape, dtype, *, like=None) -> torch.Tensor: result = self._csa2_buffers.get(cache_key) if result is None: result = ( - torch.empty_like(like, device=self.kv_lens_cuda.device) + torch.empty_like(like, device=md.kv_lens_cuda.device) if like is not None - else torch.empty(shape, dtype=dtype, device=self.kv_lens_cuda.device) + else torch.empty(shape, dtype=dtype, device=md.kv_lens_cuda.device) ) self._csa2_buffers[cache_key] = result return result @@ -1161,11 +1179,12 @@ def _swa_write_pages(starts, lengths, floors, allocated, block, table_shape): "CSA2 query/replay interval requires writable SWA pages; " "reserve the complete replay range before prepare" ) - return CSA2TrtllmMetadata._flat_page_ranges(begin, np.where(active, end - begin, 0)) + return CSA2Metadata._flat_page_ranges(begin, np.where(active, end - begin, 0)) def _check_swa_write_pages(self, starts, lengths, floors, allocated_until, page_count): """Raise unless every SWA page the query intervals write is mapped.""" - manager = self.kv_cache_manager + md = self.metadata + manager = md.kv_cache_manager write_rows, write_columns = self._swa_write_pages( starts, lengths, @@ -1185,7 +1204,8 @@ def _check_owner_pages(self, owner, owner_starts, owner_ends, allocated_until, p inside the allocated capacity are mapped.""" from .cache_manager import CSA2CacheRole - manager = self.kv_cache_manager + md = self.metadata + manager = md.kv_cache_manager block = manager.tokens_per_block ratio = manager.layout.compress_ratios[owner] page_size = block // ratio @@ -1259,6 +1279,7 @@ def _swa_table_descriptors(self, layers: int) -> torch.Tensor: def _device_descriptors(self, kind: str, rows) -> torch.Tensor: """Device int64 table of the host descriptor ``rows``, uploaded only when they change.""" + md = self.metadata key = (kind, tuple(rows)) cache = getattr(self, "_csa2_swa_table_cache", None) if cache is None or cache[0] is not self._csa2_buffers: @@ -1266,7 +1287,7 @@ def _device_descriptors(self, kind: str, rows) -> torch.Tensor: # descriptor index with it so every copy sees the same fills. cache = self._csa2_swa_table_cache = (self._csa2_buffers, {}, {}) _, names, filled = cache - if self.is_cuda_graph: + if md.is_cuda_graph: # A captured launch reads its descriptor table at replay time, so # the content must never change after the fill: every distinct # binding (one per captured batch shape) owns its own buffer. @@ -1296,11 +1317,12 @@ def _ensure_swa_slots(self) -> None: re-arm it, so a graph capture always records the launch. """ # None: not prepared by prepare_csa2. + md = self.metadata if getattr(self, "_csa2_swa_resolved", None) is not False: return from . import kernel - manager = self.kv_cache_manager + manager = md.kv_cache_manager positions, requests = self.csa2_positions, self.csa2_token_requests layers, count, window = ( # SWA layers are a prefix; a decoder-only pass binds its suffix. @@ -1410,13 +1432,14 @@ def _refresh_swa_tensor_outputs( def _bind_swa_pages(self, layer, pages, published_pages): # Select backing once for each existing shape/graph key. A direct or # draft fallback must update it, never redirect a captured consumer. + md = self.metadata if published_pages is not None and ( published_pages.dtype != pages.dtype - or published_pages.device != self.kv_lens_cuda.device + or published_pages.device != md.kv_lens_cuda.device or published_pages.shape != pages.shape ): published_pages = None - key = (f"swa_pages/{layer}", tuple(pages.shape), pages.dtype, self.is_cuda_graph) + key = (f"swa_pages/{layer}", tuple(pages.shape), pages.dtype, md.is_cuda_graph) if not hasattr(self, "_csa2_buffers"): self._csa2_buffers = {} if published_pages is not None and key not in self._csa2_buffers: @@ -1443,27 +1466,30 @@ def prepare_csa2(self, *, swa_publication=None) -> None: from .cache_manager import CSA2CacheManager, CSA2CacheRole from .compressor import CSA2CompressionBatch + md = self.metadata self._csa2_deferred_decode_outputs = False self._csa2_ready_for_kv_update = False # Every prepared forward starts a new staging epoch: the first layer of # each staging group compacts again. self._csa2_stage_epoch = getattr(self, "_csa2_stage_epoch", 0) + 1 - manager = self.kv_cache_manager + manager = md.kv_cache_manager if not isinstance(manager, CSA2CacheManager): raise TypeError("CSA2 runtime metadata requires CSA2CacheManager") self.csa2_query_layer_start = getattr(self, "_csa2_pending_query_layer_start", 0) if hasattr(self, "_csa2_pending_query_layer_start"): del self._csa2_pending_query_layer_start - if self.beam_width != 1 or self.is_spec_dec_dynamic_tree: + if md.beam_width != 1 or md.ext(SpecDecodingMetadata).is_spec_dec_dynamic_tree: raise NotImplementedError( "CSA2 supports contiguous-prefix verification, not beams or dynamic trees" ) - request_ids = list(self.request_ids) - lengths = self.seq_lens.tolist() - if manager.context_swa_layer_limit is not None and self.num_contexts != len(lengths): + request_ids = list(md.request_ids) + lengths = md.seq_lens.tolist() + if manager.context_swa_layer_limit is not None and md.num_contexts != len(lengths): raise ValueError("CSA2 context-only cache cannot prepare generation requests") - manager.validate_verification(lengths, self.num_contexts, self.is_spec_decoding_enabled) - starts = list(self.kv_cache_params.num_cached_tokens_per_seq) + manager.validate_verification( + lengths, md.num_contexts, md.ext(SpecDecodingMetadata).is_spec_decoding_enabled + ) + starts = list(md.kv_cache_params.num_cached_tokens_per_seq) if len(request_ids) != len(lengths) or len(starts) != len(lengths): raise ValueError("CSA2 requires one query length and cached length per request") replay = getattr(self, "_csa2_pending_replay", None) @@ -1489,7 +1515,7 @@ def prepare_csa2(self, *, swa_publication=None) -> None: for row, (prefix, start, length) in enumerate(zip(prefixes, starts, lengths)): if prefix is None: continue - if row >= self.num_contexts: + if row >= md.num_contexts: raise ValueError("CSA2 bounded replay applies only to context requests") if start != replay_starts[row]: raise ValueError( @@ -1521,14 +1547,14 @@ def prepare_csa2(self, *, swa_publication=None) -> None: # speculative/draft metadata keep the complete preparation contract. defer_outputs = ( getattr(self, "_csa2_defer_decode_outputs", False) - and self.is_cuda_graph - and self.num_contexts == 0 + and md.is_cuda_graph + and md.num_contexts == 0 and bool(lengths) and all(length == 1 for length in lengths) - and not self.is_spec_decoding_enabled + and not md.ext(SpecDecodingMetadata).is_spec_decoding_enabled and replay is None and not external_source - and getattr(self, "csa2_remote_tail_mode", None) is None + and getattr(md, "csa2_remote_tail_mode", None) is None ) source_lengths, source_starts = getattr(self, "_csa2_source_batch", (lengths, starts)) if external_source: @@ -1539,9 +1565,9 @@ def prepare_csa2(self, *, swa_publication=None) -> None: # step's full verification width; positions beyond the allocated # capacity stay unmapped and on_update_kv_lens narrows them later. allocated_until = np.full(len(lengths), manager.max_seq_len, dtype=np.int64) - allocated_until[self.num_contexts :] = [ + allocated_until[md.num_contexts :] = [ manager.kv_cache_map[request_id].capacity - for request_id in request_ids[self.num_contexts :] + for request_id in request_ids[md.num_contexts :] ] lengths_np = np.asarray(lengths, dtype=np.int64) starts_np = np.asarray(starts, dtype=np.int64) @@ -1583,7 +1609,7 @@ def prepare_csa2(self, *, swa_publication=None) -> None: raise ValueError("CSA2 query positions exceed the manager context capacity") self.csa2_request_start_positions = tuple(starts) self.csa2_request_lengths = tuple(lengths) - self.csa2_num_context_requests = self.num_contexts + self.csa2_num_context_requests = md.num_contexts query_ends = np.cumsum(lengths_np) query_begins = query_ends - lengths_np packed_start = int(query_ends[-1]) if query_ends.size else 0 @@ -1622,10 +1648,10 @@ def prepare_csa2(self, *, swa_publication=None) -> None: self._csa2_compression = {} self._csa2_compressed_positions = {} published_swa = manager._get_swa_publication( - swa_publication, self.kv_cache_block_offsets, request_ids, self.num_contexts, page_count + swa_publication, md.kv_cache_block_offsets, request_ids, md.num_contexts, page_count ) if published_swa is None: - manager.compute_batch_page_tables(request_ids, self.num_contexts) + manager.compute_batch_page_tables(request_ids, md.num_contexts) # Owner slots are resolved from these host KV ends: a draft prepare # runs without the generic prepare, so kv_lens_cuda may hold live # lengths that differ from the host positions. @@ -1638,7 +1664,7 @@ def prepare_csa2(self, *, swa_publication=None) -> None: ) shared_cu = copy("source_cu", np.cumsum([0, *source_lengths], dtype=np.int32)) owners = manager.layout.kv_source_layer_ids - if self.num_contexts: + if md.num_contexts: # Reserve every currently addressable logical row, including # upward live-end corrections within already published pages. # This is capacity only; selected indices retain visibility. @@ -1687,7 +1713,7 @@ def prepare_csa2(self, *, swa_publication=None) -> None: ) capacity = sum(owner_lengths) global_spec = (owner, CSA2CacheRole.GLOBAL) - if self.num_contexts: + if md.num_contexts: self._csa2_main_domain_counts[owner] = tuple( min(page * (block // ratio), manager.max_seq_len // ratio) for page in last_pages[index] @@ -1752,11 +1778,11 @@ def prepare_csa2(self, *, swa_publication=None) -> None: self.csa2_token_requests = copy("token_requests", token_requests_np) # Generation-row -> generation-request map for the sparse paged indexer kernels decode_start = ( - query_ranges[self.num_contexts][0] - if self.num_contexts < len(query_ranges) + query_ranges[md.num_contexts][0] + if md.num_contexts < len(query_ranges) else packed_start ) - decode_rows_np = (token_requests_np[decode_start:] - self.num_contexts).astype(np.int32) + decode_rows_np = (token_requests_np[decode_start:] - md.num_contexts).astype(np.int32) self.csa2_decode_row_requests = ( copy("decode_row_requests", decode_rows_np) if decode_rows_np.size else None ) @@ -1812,11 +1838,11 @@ def prepare_csa2(self, *, swa_publication=None) -> None: # Only a plain generation prepare seeds the steady path: context, # replay, source-override and decoder-boundary geometry must be rebuilt. steady_seed = ( - not self.num_contexts + not md.num_contexts and replay is None and not external_source and not self.csa2_query_layer_start - and getattr(self, "csa2_remote_tail_mode", None) is None + and getattr(md, "csa2_remote_tail_mode", None) is None ) self._finish_prepare_csa2(snapshot if steady_seed else None, defer_outputs=defer_outputs) @@ -1826,8 +1852,9 @@ def _steady_generation_step(self, snapshot) -> bool: The same generation requests, in the same order and with the same query lengths, each start where the previous query ended. """ + md = self.metadata previous = getattr(self, "_csa2_steady_snapshot", None) - if previous is None or self.num_contexts: + if previous is None or md.num_contexts: return False ids, lengths, starts, _ = snapshot if (ids, lengths) != (previous[0], previous[1]): @@ -1837,9 +1864,10 @@ def _steady_generation_step(self, snapshot) -> bool: def _steady_pages_moved(self, snapshot) -> bool: """Whether a steady step's last position, SWA window start or allocated capacity crossed a page boundary, so its page tables may have changed.""" + md = self.metadata previous = self._csa2_steady_snapshot - block = self.kv_cache_manager.tokens_per_block - window = self.kv_cache_manager.layout.window_size + block = md.kv_cache_manager.tokens_per_block + window = md.kv_cache_manager.layout.window_size _, lengths, starts, allocated = snapshot last = starts + np.asarray(lengths, dtype=np.int64) - 1 prev_last = last - (starts - previous[2]) @@ -1862,12 +1890,13 @@ def _refresh_steady_page_tables( """ from .cache_manager import CSA2CacheRole - manager = self.kv_cache_manager + md = self.metadata + manager = md.kv_cache_manager published_swa = manager._get_swa_publication( - swa_publication, self.kv_cache_block_offsets, request_ids, self.num_contexts, page_count + swa_publication, md.kv_cache_block_offsets, request_ids, md.num_contexts, page_count ) if published_swa is None: - manager.compute_batch_page_tables(request_ids, self.num_contexts) + manager.compute_batch_page_tables(request_ids, md.num_contexts) self._check_swa_write_pages(starts, lengths, floors, allocated_until, page_count) ends = [start + length for start, length in zip(starts, lengths)] for owner in self.csa2_global_page_tables: @@ -1895,7 +1924,8 @@ def refresh(target, spec, published=None): torch._foreach_copy_(targets, sources) def _finish_prepare_csa2(self, snapshot, *, defer_outputs: bool) -> None: - manager = self.kv_cache_manager + md = self.metadata + manager = md.kv_cache_manager self._csa2_prepared_manager = manager self._csa2_forward_serial = getattr(self, "_csa2_forward_serial", 0) + 1 for layer_idx in getattr(self, "_csa2_priors", {}): @@ -1918,7 +1948,7 @@ def _restore_cache_fields(self, fields: dict) -> None: for name, value in fields.items(): setattr(self, name, value) - def prepare_for_draft_forward(self) -> dict | None: + def draft_swap(self) -> dict | None: """Rebuild draft-owned fields after the existing interface swaps managers. Only identical explicit layer layouts are representable here. Virtual @@ -1926,7 +1956,7 @@ def prepare_for_draft_forward(self) -> dict | None: This hook runs before capture/replay, never inside a captured forward. """ target = getattr(self, "_csa2_prepared_manager", None) - draft = self.kv_cache_manager + draft = self.metadata.kv_cache_manager if target is None or target is draft: return None if torch.cuda.is_current_stream_capturing(): @@ -1953,7 +1983,7 @@ def prepare_for_draft_forward(self) -> dict | None: raise return saved - def restore_after_draft_forward(self, saved_state: dict | None) -> None: + def draft_restore(self, saved_state: dict | None) -> None: if saved_state is None: return draft = self._csa2_prepared_manager @@ -1964,37 +1994,37 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: for callback in getattr(self, "_csa2_indexer_resets", {}).values(): callback() - def begin_model_forward(self) -> None: + def begin_forward(self) -> None: """Reset only per-model-forward publications, retaining prepared buffers.""" self.reset_routing() self.csa2_precomputed_kv_layers = set() self.csa2_replay_query_rows = None self.csa2_decoder_capture_lens = None - def on_update_kv_lens(self) -> None: + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: """Refresh query/cache addresses after device-side acceptance correction. Host cached lengths reserve verification capacity; live device endpoints may be shorter. All address/ratio computations stay on device and retain the prepared tensor shapes for captured replay. """ - super().on_update_kv_lens() - if getattr(self, "_csa2_ready_for_kv_update", False): + if not spec_dec_step and getattr(self, "_csa2_ready_for_kv_update", False): self._refresh_live_endpoints() def _refresh_live_endpoints(self) -> None: """Recompute positions and owner slots; SWA slots follow in the next forward.""" + md = self.metadata self._csa2_shared_domains = {} self._csa2_swa_resolved = False _, _, positions = self._refresh_query_positions( - self.kv_lens_cuda, + md.kv_lens_cuda, self._csa2_query_lengths, self._csa2_query_base, self.csa2_token_requests.long(), self._csa2_query_offsets, ) self.csa2_positions.copy_(positions) - self._resolve_owner_slots(self.kv_lens_cuda) + self._resolve_owner_slots(md.kv_lens_cuda) def _resolve_owner_slots(self, kv_lens: torch.Tensor) -> None: """Every KV owner's GLOBAL write slots and compressed groups for the KV ends @@ -2034,7 +2064,8 @@ def _resolve_owner_slots(self, kv_lens: torch.Tensor) -> None: def _owner_slot_arguments(self): """``(owner, (source base, source lengths, pages, slots, compressed, ratio, page size, batch starts, batch lengths, batch cu))`` of every KV owner.""" - manager = self.kv_cache_manager + md = self.metadata + manager = md.kv_cache_manager for owner, (source_base, source_lengths) in self._csa2_source_geometry.items(): ratio = manager.layout.compress_ratios[owner] batch = self._csa2_compression.get(owner) if ratio == 2 else None @@ -2190,7 +2221,7 @@ def _map_global_slots(logical, requests, table, visible_lengths, page_size, max_ slots = torch.where(logical < visible_lengths[:, None], slots, -1) return slots - def prepare_indexer(self, layer_idx: int) -> CSA2TrtllmMetadata | None: + def prepare_indexer(self, layer_idx: int) -> CSA2Metadata | None: """Bind native paged-indexer descriptors for the real generation queries. The owner's index rows are read in place: a GLOBAL page is a run of @@ -2206,10 +2237,11 @@ def prepare_indexer(self, layer_idx: int) -> CSA2TrtllmMetadata | None: from .quantization import INDEX_PAGE_ROWS - manager = self.kv_cache_manager + md = self.metadata + manager = md.kv_cache_manager if manager is None: raise ValueError("CSA2 paged indexer requires a cache manager") - if torch.cuda.is_current_stream_capturing() and not self.is_cuda_graph: + if torch.cuda.is_current_stream_capturing() and not md.is_cuda_graph: raise RuntimeError("Set is_cuda_graph before warming CSA2 paged indexer metadata") owner = self.csa2_kv_sources[layer_idx] if owner is None: @@ -2226,7 +2258,7 @@ def prepare_indexer(self, layer_idx: int) -> CSA2TrtllmMetadata | None: count = decode_end - decode_start request_count = len(generation_ranges) ratio = manager.layout.compress_ratios[owner] - if self.is_cuda_graph: + if md.is_cuda_graph: max_positions = max(1, self.csa2_global_max_positions[owner]) else: max_positions = max( @@ -2240,7 +2272,7 @@ def prepare_indexer(self, layer_idx: int) -> CSA2TrtllmMetadata | None: ), ) native_page_size = INDEX_PAGE_ROWS - if self.is_cuda_graph: + if md.is_cuda_graph: # Reserve the largest owner geometry during warmup so another # serialized owner cannot replace a buffer captured by this graph. required_pages = max( @@ -2255,17 +2287,17 @@ def prepare_indexer(self, layer_idx: int) -> CSA2TrtllmMetadata | None: table = self.csa2_global_page_tables[owner][context_requests:] device = table.device draft_width = 1 + getattr(manager, "max_total_draft_tokens", 0) - query_capacity = self.max_num_requests * draft_width if self.is_cuda_graph else count - if self.is_cuda_graph and count > query_capacity: + query_capacity = md.max_num_requests * draft_width if md.is_cuda_graph else count + if md.is_cuda_graph and count > query_capacity: raise ValueError("CSA2 graph query count exceeds configured batch and draft width") key = ( (device, manager.layout.index_topk, True) - if self.is_cuda_graph + if md.is_cuda_graph else (device, request_count, count, manager.layout.index_topk, False) ) if not hasattr(self, "_csa2_indexer_workspaces"): self._csa2_indexer_workspaces = {} - if not self.is_cuda_graph: + if not md.is_cuda_graph: # Runtime metadata survives many batch geometries. Retain just # the latest eager arena per device; active descriptors still own # their tensors, and serialized stream work preserves their use. @@ -2327,7 +2359,7 @@ def prepare_indexer(self, layer_idx: int) -> CSA2TrtllmMetadata | None: id(manager), getattr(self, "_csa2_forward_serial", 0), count, - self.is_cuda_graph, + md.is_cuda_graph, torch.cuda.is_current_stream_capturing(), ) if scratch["owners"].get(owner) != serial: @@ -2383,6 +2415,7 @@ def prepare_sparse_indexer( """ from tensorrt_llm.deep_gemm import get_paged_sparse_mqa_logits_metadata + md = self.metadata row_requests = getattr(self, "csa2_decode_row_requests", None) table = self.csa2_indexer_block_table owner = self.csa2_kv_sources[layer_idx] @@ -2402,10 +2435,10 @@ def prepare_sparse_indexer( counts = counts[counts.shape[0] - count :].reshape(-1) blocks = blocks[blocks.shape[0] - count :] serial = ( - id(self.kv_cache_manager), + id(md.kv_cache_manager), getattr(self, "_csa2_forward_serial", 0), count, - self.is_cuda_graph, + md.is_cuda_graph, torch.cuda.is_current_stream_capturing(), ) schedules = self.__dict__.setdefault("_csa2_sparse_schedules", {}) @@ -2570,14 +2603,14 @@ def visit(value): def get_workspace_bytes(self) -> int: """Report retained GPU workspace once per storage, excluding KV pools.""" + md = self.metadata values = [ value - for name, value in vars(self).items() + for name, value in (*vars(md).items(), *vars(self).items()) if isinstance(value, torch.Tensor) and name != "position_ids" ] + values += [md.workspace, md.cuda_graph_workspace] for name in ( - "workspace", - "cuda_graph_workspace", "_csa2_buffers", "_csa2_indexer_workspaces", "_csa2_priors", @@ -2586,8 +2619,9 @@ def get_workspace_bytes(self) -> int: ): values.append(getattr(self, name, None)) for metadata in getattr(self, "_csa2_query_tiles", {}).values(): - for name, value in vars(metadata).items(): - if name != "kv_cache_manager": + tile = metadata.ext(CSA2Metadata) + for name, value in (*vars(metadata).items(), *vars(tile).items()): + if name not in ("kv_cache_manager", "metadata"): values.append(value) return self._workspace_storage_bytes(values) @@ -2598,6 +2632,7 @@ def register_indexer_reset(self, layer_idx: int, callback) -> None: self._csa2_indexer_resets[layer_idx] = callback def prepare_indexer_prior(self, layer_idx: int, topk: int) -> torch.Tensor: + md = self.metadata if not hasattr(self, "_csa2_priors"): self._csa2_priors = {} self.csa2_indexer_prior_capacity = {} @@ -2605,7 +2640,7 @@ def prepare_indexer_prior(self, layer_idx: int, topk: int) -> torch.Tensor: if record is None: if torch.cuda.is_current_stream_capturing(): raise RuntimeError("Warm up CSA2 temporal priors before graph capture") - capacity = self.max_num_tokens + capacity = md.max_num_tokens device = self.csa2_token_requests.device record = { "prior": torch.full((capacity, topk), -1, dtype=torch.int32, device=device), @@ -2625,17 +2660,18 @@ def prepare_indexer_prior(self, layer_idx: int, topk: int) -> torch.Tensor: return record["prior"][: record["count"]] def _refresh_indexer_prior(self, layer_idx: int) -> None: + md = self.metadata record = self._csa2_priors[layer_idx] previous = {key: row for row, key in enumerate(record["keys"]) if key is not None} rows, identities = [], [] publish_keys = [None] * sum(self.csa2_request_lengths) reset = self.csa2_num_context_requests > 0 for request_row, length in enumerate(self.csa2_request_lengths): - request_id = self.request_ids[request_row] + request_id = md.request_ids[request_row] start = self.csa2_request_start_positions[request_row] identity = ( - id(self.kv_cache_manager), - self.kv_cache_manager.request_epoch(request_id), + id(md.kv_cache_manager), + md.kv_cache_manager.request_epoch(request_id), request_id, ) begin, end = self.csa2_request_query_ranges[request_row] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/module.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/module.py index b0cb7654db28..0c40c918073b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/module.py @@ -11,6 +11,7 @@ from transformers import PretrainedConfig from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.attention.rotary_embedding import RotaryEmbedding from tensorrt_llm._torch.distributed import AllReduceStrategy from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode @@ -23,7 +24,7 @@ from ..deepseek_v4.module import project_sparse_attn_output from .compressor import CSA2Compressor -from .metadata import CSA2TrtllmMetadata +from .metadata import CSA2Metadata from .params import ( CSA2BackendForwardArgs, CSA2ForwardState, @@ -378,29 +379,30 @@ def global_transforms(self, compressed_positions): ) def prepare_global_cache( - self, hidden_states: torch.Tensor, metadata: CSA2TrtllmMetadata + self, hidden_states: torch.Tensor, metadata: TrtllmAttentionMetadata ) -> None: """Publish a CED owner's full encoder rows before narrowing decoder queries.""" + csa2 = metadata.ext(CSA2Metadata) layer_idx = self.layer.layer_idx if self.layer.mode != CSA2Mode.FULL or self.compressor is None: raise ValueError("Only a Full CSA2 layer can precompute global KV") - if layer_idx in metadata.csa2_precomputed_kv_layers: + if layer_idx in csa2.csa2_precomputed_kv_layers: raise ValueError("CSA2 global KV was already precomputed for this forward") - source = metadata.select_global_source(layer_idx, hidden_states, None) + source = csa2.select_global_source(layer_idx, hidden_states, None) main_kv, index_k = self._prepare_global( hidden_states, source, - metadata.get_compression_batch(layer_idx), - metadata.get_compressed_positions(layer_idx), + csa2.get_compression_batch(layer_idx), + csa2.get_compressed_positions(layer_idx), ) metadata.kv_cache_manager.write_global( self.layer.kv_source, - metadata.csa2_main_write_slots[self.layer.kv_source], + csa2.csa2_main_write_slots[self.layer.kv_source], main_kv, index_k, - *self.global_transforms(metadata.get_compressed_positions(layer_idx)), + *self.global_transforms(csa2.get_compressed_positions(layer_idx)), ) - metadata.csa2_precomputed_kv_layers.add(layer_idx) + csa2.csa2_precomputed_kv_layers.add(layer_idx) def _prepare_auxiliary( self, @@ -503,7 +505,7 @@ def forward( self, hidden_states: torch.Tensor, positions: torch.Tensor, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, *, global_hidden_states: torch.Tensor | None = None, ) -> torch.Tensor: @@ -514,17 +516,16 @@ def forward( completed group. ``global_hidden_states`` is the encoder output when decoder global KV is prepared from more rows than its SWA replay. """ - precomputed = self.layer.layer_idx in attn_metadata.csa2_precomputed_kv_layers - compression = attn_metadata.get_compression_batch(self.layer.layer_idx) + csa2 = attn_metadata.ext(CSA2Metadata) + precomputed = self.layer.layer_idx in csa2.csa2_precomputed_kv_layers + compression = csa2.get_compression_batch(self.layer.layer_idx) compressed_positions = ( - attn_metadata.get_compressed_positions(self.layer.layer_idx) + csa2.get_compressed_positions(self.layer.layer_idx) if self.layer.mode == CSA2Mode.FULL else None ) global_source = ( - attn_metadata.select_global_source( - self.layer.layer_idx, hidden_states, global_hidden_states - ) + csa2.select_global_source(self.layer.layer_idx, hidden_states, global_hidden_states) if self.compressor is not None and not precomputed else None ) @@ -621,7 +622,7 @@ def forward( if start == end: continue tile = q[start:end] - metadata = attn_metadata.get_query_tile_metadata( + metadata = csa2.get_query_tile_metadata( tile, 0 if self.layer.mode == CSA2Mode.SWA else self.layout.index_topk, query_start=start, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/csa2/params.py b/tensorrt_llm/_torch/attention/backends/sparse/csa2/params.py index 39ede554d38e..bc73aa9d79a4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/csa2/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/csa2/params.py @@ -13,7 +13,7 @@ from ..params import SparseBackendForwardArgs, SparseMetadataParams, SparseParams if TYPE_CHECKING: - from .metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata class CSA2Mode(Enum): @@ -245,7 +245,7 @@ class CSA2ForwardState: named next to them; the cache writes apply those in the quantize launch. """ - metadata: CSA2TrtllmMetadata + metadata: TrtllmAttentionMetadata swa_kv: torch.Tensor index_q: torch.Tensor | None = None index_q_scale: torch.Tensor | None = None diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/__init__.py index a822bd90bd40..4528b1529a4f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/__init__.py @@ -16,7 +16,7 @@ from .backend import DeepseekV4TrtllmAttention from .cache_manager import DeepseekV4CacheManager from .indexer import DeepseekV4Indexer -from .metadata import DeepseekV4TrtllmAttentionMetadata +from .metadata import DeepseekV4Metadata from .params import DeepseekV4AttentionType, DeepSeekV4MetadataParams, DeepSeekV4Params __all__ = [ @@ -26,5 +26,5 @@ "DeepseekV4CacheManager", "DeepseekV4Indexer", "DeepseekV4TrtllmAttention", - "DeepseekV4TrtllmAttentionMetadata", + "DeepseekV4Metadata", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py index 1337b7c747f9..7cc919ff8855 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/backend.py @@ -27,7 +27,7 @@ PositionalEmbeddingParams, merge_attention_forward_args, ) -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm.bindings import DataType from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantMode @@ -37,7 +37,7 @@ from .compressor import Compressor from .indexer import DeepseekV4Indexer from .kernels import deepseek_v4_local_to_global_indices -from .metadata import DeepseekV4TrtllmAttentionMetadata +from .metadata import DeepseekV4Metadata from .params import DeepseekV4AttentionType, DeepSeekV4Params if TYPE_CHECKING: @@ -70,7 +70,8 @@ def _alloc_nvfp4_compress_fp8_scratch( class DeepseekV4TrtllmAttention(TrtllmAttention): - Metadata = DeepseekV4TrtllmAttentionMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (DeepseekV4Metadata,) @property def uses_fp4_mla_attention(self) -> bool: @@ -203,9 +204,10 @@ def __init__( def _prepare_sparse_forward_args( self, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> None: + dsv4 = metadata.ext(DeepseekV4Metadata) attention_input_type = forward_args.attention_input_type if attention_input_type == AttentionInputType.context_only: start_idx = 0 @@ -218,17 +220,17 @@ def _prepare_sparse_forward_args( end_idx = metadata.num_tokens sparse_args = forward_args.sparse_runtime_params - sparse_args.sparse_attn_kv_lens = metadata.sparse_mla_topk_lens[self.compress_ratio][ + sparse_args.sparse_attn_kv_lens = dsv4.sparse_mla_topk_lens[self.compress_ratio][ start_idx:end_idx ] if self.compress_ratio > 1: - sparse_args.aux_kv_cache_pool_ptr = metadata.sparse_mla_base_ptrs[self.compress_ratio] + sparse_args.aux_kv_cache_pool_ptr = dsv4.sparse_mla_base_ptrs[self.compress_ratio] else: sparse_args.aux_kv_cache_pool_ptr = None metadata.num_sparse_topk = ( self.sparse_attention_config.window_size - + metadata.max_compressed_indices[self.compress_ratio] + + dsv4.max_compressed_indices[self.compress_ratio] ) def forward( @@ -236,7 +238,7 @@ def forward( q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: Optional[AttentionForwardArgs] = None, **kwargs, ): @@ -261,18 +263,19 @@ def sparse_attn_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Convert local indices (SWA + compressed) to global pool indices.""" + dsv4 = metadata.ext(DeepseekV4Metadata) layer_idx = self.layer_idx kv_cache_manager = metadata.kv_cache_manager attention_input_type = forward_args.attention_input_type - swa_pool_base_ptr = metadata.sparse_mla_base_ptrs[1] + swa_pool_base_ptr = dsv4.sparse_mla_base_ptrs[1] # Get cached buffer pointers - swa_buffer_ptr = metadata.swa_buffer_ptrs[layer_idx] + swa_buffer_ptr = dsv4.swa_buffer_ptrs[layer_idx] # Token stride index_head_dim = self.sparse_attention_config.index_head_dim @@ -298,18 +301,18 @@ def sparse_attn_predict( end_idx = metadata.num_tokens # Use global req_id directly - req_id = metadata.req_idx_per_token[start_idx:end_idx] - swa_local_indices = metadata.swa_local_indices_cuda[start_idx:end_idx] + req_id = dsv4.req_idx_per_token[start_idx:end_idx] + swa_local_indices = dsv4.swa_local_indices_cuda[start_idx:end_idx] local_layer_idx = kv_cache_manager.layer_offsets[layer_idx] - block_table_swa = metadata.sliding_block_tables[ + block_table_swa = dsv4.sliding_block_tables[ local_layer_idx, DeepseekV4AttentionType.SWA.value ] active_request_count = None if self.compress_ratio > 1: - compressed_buffer_ptr = metadata.compressed_buffer_ptrs[layer_idx] - compress_pool_base_ptr = metadata.sparse_mla_base_ptrs[self.compress_ratio] - block_table_compressed = metadata.compress_block_tables[self.compress_ratio] + compressed_buffer_ptr = dsv4.compressed_buffer_ptrs[layer_idx] + compress_pool_base_ptr = dsv4.sparse_mla_base_ptrs[self.compress_ratio] + block_table_compressed = dsv4.compress_block_tables[self.compress_ratio] if self.compress_ratio == 4: sparse_backend_args = forward_args.sparse_backend_args assert sparse_backend_args is not None, ( @@ -324,7 +327,7 @@ def sparse_attn_predict( self.sparse_attention_config.enable_kv_cache_offload and attention_input_type == AttentionInputType.generation_only ): - state = metadata.sparse_offload_state + state = dsv4.sparse_offload_state assert state is not None, "Sparse offload metadata must be initialized" # module.forward_sparse_attn has joined both producer # streams. Keep fetch before the conversion kernel's @@ -334,14 +337,14 @@ def sparse_attn_predict( layer_idx, topk_indices, req_id, - metadata.compressed_kv_lens_cuda[self.compress_ratio], + dsv4.compressed_kv_lens_cuda[self.compress_ratio], block_table_compressed, ) # Metadata counts include CUDA graph dummy rows. The active # count is refreshed in place before replay to mask them. active_request_count = state.active_request_count else: - compressed_local_indices = metadata.compressed_local_indices_cuda[start_idx:end_idx] + compressed_local_indices = dsv4.compressed_local_indices_cuda[start_idx:end_idx] else: compressed_buffer_ptr = 0 compress_pool_base_ptr = 0 @@ -401,7 +404,7 @@ def sparse_attn_predict( compress_pool_base_ptr=compress_pool_base_ptr, compressed_buffer_ptr=compressed_buffer_ptr, compress_ratio=self.compress_ratio, - num_compressed_indices=metadata.max_compressed_indices[self.compress_ratio], + num_compressed_indices=dsv4.max_compressed_indices[self.compress_ratio], **sched_kwargs, split_extra=self.use_fp8_ds_mla, active_request_count=active_request_count, @@ -413,7 +416,7 @@ def sparse_attn_predict( global_indices = result if self._uses_nvfp4_compress and self.compress_ratio > 1: - num_compressed_indices = metadata.max_compressed_indices[self.compress_ratio] + num_compressed_indices = dsv4.max_compressed_indices[self.compress_ratio] compressed_indices = global_indices[:, -num_compressed_indices:].contiguous() data_pool, scale_pool = kv_cache_manager.get_compress_pool_buffers(self.compress_ratio) num_pool_tokens = data_pool.numel() // ( @@ -422,7 +425,7 @@ def sparse_attn_predict( # The previous layer has already enqueued its attention on this stream. # Release its Python reference before allocating the next scratch so # the caching allocator can reuse the storage in stream order. - metadata.nvfp4_compress_fp8_scratch = None + dsv4.nvfp4_compress_fp8_scratch = None if attention_input_type == AttentionInputType.generation_only: num_scratch_rows = compressed_indices.numel() scratch = _alloc_nvfp4_compress_fp8_scratch( @@ -443,7 +446,7 @@ def sparse_attn_predict( if attention_input_type == AttentionInputType.context_only else metadata.num_seqs ) - compressed_kv_lengths = metadata.compressed_kv_lens_cuda[self.compress_ratio][ + compressed_kv_lengths = dsv4.compressed_kv_lens_cuda[self.compress_ratio][ :request_count ].contiguous() total_raw_kv_tokens = int(metadata.host_total_kv_lens[0].item()) @@ -478,7 +481,7 @@ def sparse_attn_predict( max_compressed_kv_tokens, num_pool_tokens, ) - metadata.nvfp4_compress_fp8_scratch = scratch + dsv4.nvfp4_compress_fp8_scratch = scratch global_indices[:, -num_compressed_indices:] = compressed_indices forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr() @@ -488,7 +491,7 @@ def sparse_kv_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: return None, None diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py index 64d90b780996..c27a6bc93a85 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py @@ -17,7 +17,7 @@ from .params import DeepseekV4AttentionType if TYPE_CHECKING: - from .metadata import DeepseekV4TrtllmAttentionMetadata + from ...trtllm import TrtllmAttentionMetadata class KVCacheDtype(IntEnum): @@ -158,7 +158,7 @@ def __init__( def forward( self, x: torch.Tensor, - metadata: "DeepseekV4TrtllmAttentionMetadata", + metadata: "TrtllmAttentionMetadata", ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Forward pass for paged KV compression. @@ -175,6 +175,10 @@ def forward( - no compressed tokens: (None, None) """ # Extract metadata + # The metadata module imports this one through the indexer. + from .metadata import DeepseekV4Metadata + + dsv4 = metadata.ext(DeepseekV4Metadata) num_contexts = metadata.num_contexts num_generations = metadata.num_generations num_ctx_tokens = metadata.num_ctx_tokens @@ -203,11 +207,11 @@ def forward( # Get block tables local_layer_idx = metadata.kv_cache_manager.layer_offsets[self.layer_idx] if self.is_indexer: - block_table = metadata.indexer_k_cache_block_offsets + block_table = dsv4.indexer_k_cache_block_offsets else: - block_table = metadata.compress_block_tables[self.compress_ratio] - block_table_kv_state = metadata.sliding_block_tables[local_layer_idx, kv_type.value] - block_table_score_state = metadata.sliding_block_tables[local_layer_idx, score_type.value] + block_table = dsv4.compress_block_tables[self.compress_ratio] + block_table_kv_state = dsv4.sliding_block_tables[local_layer_idx, kv_type.value] + block_table_score_state = dsv4.sliding_block_tables[local_layer_idx, score_type.value] # Get tokens_per_block from cache manager # state_tokens_per_block: for compressor kv/score state caches (used in compress kernels) @@ -216,11 +220,11 @@ def forward( compress_tokens_per_block = metadata.kv_cache_manager.compressed_block_sizes[self.layer_idx] # Get compression metadata - cu_new_comp_kv = metadata.cu_new_comp_kv_cuda[self.compress_ratio] + cu_new_comp_kv = dsv4.cu_new_comp_kv_cuda[self.compress_ratio] kv_lens = metadata.kv_lens_cuda_runtime - total_num_comp_tokens = metadata.num_total_compressed_tokens[self.compress_ratio] - num_comp_tokens = metadata.new_comp_kv_lens_cuda[self.compress_ratio][:bsz] - max_ctx_comp_kv_lens = metadata.max_ctx_compressed_tokens[self.compress_ratio] + total_num_comp_tokens = dsv4.num_total_compressed_tokens[self.compress_ratio] + num_comp_tokens = dsv4.new_comp_kv_lens_cuda[self.compress_ratio][:bsz] + max_ctx_comp_kv_lens = dsv4.max_ctx_compressed_tokens[self.compress_ratio] # Project input to KV and score in the checkpoint dtype. The compressor # kernels accept bf16 or fp32 kv_score and convert values to fp32 @@ -241,8 +245,8 @@ def forward( block_table_score_state[:num_contexts], kv_comp, kv_lens[:num_contexts], - metadata.cached_token_lens_cuda[:num_contexts], - metadata.cu_seq_lens_cuda, + dsv4.cached_token_lens_cuda[:num_contexts], + dsv4.cu_seq_lens_cuda, cu_new_comp_kv[: num_contexts + 1], num_contexts, state_tokens_per_block, @@ -257,8 +261,8 @@ def forward( # batches: next_n is only the batch maximum (the kernel's loop # bound) and the real per-request counts come through # new_tokens_per_seq. - next_n = metadata.num_gen_tokens_per_seq - new_tokens_per_seq = metadata.gen_new_tokens_per_seq + next_n = dsv4.num_gen_tokens_per_seq + new_tokens_per_seq = dsv4.gen_new_tokens_per_seq # Pass full kv_score (not sliced) with the generation portion of # cu_seq_lens so the kernel reads at the correct absolute offsets. torch.ops.trtllm.compressor_paged_kv_compress( @@ -270,7 +274,7 @@ def forward( block_table_score_state[num_contexts:], kv_comp, gen_kv_lens, - metadata.cu_seq_lens_cuda[num_contexts:], + dsv4.cu_seq_lens_cuda[num_contexts:], cu_new_comp_kv[num_contexts:], num_generations, state_tokens_per_block, @@ -281,7 +285,7 @@ def forward( ) # Scatter to cache with appropriate quantization (all modes fused) - start_pos = metadata.past_kv_lens_cuda[self.compress_ratio][:bsz] + start_pos = dsv4.past_kv_lens_cuda[self.compress_ratio][:bsz] total_tokens = kv_comp.shape[0] # Allocate optional returned postprocess buffers for indexer paths. @@ -308,8 +312,8 @@ def forward( total_tokens, num_scale_blocks, dtype=torch.uint8, device=kv_comp.device ) - position_ids = metadata.compressed_position_ids_cuda[self.compress_ratio][:total_tokens] - compressed_mask = metadata.compressed_mask_cuda[self.compress_ratio][:total_tokens] + position_ids = dsv4.compressed_position_ids_cuda[self.compress_ratio][:total_tokens] + compressed_mask = dsv4.compressed_mask_cuda[self.compress_ratio][:total_tokens] if self.footer_scale_cache and not self.is_indexer: self._footer_scale_postprocess_scatter( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flash_mla.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flash_mla.py index 5018ee2fdad0..95b7d730a340 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flash_mla.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flash_mla.py @@ -8,13 +8,13 @@ import torch -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._utils import nvtx_range, nvtx_range_debug from tensorrt_llm.logger import logger from . import footer_scale_kv from .kernels import deepseek_v4_local_to_global_indices -from .metadata import DeepseekV4TrtllmAttentionMetadata +from .metadata import DeepseekV4Metadata from .params import DEEPSEEK_V4_SPARSE_RATIO, DeepseekV4AttentionType _BF16_CONTEXT_CHUNK_SIZE = 512 @@ -47,7 +47,7 @@ def _prepare_q_and_cache( self, q: torch.Tensor, latent_cache: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, *, is_generation: bool, rotary_cos_sin: Optional[torch.Tensor] = None, @@ -79,37 +79,38 @@ def _prepare_q_and_cache( def _prepare_pool_indices( self, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, topk_indices: Optional[torch.Tensor], *, is_generation: bool, ) -> tuple[torch.Tensor, int, int]: """Convert local dual-pool indices to layer-pool-relative indices.""" + dsv4 = metadata.ext(DeepseekV4Metadata) kv_cache_manager = metadata.kv_cache_manager - window_size = metadata.window_size + window_size = dsv4.window_size start_idx = metadata.num_ctx_tokens if is_generation else 0 end_idx = metadata.num_tokens if is_generation else metadata.num_ctx_tokens - req_id = metadata.req_idx_per_token[start_idx:end_idx] - swa_local_indices = metadata.swa_local_indices_cuda[start_idx:end_idx] + req_id = dsv4.req_idx_per_token[start_idx:end_idx] + swa_local_indices = dsv4.swa_local_indices_cuda[start_idx:end_idx] local_layer_idx = kv_cache_manager.layer_offsets[self._layer_idx] - block_table_swa = metadata.sliding_block_tables[ + block_table_swa = dsv4.sliding_block_tables[ local_layer_idx, DeepseekV4AttentionType.SWA.value ] - swa_buffer_ptr = metadata.swa_buffer_ptrs[self._layer_idx] + swa_buffer_ptr = dsv4.swa_buffer_ptrs[self._layer_idx] block_table_compressed = None compressed_local_indices = None compressed_buffer_ptr = 0 if self._compress_ratio > 1: - compressed_buffer_ptr = metadata.compressed_buffer_ptrs[self._layer_idx] - block_table_compressed = metadata.compress_block_tables[self._compress_ratio] + compressed_buffer_ptr = dsv4.compressed_buffer_ptrs[self._layer_idx] + block_table_compressed = dsv4.compress_block_tables[self._compress_ratio] if self._compress_ratio == DEEPSEEK_V4_SPARSE_RATIO: if topk_indices is None: raise ValueError("DeepSeek-V4 ratio-4 attention requires top-k indices") compressed_local_indices = topk_indices else: - compressed_local_indices = metadata.compressed_local_indices_cuda[start_idx:end_idx] + compressed_local_indices = dsv4.compressed_local_indices_cuda[start_idx:end_idx] # Equal base/layer pointers make both output slices relative to the # tensors passed to FlashMLA instead of the shared sparse arena. @@ -126,7 +127,7 @@ def _prepare_pool_indices( compress_pool_base_ptr=compressed_buffer_ptr, compressed_buffer_ptr=compressed_buffer_ptr, compress_ratio=self._compress_ratio, - num_compressed_indices=metadata.max_compressed_indices[self._compress_ratio], + num_compressed_indices=dsv4.max_compressed_indices[self._compress_ratio], ) return global_indices, self._compress_ratio, window_size @@ -238,7 +239,7 @@ def forward_context( self, q: torch.Tensor, latent_cache: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, output: torch.Tensor, topk_indices: Optional[torch.Tensor], softmax_scale: float, @@ -263,7 +264,7 @@ def _forward_bf16( self, q: torch.Tensor, latent_cache: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, output: torch.Tensor, topk_indices: Optional[torch.Tensor], softmax_scale: float, @@ -363,7 +364,7 @@ def _forward_fp8( self, q: torch.Tensor, latent_cache: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, output: torch.Tensor, topk_indices: torch.Tensor, softmax_scale: float, @@ -464,7 +465,7 @@ def forward_generation( self, q: torch.Tensor, latent_cache: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, output: torch.Tensor, topk_indices: Optional[torch.Tensor], softmax_scale: float, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flashinfer.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flashinfer.py index efa20445a310..55539f5fd76b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/flashinfer.py @@ -10,30 +10,32 @@ from tensorrt_llm._torch.attention.rotary_embedding import RotaryEmbedding from ...interface import AttentionForwardArgs, AttentionInputType +from ...trtllm import TrtllmAttentionMetadata from ..flashinfer_utils import ( SPARSE_MLA_SPLIT_KV_TILE, allocate_sparse_mla_split_workspace, get_sparse_mla_op, ) from . import footer_scale_kv -from .metadata import DeepseekV4TrtllmAttentionMetadata +from .metadata import DeepseekV4Metadata from .params import DeepseekV4AttentionType def run_flashinfer_sparse_mla( attn, q: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, rotary_emb: RotaryEmbedding, ) -> None: """Run DeepSeek-V4 attention for ``FlashInferSparseMlaFmha``.""" + dsv4 = metadata.ext(DeepseekV4Metadata) if forward_args.enable_dsv4_epilogue_fusion: raise NotImplementedError( "DSv4 epilogue fusion is fused into the C++ attention op; the " "FlashInfer SM120 FMHA does not provide it." ) - if metadata.max_draft_tokens > 0: + if dsv4.max_draft_tokens > 0: raise NotImplementedError( "MTP / speculative decoding is not supported by the FlashInfer DSv4 FMHA yet." ) @@ -83,7 +85,7 @@ def run_flashinfer_sparse_mla( rotary_emb.is_neox, apply_rope=not forward_args.skip_mla_rope_generation, ) - positions = metadata.token_positions_cuda[start_idx:end_idx] + positions = dsv4.token_positions_cuda[start_idx:end_idx] sparse_runtime_params = forward_args.sparse_runtime_params swa_indices = sparse_runtime_params.sparse_attn_indices @@ -98,7 +100,7 @@ def run_flashinfer_sparse_mla( if extra_indices is None: raise RuntimeError("Compressed DeepSeek-V4 layers require extra sparse indices.") extra_lens = ( - metadata.sparse_mla_topk_lens[attn.compress_ratio][start_idx:end_idx] - window + dsv4.sparse_mla_topk_lens[attn.compress_ratio][start_idx:end_idx] - window ).clamp(min=0) extra_pool = footer_scale_kv.get_pool_2d( metadata, DeepseekV4AttentionType.COMPRESS, attn.compress_ratio diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/footer_scale_kv.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/footer_scale_kv.py index a1cfd85c1225..479999e4f674 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/footer_scale_kv.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/footer_scale_kv.py @@ -20,12 +20,14 @@ from tensorrt_llm.bindings import DataType from .kernels import deepseek_v4_local_to_global_indices +from .metadata import DeepseekV4Metadata from .params import DeepseekV4AttentionType if TYPE_CHECKING: - from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention - - from .metadata import DeepseekV4TrtllmAttentionMetadata + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) DIM_NOPE = 448 DIM_ROPE = 64 @@ -152,7 +154,7 @@ def quant_scatter( def get_pool_2d( - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, attn_type: DeepseekV4AttentionType, compress_ratio: int, page_size: int | None = None, @@ -210,25 +212,24 @@ def get_pool_2d( def append_swa( attn: TrtllmAttention, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, latent_rows: torch.Tensor, start_idx: int, end_idx: int, page_size: int = PAGE_SIZE, ) -> None: """Quantize new latent rows and scatter them into the canonical SWA pool.""" - positions = metadata.token_positions_cuda[start_idx:end_idx] - req_id = metadata.req_idx_per_token[start_idx:end_idx] + dsv4 = metadata.ext(DeepseekV4Metadata) + positions = dsv4.token_positions_cuda[start_idx:end_idx] + req_id = dsv4.req_idx_per_token[start_idx:end_idx] local_layer_idx = metadata.kv_cache_manager.layer_offsets[attn.layer_idx] - block_table_swa = metadata.sliding_block_tables[ - local_layer_idx, DeepseekV4AttentionType.SWA.value - ] + block_table_swa = dsv4.sliding_block_tables[local_layer_idx, DeepseekV4AttentionType.SWA.value] loc = deepseek_v4_local_to_global_indices( req_id=req_id, block_table_swa=block_table_swa, swa_local_indices=positions.unsqueeze(1).contiguous(), - swa_pool_base_ptr=metadata.sparse_mla_base_ptrs[1], - swa_buffer_ptr=metadata.swa_buffer_ptrs[attn.layer_idx], + swa_pool_base_ptr=dsv4.sparse_mla_base_ptrs[1], + swa_buffer_ptr=dsv4.swa_buffer_ptrs[attn.layer_idx], tokens_per_block=metadata.kv_cache_manager.tokens_per_block, token_stride=TOKEN_BYTES, ).view(-1) @@ -238,7 +239,7 @@ def append_swa( def apply_rope_and_append_swa( attn: TrtllmAttention, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, q: torch.Tensor, latent_rows: torch.Tensor | None, start_idx: int, @@ -250,13 +251,14 @@ def apply_rope_and_append_swa( page_size: int = PAGE_SIZE, ) -> None: """Apply MLA RoPE and append latent rows to the canonical SWA cache.""" + dsv4 = metadata.ext(DeepseekV4Metadata) num_tokens = end_idx - start_idx if num_tokens == 0: return if attn.mla_params is None: raise ValueError("Footer-scale MLA cache requires MLA parameters") - positions = metadata.token_positions_cuda[start_idx:end_idx] + positions = dsv4.token_positions_cuda[start_idx:end_idx] if positions.numel() != num_tokens: raise ValueError( "Expected one cached token position per footer-scale append, got " diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py index 995fc1c96ac8..c9043e766b7f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/indexer.py @@ -22,8 +22,7 @@ from .params import DeepSeekV4Params if TYPE_CHECKING: - from ..dsa.metadata import DSAtrtllmAttentionMetadata - from .metadata import DeepseekV4TrtllmAttentionMetadata + from ...trtllm import TrtllmAttentionMetadata class DeepseekV4Indexer(Indexer): @@ -214,7 +213,7 @@ def _update_k_cache_if_needed( self, k_fp8: Optional[torch.Tensor], k_scale: Optional[torch.Tensor], - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, ) -> None: if k_fp8 is None: return @@ -231,7 +230,7 @@ def _update_k_cache_if_needed( def precompute_aux( self, hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, start_event: Optional[torch.cuda.Event] = None, ) -> Optional[Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]]: """Pre-launch the qr-independent half of the indexer prepare phase. @@ -265,7 +264,7 @@ def _run_overlapped_indexer_prepare( self, qr: torch.Tensor, hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, position_ids: torch.Tensor, pre_aux: Optional[ Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] @@ -338,7 +337,7 @@ def _run_serial_indexer_prepare( self, qr: torch.Tensor, hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, position_ids: torch.Tensor, ) -> Tuple[ torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor @@ -357,7 +356,7 @@ def _update_k_cache( self, k_fp8: torch.Tensor, k_scale: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, ) -> None: # DSV4's indexer compressor already scatters INDEXER_COMPRESS. The # shared DSA scatter would duplicate that write. @@ -367,12 +366,15 @@ def forward( self, qr: torch.Tensor, hidden_states: torch.Tensor, - metadata: DeepseekV4TrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, position_ids: torch.Tensor, pre_aux: Optional[ Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]] ] = None, ): + from .metadata import DeepseekV4Metadata + + dsv4 = metadata.ext(DeepseekV4Metadata) use_overlapped_prepare = ( do_multi_stream() and self.aux_stream is not None @@ -394,7 +396,7 @@ def forward( # If there are no compressed tokens, return an topk indices buffer with all -1s in the tensor. if k_fp8 is None: - topk_indices = metadata.empty_topk_indices_buffer[: hidden_states.shape[0]] + topk_indices = dsv4.empty_topk_indices_buffer[: hidden_states.shape[0]] else: topk_indices = self.sparse_attn_indexer( metadata, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 4ed9a04e968b..560dc3edc9da 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -9,11 +9,10 @@ import torch -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.utils import maybe_compile from tensorrt_llm._utils import prefer_pinned -from ..dsa.metadata import DSAtrtllmAttentionMetadata +from ..dsa.metadata import DSAMetadata from .indexer import DeepseekV4Indexer from .offload import SparseOffloadState from .params import ( @@ -28,7 +27,9 @@ from .cache_manager import DeepseekV4CacheManager -class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): +class DeepseekV4Metadata(DSAMetadata): + """DeepSeek-V4 compressed and sliding-window state on top of the DSA indexer state.""" + # The set of compress ratios for the layers compress_ratio_set: Set[int] # The set of (compress ratio, attention type) for the layers @@ -38,21 +39,19 @@ class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): # The max number of context compressed tokens for each compress ratio max_ctx_compressed_tokens: Dict[int, int] - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def _invalidate_mla_scheduler_buffers(self) -> None: + def invalidate_length_derived_buffers(self) -> None: """Invalidate request- and token-prefix buffers derived from sequence lengths.""" - super()._invalidate_mla_scheduler_buffers() + super().invalidate_length_derived_buffers() self._fused_q_gen_cu_seqlens_valid = False - def __post_init__(self): - super().__post_init__() + def allocate(self, buffers) -> None: + md = self.metadata + super().allocate(buffers) self.num_total_compressed_tokens = {} self.max_ctx_compressed_tokens = {} self.nvfp4_compress_fp8_scratch: Optional[torch.Tensor] = None self._ctx_output_sizes: Optional[Dict[int, int]] = None - sparse_metadata_params = self.sparse_metadata_params + sparse_metadata_params = md.sparse_metadata_params if not isinstance(sparse_metadata_params, DeepSeekV4MetadataParams): raise ValueError("DeepSeek-V4 sparse attention metadata params are not set") self.window_size = sparse_metadata_params.window_size @@ -61,7 +60,7 @@ def __post_init__(self): f"Dual-pool sparse MLA requires window_size == 128, which equals to the" f"TileSizeKV of the FMHA kernel. (got {window_size})." ) - capture_graph = self.is_cuda_graph + capture_graph = md.is_cuda_graph self.compress_ratio_set = set(self.compress_ratios) # Cache a sorted list for deterministic iteration order across # torch.compile traces. Avoids repeated set→list conversion and @@ -98,9 +97,9 @@ def __post_init__(self): # Create buffers for the compressor # cu_seq_lens_cuda is the cumulative sequence lengths for the requests - self.cu_seq_lens_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), + self.cu_seq_lens_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences + 1,), dtype=torch.int, cache_name="cu_seq_lens_cuda", capture_graph=capture_graph, @@ -117,18 +116,18 @@ def __post_init__(self): if self.enable_ragged_verification: # The fused Q normalization op consumes request-major token # prefixes, rather than the sparse-MLA token-major row view. - self.fused_q_gen_cu_seqlens = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), + self.fused_q_gen_cu_seqlens = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences + 1,), cache_name="fused_q_gen_cu_seqlens", dtype=torch.int32, capture_graph=capture_graph, ) self.fused_q_gen_cu_seqlens[:1].zero_() # Exact tokens appended by each generation request. - self.gen_new_tokens_per_seq_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + self.gen_new_tokens_per_seq_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), dtype=torch.int, cache_name="gen_new_tokens_per_seq_cuda", capture_graph=capture_graph, @@ -136,9 +135,9 @@ def __post_init__(self): # new_comp_kv_lens_cuda is the number of new compressed tokens for the requests self.new_comp_kv_lens_cuda = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), dtype=torch.int, cache_name=f"new_comp_kv_lens_cuda_{compress_ratio}", capture_graph=capture_graph, @@ -148,9 +147,9 @@ def __post_init__(self): # cu_new_comp_kv_cuda is the cumulative number of new compressed tokens for the requests self.cu_new_comp_kv_cuda = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences + 1,), dtype=torch.int, cache_name=f"cu_new_comp_kv_cuda_{compress_ratio}", capture_graph=capture_graph, @@ -160,9 +159,9 @@ def __post_init__(self): # compressed_kv_lens_cuda is the number of compressed tokens for the requests self.compressed_kv_lens_cuda = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), dtype=torch.int, cache_name=f"compressed_kv_lens_cuda_{compress_ratio}", capture_graph=capture_graph, @@ -172,9 +171,9 @@ def __post_init__(self): # past_kv_lens_cuda is the number of past compressed tokens for the requests self.past_kv_lens_cuda = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), dtype=torch.int, cache_name=f"past_kv_lens_cuda_{compress_ratio}", capture_graph=capture_graph, @@ -184,9 +183,9 @@ def __post_init__(self): # compressed_position_ids_cuda is the compressed position ids for the requests self.compressed_position_ids_cuda = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), dtype=torch.int, cache_name=f"compressed_position_ids_cuda_{compress_ratio}", capture_graph=capture_graph, @@ -195,8 +194,8 @@ def __post_init__(self): } # Device scalar batch size for the shape-static compressed_mask graph. - self.batch_size_cuda = self.get_empty( - self.cuda_graph_buffers, + self.batch_size_cuda = md.get_empty( + md.cuda_graph_buffers, (1,), dtype=torch.int, cache_name="batch_size_cuda", @@ -206,9 +205,9 @@ def __post_init__(self): # compressed_mask_cuda: per-token bool mask for postprocess scatter. # Precomputed from new_comp_kv_lens to skip padded generation slots. self.compressed_mask_cuda = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), dtype=torch.bool, cache_name=f"compressed_mask_cuda_{compress_ratio}", capture_graph=capture_graph, @@ -223,9 +222,9 @@ def __post_init__(self): } # empty topk indices buffer with all -1s in the tensor - self.empty_topk_indices_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, self.sparse_mla_topk), + self.empty_topk_indices_buffer = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens, self.sparse_mla_topk), cache_name="empty_topk_indices_buffer", dtype=torch.int32, capture_graph=capture_graph, @@ -233,9 +232,9 @@ def __post_init__(self): self.empty_topk_indices_buffer.fill_(-1) # SWA local indices - self.swa_local_indices_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, self.window_size), + self.swa_local_indices_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens, self.window_size), cache_name="swa_local_indices_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -245,9 +244,9 @@ def __post_init__(self): # lengths and the fused output RoPE). The DSA base allocates them only # for the fp8_ds_mla layout, so provide them for every DSv4 cache mode. if getattr(self, "token_positions_cuda", None) is None: - self.token_positions_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.token_positions_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="token_positions_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -256,7 +255,7 @@ def __post_init__(self): # Compute max_compressed_indices for CUDA graph compatibility. # For ratio=4, the indexer selects index_topk compressed tokens. # For ratio=128, use max_seq_len / 128 rounded up to next power of 2 - raw_128 = math.ceil(self.max_seq_len / 128) + raw_128 = math.ceil(md.max_seq_len / 128) po2_128 = 1 << (raw_128 - 1).bit_length() if raw_128 > 0 else 1 self.max_compressed_indices = { 1: 0, # No compressed indices @@ -266,9 +265,9 @@ def __post_init__(self): # Compressed local indices for compress_ratio=128 # Note: ratio=4 uses dynamic topk_indices from indexer, so we only pre-allocate for ratio=128 - self.compressed_local_indices_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, self.max_compressed_indices[128]), + self.compressed_local_indices_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens, self.max_compressed_indices[128]), cache_name="compressed_local_indices_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -279,13 +278,13 @@ def __post_init__(self): # uses ratio-shared page indices, and INDEXER_COMPRESS keeps a separate # compatibility table for the generic DSA indexer path. block_table_shape = ( - self.kv_cache_manager.num_local_layers, + md.kv_cache_manager.num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION), - self.max_num_sequences, - self.kv_cache_manager.max_blocks_per_seq, + md.max_num_sequences, + md.kv_cache_manager.max_blocks_per_seq, ) - self.sliding_block_tables = self.get_empty( - self.cuda_graph_buffers, + self.sliding_block_tables = md.get_empty( + md.cuda_graph_buffers, block_table_shape, cache_name="sliding_block_tables", dtype=torch.int32, @@ -293,12 +292,12 @@ def __post_init__(self): ) compress_block_table_shape = ( - self.max_num_sequences, - self.kv_cache_manager.max_blocks_per_seq, + md.max_num_sequences, + md.kv_cache_manager.max_blocks_per_seq, ) self.compress_block_tables = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, + compress_ratio: md.get_empty( + md.cuda_graph_buffers, compress_block_table_shape, cache_name=f"compress_block_tables_{compress_ratio}", dtype=torch.int32, @@ -312,9 +311,9 @@ def __post_init__(self): # sparse_mla_topk_lens: actual token count per token for each compress_ratio (SWA + compressed) # Shape: [max_num_tokens] per compress_ratio self.sparse_mla_topk_lens = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name=f"sparse_mla_topk_lens_{compress_ratio}", dtype=torch.int32, capture_graph=capture_graph, @@ -325,17 +324,17 @@ def __post_init__(self): # per step next to sparse_mla_topk_lens: SWA rows live per token # (min(pos + 1, window)) and compressed rows live per compress ratio # (sparse_mla_topk_lens[ratio] - window). Shape: [max_num_tokens]. - self.sparse_mla_swa_lens = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.sparse_mla_swa_lens = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="sparse_mla_swa_lens", dtype=torch.int32, capture_graph=capture_graph, ) self.sparse_mla_extra_lens = { - compress_ratio: self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + compress_ratio: md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name=f"sparse_mla_extra_lens_{compress_ratio}", dtype=torch.int32, capture_graph=capture_graph, @@ -345,9 +344,9 @@ def __post_init__(self): } # cached_token_lens_cuda: number of tokens already cached per request - self.cached_token_lens_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + self.cached_token_lens_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), cache_name="cached_token_lens_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -365,19 +364,20 @@ def __post_init__(self): def _init_sparse_offload_state(self) -> None: """Validate fixed offload configuration and allocate persistent metadata.""" + md = self.metadata self.sparse_offload_state: SparseOffloadState | None = None - if not self.sparse_metadata_params.enable_kv_cache_offload: + if not md.sparse_metadata_params.enable_kv_cache_offload: return - manager = self.kv_cache_manager + manager = md.kv_cache_manager # Per-batch beam/draft values start at defaults. Also check the manager's # configured limits before allocating any offload buffers. if ( - self.beam_width != 1 + md.beam_width != 1 or manager.max_beam_width != 1 or self.max_draft_tokens or manager.max_draft_len or manager.max_total_draft_tokens - or self.draft_kv_cache_manager is not None + or md.draft_kv_cache_manager is not None ): raise NotImplementedError( "Sparse offload requires single-beam, non-speculative execution" @@ -392,16 +392,16 @@ def _init_sparse_offload_state(self) -> None: return if self.sparse_mla_topk <= 0: raise ValueError("Sparse offload metadata requires positive top-k") - batch = self.max_num_sequences - max_blocks = self.kv_cache_manager.max_blocks_per_seq + batch = md.max_num_sequences + max_blocks = md.kv_cache_manager.max_blocks_per_seq def allocate(name: str, shape: tuple[int, ...], value: int) -> torch.Tensor: - tensor = self.get_empty( - self.cuda_graph_buffers, + tensor = md.get_empty( + md.cuda_graph_buffers, shape, dtype=torch.int32, cache_name=f"dsv4_sparse_offload_{name}", - capture_graph=self.is_cuda_graph, + capture_graph=md.is_cuda_graph, ) tensor.fill_(value) return tensor @@ -435,13 +435,14 @@ def mla_prepare_fused_q_gen_cu_seqlens(self) -> Optional[torch.Tensor]: into graph-owned storage once per iteration; device-window graph replay therefore updates values without changing the captured address. """ - if not self.is_ragged_verify or self.num_generations <= 0: + md = self.metadata + if not self.is_ragged_verify or md.num_generations <= 0: return None - num_contexts = self.num_contexts - num_generations = self.num_generations + num_contexts = md.num_contexts + num_generations = md.num_generations cu_seqlens = getattr(self, "fused_q_gen_cu_seqlens", None) - seq_lens_cuda = getattr(self, "seq_lens_cuda", None) + seq_lens_cuda = getattr(md, "seq_lens_cuda", None) if cu_seqlens is None or seq_lens_cuda is None or num_generations + 1 > cu_seqlens.shape[0]: return None @@ -462,24 +463,26 @@ def apply_device_ragged_layout( kv_correction: torch.Tensor, ) -> None: """Install device-selected windows and invalidate derived Q prefixes.""" + md = self.metadata super().apply_device_ragged_layout(verify_lens, req_idx, kv_correction) # This call precedes the first model layer. Its fixed-address cumsum is # therefore captured once after the layout update and replayed in order. - self._invalidate_mla_scheduler_buffers() + md._invalidate_mla_scheduler_buffers() def prepare_for_indexer_k_cache(self): """Prepare the shared indexer K-cache decode table for DSA kernels.""" # INDEXER_COMPRESS uses shared page indices, so the generic DSA # indexer path only needs one 2D block table. - self.kv_cache_manager.copy_batch_indexer_compress_block_tables( + md = self.metadata + md.kv_cache_manager.copy_batch_indexer_compress_block_tables( self.host_indexer_k_cache_block_offsets, - self.request_ids, - beam_width=self.beam_width, - num_contexts=self.num_contexts, - num_seqs=self.num_seqs, + md.request_ids, + beam_width=md.beam_width, + num_contexts=md.num_contexts, + num_seqs=md.num_seqs, ) - self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[: self.num_seqs], + self.indexer_k_cache_block_offsets[: md.num_seqs].copy_( + self.host_indexer_k_cache_block_offsets[: md.num_seqs], non_blocking=True, ) # Columns beyond each sequence's allocated indexer blocks contain BAD_PAGE_INDEX (-1). @@ -489,45 +492,47 @@ def prepare_for_indexer_k_cache(self): def prepare_for_block_tables(self): """Prepare block tables for sliding-window and compressed attention.""" - self.kv_cache_manager.copy_batch_sliding_block_tables( + md = self.metadata + md.kv_cache_manager.copy_batch_sliding_block_tables( self.sliding_block_tables, - self.request_ids, - self.num_contexts, - self.num_seqs, + md.request_ids, + md.num_contexts, + md.num_seqs, ) for compress_ratio, compress_block_table in self.compress_block_tables.items(): if ( compress_ratio == DEEPSEEK_V4_SPARSE_RATIO - and self.sparse_metadata_params.enable_kv_cache_offload + and md.sparse_metadata_params.enable_kv_cache_offload ): if self.sparse_offload_state is None: compress_block_table.fill_(-1) else: - self.kv_cache_manager.prepare_sparse_offload( + md.kv_cache_manager.prepare_sparse_offload( self.sparse_offload_state, - self.request_ids, + md.request_ids, compress_block_table, - beam_width=self.beam_width, + beam_width=md.beam_width, ) continue - self.kv_cache_manager.copy_batch_compress_block_tables( + md.kv_cache_manager.copy_batch_compress_block_tables( compress_block_table, - self.request_ids, + md.request_ids, compress_ratio=compress_ratio, - beam_width=self.beam_width, - num_contexts=self.num_contexts, - num_seqs=self.num_seqs, + beam_width=md.beam_width, + num_contexts=md.num_contexts, + num_seqs=md.num_seqs, ) def prepare_for_deepseek_v4_indices(self, token_positions=None): """Prepare SWA/compressed local indices and sparse_mla_topk_lens.""" + md = self.metadata window_size = self.window_size device = self.swa_local_indices_cuda.device if token_positions is None: # Initial prepare() path — build token_positions from CPU data. - num_tokens = self.num_tokens - num_requests = self.num_seqs + num_tokens = md.num_tokens + num_requests = md.num_seqs # cu_seq_lens_cuda must already be populated before this call cu_seq_lens = self.cu_seq_lens_cuda[: num_requests + 1] @@ -640,6 +645,7 @@ def _build_cache_buffer_data_pointers( def _init_cache_buffer_data_pointers(self): # If MTP is enabled, enlarge the compress ratios by max_draft_tokens - 1. + md = self.metadata extend_compress_ratios = self.compress_ratios + [self.compress_ratios[-1]] * ( self.max_draft_tokens - 1 ) @@ -647,15 +653,16 @@ def _init_cache_buffer_data_pointers(self): self.sparse_mla_base_ptrs, self.swa_buffer_ptrs, self.compressed_buffer_ptrs, - ) = self._build_cache_buffer_data_pointers(self.kv_cache_manager, extend_compress_ratios) + ) = self._build_cache_buffer_data_pointers(md.kv_cache_manager, extend_compress_ratios) def _init_draft_sparse_buffers(self): """Initialize sparse buffers for a separate one-model MTP draft cache.""" + md = self.metadata self.draft_sliding_block_tables = None self.draft_sparse_mla_base_ptrs = None self.draft_swa_buffer_ptrs = None - draft_mgr = self.draft_kv_cache_manager + draft_mgr = md.draft_kv_cache_manager from .cache_manager import DeepseekV4CacheManager if not isinstance(draft_mgr, DeepseekV4CacheManager): @@ -673,15 +680,15 @@ def _init_draft_sparse_buffers(self): draft_block_table_shape = ( draft_mgr.num_local_layers, len(DEEPSEEK_V4_SLIDING_ATTENTION), - self.max_num_sequences, + md.max_num_sequences, draft_mgr.max_blocks_per_seq, ) - self.draft_sliding_block_tables = self.get_empty( - self.cuda_graph_buffers, + self.draft_sliding_block_tables = md.get_empty( + md.cuda_graph_buffers, draft_block_table_shape, cache_name="draft_sliding_block_tables", dtype=torch.int32, - capture_graph=self.is_cuda_graph, + capture_graph=md.is_cuda_graph, ) extend_compress_ratios = self.compress_ratios + [draft_ratio] * (self.max_draft_tokens - 1) ( @@ -696,18 +703,20 @@ def _init_draft_sparse_buffers(self): "swa_buffer_ptrs", ) - def prepare_for_draft_forward(self) -> dict | None: - """Repoint sparse fields to the draft buffers for a draft forward.""" + def draft_swap(self) -> dict | None: + """Repoint sparse fields to the draft buffers for a draft forward. + + Replaces the DSA swap: the DeepSeek-V4 indexer buffers stay bound. + """ if self.draft_sliding_block_tables is None: return None - saved_state = {field: getattr(self, field) for field in self._DRAFT_SPARSE_FIELDS} - for field in self._DRAFT_SPARSE_FIELDS: - self.swap_for_draft(field, getattr(self, f"draft_{field}")) - return saved_state + return self.swap_attrs_for_draft( + {field: getattr(self, f"draft_{field}") for field in self._DRAFT_SPARSE_FIELDS} + ) - def restore_after_draft_forward(self, saved_state: dict | None) -> None: - """No-op: restore_draft_swaps() restores the sparse fields, and DSA's hook must not run.""" - return None + def draft_restore(self, saved_state: dict | None) -> None: + if saved_state is not None: + self.restore_swapped_attrs(saved_state) def validate_sparse_offload_batch(self) -> None: """Reject unsupported layouts on the host before any model writes. @@ -717,34 +726,35 @@ def validate_sparse_offload_batch(self) -> None: the per-layer path consumes the resulting phase and device metadata. Fixed configuration is validated during offload metadata initialization. """ - if not self.sparse_metadata_params.enable_kv_cache_offload: + md = self.metadata + if not md.sparse_metadata_params.enable_kv_cache_offload: return state = self.sparse_offload_state if state is not None: state.prepared = False state.is_prefill = False - if self.num_contexts and self.num_generations: + if md.num_contexts and md.num_generations: raise NotImplementedError( "Sparse offload does not support mixed context/decode batches" ) - if self.seq_lens.device.type != "cpu" or self.seq_lens_kv.device.type != "cpu": + if md.seq_lens.device.type != "cpu" or md.seq_lens_kv.device.type != "cpu": raise ValueError("Sparse offload batch validation requires host sequence lengths") - if not torch.equal(self.seq_lens, self.seq_lens_kv): + if not torch.equal(md.seq_lens, md.seq_lens_kv): raise NotImplementedError("Sparse offload requires matching query and new-KV lengths") - if not self.num_contexts: - if not torch.all(self.seq_lens == 1): + if not md.num_contexts: + if not torch.all(md.seq_lens == 1): raise NotImplementedError("Sparse offload decode requires one query per request") return - cached = self.kv_cache_params.num_cached_tokens_per_seq + cached = md.kv_cache_params.num_cached_tokens_per_seq if ( - self.prompt_lens is None - or len(self.prompt_lens) < self.num_contexts - or len(cached) < self.num_contexts - or any(cached[: self.num_contexts]) - or self.seq_lens.tolist() != self.prompt_lens[: self.num_contexts] + md.prompt_lens is None + or len(md.prompt_lens) < md.num_contexts + or len(cached) < md.num_contexts + or any(cached[: md.num_contexts]) + or md.seq_lens.tolist() != md.prompt_lens[: md.num_contexts] or any( - self.kv_cache_manager.kv_cache_map[request_id].history_length != 0 - for request_id in self.request_ids + md.kv_cache_manager.kv_cache_map[request_id].history_length != 0 + for request_id in md.request_ids ) ): raise NotImplementedError( @@ -753,44 +763,46 @@ def validate_sparse_offload_batch(self) -> None: if state is not None: state.is_prefill = True - def prepare(self) -> None: + def pre_derive_lengths(self) -> None: + md = self.metadata self._validate_ragged_verification() - assert self.kv_cache_manager is not None - assert self.request_ids is not None + assert md.kv_cache_manager is not None + assert md.request_ids is not None self.validate_sparse_offload_batch() - self.kv_cache_manager.compute_sliding_block_tables( - self.request_ids, - self.num_contexts, + md.kv_cache_manager.compute_sliding_block_tables( + md.request_ids, + md.num_contexts, ) # Prepare the draft manager's tables before the generic prepare copies # its block offsets, and populate the dedicated DSv4 draft sparse table. - draft_mgr = self.draft_kv_cache_manager + draft_mgr = md.draft_kv_cache_manager if draft_mgr is not None and hasattr(draft_mgr, "compute_sliding_block_tables"): draft_mgr.compute_sliding_block_tables( - self.request_ids, - self.num_contexts, + md.request_ids, + md.num_contexts, ) if self.draft_sliding_block_tables is not None: draft_mgr.copy_batch_sliding_block_tables( self.draft_sliding_block_tables, - self.request_ids, - self.num_contexts, - self.num_seqs, + md.request_ids, + md.num_contexts, + md.num_seqs, ) - TrtllmAttentionMetadata.prepare(self) - - num_requests = self.num_contexts + self.num_generations - num_gen_tokens = self.num_tokens - self.num_ctx_tokens + def prepare(self) -> None: + """Derive the DeepSeek-V4 state; replaces the DSA prepare.""" + md = self.metadata + num_requests = md.num_contexts + md.num_generations + num_gen_tokens = md.num_tokens - md.num_ctx_tokens self.cached_token_lens_cpu[:num_requests] = torch.tensor( - self.kv_cache_params.num_cached_tokens_per_seq[:num_requests], + md.kv_cache_params.num_cached_tokens_per_seq[:num_requests], dtype=torch.int32, ) cached_token_lens = self.cached_token_lens_cpu - kv_lens = cached_token_lens[:num_requests] + self.seq_lens_kv[:num_requests] + kv_lens = cached_token_lens[:num_requests] + md.seq_lens_kv[:num_requests] kv_lens = self._prepare_device_window_kv_lens(kv_lens) @@ -799,7 +811,7 @@ def prepare(self) -> None: ) # Prepare cu_seq_lens early — needed by prepare_for_deepseek_v4_indices - self.cu_seq_lens[1 : num_requests + 1] = self.seq_lens.cumsum(0) + self.cu_seq_lens[1 : num_requests + 1] = md.seq_lens.cumsum(0) self.cu_seq_lens_cuda[: num_requests + 1].copy_( self.cu_seq_lens[: num_requests + 1], non_blocking=True ) @@ -824,14 +836,14 @@ def prepare(self) -> None: # Prepare metadata for indexer (only needed when sparse layers exist) if has_sparse_layers: - DeepseekV4Indexer.prepare(metadata=self) + DeepseekV4Indexer.prepare(metadata=md) # --- Per-ratio metadata --- # 1) CPU-side: compute scalar metadata (num_total_compressed_tokens, etc.) # 2) CUDA-side: fill *_cuda buffers via prepare_compressed_kv_metadata() num_gen_tokens_per_seq = self._sync_gen_tokens_per_seq(num_gen_tokens) - num_contexts = self.num_contexts - num_generations = self.num_generations + num_contexts = md.num_contexts + num_generations = md.num_generations kv_lens_slice = kv_lens[:num_requests] cached_slice = cached_token_lens[:num_requests] @@ -869,9 +881,7 @@ def prepare(self) -> None: self._ctx_output_sizes = ctx_output_sizes # 2) CUDA-side: fill *_cuda buffers on device. - kv_lens_cuda = ( - self.cached_token_lens_cuda[:num_requests] + self._seq_lens_cuda[:num_requests] - ) + kv_lens_cuda = self.cached_token_lens_cuda[:num_requests] + md._seq_lens_cuda[:num_requests] cached_tokens_cuda = self.cached_token_lens_cuda[:num_requests] self.prepare_compressed_kv_metadata(kv_lens_cuda, cached_tokens_cuda, ctx_output_sizes) @@ -894,16 +904,17 @@ def _sync_gen_tokens_per_seq(self, num_gen_tokens: int) -> int: observes device-window rewrites instead of reloading the captured host shape split. """ - if self.is_ragged_verify and self.num_generations > 0: + md = self.metadata + if self.is_ragged_verify and md.num_generations > 0: self.num_gen_tokens_per_seq = 1 + self.max_draft_tokens - num_contexts = self.num_contexts - num_generations = self.num_generations + num_contexts = md.num_contexts + num_generations = md.num_generations gen_new_tokens = self.gen_new_tokens_per_seq_cuda[:num_generations] - gen_new_tokens.copy_(self._seq_lens_cuda[num_contexts : num_contexts + num_generations]) + gen_new_tokens.copy_(md._seq_lens_cuda[num_contexts : num_contexts + num_generations]) self.gen_new_tokens_per_seq = gen_new_tokens else: self.num_gen_tokens_per_seq = ( - num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 + num_gen_tokens // md.num_generations if md.num_generations > 0 else 0 ) self.gen_new_tokens_per_seq = None return self.num_gen_tokens_per_seq @@ -928,9 +939,10 @@ def prepare_compressed_kv_metadata( reuses the cached copy unless the extend_ctx path may have mutated ctx-row kv_lens on device. """ + md = self.metadata batch_size = kv_lens.shape[0] - num_contexts = self.num_contexts - num_generations = self.num_generations + num_contexts = md.num_contexts + num_generations = md.num_generations self._compute_per_ratio_kv_lens( kv_lens, @@ -972,24 +984,25 @@ def prepare_compressed_kv_metadata( gen_output_offsets, ) - def on_update_kv_lens(self): + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: """Recompute kv-lens-dependent DeepSeek-V4 metadata on device.""" - super().on_update_kv_lens() + md = self.metadata + super().on_kv_lens_changed(spec_dec_step) - batch_size = self.num_seqs - num_tokens = self.num_tokens - kv_lens = self.kv_lens_cuda[:batch_size] - seq_lens = self._seq_lens_cuda[:batch_size] + batch_size = md.num_seqs + num_tokens = md.num_tokens + kv_lens = md.kv_lens_cuda[:batch_size] + seq_lens = md._seq_lens_cuda[:batch_size] cached_tokens = kv_lens - seq_lens - num_gen_tokens = num_tokens - self.num_ctx_tokens + num_gen_tokens = num_tokens - md.num_ctx_tokens self._sync_gen_tokens_per_seq(num_gen_tokens) # Reuse prepare()'s host-computed ctx sizes unless the extend_ctx path # (num_chunked_ctx_requests > 0) may have mutated ctx-row kv_lens on # device; every other path only changes gen rows. ctx_output_sizes = ( - self._ctx_output_sizes if getattr(self, "num_chunked_ctx_requests", 0) == 0 else None + self._ctx_output_sizes if getattr(md, "num_chunked_ctx_requests", 0) == 0 else None ) self.prepare_compressed_kv_metadata(kv_lens, cached_tokens, ctx_output_sizes) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py index dcbd878d613d..8f59c0f4a371 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py @@ -29,6 +29,7 @@ from ..hooks import MLASparseHooks, register_mla_sparse_hooks from ..params import SparseBackendForwardArgs from .flash_mla import DeepSeekV4FlashMLA +from .metadata import DeepseekV4Metadata if TYPE_CHECKING: from tensorrt_llm._torch.attention.mla import MLA @@ -660,7 +661,8 @@ def _is_dsv4_q_b_gemm_norm_rope_fusion_enabled( return False cu_q_seqlens = attn_metadata.cu_q_seqlens if cu_q_seqlens is None and num_contexts == 0: - generation_cu_seqlens = getattr(attn_metadata, "cu_seq_lens_cuda", None) + dsv4 = attn_metadata.find_ext(DeepseekV4Metadata) + generation_cu_seqlens = getattr(dsv4, "cu_seq_lens_cuda", None) if generation_cu_seqlens is not None: cu_q_seqlens = generation_cu_seqlens[: num_generations + 1] kv_cache_lengths = getattr(attn_metadata, "kv_lens_cuda_runtime", None) @@ -690,7 +692,8 @@ def _dsv4_q_b_gemm_cu_q_seqlens( if attn_metadata.num_contexts > 0: prep_ctx = getattr(attn_metadata, "mla_prepare_ctx_cu_seqlens", None) return prep_ctx() if prep_ctx is not None else None - generation_cu_seqlens = getattr(attn_metadata, "cu_seq_lens_cuda", None) + dsv4 = attn_metadata.find_ext(DeepseekV4Metadata) + generation_cu_seqlens = getattr(dsv4, "cu_seq_lens_cuda", None) if generation_cu_seqlens is None: return None return generation_cu_seqlens[: attn_metadata.num_generations + 1] @@ -794,8 +797,9 @@ def _fused_q_rope_specs( num_gen_tokens = num_tokens - num_ctx_tokens if num_gen_tokens <= 0: return None, [] - if bool(getattr(attn_metadata, "is_ragged_verify", False)): - prep_gen = getattr(attn_metadata, "mla_prepare_fused_q_gen_cu_seqlens", None) + dsv4 = attn_metadata.find_ext(DeepseekV4Metadata) + if bool(getattr(dsv4, "is_ragged_verify", False)): + prep_gen = getattr(dsv4, "mla_prepare_fused_q_gen_cu_seqlens", None) if prep_gen is None: return None, [] cu_q_seqlens = prep_gen() diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/__init__.py index d3f8152e937c..60deca5192d8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/__init__.py @@ -24,7 +24,7 @@ split_prefill_chunks, transform_local_topk_and_prepare_pool_view, ) -from .metadata import DSAtrtllmAttentionMetadata, build_req_idx_per_token +from .metadata import DSAMetadata, build_req_idx_per_token from .params import DSABackendForwardArgs, DSAMetadataParams, DSAParams __all__ = [ @@ -32,10 +32,10 @@ "DSABackendForwardArgs", "DSACacheManager", "DSACacheManagerV2", + "DSAMetadata", "DSAMetadataParams", "DSAParams", "DSATrtllmAttention", - "DSAtrtllmAttentionMetadata", "Indexer", "IndexerParams", "IndexerPrefillChunkMetadata", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py index 1c14f3421a51..0651e5c09ff6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/backend.py @@ -18,7 +18,7 @@ MLAParams, PositionalEmbeddingParams, ) -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm.models.modeling_utils import QuantConfig from .indexer import ( @@ -26,7 +26,7 @@ transform_local_topk_and_prepare_pool_view, transform_local_topk_and_prepare_pool_view_grouped, ) -from .metadata import DSAtrtllmAttentionMetadata +from .metadata import DSAMetadata from .params import DSAParams ModelConfig = tensorrt_llm.bindings.ModelConfig @@ -68,7 +68,8 @@ def _get_nvfp4_mla_kv_cache_amax() -> float: class DSATrtllmAttention(TrtllmAttention): """TRT-LLM attention layer with DSA sparse indexer for MLA models.""" - Metadata = DSAtrtllmAttentionMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (DSAMetadata,) def __init__( self, @@ -156,15 +157,16 @@ def sparse_attn_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """Run the DSA indexer and convert its local TopK to paged indices.""" + dsa = metadata.ext(DSAMetadata) is_generation = forward_args.attention_input_type == AttentionInputType.generation_only sparse_backend_args = forward_args.sparse_backend_args phase_start = metadata.num_ctx_tokens if is_generation else 0 phase_end = metadata.num_tokens if is_generation else metadata.num_ctx_tokens - shared_topk_indices = metadata.shared_topk_indices + shared_topk_indices = dsa.shared_topk_indices if self.indexer is None: topk_indices = shared_topk_indices[phase_start:phase_end] else: @@ -175,7 +177,7 @@ def sparse_attn_predict( indexer_intermediates, is_generation=is_generation, ) - preserve_mtp_topk = metadata.in_mtp_draft_loop and self.indexer.mtp_index_share + preserve_mtp_topk = dsa.in_mtp_draft_loop and self.indexer.mtp_index_share if shared_topk_indices is not None and not preserve_mtp_topk: shared_topk_indices[ phase_start : phase_start + topk_indices.shape[0], @@ -185,11 +187,11 @@ def sparse_attn_predict( local_layer_idx = self.get_local_layer_idx(metadata) kv_cache_dtype = getattr(metadata.kv_cache_manager, "dtype", None) if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4 and not is_generation: - metadata._ensure_pool_view_cached() + dsa._ensure_pool_view_cached() page_index_scale, layer_offset = ( metadata.kv_cache_manager.get_primary_pool_page_index_params(local_layer_idx) ) - total_kv_tokens = metadata.num_ctx_mla_kv_tokens + total_kv_tokens = dsa.num_ctx_mla_kv_tokens if total_kv_tokens <= 0: raise RuntimeError("NVFP4 DSA context has no active KV tokens") head_dim = metadata.kv_cache_manager.head_dim @@ -204,7 +206,7 @@ def sparse_attn_predict( # stream. Drop its Python reference before allocating the next # layer's scratch so the caching allocator can reuse that storage # in stream order instead of transiently keeping two full buffers. - metadata.nvfp4_mla_context_fp8_scratch = None + dsa.nvfp4_mla_context_fp8_scratch = None scratch = torch.empty( (scratch_capacity, 1, head_dim), dtype=torch.float8_e4m3fn, @@ -215,29 +217,29 @@ def sparse_attn_predict( metadata.host_kv_cache_pool_pointers, metadata.host_kv_cache_pool_mapping, topk_indices, - metadata._cached_req_idx_ctx, - metadata._cached_block_table_ctx, + dsa._cached_req_idx_ctx, + dsa._cached_block_table_ctx, metadata.ctx_kv_indptr[: metadata.num_contexts + 1], scratch, compact_indices, self._nvfp4_mla_kv_scale_quant_orig, local_layer_idx, total_kv_tokens, - metadata._cached_tokens_per_block, - page_index_scale * metadata._cached_tokens_per_block, + dsa._cached_tokens_per_block, + page_index_scale * dsa._cached_tokens_per_block, layer_offset, metadata.kv_cache_manager.mla_kv_cache_residual_dim, - metadata._cached_num_pool_tokens, + dsa._cached_num_pool_tokens, ) - metadata.nvfp4_mla_context_fp8_scratch = scratch + dsa.nvfp4_mla_context_fp8_scratch = scratch forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr = scratch.data_ptr() return compact_indices, None topk_indices_global = self._remap_topk_to_global(topk_indices, metadata, is_generation) if kv_cache_dtype == tensorrt_llm.bindings.DataType.NVFP4: - scratch = metadata.nvfp4_mla_fp8_scratch - compact_indices = metadata.nvfp4_mla_compact_indices + scratch = dsa.nvfp4_mla_fp8_scratch + compact_indices = dsa.nvfp4_mla_compact_indices if scratch is None or compact_indices is None: raise RuntimeError("NVFP4 MLA scratch buffers were not allocated") num_rows = topk_indices_global.shape[0] @@ -257,7 +259,7 @@ def sparse_attn_predict( self._nvfp4_mla_kv_scale_quant_orig, local_layer_idx, metadata.kv_cache_manager.mla_kv_cache_residual_dim, - metadata._cached_num_pool_tokens, + dsa._cached_num_pool_tokens, ) # Static sparse MLA treats indices as offsets from kvPtr. Feed it # compact offsets and the FP8 scratch base; TRTLLMGen is unchanged. @@ -269,7 +271,7 @@ def sparse_attn_predict( def _remap_topk_to_global( self, topk_indices: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, is_generation: bool, ) -> torch.Tensor: """Convert this layer's local top-k to global pool indices. @@ -277,6 +279,7 @@ def _remap_topk_to_global( Uses the cross-layer fan-out (grouped) remap when enabled and applicable; otherwise the per-layer single-op path (unchanged). """ + dsa = metadata.ext(DSAMetadata) local_layer_idx = self.get_local_layer_idx(metadata) # Grouping applies only to the MLA target-verify generation forward. @@ -284,7 +287,7 @@ def _remap_topk_to_global( # reuse a frozen top-k for shared layers that differs from the leader's, # so fan-out is unsafe there; context is excluded to avoid a group-sized # allocation over many tokens). - if _GROUP_REMAP and is_generation and not metadata.in_mtp_draft_loop: + if _GROUP_REMAP and is_generation and not dsa.in_mtp_draft_loop: grouped = self._grouped_remap_topk_to_global( topk_indices, metadata, local_layer_idx, is_generation ) @@ -299,7 +302,7 @@ def _remap_topk_to_global( def _grouped_remap_topk_to_global( self, topk_indices: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, local_layer_idx: int, is_generation: bool, ) -> Optional[torch.Tensor]: @@ -314,7 +317,8 @@ def _grouped_remap_topk_to_global( that equivalence is covered by unit tests (see ``tests/unittest/_torch/attention/sparse/dsa/test_cpp_custom_ops.py``). """ - struct = metadata._ensure_group_remap_struct() + dsa = metadata.ext(DSAMetadata) + struct = dsa._ensure_group_remap_struct() leader_of = struct.get("leader_of") if leader_of is None or local_layer_idx >= len(leader_of): return None @@ -332,7 +336,7 @@ def _grouped_remap_topk_to_global( struct["group_scale"][leader], is_generation, ) - metadata._group_remap_batched[leader] = batched + dsa._group_remap_batched[leader] = batched else: # Shared layer: read the leader's precomputed batch. The batch is # cleared at every step boundary, so a present entry is from this @@ -340,7 +344,7 @@ def _grouped_remap_topk_to_global( # top-k shape, and otherwise fall back to the per-layer path. This # makes any leader/follower ordering or shape mismatch a safe # per-layer fallback rather than a stale/out-of-bounds read. - batched = metadata._group_remap_batched.get(leader, None) + batched = dsa._group_remap_batched.get(leader, None) if ( batched is None or slot >= batched.shape[0] @@ -354,7 +358,7 @@ def sparse_kv_predict( self, q: torch.Tensor, k: Optional[torch.Tensor], - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: """No-op KV prediction; DSA uses indexer-based selection instead.""" @@ -364,16 +368,17 @@ def mla_rope_append_paged_kv_assign_q( self, q: torch.Tensor, latent_cache: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, is_generation: bool = False, **kwargs, ) -> None: """Apply RoPE, append latent cache to paged KV, and assign query for MLA.""" + dsa = metadata.ext(DSAMetadata) if is_generation: - cached_token_indptr = metadata.gen_cached_token_indptr - kv_indptr = metadata.gen_kv_indptr + cached_token_indptr = dsa.gen_cached_token_indptr + kv_indptr = dsa.gen_kv_indptr num_seqs = metadata.num_generations - max_seq_len = metadata.max_gen_seq_len + max_seq_len = dsa.max_gen_seq_len block_offsets = metadata.kv_cache_block_offsets[:, metadata.num_contexts :] else: cached_token_indptr = metadata.ctx_cached_token_indptr diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index a4d3c25b7bce..334ad9e46cd9 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -75,7 +75,19 @@ if TYPE_CHECKING: from tensorrt_llm.mapping import Mapping - from .metadata import DSAtrtllmAttentionMetadata + from ...trtllm import TrtllmAttentionMetadata + from .metadata import DSAMetadata + + +def _dsa_state(metadata: "TrtllmAttentionMetadata") -> "DSAMetadata": + """The DSA extension of ``metadata``. + + Imported on use: the metadata module imports this one. + """ + from .metadata import DSAMetadata + + return metadata.ext(DSAMetadata) + # Optional import: fast-hadamard-transform causes CI build issues (requires wheel+torch pre-installed) try: @@ -293,7 +305,7 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: def transform_local_topk_and_prepare_pool_view( topk_indices: torch.Tensor, - attn_metadata: "DSAtrtllmAttentionMetadata", + attn_metadata: "TrtllmAttentionMetadata", layer_idx: int, is_generation: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -302,36 +314,37 @@ def transform_local_topk_and_prepare_pool_view( Uses cached values from attn_metadata._ensure_pool_view_cached() to avoid redundant Python/CUDA overhead across layers. """ + dsa = _dsa_state(attn_metadata) assert topk_indices.dtype == torch.int32 - attn_metadata._ensure_pool_view_cached() + dsa._ensure_pool_view_cached() page_index_scale, layer_offset = ( attn_metadata.kv_cache_manager.get_primary_pool_page_index_params(layer_idx) ) if is_generation: - block_table = attn_metadata._cached_block_table_gen - req_idx = attn_metadata._cached_req_idx_gen + block_table = dsa._cached_block_table_gen + req_idx = dsa._cached_req_idx_gen else: - block_table = attn_metadata._cached_block_table_ctx - req_idx = attn_metadata._cached_req_idx_ctx + block_table = dsa._cached_block_table_ctx + req_idx = dsa._cached_req_idx_ctx global_indices = torch.ops.trtllm.convert_req_index_to_global( req_idx, block_table, topk_indices, - attn_metadata._cached_tokens_per_block, + dsa._cached_tokens_per_block, topk_indices.shape[1], - page_index_scale * attn_metadata._cached_tokens_per_block, + page_index_scale * dsa._cached_tokens_per_block, layer_offset, ) - return global_indices, attn_metadata._cached_pool_view + return global_indices, dsa._cached_pool_view def transform_local_topk_and_prepare_pool_view_grouped( topk_indices: torch.Tensor, - attn_metadata: "DSAtrtllmAttentionMetadata", + attn_metadata: "TrtllmAttentionMetadata", layer_ids: torch.Tensor, page_index_scale: int, is_generation: bool = False, @@ -351,24 +364,25 @@ def transform_local_topk_and_prepare_pool_view_grouped( is bit-identical to ``transform_local_topk_and_prepare_pool_view`` for the member with offset ``layer_ids[z]``. MLA-only path. """ + dsa = _dsa_state(attn_metadata) assert topk_indices.dtype == torch.int32 - attn_metadata._ensure_pool_view_cached() + dsa._ensure_pool_view_cached() if is_generation: - block_table = attn_metadata._cached_block_table_gen - req_idx = attn_metadata._cached_req_idx_gen + block_table = dsa._cached_block_table_gen + req_idx = dsa._cached_req_idx_gen else: - block_table = attn_metadata._cached_block_table_ctx - req_idx = attn_metadata._cached_req_idx_ctx + block_table = dsa._cached_block_table_ctx + req_idx = dsa._cached_req_idx_ctx global_indices = torch.ops.trtllm.convert_req_index_to_global_grouped( req_idx, block_table, topk_indices, - attn_metadata._cached_tokens_per_block, + dsa._cached_tokens_per_block, topk_indices.shape[1], - page_index_scale * attn_metadata._cached_tokens_per_block, + page_index_scale * dsa._cached_tokens_per_block, layer_ids, ) @@ -686,13 +700,13 @@ def _split_prefill_queries( return count * rank // size, count * (rank + 1) // size, sizes -def _prefill_q_split_threshold(metadata: DSAtrtllmAttentionMetadata) -> int: +def _prefill_q_split_threshold(metadata: TrtllmAttentionMetadata) -> int: """Query-token count from which prefill indexer queries are split across TP ranks.""" sparse_metadata_params = metadata.sparse_metadata_params return sparse_metadata_params.q_split_threshold if sparse_metadata_params is not None else 8192 -def _packed_logits_elem_budget(metadata: DSAtrtllmAttentionMetadata) -> int: +def _packed_logits_elem_budget(metadata: TrtllmAttentionMetadata) -> int: """Packed-logits cap for grouping context requests into one prefill chunk. A single request's logits tile is at most ``min(_INDEXER_MQA_LOGITS_ELEM_BUDGET, @@ -1028,7 +1042,7 @@ def post_load_weights(self) -> None: @staticmethod def prepare_one_prefill_chunk( - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, chunk_specs: List[Tuple[int, int, int, int]], ) -> IndexerPrefillChunkMetadata: """ @@ -1044,9 +1058,10 @@ def prepare_one_prefill_chunk( Note: Cached token counts are derived from metadata.host_ctx_cached_token_indptr """ - device = metadata.cu_seqlen_ks.device + dsa = _dsa_state(metadata) + device = dsa.cu_seqlen_ks.device compress_ratio = _effective_compress_ratio_divisor( - _select_indexer_compress_ratio(metadata.compress_ratios) + _select_indexer_compress_ratio(dsa.compress_ratios) ) if len(chunk_specs) == 1: # Single request or intra-request Q-block @@ -1158,25 +1173,26 @@ def prepare_one_prefill_chunk( ) @staticmethod - def build_indexer_params(metadata: DSAtrtllmAttentionMetadata) -> Optional[IndexerParams]: + def build_indexer_params(metadata: TrtllmAttentionMetadata) -> Optional[IndexerParams]: + dsa = _dsa_state(metadata) kv_cache_manager = metadata.kv_cache_manager if kv_cache_manager is None or not hasattr(kv_cache_manager, "index_head_dim"): return None - head_dim = metadata.indexer_head_dim + head_dim = dsa.indexer_head_dim data_bytes_per_token = ( head_dim // 2 if getattr(kv_cache_manager, "use_fp4", False) else head_dim ) compress_ratio = _effective_compress_ratio_divisor( - _select_indexer_compress_ratio(metadata.compress_ratios) + _select_indexer_compress_ratio(dsa.compress_ratios) ) return IndexerParams( num_contexts=metadata.num_contexts, num_generations=metadata.num_generations, num_ctx_tokens=metadata.num_ctx_tokens, head_dim=head_dim, - quant_block_size=metadata.indexer_quant_block_size, - tokens_per_block=metadata._tokens_per_block, + quant_block_size=dsa.indexer_quant_block_size, + tokens_per_block=dsa._tokens_per_block, compress_ratio=compress_ratio, request_ids=metadata.request_ids, num_past_tokens=metadata.kv_cache_params.num_cached_tokens_per_seq, @@ -1186,7 +1202,7 @@ def build_indexer_params(metadata: DSAtrtllmAttentionMetadata) -> Optional[Index @staticmethod def recompute_slot_mappings( - metadata: DSAtrtllmAttentionMetadata, indexer_params: Optional[IndexerParams] = None + metadata: TrtllmAttentionMetadata, indexer_params: Optional[IndexerParams] = None ): """Recompute slot_mapping_fp8/scale from current block offsets. @@ -1195,6 +1211,7 @@ def recompute_slot_mappings( isolation after a caller has swapped the active KV cache manager and block-offset buffers, such as during draft KV-cache replay. """ + dsa = _dsa_state(metadata) if indexer_params is None: indexer_params = Indexer.build_indexer_params(metadata) if indexer_params is None: @@ -1222,14 +1239,14 @@ def recompute_slot_mappings( # Block indices/pos for all kv tokens in the batch block_indices_in_seq = global_positions // tokens_per_block - max_blocks = metadata.host_indexer_k_cache_block_offsets.shape[1] + max_blocks = dsa.host_indexer_k_cache_block_offsets.shape[1] assert (block_indices_in_seq < max_blocks).all(), ( f"Block index out of bounds: max={max_blocks}, got indices up to {block_indices_in_seq.max().item()}" ) fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( global_positions, - metadata.host_indexer_k_cache_block_offsets, + dsa.host_indexer_k_cache_block_offsets, req_indices, head_dim, tokens_per_block, @@ -1237,25 +1254,26 @@ def recompute_slot_mappings( data_bytes_per_token=data_bytes_per_token, ) - metadata.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices - metadata.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices + dsa.host_slot_mapping_fp8[:total_new_kv_tokens] = fp8_flat_indices + dsa.host_slot_mapping_scale[:total_new_kv_tokens] = scale_flat_indices - metadata.slot_mapping_fp8[:total_new_kv_tokens].copy_( - metadata.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True + dsa.slot_mapping_fp8[:total_new_kv_tokens].copy_( + dsa.host_slot_mapping_fp8[:total_new_kv_tokens], non_blocking=True ) - metadata.slot_mapping_scale[:total_new_kv_tokens].copy_( - metadata.host_slot_mapping_scale[:total_new_kv_tokens], non_blocking=True + dsa.slot_mapping_scale[:total_new_kv_tokens].copy_( + dsa.host_slot_mapping_scale[:total_new_kv_tokens], non_blocking=True ) @staticmethod def recompute_context_kv_gather_mappings( - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, indexer_params: Optional[IndexerParams] = None, ) -> None: """Recompute context KV-cache gather mappings for the active cache manager.""" + dsa = _dsa_state(metadata) if not metadata.enable_context_mla_with_cached_kv: - metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 - metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + dsa.slot_mapping_fp8_fullkv = dsa.slot_mapping_fp8 + dsa.slot_mapping_scale_fullkv = dsa.slot_mapping_scale return if indexer_params is None: @@ -1270,11 +1288,11 @@ def recompute_context_kv_gather_mappings( # while capturing generation CUDA graphs. return - if metadata.skip_indexer_for_ctx_reqs: + if dsa.skip_indexer_for_ctx_reqs: # The dense short-context path does not gather from the indexer K # cache, so full-KV mappings would be allocated but never consumed. - metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 - metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + dsa.slot_mapping_fp8_fullkv = dsa.slot_mapping_fp8 + dsa.slot_mapping_scale_fullkv = dsa.slot_mapping_scale return cached_kv_tokens = indexer_params.cached_kv_tokens[:num_contexts] @@ -1282,8 +1300,8 @@ def recompute_context_kv_gather_mappings( # Without compression or cached tokens, the regular mappings cover # every context KV token in the same order and remain equivalent # when on_update_kv_lens() refreshes them in place. - metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 - metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + dsa.slot_mapping_fp8_fullkv = dsa.slot_mapping_fp8 + dsa.slot_mapping_scale_fullkv = dsa.slot_mapping_scale return total_kv_per_request = indexer_params.kv_lens[:num_contexts] @@ -1291,8 +1309,8 @@ def recompute_context_kv_gather_mappings( if total_kv_len == 0: # No chunk can gather indexer KV, so avoid zero-length allocation # and H2D work for short compressed contexts. - metadata.slot_mapping_fp8_fullkv = metadata.slot_mapping_fp8 - metadata.slot_mapping_scale_fullkv = metadata.slot_mapping_scale + dsa.slot_mapping_fp8_fullkv = dsa.slot_mapping_fp8 + dsa.slot_mapping_scale_fullkv = dsa.slot_mapping_scale return host_slot_mapping_fp8_fullkv = torch.empty( @@ -1314,7 +1332,7 @@ def recompute_context_kv_gather_mappings( fp8_flat_indices, scale_flat_indices = _compute_slot_mappings( kv_positions, - metadata.host_indexer_k_cache_block_offsets, + dsa.host_indexer_k_cache_block_offsets, req_indices, indexer_params.head_dim, indexer_params.tokens_per_block, @@ -1324,12 +1342,12 @@ def recompute_context_kv_gather_mappings( host_slot_mapping_fp8_fullkv.copy_(fp8_flat_indices) host_slot_mapping_scale_fullkv.copy_(scale_flat_indices) - metadata.slot_mapping_fp8_fullkv = host_slot_mapping_fp8_fullkv.cuda(non_blocking=True) - metadata.slot_mapping_scale_fullkv = host_slot_mapping_scale_fullkv.cuda(non_blocking=True) + dsa.slot_mapping_fp8_fullkv = host_slot_mapping_fp8_fullkv.cuda(non_blocking=True) + dsa.slot_mapping_scale_fullkv = host_slot_mapping_scale_fullkv.cuda(non_blocking=True) @staticmethod def prepare_for_update_k_cache( - metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams + metadata: TrtllmAttentionMetadata, indexer_params: IndexerParams ): """ Prepare indexer for the update_k_cache stage. @@ -1341,14 +1359,15 @@ def prepare_for_update_k_cache( @staticmethod def prepare_for_chunked_prefill( - metadata: DSAtrtllmAttentionMetadata, indexer_params: IndexerParams + metadata: TrtllmAttentionMetadata, indexer_params: IndexerParams ): """ Prepare indexer for the chunked prefill. """ + dsa = _dsa_state(metadata) num_contexts = indexer_params.num_contexts seq_lens = indexer_params.seq_lens - metadata.indexer_prefill_q_split = None + dsa.indexer_prefill_q_split = None # When MLA chunked prefill is active, it already bounds the new tokens per # request, so requests are not split by length here. From @@ -1362,7 +1381,7 @@ def prepare_for_chunked_prefill( ) if has_mla_chunked_prefill: compress_ratio = _effective_compress_ratio_divisor( - _select_indexer_compress_ratio(metadata.compress_ratios) + _select_indexer_compress_ratio(dsa.compress_ratios) ) q_lens = seq_lens[:num_contexts].tolist() cached_indptr = metadata.host_ctx_cached_token_indptr[: num_contexts + 1].tolist() @@ -1374,8 +1393,8 @@ def prepare_for_chunked_prefill( sum(q_lens), metadata.mapping, _prefill_q_split_threshold(metadata) ) if q_sizes is not None: - metadata.indexer_prefill_q_split = (q_begin, q_end, q_sizes) - metadata.indexer_prefill_chunks = [ + dsa.indexer_prefill_q_split = (q_begin, q_end, q_sizes) + dsa.indexer_prefill_chunks = [ Indexer.prepare_one_prefill_chunk(metadata, chunk_specs) for chunk_specs in plan_mla_chunked_prefill_chunks( q_lens, kv_lens, _packed_logits_elem_budget(metadata), q_begin, q_end @@ -1391,7 +1410,7 @@ def prepare_for_chunked_prefill( int(indexer_params.kv_lens[:num_contexts].max().item()) if num_contexts > 0 else 0 ) effective_chunk_size = select_indexer_chunk_size( - metadata.indexer_max_chunk_size, max_k_compressed + dsa.indexer_max_chunk_size, max_k_compressed ) chunk_groups = split_prefill_chunks( seq_lens[:num_contexts], @@ -1400,7 +1419,7 @@ def prepare_for_chunked_prefill( ) if len(chunk_groups) > 1 or metadata.enable_context_mla_with_cached_kv: - metadata.indexer_prefill_chunks = [ + dsa.indexer_prefill_chunks = [ Indexer.prepare_one_prefill_chunk( metadata, chunk_specs, @@ -1408,53 +1427,54 @@ def prepare_for_chunked_prefill( for chunk_specs in chunk_groups ] else: - metadata.indexer_prefill_chunks = None + dsa.indexer_prefill_chunks = None # Chunked prefill and KV-cache reuse require gather mappings that cover # all cached and newly appended indexer tokens. Indexer.recompute_context_kv_gather_mappings(metadata, indexer_params) @staticmethod - def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): + def prepare_scheduler_metadata(metadata: TrtllmAttentionMetadata): """ Prepare scheduler metadata for the DeepGEMM decode MQA kernel. """ + dsa = _dsa_state(metadata) num_contexts = metadata.num_contexts num_generations = metadata.num_generations - gen_seq_lens = metadata.get_indexer_kv_lens( + gen_seq_lens = dsa.get_indexer_kv_lens( metadata.kv_lens_cuda_runtime[num_contexts : num_contexts + num_generations] ) - metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens - if not metadata.use_expanded_buffers_for_mtp: - scheduler_context_lens = metadata.update_indexer_kv_lens_2d(gen_seq_lens) + dsa.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens + if not dsa.use_expanded_buffers_for_mtp: + scheduler_context_lens = dsa.update_indexer_kv_lens_2d(gen_seq_lens) scheduler_metadata_buffer = get_paged_mqa_logits_metadata( - gen_seq_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, metadata.num_sms + gen_seq_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, dsa.num_sms ) - metadata.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True) - if metadata.max_draft_tokens > 0: + dsa.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True) + if dsa.max_draft_tokens > 0: scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( scheduler_context_lens.contiguous(), _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms, + dsa.num_sms, ) - metadata.scheduler_metadata_buffer_full_next_n.copy_( + dsa.scheduler_metadata_buffer_full_next_n.copy_( scheduler_metadata_buffer_full_next_n, non_blocking=True ) else: # Expand schedule metadata buffer (only generation). The DeepGEMM # API requires 2D; each expanded token becomes a (1,) row. - num_tokens = sum(metadata.gen_token_repeat_list()) + num_tokens = sum(dsa.gen_token_repeat_list()) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( - metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), + dsa.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, - metadata.num_sms, + dsa.num_sms, ) - metadata.scheduler_metadata_buffer_expanded.copy_( + dsa.scheduler_metadata_buffer_expanded.copy_( scheduler_metadata_buffer_expanded, non_blocking=True ) @staticmethod - def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: + def prepare(metadata: TrtllmAttentionMetadata) -> None: """ Prepare indexer for the forward pass. This should be called during metadata.prepare() stage. @@ -1468,6 +1488,7 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: # This can happen when the metadata is being used with a draft KV cache manager # during MTP speculative decoding, which uses a regular KVCacheManager instead # of DSACacheManager. + dsa = _dsa_state(metadata) indexer_params = Indexer.build_indexer_params(metadata) if indexer_params is None: return @@ -1477,10 +1498,10 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: num_ctx_tokens = metadata.num_ctx_tokens seq_lens = metadata.seq_lens compress_ratio = _effective_compress_ratio_divisor( - _select_indexer_compress_ratio(metadata.compress_ratios) + _select_indexer_compress_ratio(dsa.compress_ratios) ) # Store compressed KV token count for context requests - metadata.num_ctx_kv_tokens = indexer_params.new_kv_tokens[:num_contexts].sum().item() + dsa.num_ctx_kv_tokens = indexer_params.new_kv_tokens[:num_contexts].sum().item() # Prepare for update_k_cache Indexer.prepare_for_update_k_cache(metadata, indexer_params) @@ -1502,10 +1523,10 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: compress_ratio, ) - metadata.cu_seqlen_ks[:num_ctx_tokens].copy_( + dsa.cu_seqlen_ks[:num_ctx_tokens].copy_( maybe_pin_memory(host_cu_seqlen_ks), non_blocking=True ) - metadata.cu_seqlen_ke[:num_ctx_tokens].copy_( + dsa.cu_seqlen_ke[:num_ctx_tokens].copy_( maybe_pin_memory(host_cu_seqlen_ke), non_blocking=True ) Indexer.prepare_for_chunked_prefill(metadata, indexer_params) @@ -1517,7 +1538,7 @@ def prepare(metadata: DSAtrtllmAttentionMetadata) -> None: Indexer.prepare_scheduler_metadata(metadata) def _update_k_cache( - self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: DSAtrtllmAttentionMetadata + self, k_fp8: torch.Tensor, k_scale: torch.Tensor, metadata: TrtllmAttentionMetadata ) -> None: """ Insert/append k values and scales into the indexer k cache using pre-computed slot mappings. @@ -1527,7 +1548,8 @@ def _update_k_cache( k_fp8: FP8 quantized k tensor, shape [total_tokens, head_dim] k_scale: Scaling factors, shape [total_tokens, head_dim // quant_block_size] """ - if metadata.kv_cache_manager is None or metadata.slot_mapping_fp8 is None: + dsa = _dsa_state(metadata) + if metadata.kv_cache_manager is None or dsa.slot_mapping_fp8 is None: return k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) @@ -1544,14 +1566,14 @@ def _update_k_cache( k_fp8, k_scale, k_cache, - metadata.slot_mapping_fp8, - metadata.slot_mapping_scale, + dsa.slot_mapping_fp8, + dsa.slot_mapping_scale, num_tokens, ) def _gather_k_cache_for_chunk( self, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, chunk: IndexerPrefillChunkMetadata, ) -> Tuple[torch.Tensor, torch.Tensor]: """ @@ -1568,7 +1590,8 @@ def _gather_k_cache_for_chunk( k_fp8: FP8 quantized k tensor, shape [num_k_tokens, head_dim] k_scale: Scaling factors, shape [num_k_tokens, 1] """ - assert metadata.slot_mapping_fp8_fullkv is not None, ( + dsa = _dsa_state(metadata) + assert dsa.slot_mapping_fp8_fullkv is not None, ( "_gather_k_cache_for_chunk requires extended slot mappings (only available with cached tokens)" ) @@ -1583,8 +1606,8 @@ def _gather_k_cache_for_chunk( k_token_end = chunk.k_token_end num_k_tokens = k_token_end - k_token_start - slot_mapping_fp8_chunk = metadata.slot_mapping_fp8_fullkv[k_token_start:k_token_end] - slot_mapping_scale_chunk = metadata.slot_mapping_scale_fullkv[k_token_start:k_token_end] + slot_mapping_fp8_chunk = dsa.slot_mapping_fp8_fullkv[k_token_start:k_token_end] + slot_mapping_scale_chunk = dsa.slot_mapping_scale_fullkv[k_token_start:k_token_end] # Vectorized gather using pre-computed slot mappings # Gather FP8 data @@ -1779,7 +1802,7 @@ def _call_paged_mqa_logits( def sparse_attn_indexer( self, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, hidden_states: torch.Tensor, q_fp8: torch.Tensor, k_fp8: torch.Tensor, @@ -1799,6 +1822,7 @@ def sparse_attn_indexer( # 128 (FP8 path) or 32 (MXFP4 path); both round-trip to the same # 4-byte scale word per token at index_head_dim=128, so the slot # mapping arithmetic is unchanged in either mode. + dsa = _dsa_state(metadata) assert metadata.kv_cache_manager is None or metadata.kv_cache_manager.quant_block_size in ( 32, 128, @@ -1814,9 +1838,9 @@ def sparse_attn_indexer( num_gen_tokens = num_tokens - num_ctx_tokens gvr_prior_indices = None if self.top_k.needs_gvr_prior: - assert metadata.gvr_prior_indices is not None + assert dsa.gvr_prior_indices is not None local_layer = metadata.kv_cache_manager.layer_offsets[self.layer_idx] - gvr_prior_indices = metadata.gvr_prior_indices[local_layer] + gvr_prior_indices = dsa.gvr_prior_indices[local_layer] if is_generation is None: has_prefill = num_contexts > 0 has_decode = num_generations > 0 @@ -1843,13 +1867,13 @@ def sparse_attn_indexer( capture_graph=metadata.is_cuda_graph, ) - if has_prefill and not metadata.skip_indexer_for_ctx_reqs: + if has_prefill and not dsa.skip_indexer_for_ctx_reqs: # Use chunked prefill to reduce memory footprint - if metadata.indexer_prefill_chunks is not None: + if dsa.indexer_prefill_chunks is not None: q_split_threshold = _prefill_q_split_threshold(metadata) # Chunks planned for this rank's share of the queries only are # never split again; their top-k rows are gathered after the loop. - batch_q_split = metadata.indexer_prefill_q_split + batch_q_split = dsa.indexer_prefill_q_split k_cache_4d = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) # FP4 packs two codes per byte so the gathered row holds half # as many bytes as in the FP8 path. The scale (4 bytes) is the @@ -1857,7 +1881,7 @@ def sparse_attn_indexer( # into one int32 to match FP8's float32 scale width. gather_head_dim = self.head_dim // 2 if self.use_fp4 else self.head_dim - for chunk in metadata.indexer_prefill_chunks: + for chunk in dsa.indexer_prefill_chunks: # Skip chunks with no compressed KV tokens (e.g., warmup # sequences shorter than compress_ratio produce zero KV). if chunk.k_token_start >= chunk.k_token_end: @@ -1866,8 +1890,8 @@ def sparse_attn_indexer( num_k_tokens = chunk.k_token_end - chunk.k_token_start chunk_k_fp8, chunk_k_scale = torch.ops.trtllm.indexer_k_cache_gather_op( k_cache_4d, - metadata.slot_mapping_fp8_fullkv, - metadata.slot_mapping_scale_fullkv, + dsa.slot_mapping_fp8_fullkv, + dsa.slot_mapping_scale_fullkv, chunk.k_token_start, num_k_tokens, gather_head_dim, @@ -1935,14 +1959,14 @@ def sparse_attn_indexer( dim=0, sizes=q_sizes, ) - elif metadata.num_ctx_kv_tokens == 0: + elif dsa.num_ctx_kv_tokens == 0: # No compressed KV tokens — fill with -1 (no valid indices) topk_indices_buffer[:num_ctx_tokens, :].fill_(-1) else: # Fallback: single-pass indexer prefill (TODO: remove this once chunked prefill is fully tested) - num_ctx_kv_tokens = metadata.num_ctx_kv_tokens - cu_seqlen_ks = metadata.cu_seqlen_ks[:num_ctx_tokens] - cu_seqlen_ke = metadata.cu_seqlen_ke[:num_ctx_tokens] + num_ctx_kv_tokens = dsa.num_ctx_kv_tokens + cu_seqlen_ks = dsa.cu_seqlen_ks[:num_ctx_tokens] + cu_seqlen_ke = dsa.cu_seqlen_ke[:num_ctx_tokens] ctx_q_scale = q_scale[:num_ctx_tokens, ...] if self.use_fp4 else None self.forward_prepared( @@ -1955,11 +1979,9 @@ def sparse_attn_indexer( topk_indices_buffer[:num_ctx_tokens, :], ctx_q_scale, ) - elif has_prefill and metadata.skip_indexer_for_ctx_reqs: + elif has_prefill and dsa.skip_indexer_for_ctx_reqs: # Fill topk_indices_buffer with pre-defined dense topk indices - topk_indices_buffer[:num_ctx_tokens, :] = metadata.topk_indices_buffer[ - :num_ctx_tokens, : - ] + topk_indices_buffer[:num_ctx_tokens, :] = dsa.topk_indices_buffer[:num_ctx_tokens, :] # Dense skip outputs remain valid priors if a later step enters GVR. if has_prefill: @@ -1981,19 +2003,19 @@ def sparse_attn_indexer( reuse_topk = ( self.mtp_index_share - and metadata.in_mtp_draft_loop - and metadata.indexer_skip_topk - and metadata.shared_topk_indices is not None + and dsa.in_mtp_draft_loop + and dsa.indexer_skip_topk + and dsa.shared_topk_indices is not None ) - if has_decode and not metadata.skip_indexer_for_gen_reqs and reuse_topk: + if has_decode and not dsa.skip_indexer_for_gen_reqs and reuse_topk: topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = ( - metadata.shared_topk_indices[:num_generations, :] + dsa.shared_topk_indices[:num_generations, :] ) - elif has_decode and not metadata.skip_indexer_for_gen_reqs: + elif has_decode and not dsa.skip_indexer_for_gen_reqs: # Get decode lengths per request (from seq_lens) for validation gen_seq_lens = metadata.seq_lens[num_contexts : num_contexts + num_generations] - if not metadata.is_ragged_verify: + if not dsa.is_ragged_verify: max_decode_len = gen_seq_lens.max().item() min_decode_len = gen_seq_lens.min().item() assert max_decode_len == min_decode_len, ( @@ -2007,17 +2029,17 @@ def sparse_attn_indexer( # Because fp8_paged_mqa_logits can only support next_n == 1/2/4 on sm100, and # next_n == 1/2 on sm90, for other next_n, we need to flatten the q_decode tensor # and expand the corresponding metadata. - if metadata.is_ragged_verify: + if dsa.is_ragged_verify: take_strided_path = False else: - take_strided_path = not metadata.use_expanded_buffers_for_mtp or next_n == 1 + take_strided_path = not dsa.use_expanded_buffers_for_mtp or next_n == 1 if take_strided_path: q_decode = q_decode.view(num_generations, -1, *q_fp8.shape[1:]) # 2D context_lens slice from the pre-allocated buffer; matches # q_decode's (batch, next_n) layout required by the new # DeepGEMM paged MQA logits API. - context_lens = metadata.kv_lens_cuda_2d[:num_generations, :next_n].contiguous() - block_table = metadata.indexer_k_cache_block_offsets[ + context_lens = dsa.kv_lens_cuda_2d[:num_generations, :next_n].contiguous() + block_table = dsa.indexer_k_cache_block_offsets[ num_contexts : num_contexts + num_generations ] # The 2D-context_lens metadata kernel encodes next_n into the @@ -2028,21 +2050,21 @@ def sparse_attn_indexer( # uses its own schedule buffer (built with num_next_n_atoms=1 # via a (num_gen, 1) input shape) and overrides this below. if next_n == 1: - scheduler_metadata_buffer = metadata.scheduler_metadata_buffer + scheduler_metadata_buffer = dsa.scheduler_metadata_buffer else: - scheduler_metadata_buffer = metadata.scheduler_metadata_buffer_full_next_n + scheduler_metadata_buffer = dsa.scheduler_metadata_buffer_full_next_n else: q_decode = q_decode.view(-1, 1, *q_fp8.shape[1:]) num_tokens = q_decode.shape[0] # New API requires 2D; each expanded token becomes a (1,) row. - context_lens = metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1) - block_table = metadata.block_table_expanded[:num_tokens] - scheduler_metadata_buffer = metadata.scheduler_metadata_buffer_expanded + context_lens = dsa.kv_lens_expanded_cuda[:num_tokens].view(-1, 1) + block_table = dsa.block_table_expanded[:num_tokens] + scheduler_metadata_buffer = dsa.scheduler_metadata_buffer_expanded - if metadata.is_ragged_verify: - assert num_gen_tokens == sum(metadata.ragged_verify_lens), ( + if dsa.is_ragged_verify: + assert num_gen_tokens == sum(dsa.ragged_verify_lens), ( f"ragged gen tokens {num_gen_tokens} != " - f"sum(verify_lens) {sum(metadata.ragged_verify_lens)}" + f"sum(verify_lens) {sum(dsa.ragged_verify_lens)}" ) else: assert num_gen_tokens == batch_size * next_n @@ -2052,9 +2074,9 @@ def sparse_attn_indexer( # with prepared decode metadata. # [num_blocks, tokens_per_block, 1, head_dim + scale_size] k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) - indexer_max_seq_len = metadata.get_indexer_max_seq_len() + indexer_max_seq_len = dsa.get_indexer_max_seq_len() - if self.use_cute_dsl_paged_mqa_logits and not metadata.is_ragged_verify: + if self.use_cute_dsl_paged_mqa_logits and not dsa.is_ragged_verify: # DSL kernel design: 1 atom per q (atom = real next_n positions), # kNumNextNAtoms = 1 for any real next_n. The matching schedule # is `scheduler_metadata_buffer` — built in `Indexer.prepare()` @@ -2065,7 +2087,7 @@ def sparse_attn_indexer( # the same KV length on this path (kv_lens_cuda_2d broadcasts), # so passing the 1D contiguous kv_lens slice for context_lens # avoids materializing a 2D contiguous tensor per call. - dsl_context_lens = metadata.gen_indexer_kv_lens_cuda_runtime + dsl_context_lens = dsa.gen_indexer_kv_lens_cuda_runtime assert dsl_context_lens is not None # Wave-aware atom-split: the picker in `_pick_dsl_expand` caches # (factor, atom) on metadata with invariant @@ -2078,19 +2100,16 @@ def sparse_attn_indexer( # tokens; gate here so i≥1 draft calls fall back to the # kernel-native next_n=1 path. dsl_atom_split = ( - metadata.dsl_expand_factor > 1 - and next_n == metadata.dsl_expand_factor * metadata.dsl_atom + dsa.dsl_expand_factor > 1 and next_n == dsa.dsl_expand_factor * dsa.dsl_atom ) # In-kernel uneven split (next_n with no exact atom divisor): # native Q / context_lens / block_table, schedule over the # staggered per-atom rows. Same full-window gate as above. - dsl_kernel_split = ( - metadata.dsl_kernel_atoms > 1 and next_n == metadata.dsl_kernel_next_n - ) + dsl_kernel_split = dsa.dsl_kernel_atoms > 1 and next_n == dsa.dsl_kernel_next_n dsl_split_kwargs: dict = {} if dsl_kernel_split: dsl_split_kwargs = dict( - num_next_n_atoms=metadata.dsl_kernel_atoms, atom=metadata.dsl_atom + num_next_n_atoms=dsa.dsl_kernel_atoms, atom=dsa.dsl_atom ) gvr_emit_kwargs: dict = {} # emitting for a step the Top-K cannot consume only churns @@ -2104,7 +2123,7 @@ def sparse_attn_indexer( and num_gen_tokens <= 256 # ext tiers are single-CTA/sort-path only; row reordering # routes the Top-K through order_row, which excludes them - and metadata.kv_lens_row_reorder is None + and dsa.kv_lens_row_reorder is None ): gvr_emit_kwargs = self.top_k.prepare_gvr_emission( num_generations, @@ -2124,18 +2143,18 @@ def sparse_attn_indexer( ) dsl_q = q_decode.view(torch.uint8) dsl_block_table = block_table - dsl_schedule_meta = metadata.scheduler_metadata_buffer + dsl_schedule_meta = dsa.scheduler_metadata_buffer if dsl_atom_split: - factor = metadata.dsl_expand_factor - eff_next_n = metadata.dsl_atom + factor = dsa.dsl_expand_factor + eff_next_n = dsa.dsl_atom exp_B = num_generations * factor dsl_q = dsl_q.reshape(exp_B, eff_next_n, self.n_heads, self.head_dim // 2) decode_q_scale = decode_q_scale.reshape(exp_B, eff_next_n, self.n_heads) - dsl_context_lens = metadata.kv_lens_expanded_cuda[:exp_B] - dsl_block_table = metadata.block_table_expanded[:exp_B] - dsl_schedule_meta = metadata.scheduler_metadata_buffer_expanded + dsl_context_lens = dsa.kv_lens_expanded_cuda[:exp_B] + dsl_block_table = dsa.block_table_expanded[:exp_B] + dsl_schedule_meta = dsa.scheduler_metadata_buffer_expanded elif dsl_kernel_split: - dsl_schedule_meta = metadata.scheduler_metadata_buffer_expanded + dsl_schedule_meta = dsa.scheduler_metadata_buffer_expanded logits_decode = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits( dsl_q, @@ -2158,17 +2177,17 @@ def sparse_attn_indexer( dsl_q = q_decode fp8_ctx_lens = dsl_context_lens fp8_block_table = block_table - fp8_schedule_meta = metadata.scheduler_metadata_buffer + fp8_schedule_meta = dsa.scheduler_metadata_buffer if dsl_atom_split: - factor = metadata.dsl_expand_factor - atom = metadata.dsl_atom + factor = dsa.dsl_expand_factor + atom = dsa.dsl_atom exp_B = num_generations * factor dsl_q = q_decode.reshape(exp_B, atom, self.n_heads, self.head_dim) - fp8_ctx_lens = metadata.kv_lens_expanded_cuda[:exp_B] - fp8_block_table = metadata.block_table_expanded[:exp_B] - fp8_schedule_meta = metadata.scheduler_metadata_buffer_expanded + fp8_ctx_lens = dsa.kv_lens_expanded_cuda[:exp_B] + fp8_block_table = dsa.block_table_expanded[:exp_B] + fp8_schedule_meta = dsa.scheduler_metadata_buffer_expanded elif dsl_kernel_split: - fp8_schedule_meta = metadata.scheduler_metadata_buffer_expanded + fp8_schedule_meta = dsa.scheduler_metadata_buffer_expanded logits_decode = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( dsl_q, k_cache, @@ -2208,47 +2227,45 @@ def sparse_attn_indexer( gen_kv_lens_cuda = metadata.kv_lens_cuda_runtime[ num_contexts : num_contexts + num_generations ] - scan_lengths = metadata.gen_indexer_kv_lens_cuda_runtime + scan_lengths = dsa.gen_indexer_kv_lens_cuda_runtime assert scan_lengths is not None gvr_ext_kwargs = ( { "gvr_prior_indices": gvr_prior_indices[:num_generations], - "gvr_row_order": metadata.kv_lens_row_reorder, + "gvr_row_order": dsa.kv_lens_row_reorder, } if gvr_prior_indices is not None else None ) - assert metadata.radix_aux_indices is not None - assert metadata.radix_aux_logits is not None + assert dsa.radix_aux_indices is not None + assert dsa.radix_aux_logits is not None self.top_k( logits_decode, topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :], is_prefill=False, sequence_lengths=gen_kv_lens_cuda, scan_lengths=scan_lengths, - row_kv_lens=metadata.ragged_row_kv_lens(num_gen_tokens), - next_n=1 if metadata.is_ragged_verify else next_n, + row_kv_lens=dsa.ragged_row_kv_lens(num_gen_tokens), + next_n=1 if dsa.is_ragged_verify else next_n, max_seq_len=indexer_max_seq_len, gvr_ext_kwargs=gvr_ext_kwargs, - radix_aux_indices=metadata.radix_aux_indices, - radix_aux_logits=metadata.radix_aux_logits, + radix_aux_indices=dsa.radix_aux_indices, + radix_aux_logits=dsa.radix_aux_logits, ) - elif has_decode and metadata.skip_indexer_for_gen_reqs: + elif has_decode and dsa.skip_indexer_for_gen_reqs: # Fill topk_indices_buffer with pre-defined dense topk indices topk_indices_buffer[token_offset : token_offset + num_gen_tokens, :] = ( - metadata.topk_indices_buffer[num_ctx_tokens:num_tokens, :] + dsa.topk_indices_buffer[num_ctx_tokens:num_tokens, :] ) # Keep the GVR prior current for computed, reused, and dense-skip TopK. if gvr_prior_indices is not None and has_decode: decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] - if metadata.is_ragged_verify: + if dsa.is_ragged_verify: last_rows = ( - torch.cumsum(metadata.gen_token_repeats_cuda[:num_generations], dim=0).to( - torch.long - ) + torch.cumsum(dsa.gen_token_repeats_cuda[:num_generations], dim=0).to(torch.long) - 1 ) last_mtp_topk = decode_topk[last_rows] @@ -2267,7 +2284,7 @@ def sparse_attn_indexer( else: prev_topk_dst.copy_(last_mtp_topk) - if self.mtp_index_share and metadata.in_mtp_draft_loop and not reuse_topk: + if self.mtp_index_share and dsa.in_mtp_draft_loop and not reuse_topk: rows = None if has_decode: next_n = num_gen_tokens // num_generations @@ -2282,9 +2299,9 @@ def sparse_attn_indexer( ctx_rows = topk_indices_buffer[ctx_last] rows = ctx_rows if rows is None else torch.cat([ctx_rows, rows]) if rows is not None: - shared_topk_indices = metadata.shared_topk_indices + shared_topk_indices = dsa.shared_topk_indices if shared_topk_indices is None: - metadata.shared_topk_indices = rows.contiguous() + dsa.shared_topk_indices = rows.contiguous() else: row_start = num_contexts if is_generation else 0 shared_topk_indices[ @@ -2295,13 +2312,14 @@ def sparse_attn_indexer( def _mtp_last_accepted_rows( self, gen_topk: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, num_contexts: int, num_generations: int, next_n: int, ) -> torch.Tensor: """Select each generation request's last accepted TopK row.""" - num_accepted = metadata.mtp_num_accepted + dsa = _dsa_state(metadata) + num_accepted = dsa.mtp_num_accepted if num_accepted is None: return gen_topk[next_n - 1 :: next_n] gen_num_accepted = num_accepted[num_contexts : num_contexts + num_generations] @@ -2430,7 +2448,7 @@ def pre_indexer_proj( def forward_from_projected( self, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, hidden_states: torch.Tensor, indexer_intermediates: List[torch.Tensor], is_generation: Optional[bool] = None, @@ -2472,7 +2490,7 @@ def forward( self, qr: torch.Tensor, hidden_states: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, position_ids: torch.Tensor, ): q_fp8, k_fp8, k_scale, weights, q_scale = self.pre_indexer_proj( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index e0773630dcb3..99a2d8d06bf4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -6,13 +6,14 @@ import functools import os -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar, FrozenSet, List, Optional, Tuple import torch import tensorrt_llm import tensorrt_llm.bindings +from tensorrt_llm._torch.attention.backends.features import HelixMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._torch.utils import maybe_compile @@ -20,7 +21,8 @@ from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata from tensorrt_llm.logger import logger -from ..params import use_self_sampling_gvr +from ...features.spec_decoding import SpecDecodingMetadata +from ..params import MTPIndexShareExtension, SparseAlgorithmExtension, use_self_sampling_gvr from .cache_manager import is_dsa_cache_manager from .indexer import ( _DG_SCHEDULE_BLOCK_KV, @@ -44,7 +46,6 @@ _MAX_RADIX_BLOCKS_PER_ROW = 10 if TYPE_CHECKING: - from tensorrt_llm._torch.speculative.interface import SpecMetadata from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager @@ -90,94 +91,118 @@ class TokenMajorGenView: max_num_rows: int -@dataclass(init=False) -class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): - """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" - - sparse_metadata_params: Optional[DSAMetadataParams] = None - use_fp8_ds_mla: bool = field(default=False, init=False) - # Store reference to indexer for preparation stage - indexer: Optional["Indexer"] = None - # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed) - indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None - # (local q_begin, local q_end, per-rank sizes) when indexer_prefill_chunks - # cover only this TP rank's share of the context query tokens; the top-k - # results are then all-gathered once after the chunks. - indexer_prefill_q_split: Optional[Tuple[int, int, List[int]]] = None - # Max chunk size for two-level chunking: - # 1. Request-level: Pack multiple small requests into one chunk (up to indexer_max_chunk_size) - # 2. Intra-request: Split large requests into Q-blocks when seq_len > max_chunk_size - indexer_max_chunk_size: int - # TopK for static token sparse attention - num_sparse_topk: int - # TopK for dynamic sparse MLA - sparse_mla_topk: int - radix_aux_indices: Optional[torch.Tensor] = field(default=None, init=False) - radix_aux_logits: Optional[torch.Tensor] = field(default=None, init=False) - # max number of draft tokens - max_draft_tokens: int = 0 - # Indexer head dimension - indexer_head_dim: int = 128 - # Elements per four-byte scale word, matching Indexer.prepare(). - indexer_quant_block_size: int = 128 - # Enable indexer skip for short sequences - enable_indexer_skip: bool = False - shared_topk_indices: Optional[torch.Tensor] = None - indexer_skip_topk: bool = False - in_mtp_draft_loop: bool = False - mtp_num_accepted: Optional[torch.Tensor] = None - # Whether skip the indexer for context requests - skip_indexer_for_ctx_reqs: bool = False - # Whether skip the indexer for generation requests - skip_indexer_for_gen_reqs: bool = False - # Whether to use the expanded buffers for MTP support - use_expanded_buffers_for_mtp: bool = False - # Whether to reshape the DSL paged MQA logits Q tensor into a kernel- - # supported `effective_next_n` via caller-side atom-split (FP4: {1,2,3}; - # FP8: {1,2,3,4}; see `_pick_dsl_expand`). Reuses - # `kv_lens_expanded_cuda` / `block_table_expanded` / - # `scheduler_metadata_buffer_expanded`; runtime mutually exclusive with - # `use_expanded_buffers_for_mtp` (the latter requires `not _use_dsl`). - expand_for_dsl: bool = False - # Cached (expand_factor, atom) decision from the wave-aware picker. Set at - # `prepare()` time and read by forward call sites — avoids re-running the - # picker per call and guarantees prepare/forward use the SAME decision - # (otherwise the populated buffers would mismatch the kernel reshape). - dsl_expand_factor: int = 1 - dsl_atom: int = 1 - # In-kernel uneven split (``_pick_dsl_split``): atom count, the window it - # was picked for, and the per-atom causal stagger on device. Its schedule - # reuses ``kv_lens_expanded_cuda`` / ``scheduler_metadata_buffer_expanded``. - dsl_kernel_atoms: int = 1 - dsl_kernel_next_n: int = 0 - dsl_atom_ctx_offsets: Optional[torch.Tensor] = None - dsl_atom_ctx_offsets_key: Tuple[int, int] = (0, 0) - # Compression ratio for KV tokens - compress_ratios: List[int] = field(default_factory=lambda: [1]) - # Number of compressed KV tokens for context requests - num_ctx_kv_tokens: int = 0 - gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None +class DSAMetadata(SparseAlgorithmExtension, MTPIndexShareExtension): + """Indexer state of DSA (Dense Sparse Attention) on a TRTLLM metadata.""" + # Rebound by on_update_kv_lens(); the decode graph recomputes it. graph_temporary_attrs: ClassVar[FrozenSet[str]] = frozenset( {"gen_indexer_kv_lens_cuda_runtime"} ) - # Temporal-GVR prior state: allocated only when the two-level dispatch - # selects the temporal engine (the self-sampling engine keeps no - # cross-step state). - needs_gvr_prior: bool = field(default=False, init=False) - gvr_prior_indices: Optional[torch.Tensor] = field(default=None, init=False) - # Query tokens per generation request this step, in batch order. None - # denotes the uniform runtime stride. - ragged_verify_lens: Optional[List[int]] = None - # Allocate the row-major and token-major buffers used only by ragged - # verification. This is fixed before CUDA-graph metadata is created. - enable_ragged_verification: bool = False - # In device-window mode the host split determines only the captured shape; - # the true windows are installed before replay. - device_windows_mode: bool = False - # Query tokens per generation request for a uniform runtime tier. Zero - # falls back to the static draft-token ceiling. - runtime_tokens_per_gen_step: int = 0 + + def __init__(self, metadata: TrtllmAttentionMetadata) -> None: + super().__init__(metadata) + self.use_fp8_ds_mla = False + # Store reference to indexer for preparation stage + self.indexer: Optional["Indexer"] = None + # Chunked prefill metadata for indexer (prefill-only, no CUDA graph needed) + self.indexer_prefill_chunks: Optional[List[IndexerPrefillChunkMetadata]] = None + # (local q_begin, local q_end, per-rank sizes) when indexer_prefill_chunks + # cover only this TP rank's share of the context query tokens; the top-k + # results are then all-gathered once after the chunks. + self.indexer_prefill_q_split: Optional[Tuple[int, int, List[int]]] = None + # Max chunk size for two-level chunking: + # 1. Request-level: Pack multiple small requests into one chunk (up to indexer_max_chunk_size) + # 2. Intra-request: Split large requests into Q-blocks when seq_len > max_chunk_size + self.indexer_max_chunk_size = 0 + # TopK for dynamic sparse MLA + self.sparse_mla_topk = 0 + self.radix_aux_indices: Optional[torch.Tensor] = None + self.radix_aux_logits: Optional[torch.Tensor] = None + # max number of draft tokens + self.max_draft_tokens = 0 + # Indexer head dimension + self.indexer_head_dim = 128 + # Elements per four-byte scale word, matching Indexer.prepare(). + self.indexer_quant_block_size = 128 + # Enable indexer skip for short sequences + self.enable_indexer_skip = False + self.shared_topk_indices: Optional[torch.Tensor] = None + self.indexer_skip_topk = False + self.in_mtp_draft_loop = False + self.mtp_num_accepted: Optional[torch.Tensor] = None + # Whether skip the indexer for context requests + self.skip_indexer_for_ctx_reqs = False + # Whether skip the indexer for generation requests + self.skip_indexer_for_gen_reqs = False + # Whether to use the expanded buffers for MTP support + self.use_expanded_buffers_for_mtp = False + # Whether to reshape the DSL paged MQA logits Q tensor into a kernel- + # supported `effective_next_n` via caller-side atom-split (FP4: {1,2,3}; + # FP8: {1,2,3,4}; see `_pick_dsl_expand`). Reuses + # `kv_lens_expanded_cuda` / `block_table_expanded` / + # `scheduler_metadata_buffer_expanded`; runtime mutually exclusive with + # `use_expanded_buffers_for_mtp` (the latter requires `not _use_dsl`). + self.expand_for_dsl = False + # Cached (expand_factor, atom) decision from the wave-aware picker. Set at + # `prepare()` time and read by forward call sites — avoids re-running the + # picker per call and guarantees prepare/forward use the SAME decision + # (otherwise the populated buffers would mismatch the kernel reshape). + self.dsl_expand_factor = 1 + self.dsl_atom = 1 + # In-kernel uneven split (``_pick_dsl_split``): atom count, the window it + # was picked for, and the per-atom causal stagger on device. Its schedule + # reuses ``kv_lens_expanded_cuda`` / ``scheduler_metadata_buffer_expanded``. + self.dsl_kernel_atoms = 1 + self.dsl_kernel_next_n = 0 + self.dsl_atom_ctx_offsets: Optional[torch.Tensor] = None + self.dsl_atom_ctx_offsets_key: Tuple[int, int] = (0, 0) + # Compression ratio for KV tokens + self.compress_ratios: List[int] = [1] + # Number of compressed KV tokens for context requests + self.num_ctx_kv_tokens = 0 + self.gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None + # Temporal-GVR prior state: allocated only when the two-level dispatch + # selects the temporal engine (the self-sampling engine keeps no + # cross-step state). + self.needs_gvr_prior = False + self.gvr_prior_indices: Optional[torch.Tensor] = None + # Query tokens per generation request this step, in batch order. None + # denotes the uniform runtime stride. + self.ragged_verify_lens: Optional[List[int]] = None + # Allocate the row-major and token-major buffers used only by ragged + # verification. This is fixed before CUDA-graph metadata is created. + self.enable_ragged_verification = False + # In device-window mode the host split determines only the captured shape; + # the true windows are installed before replay. + self.device_windows_mode = False + # Query tokens per generation request for a uniform runtime tier. Zero + # falls back to the static draft-token ceiling. + self.runtime_tokens_per_gen_step = 0 + self.num_sms = tensorrt_llm.deep_gemm.get_num_sms() + # Cached step-invariant values for transform_local_topk_and_prepare_pool_view. + # These are recomputed once per step in _ensure_pool_view_cached() and + # reused across all layers to avoid redundant Python/CUDA overhead. + self._pool_cache_valid = False + self._cached_kv_mgr_id = 0 + self._cached_pool_view = None + self._cached_num_pool_tokens = 0 + self._cached_tokens_per_block = 0 + self._cached_block_table_ctx = None + self._cached_block_table_gen = None + self._cached_req_idx_ctx = None + self._cached_req_idx_gen = None + # Cross-layer fan-out of the DSA index remap (TRTLLM_DISABLE_DSA_GROUP_REMAP). + # `_group_remap_struct` is the (static) full+shared indexer + # group layout; `_group_remap_batched` holds the current forward's + # per-group batched remap output (leader writes, shared members read). + self._group_remap_struct = None + self._group_remap_struct_kv_id = 0 + self._group_remap_batched = {} + self.num_ctx_mla_kv_tokens = 0 + self.max_gen_seq_len = 0 + self.nvfp4_mla_context_fp8_scratch = None + self._radix_rows_per_sequence = 1 + self._ragged_mla_buffers_valid = False @property def is_ragged_verify(self) -> bool: @@ -198,20 +223,21 @@ def gen_token_stride(self) -> int: def update_indexer_kv_lens_2d(self, gen_kv_lens: torch.Tensor) -> torch.Tensor: """Fill and return the active query-width view of persistent storage.""" + md = self.metadata width = min(self.gen_token_stride, self.kv_lens_cuda_2d.shape[1]) - context_lens = self.kv_lens_cuda_2d[: self.num_generations, :width] + context_lens = self.kv_lens_cuda_2d[: md.num_generations, :width] context_lens.copy_(gen_kv_lens.unsqueeze(-1).expand(-1, width)) return context_lens - def mla_prepare_scheduler_buffers(self, num_heads: int) -> tuple[torch.Tensor, torch.Tensor]: + def mla_scheduler_buffers(self, num_heads: int) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: """Precompute prefixes for the same rows presented to MLA generation.""" view = self.token_major_gen_view() if view is None: - return super().mla_prepare_scheduler_buffers(num_heads) + return None - num_ctx = self.num_contexts + num_ctx = self.metadata.num_contexts n_gen = view.num_rows - num_ctx - if not self._mla_scheduler_buffers_valid: + if not self._ragged_mla_buffers_valid: # RoPE presents one query token per generation row. The device # extents are causal, pre-extra-KV lengths, just like the # cache_seq_lens scanned by the non-precomputed kernel path. @@ -228,15 +254,16 @@ def mla_prepare_scheduler_buffers(self, num_heads: int) -> tuple[torch.Tensor, t dtype=torch.int32, out=self.ragged_mla_cu_kv_seqlens[1 : n_gen + 1], ) - self._mla_scheduler_buffers_valid = True + self._ragged_mla_buffers_valid = True return ( self.ragged_mla_cu_q_rows[: n_gen + 1], self.ragged_mla_cu_kv_seqlens[: n_gen + 1], ) def gen_token_repeat_list(self) -> List[int]: + md = self.metadata if not self.is_ragged_verify: - return [self.gen_token_stride] * self.num_generations + return [self.gen_token_stride] * md.num_generations return list(self.ragged_verify_lens) def gen_token_repeats(self, device: Optional[torch.device] = None) -> torch.Tensor: @@ -244,12 +271,13 @@ def gen_token_repeats(self, device: Optional[torch.device] = None) -> torch.Tens def expand_per_gen_token(self, values: torch.Tensor, dim: int = 0) -> Tuple[torch.Tensor, int]: """Repeat request values once per generation query token.""" + md = self.metadata repeats = self.gen_token_repeat_list() num_tokens = sum(repeats) if not self.is_ragged_verify: return values.repeat_interleave(self.gen_token_stride, dim=dim), num_tokens if values.device == self.gen_token_repeats_cuda.device: - repeats_dev = self.gen_token_repeats_cuda[: self.num_generations] + repeats_dev = self.gen_token_repeats_cuda[: md.num_generations] else: repeats_dev = torch.tensor(repeats, dtype=torch.long, device=values.device) return ( @@ -257,80 +285,41 @@ def expand_per_gen_token(self, values: torch.Tensor, dim: int = 0) -> Tuple[torc num_tokens, ) - def __init__(self, *args, **kwargs): - """Initialize DSA metadata with SM count and indexer chunk size.""" - self.enable_ragged_verification = bool(kwargs.pop("enable_ragged_verification", False)) - sparse_attention_config = kwargs.pop("sparse_attention_config", None) - if ( - kwargs.get("sparse_metadata_params") is None - and sparse_attention_config is not None - and hasattr(sparse_attention_config, "to_sparse_metadata_params") - ): - kwargs["sparse_metadata_params"] = sparse_attention_config.to_sparse_metadata_params() - self.num_sms = tensorrt_llm.deep_gemm.get_num_sms() - # Cached step-invariant values for transform_local_topk_and_prepare_pool_view. - # These are recomputed once per step in _ensure_pool_view_cached() and - # reused across all layers to avoid redundant Python/CUDA overhead. - # Initialized here as plain instance attributes (not class-level - # annotations) to stay invisible to dataclass/torch.compile introspection. - self._pool_cache_valid = False - self._cached_kv_mgr_id = 0 - self._cached_pool_view = None - self._cached_num_pool_tokens = 0 - self._cached_tokens_per_block = 0 - self._cached_block_table_ctx = None - self._cached_block_table_gen = None - self._cached_req_idx_ctx = None - self._cached_req_idx_gen = None - # Cross-layer fan-out of the DSA index remap (TRTLLM_DISABLE_DSA_GROUP_REMAP). - # `_group_remap_struct` is the (static) full+shared indexer - # group layout; `_group_remap_batched` holds the current forward's - # per-group batched remap output (leader writes, shared members read). - self._group_remap_struct = None - self._group_remap_struct_kv_id = 0 - self._group_remap_batched = {} - self.num_ctx_mla_kv_tokens = 0 - self.nvfp4_mla_context_fp8_scratch = None - self._radix_rows_per_sequence = 1 - super().__init__(*args, **kwargs) - sparse_metadata_params = self.sparse_metadata_params - if not isinstance(sparse_metadata_params, DSAMetadataParams): - raise ValueError("DSA sparse attention metadata params are not set") - self.indexer_max_chunk_size = sparse_metadata_params.indexer_max_chunk_size - @property def use_cute_dsl_paged_mqa_logits(self) -> bool: """The Indexer's decision: the config flag gated on SM100 and cutlass-dsl.""" + md = self.metadata return use_cute_dsl_paged_mqa_logits( - self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits + md.sparse_metadata_params.use_cute_dsl_paged_mqa_logits ) - def __post_init__(self): + def allocate(self, buffers) -> None: """Allocate indexer K-cache buffers and heuristic TopK metadata.""" - super().__post_init__() - if self.is_cross: + md = self.metadata + if md.is_cross: # Ragged verification is a target self-attention capability. A # cross-attention metadata object may be a shallow copy of one, but # must not inherit its verifier-only buffers or per-step layout. self.enable_ragged_verification = False self.ragged_verify_lens = None self._validate_ragged_verification() - if not is_dsa_cache_manager(self.kv_cache_manager): + if not is_dsa_cache_manager(md.kv_cache_manager): has_deepseek_v4_cache_interface = all( - hasattr(self.kv_cache_manager, attr) + hasattr(md.kv_cache_manager, attr) for attr in ("compressed_block_sizes", "get_cache_indices") ) assert has_deepseek_v4_cache_interface, ( - "DSAtrtllmAttentionMetadata requires DSACacheManager-compatible " - f"cache manager, got {type(self.kv_cache_manager)}" + "DSAMetadata requires DSACacheManager-compatible " + f"cache manager, got {type(md.kv_cache_manager)}" ) - self.use_fp8_ds_mla = getattr(self.kv_cache_manager, "use_fp8_ds_mla", False) + self.use_fp8_ds_mla = getattr(md.kv_cache_manager, "use_fp8_ds_mla", False) - sparse_metadata_params = self.sparse_metadata_params + sparse_metadata_params = md.sparse_metadata_params if not isinstance(sparse_metadata_params, DSAMetadataParams): raise ValueError("DSA sparse attention metadata params are not set") - self.num_sparse_topk = sparse_metadata_params.max_sparse_topk - self.sparse_mla_topk = self.num_sparse_topk + self.indexer_max_chunk_size = sparse_metadata_params.indexer_max_chunk_size + md.num_sparse_topk = sparse_metadata_params.max_sparse_topk + self.sparse_mla_topk = md.num_sparse_topk self.indexer_head_dim = sparse_metadata_params.index_head_dim self.indexer_quant_block_size = 128 self.enable_indexer_skip = sparse_metadata_params.enable_indexer_skip @@ -341,7 +330,7 @@ def __post_init__(self): sparse_metadata_params.enable_heuristic_topk and get_sm_version() >= 100 ) self.kv_lens_row_reorder = None - capture_graph = self.is_cuda_graph + capture_graph = md.is_cuda_graph # Plain DSA has no compression and uses the default [1]. DeepSeek-V4's # metadata params carry the model-specific compression ratios. self.compress_ratios = getattr(sparse_metadata_params, "compress_ratios", [1]) @@ -350,9 +339,9 @@ def __post_init__(self): # DeepSeek-V4's indexer cache uses layer-dependent compressed block sizes # (tokens_per_block // compress_ratio), so slot mappings must be built # against that stride — not kv_cache_manager.tokens_per_block directly. - tpb = self.kv_cache_manager.tokens_per_block + tpb = md.kv_cache_manager.tokens_per_block self._indexer_compress_ratio = _select_indexer_compress_ratio(self.compress_ratios) - if hasattr(self.kv_cache_manager, "compressed_block_sizes"): + if hasattr(md.kv_cache_manager, "compressed_block_sizes"): tpb = tpb // _effective_compress_ratio_divisor(self._indexer_compress_ratio) self._tokens_per_block = tpb # Mirror the indexer's two-level GVR decision. The representative @@ -361,7 +350,7 @@ def __post_init__(self): self.use_self_sampling_topk = use_self_sampling_gvr( enable_heuristic_topk=self.enable_gvr_topk, use_self_sampling_topk=sparse_metadata_params.use_self_sampling_topk, - index_topk=self.num_sparse_topk, + index_topk=md.num_sparse_topk, compress_ratio=self._indexer_compress_ratio, is_cute_dsl_available=IS_CUTLASS_DSL_AVAILABLE, sm_version=get_sm_version(), @@ -377,27 +366,30 @@ def __post_init__(self): self.create_buffers_for_indexer(capture_graph=capture_graph) self._create_nvfp4_mla_generation_buffers(capture_graph=capture_graph) - def prepare(self): + def pre_derive_lengths(self) -> None: self._validate_ragged_verification() - super().prepare() + + def prepare(self) -> None: + md = self.metadata self._invalidate_pool_view_cache() # Get kv lengths - assert self.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" + assert md.kv_cache_params.use_cache is True, "DSA requires use_cache to be True" cached_token_lens = torch.tensor( - self.kv_cache_params.num_cached_tokens_per_seq, + md.kv_cache_params.num_cached_tokens_per_seq, dtype=torch.int, device="cpu", ) - if self.enable_helix: + helix = md.find_ext(HelixMetadata) + if helix is not None: # For Helix CP, inactive ranks only attend to previously cached # tokens (no new token appended), while active ranks add new tokens. # This mirrors the kv_lens logic in TrtllmAttentionMetadata.prepare(). - active_rank = ~self.helix_is_inactive_rank_cpu[: self.num_seqs] + active_rank = ~helix.is_inactive_rank_cpu[: md.num_seqs] kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.seq_lens_kv[active_rank] + kv_lens[active_rank] += md.seq_lens_kv[active_rank] else: - kv_lens = cached_token_lens + self.seq_lens_kv + kv_lens = cached_token_lens + md.seq_lens_kv kv_lens = self._prepare_device_window_kv_lens(kv_lens) @@ -417,68 +409,61 @@ def prepare(self): self.prepare_for_spec_decode(kv_lens) # Prepare metadata for indexer - Indexer.prepare(metadata=self) + Indexer.prepare(metadata=md) def _prepare_device_window_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: """Bound host-side KV extents until device-selected windows are installed.""" + md = self.metadata if self.device_windows_mode and self.is_ragged_verify: - nc, ns = self.num_contexts, self.num_seqs + nc, ns = md.num_contexts, md.num_seqs kv_lens = kv_lens.clone() - kv_lens[nc:ns] += (1 + self.max_draft_tokens) - self.seq_lens_kv[nc:ns] - self.kv_lens[nc:ns] = kv_lens[nc:ns] + kv_lens[nc:ns] += (1 + self.max_draft_tokens) - md.seq_lens_kv[nc:ns] + md.kv_lens[nc:ns] = kv_lens[nc:ns] return kv_lens - def prepare_for_draft_forward(self) -> dict | None: + # Indexer buffers rebound to the draft manager's for a draft forward. + _DRAFT_SWAPPED_FIELDS = { + "host_indexer_k_cache_block_offsets": "host_draft_indexer_k_cache_block_offsets", + "indexer_k_cache_block_offsets": "draft_indexer_k_cache_block_offsets", + "host_slot_mapping_fp8": "host_draft_slot_mapping_fp8", + "slot_mapping_fp8": "draft_slot_mapping_fp8", + "host_slot_mapping_scale": "host_draft_slot_mapping_scale", + "slot_mapping_scale": "draft_slot_mapping_scale", + "block_table": "draft_block_table", + "block_table_expanded": "draft_block_table_expanded", + "host_block_table_expanded": "host_draft_block_table_expanded", + } + + def draft_swap(self) -> dict | None: """Select native DSA indexer metadata for a draft forward.""" - # DeepSeek-V4 metadata inherits DSA metadata, but its cache manager uses a - # different dual-pool layout. Only native DSA cache managers use the DSA + # DeepSeek-V4 extends DSA, but its cache manager uses a different + # dual-pool layout. Only native DSA cache managers use the DSA # draft-replay buffers below. - if not is_dsa_cache_manager(self.kv_cache_manager): + md = self.metadata + if not is_dsa_cache_manager(md.kv_cache_manager): return None - saved_state = { - "host_indexer_k_cache_block_offsets": self.host_indexer_k_cache_block_offsets, - "indexer_k_cache_block_offsets": self.indexer_k_cache_block_offsets, - "host_slot_mapping_fp8": self.host_slot_mapping_fp8, - "host_slot_mapping_scale": self.host_slot_mapping_scale, - "slot_mapping_fp8": self.slot_mapping_fp8, - "slot_mapping_scale": self.slot_mapping_scale, - "block_table": self.block_table, - "block_table_expanded": self.block_table_expanded, - "host_block_table_expanded": self.host_block_table_expanded, - } - # The cached-KV feature owns these references even when an optimized - # path aliases them to slot_mapping_*. With the feature disabled, the - # aliases are lazy and may not exist on the first generation replay. - if self.enable_context_mla_with_cached_kv: - saved_state.update( - { - "slot_mapping_fp8_fullkv": self.slot_mapping_fp8_fullkv, - "slot_mapping_scale_fullkv": self.slot_mapping_scale_fullkv, - } - ) - # Rebind to the draft manager's dedicated buffers instead of # overwriting the target tensors in place. Rebinding is invisible to # CUDA graph capture, so the target and draft segments of the graph - # bake distinct addresses (like draft_kv_cache_block_offsets) and no + # bake distinct addresses (like the draft KV block offsets) and no # graph-recorded copy from a transient host buffer is needed. - self.swap_for_draft( - "host_indexer_k_cache_block_offsets", self.host_draft_indexer_k_cache_block_offsets - ) - self.swap_for_draft( - "indexer_k_cache_block_offsets", self.draft_indexer_k_cache_block_offsets - ) - self.swap_for_draft("host_slot_mapping_fp8", self.host_draft_slot_mapping_fp8) - self.swap_for_draft("slot_mapping_fp8", self.draft_slot_mapping_fp8) - self.swap_for_draft("host_slot_mapping_scale", self.host_draft_slot_mapping_scale) - self.swap_for_draft("slot_mapping_scale", self.draft_slot_mapping_scale) - self.swap_for_draft("block_table", self.draft_block_table) - self.swap_for_draft("block_table_expanded", self.draft_block_table_expanded) - self.swap_for_draft("host_block_table_expanded", self.host_draft_block_table_expanded) - # recompute_context_kv_gather_mappings() rebinds the aliases below. - self.record_draft_swap("slot_mapping_fp8_fullkv") - self.record_draft_swap("slot_mapping_scale_fullkv") + saved_state = self.swap_attrs_for_draft( + { + name: getattr(self, draft_name) + for name, draft_name in self._DRAFT_SWAPPED_FIELDS.items() + } + ) + # recompute_context_kv_gather_mappings() rebinds the full-KV aliases, + # so they are restored as well. + saved_state.update( + self.swap_attrs_for_draft( + { + name: getattr(self, name, None) + for name in ("slot_mapping_fp8_fullkv", "slot_mapping_scale_fullkv") + } + ) + ) self._invalidate_pool_view_cache() # Recording a capture executes no kernels, so the draft mappings only @@ -486,23 +471,22 @@ def prepare_for_draft_forward(self) -> dict | None: # (warmup) and the pre-replay call from model_engine. The per-step # advance inside the captured graph re-derives slot mappings on # device from the rebound block-offset buffer. - # kv_cache_manager was already swapped to the draft manager above. + # kv_cache_manager was already swapped to the draft manager. if not torch.cuda.is_current_stream_capturing(): self.prepare_for_indexer_k_cache() self._refresh_expanded_block_table() - Indexer.recompute_slot_mappings(self) - Indexer.recompute_context_kv_gather_mappings(self) + Indexer.recompute_slot_mappings(md) + Indexer.recompute_context_kv_gather_mappings(md) return saved_state - def restore_after_draft_forward(self, saved_state: dict | None) -> None: + def draft_restore(self, saved_state: dict | None) -> None: """Restore native DSA indexer metadata after a draft forward.""" if saved_state is None: return - - # restore_draft_swaps() has already rebound the swapped fields. + self.restore_swapped_attrs(saved_state) self._invalidate_pool_view_cache() - if "slot_mapping_fp8_fullkv" not in saved_state: + if not self.metadata.enable_context_mla_with_cached_kv: # Without cached-KV context MLA the aliases track the target mappings. self.slot_mapping_fp8_fullkv = self.slot_mapping_fp8 self.slot_mapping_scale_fullkv = self.slot_mapping_scale @@ -513,18 +497,20 @@ def get_indexer_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: return kv_lens // self._indexer_compress_ratio def get_indexer_max_seq_len(self) -> int: + md = self.metadata if self._indexer_compress_ratio <= 1: - return self.kv_cache_manager.max_seq_len - return max(1, self.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) + return md.kv_cache_manager.max_seq_len + return max(1, md.kv_cache_manager.max_seq_len // self._indexer_compress_ratio) def warmup_cute_dsl_radix_topk(self, next_n: int) -> None: """Pre-compile CuTe DSL radix variants not covered by engine warmup.""" - sparse_params = self.sparse_metadata_params + md = self.metadata + sparse_params = md.sparse_metadata_params if not self.use_cute_dsl_topk or ( sparse_params.enable_heuristic_topk and get_sm_version() >= 100 ): return - if self.kv_cache_manager is None: + if md.kv_cache_manager is None: return top_k = self.sparse_mla_topk if not top_k: @@ -564,9 +550,10 @@ def warmup_selfsampling_topk( """ # same two-level dispatch and hardware gates as the indexer __init__: # never compile these kernels on unsupported stacks during warmup + md = self.metadata if not IS_CUTLASS_DSL_AVAILABLE or get_sm_version() not in (100, 103): return - if not self.enable_gvr_topk or self.kv_cache_manager is None: + if not self.enable_gvr_topk or md.kv_cache_manager is None: return if not self.use_self_sampling_topk: return @@ -596,7 +583,7 @@ def warmup_selfsampling_topk( # forwarded only when enabled (keeps the default warmup calls unchanged) det_kw = ( {"deterministic": True} - if getattr(self.sparse_metadata_params, "gvr_deterministic_topk", False) + if getattr(md.sparse_metadata_params, "gvr_deterministic_topk", False) else {} ) try: @@ -638,7 +625,19 @@ def warmup_selfsampling_topk( "will JIT-compile lazily on first touch instead." ) - def on_update_kv_lens(self) -> None: + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: + if spec_dec_step: + md = self.metadata + # host + md.max_ctx_kv_len = 0 + md.num_ctx_cached_tokens = 0 + self.max_gen_seq_len = 1 + self._refresh_after_kv_lens_change() + + def invalidate_length_derived_buffers(self) -> None: + self._ragged_mla_buffers_valid = False + + def _refresh_after_kv_lens_change(self) -> None: # After changing the kv_lens/kv_lens_cuda, we may need to update other metadatas. # Especially for the changes in DecoderInputPreparer.preprocess_inputs(). # @@ -648,8 +647,7 @@ def on_update_kv_lens(self) -> None: # slot_mapping_* buffers also depend on these effective cached lengths. If we do not # refresh slot mappings here, indexer K-cache updates can be written with stale offsets. - super().on_update_kv_lens() - + md = self.metadata # preprocess_inputs() also uses this as a general hook to "invalidate per-forward-pass # caches so they are recomputed (and captured) on every _forward_step". Invalidate the # pool_view cache here so it is recomputed on the next @@ -669,17 +667,17 @@ def on_update_kv_lens(self) -> None: # it stale; those configs keep the eager chain. fused_eligible = ( _fused_dsa_meta_enabled() - and self.kv_cache_manager is not None - and self.num_tokens > 0 - and self.num_generations > 0 - and self.num_contexts == 0 + and md.kv_cache_manager is not None + and md.num_tokens > 0 + and md.num_generations > 0 + and md.num_contexts == 0 and not self.use_fp8_ds_mla ) # The draft loop can rewrite seq_lens after prepare() built this map. - if self.num_tokens > 0 and not fused_eligible: - self.req_idx_per_token[: self.num_tokens] = build_req_idx_per_token( - self.seq_lens_cuda[: self.num_seqs], self.num_tokens + if md.num_tokens > 0 and not fused_eligible: + self.req_idx_per_token[: md.num_tokens] = build_req_idx_per_token( + md.seq_lens_cuda[: md.num_seqs], md.num_tokens ).to(self.req_idx_per_token.dtype) # The overlap correction may have changed request KV lengths after @@ -689,25 +687,25 @@ def on_update_kv_lens(self) -> None: self.refresh_token_major_gen_rows() # Uncached DSA prefill aliases these scatter maps as gather maps. - if self.kv_cache_manager is not None and self.num_tokens > 0 and not fused_eligible: - seq_lens = self.seq_lens_cuda[: self.num_seqs] + if md.kv_cache_manager is not None and md.num_tokens > 0 and not fused_eligible: + seq_lens = md.seq_lens_cuda[: md.num_seqs] # Runtime cached lengths after overlap/spec-dec correction. - start_positions = self.kv_lens_cuda[: self.num_seqs] - seq_lens + start_positions = md.kv_lens_cuda[: md.num_seqs] - seq_lens - req_indices = self.req_idx_per_token[: self.num_tokens].to(dtype=torch.int64) + req_indices = self.req_idx_per_token[: md.num_tokens].to(dtype=torch.int64) seq_starts = torch.cumsum(seq_lens, dim=0, dtype=torch.int64) - seq_lens.to(torch.int64) token_offsets = ( - torch.arange(self.num_tokens, device=seq_lens.device, dtype=torch.int64) + torch.arange(md.num_tokens, device=seq_lens.device, dtype=torch.int64) - seq_starts[req_indices] ) global_positions = start_positions[req_indices] + token_offsets if self.use_fp8_ds_mla: - self.token_positions_cuda[: self.num_tokens] = global_positions.to(torch.int32) + self.token_positions_cuda[: md.num_tokens] = global_positions.to(torch.int32) # Honor MXFP4 indexer K cache layout (½ byte per value vs FP8's # 1 byte) when the cache manager exposes a use_fp4 flag. - index_head_dim = self.kv_cache_manager.index_head_dim - use_fp4 = getattr(self.kv_cache_manager, "use_fp4", False) + index_head_dim = md.kv_cache_manager.index_head_dim + use_fp4 = getattr(md.kv_cache_manager, "use_fp4", False) data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim fp8_indices, scale_indices = _compute_slot_mappings( global_positions, @@ -718,29 +716,27 @@ def on_update_kv_lens(self) -> None: self.indexer_quant_block_size, data_bytes_per_token=data_bytes_per_token, ) - self.slot_mapping_fp8[: self.num_tokens] = fp8_indices - self.slot_mapping_scale[: self.num_tokens] = scale_indices + self.slot_mapping_fp8[: md.num_tokens] = fp8_indices + self.slot_mapping_scale[: md.num_tokens] = scale_indices - if self.num_generations > 0: + if md.num_generations > 0: if not fused_eligible: torch.cumsum( - self.kv_lens_cuda[ - self.num_contexts : self.num_seqs - ], # num_contexts should be 0 + md.kv_lens_cuda[md.num_contexts : md.num_seqs], # num_contexts should be 0 dim=0, dtype=torch.int64, - out=self.gen_kv_indptr[1 : self.num_generations + 1], + out=self.gen_kv_indptr[1 : md.num_generations + 1], ) torch.cumsum( ( - self.kv_lens_cuda[self.num_contexts : self.num_seqs] - - self.seq_lens_cuda[self.num_contexts : self.num_seqs] + md.kv_lens_cuda[md.num_contexts : md.num_seqs] + - md.seq_lens_cuda[md.num_contexts : md.num_seqs] ), dim=0, dtype=torch.int64, - out=self.gen_cached_token_indptr[1 : self.num_generations + 1], + out=self.gen_cached_token_indptr[1 : md.num_generations + 1], ) - gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] + gen_kv_lens = md.kv_lens_cuda[md.num_contexts : md.num_seqs] gen_indexer_kv_lens = self.get_indexer_kv_lens(gen_kv_lens) self.gen_indexer_kv_lens_cuda_runtime = gen_indexer_kv_lens scheduler_context_lens = self.update_indexer_kv_lens_2d(gen_indexer_kv_lens) @@ -769,7 +765,7 @@ def on_update_kv_lens(self) -> None: scheduler_metadata_buffer_expanded, non_blocking=True ) if self.expand_for_dsl and self.dsl_kernel_atoms > 1: - num_rows = self.num_generations * self.dsl_kernel_atoms + num_rows = md.num_generations * self.dsl_kernel_atoms atom_lens = gen_indexer_kv_lens.unsqueeze(1) - self.dsl_atom_ctx_offsets self.kv_lens_expanded_cuda[:num_rows].copy_(atom_lens.reshape(-1)) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( @@ -782,7 +778,7 @@ def on_update_kv_lens(self) -> None: ) if self.expand_for_dsl and self.dsl_expand_factor > 1: expand_factor = self.dsl_expand_factor - num_tokens = self.num_generations * expand_factor + num_tokens = md.num_generations * expand_factor gen_kv_lens_expanded = gen_indexer_kv_lens.repeat_interleave(expand_factor) self.kv_lens_expanded_cuda[:num_tokens].copy_(gen_kv_lens_expanded) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( @@ -798,37 +794,38 @@ def on_update_kv_lens(self) -> None: if not getattr(self, "_fused_dsa_meta_armed", False): logger.info( "Fused DSA decode-metadata " - f"kernel armed (num_seqs={self.num_seqs}, " - f"num_tokens={self.num_tokens})." + f"kernel armed (num_seqs={md.num_seqs}, " + f"num_tokens={md.num_tokens})." ) self._fused_dsa_meta_armed = True self._run_fused_dsa_decode_metadata() self._compute_kv_lens_row_reorder() - self.prepare_dense_topk_indices(self.kv_lens_cuda, device=True) + self.prepare_dense_topk_indices(md.kv_lens_cuda, device=True) def _run_fused_dsa_decode_metadata(self): """Fill req_idx_per_token + slot mappings + gen indptrs via one Triton launch (see kernels.py), replacing the eager chain for the pure-decode step. Capture-safe: the launch is pre-compiled at warmup and replays inside the decode CUDA graph.""" - num_tokens = self.num_tokens - num_seqs = self.num_seqs + md = self.metadata + num_tokens = md.num_tokens + num_seqs = md.num_seqs # The fused header writes gen_*_indptr[1 : num_seqs + 1] while the eager # chain writes [1 : num_generations + 1]; they coincide only because the # eligibility gate requires num_contexts == 0 (=> num_seqs == # num_generations). Pin that coupling down explicitly. - assert num_seqs == self.num_generations, ( + assert num_seqs == md.num_generations, ( f"fused DSA metadata expects num_seqs ({num_seqs}) == " - f"num_generations ({self.num_generations})" + f"num_generations ({md.num_generations})" ) - index_head_dim = self.kv_cache_manager.index_head_dim - use_fp4 = getattr(self.kv_cache_manager, "use_fp4", False) + index_head_dim = md.kv_cache_manager.index_head_dim + use_fp4 = getattr(md.kv_cache_manager, "use_fp4", False) data_bytes_per_token = index_head_dim // 2 if use_fp4 else index_head_dim fused_dsa_decode_metadata( - self.seq_lens_cuda[:num_seqs], - self.kv_lens_cuda[:num_seqs], + md.seq_lens_cuda[:num_seqs], + md.kv_lens_cuda[:num_seqs], self.indexer_k_cache_block_offsets[:num_seqs], self.req_idx_per_token[:num_tokens], self.slot_mapping_fp8[:num_tokens], @@ -845,58 +842,50 @@ def _run_fused_dsa_decode_metadata(self): def _compute_kv_lens_row_reorder(self) -> None: """Prepare the longest-job-first GVR row order once per forward step.""" + md = self.metadata next_n = self.gen_token_stride - if self.needs_gvr_prior and self.num_generations * next_n >= 2 * self.num_sms: - gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] + if self.needs_gvr_prior and md.num_generations * next_n >= 2 * self.num_sms: + gen_kv_lens = md.kv_lens_cuda[md.num_contexts : md.num_seqs] order = torch.argsort(gen_kv_lens, descending=True).to(torch.int32) - self.kv_lens_row_reorder_buffer[: self.num_generations].copy_(order) - self.kv_lens_row_reorder = self.kv_lens_row_reorder_buffer[: self.num_generations] + self.kv_lens_row_reorder_buffer[: md.num_generations].copy_(order) + self.kv_lens_row_reorder = self.kv_lens_row_reorder_buffer[: md.num_generations] else: self.kv_lens_row_reorder = None - def update_for_spec_dec(self): - super().update_for_spec_dec() - # host - self.max_ctx_kv_len = 0 - self.num_ctx_cached_tokens = 0 - self.max_gen_seq_len = 1 - - # device - self.on_update_kv_lens() - # Create buffers for mla_rope_append_paged_kv_assign_q def create_buffers_for_mla_rope_append(self, capture_graph=False): # New context buffers for dsa - if not self.enable_context_mla_with_cached_kv: - self.ctx_cached_token_indptr = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_requests + 1,), + md = self.metadata + if not md.enable_context_mla_with_cached_kv: + md.ctx_cached_token_indptr = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_requests + 1,), cache_name="ctx_cached_token_indptr", dtype=torch.int64, capture_graph=capture_graph, ) - self.host_ctx_cached_token_indptr = torch.zeros_like( - self.ctx_cached_token_indptr, + md.host_ctx_cached_token_indptr = torch.zeros_like( + md.ctx_cached_token_indptr, device="cpu", pin_memory=prefer_pinned(), ) - self.ctx_kv_indptr = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_requests + 1,), + md.ctx_kv_indptr = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_requests + 1,), cache_name="ctx_kv_indptr", dtype=torch.int64, capture_graph=capture_graph, ) - self.host_ctx_kv_indptr = torch.zeros_like( - self.ctx_kv_indptr, + md.host_ctx_kv_indptr = torch.zeros_like( + md.ctx_kv_indptr, device="cpu", pin_memory=prefer_pinned(), ) # New generation buffers for dsa - self.gen_cached_token_indptr = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_requests + 1,), + self.gen_cached_token_indptr = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_requests + 1,), cache_name="gen_cached_token_indptr", dtype=torch.int64, capture_graph=capture_graph, @@ -906,9 +895,9 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - self.gen_kv_indptr = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_requests + 1,), + self.gen_kv_indptr = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_requests + 1,), cache_name="gen_kv_indptr", dtype=torch.int64, capture_graph=capture_graph, @@ -921,9 +910,10 @@ def create_buffers_for_mla_rope_append(self, capture_graph=False): def _create_nvfp4_mla_generation_buffers(self, capture_graph=False): """Allocate stable NVFP4 gather outputs for all MTP generation rows.""" + md = self.metadata self.nvfp4_mla_fp8_scratch = None self.nvfp4_mla_compact_indices = None - if getattr(self.kv_cache_manager, "dtype", None) != tensorrt_llm.bindings.DataType.NVFP4: + if getattr(md.kv_cache_manager, "dtype", None) != tensorrt_llm.bindings.DataType.NVFP4: return # MTP verifies the current token and max_draft_tokens draft tokens for @@ -931,38 +921,39 @@ def _create_nvfp4_mla_generation_buffers(self, capture_graph=False): # so the dequantized FP8 rows and compact offsets must cover the full # verification window. The buffers have stable addresses for CUDA # graphs and are reused sequentially by every attention layer. - max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) - self.nvfp4_mla_fp8_scratch = self.get_empty( - self.cuda_graph_buffers, + max_gen_tokens = md.max_num_sequences * (1 + self.max_draft_tokens) + self.nvfp4_mla_fp8_scratch = md.get_empty( + md.cuda_graph_buffers, ( max_gen_tokens, - self.num_sparse_topk, - self.kv_cache_manager.head_dim, + md.num_sparse_topk, + md.kv_cache_manager.head_dim, ), cache_name="nvfp4_mla_fp8_scratch", dtype=torch.float8_e4m3fn, capture_graph=capture_graph, ) - self.nvfp4_mla_compact_indices = self.get_empty( - self.cuda_graph_buffers, - (max_gen_tokens, self.num_sparse_topk), + self.nvfp4_mla_compact_indices = md.get_empty( + md.cuda_graph_buffers, + (max_gen_tokens, md.num_sparse_topk), cache_name="nvfp4_mla_compact_indices", dtype=torch.int32, capture_graph=capture_graph, ) def _create_radix_aux_buffers(self, capture_graph: bool) -> None: - max_num_rows = self.max_num_sequences * self._radix_rows_per_sequence - radix_shape = (max_num_rows, _MAX_RADIX_BLOCKS_PER_ROW, self.num_sparse_topk) - self.radix_aux_indices = self.get_empty( - self.cuda_graph_buffers, + md = self.metadata + max_num_rows = md.max_num_sequences * self._radix_rows_per_sequence + radix_shape = (max_num_rows, _MAX_RADIX_BLOCKS_PER_ROW, md.num_sparse_topk) + self.radix_aux_indices = md.get_empty( + md.cuda_graph_buffers, radix_shape, cache_name="radix_aux_indices", dtype=torch.int32, capture_graph=capture_graph, ) - self.radix_aux_logits = self.get_empty( - self.cuda_graph_buffers, + self.radix_aux_logits = md.get_empty( + md.cuda_graph_buffers, radix_shape, cache_name="radix_aux_logits", dtype=torch.float32, @@ -970,12 +961,13 @@ def _create_radix_aux_buffers(self, capture_graph: bool) -> None: ) def create_buffers_for_indexer(self, capture_graph=False): - sparse_metadata_params = self.sparse_metadata_params + md = self.metadata + sparse_metadata_params = md.sparse_metadata_params if not isinstance(sparse_metadata_params, DSAMetadataParams): raise ValueError("DSA sparse attention metadata params are not set") - self.indexer_k_cache_block_offsets = self.get_empty( - self.cuda_graph_buffers, - [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], + self.indexer_k_cache_block_offsets = md.get_empty( + md.cuda_graph_buffers, + [md.max_num_sequences, md.kv_cache_manager.max_blocks_per_seq], cache_name="indexer_k_cache_block_offsets", dtype=torch.int32, capture_graph=capture_graph, @@ -990,9 +982,9 @@ def create_buffers_for_indexer(self, capture_graph=False): sparse_metadata_params.has_shared_indexer_layers or sparse_metadata_params.mtp_index_share ): - self.shared_topk_indices = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, self.num_sparse_topk), + self.shared_topk_indices = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens, md.num_sparse_topk), cache_name="shared_topk_indices", dtype=torch.int32, capture_graph=capture_graph, @@ -1006,9 +998,9 @@ def create_buffers_for_indexer(self, capture_graph=False): # Indexer metadata # Separate slot mappings for non-interleaved layout (flat byte indices) - self.slot_mapping_fp8 = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.slot_mapping_fp8 = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="slot_mapping_fp8", dtype=torch.int64, capture_graph=capture_graph, @@ -1018,9 +1010,9 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - self.slot_mapping_scale = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.slot_mapping_scale = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="slot_mapping_scale", dtype=torch.int64, capture_graph=capture_graph, @@ -1032,9 +1024,9 @@ def create_buffers_for_indexer(self, capture_graph=False): ) self.token_positions_cuda = None if self.use_fp8_ds_mla: - self.token_positions_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.token_positions_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="token_positions_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -1044,10 +1036,10 @@ def create_buffers_for_indexer(self, capture_graph=False): # the draft-replay context swaps these in by rebinding, so CUDA graph # capture bakes distinct addresses for the target and draft segments # and both sides can be refreshed eagerly outside the graph. - if is_dsa_cache_manager(self.draft_kv_cache_manager): - self.draft_indexer_k_cache_block_offsets = self.get_empty( - self.cuda_graph_buffers, - [self.max_num_sequences, self.draft_kv_cache_manager.max_blocks_per_seq], + if is_dsa_cache_manager(md.draft_kv_cache_manager): + self.draft_indexer_k_cache_block_offsets = md.get_empty( + md.cuda_graph_buffers, + [md.max_num_sequences, md.draft_kv_cache_manager.max_blocks_per_seq], cache_name="draft_indexer_k_cache_block_offsets", dtype=torch.int32, capture_graph=capture_graph, @@ -1057,9 +1049,9 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - self.draft_slot_mapping_fp8 = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.draft_slot_mapping_fp8 = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="draft_slot_mapping_fp8", dtype=torch.int64, capture_graph=capture_graph, @@ -1069,9 +1061,9 @@ def create_buffers_for_indexer(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - self.draft_slot_mapping_scale = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.draft_slot_mapping_scale = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="draft_slot_mapping_scale", dtype=torch.int64, capture_graph=capture_graph, @@ -1084,13 +1076,13 @@ def create_buffers_for_indexer(self, capture_graph=False): # Only when MLA chunked prefill is enabled, we need to gather the full KV for indexer's logit computation. # Allocate these buffers dynamically in Indexer.prepare() # based on the actual total_kv_len to save memory. - if self.enable_context_mla_with_cached_kv: + if md.enable_context_mla_with_cached_kv: self.slot_mapping_fp8_fullkv = None self.slot_mapping_scale_fullkv = None # Per-token request index buffer for topk_indices conversion - self.req_idx_per_token = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.req_idx_per_token = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="req_idx_per_token", dtype=torch.int32, capture_graph=capture_graph, @@ -1099,26 +1091,26 @@ def create_buffers_for_indexer(self, capture_graph=False): self.req_idx_per_token, device="cpu", pin_memory=prefer_pinned() ) # Block table for topk_indices conversion (shared for context and generation) - self.block_table = self.get_empty( - self.cuda_graph_buffers, - [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq], + self.block_table = md.get_empty( + md.cuda_graph_buffers, + [md.max_num_sequences, md.kv_cache_manager.max_blocks_per_seq], cache_name="block_table", dtype=torch.int32, capture_graph=capture_graph, ) - if is_dsa_cache_manager(self.draft_kv_cache_manager): - self.draft_block_table = self.get_empty( - self.cuda_graph_buffers, + if is_dsa_cache_manager(md.draft_kv_cache_manager): + self.draft_block_table = md.get_empty( + md.cuda_graph_buffers, [ - self.max_num_sequences, - self.draft_kv_cache_manager.max_blocks_per_seq, + md.max_num_sequences, + md.draft_kv_cache_manager.max_blocks_per_seq, ], cache_name="draft_block_table", dtype=torch.int32, capture_graph=capture_graph, ) - self.scheduler_metadata_buffer = self.get_empty( - self.cuda_graph_buffers, + self.scheduler_metadata_buffer = md.get_empty( + md.cuda_graph_buffers, (self.num_sms + 1, 2), cache_name="scheduler_metadata_buffer", dtype=torch.int32, @@ -1133,8 +1125,8 @@ def create_buffers_for_indexer(self, capture_graph=False): # path keeps using `scheduler_metadata_buffer`. Always allocate (a # few KB) so transitions in `update_spec_dec_param` don't have to # special-case its existence. - self.scheduler_metadata_buffer_full_next_n = self.get_empty( - self.cuda_graph_buffers, + self.scheduler_metadata_buffer_full_next_n = md.get_empty( + md.cuda_graph_buffers, (self.num_sms + 1, 2), cache_name="scheduler_metadata_buffer_full_next_n", dtype=torch.int32, @@ -1146,25 +1138,25 @@ def create_buffers_for_indexer(self, capture_graph=False): # slice per forward. Avoids per-forward .expand().contiguous() # allocations that would break CUDA graphs. self._create_kv_lens_2d_buffer(capture_graph=capture_graph) - self.cu_seqlen_ks = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.cu_seqlen_ks = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="cu_seqlen_ks", dtype=torch.int32, capture_graph=capture_graph, ) - self.cu_seqlen_ke = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens,), + self.cu_seqlen_ke = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens,), cache_name="cu_seqlen_ke", dtype=torch.int32, capture_graph=capture_graph, ) # Topk indices buffer to support skip indexer for requests with short sequence lengths if self.enable_indexer_skip: - self.topk_indices_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_tokens, self.num_sparse_topk), + self.topk_indices_buffer = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_tokens, md.num_sparse_topk), cache_name="topk_indices_buffer", dtype=torch.int32, capture_graph=capture_graph, @@ -1175,21 +1167,21 @@ def create_buffers_for_indexer(self, capture_graph=False): pin_memory=prefer_pinned(), ) if self.needs_gvr_prior: - self.gvr_prior_indices = self.get_empty( - self.cuda_graph_buffers, + self.gvr_prior_indices = md.get_empty( + md.cuda_graph_buffers, ( - self.kv_cache_manager.num_local_layers, - self.max_num_sequences, - self.num_sparse_topk, + md.kv_cache_manager.num_local_layers, + md.max_num_sequences, + md.num_sparse_topk, ), cache_name="gvr_prior_indices", dtype=torch.int32, capture_graph=capture_graph, ) self.gvr_prior_indices.zero_() - self.kv_lens_row_reorder_buffer = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + self.kv_lens_row_reorder_buffer = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), cache_name="kv_lens_row_reorder_buffer", dtype=torch.int32, capture_graph=capture_graph, @@ -1213,9 +1205,10 @@ def _create_kv_lens_2d_buffer(self, capture_graph=False): in place. Scheduling receives a contiguous view at the active runtime query width, which may be smaller than the allocation ceiling. """ - self.kv_lens_cuda_2d = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences, 1 + self._draft_sizing_cap), + md = self.metadata + self.kv_lens_cuda_2d = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences, 1 + self._draft_sizing_cap), cache_name="kv_lens_cuda_2d", dtype=torch.int32, capture_graph=capture_graph, @@ -1224,9 +1217,10 @@ def _create_kv_lens_2d_buffer(self, capture_graph=False): # TODO: remove these expanded buffers when fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. def create_expanded_buffers(self, capture_graph=False): """Create expanded KV-length and block-table buffers for speculative decoding.""" - self.kv_lens_expanded_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences * (1 + self._draft_sizing_cap),), + md = self.metadata + self.kv_lens_expanded_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences * (1 + self._draft_sizing_cap),), cache_name="kv_lens_expanded_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -1260,11 +1254,11 @@ def create_expanded_buffers(self, capture_graph=False): self.ragged_mla_cu_kv_seqlens = None if self.enable_ragged_verification: self._create_ragged_verification_buffers(capture_graph) - self.block_table_expanded = self.get_empty( - self.cuda_graph_buffers, + self.block_table_expanded = md.get_empty( + md.cuda_graph_buffers, [ - self.max_num_sequences * (1 + self._draft_sizing_cap), - self.kv_cache_manager.max_blocks_per_seq, + md.max_num_sequences * (1 + self._draft_sizing_cap), + md.kv_cache_manager.max_blocks_per_seq, ], cache_name="block_table_expanded", dtype=torch.int32, @@ -1275,12 +1269,12 @@ def create_expanded_buffers(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - if is_dsa_cache_manager(self.draft_kv_cache_manager): - self.draft_block_table_expanded = self.get_empty( - self.cuda_graph_buffers, + if is_dsa_cache_manager(md.draft_kv_cache_manager): + self.draft_block_table_expanded = md.get_empty( + md.cuda_graph_buffers, [ - self.max_num_sequences * (1 + self._draft_sizing_cap), - self.draft_kv_cache_manager.max_blocks_per_seq, + md.max_num_sequences * (1 + self._draft_sizing_cap), + md.draft_kv_cache_manager.max_blocks_per_seq, ], cache_name="draft_block_table_expanded", dtype=torch.int32, @@ -1291,8 +1285,8 @@ def create_expanded_buffers(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) - self.scheduler_metadata_buffer_expanded = self.get_empty( - self.cuda_graph_buffers, + self.scheduler_metadata_buffer_expanded = md.get_empty( + md.cuda_graph_buffers, (self.num_sms + 1, 2), cache_name="scheduler_metadata_buffer_expanded", dtype=torch.int32, @@ -1302,9 +1296,10 @@ def create_expanded_buffers(self, capture_graph=False): def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: """Allocate persistent row layouts required only by ragged verification.""" # Stable-address repeat vector used by in-graph ragged expansions. - self.gen_token_repeats_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), + md = self.metadata + self.gen_token_repeats_cuda = md.get_empty( + md.cuda_graph_buffers, + (md.max_num_sequences,), cache_name="gen_token_repeats_cuda", dtype=torch.int64, capture_graph=capture_graph, @@ -1314,10 +1309,10 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: device="cpu", pin_memory=prefer_pinned(), ) - row_cap = self.max_num_sequences * (1 + self._draft_sizing_cap) + row_cap = md.max_num_sequences * (1 + self._draft_sizing_cap) for name in ("ragged_mla_cu_q_rows", "ragged_mla_cu_kv_seqlens"): - prefix = self.get_empty( - self.cuda_graph_buffers, + prefix = md.get_empty( + md.cuda_graph_buffers, (row_cap + 1,), cache_name=name, dtype=torch.int32, @@ -1325,8 +1320,8 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: ) prefix[:1].zero_() setattr(self, name, prefix) - self.row_kv_lens_cuda = self.get_empty( - self.cuda_graph_buffers, + self.row_kv_lens_cuda = md.get_empty( + md.cuda_graph_buffers, (row_cap,), cache_name="row_kv_lens_cuda", dtype=torch.int32, @@ -1335,8 +1330,8 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: self.row_kv_lens_host = torch.zeros_like( self.row_kv_lens_cuda, device="cpu", pin_memory=prefer_pinned() ) - self.row_kv_correction_cuda = self.get_empty( - self.cuda_graph_buffers, + self.row_kv_correction_cuda = md.get_empty( + md.cuda_graph_buffers, (row_cap,), cache_name="row_kv_correction_cuda", dtype=torch.int32, @@ -1345,8 +1340,8 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: self.row_kv_correction_host = torch.zeros_like( self.row_kv_correction_cuda, device="cpu", pin_memory=prefer_pinned() ) - self.row_req_idx_cuda = self.get_empty( - self.cuda_graph_buffers, + self.row_req_idx_cuda = md.get_empty( + md.cuda_graph_buffers, (row_cap,), cache_name="row_req_idx_cuda", dtype=torch.long, @@ -1358,23 +1353,23 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: # MLA RoPE and sparse-MLA generation consume one row per query token; # every other consumer keeps using request-major metadata. - attn_rows_cap = self.max_num_sequences * (2 + self._draft_sizing_cap) - self.attn_row_kv_lens_cuda = self.get_empty( - self.cuda_graph_buffers, + attn_rows_cap = md.max_num_sequences * (2 + self._draft_sizing_cap) + self.attn_row_kv_lens_cuda = md.get_empty( + md.cuda_graph_buffers, (attn_rows_cap,), cache_name="attn_row_kv_lens_cuda", dtype=torch.int32, capture_graph=capture_graph, ) - self.attn_row_kv_correction_cuda = self.get_empty( - self.cuda_graph_buffers, + self.attn_row_kv_correction_cuda = md.get_empty( + md.cuda_graph_buffers, (attn_rows_cap,), cache_name="attn_row_kv_correction_cuda", dtype=torch.int32, capture_graph=capture_graph, ) - self.attn_row_req_idx_cuda = self.get_empty( - self.cuda_graph_buffers, + self.attn_row_req_idx_cuda = md.get_empty( + md.cuda_graph_buffers, (attn_rows_cap,), cache_name="attn_row_req_idx_cuda", dtype=torch.long, @@ -1392,8 +1387,8 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: self.attn_row_request_types_host = torch.ones( (attn_rows_cap,), dtype=torch.int32, pin_memory=prefer_pinned() ) - self.attn_row_prompt_lens_cuda = self.get_empty( - self.cuda_graph_buffers, + self.attn_row_prompt_lens_cuda = md.get_empty( + md.cuda_graph_buffers, (attn_rows_cap,), cache_name="attn_row_prompt_lens_cuda", dtype=torch.int32, @@ -1403,17 +1398,17 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: (attn_rows_cap,), dtype=torch.int32, pin_memory=prefer_pinned() ) num_attention_op_pools = getattr( - self.kv_cache_manager, + md.kv_cache_manager, "num_attention_op_pools", - self.kv_cache_manager.num_pools, + md.kv_cache_manager.num_pools, ) - self.attn_row_block_offsets = self.get_empty( - self.cuda_graph_buffers, + self.attn_row_block_offsets = md.get_empty( + md.cuda_graph_buffers, ( num_attention_op_pools, attn_rows_cap, 2, - self.kv_cache_manager.max_blocks_per_seq, + md.kv_cache_manager.max_blocks_per_seq, ), cache_name="attn_row_block_offsets", dtype=torch.int32, @@ -1422,33 +1417,17 @@ def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: # This function is only used to create the expanded buffers when the max_draft_tokens is changed. # TODO: remove this function once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. - def update_spec_dec_param( + def on_spec_dec_param( self, - batch_size, - is_spec_decoding_enabled, - is_spec_dec_tree, - is_spec_dec_dynamic_tree, - max_draft_len, - max_total_draft_tokens, - spec_metadata: Optional["SpecMetadata"] = None, - spec_tree_manager: Optional["SpecTreeManager"] = None, - num_contexts: int = 0, - ): + max_draft_len: int, + max_total_draft_tokens: int, + spec_tree_manager: Optional["SpecTreeManager"], + ) -> None: """Update speculative decoding parameters and create expanded buffers.""" - super().update_spec_dec_param( - batch_size, - is_spec_decoding_enabled, - is_spec_dec_tree, - is_spec_dec_dynamic_tree, - max_draft_len, - max_total_draft_tokens, - spec_metadata, - spec_tree_manager, - num_contexts=num_contexts, - ) + md = self.metadata self.max_draft_tokens = max_draft_len rows_per_sequence = max(1 + max_draft_len, 1 + (max_total_draft_tokens or 0)) - if is_spec_dec_dynamic_tree and spec_tree_manager is not None: + if md.ext(SpecDecodingMetadata).is_spec_dec_dynamic_tree and spec_tree_manager is not None: rows_per_sequence = max(rows_per_sequence, spec_tree_manager._internal_buf_dim) resize_radix_workspace = rows_per_sequence != self._radix_rows_per_sequence self._radix_rows_per_sequence = rows_per_sequence @@ -1456,8 +1435,8 @@ def update_spec_dec_param( int(max_total_draft_tokens), int(getattr(self, "_draft_alloc_cap", 0) or 0), ) - capture_graph = self.is_cuda_graph - max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) + capture_graph = md.is_cuda_graph + max_gen_tokens = md.max_num_sequences * (1 + self.max_draft_tokens) if ( self.nvfp4_mla_fp8_scratch is not None and self.nvfp4_mla_fp8_scratch.shape[0] != max_gen_tokens @@ -1474,20 +1453,21 @@ def update_spec_dec_param( def _update_indexer_k_cache_block_offsets(self) -> torch.Tensor: """Refresh INDEX_KEY offsets and return their physical pool slots.""" - cache_manager = self.kv_cache_manager + md = self.metadata + cache_manager = md.kv_cache_manager # Raw block IDs can exceed the primary pool after host-cache offload; # the manager resolves their current physical slots. pool_indices = cache_manager.get_pool_block_indices( - self.num_seqs, - request_ids=self.request_ids, - num_contexts=self.num_contexts, - beam_width=self.beam_width, + md.num_seqs, + request_ids=md.request_ids, + num_contexts=md.num_contexts, + beam_width=md.beam_width, ) page_indices = pool_indices * cache_manager.indexer_k_cache_page_scale num_blocks = page_indices.shape[1] - self.host_indexer_k_cache_block_offsets[: self.num_seqs, :num_blocks].copy_(page_indices) - self.indexer_k_cache_block_offsets[: self.num_seqs].copy_( - self.host_indexer_k_cache_block_offsets[: self.num_seqs], non_blocking=True + self.host_indexer_k_cache_block_offsets[: md.num_seqs, :num_blocks].copy_(page_indices) + self.indexer_k_cache_block_offsets[: md.num_seqs].copy_( + self.host_indexer_k_cache_block_offsets[: md.num_seqs], non_blocking=True ) # Sanitize graph-padding entries that may be stale after cache # eviction or host-cache onboarding. @@ -1531,11 +1511,12 @@ def _ensure_pool_view_cached(self): start of every step (prepare() and on_update_kv_lens()), so the boolean flag is always cleared before the first per-layer call within a step. """ - if self._pool_cache_valid and self._cached_kv_mgr_id == id(self.kv_cache_manager): + md = self.metadata + if self._pool_cache_valid and self._cached_kv_mgr_id == id(md.kv_cache_manager): return - pool = self.kv_cache_manager.get_unique_primary_pool() - kv_cache_manager = self.kv_cache_manager + pool = md.kv_cache_manager.get_unique_primary_pool() + kv_cache_manager = md.kv_cache_manager num_blocks, num_layers = pool.shape[:2] self._cached_tokens_per_block = kv_cache_manager.tokens_per_block head_dim = kv_cache_manager.head_dim @@ -1547,11 +1528,11 @@ def _ensure_pool_view_cached(self): self._cached_pool_view = None else: self._cached_pool_view = pool.squeeze(2).view(-1, 1, head_dim) - self._cached_block_table_ctx = self.block_table[: self.num_contexts] - self._cached_block_table_gen = self.block_table[self.num_contexts : self.num_seqs] - self._cached_req_idx_ctx = self.req_idx_per_token[: self.num_ctx_tokens] + self._cached_block_table_ctx = self.block_table[: md.num_contexts] + self._cached_block_table_gen = self.block_table[md.num_contexts : md.num_seqs] + self._cached_req_idx_ctx = self.req_idx_per_token[: md.num_ctx_tokens] self._cached_req_idx_gen = ( - self.req_idx_per_token[self.num_ctx_tokens : self.num_tokens] - self.num_contexts + self.req_idx_per_token[md.num_ctx_tokens : md.num_tokens] - md.num_contexts ) self._cached_kv_mgr_id = id(kv_cache_manager) self._pool_cache_valid = True @@ -1572,7 +1553,8 @@ def _ensure_group_remap_struct(self): stride_factor). Singleton / non-uniform groups fall back to the per-layer single-op path (identical to flag-off behavior). """ - kvm = self.kv_cache_manager + md = self.metadata + kvm = md.kv_cache_manager if self._group_remap_struct is not None and self._group_remap_struct_kv_id == id(kvm): return self._group_remap_struct @@ -1646,6 +1628,7 @@ def _ensure_group_remap_struct(self): @maybe_compile(dynamic=True) def _get_dense_topk_indices(self, seq_lens, kv_lens, num_tokens): + md = self.metadata device = kv_lens.device past_kv_lens = kv_lens - seq_lens # get position ids @@ -1657,50 +1640,51 @@ def _get_dense_topk_indices(self, seq_lens, kv_lens, num_tokens): repeated_offsets = per_seq_offsets[batch_indices] position_ids = global_indices + repeated_offsets # get the dense topk indices with causal mask - range_row = torch.arange(self.num_sparse_topk, device=device) + range_row = torch.arange(md.num_sparse_topk, device=device) mask = range_row <= position_ids.unsqueeze(1) return torch.where(mask, range_row, -1) def prepare_dense_topk_indices(self, kv_lens, device=False): # device=False means use CPU """Prepare dense TopK indices for short sequences that skip the indexer.""" - if self.num_contexts > 0 and self.skip_indexer_for_ctx_reqs: - ctx_range = slice(self.num_ctx_tokens) + md = self.metadata + if md.num_contexts > 0 and self.skip_indexer_for_ctx_reqs: + ctx_range = slice(md.num_ctx_tokens) if device: self.topk_indices_buffer[ctx_range, :].copy_( self._get_dense_topk_indices( - self.seq_lens_cuda[: self.num_contexts], - kv_lens[: self.num_contexts], - self.num_ctx_tokens, + md.seq_lens_cuda[: md.num_contexts], + kv_lens[: md.num_contexts], + md.num_ctx_tokens, ), non_blocking=True, ) else: self.host_topk_indices_buffer[ctx_range, :] = self._get_dense_topk_indices( - self.seq_lens[: self.num_contexts], - kv_lens[: self.num_contexts], - self.num_ctx_tokens, + md.seq_lens[: md.num_contexts], + kv_lens[: md.num_contexts], + md.num_ctx_tokens, ) self.topk_indices_buffer[ctx_range, :].copy_( self.host_topk_indices_buffer[ctx_range, :], non_blocking=True ) - if self.num_generations > 0 and self.skip_indexer_for_gen_reqs: - gen_range = slice(self.num_ctx_tokens, self.num_tokens) + if md.num_generations > 0 and self.skip_indexer_for_gen_reqs: + gen_range = slice(md.num_ctx_tokens, md.num_tokens) if device: self.topk_indices_buffer[gen_range, :].copy_( self._get_dense_topk_indices( - self.seq_lens_cuda[self.num_contexts : self.num_seqs], - kv_lens[self.num_contexts : self.num_seqs], - self.num_tokens - self.num_ctx_tokens, + md.seq_lens_cuda[md.num_contexts : md.num_seqs], + kv_lens[md.num_contexts : md.num_seqs], + md.num_tokens - md.num_ctx_tokens, ), non_blocking=True, ) else: self.host_topk_indices_buffer[gen_range, :] = self._get_dense_topk_indices( - self.seq_lens[self.num_contexts : self.num_seqs], - kv_lens[self.num_contexts : self.num_seqs], - self.num_tokens - self.num_ctx_tokens, + md.seq_lens[md.num_contexts : md.num_seqs], + kv_lens[md.num_contexts : md.num_seqs], + md.num_tokens - md.num_ctx_tokens, ) self.topk_indices_buffer[gen_range, :].copy_( self.host_topk_indices_buffer[gen_range, :], non_blocking=True @@ -1716,6 +1700,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # TODO: # - Drop this sm90 branch (and the expanded buffers) once # fp8_paged_mqa_logits supports an arbitrary next_n on sm90 too. + md = self.metadata use_dsl = self.use_cute_dsl_paged_mqa_logits if not self.is_ragged_verify: self._ragged_num_rows = 0 @@ -1724,7 +1709,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # The DSL paged-MQA path reconstructs rows from one scalar window # and cannot express a per-request split. use_dsl = False - assert not self.is_spec_decoding_enabled, ( + assert not md.ext(SpecDecodingMetadata).is_spec_decoding_enabled, ( "ragged verification requires dense spec-decoding metadata to be disabled" ) @@ -1733,7 +1718,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): ) if self.use_expanded_buffers_for_mtp: if self.is_ragged_verify: - n_gen = self.num_generations + n_gen = md.num_generations self.host_gen_token_repeats[:n_gen].copy_( torch.tensor(self.gen_token_repeat_list(), dtype=torch.int64) ) @@ -1741,17 +1726,17 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.host_gen_token_repeats[:n_gen], non_blocking=True ) - gen_kv_lens = self.get_indexer_kv_lens(kv_lens[self.num_contexts : self.num_seqs]) + gen_kv_lens = self.get_indexer_kv_lens(kv_lens[md.num_contexts : md.num_seqs]) gen_kv_lens_expanded, num_tokens = self.expand_per_gen_token(gen_kv_lens) self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) self.kv_lens_expanded_cuda[:num_tokens].copy_( self.kv_lens_expanded_host[:num_tokens], non_blocking=True ) - if self.kv_cache_manager is not None and self.num_generations > 0: + if md.kv_cache_manager is not None and md.num_generations > 0: max_len = self.host_indexer_k_cache_block_offsets.shape[1] gen_blocks = self.host_indexer_k_cache_block_offsets[ - self.num_contexts : self.num_seqs, :max_len + md.num_contexts : md.num_seqs, :max_len ] expanded_blocks, _ = self.expand_per_gen_token(gen_blocks, dim=0) self.host_block_table_expanded[:num_tokens, :max_len].copy_(expanded_blocks) @@ -1765,16 +1750,16 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self._prepare_token_major_gen_rows(kv_lens) self.expand_for_dsl = ( - use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1 + use_dsl and md.kv_cache_manager is not None and self.max_draft_tokens >= 1 ) - if self.expand_for_dsl and self.num_generations > 0: + if self.expand_for_dsl and md.num_generations > 0: next_n = self.gen_token_stride - kernel_atoms = (1, 2, 3) if self.kv_cache_manager.use_fp4 else (1, 2, 3, 4) - gen_kv_lens = self.get_indexer_kv_lens(kv_lens[self.num_contexts : self.num_seqs]) + kernel_atoms = (1, 2, 3) if md.kv_cache_manager.use_fp4 else (1, 2, 3, 4) + gen_kv_lens = self.get_indexer_kv_lens(kv_lens[md.num_contexts : md.num_seqs]) max_ctx = int(gen_kv_lens.max().item()) if gen_kv_lens.numel() else 0 expand_factor, atom, kernel_split_atoms = _pick_dsl_split( next_n, - batch_size=self.num_generations, + batch_size=md.num_generations, max_ctx=max_ctx, num_sms=self.num_sms, kernel_atoms=kernel_atoms, @@ -1791,7 +1776,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): offsets, dtype=torch.int32, device="cuda" ) self.dsl_atom_ctx_offsets_key = (next_n, kernel_split_atoms) - num_rows = self.num_generations * kernel_split_atoms + num_rows = md.num_generations * kernel_split_atoms atom_lens = gen_kv_lens.unsqueeze(1) - torch.tensor( offsets, dtype=gen_kv_lens.dtype, device=gen_kv_lens.device ) @@ -1801,7 +1786,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.kv_lens_expanded_host[:num_rows], non_blocking=True ) if expand_factor > 1: - num_tokens = self.num_generations * expand_factor + num_tokens = md.num_generations * expand_factor gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(expand_factor) self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) self.kv_lens_expanded_cuda[:num_tokens].copy_( @@ -1815,14 +1800,15 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: """Populate the causal KV extent of each ragged generation row.""" + md = self.metadata verify_lens = self.ragged_verify_lens if not verify_lens: return - gen_kv_lens = kv_lens[self.num_contexts : self.num_seqs].tolist() + gen_kv_lens = kv_lens[md.num_contexts : md.num_seqs].tolist() assert len(gen_kv_lens) == len(verify_lens), ( f"ragged verify lengths {len(verify_lens)} != generation requests {len(gen_kv_lens)}" ) - expected_gen_tokens = self.num_tokens - self.num_ctx_tokens + expected_gen_tokens = md.num_tokens - md.num_ctx_tokens assert sum(verify_lens) == expected_gen_tokens, ( f"ragged verify windows sum to {sum(verify_lens)} but this step has " f"{expected_gen_tokens} generation tokens" @@ -1852,10 +1838,11 @@ def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: def refresh_ragged_row_kv_lens(self) -> None: """Refresh ragged extents after overlap changes ``kv_lens_cuda``.""" + md = self.metadata num_rows = self._ragged_num_rows if not self.is_ragged_verify or num_rows <= 0: return - gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] + gen_kv_lens = md.kv_lens_cuda[md.num_contexts : md.num_seqs] row_kv_lens = self.row_kv_lens_cuda[:num_rows] torch.index_select( gen_kv_lens, @@ -1867,13 +1854,14 @@ def refresh_ragged_row_kv_lens(self) -> None: def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: """Build the temporary one-row-per-query-token attention view.""" + md = self.metadata verify_lens = self.ragged_verify_lens num_gen_rows = self._ragged_num_rows if not verify_lens or num_gen_rows <= 0: self._attn_num_rows = 0 return - num_contexts = self.num_contexts + num_contexts = md.num_contexts num_rows = num_contexts + num_gen_rows capacity = self.attn_row_kv_lens_cuda.shape[0] assert num_rows <= capacity, ( @@ -1887,7 +1875,7 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: torch.arange(num_contexts, dtype=torch.long) ) self.attn_row_prompt_lens_cpu[:num_contexts].copy_( - self.prompt_lens_cpu[:num_contexts].to(torch.int32) + md.prompt_lens_cpu[:num_contexts].to(torch.int32) ) self.attn_row_kv_lens_host[num_contexts:num_rows].copy_( @@ -1923,11 +1911,12 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: def refresh_token_major_gen_rows(self) -> None: """Refresh token-major KV extents after overlap correction.""" + md = self.metadata num_rows = self._attn_num_rows if not self.is_ragged_verify or num_rows <= 0: return torch.index_select( - self.kv_lens_cuda[: self.num_seqs], + md.kv_lens_cuda[: md.num_seqs], 0, self.attn_row_req_idx_cuda[:num_rows], out=self.attn_row_kv_lens_cuda[:num_rows], @@ -1936,10 +1925,11 @@ def refresh_token_major_gen_rows(self) -> None: def refresh_token_major_block_table(self) -> None: """Expand the main attention block table along its sequence axis.""" + md = self.metadata num_rows = self._attn_num_rows - if num_rows <= 0 or self.kv_cache_block_offsets is None: + if num_rows <= 0 or md.kv_cache_block_offsets is None: return - source = self.kv_cache_block_offsets + source = md.kv_cache_block_offsets num_pools = min(source.shape[0], self.attn_row_block_offsets.shape[0]) width = min(source.shape[-1], self.attn_row_block_offsets.shape[-1]) torch.index_select( @@ -1981,13 +1971,14 @@ def apply_device_ragged_layout( kv_correction: torch.Tensor, ) -> None: """Install device-selected windows into stable captured buffers.""" + md = self.metadata self._validate_ragged_verification() if not self.is_ragged_verify: raise RuntimeError("device ragged layout requires enabled ragged verification") - num_contexts = self.num_contexts + num_contexts = md.num_contexts num_rows = self._ragged_num_rows - num_generations = self.num_generations - self.seq_lens_cuda[num_contexts : self.num_seqs] = verify_lens.to(self.seq_lens_cuda.dtype) + num_generations = md.num_generations + md.seq_lens_cuda[num_contexts : md.num_seqs] = verify_lens.to(md.seq_lens_cuda.dtype) self.gen_token_repeats_cuda[:num_generations] = verify_lens.to(torch.int64) self.row_req_idx_cuda[:num_rows] = req_idx self.row_kv_correction_cuda[:num_rows] = kv_correction @@ -1995,7 +1986,7 @@ def apply_device_ragged_layout( attn_row_end = attn_row_start + num_rows self.attn_row_req_idx_cuda[attn_row_start:attn_row_end] = req_idx + num_contexts self.attn_row_kv_correction_cuda[attn_row_start:attn_row_end] = kv_correction - token_start = self.num_ctx_tokens + token_start = md.num_ctx_tokens token_end = token_start + num_rows self.req_idx_per_token[token_start:token_end] = (req_idx + num_contexts).to( self.req_idx_per_token.dtype @@ -2006,14 +1997,15 @@ def apply_device_ragged_layout( self.block_table_expanded.shape[-1], ) self.block_table_expanded[:num_rows, :width] = ( - self.indexer_k_cache_block_offsets[num_contexts : self.num_seqs, :width] + self.indexer_k_cache_block_offsets[num_contexts : md.num_seqs, :width] .index_select(0, req_idx) .clamp_(min=0) ) def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None): """Refresh the active cache's expanded INDEX_KEY page table.""" - if self.kv_cache_manager is None or self.num_generations == 0: + md = self.metadata + if md.kv_cache_manager is None or md.num_generations == 0: return use_runtime_layout = False if repeat_factor is None: @@ -2026,12 +2018,12 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None): max_len = self.host_indexer_k_cache_block_offsets.shape[1] gen_block_tensor = self.host_indexer_k_cache_block_offsets[ - self.num_contexts : self.num_seqs, :max_len + md.num_contexts : md.num_seqs, :max_len ] if use_runtime_layout: expanded_blocks, num_tokens = self.expand_per_gen_token(gen_block_tensor, dim=0) else: - num_tokens = self.num_generations * repeat_factor + num_tokens = md.num_generations * repeat_factor expanded_blocks = gen_block_tensor.repeat_interleave(repeat_factor, dim=0) self.host_block_table_expanded[:num_tokens, :max_len].copy_( expanded_blocks, non_blocking=True @@ -2042,26 +2034,28 @@ def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None): self.block_table_expanded.clamp_(min=0) def prepare_for_indexer_k_cache(self): - if self.kv_cache_manager is None: + md = self.metadata + if md.kv_cache_manager is None: return # Keep physical slots for the primary-pool TopK conversion below. pool_indices = self._update_indexer_k_cache_block_offsets() # Build block_table for topk_indices conversion (actual block allocation) cached_token_lens = torch.tensor( - self.kv_cache_params.num_cached_tokens_per_seq, + md.kv_cache_params.num_cached_tokens_per_seq, dtype=torch.int, device="cpu", ) - if self.enable_helix: - active_rank = ~self.helix_is_inactive_rank_cpu[: self.num_seqs] + helix = md.find_ext(HelixMetadata) + if helix is not None: + active_rank = ~helix.is_inactive_rank_cpu[: md.num_seqs] kv_lens = cached_token_lens.clone() - kv_lens[active_rank] += self.seq_lens_kv[active_rank] + kv_lens[active_rank] += md.seq_lens_kv[active_rank] else: - kv_lens = cached_token_lens + self.seq_lens_kv - tokens_per_block = self.kv_cache_manager.tokens_per_block - num_blocks_per_seq = (kv_lens[: self.num_seqs] + tokens_per_block - 1) // tokens_per_block - max_blocks_used = num_blocks_per_seq.max().item() if self.num_seqs > 0 else 1 + kv_lens = cached_token_lens + md.seq_lens_kv + tokens_per_block = md.kv_cache_manager.tokens_per_block + num_blocks_per_seq = (kv_lens[: md.num_seqs] + tokens_per_block - 1) // tokens_per_block + max_blocks_used = num_blocks_per_seq.max().item() if md.num_seqs > 0 else 1 # pool_indices already has correct values; set padding to -1. # Stage through a fresh pinned buffer: an async H2D from pageable # memory would block the host behind the busy execution stream. @@ -2073,89 +2067,88 @@ def prepare_for_indexer_k_cache(self): host_block_table.copy_(pool_indices[:, :max_blocks_used]) pad_cols = torch.arange(max_blocks_used, dtype=num_blocks_per_seq.dtype) host_block_table.masked_fill_( - pad_cols.unsqueeze(0) >= num_blocks_per_seq[: self.num_seqs].unsqueeze(1), -1 + pad_cols.unsqueeze(0) >= num_blocks_per_seq[: md.num_seqs].unsqueeze(1), -1 ) # Copy to GPU - self.block_table[: self.num_seqs, :max_blocks_used].copy_( - host_block_table, non_blocking=True - ) + self.block_table[: md.num_seqs, :max_blocks_used].copy_(host_block_table, non_blocking=True) def prepare_for_skip_indexer(self, kv_lens: torch.Tensor): - num_extra_kv_tokens = self.kv_cache_params.num_extra_kv_tokens - if self.num_contexts > 0 and self.enable_indexer_skip: + md = self.metadata + num_extra_kv_tokens = md.kv_cache_params.num_extra_kv_tokens + if md.num_contexts > 0 and self.enable_indexer_skip: # Minus the number of extra KV tokens because when using one-model MTP, the # draft layers needs more KV tokens for the next draft forwards. self.skip_indexer_for_ctx_reqs = ( - kv_lens[: self.num_contexts].max().item() - <= self.num_sparse_topk - num_extra_kv_tokens + kv_lens[: md.num_contexts].max().item() <= md.num_sparse_topk - num_extra_kv_tokens ) else: self.skip_indexer_for_ctx_reqs = False - if self.num_generations > 0 and self.enable_indexer_skip: + if md.num_generations > 0 and self.enable_indexer_skip: # Minus the number of extra KV tokens because when using one-model MTP, the # draft layers needs more KV tokens for the next draft forwards. self.skip_indexer_for_gen_reqs = ( - kv_lens[self.num_contexts : self.num_seqs].max().item() - <= self.num_sparse_topk - num_extra_kv_tokens + kv_lens[md.num_contexts : md.num_seqs].max().item() + <= md.num_sparse_topk - num_extra_kv_tokens ) else: self.skip_indexer_for_gen_reqs = False self.prepare_dense_topk_indices(kv_lens) def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens: torch.Tensor): - if self.num_contexts > 0: - self.num_ctx_cached_tokens = cached_token_lens[: self.num_contexts].sum().item() - self.num_ctx_mla_kv_tokens = kv_lens[: self.num_contexts].sum().item() - self.max_ctx_kv_len = kv_lens[: self.num_contexts].max().item() - self.max_ctx_seq_len = self.seq_lens[: self.num_contexts].max().item() + md = self.metadata + if md.num_contexts > 0: + md.num_ctx_cached_tokens = cached_token_lens[: md.num_contexts].sum().item() + self.num_ctx_mla_kv_tokens = kv_lens[: md.num_contexts].sum().item() + md.max_ctx_kv_len = kv_lens[: md.num_contexts].max().item() + md.max_ctx_seq_len = md.seq_lens[: md.num_contexts].max().item() # context cached token indptr torch.cumsum( - cached_token_lens[: self.num_contexts], + cached_token_lens[: md.num_contexts], dim=0, dtype=torch.int64, - out=self.host_ctx_cached_token_indptr[1 : self.num_contexts + 1], + out=md.host_ctx_cached_token_indptr[1 : md.num_contexts + 1], ) - self.ctx_cached_token_indptr[: self.num_contexts + 1].copy_( - self.host_ctx_cached_token_indptr[: self.num_contexts + 1], non_blocking=True + md.ctx_cached_token_indptr[: md.num_contexts + 1].copy_( + md.host_ctx_cached_token_indptr[: md.num_contexts + 1], non_blocking=True ) # context kv indptr torch.cumsum( - kv_lens[: self.num_contexts], + kv_lens[: md.num_contexts], dim=0, dtype=torch.int64, - out=self.host_ctx_kv_indptr[1 : self.num_contexts + 1], + out=md.host_ctx_kv_indptr[1 : md.num_contexts + 1], ) - self.ctx_kv_indptr[: self.num_contexts + 1].copy_( - self.host_ctx_kv_indptr[: self.num_contexts + 1], non_blocking=True + md.ctx_kv_indptr[: md.num_contexts + 1].copy_( + md.host_ctx_kv_indptr[: md.num_contexts + 1], non_blocking=True ) else: - self.num_ctx_cached_tokens = 0 + md.num_ctx_cached_tokens = 0 self.num_ctx_mla_kv_tokens = 0 - self.max_ctx_kv_len = 0 - self.max_ctx_seq_len = 0 + md.max_ctx_kv_len = 0 + md.max_ctx_seq_len = 0 - if self.num_generations > 0: - self.max_gen_seq_len = self.seq_lens[self.num_contexts : self.num_seqs].max().item() + if md.num_generations > 0: + self.max_gen_seq_len = md.seq_lens[md.num_contexts : md.num_seqs].max().item() # generation cached token indptr torch.cumsum( - cached_token_lens[self.num_contexts : self.num_seqs], + cached_token_lens[md.num_contexts : md.num_seqs], dim=0, dtype=torch.int64, - out=self.host_gen_cached_token_indptr[1 : self.num_generations + 1], + out=self.host_gen_cached_token_indptr[1 : md.num_generations + 1], ) - self.gen_cached_token_indptr[: self.num_generations + 1].copy_( - self.host_gen_cached_token_indptr[: self.num_generations + 1], non_blocking=True + self.gen_cached_token_indptr[: md.num_generations + 1].copy_( + self.host_gen_cached_token_indptr[: md.num_generations + 1], non_blocking=True ) # generation kv indptr torch.cumsum( - kv_lens[self.num_contexts : self.num_seqs], + kv_lens[md.num_contexts : md.num_seqs], dim=0, dtype=torch.int64, - out=self.host_gen_kv_indptr[1 : self.num_generations + 1], + out=self.host_gen_kv_indptr[1 : md.num_generations + 1], ) - self.gen_kv_indptr[: self.num_generations + 1].copy_( - self.host_gen_kv_indptr[: self.num_generations + 1], non_blocking=True + self.gen_kv_indptr[: md.num_generations + 1].copy_( + self.host_gen_kv_indptr[: md.num_generations + 1], non_blocking=True ) else: self.max_gen_seq_len = 0 @@ -2163,12 +2156,13 @@ def prepare_for_mla_rope_append(self, cached_token_lens: torch.Tensor, kv_lens: def prepare_for_indices_conversion(self): # Build req_idx_per_token for topk_indices conversion # Use pinned staging buffer to avoid pageable H2D memcpy - self.host_req_idx_per_token[: self.num_tokens] = torch.repeat_interleave( - torch.arange(self.num_seqs, dtype=torch.int32), - self.seq_lens, + md = self.metadata + self.host_req_idx_per_token[: md.num_tokens] = torch.repeat_interleave( + torch.arange(md.num_seqs, dtype=torch.int32), + md.seq_lens, dim=0, ) - self.req_idx_per_token[: self.num_tokens].copy_( - self.host_req_idx_per_token[: self.num_tokens], + self.req_idx_per_token[: md.num_tokens].copy_( + self.host_req_idx_per_token[: md.num_tokens], non_blocking=True, ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py index f1ca16154b73..ed3a00547021 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/module.py @@ -20,7 +20,7 @@ from tensorrt_llm.logger import logger from ..hooks import MLASparseHooks, register_mla_sparse_hooks -from .metadata import DSAtrtllmAttentionMetadata +from .metadata import DSAMetadata from .params import DSABackendForwardArgs if TYPE_CHECKING: @@ -447,7 +447,7 @@ def forward_sparse_mla_kvcache_bf16( self, q: torch.Tensor, latent_cache: torch.Tensor, - attn_metadata: DSAtrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, output: torch.Tensor, indexer_intermediates: Optional[List[torch.Tensor]], is_generation: bool = False, @@ -460,9 +460,7 @@ def forward_sparse_mla_kvcache_bf16( 2. Load full kv cache from paged memory (with k rope applied) 3. Call FlashMLA sparse attention kernel for sparse prefill/decode """ - assert isinstance(attn_metadata, DSAtrtllmAttentionMetadata), ( - "DSA requires DSAtrtllmAttentionMetadata" - ) + dsa = attn_metadata.ext(DSAMetadata) attention_input_type = ( AttentionInputType.generation_only if is_generation else AttentionInputType.context_only ) @@ -476,8 +474,8 @@ def forward_sparse_mla_kvcache_bf16( ), ) assert topk_indices_pool is not None - attn_metadata._ensure_pool_view_cached() - kv_cache_pool = attn_metadata._cached_pool_view + dsa._ensure_pool_view_cached() + kv_cache_pool = dsa._cached_pool_view assert kv_cache_pool is not None # Append current tokens to paged cache and apply RoPE to q diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa_flashinfer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa_flashinfer.py index 3fc20e093d91..d99e1285a8cb 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa_flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa_flashinfer.py @@ -10,8 +10,9 @@ from tensorrt_llm._torch.attention.rotary_embedding import RotaryEmbedding from ..interface import AttentionForwardArgs, AttentionInputType +from ..trtllm import TrtllmAttentionMetadata from . import inline_scale_kv -from .dsa import DSAtrtllmAttentionMetadata +from .dsa import DSAMetadata from .flashinfer_utils import ( SPARSE_MLA_SPLIT_KV_TILE, allocate_sparse_mla_split_workspace, @@ -19,7 +20,7 @@ ) -def _inline_scale_pool_paged(metadata: DSAtrtllmAttentionMetadata) -> torch.Tensor: +def _inline_scale_pool_paged(metadata: TrtllmAttentionMetadata) -> torch.Tensor: """Return the layer-interleaved primary pool as inline-scale pages.""" manager = metadata.kv_cache_manager cached = getattr(metadata, "_inline_scale_pool", None) @@ -37,31 +38,32 @@ def _inline_scale_pool_paged(metadata: DSAtrtllmAttentionMetadata) -> torch.Tens def _latent_append( attn, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, latent_rows: torch.Tensor, start_idx: int, end_idx: int, is_generation: bool, ) -> None: """Quantize the new latent rows and scatter them into the main pool.""" - metadata._ensure_pool_view_cached() + dsa = metadata.ext(DSAMetadata) + dsa._ensure_pool_view_cached() page_index_scale, layer_offset = metadata.kv_cache_manager.get_primary_pool_page_index_params( attn.get_local_layer_idx(metadata) ) - positions = metadata.token_positions_cuda[start_idx:end_idx] + positions = dsa.token_positions_cuda[start_idx:end_idx] if is_generation: - block_table = metadata._cached_block_table_gen - req_idx = metadata._cached_req_idx_gen + block_table = dsa._cached_block_table_gen + req_idx = dsa._cached_req_idx_gen else: - block_table = metadata._cached_block_table_ctx - req_idx = metadata._cached_req_idx_ctx + block_table = dsa._cached_block_table_ctx + req_idx = dsa._cached_req_idx_ctx loc = torch.ops.trtllm.convert_req_index_to_global( req_idx, block_table, positions.unsqueeze(1).contiguous(), - metadata._cached_tokens_per_block, + dsa._cached_tokens_per_block, 1, - page_index_scale * metadata._cached_tokens_per_block, + page_index_scale * dsa._cached_tokens_per_block, layer_offset, ).view(-1) pool = _inline_scale_pool_paged(metadata) @@ -69,19 +71,20 @@ def _latent_append( pool.view(pool.shape[0], -1), loc, latent_rows, - page_size=metadata._cached_tokens_per_block, + page_size=dsa._cached_tokens_per_block, ) def run_flashinfer_sparse_mla( attn, q: torch.Tensor, - metadata: DSAtrtllmAttentionMetadata, + metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, rotary_emb: RotaryEmbedding, ) -> None: """Run DSA attention for ``FlashInferSparseMlaFmha``.""" - if metadata.max_draft_tokens > 0: + dsa = metadata.ext(DSAMetadata) + if dsa.max_draft_tokens > 0: raise NotImplementedError( "MTP / speculative decoding is not supported by the FlashInfer DSA FMHA yet." ) @@ -114,7 +117,7 @@ def run_flashinfer_sparse_mla( "paths must stay disabled for this FMHA." ) - positions = metadata.token_positions_cuda[start_idx:end_idx] + positions = dsa.token_positions_cuda[start_idx:end_idx] latent_cache = forward_args.latent_cache latent_rows = ( latent_cache.view(num_tokens, attn.head_dim) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py index 13937deed749..f4237bbece7f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py @@ -33,6 +33,7 @@ from tensorrt_llm._utils import get_sm_version +from ...features import RecurrentMetadata from ...interface import ( AttentionForwardArgs, AttentionInputType, @@ -241,7 +242,7 @@ def _cache_state(self, metadata) -> _GlmKpoolCacheState: "glm_kpool metadata has no kv_cache_manager; the hybrid " "KVCacheManagerV2 owns the latent/indexer pools" ) - mamba_metadata = metadata.mamba_metadata + mamba_metadata = metadata.ext(RecurrentMetadata).mamba_metadata assert isinstance(mamba_metadata, Glm5NextMamba2Metadata), ( "glm_kpool requires Glm5NextMamba2Metadata" ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py index 29398d8e5f3b..3bf6daf7d630 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/__init__.py @@ -7,8 +7,8 @@ * :mod:`.triton_metadata` -- ``MiniMaxM3TritonSparseAttentionMetadata`` dataclass, CUDA-graph-stable buffer allocator + builder, and the - :class:`AttentionMetadata` subclass - factory for the Triton reference path. + :class:`AttentionMetadata` extension + for the Triton reference path. * :mod:`.cache_manager` -- standalone side index cache used by tests and the :class:`KVCacheManagerV2` subclass factory. Shared by both backends. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py index 774c21d88126..07b37a1f8109 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/msa_utils.py @@ -170,6 +170,8 @@ def write_msa_phase_kv( MiniMaxM3MsaSparseAttention.write_layer_caches ahead of its indexer, tells both libraries the cache is already resident. """ + from ..msa_backend import MiniMaxM3MsaMetadata + if attention_input_type != AttentionInputType.mixed: raise NotImplementedError( "MiniMax-M3 MSA attention requires the mixed attention input type, " @@ -178,13 +180,14 @@ def write_msa_phase_kv( ) if k is None or v is None: return + msa = metadata.ext(MiniMaxM3MsaMetadata) num_tokens = int(k.shape[0]) if num_tokens == 0: return write_msa_main_kv( metadata.kv_cache_manager, attn.layer_idx, - metadata.msa_out_cache_loc[token_offset : token_offset + num_tokens], + msa.msa_out_cache_loc[token_offset : token_offset + num_tokens], k, v, # k and v are this phase's own token slice, so every row owns a slot. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index dca8a6abd44c..1cccd159bce0 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -10,38 +10,39 @@ * The indexer calls fmha_sm100 directly to produce the per-query selected block indices, which the model layer threads through forward_args.sparse_backend_args. - * MiniMaxM3MsaSparseAttentionMetadata subclasses TrtllmAttentionMetadata and - stores its per-forward MSA tensors in CUDA-graph-stable buffers. - The buffers are allocated once in __post_init__ via - get_empty(capture_graph=...), and prepare() copies the per-step values - into them. The standard CUDAGraphRunner clones one metadata per graph - batch size (create_cuda_graph_metadata), so no per-batch-size cache is - needed here. + * MiniMaxM3MsaMetadata extends TrtllmAttentionMetadata and stores its + per-forward MSA tensors in CUDA-graph-stable buffers. The buffers are + allocated once per metadata in allocate() via get_empty(capture_graph=...), + and prepare() copies the per-step values into them. The standard + CUDAGraphRunner clones one metadata per graph batch size + (create_cuda_graph_metadata), so no per-batch-size cache is needed here. * With Eagle3 a decode row has 1 + draft_len query tokens. Slots and valid-block counts are per token, and the decode kernels take that uniform - query length from msa_decode_span. on_update_kv_lens re-derives the + query length from msa_decode_span. on_kv_lens_changed re-derives the per-request lengths, the slots and the counts after the overlap scheduler corrects kv_lens on device. -The classes subclass TrtllmAttention and TrtllmAttentionMetadata, imported at +The classes build on TrtllmAttention and TrtllmAttentionMetadata, imported at module scope. That is cycle-free only because the dependency runs one way: the -two FMHA libraries reach the kernels through ...minimax_m3.kernels and never -load this module, so trtllm's import chain does not come back here. +two FMHA libraries reach the kernels through ...minimax_m3.kernels and load +this module only lazily, at call time, to look up MiniMaxM3MsaMetadata, so +trtllm's import chain does not come back here. """ from __future__ import annotations -from dataclasses import dataclass from typing import NamedTuple, Optional, Tuple import torch from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention.backends.sparse.params import SparseAlgorithmExtension from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType from tensorrt_llm.models.modeling_utils import QuantConfig +from ...features.spec_decoding import SpecDecodingMetadata from .common import ( MiniMaxM3SparseConfig, MiniMaxM3SparseMetadataParams, @@ -115,9 +116,8 @@ class MsaDecodeSpan(NamedTuple): query_len: int -@dataclass(init=False) -class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): - """TrtllmAttentionMetadata for MiniMax-M3 MSA sparse layers. +class MiniMaxM3MsaMetadata(SparseAlgorithmExtension): + """TrtllmAttentionMetadata extension for MiniMax-M3 MSA sparse layers. A step is prepared for a fixed division of labour: its context rows run on fmha_sm100 through MsaPrefillFmha, and its generation rows on MiniMax-M3's @@ -144,7 +144,7 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): query_len is that count, slots and valid-block counts are per token, and the overlap scheduler corrects kv_lens on device after prepare(). The decode kernels read their lengths from msa_seq_lens_cuda, so - on_update_kv_lens patches that buffer, the slots and the counts. The + on_kv_lens_changed patches that buffer, the slots and the counts. The fmha_sm100 plans cover context rows only, whose lengths the correction never touches, so they need no patch. """ @@ -175,7 +175,7 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # factor, or 0 where the pool has no single one; see msa_subpage_rows. msa_subpage_block_table: Optional[torch.Tensor] = None _msa_subpages_per_slot: int = 0 - # Inputs for on_update_kv_lens: each query token's request row and offset + # Inputs for on_kv_lens_changed: each query token's request row and offset # within the request, and the kv_lens prepare() staged (the upper bound the # correction is clamped to). The write slots are re-derived from # msa_block_table, so no per-token slot table is kept. @@ -215,13 +215,13 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): _msa_max_kv_len: int = 0 # See msa_worst_case_max_k_tiles. _msa_worst_case_max_k_tiles: int = 0 - # True only with speculative decoding; otherwise on_update_kv_lens and its + # True only with speculative decoding; otherwise on_kv_lens_changed and its # staging are skipped and non-speculative steps run as before. _msa_kv_lens_dynamic: bool = False - def __post_init__(self) -> None: - super().__post_init__() - params = self.sparse_metadata_params + def allocate(self, buffers) -> None: + md = self.metadata + params = md.sparse_metadata_params self._msa_params = params if isinstance(params, MiniMaxM3SparseMetadataParams) else None self._create_msa_buffers() self._validate_decode_kernel_support() @@ -249,18 +249,19 @@ def _stage_host_lengths(self) -> None: what lets the planners stage them with non-blocking copies instead of copying out of pageable memory, so it is worth paying for once a step. """ + md = self.metadata def as_pinned_int32(lens: torch.Tensor) -> torch.Tensor: - rows = lens[: self.num_seqs] + rows = lens[: md.num_seqs] return maybe_pin_memory(rows.to(torch.int32) if rows.dtype != torch.int32 else rows) - seq_lens = self.seq_lens - kv_lens = getattr(self, "kv_lens", None) + seq_lens = md.seq_lens + kv_lens = getattr(md, "kv_lens", None) # The host kv_lens carries num_extra_kv_tokens (speculative draft # slots); MSA needs the attended length, which is what kv_lens_cuda # holds. Every consumer below, msa_seq_lens_cuda included, inherits # the correction from here. - params = self.kv_cache_params + params = md.kv_cache_params extra = int(params.num_extra_kv_tokens) if params is not None else 0 if kv_lens is not None and extra: kv_lens = kv_lens - extra @@ -359,7 +360,8 @@ def _msa_main_kv_is_fp8(self) -> bool: sparse-prefill path selects the FP8 AOT kernels. Mirrors the k_paged dtype check in run_msa_prefill_gqa. """ - kv_cache_manager = self.kv_cache_manager + md = self.metadata + kv_cache_manager = md.kv_cache_manager return kv_cache_manager is not None and kv_cache_manager.dtype in ( DataType.FP8, DataType.NVFP4, @@ -375,8 +377,9 @@ def _validate_decode_kernel_support(self) -> None: not the layer, fixes the page size, the index dtype and the sub-page pool. """ + md = self.metadata params = self._msa_params - kv_cache_manager = self.kv_cache_manager + kv_cache_manager = md.kv_cache_manager if params is None or kv_cache_manager is None: # No MSA geometry to check, as for a structural test's metadata. return @@ -386,7 +389,7 @@ def _validate_decode_kernel_support(self) -> None: # num_index_heads * query tokens into one Q block, which bounds the # draft length it can verify. decode_query_len = self._msa_max_decode_query_len() - num_index_heads = params.sharded_index_head_count(self.mapping) + num_index_heads = params.sharded_index_head_count(md.mapping) if not self._cutedsl_indexer_supported( num_index_heads=num_index_heads, page_size=page_size, @@ -413,7 +416,8 @@ def _create_msa_buffers(self) -> None: max_num_tokens for cache slots, max_num_sequences * max_blocks_per_seq for the page table, and worst-case max_k_tiles for proxy scratch. """ - kv_cache_manager = self.kv_cache_manager + md = self.metadata + kv_cache_manager = md.kv_cache_manager self._msa_buffers_ready = False if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): return @@ -430,35 +434,35 @@ def _create_msa_buffers(self) -> None: for layer_idx in getattr(kv_cache_manager, "sparse_layer_ids", ()) if layer_idx in kv_cache_manager.layer_offsets } - capture_graph = self.is_cuda_graph - buffers = self.cuda_graph_buffers - max_num_sequences = int(self.max_num_sequences) + capture_graph = md.is_cuda_graph + buffers = md.cuda_graph_buffers + max_num_sequences = int(md.max_num_sequences) max_blocks_per_seq = int(kv_cache_manager.max_blocks_per_seq) max_total_pages = max_num_sequences * max_blocks_per_seq - max_num_tokens = int(self.max_num_tokens) + max_num_tokens = int(md.max_num_tokens) - self.msa_out_cache_loc = self.get_empty( + self.msa_out_cache_loc = md.get_empty( buffers, (max_num_tokens,), cache_name="msa_out_cache_loc", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_kv_indices = self.get_empty( + self.msa_kv_indices = md.get_empty( buffers, (max_total_pages,), cache_name="msa_kv_indices", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_block_table = self.get_empty( + self.msa_block_table = md.get_empty( buffers, (max_num_sequences, max_blocks_per_seq), cache_name="msa_block_table", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_seq_lens_cuda = self.get_empty( + self.msa_seq_lens_cuda = md.get_empty( buffers, (max_num_sequences,), cache_name="msa_seq_lens_cuda", @@ -466,21 +470,21 @@ def _create_msa_buffers(self) -> None: capture_graph=capture_graph, ) if getattr(kv_cache_manager, "dtype", None) == DataType.NVFP4: - self.msa_qo_lens_cuda = self.get_empty( + self.msa_qo_lens_cuda = md.get_empty( buffers, (max_num_sequences,), cache_name="msa_qo_lens_cuda", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_cu_q_lens = self.get_empty( + self.msa_cu_q_lens = md.get_empty( buffers, (max_num_sequences + 1,), cache_name="msa_cu_q_lens", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_cu_kv_lens = self.get_empty( + self.msa_cu_kv_lens = md.get_empty( buffers, (max_num_sequences + 1,), cache_name="msa_cu_kv_lens", @@ -491,29 +495,29 @@ def _create_msa_buffers(self) -> None: # pool's layout for the life of the manager. self._msa_subpages_per_slot = uniform_subpages_per_slot(kv_cache_manager) if self._msa_subpages_per_slot > 0: - self.msa_subpage_block_table = self.get_empty( + self.msa_subpage_block_table = md.get_empty( buffers, (max_num_sequences, 2, max_blocks_per_seq), cache_name="msa_subpage_block_table", dtype=torch.int32, capture_graph=capture_graph, ) - # Inputs for on_update_kv_lens. - self.msa_q_batch_row = self.get_empty( + # Inputs for on_kv_lens_changed. + self.msa_q_batch_row = md.get_empty( buffers, (max_num_tokens,), cache_name="msa_q_batch_row", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_q_intra = self.get_empty( + self.msa_q_intra = md.get_empty( buffers, (max_num_tokens,), cache_name="msa_q_intra", dtype=torch.int32, capture_graph=capture_graph, ) - self.msa_kv_lens_staged = self.get_empty( + self.msa_kv_lens_staged = md.get_empty( buffers, (max_num_sequences,), cache_name="msa_kv_lens_staged", @@ -529,13 +533,13 @@ def _create_msa_buffers(self) -> None: fmha_sm100 = require_msa_module() max_k_tiles = _worst_case_proxy_max_k_tiles( fmha_sm100, - num_index_heads=params.sharded_index_head_count(self.mapping), + num_index_heads=params.sharded_index_head_count(md.mapping), kv_cache_manager=kv_cache_manager, max_batch=max_num_sequences, ) self._msa_worst_case_max_k_tiles = int(max_k_tiles) self._alloc_msa_proxy_scratch( - num_index_heads=params.sharded_index_head_count(self.mapping), + num_index_heads=params.sharded_index_head_count(md.mapping), max_tokens=self._msa_max_decode_tokens(), max_k_tiles=max_k_tiles, capture_graph=capture_graph, @@ -550,7 +554,8 @@ def _msa_max_decode_query_len(self) -> int: one source for both the up-front kernel validation and the scratch sizing. A run without speculative decoding reports 1. """ - draft_len = int(getattr(self.kv_cache_manager, "max_total_draft_tokens", 0) or 0) + md = self.metadata + draft_len = int(getattr(md.kv_cache_manager, "max_total_draft_tokens", 0) or 0) return 1 + max(0, draft_len) def _msa_max_decode_tokens(self) -> int: @@ -561,9 +566,10 @@ def _msa_max_decode_tokens(self) -> int: 16384, which also keeps a whole-batch fmha_sm100 proxy plan under its 65536 total_q * num_qo_heads limit with 4 index heads. """ - max_seqs = int(self.max_num_sequences) + md = self.metadata + max_seqs = int(md.max_num_sequences) tokens = max_seqs * self._msa_max_decode_query_len() - max_toks = int(self.max_num_tokens or 0) + max_toks = int(md.max_num_tokens or 0) if max_toks > 0: tokens = min(tokens, max_toks) return max(max_seqs, min(tokens, 16384)) @@ -582,15 +588,16 @@ def _alloc_msa_proxy_scratch( the batch size under speculative verify), so one allocation serves every decode step. msa_proxy_max_score_view slices the per-step shape. """ - buffers = self.cuda_graph_buffers - self.msa_max_score = self.get_empty( + md = self.metadata + buffers = md.cuda_graph_buffers + self.msa_max_score = md.get_empty( buffers, (num_index_heads * max_k_tiles * max_tokens,), cache_name="msa_max_score", dtype=torch.float32, capture_graph=capture_graph, ) - self.msa_n_valid_blocks = self.get_empty( + self.msa_n_valid_blocks = md.get_empty( buffers, (max_tokens,), cache_name="msa_n_valid_blocks", @@ -607,6 +614,7 @@ def _ensure_msa_decode_scratch_buffers( required_max_k_tiles: int, ) -> None: """Ensure proxy scratch buffers exist and cover the current plan.""" + md = self.metadata max_tokens = max(int(max_batch), self._msa_max_decode_tokens()) required_numel = num_index_heads * required_max_k_tiles * max_tokens if self.msa_max_score is not None: @@ -619,7 +627,7 @@ def _ensure_msa_decode_scratch_buffers( ) return - kv_cache_manager = self.kv_cache_manager + kv_cache_manager = md.kv_cache_manager if kv_cache_manager is None: return @@ -652,13 +660,12 @@ def _ensure_prefill_n_valid_buffer(self, total_q: int, device: torch.device) -> """ buf = self._msa_prefill_n_valid_buf if buf is None or buf.numel() < total_q or buf.device != device: - cap = max(int(total_q), int(getattr(self, "max_num_tokens", 0) or 0), 1) + cap = max(int(total_q), int(getattr(self.metadata, "max_num_tokens", 0) or 0), 1) buf = torch.empty(cap, dtype=torch.int32, device=device) self._msa_prefill_n_valid_buf = buf return buf def prepare(self) -> None: - super().prepare() self._check_beam_width() # Everything below reads these. self._stage_host_lengths() @@ -676,10 +683,11 @@ def _check_beam_width(self) -> None: beam_width rows, so the block table and lengths handed to them would cover only the first 1 / beam_width of the batch. """ - if self.beam_width != 1: + md = self.metadata + if md.beam_width != 1: raise NotImplementedError( "MiniMax-M3 MSA attention does not support beam search, but this " - f"step has beam_width={self.beam_width}. Use beam_width=1 or the " + f"step has beam_width={md.beam_width}. Use beam_width=1 or the " "non-MSA MiniMax-M3 backend." ) @@ -700,13 +708,14 @@ def _set_decode_span(self) -> None: and may not exceed what _validate_decode_kernel_support settled the scorer for. """ + md = self.metadata self._msa_decode_span = None self._msa_max_kv_len = 0 qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu if qo_lens_cpu is None or kv_lens_cpu is None: return - row_first = int(self.num_contexts or 0) + row_first = int(md.num_contexts or 0) row_last = int(qo_lens_cpu.shape[0]) if row_first >= row_last: # Pure prefill: no generation row to describe. @@ -742,10 +751,11 @@ def _check_capture_is_pure_decode(self) -> None: fmha_sm100 against the addresses of a step that has passed. Decode needs no such plan, which is why capture is confined to it. """ - if self.is_cuda_graph and int(self.num_contexts or 0) > 0: + md = self.metadata + if md.is_cuda_graph and int(md.num_contexts or 0) > 0: raise RuntimeError( "MiniMax-M3 MSA attention captured a CUDA graph for a step with " - f"{int(self.num_contexts)} context requests. Only pure-decode steps " + f"{int(md.num_contexts)} context requests. Only pure-decode steps " "are graph-safe here; see _build_step_plans." ) @@ -771,9 +781,10 @@ def _msa_proxy_plan_rows(self) -> Optional[Tuple[int, int]]: * mixed: the context prefix only; * pure decode: None, no rows left for the proxy. """ + md = self.metadata span = self._msa_decode_span if span is None: - return (0, int(self.num_seqs)) + return (0, int(md.num_seqs)) return (0, span.row_first) if span.row_first > 0 else None def _msa_attn_plan_rows(self) -> Optional[Tuple[int, int]]: @@ -782,7 +793,8 @@ def _msa_attn_plan_rows(self) -> Optional[Tuple[int, int]]: The context prefix, which is the whole of what fmha_sm100 attends: the generation rows are the decode kernels' and are never planned for it. """ - num_contexts = int(self.num_contexts or 0) + md = self.metadata + num_contexts = int(md.num_contexts or 0) return (0, num_contexts) if num_contexts > 0 else None def _msa_index_kv_dtype(self) -> torch.dtype: @@ -791,7 +803,8 @@ def _msa_index_kv_dtype(self) -> torch.dtype: The CuTe DSL scorer requires index Q and K to match, and run_indexer casts Q to the cache dtype, so the cache decides what the scorer sees. """ - indexer_kv_dtype = str(getattr(self.kv_cache_manager, "indexer_kv_dtype", "bf16")) + md = self.metadata + indexer_kv_dtype = str(getattr(md.kv_cache_manager, "indexer_kv_dtype", "bf16")) return torch.float8_e4m3fn if indexer_kv_dtype == "fp8" else torch.bfloat16 def _cutedsl_indexer_supported( @@ -818,17 +831,19 @@ def _msa_kv_lens_may_change(self) -> bool: max_total_draft_tokens is set by the engine's update_spec_dec_param; the other checks mirror TrtllmAttentionMetadata's spec_active. """ - params = self.kv_cache_params - runtime_features = self.runtime_features + md = self.metadata + spec = md.find_ext(SpecDecodingMetadata) + params = md.kv_cache_params + runtime_features = md.runtime_features return bool( - self.max_total_draft_tokens - or self.draft_kv_cache_manager is not None - or self.is_spec_decoding_enabled + spec.max_total_draft_tokens + or md.draft_kv_cache_manager is not None + or spec.is_spec_decoding_enabled or (params is not None and params.num_extra_kv_tokens) or (runtime_features is not None and runtime_features.has_speculative_draft_tokens) ) - def on_update_kv_lens(self) -> None: + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: """Re-derive lengths, slots and valid-block counts from the corrected kv_lens_cuda. The overlap scheduler shortens kv_lens on device after prepare() staged @@ -847,23 +862,23 @@ def on_update_kv_lens(self) -> None: extend_ctx promotes speculative generation rows into the context prefix. Staged host bounds remain valid upper bounds as lengths shrink. """ - super().on_update_kv_lens() - if not self._msa_fields_ready or not self._msa_kv_lens_dynamic: + if spec_dec_step or not self._msa_fields_ready or not self._msa_kv_lens_dynamic: return - batch = int(self.num_seqs) - total_q = int(self.num_tokens) + md = self.metadata + batch = int(md.num_seqs) + total_q = int(md.num_tokens) if batch <= 0 or total_q <= 0: return # kv_lens_cuda is the attended length (no num_extra_kv_tokens), the same # domain as the staged bound and msa_seq_lens_cuda. - kv_true = torch.minimum(self.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch]) + kv_true = torch.minimum(md.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch]) self.msa_seq_lens_cuda[:batch].copy_(kv_true) if self.msa_cu_kv_lens is not None: self.msa_cu_kv_lens[0].zero_() torch.cumsum(kv_true, 0, out=self.msa_cu_kv_lens[1 : batch + 1]) qbr = self.msa_q_batch_row[:total_q].to(torch.long) - qo_dev = self.seq_lens_cuda[:batch] + qo_dev = md.seq_lens_cuda[:batch] # Position each query token attends up to. pos = kv_true[qbr] - qo_dev[qbr] + self.msa_q_intra[:total_q] @@ -872,7 +887,7 @@ def on_update_kv_lens(self) -> None: # pos % page. Positions only move down, so the page is one prepare() # staged for that request; a padding row (pos -1) keeps its own first # slot, as on the host. - page = int(self.kv_cache_manager.tokens_per_block) + page = int(md.kv_cache_manager.tokens_per_block) table = self.msa_block_table pos_slot = pos.clamp_min(0) block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( @@ -910,6 +925,7 @@ def _build_step_plans(self) -> None: (_check_capture_is_pure_decode), and so lets them be plain tuples of per-step tensors. """ + md = self.metadata self._msa_prefill_proxy_plan = None self._msa_prefill_proxy_kv_indices = None self._msa_prefill_gqa_plan = None @@ -921,21 +937,21 @@ def _build_step_plans(self) -> None: params = self._msa_params if params is None: return - num_index_heads = params.sharded_index_head_count(self.mapping) + num_index_heads = params.sharded_index_head_count(md.mapping) qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu if qo_lens_cpu is None or kv_lens_cpu is None or qo_offset_cpu is None: return batch = int(qo_lens_cpu.shape[0]) - page_size = int(self.kv_cache_manager.tokens_per_block) + page_size = int(md.kv_cache_manager.tokens_per_block) if self._msa_runs_no_fmha(): # Pure decode: no plan to build, but the scorer still writes its # scores into the proxy scratch and reads the valid-block count. self._ensure_msa_decode_scratch_buffers( num_index_heads=num_index_heads, - max_batch=int(self.max_num_sequences), - capture_graph=self.is_cuda_graph, + max_batch=int(md.max_num_sequences), + capture_graph=md.is_cuda_graph, # No proxy plan, so the worst case is the only bound available. required_max_k_tiles=self._msa_worst_case_max_k_tiles, ) @@ -949,7 +965,7 @@ def _build_step_plans(self) -> None: return fmha_sm100 = require_msa_module() - num_q_heads, num_kv_heads = params.sharded_head_counts(self.mapping) + num_q_heads, num_kv_heads = params.sharded_head_counts(md.mapping) # The main-attention GQA and dense plans need use_fp8_kvcache so the # inline sparse-prefill kernel selection matches an FP8 paged cache. # The proxy runs over the bf16 index-K cache, so it never needs the @@ -1001,7 +1017,7 @@ def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: # Sparse layers: kv_block_num=topk limits attention to top-k blocks. # NVFP4 sparse layers run the MSA CSR kernel, which takes no plan. self._msa_prefill_gqa_plan = plan_for( - None if self.kv_cache_manager.dtype == DataType.NVFP4 else attn_rows, + None if md.kv_cache_manager.dtype == DataType.NVFP4 else attn_rows, num_qo_heads=num_q_heads, num_kv_heads=num_kv_heads, kv_block_num=params.topk, @@ -1021,7 +1037,7 @@ def plan_for(rows: Optional[Tuple[int, int]], **plan_kwargs) -> Optional[tuple]: ) total_q = int(n_valid_host.shape[0]) if total_q > 0: - dev_buf = self._ensure_prefill_n_valid_buffer(total_q, _cache_device(self)) + dev_buf = self._ensure_prefill_n_valid_buffer(total_q, _cache_device(md)) dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True) self._msa_prefill_n_valid_blocks = dev_buf[:total_q] @@ -1031,12 +1047,13 @@ def _build_msa_fields(self) -> None: The page table and per-new-token cache slots are derived via the build_paged_kv_slot_mapping helper, then copied into the persistent buffers. The transient builder tensors are discarded. With speculative - decoding the inputs for on_update_kv_lens are staged as well. + decoding the inputs for on_kv_lens_changed are staged as well. """ + md = self.metadata self._msa_fields_ready = False if not self._msa_buffers_ready: return - request_ids = self.request_ids + request_ids = md.request_ids qo_lens_cpu = self.msa_qo_lens_cpu kv_lens_cpu = self.msa_kv_lens_cpu qo_offset_cpu = self.msa_qo_offset_cpu @@ -1048,8 +1065,8 @@ def _build_msa_fields(self) -> None: self.msa_out_cache_loc.fill_(-1) return - kv_cache_manager = self.kv_cache_manager - cache_device = _cache_device(self) + kv_cache_manager = md.kv_cache_manager + cache_device = _cache_device(md) page_size = int(kv_cache_manager.tokens_per_block) # Built in prepare() (outside capture), so these transients are @@ -1085,10 +1102,10 @@ def _build_msa_fields(self) -> None: # have to describe the same rows. The mapping emits one slot per new # token, so they agree unless a caller staged lengths this metadata's # seq_lens does not match. - if total_new_tokens != int(self.num_tokens): + if total_new_tokens != int(md.num_tokens): raise ValueError( f"MSA slot mapping covers {total_new_tokens} new tokens, but the " - f"step's token count is {int(self.num_tokens)}." + f"step's token count is {int(md.num_tokens)}." ) if kv_indices is not None and int(kv_indices.shape[0]) > self.msa_kv_indices.shape[0]: raise ValueError( @@ -1154,7 +1171,7 @@ def _build_msa_fields(self) -> None: # NVFP4 sparse layer's CSR kernel covers once the ported decode kernels # take the generation suffix. Sizing that call from the whole batch # would let a long generation row inflate its worklist. - context_rows = min(int(self.num_contexts or 0), batch_size) + context_rows = min(int(md.num_contexts or 0), batch_size) if context_rows > 0: prefix_kv = kv_lens_cpu[:context_rows].to(torch.int64) self._msa_context_prefix_bounds = ( @@ -1184,7 +1201,7 @@ def _build_msa_fields(self) -> None: self._msa_fields_ready = True return - # Inputs for on_update_kv_lens: each query token's (request row, offset + # Inputs for on_kv_lens_changed: each query token's (request row, offset # in request). Pinned and non-blocking so the copies do not synchronize # the stream. The slots themselves come from msa_block_table above. qo_long = qo_lens_cpu.to(torch.long) @@ -1200,7 +1217,7 @@ def _build_msa_fields(self) -> None: maybe_pin_memory(batch_row_cpu), non_blocking=True ) self.msa_q_intra[:total_new_tokens].copy_(maybe_pin_memory(intra_cpu), non_blocking=True) - # The staged lens are the upper bound on_update_kv_lens clamps to: the + # The staged lens are the upper bound on_kv_lens_changed clamps to: the # same attended lengths msa_seq_lens_cuda was just filled from, kept # apart so the clamp stays fixed however often the hook runs. self.msa_kv_lens_staged[:batch_size].copy_(kv_lens_cpu, non_blocking=True) @@ -1208,10 +1225,12 @@ def _build_msa_fields(self) -> None: def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor: """Return the paged index-K cache in the HND layout MSA consumes.""" - return self.kv_cache_manager.get_index_k_buffer(layer_idx) + md = self.metadata + return md.kv_cache_manager.get_index_k_buffer(layer_idx) def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: """Write the new-token index-K into the side cache at out_cache_loc.""" + md = self.metadata cache = self.msa_idx_k_cache(layer_idx) sparse_index_dim = int(cache.shape[-1]) num_tokens = int(idx_k.shape[0]) @@ -1221,7 +1240,7 @@ def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: idx_k.reshape(num_tokens, 1, sparse_index_dim), # idx_k arrives over the padded token extent; the live prefix is # where msa_out_cache_loc stops holding real slots. - int(self.num_tokens), + int(md.num_tokens), layout="HND", ) @@ -1258,7 +1277,8 @@ def msa_proxy_max_score_view( class MiniMaxM3MsaSparseAttention(TrtllmAttention): """MSA-backed MiniMax-M3 sparse attention.""" - Metadata = MiniMaxM3MsaSparseAttentionMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (MiniMaxM3MsaMetadata,) def __init__( self, @@ -1370,6 +1390,7 @@ def write_layer_caches( """ from .kernels.msa_scatter import fused_write_layer_caches, fused_write_layer_caches_nvfp4 + msa = metadata.ext(MiniMaxM3MsaMetadata) layer_idx = self.layer_idx manager = metadata.kv_cache_manager if getattr(manager, "is_nvfp4_layer", lambda _: False)(layer_idx): @@ -1377,14 +1398,14 @@ def write_layer_caches( raise RuntimeError("NVFP4 cache writes require quantization scales") buffers = manager.get_buffers(layer_idx, "HND") scales = manager.get_block_scale_buffers(layer_idx, "HND") - idx_cache = metadata.msa_idx_k_cache(layer_idx) if idx_k is not None else None + idx_cache = msa.msa_idx_k_cache(layer_idx) if idx_k is not None else None if not fused_write_layer_caches_nvfp4( buffers[:, 0], buffers[:, 1], scales[:, 0], scales[:, 1], idx_cache, - metadata.msa_out_cache_loc[: k.shape[0]], + msa.msa_out_cache_loc[: k.shape[0]], k, v, idx_k, @@ -1396,9 +1417,9 @@ def write_layer_caches( return buffers = metadata.kv_cache_manager.get_buffers(layer_idx, kv_layout="HND") k_view, v_view = buffers[:, 0], buffers[:, 1] - idx_cache = metadata.msa_idx_k_cache(layer_idx) if idx_k is not None else None + idx_cache = msa.msa_idx_k_cache(layer_idx) if idx_k is not None else None num_tokens = int(k.shape[0]) - out_cache_loc = metadata.msa_out_cache_loc[:num_tokens] + out_cache_loc = msa.msa_out_cache_loc[:num_tokens] if fused_write_layer_caches(k_view, v_view, idx_cache, out_cache_loc, k, v, idx_k): return num_kv_heads = int(k_view.shape[1]) @@ -1447,16 +1468,17 @@ def run_indexer( `idx_k_prewritten` marks that the fused per-layer cache write (write_layer_caches) already stored this layer's index-K. """ + msa = metadata.ext(MiniMaxM3MsaMetadata) config = self.m3_config idx_sm_scale = idx_sm_scale if idx_sm_scale is not None else config.sparse_index_dim**-0.5 num_tokens = int(idx_q.shape[0]) # The span says how the scoring is split between the CuTe DSL kernel # and the fmha_sm100 proxy pass. - span = metadata.msa_decode_span + span = msa.msa_decode_span # Preserve split column views without allowing an implicit copy. The # scorer and cache writer below both honor their source strides. idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) - idx_k_cache = metadata.msa_idx_k_cache(self.layer_idx) + idx_k_cache = msa.msa_idx_k_cache(self.layer_idx) configured_for_fp8 = self.indexer_kv_dtype == "fp8" expected_cache_dtype = torch.float8_e4m3fn if configured_for_fp8 else torch.bfloat16 if idx_k_cache.dtype != expected_cache_dtype: @@ -1483,7 +1505,7 @@ def run_indexer( # ahead of the proxy pass; write it here only when it did not. if not idx_k_prewritten: idx_k_view = idx_k.view(num_tokens, 1, config.sparse_index_dim) - metadata.msa_write_idx_k(self.layer_idx, idx_k_view) + msa.msa_write_idx_k(self.layer_idx, idx_k_view) # The FP8 indexer mirrors vLLM's unscaled E4M3 contract: normalized # index Q/K are cast directly and the proxy accumulates their QK scores # in FP32. Block ordering is invariant to the omitted positive scale. @@ -1509,25 +1531,25 @@ def run_indexer( # PhasedFmha derives the attention phase's token offset, so the # scorer and the decode kernels agree on the boundary. gen_first = num_tokens - decode_query_len * (row_last - ctx_rows) - block_table = metadata.msa_block_table[ctx_rows:row_last] - seq_lens_cuda = metadata.msa_seq_lens_cuda[ctx_rows:row_last] + block_table = msa.msa_block_table[ctx_rows:row_last] + seq_lens_cuda = msa.msa_seq_lens_cuda[ctx_rows:row_last] # One selection path, over the scratch prepare() staged for whichever # scorer owns the rows: a span starting at row 0 leaves the CuTe DSL # kernel every row and the graph-stable valid-block buffer, anything # else keeps a proxy plan and the per-step count. When neither is # present (a standalone test that skips prepare) select_blocks plans # inline and computes the valid-block count itself. - proxy_plan = metadata.msa_prefill_proxy_plan + proxy_plan = msa.msa_prefill_proxy_plan if proxy_plan is None and span is not None and span.row_first == 0: # No proxy plan to read max_k_tiles from for the contiguous score # view, so it is shaped to the worst case, which the scorer # accepts: it takes every score stride at runtime. - max_score = metadata.msa_proxy_max_score_view( - config.num_index_heads, metadata.msa_worst_case_max_k_tiles, num_tokens + max_score = msa.msa_proxy_max_score_view( + config.num_index_heads, msa.msa_worst_case_max_k_tiles, num_tokens ) - n_valid_blocks = metadata.msa_n_valid_blocks[:num_tokens] + n_valid_blocks = msa.msa_n_valid_blocks[:num_tokens] else: - n_valid_blocks = metadata.msa_prefill_n_valid_blocks + n_valid_blocks = msa.msa_prefill_n_valid_blocks if n_valid_blocks is not None: n_valid_blocks = n_valid_blocks[:num_tokens] # The scorer fills the buffer it is handed, shaped to the span's @@ -1536,9 +1558,9 @@ def run_indexer( # so cannot take a slice of this one. The span's tokens are at most # a decode step's worth, which is what the store was sized for. max_score = ( - metadata.msa_proxy_max_score_view( + msa.msa_proxy_max_score_view( config.num_index_heads, - metadata.msa_worst_case_max_k_tiles, + msa.msa_worst_case_max_k_tiles, num_tokens - gen_first, ) if span is not None @@ -1548,10 +1570,10 @@ def run_indexer( idx_q_view, idx_k_cache, idx_sm_scale=idx_sm_scale, - kv_indices=metadata.msa_prefill_proxy_kv_indices, - qo_lens_cpu=metadata.msa_qo_lens_cpu, - kv_lens_cpu=metadata.msa_kv_lens_cpu, - qo_offset_cpu=metadata.msa_qo_offset_cpu, + kv_indices=msa.msa_prefill_proxy_kv_indices, + qo_lens_cpu=msa.msa_qo_lens_cpu, + kv_lens_cpu=msa.msa_kv_lens_cpu, + qo_offset_cpu=msa.msa_qo_offset_cpu, proxy_plan=proxy_plan, max_score=max_score, n_valid_blocks=n_valid_blocks, @@ -1588,5 +1610,5 @@ def sparse_kv_predict( __all__ = [ "MiniMaxM3MsaSparseAttention", - "MiniMaxM3MsaSparseAttentionMetadata", + "MiniMaxM3MsaMetadata", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py index fe53b4de0fa0..04c88fe65ec0 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_backend.py @@ -68,10 +68,11 @@ AttentionMetadata, merge_attention_forward_args, ) +from ...trtllm import TrtllmAttentionMetadata from .common import _INIT_SCORE, _LOCAL_SCORE, MiniMaxM3SparseConfig, write_kv_slots from .triton_kernels import triton_block_max_score, triton_sparse_softmax from .triton_metadata import ( - MiniMaxM3AttentionMetadata, + MiniMaxM3TritonMetadata, MiniMaxM3TritonSparseAttentionMetadata, ensure_metadata_on_device, ) @@ -121,7 +122,7 @@ def _gather_paged_batched( ``req_to_token`` and ``slot_ids`` must already live on the cache device — the production - :class:`MiniMaxM3AttentionMetadata.prepare` builds them there, and + :class:`MiniMaxM3TritonMetadata.prepare` builds them there, and test callers do the same. ``.to(dtype=torch.long)`` below is a same-device dtype conversion (capture-safe). """ @@ -918,7 +919,8 @@ class MiniMaxM3SparseRuntimeBackend(AttentionBackend[AttentionMetadata]): :meth:`forward_sparse` when ``idx_q`` is supplied. """ - Metadata = MiniMaxM3AttentionMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (MiniMaxM3TritonMetadata,) def __init__( self, @@ -1014,7 +1016,7 @@ def forward_sparse( idx_k_view = idx_k.view(num_tokens, 1, sparse_index_dim) # Production paths build the M3 metadata on the cache - # device in ``MiniMaxM3AttentionMetadata.prepare`` (called + # device in ``MiniMaxM3TritonMetadata.prepare`` (called # outside any CUDA-graph capture window), and test paths # construct it directly on the desired device. So all # metadata tensors should already live on ``k_cache.device`` diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py index a1bea4f942fc..9762af4a5a89 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py @@ -8,8 +8,8 @@ dataclass with a CUDA-graph-safe :meth:`prepare`. * Helpers to migrate metadata across devices, build it from a real :class:`KVCacheManagerV2`, and pre-allocate CUDA-graph-stable buffers. - * :class:`MiniMaxM3AttentionMetadata` -- the :class:`AttentionMetadata` - subclass the pyexecutor wires into the M3 sparse layer's forward path. + * :class:`MiniMaxM3TritonMetadata` -- the :class:`AttentionMetadata` + extension the pyexecutor wires into the M3 sparse layer's forward path. """ from __future__ import annotations @@ -20,7 +20,7 @@ import torch -from ...trtllm import TrtllmAttentionMetadata +from ..params import SparseAlgorithmExtension from .common import build_paged_kv_slot_mapping @@ -435,7 +435,7 @@ def build_runtime_metadata_from_kv_manager( ``device`` selects the placement of the returned tensors. Default is ``seq_lens.device`` so existing callers (focused tests) keep their behaviour; the production - :class:`MiniMaxM3AttentionMetadata.prepare` path passes the cache + :class:`MiniMaxM3TritonMetadata.prepare` path passes the cache device explicitly so the forward never needs a CPU->GPU copy. ``static_buffers`` is the CUDA-graph-stable buffer dict produced by @@ -664,28 +664,24 @@ def build_runtime_metadata_from_kv_manager( return meta, out_cache_loc -class MiniMaxM3AttentionMetadata(TrtllmAttentionMetadata): - """:class:`TrtllmAttentionMetadata` that pre-builds MiniMax-M3 metadata. +class MiniMaxM3TritonMetadata(SparseAlgorithmExtension): + """:class:`TrtllmAttentionMetadata` extension that pre-builds MiniMax-M3 metadata. - Subclasses :class:`TrtllmAttentionMetadata` (like - ``DSAtrtllmAttentionMetadata``) so one-model Eagle3 draft layers can run - :class:`TrtllmAttention` on this metadata. + Extends :class:`TrtllmAttentionMetadata` so one-model Eagle3 draft layers + can run :class:`TrtllmAttention` on the same metadata. - Overrides :meth:`prepare` so the M3-sparse + :meth:`prepare` builds the M3-sparse :class:`MiniMaxM3TritonSparseAttentionMetadata` and the per-new-token - ``out_cache_loc`` are built **once per scheduler step**, on the - cache device, before the model forward runs. The result is - stored as ``self.minimax_m3 = {"metadata": m3_meta, - "out_cache_loc": out_cache_loc}`` so the model layer's - ``_dense_forward`` and ``_sparse_forward`` can read it without - any device migration. - - Test paths that build their own metadata can short-circuit by - attaching ``attn_metadata.minimax_m3`` directly before calling - the forward; those paths do not go through :meth:`prepare`. + ``out_cache_loc`` **once per scheduler step**, on the cache device, + before the model forward runs, as ``sparse_metadata`` and + ``out_cache_loc``, so the model layer's ``_dense_forward`` and + ``_sparse_forward`` can read them without any device migration. """ - minimax_m3: Optional[dict] = None + # This step's M3-sparse metadata and per-new-token cache slots; None until + # prepare() builds them. + sparse_metadata: Optional[MiniMaxM3TritonSparseAttentionMetadata] = None + out_cache_loc: Optional[torch.Tensor] = None # Lazily allocated dict of persistent device buffers used to keep # ``MiniMaxM3TritonSparseAttentionMetadata`` tensor addresses stable # across CUDA-graph capture/replay. None until the first @@ -703,7 +699,7 @@ def _maybe_get_m3_static_buffers( needed and caches it on ``self._m3_static_buffers``. We allocate the buffers under two conditions: - * ``self.is_cuda_graph`` is True -- the captured graph + * the metadata's ``is_cuda_graph`` is True -- the captured graph requires stable ``data_ptr()`` across replays; OR * the previous prepare() already allocated buffers -- we keep using them so the algorithm sees the same @@ -714,9 +710,8 @@ def _maybe_get_m3_static_buffers( Returns ``None`` when no static buffers should be used (e.g. eager-only test paths that rely on per-call allocations). """ - need_static = ( - bool(getattr(self, "is_cuda_graph", False)) or self._m3_static_buffers is not None - ) + md = self.metadata + need_static = bool(md.is_cuda_graph) or self._m3_static_buffers is not None if not need_static: return None if self._m3_static_buffers is not None: @@ -740,20 +735,16 @@ def _maybe_get_m3_static_buffers( # Caller-provided hints used by the lazy allocator below # when it sizes the persistent buffers on the first real # prepare() call. - "max_num_sequences_hint": int( - getattr(self, "max_num_sequences", None) or self.max_num_requests - ), + "max_num_sequences_hint": int(md.max_num_sequences or md.max_num_requests), "max_num_tokens_hint": int( - getattr(self, "max_num_tokens", None) - or (int(getattr(self, "max_num_sequences", None) or self.max_num_requests)) + md.max_num_tokens or md.max_num_sequences or md.max_num_requests ), } self._m3_static_buffers = placeholder return placeholder def prepare(self) -> None: - super().prepare() - + md = self.metadata # Always rebuild the M3 metadata block on each prepare() # call so it reflects the current scheduler step's seq_lens # / request_ids / num_cached_tokens. Production @@ -768,22 +759,23 @@ def prepare(self) -> None: # over ``req_to_token``/``slot_ids`` reads from freed warmup # memory and either produces wrong tokens or fires # ``Indexing.cu:1515`` ``srcIndex < srcSelectDimSize``. - self.minimax_m3 = None + self.sparse_metadata = None + self.out_cache_loc = None # Production path: build the M3 metadata from the standard # AttentionMetadata fields. Requires kv_cache_manager + the # M3 sparse-cache contract. - kv_cache_manager = getattr(self, "kv_cache_manager", None) + kv_cache_manager = md.kv_cache_manager if kv_cache_manager is None or not hasattr(kv_cache_manager, "get_index_k_buffer"): # Not an M3 KV cache manager: nothing to build. The # forward path will raise a clear error if the M3 # backend ends up dispatched without the M3 cache. return - request_ids = getattr(self, "request_ids", None) - seq_lens = self.seq_lens + request_ids = md.request_ids + seq_lens = md.seq_lens if request_ids is None or seq_lens is None: return - num_contexts = int(getattr(self, "num_contexts", 0) or 0) + num_contexts = int(md.num_contexts or 0) batch_size = int(seq_lens.shape[0]) if batch_size == 0: return @@ -797,14 +789,14 @@ def prepare(self) -> None: cache_device = torch.device(f"cuda:{torch.cuda.current_device()}") seq_lens_cpu = ( - getattr(self, "seq_lens_cpu", None) - if hasattr(self, "seq_lens_cpu") + getattr(md, "seq_lens_cpu", None) + if hasattr(md, "seq_lens_cpu") else seq_lens.detach().to("cpu") ) if seq_lens_cpu is None: seq_lens_cpu = seq_lens.detach().to("cpu") - kv_cache_params = getattr(self, "kv_cache_params", None) + kv_cache_params = md.kv_cache_params num_cached_per_seq = ( kv_cache_params.num_cached_tokens_per_seq if kv_cache_params is not None @@ -860,29 +852,25 @@ def prepare(self) -> None: static_buffers=static_buffers, ) - self.minimax_m3 = { - "metadata": m3_meta, - "out_cache_loc": out_cache_loc, - } + self.sparse_metadata = m3_meta + self.out_cache_loc = out_cache_loc - def on_update_kv_lens(self) -> None: + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: """Re-derive the M3 attachment from the corrected ``kv_lens_cuda``. With the overlap scheduler and speculative decoding, prepare() runs with optimistic lengths and the engine corrects ``kv_lens_cuda`` on - device before calling this (as DSAtrtllmAttentionMetadata does). + device before calling ``on_update_kv_lens``. Device-only and idempotent. ``seq_lens_cpu`` / ``max_seqlen_k`` keep the optimistic values; they only bound widths the kernels mask by ``seq_lens``. """ - super().on_update_kv_lens() - attachment = self.minimax_m3 - if not attachment: + meta = self.sparse_metadata + if spec_dec_step or meta is None: return - meta = attachment["metadata"] - out_cache_loc = attachment["out_cache_loc"] + out_cache_loc = self.out_cache_loc batch = int(meta.slot_ids.shape[0]) - kv_lens = self.kv_lens_cuda[:batch] + kv_lens = self.metadata.kv_lens_cuda[:batch] meta.seq_lens[:batch].copy_(kv_lens) if meta.is_prefill: # Only the K-side prefix moves with rejections; the Q-side layout @@ -904,7 +892,7 @@ def on_update_kv_lens(self) -> None: __all__ = [ - "MiniMaxM3AttentionMetadata", + "MiniMaxM3TritonMetadata", "MiniMaxM3TritonSparseAttentionMetadata", "allocate_minimax_m3_static_buffers", "build_runtime_metadata_from_kv_manager", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 065ac68a503e..bcffc83f572f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -16,10 +16,12 @@ import os from dataclasses import dataclass -from typing import Literal, Optional, Protocol, runtime_checkable +from typing import Literal, Optional import torch +from ..extensions import MetadataExtension + _INDEXER_MQA_LOGITS_DEFAULT_ELEM_BUDGET = 1 << 31 _INDEXER_MQA_LOGITS_BYTES_PER_ELEMENT = 4 @@ -43,22 +45,30 @@ def get_indexer_mqa_logits_workspace_bytes( return elem_budget * _INDEXER_MQA_LOGITS_BYTES_PER_ELEMENT -@runtime_checkable -class MTPIndexShareMetadata(Protocol): - """Draft-loop state a sparse backend exposes so MTP can share one selection. +class MTPIndexShareExtension(MetadataExtension): + """Draft-loop state a sparse algorithm exposes so MTP can share one selection. - The three calls are used together for one draft loop, so a backend that + The three calls are used together for one draft loop, so an algorithm that implements only some of them cannot serve the reuse at all. """ def set_in_mtp_draft_loop(self, active: bool) -> None: """Mark whether a draft loop is running.""" + raise NotImplementedError def set_mtp_num_accepted(self, num_accepted: Optional[torch.Tensor]) -> None: """Supply the per-request accepted counts used to pick the capture row.""" + raise NotImplementedError def set_skip_topk(self, skip: bool) -> None: """Ask the indexer to reuse its captured selection instead of scoring.""" + raise NotImplementedError + + def reset_mtp_index_share(self) -> None: + """Clear the draft-loop state the indexer reads.""" + self.set_skip_topk(False) + self.set_in_mtp_draft_loop(False) + self.set_mtp_num_accepted(None) def use_self_sampling_gvr( @@ -96,6 +106,10 @@ class SparseMetadataParams: """Base parameters for sparse attention metadata.""" +class SparseAlgorithmExtension(MetadataExtension): + """Base of a sparse algorithm's metadata state; a metadata carries at most one.""" + + @dataclass(kw_only=True, slots=True) class SparseBackendForwardArgs: """Sparse inputs passed from an attention module to its backend.""" diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py index bb11190acf00..362a5e34a6be 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py @@ -7,7 +7,7 @@ _EXPORTS = { "QSATrtllmAttention": "backend", - "QSAAttentionMetadata": "metadata", + "QSAMetadata": "metadata", "QSASparseMetadataParams": "params", "QSASparseParams": "params", } diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/backend.py index 3bb75521893c..aeb272ab7274 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/backend.py @@ -2,9 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """TRT-LLM backend wrapper for QSA sparse metadata.""" -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata -from .metadata import QSAAttentionMetadata +from .metadata import QSAMetadata class QSATrtllmAttention(TrtllmAttention): @@ -14,7 +14,8 @@ class QSATrtllmAttention(TrtllmAttention): this class only installs metadata required by both paths. """ - Metadata = QSAAttentionMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (QSAMetadata,) __all__ = ["QSATrtllmAttention"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/indexer.py index dd840eb93944..f93f9497ec2c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/indexer.py @@ -22,12 +22,12 @@ QSA_INDEX_K_CACHE_DTYPE, QSA_POSITION_COORDINATE_AXES, ) +from .metadata import QSAMetadata from .params import QSASparseParams if TYPE_CHECKING: from tensorrt_llm._torch.attention.attention import Attention - - from .metadata import QSAAttentionMetadata + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata _GROUP_START_MEMBER = 0 @@ -79,17 +79,21 @@ def _position_coordinates(position_ids: torch.Tensor, num_tokens: int) -> torch. def _logical_to_pages( - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", request_indices: torch.Tensor, logical_positions: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: tokens_per_block = metadata.kv_cache_manager.tokens_per_block logical_positions = logical_positions.to(torch.long) page_columns = logical_positions // tokens_per_block - pages = metadata.qsa_block_table[ - request_indices.to(torch.long), - page_columns, - ].to(torch.long) + pages = ( + metadata.ext(QSAMetadata) + .block_table[ + request_indices.to(torch.long), + page_columns, + ] + .to(torch.long) + ) within = logical_positions % tokens_per_block return pages, within @@ -281,7 +285,7 @@ def project_and_update_cache( hidden_states: torch.Tensor, position_ids: torch.Tensor, layer_idx: int, - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", ) -> torch.Tensor: """Project the index Q/K and update its paged side cache. @@ -290,6 +294,7 @@ def project_and_update_cache( speculative verification, CPU execution, and unsupported RoPE layouts retain the reference path below. """ + qsa = metadata.find_ext(QSAMetadata) num_tokens = hidden_states.shape[0] qk = self.index_qk_proj(hidden_states) q_width = self.params.index_n_heads * self.params.index_head_dim @@ -337,9 +342,9 @@ def project_and_update_cache( q=q_raw, token_k=token_k, position_coordinates=coordinates, - request_indices=metadata.qsa_req_idx_per_token[:num_tokens], - logical_positions=metadata.qsa_logical_positions[:num_tokens], - block_table=metadata.qsa_block_table, + request_indices=qsa.req_idx_per_token[:num_tokens], + logical_positions=qsa.logical_positions[:num_tokens], + block_table=qsa.block_table, index_cache=index_cache, position_cache=position_cache, q_norm_weight=self.q_layernorm.weight, @@ -366,7 +371,7 @@ def update_cache_and_compress( layer_idx: int, token_k: torch.Tensor, position_coordinates: torch.Tensor, - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", ) -> None: """Store raw index keys and replace completed group anchors. @@ -376,9 +381,10 @@ def update_cache_and_compress( index side cache is overwritten; the main K/V cache retains every raw token required by exact sparse attention. """ + qsa = metadata.find_ext(QSAMetadata) num_tokens = token_k.shape[0] - req_idx = metadata.qsa_req_idx_per_token[:num_tokens] - logical = metadata.qsa_logical_positions[:num_tokens] + req_idx = qsa.req_idx_per_token[:num_tokens] + logical = qsa.logical_positions[:num_tokens] pages, within = _logical_to_pages(metadata, req_idx, logical) index_cache = metadata.kv_cache_manager.get_index_k_buffer(layer_idx) position_cache = metadata.kv_cache_manager.get_qsa_position_buffer() @@ -418,7 +424,7 @@ def update_cache_and_compress( triton_qsa_prefill_compress( logical_positions=logical, request_indices=req_idx, - block_table=metadata.qsa_block_table, + block_table=qsa.block_table, index_cache=index_cache, position_cache=position_cache, k_norm_weight=self.k_layernorm.weight, @@ -512,7 +518,7 @@ def _capture_speculative_cache_state( within: torch.Tensor, index_cache: torch.Tensor, position_cache: torch.Tensor, - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", ) -> None: """Snapshot QSA side-cache entries that target verification may reject.""" self._pending_speculative_cache = None @@ -524,7 +530,7 @@ def _capture_speculative_cache_state( return # Mixed-IFB forwards are eager. Derive this branch once from the host # query lengths rather than synchronizing a device reduction per layer. - if not metadata.is_cuda_graph and not metadata.qsa_needs_speculative_snapshot: + if not metadata.is_cuda_graph and not metadata.ext(QSAMetadata).needs_speculative_snapshot: return seq_lens = metadata.seq_lens_cuda[:num_seqs].to(torch.long) generation_lens = seq_lens[num_contexts:] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index f689b1af5320..3f65792ec96c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -1,164 +1,99 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Attention metadata for QSA sparse attention.""" +"""Attention metadata state for QSA sparse attention.""" -from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Optional import torch -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.features import DraftKVMetadata +from tensorrt_llm._torch.attention.backends.sparse.params import ( + MTPIndexShareExtension, + SparseAlgorithmExtension, +) from .constants import QSA_KEY_ROLE_INDEX from .kernels import triton_qsa_unscale_block_table from .params import QSASparseMetadataParams -# Keep inherited dataclass fields while intercepting one QSA-only constructor arg. -@dataclass(init=False) -class QSAAttentionMetadata(TrtllmAttentionMetadata): - """Extend TRTLLM metadata with graph-stable QSA page and row mappings.""" - - sparse_metadata_params: Optional[QSASparseMetadataParams] = None - qsa_needs_speculative_snapshot: bool = field( - init=False, - default=False, - repr=False, - ) - qsa_has_local_layers: bool = field( - init=False, - default=False, - repr=False, - ) - - def __init__(self, *args: object, **kwargs: object) -> None: - # Forwards the parent dataclass's full field set, so the signature - # stays open rather than restating it. +class QSAMetadata(SparseAlgorithmExtension, MTPIndexShareExtension): + """Graph-stable QSA page and row mappings of a TRTLLM metadata.""" + + def allocate(self, buffers) -> None: """Allocate the fixed-shape QSA buffers updated before each graph replay.""" - sparse_attention_config = kwargs.pop("sparse_attention_config", None) - if kwargs.get("sparse_metadata_params") is None and sparse_attention_config is not None: - kwargs["sparse_metadata_params"] = sparse_attention_config.to_sparse_metadata_params() - # The generated parent initializer dispatches to this class's - # ``__post_init__``; keep QSA validation there so both constructor forms - # (lowered params or the legacy user config) enforce the same contract. - super().__init__(*args, **kwargs) - - def __post_init__(self) -> None: from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager - super().__post_init__() + md = self.metadata + self.needs_speculative_snapshot = False # Draft-loop state; set by the speculative driver, read by the indexer. - self.qsa_in_mtp_draft_loop = False - self.qsa_indexer_skip_topk = False - self.qsa_mtp_num_accepted: Optional[torch.Tensor] = None - self.qsa_shared_topk_indices: Optional[torch.Tensor] = None - self.qsa_shared_topk_visible_blocks: Optional[torch.Tensor] = None - self.qsa_shared_topk_captured_slots: set[int] = set() - self.qsa_draft_block_table: Optional[torch.Tensor] = None - if not isinstance(self.kv_cache_manager, Qwen4ExpHybridCacheManager): + self.in_mtp_draft_loop = False + self.indexer_skip_topk = False + self.mtp_num_accepted: Optional[torch.Tensor] = None + self.shared_topk_indices: Optional[torch.Tensor] = None + self.shared_topk_visible_blocks: Optional[torch.Tensor] = None + self.shared_topk_captured_slots: set[int] = set() + self.draft_block_table: Optional[torch.Tensor] = None + if not isinstance(md.kv_cache_manager, Qwen4ExpHybridCacheManager): raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManager") - if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): + params = md.sparse_metadata_params + if not isinstance(params, QSASparseMetadataParams): raise TypeError("QSA sparse attention requires QSASparseMetadataParams") - self.qsa_has_local_layers = self.kv_cache_manager.qsa_position_layer_id is not None - if not self.qsa_has_local_layers: + self.params = params + self.has_local_layers = md.kv_cache_manager.qsa_position_layer_id is not None + if not self.has_local_layers: # A PP rank may own only recurrent layers. It still uses the hybrid - # manager and parent metadata, but needs no QSA page views. + # manager and core metadata, but needs no QSA page views. return - capture_graph = self.is_cuda_graph - buffers = self.cuda_graph_buffers - # The sparse kernels index raw V2 pool pages, while the parent metadata + # The sparse kernels index raw V2 pool pages, while the core metadata # stores attention-operator encoded offsets. Keep an unscaled table. - self.qsa_block_table = self.get_empty( - buffers, - ( - self.max_num_sequences, - self.kv_cache_manager.max_blocks_per_seq, - ), - cache_name="qsa_block_table", - dtype=torch.int32, - capture_graph=capture_graph, + self.block_table = md.get_buffer( + (md.max_num_sequences, md.kv_cache_manager.max_blocks_per_seq), + torch.int32, + "qsa_block_table", ) - self._qsa_attention_pool_id, self._qsa_page_index_scale = ( - self.kv_cache_manager.get_qsa_attention_pool_layout() + self._attention_pool_id, self._page_index_scale = ( + md.kv_cache_manager.get_qsa_attention_pool_layout() ) - self._init_qsa_draft_block_table() + self._init_draft_block_table() # Packed-row mapping shared by compression, selection, and sparse GQA. - self.qsa_req_idx_per_token = self.get_empty( - buffers, - (self.max_num_tokens,), - cache_name="qsa_req_idx_per_token", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.qsa_logical_positions = self.get_empty( - buffers, - (self.max_num_tokens,), - cache_name="qsa_logical_positions", - dtype=torch.int64, - capture_graph=capture_graph, + self.req_idx_per_token = md.get_buffer( + (md.max_num_tokens,), torch.int32, "qsa_req_idx_per_token" ) - self.qsa_sequence_lengths = self.get_empty( - buffers, - (self.max_num_tokens,), - cache_name="qsa_sequence_lengths", - dtype=torch.int32, - capture_graph=capture_graph, + self.logical_positions = md.get_buffer( + (md.max_num_tokens,), torch.int64, "qsa_logical_positions" ) - self.qsa_visible_blocks = self.get_empty( - buffers, - (self.max_num_tokens,), - cache_name="qsa_visible_blocks", - dtype=torch.int32, - capture_graph=capture_graph, + self.sequence_lengths = md.get_buffer( + (md.max_num_tokens,), torch.int32, "qsa_sequence_lengths" ) + self.visible_blocks = md.get_buffer((md.max_num_tokens,), torch.int32, "qsa_visible_blocks") # Reuse graph-stable row IDs instead of allocating arange each step. - self._qsa_token_arange = self.get_empty( - buffers, - (self.max_num_tokens,), - cache_name="qsa_token_arange", - dtype=torch.int64, - capture_graph=capture_graph, - ) - self._qsa_token_arange.copy_( + self._token_arange = md.get_buffer((md.max_num_tokens,), torch.int64, "qsa_token_arange") + self._token_arange.copy_( torch.arange( - self.max_num_tokens, - device=self._qsa_token_arange.device, + md.max_num_tokens, + device=self._token_arange.device, dtype=torch.int64, ) ) # The leading zero turns request lengths into packed-row boundaries. - self._qsa_cu_seq_lens = self.get_empty( - buffers, - (self.max_num_sequences + 1,), - cache_name="qsa_cu_seq_lens", - dtype=torch.int64, - capture_graph=capture_graph, + self._cu_seq_lens = md.get_buffer( + (md.max_num_sequences + 1,), torch.int64, "qsa_cu_seq_lens" ) - self._qsa_cu_seq_lens.zero_() + self._cu_seq_lens.zero_() # TopK writes in place; caller ownership keeps the address graph-stable. - self.qsa_topk_indices = self.get_empty( - buffers, - ( - self.max_num_tokens, - self.sparse_metadata_params.block_topk, - ), - cache_name="qsa_topk_indices", - dtype=torch.int32, - capture_graph=capture_graph, + self.topk_indices = md.get_buffer( + (md.max_num_tokens, params.block_topk), torch.int32, "qsa_topk_indices" ) # Per-row Top-K lower bounds; QSA searches every row from zero. - self.qsa_topk_row_starts = self.get_empty( - buffers, - (self.max_num_tokens,), - cache_name="qsa_topk_row_starts", - dtype=torch.int32, - capture_graph=capture_graph, + self.topk_row_starts = md.get_buffer( + (md.max_num_tokens,), torch.int32, "qsa_topk_row_starts" ) - self.qsa_topk_row_starts.zero_() + self.topk_row_starts.zero_() # One captured selection per sparse layer and request, reused by the # MTP draft loop. The layer axis is required because each layer runs @@ -166,32 +101,24 @@ def __post_init__(self) -> None: # selections differ. ``shared_topk_visible_blocks`` records how many # compressed groups were complete when the row was captured, so the # reuse can tell which groups formed afterwards. - if self.sparse_metadata_params.mtp_index_share: + if params.mtp_index_share: # Slots are numbered by the bound manager; the draft loop binds the # separate draft manager. - num_slots = len(self.kv_cache_manager.qsa_shared_topk_slots) - if self.qsa_draft_block_table is not None: - num_slots = max(num_slots, len(self.draft_kv_cache_manager.qsa_shared_topk_slots)) - self.qsa_shared_topk_indices = self.get_empty( - buffers, - ( - num_slots, - self.max_num_sequences, - self.sparse_metadata_params.block_topk, - ), - cache_name="qsa_shared_topk_indices", - dtype=torch.int32, - capture_graph=capture_graph, + num_slots = len(md.kv_cache_manager.qsa_shared_topk_slots) + if self.draft_block_table is not None: + num_slots = max(num_slots, len(md.draft_kv_cache_manager.qsa_shared_topk_slots)) + self.shared_topk_indices = md.get_buffer( + (num_slots, md.max_num_sequences, params.block_topk), + torch.int32, + "qsa_shared_topk_indices", ) - self.qsa_shared_topk_visible_blocks = self.get_empty( - buffers, - (num_slots, self.max_num_sequences), - cache_name="qsa_shared_topk_visible_blocks", - dtype=torch.int32, - capture_graph=capture_graph, + self.shared_topk_visible_blocks = md.get_buffer( + (num_slots, md.max_num_sequences), + torch.int32, + "qsa_shared_topk_visible_blocks", ) - def _init_qsa_draft_block_table(self) -> None: + def _init_draft_block_table(self) -> None: """Allocate the page table of a separate one-model draft KV cache. The draft layers' pages live in the draft manager's pools, which the @@ -202,85 +129,83 @@ def _init_qsa_draft_block_table(self) -> None: """ from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager - draft_manager = self.draft_kv_cache_manager + md = self.metadata + draft_manager = md.draft_kv_cache_manager if ( not isinstance(draft_manager, Qwen4ExpHybridCacheManager) or draft_manager.qsa_position_layer_id is None ): return - self.qsa_draft_block_table = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences, draft_manager.max_blocks_per_seq), - cache_name="qsa_draft_block_table", - dtype=torch.int32, - capture_graph=self.is_cuda_graph, + self.draft_block_table = md.get_buffer( + (md.max_num_sequences, draft_manager.max_blocks_per_seq), + torch.int32, + "qsa_draft_block_table", ) - self._qsa_draft_attention_pool_id, self._qsa_draft_page_index_scale = ( + self._draft_attention_pool_id, self._draft_page_index_scale = ( draft_manager.get_qsa_attention_pool_layout() ) def set_skip_topk(self, skip: bool) -> None: """Mark a draft step that reuses the captured selection.""" - self.qsa_indexer_skip_topk = skip + self.indexer_skip_topk = skip def set_in_mtp_draft_loop(self, active: bool) -> None: """Bracket the MTP draft loop so the indexer captures once and reuses.""" - self.qsa_in_mtp_draft_loop = active + self.in_mtp_draft_loop = active # Readiness is per layer slot, not global: a layer may reuse only the # row it captured itself, and the layers do not all capture in the same # pass. A loop whose first step is a context pass captures nothing, so # the next step captures layer by layer. - self.qsa_shared_topk_captured_slots.clear() + self.shared_topk_captured_slots.clear() def set_mtp_num_accepted(self, num_accepted: Optional[torch.Tensor]) -> None: """Supply the per-request accepted counts used to pick the capture row.""" - self.qsa_mtp_num_accepted = num_accepted + self.mtp_num_accepted = num_accepted - def _refresh_qsa_token_mapping(self) -> None: + def _refresh_token_mapping(self) -> None: """Map packed rows to requests, logical positions, and causal limits.""" - num_seqs = self.num_seqs - num_tokens = self.num_tokens + md = self.metadata + num_seqs = md.num_seqs + num_tokens = md.num_tokens if num_seqs <= 0 or num_tokens <= 0: return - seq_lens = self.seq_lens_cuda[:num_seqs] - if self.num_contexts == 0 and num_tokens == num_seqs: + seq_lens = md.seq_lens_cuda[:num_seqs] + if md.num_contexts == 0 and num_tokens == num_seqs: from .kernels import triton_qsa_decode_token_mapping triton_qsa_decode_token_mapping( - kv_lens=self.kv_lens_cuda_runtime[:num_seqs], + kv_lens=md.kv_lens_cuda_runtime[:num_seqs], seq_lens=seq_lens, - request_indices=self.qsa_req_idx_per_token[:num_tokens], - logical_positions=self.qsa_logical_positions[:num_tokens], - sequence_lengths=self.qsa_sequence_lengths[:num_tokens], - visible_blocks=self.qsa_visible_blocks[:num_tokens], - compress_ratio=self.sparse_metadata_params.compress_ratio, + request_indices=self.req_idx_per_token[:num_tokens], + logical_positions=self.logical_positions[:num_tokens], + sequence_lengths=self.sequence_lengths[:num_tokens], + visible_blocks=self.visible_blocks[:num_tokens], + compress_ratio=self.params.compress_ratio, ) return torch.cumsum( seq_lens, dim=0, dtype=torch.int64, - out=self._qsa_cu_seq_lens[1 : num_seqs + 1], + out=self._cu_seq_lens[1 : num_seqs + 1], ) - token_ids = self._qsa_token_arange[:num_tokens] + token_ids = self._token_arange[:num_tokens] req_idx = torch.searchsorted( - self._qsa_cu_seq_lens[1 : num_seqs + 1], + self._cu_seq_lens[1 : num_seqs + 1], token_ids, right=True, ) - self.qsa_req_idx_per_token[:num_tokens].copy_(req_idx.to(torch.int32)) - seq_starts = self._qsa_cu_seq_lens[:num_seqs] - cached_lens = self.kv_lens_cuda_runtime[:num_seqs].to(torch.int64) - seq_lens.to( - torch.int64 - ) + self.req_idx_per_token[:num_tokens].copy_(req_idx.to(torch.int32)) + seq_starts = self._cu_seq_lens[:num_seqs] + cached_lens = md.kv_lens_cuda_runtime[:num_seqs].to(torch.int64) - seq_lens.to(torch.int64) logical = cached_lens[req_idx] + token_ids - seq_starts[req_idx] - self.qsa_logical_positions[:num_tokens].copy_(logical) - self.qsa_sequence_lengths[:num_tokens].copy_(self.kv_lens_cuda_runtime[:num_seqs][req_idx]) - self.qsa_visible_blocks[:num_tokens].copy_( - ((logical + 1) // self.sparse_metadata_params.compress_ratio).to(torch.int32) + self.logical_positions[:num_tokens].copy_(logical) + self.sequence_lengths[:num_tokens].copy_(md.kv_lens_cuda_runtime[:num_seqs][req_idx]) + self.visible_blocks[:num_tokens].copy_( + ((logical + 1) // self.params.compress_ratio).to(torch.int32) ) - def _unscale_qsa_block_table( + def _unscale_block_table( self, block_offsets: torch.Tensor, pool_id: int, @@ -288,7 +213,7 @@ def _unscale_qsa_block_table( block_table: torch.Tensor, ) -> None: """Decode one manager's V2 attention offsets into raw page IDs.""" - num_seqs = self.num_seqs + num_seqs = self.metadata.num_seqs if num_seqs <= 0: return triton_qsa_unscale_block_table( @@ -297,94 +222,92 @@ def _unscale_qsa_block_table( page_index_scale=page_index_scale, ) - def _refresh_qsa_block_table(self) -> None: + def _refresh_block_table(self) -> None: """Decode the bound manager's offsets into the table QSA kernels read. - While ``prepare_for_draft_forward()`` is in effect every operand here is - the draft manager's, so this rebuilds the draft table. + While a draft forward is in effect every operand here is the draft + manager's, so this rebuilds the draft table. """ - self._unscale_qsa_block_table( - self.kv_cache_block_offsets, - self._qsa_attention_pool_id, - self._qsa_page_index_scale, - self.qsa_block_table, + self._unscale_block_table( + self.metadata.kv_cache_block_offsets, + self._attention_pool_id, + self._page_index_scale, + self.block_table, ) - def _refresh_qsa_draft_block_table(self) -> None: + def _refresh_draft_block_table(self) -> None: """Decode the separate draft manager's offsets into the draft table.""" - if self.qsa_draft_block_table is None: + if self.draft_block_table is None: return - self._unscale_qsa_block_table( - self.draft_kv_cache_block_offsets, - self._qsa_draft_attention_pool_id, - self._qsa_draft_page_index_scale, - self.qsa_draft_block_table, + self._unscale_block_table( + self.metadata.ext(DraftKVMetadata).kv_cache_block_offsets, + self._draft_attention_pool_id, + self._draft_page_index_scale, + self.draft_block_table, ) - def _refresh_qsa_speculative_snapshot_state(self) -> None: + def _refresh_speculative_snapshot_state(self) -> None: """Cache whether eager generation contains multi-token verification rows. Host lengths are already available during prepare; using them avoids a device reduction and synchronization in every QSA layer. """ - generation_lens = self.seq_lens[self.num_contexts : self.num_seqs] - self.qsa_needs_speculative_snapshot = any(length > 1 for length in generation_lens.tolist()) + md = self.metadata + generation_lens = md.seq_lens[md.num_contexts : md.num_seqs] + self.needs_speculative_snapshot = any(length > 1 for length in generation_lens.tolist()) - def _refresh_qsa_device_state(self) -> None: + def _refresh_device_state(self) -> None: """Rebuild the device-side block table and token mapping for this step.""" - if not self.qsa_has_local_layers: + if not self.has_local_layers: return - self._refresh_qsa_block_table() - self._refresh_qsa_draft_block_table() - self._refresh_qsa_token_mapping() + self._refresh_block_table() + self._refresh_draft_block_table() + self._refresh_token_mapping() - def prepare_for_draft_forward(self) -> None: + def draft_swap(self) -> Any: """Point the QSA page table at the separate draft manager's pages. The caller has already swapped ``kv_cache_manager`` and ``kv_cache_block_offsets`` to the draft manager's. The draft table was decoded in ``prepare()``, so this only rebinds fields and launches - nothing; ``restore_draft_swaps()`` restores the target bindings. The - row mapping is manager-independent and stays as it is. + nothing. The row mapping is manager-independent and stays as it is. """ - if self.qsa_draft_block_table is None: - return - if self.kv_cache_manager is not self.draft_kv_cache_manager: + if self.draft_block_table is None: + return None + md = self.metadata + if md.kv_cache_manager is not md.draft_kv_cache_manager: raise RuntimeError( "QSA draft forward expects the metadata's draft KV cache manager " - f"to be bound; got {type(self.kv_cache_manager).__name__}" + f"to be bound; got {type(md.kv_cache_manager).__name__}" ) - self.swap_for_draft("qsa_block_table", self.qsa_draft_block_table) - self.swap_for_draft("_qsa_attention_pool_id", self._qsa_draft_attention_pool_id) - self.swap_for_draft("_qsa_page_index_scale", self._qsa_draft_page_index_scale) + return self.swap_attrs_for_draft( + { + "block_table": self.draft_block_table, + "_attention_pool_id": self._draft_attention_pool_id, + "_page_index_scale": self._draft_page_index_scale, + } + ) + + def draft_restore(self, saved: Any) -> None: + if saved is not None: + self.restore_swapped_attrs(saved) def prepare(self) -> None: """Refresh host decisions, then rebuild graph-stable device mappings.""" - super().prepare() - if not self.qsa_has_local_layers: - return - self._refresh_qsa_speculative_snapshot_state() - self._refresh_qsa_device_state() - - def on_update_kv_lens(self) -> None: - """Refresh row limits after speculative decoding advances KV lengths.""" - super().on_update_kv_lens() - if not self.qsa_has_local_layers: + if not self.has_local_layers: return - # Page allocation is unchanged within the step, so only the row mapping - # depends on the updated lengths. - self._refresh_qsa_token_mapping() + self._refresh_speculative_snapshot_state() + self._refresh_device_state() - def update_for_spec_dec(self) -> None: - """Refresh the row mapping after a draft step repacks the batch. + def on_kv_lens_changed(self, spec_dec_step: bool) -> None: + """Refresh row limits after speculative decoding advances KV lengths. - The first draft step rewrites seq_lens from one row per accepted token - to a single row per request and advances kv_lens_cuda in place, so the - packed-row mapping built during prepare() no longer describes the rows - the next draft step attends over. + Page allocation is unchanged within the step, so only the row mapping + depends on the updated lengths. A draft step also repacks the batch + from one row per accepted token to a single row per request. """ - super().update_for_spec_dec() - self.on_update_kv_lens() + if self.has_local_layers: + self._refresh_token_mapping() -__all__ = ["QSAAttentionMetadata"] +__all__ = ["QSAMetadata"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py index c1c127aa35df..5ad51ebcbeb1 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/module.py @@ -19,11 +19,12 @@ QSA_VALUE_ROLE_INDEX, ) from .indexer import _is_power_of_two, _logical_to_pages -from .metadata import QSAAttentionMetadata +from .metadata import QSAMetadata from .params import QSASparseParams if TYPE_CHECKING: from tensorrt_llm._torch.attention.attention import Attention + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata _DEFAULT_SCORE_WORKSPACE_BYTES = 128 * 1024 * 1024 @@ -158,7 +159,7 @@ def qsa_sparse_gqa( selected_tokens: torch.Tensor, request_idx: int | None = None, request_indices: torch.Tensor | None = None, - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", softmax_scale: float, query_positions: torch.Tensor | None = None, compress_ratio: int | None = None, @@ -190,7 +191,7 @@ def qsa_sparse_gqa( q=q, k_cache=k_cache, v_cache=v_cache, - block_table=metadata.qsa_block_table, + block_table=metadata.ext(QSAMetadata).block_table, selected_tokens=selected_tokens, request_indices=request_indices.contiguous(), tokens_per_block=metadata.kv_cache_manager.tokens_per_block, @@ -221,7 +222,7 @@ def qsa_sparse_gqa_reference( v_cache: torch.Tensor, selected_tokens: torch.Tensor, request_indices: torch.Tensor, - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", softmax_scale: float, ) -> torch.Tensor: """Torch reference sparse GQA over V2 paged K/V.""" @@ -259,7 +260,7 @@ def qsa_sparse_gqa_reference( def _capture_mtp_shared_topk( - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", slot: int, indices: torch.Tensor, visible_blocks: torch.Tensor, @@ -272,7 +273,8 @@ def _capture_mtp_shared_topk( The row kept for each request is its last accepted one, so the reuse starts from the ranking the target actually committed to. """ - shared = metadata.qsa_shared_topk_indices + qsa = metadata.find_ext(QSAMetadata) + shared = qsa.shared_topk_indices if shared is None: return num_seqs = metadata.num_seqs @@ -285,7 +287,7 @@ def _capture_mtp_shared_topk( # reads back through ``request_indices``. seq_lens = metadata.seq_lens_cuda[:num_seqs].to(torch.int64) row_starts = torch.cumsum(seq_lens, dim=0) - seq_lens - num_accepted = metadata.qsa_mtp_num_accepted + num_accepted = qsa.mtp_num_accepted if num_accepted is None: offsets = seq_lens - 1 else: @@ -293,10 +295,8 @@ def _capture_mtp_shared_topk( rows = (row_starts + offsets).clamp_(0, num_rows - 1) shared[slot, :num_seqs].copy_(indices.index_select(0, rows)) - metadata.qsa_shared_topk_visible_blocks[slot, :num_seqs].copy_( - visible_blocks.index_select(0, rows) - ) - metadata.qsa_shared_topk_captured_slots.add(slot) + qsa.shared_topk_visible_blocks[slot, :num_seqs].copy_(visible_blocks.index_select(0, rows)) + qsa.shared_topk_captured_slots.add(slot) def select_qsa_tokens( @@ -365,7 +365,7 @@ def select_qsa_paged_tokens( query_positions: torch.Tensor, sequence_lengths: torch.Tensor, request_indices: torch.Tensor, - metadata: "QSAAttentionMetadata", + metadata: "TrtllmAttentionMetadata", params: QSASparseParams, *, top_k: TopK | None = None, @@ -389,19 +389,20 @@ def select_qsa_paged_tokens( """ from .kernels import triton_qsa_mtp_reuse_topk, triton_qsa_paged_index_scores + qsa = metadata.find_ext(QSAMetadata) if visible_blocks is None: visible_blocks = ((query_positions + 1) // params.compress_ratio).to(torch.int32) if ( mtp_share_slot is not None - and metadata.qsa_indexer_skip_topk - and mtp_share_slot in metadata.qsa_shared_topk_captured_slots + and qsa.indexer_skip_topk + and mtp_share_slot in qsa.shared_topk_captured_slots and top_k_output is not None and q.is_cuda ): indices = triton_qsa_mtp_reuse_topk( - shared_indices=metadata.qsa_shared_topk_indices[mtp_share_slot], - shared_visible_blocks=metadata.qsa_shared_topk_visible_blocks[mtp_share_slot], + shared_indices=qsa.shared_topk_indices[mtp_share_slot], + shared_visible_blocks=qsa.shared_topk_visible_blocks[mtp_share_slot], request_indices=request_indices, visible_blocks=visible_blocks, output=top_k_output[: q.shape[0]], @@ -417,7 +418,7 @@ def select_qsa_paged_tokens( logits = triton_qsa_paged_index_scores( q=q, index_cache=index_cache, - block_table=metadata.qsa_block_table, + block_table=qsa.block_table, query_positions=query_positions, request_indices=request_indices, tokens_per_block=metadata.kv_cache_manager.tokens_per_block, @@ -482,7 +483,7 @@ def forward( q: torch.Tensor, k: Optional[torch.Tensor], v: Optional[torch.Tensor], - attn_metadata: QSAAttentionMetadata, + attn_metadata: "TrtllmAttentionMetadata", attention_mask: AttentionMask, attention_window_size: Optional[int], attention_mask_data: Optional[torch.Tensor], @@ -494,6 +495,7 @@ def forward( output_gate: Optional[torch.Tensor], **kwargs: object, ) -> Optional[torch.Tensor]: + qsa = attn_metadata.find_ext(QSAMetadata) if attention_mask is not PredefinedAttentionMask.CAUSAL: raise NotImplementedError("QSA sparse attention supports only a causal mask") unsupported = { @@ -514,7 +516,7 @@ def forward( # The model's QKNormRoPEAttention preprocesses q/k before this hook; # mrope_config is therefore backend state, not work for the QSA kernel. del mrope_config - if not isinstance(attn_metadata, QSAAttentionMetadata): + if qsa is None: raise TypeError("QSA sparse attention received incompatible metadata") num_tokens = attn_metadata.num_tokens if num_tokens == 0: @@ -587,8 +589,8 @@ def forward( raise RuntimeError("QSA sparse attention requires adjacent paged HND K/V buffers") k_cache = kv_pool[:, QSA_KEY_ROLE_INDEX] v_cache = kv_pool[:, QSA_VALUE_ROLE_INDEX] - req_idx = attn_metadata.qsa_req_idx_per_token[:num_tokens] - logical = attn_metadata.qsa_logical_positions[:num_tokens] + req_idx = qsa.req_idx_per_token[:num_tokens] + logical = qsa.logical_positions[:num_tokens] tokens_per_block = attn_metadata.kv_cache_manager.tokens_per_block # The direct row-wise kernel removes launch overhead for small decode # batches. Larger batches and prefill chunks have enough work for @@ -614,7 +616,7 @@ def forward( v_cache=v_cache, request_indices=req_idx, logical_positions=logical, - block_table=attn_metadata.qsa_block_table, + block_table=qsa.block_table, tokens_per_block=tokens_per_block, ) else: @@ -625,7 +627,7 @@ def forward( v_cache, req_idx, logical, - attn_metadata.qsa_block_table, + qsa.block_table, tokens_per_block, ) @@ -642,9 +644,9 @@ def forward( # KV-length mirror: speculative decoding advances the authoritative # lengths on device and the host mirror is not updated between its # sub-steps. - sequence_lengths = attn_metadata.qsa_sequence_lengths[:num_tokens] + sequence_lengths = qsa.sequence_lengths[:num_tokens] mtp_share_slot = None - if params.mtp_index_share and attn_metadata.qsa_in_mtp_draft_loop: + if params.mtp_index_share and qsa.in_mtp_draft_loop: mtp_share_slot = attn_metadata.kv_cache_manager.qsa_shared_topk_slots[ attention.layer_idx ] @@ -657,9 +659,9 @@ def forward( attn_metadata, params, top_k=attention.indexer.top_k, - top_k_output=attn_metadata.qsa_topk_indices, - top_k_row_starts=attn_metadata.qsa_topk_row_starts, - visible_blocks=attn_metadata.qsa_visible_blocks[:num_tokens], + top_k_output=qsa.topk_indices, + top_k_row_starts=qsa.topk_row_starts, + visible_blocks=qsa.visible_blocks[:num_tokens], mtp_share_slot=mtp_share_slot, ) fuse_output_gate = False @@ -695,10 +697,8 @@ def forward( index_cache = attn_metadata.kv_cache_manager.get_index_k_buffer(attention.layer_idx) if index_cache is None: raise RuntimeError(f"QSA index cache is unavailable for layer {attention.layer_idx}") - sequence_lengths = attn_metadata.qsa_sequence_lengths[:num_tokens] - score_columns = ( - attn_metadata.qsa_block_table.shape[1] * tokens_per_block // params.compress_ratio - ) + sequence_lengths = qsa.sequence_lengths[:num_tokens] + score_columns = qsa.block_table.shape[1] * tokens_per_block // params.compress_ratio chunk_size = _query_chunk_size(num_tokens, score_columns) output = torch.empty_like(q) for start in range(0, num_tokens, chunk_size): @@ -713,9 +713,9 @@ def forward( attn_metadata, params, top_k=attention.indexer.top_k, - top_k_output=attn_metadata.qsa_topk_indices[packed_slice], - top_k_row_starts=attn_metadata.qsa_topk_row_starts[packed_slice], - visible_blocks=attn_metadata.qsa_visible_blocks[packed_slice], + top_k_output=qsa.topk_indices[packed_slice], + top_k_row_starts=qsa.topk_row_starts[packed_slice], + visible_blocks=qsa.visible_blocks[packed_slice], context_rows=True, ) output[packed_slice] = qsa_sparse_gqa( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/rocket/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/rocket/__init__.py index 85c3c9e33f3b..8f97833017df 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/rocket/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/rocket/__init__.py @@ -5,15 +5,14 @@ from .backend import RocketTrtllmAttention, RocketVanillaAttention from .cache_manager import RocketKVCacheManager -from .metadata import RocketTrtllmAttentionMetadata, RocketVanillaAttentionMetadata +from .metadata import RocketMetadata from .params import RocketKVMetadataParams, RocketKVParams __all__ = [ "RocketKVCacheManager", "RocketKVMetadataParams", "RocketKVParams", + "RocketMetadata", "RocketTrtllmAttention", - "RocketTrtllmAttentionMetadata", "RocketVanillaAttention", - "RocketVanillaAttentionMetadata", ] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/rocket/backend.py b/tensorrt_llm/_torch/attention/backends/sparse/rocket/backend.py index 836ced6425f1..371e7779bc23 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/rocket/backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/rocket/backend.py @@ -13,7 +13,11 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata -from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttention, repeat_kv +from tensorrt_llm._torch.attention.backends.vanilla import ( + VanillaAttention, + VanillaAttentionMetadata, + repeat_kv, +) from tensorrt_llm.models.modeling_utils import QuantConfig if TYPE_CHECKING: @@ -32,14 +36,15 @@ triton_softmax, triton_topk, ) -from .metadata import RocketTrtllmAttentionMetadata, RocketVanillaAttentionMetadata +from .metadata import RocketMetadata from .params import RocketKVParams ModelConfig = tensorrt_llm.bindings.ModelConfig class RocketTrtllmAttention(TrtllmAttention): - Metadata = RocketTrtllmAttentionMetadata + Metadata = TrtllmAttentionMetadata + metadata_extensions = (RocketMetadata,) # Access type for different dtype sizes _access_type = {1: torch.int8, 2: torch.int16, 4: torch.int32, 8: torch.int64} @@ -91,37 +96,38 @@ def sparse_kv_predict( num_ctx_tokens = metadata.num_ctx_tokens if num_ctx_tokens == 0: return None, None + rocket = metadata.ext(RocketMetadata) if k is None: qkv_input = q[:num_ctx_tokens] else: qkv_input = torch.cat([q, k], dim=1) - if metadata.valid_batch_size > 0: + if rocket.valid_batch_size > 0: q_window, k_context = triton_rocket_qk_split( qkv_input, metadata.prompt_lens_cuda, - metadata.context_cumsum_cuda, - metadata.valid_seq_indices_cuda, - metadata.k_cu_seqlens_cuda, - metadata.total_rocket_k_ctx_tokens, + rocket.context_cumsum_cuda, + rocket.valid_seq_indices_cuda, + rocket.k_cu_seqlens_cuda, + rocket.total_rocket_k_ctx_tokens, self.num_heads, self.num_kv_heads, self.head_dim, self.window_size, - metadata.valid_batch_size, + rocket.valid_batch_size, ) scores = triton_bmm( q_window, k_context, - metadata.q_cu_seqlens_cuda, - metadata.k_cu_seqlens_cuda, - metadata.valid_batch_size, + rocket.q_cu_seqlens_cuda, + rocket.k_cu_seqlens_cuda, + rocket.valid_batch_size, causal=False, ) - scores = triton_softmax(scores, metadata.k_cu_seqlens_cuda, metadata.valid_batch_size) + scores = triton_softmax(scores, rocket.k_cu_seqlens_cuda, rocket.valid_batch_size) # scores: [num_heads, window_size, total_k_tokens] -> [num_kv_heads, total_k_tokens] scores = scores.view( @@ -132,9 +138,9 @@ def sparse_kv_predict( # scores: [num_kv_heads, total_k_tokens] -> [valid_batch_size, num_kv_heads, padding_size] scores = triton_flatten_to_batch( scores, - metadata.k_cu_seqlens_cuda, - metadata.valid_batch_size, - metadata.max_rocket_k_ctx_len, + rocket.k_cu_seqlens_cuda, + rocket.valid_batch_size, + rocket.max_rocket_k_ctx_len, ) scores = torch.nn.functional.max_pool1d( @@ -142,7 +148,7 @@ def sparse_kv_predict( ) # Use indexer topk prefill to select topk prefix indices - total_tasks = metadata.valid_batch_size * self.num_kv_heads + total_tasks = rocket.valid_batch_size * self.num_kv_heads selected_prefix_indices = torch.empty( (total_tasks, self.prompt_budget - self.window_size), @@ -152,10 +158,10 @@ def sparse_kv_predict( scores = scores.view(total_tasks, -1) - row_starts = metadata.k_context_start_cuda[ - : metadata.valid_batch_size - ].repeat_interleave(self.num_kv_heads) - row_ends = metadata.k_context_lens_cuda[: metadata.valid_batch_size].repeat_interleave( + row_starts = rocket.k_context_start_cuda[: rocket.valid_batch_size].repeat_interleave( + self.num_kv_heads + ) + row_ends = rocket.k_context_lens_cuda[: rocket.valid_batch_size].repeat_interleave( self.num_kv_heads ) torch.ops.trtllm.indexer_topk_prefill( @@ -175,16 +181,16 @@ def sparse_kv_predict( dtype=torch.int32, ) - sparse_kv_offsets = metadata.sparse_offsets_ctx_cuda[: metadata.num_contexts + 1] + sparse_kv_offsets = rocket.sparse_offsets_ctx_cuda[: metadata.num_contexts + 1] # Flatten sparse indices sparse_kv_indices = triton_rocket_batch_to_flatten( selected_prefix_indices, metadata.prompt_lens_cuda, - metadata.valid_seq_indices_cuda, + rocket.valid_seq_indices_cuda, sparse_kv_offsets, metadata.num_contexts, - metadata.total_sparse_ctx_indices, + rocket.total_sparse_ctx_indices, self.window_size, self.prompt_budget, self.num_kv_heads, @@ -196,8 +202,8 @@ def sparse_kv_predict( triton_rocket_update_kt_cache_ctx( qkv_input.contiguous(), kt_cache_tensor, - metadata.kt_cache_block_offsets[: metadata.num_contexts], - metadata.context_cumsum_cuda[: metadata.num_contexts + 1], + rocket.kt_cache_block_offsets[: metadata.num_contexts], + rocket.context_cumsum_cuda[: metadata.num_contexts + 1], sparse_kv_indices, sparse_kv_offsets, self.num_heads, @@ -205,19 +211,19 @@ def sparse_kv_predict( self.head_dim, self.page_size, self.prompt_budget, - metadata.kt_tokens_per_block, + rocket.kt_tokens_per_block, metadata.kv_cache_manager.max_kt_blocks_per_seq, ) # Reduce overhead of post processing - if metadata.valid_batch_size == 0: + if rocket.valid_batch_size == 0: return None, None return sparse_kv_indices, sparse_kv_offsets @torch.compile(dynamic=True) def preprocess_for_gen( - self, q: torch.Tensor, k: Optional[torch.Tensor], metadata: RocketTrtllmAttentionMetadata + self, q: torch.Tensor, k: Optional[torch.Tensor], metadata: TrtllmAttentionMetadata ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if k is None: qkv_input = q[metadata.num_ctx_tokens :] @@ -249,6 +255,7 @@ def sparse_attn_predict( ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: if metadata.num_generations == 0: return None, None + rocket = metadata.ext(RocketMetadata) q, k = self.preprocess_for_gen(q, k, metadata) @@ -261,10 +268,10 @@ def sparse_attn_predict( triton_rocket_update_kt_cache_gen( k, kt_cache_tensor, - metadata.kt_cache_block_offsets[metadata.num_contexts :], + rocket.kt_cache_block_offsets[metadata.num_contexts :], metadata.kv_lens_cuda_runtime[metadata.num_contexts :], - metadata.page_size, - metadata.kt_tokens_per_block, + rocket.page_size, + rocket.kt_tokens_per_block, metadata.kv_cache_manager.max_kt_blocks_per_seq, self.num_kv_heads, self.head_dim, @@ -274,34 +281,34 @@ def sparse_attn_predict( scores = triton_rocket_paged_kt_cache_bmm( q, kt_cache_tensor, - metadata.kt_cache_block_offsets[metadata.num_contexts :], + rocket.kt_cache_block_offsets[metadata.num_contexts :], metadata.kv_lens_cuda_runtime[metadata.num_contexts :], - metadata.cum_kt_lens_cuda, - metadata.page_size, - metadata.kt_tokens_per_block, + rocket.cum_kt_lens_cuda, + rocket.page_size, + rocket.kt_tokens_per_block, metadata.kv_cache_manager.max_kt_blocks_per_seq, - metadata.total_kt_tokens, + rocket.total_kt_tokens, ) - scores = triton_softmax(scores, metadata.cum_kt_lens_cuda, metadata.num_generations) + scores = triton_softmax(scores, rocket.cum_kt_lens_cuda, metadata.num_generations) # Mean over num_heads_per_kv for each batch separately scores = triton_rocket_reduce_scores( scores, - metadata.cum_kt_lens_cuda, + rocket.cum_kt_lens_cuda, metadata.num_generations, self.num_kv_heads, self.num_heads // self.num_kv_heads, ) - sparse_attn_offsets = metadata.sparse_offsets_gen_cuda[: metadata.num_generations + 1] + sparse_attn_offsets = rocket.sparse_offsets_gen_cuda[: metadata.num_generations + 1] selected_indices = triton_topk( scores, - metadata.cum_kt_lens_cuda, + rocket.cum_kt_lens_cuda, sparse_attn_offsets, - metadata.total_sparse_gen_indices, - metadata.topk, + rocket.total_sparse_gen_indices, + rocket.topk, ) return selected_indices, sparse_attn_offsets @@ -314,7 +321,8 @@ class RocketVanillaAttention(VanillaAttention): - Generation phase: only support sparse attention computation """ - Metadata = RocketVanillaAttentionMetadata + Metadata = VanillaAttentionMetadata + metadata_extensions = (RocketMetadata,) def __init__( self, @@ -354,7 +362,7 @@ def _single_request_sparse_kv_predict( q: Optional[Tensor], k: Optional[Tensor], v: Optional[Tensor], - metadata: RocketVanillaAttentionMetadata, + metadata: VanillaAttentionMetadata, past_seen_token: int, sample_idx: int, **kwargs, @@ -392,7 +400,7 @@ def _single_request_sparse_kv_predict( self._single_request_update_kt_cache( k_snap, kt_cache_tensor, - metadata.kt_cache_block_offsets[sample_idx], + metadata.ext(RocketMetadata).kt_cache_block_offsets[sample_idx], target_seq_len, kt_cache_position, update=False, @@ -405,7 +413,7 @@ def _single_request_sparse_attn_predict( k: Optional[Tensor], v: Optional[Tensor], kv_cache_tensor: Tensor, - metadata: RocketVanillaAttentionMetadata, + metadata: VanillaAttentionMetadata, past_seen_token: int, sample_idx: int, **kwargs, @@ -530,7 +538,7 @@ def _rocketkv_selection( self, q: Tensor, k: Tensor, - metadata: RocketVanillaAttentionMetadata, + metadata: VanillaAttentionMetadata, past_seen_token: int, sample_idx: int, ) -> Tensor: @@ -560,7 +568,7 @@ def _scaled_softmax(x: Tensor, divscale: Tensor | float, dim: int) -> Tensor: kt_states = self._single_request_update_kt_cache( k, kt_cache_tensor, - metadata.kt_cache_block_offsets[sample_idx], + metadata.ext(RocketMetadata).kt_cache_block_offsets[sample_idx], target_seq_len, kt_cache_position, ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/rocket/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/rocket/metadata.py index 054c69855dbb..b546511b4127 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/rocket/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/rocket/metadata.py @@ -10,8 +10,8 @@ import tensorrt_llm import tensorrt_llm.bindings +from tensorrt_llm._torch.attention.backends.sparse.params import SparseAlgorithmExtension from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata from tensorrt_llm._utils import prefer_pinned from .params import RocketKVMetadataParams @@ -19,108 +19,99 @@ ModelConfig = tensorrt_llm.bindings.ModelConfig -class RocketTrtllmAttentionMetadata(TrtllmAttentionMetadata): - def __post_init__(self): - super().__post_init__() - sparse_metadata_params = self.sparse_metadata_params - if not isinstance(sparse_metadata_params, RocketKVMetadataParams): +class RocketMetadata(SparseAlgorithmExtension): + """RocketKV state: the prompt budget, the KT cache and the index layouts. + + Serves the TRTLLM and Vanilla backends. On both, a generation request + attends to at most ``prompt_budget`` prompt tokens, so its cached length is + reduced before the KV lengths are derived. The TRTLLM backend additionally + predicts sparse indices from the per-step context and generation layouts + prepared here; the Vanilla backend predicts per request in Python. + """ + + def __init__(self, metadata) -> None: + super().__init__(metadata) + params = metadata.sparse_metadata_params + if not isinstance(params, RocketKVMetadataParams): raise ValueError("RocketKV sparse attention metadata params are not set") - self.prompt_budget = sparse_metadata_params.prompt_budget - self.window_size = sparse_metadata_params.window_size - self.page_size = sparse_metadata_params.page_size - self.topk = sparse_metadata_params.topk + self.prompt_budget = params.prompt_budget + self.window_size = params.window_size + self.page_size = params.page_size + self.topk = params.topk + + @property + def _predicts_on_device(self) -> bool: + return isinstance(self.metadata, TrtllmAttentionMetadata) + @property + def kt_tokens_per_block(self) -> Optional[int]: + """The number of KT tokens per block from the KV cache manager.""" + kv_cache_manager = self.metadata.kv_cache_manager + return kv_cache_manager.kt_tokens_per_block if kv_cache_manager is not None else None + + def allocate(self, buffers) -> None: + md = self.metadata + if not self._predicts_on_device: + self.kt_cache_block_offsets = torch.empty( + [md.max_num_sequences, md.kv_cache_manager.max_kt_blocks_per_seq], + dtype=torch.int32, + device="cuda", + ) + self.host_kt_cache_block_offsets = torch.zeros_like( + self.kt_cache_block_offsets, + device="cpu", + pin_memory=prefer_pinned(), + ) + return + + # The TRTLLM kernels read the RocketKV page size as the cache page size. + md.page_size = self.page_size assert self.page_size == next_power_of_2(self.page_size), "Page size must be a power of 2" - capture_graph = self.is_cuda_graph + capture_graph = md.is_cuda_graph + max_num_sequences = md.max_num_sequences - # Cumulative valid sequence lengths for query and key - self.q_cu_seqlens_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), - dtype=torch.int32, - cache_name="q_cu_seqlens_cuda", - capture_graph=capture_graph, - ) + def device_buffer(shape, name, pool=buffers): + return md.get_empty( + pool, shape, dtype=torch.int32, cache_name=name, capture_graph=capture_graph + ) - self.q_cu_seqlens = torch.zeros_like( - self.q_cu_seqlens_cuda, device="cpu", dtype=torch.int32 - ) + def host_mirror(device): + return torch.zeros_like(device, device="cpu", dtype=torch.int32) - self.k_cu_seqlens_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), - dtype=torch.int32, - cache_name="k_cu_seqlens_cuda", - capture_graph=capture_graph, - ) - self.k_cu_seqlens = torch.zeros_like( - self.k_cu_seqlens_cuda, device="cpu", dtype=torch.int32 - ) + # Cumulative valid sequence lengths for query and key + self.q_cu_seqlens_cuda = device_buffer((max_num_sequences + 1,), "q_cu_seqlens_cuda") + self.q_cu_seqlens = host_mirror(self.q_cu_seqlens_cuda) + self.k_cu_seqlens_cuda = device_buffer((max_num_sequences + 1,), "k_cu_seqlens_cuda") + self.k_cu_seqlens = host_mirror(self.k_cu_seqlens_cuda) # Context length of RocketKV key for each valid sequence - self.k_context_lens_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), - dtype=torch.int32, - cache_name="k_context_lens_cuda", - capture_graph=capture_graph, - ) - self.k_context_lens = torch.zeros_like( - self.k_context_lens_cuda, device="cpu", dtype=torch.int32 - ) + self.k_context_lens_cuda = device_buffer((max_num_sequences,), "k_context_lens_cuda") + self.k_context_lens = host_mirror(self.k_context_lens_cuda) # Start index of RocketKV key for each valid sequence - self.k_context_start_cuda = self.get_empty( - None, - (self.max_num_sequences,), - dtype=torch.int32, - cache_name="k_context_start_cuda", - capture_graph=capture_graph, + self.k_context_start_cuda = device_buffer( + (max_num_sequences,), "k_context_start_cuda", pool=None ) # Cumulative context lengths for each sequence - self.context_cumsum_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), - dtype=torch.int32, - cache_name="context_cumsum_cuda", - capture_graph=capture_graph, - ) - self.context_cumsum = torch.zeros_like( - self.context_cumsum_cuda, device="cpu", dtype=torch.int32 - ) + self.context_cumsum_cuda = device_buffer((max_num_sequences + 1,), "context_cumsum_cuda") + self.context_cumsum = host_mirror(self.context_cumsum_cuda) # Sparse kv indices offsets for each sequence in context phase - self.sparse_offsets_ctx_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), - dtype=torch.int32, - cache_name="sparse_offsets_ctx_cuda", - capture_graph=capture_graph, - ) - self.sparse_offsets_ctx = torch.zeros_like( - self.sparse_offsets_ctx_cuda, device="cpu", dtype=torch.int32 + self.sparse_offsets_ctx_cuda = device_buffer( + (max_num_sequences + 1,), "sparse_offsets_ctx_cuda" ) + self.sparse_offsets_ctx = host_mirror(self.sparse_offsets_ctx_cuda) # Valid sequence indices used in sparse kv indices prediction - self.valid_seq_indices_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences,), - dtype=torch.int32, - cache_name="valid_seq_indices_cuda", - capture_graph=capture_graph, - ) + self.valid_seq_indices_cuda = device_buffer((max_num_sequences,), "valid_seq_indices_cuda") # KT cache block offsets used in KT cache related kernels - self.kt_cache_block_offsets = self.get_empty( - self.cuda_graph_buffers, - [self.max_num_sequences, self.kv_cache_manager.max_kt_blocks_per_seq], - dtype=torch.int32, - cache_name="kt_cache_block_offsets", - capture_graph=capture_graph, + self.kt_cache_block_offsets = device_buffer( + [max_num_sequences, md.kv_cache_manager.max_kt_blocks_per_seq], + "kt_cache_block_offsets", ) - self.host_kt_cache_block_offsets = torch.zeros_like( self.kt_cache_block_offsets, device="cpu", @@ -128,95 +119,75 @@ def __post_init__(self): ) # Number of KT tokens for each sequence - self.num_kt_tokens = torch.empty( - self.max_num_sequences, - device="cpu", - dtype=torch.int32, - ) + self.num_kt_tokens = torch.empty(max_num_sequences, device="cpu", dtype=torch.int32) # Cumulative KT lengths for each sequence - self.cum_kt_lens_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), - dtype=torch.int32, - cache_name="cum_kt_lens_cuda", - capture_graph=capture_graph, - ) - self.cum_kt_lens = torch.zeros_like(self.cum_kt_lens_cuda, device="cpu", dtype=torch.int32) + self.cum_kt_lens_cuda = device_buffer((max_num_sequences + 1,), "cum_kt_lens_cuda") + self.cum_kt_lens = host_mirror(self.cum_kt_lens_cuda) # Sparse attn indices offsets for each sequence in generation phase - self.sparse_offsets_gen_cuda = self.get_empty( - self.cuda_graph_buffers, - (self.max_num_sequences + 1,), - dtype=torch.int32, - cache_name="sparse_offsets_gen_cuda", - capture_graph=capture_graph, - ) - self.sparse_offsets_gen = torch.zeros_like( - self.sparse_offsets_gen_cuda, device="cpu", dtype=torch.int32 + self.sparse_offsets_gen_cuda = device_buffer( + (max_num_sequences + 1,), "sparse_offsets_gen_cuda" ) + self.sparse_offsets_gen = host_mirror(self.sparse_offsets_gen_cuda) # Maximum number of KT tokens - self.max_kt_tokens = (self.max_seq_len + self.page_size - 1) // self.page_size - - @property - def kt_tokens_per_block(self) -> Optional[int]: - """ - Returns the number of kt tokens per block from the KV cache manager. - """ - return ( - self.kv_cache_manager.kt_tokens_per_block if self.kv_cache_manager is not None else None - ) - - def prepare(self): - if self.kv_cache_manager is not None: - num_contexts = self.num_contexts - num_generations = self.num_generations - num_requests = num_contexts + num_generations - - for i in range(num_requests): - if i < num_contexts: - self.kv_cache_params.num_cached_tokens_per_seq[i] = 0 - else: - if self.prompt_lens[i] > self.prompt_budget: - self.kv_cache_params.num_cached_tokens_per_seq[i] += ( - self.prompt_budget - self.prompt_lens[i] - ) - - super().prepare() - - # Update prompt lens for sparse attention - if self.kv_cache_manager is not None: - _prompt_lens = self.prompt_lens.copy() - for i in range(num_requests): - if i >= num_contexts: - _prompt_lens[i] = min(_prompt_lens[i], self.prompt_budget) - _prompt_lens = torch.tensor(_prompt_lens, dtype=torch.int, device="cpu") - self.prompt_lens_cpu[: self.num_seqs].copy_(_prompt_lens) - self.prompt_lens_cuda[: self.num_seqs].copy_( - self.prompt_lens_cpu[: self.num_seqs], non_blocking=True - ) - self.prompt_lens_cuda_runtime = self.prompt_lens_cuda[: self.num_seqs] - self.prompt_lens_cpu_runtime = self.prompt_lens_cpu[: self.num_seqs] + self.max_kt_tokens = (md.max_seq_len + self.page_size - 1) // self.page_size + + def pre_derive_lengths(self) -> None: + """Cap each generation request's cached length at the prompt budget.""" + md = self.metadata + if self._predicts_on_device and md.kv_cache_manager is None: + return + num_contexts = md.num_contexts + num_requests = num_contexts + md.num_generations + cached = md.kv_cache_params.num_cached_tokens_per_seq + for i in range(num_requests): + if i < num_contexts: + cached[i] = 0 + elif md.prompt_lens[i] > self.prompt_budget: + cached[i] += self.prompt_budget - md.prompt_lens[i] - # for kt cache - self.kv_cache_manager.copy_kt_block_offsets( - self.request_ids, self.host_kt_cache_block_offsets + def prepare(self) -> None: + md = self.metadata + if md.kv_cache_manager is None: + return + md.kv_cache_manager.copy_kt_block_offsets(md.request_ids, self.host_kt_cache_block_offsets) + self.kt_cache_block_offsets[: md.num_seqs].copy_( + self.host_kt_cache_block_offsets[: md.num_seqs], non_blocking=True + ) + + def post_derive(self) -> None: + if not self._predicts_on_device: + return + md = self.metadata + num_contexts = md.num_contexts + num_generations = md.num_generations + + if md.kv_cache_manager is not None: + # Generation requests attend to at most prompt_budget prompt tokens. + prompt_lens = md.prompt_lens.copy() + for i in range(num_contexts, num_contexts + num_generations): + prompt_lens[i] = min(prompt_lens[i], self.prompt_budget) + md.prompt_lens_cpu[: md.num_seqs].copy_( + torch.tensor(prompt_lens, dtype=torch.int, device="cpu") ) - self.kt_cache_block_offsets[: self.num_seqs].copy_( - self.host_kt_cache_block_offsets[: self.num_seqs], non_blocking=True + md.prompt_lens_cuda[: md.num_seqs].copy_( + md.prompt_lens_cpu[: md.num_seqs], non_blocking=True ) + md.prompt_lens_cuda_runtime = md.prompt_lens_cuda[: md.num_seqs] + md.prompt_lens_cpu_runtime = md.prompt_lens_cpu[: md.num_seqs] # -------------------------------- Context phase -------------------------------- - self.context_cumsum[1 : self.num_contexts + 1] = torch.cumsum( - self.prompt_lens_cpu[: self.num_contexts], dim=0 + self.context_cumsum[1 : num_contexts + 1] = torch.cumsum( + md.prompt_lens_cpu[:num_contexts], dim=0 ) - self.context_cumsum_cuda[: self.num_contexts + 1].copy_( - self.context_cumsum[: self.num_contexts + 1], non_blocking=True + self.context_cumsum_cuda[: num_contexts + 1].copy_( + self.context_cumsum[: num_contexts + 1], non_blocking=True ) # We need to filter out sequences that are too short to skip sparse kv indices prediction - valid_mask = self.prompt_lens_cpu[: self.num_contexts] >= self.prompt_budget + valid_mask = md.prompt_lens_cpu[:num_contexts] >= self.prompt_budget valid_seq_indices = torch.where(valid_mask)[0] invalid_seq_indices = torch.where(~valid_mask)[0] valid_batch_size = len(valid_seq_indices) @@ -224,19 +195,19 @@ def prepare(self): # Only consider sequences that are long enough for sparse kv indices prediction in context phase self.k_context_lens[:valid_batch_size] = ( - self.prompt_lens_cpu[valid_seq_indices] - self.window_size + md.prompt_lens_cpu[valid_seq_indices] - self.window_size ) self.k_context_lens_cuda[:valid_batch_size].copy_( self.k_context_lens[:valid_batch_size], non_blocking=True ) - sparse_counts_ctx = torch.zeros(self.num_contexts, dtype=torch.int32, device="cpu") + sparse_counts_ctx = torch.zeros(num_contexts, dtype=torch.int32, device="cpu") sparse_counts_ctx[valid_seq_indices] = self.prompt_budget - sparse_counts_ctx[invalid_seq_indices] = self.prompt_lens_cpu[invalid_seq_indices] + sparse_counts_ctx[invalid_seq_indices] = md.prompt_lens_cpu[invalid_seq_indices] - self.sparse_offsets_ctx[1 : self.num_contexts + 1] = torch.cumsum(sparse_counts_ctx, dim=0) - self.sparse_offsets_ctx_cuda[: self.num_contexts + 1].copy_( - self.sparse_offsets_ctx[: self.num_contexts + 1], non_blocking=True + self.sparse_offsets_ctx[1 : num_contexts + 1] = torch.cumsum(sparse_counts_ctx, dim=0) + self.sparse_offsets_ctx_cuda[: num_contexts + 1].copy_( + self.sparse_offsets_ctx[: num_contexts + 1], non_blocking=True ) self.q_cu_seqlens[: valid_batch_size + 1] = ( @@ -262,76 +233,33 @@ def prepare(self): self.total_rocket_k_ctx_tokens = 0 self.valid_batch_size = valid_batch_size - self.total_sparse_ctx_indices = self.sparse_offsets_ctx[self.num_contexts].item() + self.total_sparse_ctx_indices = self.sparse_offsets_ctx[num_contexts].item() # -------------------------------- Generation phase -------------------------------- - self.num_kt_tokens[: self.num_generations] = ( - self.kv_lens[self.num_contexts : self.num_seqs] + self.page_size - 1 + self.num_kt_tokens[:num_generations] = ( + md.kv_lens[num_contexts : md.num_seqs] + self.page_size - 1 ) // self.page_size - self.cum_kt_lens[1 : self.num_generations + 1] = torch.cumsum( - self.num_kt_tokens[: self.num_generations], dim=0 + self.cum_kt_lens[1 : num_generations + 1] = torch.cumsum( + self.num_kt_tokens[:num_generations], dim=0 ) - self.cum_kt_lens_cuda[: self.num_generations + 1].copy_( - self.cum_kt_lens[: self.num_generations + 1], non_blocking=True + self.cum_kt_lens_cuda[: num_generations + 1].copy_( + self.cum_kt_lens[: num_generations + 1], non_blocking=True ) - self.total_kt_tokens = self.num_generations * self.max_kt_tokens + self.total_kt_tokens = num_generations * self.max_kt_tokens topk_tensor = torch.tensor(self.topk, dtype=torch.int32) # Some sequences may have less than topk KT tokens # We need to use the minimum of topk and the number of KT tokens - sparse_counts_gen = torch.minimum(topk_tensor, self.num_kt_tokens[: self.num_generations]) - - self.sparse_offsets_gen[1 : self.num_generations + 1] = torch.cumsum( - sparse_counts_gen[: self.num_generations], dim=0 - ) - self.sparse_offsets_gen_cuda[: self.num_generations + 1].copy_( - self.sparse_offsets_gen[: self.num_generations + 1], non_blocking=True - ) + sparse_counts_gen = torch.minimum(topk_tensor, self.num_kt_tokens[:num_generations]) - self.total_sparse_gen_indices = self.topk * self.num_generations - - -class RocketVanillaAttentionMetadata(VanillaAttentionMetadata): - def __post_init__(self): - super().__post_init__() - sparse_metadata_params = self.sparse_metadata_params - if not isinstance(sparse_metadata_params, RocketKVMetadataParams): - raise ValueError("RocketKV sparse attention metadata params are not set") - self.prompt_budget = sparse_metadata_params.prompt_budget - self.kt_cache_block_offsets = torch.empty( - [self.max_num_sequences, self.kv_cache_manager.max_kt_blocks_per_seq], - dtype=torch.int32, - device="cuda", + self.sparse_offsets_gen[1 : num_generations + 1] = torch.cumsum( + sparse_counts_gen[:num_generations], dim=0 ) - self.host_kt_cache_block_offsets = torch.zeros_like( - self.kt_cache_block_offsets, - device="cpu", - pin_memory=prefer_pinned(), + self.sparse_offsets_gen_cuda[: num_generations + 1].copy_( + self.sparse_offsets_gen[: num_generations + 1], non_blocking=True ) - def prepare(self) -> None: - super().prepare() - num_contexts = self.num_contexts - num_generations = self.num_generations - num_requests = num_contexts + num_generations - - for i in range(num_requests): - if i < num_contexts: - self.kv_cache_params.num_cached_tokens_per_seq[i] = 0 - else: - if self.prompt_lens[i] > self.prompt_budget: - self.kv_cache_params.num_cached_tokens_per_seq[i] += ( - self.prompt_budget - self.prompt_lens[i] - ) - - if self.kv_cache_manager is not None: - # for kt cache - self.kv_cache_manager.copy_kt_block_offsets( - self.request_ids, self.host_kt_cache_block_offsets - ) - self.kt_cache_block_offsets[: self.num_seqs].copy_( - self.host_kt_cache_block_offsets[: self.num_seqs], non_blocking=True - ) + self.total_sparse_gen_indices = self.topk * num_generations diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 9b5ac32120db..c2c43f0b7b96 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -14,7 +14,6 @@ # limitations under the License. import contextlib -import functools import math import os import weakref @@ -42,11 +41,12 @@ from ...pyexecutor.config_utils import is_mla from ...utils import (compute_swizzled_sf_shape, get_global_attrs, - get_model_extra_attrs, helix_local_len_tensor) + get_model_extra_attrs) from ..kernels.sage_quant import SageQuantized, sage_quantize +from .features import DraftKVMetadata, FlashMlaMetadata, HelixMetadata +from .features.spec_decoding import SpecDecodingMetadata from .fmha.manager import FmhaManager -from .fp4_mla import (can_fuse_fp4_mla_q_quant, populate_helix_spec_buffers, - scatter_fp4_mla_kv_cache) +from .fp4_mla import can_fuse_fp4_mla_q_quant, scatter_fp4_mla_kv_cache from .fp4_mla.state import Fp4MlaState from .interface import (AttentionBackend, AttentionForwardArgs, AttentionInputType, AttentionMask, AttentionMetadata, @@ -97,33 +97,6 @@ def _resolve_uses_spcompress(sparse_params: Optional[SparseParams], return False -@functools.cache -def generate_spec_decoding_position_offsets(max_num_requests: int, - draft_len: int) -> torch.Tensor: - width = draft_len + 1 - row = torch.arange(width, dtype=torch.int, device='cuda') - return row.unsqueeze(0).expand(max_num_requests, -1).contiguous() - - -@functools.cache -def generate_spec_decoding_packed_mask(max_num_requests: int, - draft_len: int) -> torch.Tensor: - width = draft_len + 1 - num_blocks = math.ceil(width / 32) - mask = torch.zeros([max_num_requests, width, num_blocks], - dtype=torch.int, - device='cuda') - remaining = width - for blk in range(num_blocks): - if remaining <= 0: - break - n = min(32, remaining) - vals = (torch.pow(2, torch.arange(n) + 1) - 1).int() - mask[:, blk * 32:blk * 32 + n, blk] = vals - remaining -= 32 - return mask - - @dataclass(kw_only=True) class TrtllmAttentionMetadata(AttentionMetadata): workspace: Optional[torch.Tensor] = None @@ -173,72 +146,16 @@ def effective_beam_width(self) -> int: # EncoderCUDAGraphRunner; None elsewhere. max_context_q_len_override: Optional[int] = None - # Flags to enable spec-dec mode (multi-query mode) in TRTLLM XQA Kernels - # spec decoding mode can be enabled for non-TRTLLM-gen kernels (pre-Blackwell XQA kernels) - # is_spec_decoding_enabled specifies if spec-dec mode is supported for the entire runtime. - is_spec_decoding_enabled: bool = False - # use_spec_decoding determines if the attention layer should be run in spec-dec mode at the specific step / layer. - use_spec_decoding: bool = False - - # if spec-dec tree is a tree or a chain (linear tree) - is_spec_dec_tree: bool = False - # if spec-dec tree wouldn't be changed at all, the mask won't be computed every step. - is_spec_dec_dynamic_tree: bool = False - force_prepare_spec_dec_tree_mask: bool = False - - # parameters required for spec-dec mode - max_total_draft_tokens: Optional[int] = None - spec_decoding_position_offsets: Optional[torch.Tensor] = None - # C++ attention op requires a 2-D position_offsets tensor and reads - # sizes()[1] as the generation length / packed-mask row stride. - spec_decoding_position_offsets_cpp: Optional[torch.Tensor] = None - # Compact Hopper C++ row stride for 1D dynamic-tree offsets. - position_offsets_stride: int = 0 - spec_decoding_packed_mask: Optional[torch.Tensor] = None - spec_decoding_generation_lengths: Optional[torch.Tensor] = None - spec_decoding_bl_tree_mask_offset: Optional[torch.Tensor] = None - spec_decoding_bl_tree_mask: Optional[torch.Tensor] = None - spec_bl_tree_first_sparse_mask_offset_kv: Optional[torch.Tensor] = None - # TRTLLM-Gen FMHA JIT warmup controls. trtllm_gen_jit_warmup: bool = False - # Flag to enable helix parallelism. - enable_helix: bool = False - - # Global position ids of tokens for each sequence in the batch. Given - # each helix rank has only a subset of tokens for a sequence, we compute - # a global position id for each token here. - helix_position_offsets: Optional[torch.Tensor] = None - helix_position_offsets_cpu: Optional[torch.Tensor] = None - - # Whether the current rank is inactive for helix parallelism. - # In helix parallelism, only the active rank appends KV cache for the query token - # and attends to the previously cached tokens as well as the query token. Inactive - # ranks do not append KV cache for the query token and attend to the previously - # cached tokens only. - helix_is_inactive_rank: Optional[torch.Tensor] = None - helix_is_inactive_rank_cpu: Optional[torch.Tensor] = None - - # Per-token helix state for speculative verify groups (a 1 + draft_len - # group may straddle a ledger-page boundary onto two CP ranks, so the - # per-sequence boolean above is insufficient there). See - # recompute_helix_spec_buffers for the derivation. - helix_local_slots: Optional[torch.Tensor] = None - helix_kv_bounds: Optional[torch.Tensor] = None - helix_owned_new_tokens_cpu: Optional[torch.Tensor] = None - _helix_spec_tokens_valid: bool = False - # Block offsets for the target and draft KV caches kv_cache_block_offsets: Optional[torch.Tensor] = None host_kv_cache_block_offsets: Optional[torch.Tensor] = None - draft_kv_cache_block_offsets: Optional[torch.Tensor] = None - # Active block-ID buffers, defaulting to the target cache. Separate draft - # storage lets CUDA graphs bind stable target and draft addresses. + # Per-sequence block IDs: generation rows (compacted) and every row. None + # unless an extension that needs them (FlashMLA) provides them. block_ids_per_seq: Optional[torch.Tensor] = None kv_block_ids_per_seq: Optional[torch.Tensor] = None - draft_block_ids_per_seq: Optional[torch.Tensor] = None - draft_kv_block_ids_per_seq: Optional[torch.Tensor] = None # Originals of attributes swapped for draft replay, recorded by # swap_for_draft; the strict buffer check validates these instead. draft_replay_swapped_attrs: Dict[str, Any] = field(default_factory=dict, @@ -249,20 +166,6 @@ def effective_beam_width(self) -> int: # buffer check skips them. graph_temporary_attrs: ClassVar[FrozenSet[str]] = frozenset() - # Batch-shared FP4 state; other attention paths allocate none of it. - fp4_mla_state: Optional[Fp4MlaState] = field(init=False, - default=None, - repr=False, - compare=False) - - # Pre-computed FlashMLA tile-scheduler metadata and num_splits. - # Computed once per forward pass in TrtllmAttention.forward() and reused across layers. - flash_mla_tile_scheduler_metadata: Optional[torch.Tensor] = None - flash_mla_num_splits: Optional[torch.Tensor] = None - _flash_mla_metadata_valid: bool = field(default=False, - init=False, - repr=False) - use_paged_context_fmha: bool = field(init=False, default=False, repr=False) # FMHA prologue buffers for the MLA generation path, hoisted out of the per-layer @@ -286,7 +189,7 @@ def effective_beam_width(self) -> int: default=None, repr=False) - # `DSAtrtllmAttentionMetadata` overrides this; the dense path keeps 0. + # The DSA extension sets this; the dense path keeps 0. num_sparse_topk: int = 0 @property @@ -294,19 +197,6 @@ def effective_workspace(self) -> Optional[torch.Tensor]: """Attention-kernel workspace, switching to the CUDA-graph copy under capture.""" return self.cuda_graph_workspace if self.is_cuda_graph else self.workspace - @property - def spec_decoding_position_offsets_for_cpp(self) -> Optional[torch.Tensor]: - """``spec_decoding_position_offsets`` reshaped to the 2D layout the C++ - kernel expects.""" - offsets = self.spec_decoding_position_offsets - if offsets is not None and offsets.dim() == 1: - if (self.spec_decoding_position_offsets_cpp is not None - and not self.is_sm_version_trtllm_gen_kernel( - sm=get_sm_version())): - return self.spec_decoding_position_offsets_cpp - return offsets.view(self.max_num_requests, -1) - return offsets - @property def max_context_length(self) -> int: """ @@ -340,6 +230,9 @@ def tokens_per_block(self) -> Optional[int]: """ Returns the number of tokens per block from the KV cache manager. """ + pools = self._extension_kv_cache_pools() + if pools is not None: + return pools[2] return self.kv_cache_manager.tokens_per_block if self.kv_cache_manager is not None else None @property @@ -362,6 +255,9 @@ def host_kv_cache_pool_pointers(self) -> Optional[torch.Tensor]: """ Returns the host KV cache pool pointers from the KV cache manager if KV cache manager is not None. """ + pools = self._extension_kv_cache_pools() + if pools is not None: + return pools[0] return self.kv_cache_manager.kv_cache_pool_pointers if self.kv_cache_manager is not None else None @property @@ -369,12 +265,34 @@ def host_kv_cache_pool_mapping(self) -> Optional[torch.Tensor]: """ Returns the host KV cache pool mapping from the KV cache manager if KV cache manager is not None. """ + pools = self._extension_kv_cache_pools() + if pools is not None: + return pools[1] return self.kv_cache_manager.kv_cache_pool_mapping if self.kv_cache_manager is not None else None - def __post_init__(self) -> None: - super().__post_init__() - self.enable_helix = self.mapping.has_cp_helix( - ) if self.mapping is not None else False + def _extension_kv_cache_pools( + self) -> Optional[Tuple[torch.Tensor, torch.Tensor, int]]: + for extension in self._attached_extensions(): + pools = extension.kv_cache_pools() + if pools is not None: + return pools + return None + + def _default_extension_types(self) -> list: + types = [*super()._default_extension_types(), SpecDecodingMetadata] + if self.mapping is not None and self.mapping.has_cp_helix(): + types.append(HelixMetadata) + if (self.kv_cache_manager is not None + and self.draft_kv_cache_manager is not None): + types.append(DraftKVMetadata) + if callable( + getattr(self.kv_cache_manager, "get_fp4_mla_page_table_spec", + None)): + types.append(Fp4MlaState) + return types + + def _post_init_core(self) -> None: + super()._post_init_core() self.use_paged_context_fmha = ( self.runtime_features.chunked_prefill or self.runtime_features.cache_reuse @@ -387,23 +305,6 @@ def __post_init__(self) -> None: self.draft_replay_swapped_attrs = {} self._post_init_with_buffers(self.cuda_graph_buffers) - def update_position_offsets_for_cpp(self, query_len: int) -> None: - """Refresh the C++ view of spec-dec position offsets.""" - offsets = self.spec_decoding_position_offsets - if offsets is None or offsets.dim() != 1: - self.spec_decoding_position_offsets_cpp = offsets - self.position_offsets_stride = 0 - return - - if self.max_num_requests > 0 and query_len > 0: - self.position_offsets_stride = query_len - total = self.max_num_requests * query_len - self.spec_decoding_position_offsets_cpp = offsets[:total].view( - self.max_num_requests, query_len) - else: - self.spec_decoding_position_offsets_cpp = offsets - self.position_offsets_stride = 0 - def _post_init_with_buffers(self, buffers) -> None: # Set a default value, as max_num_sequences is not always set. @@ -498,86 +399,7 @@ def _post_init_with_buffers(self, buffers) -> None: self.host_kv_cache_block_offsets = self.kv_cache_manager.host_kv_cache_block_offsets self.block_ids_per_seq = None self.kv_block_ids_per_seq = None - self.draft_block_ids_per_seq = None - self.draft_kv_block_ids_per_seq = None - - # Allocate separate block offset tensors for draft KV cache manager - # Used in one-model speculative decoding with different KV cache layouts - if self.draft_kv_cache_manager is not None: - num_draft_attention_op_pools = getattr( - self.draft_kv_cache_manager, "num_attention_op_pools", - self.draft_kv_cache_manager.num_pools) - self.draft_kv_cache_block_offsets = self.get_empty( - buffers, - [ - num_draft_attention_op_pools, self.max_num_sequences, 2, - self.draft_kv_cache_manager.max_blocks_per_seq - ], - cache_name="draft_kv_cache_block_offsets", - dtype=torch.int32, - capture_graph=capture_graph, - ) - if self.enable_flash_mla: - self.block_ids_per_seq = self.get_empty( - buffers, - [ - self.kv_cache_manager.max_batch_size, - self.kv_cache_manager.max_blocks_per_seq - ], - cache_name="block_ids_per_seq", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.kv_block_ids_per_seq = self.get_empty( - buffers, - [ - self.kv_cache_manager.max_batch_size, - self.kv_cache_manager.max_blocks_per_seq - ], - cache_name="kv_block_ids_per_seq", - dtype=torch.int32, - capture_graph=capture_graph, - ) - if self.draft_kv_cache_manager is not None: - self.draft_block_ids_per_seq = self.get_empty( - buffers, - [ - self.draft_kv_cache_manager.max_batch_size, - self.draft_kv_cache_manager.max_blocks_per_seq - ], - cache_name="draft_block_ids_per_seq", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.draft_kv_block_ids_per_seq = self.get_empty( - buffers, - [ - self.draft_kv_cache_manager.max_batch_size, - self.draft_kv_cache_manager.max_blocks_per_seq - ], - cache_name="draft_kv_block_ids_per_seq", - dtype=torch.int32, - capture_graph=capture_graph, - ) - # Allocate fixed-size buffers for pre-computed FlashMLA metadata. - # These are pre-allocated so their GPU addresses are stable across CUDA graph captures. - sm_count = torch.cuda.get_device_properties( - torch.cuda.current_device()).multi_processor_count - self.flash_mla_tile_scheduler_metadata = self.get_empty( - buffers, - [sm_count * 8], # TileSchedulerMetaDataSize = 8 - cache_name="flash_mla_tile_scheduler_metadata", - dtype=torch.int32, - capture_graph=capture_graph, - ) - self.flash_mla_num_splits = self.get_empty( - buffers, - [self.kv_cache_manager.max_batch_size + 1], - cache_name="flash_mla_num_splits", - dtype=torch.int32, - capture_graph=capture_graph, - ) if self.enable_context_mla_with_cached_kv: # for kv cache reuse/chunked context in MLA self.ctx_cached_token_indptr = self.get_empty( @@ -618,89 +440,14 @@ def _post_init_with_buffers(self, buffers) -> None: pin_memory=prefer_pinned(), ) - if callable( - getattr(self.kv_cache_manager, "get_fp4_mla_page_table_spec", - None)): - self.fp4_mla_state = Fp4MlaState.create(self, buffers) - - # Allocate static buffers for helix parallelism support. - if self.enable_helix: - self.helix_position_offsets = self.get_empty( - buffers, - (self.max_num_tokens, ), - cache_name="helix_position_offsets", - dtype=torch.int, - capture_graph=capture_graph, - ) - self.helix_position_offsets_cpu = torch.empty_like( - self.helix_position_offsets, - device='cpu', - pin_memory=prefer_pinned(), - ) - self.helix_is_inactive_rank = self.get_empty( - buffers, - (self.max_num_sequences, ), - cache_name="helix_is_inactive_rank", - dtype=torch.bool, - capture_graph=capture_graph, - ) - self.helix_is_inactive_rank_cpu = torch.empty_like( - self.helix_is_inactive_rank, - device='cpu', - pin_memory=prefer_pinned(), - ) - # Per-token buffers for speculative verify groups under helix. - # A group of 1 + draft_len tokens can straddle a ledger-page - # boundary, splitting ownership between two CP ranks, so the - # per-sequence flag above is not expressive enough: - # helix_local_slots[t]: rank-local KV write slot of gen token t - # on this rank, or -1 when another rank owns its position - # (consumed by the mla_rope_generation append kernel). - # helix_kv_bounds[t]: number of rank-local KV entries token t - # may attend to, i.e. local_len(pos_t + 1) (consumed by the - # CuTe DSL MLA decode mask and the helix stats identity). - # Filled by recompute_helix_spec_buffers() on the spec path only. - self.helix_local_slots = self.get_empty( - buffers, - (self.max_num_tokens, ), - cache_name="helix_local_slots", - dtype=torch.int, - capture_graph=capture_graph, - ) - self.helix_kv_bounds = self.get_empty( - buffers, - (self.max_num_tokens, ), - cache_name="helix_kv_bounds", - dtype=torch.int, - capture_graph=capture_graph, - ) - # Host-side per-sequence count of this step's new tokens owned by - # this rank (spec path; single-token path derives it from the - # boolean flag). Consumed by prepare()'s helix kv_lens branch. - self.helix_owned_new_tokens_cpu = torch.zeros( - (self.max_num_sequences, ), - device='cpu', - dtype=torch.int, - pin_memory=prefer_pinned(), - ) - self._helix_spec_tokens_valid = False - def on_update_kv_lens(self): # KV lengths can change between speculative decoding sub-steps. - if self.enable_flash_mla: - self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() - state = getattr(self, "fp4_mla_state", None) - if state is not None: - state.on_update_kv_lens(self) + super().on_update_kv_lens() def update_for_spec_dec(self) -> None: - if self.enable_flash_mla: - self._flash_mla_metadata_valid = False self._invalidate_mla_scheduler_buffers() - state = getattr(self, "fp4_mla_state", None) - if state is not None: - state.update_for_spec_dec(self) + super().update_for_spec_dec() def _invalidate_mla_scheduler_buffers(self) -> None: # Spec-dec rewrites q_lens and kv_lens between sub-steps, so the cumulative @@ -708,116 +455,8 @@ def _invalidate_mla_scheduler_buffers(self) -> None: self._mla_scheduler_buffers_valid = False self._mla_ctx_cu_seqlens_valid = False self._cute_dsl_mla_split_plan_key = None - - def restore_from_spec_dec(self) -> None: - super().restore_from_spec_dec() - if self.fp4_mla_state is not None: - self.fp4_mla_state.restore_from_spec_dec(self) - - def update_helix_param( - self, - helix_position_offsets: List[int], - helix_is_inactive_rank: List[bool], - helix_owned_new_tokens: Optional[List[int]] = None, - ) -> None: - """ - Update helix parameters by copying into static buffers for CUDA graph compatibility. - - Args: - helix_position_offsets: Position offsets for helix parallelism with shape (num_tokens,). - helix_is_inactive_rank: Whether the current rank is inactive with shape (batch_size,). - helix_owned_new_tokens: Per-sequence count of this step's new - tokens owned by this rank (speculative verify groups; one - group may straddle a page boundary onto two ranks). None on - the single-token path, where the boolean flag carries it. - """ - if helix_position_offsets is not None and self.helix_position_offsets is not None: - num_tokens = len(helix_position_offsets) - self.helix_position_offsets_cpu[:num_tokens].copy_( - torch.tensor(helix_position_offsets, dtype=torch.int)) - self.helix_position_offsets[:num_tokens].copy_( - self.helix_position_offsets_cpu[:num_tokens], non_blocking=True) - - if helix_is_inactive_rank is not None and self.helix_is_inactive_rank is not None: - batch_size = len(helix_is_inactive_rank) - self.helix_is_inactive_rank_cpu[:batch_size].copy_( - torch.tensor(helix_is_inactive_rank, dtype=torch.bool)) - self.helix_is_inactive_rank[:batch_size].copy_( - self.helix_is_inactive_rank_cpu[:batch_size], non_blocking=True) - - self._helix_spec_tokens_valid = False - if helix_owned_new_tokens is not None: - batch_size = len(helix_owned_new_tokens) - self.helix_owned_new_tokens_cpu[:batch_size].copy_( - torch.tensor(helix_owned_new_tokens, dtype=torch.int)) - self._helix_spec_tokens_valid = True - - def helix_local_len_vec(self, global_lens: torch.Tensor) -> torch.Tensor: - """Vectorized rank-local prefix length for helix round-robin pages. - - For each global sequence length g, returns the number of the first g - tokens whose ledger page lives on this CP rank (page b -> rank - b % cp_size). The rule and its scalar twin live in - ``_torch.utils``; KVCacheManagerV2._helix_local_len and the host - packing in model_engine use the same definition. - """ - return helix_local_len_tensor(global_lens, - self.kv_cache_manager.tokens_per_block, - self.mapping.cp_size, - self.mapping.cp_rank) - - def recompute_helix_spec_buffers(self, num_gen_tokens: int, - tokens_per_gen_seq: int) -> None: - """Derive per-token helix buffers from (corrected) global positions. - - Called after the overlap-scheduler device correction has been applied - to helix_position_offsets, so every derived quantity reflects the - real committed length even though the host packed provisional values. - Static shapes only; safe under CUDA graph capture. - - ``tokens_per_gen_seq`` is the uniform verify-group width. Non-uniform - groups are rejected rather than silently mis-sliced: without the - overlap scheduler the extend loop packs a per-request - ``1 + get_draft_token_length(request)`` and a request entering with no - draft tokens is packed as a single-token generation row instead, so a - batch can arrive whose total happens to divide but whose rows do not - line up. - - Two index bases meet here, and they are not the same: - * helix_position_offsets / helix_local_slots / helix_kv_bounds are - GENERATION-RELATIVE -- the packing loops only append for extend - and generation rows, so token 0 is the first generation token. - * kv_lens_cuda is BATCH-indexed, hence the num_contexts offset on - the write below. - """ - assert num_gen_tokens % tokens_per_gen_seq == 0, ( - f"helix spec expects uniform verify groups: {num_gen_tokens} gen " - f"tokens not divisible by group size {tokens_per_gen_seq}") - num_gen_seqs = num_gen_tokens // tokens_per_gen_seq - # Divisibility alone does not imply uniformity: a batch of mixed group - # widths can still divide and would then write the wrong number of - # kv_lens_cuda rows with values taken from the wrong tokens. - assert num_gen_seqs == self.num_generations, ( - f"helix spec expects uniform verify groups: {num_gen_tokens} gen " - f"tokens over {self.num_generations} generation rows do not all " - f"have width {tokens_per_gen_seq}") - # One launch derives the write slots, the attention bounds and the - # per-sequence rank-local kv length (the bound of the sequence's last - # token: committed + owned in-flight tokens) from the positions; it - # replaces a chain of elementwise ops on the step's critical path. - # Static shapes and allocation-free, so it is safe under CUDA graph - # capture. - populate_helix_spec_buffers( - self.helix_position_offsets[:num_gen_tokens], - self.helix_local_slots[:num_gen_tokens], - self.helix_kv_bounds[:num_gen_tokens], - self.kv_lens_cuda[self.num_contexts:self.num_contexts + - num_gen_seqs], - num_tokens=num_gen_tokens, - tokens_per_seq=tokens_per_gen_seq, - tokens_per_block=self.kv_cache_manager.tokens_per_block, - cp_size=self.mapping.cp_size, - cp_rank=self.mapping.cp_rank) + for extension in self._attached_extensions(): + extension.invalidate_length_derived_buffers() def _bind_runtime_views( self, @@ -838,6 +477,8 @@ def _bind_runtime_views( self.prompt_lens_cuda_runtime = prompt_lens_cuda self.prompt_lens_cpu_runtime = prompt_lens_cpu self.host_request_types_runtime = host_request_types + for extension in self._attached_extensions(): + extension.on_runtime_views_bound() def mla_prepare_scheduler_buffers( self, num_heads: int) -> tuple[torch.Tensor, torch.Tensor]: @@ -856,6 +497,10 @@ def mla_prepare_scheduler_buffers( touches the device copy. Results land in fixed-address buffers so a captured CUDA graph keeps reading the same pointers, and no D2H sync is introduced. """ + for extension in self._attached_extensions(): + buffers = extension.mla_scheduler_buffers(num_heads) + if buffers is not None: + return buffers num_ctx = self.num_contexts n_gen = self.num_seqs - num_ctx if not self._mla_scheduler_buffers_valid: @@ -912,20 +557,12 @@ def restore_draft_swaps(self) -> None: setattr(self, name, original) self.draft_replay_swapped_attrs = {} - def prepare_for_draft_forward(self) -> dict | None: - """Prepare backend state shared by draft-forward execution paths.""" - return None - - def restore_after_draft_forward(self, saved_state: dict | None) -> None: - """Restore backend state modified for draft-forward execution.""" - return None - def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: - """The token-major generation presentation, or None. - - None on every backend but DSA, and on every uniform DSA batch. See - ``DSAtrtllmAttentionMetadata.token_major_gen_view``. - """ + """The token-major generation presentation an extension provides, or None.""" + for extension in self._attached_extensions(): + view = extension.token_major_gen_view() + if view is not None: + return view return None @contextlib.contextmanager @@ -983,8 +620,8 @@ def presented_token_major( self.max_num_requests, self.prompt_lens_cuda_runtime, self.host_request_types_runtime) = saved - def prepare(self) -> None: - super().prepare() + def _prepare_core(self) -> None: + super()._prepare_core() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. self._invalidate_mla_scheduler_buffers() extra_attrs = get_model_extra_attrs() @@ -1021,45 +658,10 @@ def prepare(self) -> None: device='cpu', ) if self.kv_cache_params.use_cache else None - if self.enable_flash_mla: - self.prepare_flash_mla() - # number of tokens needed in the kv cache for each sequence after the next pass. - if self.enable_helix: - # If helix is inactive, attend to the previously cached tokens only. - assert cached_token_lens is not None, "cached_token_lens should be set for helix" - # The helix per-sequence buffers are GENERATION-relative: the - # packing loops in model_engine append only for extend and plain - # generation rows, so update_helix_param writes exactly - # [0, num_generations). Every device consumer indexes them the - # same way (the MLA rope generation kernel, the XQA preprocessing - # kernels, the FP4 MLA generation kernel), so the host read cannot - # slice them from 0 against a contexts-first cached_token_lens: - # that both shifts every pairing by num_contexts and reads past - # the written region, which for helix_is_inactive_rank_cpu is - # uninitialized memory. Pair them with the generation slice of the - # batch-indexed tensors instead. - num_gen = self.num_generations - gen = slice(self.num_contexts, self.num_seqs) - # Context rows are not part of a verify group and are not packed - # into the helix buffers at all; they append every one of their - # tokens, exactly like the non-helix path below. - kv_lens = cached_token_lens + self.seq_lens_kv - if self._helix_spec_tokens_valid: - # Speculative verify groups: a group may straddle a page - # boundary, so ownership of this step's new tokens is a - # per-sequence COUNT, not a boolean. Provisional host values; - # recompute_helix_spec_buffers overrides the device copy - # after the overlap correction. - kv_lens[gen] = (cached_token_lens[gen] + - self.helix_owned_new_tokens_cpu[:num_gen]) - else: - inactive_rank = self.helix_is_inactive_rank_cpu[:num_gen] - kv_lens[gen] = torch.where(inactive_rank, - cached_token_lens[gen], kv_lens[gen]) - else: - kv_lens = cached_token_lens + \ - self.seq_lens_kv if cached_token_lens is not None else self.seq_lens_kv + kv_lens = cached_token_lens + \ + self.seq_lens_kv if cached_token_lens is not None else self.seq_lens_kv + kv_lens = self._override_kv_lens(kv_lens) # self.kv_lens is the valid kv cache length, while the self.kv_lens_cuda is # the sequence length including the cached tokens and the input tokens. @@ -1096,11 +698,12 @@ def prepare(self) -> None: # kv_lens_cuda on device past the host snapshot, and their # kernels dereference block columns a host-derived cap would # leave unstaged (uninitialized in this buffer). - spec_active = (self.draft_kv_cache_manager is not None - or self.is_spec_decoding_enabled - or bool(self.kv_cache_params.num_extra_kv_tokens) or - (self.runtime_features is not None and - self.runtime_features.has_speculative_draft_tokens)) + spec_active = ( + self.ext(SpecDecodingMetadata).is_spec_decoding_enabled + or bool(self.kv_cache_params.num_extra_kv_tokens) + or (self.runtime_features is not None + and self.runtime_features.has_speculative_draft_tokens) + or self._needs_full_block_table_width()) max_blocks = None if not spec_active and self.kv_cache_manager.tokens_per_block: max_blocks = ceil_div(max_kv_len, @@ -1113,43 +716,22 @@ def prepare(self) -> None: self.num_seqs, max_blocks=max_blocks) - # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists - if self.draft_kv_cache_manager is not None: - # Use the wrapper method which works for both V1 and V2 - self.draft_kv_cache_manager.copy_batch_block_offsets( - self.draft_kv_cache_block_offsets, - self.request_ids, - self.beam_width, - self.num_contexts, - self.num_seqs, - max_blocks=max_blocks) + draft_kv = self.find_ext(DraftKVMetadata) + if draft_kv is not None: + draft_kv.stage_block_offsets(max_blocks) # Don't pass self.kv_lens as kv_lens here because it includes extra # tokens. Use the actual KV length (without extra tokens) for # kv_lens_runtime, which becomes host_past_key_value_lengths and # eventually mMaxSeqLenKv. - if self.fp4_mla_state is not None: - # FP4 MLA needs the per-forward append lengths rather than the - # original prompt lengths. Context scratch-cache metadata and MTP - # generation both consume these runtime views. - kv_lens_runtime = kv_lens[:self.num_seqs] - prompt_lens_cuda_runtime = self.seq_lens_kv_cuda[:self.num_seqs] - prompt_lens_cpu_runtime = self.seq_lens_kv[:self.num_seqs] - else: - kv_lens_runtime = kv_lens[:self.num_seqs] - prompt_lens_cuda_runtime = self.prompt_lens_cuda[:self.num_seqs] - prompt_lens_cpu_runtime = self.prompt_lens_cpu[:self.num_seqs] self._bind_runtime_views( kv_lens_cuda=self.kv_lens_cuda[:self.num_seqs], - kv_lens=kv_lens_runtime, - prompt_lens_cuda=prompt_lens_cuda_runtime, - prompt_lens_cpu=prompt_lens_cpu_runtime, + kv_lens=kv_lens[:self.num_seqs], + prompt_lens_cuda=self.prompt_lens_cuda[:self.num_seqs], + prompt_lens_cpu=self.prompt_lens_cpu[:self.num_seqs], host_request_types=self.host_request_types[:self.num_seqs], ) - if self.fp4_mla_state is not None: - self.fp4_mla_state.prepare(self, kv_lens) - if self._has_mixed_precision_kv(): # Native BF16 paged-context FMHA reads 32-token pages. self.prepare_mixed_precision_kv( @@ -1197,17 +779,17 @@ def prepare_encoder_decoder_from_precomputed_lengths( context_kv_tokens: int, generation_kv_tokens: int, max_kv_len: int) -> None: """Prepare encoder-decoder attention from precomputed lengths.""" - super().prepare() + AttentionMetadata._prepare_core(self) extra_attrs = get_model_extra_attrs() if extra_attrs is None: get_global_attrs().attention_metadata = weakref.ref(self) assert self.kv_cache_manager is not None assert self.draft_kv_cache_manager is None - assert not self.is_spec_decoding_enabled + assert not self.ext(SpecDecodingMetadata).is_spec_decoding_enabled assert self.kv_cache_params.num_extra_kv_tokens == 0 - assert not self.enable_flash_mla - assert not self.enable_helix + assert self.find_ext(FlashMlaMetadata) is None + assert self.find_ext(HelixMetadata) is None assert not self.enable_context_mla_with_cached_kv assert self.request_ids is not None assert max_kv_len <= self.kv_cache_manager.max_seq_len, ( @@ -1301,23 +883,6 @@ def prepare_encoder_cuda_graph_replay(self, seq_lens: List[int], self._num_ctx_tokens = padded_num_tokens self.host_total_kv_lens[0] = padded_num_tokens - def prepare_flash_mla(self) -> None: - self._flash_mla_metadata_valid = False - # Request-specific fills and H2D copies must happen before replay, not - # become fixed operations in the captured graph. - if torch.cuda.is_current_stream_capturing(): - return - - block_ids_per_seq = maybe_pin_memory( - self.kv_cache_manager.get_block_ids_per_seq(self.request_ids)) - num_blocks = block_ids_per_seq.shape[1] - self.kv_block_ids_per_seq.fill_(0) - self.kv_block_ids_per_seq[:self.num_seqs, :num_blocks].copy_( - block_ids_per_seq, non_blocking=True) - self.block_ids_per_seq.fill_(0) - self.block_ids_per_seq[:self.num_generations, :num_blocks].copy_( - block_ids_per_seq[self.num_contexts:], non_blocking=True) - def pre_process_for_chunked_prefill( self, chunked_seq_len: torch.Tensor, @@ -1497,46 +1062,6 @@ def prepare_context_mla_with_cached_kv(self, self.ctx_kv_indptr[:self.num_contexts + 1].copy_( self.host_ctx_kv_indptr[:self.num_contexts + 1], non_blocking=True) - def spec_decoding_param_prepare_for_blackwell(self) -> None: - """ - Prepare the blackwell parameters for the speculative decoding (Medusa and Eagle) generation-phase attention kernels. - Uses persistent buffers (only allocate if None) for CUDA graph compatibility. - """ - if self.spec_decoding_bl_tree_mask_offset is None: - self.spec_decoding_bl_tree_mask_offset = torch.zeros( - [self.max_num_requests], - dtype=torch.int64, - device='cuda', - ) - - if self.spec_bl_tree_first_sparse_mask_offset_kv is None: - self.spec_bl_tree_first_sparse_mask_offset_kv = torch.zeros( - [self.max_num_requests], - dtype=torch.int32, - device='cuda', - ) - - if self.spec_decoding_bl_tree_mask is None: - seqLenQ = self.spec_decoding_packed_mask.shape[ - 1] if self.spec_decoding_packed_mask is not None else 1 - mask_size = torch.ops.trtllm.spec_dec_tree_custom_mask_capacity( - self.max_num_requests, seqLenQ, self.max_num_heads_per_kv) - self.spec_decoding_bl_tree_mask = torch.zeros( - mask_size, - dtype=torch.uint32, - device='cuda', - ) - - def update_blackwell_first_sparse_mask_offset(self) -> None: - """Fill gen slots [0:ng) with (kv_lens_cuda - seq_lens); in-place for CUDA graph.""" - nc = self.num_contexts - n = self.num_seqs - ng = n - nc - if ng > 0: - gen_offset = (self.kv_lens_cuda[nc:n] - - self._seq_lens_cuda[nc:n]).to(torch.int32) - self.spec_bl_tree_first_sparse_mask_offset_kv[:ng].copy_(gen_offset) - def update_spec_dec_param( self, batch_size, @@ -1549,156 +1074,15 @@ def update_spec_dec_param( spec_tree_manager: Optional['SpecTreeManager'] = None, num_contexts: int = 0, ) -> None: - ''' - Update the spec-dec parameters for the TRTLLM attention layer. - Args: - batch_size: int, the number of requests in the batch. - is_spec_decoding_enabled: bool, whether the attention need to be spec_decoding mode, which is determined by attention_need_spec_dec_mode() function. - is_spec_dec_tree: bool, whether the spec-dec mode is a tree, i.e., dynamic tree. For linear-tree, it is always False. - is_spec_dec_dynamic_tree: bool, whether using dynamic tree. - max_draft_len: int, the number of the draft layers. - max_total_draft_tokens: int, the number of all nodes in the tree (except the root). - spec_metadata: Optional['SpecMetadata'] = None, the metadata of the spec-dec. - spec_tree_manager: Optional['SpecTreeManager'] = None, the spec_tree_manager for draft token tree. - num_contexts: int = 0, the number of context (prefill) requests in the - batch. The slot-storage layout is ``[ctx | gen]``; dynamic-tree - metadata only describes the gen rows, so we must source from - ``[num_contexts:batch_size]`` rather than ``[:batch_size]``. - ''' - - # Blackwell trtllm-gen spec-dec is enabled only for dynamic-tree masks. - self.is_spec_decoding_enabled = is_spec_decoding_enabled and ( - not self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()) - or is_spec_dec_dynamic_tree) - - # use_spec_decoding is default to true by default, change in runtime by layers / requests - self.use_spec_decoding = self.is_spec_decoding_enabled - self.is_spec_dec_tree = is_spec_dec_tree - self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree - # A hybrid model's first executed attention layer can have a nonzero - # cache-local index because recurrent layers precede it. Do not rely - # on the C++ ``layer_idx == 0`` fallback to rebuild the target mask: - # the dynamic draft loop clears that mask before the next target step. - self.force_prepare_spec_dec_tree_mask = is_spec_dec_dynamic_tree - # Forward the tree length to FMHA kernel selection. - self.max_total_draft_tokens = max_total_draft_tokens - - # Parameters can be fixed and not changed during runtime if the - if self.is_spec_decoding_enabled: - # Skip pre-allocating position_offsets and packed_mask when dynamic draft length is enabled. - # We will use per-draft-len cached tensors for position_offsets and packed_mask instead. - # Currently dynamic draft length is only supported for linear tree (not is_spec_dec_tree). - - # These buffers are accessed more like removing input padding, - # rather than using max_total_draft_tokens + 1 as the offset between different requests. - if is_spec_dec_tree and self.spec_decoding_position_offsets is None: - # Dynamic tree: use _internal_buf_dim which may be larger - # than max_total_draft_tokens+1 to accommodate K*max_draft_len. - # 1D layout for flexible view() in the drafting loop. - buf_dim = spec_tree_manager._internal_buf_dim - self.spec_decoding_position_offsets = torch.empty( - (self.max_num_requests * buf_dim, ), - dtype=torch.int, - device='cuda', - ) - if is_spec_dec_tree and self.spec_decoding_packed_mask is None: - buf_dim = spec_tree_manager._internal_buf_dim - # Zero-init: dynamic-tree dst has inner dim - # ceil(buf_dim/32) but only ceil((max_total_draft_tokens+1)/32) - # is written each step. Unwritten cols would otherwise feed the - # C++ XQA kernel as live mask bits. - self.spec_decoding_packed_mask = torch.zeros( - [self.max_num_requests, buf_dim, - math.ceil(buf_dim / 32)], - dtype=torch.int, - device='cuda', - ) - - if self.spec_decoding_generation_lengths is None: - self.spec_decoding_generation_lengths = torch.empty( - [self.max_num_requests], - dtype=torch.int, - device='cuda', - ) - - if self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()): - self.spec_decoding_param_prepare_for_blackwell() - else: - self.spec_decoding_bl_tree_mask_offset = None - self.spec_decoding_bl_tree_mask = None - self.spec_bl_tree_first_sparse_mask_offset_kv = None - - cpp_query_len = 0 - - # Case 1: dynamic tree — copy per-request params from spec_tree_manager. - if self.is_spec_dec_dynamic_tree: - assert spec_tree_manager is not None, "spec_tree_manager is required for dynamic tree" - n_dt = spec_tree_manager.max_total_draft_tokens + 1 - buf_dim = spec_tree_manager._internal_buf_dim - - assert buf_dim >= n_dt, ( - f"Dynamic-tree copy requires buf_dim >= n_dt; got " - f"buf_dim={buf_dim}, n_dt={n_dt}.") - - slot_storage = spec_tree_manager.slot_storage - - num_gens = batch_size - num_contexts - if num_gens > 0: - slot_ids = slot_storage.all_ids_buf[num_contexts:batch_size] - - pos_src = torch.index_select(slot_storage.position_offsets, - 0, slot_ids)[:, :n_dt] - pos_dst = self.spec_decoding_position_offsets[:num_gens * - n_dt].view( - num_gens, - n_dt) - pos_dst.copy_(pos_src, non_blocking=True) - - actual_mask_width = math.ceil(n_dt / 32) - mask_src = torch.index_select( - slot_storage.packed_mask, 0, - slot_ids)[:, :, :actual_mask_width] - if self.is_sm_version_trtllm_gen_kernel( - sm=get_sm_version()): - # Blackwell reads the padded 3D mask layout. - self.spec_decoding_packed_mask[:num_gens, :n_dt, : - actual_mask_width].copy_( - mask_src, - non_blocking=True) - else: - # Hopper XQA reads a compact flat prefix. - total = num_gens * n_dt * actual_mask_width - self.spec_decoding_packed_mask.view(-1)[:total].copy_( - mask_src.reshape(-1), non_blocking=True) - - self.spec_decoding_generation_lengths[:batch_size].fill_(n_dt) - cpp_query_len = n_dt - - # Case 2: linear tree - else: - # Currently dynamic draft length is only supported for linear tree - # Dynamic draft length needs position offsets and packed mask to be shaped for each runtime draft length. - # So we create cache for position offsets and packed mask for each draft length to avoid reallocation. - assert max_draft_len == max_total_draft_tokens, "max_draft_len should be equal to max_total_draft_tokens for linear tree" - # For algos other than PARD, this equals runtime_draft_len (K); for PARD it's 2K-1. - runtime_draft_token_buffer_width = ( - spec_metadata.runtime_tokens_per_gen_step - - 1 if spec_metadata is not None else max_draft_len) - self.generate_spec_decoding_generation_length( - runtime_draft_len=runtime_draft_token_buffer_width) - self.spec_decoding_position_offsets = generate_spec_decoding_position_offsets( - self.max_num_requests, runtime_draft_token_buffer_width) - self.spec_decoding_packed_mask = generate_spec_decoding_packed_mask( - self.max_num_requests, runtime_draft_token_buffer_width) - - self.update_position_offsets_for_cpp(cpp_query_len) - - def generate_spec_decoding_generation_length(self, runtime_draft_len): - self.spec_decoding_generation_lengths[:self.max_num_requests].fill_( - runtime_draft_len + 1) - - def is_sm_version_trtllm_gen_kernel(self, sm): - return not (sm < 100 or sm in [120, 121]) + self.ext(SpecDecodingMetadata).update( + batch_size, is_spec_decoding_enabled, is_spec_dec_tree, + is_spec_dec_dynamic_tree, max_draft_len, max_total_draft_tokens, + spec_metadata, spec_tree_manager, num_contexts) + super().update_spec_dec_param(batch_size, is_spec_decoding_enabled, + is_spec_dec_tree, + is_spec_dec_dynamic_tree, max_draft_len, + max_total_draft_tokens, spec_metadata, + spec_tree_manager, num_contexts) class TrtllmAttention(AttentionBackend[TrtllmAttentionMetadata]): @@ -2039,31 +1423,6 @@ def get_quantize_output_dtype( return torch.float8_e4m3fn return None - def _compute_flash_mla_metadata(self, - metadata: TrtllmAttentionMetadata) -> None: - num_generations = metadata.num_generations - if num_generations <= 0: - return - - generation_slice = slice(metadata.num_contexts, - metadata.num_contexts + num_generations) - generation_kv_lens = metadata.kv_lens_cuda_runtime[generation_slice] - generation_num_splits = metadata.flash_mla_num_splits[:num_generations + - 1] - - s_q = int(metadata.seq_lens[metadata.num_contexts].item()) - - thop.compute_flash_mla_metadata( - generation_kv_lens, - metadata.flash_mla_tile_scheduler_metadata, - generation_num_splits, - num_generations, - s_q, - self.num_heads, - 1, - self.kv_lora_rank, - ) - def _quantize_sage_inputs( self, q: torch.Tensor, @@ -2259,7 +1618,8 @@ def forward( token_major_view = (metadata.token_major_gen_view() if forward_args.attention_input_type == AttentionInputType.generation_only else None) - if metadata.enable_flash_mla and token_major_view is not None: + if (metadata.find_ext(FlashMlaMetadata) is not None + and token_major_view is not None): raise ValueError( "FlashMLA does not support ragged token-major generation; " "disable FlashMLA for this layout.") @@ -2417,18 +1777,19 @@ def forward( # The flag is invalidated whenever FlashMLA inputs change. The metadata # must always match the compacted generation sub-batch, which is also # the layout used by block_ids_per_seq. - if (metadata.enable_flash_mla and forward_args.attention_input_type + flash_mla = metadata.find_ext(FlashMlaMetadata) + if (flash_mla is not None and forward_args.attention_input_type != AttentionInputType.context_only - and metadata.num_generations > 0 - and not metadata._flash_mla_metadata_valid): - self._compute_flash_mla_metadata(metadata) - metadata._flash_mla_metadata_valid = True + and metadata.num_generations > 0): + flash_mla.compute_scheduler_metadata(self.num_heads, + self.kv_lora_rank) # Blackwell first_sparse: refresh at layer 0 before kernel launch. + spec = metadata.ext(SpecDecodingMetadata) if self.get_local_layer_idx(metadata) == 0 and ( - metadata.spec_bl_tree_first_sparse_mask_offset_kv is not None + spec.bl_tree_first_sparse_mask_offset_kv is not None and metadata._seq_lens_cuda is not None): - metadata.update_blackwell_first_sparse_mask_offset() + spec.update_blackwell_first_sparse_mask_offset() if metadata.is_cross: if k is not None and v is not None: @@ -2503,8 +1864,9 @@ def forward( # populated int rather than the `None` placeholder. # The call is a fast cache hit after the first forward. self.local_layer_idx = self.get_local_layer_idx(metadata) - if metadata.spec_decoding_bl_tree_mask is not None and self.local_layer_idx == 0: - metadata.spec_decoding_bl_tree_mask.zero_() + spec = metadata.ext(SpecDecodingMetadata) + if spec.bl_tree_mask is not None and self.local_layer_idx == 0: + spec.bl_tree_mask.zero_() if self.print_skip_softmax_stat: self.skip_softmax_stat.zero_() @@ -2886,14 +2248,16 @@ def mla_rope_generation( ) return + helix = metadata.find_ext(HelixMetadata) helix_tensor_params = [ - metadata.helix_position_offsets, metadata.helix_is_inactive_rank + helix.position_offsets if helix is not None else None, + helix.is_inactive_rank if helix is not None else None, ] - if metadata._helix_spec_tokens_valid: + if helix is not None and helix.spec_tokens_valid: # Speculative verify groups: per-token KV write slots (-1 = this # rank does not own the token's position). The append kernel then # gates and addresses per token instead of per sequence. - helix_tensor_params.append(metadata.helix_local_slots) + helix_tensor_params.append(helix.local_slots) # Same presentation as the attention dispatch above: the # rope kernel derives batch_idx = tok / seq_len and the KV write offset @@ -3029,6 +2393,7 @@ def _fp4_mla_rope_generation( "FP4 MLA generation requires fused BF16 RoPE/cache update " "with an FP32 rotary table.") + helix = metadata.find_ext(HelixMetadata) hp_pool_updated = scatter_fp4_mla_kv_cache( metadata, latent_cache, @@ -3041,8 +2406,10 @@ def _fp4_mla_rope_generation( q_pe=q_pe, q_rope_out=fused_q[..., self.kv_lora_rank:], q_quant_input=fused_q, - helix_position_offsets=metadata.helix_position_offsets, - helix_is_inactive_rank=metadata.helix_is_inactive_rank, + helix_position_offsets=helix.position_offsets + if helix is not None else None, + helix_is_inactive_rank=helix.is_inactive_rank + if helix is not None else None, ) if not hp_pool_updated: raise RuntimeError( diff --git a/tensorrt_llm/_torch/attention/backends/vanilla.py b/tensorrt_llm/_torch/attention/backends/vanilla.py index 2477c271b1e2..3281c006d2ff 100644 --- a/tensorrt_llm/_torch/attention/backends/vanilla.py +++ b/tensorrt_llm/_torch/attention/backends/vanilla.py @@ -63,12 +63,12 @@ def generate_sliding_window_mask(batch_size: int, target_length: int, class VanillaAttentionMetadata(AttentionMetadata): - def __post_init__(self) -> None: - super().__post_init__() + def _post_init_core(self) -> None: + super()._post_init_core() self.kv_layout = "NHD" - def prepare(self) -> None: - super().prepare() + def _prepare_core(self) -> None: + super()._prepare_core() # indices of used cache blocks for each sequence assert self.request_ids is not None self.block_ids_per_seq = self.kv_cache_manager.get_batch_cache_indices( diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index 328bb37accbf..40106b8f9449 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -60,8 +60,10 @@ TrtllmAttention, TrtllmAttentionMetadata, ) +from .backends.features import HelixMetadata from .backends.interface import AttentionBackend, PositionalEmbeddingParams, PredefinedAttentionMask from .backends.sparse.hooks import get_sparse_mla_hooks +from .backends.sparse.params import SparseAlgorithmExtension from .backends.utils import create_attention from .rotary_embedding import RotaryEmbedding @@ -843,10 +845,9 @@ def _attn_forward_gen( kv_lora_rank = partial_o.shape[-1] // self.num_heads_tp assert self.kv_lora_rank == kv_lora_rank - helix_kv_bounds = getattr(attn_metadata, "helix_kv_bounds", None) - # helix_kv_bounds is non-None only on TrtllmAttentionMetadata, - # where _helix_spec_tokens_valid is a declared field. - if helix_kv_bounds is not None and attn_metadata._helix_spec_tokens_valid: + helix = attn_metadata.find_ext(HelixMetadata) + helix_kv_bounds = helix.kv_bounds if helix is not None else None + if helix_kv_bounds is not None and helix.spec_tokens_valid: # Speculative verify groups: KV ownership is per-TOKEN. A rank # owning only the tail page of a group has zero visible KV for # the group's leading tokens while its per-sequence kv_len is @@ -1621,7 +1622,7 @@ def forward_absorption_generation( # allocation rather than scaling these temporary buffers by K. cu_seqlens_rows = ( max(num_seqs, num_tokens) + 1 - if getattr(attn_metadata, "is_ragged_verify", False) + if getattr(attn_metadata.find_ext(SparseAlgorithmExtension), "is_ragged_verify", False) else num_seqs + 1 ) cu_q_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 4892807e0c6e..7c40623ad9d2 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -53,8 +53,7 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo -from ..attention.backends.interface import PositionalEmbeddingParams, RopeParams -from ..attention.backends.sparse.deepseek_v4 import DeepseekV4TrtllmAttentionMetadata +from ..attention.backends.interface import AttentionMetadata, PositionalEmbeddingParams, RopeParams from ..attention.mla import MLA from ..distributed import ( AllReduce, @@ -1309,7 +1308,7 @@ def forward( self, hidden_states: torch.Tensor, lm_head: Linear, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, return_context_logits: bool = False, ) -> torch.Tensor: if not return_context_logits: @@ -1372,7 +1371,7 @@ def forward( self, hidden_states: torch.Tensor, lm_head: Linear, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, return_context_logits: bool = False, ) -> torch.Tensor: if not self.model_config.mapping.is_last_pp_rank(): @@ -2128,7 +2127,7 @@ def forward( self, position_ids: torch.IntTensor, hc_state, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, spec_metadata: Optional[SpecMetadata] = None, input_ids: Optional[torch.IntTensor] = None, engram_embeddings=None, @@ -2295,7 +2294,7 @@ def _entry_boundary(self, hc_state, engram_embeddings, has_engram): def forward_MoE( self, hidden_states: torch.Tensor, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, spec_metadata: Optional[SpecMetadata] = None, input_ids: Optional[torch.IntTensor] = None, *, @@ -2484,7 +2483,7 @@ def forward( position_ids: torch.IntTensor, hidden_states: torch.Tensor, embed_tokens: Embedding, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, all_rank_num_tokens: Optional[List[int]] = None, spec_metadata: Optional[SpecMetadata] = None, **kwargs, @@ -2636,7 +2635,7 @@ def __pp_init__(self): def forward( self, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, input_ids: Optional[torch.IntTensor] = None, position_ids: Optional[torch.IntTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, @@ -2790,7 +2789,7 @@ def _compute_engram_hashes( self, input_ids: torch.Tensor, position_ids: Optional[torch.Tensor], - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, *, token_mask: Optional[torch.Tensor] = None, ) -> Dict[int, torch.Tensor]: @@ -2810,7 +2809,7 @@ def _precompute_engram( layer_id: int, hash_indices: torch.Tensor, dtype: torch.dtype, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, ) -> torch.Tensor: """Launch a lookup; sharded variants may first redistribute hash rows.""" return self.layers[layer_id].engram.precompute(hash_indices, dtype=dtype) @@ -3092,7 +3091,7 @@ def _refresh( def forward( self, - attn_metadata: DeepseekV4TrtllmAttentionMetadata, + attn_metadata: AttentionMetadata, input_ids: torch.IntTensor = None, position_ids: Optional[torch.IntTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv41.py b/tensorrt_llm/_torch/models/modeling_deepseekv41.py index 07628c7456b1..41c915f98795 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv41.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv41.py @@ -38,9 +38,10 @@ plan_decoder_replay, scatter_replayed_rows, ) -from ..attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from ..attention.backends.sparse.csa2.metadata import CSA2Metadata from ..attention.backends.sparse.csa2.module import DeepseekV41Attention as CSA2Attention from ..attention.backends.sparse.csa2.params import CSA2Mode +from ..attention.backends.trtllm import TrtllmAttentionMetadata from ..configs.deepseek_v41 import ( DeepseekV41QuantLayout, assert_weight_layout, @@ -622,7 +623,9 @@ def forward( if capture_this_layer and spec_metadata.spec_dec_mode.is_dspark(): # Capture the post-Engram entry state, not V4's post-MoE state. # DSpark averages the mult axis of the flat [N, mult * hidden] stream. - capture_rows = getattr(attn_metadata, "csa2_replay_query_rows", None) + capture_rows = getattr( + attn_metadata.find_ext(CSA2Metadata), "csa2_replay_query_rows", None + ) capture_input = residual.reshape(residual.shape[0], -1) if capture_rows is None: spec_metadata.maybe_capture_hidden_states(self.layer_idx, capture_input, None) @@ -778,7 +781,7 @@ def _precompute_engram_kv( def forward(self, attn_metadata, *args, **kwargs): from ..pyexecutor.ced_replay import uses_private_decoder_cache - attn_metadata.begin_model_forward() + attn_metadata.begin_forward() spec_metadata = kwargs.get("spec_metadata") if isinstance(spec_metadata, DSparkSpecMetadata): spec_metadata.context_capture_lens = None @@ -802,7 +805,9 @@ def forward(self, attn_metadata, *args, **kwargs): requests = None output = super().forward(attn_metadata, *args, **kwargs) if isinstance(spec_metadata, DSparkSpecMetadata): - spec_metadata.context_capture_lens = attn_metadata.csa2_decoder_capture_lens + spec_metadata.context_capture_lens = attn_metadata.ext( + CSA2Metadata + ).csa2_decoder_capture_lens if requests and self.encoder_replay_enabled: from ..attention.backends.sparse.csa2.ced import complete_encoder_replay @@ -813,7 +818,7 @@ def _compute_engram_hashes( self, input_ids: torch.Tensor, position_ids: torch.Tensor | None, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, *, token_mask: torch.Tensor | None = None, padded_num_tokens: int | None = None, @@ -846,7 +851,7 @@ def _precompute_engram( layer_id: int, hash_indices: torch.Tensor, dtype: torch.dtype, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, ) -> torch.Tensor | _AttentionDPEngramEmbeddings: return self.layers[layer_id].engram.precompute( hash_indices, dtype=dtype, all_rank_num_tokens=attn_metadata.all_rank_num_tokens @@ -986,7 +991,7 @@ def refuse(reason: str) -> Tuple[None, int]: def prepare_adp_inputs( self, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, *, all_token_states_required: bool, requests: list["LlmRequest"] | None = None, @@ -996,13 +1001,13 @@ def prepare_adp_inputs( Direct ADP callers use this before the normal input count exchange too; full-state requirements must not change after counts are exchanged. """ - attn_metadata.decoder_replay_plan = self._make_bounded_replay_plan( + attn_metadata.ext(CSA2Metadata).decoder_replay_plan = self._make_bounded_replay_plan( attn_metadata, all_token_states_required, requests ) def _plan_bounded_replay( self, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, all_token_states_required: bool, requests: list["LlmRequest"] | None = None, ) -> tuple[int | None, DecoderReplayPlan | None]: @@ -1014,7 +1019,7 @@ def _plan_bounded_replay( # Captured forwards retain the full, padded input on every rank. if attn_metadata.is_cuda_graph or get_per_request_prefill_cuda_graph_flag(): return None, None - plan = getattr(attn_metadata, "decoder_replay_plan", None) + plan = attn_metadata.ext(CSA2Metadata).decoder_replay_plan if plan is not None: if plan.replay_all_rank_num_tokens is None: raise ValueError("Exchange ADP token counts after prepare_adp_inputs().") @@ -1050,7 +1055,7 @@ def _plan_bounded_replay( def _make_bounded_replay_plan( self, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, all_token_states_required: bool, requests: list["LlmRequest"] | None = None, ) -> DecoderReplayPlan | None: @@ -1133,7 +1138,7 @@ def _enter_remote_tail_boundary( def _enter_bounded_replay( self, plan: DecoderReplayPlan, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, position_ids: torch.Tensor | None, input_ids: torch.Tensor | None, hc_state: HCState, @@ -1154,7 +1159,7 @@ def _enter_bounded_replay( if uses_private_decoder_cache(self, attn_metadata.kv_cache_manager) else None, ) - attn_metadata.csa2_replay_query_rows = plan.rows + attn_metadata.ext(CSA2Metadata).csa2_replay_query_rows = plan.rows return ( gather_replayed_rows(position_ids, plan), gather_replayed_rows(input_ids, plan), @@ -1168,10 +1173,11 @@ def _exit_bounded_replay(self, plan, attn_metadata, hidden_states): """Restore the output layout and metadata after successful Decoder execution.""" hidden_states = scatter_replayed_rows(hidden_states, plan) exit_decoder_replay(attn_metadata, plan) - attn_metadata.csa2_decoder_capture_lens = tuple( + csa2 = attn_metadata.ext(CSA2Metadata) + csa2.csa2_decoder_capture_lens = tuple( plan.replay_seq_lens[: attn_metadata.num_contexts].tolist() ) - attn_metadata.csa2_precomputed_kv_layers.clear() + csa2.csa2_precomputed_kv_layers.clear() return hidden_states def _resolve_hc_state(self, hc_state: HCState) -> HCState: @@ -2265,7 +2271,7 @@ def forward( self, hidden_states: torch.Tensor, lm_head: nn.Module, - attn_metadata: CSA2TrtllmMetadata, + attn_metadata: TrtllmAttentionMetadata, return_context_logits: bool = False, ) -> torch.Tensor: if return_context_logits: diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next.py b/tensorrt_llm/_torch/models/modeling_glm5_next.py index dcc02d2098ab..4b942fe7c31b 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next.py @@ -36,6 +36,7 @@ from ...mapping import Mapping from ...models.modeling_utils import QuantConfig from ...quantization.mode import QuantAlgo +from ..attention.backends.features import RecurrentMetadata from ..attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, @@ -352,7 +353,7 @@ def build_glm5_next_runtime_context(attn_metadata: AttentionMetadata) -> Glm5Nex """ if attn_metadata.kv_cache_manager is None: raise ValueError("glm5_next requires a kv cache manager; got None") - mamba_metadata = attn_metadata.mamba_metadata + mamba_metadata = attn_metadata.ext(RecurrentMetadata).mamba_metadata assert isinstance(mamba_metadata, Glm5NextMamba2Metadata), ( "glm5_next requires Glm5NextMamba2Metadata" ) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 311600ab67ca..c5f8d8e8a26c 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -56,6 +56,8 @@ _write_main_kv_slots_to_pool, ) from ..attention.backends.sparse.minimax_m3.common import index_head_range +from ..attention.backends.sparse.minimax_m3.msa_backend import MiniMaxM3MsaMetadata +from ..attention.backends.sparse.minimax_m3.triton_metadata import MiniMaxM3TritonMetadata from ..attention.backends.sparse.params import SparseBackendForwardArgs from ..attention.backends.trtllm import TrtllmAttention from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams, MiniMaxAllReduceRMS @@ -395,7 +397,7 @@ def _validate_sparse_attention_runtime_config( if sparse_config is None or sparse_config.algorithm != "minimax_m3": raise ValueError( "MiniMax-M3 requires sparse_attention_config.algorithm='minimax_m3' " - "to create its KV-cache manager and prepare attn_metadata.minimax_m3. " + "to create its KV-cache manager and prepare its attention metadata. " "Set the following in the LLM API configuration:\n" "sparse_attention_config:\n algorithm: minimax_m3" ) @@ -1537,7 +1539,8 @@ def _fused_fp8_index_qk_norm_rope( ) rotary_dim = int(self.pos_embd_params.rope.dim) - index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + msa = attn_metadata.ext(MiniMaxM3MsaMetadata) + index_k_cache = msa.msa_idx_k_cache(self.layer_idx) if index_k_cache.dtype != torch.float8_e4m3fn: raise ValueError( "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " @@ -1547,7 +1550,7 @@ def _fused_fp8_index_qk_norm_rope( return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( idx_qk.contiguous(), index_k_cache, - attn_metadata.msa_out_cache_loc[:num_tokens], + msa.msa_out_cache_loc[:num_tokens], self.sparse_num_index_heads, self.sparse_index_dim, rotary_dim, @@ -1618,14 +1621,15 @@ def _fused_fp8_qkv_indexer_norm_rope_kv_insert( if cache_tensors is None: if kv_cache_manager is None: return None + msa = attn_metadata.ext(MiniMaxM3MsaMetadata) buffers = kv_cache_manager.get_buffers(self.layer_idx, kv_layout="HND") - index_k_cache = attn_metadata.msa_idx_k_cache(self.layer_idx) + index_k_cache = msa.msa_idx_k_cache(self.layer_idx) scale_buffers = ( kv_cache_manager.get_block_scale_buffers(self.layer_idx, kv_layout="HND") if kv_cache_manager.is_nvfp4_layer(self.layer_idx) else None ) - out_cache_loc = getattr(attn_metadata, "msa_out_cache_loc", None) + out_cache_loc = msa.msa_out_cache_loc else: buffers, index_k_cache, scale_buffers, out_cache_loc = cache_tensors num_tokens = int(packed.shape[0]) @@ -1897,19 +1901,19 @@ def _sdpa_dense_attention_core( v_cache_view = kv_pool[:, 1] # 5. Read the pre-built M3 metadata. Building is the - # responsibility of ``MiniMaxM3AttentionMetadata.prepare`` (or + # responsibility of ``MiniMaxM3TritonMetadata.prepare`` (or # the test path), so the forward stays CUDA-graph-capture safe. - m3_attachment = getattr(attn_metadata, "minimax_m3", None) - if m3_attachment is None: + m3 = attn_metadata.find_ext(MiniMaxM3TritonMetadata) + if m3 is None or m3.sparse_metadata is None: raise RuntimeError( f"MiniMax-M3 dense forward (layer {self.layer_idx}) requires " - "attn_metadata.minimax_m3 to be pre-built by " - "MiniMaxM3AttentionMetadata.prepare(); the model_engine " - "configures the M3 backend's Metadata class so this happens " - "automatically. Test callers must attach minimax_m3 manually." + "the MiniMaxM3TritonMetadata extension to be prepared; the " + "model_engine attaches it through the M3 backend's " + "metadata_extensions so this happens automatically. Test " + "callers must attach and fill it manually." ) - m3_meta = m3_attachment["metadata"] - out_cache_loc = m3_attachment["out_cache_loc"] + m3_meta = m3.sparse_metadata + out_cache_loc = m3.out_cache_loc # 6. Write the new tokens' K/V to the pool. All inputs come # from prepare() which lands them on the cache device, so the @@ -1955,7 +1959,7 @@ def _sdpa_dense_attention_core( # prefill) preserves causality. ``q_positions`` from the prefill # metadata names each Q token's K-side position. # The metadata tensors are produced by - # :meth:`MiniMaxM3AttentionMetadata.prepare` on the cache + # :meth:`MiniMaxM3TritonMetadata.prepare` on the cache # device, so ``.to(dtype=torch.long)`` is a same-device dtype # conversion (capture-safe). seq_lens_dev = m3_meta.seq_lens.to(dtype=torch.long) @@ -2213,11 +2217,11 @@ def _sparse_forward( 8. Apply ``o_proj`` and return. Production callers (the LLM API path) drive - :meth:`MiniMaxM3AttentionMetadata.prepare` outside any - CUDA-graph capture window; that method attaches a pre-built + :meth:`MiniMaxM3TritonMetadata.prepare` outside any + CUDA-graph capture window; that method stores a pre-built :class:`MiniMaxM3TritonSparseAttentionMetadata` and an - ``out_cache_loc`` tensor as ``attn_metadata.minimax_m3``. Test - callers attach the same dict manually. This forward path + ``out_cache_loc`` tensor on the extension. Test callers fill + the same extension manually. This forward path always reads the pre-built attachment and never builds metadata itself — both would trigger ``cudaErrorStreamCaptureUnsupported`` for the @@ -2246,7 +2250,8 @@ def _sparse_forward( ): # Metadata stages these zero-copy views before compilation; # tracing the cache manager's native pointer access is unsafe. - kv_cache, index_k_cache, kv_scale_cache = attn_metadata.msa_layer_cache_tensors[ + msa = attn_metadata.ext(MiniMaxM3MsaMetadata) + kv_cache, index_k_cache, kv_scale_cache = msa.msa_layer_cache_tensors[ self.layer_idx ] q, idx_q = torch.ops.trtllm.minimax_m3_fused_sparse_qkv_producer( @@ -2255,7 +2260,7 @@ def _sparse_forward( kv_cache, index_k_cache, kv_scale_cache, - attn_metadata.msa_out_cache_loc, + msa.msa_out_cache_loc, self.layer_idx_str, ) o = self._forward_attention_core(q, None, None, idx_q, None, attn_metadata) @@ -2438,19 +2443,19 @@ def _triton_sparse_attention_core( ) # 5. Read the pre-built M3 metadata. Building is the - # responsibility of ``MiniMaxM3AttentionMetadata.prepare`` (or + # responsibility of ``MiniMaxM3TritonMetadata.prepare`` (or # the test path), so the forward stays CUDA-graph-capture safe. - m3_attachment = getattr(attn_metadata, "minimax_m3", None) - if m3_attachment is None: + m3 = attn_metadata.find_ext(MiniMaxM3TritonMetadata) + if m3 is None or m3.sparse_metadata is None: raise RuntimeError( f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires " - "attn_metadata.minimax_m3 to be pre-built by " - "MiniMaxM3AttentionMetadata.prepare(); the model_engine " - "configures the M3 backend's Metadata class so this happens " - "automatically. Test callers must attach minimax_m3 manually." + "the MiniMaxM3TritonMetadata extension to be prepared; the " + "model_engine attaches it through the M3 backend's " + "metadata_extensions so this happens automatically. Test " + "callers must attach and fill it manually." ) - m3_meta = m3_attachment["metadata"] - out_cache_loc = m3_attachment["out_cache_loc"] + m3_meta = m3.sparse_metadata + out_cache_loc = m3.out_cache_loc if not self.disable_index_value and idx_v_cache is not None: # The shared idx_v_proj is not part of the M3 checkpoint diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 2917e7baea77..d75f8094aa0f 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -34,6 +34,7 @@ from ..attention.attention import Attention from ..attention.backends import AttentionMetadata +from ..attention.backends.features import RecurrentMetadata from ..distributed import AllReduce, AllReduceFusionOp, AllReduceParams from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer @@ -751,7 +752,7 @@ def forward( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) - mamba_metadata = attn_metadata.mamba_metadata + mamba_metadata = attn_metadata.ext(RecurrentMetadata).mamba_metadata if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -1250,7 +1251,8 @@ def forward( attn_metadata=attn_metadata, all_rank_num_tokens=all_rank_num_tokens, spec_metadata=spec_metadata, - mamba_metadata=attn_metadata.mamba_metadata, + mamba_metadata=attn_metadata.ext( + RecurrentMetadata).mamba_metadata, lora_params=lora_params, ) return hidden_states diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 1546d5189d49..580bb0e0e418 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -45,6 +45,7 @@ from ...logger import logger from ..attention.backends import AttentionMetadata +from ..attention.backends.features import RecurrentMetadata from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, MoEAllReduce, MoEAllReduceParams, allgather) from ..model_config import ModelConfig @@ -1104,10 +1105,11 @@ def forward( "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" ) - mamba_metadata = attn_metadata.mamba_metadata - if mamba_metadata.max_batch_size != attn_metadata.max_num_requests: - attn_metadata.mamba_metadata = Mamba2Metadata( + recurrent = attn_metadata.ext(RecurrentMetadata) + if recurrent.mamba_metadata.max_batch_size != attn_metadata.max_num_requests: + recurrent.mamba_metadata = Mamba2Metadata( attn_metadata.max_num_requests, chunk_size=128) + mamba_metadata = recurrent.mamba_metadata if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) diff --git a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py index c2e54cba74ec..1bfef4566f45 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py +++ b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py @@ -31,6 +31,7 @@ register_input_processor, ) from ..attention.backends import AttentionMetadata +from ..attention.backends.features import RecurrentMetadata from ..attention.backends.sparse.qsa.indexer import QSAIndexer from ..attention.backends.sparse.qsa.params import QSASparseParams from ..distributed import AllReduce, AllReduceParams, allgather @@ -791,15 +792,16 @@ def forward( ) # Refresh mamba metadata for the batch size, mirroring Qwen3NextModel. - mamba_metadata = getattr(attn_metadata, "mamba_metadata", None) + recurrent = attn_metadata.find_ext(RecurrentMetadata) + mamba_metadata = recurrent.mamba_metadata if recurrent is not None else None if ( mamba_metadata is not None and getattr(mamba_metadata, "max_batch_size", None) != attn_metadata.max_num_requests ): - attn_metadata.mamba_metadata = Mamba2Metadata( + recurrent.mamba_metadata = Mamba2Metadata( attn_metadata.max_num_requests, chunk_size=128 ) - mamba_metadata = attn_metadata.mamba_metadata + mamba_metadata = recurrent.mamba_metadata if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) @@ -1119,7 +1121,7 @@ def norm_hidden() -> torch.Tensor: position_ids=position_ids, hidden_states=hidden_states.flatten(start_dim=-2), attn_metadata=attn_metadata, - mamba_metadata=attn_metadata.mamba_metadata, + mamba_metadata=attn_metadata.ext(RecurrentMetadata).mamba_metadata, spec_metadata=spec_metadata, **kwargs, ) diff --git a/tensorrt_llm/_torch/models/modeling_step3p7.py b/tensorrt_llm/_torch/models/modeling_step3p7.py index 896205ef55de..547889cfc6a8 100644 --- a/tensorrt_llm/_torch/models/modeling_step3p7.py +++ b/tensorrt_llm/_torch/models/modeling_step3p7.py @@ -42,6 +42,7 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType from ..attention.attention import Attention @@ -597,12 +598,14 @@ def forward( attention_window_size if attention_window_size is not None else self.sliding_window ) + spec = attn_metadata.find_ext(SpecDecodingMetadata) if ( not self.rope_fusion - and getattr(attn_metadata, "is_spec_dec_dynamic_tree", False) - and getattr(attn_metadata, "use_spec_decoding", False) - and getattr(attn_metadata, "spec_decoding_position_offsets", None) is not None - and attn_metadata.spec_decoding_position_offsets.dim() == 1 + and spec is not None + and spec.is_spec_dec_dynamic_tree + and spec.use_spec_decoding + and spec.position_offsets is not None + and spec.position_offsets.dim() == 1 and position_ids is not None ): position_ids = self._adjust_position_ids_for_spec_dec( diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index e40fd30a500b..3bf7205fe2f6 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -540,7 +540,9 @@ def commit_all_layers(self, attention_metadata: object, num_decodes: int) -> Non "GDN cached replay all-layer commit requires V2 replay state buffers." ) - mamba_metadata = attention_metadata.mamba_metadata + from ...attention.backends.features import RecurrentMetadata + + mamba_metadata = attention_metadata.ext(RecurrentMetadata).mamba_metadata if mamba_metadata.replay_num_decodes != num_decodes: raise RuntimeError( "GDN replay metadata contains " diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 33f388bad319..9506252728bc 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -30,6 +30,7 @@ from ....models.modeling_utils import QuantConfig from ....quantization import QuantAlgo from ...attention.backends import AttentionMetadata +from ...attention.backends.features import RecurrentMetadata from ...distributed import AllReduce, AllReduceStrategy from ...model_config import ModelConfig from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph @@ -489,7 +490,7 @@ def _forward_impl( verify-only batch without ``output`` may instead return the fused verify epilogue's MXFP8 core (see ``forward_verify_fused``). """ - mamba_metadata = attn_metadata.mamba_metadata + mamba_metadata = attn_metadata.ext(RecurrentMetadata).mamba_metadata num_prefills = attn_metadata.num_contexts num_ctx_tokens = attn_metadata.num_ctx_tokens num_tokens = attn_metadata.num_tokens diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 8bd54a3e073a..3c2d8fd52e02 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -31,6 +31,7 @@ from tensorrt_llm.mapping import Mapping from ...attention.backends import AttentionMetadata +from ...attention.backends.features import RecurrentMetadata from ...distributed import AllReduceParams from ...model_config import ModelConfig from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph @@ -127,7 +128,7 @@ def gdn_custom_op_inplace( a[:num_tokens], b[:num_tokens], attn_metadata, - attn_metadata.mamba_metadata, + attn_metadata.ext(RecurrentMetadata).mamba_metadata, spec_metadata=spec_metadata, output=output[:, :num_tokens, :, :], ) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 9c41b9227e1e..ba2bbf749f16 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -30,6 +30,7 @@ from tensorrt_llm.mapping import Mapping from ...attention.backends import AttentionMetadata +from ...attention.backends.features import RecurrentMetadata from ...model_config import ModelConfig from ...peft.lora.layer import LoraLayer, LoraModuleType from ...pyexecutor.breakable_cuda_graph import (eager_on_graph, @@ -83,9 +84,10 @@ def mamba2_custom_op_inplace(zxbcdt: torch.Tensor, layer_idx: str, # every surrounding piece stays uniform in the num_tokens dim. attn_metadata, mamba_layer, spec_metadata = _extract_mamba2_extra_attrs( layer_idx) - mamba_layer.forward_core(zxbcdt, attn_metadata, - attn_metadata.mamba_metadata, spec_metadata, - ssm_out) + mamba_layer.forward_core( + zxbcdt, attn_metadata, + attn_metadata.ext(RecurrentMetadata).mamba_metadata, spec_metadata, + ssm_out) maybe_bcg_mamba2_custom_op_inplace = eager_on_graph(mamba2_custom_op_inplace) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 235dfdbb26cd..db541956be71 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -959,13 +959,26 @@ def _snapshot_graph_tensor_ptrs(obj: Any) -> Dict[str, int]: in; any of them changing before a later replay means the graph would read stale (or freed) memory. obj.graph_temporary_attrs are skipped. """ - graph_temporary_attrs = getattr(obj, "graph_temporary_attrs", ()) - return { - name: value.data_ptr() - for name, value in vars(obj).items() - if isinstance(value, torch.Tensor) and value.is_cuda - and name not in graph_temporary_attrs - } + ptrs = {} + for prefix, owner in CUDAGraphRunner._graph_tensor_owners(obj): + graph_temporary_attrs = getattr(owner, "graph_temporary_attrs", ()) + ptrs.update({ + prefix + name: value.data_ptr() + for name, value in vars(owner).items() + if isinstance(value, torch.Tensor) and value.is_cuda + and name not in graph_temporary_attrs + }) + return ptrs + + @staticmethod + def _graph_tensor_owners(obj: Any) -> List[Tuple[str, Any]]: + """obj and, for attention metadata, its extensions, with key prefixes.""" + owners = [("", obj)] + attached = obj.__dict__.get("_extensions") if hasattr( + obj, "__dict__") else None + for extension in attached or (): + owners.append((f"{type(extension).__name__}.", extension)) + return owners @staticmethod def _assert_graph_tensor_ptrs_stable(obj: Any, @@ -977,20 +990,23 @@ def _assert_graph_tensor_ptrs_stable(obj: Any, live value is not the captured one. The target value saved before the swap is checked in its place. """ - swapped_attrs = getattr(obj, "draft_replay_swapped_attrs", {}) - for name, expected_ptr in expected_ptrs.items(): + owners = dict(CUDAGraphRunner._graph_tensor_owners(obj)) + for key, expected_ptr in expected_ptrs.items(): + prefix, _, name = key.rpartition(".") + owner = owners.get(f"{prefix}." if prefix else "") + swapped_attrs = getattr(owner, "draft_replay_swapped_attrs", {}) value = (swapped_attrs[name] if name in swapped_attrs else getattr( - obj, name, None)) + owner, name, None)) if not isinstance(value, torch.Tensor): raise RuntimeError( f"CUDA graph metadata attribute `{type(obj).__name__}." - f"{name}` was a tensor at capture time but is now " + f"{key}` was a tensor at capture time but is now " f"{type(value).__name__}. Graph-visible buffers must be " "updated in place, not rebound.") if value.data_ptr() != expected_ptr: raise RuntimeError( f"CUDA graph metadata tensor `{type(obj).__name__}." - f"{name}` was reallocated between capture and replay " + f"{key}` was reallocated between capture and replay " f"(data_ptr 0x{expected_ptr:x} -> 0x{value.data_ptr():x}). " "Graph-visible buffers must be updated in place; " "rebinding invalidates addresses baked into the captured " diff --git a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py index 8b130ab9877c..6027b37c95e8 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/metadata.py @@ -9,11 +9,13 @@ import torch +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionBackend, AttentionMetadata, AttentionRuntimeFeatures, ) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig, _get_kv_cache_layer_specs from tensorrt_llm.mapping import Mapping @@ -75,7 +77,14 @@ def build_attention_metadata( if sparse_attention_config is not None else None ) - return attention_backend.Metadata( + # FlashMLA kernels run on the TRTLLM backend only. + extension_types = ( + (FlashMlaMetadata,) + if model_config.enable_flash_mla + and issubclass(attention_backend.Metadata, TrtllmAttentionMetadata) + else () + ) + return attention_backend.create_metadata( max_num_requests=max_batch_size, max_num_tokens=max_num_tokens, max_num_sequences=max_batch_size * max_beam_width, @@ -83,9 +92,9 @@ def build_attention_metadata( draft_kv_cache_manager=draft_kv_cache_manager, mapping=mapping, runtime_features=attention_runtime_features, - enable_flash_mla=model_config.enable_flash_mla, enable_context_mla_with_cached_kv=enable_context_mla_with_cached_kv, cache_indirection=cache_indirection, max_num_heads_per_kv=max_num_heads_per_kv, sparse_metadata_params=sparse_metadata_params, + extension_types=extension_types, ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index 54fc999d19ad..85d8a7fcb94f 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -11,6 +11,7 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch.attention.backends.features import HelixMetadata, RecurrentMetadata from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.distributed import Distributed @@ -671,8 +672,11 @@ def _apply_steady_gen_fast_prepare( num_extra_kv_tokens=get_num_extra_kv_tokens(None), ) attn_metadata.kv_cache_manager = kv_cache_manager - if hasattr(self.model.model_config.pretrained_config, "chunk_size"): - attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size + recurrent = attn_metadata.find_ext(RecurrentMetadata) + if recurrent is not None and hasattr( + self.model.model_config.pretrained_config, "chunk_size" + ): + recurrent.chunk_size = self.model.model_config.pretrained_config.chunk_size with nvtx_range("steady_gen_metadata_prepare"): attn_metadata.prepare() @@ -2064,13 +2068,15 @@ def _prepare_attn_metadata( # arm the spec path for ordinary helix generation and send its # consumers to uninitialized buffers. helix_spec_active = bool(enable_spec_decode and host.lists.helix_owned_new_tokens) - attn_metadata.update_helix_param( - helix_position_offsets=host.lists.helix_position_offsets, - helix_is_inactive_rank=host.lists.helix_is_inactive_rank, - helix_owned_new_tokens=( - host.lists.helix_owned_new_tokens if helix_spec_active else None - ), - ) + helix = attn_metadata.find_ext(HelixMetadata) + if helix is not None: + helix.update( + position_offsets=host.lists.helix_position_offsets, + is_inactive_rank=host.lists.helix_is_inactive_rank, + owned_new_tokens=( + host.lists.helix_owned_new_tokens if helix_spec_active else None + ), + ) if not attn_metadata.is_cuda_graph: # Assumes seq lens do not change between CUDA graph invocations. This applies @@ -2147,22 +2153,23 @@ def _prepare_attn_metadata( ) attn_metadata.kv_cache_manager = kv_cache_manager - if hasattr(self.model.model_config.pretrained_config, "chunk_size"): - attn_metadata.mamba_chunk_size = self.model.model_config.pretrained_config.chunk_size + recurrent = attn_metadata.find_ext(RecurrentMetadata) + if recurrent is not None and hasattr( + self.model.model_config.pretrained_config, "chunk_size" + ): + recurrent.chunk_size = self.model.model_config.pretrained_config.chunk_size # Some sparse backends (RocketKV) clamp # kv_cache_params.num_cached_tokens_per_seq in place during prepare(), # and KVCacheParams holds the list by reference. Snapshot the true # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) - prepare_context_replay = getattr(attn_metadata, "prepare_context_replay", None) - if prepare_context_replay is not None: - prepare_context_replay( - scheduled_requests.context_requests, - allow_final_window=( - self._config.prefill_cuda_graph_backend == PrefillCudaGraphBackend.DISABLED - ), - ) + attn_metadata.bind_context_requests( + scheduled_requests.context_requests, + prefill_cuda_graphs=( + self._config.prefill_cuda_graph_backend != PrefillCudaGraphBackend.DISABLED + ), + ) attn_metadata.prepare() host.num_cached_tokens_snapshot = num_cached_tokens_snapshot host.remote_tail_modes = remote_tail_modes @@ -2484,20 +2491,21 @@ def preprocess_inputs( # first apply the same accepted-count correction position_ids # got above; without overlap the host values are already exact. md = inputs.get("attn_metadata") - if ( - md is not None - and md.kv_cache_manager is not None - and getattr(md, "_helix_spec_tokens_valid", False) - ): + helix = ( + md.find_ext(HelixMetadata) + if md is not None and self.mapping.has_cp_helix() + else None + ) + if helix is not None and md.kv_cache_manager is not None and helix.spec_tokens_valid: helix_gen_tokens = inputs["input_ids"].shape[0] - md.num_ctx_tokens if not self._config.disable_overlap_scheduler: # The kv_lens override in the recompute supersedes the # generic previous_kv_lens_offsets adjustment above, # which is not ownership-aware. - md.helix_position_offsets[:helix_gen_tokens] += ( + helix.position_offsets[:helix_gen_tokens] += ( self._buffers.previous_pos_id_offsets_cuda[:helix_gen_tokens] ) - md.recompute_helix_spec_buffers( + helix.recompute_spec_buffers( helix_gen_tokens, self.get_runtime_tokens_per_gen_step(runtime_draft_len) ) md.on_update_kv_lens() @@ -2556,14 +2564,17 @@ def postprocess_inputs( restore=True, ) - if self.mapping.has_cp_helix() and getattr( - inputs["attn_metadata"], "_helix_spec_tokens_valid", False - ): + helix = ( + inputs["attn_metadata"].find_ext(HelixMetadata) + if self.mapping.has_cp_helix() + else None + ) + if helix is not None and helix.spec_tokens_valid: # Mirror of the helix position correction in # preprocess_inputs (capture symmetry, like position_ids # above). The recompute's OVERWRITES (slots/bounds/ # kv_lens) need no reversal: every consumer buffer is # rewritten from host state at the next step's prepare. - inputs["attn_metadata"].helix_position_offsets[:previous_batch_tokens] -= ( + helix.position_offsets[:previous_batch_tokens] -= ( self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens] ) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py index 11797542225c..40093ffc9090 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py @@ -9,10 +9,10 @@ import torch +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import BaseMambaCacheManager from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, KVCacheManager, @@ -97,11 +97,10 @@ def setup_attention( kv_cache_manager=kv_cache_manager, draft_kv_cache_manager=draft_kv_cache_manager, ) - if isinstance(kv_cache_manager, BaseMambaCacheManager): - self.attn_metadata.mamba_chunk_size = getattr( - config, "chunk_size", self.attn_metadata.mamba_chunk_size - ) - self.attn_metadata.mamba_metadata_cls = resolve_mamba_metadata_cls(self.model) + recurrent = self.attn_metadata.find_ext(RecurrentMetadata) + if recurrent is not None: + recurrent.chunk_size = getattr(config, "chunk_size", recurrent.chunk_size) + recurrent.mamba_metadata_cls = resolve_mamba_metadata_cls(self.model) return self.attn_metadata diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 7dbcebefe67c..ae8c56de82f2 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -967,7 +967,7 @@ def _cuda_graph_metadata_prepare_scope( consumers. Non-speculative overlap has the same ordering; speculative acceptance corrections and draft-manager swaps keep full preparation. """ - defer = getattr(attn_metadata, "defer_cuda_graph_decode_outputs", None) + defer = attn_metadata.defer_decode_outputs() if ( defer is not None and isinstance(key, KeyType) @@ -982,7 +982,7 @@ def _cuda_graph_metadata_prepare_scope( and self._deferred_decode_outputs_supported and self._mapping.cp_size == 1 ): - return defer() + return defer return contextlib.nullcontext() @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index b9146c9c6bc3..ec445b27da10 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -593,11 +593,12 @@ def _warmup_dg_paged_mqa_logits_metadata(self) -> None: try: from tensorrt_llm._torch.attention.backends.sparse.dsa import ( _DG_SCHEDULE_BLOCK_KV, - DSAtrtllmAttentionMetadata, + DSAMetadata, ) except ImportError: return - if not isinstance(attn_meta, DSAtrtllmAttentionMetadata): + dsa = attn_meta.find_ext(DSAMetadata) + if dsa is None: return try: from tensorrt_llm.deep_gemm import get_paged_mqa_logits_metadata @@ -608,7 +609,7 @@ def _warmup_dg_paged_mqa_logits_metadata(self) -> None: ) return - num_sms = attn_meta.num_sms + num_sms = dsa.num_sms max_bs = max(1, int(self._config.max_batch_size)) beam_width = max(1, self._config.max_beam_width) # Static upper bound on the row-count multiplier applied to @@ -616,10 +617,8 @@ def _warmup_dg_paged_mqa_logits_metadata(self) -> None: # are bounded above by `(1 + max_draft_tokens)`; see the # docstring for why we don't read the runtime `dsl_expand_factor` # here. - max_draft_tokens = int(getattr(attn_meta, "max_draft_tokens", 0) or 0) - expands_batch = getattr(attn_meta, "use_expanded_buffers_for_mtp", False) or getattr( - attn_meta, "expand_for_dsl", False - ) + max_draft_tokens = int(dsa.max_draft_tokens or 0) + expands_batch = dsa.use_expanded_buffers_for_mtp or dsa.expand_for_dsl expand_factor = 1 + max_draft_tokens if expands_batch else 1 max_aligned = ((max_bs * beam_width * expand_factor + 31) // 32) * 32 buckets = list(range(32, max_aligned + 32, 32)) @@ -734,11 +733,11 @@ def _ensure_dsa_attn_metadata_for_warmup(self, resource_manager: ResourceManager if self._metadata.attn_metadata is not None: return try: - from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata + from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata except ImportError: return - metadata_cls = getattr(self._config.attention_backend, "Metadata", None) - if metadata_cls is None or not issubclass(metadata_cls, DSAtrtllmAttentionMetadata): + extension_types = getattr(self._config.attention_backend, "metadata_extensions", ()) + if not any(issubclass(cls, DSAMetadata) for cls in extension_types): return kv_cache_manager = resource_manager.get_resource_manager( ResourceManagerType.KV_CACHE_MANAGER @@ -763,16 +762,14 @@ def _warmup_cute_dsl_radix_topk(self) -> None: if attn_meta is None: return try: - from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata + from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata except ImportError: return - if isinstance(attn_meta, DSAtrtllmAttentionMetadata): + dsa = attn_meta.find_ext(DSAMetadata) + if dsa is not None: next_n = 1 + self._config.original_max_draft_len - attn_meta.warmup_cute_dsl_radix_topk(next_n) - if hasattr(attn_meta, "warmup_selfsampling_topk"): - attn_meta.warmup_selfsampling_topk( - next_n, batch_sizes=self._config.cuda_graph_batch_sizes - ) + dsa.warmup_cute_dsl_radix_topk(next_n) + dsa.warmup_selfsampling_topk(next_n, batch_sizes=self._config.cuda_graph_batch_sizes) def _general_warmup( self, resource_manager: ResourceManager, warmup_requests_configs: list[tuple[int, int]] diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py index 883123303745..ebc9d863c007 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py @@ -16,6 +16,7 @@ AttentionMetadata, AttentionRuntimeFeatures, ) +from tensorrt_llm._torch.attention.backends.sparse.params import SparseAlgorithmExtension from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata from tensorrt_llm._torch.distributed import Distributed @@ -562,6 +563,7 @@ def _prepare_inputs_fast_path( if ( not promoted_context_request_ids and type(attn_metadata) is TrtllmAttentionMetadata + and attn_metadata.find_ext(SparseAlgorithmExtension) is None and self._can_use_fast_path( scheduled_requests, new_tokens_device, next_draft_tokens_device, enable_spec_decode ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py index e3bfeec169bb..7eeba3b8b51f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py @@ -1980,7 +1980,9 @@ def _commit_gdn_cached_replay_history_layers( ): raise RuntimeError("GDN cached replay all-layer commit requires replay state buffers.") - mamba_metadata = attn_metadata.mamba_metadata + from ....attention.backends.features import RecurrentMetadata + + mamba_metadata = attn_metadata.ext(RecurrentMetadata).mamba_metadata if mamba_metadata.replay_num_decodes != num_decodes: raise RuntimeError( "GDN replay metadata contains " @@ -2118,7 +2120,7 @@ def update_mamba_states( position_source_is_token_count = False # Match the API of MambaCacheManager.update_mamba_states: callers # may pass per-request state slot indices explicitly (e.g. MTP via - # attn_metadata.mamba_metadata.state_indices). Fall back to this + # RecurrentMetadata.mamba_metadata.state_indices). Fall back to this # manager's own slot mapping when not provided. if state_indices is None: state_indices = self.get_state_indices() diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 6b7e6b517876..6ceb475f990a 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -30,6 +30,7 @@ from tensorrt_llm.mapping import Mapping from ..attention.backends import AttentionMetadata +from ..attention.backends.features import DraftKVMetadata, RecurrentMetadata from ..modules.logits_processor import LogitsProcessor from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID @@ -1165,14 +1166,15 @@ def _refresh_ctx_block_tables(self, attn_metadata, num_seqs: int) -> bool: """ if self._ctx_block_tables is None or num_seqs <= 0: return False - src = getattr(attn_metadata, "draft_kv_cache_block_offsets", None) + draft_kv = attn_metadata.find_ext(DraftKVMetadata) + src = draft_kv.kv_cache_block_offsets if draft_kv is not None else None if src is None: # The table is bound but unfillable. Leaving it at zeros would send # every request's every layer to pool block 0 -- silently. raise RuntimeError( "DFlash: bound to the draft KV pool but " f"{type(attn_metadata).__name__} carries no " - "draft_kv_cache_block_offsets this iteration." + "draft KV block offsets this iteration." ) # Field 0 of the K/V pair: both name the same block for a kv_factor=2 # pool, and the drafter addresses blocks, not planes. The block count @@ -1191,12 +1193,10 @@ def _lazy_init_ctx_buffers( self, draft_model, spec_metadata, attn_metadata, draft_kv_cache_manager=None ): self._validate_draft_attention_backend(draft_model) - # Only TrtllmAttentionMetadata builds draft_kv_cache_block_offsets. - # Without it the table stays all-zero and every request would read and - # write pool block 0 -- silently, at acceptance 1.0. - if draft_kv_cache_manager is not None and not hasattr( - attn_metadata, "draft_kv_cache_block_offsets" - ): + # Only TrtllmAttentionMetadata attaches the draft KV block offsets. + # Without them the table stays all-zero and every request would read + # and write pool block 0 -- silently, at acceptance 1.0. + if draft_kv_cache_manager is not None and attn_metadata.find_ext(DraftKVMetadata) is None: logger.warning( f"DFlash: {type(attn_metadata).__name__} carries no draft KV block offsets; " "using the private ctx arena." @@ -1868,7 +1868,7 @@ def _forward_impl( attn_metadata.kv_cache_manager.update_mamba_states( attn_metadata=attn_metadata, num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices, + state_indices=attn_metadata.ext(RecurrentMetadata).mamba_metadata.state_indices, ) self._prepare_attn_metadata_for_dflash(attn_metadata, spec_metadata) diff --git a/tensorrt_llm/_torch/speculative/draft_target.py b/tensorrt_llm/_torch/speculative/draft_target.py index 1b853b0bf1fb..5e2204e6ac2d 100644 --- a/tensorrt_llm/_torch/speculative/draft_target.py +++ b/tensorrt_llm/_torch/speculative/draft_target.py @@ -26,6 +26,7 @@ import torch from torch import nn +from tensorrt_llm._torch.attention.backends.features.spec_decoding import set_use_spec_decoding from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.mapping import Mapping @@ -241,7 +242,7 @@ def _forward_impl( hidden_states = hidden_states[0] # Disable spec-dec mode for chained draft steps - attn_metadata.use_spec_decoding = False + set_use_spec_decoding(attn_metadata, False) logits = draft_model.logits_processor( hidden_states[gather_ids], draft_model.lm_head, attn_metadata, True @@ -298,7 +299,7 @@ def _forward_impl( num_accepted_tokens, ) - attn_metadata.use_spec_decoding = True + set_use_spec_decoding(attn_metadata, True) return { "logits": raw_logits, diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index bac560ef8f36..d495a9048335 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -9,13 +9,16 @@ import torch.nn.functional as F from torch import nn +from tensorrt_llm._torch.attention.backends.features.spec_decoding import ( + SpecDecodingMetadata, set_use_spec_decoding) from tensorrt_llm._torch.custom_ops import inplace_slice_copy from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.mapping import Mapping from ..attention.backends import AttentionMetadata +from ..attention.backends.features import FlashMlaMetadata, RecurrentMetadata from ..attention.backends.flashinfer import FlashInferAttentionMetadata -from ..attention.backends.sparse.params import MTPIndexShareMetadata +from ..attention.backends.sparse.params import MTPIndexShareExtension from ..model_config import ModelConfig from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest @@ -31,13 +34,6 @@ from ...llmapi.llm_args import EagleDecodingConfig -def _reset_mtp_index_share(attn_metadata: MTPIndexShareMetadata) -> None: - """Clear the draft-loop state a sparse indexer reads.""" - attn_metadata.set_skip_topk(False) - attn_metadata.set_in_mtp_draft_loop(False) - attn_metadata.set_mtp_num_accepted(None) - - class Eagle3ResourceManager(BaseResourceManager): """ Eagle3 needs to save the hidden states for the draft model. Use this class @@ -494,43 +490,39 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): # Without this, CUDA graph warmup's second iteration would run # the target model attention with stale draft-layer masks # instead of the correct target-tree masks. - if attn_metadata.spec_decoding_packed_mask is not None: - self._saved_packed_mask = attn_metadata.spec_decoding_packed_mask[: - batch_size].clone( - ) + spec = attn_metadata.find_ext(SpecDecodingMetadata) + if spec is not None and spec.packed_mask is not None: + self._saved_packed_mask = spec.packed_mask[:batch_size].clone() else: self._saved_packed_mask = None - if attn_metadata.spec_decoding_position_offsets is not None: - self._saved_position_offsets = attn_metadata.spec_decoding_position_offsets.clone( - ) - self._saved_position_offsets_cpp = ( - attn_metadata.spec_decoding_position_offsets_cpp) + if spec is not None and spec.position_offsets is not None: + self._saved_position_offsets = spec.position_offsets.clone() + self._saved_position_offsets_cpp = spec.position_offsets_cpp else: self._saved_position_offsets = None self._saved_position_offsets_cpp = None - if attn_metadata.spec_decoding_generation_lengths is not None: - self._saved_generation_lengths = attn_metadata.spec_decoding_generation_lengths[:batch_size].clone( - ) + if spec is not None and spec.generation_lengths is not None: + self._saved_generation_lengths = spec.generation_lengths[: + batch_size].clone( + ) else: self._saved_generation_lengths = None def _restore_attn_metadata_from_spec_dec(self, attn_metadata): + spec = attn_metadata.find_ext(SpecDecodingMetadata) super()._restore_attn_metadata_from_spec_dec(attn_metadata) if self._saved_packed_mask is not None: batch_size = self._saved_packed_mask.shape[0] - attn_metadata.spec_decoding_packed_mask[:batch_size].copy_( - self._saved_packed_mask) + spec.packed_mask[:batch_size].copy_(self._saved_packed_mask) self._saved_packed_mask = None if self._saved_position_offsets is not None: - attn_metadata.spec_decoding_position_offsets.copy_( - self._saved_position_offsets) - attn_metadata.spec_decoding_position_offsets_cpp = ( - self._saved_position_offsets_cpp) + spec.position_offsets.copy_(self._saved_position_offsets) + spec.position_offsets_cpp = (self._saved_position_offsets_cpp) self._saved_position_offsets = None self._saved_position_offsets_cpp = None if self._saved_generation_lengths is not None: batch_size = self._saved_generation_lengths.shape[0] - attn_metadata.spec_decoding_generation_lengths[:batch_size].copy_( + spec.generation_lengths[:batch_size].copy_( self._saved_generation_lengths) self._saved_generation_lengths = None @@ -547,6 +539,7 @@ def _forward_impl(self, draft_model, resource_manager=None): + recurrent = attn_metadata.find_ext(RecurrentMetadata) runtime_draft_len = spec_metadata.runtime_draft_len # skip the draft forward if the runtime draft length is 0 if runtime_draft_len == 0: @@ -579,7 +572,7 @@ def _forward_impl(self, attn_metadata.kv_cache_manager.update_mamba_states( attn_metadata=attn_metadata, num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices) + state_indices=recurrent.mamba_metadata.state_indices) sa_manager = getattr(spec_metadata.spec_resource_manager, 'sa_manager', None) @@ -647,7 +640,7 @@ def _forward_impl(self, accepted_tokens, next_draft_tokens, spec_metadata.batch_indices_cuda, batch_size, num_accepted_tokens) - attn_metadata.use_spec_decoding = True + set_use_spec_decoding(attn_metadata, True) # Commit model-owned target state only after draft preparation has # succeeded; SpecWorkerBase aborts the still-pending snapshots on any @@ -655,7 +648,7 @@ def _forward_impl(self, if self._auxiliary_state_handlers: self.commit_auxiliary_speculative_states( num_accepted_tokens, - attn_metadata.mamba_metadata.state_indices[:batch_size], + recurrent.mamba_metadata.state_indices[:batch_size], num_contexts, ) @@ -706,7 +699,7 @@ def _forward_external_shared_target_kv_draft_loop( draft_metadata.update_shared_kv_draft_lengths(attn_metadata, num_accepted_tokens, num_contexts) - draft_metadata.use_spec_decoding = False + set_use_spec_decoding(draft_metadata, False) draft_metadata.padded_num_tokens = None draft_metadata.all_rank_num_tokens = ( spec_metadata.subseq_all_rank_num_tokens) @@ -779,20 +772,20 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # Sparse backends that can reuse one indexer selection across the draft # loop expose this setter; each indexer still decides whether its own # config turns the reuse on. - uses_mtp_index_share = self.is_mtp_eagle and isinstance( - attn_metadata, MTPIndexShareMetadata) + index_share = (attn_metadata.find_ext(MTPIndexShareExtension) + if self.is_mtp_eagle else None) with contextlib.ExitStack() as draft_scope: - if uses_mtp_index_share: - attn_metadata.set_in_mtp_draft_loop(True) + if index_share is not None: + index_share.set_in_mtp_draft_loop(True) # Accepted counts let the indexer stash each gen's last-accepted row. - attn_metadata.set_mtp_num_accepted(num_accepted_tokens) + index_share.set_mtp_num_accepted(num_accepted_tokens) # The metadata outlives this call, so clear the state even when a # draft step raises: a stale draft-loop flag would make the next # target forward reuse these selections. The loop writes - # set_skip_topk on the metadata yielded below, which is this same - # object: draft_kv_cache_context swaps buffers in place for every - # TrtllmAttentionMetadata, and sparse metadata all derive from it. - draft_scope.callback(_reset_mtp_index_share, attn_metadata) + # set_skip_topk on the extension of the metadata yielded below, + # which is this same object: draft_kv_cache_context swaps buffers + # in place for every TrtllmAttentionMetadata. + draft_scope.callback(index_share.reset_mtp_index_share) draft_attn_metadata = draft_scope.enter_context( self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager)) @@ -807,8 +800,8 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # writes the separate DSA indexer cache at the same positions. attn_metadata.on_update_kv_lens() for i in range(runtime_draft_len): - if uses_mtp_index_share: - attn_metadata.set_skip_topk(i > 0) + if index_share is not None: + index_share.set_skip_topk(i > 0) # Run draft model (mode-specific via helper). The helper # passes ``all_rank_num_tokens`` as a kwarg so the draft model # handles save/restore internally (Eagle3DraftModel.forward @@ -1014,7 +1007,7 @@ def _forward_linear_draft_loop(self, inputs, attn_metadata, spec_metadata, # mode), then 1 token per step in subsequent iterations. # Disable spec_decoding here so the masks/positions stay # correct on subsequent iters. - attn_metadata.use_spec_decoding = False + set_use_spec_decoding(attn_metadata, False) else: if hasattr(attn_metadata, 'kv_lens_cuda'): attn_metadata.kv_lens_cuda[:batch_size] += 1 @@ -1081,7 +1074,8 @@ def _prepare_flash_mla_generation_layout(self, attn_metadata, num_contexts, With no context rows ``prepare_flash_mla`` already staged this exact layout, so the early return merely skips a redundant copy. """ - if num_contexts <= 0 or not attn_metadata.enable_flash_mla: + if (num_contexts <= 0 + or attn_metadata.find_ext(FlashMlaMetadata) is None): return attn_metadata.block_ids_per_seq[:batch_size].copy_( attn_metadata.kv_block_ids_per_seq[:batch_size], non_blocking=True) diff --git a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py index 760b7e1f8c8a..d2e22b98c3ee 100644 --- a/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/eagle3_dynamic_tree.py @@ -21,6 +21,7 @@ import triton import triton.language as tl +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._utils import get_sm_version, nvtx_range from ..attention.backends import AttentionMetadata @@ -340,8 +341,8 @@ def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): def _apply_spec_metadata(self, attn_metadata, batch_size, query_len): """Set spec-dec gen lengths and refresh the C++ position-offset view.""" - attn_metadata.spec_decoding_generation_lengths[:batch_size] = query_len - attn_metadata.update_position_offsets_for_cpp(query_len) + attn_metadata.ext(SpecDecodingMetadata).generation_lengths[:batch_size] = query_len + attn_metadata.ext(SpecDecodingMetadata).update_position_offsets_for_cpp(query_len) @nvtx_range("eagle3_dyn._ensure_spec_tree_manager") def _ensure_spec_tree_manager(self, resource_manager): @@ -619,16 +620,18 @@ def _forward_draft_loop( attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens # Step-0 causal spec-dec (None in prefill-only warmup). - if attn_metadata.spec_decoding_generation_lengths is not None: + if attn_metadata.ext(SpecDecodingMetadata).generation_lengths is not None: total = num_gens * num_step0_tokens - dst = attn_metadata.spec_decoding_position_offsets[:total].view( - num_gens, num_step0_tokens + dst = ( + attn_metadata.ext(SpecDecodingMetadata) + .position_offsets[:total] + .view(num_gens, num_step0_tokens) ) dst.copy_(self._causal_offs[:num_step0_tokens].unsqueeze(0).expand(num_gens, -1)) self._apply_spec_metadata(attn_metadata, num_gens, num_step0_tokens) # Hopper XQA needs packed 1D mask; Blackwell expects padded 3D. - packed_mask = attn_metadata.spec_decoding_packed_mask + packed_mask = attn_metadata.ext(SpecDecodingMetadata).packed_mask packed_mask[:num_gens].zero_() packed_mask[:num_gens, :num_step0_tokens, 0] = self._step0_causal_mask[ :num_step0_tokens @@ -636,7 +639,7 @@ def _forward_draft_loop( if self._needs_mask_repack: self._repack_mask_padded_to_packed(packed_mask, num_gens, num_step0_tokens) - attn_metadata.use_spec_decoding = num_gens > 0 + attn_metadata.ext(SpecDecodingMetadata).use_spec_decoding = num_gens > 0 # KV correction: target processed tokens_per_gen_step per request, # but step 0 should only attend to max_draft_len + 1 accepted-path tokens. @@ -1122,13 +1125,14 @@ def prepare_tree_mask_and_position_offset( selected_parents: torch.Tensor = None, ): """Prepare the mask and position offsets for the next layer.""" - if attn_metadata.spec_decoding_packed_mask is None: + spec = attn_metadata.find_ext(SpecDecodingMetadata) + if spec.packed_mask is None: return spec_tree_manager = self.spec_tree_manager batch_size = attn_metadata.num_seqs num_tokens_current_layer = self.K * (cur_draft_idx + 1) num_tokens_previous_layer = self.K * cur_draft_idx - packed_mask = attn_metadata.spec_decoding_packed_mask + packed_mask = spec.packed_mask if cur_draft_idx == 0: spec_tree_manager.compute_spec_dec_packed_mask( self.tree_mask_init_buffer[:batch_size], @@ -1137,7 +1141,7 @@ def prepare_tree_mask_and_position_offset( self.tree_mask_buffer[ : batch_size * num_tokens_current_layer * num_tokens_current_layer ].copy_(self.tree_mask_init_buffer[:batch_size].view(-1)) - attn_metadata.spec_decoding_position_offsets.fill_(0) + spec.position_offsets.fill_(0) self._apply_spec_metadata(attn_metadata, batch_size, num_tokens_current_layer) else: num_parent_mask = batch_size * cur_draft_idx * self.K * cur_draft_idx * self.K @@ -1146,9 +1150,9 @@ def prepare_tree_mask_and_position_offset( ) prev_total = batch_size * num_tokens_previous_layer - previous_position_offsets = attn_metadata.spec_decoding_position_offsets[ - :prev_total - ].view(batch_size, num_tokens_previous_layer) + previous_position_offsets = spec.position_offsets[:prev_total].view( + batch_size, num_tokens_previous_layer + ) # Fused: gather parent masks + build current mask + compute positions current_mask, new_positions = _build_mask_and_position( @@ -1167,7 +1171,7 @@ def prepare_tree_mask_and_position_offset( ].copy_(current_mask.reshape(-1)) cur_total = batch_size * num_tokens_current_layer - attn_metadata.spec_decoding_position_offsets[:cur_total] = new_positions.reshape(-1) + spec.position_offsets[:cur_total] = new_positions.reshape(-1) self._apply_spec_metadata(attn_metadata, batch_size, num_tokens_current_layer) # Hopper XQA needs packed 1D mask; Blackwell expects padded 3D. @@ -1211,7 +1215,7 @@ def prepare_for_generation( ) - num_accepted_tokens[num_contexts:batch_size] attn_metadata.kv_lens_cuda[:batch_size] += self.K - attn_metadata.use_spec_decoding = True + attn_metadata.ext(SpecDecodingMetadata).use_spec_decoding = True else: num_tokens_previous_layer = cur_draft_idx * self.K diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 407e8ee100be..f08915b4075a 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -28,6 +28,7 @@ from tensorrt_llm.logger import logger from ..._utils import get_sm_version, prefer_pinned +from ..attention.backends.features import DraftKVMetadata from ..attention.backends.interface import AttentionMetadata from ..attention.backends.trtllm import (AttentionBackend, TrtllmAttention, TrtllmAttentionMetadata) @@ -181,37 +182,22 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata, """ if draft_kv_cache_manager is None: return None - if not isinstance(attn_metadata, TrtllmAttentionMetadata): - return None - draft_block_offsets = getattr(attn_metadata, 'draft_kv_cache_block_offsets', - None) - if draft_block_offsets is None: + draft_kv = attn_metadata.find_ext(DraftKVMetadata) + if draft_kv is None: return None - # Marks that a swap happened and carries the backend's state for its - # restore hook; swap_for_draft records the originals on attn_metadata. + # Marks that a swap happened and carries the extensions' state for their + # restore hooks; swap_for_draft records the originals on attn_metadata. saved = {} - if attn_metadata.enable_flash_mla: - if (attn_metadata.draft_block_ids_per_seq is None - or attn_metadata.draft_kv_block_ids_per_seq is None): - raise RuntimeError( - "FlashMLA separate draft KV cache requires dedicated draft block-ID buffers" - ) - attn_metadata.swap_for_draft('block_ids_per_seq', - attn_metadata.draft_block_ids_per_seq) - attn_metadata.swap_for_draft('kv_block_ids_per_seq', - attn_metadata.draft_kv_block_ids_per_seq) attn_metadata.swap_for_draft('kv_cache_manager', draft_kv_cache_manager) attn_metadata.swap_for_draft('kv_cache_block_offsets', - attn_metadata.draft_kv_cache_block_offsets) + draft_kv.kv_cache_block_offsets) attn_metadata.swap_for_draft( 'host_kv_cache_block_offsets', draft_kv_cache_manager.host_kv_cache_block_offsets) - if attn_metadata.enable_flash_mla: - attn_metadata.prepare_flash_mla() - # Backends select any additional draft-forward state, such as native DSA - # indexer buffers or DeepSeek-V4 sparse tables and pool pointers. + # Extensions select any additional draft-forward state, such as FlashMLA + # block IDs, native DSA indexer buffers or DeepSeek-V4 sparse tables. backend_saved = attn_metadata.prepare_for_draft_forward() if backend_saved is not None: saved['saved_backend_state'] = backend_saved @@ -223,11 +209,6 @@ def restore_attn_metadata_after_draft_replay(attn_metadata, saved_state): if saved_state is None: return attn_metadata.restore_draft_swaps() - if attn_metadata.enable_flash_mla: - # Target and draft block-ID buffers are independent. Restoring only - # needs to invalidate the scheduler metadata; refreshing the unchanged - # target buffers would repeat request-specific H2D work. - attn_metadata._flash_mla_metadata_valid = False attn_metadata.restore_after_draft_forward( saved_state.get('saved_backend_state')) diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 2107fbed26c8..92d4bb84b625 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -9,6 +9,8 @@ from tensorrt_llm._utils import prefer_pinned from ..attention.backends import AttentionMetadata +from ..attention.backends.features import RecurrentMetadata +from ..attention.backends.fp4_mla.state import Fp4MlaState from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager @@ -299,7 +301,8 @@ def _commit_target_mamba_states( attn_metadata.kv_cache_manager.update_mamba_states( attn_metadata=attn_metadata, num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices, + state_indices=attn_metadata.ext( + RecurrentMetadata).mamba_metadata.state_indices, ) def _forward_impl( @@ -437,7 +440,8 @@ def _forward_impl( if self._auxiliary_state_handlers: self.commit_auxiliary_speculative_states( num_accepted_tokens, - attn_metadata.mamba_metadata.state_indices[:batch_size], + attn_metadata.ext(RecurrentMetadata).mamba_metadata. + state_indices[:batch_size], attn_metadata.num_contexts, ) @@ -951,7 +955,7 @@ def change_attn_metadata(self, num_accepted_tokens: torch.Tensor, attn_metadata.kv_lens_cuda[num_contexts:batch_size].clamp_( min=runtime_draft_len) attn_metadata.on_update_kv_lens() - if getattr(attn_metadata, "fp4_mla_state", None) is not None: + if attn_metadata.find_ext(Fp4MlaState) is not None: attn_metadata.update_for_spec_dec() if attn_metadata.kv_cache_params is not None and not attn_metadata.is_cuda_graph: diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 91bb857fe279..1529f6c40b38 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -20,10 +20,12 @@ import torch import triton +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from tensorrt_llm._utils import get_sm_version, nvtx_range from tensorrt_llm.mapping import Mapping +from ..attention.backends.features import RecurrentMetadata from ..distributed.ops import allgather from ..model_config import ModelConfig from ..pyexecutor.llm_request import LlmRequest @@ -178,6 +180,7 @@ def __init__( self._needs_mask_repack = sm < 100 or sm in (120, 121) def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): + spec = attn_metadata.find_ext(SpecDecodingMetadata) super()._prepare_attn_metadata_for_spec_dec(attn_metadata) batch_size = attn_metadata.num_seqs @@ -189,20 +192,18 @@ def _prepare_attn_metadata_for_spec_dec(self, attn_metadata): self._saved_kv_lens_cuda = None # Restore verify metadata after the draft loop mutates it. - if attn_metadata.spec_decoding_packed_mask is not None: - self._saved_packed_mask = attn_metadata.spec_decoding_packed_mask[:batch_size].clone() + if spec.packed_mask is not None: + self._saved_packed_mask = spec.packed_mask[:batch_size].clone() else: self._saved_packed_mask = None - if attn_metadata.spec_decoding_position_offsets is not None: - self._saved_position_offsets = attn_metadata.spec_decoding_position_offsets.clone() - self._saved_position_offsets_cpp = attn_metadata.spec_decoding_position_offsets_cpp + if spec.position_offsets is not None: + self._saved_position_offsets = spec.position_offsets.clone() + self._saved_position_offsets_cpp = spec.position_offsets_cpp else: self._saved_position_offsets = None self._saved_position_offsets_cpp = None - if attn_metadata.spec_decoding_generation_lengths is not None: - self._saved_generation_lengths = attn_metadata.spec_decoding_generation_lengths[ - :batch_size - ].clone() + if spec.generation_lengths is not None: + self._saved_generation_lengths = spec.generation_lengths[:batch_size].clone() else: self._saved_generation_lengths = None @@ -212,6 +213,7 @@ def prepare_position_ids_and_last_tokens(self, position_ids, seq_lens_cuda): return position_ids, last_tokens_idx def _restore_attn_metadata_from_spec_dec(self, attn_metadata): + spec = attn_metadata.find_ext(SpecDecodingMetadata) super()._restore_attn_metadata_from_spec_dec(attn_metadata) if self._saved_kv_lens_cuda is not None: @@ -221,18 +223,16 @@ def _restore_attn_metadata_from_spec_dec(self, attn_metadata): if self._saved_packed_mask is not None: batch_size = self._saved_packed_mask.shape[0] - attn_metadata.spec_decoding_packed_mask[:batch_size].copy_(self._saved_packed_mask) + spec.packed_mask[:batch_size].copy_(self._saved_packed_mask) self._saved_packed_mask = None if self._saved_position_offsets is not None: - attn_metadata.spec_decoding_position_offsets.copy_(self._saved_position_offsets) - attn_metadata.spec_decoding_position_offsets_cpp = self._saved_position_offsets_cpp + spec.position_offsets.copy_(self._saved_position_offsets) + spec.position_offsets_cpp = self._saved_position_offsets_cpp self._saved_position_offsets = None self._saved_position_offsets_cpp = None if self._saved_generation_lengths is not None: batch_size = self._saved_generation_lengths.shape[0] - attn_metadata.spec_decoding_generation_lengths[:batch_size].copy_( - self._saved_generation_lengths - ) + spec.generation_lengths[:batch_size].copy_(self._saved_generation_lengths) self._saved_generation_lengths = None def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): @@ -245,7 +245,7 @@ def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): return ( attn_metadata.all_rank_num_tokens, - attn_metadata.force_prepare_spec_dec_tree_mask, + attn_metadata.ext(SpecDecodingMetadata).force_prepare_tree_mask, ) = self._saved_extra_state self._saved_extra_state = None @@ -254,24 +254,18 @@ def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): # ------------------------------------------------------------------ # def _apply_spec_metadata(self, attn_metadata, batch_size, query_len): """Set spec-dec gen lengths and refresh the C++ position-offset view.""" - attn_metadata.spec_decoding_generation_lengths[:batch_size] = query_len - attn_metadata.update_position_offsets_for_cpp(query_len) + attn_metadata.ext(SpecDecodingMetadata).generation_lengths[:batch_size] = query_len + attn_metadata.ext(SpecDecodingMetadata).update_position_offsets_for_cpp(query_len) def _refresh_blackwell_tree_mask_metadata(self, attn_metadata): - if not getattr(attn_metadata, "use_spec_decoding", False): - return - if not getattr(attn_metadata, "is_spec_dec_dynamic_tree", False): - return - - first_sparse = getattr(attn_metadata, "spec_bl_tree_first_sparse_mask_offset_kv", None) - bl_tree_mask = getattr(attn_metadata, "spec_decoding_bl_tree_mask", None) - if first_sparse is None and bl_tree_mask is None: + spec = attn_metadata.find_ext(SpecDecodingMetadata) + if spec is None or not spec.use_spec_decoding or not spec.is_spec_dec_dynamic_tree: return - if bl_tree_mask is not None: - bl_tree_mask.zero_() - if first_sparse is not None: - attn_metadata.update_blackwell_first_sparse_mask_offset() + if spec.bl_tree_mask is not None: + spec.bl_tree_mask.zero_() + if spec.bl_tree_first_sparse_mask_offset_kv is not None: + spec.update_blackwell_first_sparse_mask_offset() def _repack_mask_padded_to_packed(self, mask_buf, n_req, n_tok): """Compact padded masks into the flat prefix XQA expects.""" @@ -401,13 +395,14 @@ def prepare_tree_mask_and_position_offset( self, cur_draft_idx, attn_metadata, selected_parents=None ): """Prepare mask and position offsets for the next draft layer.""" - if attn_metadata.spec_decoding_packed_mask is None: + spec = attn_metadata.find_ext(SpecDecodingMetadata) + if spec.packed_mask is None: return spec_tree_manager = self.spec_tree_manager batch_size = attn_metadata.num_seqs num_tokens_current_layer = self.K * (cur_draft_idx + 1) num_tokens_previous_layer = self.K * cur_draft_idx - packed_mask = attn_metadata.spec_decoding_packed_mask + packed_mask = spec.packed_mask if cur_draft_idx == 0: spec_tree_manager.compute_spec_dec_packed_mask( self.tree_mask_init_buffer[:batch_size], @@ -416,7 +411,7 @@ def prepare_tree_mask_and_position_offset( self.tree_mask_buffer[ : batch_size * num_tokens_current_layer * num_tokens_current_layer ].copy_(self.tree_mask_init_buffer[:batch_size].view(-1)) - attn_metadata.spec_decoding_position_offsets.fill_(0) + spec.position_offsets.fill_(0) self._apply_spec_metadata(attn_metadata, batch_size, num_tokens_current_layer) else: num_parent_mask = batch_size * cur_draft_idx * self.K * cur_draft_idx * self.K @@ -425,9 +420,9 @@ def prepare_tree_mask_and_position_offset( ) prev_total = batch_size * num_tokens_previous_layer - previous_position_offsets = attn_metadata.spec_decoding_position_offsets[ - :prev_total - ].view(batch_size, num_tokens_previous_layer) + previous_position_offsets = spec.position_offsets[:prev_total].view( + batch_size, num_tokens_previous_layer + ) current_mask, new_positions = _build_mask_and_position( parent_mask, @@ -445,7 +440,7 @@ def prepare_tree_mask_and_position_offset( ].copy_(current_mask.reshape(-1)) cur_total = batch_size * num_tokens_current_layer - attn_metadata.spec_decoding_position_offsets[:cur_total] = new_positions.reshape(-1) + spec.position_offsets[:cur_total] = new_positions.reshape(-1) self._apply_spec_metadata(attn_metadata, batch_size, num_tokens_current_layer) if self._needs_mask_repack: @@ -703,7 +698,7 @@ def _forward_impl( attn_metadata.kv_cache_manager.update_mamba_states( attn_metadata=attn_metadata, num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices, + state_indices=attn_metadata.ext(RecurrentMetadata).mamba_metadata.state_indices, accepted_leaf_positions=accepted_leaf_positions, ) @@ -712,10 +707,10 @@ def _forward_impl( # draft loop raises (see SpecWorkerBase.forward try/finally). self._saved_extra_state = ( attn_metadata.all_rank_num_tokens, - attn_metadata.force_prepare_spec_dec_tree_mask, + attn_metadata.ext(SpecDecodingMetadata).force_prepare_tree_mask, ) self._prepare_attn_metadata_for_spec_dec(attn_metadata) - attn_metadata.force_prepare_spec_dec_tree_mask = True + attn_metadata.ext(SpecDecodingMetadata).force_prepare_tree_mask = True # (c) Run the MTP draft tree loop -> build + store the next tree. draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) @@ -738,10 +733,10 @@ def _forward_impl( self._restore_attn_metadata_from_spec_dec(attn_metadata) ( attn_metadata.all_rank_num_tokens, - attn_metadata.force_prepare_spec_dec_tree_mask, + attn_metadata.ext(SpecDecodingMetadata).force_prepare_tree_mask, ) = self._saved_extra_state self._saved_extra_state = None - attn_metadata.use_spec_decoding = True + attn_metadata.ext(SpecDecodingMetadata).use_spec_decoding = True # (d) Prepare next_new_tokens for overlap scheduler. next_new_tokens = self._prepare_next_new_tokens( @@ -876,6 +871,7 @@ def _forward_draft_loop( batch_size, ): """Draft the next dynamic tree with growing context.""" + spec = attn_metadata.find_ext(SpecDecodingMetadata) spec_tree_manager = self.spec_tree_manager assert batch_size <= self._max_batch_size, ( @@ -898,21 +894,19 @@ def _forward_draft_loop( # Reset verify-time tree metadata to accepted-path width. num_step0_tokens = self._max_path_len - if attn_metadata.spec_decoding_generation_lengths is not None: + if spec.generation_lengths is not None: total = num_gens * num_step0_tokens - dst = attn_metadata.spec_decoding_position_offsets[:total].view( - num_gens, num_step0_tokens - ) + dst = spec.position_offsets[:total].view(num_gens, num_step0_tokens) dst.copy_(self._causal_offs[:num_step0_tokens].unsqueeze(0).expand(num_gens, -1)) self._apply_spec_metadata(attn_metadata, num_gens, num_step0_tokens) - packed_mask = attn_metadata.spec_decoding_packed_mask + packed_mask = spec.packed_mask packed_mask[:num_gens].zero_() packed_mask[:num_gens, :num_step0_tokens, 0] = self._step0_causal_mask[ :num_step0_tokens ] if self._needs_mask_repack: self._repack_mask_padded_to_packed(packed_mask, num_gens, num_step0_tokens) - attn_metadata.use_spec_decoding = num_gens > 0 + spec.use_spec_decoding = num_gens > 0 if num_gens > 0 and hasattr(attn_metadata, "kv_lens_cuda"): attn_metadata.kv_lens_cuda[num_contexts:batch_size] -= self._kv_correction self._refresh_blackwell_tree_mask_metadata(attn_metadata) @@ -1073,7 +1067,7 @@ def _prepare_draft_layer_metadata( self._max_path_len ) - num_accepted_tokens[num_contexts:batch_size] attn_metadata.kv_lens_cuda[:batch_size] += self.K - attn_metadata.use_spec_decoding = True + attn_metadata.ext(SpecDecodingMetadata).use_spec_decoding = True self._refresh_blackwell_tree_mask_metadata(attn_metadata) else: num_tokens_previous_layer = cur_draft_idx * self.K diff --git a/tensorrt_llm/_torch/speculative/sa_worker.py b/tensorrt_llm/_torch/speculative/sa_worker.py index 0f508333fb22..e48954a7482d 100644 --- a/tensorrt_llm/_torch/speculative/sa_worker.py +++ b/tensorrt_llm/_torch/speculative/sa_worker.py @@ -30,6 +30,7 @@ from tensorrt_llm._utils import prefer_pinned +from ..attention.backends.features import RecurrentMetadata from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from .interface import SpecMetadata, SpecWorkerBase from .suffix_automaton import SuffixAutomatonManager @@ -189,7 +190,7 @@ def _forward_impl( attn_metadata.kv_cache_manager.update_mamba_states( attn_metadata=attn_metadata, num_accepted_tokens=num_accepted_tokens, - state_indices=attn_metadata.mamba_metadata.state_indices, + state_indices=attn_metadata.ext(RecurrentMetadata).mamba_metadata.state_indices, ) # Step 3-4: Extend SA and generate next draft tokens using GPU kernel diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index 4382827aad45..cf0850fa54d8 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -16,6 +16,7 @@ import tensorrt_llm._torch.model_config import tensorrt_llm._torch.pyexecutor.config_utils import tensorrt_llm.bindings +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata from tensorrt_llm._torch.attention.backends.utils import get_attention_backend from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper from tensorrt_llm._torch.metadata import KVCacheParams @@ -662,13 +663,12 @@ def create_run_pack( AttentionCls = get_attention_backend( self.model_config.attn_backend, sparse_params=sparse_params ) - metadata_cls = AttentionCls.Metadata sparse_metadata_params = ( sparse_attention_config.to_sparse_metadata_params(pretrained_config=pretrained_config) if sparse_attention_config is not None else None ) - attn_metadata = metadata_cls( + attn_metadata = AttentionCls.create_metadata( seq_lens=torch.tensor([seq_len_q] * batch_size, dtype=torch.int), request_ids=list(range(request_id_begin, request_id_begin + batch_size)), max_num_requests=kv_cache_manager.max_batch_size, @@ -704,12 +704,12 @@ def create_run_pack( # buffer stays one column wide and DeepGEMM aborts on the next_n mismatch. # Shapes only -- spec-dec masking stays off. # - # Gate on kv_lens_cuda_2d, not on the method: update_spec_dec_param() is on + # Gate on the DSA extension, not on the method: update_spec_dec_param() is on # the base metadata class, so hasattr() would let every backend in, and the # base sets max_total_draft_tokens unconditionally -- which reaches the - # attention op cache key and FMHA kernel selection. kv_lens_cuda_2d exists - # only on DSA metadata, which is the backend that needs this. - if run_type == "GEN" and seq_len_q > 1 and hasattr(attn_metadata, "kv_lens_cuda_2d"): + # attention op cache key and FMHA kernel selection. DSA is the backend that + # needs this. + if run_type == "GEN" and seq_len_q > 1 and attn_metadata.find_ext(DSAMetadata) is not None: attn_metadata.update_spec_dec_param( batch_size=batch_size, is_spec_decoding_enabled=False, diff --git a/tests/microbenchmarks/attention_perf/attention_perf_harness.py b/tests/microbenchmarks/attention_perf/attention_perf_harness.py index bc7ccc641f46..a2f65647d41a 100644 --- a/tests/microbenchmarks/attention_perf/attention_perf_harness.py +++ b/tests/microbenchmarks/attention_perf/attention_perf_harness.py @@ -52,6 +52,7 @@ import tensorrt_llm from tensorrt_llm._torch.attention.backends import TrtllmAttention +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionInputType, AttentionRuntimeFeatures, @@ -548,7 +549,7 @@ def _build_mla_gen_metadata(case: AttnCase, kv_mgr: KVCacheManagerV2): mapping=Mapping(world_size=1, tp_size=1, rank=0), # FlashMLA decode kernel is sm90-only — that is where the DeepSeek-R1 # gen regressions live; on other archs this routes to the generic kernel. - enable_flash_mla=(torch.cuda.get_device_capability() == (9, 0)), + extension_types=(FlashMlaMetadata,) if torch.cuda.get_device_capability() == (9, 0) else (), ) metadata = TrtllmAttention.Metadata(**kwargs) metadata.prepare() @@ -600,7 +601,7 @@ def _make_mla_gen_inputs(case: AttnCase, device: torch.device): # which owns an Indexer submodule; its projections feed _forward_dsa_attn, # which updates the indexer cache and dispatches sparse MLA attention. # Unlike the GQA/MLA paths this needs the FULL MLA module (weights + indexer), -# a DSACacheManagerV2 (indexer-K pool), and DSAtrtllmAttentionMetadata. +# a DSACacheManagerV2 (indexer-K pool), and metadata with the DSA extension. # # ARCH: the indexer's DeepGEMM MQA-logits kernels exist only for sm90/sm100/sm103. # On sm120 (and any arch without DeepGEMM) construction/forward raises; the caller @@ -722,7 +723,7 @@ def _build_dsa_ctx(case: AttnCase, device: torch.device): assert requests is not None AttentionCls = get_attention_backend("TRTLLM", sparse_config) - metadata = AttentionCls.Metadata( + metadata = AttentionCls.create_metadata( seq_lens=torch.tensor([case.seq_len] * case.batch_size, dtype=torch.int), request_ids=request_ids, max_num_requests=case.batch_size, @@ -734,7 +735,7 @@ def _build_dsa_ctx(case: AttnCase, device: torch.device): use_cache=True, num_cached_tokens_per_seq=[0] * case.batch_size ), mapping=Mapping(world_size=1, tp_size=1, rank=0), - sparse_attention_config=sparse_config, + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), ) metadata.prepare() diff --git a/tests/microbenchmarks/bench_fp4_mla_decode.py b/tests/microbenchmarks/bench_fp4_mla_decode.py index 32466a6c50ca..ec42f0591dda 100644 --- a/tests/microbenchmarks/bench_fp4_mla_decode.py +++ b/tests/microbenchmarks/bench_fp4_mla_decode.py @@ -53,7 +53,7 @@ scatter_fp4_mla_kv_cache, ) from tensorrt_llm._torch.attention.backends.fp4_mla.cache_manager import Fp4MlaKVCacheManagerV2 -from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState +from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState, find_fp4_mla from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig from tensorrt_llm.mapping import Mapping @@ -137,6 +137,13 @@ def _expand_qk_residual_terms( ) +def _fp4_mla_metadata(fp4_mla_state: Fp4MlaState, **fields) -> SimpleNamespace: + """A metadata stub carrying ``fp4_mla_state`` as its FP4 MLA extension.""" + metadata = SimpleNamespace(**fields) + metadata.find_ext = lambda cls: fp4_mla_state if cls is Fp4MlaState else None + return metadata + + def _build_multi_seq_metadata( kv_cache_manager: Fp4MlaKVCacheManagerV2, *, seq_lens: list[int], page_size: int, layer_idx: int ) -> SimpleNamespace: @@ -194,7 +201,7 @@ def _build_multi_seq_metadata( kv_global_scale = torch.tensor([_BENCH_KV_GLOBAL_SCALE], dtype=torch.float32, device=device) q_global_scale = torch.tensor([FP4_MLA_Q_GLOBAL_SCALE], dtype=torch.float32, device=device) - return SimpleNamespace( + return _fp4_mla_metadata( kv_cache_manager=kv_cache_manager, page_size=page_size, num_contexts=num_seqs, @@ -230,15 +237,16 @@ def _build_multi_seq_metadata( def _materialize_reference_cache_storage( metadata: SimpleNamespace, layer_idx: int, head_dim: int ) -> torch.Tensor: + fp4_mla_state = find_fp4_mla(metadata) kv_cache, sf_cache = metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) sf_cache = sf_cache.view(torch.float8_e4m3fn) storage_head_dim = kv_cache.shape[-1] * 2 static_global_scale = float(_get_fp4_mla_global_scale(metadata, kv_cache.device).item()) pages = [] dequantized_pages = {} - page_rows = metadata.fp4_mla_state.paged_kv_indices.view( + page_rows = fp4_mla_state.paged_kv_indices.view( metadata.num_seqs, - metadata.fp4_mla_state.page_table_stride, + fp4_mla_state.page_table_stride, ) # Preserve the fixed row stride used by paged_kv_indptr_decode, including # padding for shorter requests. The reference masks padding by KV length. @@ -341,11 +349,11 @@ def _build_fp4_mla_attention_decode_case( page_size=page_size, layer_idx=layer_idx, ) - assert metadata.fp4_mla_state.v_scale_pool is not None + assert find_fp4_mla(metadata).v_scale_pool is not None persistent_pool_base = kv_cache_manager.get_mla_v_packed_pool_base() if persistent_pool_base is not None: persistent_pool_base.zero_() - metadata.fp4_mla_state.v_scale_pool.zero_() + find_fp4_mla(metadata).v_scale_pool.zero_() latent = ( torch.randn(sum(seq_lens), head_dim, dtype=torch.bfloat16, device=device) * 0.25 @@ -360,13 +368,13 @@ def _build_fp4_mla_attention_decode_case( v_head_dim=kv_lora_rank, ) # Decode reference exercises the quantized cache without an HP overlay. - metadata.fp4_mla_state.hp_pool.zero_() + find_fp4_mla(metadata).hp_pool.zero_() torch.cuda.synchronize() metadata.num_contexts = 0 - metadata.fp4_mla_state.num_context_blocks = 0 - metadata.fp4_mla_state.num_generation_blocks = ( - len(seq_lens) * metadata.fp4_mla_state.page_table_stride + find_fp4_mla(metadata).num_context_blocks = 0 + find_fp4_mla(metadata).num_generation_blocks = ( + len(seq_lens) * find_fp4_mla(metadata).page_table_stride ) metadata.prompt_lens_cuda_runtime = torch.full( (len(seq_lens),), query_len_per_seq, dtype=torch.int32, device=device @@ -404,6 +412,7 @@ def _fp4_mla_attention_decode_reference( layer_idx: int = 0, local_layer: int = 0, ) -> tuple[torch.Tensor, list[torch.Tensor], list[torch.Tensor]]: + fp4_mla_state = find_fp4_mla(metadata) head_dim = kv_lora_rank + qk_rope_head_dim dequant_cache = _materialize_reference_cache(metadata, layer_idx, head_dim) dequant_k_residual = None @@ -415,7 +424,7 @@ def _fp4_mla_attention_decode_reference( ) num_heads = q_nope.shape[1] q_full = torch.cat((q_nope, q_pe), dim=-1).reshape(-1, head_dim) - global_scale = metadata.fp4_mla_state.q_global_scale + global_scale = fp4_mla_state.q_global_scale q_fp4, q_sf = torch.ops.trtllm.fp4_quantize_with_residual( q_full, global_scale, @@ -432,16 +441,16 @@ def _fp4_mla_attention_decode_reference( ) p_dequant = None - if "_fp4_mla_attention_p_buf" in metadata.fp4_mla_state.workspaces: + if "_fp4_mla_attention_p_buf" in fp4_mla_state.workspaces: p_dequant = _dequant_fp4_swizzled( - metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_buf"], - metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_sf_buf"], + fp4_mla_state.workspaces["_fp4_mla_attention_p_buf"], + fp4_mla_state.workspaces["_fp4_mla_attention_p_sf_buf"], logical_dim=metadata.page_size, sf_per_token=metadata.page_size // FP4_BLOCK_SIZE, global_scale=FP4_MLA_P_GLOBAL_SCALE, ) - indptr = metadata.fp4_mla_state.paged_kv_indptr_decode.cpu().tolist() + indptr = fp4_mla_state.paged_kv_indptr_decode.cpu().tolist() kv_lens = metadata.kv_lens_cuda_runtime.cpu().tolist() num_seqs = metadata.num_seqs - metadata.num_contexts query_len_per_seq = q_nope.shape[0] // num_seqs @@ -1285,15 +1294,15 @@ def run_one( # The synthetic benchmark never changes its sequence or append lengths. # Bind the authoritative runtime tensors as an already-populated result so # timing excludes the production MTP length-correction helper. - metadata.fp4_mla_state.generation_kv_lens = metadata.kv_lens_cuda_runtime[:batch] - metadata.fp4_mla_state.generation_append_lens = metadata.prompt_lens_cuda_runtime[:batch] - metadata.fp4_mla_state.generation_lengths_num_tokens = batch * q_len - metadata.fp4_mla_state.generation_lengths_num_seqs = batch - metadata.fp4_mla_state.generation_lengths_num_contexts = metadata.num_contexts - metadata.fp4_mla_state.generation_lengths_capture_recorded = True + find_fp4_mla(metadata).generation_kv_lens = metadata.kv_lens_cuda_runtime[:batch] + find_fp4_mla(metadata).generation_append_lens = metadata.prompt_lens_cuda_runtime[:batch] + find_fp4_mla(metadata).generation_lengths_num_tokens = batch * q_len + find_fp4_mla(metadata).generation_lengths_num_seqs = batch + find_fp4_mla(metadata).generation_lengths_num_contexts = metadata.num_contexts + find_fp4_mla(metadata).generation_lengths_capture_recorded = True if use_cuda_graph: # Synthetic benchmark page tables are immutable and never grow on replay. - metadata.fp4_mla_state._paged_kv_indices = metadata.fp4_mla_state.paged_kv_indices + find_fp4_mla(metadata)._paged_kv_indices = find_fp4_mla(metadata).paged_kv_indices generation_latent = None generation_rotary_cos_sin = None persistent_v_pack = False @@ -1304,10 +1313,10 @@ def run_one( device=q_nope.device, ) metadata.num_ctx_tokens = 0 - metadata.fp4_mla_state.batch_indices = torch.arange( + find_fp4_mla(metadata).batch_indices = torch.arange( batch, dtype=torch.int32, device=q_nope.device ).repeat_interleave(q_len) - metadata.fp4_mla_state.positions = ( + find_fp4_mla(metadata).positions = ( context_lens[:, None] + torch.arange(q_len, dtype=torch.int32, device=q_nope.device)[None, :] ).reshape(-1) @@ -1405,16 +1414,16 @@ def record_launcher(*args: object, **kwargs: object) -> object: q = torch.cat((q_nope, q_pe), dim=-1).contiguous() baseline_q_fp4, baseline_q_sf = torch.ops.trtllm.fp4_quantize_with_residual( q.view(-1, q.shape[-1]), - metadata.fp4_mla_state.q_global_scale, + find_fp4_mla(metadata).q_global_scale, FP4_MLA_Q_RESIDUAL_DIM, is_act=True, ) baseline_q_sf = baseline_q_sf.view(torch.float8_e4m3fn).reshape(-1) def decode() -> None: - q_fp4 = getattr(metadata.fp4_mla_state, "prequantized_q", None) - q_sf = getattr(metadata.fp4_mla_state, "prequantized_q_sf", None) - q_batch_capacity = getattr(metadata.fp4_mla_state, "q_batch_capacity", None) + q_fp4 = getattr(find_fp4_mla(metadata), "prequantized_q", None) + q_sf = getattr(find_fp4_mla(metadata), "prequantized_q_sf", None) + q_batch_capacity = getattr(find_fp4_mla(metadata), "q_batch_capacity", None) if q_fp4 is None: q_fp4 = baseline_q_fp4 q_sf = baseline_q_sf diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index d5ae45ecf748..d7246a7447cd 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -510,7 +510,7 @@ def _forward(metadata, q, q_pe): return out[0] if isinstance(out, tuple) else out def _create_metadata(AttentionCls, case, mgr): - return AttentionCls.Metadata( + return AttentionCls.create_metadata( num_contexts=0, kv_cache_params=KVCacheParams( use_cache=True, num_cached_tokens_per_seq=case.num_cached_tokens @@ -699,7 +699,7 @@ def _run_mla_context_backend( ) expected_latents = _split_packed_tokens(expected_latent_cache, case.seq_lens) cache_atol, cache_rtol = _tolerances(case, kv_dtype) if fuse_rope else (0.0, 0.0) - metadata = AttentionCls.Metadata( + metadata = AttentionCls.create_metadata( num_contexts=case.num_contexts, kv_cache_params=KVCacheParams( use_cache=True, num_cached_tokens_per_seq=case.num_cached_tokens @@ -936,7 +936,7 @@ def run_backend( mgr = None if case.cache == "none": - metadata = AttentionCls.Metadata( + metadata = AttentionCls.create_metadata( max_num_requests=case.num_seqs, max_num_tokens=case.max_num_tokens, kv_cache_manager=None, @@ -958,7 +958,7 @@ def run_backend( def create_metadata(AttentionCls, case, mgr, *, num_contexts: int = 0): seq_lens_kv = torch.tensor(case.seq_lens_kv, dtype=torch.int) if case.is_cross else None - return AttentionCls.Metadata( + return AttentionCls.create_metadata( num_contexts=num_contexts, kv_cache_params=KVCacheParams( use_cache=True, num_cached_tokens_per_seq=case.num_cached_tokens diff --git a/tests/unittest/_torch/attention/fmha_test_utils.py b/tests/unittest/_torch/attention/fmha_test_utils.py index a20a99dbf156..6f734b781033 100644 --- a/tests/unittest/_torch/attention/fmha_test_utils.py +++ b/tests/unittest/_torch/attention/fmha_test_utils.py @@ -13,10 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace from typing import Callable import torch +from tensorrt_llm._torch.attention.backends.features import HelixMetadata +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha, FmhaPhase from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams, PhasedFmha from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs @@ -145,3 +148,44 @@ def forward( forward_args: AttentionForwardArgs, ) -> None: self._events.append(("forward", self._name)) + + +_SPEC_FIELDS = ( + "is_spec_decoding_enabled", + "use_spec_decoding", + "is_spec_dec_tree", + "is_spec_dec_dynamic_tree", + "max_total_draft_tokens", +) +_HELIX_FIELDS = { + "helix_position_offsets": "position_offsets", + "helix_is_inactive_rank": "is_inactive_rank", + "helix_kv_bounds": "kv_bounds", + "helix_local_slots": "local_slots", + "_helix_spec_tokens_valid": "spec_tokens_valid", +} + + +def metadata_stub(**fields) -> SimpleNamespace: + """A duck-typed metadata serving its spec-dec and Helix fields through extensions. + + The spec-dec extension is always attached, as on TRTLLM metadata; Helix only + when one of its fields is set. + """ + spec = SpecDecodingMetadata(None) + for name in _SPEC_FIELDS: + if name in fields: + setattr(spec, name, fields.pop(name)) + extensions = {SpecDecodingMetadata: spec} + helix_fields = { + _HELIX_FIELDS[name]: fields.pop(name) for name in list(fields) if name in _HELIX_FIELDS + } + if any(value is not None and value is not False for value in helix_fields.values()): + helix = HelixMetadata(None) + for name, value in helix_fields.items(): + setattr(helix, name, value) + extensions[HelixMetadata] = helix + metadata = SimpleNamespace(**fields) + metadata.find_ext = extensions.get + metadata.ext = extensions.__getitem__ + return metadata diff --git a/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py b/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py index 1027b6b97909..513816319d74 100644 --- a/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py +++ b/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py @@ -3,58 +3,65 @@ """CPU coverage of ragged preparation, update and backend dispatch contracts.""" from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, patch import pytest import torch +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import AttentionInputType from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import ( forward_generation_sparse_attn, ) from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer -from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata pytestmark = pytest.mark.cpu_only -def _metadata(seq_lens: list[int], num_contexts: int = 0) -> DSAtrtllmAttentionMetadata: - metadata = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) +def _metadata(seq_lens: list[int], num_contexts: int = 0) -> TrtllmAttentionMetadata: + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) metadata._seq_lens = torch.tensor(seq_lens, dtype=torch.int32) metadata._seq_lens_cuda = metadata._seq_lens.clone() metadata._num_contexts = num_contexts metadata._num_generations = len(seq_lens) - num_contexts metadata._num_ctx_tokens = sum(seq_lens[:num_contexts]) metadata._num_tokens = sum(seq_lens) - metadata.max_draft_tokens = 5 - metadata.runtime_tokens_per_gen_step = 0 - metadata._indexer_compress_ratio = 1 metadata.kv_cache_manager = None - metadata.num_sms = 1 + with patch( + "tensorrt_llm.deep_gemm.get_num_sms", + return_value=1, + ): + dsa = DSAMetadata(metadata) + dsa.max_draft_tokens = 5 + dsa._indexer_compress_ratio = 1 + metadata.replace_ext(dsa) + metadata.replace_ext(SpecDecodingMetadata(metadata)) return metadata @pytest.mark.parametrize("runtime_width", [1, 2, 6]) def test_schedule_update_keeps_runtime_query_width(monkeypatch, runtime_width: int) -> None: metadata = _metadata([runtime_width, runtime_width]) - metadata.runtime_tokens_per_gen_step = runtime_width + metadata.ext(DSAMetadata).runtime_tokens_per_gen_step = runtime_width + dsa = metadata.ext(DSAMetadata) metadata.kv_lens_cuda = torch.tensor([102, 204], dtype=torch.int32) metadata.kv_lens_cuda_runtime = metadata.kv_lens_cuda - metadata.kv_lens_cuda_2d = torch.full((2, 6), -1, dtype=torch.int32) - metadata.scheduler_metadata_buffer = torch.empty(1, dtype=torch.int32) - metadata.scheduler_metadata_buffer_full_next_n = torch.empty(1, dtype=torch.int32) - metadata.req_idx_per_token = torch.empty(metadata.num_tokens, dtype=torch.int32) - metadata.gen_kv_indptr = torch.zeros(3, dtype=torch.int64) - metadata.gen_cached_token_indptr = torch.zeros(3, dtype=torch.int64) - metadata.enable_ragged_verification = False - metadata.enable_flash_mla = False - metadata.use_fp8_ds_mla = False - metadata.use_expanded_buffers_for_mtp = False - metadata.expand_for_dsl = False - metadata._invalidate_pool_view_cache = Mock() - metadata._compute_kv_lens_row_reorder = Mock() - metadata.prepare_dense_topk_indices = Mock() + dsa.kv_lens_cuda_2d = torch.full((2, 6), -1, dtype=torch.int32) + dsa.scheduler_metadata_buffer = torch.empty(1, dtype=torch.int32) + dsa.scheduler_metadata_buffer_full_next_n = torch.empty(1, dtype=torch.int32) + dsa.req_idx_per_token = torch.empty(metadata.num_tokens, dtype=torch.int32) + dsa.gen_kv_indptr = torch.zeros(3, dtype=torch.int64) + dsa.gen_cached_token_indptr = torch.zeros(3, dtype=torch.int64) + dsa.enable_ragged_verification = False + dsa.use_fp8_ds_mla = False + dsa.use_expanded_buffers_for_mtp = False + dsa.expand_for_dsl = False + dsa._invalidate_pool_view_cache = Mock() + dsa._compute_kv_lens_row_reorder = Mock() + dsa.prepare_dense_topk_indices = Mock() scheduler_inputs = [] def schedule(context_lens, block_kv, num_sms): @@ -91,53 +98,53 @@ def test_prepare_initializes_row_extents_and_refresh_reuses_storage( monkeypatch, num_contexts: int ) -> None: metadata = _metadata(([3] if num_contexts else []) + [2, 4], num_contexts) - metadata.ragged_verify_lens = [2, 4] - metadata.enable_ragged_verification = True - metadata.is_spec_decoding_enabled = False + metadata.ext(DSAMetadata).ragged_verify_lens = [2, 4] + dsa = metadata.ext(DSAMetadata) + dsa.enable_ragged_verification = True metadata.sparse_metadata_params = SimpleNamespace(use_cute_dsl_paged_mqa_logits=False) - metadata.gen_token_repeats_cuda = torch.empty(2, dtype=torch.int64) - metadata.host_gen_token_repeats = torch.empty(2, dtype=torch.int64) - metadata.kv_lens_expanded_host = torch.empty(6, dtype=torch.int32) - metadata.kv_lens_expanded_cuda = torch.empty(6, dtype=torch.int32) + dsa.gen_token_repeats_cuda = torch.empty(2, dtype=torch.int64) + dsa.host_gen_token_repeats = torch.empty(2, dtype=torch.int64) + dsa.kv_lens_expanded_host = torch.empty(6, dtype=torch.int32) + dsa.kv_lens_expanded_cuda = torch.empty(6, dtype=torch.int32) for name in ("row_kv_lens", "row_kv_correction", "row_req_idx"): dtype = torch.long if name == "row_req_idx" else torch.int32 for suffix in ("host", "cuda"): - setattr(metadata, f"{name}_{suffix}", torch.full((6,), -777, dtype=dtype)) + setattr(dsa, f"{name}_{suffix}", torch.full((6,), -777, dtype=dtype)) for name in ("attn_row_kv_lens", "attn_row_kv_correction", "attn_row_req_idx"): dtype = torch.long if name == "attn_row_req_idx" else torch.int32 for suffix in ("host", "cuda"): - setattr(metadata, f"{name}_{suffix}", torch.full((7,), -777, dtype=dtype)) + setattr(dsa, f"{name}_{suffix}", torch.full((7,), -777, dtype=dtype)) metadata.prompt_lens_cpu = torch.tensor(([3] if num_contexts else []) + [100, 200]) - metadata.attn_row_prompt_lens_cpu = torch.empty(7, dtype=torch.int32) - metadata.attn_row_prompt_lens_cuda = torch.empty(7, dtype=torch.int32) - metadata.attn_row_request_types_host = torch.empty(7, dtype=torch.int32) - metadata.attn_row_block_offsets = torch.empty(1, 7, 2, 1, dtype=torch.int32) + dsa.attn_row_prompt_lens_cpu = torch.empty(7, dtype=torch.int32) + dsa.attn_row_prompt_lens_cuda = torch.empty(7, dtype=torch.int32) + dsa.attn_row_request_types_host = torch.empty(7, dtype=torch.int32) + dsa.attn_row_block_offsets = torch.empty(1, 7, 2, 1, dtype=torch.int32) metadata.kv_cache_block_offsets = None monkeypatch.setattr( "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_sm_version", lambda: 100, ) kv_lens = torch.tensor(([3] if num_contexts else []) + [102, 204], dtype=torch.int32) - metadata.prepare_for_spec_decode(kv_lens) + dsa.prepare_for_spec_decode(kv_lens) expected = torch.tensor([101, 102, 201, 202, 203, 204], dtype=torch.int32) - torch.testing.assert_close(metadata.ragged_row_kv_lens(6), expected) + torch.testing.assert_close(dsa.ragged_row_kv_lens(6), expected) view = metadata.token_major_gen_view() torch.testing.assert_close(view.sequence_length[num_contexts:], expected) if num_contexts: assert view.sequence_length[0] == 3 - row_ptr = metadata.row_kv_lens_cuda.data_ptr() - attention_ptr = metadata.attn_row_kv_lens_cuda.data_ptr() + row_ptr = dsa.row_kv_lens_cuda.data_ptr() + attention_ptr = dsa.attn_row_kv_lens_cuda.data_ptr() metadata.kv_lens_cuda = kv_lens.clone() metadata.kv_lens_cuda[num_contexts:] -= torch.tensor([3, 6]) - metadata.refresh_ragged_row_kv_lens() - metadata.refresh_token_major_gen_rows() + dsa.refresh_ragged_row_kv_lens() + dsa.refresh_token_major_gen_rows() expected -= torch.tensor([3, 3, 6, 6, 6, 6]) - torch.testing.assert_close(metadata.ragged_row_kv_lens(6), expected) + torch.testing.assert_close(dsa.ragged_row_kv_lens(6), expected) torch.testing.assert_close(view.sequence_length[num_contexts:], expected) - assert metadata.row_kv_lens_cuda.data_ptr() == row_ptr - assert metadata.attn_row_kv_lens_cuda.data_ptr() == attention_ptr + assert dsa.row_kv_lens_cuda.data_ptr() == row_ptr + assert dsa.attn_row_kv_lens_cuda.data_ptr() == attention_ptr assert metadata.num_seqs == num_contexts + 2 @@ -148,8 +155,9 @@ class _StopBeforeDispatch(Exception): def _prefix_metadata(monkeypatch, num_contexts=0, ragged=True): seq_lens = ([3] if num_contexts else []) + [1, 2, 4] metadata = _metadata(seq_lens, num_contexts) - metadata.enable_ragged_verification = ragged - metadata.ragged_verify_lens = [1, 2, 4] if ragged else None + metadata.ext(DSAMetadata).enable_ragged_verification = ragged + dsa = metadata.ext(DSAMetadata) + dsa.ragged_verify_lens = [1, 2, 4] if ragged else None metadata.max_num_sequences = len(seq_lens) metadata.cuda_graph_buffers = {} metadata.draft_kv_cache_manager = None @@ -171,7 +179,7 @@ def allocate(buffers, shape, *, cache_name, dtype, capture_graph): "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.is_dsa_cache_manager", lambda value: False, ) - metadata.create_expanded_buffers() + dsa.create_expanded_buffers() metadata.kv_cache_manager = None metadata.mla_cu_q_rows = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) metadata.mla_cu_kv_seqlens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) @@ -180,7 +188,6 @@ def allocate(buffers, shape, *, cache_name, dtype, capture_graph): metadata._mla_scheduler_buffers_valid = False metadata.prompt_lens_cpu = metadata.seq_lens.clone() metadata.kv_cache_block_offsets = None - metadata.is_spec_decoding_enabled = False metadata.sparse_metadata_params = SimpleNamespace(use_cute_dsl_paged_mqa_logits=False) monkeypatch.setattr( "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_sm_version", @@ -191,7 +198,7 @@ def allocate(buffers, shape, *, cache_name, dtype, capture_graph): # Reserved extra tokens must not leak into the kernel's causal KV scan. metadata.kv_lens = kv_lens + 9 if ragged: - metadata.prepare_for_spec_decode(kv_lens) + dsa.prepare_for_spec_decode(kv_lens) return metadata, allocation_names @@ -208,7 +215,7 @@ def test_feature_off_keeps_request_major_prefixes_and_allocations(monkeypatch): def test_uniform_step_with_opt_in_keeps_request_major_prefixes(monkeypatch): metadata, _ = _prefix_metadata(monkeypatch) - metadata.ragged_verify_lens = None + metadata.ext(DSAMetadata).ragged_verify_lens = None metadata._seq_lens.fill_(2) metadata._seq_lens_cuda.copy_(metadata._seq_lens) q_prefix, kv_prefix = metadata.mla_prepare_scheduler_buffers(2) @@ -278,25 +285,25 @@ def launch(): # Exercise the actual overlap hook: refresh device rows and invalidate the # cached prefix before the next module-level dispatch, without host refresh. metadata.kv_cache_manager = None - metadata.enable_flash_mla = False - metadata.use_fp8_ds_mla = False - metadata.expand_for_dsl = False - metadata.kv_lens_cuda_2d = torch.empty(metadata.num_generations, 6, dtype=torch.int32) - metadata.scheduler_metadata_buffer = torch.empty(1, dtype=torch.int32) - metadata.scheduler_metadata_buffer_expanded = torch.empty(1, dtype=torch.int32) - metadata.req_idx_per_token = torch.empty(metadata.num_tokens, dtype=torch.int32) - metadata.gen_kv_indptr = torch.zeros(4, dtype=torch.int64) - metadata.gen_cached_token_indptr = torch.zeros(4, dtype=torch.int64) - metadata._invalidate_pool_view_cache = Mock() - metadata._compute_kv_lens_row_reorder = Mock() - metadata.prepare_dense_topk_indices = Mock() + metadata.ext(DSAMetadata).use_fp8_ds_mla = False + dsa = metadata.ext(DSAMetadata) + dsa.expand_for_dsl = False + dsa.kv_lens_cuda_2d = torch.empty(metadata.num_generations, 6, dtype=torch.int32) + dsa.scheduler_metadata_buffer = torch.empty(1, dtype=torch.int32) + dsa.scheduler_metadata_buffer_expanded = torch.empty(1, dtype=torch.int32) + dsa.req_idx_per_token = torch.empty(metadata.num_tokens, dtype=torch.int32) + dsa.gen_kv_indptr = torch.zeros(4, dtype=torch.int64) + dsa.gen_cached_token_indptr = torch.zeros(4, dtype=torch.int64) + dsa._invalidate_pool_view_cache = Mock() + dsa._compute_kv_lens_row_reorder = Mock() + dsa.prepare_dense_topk_indices = Mock() monkeypatch.setattr( "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_paged_mqa_logits_metadata", lambda *unused: torch.zeros(1, dtype=torch.int32), ) metadata.kv_lens_cuda[num_contexts:] -= torch.tensor([2, 3, 6]) metadata.on_update_kv_lens() - assert metadata._mla_scheduler_buffers_valid is False + assert dsa._ragged_mla_buffers_valid is False launch() expected_kv -= torch.tensor([2, 3, 3, 6, 6, 6, 6]) torch.testing.assert_close(observed[-1][0], expected_q) @@ -308,8 +315,9 @@ def launch(): @pytest.mark.parametrize("metadata_valid", [False, True]) def test_flash_mla_rejects_token_major_generation_before_dispatch(metadata_valid: bool) -> None: metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) - metadata.enable_flash_mla = True - metadata._flash_mla_metadata_valid = metadata_valid + flash_mla = FlashMlaMetadata(metadata) + flash_mla.metadata_valid = metadata_valid + metadata.replace_ext(flash_mla) metadata.token_major_gen_view = Mock(return_value=object()) attention = SimpleNamespace( sparse_params=None, @@ -335,7 +343,8 @@ def test_uniform_flash_mla_and_ragged_native_pass_the_guard( flash_mla: bool, token_major: bool ) -> None: metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) - metadata.enable_flash_mla = flash_mla + if flash_mla: + metadata.replace_ext(FlashMlaMetadata(metadata)) metadata.token_major_gen_view = Mock(return_value=object() if token_major else None) metadata.use_paged_context_fmha = False attention = SimpleNamespace( diff --git a/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py b/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py index 73121f364609..745a15609516 100644 --- a/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py +++ b/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py @@ -7,18 +7,22 @@ import pytest +from tensorrt_llm._torch.attention.backends.features import HelixMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention @pytest.mark.parametrize("helix_spec_tokens_valid", [False, True]) def test_mla_rope_generation_forwards_kernel_controls(helix_spec_tokens_valid: bool) -> None: + helix = SimpleNamespace( + position_offsets=object(), + is_inactive_rank=object(), + spec_tokens_valid=helix_spec_tokens_valid, + local_slots=object(), + ) metadata = SimpleNamespace( kv_cache_manager=object(), max_seq_len=1024, - helix_position_offsets=object(), - helix_is_inactive_rank=object(), - _helix_spec_tokens_valid=helix_spec_tokens_valid, - helix_local_slots=object(), + find_ext=lambda cls: helix if cls is HelixMetadata else None, ) dispatch = Mock() attention = SimpleNamespace( @@ -53,8 +57,8 @@ def test_mla_rope_generation_forwards_kernel_controls(helix_spec_tokens_valid: b dispatch.assert_called_once() assert dispatch.call_args.args[0] is metadata assert dispatch.call_args.args[1:10] == tuple(tensors) - helix_tensor_params = [metadata.helix_position_offsets, metadata.helix_is_inactive_rank] + helix_tensor_params = [helix.position_offsets, helix.is_inactive_rank] if helix_spec_tokens_valid: - helix_tensor_params.append(metadata.helix_local_slots) + helix_tensor_params.append(helix.local_slots) assert dispatch.call_args.args[11] == helix_tensor_params assert dispatch.call_args.args[-len(controls) :] == tuple(controls.values()) diff --git a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py index 74b64b666742..593bf0305782 100644 --- a/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py +++ b/tests/unittest/_torch/attention/kernels/parallel_hw_agnostic/test_helix_postprocess.py @@ -23,6 +23,7 @@ import tensorrt_llm from tensorrt_llm._torch.attention.attention import _helix_sanitize_empty_kv, _helix_zero_kv_mask +from tensorrt_llm._torch.attention.backends.features import HelixMetadata class _FakeAttnMetadata: @@ -522,11 +523,11 @@ def test_mla_skips_generic_mask_when_per_token_bounds_are_valid(self): num_heads = 2 kv_lora_rank = 4 helix_kv_bounds = torch.tensor([0, 5, 0], dtype=torch.int32) + helix = SimpleNamespace(kv_bounds=helix_kv_bounds, spec_tokens_valid=True) attn_metadata = SimpleNamespace( - helix_kv_bounds=helix_kv_bounds, - _helix_spec_tokens_valid=True, num_contexts=0, num_generations=1, + find_ext=lambda cls: helix if cls is HelixMetadata else None, ) attn_backend = Mock() attn_backend.forward.return_value = torch.empty(num_tokens, num_heads * kv_lora_rank) diff --git a/tests/unittest/_torch/attention/multi_gpu/helix_test_utils.py b/tests/unittest/_torch/attention/multi_gpu/helix_test_utils.py index d2d81f93b4d0..4eeca4219b0d 100644 --- a/tests/unittest/_torch/attention/multi_gpu/helix_test_utils.py +++ b/tests/unittest/_torch/attention/multi_gpu/helix_test_utils.py @@ -20,6 +20,7 @@ from mpi4py.futures import MPIPoolExecutor import tensorrt_llm +from tensorrt_llm._torch.attention.backends.features import HelixMetadata from tensorrt_llm._torch.attention.backends.interface import KVCacheParams from tensorrt_llm._torch.attention.backends.utils import get_attention_backend from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 @@ -81,17 +82,12 @@ def split_inputs_for_rank(input_ctx, position_ids_ctx, scenario, rank, world_siz def activate_all_ranks_for_context(attn_metadata, position_ids): """Override helix metadata so all ranks participate in the context step.""" - attn_metadata.helix_position_offsets = position_ids - if ( - hasattr(attn_metadata, "helix_is_inactive_rank") - and attn_metadata.helix_is_inactive_rank is not None - ): - attn_metadata.helix_is_inactive_rank.fill_(False) - if ( - hasattr(attn_metadata, "helix_is_inactive_rank_cpu") - and attn_metadata.helix_is_inactive_rank_cpu is not None - ): - attn_metadata.helix_is_inactive_rank_cpu.fill_(False) + helix = attn_metadata.add_ext(HelixMetadata) + helix.position_offsets = position_ids + if helix.is_inactive_rank is not None: + helix.is_inactive_rank.fill_(False) + if helix.is_inactive_rank_cpu is not None: + helix.is_inactive_rank_cpu.fill_(False) def create_helix_gen_metadata( @@ -125,19 +121,11 @@ def create_helix_gen_metadata( ), **kwargs, ) - attn_metadata.enable_helix = True - attn_metadata.helix_is_inactive_rank = torch.tensor( - helix_is_inactive_rank, dtype=torch.bool, device="cuda" - ) - attn_metadata.helix_is_inactive_rank_cpu = attn_metadata.helix_is_inactive_rank.to( - device="cpu" - ).pin_memory() - attn_metadata.helix_position_offsets = torch.tensor( - position_ids_gen, dtype=torch.int, device="cuda" - ) - attn_metadata.helix_position_offsets_cpu = attn_metadata.helix_position_offsets.to( - device="cpu" - ).pin_memory() + helix = attn_metadata.add_ext(HelixMetadata) + helix.is_inactive_rank = torch.tensor(helix_is_inactive_rank, dtype=torch.bool, device="cuda") + helix.is_inactive_rank_cpu = helix.is_inactive_rank.to(device="cpu").pin_memory() + helix.position_offsets = torch.tensor(position_ids_gen, dtype=torch.int, device="cuda") + helix.position_offsets_cpu = helix.position_offsets.to(device="cpu").pin_memory() attn_metadata.prepare() return attn_metadata diff --git a/tests/unittest/_torch/attention/sparse/csa2/_utils.py b/tests/unittest/_torch/attention/sparse/csa2/_utils.py index f105e24ee496..ae7b44d1783e 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/_utils.py +++ b/tests/unittest/_torch/attention/sparse/csa2/_utils.py @@ -7,8 +7,9 @@ import torch +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.indexer import CSA2Indexer -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import ( CSA2ForwardState, CSA2Layout, @@ -24,6 +25,15 @@ store_index_rows, store_rows, ) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + + +def _bare_csa2_metadata() -> TrtllmAttentionMetadata: + """Metadata built without ``__init__``, carrying a bare CSA2 extension.""" + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(CSA2Metadata(metadata)) + metadata.replace_ext(SpecDecodingMetadata(metadata)) + return metadata def _index_pages(rows: int, device="cpu") -> torch.Tensor: @@ -80,13 +90,13 @@ def _attention_reference(q, kv, valid, sink): def _page_metadata(table, requests, page_size, max_positions): # Exercise device-independent metadata transforms without allocating native # compute buffers. Constructor/lifecycle coverage uses real V2 managers. - meta = object.__new__(CSA2TrtllmMetadata) - meta.csa2_global_page_tables = {0: table} - meta.csa2_token_requests = requests - meta.csa2_global_page_sizes = {0: page_size} - meta.csa2_global_max_positions = {0: max_positions} - meta.csa2_kv_sources = {0: 0} - meta.reset_routing() + meta = _bare_csa2_metadata() + meta.ext(CSA2Metadata).csa2_global_page_tables = {0: table} + meta.ext(CSA2Metadata).csa2_token_requests = requests + meta.ext(CSA2Metadata).csa2_global_page_sizes = {0: page_size} + meta.ext(CSA2Metadata).csa2_global_max_positions = {0: max_positions} + meta.ext(CSA2Metadata).csa2_kv_sources = {0: 0} + meta.ext(CSA2Metadata).reset_routing() return meta @@ -104,21 +114,29 @@ def _model_metadata( ): meta = _page_metadata(table, requests, page_size, max_positions) meta.kv_cache_manager = manager - meta.csa2_swa_indices = {i: swa_reads for i in range(len(layout.compress_ratios))} - meta.csa2_swa_write_slots = {i: swa_writes for i in range(len(layout.compress_ratios))} - meta.csa2_visible_lengths = {i: visible for i in range(len(layout.compress_ratios))} - meta.csa2_kv_sources = { + meta.ext(CSA2Metadata).csa2_swa_indices = { + i: swa_reads for i in range(len(layout.compress_ratios)) + } + meta.ext(CSA2Metadata).csa2_swa_write_slots = { + i: swa_writes for i in range(len(layout.compress_ratios)) + } + meta.ext(CSA2Metadata).csa2_visible_lengths = { + i: visible for i in range(len(layout.compress_ratios)) + } + meta.ext(CSA2Metadata).csa2_kv_sources = { i: layout.layer(i).kv_source for i in range(len(layout.compress_ratios)) } - meta.csa2_main_write_slots = {0: main_writes} + meta.ext(CSA2Metadata).csa2_main_write_slots = {0: main_writes} # These prediction fixtures allow arbitrary visibility per query. Treat # each row as a context chunk; real packed request phases are tested with # the manager-backed runtime metadata. - meta.csa2_request_query_ranges = tuple((i, i + 1) for i in range(len(requests))) - meta.csa2_request_start_positions = tuple( + meta.ext(CSA2Metadata).csa2_request_query_ranges = tuple( + (i, i + 1) for i in range(len(requests)) + ) + meta.ext(CSA2Metadata).csa2_request_start_positions = tuple( max_positions * layout.compress_ratios[0] - 1 for _ in requests ) - meta.csa2_num_context_requests = len(requests) + meta.ext(CSA2Metadata).csa2_num_context_requests = len(requests) meta.mapping = None meta.is_cuda_graph = False return meta @@ -128,23 +146,26 @@ def _run_indexer(layout, layer_idx, state): meta = state.metadata manager = meta.kv_cache_manager layer = layout.layer(layer_idx) - meta.enter_layer(layer) + meta.ext(CSA2Metadata).enter_layer(layer) manager.write_swa( - layer_idx, meta.csa2_swa_write_slots[layer_idx], state.swa_kv, state.swa_transform + layer_idx, + meta.ext(CSA2Metadata).csa2_swa_write_slots[layer_idx], + state.swa_kv, + state.swa_transform, ) if layer.mode == CSA2Mode.FULL: manager.write_global( layer.kv_source, - meta.csa2_main_write_slots[layer.kv_source], + meta.ext(CSA2Metadata).csa2_main_write_slots[layer.kv_source], state.main_kv, state.index_k, state.main_transform, state.index_transform, ) if layer.mode == CSA2Mode.REUSE: - if layer.index_source not in meta.csa2_indices: + if layer.index_source not in meta.ext(CSA2Metadata).csa2_indices: raise ValueError("CSA2 index source did not run") - return meta.csa2_indices[layer.index_source] + return meta.ext(CSA2Metadata).csa2_indices[layer.index_source] if layer_idx not in manager.indexers: manager.indexers[layer_idx] = CSA2Indexer( layout, layer_idx, state.index_q.shape[1], state.index_q.shape[-1] @@ -161,10 +182,10 @@ def _prediction_reference(layout, layer, q, swa, sink, meta, **kwargs): state = CSA2ForwardState(metadata=meta, swa_kv=swa, **kwargs) logical = _run_indexer(layout, layer, state) manager = meta.kv_cache_manager - swa_slots = meta.csa2_swa_indices[layer] + swa_slots = meta.ext(CSA2Metadata).csa2_swa_indices[layer] kv = gather_rows(manager.get_swa_buffer(layer), swa_slots, q.shape[-1], "swa") valid = swa_slots >= 0 - slots = meta.global_slot_tile(layer, 0, q.shape[0], logical) + slots = meta.ext(CSA2Metadata).global_slot_tile(layer, 0, q.shape[0], logical) kv = torch.cat( ( kv, @@ -265,7 +286,8 @@ def _run_modes(device, graph_replay=False): def run(): meta = copy(base) - meta.reset_routing() + meta.replace_ext(base.ext(CSA2Metadata).clone_for(meta)) + meta.ext(CSA2Metadata).reset_routing() outputs = [] for i in range(3): kwargs = dict(index_q=iq if i == 0 else -iq, index_weights=weights) if i != 1 else {} @@ -288,20 +310,22 @@ def run(): (1.0, 1, False), ): iq.fill_(sign) - base.csa2_visible_lengths[0].fill_(visible) - page_slots = base.csa2_global_page_tables[0] + base.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(visible) + page_slots = base.ext(CSA2Metadata).csa2_global_page_tables[0] page_slots.copy_(torch.tensor([[1, 0] if reverse else [0, 1]], device=device)) graph.replay() actual = [o.clone() for o in outputs] - actual_indices = routing.csa2_indices[0].clone() + actual_indices = routing.ext(CSA2Metadata).csa2_indices[0].clone() expected, expected_routing = run() for a, e in zip(actual, expected): torch.testing.assert_close(a, e, atol=0, rtol=0) - torch.testing.assert_close(actual_indices, expected_routing.csa2_indices[0], atol=0, rtol=0) + torch.testing.assert_close( + actual_indices, expected_routing.ext(CSA2Metadata).csa2_indices[0], atol=0, rtol=0 + ) if graph_replay: for request_id in (-1, 1, 0): - base.csa2_token_requests.fill_(request_id) - base.csa2_global_page_tables[0][0, 1] = -1 + base.ext(CSA2Metadata).csa2_token_requests.fill_(request_id) + base.ext(CSA2Metadata).csa2_global_page_tables[0][0, 1] = -1 graph.replay() actual = [o.clone() for o in outputs] expected, _ = run() @@ -332,7 +356,7 @@ def __init__(self, num_cached): class _FakeMetadata: """The attributes ``plan_decoder_replay`` reads, plus the SWA floors and slot resolution. - A real ``CSA2TrtllmMetadata`` needs a KV cache manager, pools + A real CSA2 metadata needs a KV cache manager, pools and a device; the planner is pure host arithmetic over these fields, so standing them up here keeps this a CPU test and keeps the refusals -- which are the part worth pinning -- independent of cache-manager plumbing. @@ -341,7 +365,18 @@ class _FakeMetadata: what it clears; see the method. """ - reset_routing = CSA2TrtllmMetadata.reset_routing + reset_routing = CSA2Metadata.reset_routing + decoder_context_ends = CSA2Metadata.decoder_context_ends + + @property + def metadata(self) -> "_FakeMetadata": + # The fake serves as both the core metadata and its CSA2 extension. + return self + + def ext(self, cls: type) -> "_FakeMetadata": + return self + + find_ext = ext def __init__( self, diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_backend.py b/tests/unittest/_torch/attention/sparse/csa2/test_backend.py index 3aa1b80d606e..a802de86f4fc 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_backend.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_backend.py @@ -7,7 +7,12 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2FlashInfer, CSA2FlashMLA +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import ( + CSA2FlashInfer, + CSA2FlashMLA, + CSA2TrtllmAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Params skip_pre_blackwell = pytest.mark.skipif( @@ -172,9 +177,8 @@ def record(attn, *args, **kwargs): extra_valid[0, 1::2] = False # holes before later valid entries extra_valid[1] = False sink = torch.randn(heads, device="cuda") - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - metadata = CSA2TrtllmMetadata.for_query_tile(q, extra_width) + metadata = CSA2Metadata.for_query_tile(q, extra_width) args = _inputs(q, swa, extra, swa_valid, extra_valid, sink) out = attn.forward(q.flatten(1), None, None, metadata, forward_args=args).view_as(q) torch.cuda.synchronize() @@ -196,10 +200,9 @@ def test_trtllm_graph_replay_resets_sparse_state(): extra_valid = torch.ones(2, 17, device="cuda", dtype=torch.bool) sink = torch.zeros(64, device="cuda") - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import pack_rows - frame = CSA2TrtllmMetadata.for_query_tile(q, 17) + frame = CSA2Metadata.for_query_tile(q, 17) def run(): frame.is_cuda_graph = torch.cuda.is_current_stream_capturing() @@ -208,7 +211,7 @@ def run(): for _ in range(3): run() - pointers = frame.pool_pointers.clone() + pointers = frame.ext(CSA2Metadata).pool_pointers.clone() workspace_ptr = frame.workspace.data_ptr() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): @@ -216,7 +219,7 @@ def run(): # Larger eager queries use independent caller-owned metadata. No recursive # warmup or backend-private frame allocation is needed. large_q = q.repeat(8, 1, 1) - large_frame = CSA2TrtllmMetadata.for_query_tile(large_q, 17) + large_frame = CSA2Metadata.for_query_tile(large_q, 17) large_args = _inputs( large_q, swa.repeat(8, 1, 1), @@ -244,7 +247,7 @@ def run(): atol=0.03, rtol=0.03, ) - torch.testing.assert_close(frame.pool_pointers, pointers, atol=0, rtol=0) + torch.testing.assert_close(frame.ext(CSA2Metadata).pool_pointers, pointers, atol=0, rtol=0) assert frame.workspace.data_ptr() == workspace_ptr assert large_frame.workspace.data_ptr() != workspace_ptr @@ -264,7 +267,6 @@ def test_standard_backend_multiple_tiles_and_reuse(monkeypatch): CSA2CacheRole, ) from tensorrt_llm._torch.attention.backends.sparse.csa2.indexer import CSA2Indexer - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import ( CSA2BackendForwardArgs, CSA2ForwardState, @@ -302,7 +304,7 @@ def predict(indexer, metadata, hidden_states, *args, **kwargs): request = manager._create_kv_cache(100, None, []) assert manager._resume_and_restore(100, request) assert request.resize(count) - runtime = CSA2TrtllmMetadata( + runtime = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=count, kv_cache_manager=manager ) positions = torch.arange(count, device="cuda") @@ -313,30 +315,32 @@ def predict(indexer, metadata, hidden_states, *args, **kwargs): index_q = torch.randn(count, 2, 128, device="cuda", dtype=torch.bfloat16) weights = torch.ones(count, 2, device="cuda", dtype=torch.bfloat16) sink = torch.randn(heads, device="cuda") - runtime.reset_routing() + runtime.ext(CSA2Metadata).reset_routing() global_base = manager.get_cache_indices(100, 0, CSA2CacheRole.GLOBAL)[0] * 128 global_slots = global_base + torch.arange(6, device="cuda") - runtime.csa2_token_requests = torch.zeros(count, dtype=torch.int64, device="cuda") - runtime.csa2_request_query_ranges = ((0, count),) - runtime.csa2_request_start_positions = (0,) - runtime.csa2_num_context_requests = 1 - runtime.csa2_global_page_tables = { + runtime.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + count, dtype=torch.int64, device="cuda" + ) + runtime.ext(CSA2Metadata).csa2_request_query_ranges = ((0, count),) + runtime.ext(CSA2Metadata).csa2_request_start_positions = (0,) + runtime.ext(CSA2Metadata).csa2_num_context_requests = 1 + runtime.ext(CSA2Metadata).csa2_global_page_tables = { 0: torch.tensor([[global_base // 128]], dtype=torch.int32, device="cuda") } - runtime.csa2_global_page_sizes = {0: 128} - runtime.csa2_global_max_positions = {0: 6} - runtime.csa2_main_write_slots = {0: global_slots} - runtime.csa2_kv_sources = {0: 0, 1: 0, 2: 0} - runtime.csa2_swa_indices = {} - runtime.csa2_swa_write_slots = {} - runtime.csa2_visible_lengths = {} + runtime.ext(CSA2Metadata).csa2_global_page_sizes = {0: 128} + runtime.ext(CSA2Metadata).csa2_global_max_positions = {0: 6} + runtime.ext(CSA2Metadata).csa2_main_write_slots = {0: global_slots} + runtime.ext(CSA2Metadata).csa2_kv_sources = {0: 0, 1: 0, 2: 0} + runtime.ext(CSA2Metadata).csa2_swa_indices = {} + runtime.ext(CSA2Metadata).csa2_swa_write_slots = {} + runtime.ext(CSA2Metadata).csa2_visible_lengths = {} for layer in range(3): swa_base = manager.get_cache_indices(100, layer, CSA2CacheRole.SWA)[0] * 128 local = positions[:, None] - torch.arange(4, device="cuda")[None, :] swa_slots = torch.where(local >= 0, local + swa_base, -1) - runtime.csa2_swa_indices[layer] = swa_slots - runtime.csa2_swa_write_slots[layer] = positions + swa_base - runtime.csa2_visible_lengths[layer] = positions.remainder(6) + 1 + runtime.ext(CSA2Metadata).csa2_swa_indices[layer] = swa_slots + runtime.ext(CSA2Metadata).csa2_swa_write_slots[layer] = positions + swa_base + runtime.ext(CSA2Metadata).csa2_visible_lengths[layer] = positions.remainder(6) + 1 args_dict = {} if layer != 1: args_dict.update(index_q=index_q * (-1 if layer else 1), index_weights=weights) @@ -347,7 +351,7 @@ def predict(indexer, metadata, hidden_states, *args, **kwargs): outputs = [] for start in range(0, count, 16): tile = q[start : start + 16] - frame = runtime.get_query_tile_metadata(tile, layout.index_topk) + frame = runtime.ext(CSA2Metadata).get_query_tile_metadata(tile, layout.index_topk) args = AttentionForwardArgs( attention_input_type=AttentionInputType.generation_only, attention_sinks=sink, @@ -359,13 +363,13 @@ def predict(indexer, metadata, hidden_states, *args, **kwargs): ) ) actual = torch.cat(outputs) - logical = runtime.csa2_indices[0 if layer == 1 else layer] - slots = runtime.global_slot_tile(layer, 0, count, logical) + logical = runtime.ext(CSA2Metadata).csa2_indices[0 if layer == 1 else layer] + slots = runtime.ext(CSA2Metadata).global_slot_tile(layer, 0, count, logical) swa_values = gather_rows(manager.get_swa_buffer(layer), swa_slots, 512, "swa") main_values = gather_rows(manager.get_main_buffer(0), slots, 512, "main") expected = _reference(q, swa_values, main_values, swa_slots >= 0, slots >= 0, sink) torch.testing.assert_close(actual, expected, atol=0.03, rtol=0.03) - frame = runtime.get_query_tile_metadata(q[-3:], 4) + frame = runtime.ext(CSA2Metadata).get_query_tile_metadata(q[-3:], 4) assert frame.host_total_kv_lens.tolist() == [0, 3 * 256] assert predictions == [(0, count), (2, count)] finally: @@ -378,15 +382,13 @@ def predict(indexer, metadata, hidden_states, *args, **kwargs): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_shared_metadata_resets_global_to_swa_inputs(): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - q = torch.randn(2, 8, 512, device="cuda", dtype=torch.bfloat16) swa = torch.randn(2, 4, 512, device="cuda", dtype=torch.bfloat16) main = torch.randn(2, 17, 512, device="cuda", dtype=torch.bfloat16) swa_valid = torch.ones(2, 4, device="cuda", dtype=torch.bool) main_valid = torch.ones(2, 17, device="cuda", dtype=torch.bool) sink = torch.zeros(8, device="cuda") - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17) + metadata = CSA2Metadata.for_query_tile(q, 17) first, second = _backend(8, 20), _backend(8, 21) args = _inputs(q, swa, main, swa_valid, main_valid, sink) first.forward(q.flatten(1), None, None, metadata, forward_args=args) @@ -397,7 +399,10 @@ def test_shared_metadata_resets_global_to_swa_inputs(): args.output = None output = second.forward(q.flatten(1), None, None, metadata, forward_args=args).view_as(q) assert args.sparse_runtime_params.aux_kv_cache_pool_ptr is None - torch.testing.assert_close(metadata.prepared_lens, torch.full_like(metadata.prepared_lens, 4)) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_lens, + torch.full_like(metadata.ext(CSA2Metadata).prepared_lens, 4), + ) torch.testing.assert_close( output, _decoded_reference(q, args, swa_valid, None, sink), atol=0.03, rtol=0.03 ) @@ -407,14 +412,12 @@ def test_shared_metadata_resets_global_to_swa_inputs(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_cold_metadata_rejects_capture(): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - attn = _backend(8) q = torch.zeros(2, 8, 512, device="cuda", dtype=torch.bfloat16) swa = torch.zeros(2, 1, 512, device="cuda", dtype=torch.bfloat16) valid = torch.ones(2, 1, device="cuda", dtype=torch.bool) sink = torch.zeros(8, device="cuda") - metadata = CSA2TrtllmMetadata.for_query_tile(q, 0) + metadata = CSA2Metadata.for_query_tile(q, 0) args = _inputs(q, swa, None, valid, None, sink) metadata.is_cuda_graph = True graph = torch.cuda.CUDAGraph() @@ -430,7 +433,6 @@ def test_cold_metadata_rejects_capture(): @torch.inference_mode() def test_alternative_helper_dispatch(implementation, context, monkeypatch): from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata def unexpected_fallback(*args, **kwargs): raise AssertionError("Selected CSA2 implementation fell through to native fallback") @@ -446,12 +448,12 @@ def unexpected_fallback(*args, **kwargs): swa_valid[1] = False main_valid[1] = False sink = torch.linspace(-2, 2, 8, device="cuda") - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17) + metadata = CSA2Metadata.for_query_tile(q, 17) args = _inputs(q, swa, main, swa_valid, main_valid, sink) if context: from tensorrt_llm._torch.attention.backends.interface import AttentionInputType - metadata._bind_context_tile([q.shape[0]]) + metadata.ext(CSA2Metadata)._bind_context_tile([q.shape[0]]) args.attention_input_type = AttentionInputType.context_only args.output = torch.empty_like(q).flatten(1) actual = _helper_forward(attn, q.flatten(1), metadata, args).view_as(q) @@ -469,7 +471,6 @@ def unexpected_fallback(*args, **kwargs): @torch.inference_mode() def test_alternative_backend_rejects_unsupported_options(option): from tensorrt_llm._torch.attention.backends.interface import CustomAttentionMask - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata attn = _backend(8, compute_backend="flashinfer") q = torch.zeros(1, 8, 512, device="cuda", dtype=torch.bfloat16) @@ -487,7 +488,7 @@ def test_alternative_backend_rejects_unsupported_options(option): else: args.output = torch.empty_like(q).flatten(1) args.output_sf = torch.empty(1, device="cuda", dtype=torch.uint8) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 0) + metadata = CSA2Metadata.for_query_tile(q, 0) with pytest.raises(ValueError, match="(do not support.*mask/output format|output must match)"): _helper_forward(attn, q.flatten(1), metadata, args) @@ -496,14 +497,12 @@ def test_alternative_backend_rejects_unsupported_options(option): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_alternative_helper_revalidates_output_scale(): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - attn = _backend(8, compute_backend="flashinfer") q = torch.zeros(1, 8, 512, device="cuda", dtype=torch.bfloat16) swa = torch.zeros(1, 1, 512, device="cuda", dtype=torch.bfloat16) valid = torch.ones(1, 1, device="cuda", dtype=torch.bool) args = _inputs(q, swa, None, valid, None, torch.zeros(8, device="cuda")) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 0) + metadata = CSA2Metadata.for_query_tile(q, 0) _helper_forward(attn, q.flatten(1), metadata, args) torch.cuda.synchronize() args.out_scale = torch.ones(1, device="cuda") @@ -520,7 +519,6 @@ def test_alternative_helper_revalidates_output_scale(): def test_native_context_preserves_real_query_groups(extra_width, prefix, tile_start, monkeypatch): from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha from tensorrt_llm._torch.attention.backends.interface import AttentionInputType - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata torch.manual_seed(459) heads, count = 64, 5 @@ -542,14 +540,21 @@ def test_native_context_preserves_real_query_groups(extra_width, prefix, tile_st extra_valid[1, ::3] = False extra_valid[2] = False sink = torch.randn(heads, device="cuda") - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=count + tile_start) + source = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=count + tile_start + ) source._num_contexts = 2 source._num_ctx_tokens = source._num_tokens = count + tile_start - source.csa2_num_context_requests = 2 - source.csa2_request_query_ranges = ((0, 3 + tile_start), (3 + tile_start, count + tile_start)) - source.csa2_request_start_positions = (prefix, prefix * 2) - source.csa2_request_lengths = (3 + tile_start, 2) - metadata = source.get_query_tile_metadata(q, extra_width, query_start=tile_start) + source.ext(CSA2Metadata).csa2_num_context_requests = 2 + source.ext(CSA2Metadata).csa2_request_query_ranges = ( + (0, 3 + tile_start), + (3 + tile_start, count + tile_start), + ) + source.ext(CSA2Metadata).csa2_request_start_positions = (prefix, prefix * 2) + source.ext(CSA2Metadata).csa2_request_lengths = (3 + tile_start, 2) + metadata = source.ext(CSA2Metadata).get_query_tile_metadata( + q, extra_width, query_start=tile_start + ) assert metadata.num_contexts == 2 assert metadata.num_ctx_tokens == count assert metadata.num_generations == 0 @@ -584,7 +589,7 @@ def record(provider, query, key, value, meta, forward_args): ) monkeypatch.setattr(FallbackFmha, "forward", original) - generation = CSA2TrtllmMetadata.for_query_tile(q, extra_width) + generation = CSA2Metadata.for_query_tile(q, extra_width) generation_args = _inputs(q, swa, extra, swa_valid, extra_valid, sink) reference_generation = attn.forward( q.flatten(1), None, None, generation, forward_args=generation_args @@ -601,7 +606,6 @@ def test_packed_helper_skips_bf16_staging(extra_width, monkeypatch): pytest.skip("SM100 packed attention") from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata torch.manual_seed(460) heads, count = 32, 3 @@ -629,9 +633,11 @@ def test_packed_helper_skips_bf16_staging(extra_width, monkeypatch): records = torch.full((packed.shape[0], 356), 73, dtype=torch.uint8, device="cuda") records[:, :288].copy_(packed) args.sparse_backend_args.main_pool = records[:, :288] - metadata = CSA2TrtllmMetadata.for_query_tile(q, extra_width) + metadata = CSA2Metadata.for_query_tile(q, extra_width) monkeypatch.setattr( - metadata, "stage_selected", lambda *a: pytest.fail("Packed FMHA staged BF16 rows") + metadata.ext(CSA2Metadata), + "stage_selected", + lambda *a: pytest.fail("Packed FMHA staged BF16 rows"), ) output = _helper_forward(attn, q.flatten(1), metadata, args).view_as(q) torch.cuda.synchronize() @@ -688,7 +694,6 @@ def test_native_backend_rejects_disabled_fallback(monkeypatch): @torch.inference_mode() def test_context_graph_uses_fixed_generation_frame(): from tensorrt_llm._torch.attention.backends.interface import AttentionInputType - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata torch.manual_seed(461) heads, count, extra_width = 64, 5, 17 @@ -699,15 +704,15 @@ def test_context_graph_uses_fixed_generation_frame(): swa_valid = torch.ones(count, 128, device="cuda", dtype=torch.bool) extra_valid = torch.ones(count, extra_width, device="cuda", dtype=torch.bool) sink = torch.randn(heads, device="cuda") - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=count) + source = CSA2TrtllmAttention.create_metadata(max_num_requests=2, max_num_tokens=count) source._num_contexts = 2 source._num_ctx_tokens = source._num_tokens = count - source.csa2_num_context_requests = 2 - source.csa2_request_query_ranges = ((0, 3), (3, 5)) - source.csa2_request_start_positions = (1024, 2048) - source.csa2_request_lengths = (3, 2) + source.ext(CSA2Metadata).csa2_num_context_requests = 2 + source.ext(CSA2Metadata).csa2_request_query_ranges = ((0, 3), (3, 5)) + source.ext(CSA2Metadata).csa2_request_start_positions = (1024, 2048) + source.ext(CSA2Metadata).csa2_request_lengths = (3, 2) source.is_cuda_graph = True - metadata = source.get_query_tile_metadata(q, extra_width, query_start=0) + metadata = source.ext(CSA2Metadata).get_query_tile_metadata(q, extra_width, query_start=0) assert metadata.num_contexts == 0 assert metadata.num_generations == count args = _inputs(q, swa, extra, swa_valid, extra_valid, sink) @@ -719,13 +724,13 @@ def test_context_graph_uses_fixed_generation_frame(): for _ in range(3): attn.forward(q.flatten(1), None, None, metadata, forward_args=args) pointers = ( - metadata.swa_pool.data_ptr(), - metadata.extra_pool.data_ptr(), + metadata.ext(CSA2Metadata).swa_pool.data_ptr(), + metadata.ext(CSA2Metadata).extra_pool.data_ptr(), metadata.workspace.data_ptr(), ) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - captured = source.get_query_tile_metadata(q, extra_width, query_start=0) + captured = source.ext(CSA2Metadata).get_query_tile_metadata(q, extra_width, query_start=0) assert captured is metadata output = attn.forward(q.flatten(1), None, None, captured, forward_args=args).view_as(q) swa_ids = torch.arange(count * 128, device="cuda").reshape(count, 128) @@ -745,8 +750,8 @@ def test_context_graph_uses_fixed_generation_frame(): output, _decoded_reference(q, args, swa_valid, extra_valid, sink), atol=0.03, rtol=0.03 ) assert pointers == ( - metadata.swa_pool.data_ptr(), - metadata.extra_pool.data_ptr(), + metadata.ext(CSA2Metadata).swa_pool.data_ptr(), + metadata.ext(CSA2Metadata).extra_pool.data_ptr(), metadata.workspace.data_ptr(), ) @@ -756,13 +761,11 @@ def test_context_graph_uses_fixed_generation_frame(): @pytest.mark.parametrize("implementation,packed", [("flashinfer", False), ("trtllm", True)]) @torch.inference_mode() def test_direct_compute_rejects_native_forward_misuse(implementation, packed): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - attn = _backend(32, compute_backend=implementation, use_packed=packed) q = torch.zeros(1, 32, 512, device="cuda", dtype=torch.bfloat16) values = torch.zeros(1, 1, 512, device="cuda", dtype=torch.bfloat16) valid = torch.ones(1, 1, device="cuda", dtype=torch.bool) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 0) + metadata = CSA2Metadata.for_query_tile(q, 0) args = _inputs(q, values, None, valid, None, torch.zeros(32, device="cuda")) with pytest.raises(ValueError, match="explicit module helper"): attn.forward(q.flatten(1), None, None, metadata, forward_args=args) @@ -905,14 +908,16 @@ def test_global_slot_graph_mapping_folds_visibility(monkeypatch): logical = storage[:, ::2] logical.copy_(torch.arange(512, device="cuda")) visible = torch.full((8,), 512, dtype=torch.int64, device="cuda") - metadata.global_slot_tile(0, 0, 8, logical, visible_lengths=visible) + metadata.ext(CSA2Metadata).global_slot_tile(0, 0, 8, logical, visible_lengths=visible) metadata.is_cuda_graph = True - metadata.global_slot_tile(0, 0, 8, logical) + metadata.ext(CSA2Metadata).global_slot_tile(0, 0, 8, logical) for _ in range(3): - metadata.global_slot_tile(0, 0, 8, logical, visible_lengths=visible) + metadata.ext(CSA2Metadata).global_slot_tile(0, 0, 8, logical, visible_lengths=visible) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - slots = metadata.global_slot_tile(0, 0, 8, logical, visible_lengths=visible) + slots = metadata.ext(CSA2Metadata).global_slot_tile( + 0, 0, 8, logical, visible_lengths=visible + ) for step in range(3): table[0, 1] = -1 if step == 0 else 1 requests.copy_( @@ -948,7 +953,6 @@ def test_global_slot_graph_mapping_folds_visibility(monkeypatch): def test_native_fp8_quantized_attention(phase, extra_width, heads, monkeypatch, record_property): from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha from tensorrt_llm._torch.attention.backends.interface import AttentionInputType - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import unpack_rows torch.manual_seed(482) @@ -979,20 +983,18 @@ def test_native_fp8_quantized_attention(phase, extra_width, heads, monkeypatch, args = _inputs(q, swa, extra, sv, ev, sink) # Staging derives both scales from the packed pools; the oracle below reads # them back off the metadata and the forward args. - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=count) + source = CSA2TrtllmAttention.create_metadata(max_num_requests=2, max_num_tokens=count) if phase == "generation": - metadata = CSA2TrtllmMetadata.for_query_tile( - q, extra_width, staging_dtype=attn.staging_dtype - ) + metadata = CSA2Metadata.for_query_tile(q, extra_width, staging_dtype=attn.staging_dtype) else: source._num_contexts = 2 source._num_ctx_tokens = source._num_tokens = count - source.csa2_num_context_requests = 2 - source.csa2_request_query_ranges = ((0, 3), (3, 5)) - source.csa2_request_start_positions = (1024, 2048) - source.csa2_request_lengths = (3, 2) + source.ext(CSA2Metadata).csa2_num_context_requests = 2 + source.ext(CSA2Metadata).csa2_request_query_ranges = ((0, 3), (3, 5)) + source.ext(CSA2Metadata).csa2_request_start_positions = (1024, 2048) + source.ext(CSA2Metadata).csa2_request_lengths = (3, 2) source.is_cuda_graph = phase == "captured_context" - metadata = source.get_query_tile_metadata( + metadata = source.ext(CSA2Metadata).get_query_tile_metadata( q, extra_width, query_start=0, staging_dtype=attn.staging_dtype ) args.attention_input_type = ( @@ -1008,7 +1010,11 @@ def test_native_fp8_quantized_attention(phase, extra_width, heads, monkeypatch, def observe(provider, query, k, v, meta, forward_args): calls.append(forward_args.attention_input_type) assert query.dtype == torch.bfloat16 - assert meta.swa_pool.dtype == meta.extra_pool.dtype == torch.float8_e4m3fn + assert ( + meta.ext(CSA2Metadata).swa_pool.dtype + == meta.ext(CSA2Metadata).extra_pool.dtype + == torch.float8_e4m3fn + ) if meta.num_contexts: assert forward_args.quant_q_buffer is None assert forward_args.quant_scale_qkv is forward_args.kv_scale_orig_quant @@ -1024,7 +1030,7 @@ def observe(provider, query, k, v, meta, forward_args): def run(): if phase == "captured_context": - captured = source.get_query_tile_metadata( + captured = source.ext(CSA2Metadata).get_query_tile_metadata( q, extra_width, query_start=0, staging_dtype=attn.staging_dtype ) assert captured is metadata @@ -1042,8 +1048,8 @@ def run(): graph = None if phase == "captured_context": pointers = ( - metadata.swa_pool.data_ptr(), - metadata.extra_pool.data_ptr(), + metadata.ext(CSA2Metadata).swa_pool.data_ptr(), + metadata.ext(CSA2Metadata).extra_pool.data_ptr(), metadata.workspace.data_ptr(), ) graph = torch.cuda.CUDAGraph() @@ -1056,7 +1062,7 @@ def run(): sv[0, : 3 * (replay + 1)] = False graph.replay() torch.cuda.synchronize() - kv_dq = metadata.stage_kv_scales[1] + kv_dq = metadata.ext(CSA2Metadata).stage_kv_scales[1] scale_kv = kv_dq.item() # A power of two costs no relative precision: it only moves the exponent. assert scale_kv > 0.0 and math.log2(scale_kv).is_integer() @@ -1068,7 +1074,7 @@ def run(): # The C++ op points dequant_scale_q at dequant_scale_kv, so context # cannot give Q a scale of its own: it gets the staging pair as is. orig_quant, q_dq = seen[-1] - assert (orig_quant, q_dq) == metadata.stage_kv_scales + assert (orig_quant, q_dq) == metadata.ext(CSA2Metadata).stage_kv_scales else: # Q takes the reciprocal, which keeps BMM1's dequant product at one. # The in-tree split-KV reduction no longer requires that, but the same @@ -1086,12 +1092,14 @@ def run(): ) scale = q_dq.item() * scale_kv * 512**-0.5 torch.testing.assert_close( - metadata.mla_bmm1_scale, + metadata.ext(CSA2Metadata).mla_bmm1_scale, torch.tensor([scale, scale * 1.4426950408889634], device="cuda"), atol=0, rtol=1e-6, ) - torch.testing.assert_close(metadata.mla_bmm2_scale, kv_dq, atol=0, rtol=0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).mla_bmm2_scale, kv_dq, atol=0, rtol=0 + ) inp = args.sparse_backend_args decoded_swa = unpack_rows(inp.swa_pool, 512, "swa").reshape(count, 128, 512) qswa, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(decoded_swa, kv_dq) @@ -1119,8 +1127,8 @@ def run(): ) if graph: assert pointers == ( - metadata.swa_pool.data_ptr(), - metadata.extra_pool.data_ptr(), + metadata.ext(CSA2Metadata).swa_pool.data_ptr(), + metadata.ext(CSA2Metadata).extra_pool.data_ptr(), metadata.workspace.data_ptr(), ) @@ -1147,7 +1155,6 @@ def test_native_fp8_staging_keeps_rows_beyond_the_e4m3_range(phase, extra_width) from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import ( _CONTEXT_STAGE_SCALE_MAX_EXPONENT, ) - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import unpack_rows torch.manual_seed(482) @@ -1165,18 +1172,16 @@ def test_native_fp8_staging_keeps_rows_beyond_the_e4m3_range(phase, extra_width) sink = torch.randn(heads, device="cuda") args = _inputs(q, swa, extra, sv, ev, sink) if phase == "generation": - metadata = CSA2TrtllmMetadata.for_query_tile( - q, extra_width, staging_dtype=attn.staging_dtype - ) + metadata = CSA2Metadata.for_query_tile(q, extra_width, staging_dtype=attn.staging_dtype) else: - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=count) + source = CSA2TrtllmAttention.create_metadata(max_num_requests=2, max_num_tokens=count) source._num_contexts = 2 source._num_ctx_tokens = source._num_tokens = count - source.csa2_num_context_requests = 2 - source.csa2_request_query_ranges = ((0, 3), (3, 5)) - source.csa2_request_start_positions = (1024, 2048) - source.csa2_request_lengths = (3, 2) - metadata = source.get_query_tile_metadata( + source.ext(CSA2Metadata).csa2_num_context_requests = 2 + source.ext(CSA2Metadata).csa2_request_query_ranges = ((0, 3), (3, 5)) + source.ext(CSA2Metadata).csa2_request_start_positions = (1024, 2048) + source.ext(CSA2Metadata).csa2_request_lengths = (3, 2) + metadata = source.ext(CSA2Metadata).get_query_tile_metadata( q, extra_width, query_start=0, staging_dtype=attn.staging_dtype ) args.attention_input_type = ( @@ -1185,7 +1190,7 @@ def test_native_fp8_staging_keeps_rows_beyond_the_e4m3_range(phase, extra_width) else AttentionInputType.generation_only ) output = attn.forward(q.flatten(1), None, None, metadata, forward_args=args).view_as(q) - kv_dq = metadata.stage_kv_scales[1] + kv_dq = metadata.ext(CSA2Metadata).stage_kv_scales[1] scale_kv = kv_dq.item() inp = args.sparse_backend_args decoded_swa = unpack_rows(inp.swa_pool, 512, "swa").reshape(count, 128, 512).float() @@ -1309,8 +1314,6 @@ def test_default_staging_follows_the_backend_that_can_serve_it(kv_dtype): @skip_pre_blackwell @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_native_fp8_rejects_incompatible_options(): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - # Packed attention reads persistent rows directly and has no FP8 staging. with pytest.raises(ValueError): _backend(64, kv_cache_dtype="fp8", use_packed=True, fp8_staging=True) @@ -1321,9 +1324,7 @@ def test_native_fp8_rejects_incompatible_options(): valid = torch.ones(1, 128, device="cuda", dtype=torch.bool) args = _inputs(q, swa, None, valid, None, torch.zeros(64, device="cuda")) with pytest.raises(ValueError): - attn.forward( - q.flatten(1), None, None, CSA2TrtllmMetadata.for_query_tile(q, 0), forward_args=args - ) + attn.forward(q.flatten(1), None, None, CSA2Metadata.for_query_tile(q, 0), forward_args=args) @skip_pre_blackwell @@ -1340,10 +1341,7 @@ def test_shared_native_context_matches_selected_rows( AttentionInputType, ) from tensorrt_llm._torch.attention.backends.sparse.csa2 import kernel - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import ( - CSA2SharedKVPlan, - CSA2TrtllmMetadata, - ) + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2SharedKVPlan from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2BackendForwardArgs from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( gather_rows, @@ -1405,7 +1403,7 @@ def observe(*args, **kwargs): monkeypatch.setattr(kernel, "stage_shared_rows", observe) outputs = [] for shared in (False, True): - metadata = CSA2TrtllmMetadata.for_query_tile( + metadata = CSA2Metadata.for_query_tile( q, width, staging_dtype=attn.staging_dtype, @@ -1413,10 +1411,14 @@ def observe(*args, **kwargs): shared_rows=1289 if shared else None, ) if shared: - metadata.shared_plan = plan - assert metadata.swa_pool is metadata.extra_pool is metadata.shared_pool - assert metadata.shared_pool.shape == (1289, 512) - assert metadata.shared_pool.numel() < count * (128 + width) * 512 + metadata.ext(CSA2Metadata).shared_plan = plan + assert ( + metadata.ext(CSA2Metadata).swa_pool + is metadata.ext(CSA2Metadata).extra_pool + is metadata.ext(CSA2Metadata).shared_pool + ) + assert metadata.ext(CSA2Metadata).shared_pool.shape == (1289, 512) + assert metadata.ext(CSA2Metadata).shared_pool.numel() < count * (128 + width) * 512 args = AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, attention_sinks=sink, @@ -1525,6 +1527,7 @@ def phase(n, contexts): selection[0] = torch.arange(n * width, dtype=torch.int32).reshape(n, width) % 7 query = torch.zeros(n, heads * 512, dtype=torch.bfloat16) metadata = SimpleNamespace(num_contexts=contexts, shared_plan=True) + metadata.ext = lambda cls: metadata if contexts: args = AttentionForwardArgs(attention_input_type=AttentionInputType.context_only) return helper.forward_context(query, metadata, args) diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py b/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py index 958c004d745b..ac1c0d64ec40 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_cache_manager.py @@ -7,11 +7,12 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import ( CSA2CacheManager, CSA2CacheRole, ) -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( read_index_rows, @@ -401,7 +402,7 @@ def _request(request_id, tokens): def _metadata(manager, request, metadata=None): - metadata = metadata or CSA2TrtllmMetadata( + metadata = metadata or CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=512, kv_cache_manager=manager ) metadata.request_ids = [request.py_request_id] @@ -494,7 +495,7 @@ def _spec_manager(block): def _spec_metadata(manager): - return CSA2TrtllmMetadata( + return CSA2TrtllmAttention.create_metadata( max_num_requests=manager.max_batch_size, max_num_tokens=manager.max_num_tokens, kv_cache_manager=manager, @@ -559,16 +560,16 @@ def test_first_verification_step_tolerates_the_overlap_tail(block): metadata = _spec_prepare(manager, requests, cached) owner = manager.layout.kv_source_layer_ids[-1] assert manager.layout.compress_ratios[owner] == 1 - metadata._ensure_swa_slots() + metadata.ext(CSA2Metadata)._ensure_swa_slots() for row, request in enumerate(requests): cache = manager.kv_cache_map[request.py_request_id] assert cache.capacity == request.prompt_len + 2 * DRAFT_LEN - tail = metadata.csa2_request_query_ranges[row][1] - 1 + tail = metadata.ext(CSA2Metadata).csa2_request_query_ranges[row][1] - 1 for layer in range(len(manager.layout.compress_ratios)): - writes = metadata.csa2_swa_write_slots[layer] + writes = metadata.ext(CSA2Metadata).csa2_swa_write_slots[layer] assert int(writes[tail]) == -1 assert all(int(v) >= 0 for v in writes[tail - DRAFT_LEN : tail]) - assert int(metadata.csa2_main_write_slots[owner][tail]) == -1 + assert int(metadata.ext(CSA2Metadata).csa2_main_write_slots[owner][tail]) == -1 finally: _spec_cleanup(manager) diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_context_only_metadata.py b/tests/unittest/_torch/attention/sparse/csa2/test_context_only_metadata.py index 02da5bd7eb48..0063bbb2b0a9 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_context_only_metadata.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_context_only_metadata.py @@ -6,11 +6,12 @@ import torch from utils.util import skip_pre_blackwell +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import ( CSA2CacheManager, CSA2CacheRole, ) -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -58,7 +59,7 @@ def test_context_metadata_omits_decoder_swa_and_preserves_boundary_global(length assert manager.prepare_context(request) assert manager.resize_context(request, length) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=8192, kv_cache_manager=manager ) metadata.request_ids = [71] @@ -72,16 +73,16 @@ def test_context_metadata_omits_decoder_swa_and_preserves_boundary_global(length # live metadata, then its device-side address refresh. for direct in (False, True): if direct: - metadata.prepare_csa2() + metadata.ext(CSA2Metadata).prepare_csa2() metadata.on_update_kv_lens() - metadata._ensure_swa_slots() - assert set(metadata.csa2_swa_indices) == set(range(40)) - assert set(metadata._csa2_swa_page_tables) == set(range(20)) - assert set(metadata.csa2_global_page_tables) == {2, 8, 14, 20} + metadata.ext(CSA2Metadata)._ensure_swa_slots() + assert set(metadata.ext(CSA2Metadata).csa2_swa_indices) == set(range(40)) + assert set(metadata.ext(CSA2Metadata)._csa2_swa_page_tables) == set(range(20)) + assert set(metadata.ext(CSA2Metadata).csa2_global_page_tables) == {2, 8, 14, 20} positions = torch.arange(length, dtype=torch.int64, device="cuda") for layer in range(40): - reads = metadata.csa2_swa_indices[layer] - writes = metadata.csa2_swa_write_slots[layer] + reads = metadata.ext(CSA2Metadata).csa2_swa_indices[layer] + writes = metadata.ext(CSA2Metadata).csa2_swa_write_slots[layer] if direct: assert initial_pointers[layer] == (reads.data_ptr(), writes.data_ptr()) else: @@ -100,12 +101,14 @@ def test_context_metadata_omits_decoder_swa_and_preserves_boundary_global(length expected_visible = ( (positions + 1) // ratio if ratio else torch.zeros_like(positions) ) - torch.testing.assert_close(metadata.csa2_visible_lengths[layer], expected_visible) - assert torch.all(metadata.csa2_main_write_slots[20] >= 0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_visible_lengths[layer], expected_visible + ) + assert torch.all(metadata.ext(CSA2Metadata).csa2_main_write_slots[20] >= 0) assert manager.get_main_buffer(20).data_ptr() == manager.get_main_buffer(39).data_ptr() metadata.num_contexts = 0 with pytest.raises(ValueError, match="context-only cache cannot prepare generation"): - metadata.prepare_csa2() + metadata.ext(CSA2Metadata).prepare_csa2() finally: manager.free_resources(request) manager.shutdown() diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_indexer.py b/tests/unittest/_torch/attention/sparse/csa2/test_indexer.py index 09197c391c13..60e873b5f19c 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_indexer.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_indexer.py @@ -6,6 +6,7 @@ import torch import torch.nn.functional as F +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager from tensorrt_llm._torch.attention.backends.sparse.csa2.compressor import CSA2CompressionBatch from tensorrt_llm._torch.attention.backends.sparse.csa2.indexer import ( @@ -14,7 +15,7 @@ _ChunkInputs, _QueryChunk, ) -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import ( CSA2ForwardState, CSA2Layout, @@ -862,29 +863,31 @@ def test_native_paged_owner_indexer_and_graph(monkeypatch, ratio): # One request carries two verification queries; declare that arena width. manager.max_total_draft_tokens = 1 q, _, weights = _projected(3) - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=3) + metadata = CSA2TrtllmAttention.create_metadata(max_num_requests=2, max_num_tokens=3) metadata.kv_cache_manager = manager metadata.mapping = Mapping() metadata.is_cuda_graph = True - metadata.csa2_kv_sources = {0: 0} - metadata.csa2_num_context_requests = 0 - metadata.csa2_request_query_ranges = [(0, 2), (2, 3)] - metadata.csa2_request_start_positions = [125, 127] - metadata.csa2_request_lengths = [2, 1] - metadata.csa2_global_max_positions = {0: 128} + metadata.ext(CSA2Metadata).csa2_kv_sources = {0: 0} + metadata.ext(CSA2Metadata).csa2_num_context_requests = 0 + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, 2), (2, 3)] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [125, 127] + metadata.ext(CSA2Metadata).csa2_request_lengths = [2, 1] + metadata.ext(CSA2Metadata).csa2_global_max_positions = {0: 128} page_size = 128 // ratio - metadata.csa2_global_page_sizes = {0: page_size} + metadata.ext(CSA2Metadata).csa2_global_page_sizes = {0: page_size} table = torch.tensor([[0, 1], [2, 3]], device="cuda", dtype=torch.int32) visible = torch.tensor([63, 63, 64], device="cuda", dtype=torch.int32) - metadata.csa2_global_page_tables = {0: table} - metadata.csa2_visible_lengths = {0: visible} - metadata.csa2_token_requests = torch.tensor([0, 0, 1], device="cuda", dtype=torch.int64) - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} - metadata.indexer_max_chunk_size = 16 - metadata.indexer_q_split_threshold = -1 + metadata.ext(CSA2Metadata).csa2_global_page_tables = {0: table} + metadata.ext(CSA2Metadata).csa2_visible_lengths = {0: visible} + metadata.ext(CSA2Metadata).csa2_token_requests = torch.tensor( + [0, 0, 1], device="cuda", dtype=torch.int64 + ) + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).indexer_max_chunk_size = 16 + metadata.ext(CSA2Metadata).indexer_q_split_threshold = -1 state = CSA2ForwardState( metadata=metadata, swa_kv=torch.zeros(3, 512, device="cuda", dtype=torch.bfloat16), @@ -916,8 +919,10 @@ def paged(*args, **kwargs): assert native_calls and all(shape[:2] == (3, 1) for shape in native_calls) assert (32, False) in topk_calls if ratio == 1: - assert metadata.csa2_candidates[0].shape[0] == 3 # Candidate blocks were published. - bound = metadata.prepare_indexer(0) + assert ( + metadata.ext(CSA2Metadata).csa2_candidates[0].shape[0] == 3 + ) # Candidate blocks were published. + bound = metadata.ext(CSA2Metadata).prepare_indexer(0) ptrs = ( bound.csa2_indexer_k_cache.data_ptr(), bound.csa2_indexer_block_table.data_ptr(), @@ -937,14 +942,14 @@ def paged(*args, **kwargs): graph.replay() torch.cuda.synchronize() selected = output.clone() - refreshed = metadata.prepare_indexer(0) + refreshed = metadata.ext(CSA2Metadata).prepare_indexer(0) assert ptrs == ( refreshed.csa2_indexer_k_cache.data_ptr(), refreshed.csa2_indexer_block_table.data_ptr(), refreshed.csa2_indexer_scheduler_metadata.data_ptr(), ) positions = torch.arange(128, device="cuda").expand(3, -1) - req = metadata.csa2_token_requests + req = metadata.ext(CSA2Metadata).csa2_token_requests physical = table[req[:, None], positions // page_size] valid = (physical >= 0) & (positions < visible[:, None]) rows = read_index_rows( @@ -990,38 +995,47 @@ def test_multi_rank_candidate_publication_and_reindex(count, chunk_size, mixed, manager = _StridedOwner(layout, keys) total = count + int(mixed) q, _, weights = _projected(total) - metadata = CSA2TrtllmMetadata(max_num_requests=1 + int(mixed), max_num_tokens=total) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1 + int(mixed), max_num_tokens=total + ) metadata.kv_cache_manager = manager metadata.mapping = Mapping(world_size=world_size, rank=mpi_rank(), tp_size=world_size) - metadata.csa2_kv_sources = {0: 0, 1: 0} - metadata.csa2_num_context_requests = 1 - metadata.csa2_request_query_ranges = [(0, count)] - metadata.csa2_request_start_positions = [59] - metadata.csa2_request_lengths = [count] - metadata.csa2_global_max_positions = {0: 128} - metadata.csa2_global_page_sizes = {0: 128} - metadata.csa2_global_page_tables = {0: torch.zeros(1, 1, device="cuda", dtype=torch.int32)} + metadata.ext(CSA2Metadata).csa2_kv_sources = {0: 0, 1: 0} + metadata.ext(CSA2Metadata).csa2_num_context_requests = 1 + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, count)] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [59] + metadata.ext(CSA2Metadata).csa2_request_lengths = [count] + metadata.ext(CSA2Metadata).csa2_global_max_positions = {0: 128} + metadata.ext(CSA2Metadata).csa2_global_page_sizes = {0: 128} + metadata.ext(CSA2Metadata).csa2_global_page_tables = { + 0: torch.zeros(1, 1, device="cuda", dtype=torch.int32) + } visible = torch.arange(60, 60 + count, device="cuda", dtype=torch.int32) - metadata.csa2_visible_lengths = {0: visible, 1: visible} - metadata.csa2_token_requests = torch.zeros(count, device="cuda", dtype=torch.int64) + metadata.ext(CSA2Metadata).csa2_visible_lengths = {0: visible, 1: visible} + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + count, device="cuda", dtype=torch.int64 + ) if mixed: - metadata.csa2_request_query_ranges.append((count, total)) - metadata.csa2_request_start_positions.append(42) - metadata.csa2_request_lengths.append(1) - metadata.csa2_global_page_tables[0] = torch.zeros( + metadata.ext(CSA2Metadata).csa2_request_query_ranges.append((count, total)) + metadata.ext(CSA2Metadata).csa2_request_start_positions.append(42) + metadata.ext(CSA2Metadata).csa2_request_lengths.append(1) + metadata.ext(CSA2Metadata).csa2_global_page_tables[0] = torch.zeros( 2, 1, device="cuda", dtype=torch.int32 ) visible = torch.cat((visible, torch.tensor([43], device="cuda", dtype=torch.int32))) - metadata.csa2_visible_lengths = {0: visible, 1: visible} - metadata.csa2_token_requests = torch.cat( - (metadata.csa2_token_requests, torch.ones(1, device="cuda", dtype=torch.int64)) + metadata.ext(CSA2Metadata).csa2_visible_lengths = {0: visible, 1: visible} + metadata.ext(CSA2Metadata).csa2_token_requests = torch.cat( + ( + metadata.ext(CSA2Metadata).csa2_token_requests, + torch.ones(1, device="cuda", dtype=torch.int64), + ) ) - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} - metadata.indexer_max_chunk_size = chunk_size - metadata.indexer_q_split_threshold = 0 + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).indexer_max_chunk_size = chunk_size + metadata.ext(CSA2Metadata).indexer_q_split_threshold = 0 state = CSA2ForwardState( metadata=metadata, swa_kv=torch.zeros(total, 512, device="cuda", dtype=torch.bfloat16), @@ -1030,9 +1044,9 @@ def test_multi_rank_candidate_publication_and_reindex(count, chunk_size, mixed, ) source = CSA2Indexer(layout, 0, 32, 128) source_result = source(state, 0, total).clone() - gathered_candidates = metadata.csa2_candidates[0].clone() - gathered_blocks = metadata.csa2_candidate_blocks[0].clone() - gathered_counts = metadata.csa2_candidate_counts[0].clone() + gathered_candidates = metadata.ext(CSA2Metadata).csa2_candidates[0].clone() + gathered_blocks = metadata.ext(CSA2Metadata).csa2_candidate_blocks[0].clone() + gathered_counts = metadata.ext(CSA2Metadata).csa2_candidate_counts[0].clone() consumer = CSA2Indexer(layout, 1, 32, 128) state.index_q = -q shared_calls = [] @@ -1055,15 +1069,15 @@ def record_shared(*args, **kwargs): # Execute the same requests without query splitting through the same # scoring path, including the rank with zero local Q. metadata.mapping = Mapping() - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} state.index_q = q expected_source = source(state, 0, total).clone() - expected_candidates = metadata.csa2_candidates[0].clone() - expected_blocks = metadata.csa2_candidate_blocks[0].clone() - expected_counts = metadata.csa2_candidate_counts[0].clone() + expected_candidates = metadata.ext(CSA2Metadata).csa2_candidates[0].clone() + expected_blocks = metadata.ext(CSA2Metadata).csa2_candidate_blocks[0].clone() + expected_counts = metadata.ext(CSA2Metadata).csa2_candidate_counts[0].clone() state.index_q = -q expected_consumer = consumer(state, 0, total).clone() torch.testing.assert_close(source_result, expected_source, atol=0, rtol=0) @@ -1090,35 +1104,37 @@ def _single_request_state(heads=8, two_owners=False): 1: _index_pages_from_rows(pack_rows(-keys, "index")), } manager.get_index_pages = lambda layer: manager.owner_buffers[layout.layer(layer).kv_source] - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=1) + metadata = CSA2TrtllmAttention.create_metadata(max_num_requests=1, max_num_tokens=1) metadata.kv_cache_manager = manager metadata.mapping = Mapping() - metadata.csa2_kv_sources = {i: i for i in layout.kv_source_layer_ids} - metadata.csa2_num_context_requests = 0 - metadata.csa2_request_query_ranges = [(0, 1)] - metadata.csa2_request_start_positions = [3] - metadata.csa2_request_lengths = [1] - metadata.csa2_global_max_positions = { + metadata.ext(CSA2Metadata).csa2_kv_sources = {i: i for i in layout.kv_source_layer_ids} + metadata.ext(CSA2Metadata).csa2_num_context_requests = 0 + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, 1)] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [3] + metadata.ext(CSA2Metadata).csa2_request_lengths = [1] + metadata.ext(CSA2Metadata).csa2_global_max_positions = { i: 256 // layout.compress_ratios[i] for i in layout.kv_source_layer_ids } - metadata.csa2_global_page_sizes = { + metadata.ext(CSA2Metadata).csa2_global_page_sizes = { i: 128 // layout.compress_ratios[i] for i in layout.kv_source_layer_ids } - metadata.csa2_global_page_tables = { + metadata.ext(CSA2Metadata).csa2_global_page_tables = { i: torch.tensor([[0, 1]], device="cuda", dtype=torch.int32) for i in layout.kv_source_layer_ids } - metadata.csa2_visible_lengths = { + metadata.ext(CSA2Metadata).csa2_visible_lengths = { i: torch.tensor([4 // layout.compress_ratios[i]], device="cuda", dtype=torch.int32) for i in layout.kv_source_layer_ids } - metadata.csa2_token_requests = torch.zeros(1, device="cuda", dtype=torch.int64) - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} - metadata.indexer_max_chunk_size = 16 - metadata.indexer_q_split_threshold = -1 + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + 1, device="cuda", dtype=torch.int64 + ) + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).indexer_max_chunk_size = 16 + metadata.ext(CSA2Metadata).indexer_q_split_threshold = -1 state = CSA2ForwardState( metadata=metadata, swa_kv=torch.zeros(1, 512, device="cuda", dtype=torch.bfloat16), @@ -1131,10 +1147,12 @@ def _single_request_state(heads=8, two_owners=False): def _assert_state_selection(output, state, owner): metadata = state.metadata manager = metadata.kv_cache_manager - width = metadata.csa2_global_max_positions[owner] + width = metadata.ext(CSA2Metadata).csa2_global_max_positions[owner] positions = torch.arange(width, device="cuda").unsqueeze(0) - slots = metadata.global_slot_tile(owner, 0, 1, positions) - valid = (slots >= 0) & (positions < metadata.csa2_visible_lengths[owner][:, None]) + slots = metadata.ext(CSA2Metadata).global_slot_tile(owner, 0, 1, positions) + valid = (slots >= 0) & ( + positions < metadata.ext(CSA2Metadata).csa2_visible_lengths[owner][:, None] + ) rows = read_index_rows(manager.get_index_pages(owner), slots.clamp_min(0)) keys = unpack_rows(rows, 128, "index").float() query = unpack_rows(pack_rows(state.index_q, "index"), 128, "index").float() @@ -1160,7 +1178,7 @@ def gather(owner, slots): monkeypatch.setattr(manager, "gather_indexer_keys", gather) # Exercise the bounded fallback independently of native head coverage. - monkeypatch.setattr(metadata, "prepare_indexer", lambda layer: None) + monkeypatch.setattr(metadata.ext(CSA2Metadata), "prepare_indexer", lambda layer: None) monkeypatch.setattr( indexer, "_call_paged_mqa_logits", @@ -1176,8 +1194,8 @@ def gather(owner, slots): with torch.cuda.graph(graph): output = indexer(state, 0, 1) for length, second_page in ((96, 1), (96, 2), (0, -1), (65, 1)): - metadata.csa2_visible_lengths[0].fill_(length) - metadata.csa2_global_page_tables[0][0, 1] = second_page + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(length) + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 1] = second_page graph.replay() torch.cuda.synchronize() _assert_state_selection(output, state, 0) @@ -1195,12 +1213,12 @@ def test_paged_metadata_eager_growth_reuses_bounded_arena(): previous = None previous_capacity = 0 for visible in (1, 2, 31, 32, 63, 64, 65, 66, 96, 127, 128): - metadata.csa2_request_start_positions = [visible * 2 - 1] - metadata.csa2_visible_lengths[0].fill_(visible) - metadata.csa2_visible_lengths[1].fill_(visible * 2) + metadata.ext(CSA2Metadata).csa2_request_start_positions = [visible * 2 - 1] + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(visible) + metadata.ext(CSA2Metadata).csa2_visible_lengths[1].fill_(visible * 2) for owner in layout.kv_source_layer_ids: - binding = metadata.prepare_indexer(owner) - assert len(metadata._csa2_indexer_workspaces) == 1 + binding = metadata.ext(CSA2Metadata).prepare_indexer(owner) + assert len(metadata.ext(CSA2Metadata)._csa2_indexer_workspaces) == 1 # Index pages are read in place; only descriptors are retained, # and not per exact sequence length or per owner. assert ( @@ -1213,16 +1231,18 @@ def test_paged_metadata_eager_growth_reuses_bounded_arena(): previous, previous_capacity = pointer, capacity assert previous_capacity <= 4 # A changed packed-query geometry must replace its eager arena too. - metadata.csa2_request_query_ranges = [(0, 2)] - metadata.csa2_request_lengths = [2] - metadata.csa2_request_start_positions = [0] - metadata.csa2_token_requests = torch.zeros(2, device="cuda", dtype=torch.int64) - metadata.csa2_visible_lengths = { + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, 2)] + metadata.ext(CSA2Metadata).csa2_request_lengths = [2] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [0] + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + 2, device="cuda", dtype=torch.int64 + ) + metadata.ext(CSA2Metadata).csa2_visible_lengths = { owner: torch.ones(2, device="cuda", dtype=torch.int32) for owner in layout.kv_source_layer_ids } - metadata.prepare_indexer(0) - assert len(metadata._csa2_indexer_workspaces) == 1 + metadata.ext(CSA2Metadata).prepare_indexer(0) + assert len(metadata.ext(CSA2Metadata)._csa2_indexer_workspaces) == 1 @skip_pre_blackwell @@ -1244,12 +1264,12 @@ def run(): run() pointers = [] for owner in layout.kv_source_layer_ids: - metadata.prepare_indexer(owner) + metadata.ext(CSA2Metadata).prepare_indexer(owner) pointers.append( ( - metadata.csa2_indexer_k_cache.data_ptr(), - metadata.csa2_indexer_block_table.data_ptr(), - metadata.csa2_indexer_scheduler_metadata.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_k_cache.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_block_table.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_scheduler_metadata.data_ptr(), ) ) assert pointers[0][1] == pointers[1][1] @@ -1259,20 +1279,20 @@ def run(): with torch.cuda.graph(graph): outputs = run() for first, second in ((65, 129), (1, 2), (0, 0), (96, 191)): - metadata.csa2_visible_lengths[0].fill_(first) - metadata.csa2_visible_lengths[1].fill_(second) + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(first) + metadata.ext(CSA2Metadata).csa2_visible_lengths[1].fill_(second) state.index_q.neg_() graph.replay() torch.cuda.synchronize() for owner, output in enumerate(outputs): _assert_state_selection(output, state, owner) - binding = metadata.prepare_indexer(owner) + binding = metadata.ext(CSA2Metadata).prepare_indexer(owner) assert pointers[owner] == ( binding.csa2_indexer_k_cache.data_ptr(), binding.csa2_indexer_block_table.data_ptr(), binding.csa2_indexer_scheduler_metadata.data_ptr(), ) - assert len(metadata._csa2_indexer_workspaces) == 1 + assert len(metadata.ext(CSA2Metadata)._csa2_indexer_workspaces) == 1 @skip_pre_blackwell @@ -1282,19 +1302,21 @@ def test_paged_descriptors_refresh_when_the_manager_changes(): """Target and draft managers share the descriptor arena and advance serials in lockstep.""" layout, manager, metadata, state = _single_request_state() metadata.is_cuda_graph = True - metadata.csa2_visible_lengths[0].fill_(96) - assert metadata.prepare_indexer(0) is metadata - assert metadata.csa2_indexer_block_table.tolist() == [[0, 1]] + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(96) + assert metadata.ext(CSA2Metadata).prepare_indexer(0) is metadata.ext(CSA2Metadata) + assert metadata.ext(CSA2Metadata).csa2_indexer_block_table.tolist() == [[0, 1]] keys = torch.ones(512, 128, device="cuda", dtype=torch.bfloat16) draft = _StridedOwner(layout, pack_rows(keys, "index")) metadata.kv_cache_manager = draft - metadata.csa2_global_page_tables[0].copy_( + metadata.ext(CSA2Metadata).csa2_global_page_tables[0].copy_( torch.tensor([[3, 2]], device="cuda", dtype=torch.int32) ) # Same forward serial, different manager: descriptors must follow the draft. - assert metadata.prepare_indexer(0) is metadata - assert metadata.csa2_indexer_k_cache.data_ptr() == draft.index_pages.data_ptr() - assert metadata.csa2_indexer_block_table.tolist() == [[3, 2]] + assert metadata.ext(CSA2Metadata).prepare_indexer(0) is metadata.ext(CSA2Metadata) + assert ( + metadata.ext(CSA2Metadata).csa2_indexer_k_cache.data_ptr() == draft.index_pages.data_ptr() + ) + assert metadata.ext(CSA2Metadata).csa2_indexer_block_table.tolist() == [[3, 2]] @skip_pre_blackwell @@ -1307,7 +1329,7 @@ def test_native_paged_head_coverage_and_optional_dsl(monkeypatch, heads, use_dsl pytest.skip("Native FP4 paged variants require SM100") layout, _, metadata, state = _single_request_state(heads=heads) metadata.is_cuda_graph = True - metadata.csa2_visible_lengths[0].fill_(96) + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(96) indexer = CSA2Indexer( layout, 0, heads, 128, options=CSA2Params(use_cute_dsl_paged_mqa_logits=use_dsl) ) @@ -1327,8 +1349,8 @@ def record(*args, **kwargs): with torch.cuda.graph(graph): output = indexer(state, 0, 1) for length, page in ((65, 1), (96, 2), (0, -1)): - metadata.csa2_visible_lengths[0].fill_(length) - metadata.csa2_global_page_tables[0][0, 1] = page + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(length) + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 1] = page graph.replay() _assert_state_selection(output, state, 0) @@ -1351,7 +1373,7 @@ def _temporal_gvr_case(monkeypatch, emission, *, page_holes=False, candidate_sou ) manager.epoch = 1 manager.request_epoch = lambda request: manager.epoch - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=1) + metadata = CSA2TrtllmAttention.create_metadata(max_num_requests=1, max_num_tokens=1) metadata.kv_cache_manager = manager metadata.mapping = Mapping() metadata.request_ids = [13] @@ -1359,27 +1381,31 @@ def _temporal_gvr_case(monkeypatch, emission, *, page_holes=False, candidate_sou metadata._num_contexts = 0 metadata._num_tokens = 1 metadata._num_generations = 1 - metadata.csa2_kv_sources = {0: 0} - metadata.csa2_num_context_requests = 0 - metadata.csa2_request_query_ranges = [(0, 1)] - metadata.csa2_request_start_positions = [32767] - metadata.csa2_request_lengths = [1] - metadata.csa2_global_max_positions = {0: width} - metadata.csa2_global_page_sizes = {0: 128} - metadata.csa2_global_page_tables = { + metadata.ext(CSA2Metadata).csa2_kv_sources = {0: 0} + metadata.ext(CSA2Metadata).csa2_num_context_requests = 0 + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, 1)] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [32767] + metadata.ext(CSA2Metadata).csa2_request_lengths = [1] + metadata.ext(CSA2Metadata).csa2_global_max_positions = {0: width} + metadata.ext(CSA2Metadata).csa2_global_page_sizes = {0: 128} + metadata.ext(CSA2Metadata).csa2_global_page_tables = { 0: torch.arange(width // 128, device="cuda", dtype=torch.int32)[None, :] } if page_holes: - metadata.csa2_global_page_tables[0][0, 2] = -1 - metadata.csa2_visible_lengths = {0: torch.tensor([32768], device="cuda", dtype=torch.int32)} - metadata.csa2_token_requests = torch.zeros(1, device="cuda", dtype=torch.int64) - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} - metadata.indexer_max_chunk_size = 16 - metadata.indexer_q_split_threshold = -1 - metadata._csa2_forward_serial = 0 + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 2] = -1 + metadata.ext(CSA2Metadata).csa2_visible_lengths = { + 0: torch.tensor([32768], device="cuda", dtype=torch.int32) + } + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + 1, device="cuda", dtype=torch.int64 + ) + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).indexer_max_chunk_size = 16 + metadata.ext(CSA2Metadata).indexer_q_split_threshold = -1 + metadata.ext(CSA2Metadata)._csa2_forward_serial = 0 state = CSA2ForwardState( metadata=metadata, swa_kv=torch.zeros(1, 512, device="cuda", dtype=torch.bfloat16), @@ -1426,24 +1452,24 @@ def logits(*args, **kwargs): output = indexer(state, 0, 1) graph.replay() for start, epoch, expect_hint in ((32768, 1, True), (32769, 2, False), (511, 2, False)): - metadata.csa2_request_start_positions = [start] - metadata.csa2_visible_lengths[0].fill_(start + 1) + metadata.ext(CSA2Metadata).csa2_request_start_positions = [start] + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(start + 1) manager.epoch = epoch - metadata._csa2_forward_serial += 1 + metadata.ext(CSA2Metadata)._csa2_forward_serial += 1 resets_before = len(resets) - prior = metadata.prepare_indexer_prior(0, 512) + prior = metadata.ext(CSA2Metadata).prepare_indexer_prior(0, 512) assert bool((prior >= 0).any()) == expect_hint if not expect_hint: assert len(resets) > resets_before state.index_q.copy_(torch.randn_like(state.index_q)) if page_holes: - metadata.csa2_global_page_tables[0][0, 2] = 2 if expect_hint else -1 + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 2] = 2 if expect_hint else -1 graph.replay() _assert_state_selection(output, state, 0) if candidate_source: - candidates = metadata.csa2_candidates[0] + candidates = metadata.ext(CSA2Metadata).csa2_candidates[0] assert candidates.shape == (1, 32) - slots = metadata.global_slot_tile(0, 0, 1, candidates) + slots = metadata.ext(CSA2Metadata).global_slot_tile(0, 0, 1, candidates) assert bool(((candidates < 0) | (slots >= 0)).all()) @@ -1471,7 +1497,7 @@ def test_packed_query_state_avoids_requantization(monkeypatch): pytest.skip("Packed-query native integration requires SM100") layout, _, metadata, state = _single_request_state(heads=64) metadata.is_cuda_graph = True - metadata.csa2_visible_lengths[0].fill_(96) + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(96) indexer = CSA2Indexer(layout, 0, 64, 128) expected = indexer(state, 0, 1) packed = pack_rows(state.index_q, "index") @@ -1509,9 +1535,9 @@ def test_decode_short_skip_respects_graph_admission(monkeypatch, mode): layout, manager, metadata, state = _single_request_state(heads=64) metadata.is_cuda_graph = mode != "eager" if mode == "bounded_graph": - metadata.csa2_global_max_positions[0] = 16 - metadata.csa2_visible_lengths[0].fill_(12) - metadata.csa2_request_start_positions = [23] + metadata.ext(CSA2Metadata).csa2_global_max_positions[0] = 16 + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(12) + metadata.ext(CSA2Metadata).csa2_request_start_positions = [23] indexer = CSA2Indexer(layout, 0, 64, 128) calls = [] native = indexer._paged_mqa_logits @@ -1534,14 +1560,14 @@ def unexpected(*args, **kwargs): output = indexer(state, 0, 1) _assert_state_selection(output, state, 0) if mode == "eager": - metadata.csa2_global_page_tables[0][0, 0] = -1 + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 0] = -1 output = indexer(state, 0, 1) assert bool((output == -1).all()) return graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): output = indexer(state, 0, 1) - metadata.csa2_visible_lengths[0].fill_(96 if mode == "growing_graph" else 16) + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(96 if mode == "growing_graph" else 16) graph.replay() _assert_state_selection(output, state, 0) assert bool(calls) == (mode == "growing_graph") @@ -1564,10 +1590,10 @@ def test_short_candidate_source_still_computes_block_scores(monkeypatch): index_topk=32, ) manager.layout = layout - metadata.csa2_global_page_sizes[0] = 128 - metadata.csa2_global_max_positions[0] = 16 - metadata.csa2_visible_lengths[0].fill_(4) - metadata.csa2_request_start_positions = [3] + metadata.ext(CSA2Metadata).csa2_global_page_sizes[0] = 128 + metadata.ext(CSA2Metadata).csa2_global_max_positions[0] = 16 + metadata.ext(CSA2Metadata).csa2_visible_lengths[0].fill_(4) + metadata.ext(CSA2Metadata).csa2_request_start_positions = [3] indexer = CSA2Indexer(layout, 0, 64, 128) calls = [] native = indexer._paged_mqa_logits @@ -1579,7 +1605,7 @@ def record(*args, **kwargs): monkeypatch.setattr(indexer, "_paged_mqa_logits", record) output = indexer(state, 0, 1) assert len(calls) == 1 - assert metadata.csa2_candidates[0].shape == (1, 8) + assert metadata.ext(CSA2Metadata).csa2_candidates[0].shape == (1, 8) _assert_state_selection(output, state, 0) @@ -1602,38 +1628,58 @@ def test_live_acceptance_graph_refreshes_native_paged_indexer(monkeypatch, query manager.max_total_draft_tokens = query_width - 1 metadata.is_cuda_graph = True metadata._num_tokens = query_width - metadata.csa2_request_query_ranges = [(0, query_width)] - metadata.csa2_request_lengths = [query_width] - metadata.csa2_request_start_positions = [127] + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, query_width)] + metadata.ext(CSA2Metadata).csa2_request_lengths = [query_width] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [127] device = torch.device("cuda") - metadata._csa2_ready_for_kv_update = True - metadata._csa2_query_base = torch.tensor([127], dtype=torch.int32, device=device) - metadata._csa2_query_lengths = torch.tensor([query_width], dtype=torch.int32, device=device) - metadata._csa2_query_offsets = torch.arange(query_width, dtype=torch.int32, device=device) - metadata.csa2_token_requests = torch.zeros(query_width, dtype=torch.int64, device=device) - metadata.csa2_positions = torch.empty(query_width, dtype=torch.int32, device=device) - metadata.csa2_replay_start_positions = torch.zeros(1, dtype=torch.int64, device=device) - metadata._csa2_swa_page_tables = { + metadata.ext(CSA2Metadata)._csa2_ready_for_kv_update = True + metadata.ext(CSA2Metadata)._csa2_query_base = torch.tensor( + [127], dtype=torch.int32, device=device + ) + metadata.ext(CSA2Metadata)._csa2_query_lengths = torch.tensor( + [query_width], dtype=torch.int32, device=device + ) + metadata.ext(CSA2Metadata)._csa2_query_offsets = torch.arange( + query_width, dtype=torch.int32, device=device + ) + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + query_width, dtype=torch.int64, device=device + ) + metadata.ext(CSA2Metadata).csa2_positions = torch.empty( + query_width, dtype=torch.int32, device=device + ) + metadata.ext(CSA2Metadata).csa2_replay_start_positions = torch.zeros( + 1, dtype=torch.int64, device=device + ) + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = { owner: torch.tensor([[0, 1]], dtype=torch.int32, device=device) for owner in (0, 1) } - metadata._csa2_buffers = {} - metadata._csa2_swa_descriptors = metadata._swa_table_descriptors(2) - metadata._csa2_swa_ratios = metadata._swa_ratio_tensor(layout.compress_ratios) + metadata.ext(CSA2Metadata)._csa2_buffers = {} + metadata.ext(CSA2Metadata)._csa2_swa_descriptors = metadata.ext( + CSA2Metadata + )._swa_table_descriptors(2) + metadata.ext(CSA2Metadata)._csa2_swa_ratios = metadata.ext(CSA2Metadata)._swa_ratio_tensor( + layout.compress_ratios + ) # The forward resolves the SWA slots; keep them for the checks. kept = { owner: tuple(torch.empty(query_width, dtype=torch.int64, device=device) for _ in range(2)) for owner in (0, 1) } - metadata._csa2_source_geometry = { - owner: (metadata._csa2_query_base, metadata._csa2_query_lengths) for owner in (0, 1) + metadata.ext(CSA2Metadata)._csa2_source_geometry = { + owner: ( + metadata.ext(CSA2Metadata)._csa2_query_base, + metadata.ext(CSA2Metadata)._csa2_query_lengths, + ) + for owner in (0, 1) } - metadata.csa2_main_write_slots = { + metadata.ext(CSA2Metadata).csa2_main_write_slots = { owner: torch.empty(query_width, dtype=torch.int64, device=device) for owner in (0, 1) } - metadata._csa2_compressed_positions = { + metadata.ext(CSA2Metadata)._csa2_compressed_positions = { owner: torch.empty(query_width, dtype=torch.int32, device=device) for owner in (0, 1) } - metadata._csa2_compression = { + metadata.ext(CSA2Metadata)._csa2_compression = { 0: CSA2CompressionBatch( torch.empty(0, device=device), torch.empty(0, device=device), @@ -1648,7 +1694,9 @@ def test_live_acceptance_graph_refreshes_native_paged_indexer(monkeypatch, query query_width, ) } - metadata._csa2_owner_descriptors = metadata._owner_slot_descriptors() + metadata.ext(CSA2Metadata)._csa2_owner_descriptors = metadata.ext( + CSA2Metadata + )._owner_slot_descriptors() state.index_q = state.index_q.expand(query_width, -1, -1).clone() state.index_weights = state.index_weights.expand(query_width, -1).clone() state.swa_kv = state.swa_kv.expand(query_width, -1).clone() @@ -1662,19 +1710,19 @@ def test_live_acceptance_graph_refreshes_native_paged_indexer(monkeypatch, query def run(): metadata.on_update_kv_lens() - metadata._ensure_swa_slots() # As entering the first layer does. + metadata.ext(CSA2Metadata)._ensure_swa_slots() # As entering the first layer does. for owner in (0, 1): - kept[owner][0].copy_(metadata.csa2_swa_write_slots[owner]) - kept[owner][1].copy_(metadata.csa2_visible_lengths[owner]) + kept[owner][0].copy_(metadata.ext(CSA2Metadata).csa2_swa_write_slots[owner]) + kept[owner][1].copy_(metadata.ext(CSA2Metadata).csa2_visible_lengths[owner]) return tuple(indexer(state, 0, query_width) for indexer in indexers) metadata.kv_lens_cuda[0] = 130 for _ in range(3): run() pointers = ( - metadata.csa2_positions.data_ptr(), - metadata.csa2_main_write_slots[0].data_ptr(), - metadata.csa2_indexer_scheduler_metadata.data_ptr(), + metadata.ext(CSA2Metadata).csa2_positions.data_ptr(), + metadata.ext(CSA2Metadata).csa2_main_write_slots[0].data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_scheduler_metadata.data_ptr(), ) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): @@ -1685,19 +1733,19 @@ def run(): graph.replay() torch.cuda.synchronize() positions = torch.arange(endpoint - query_width, endpoint, device=device) - torch.testing.assert_close(metadata.csa2_positions.long(), positions) + torch.testing.assert_close(metadata.ext(CSA2Metadata).csa2_positions.long(), positions) torch.testing.assert_close(kept[0][0], positions) for owner, selected in enumerate(output): ratio = layout.compress_ratios[owner] visible = (positions + 1) // ratio torch.testing.assert_close(kept[owner][1], visible) groups = list(range((endpoint - query_width) // ratio, endpoint // ratio)) - assert metadata.csa2_main_write_slots[owner].tolist() == groups + [-1] * ( - query_width - len(groups) - ) - width = metadata.csa2_global_max_positions[owner] + assert metadata.ext(CSA2Metadata).csa2_main_write_slots[owner].tolist() == groups + [ + -1 + ] * (query_width - len(groups)) + width = metadata.ext(CSA2Metadata).csa2_global_max_positions[owner] logical = torch.arange(width, device=device).expand(query_width, -1) - slots = metadata.global_slot_tile(owner, 0, query_width, logical) + slots = metadata.ext(CSA2Metadata).global_slot_tile(owner, 0, query_width, logical) keys = unpack_rows( read_index_rows(manager.get_index_pages(owner), slots.clamp_min(0)), 128, "index" ).float() @@ -1706,7 +1754,7 @@ def run(): torch.einsum("qhd,qkd->qhk", query, keys).relu() * state.index_weights[..., None] ).sum(1) _assert_selection(selected, scores, (slots >= 0) & (logical < visible[:, None])) - batch = metadata._csa2_compression[0] + batch = metadata.ext(CSA2Metadata)._csa2_compression[0] assert batch.start_positions.item() == endpoint - query_width assert batch.kv_lengths.item() == endpoint assert batch.cu_compressed_lengths.tolist() == [ @@ -1716,17 +1764,21 @@ def run(): # The serialized ratio-one owner ran last, so its native scheduling # inputs and refreshed schedule occupy the shared arena now. expected_context = (positions + 1).int().clamp_min(1)[:, None] - torch.testing.assert_close(metadata.csa2_indexer_context_lengths, expected_context) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_indexer_context_lengths, expected_context + ) expected_schedule = get_paged_mqa_logits_metadata( expected_context, 64, torch.cuda.get_device_properties(device).multi_processor_count, ) - torch.testing.assert_close(metadata.csa2_indexer_scheduler_metadata, expected_schedule) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_indexer_scheduler_metadata, expected_schedule + ) assert pointers == ( - metadata.csa2_positions.data_ptr(), - metadata.csa2_main_write_slots[0].data_ptr(), - metadata.csa2_indexer_scheduler_metadata.data_ptr(), + metadata.ext(CSA2Metadata).csa2_positions.data_ptr(), + metadata.ext(CSA2Metadata).csa2_main_write_slots[0].data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_scheduler_metadata.data_ptr(), ) @@ -1737,10 +1789,10 @@ def test_candidates_pin_latest_block_and_mask_future(): ) layout, state = _selection_state(1, keys, torch.tensor([5])) _run_indexer(layout, 0, state) - assert state.metadata.csa2_candidates[0].tolist() == [[4, -1]] + assert state.metadata.ext(CSA2Metadata).csa2_candidates[0].tolist() == [[4, -1]] state.main_kv = state.index_k = None _run_indexer(layout, 1, state) - assert state.metadata.csa2_indices[1].tolist() == [[4, -1, -1]] + assert state.metadata.ext(CSA2Metadata).csa2_indices[1].tolist() == [[4, -1, -1]] @pytest.mark.cpu_only @@ -1752,16 +1804,16 @@ def test_empty_visibility(queries, width): torch.zeros(queries, dtype=torch.int32), ) _run_indexer(layout, 0, state) - assert state.metadata.csa2_indices[0].shape == (queries, 3) - assert torch.all(state.metadata.csa2_indices[0] == -1) - assert torch.all(state.metadata.csa2_candidates[0] == -1) + assert state.metadata.ext(CSA2Metadata).csa2_indices[0].shape == (queries, 3) + assert torch.all(state.metadata.ext(CSA2Metadata).csa2_indices[0] == -1) + assert torch.all(state.metadata.ext(CSA2Metadata).csa2_candidates[0] == -1) @pytest.mark.cpu_only def test_full_reuse_reindex_private_swa_and_shared_sink(): outputs, routing, cache, batch = _run_modes("cpu") - assert routing.csa2_indices[0].tolist() == [[0]] - assert routing.csa2_indices[2].tolist() == [[1]] + assert routing.ext(CSA2Metadata).csa2_indices[0].tolist() == [[0]] + assert routing.ext(CSA2Metadata).csa2_indices[2].tolist() == [[1]] # q=0 means selected global KV, private SWA, and sink each get 1/3. for actual, expected in zip(outputs, [2.0, 3.0, 5.0]): torch.testing.assert_close( @@ -1769,10 +1821,10 @@ def test_full_reuse_reindex_private_swa_and_shared_sink(): ) layout = CSA2Layout((1, 1), (0,), (0,)) with pytest.raises(ValueError, match="across forwards"): - routing.enter_layer(layout.layer(0)) + routing.ext(CSA2Metadata).enter_layer(layout.layer(0)) q = torch.zeros(1, 1, 128, dtype=torch.bfloat16) with pytest.raises(ValueError, match="source did not run"): - batch.reset_routing() + batch.ext(CSA2Metadata).reset_routing() _prediction_reference(layout, 1, q, q[:, 0], torch.zeros(1), batch) @@ -1782,7 +1834,7 @@ def test_all_modes_cuda_graph_changed_routing(): # Preserve the independent analytical oracle before exercising replay changes. outputs, routing, _, _ = _run_modes("cuda") torch.cuda.synchronize() - assert routing.csa2_indices[2].tolist() == [[1]] + assert routing.ext(CSA2Metadata).csa2_indices[2].tolist() == [[1]] torch.testing.assert_close(outputs[-1], torch.full_like(outputs[-1], 5.0), atol=0.02, rtol=0) _run_modes("cuda", graph_replay=True) @@ -1802,13 +1854,13 @@ def _publish_manual_candidates(metadata, layer, block_ids, block_size=8): column count. Existing publications are updated in place so captured graphs keep reading the same buffers. Also binds the generation-row -> request map. """ - visible = metadata.csa2_visible_lengths[layer + 1].to(torch.int64) + visible = metadata.ext(CSA2Metadata).csa2_visible_lengths[layer + 1].to(torch.int64) count = visible.shape[0] device = visible.device - max_positions = metadata.csa2_global_max_positions[0] - table = metadata.csa2_global_page_tables[0] - page_size = metadata.csa2_global_page_sizes[0] - requests = metadata.csa2_token_requests.long() + max_positions = metadata.ext(CSA2Metadata).csa2_global_max_positions[0] + table = metadata.ext(CSA2Metadata).csa2_global_page_tables[0] + page_size = metadata.ext(CSA2Metadata).csa2_global_page_sizes[0] + requests = metadata.ext(CSA2Metadata).csa2_token_requests.long() width = -(-len(block_ids) // 4) * 4 * block_size ids = torch.tensor(sorted(set(block_ids)), dtype=torch.int64, device=device) first = ids * block_size @@ -1830,24 +1882,24 @@ def _publish_manual_candidates(metadata, layer, block_ids, block_size=8): blocks = torch.cat((blocks, last.expand(-1, width // block_size - blocks.shape[1])), 1).int() counts = (candidates >= 0).sum(1, keepdim=True).int() for store, value in ( - (metadata.csa2_candidates, candidates), - (metadata.csa2_candidate_blocks, blocks), - (metadata.csa2_candidate_counts, counts), + (metadata.ext(CSA2Metadata).csa2_candidates, candidates), + (metadata.ext(CSA2Metadata).csa2_candidate_blocks, blocks), + (metadata.ext(CSA2Metadata).csa2_candidate_counts, counts), ): if layer in store and store[layer].shape == value.shape: store[layer].copy_(value) else: store[layer] = value.clone() - contexts = metadata.csa2_num_context_requests - ranges = metadata.csa2_request_query_ranges + contexts = metadata.ext(CSA2Metadata).csa2_num_context_requests + ranges = metadata.ext(CSA2Metadata).csa2_request_query_ranges decode_start = ranges[contexts][0] if contexts < len(ranges) else count rows = (requests[decode_start:] - contexts).int() - existing = getattr(metadata, "csa2_decode_row_requests", None) + existing = getattr(metadata.ext(CSA2Metadata), "csa2_decode_row_requests", None) # Captured graphs read this map by address as well: update it in place. if existing is not None and existing.shape == rows.shape: existing.copy_(rows) else: - metadata.csa2_decode_row_requests = rows + metadata.ext(CSA2Metadata).csa2_decode_row_requests = rows def _candidate_decode_fixture(*, mixed=False, random=False): @@ -1864,33 +1916,43 @@ def _candidate_decode_fixture(*, mixed=False, random=False): manager = _StridedOwner(layout, pack_rows(keys, "index")) manager.max_total_draft_tokens = 1 count = 4 if mixed else 3 - metadata = CSA2TrtllmMetadata(max_num_requests=3 if mixed else 2, max_num_tokens=count) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=3 if mixed else 2, max_num_tokens=count + ) metadata.kv_cache_manager = manager metadata.mapping = Mapping() metadata.is_cuda_graph = True - metadata.csa2_kv_sources = {0: 0, 1: 0} - metadata.csa2_num_context_requests = int(mixed) - metadata.csa2_request_query_ranges = [(0, 1), (1, 3), (3, 4)] if mixed else [(0, 2), (2, 3)] - metadata.csa2_request_start_positions = [191, 191, 192] if mixed else [191, 192] - metadata.csa2_request_lengths = [1, 2, 1] if mixed else [2, 1] - metadata.csa2_global_max_positions = {0: 256} - metadata.csa2_global_page_sizes = {0: 128} - metadata.csa2_global_page_tables = { + metadata.ext(CSA2Metadata).csa2_kv_sources = {0: 0, 1: 0} + metadata.ext(CSA2Metadata).csa2_num_context_requests = int(mixed) + metadata.ext(CSA2Metadata).csa2_request_query_ranges = ( + [(0, 1), (1, 3), (3, 4)] if mixed else [(0, 2), (2, 3)] + ) + metadata.ext(CSA2Metadata).csa2_request_start_positions = ( + [191, 191, 192] if mixed else [191, 192] + ) + metadata.ext(CSA2Metadata).csa2_request_lengths = [1, 2, 1] if mixed else [2, 1] + metadata.ext(CSA2Metadata).csa2_global_max_positions = {0: 256} + metadata.ext(CSA2Metadata).csa2_global_page_sizes = {0: 128} + metadata.ext(CSA2Metadata).csa2_global_page_tables = { 0: torch.tensor( [[3, 2], [1, 0], [2, 3]] if mixed else [[1, 0], [2, 3]], dtype=torch.int32, device="cuda", ) } - metadata.csa2_visible_lengths = {1: torch.full((count,), 192, dtype=torch.int64, device="cuda")} - metadata.csa2_token_requests = torch.tensor([0, 1, 1, 2] if mixed else [0, 0, 1], device="cuda") - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).csa2_visible_lengths = { + 1: torch.full((count,), 192, dtype=torch.int64, device="cuda") + } + metadata.ext(CSA2Metadata).csa2_token_requests = torch.tensor( + [0, 1, 1, 2] if mixed else [0, 0, 1], device="cuda" + ) + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} _publish_manual_candidates(metadata, 0, _MANUAL_BLOCKS) - metadata.indexer_max_chunk_size = 16 - metadata.indexer_q_split_threshold = -1 + metadata.ext(CSA2Metadata).indexer_max_chunk_size = 16 + metadata.ext(CSA2Metadata).indexer_q_split_threshold = -1 q = torch.zeros(count, 32, 128, dtype=torch.bfloat16, device="cuda") q[:, :9, :9] = torch.eye(9, dtype=torch.bfloat16, device="cuda") weights = torch.zeros(count, 32, device="cuda") @@ -1925,8 +1987,8 @@ def _candidate_decode_run(indexer, state, monkeypatch, *, old): select = indexer._select_mapped_logits dense = indexer._call_mqa_logits sparse = indexer._sparse_candidate_logits - candidates = metadata.csa2_candidates[0] - width = max(metadata.csa2_global_max_positions[0], int(candidates.max()) + 1) + candidates = metadata.ext(CSA2Metadata).csa2_candidates[0] + width = max(metadata.ext(CSA2Metadata).csa2_global_max_positions[0], int(candidates.max()) + 1) rows_done = [0] def collect(logits, starts, ends, positions, visible, *args, **kwargs): @@ -2000,29 +2062,33 @@ def test_candidate_decode_batched_exact_and_mixed(monkeypatch, mixed): pytest.skip("Native paged FP4 requires SM100 family") indexer, state = _candidate_decode_fixture(mixed=mixed) # Keep active queries on both requests, plus an empty first query and a hole. - state.metadata.csa2_visible_lengths[1][0] = 0 - state.metadata.csa2_global_page_tables[0][0, 0] = -1 + state.metadata.ext(CSA2Metadata).csa2_visible_lengths[1][0] = 0 + state.metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 0] = -1 _publish_manual_candidates(state.metadata, 0, _MANUAL_BLOCKS) - candidate = state.metadata.csa2_candidates[0].clone() + candidate = state.metadata.ext(CSA2Metadata).csa2_candidates[0].clone() if not _sparse_kernels_available(): pytest.skip("DeepGEMM sparse MQA logits required") old, scores, control = _candidate_decode_run(indexer, state, monkeypatch, old=True) actual, new_scores, calls = _candidate_decode_run(indexer, state, monkeypatch, old=False) assert control == { - "dense": len(state.metadata.csa2_request_query_ranges), + "dense": len(state.metadata.ext(CSA2Metadata).csa2_request_query_ranges), "sparse_prefill": 0, "sparse_paged": 0, } assert calls == {"dense": 0, "sparse_prefill": int(mixed), "sparse_paged": 1} # The sparse kernels emit bf16 logits; the per-query gather reference is fp32. torch.testing.assert_close(new_scores, scores, atol=0, rtol=2**-7) - torch.testing.assert_close(state.metadata.csa2_candidates[0], candidate, atol=0, rtol=0) + torch.testing.assert_close( + state.metadata.ext(CSA2Metadata).csa2_candidates[0], candidate, atol=0, rtol=0 + ) _check_candidate_multiset(actual, scores, candidate, _bf16_tolerance(scores)) if mixed: # Without paged sparse descriptors the decode rows fall back to the # per-query candidate gathers without row-offset loss. with monkeypatch.context() as context: - context.setattr(state.metadata, "prepare_sparse_indexer", lambda *a: None) + context.setattr( + state.metadata.ext(CSA2Metadata), "prepare_sparse_indexer", lambda *a: None + ) fallback, fallback_scores, fallback_calls = _candidate_decode_run( indexer, state, monkeypatch, old=False ) @@ -2031,7 +2097,7 @@ def test_candidate_decode_batched_exact_and_mixed(monkeypatch, mixed): torch.testing.assert_close(fallback_scores, scores, atol=0, rtol=2**-7) # An admitted domain wider than the candidate pool changes nothing: only # the published candidate blocks are scored. - state.metadata.csa2_global_max_positions[0] = 257 + state.metadata.ext(CSA2Metadata).csa2_global_max_positions[0] = 257 wider, wider_scores, wider_calls = _candidate_decode_run(indexer, state, monkeypatch, old=False) assert wider_calls == calls _check_candidate_multiset(wider, scores, candidate, _bf16_tolerance(scores)) @@ -2052,12 +2118,14 @@ def test_candidate_decode_batched_graph_refresh(monkeypatch): with torch.cuda.graph(graph): actual = indexer(state, 0, 3) pointers = ( - metadata.csa2_indexer_k_cache.data_ptr(), - metadata.csa2_indexer_block_table.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_k_cache.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_block_table.data_ptr(), ) for step, lengths in enumerate(([192, 193, 190], [0, 0, 0], [129, 191, 192])): - metadata.csa2_visible_lengths[1].copy_(torch.tensor(lengths, device="cuda")) - metadata.csa2_global_page_tables[0][0, 1] = -1 if step == 0 else 0 + metadata.ext(CSA2Metadata).csa2_visible_lengths[1].copy_( + torch.tensor(lengths, device="cuda") + ) + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 1] = -1 if step == 0 else 0 # A source republishes candidates for the new visibility/pages every step _publish_manual_candidates( metadata, 0, _MANUAL_BLOCKS[step:] + [40 + i for i in range(step)] @@ -2067,11 +2135,11 @@ def test_candidate_decode_batched_graph_refresh(monkeypatch): graph.replay() torch.cuda.synchronize() _check_candidate_multiset( - actual, scores, metadata.csa2_candidates[0], _bf16_tolerance(scores) + actual, scores, metadata.ext(CSA2Metadata).csa2_candidates[0], _bf16_tolerance(scores) ) assert pointers == ( - metadata.csa2_indexer_k_cache.data_ptr(), - metadata.csa2_indexer_block_table.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_k_cache.data_ptr(), + metadata.ext(CSA2Metadata).csa2_indexer_block_table.data_ptr(), ) @@ -2098,7 +2166,10 @@ def test_candidate_batched_random_logits(monkeypatch, phase): # Existing native selection tests allow only 0.005 absolute cutoff error. torch.testing.assert_close(new_scores[valid], scores[valid], atol=0.005, rtol=2**-7) _check_candidate_multiset( - actual, scores, state.metadata.csa2_candidates[0], 0.005 + _bf16_tolerance(scores) + actual, + scores, + state.metadata.ext(CSA2Metadata).csa2_candidates[0], + 0.005 + _bf16_tolerance(scores), ) separated = gaps > 2 * difference.abs().amax(1) torch.testing.assert_close(actual[separated], old[separated], atol=0, rtol=0) @@ -2113,23 +2184,25 @@ def _candidate_prefill_fixture(*, random=False): state.index_weights = state.index_weights[rows] state.swa_kv = torch.zeros(count, 512, dtype=torch.bfloat16, device="cuda") metadata.is_cuda_graph = False - metadata.csa2_num_context_requests = 2 - metadata.csa2_request_query_ranges = [(0, 137), (137, 204), (204, 206)] - metadata.csa2_request_start_positions = [32, 144, 192] - metadata.csa2_request_lengths = [137, 67, 2] - metadata.csa2_token_requests = torch.tensor([0] * 137 + [1] * 67 + [2] * 2, device="cuda") - metadata.csa2_global_page_tables = { + metadata.ext(CSA2Metadata).csa2_num_context_requests = 2 + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, 137), (137, 204), (204, 206)] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [32, 144, 192] + metadata.ext(CSA2Metadata).csa2_request_lengths = [137, 67, 2] + metadata.ext(CSA2Metadata).csa2_token_requests = torch.tensor( + [0] * 137 + [1] * 67 + [2] * 2, device="cuda" + ) + metadata.ext(CSA2Metadata).csa2_global_page_tables = { 0: torch.tensor([[1, 0], [2, 3], [0, 1]], device="cuda", dtype=torch.int32) } - metadata.csa2_visible_lengths[1] = torch.cat( + metadata.ext(CSA2Metadata).csa2_visible_lengths[1] = torch.cat( ( torch.arange(33, 170, device="cuda"), torch.arange(145, 212, device="cuda"), torch.tensor([193, 194], device="cuda"), ) ) - metadata.csa2_visible_lengths[1][5] = 0 - metadata.csa2_global_page_tables[0][0, 1] = -1 + metadata.ext(CSA2Metadata).csa2_visible_lengths[1][5] = 0 + metadata.ext(CSA2Metadata).csa2_global_page_tables[0][0, 1] = -1 _publish_manual_candidates(metadata, 0, _MANUAL_BLOCKS) return indexer, state @@ -2163,7 +2236,7 @@ def test_candidate_shared_prefill_exact_and_mixed(monkeypatch): if not _sparse_kernels_available(): pytest.skip("DeepGEMM sparse MQA logits required") indexer, state = _candidate_prefill_fixture() - candidates = state.metadata.csa2_candidates[0].clone() + candidates = state.metadata.ext(CSA2Metadata).csa2_candidates[0].clone() old, scores, old_calls, _, old_gathers = _candidate_prefill_run( indexer, state, monkeypatch, True ) @@ -2177,7 +2250,9 @@ def test_candidate_shared_prefill_exact_and_mixed(monkeypatch): assert gathers == [169, 211] torch.testing.assert_close(new_scores, scores, atol=0, rtol=2**-7) _check_candidate_multiset(actual, scores, candidates, _bf16_tolerance(scores)) - torch.testing.assert_close(state.metadata.csa2_candidates[0], candidates, atol=0, rtol=0) + torch.testing.assert_close( + state.metadata.ext(CSA2Metadata).csa2_candidates[0], candidates, atol=0, rtol=0 + ) @pytest.mark.cpu_only @@ -2442,29 +2517,33 @@ def test_sparse_candidate_paged_decode_and_graph(monkeypatch): manager = _StridedOwner(layout, packed_keys) manager.max_total_draft_tokens = 1 q, _, weights = _projected(3) - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=3) + metadata = CSA2TrtllmAttention.create_metadata(max_num_requests=2, max_num_tokens=3) metadata.kv_cache_manager = manager metadata.mapping = Mapping() metadata.is_cuda_graph = True - metadata.csa2_kv_sources = {0: 0, 1: 0} - metadata.csa2_num_context_requests = 0 - metadata.csa2_request_query_ranges = [(0, 2), (2, 3)] - metadata.csa2_request_start_positions = [125, 127] - metadata.csa2_request_lengths = [2, 1] - metadata.csa2_global_max_positions = {0: 128} - metadata.csa2_global_page_sizes = {0: 128} + metadata.ext(CSA2Metadata).csa2_kv_sources = {0: 0, 1: 0} + metadata.ext(CSA2Metadata).csa2_num_context_requests = 0 + metadata.ext(CSA2Metadata).csa2_request_query_ranges = [(0, 2), (2, 3)] + metadata.ext(CSA2Metadata).csa2_request_start_positions = [125, 127] + metadata.ext(CSA2Metadata).csa2_request_lengths = [2, 1] + metadata.ext(CSA2Metadata).csa2_global_max_positions = {0: 128} + metadata.ext(CSA2Metadata).csa2_global_page_sizes = {0: 128} table = torch.tensor([[0, 1], [2, 3]], device="cuda", dtype=torch.int32) visible = torch.tensor([63, 63, 64], device="cuda", dtype=torch.int32) - metadata.csa2_global_page_tables = {0: table} - metadata.csa2_visible_lengths = {0: visible, 1: visible} - metadata.csa2_token_requests = torch.tensor([0, 0, 1], device="cuda", dtype=torch.int64) - metadata.csa2_decode_row_requests = torch.tensor([0, 0, 1], device="cuda", dtype=torch.int32) - metadata.csa2_indices = {} - metadata.csa2_candidates = {} - metadata.csa2_candidate_blocks = {} - metadata.csa2_candidate_counts = {} - metadata.indexer_max_chunk_size = 16 - metadata.indexer_q_split_threshold = -1 + metadata.ext(CSA2Metadata).csa2_global_page_tables = {0: table} + metadata.ext(CSA2Metadata).csa2_visible_lengths = {0: visible, 1: visible} + metadata.ext(CSA2Metadata).csa2_token_requests = torch.tensor( + [0, 0, 1], device="cuda", dtype=torch.int64 + ) + metadata.ext(CSA2Metadata).csa2_decode_row_requests = torch.tensor( + [0, 0, 1], device="cuda", dtype=torch.int32 + ) + metadata.ext(CSA2Metadata).csa2_indices = {} + metadata.ext(CSA2Metadata).csa2_candidates = {} + metadata.ext(CSA2Metadata).csa2_candidate_blocks = {} + metadata.ext(CSA2Metadata).csa2_candidate_counts = {} + metadata.ext(CSA2Metadata).indexer_max_chunk_size = 16 + metadata.ext(CSA2Metadata).indexer_q_split_threshold = -1 state = CSA2ForwardState( metadata=metadata, swa_kv=torch.zeros(3, 512, device="cuda", dtype=torch.bfloat16), @@ -2489,19 +2568,23 @@ def sparse(*args, **kwargs): ) def forward(): - metadata._csa2_forward_serial = getattr(metadata, "_csa2_forward_serial", 0) + 1 + metadata.ext(CSA2Metadata)._csa2_forward_serial = ( + getattr(metadata.ext(CSA2Metadata), "_csa2_forward_serial", 0) + 1 + ) source(state, 0, 3) return consumer(state, 0, 3) for _ in range(2): output = forward() assert sparse_calls and all(shape == (3, 32, 64) for shape in sparse_calls) - descriptors = metadata.prepare_sparse_indexer(1, consumer.sparse_block) + descriptors = metadata.ext(CSA2Metadata).prepare_sparse_indexer(1, consumer.sparse_block) assert descriptors is not None pages, block_table, context_lens, rows, schedule = descriptors assert schedule is not None and schedule.is_cuda # The schedule is shared by every consumer layer of a forward - assert metadata.prepare_sparse_indexer(1, consumer.sparse_block)[4] is schedule + assert ( + metadata.ext(CSA2Metadata).prepare_sparse_indexer(1, consumer.sparse_block)[4] is schedule + ) assert pages.shape[1:] == (64, 1, 68) assert block_table.shape[0] == 3 and context_lens.shape == (3,) assert rows.tolist() == [0, 0, 1] @@ -2519,9 +2602,9 @@ def forward(): graph.replay() torch.cuda.synchronize() selected = output.clone() - candidates = metadata.csa2_candidates[0].clone() + candidates = metadata.ext(CSA2Metadata).csa2_candidates[0].clone() positions = torch.arange(128, device="cuda").expand(3, -1) - req = metadata.csa2_token_requests + req = metadata.ext(CSA2Metadata).csa2_token_requests physical = table[req[:, None], positions // 128] valid = (physical >= 0) & (positions < visible[:, None]) rows_ = read_index_rows( @@ -2553,35 +2636,37 @@ def forward(): source = CSA2Indexer(layout, 0, 32, 128) consumer = CSA2Indexer(layout, 1, 32, 128) metadata.is_cuda_graph = False - metadata.csa2_global_max_positions[0] = 8192 - metadata.csa2_request_start_positions = [3246, 0] + metadata.ext(CSA2Metadata).csa2_global_max_positions[0] = 8192 + metadata.ext(CSA2Metadata).csa2_request_start_positions = [3246, 0] table = torch.full((2, 64), -1, dtype=torch.int32, device="cuda") table[:, 0] = torch.tensor([0, 2], device="cuda") table[:, 25] = torch.tensor([1, 3], device="cuda") - metadata.csa2_global_page_tables[0] = table + metadata.ext(CSA2Metadata).csa2_global_page_tables[0] = table visible.copy_(torch.tensor([3247, 3248, 0], dtype=torch.int32, device="cuda")) def check_holey(actual: torch.Tensor, counts: list[int]) -> None: - assert metadata.csa2_candidate_counts[0].flatten().tolist() == counts + assert metadata.ext(CSA2Metadata).csa2_candidate_counts[0].flatten().tolist() == counts _, scores, _ = _candidate_decode_run(consumer, state, monkeypatch, old=True) _check_candidate_multiset( - actual, scores, metadata.csa2_candidates[0], _bf16_tolerance(scores) + actual, scores, metadata.ext(CSA2Metadata).csa2_candidates[0], _bf16_tolerance(scores) ) for num_contexts, decode_start in ((0, 0), (1, 2), (2, 3)): - metadata.csa2_num_context_requests = num_contexts - metadata.csa2_decode_row_requests = ( - metadata.csa2_token_requests[decode_start:].int() - num_contexts + metadata.ext(CSA2Metadata).csa2_num_context_requests = num_contexts + metadata.ext(CSA2Metadata).csa2_decode_row_requests = ( + metadata.ext(CSA2Metadata).csa2_token_requests[decode_start:].int() - num_contexts ) actual = forward().clone() torch.cuda.synchronize() check_holey(actual, [175, 176, 0]) if num_contexts < 2: - counts = metadata.prepare_sparse_indexer(1, consumer.sparse_block)[2] + counts = metadata.ext(CSA2Metadata).prepare_sparse_indexer(1, consumer.sparse_block)[2] assert counts.tolist() == [175, 176, 0][decode_start:] - metadata.csa2_num_context_requests = 0 - metadata.csa2_decode_row_requests = metadata.csa2_token_requests.int() + metadata.ext(CSA2Metadata).csa2_num_context_requests = 0 + metadata.ext(CSA2Metadata).csa2_decode_row_requests = metadata.ext( + CSA2Metadata + ).csa2_token_requests.int() metadata.is_cuda_graph = True for _ in range(2): forward() diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_indexer_descriptors.py b/tests/unittest/_torch/attention/sparse/csa2/test_indexer_descriptors.py index 94aea2804df4..85112a9ca473 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_indexer_descriptors.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_indexer_descriptors.py @@ -12,11 +12,12 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import INDEX_PAGE_ROWS -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from . import test_metadata +from ._utils import _bare_csa2_metadata from .test_metadata import _prepare_decode_metadata skip_pre_blackwell = pytest.mark.skipif( @@ -153,8 +154,8 @@ def test_indexer_descriptors_follow_live_graph_inputs( ) -> None: """Replay changed contents against eager and scalar references with exact integer parity.""" monkeypatch.setattr("tensorrt_llm._torch.utils.is_piecewise_running", lambda: False) - eager = CSA2TrtllmMetadata._fill_indexer_descriptors - captured = CSA2TrtllmMetadata._fill_indexer_descriptors + eager = CSA2Metadata._fill_indexer_descriptors + captured = CSA2Metadata._fill_indexer_descriptors host_inputs = _inputs(query_count, request_count, source_pages, page_capacity, 0) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -203,16 +204,13 @@ def test_indexer_descriptors_follow_live_graph_inputs( @pytest.mark.cpu_only -def test_graph_clone_post_init_clears_decode_deferral_state( - monkeypatch: pytest.MonkeyPatch, -) -> None: - metadata = object.__new__(CSA2TrtllmMetadata) - metadata._csa2_deferred_decode_outputs = True +def test_graph_clone_post_init_clears_decode_deferral_state() -> None: + metadata = _bare_csa2_metadata() + metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs = True clone = copy(metadata) - monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) - clone.__post_init__() - assert not clone._csa2_deferred_decode_outputs - assert metadata._csa2_deferred_decode_outputs + clone._attach_extensions() + assert not clone.ext(CSA2Metadata)._csa2_deferred_decode_outputs + assert metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs @pytest.mark.cpu_only @@ -220,8 +218,8 @@ def test_graph_clone_post_init_clears_decode_deferral_state( def test_failed_prepare_clears_decode_deferral_state( monkeypatch: pytest.MonkeyPatch, entrypoint: str ) -> None: - metadata = object.__new__(CSA2TrtllmMetadata) - metadata._csa2_deferred_decode_outputs = True + metadata = _bare_csa2_metadata() + metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs = True metadata.kv_cache_manager = SimpleNamespace() metadata.is_cuda_graph = False metadata.kv_lens_cuda = torch.empty(0) @@ -234,8 +232,9 @@ def fail_prepare() -> None: error = ValueError if entrypoint == "prepare" else TypeError message = "injected prepare failure" if entrypoint == "prepare" else "CSA2CacheManager" with pytest.raises(error, match=message): - getattr(metadata, entrypoint)() - assert not metadata._csa2_deferred_decode_outputs + target = metadata if entrypoint == "prepare" else metadata.ext(CSA2Metadata) + getattr(target, entrypoint)() + assert not metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs @skip_pre_blackwell @@ -246,23 +245,25 @@ def test_selected_decode_prepare_controls_indexer_dispatch( monkeypatch: pytest.MonkeyPatch, ) -> None: manager, requests = manager_requests - base = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + base = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) # A shallow graph clone must discard any eligibility inherited from its source. - base._csa2_deferred_decode_outputs = True + base.ext(CSA2Metadata)._csa2_deferred_decode_outputs = True metadata = base.create_cuda_graph_metadata(2) - assert not metadata._csa2_deferred_decode_outputs - assert base._csa2_deferred_decode_outputs + assert not metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs + assert base.ext(CSA2Metadata)._csa2_deferred_decode_outputs starts = [127, 128] calls = [] - eager = CSA2TrtllmMetadata._fill_indexer_descriptors + eager = CSA2Metadata._fill_indexer_descriptors def record_refresh(*args: torch.Tensor | int) -> None: - calls.append(metadata._csa2_deferred_decode_outputs) + calls.append(metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs) return eager(*args) # The same fused dispatcher serves every preparation mode. Graph parity # is checked above; isolate per-owner reuse and preparation state here. - monkeypatch.setattr(metadata, "_fill_indexer_descriptors", record_refresh) + monkeypatch.setattr(metadata.ext(CSA2Metadata), "_fill_indexer_descriptors", record_refresh) monkeypatch.setattr("tensorrt_llm._torch.utils.is_piecewise_running", lambda: True) monkeypatch.setattr( "tensorrt_llm.deep_gemm.get_paged_mqa_logits_metadata", @@ -272,28 +273,28 @@ def record_refresh(*args: torch.Tensor | int) -> None: ) _prepare_decode_metadata(metadata, requests, starts, deferred=True) - assert metadata._csa2_deferred_decode_outputs - assert not metadata._csa2_defer_decode_outputs + assert metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs + assert not metadata.ext(CSA2Metadata)._csa2_defer_decode_outputs metadata.on_update_kv_lens() - assert metadata.prepare_indexer(1) is metadata + assert metadata.ext(CSA2Metadata).prepare_indexer(1) is metadata.ext(CSA2Metadata) assert calls == [True] # Layers 1 and 2 share an owner. One serial refresh serves both layers. - metadata.prepare_indexer(1) - metadata.prepare_indexer(2) + metadata.ext(CSA2Metadata).prepare_indexer(1) + metadata.ext(CSA2Metadata).prepare_indexer(2) assert calls == [True] _prepare_decode_metadata(metadata, requests, starts) - assert not metadata._csa2_deferred_decode_outputs + assert not metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs metadata.on_update_kv_lens() - metadata.prepare_indexer(1) + metadata.ext(CSA2Metadata).prepare_indexer(1) assert calls == [True, False] _prepare_decode_metadata(metadata, requests, starts, deferred=True) - assert metadata._csa2_deferred_decode_outputs - metadata.set_source_batch([1, 1], starts) + assert metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs + metadata.ext(CSA2Metadata).set_source_batch([1, 1], starts) _prepare_decode_metadata(metadata, requests, starts, deferred=True) - assert not hasattr(metadata, "_csa2_source_batch") - assert not metadata._csa2_deferred_decode_outputs + assert not hasattr(metadata.ext(CSA2Metadata), "_csa2_source_batch") + assert not metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs metadata.on_update_kv_lens() - metadata.prepare_indexer(1) + metadata.ext(CSA2Metadata).prepare_indexer(1) assert calls == [True, False, False] diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_kernel.py b/tests/unittest/_torch/attention/sparse/csa2/test_kernel.py index 71a1f1803c1d..6a367794b824 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_kernel.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_kernel.py @@ -14,6 +14,7 @@ packed_sparse_attention, quantize_scatter_rows, ) +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( INDEX_PAGE_BYTES, INDEX_PAGE_ROWS, @@ -673,12 +674,11 @@ def test_quantization_cuda_graph_changed_values(): def _stage_fixture(queries, swa_width, main_width, *, empty_pool=False, extra_width=None): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2BackendForwardArgs torch.manual_seed(711) q = torch.empty(queries, 1, 512, dtype=torch.bfloat16, device="cuda") - metadata = CSA2TrtllmMetadata.for_query_tile( + metadata = CSA2Metadata.for_query_tile( q, (main_width or 0) if extra_width is None else extra_width ) rows = 0 if empty_pool else 23 @@ -726,7 +726,10 @@ def _stage_oracle(metadata, inputs): packed = rows.gather(1, order[..., None].expand(-1, -1, 512)) counts = valid.sum(1, dtype=torch.int32) packed = torch.where((positions < counts[:, None])[..., None], packed, 0) - primary, extra = torch.zeros_like(metadata.swa_pool), torch.zeros_like(metadata.extra_pool) + primary, extra = ( + torch.zeros_like(metadata.ext(CSA2Metadata).swa_pool), + torch.zeros_like(metadata.ext(CSA2Metadata).extra_pool), + ) primary[:, : min(rows.shape[1], 128)] = packed[:, :128] if rows.shape[1] > 128: extra[:, : rows.shape[1] - 128] = packed[:, 128:] @@ -744,13 +747,18 @@ def _stage_oracle(metadata, inputs): def _check_stage(metadata, expected): for actual, reference in zip( - (metadata.swa_pool, metadata.extra_pool, metadata.prepared_lens, metadata.prepared_indices), + ( + metadata.ext(CSA2Metadata).swa_pool, + metadata.ext(CSA2Metadata).extra_pool, + metadata.ext(CSA2Metadata).prepared_lens, + metadata.ext(CSA2Metadata).prepared_indices, + ), expected, ): if actual.dtype == torch.bfloat16: actual, reference = actual.view(torch.int16), reference.view(torch.int16) torch.testing.assert_close(actual, reference, atol=0, rtol=0) - assert metadata.prepared_counter.item() == 0 + assert metadata.ext(CSA2Metadata).prepared_counter.item() == 0 @skip_pre_blackwell @@ -782,10 +790,10 @@ def stage(*args, **kwargs): queries, swa_width, main_width, empty_pool=empty_pool, extra_width=extra_width ) expected = _stage_oracle(metadata, inputs) - metadata.swa_pool.fill_(13) - metadata.extra_pool.fill_(17) - metadata.prepared_counter.fill_(19) - metadata.stage_selected(inputs) + metadata.ext(CSA2Metadata).swa_pool.fill_(13) + metadata.ext(CSA2Metadata).extra_pool.fill_(17) + metadata.ext(CSA2Metadata).prepared_counter.fill_(19) + metadata.ext(CSA2Metadata).stage_selected(inputs) _check_stage(metadata, expected) assert bool(calls) == (metadata.num_sparse_topk <= 4096) @@ -817,27 +825,33 @@ def test_staging_group_reuse_validates_each_layers_own_slots(): } # The first layer of the group resolves every column, so it tags all of them. stage_selected_rows( - inputs.swa_pool, resolved, inputs.swa_pool, None, metadata, scratch, compact=True + inputs.swa_pool, + resolved, + inputs.swa_pool, + None, + metadata.ext(CSA2Metadata), + scratch, + compact=True, ) assert scratch["counts"].tolist() == [width] * queries # The later layer's private pool starts one row into a sentinel buffer, so the # row an unvalidated negative slot reaches decodes to +1.0 instead of zero. sentinel = torch.empty(rows + 1, row_bytes, dtype=torch.uint8, device="cuda") sentinel[:, :512], sentinel[:, 512:] = 56, 127 # E4M3 +1.0 with a 2**0 scale - metadata.swa_pool.fill_(13) - metadata.extra_pool.fill_(17) + metadata.ext(CSA2Metadata).swa_pool.fill_(13) + metadata.ext(CSA2Metadata).extra_pool.fill_(17) stage_selected_rows( sentinel[1:], torch.full((queries, width), -1, device="cuda", dtype=torch.int64), sentinel[1:], None, - metadata, + metadata.ext(CSA2Metadata), scratch, compact=False, ) # Every tagged column is empty for this layer, so every staged row stays zero. - assert torch.count_nonzero(metadata.swa_pool) == 0 - assert torch.count_nonzero(metadata.extra_pool) == 0 + assert torch.count_nonzero(metadata.ext(CSA2Metadata).swa_pool) == 0 + assert torch.count_nonzero(metadata.ext(CSA2Metadata).extra_pool) == 0 @skip_pre_blackwell @@ -845,26 +859,28 @@ def test_staging_group_reuse_validates_each_layers_own_slots(): def test_native_staging_graph_refresh(): _sm100() metadata, inputs = _stage_fixture(8, 128, 17) - metadata.stage_selected(inputs) + metadata.ext(CSA2Metadata).stage_selected(inputs) tensor_names = ("swa_pool", "extra_pool", "prepared_indices", "prepared_lens") - pointers = [getattr(metadata, name).data_ptr() for name in tensor_names] - retained = metadata.get_workspace_bytes() + pointers = [getattr(metadata.ext(CSA2Metadata), name).data_ptr() for name in tensor_names] + retained = metadata.ext(CSA2Metadata).get_workspace_bytes() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - metadata.stage_selected(inputs) + metadata.ext(CSA2Metadata).stage_selected(inputs) for selection, payload in ((0, 56), (-1, 128), (2, 32), (2, 184)): inputs.swa_indices.fill_(selection) inputs.topk_indices.fill_(selection) inputs.swa_pool[2, :512].fill_(payload) inputs.main_pool[2, :256].fill_(payload) expected = _stage_oracle(metadata, inputs) - metadata.swa_pool.fill_(13) - metadata.extra_pool.fill_(17) - metadata.prepared_counter.fill_(19) + metadata.ext(CSA2Metadata).swa_pool.fill_(13) + metadata.ext(CSA2Metadata).extra_pool.fill_(17) + metadata.ext(CSA2Metadata).prepared_counter.fill_(19) graph.replay() _check_stage(metadata, expected) - assert pointers == [getattr(metadata, name).data_ptr() for name in tensor_names] - assert retained == metadata.get_workspace_bytes() + assert pointers == [ + getattr(metadata.ext(CSA2Metadata), name).data_ptr() for name in tensor_names + ] + assert retained == metadata.ext(CSA2Metadata).get_workspace_bytes() @skip_pre_blackwell @@ -876,10 +892,10 @@ def test_native_staging_validation_and_fallback(monkeypatch): metadata, inputs = _stage_fixture(1, 17, None) # An unused malformed main pool is ignored; a supplied empty selection is validated. inputs.main_pool = torch.zeros(1, device="cuda") - metadata.stage_selected(inputs) + metadata.ext(CSA2Metadata).stage_selected(inputs) inputs.topk_indices = inputs.swa_indices[:, :0] with pytest.raises(ValueError, match="wrong row shape or dtype"): - metadata.stage_selected(inputs) + metadata.ext(CSA2Metadata).stage_selected(inputs) inputs.topk_indices = None original_slots = inputs.swa_indices backing = torch.empty(inputs.swa_pool.shape[0], 1056, dtype=torch.uint8, device="cuda") @@ -889,14 +905,14 @@ def test_native_staging_validation_and_fallback(monkeypatch): monkeypatch.setattr( kernel, "stage_selected_rows", lambda *args: pytest.fail("Expected fallback") ) - metadata.stage_selected(inputs) + metadata.ext(CSA2Metadata).stage_selected(inputs) _check_stage(metadata, expected) # The unsupported zero-query direct invocation still resets the counter. metadata._num_tokens = 0 inputs.swa_indices = original_slots[:0] - metadata.prepared_counter.fill_(19) - metadata.stage_selected(inputs) - assert metadata.prepared_counter.item() == 0 + metadata.ext(CSA2Metadata).prepared_counter.fill_(19) + metadata.ext(CSA2Metadata).stage_selected(inputs) + assert metadata.ext(CSA2Metadata).prepared_counter.item() == 0 @skip_pre_blackwell @@ -931,7 +947,6 @@ def test_fused_store_changed_row_counts(cache_format, dim): @torch.inference_mode() def test_native_fp8_staging_exact_bytes(dequant, q_dequant, monkeypatch): from tensorrt_llm._torch.attention.backends.sparse.csa2 import kernel - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() control, inputs = _stage_fixture(3, 128, 17) @@ -951,7 +966,7 @@ def test_native_fp8_staging_exact_bytes(dequant, q_dequant, monkeypatch): inputs.topk_indices[2, :2] = 2 expected = _stage_oracle(control, inputs) q = torch.zeros(3, 64, 512, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) + metadata = CSA2Metadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) dq = torch.tensor([dequant], device="cuda", dtype=torch.float32) q_dq = torch.tensor([q_dequant], device="cuda", dtype=torch.float32) multiplier = dq.reciprocal() @@ -963,10 +978,10 @@ def stage(*args, **kwargs): return original(*args, **kwargs) monkeypatch.setattr(kernel, "stage_selected_rows", stage) - metadata.prepared_counter.fill_(19) - metadata.mla_bmm1_scale.fill_(float("nan")) - metadata.mla_bmm2_scale.fill_(float("nan")) - metadata.stage_selected( + metadata.ext(CSA2Metadata).prepared_counter.fill_(19) + metadata.ext(CSA2Metadata).mla_bmm1_scale.fill_(float("nan")) + metadata.ext(CSA2Metadata).mla_bmm2_scale.fill_(float("nan")) + metadata.ext(CSA2Metadata).stage_selected( inputs, kv_scale_orig_quant=multiplier, kv_scale_quant_orig=dq, @@ -974,23 +989,29 @@ def stage(*args, **kwargs): softmax_scale=512**-0.5, ) assert calls == [True] - for actual, decoded in zip((metadata.swa_pool, metadata.extra_pool), expected[:2]): + for actual, decoded in zip( + (metadata.ext(CSA2Metadata).swa_pool, metadata.ext(CSA2Metadata).extra_pool), expected[:2] + ): quantized, _ = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(decoded, dq) torch.testing.assert_close( actual.view(torch.uint8), quantized.view(torch.uint8), atol=0, rtol=0 ) - torch.testing.assert_close(metadata.prepared_lens, expected[2], atol=0, rtol=0) - torch.testing.assert_close(metadata.prepared_indices, expected[3], atol=0, rtol=0) - assert metadata.prepared_counter.item() == 0 + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_lens, expected[2], atol=0, rtol=0 + ) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_indices, expected[3], atol=0, rtol=0 + ) + assert metadata.ext(CSA2Metadata).prepared_counter.item() == 0 # BMM1 folds the two independent dequant scales, not one shared scale. scale = q_dequant * dequant * 512**-0.5 torch.testing.assert_close( - metadata.mla_bmm1_scale, + metadata.ext(CSA2Metadata).mla_bmm1_scale, torch.tensor([scale, scale * 1.4426950408889634], device="cuda"), atol=1e-7, rtol=1e-6, ) - torch.testing.assert_close(metadata.mla_bmm2_scale, dq, atol=0, rtol=0) + torch.testing.assert_close(metadata.ext(CSA2Metadata).mla_bmm2_scale, dq, atol=0, rtol=0) def _stage_scale_reference(inputs): @@ -1030,7 +1051,6 @@ def _stage_scale_reference(inputs): @torch.inference_mode() def test_stage_kv_scale_derivation(case, minimum): """The derived staging scale covers every selected row without saturating.""" - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() _, inputs = _stage_fixture(3, 128, 17) @@ -1052,9 +1072,11 @@ def test_stage_kv_scale_derivation(case, minimum): inputs.swa_pool[2, 512:].fill_(254) # 2**127, past every clamp inputs.swa_indices[2, 0] = 2 q = torch.zeros(3, 64, 512, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) - quant, dequant = metadata.derive_stage_kv_scales(inputs, minimum=minimum, maximum=126) - assert (quant, dequant) == metadata.stage_kv_scales + metadata = CSA2Metadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) + quant, dequant = metadata.ext(CSA2Metadata).derive_stage_kv_scales( + inputs, minimum=minimum, maximum=126 + ) + assert (quant, dequant) == metadata.ext(CSA2Metadata).stage_kv_scales scale = dequant.item() # A power of two costs no relative precision: it only moves the exponent. assert scale > 0.0 and math.log2(scale).is_integer() @@ -1093,7 +1115,6 @@ def test_stage_kv_scale_respects_q_headroom(amax, expected): its codes land at ``scale * amax(|Q|)``: a scale wide enough to rescue the KV rows must still leave Q inside the E4M3 range. """ - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() _, inputs = _stage_fixture(3, 128, 17) @@ -1101,8 +1122,8 @@ def test_stage_kv_scale_respects_q_headroom(amax, expected): inputs.swa_pool[2, 512:].fill_(254) inputs.swa_indices[2, 0] = 2 q = torch.full((3, 64, 512), amax, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) - quant, dequant = metadata.derive_stage_kv_scales(inputs, q=q) + metadata = CSA2Metadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) + quant, dequant = metadata.ext(CSA2Metadata).derive_stage_kv_scales(inputs, q=q) assert dequant.item() == expected assert quant.item() == 1.0 / expected if 0.0 < amax <= 448.0: # A Q that saturates on its own bounds nothing useful. @@ -1114,7 +1135,6 @@ def test_stage_kv_scale_respects_q_headroom(amax, expected): @torch.inference_mode() def test_stage_kv_scale_rescues_saturated_rows(): """Rows decoding past 448 survive staging instead of clamping at the E4M3 max.""" - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() control, inputs = _stage_fixture(4, 128, 17) @@ -1130,15 +1150,26 @@ def test_stage_kv_scale_rescues_saturated_rows(): ) expected = _stage_oracle(control, inputs) q = torch.zeros(4, 64, 512, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) - unit = metadata.native_unit_scale - metadata.stage_selected(inputs, kv_scale_orig_quant=unit, kv_scale_quant_orig=unit) - clamped = [pool.float().clone() for pool in (metadata.swa_pool, metadata.extra_pool)] - quant, dequant = metadata.derive_stage_kv_scales(inputs) + metadata = CSA2Metadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) + unit = metadata.ext(CSA2Metadata).native_unit_scale + metadata.ext(CSA2Metadata).stage_selected( + inputs, kv_scale_orig_quant=unit, kv_scale_quant_orig=unit + ) + clamped = [ + pool.float().clone() + for pool in (metadata.ext(CSA2Metadata).swa_pool, metadata.ext(CSA2Metadata).extra_pool) + ] + quant, dequant = metadata.ext(CSA2Metadata).derive_stage_kv_scales(inputs) assert dequant.item() > 1.0 - metadata.stage_selected(inputs, kv_scale_orig_quant=quant, kv_scale_quant_orig=dequant) + metadata.ext(CSA2Metadata).stage_selected( + inputs, kv_scale_orig_quant=quant, kv_scale_quant_orig=dequant + ) saturated = 0 - for pool, before, reference in zip((metadata.swa_pool, metadata.extra_pool), clamped, expected): + for pool, before, reference in zip( + (metadata.ext(CSA2Metadata).swa_pool, metadata.ext(CSA2Metadata).extra_pool), + clamped, + expected, + ): over = reference.float().abs() > 448.0 saturated += int(over.sum()) # A unit scale clamped these rows at 448, losing every bit of magnitude. @@ -1171,7 +1202,6 @@ def test_stage_kv_scale_respects_q_resolution(amax, expected): bound therefore holds Q's own amax at unit magnitude or above, which leaves it the whole E4M3 normal range below one plus the subnormals under that. """ - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() _, inputs = _stage_fixture(3, 128, 17) @@ -1179,8 +1209,8 @@ def test_stage_kv_scale_respects_q_resolution(amax, expected): inputs.swa_pool[2, 512:].fill_(254) inputs.swa_indices[2, 0] = 2 q = torch.full((3, 64, 512), amax, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) - quant, dequant = metadata.derive_stage_kv_scales(inputs, q=q, q_shared=True) + metadata = CSA2Metadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) + quant, dequant = metadata.ext(CSA2Metadata).derive_stage_kv_scales(inputs, q=q, q_shared=True) assert dequant.item() == expected assert quant.item() == 1.0 / expected # A zero Q has no magnitude to keep; otherwise Q stays at unit magnitude, or @@ -1203,7 +1233,6 @@ def test_stage_kv_scale_rescues_context_rows(): from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import ( _CONTEXT_STAGE_SCALE_MAX_EXPONENT, ) - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() control, inputs = _stage_fixture(4, 128, 17) @@ -1222,20 +1251,31 @@ def test_stage_kv_scale_rescues_context_rows(): inputs.main_pool[0, 256:].fill_(126) expected = _stage_oracle(control, inputs) q = torch.full((4, 64, 512), 64.0, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) - unit = metadata.native_unit_scale - metadata.stage_selected(inputs, kv_scale_orig_quant=unit, kv_scale_quant_orig=unit) - clamped = [pool.float().clone() for pool in (metadata.swa_pool, metadata.extra_pool)] - quant, dequant = metadata.derive_stage_kv_scales( + metadata = CSA2Metadata.for_query_tile(q, 17, staging_dtype=torch.float8_e4m3fn) + unit = metadata.ext(CSA2Metadata).native_unit_scale + metadata.ext(CSA2Metadata).stage_selected( + inputs, kv_scale_orig_quant=unit, kv_scale_quant_orig=unit + ) + clamped = [ + pool.float().clone() + for pool in (metadata.ext(CSA2Metadata).swa_pool, metadata.ext(CSA2Metadata).extra_pool) + ] + quant, dequant = metadata.ext(CSA2Metadata).derive_stage_kv_scales( inputs, q=q, q_shared=True, maximum=_CONTEXT_STAGE_SCALE_MAX_EXPONENT ) # 2688 mapped onto 448 is three binades, which is exactly the cap: context # never has to clip what the main pool asks for. assert dequant.item() == 8.0 == 2.0**_CONTEXT_STAGE_SCALE_MAX_EXPONENT assert 64.0 / dequant.item() >= 1.0 # Q gave up resolution, not magnitude. - metadata.stage_selected(inputs, kv_scale_orig_quant=quant, kv_scale_quant_orig=dequant) + metadata.ext(CSA2Metadata).stage_selected( + inputs, kv_scale_orig_quant=quant, kv_scale_quant_orig=dequant + ) saturated = 0 - for pool, before, reference in zip((metadata.swa_pool, metadata.extra_pool), clamped, expected): + for pool, before, reference in zip( + (metadata.ext(CSA2Metadata).swa_pool, metadata.ext(CSA2Metadata).extra_pool), + clamped, + expected, + ): over = reference.float().abs() > 448.0 saturated += int(over.sum()) # A unit scale clamped these rows at 448, losing every bit of magnitude. @@ -1253,7 +1293,6 @@ def test_stage_kv_scale_rescues_context_rows(): def test_unfused_fp8_staging_matches_fused_bytes(monkeypatch): """The eager staging fallback reproduces the fused route byte for byte.""" from tensorrt_llm._torch.attention.backends.sparse.csa2 import kernel - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata _sm100() calls = [] @@ -1272,35 +1311,35 @@ def stage(*args, **kwargs): inputs.main_pool[1, 256:].fill_(126) calls.clear() q = torch.zeros(2, 64, 512, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile( - q, extra_width, staging_dtype=torch.float8_e4m3fn - ) - quant, dequant = metadata.derive_stage_kv_scales(inputs) + metadata = CSA2Metadata.for_query_tile(q, extra_width, staging_dtype=torch.float8_e4m3fn) + quant, dequant = metadata.ext(CSA2Metadata).derive_stage_kv_scales(inputs) assert dequant.item() == 8.0 # 2688 mapped onto 448. - metadata.mla_bmm1_scale.fill_(float("nan")) - metadata.mla_bmm2_scale.fill_(float("nan")) - metadata.stage_selected( + metadata.ext(CSA2Metadata).mla_bmm1_scale.fill_(float("nan")) + metadata.ext(CSA2Metadata).mla_bmm2_scale.fill_(float("nan")) + metadata.ext(CSA2Metadata).stage_selected( inputs, kv_scale_orig_quant=quant, kv_scale_quant_orig=dequant, - q_scale_quant_orig=metadata.native_unit_scale, + q_scale_quant_orig=metadata.ext(CSA2Metadata).native_unit_scale, softmax_scale=512**-0.5, ) assert bool(calls) == fused # The unfused route must publish the same folded BMM scales. scale = dequant * 512**-0.5 torch.testing.assert_close( - metadata.mla_bmm1_scale, + metadata.ext(CSA2Metadata).mla_bmm1_scale, torch.cat((scale, scale * 1.4426950408889634)), atol=0, rtol=1e-6, ) - torch.testing.assert_close(metadata.mla_bmm2_scale, dequant, atol=0, rtol=0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).mla_bmm2_scale, dequant, atol=0, rtol=0 + ) staged.append( ( - metadata.swa_pool.view(torch.uint8).clone(), - metadata.extra_pool[:, :17].view(torch.uint8).clone(), - metadata.prepared_lens.clone(), + metadata.ext(CSA2Metadata).swa_pool.view(torch.uint8).clone(), + metadata.ext(CSA2Metadata).extra_pool[:, :17].view(torch.uint8).clone(), + metadata.ext(CSA2Metadata).prepared_lens.clone(), ) ) for fused_tensor, unfused in zip(*staged): @@ -1347,8 +1386,7 @@ def strided(values): ) bank = torch.full((25, 512), 7, device="cuda", dtype=dtype) metadata = SimpleNamespace( - num_tokens=5, - num_sparse_topk=128, + metadata=SimpleNamespace(num_tokens=5, num_sparse_topk=128), shared_pool=bank, _csa2_shared_selected=torch.empty(25, device="cuda", dtype=torch.int32), prepared_indices=torch.empty(5, 128, device="cuda", dtype=torch.int32), @@ -1357,6 +1395,7 @@ def strided(values): mla_bmm1_scale=torch.empty(2, device="cuda"), mla_bmm2_scale=torch.empty(1, device="cuda"), ) + metadata.ext = lambda cls: metadata unit = torch.ones(1, device="cuda") stage_shared_rows( swa, @@ -1371,20 +1410,24 @@ def strided(values): q_scale_quant_orig=unit, ) expected_ids = [[1, 3, 4, 17, 17, 19], [2, 3, 4, 5, 18, 20], [0], [10, 11, 12, 13, 22, 24], [0]] - expected_indices = torch.full_like(metadata.prepared_indices, -1) + expected_indices = torch.full_like(metadata.ext(CSA2Metadata).prepared_indices, -1) for row, ids in enumerate(expected_ids): expected_indices[row, : len(ids)] = torch.tensor(ids, device="cuda") - torch.testing.assert_close(metadata.prepared_indices, expected_indices, atol=0, rtol=0) torch.testing.assert_close( - metadata.prepared_lens, + metadata.ext(CSA2Metadata).prepared_indices, expected_indices, atol=0, rtol=0 + ) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_lens, torch.tensor([len(ids) for ids in expected_ids], device="cuda", dtype=torch.int32), atol=0, rtol=0, ) selected_rows = sorted({index for ids in expected_ids for index in ids if index}) - expected_mask = torch.zeros_like(metadata._csa2_shared_selected) + expected_mask = torch.zeros_like(metadata.ext(CSA2Metadata)._csa2_shared_selected) expected_mask[selected_rows] = 1 - torch.testing.assert_close(metadata._csa2_shared_selected, expected_mask, atol=0, rtol=0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata)._csa2_shared_selected, expected_mask, atol=0, rtol=0 + ) expected_bank = torch.full((25, 512), 7, device="cuda", dtype=torch.bfloat16) expected_bank[0].zero_() decoded_swa, decoded_main = unpack_rows(swa, 512, "swa"), unpack_rows(main, 512, "main") @@ -1422,17 +1465,19 @@ def strided(values): kv_scale_quant_orig=unit, q_scale_quant_orig=unit, ) - assert not metadata._csa2_shared_selected.any() - assert (metadata.prepared_lens == 1).all() - assert (metadata.prepared_indices[:, 0] == 0).all() - assert (metadata.prepared_indices[:, 1:] == -1).all() + assert not metadata.ext(CSA2Metadata)._csa2_shared_selected.any() + assert (metadata.ext(CSA2Metadata).prepared_lens == 1).all() + assert (metadata.ext(CSA2Metadata).prepared_indices[:, 0] == 0).all() + assert (metadata.ext(CSA2Metadata).prepared_indices[:, 1:] == -1).all() torch.testing.assert_close(bank.view(torch.uint8), previous.view(torch.uint8), atol=0, rtol=0) # Empty request domains and no MAIN pool use the same offset upper bound. plan.swa_offsets = torch.tensor([1, 1, 9], device="cuda", dtype=torch.int64) plan.main_offsets = torch.tensor([9, 9, 9], device="cuda", dtype=torch.int64) plan.swa_rows, plan.main_rows, plan.main_pages = 8, 0, None - metadata.shared_pool = torch.full((9, 512), 7, device="cuda", dtype=dtype) - metadata._csa2_shared_selected = torch.empty(9, device="cuda", dtype=torch.int32) + metadata.ext(CSA2Metadata).shared_pool = torch.full((9, 512), 7, device="cuda", dtype=dtype) + metadata.ext(CSA2Metadata)._csa2_shared_selected = torch.empty( + 9, device="cuda", dtype=torch.int32 + ) swa_slots.copy_(torch.tensor(swa_values, device="cuda")) stage_shared_rows( swa, @@ -1449,9 +1494,11 @@ def strided(values): expected_indices.fill_(-1) expected_indices[:, 0] = 0 expected_indices[3, :4] = torch.tensor([2, 3, 4, 5], device="cuda") - torch.testing.assert_close(metadata.prepared_indices, expected_indices, atol=0, rtol=0) - assert metadata._csa2_shared_selected.tolist() == [0, 0, 1, 1, 1, 1, 0, 0, 0] - assert metadata.prepared_lens.tolist() == [1, 1, 1, 4, 1] + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_indices, expected_indices, atol=0, rtol=0 + ) + assert metadata.ext(CSA2Metadata)._csa2_shared_selected.tolist() == [0, 0, 1, 1, 1, 1, 0, 0, 0] + assert metadata.ext(CSA2Metadata).prepared_lens.tolist() == [1, 1, 1, 4, 1] expected_bank = torch.full((9, 512), 7, device="cuda", dtype=torch.bfloat16) expected_bank[0].zero_() expected_bank[2:6].copy_(decoded_swa[37:41]) @@ -1460,7 +1507,10 @@ def strided(values): 0 ] torch.testing.assert_close( - metadata.shared_pool.view(torch.uint8), expected_bank.view(torch.uint8), atol=0, rtol=0 + metadata.ext(CSA2Metadata).shared_pool.view(torch.uint8), + expected_bank.view(torch.uint8), + atol=0, + rtol=0, ) @@ -1481,8 +1531,6 @@ def _glue(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_owner_refresh_kernel_matches_reference(): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata - kernel = _glue() torch.manual_seed(1201) request_count = 3 @@ -1526,7 +1574,7 @@ def test_owner_refresh_kernel_matches_reference(): for source_base, source_lengths, pages, slots, compressed, ratio, page_size, *batch in owners: outputs = [t for t in (slots, compressed, *batch) if t is not None] expected = [torch.empty_like(t) for t in outputs] - CSA2TrtllmMetadata._refresh_owner_slots( + CSA2Metadata._refresh_owner_slots( source_base, source_lengths, delta, @@ -1715,8 +1763,7 @@ def strided_pages(values: list[list[int]]) -> torch.Tensor: ) bank = torch.full((main_offsets[-1], 512), 7, device="cuda", dtype=dtype) metadata = SimpleNamespace( - num_tokens=len(query_requests), - num_sparse_topk=128, + metadata=SimpleNamespace(num_tokens=len(query_requests), num_sparse_topk=128), shared_pool=bank, _csa2_shared_selected=torch.empty(bank.shape[0], device="cuda", dtype=torch.int32), prepared_indices=torch.empty(len(query_requests), 128, device="cuda", dtype=torch.int32), @@ -1725,6 +1772,7 @@ def strided_pages(values: list[list[int]]) -> torch.Tensor: mla_bmm1_scale=torch.empty(2, device="cuda"), mla_bmm2_scale=torch.empty(1, device="cuda"), ) + metadata.ext = lambda cls: metadata dequant_scale = torch.tensor([0.3], device="cuda", dtype=torch.float32) stage_shared_rows( swa, @@ -1738,15 +1786,19 @@ def strided_pages(values: list[list[int]]) -> torch.Tensor: kv_scale_quant_orig=dequant_scale, q_scale_quant_orig=dequant_scale, ) - expected_indices = torch.full_like(metadata.prepared_indices, -1) + expected_indices = torch.full_like(metadata.ext(CSA2Metadata).prepared_indices, -1) for query, ids in enumerate(expected_ids): expected_indices[query, : len(ids)] = torch.tensor(ids, device="cuda") - torch.testing.assert_close(metadata.prepared_indices, expected_indices, atol=0, rtol=0) - assert metadata.prepared_lens.tolist() == [len(ids) for ids in expected_ids] + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_indices, expected_indices, atol=0, rtol=0 + ) + assert metadata.ext(CSA2Metadata).prepared_lens.tolist() == [len(ids) for ids in expected_ids] selected_rows = {row for ids in expected_ids for row in ids if row} - expected_selected = torch.zeros_like(metadata._csa2_shared_selected) + expected_selected = torch.zeros_like(metadata.ext(CSA2Metadata)._csa2_shared_selected) expected_selected[list(selected_rows)] = 1 - torch.testing.assert_close(metadata._csa2_shared_selected, expected_selected, atol=0, rtol=0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata)._csa2_shared_selected, expected_selected, atol=0, rtol=0 + ) # Build the numerical oracle request by request, without a flat-row owner lookup. expected_bank = torch.full_like(bank, 7) expected_bank[0].zero_() @@ -1821,8 +1873,7 @@ def test_shared_active_rows_reuse_request_count_specialization( ) bank = torch.empty(1 + 2 * domain_rows, 512, device="cuda", dtype=dtype) metadata = SimpleNamespace( - num_tokens=queries, - num_sparse_topk=128, + metadata=SimpleNamespace(num_tokens=queries, num_sparse_topk=128), shared_pool=bank, _csa2_shared_selected=torch.empty(bank.shape[0], device="cuda", dtype=torch.int32), prepared_indices=torch.empty(queries, 128, device="cuda", dtype=torch.int32), @@ -1831,6 +1882,7 @@ def test_shared_active_rows_reuse_request_count_specialization( mla_bmm1_scale=torch.empty(2, device="cuda"), mla_bmm2_scale=torch.empty(1, device="cuda"), ) + metadata.ext = lambda cls: metadata swa_slots = torch.empty(queries, 1, device="cuda", dtype=torch.int64) main_slots = torch.empty(queries, 2, device="cuda", dtype=torch.int64) main_logical = torch.empty_like(main_slots) @@ -1898,16 +1950,20 @@ def test_shared_active_rows_reuse_request_count_specialization( else: assert staging.keys() == compiled_staging.keys() assert all(value is compiled_staging[key] for key, value in staging.items()) - expected_indices = torch.full_like(metadata.prepared_indices, -1) + expected_indices = torch.full_like(metadata.ext(CSA2Metadata).prepared_indices, -1) for query, ids in enumerate(expected_ids): expected_indices[query, : len(ids)] = torch.tensor(ids, device="cuda") - torch.testing.assert_close(metadata.prepared_indices, expected_indices, atol=0, rtol=0) - assert metadata.prepared_lens.tolist() == [len(ids) for ids in expected_ids] + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepared_indices, expected_indices, atol=0, rtol=0 + ) + assert metadata.ext(CSA2Metadata).prepared_lens.tolist() == [ + len(ids) for ids in expected_ids + ] selected_rows = sorted(expected_rows) - expected_selected = torch.zeros_like(metadata._csa2_shared_selected) + expected_selected = torch.zeros_like(metadata.ext(CSA2Metadata)._csa2_shared_selected) expected_selected[selected_rows] = 1 torch.testing.assert_close( - metadata._csa2_shared_selected, expected_selected, atol=0, rtol=0 + metadata.ext(CSA2Metadata)._csa2_shared_selected, expected_selected, atol=0, rtol=0 ) expected_bank = torch.full_like(bank, 7) expected_bank[0].zero_() diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py b/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py index 2b9d71afc0d7..66d6801639ba 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py @@ -9,6 +9,8 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager from tensorrt_llm._torch.attention.backends.sparse.csa2.compressor import CSA2CompressionBatch from tensorrt_llm._torch.attention.backends.sparse.csa2.decoder_replay import ( @@ -17,17 +19,16 @@ exit_decoder_replay, plan_decoder_replay, ) -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import pack_rows -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping -from ._utils import _FakeMetadata, _page_metadata +from ._utils import _bare_csa2_metadata, _FakeMetadata, _page_metadata skip_pre_blackwell = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), @@ -41,7 +42,7 @@ def test_decoder_swa_slots_skip_reclaimed_encoder_pages(device): if device == "cuda" and not (torch.cuda.is_available() and kernel.dsl_available()): pytest.skip("CuTe DSL SWA refresh requires CUDA") - metadata = object.__new__(CSA2TrtllmMetadata) + metadata = _bare_csa2_metadata() metadata.is_cuda_graph = False metadata.kv_lens_cuda = torch.zeros(1, device=device) metadata.kv_cache_manager = SimpleNamespace( @@ -50,20 +51,30 @@ def test_decoder_swa_slots_skip_reclaimed_encoder_pages(device): ) # Encoder pages have already been released, so a decoder-only query pass # binds only the decoder layers' tables. - metadata.csa2_query_layer_start = 2 - metadata._csa2_buffers = {} - metadata._csa2_swa_page_tables = { + metadata.ext(CSA2Metadata).csa2_query_layer_start = 2 + metadata.ext(CSA2Metadata)._csa2_buffers = {} + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = { 2: torch.tensor([[2, 3]], dtype=torch.int32, device=device), 3: torch.tensor([[4, 5]], dtype=torch.int32, device=device), } - metadata.csa2_positions = torch.tensor([128, 129], dtype=torch.int32, device=device) - metadata.csa2_token_requests = torch.zeros(2, dtype=torch.int64, device=device) - metadata.csa2_replay_start_positions = torch.zeros(1, dtype=torch.int64, device=device) + metadata.ext(CSA2Metadata).csa2_positions = torch.tensor( + [128, 129], dtype=torch.int32, device=device + ) + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros( + 2, dtype=torch.int64, device=device + ) + metadata.ext(CSA2Metadata).csa2_replay_start_positions = torch.zeros( + 1, dtype=torch.int64, device=device + ) if device == "cuda": - metadata._csa2_swa_descriptors = metadata._swa_table_descriptors(4) - metadata._csa2_swa_ratios = metadata._swa_ratio_tensor((0, 2, 1, 1)) - metadata._csa2_swa_resolved = False - metadata._ensure_swa_slots() + metadata.ext(CSA2Metadata)._csa2_swa_descriptors = metadata.ext( + CSA2Metadata + )._swa_table_descriptors(4) + metadata.ext(CSA2Metadata)._csa2_swa_ratios = metadata.ext(CSA2Metadata)._swa_ratio_tensor( + (0, 2, 1, 1) + ) + metadata.ext(CSA2Metadata)._csa2_swa_resolved = False + metadata.ext(CSA2Metadata)._ensure_swa_slots() for layer, pages in ((2, (2, 3)), (3, (4, 5))): expected = torch.tensor( [ @@ -71,15 +82,19 @@ def test_decoder_swa_slots_skip_reclaimed_encoder_pages(device): for end in (128, 129) ] ) - torch.testing.assert_close(metadata.csa2_swa_indices[layer].cpu(), expected) - torch.testing.assert_close(metadata.csa2_swa_write_slots[layer].cpu(), expected[:, -1]) torch.testing.assert_close( - metadata.csa2_visible_lengths[layer].cpu(), torch.tensor([129, 130]) + metadata.ext(CSA2Metadata).csa2_swa_indices[layer].cpu(), expected + ) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_swa_write_slots[layer].cpu(), expected[:, -1] + ) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_visible_lengths[layer].cpu(), torch.tensor([129, 130]) ) # Layers before the boundary bind no pages and read or write nothing. for layer in (0, 1): - assert (metadata.csa2_swa_indices[layer] == -1).all() - assert (metadata.csa2_swa_write_slots[layer] == -1).all() + assert (metadata.ext(CSA2Metadata).csa2_swa_indices[layer] == -1).all() + assert (metadata.ext(CSA2Metadata).csa2_swa_write_slots[layer] == -1).all() @pytest.fixture @@ -157,12 +172,12 @@ def _prepare(metadata, requests, starts, lengths): def _swa(metadata, layer): """Resolve the SWA slots from the current inputs, as a forward does, and return one layer's.""" - metadata._csa2_swa_resolved = False - metadata._ensure_swa_slots() + metadata.ext(CSA2Metadata)._csa2_swa_resolved = False + metadata.ext(CSA2Metadata)._ensure_swa_slots() return ( - metadata.csa2_swa_indices[layer], - metadata.csa2_swa_write_slots[layer], - metadata.csa2_visible_lengths[layer], + metadata.ext(CSA2Metadata).csa2_swa_indices[layer], + metadata.ext(CSA2Metadata).csa2_swa_write_slots[layer], + metadata.ext(CSA2Metadata).csa2_visible_lengths[layer], ) @@ -171,10 +186,12 @@ def _swa(metadata, layer): @torch.inference_mode() def test_manager_metadata_scratch_and_source_capacity(manager_requests): from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheRole - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) _prepare(metadata, requests, [0, 0], [257, 257]) _assert_swa_cpu_reference(metadata, manager, requests, [0, 0], [257, 257], [0, 0]) reads, writes, _ = _swa(metadata, 1) @@ -183,30 +200,43 @@ def test_manager_metadata_scratch_and_source_capacity(manager_requests): # longer than the 128-token window. torch.testing.assert_close(reads[128], writes[1:129]) assert _swa(metadata, 0)[2].count_nonzero() == 0 - assert metadata.get_compression_batch(3) is None + assert metadata.ext(CSA2Metadata).get_compression_batch(3) is None assert manager.get_main_buffer(1).stride(0) == 288 assert manager.get_index_pages(1).shape[1:] == (64, 1, 68) - metadata.set_source_batch([5, 3], [0, 1]) + metadata.ext(CSA2Metadata).set_source_batch([5, 3], [0, 1]) _prepare(metadata, requests, [4, 3], [1, 1]) - compression = metadata.get_compression_batch(1) + compression = metadata.ext(CSA2Metadata).get_compression_batch(1) assert compression.output_rows == 8 assert compression.cu_seq_lengths.tolist() == [0, 5, 8] assert compression.cu_compressed_lengths.tolist() == [0, 2, 4] - assert metadata.get_compressed_positions(1).tolist() == [0, 2, 0, 2, 0, 0, 0, 0] - assert metadata.csa2_main_write_slots[1][4:].tolist() == [-1] * 4 + assert metadata.ext(CSA2Metadata).get_compressed_positions(1).tolist() == [ + 0, + 2, + 0, + 2, + 0, + 0, + 0, + 0, + ] + assert metadata.ext(CSA2Metadata).csa2_main_write_slots[1][4:].tolist() == [-1] * 4 _prepare(metadata, requests, [0, 1], [1, 1]) - assert metadata.get_compression_batch(1).cu_compressed_lengths.tolist() == [0, 0, 1] - completed = metadata.csa2_main_write_slots[1].tolist() + assert metadata.ext(CSA2Metadata).get_compression_batch(1).cu_compressed_lengths.tolist() == [ + 0, + 0, + 1, + ] + completed = metadata.ext(CSA2Metadata).csa2_main_write_slots[1].tolist() expected_page = manager.get_cache_indices(97, 1, CSA2CacheRole.GLOBAL)[0] assert completed == [expected_page * 64, -1] # Zero completed ratio-2 groups still prepare partial state and padded slots. _prepare(metadata, requests, [0, 0], [1, 1]) - compression = metadata.get_compression_batch(1) + compression = metadata.ext(CSA2Metadata).get_compression_batch(1) assert compression.output_rows == 2 assert compression.cu_seq_lengths.tolist() == [0, 1, 2] assert compression.cu_compressed_lengths.tolist() == [0, 0, 0] - assert metadata.csa2_main_write_slots[1].tolist() == [-1, -1] - assert metadata.get_compressed_positions(1).tolist() == [0, 0] + assert metadata.ext(CSA2Metadata).csa2_main_write_slots[1].tolist() == [-1, -1] + assert metadata.ext(CSA2Metadata).get_compressed_positions(1).tolist() == [0, 0] # Consumer and source resolve the same owner pages, without owning extra KV. assert manager.get_cache_indices(41, 1, CSA2CacheRole.GLOBAL) == manager.get_cache_indices( 41, 2, CSA2CacheRole.GLOBAL @@ -222,22 +252,24 @@ def test_swa_slots_resolve_once_per_forward(manager_requests, monkeypatch): if not kernel.dsl_available(): pytest.skip("Counting native resolutions requires the CuTe DSL") manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) launches = [] refresh = kernel.refresh_swa_slots monkeypatch.setattr(kernel, "refresh_swa_slots", lambda *a: launches.append(refresh(*a))) layers = [manager.layout.layer(i) for i in range(len(manager.layout.compress_ratios))] def forward(): - metadata.begin_model_forward() + metadata.begin_forward() for layer in layers: - metadata.enter_layer(layer) + metadata.ext(CSA2Metadata).enter_layer(layer) _prepare(metadata, requests, [0, 0], [257, 257]) forward() forward() metadata.on_update_kv_lens() - metadata.prepare_indexer(1) + metadata.ext(CSA2Metadata).prepare_indexer(1) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): forward() @@ -250,14 +282,16 @@ def forward(): @torch.inference_mode() def test_manager_partial_compression_graph_and_strided_publication(manager_requests): from tensorrt_llm._torch.attention.backends.sparse.csa2.compressor import CSA2Compressor - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( pack_rows, store_rows, ) manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) metadata = metadata.create_cuda_graph_metadata(2) compressor = CSA2Compressor(32, 512, 2, 1e-6).cuda() compressor.wkv_gate.weight.normal_(std=0.1) @@ -266,10 +300,10 @@ def test_manager_partial_compression_graph_and_strided_publication(manager_reque _prepare(metadata, requests, [0, 0], [1, 1]) manager.get_main_buffer(1).fill_(73) manager.get_index_pages(1).fill_(59) - compression = metadata.get_compression_batch(1) - write_slots = metadata.csa2_main_write_slots[1] - global_page_table = metadata.csa2_global_page_tables[1] - write_slot_mapping = metadata.csa2_main_write_slots + compression = metadata.ext(CSA2Metadata).get_compression_batch(1) + write_slots = metadata.ext(CSA2Metadata).csa2_main_write_slots[1] + global_page_table = metadata.ext(CSA2Metadata).csa2_global_page_tables[1] + write_slot_mapping = metadata.ext(CSA2Metadata).csa2_main_write_slots def forward(): output = compressor(inputs, compression) @@ -291,13 +325,13 @@ def forward(): inputs.copy_(torch.randn_like(inputs)) ordered = requests if step < 2 else requests[::-1] _prepare(metadata, ordered, [step, step], [1, 1]) - refreshed = metadata.get_compression_batch(1) - assert metadata.csa2_main_write_slots is not write_slot_mapping - write_slot_mapping = metadata.csa2_main_write_slots + refreshed = metadata.ext(CSA2Metadata).get_compression_batch(1) + assert metadata.ext(CSA2Metadata).csa2_main_write_slots is not write_slot_mapping + write_slot_mapping = metadata.ext(CSA2Metadata).csa2_main_write_slots assert ptrs == ( - metadata.csa2_main_write_slots[1].data_ptr(), + metadata.ext(CSA2Metadata).csa2_main_write_slots[1].data_ptr(), refreshed.start_positions.data_ptr(), - metadata.csa2_global_page_tables[1].data_ptr(), + metadata.ext(CSA2Metadata).csa2_global_page_tables[1].data_ptr(), ) manager.get_main_buffer(1).fill_(73) values, gates = compressor.wkv_gate(inputs).float().split(512, dim=-1) @@ -320,7 +354,7 @@ def forward(): expected *= torch.rsqrt(expected.square().mean(-1, keepdim=True) + 1e-6) expected = (expected * compressor.norm.weight.float()).to(torch.bfloat16) torch.testing.assert_close(output, expected, atol=0.02, rtol=0.02) - slots = metadata.csa2_main_write_slots[1].long() + slots = metadata.ext(CSA2Metadata).csa2_main_write_slots[1].long() torch.testing.assert_close(manager.get_main_buffer(1)[slots], pack_rows(output, "main")) # Main stores must never overwrite the adjacent index bytes. assert torch.all(manager.get_index_pages(1) == 59) @@ -337,7 +371,7 @@ def test_odd_prefix_reuse_compression_matches_fresh(): CSA2CacheRole, ) from tensorrt_llm._torch.attention.backends.sparse.csa2.compressor import CSA2Compressor - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType @@ -366,7 +400,9 @@ def test_odd_prefix_reuse_compression_matches_fresh(): dtype=DataType.BF16, vocab_size=8192, ) - metadata = CSA2TrtllmMetadata(max_num_requests=3, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=3, max_num_tokens=1024, kv_cache_manager=manager + ) compressor = CSA2Compressor(32, 512, 2, 1e-6).cuda() compressor.wkv_gate.weight.normal_(std=0.1) torch.manual_seed(671) @@ -383,9 +419,11 @@ def allocate(request_id, prefix, capacity): def compress(request_id, start, values): _prepare(metadata, [SimpleNamespace(py_request_id=request_id)], [start], [len(values)]) - output = compressor(values, metadata.get_compression_batch(0)) + output = compressor(values, metadata.ext(CSA2Metadata).get_compression_batch(0)) manager.write_swa(0, _swa(metadata, 0)[1], _values(compressor, values).bfloat16()) - manager.write_global(0, metadata.csa2_main_write_slots[0], output, output[:, :128]) + manager.write_global( + 0, metadata.ext(CSA2Metadata).csa2_main_write_slots[0], output, output[:, :128] + ) return output try: @@ -437,7 +475,7 @@ def compress(request_id, start, values): def test_chain_rewind_preserves_compressor_and_swa(start, accepted_drafts): from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager from tensorrt_llm._torch.attention.backends.sparse.csa2.compressor import CSA2Compressor - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( gather_rows, @@ -474,8 +512,10 @@ def test_chain_rewind_preserves_compressor_and_swa(start, accepted_drafts): sampling_config=SamplingConfig(), is_streaming=False, ) - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=512, kv_cache_manager=manager) - metadata.is_spec_decoding_enabled = True + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=512, kv_cache_manager=manager + ) + metadata.ext(SpecDecodingMetadata).is_spec_decoding_enabled = True compressor = CSA2Compressor(32, 512, 2, 1e-6).cuda() compressor.wkv_gate.weight.normal_(std=0.1) torch.manual_seed(719 + start + accepted_drafts) @@ -491,11 +531,13 @@ def compress(position, values, context=False): use_cache=True, num_cached_tokens_per_seq=[position] ) metadata.prepare() - output = compressor(values, metadata.get_compression_batch(0)) + output = compressor(values, metadata.ext(CSA2Metadata).get_compression_batch(0)) swa = _values(compressor, values).bfloat16() projected_swa.append(swa.clone()) manager.write_swa(0, _swa(metadata, 0)[1], swa) - manager.write_global(0, metadata.csa2_main_write_slots[0], output, output[:, :128]) + manager.write_global( + 0, metadata.ext(CSA2Metadata).csa2_main_write_slots[0], output, output[:, :128] + ) return output try: @@ -556,14 +598,16 @@ def compress(position, values, context=False): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_temporal_prior_identity_rewind_and_replay(manager_requests): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.metadata import KVCacheParams manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager + ) metadata.is_cuda_graph = True resets = [] - metadata.register_indexer_reset(1, lambda: resets.append(True)) + metadata.ext(CSA2Metadata).register_indexer_reset(1, lambda: resets.append(True)) def prepare(ordered, positions): metadata.request_ids = [request.py_request_id for request in ordered] @@ -574,15 +618,15 @@ def prepare(ordered, positions): use_cache=True, num_cached_tokens_per_seq=positions ) metadata.prepare() - return metadata.prepare_indexer_prior(1, 4) + return metadata.ext(CSA2Metadata).prepare_indexer_prior(1, 4) prior = prepare(requests, [0, 0]) assert torch.all(prior == -1) - assert metadata.csa2_indexer_prior_capacity[1].shape == (8, 4) + assert metadata.ext(CSA2Metadata).csa2_indexer_prior_capacity[1].shape == (8, 4) selected = torch.arange(8, device="cuda", dtype=torch.int32).view(2, 4) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - metadata.publish_indexer_prior(1, selected) + metadata.ext(CSA2Metadata).publish_indexer_prior(1, selected) graph.replay() prior = prepare(requests, [1, 1]) torch.testing.assert_close(prior, selected) @@ -607,28 +651,35 @@ def prepare(ordered, positions): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_retained_metadata_buffers_and_workspace_accounting(manager_requests): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) for length in (1, 2, 3, 4): _prepare(metadata, requests, [0, 0], [length, length]) - eager_keys = [key for key in metadata._csa2_buffers if not key[-1]] + eager_keys = [key for key in metadata.ext(CSA2Metadata)._csa2_buffers if not key[-1]] assert len({key[0] for key in eager_keys}) == len(eager_keys) - frame = metadata.get_query_tile_metadata( + frame = metadata.ext(CSA2Metadata).get_query_tile_metadata( torch.zeros(2, 8, 512, device="cuda", dtype=torch.bfloat16), 17 ) - before = metadata.get_workspace_bytes() - metadata._csa2_buffers["alias", (), torch.bfloat16, False] = frame.swa_pool.view(-1) - assert metadata.get_workspace_bytes() == before + before = metadata.ext(CSA2Metadata).get_workspace_bytes() + metadata.ext(CSA2Metadata)._csa2_buffers["alias", (), torch.bfloat16, False] = frame.ext( + CSA2Metadata + ).swa_pool.view(-1) + assert metadata.ext(CSA2Metadata).get_workspace_bytes() == before extra = torch.empty(129, device="cuda", dtype=torch.uint8) - metadata._csa2_buffers["extra", (129,), torch.uint8, False] = extra - assert metadata.get_workspace_bytes() == before + extra.untyped_storage().nbytes() + metadata.ext(CSA2Metadata)._csa2_buffers["extra", (129,), torch.uint8, False] = extra + assert ( + metadata.ext(CSA2Metadata).get_workspace_bytes() + == before + extra.untyped_storage().nbytes() + ) # Exact sum of the graph descriptor allocations: block table, contexts, # visible lengths, position validity, radix scratch and one schedule. Index # rows are read from the owner's cache pages in place, so no page copies are retained. expected = 4 * (1 * (4 + 64) + 8 + 80 * 32) + (148 + 1) * 8 - assert metadata.workspace_reservation_bytes(2, 4, 64, 32, 148) == expected + assert metadata.ext(CSA2Metadata).workspace_reservation_bytes(2, 4, 64, 32, 148) == expected @skip_pre_blackwell @@ -637,13 +688,15 @@ def test_retained_metadata_buffers_and_workspace_accounting(manager_requests): def test_pageable_metadata_uploads_feed_captured_consumer(manager_requests): dtype = torch.int64 manager, _ = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager + ) metadata.is_cuda_graph = True width, steps = 257, 64 stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - buffer = metadata._copy_csa2_tensor( + buffer = metadata.ext(CSA2Metadata)._copy_csa2_tensor( "pageable_lifetime", torch.zeros(width, dtype=dtype, device="cpu") ) address = buffer.data_ptr() @@ -668,7 +721,7 @@ def test_pageable_metadata_uploads_feed_captured_consumer(manager_requests): else: value = backing assert not value.is_pinned() - updated = metadata._copy_csa2_tensor("pageable_lifetime", value) + updated = metadata.ext(CSA2Metadata)._copy_csa2_tensor("pageable_lifetime", value) assert updated.data_ptr() == address del value, backing # Encourage immediate pageable allocation reuse before the consumer. @@ -688,7 +741,7 @@ def test_pageable_metadata_uploads_feed_captured_consumer(manager_requests): @torch.inference_mode() def test_representable_draft_switch_restores_target_fields(manager_requests): from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.speculative.interface import ( prepare_attn_metadata_for_draft_replay, restore_attn_metadata_after_draft_replay, @@ -718,7 +771,7 @@ def test_representable_draft_switch_restores_target_fields(manager_requests): assert draft.prepare_context(request) assert draft.resize_context(request, request.context_chunk_size) draft._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=2, max_num_tokens=1024, kv_cache_manager=target, @@ -730,18 +783,18 @@ def test_representable_draft_switch_restores_target_fields(manager_requests): layer: _swa(metadata, layer)[0].clone() for layer in range(len(target.layout.compress_ratios)) } - target_pages = dict(metadata._csa2_swa_page_tables) + target_pages = dict(metadata.ext(CSA2Metadata)._csa2_swa_page_tables) target_page_values = {layer: pages.clone() for layer, pages in target_pages.items()} - target_positions = metadata.csa2_positions - target_state = metadata.get_compression_batch(1).kv_state + target_positions = metadata.ext(CSA2Metadata).csa2_positions + target_state = metadata.ext(CSA2Metadata).get_compression_batch(1).kv_state resets = [] - metadata.register_indexer_reset(1, lambda: resets.append("target")) + metadata.ext(CSA2Metadata).register_indexer_reset(1, lambda: resets.append("target")) saved = prepare_attn_metadata_for_draft_replay(metadata, draft) assert resets == ["target"] - metadata.register_indexer_reset(1, lambda: resets.append("draft")) + metadata.ext(CSA2Metadata).register_indexer_reset(1, lambda: resets.append("draft")) try: assert metadata.kv_cache_manager is draft - for layer, pages in metadata._csa2_swa_page_tables.items(): + for layer, pages in metadata.ext(CSA2Metadata)._csa2_swa_page_tables.items(): torch.testing.assert_close( target_pages[layer], target_page_values[layer], atol=0, rtol=0 ) @@ -752,8 +805,8 @@ def test_representable_draft_switch_restores_target_fields(manager_requests): _assert_swa_cpu_reference(metadata, target, requests, [0, 0], [1, 1], [0, 0]) for layer, expected in target_reads.items(): torch.testing.assert_close(_swa(metadata, layer)[0], expected, atol=0, rtol=0) - assert metadata.csa2_positions is target_positions - assert metadata.get_compression_batch(1).kv_state is target_state + assert metadata.ext(CSA2Metadata).csa2_positions is target_positions + assert metadata.ext(CSA2Metadata).get_compression_batch(1).kv_state is target_state assert resets == ["target", "target"] # A later draft switch restores its own cached callback/prior state; # contiguous identity must not suppress reset of shared emission state. @@ -845,30 +898,37 @@ def test_batch_page_tables_match_per_request_converters(manager_requests): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_temporal_prefill_seeds_first_decode(manager_requests, monkeypatch): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.metadata import KVCacheParams manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager + ) _prepare(metadata, requests, [0, 0], [2, 2]) - assert metadata.prepare_indexer_prior(1, 4).shape == (0, 4) + assert metadata.ext(CSA2Metadata).prepare_indexer_prior(1, 4).shape == (0, 4) selected = torch.arange(16, dtype=torch.int32, device="cuda").view(4, 4) - metadata.publish_indexer_prior(1, selected) + metadata.ext(CSA2Metadata).publish_indexer_prior(1, selected) metadata.seq_lens = torch.ones(2, dtype=torch.int32) metadata.num_contexts = 0 metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[2, 2]) metadata.prepare() - torch.testing.assert_close(metadata.prepare_indexer_prior(1, 4), selected[[1, 3]]) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).prepare_indexer_prior(1, 4), selected[[1, 3]] + ) if torch.cuda.get_device_capability()[0] != 10: return # Native FP4 paged descriptor coverage is validated on SM100. manager.get_index_pages(1).zero_() - assert metadata.prepare_indexer(1) is metadata + assert metadata.ext(CSA2Metadata).prepare_indexer(1) is metadata.ext(CSA2Metadata) # Index pages are read in place through the owner's block table. - assert metadata.csa2_indexer_k_cache.data_ptr() == manager.get_index_pages(1).data_ptr() - valid = metadata.csa2_indexer_valid_positions + assert ( + metadata.ext(CSA2Metadata).csa2_indexer_k_cache.data_ptr() + == manager.get_index_pages(1).data_ptr() + ) + valid = metadata.ext(CSA2Metadata).csa2_indexer_valid_positions offsets = torch.arange(valid.shape[1], device=valid.device) torch.testing.assert_close( - valid, offsets[None, :] < metadata.csa2_indexer_visible_lengths[:, None] + valid, offsets[None, :] < metadata.ext(CSA2Metadata).csa2_indexer_visible_lengths[:, None] ) @@ -876,13 +936,15 @@ def test_temporal_prefill_seeds_first_decode(manager_requests, monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_graph_workspace_is_independent_of_context_gather_cap(manager_requests): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.metadata import KVCacheParams if torch.cuda.get_device_capability()[0] != 10: pytest.skip("Native FP4 paged descriptors require SM100-family GPUs") manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=8, kv_cache_manager=manager + ) metadata.is_cuda_graph = True metadata.request_ids = [request.py_request_id for request in requests] metadata.seq_lens = torch.ones(2, dtype=torch.int32) @@ -891,12 +953,12 @@ def test_graph_workspace_is_independent_of_context_gather_cap(manager_requests): metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0, 0]) metadata.prepare() manager.fp8_ctx_mla_kv_len_cap = 1 - metadata.prepare_indexer(1) - metadata.prepare_indexer(3) - arenas = list(metadata._csa2_indexer_workspaces.values()) + metadata.ext(CSA2Metadata).prepare_indexer(1) + metadata.ext(CSA2Metadata).prepare_indexer(3) + arenas = list(metadata.ext(CSA2Metadata)._csa2_indexer_workspaces.values()) assert len(arenas) == 1 assert arenas[0]["block_table"].shape[0] == metadata.max_num_requests - expected = metadata._workspace_storage_bytes(arenas) + expected = metadata.ext(CSA2Metadata)._workspace_storage_bytes(arenas) assert expected > 0 from tensorrt_llm._torch.pyexecutor._util import ( get_attention_workspace_fixed_bytes, @@ -936,7 +998,9 @@ def replay_metadata(): assert manager._resume_and_restore(request_id, cache) assert cache.resize(32) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=64, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=64, kv_cache_manager=manager + ) yield manager, metadata for request_id in list(manager.kv_cache_map): manager.free_resources(SimpleNamespace(py_request_id=request_id)) @@ -969,7 +1033,7 @@ def _prepare_replay(metadata, prefixes, suffixes=None, *, decoder=False, ids=Non metadata.kv_cache_params = KVCacheParams( use_cache=True, num_cached_tokens_per_seq=[start for start, _ in ranges] ) - metadata.set_swa_bounded_replay(prefixes, decoder=decoder) + metadata.ext(CSA2Metadata).set_swa_bounded_replay(prefixes, decoder=decoder) metadata.prepare() return ranges @@ -981,34 +1045,40 @@ def test_owner_source_selection_and_one_shot_reset(replay_metadata): manager, metadata = replay_metadata assert _prepare_replay(metadata, [5, 6], [3, 2]) == ((1, 8), (2, 8)) # Ratio2 replays one raw odd tail plus suffix; ratio1 only projects suffix. - assert metadata.csa2_global_source_indices[1].tolist() == [3, 4, 5, 6, 11, 12] - assert metadata.csa2_global_source_indices[4].tolist() == [4, 5, 6, 11, 12] - compression = metadata.get_compression_batch(1) + assert metadata.ext(CSA2Metadata).csa2_global_source_indices[1].tolist() == [3, 4, 5, 6, 11, 12] + assert metadata.ext(CSA2Metadata).csa2_global_source_indices[4].tolist() == [4, 5, 6, 11, 12] + compression = metadata.ext(CSA2Metadata).get_compression_batch(1) assert compression.start_positions.tolist() == [4, 6] assert compression.cu_seq_lengths.tolist() == [0, 4, 6] assert compression.cu_compressed_lengths.tolist() == [0, 2, 3] - assert set(metadata.csa2_global_source_indices) == {1, 4} - assert [metadata.csa2_kv_sources[layer] for layer in (0, 1, 2, 3, 4)] == [None, 1, 1, 1, 4] + assert set(metadata.ext(CSA2Metadata).csa2_global_source_indices) == {1, 4} + assert [metadata.ext(CSA2Metadata).csa2_kv_sources[layer] for layer in (0, 1, 2, 3, 4)] == [ + None, + 1, + 1, + 1, + 4, + ] for layer in range(5): reads = _swa(metadata, layer)[0] assert reads[0, :-1].tolist() == [-1, -1, -1] assert reads[7, :-1].tolist() == [-1, -1, -1] hidden = torch.randn(13, 32, device="cuda", dtype=torch.bfloat16) torch.testing.assert_close( - metadata.select_global_source(1, hidden), hidden[[3, 4, 5, 6, 11, 12]] + metadata.ext(CSA2Metadata).select_global_source(1, hidden), hidden[[3, 4, 5, 6, 11, 12]] ) with pytest.raises(ValueError, match="external source batch"): - metadata.select_global_source(1, hidden, hidden) + metadata.ext(CSA2Metadata).select_global_source(1, hidden, hidden) # A normal subsequent prepare consumes no stale replay floor or selector. metadata.seq_lens = torch.ones(2, dtype=torch.int32) metadata.num_contexts = 0 metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[8, 8]) metadata.prepare() - assert metadata.csa2_replay_mode is None - assert metadata.csa2_global_source_indices == {} - assert metadata.csa2_replay_start_positions.tolist() == [0, 0] + assert metadata.ext(CSA2Metadata).csa2_replay_mode is None + assert metadata.ext(CSA2Metadata).csa2_global_source_indices == {} + assert metadata.ext(CSA2Metadata).csa2_replay_start_positions.tolist() == [0, 0] next_hidden = hidden[:2] - assert metadata.select_global_source(1, next_hidden) is next_hidden + assert metadata.ext(CSA2Metadata).select_global_source(1, next_hidden) is next_hidden @skip_pre_blackwell @@ -1020,20 +1090,20 @@ def test_replay_graph_updates_values_but_rejects_shape_or_mode_change(replay_met _prepare_replay(metadata, [5], [3]) hidden = torch.randn(7, 32, device="cuda", dtype=torch.bfloat16) q = torch.zeros(7, 8, 512, device="cuda", dtype=torch.bfloat16) - metadata.get_query_tile_metadata(q, 4, query_start=0) + metadata.ext(CSA2Metadata).get_query_tile_metadata(q, 4, query_start=0) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - metadata.get_query_tile_metadata(q, 4, query_start=0) - selected = metadata.select_global_source(1, hidden) - positions = metadata.csa2_positions.clone() - write_slots = metadata.csa2_main_write_slots[1].clone() + metadata.ext(CSA2Metadata).get_query_tile_metadata(q, 4, query_start=0) + selected = metadata.ext(CSA2Metadata).select_global_source(1, hidden) + positions = metadata.ext(CSA2Metadata).csa2_positions.clone() + write_slots = metadata.ext(CSA2Metadata).csa2_main_write_slots[1].clone() # Same source/query shape at a different absolute hit position is valid. _prepare_replay(metadata, [7], [3], ids=[97]) hidden.neg_() graph.replay() torch.testing.assert_close(selected, hidden[3:]) torch.testing.assert_close(positions, torch.arange(3, 10, dtype=torch.int32, device="cuda")) - torch.testing.assert_close(write_slots, metadata.csa2_main_write_slots[1]) + torch.testing.assert_close(write_slots, metadata.ext(CSA2Metadata).csa2_main_write_slots[1]) with pytest.raises(ValueError, match="fresh graph metadata"): metadata.prepare() with pytest.raises(ValueError, match="fresh graph metadata"): @@ -1056,7 +1126,7 @@ def test_replay_requires_first_window_row_to_be_physically_writable(): # requires127..254. Page0 is genuinely absent in the real allocator. assert cache.resize(255, 255) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=128, kv_cache_manager=manager ) assert manager.get_swa_replay_ranges([255]) == ((127, 255),) @@ -1072,32 +1142,41 @@ def test_replay_requires_first_window_row_to_be_physically_writable(): @pytest.mark.parametrize("ratio", [1, 2]) @pytest.mark.parametrize("accepted_end", [128, 129, 130]) def test_live_endpoint_refreshes_query_slots_and_compressed_groups( - monkeypatch, ratio, accepted_end, query_width + ratio, accepted_end, query_width ): # Base scheduler invalidation is orthogonal; exercise the actual CSA2 tensor # refresh without native allocations, including an odd ratio-two boundary. - monkeypatch.setattr(TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None) - metadata = CSA2TrtllmMetadata.__new__(CSA2TrtllmMetadata) - metadata._csa2_ready_for_kv_update = True + metadata = _bare_csa2_metadata() + metadata.ext(CSA2Metadata)._csa2_ready_for_kv_update = True metadata.kv_cache_manager = SimpleNamespace( tokens_per_block=128, layout=SimpleNamespace(window_size=128, compress_ratios=(ratio,)) ) - metadata._csa2_query_lengths = torch.tensor([query_width], dtype=torch.int32) - metadata._csa2_query_base = torch.tensor([129], dtype=torch.int32) - metadata._csa2_query_offsets = torch.arange(query_width, dtype=torch.int32) + metadata.ext(CSA2Metadata)._csa2_query_lengths = torch.tensor([query_width], dtype=torch.int32) + metadata.ext(CSA2Metadata)._csa2_query_base = torch.tensor([129], dtype=torch.int32) + metadata.ext(CSA2Metadata)._csa2_query_offsets = torch.arange(query_width, dtype=torch.int32) metadata.kv_lens_cuda = torch.tensor([accepted_end], dtype=torch.int32) - metadata.csa2_token_requests = torch.zeros(query_width, dtype=torch.int64) - metadata.csa2_positions = torch.empty(query_width, dtype=torch.int32) - metadata.csa2_replay_start_positions = torch.zeros(1, dtype=torch.int64) - metadata._csa2_swa_page_tables = {0: torch.tensor([[4, 7]], dtype=torch.int32)} + metadata.ext(CSA2Metadata).csa2_token_requests = torch.zeros(query_width, dtype=torch.int64) + metadata.ext(CSA2Metadata).csa2_positions = torch.empty(query_width, dtype=torch.int32) + metadata.ext(CSA2Metadata).csa2_replay_start_positions = torch.zeros(1, dtype=torch.int64) + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = { + 0: torch.tensor([[4, 7]], dtype=torch.int32) + } metadata.is_cuda_graph = False - metadata._csa2_source_geometry = {0: (torch.tensor([129]), torch.tensor([query_width]))} - metadata.csa2_main_write_slots = {0: torch.empty(query_width, dtype=torch.int64)} - metadata._csa2_compressed_positions = {0: torch.empty(query_width, dtype=torch.int32)} - metadata.csa2_global_page_sizes = {0: 128 // ratio} - metadata.csa2_global_page_tables = {0: torch.tensor([[9, 12]], dtype=torch.int32)} + metadata.ext(CSA2Metadata)._csa2_source_geometry = { + 0: (torch.tensor([129]), torch.tensor([query_width])) + } + metadata.ext(CSA2Metadata).csa2_main_write_slots = { + 0: torch.empty(query_width, dtype=torch.int64) + } + metadata.ext(CSA2Metadata)._csa2_compressed_positions = { + 0: torch.empty(query_width, dtype=torch.int32) + } + metadata.ext(CSA2Metadata).csa2_global_page_sizes = {0: 128 // ratio} + metadata.ext(CSA2Metadata).csa2_global_page_tables = { + 0: torch.tensor([[9, 12]], dtype=torch.int32) + } if ratio == 2: - metadata._csa2_compression = { + metadata.ext(CSA2Metadata)._csa2_compression = { 0: CSA2CompressionBatch( torch.empty(0), torch.empty(0), @@ -1112,30 +1191,33 @@ def test_live_endpoint_refreshes_query_slots_and_compressed_groups( query_width, ) } - identities = (metadata.csa2_positions.data_ptr(), metadata.csa2_main_write_slots[0].data_ptr()) - metadata.on_update_kv_lens() + identities = ( + metadata.ext(CSA2Metadata).csa2_positions.data_ptr(), + metadata.ext(CSA2Metadata).csa2_main_write_slots[0].data_ptr(), + ) + metadata.ext(CSA2Metadata).on_kv_lens_changed(spec_dec_step=False) positions = list(range(accepted_end - query_width, accepted_end)) - assert metadata.csa2_positions.tolist() == positions + assert metadata.ext(CSA2Metadata).csa2_positions.tolist() == positions # The layer resolves its SWA slots from the refreshed positions when entered. _, writes, visible = _swa(metadata, 0) assert writes.tolist() == [[4, 7][p // 128] * 128 + p % 128 for p in positions] assert visible.tolist() == [(p + 1) // ratio for p in positions] groups = list(range((accepted_end - query_width) // ratio, accepted_end // ratio)) page_size = 128 // ratio - assert metadata.csa2_main_write_slots[0].tolist() == [ + assert metadata.ext(CSA2Metadata).csa2_main_write_slots[0].tolist() == [ [9, 12][g // page_size] * page_size + g % page_size for g in groups ] + [-1] * (query_width - len(groups)) - assert metadata._csa2_compressed_positions[0].tolist() == [g * ratio for g in groups] + [0] * ( - query_width - len(groups) - ) + assert metadata.ext(CSA2Metadata)._csa2_compressed_positions[0].tolist() == [ + g * ratio for g in groups + ] + [0] * (query_width - len(groups)) if ratio == 2: - batch = metadata._csa2_compression[0] + batch = metadata.ext(CSA2Metadata)._csa2_compression[0] assert batch.start_positions.tolist() == [accepted_end - query_width] assert batch.kv_lengths.tolist() == [accepted_end] assert batch.cu_compressed_lengths.tolist() == [0, len(groups)] assert identities == ( - metadata.csa2_positions.data_ptr(), - metadata.csa2_main_write_slots[0].data_ptr(), + metadata.ext(CSA2Metadata).csa2_positions.data_ptr(), + metadata.ext(CSA2Metadata).csa2_main_write_slots[0].data_ptr(), ) @@ -1143,43 +1225,56 @@ def test_paged_owner_mapping_and_request_reordering(): pages = _page_metadata( torch.tensor([[2, 0], [1, -1]], dtype=torch.int32), torch.tensor([1, 0]), 2, 4 ) - assert pages.global_slot_tile(0, 0, 2).tolist() == [[2, 3, -1, -1], [4, 5, 0, 1]] + assert pages.ext(CSA2Metadata).global_slot_tile(0, 0, 2).tolist() == [ + [2, 3, -1, -1], + [4, 5, 0, 1], + ] logical = torch.tensor([[1, 3, -1], [2, 0, 4]], dtype=torch.int32) - assert pages.global_slot_tile(0, 0, 2, logical).tolist() == [[3, -1, -1], [0, 4, -1]] + assert pages.ext(CSA2Metadata).global_slot_tile(0, 0, 2, logical).tolist() == [ + [3, -1, -1], + [0, 4, -1], + ] def test_invalid_paged_request_does_not_alias_real_request(): pages = _page_metadata(torch.tensor([[3], [7]]), torch.tensor([-1, 2, 1]), 2, 2) - assert pages.global_slot_tile(0, 0, 3).tolist() == [[-1, -1], [-1, -1], [14, 15]] + assert pages.ext(CSA2Metadata).global_slot_tile(0, 0, 3).tolist() == [ + [-1, -1], + [-1, -1], + [14, 15], + ] def test_routing_reset_detaches_shallow_clone(): metadata = _page_metadata(torch.tensor([[0]]), torch.tensor([0]), 1, 1) layout = CSA2Layout((1, 1), (0,), (0,)) - metadata.enter_layer(layout.layer(0)) + metadata.ext(CSA2Metadata).enter_layer(layout.layer(0)) indices = torch.tensor([[0]]) - metadata.csa2_indices[0] = indices + metadata.ext(CSA2Metadata).csa2_indices[0] = indices clone = copy(metadata) - clone.reset_routing() - assert clone.csa2_indices == {} and clone.csa2_candidates == {} - assert metadata.csa2_indices[0] is indices - clone.enter_layer(layout.layer(0)) + clone.replace_ext(metadata.ext(CSA2Metadata).clone_for(clone)) + clone.ext(CSA2Metadata).reset_routing() + assert ( + clone.ext(CSA2Metadata).csa2_indices == {} and clone.ext(CSA2Metadata).csa2_candidates == {} + ) + assert metadata.ext(CSA2Metadata).csa2_indices[0] is indices + clone.ext(CSA2Metadata).enter_layer(layout.layer(0)) with pytest.raises(ValueError, match="across forwards"): - metadata.enter_layer(layout.layer(0)) - metadata.reset_routing() - assert metadata.csa2_indices == {} + metadata.ext(CSA2Metadata).enter_layer(layout.layer(0)) + metadata.ext(CSA2Metadata).reset_routing() + assert metadata.ext(CSA2Metadata).csa2_indices == {} @skip_pre_blackwell @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_staging_excludes_out_of_capacity_rows() -> None: - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2BackendForwardArgs if torch.cuda.get_device_capability()[0] != 10: pytest.skip("Fused CSA2 gather requires SM100-family validation") q = torch.empty((2, 32, 512), dtype=torch.bfloat16, device="cuda") - metadata = CSA2TrtllmMetadata.for_query_tile(q, 17) + metadata = CSA2Metadata.for_query_tile(q, 17) swa = pack_rows(torch.randn(3, 512, dtype=torch.bfloat16, device="cuda"), "swa") main = pack_rows(torch.randn(3, 512, dtype=torch.bfloat16, device="cuda"), "main") swa_slots = torch.tensor([[0, 3, 2**40, -1], [3, -1, -1, -1]], device="cuda") @@ -1187,8 +1282,8 @@ def test_staging_excludes_out_of_capacity_rows() -> None: inputs = CSA2BackendForwardArgs( swa_pool=swa, swa_indices=swa_slots, main_pool=main, topk_indices=main_slots ) - metadata.stage_selected(inputs) - assert metadata.prepared_lens.tolist() == [2, 1] + metadata.ext(CSA2Metadata).stage_selected(inputs) + assert metadata.ext(CSA2Metadata).prepared_lens.tolist() == [2, 1] expected = torch.cat( ( _staging_gather_reference(swa, swa_slots[:, :1], 512, "swa"), @@ -1196,8 +1291,8 @@ def test_staging_excludes_out_of_capacity_rows() -> None: ), dim=1, ) - _assert_staging_bf16_decode(metadata.swa_pool[0, :2], expected[0]) - assert torch.count_nonzero(metadata.swa_pool[1, 0]) == 0 + _assert_staging_bf16_decode(metadata.ext(CSA2Metadata).swa_pool[0, :2], expected[0]) + assert torch.count_nonzero(metadata.ext(CSA2Metadata).swa_pool[1, 0]) == 0 def _staging_gather_reference(pool, slots, dim, cache_format): @@ -1223,9 +1318,9 @@ def _assert_staging_bf16_decode(actual, expected): class _LifecycleMetadata(_FakeMetadata): - _get_csa2_buffer = CSA2TrtllmMetadata._get_csa2_buffer - _copy_host = CSA2TrtllmMetadata._copy_host - _copy_csa2_tensor = CSA2TrtllmMetadata._copy_csa2_tensor + _get_csa2_buffer = CSA2Metadata._get_csa2_buffer + _copy_host = CSA2Metadata._copy_host + _copy_csa2_tensor = CSA2Metadata._copy_csa2_tensor @property def num_generations(self) -> int: @@ -1264,7 +1359,7 @@ def test_replay_preserves_live_speculative_kv_lifecycle( metadata.kv_lens_cuda[2] -= 3 metadata.on_update_kv_lens() expected_kv = metadata.kv_lens_cuda.clone() - positions = metadata.csa2_positions.clone() + positions = metadata.ext(CSA2Metadata).csa2_positions.clone() untrimmed_swa = _swa(metadata, 0)[0][400:].clone() requests = [SimpleNamespace(py_ced_replay=None) for _ in range(2)] if pass_requests else None plan = plan_decoder_replay(metadata, WINDOW, requests, private_decoder=private_decoder) @@ -1281,22 +1376,28 @@ def test_replay_preserves_live_speculative_kv_lifecycle( # Encoder handoffs are published after planning and cleared by prepare(). topk = torch.arange(413 * 3, dtype=torch.int32).reshape(413, 3) candidates = torch.arange(413 * 2, dtype=torch.int32).reshape(413, 2) - metadata.csa2_indices[20] = topk - metadata.csa2_candidates[20] = candidates - metadata.csa2_precomputed_kv_layers = {20} - metadata._csa2_last_layer = 20 + metadata.ext(CSA2Metadata).csa2_indices[20] = topk + metadata.ext(CSA2Metadata).csa2_candidates[20] = candidates + metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers = {20} + metadata.ext(CSA2Metadata)._csa2_last_layer = 20 saved_prompt = metadata.prompt_lens original_buffer = metadata.kv_lens_cuda enter_decoder_replay(metadata, plan) assert metadata.prepare_calls == 1 assert metadata.kv_lens_cuda is original_buffer torch.testing.assert_close(metadata.kv_lens_cuda, expected_kv, atol=0, rtol=0) - torch.testing.assert_close(metadata.csa2_positions, positions[plan.rows], atol=0, rtol=0) - torch.testing.assert_close(metadata.csa2_indices[20], topk[plan.rows], atol=0, rtol=0) - torch.testing.assert_close(metadata.csa2_candidates[20], candidates[plan.rows], atol=0, rtol=0) - assert metadata.csa2_precomputed_kv_layers == {20} - assert metadata._csa2_last_layer == 20 - assert metadata.csa2_replay_query_rows is plan.rows + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_positions, positions[plan.rows], atol=0, rtol=0 + ) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_indices[20], topk[plan.rows], atol=0, rtol=0 + ) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_candidates[20], candidates[plan.rows], atol=0, rtol=0 + ) + assert metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers == {20} + assert metadata.ext(CSA2Metadata)._csa2_last_layer == 20 + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is plan.rows indices = _swa(metadata, 0)[0] torch.testing.assert_close((indices[:WINDOW] >= 0).sum(-1), torch.arange(1, WINDOW + 1)) torch.testing.assert_close(indices[WINDOW - 1], torch.arange(400, 528)) @@ -1309,26 +1410,30 @@ def test_replay_preserves_live_speculative_kv_lifecycle( assert metadata.seq_lens is plan.saved_seq_lens assert metadata.seq_lens.tolist() == [400, 7, 6] assert metadata.kv_cache_params.num_cached_tokens_per_seq == [128, 700, 900] - assert metadata.csa2_replay_query_rows is None + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None torch.testing.assert_close(metadata.kv_lens_cuda, expected_kv, atol=0, rtol=0) def test_remote_tail_decoder_masks_only_pre_handoff_context_swa() -> None: metadata = _LifecycleMetadata([4, 3], 1, [100, 200]) metadata.prepare() - metadata._ensure_swa_slots() - generation_indices = metadata.csa2_swa_indices[0][4:].clone() - write_slots = metadata.csa2_swa_write_slots[0].clone() + metadata.ext(CSA2Metadata)._ensure_swa_slots() + generation_indices = metadata.ext(CSA2Metadata).csa2_swa_indices[0][4:].clone() + write_slots = metadata.ext(CSA2Metadata).csa2_swa_write_slots[0].clone() enter_remote_tail_decoder(metadata, [102]) # The next decoder layer resolves the slots against the raised floors. - metadata._ensure_swa_slots() + metadata.ext(CSA2Metadata)._ensure_swa_slots() - context_indices = metadata.csa2_swa_indices[0][:4] + context_indices = metadata.ext(CSA2Metadata).csa2_swa_indices[0][:4] assert torch.all(context_indices[context_indices >= 0] >= 102) - torch.testing.assert_close(metadata.csa2_swa_indices[0][4:], generation_indices, atol=0, rtol=0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_swa_indices[0][4:], generation_indices, atol=0, rtol=0 + ) # The floor hides reads only: rows below it still store their SWA. - torch.testing.assert_close(metadata.csa2_swa_write_slots[0], write_slots, atol=0, rtol=0) + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_swa_write_slots[0], write_slots, atol=0, rtol=0 + ) def _assert_swa_cpu_reference(metadata, manager, requests, starts, lengths, floors, observed=None): @@ -1375,7 +1480,7 @@ def _assert_swa_cpu_reference(metadata, manager, requests, starts, lengths, floo def test_coalesced_uploads_mixed_dtypes_empty_values_and_slab_growth(): from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import _CoalescedUploads - metadata = CSA2TrtllmMetadata.__new__(CSA2TrtllmMetadata) + metadata = _bare_csa2_metadata() vars(metadata).update( is_cuda_graph=False, kv_lens_cuda=torch.empty(0, dtype=torch.int32, device="cuda") ) @@ -1386,7 +1491,7 @@ def test_coalesced_uploads_mixed_dtypes_empty_values_and_slab_growth(): "empty": np.empty(0, dtype=np.int64), "scalar": torch.tensor(size, dtype=torch.int32), } - uploads = _CoalescedUploads(metadata) + uploads = _CoalescedUploads(metadata.ext(CSA2Metadata)) outputs = {key: uploads.copy(key, value) for key, value in values.items()} uploads.flush() for key, value in values.items(): @@ -1397,10 +1502,10 @@ def test_coalesced_uploads_mixed_dtypes_empty_values_and_slab_growth(): @pytest.mark.cpu_only def test_swa_write_interval_validation_preserves_errors(): def check(pages, starts, lengths, floors, allocated, block): - rows, columns = CSA2TrtllmMetadata._swa_write_pages( + rows, columns = CSA2Metadata._swa_write_pages( starts, lengths, floors, allocated, block, pages.shape ) - CSA2TrtllmMetadata._check_written_pages(pages.numpy()[rows, columns] >= 0) + CSA2Metadata._check_written_pages(pages.numpy()[rows, columns] >= 0) full = [1 << 30] check(torch.tensor([[-1, 7]]), [128], [2], [0], full, 128) # Historical-only hole. @@ -1455,7 +1560,7 @@ def test_forward_swa_resolution_matches_cpu_reference(queries): assert manager.prepare_context(request) assert manager.resize_context(request, request.context_chunk_size) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=8192, kv_cache_manager=manager ) metadata.request_ids = [request.py_request_id] @@ -1465,7 +1570,7 @@ def test_forward_swa_resolution_matches_cpu_reference(queries): metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]) metadata.prepare() _assert_swa_cpu_reference(metadata, manager, [request], [0], [queries], [0]) - assert metadata._csa2_ready_for_kv_update + assert metadata.ext(CSA2Metadata)._csa2_ready_for_kv_update finally: manager.free_resources(request) manager.shutdown() @@ -1495,15 +1600,17 @@ def test_swa_resolution_integer_outputs_and_graph(monkeypatch, native): pages_host = [row[:table_width] for row in pages_host] positions = torch.tensor(positions_host, dtype=torch.int32, device="cuda") requests = torch.tensor(requests_host, dtype=torch.int64, device="cuda") - metadata = CSA2TrtllmMetadata.__new__(CSA2TrtllmMetadata) + metadata = _bare_csa2_metadata() vars(metadata).update( is_cuda_graph=False, - _csa2_buffers={}, kv_lens_cuda=positions, kv_cache_manager=SimpleNamespace( tokens_per_block=block, layout=SimpleNamespace(window_size=window, compress_ratios=ratios), ), + ) + vars(metadata.ext(CSA2Metadata)).update( + _csa2_buffers={}, csa2_positions=positions, csa2_token_requests=requests, csa2_replay_start_positions=torch.tensor(floors_host, dtype=torch.int64, device="cuda"), @@ -1518,8 +1625,12 @@ def test_swa_resolution_integer_outputs_and_graph(monkeypatch, native): def bind(): if native: - metadata._csa2_swa_descriptors = metadata._swa_table_descriptors(3) - metadata._csa2_swa_ratios = metadata._swa_ratio_tensor(ratios) + metadata.ext(CSA2Metadata)._csa2_swa_descriptors = metadata.ext( + CSA2Metadata + )._swa_table_descriptors(3) + metadata.ext(CSA2Metadata)._csa2_swa_ratios = metadata.ext( + CSA2Metadata + )._swa_ratio_tensor(ratios) outputs = [ tuple( @@ -1531,13 +1642,13 @@ def bind(): def resolve_all(force=True): if force: - metadata._csa2_swa_resolved = False - metadata._ensure_swa_slots() + metadata.ext(CSA2Metadata)._csa2_swa_resolved = False + metadata.ext(CSA2Metadata)._ensure_swa_slots() for layer, kept in enumerate(outputs): for name, out in zip( ("csa2_swa_indices", "csa2_swa_write_slots", "csa2_visible_lengths"), kept ): - out.copy_(getattr(metadata, name)[layer]) + out.copy_(getattr(metadata.ext(CSA2Metadata), name)[layer]) def check(): # Independent scalar CPU oracle, including invalid request domains. @@ -1579,7 +1690,7 @@ def check(): with torch.cuda.graph(graph): # Eager warmup resolved these inputs; the captured forward clears # the flag at its start and records the launch. - metadata.reset_routing() + metadata.ext(CSA2Metadata).reset_routing() resolve_all(force=False) for page in (-1, 13): pages_host[0][0] = page @@ -1588,11 +1699,11 @@ def check(): positions_host[1] += 1 positions.copy_(torch.tensor(positions_host, dtype=torch.int32, device="cuda")) requests.copy_(torch.tensor(requests_host, dtype=torch.int64, device="cuda")) - metadata.csa2_replay_start_positions.copy_( + metadata.ext(CSA2Metadata).csa2_replay_start_positions.copy_( torch.tensor(floors_host, dtype=torch.int64, device="cuda") ) for layer in range(3): - metadata._csa2_swa_page_tables[layer][0, 0] = page + metadata.ext(CSA2Metadata)._csa2_swa_page_tables[layer][0, 0] = page for kept in outputs: for out in kept: out.fill_(94) @@ -1651,13 +1762,13 @@ def test_indexer_descriptors_expand_pages_and_mask_missing_pages(): expected_valid, ) with torch.inference_mode(): - CSA2TrtllmMetadata._fill_indexer_descriptors( + CSA2Metadata._fill_indexer_descriptors( table, requests, decode_visible, 1, factor, 64, *outputs ) if device == "cuda" and page_size == 128 and columns: graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - CSA2TrtllmMetadata._fill_indexer_descriptors( + CSA2Metadata._fill_indexer_descriptors( table, requests, decode_visible, 1, factor, 64, *outputs ) for value in outputs: @@ -1675,23 +1786,25 @@ def test_indexer_descriptors_expand_pages_and_mask_missing_pages(): @pytest.mark.cpu_only @pytest.mark.parametrize("published_first", [False, True]) def test_swa_table_backing_stays_fixed_across_publication_and_fallback(published_first): - metadata = CSA2TrtllmMetadata.__new__(CSA2TrtllmMetadata) - metadata._csa2_buffers = {} + metadata = _bare_csa2_metadata() + metadata.ext(CSA2Metadata)._csa2_buffers = {} metadata.kv_lens_cuda = torch.empty(0, dtype=torch.int32) metadata.is_cuda_graph = True outer = torch.arange(16, dtype=torch.int32).reshape(2, 2, 4) published = outer[:, 0, :3] host = torch.full((2, 3), 17, dtype=torch.int32) - first = metadata._bind_swa_pages(0, host, published if published_first else None) + first = metadata.ext(CSA2Metadata)._bind_swa_pages( + 0, host, published if published_first else None + ) pointer, stride = first.data_ptr(), first.stride() assert (pointer == published.data_ptr()) == published_first # Direct fallback refreshes captured backing, even when that backing is an # outer-table view rather than the old private contiguous allocation. - fallback = metadata._bind_swa_pages(0, host, None) + fallback = metadata.ext(CSA2Metadata)._bind_swa_pages(0, host, None) assert (fallback.data_ptr(), fallback.stride()) == (pointer, stride) torch.testing.assert_close(fallback, host, atol=0, rtol=0) other = torch.full((2, 2, 4), 29, dtype=torch.int32)[:, 0, :3] - normal = metadata._bind_swa_pages(0, host, other) + normal = metadata.ext(CSA2Metadata)._bind_swa_pages(0, host, other) assert (normal.data_ptr(), normal.stride()) == (pointer, stride) torch.testing.assert_close(normal, other, atol=0, rtol=0) @@ -1702,18 +1815,18 @@ def rows(tables): small = [torch.zeros((1, 4), dtype=torch.int32) for _ in range(2)] large = [torch.zeros((8, 4), dtype=torch.int32) for _ in range(2)] - metadata._csa2_swa_page_tables = dict(enumerate(small)) - first = metadata._swa_table_descriptors(2) - metadata._csa2_swa_page_tables = dict(enumerate(large)) - second = metadata._swa_table_descriptors(2) + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = dict(enumerate(small)) + first = metadata.ext(CSA2Metadata)._swa_table_descriptors(2) + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = dict(enumerate(large)) + second = metadata.ext(CSA2Metadata)._swa_table_descriptors(2) assert first.data_ptr() != second.data_ptr() and second.tolist() == rows(large) - metadata._csa2_swa_page_tables = dict(enumerate(small)) - assert metadata._swa_table_descriptors(2).data_ptr() == first.data_ptr() + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = dict(enumerate(small)) + assert metadata.ext(CSA2Metadata)._swa_table_descriptors(2).data_ptr() == first.data_ptr() # Eager metadata refreshes a single buffer in place before every launch. metadata.is_cuda_graph = False - eager = metadata._swa_table_descriptors(2) - metadata._csa2_swa_page_tables = dict(enumerate(large)) - assert metadata._swa_table_descriptors(2).data_ptr() == eager.data_ptr() + eager = metadata.ext(CSA2Metadata)._swa_table_descriptors(2) + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = dict(enumerate(large)) + assert metadata.ext(CSA2Metadata)._swa_table_descriptors(2).data_ptr() == eager.data_ptr() assert eager.tolist() == rows(large) and first.tolist() == rows(small) @@ -1726,7 +1839,9 @@ def test_swa_fresh_publication_reuses_tables_and_keeps_captured_backing( from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheRole manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) metadata.is_cuda_graph = True # Physical publication order need not equal logical model-layer order. order = [2, 0, 3, 1] @@ -1758,9 +1873,9 @@ def remap(shift): assert len(conversions) == 1 pointers = { layer: (value.data_ptr(), value.stride()) - for layer, value in metadata._csa2_swa_page_tables.items() + for layer, value in metadata.ext(CSA2Metadata)._csa2_swa_page_tables.items() } - for layer, value in metadata._csa2_swa_page_tables.items(): + for layer, value in metadata.ext(CSA2Metadata)._csa2_swa_page_tables.items(): assert ( value.data_ptr() == metadata.kv_cache_block_offsets[manager.layer_offsets[layer], :2, 0, :4].data_ptr() @@ -1787,14 +1902,14 @@ def refresh_and_resolve(): remap(shift) conversions.clear() if direct: - metadata.prepare_csa2() + metadata.ext(CSA2Metadata).prepare_csa2() else: _prepare(metadata, order, [130, 131], [2, 1]) assert len(conversions) == 1 assert manager._swa_publication is manager._swa_publication_token is None assert pointers == { layer: (value.data_ptr(), value.stride()) - for layer, value in metadata._csa2_swa_page_tables.items() + for layer, value in metadata.ext(CSA2Metadata)._csa2_swa_page_tables.items() } assert output_pointers == tuple(value.data_ptr() for value in _swa(metadata, 0)) graph.replay() @@ -1809,7 +1924,7 @@ def refresh_and_resolve(): assert len(conversions) == 2 assert pointers == { layer: (value.data_ptr(), value.stride()) - for layer, value in metadata._csa2_swa_page_tables.items() + for layer, value in metadata.ext(CSA2Metadata)._csa2_swa_page_tables.items() } _assert_swa_cpu_reference(metadata, manager, requests, [130, 131], [2, 1], [0, 0]) finally: @@ -1834,15 +1949,17 @@ def test_swa_layer_reference_matches_scalar_oracle(rows, window): if rows > 256: # Request 1 starts at floor 127: page 0 is history below the floor, never a write target. tables[1][1, 0] = -1 - metadata = CSA2TrtllmMetadata.__new__(CSA2TrtllmMetadata) + metadata = _bare_csa2_metadata() vars(metadata).update( is_cuda_graph=False, - _csa2_buffers={}, kv_lens_cuda=torch.empty(0, dtype=torch.int32), kv_cache_manager=SimpleNamespace( tokens_per_block=block, layout=SimpleNamespace(window_size=window, compress_ratios=ratios), ), + ) + vars(metadata.ext(CSA2Metadata)).update( + _csa2_buffers={}, csa2_positions=torch.tensor(positions, dtype=torch.int32), csa2_token_requests=requests_host, csa2_replay_start_positions=torch.tensor(floors, dtype=torch.int64), @@ -1905,7 +2022,7 @@ def test_swa_refresh_preserves_per_layer_page_domains(device: str) -> None: ) for pages, shape, ratio in zip(pages_host, shapes, ratios) ) - CSA2TrtllmMetadata._refresh_swa_tensor_outputs( + CSA2Metadata._refresh_swa_tensor_outputs( positions, requests, floors, block, window, layer_inputs ) for (_, reads, writes, visible, ratio), pages, (table_rows, width) in zip( @@ -1938,51 +2055,52 @@ def test_native_staging_pools_cover_tma_fetch_window( dtype: torch.dtype, shared_rows: int | None ) -> None: q = torch.empty(1, 8, 512, device="cuda", dtype=torch.bfloat16) - metadata = CSA2TrtllmMetadata.for_query_tile( + metadata = CSA2Metadata.for_query_tile( q, 0, staging_dtype=dtype, context_lengths=[1] if shared_rows is not None else None, shared_rows=shared_rows, ) + csa2 = metadata.ext(CSA2Metadata) if shared_rows is None: - assert metadata.swa_pool.shape == (1, 128, 512) - assert metadata.extra_pool.shape == (1, 1, 512) + assert csa2.swa_pool.shape == (1, 128, 512) + assert csa2.extra_pool.shape == (1, 1, 512) else: - assert metadata.shared_pool.shape == (shared_rows, 512) - assert metadata.swa_pool is metadata.extra_pool is metadata.shared_pool - for pool in (metadata.swa_pool, metadata.extra_pool): + assert csa2.shared_pool.shape == (shared_rows, 512) + assert csa2.swa_pool is csa2.extra_pool is csa2.shared_pool + for pool in (csa2.swa_pool, csa2.extra_pool): assert pool.dtype == dtype and pool.is_contiguous() # TMA's encoded row domain is larger than these small logical pools. # Padding must be mapped memory retained by the view itself. backing = pool.untyped_storage() available = backing.nbytes() - (pool.data_ptr() - backing.data_ptr()) assert available >= 128 * 1024 - assert metadata.host_kv_cache_pool_pointers[0, 0].item() == metadata.swa_pool.data_ptr() + assert metadata.host_kv_cache_pool_pointers[0, 0].item() == csa2.swa_pool.data_ptr() @skip_pre_blackwell @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_context_tile_cache_retains_latest_and_preserves_graph_and_dtype(): - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=64) + source = CSA2TrtllmAttention.create_metadata(max_num_requests=2, max_num_tokens=64) q = torch.zeros(64, 8, 512, dtype=torch.bfloat16, device="cuda") - generation = source.get_query_tile_metadata(q[:2], 0) - generation.swa_pool.zero_() + generation = source.ext(CSA2Metadata).get_query_tile_metadata(q[:2], 0) + generation.ext(CSA2Metadata).swa_pool.zero_() generation.workspace.resize_(32).fill_(47) generation_workspace = generation.workspace generation_workspace_pointer = generation.workspace.data_ptr() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - assert source.get_query_tile_metadata(q[:2], 0) is generation - captured = generation.swa_pool.view(torch.uint8).clone() + assert source.ext(CSA2Metadata).get_query_tile_metadata(q[:2], 0) is generation + captured = generation.ext(CSA2Metadata).swa_pool.view(torch.uint8).clone() captured_workspace = generation.workspace.clone() - generation_pointer = generation.swa_pool.data_ptr() + generation_pointer = generation.ext(CSA2Metadata).swa_pool.data_ptr() source._num_ctx_tokens = source._num_tokens = 64 source._num_contexts = 2 - source.csa2_request_query_ranges = ((0, 32), (32, 64)) - source.csa2_request_start_positions = (0, 512) - source.csa2_request_lengths = (32, 32) + source.ext(CSA2Metadata).csa2_request_query_ranges = ((0, 32), (32, 64)) + source.ext(CSA2Metadata).csa2_request_start_positions = (0, 512) + source.ext(CSA2Metadata).csa2_request_lengths = (32, 32) protected = {} for heads, extra, dtype in ( (8, 0, torch.bfloat16), @@ -1994,10 +2112,10 @@ def test_context_tile_cache_retains_latest_and_preserves_graph_and_dtype(): workspace = None high_water = 0 for count in (16, 7, 12, 16, 5, 48, 9, 48): - frame = source.get_query_tile_metadata( + frame = source.ext(CSA2Metadata).get_query_tile_metadata( queries[:count], extra, query_start=0, staging_dtype=dtype ) - assert frame.swa_pool.dtype == dtype + assert frame.ext(CSA2Metadata).swa_pool.dtype == dtype assert frame.num_tokens == count and frame.num_ctx_tokens == count assert frame.num_contexts == (2 if count > 32 else 1) if workspace is None: @@ -2007,27 +2125,29 @@ def test_context_tile_cache_retains_latest_and_preserves_graph_and_dtype(): high_water = max(high_water, 32 + count) frame.workspace.resize_(high_water).fill_(count) for previous in protected.values(): - assert any(tile is previous for tile in source._csa2_query_tiles.values()) + assert any( + tile is previous for tile in source.ext(CSA2Metadata)._csa2_query_tiles.values() + ) assert previous.workspace is not frame.workspace # Only the latest shape of an eager context family stays resident. assert ( sum( 1 - for key, tile in source._csa2_query_tiles.items() + for key, tile in source.ext(CSA2Metadata)._csa2_query_tiles.items() if tile.num_contexts - and tile.swa_pool.dtype == dtype + and tile.ext(CSA2Metadata).swa_pool.dtype == dtype and key[1] == heads and key[3] == (extra + 127) // 128 * 128 ) == 1 ) - assert generation.swa_pool.data_ptr() == generation_pointer + assert generation.ext(CSA2Metadata).swa_pool.data_ptr() == generation_pointer assert generation.workspace is generation_workspace assert generation.workspace.data_ptr() == generation_workspace_pointer assert frame.workspace is not generation_workspace - assert source.get_query_tile_metadata(q[:2], 0) is generation - generation.swa_pool.fill_(count) - expected = generation.swa_pool.view(torch.uint8).clone() + assert source.ext(CSA2Metadata).get_query_tile_metadata(q[:2], 0) is generation + generation.ext(CSA2Metadata).swa_pool.fill_(count) + expected = generation.ext(CSA2Metadata).swa_pool.view(torch.uint8).clone() graph.replay() torch.testing.assert_close(captured, expected, atol=0, rtol=0) torch.testing.assert_close( @@ -2042,39 +2162,52 @@ def test_context_tile_cache_retains_latest_and_preserves_graph_and_dtype(): @torch.inference_mode() def test_shared_context_domains_follow_owner_source_batch_and_replay(manager_requests): manager, requests = manager_requests - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) - source.set_source_batch([200, 200], [0, 0]) + source = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) + source.ext(CSA2Metadata).set_source_batch([200, 200], [0, 0]) _prepare(source, requests, [130, 131], [2, 1]) q = torch.zeros(3, 8, 512, device="cuda", dtype=torch.bfloat16) - child = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) + child = source.ext(CSA2Metadata).get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) workspace = child.workspace workspace.resize_(64).fill_(11) - other_owner = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=3) + other_owner = source.ext(CSA2Metadata).get_query_tile_metadata( + q, 512, query_start=0, layer_idx=3 + ) other_workspace = other_owner.workspace other_workspace.resize_(96).fill_(17) assert other_workspace is not workspace - first = child.shared_plan + first = child.ext(CSA2Metadata).shared_plan assert first is not None - assert first.main_rows == sum(source._csa2_main_domain_counts[1]) + assert first.main_rows == sum(source.ext(CSA2Metadata)._csa2_main_domain_counts[1]) assert first.main_rows >= 200 # source visibility, plus admitted page capacity assert first.swa_rows == 257 assert first.swa_starts.tolist() == [3, 4] assert child.max_num_requests == 2 and child.num_tokens == 3 - assert child.swa_pool.data_ptr() == child.extra_pool.data_ptr() == child.shared_pool.data_ptr() - assert child.shared_pool.shape == (1 + first.swa_rows + first.main_rows, 512) - reuse = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=2) + assert ( + child.ext(CSA2Metadata).swa_pool.data_ptr() + == child.ext(CSA2Metadata).extra_pool.data_ptr() + == child.ext(CSA2Metadata).shared_pool.data_ptr() + ) + assert child.ext(CSA2Metadata).shared_pool.shape == (1 + first.swa_rows + first.main_rows, 512) + reuse = source.ext(CSA2Metadata).get_query_tile_metadata(q, 512, query_start=0, layer_idx=2) assert reuse is child - assert reuse.shared_plan.main_pages.data_ptr() == first.main_pages.data_ptr() - assert reuse.shared_plan.swa_pages.data_ptr() != first.swa_pages.data_ptr() + assert reuse.ext(CSA2Metadata).shared_plan.main_pages.data_ptr() == first.main_pages.data_ptr() + assert reuse.ext(CSA2Metadata).shared_plan.swa_pages.data_ptr() != first.swa_pages.data_ptr() # Pure decoder replay has no newly projected GLOBAL rows, but reads cached C. - source.set_swa_bounded_replay([130, 130], decoder=True) + source.ext(CSA2Metadata).set_swa_bounded_replay([130, 130], decoder=True) _prepare(source, requests, [2, 2], [128, 128]) q = torch.zeros(256, 8, 512, device="cuda", dtype=torch.bfloat16) - child = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) - assert child.shared_plan is not None and child.shared_plan.main_rows >= 130 + child = source.ext(CSA2Metadata).get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) + assert ( + child.ext(CSA2Metadata).shared_plan is not None + and child.ext(CSA2Metadata).shared_plan.main_rows >= 130 + ) assert child.workspace is child.cuda_graph_workspace is workspace assert child.workspace.numel() == 64 - other_owner = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=3) + other_owner = source.ext(CSA2Metadata).get_query_tile_metadata( + q, 512, query_start=0, layer_idx=3 + ) assert other_owner.workspace is other_owner.cuda_graph_workspace is other_workspace assert other_workspace.numel() == 96 torch.testing.assert_close(workspace, torch.full_like(workspace, 11), atol=0, rtol=0) @@ -2082,12 +2215,14 @@ def test_shared_context_domains_follow_owner_source_batch_and_replay(manager_req other_workspace, torch.full_like(other_workspace, 17), atol=0, rtol=0 ) assert child.max_num_requests == 2 and child.num_tokens == 256 - assert child.shared_plan.swa_starts.tolist() == [2, 2] - assert child.shared_pool.shape[0] < q.shape[0] * 640 + assert child.ext(CSA2Metadata).shared_plan.swa_starts.tolist() == [2, 2] + assert child.ext(CSA2Metadata).shared_pool.shape[0] < q.shape[0] * 640 # Captured context keeps its existing independent-generation frame. source.is_cuda_graph = True - selected = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) - assert selected.shared_plan is None and selected.swa_pool.shape == (256, 128, 512) + selected = source.ext(CSA2Metadata).get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) + assert selected.ext(CSA2Metadata).shared_plan is None and selected.ext( + CSA2Metadata + ).swa_pool.shape == (256, 128, 512) assert selected.num_contexts == 0 assert selected.workspace is not workspace and selected.workspace is not other_workspace @@ -2105,13 +2240,15 @@ def test_shared_context_covers_upward_live_endpoint_inside_published_pages(manag from .test_backend import _backend manager, requests = manager_requests - source = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) - source.set_source_batch([200, 200], [0, 0]) + source = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) + source.ext(CSA2Metadata).set_source_batch([200, 200], [0, 0]) _prepare(source, requests, [130, 131], [2, 1]) source.kv_lens_cuda.fill_(260) source.on_update_kv_lens() logical = torch.tensor([[110], [111], [112]], dtype=torch.int64, device="cuda") - slots = source.global_slot_tile(1, 0, 3, logical) + slots = source.ext(CSA2Metadata).global_slot_tile(1, 0, 3, logical) assert torch.all(slots >= 0) manager.get_swa_buffer(1).zero_() manager.get_main_buffer(1).zero_() @@ -2122,10 +2259,10 @@ def test_shared_context_covers_upward_live_endpoint_inside_published_pages(manag torch.zeros(3, 128, device="cuda", dtype=torch.bfloat16), ) q = torch.zeros(3, 8, 512, device="cuda", dtype=torch.bfloat16) - shared = source.get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) - assert shared.shared_plan is not None - assert shared.shared_plan.main_rows > 2 * (200 // 2) - selected = CSA2TrtllmMetadata.for_query_tile(q, 512, context_lengths=[2, 1]) + shared = source.ext(CSA2Metadata).get_query_tile_metadata(q, 512, query_start=0, layer_idx=1) + assert shared.ext(CSA2Metadata).shared_plan is not None + assert shared.ext(CSA2Metadata).shared_plan.main_rows > 2 * (200 // 2) + selected = CSA2Metadata.for_query_tile(q, 512, context_lengths=[2, 1]) backend = _backend(8) def arguments(): @@ -2156,12 +2293,14 @@ def test_compressor_readiness_uses_current_rows(manager_requests): from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheRole manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) base = manager.host_kv_cache_block_offsets row = int(manager.index_mapper.get_copy_index([requests[0].py_request_id], 1, 1)[0]) def prepare_source(length): - metadata.set_source_batch([length, 0], [127, 0]) + metadata.ext(CSA2Metadata).set_source_batch([length, 0], [127, 0]) _prepare(metadata, requests, [4, 4], [1, 1]) prepare_source(2) # Source spans logical pages 0 and 1. @@ -2185,19 +2324,22 @@ def prepare_source(length): @torch.inference_mode() def test_first_decode_after_context_takes_full_prepare(manager_requests): manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) # A one-token context chunk (a full prefix hit) and then its first decode # step: the context geometry must not be reused for the generation rows. _prepare(metadata, requests, [126, 129], [1, 1]) _prepare_decode_metadata(metadata, requests, [127, 130]) - assert metadata.csa2_num_context_requests == 0 - assert metadata.csa2_decode_row_requests is not None + assert metadata.ext(CSA2Metadata).csa2_num_context_requests == 0 + assert metadata.ext(CSA2Metadata).csa2_decode_row_requests is not None def _decode_derived_outputs(metadata): # SWA outputs are resolved when a forward enters its first layer. - metadata._ensure_swa_slots() - outputs = {"positions": metadata.csa2_positions} + csa2 = metadata.ext(CSA2Metadata) + csa2._ensure_swa_slots() + outputs = {"positions": csa2.csa2_positions} for name in ( "csa2_swa_indices", "csa2_swa_write_slots", @@ -2205,8 +2347,8 @@ def _decode_derived_outputs(metadata): "csa2_main_write_slots", "_csa2_compressed_positions", ): - outputs.update({(name, layer): tensor for layer, tensor in getattr(metadata, name).items()}) - for owner, batch in metadata._csa2_compression.items(): + outputs.update({(name, layer): tensor for layer, tensor in getattr(csa2, name).items()}) + for owner, batch in csa2._csa2_compression.items(): for name in ("start_positions", "kv_lengths", "cu_compressed_lengths"): outputs[owner, name] = getattr(batch, name) return outputs @@ -2220,9 +2362,9 @@ def _prepare_decode_metadata(metadata, requests, starts, *, deferred=False): metadata.num_contexts = 0 metadata.prompt_lens = [request.prompt_len for request in requests] metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=starts) - with metadata.defer_cuda_graph_decode_outputs() if deferred else nullcontext(): + with metadata.ext(CSA2Metadata).defer_decode_outputs() if deferred else nullcontext(): metadata.prepare() - assert not metadata._csa2_defer_decode_outputs + assert not metadata.ext(CSA2Metadata)._csa2_defer_decode_outputs @skip_pre_blackwell @@ -2230,11 +2372,13 @@ def _prepare_decode_metadata(metadata, requests, starts, *, deferred=False): @torch.inference_mode() def test_steady_decode_refreshes_page_tables_across_pages(manager_requests, monkeypatch): manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) refreshed = [] - refresh = metadata._refresh_steady_page_tables + refresh = metadata.ext(CSA2Metadata)._refresh_steady_page_tables monkeypatch.setattr( - metadata, + metadata.ext(CSA2Metadata), "_refresh_steady_page_tables", lambda *args: refreshed.append(args[2][0]) or refresh(*args), ) @@ -2243,15 +2387,17 @@ def test_steady_decode_refreshes_page_tables_across_pages(manager_requests, monk for start in range(125, 131): starts = [start, start + 3] _prepare_decode_metadata(metadata, requests, starts) - reference = CSA2TrtllmMetadata( + reference = CSA2TrtllmAttention.create_metadata( max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager ) _prepare_decode_metadata(reference, requests, starts) observed = _decode_derived_outputs(metadata) for name, tensor in _decode_derived_outputs(reference).items(): torch.testing.assert_close(observed[name], tensor) - for owner, table in reference.csa2_global_page_tables.items(): - torch.testing.assert_close(metadata.csa2_global_page_tables[owner], table) + for owner, table in reference.ext(CSA2Metadata).csa2_global_page_tables.items(): + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_global_page_tables[owner], table + ) assert refreshed == [128] @@ -2268,23 +2414,25 @@ def test_decode_graph_defers_derived_uploads_and_refreshes_all_outputs( from tensorrt_llm.bindings import SamplingConfig manager, requests = manager_requests - reference = CSA2TrtllmMetadata( + reference = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) + base = CSA2TrtllmAttention.create_metadata( max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager ) - base = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) metadata = base.create_cuda_graph_metadata(batch_size) copied = [] - copy_host = metadata._copy_host + copy_host = metadata.ext(CSA2Metadata)._copy_host def record_copy(key, destination, value): copied.append(key) return copy_host(key, destination, value) - monkeypatch.setattr(metadata, "_copy_host", record_copy) + monkeypatch.setattr(metadata.ext(CSA2Metadata), "_copy_host", record_copy) selected = requests[:batch_size] starts = [127] * batch_size _prepare_decode_metadata(metadata, selected, starts, deferred=True) - assert metadata._csa2_deferred_decode_outputs + assert metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs inputs = object.__new__(DecoderInputPreparer) inputs.guided_decoder = None pointers = { @@ -2319,7 +2467,7 @@ def forward(): } copied.clear() _prepare_decode_metadata(metadata, selected, starts, deferred=True) - assert metadata._csa2_deferred_decode_outputs + assert metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs if steady_start > 130: assert not copied for tensor in _decode_derived_outputs(metadata).values(): @@ -2383,7 +2531,7 @@ def forward(): # Graph metadata used by a direct caller must still be fully initialized. _prepare_decode_metadata(metadata, selected, starts) - assert not metadata._csa2_deferred_decode_outputs + assert not metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs for name, tensor in _decode_derived_outputs(metadata).items(): torch.testing.assert_close(tensor, expected[name], atol=0, rtol=0) @@ -2404,7 +2552,9 @@ def forward(): @torch.inference_mode() def test_decode_graph_deferral_keeps_unsupported_prepare_paths_complete(manager_requests, fallback): manager, requests = manager_requests - metadata = CSA2TrtllmMetadata(max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=2, max_num_tokens=1024, kv_cache_manager=manager + ) metadata.is_cuda_graph = fallback != "eager" metadata.request_ids = [request.py_request_id for request in requests] metadata.seq_lens = torch.ones(2, dtype=torch.int32) @@ -2412,23 +2562,23 @@ def test_decode_graph_deferral_keeps_unsupported_prepare_paths_complete(manager_ metadata.prompt_lens = [257, 257] metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[127, 128]) if fallback == "source": - metadata.set_source_batch([1, 1], [127, 128]) + metadata.ext(CSA2Metadata).set_source_batch([1, 1], [127, 128]) if fallback == "remote_tail": metadata.csa2_remote_tail_mode = "decode" if fallback == "speculative": - metadata.is_spec_decoding_enabled = True - with metadata.defer_cuda_graph_decode_outputs(): + metadata.ext(SpecDecodingMetadata).is_spec_decoding_enabled = True + with metadata.ext(CSA2Metadata).defer_decode_outputs(): metadata.prepare() - assert not metadata._csa2_deferred_decode_outputs - assert metadata.csa2_positions.tolist() == [127, 128] + assert not metadata.ext(CSA2Metadata)._csa2_deferred_decode_outputs + assert metadata.ext(CSA2Metadata).csa2_positions.tolist() == [127, 128] _assert_swa_cpu_reference(metadata, manager, requests, [127, 128], [1, 1], [0, 0]) @pytest.mark.cpu_only def test_decode_graph_prepare_scope_restores_proof_after_exception(): - metadata = object.__new__(CSA2TrtllmMetadata) + metadata = _bare_csa2_metadata() with pytest.raises(RuntimeError, match="prepare failed"): - with metadata.defer_cuda_graph_decode_outputs(): - assert metadata._csa2_defer_decode_outputs + with metadata.ext(CSA2Metadata).defer_decode_outputs(): + assert metadata.ext(CSA2Metadata)._csa2_defer_decode_outputs raise RuntimeError("prepare failed") - assert not metadata._csa2_defer_decode_outputs + assert not metadata.ext(CSA2Metadata)._csa2_defer_decode_outputs diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_metadata_uploads.py b/tests/unittest/_torch/attention/sparse/csa2/test_metadata_uploads.py index 4345bc477935..5804e0d8a133 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_metadata_uploads.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_metadata_uploads.py @@ -9,7 +9,9 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata skip_pre_blackwell = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() < (10, 0), @@ -28,17 +30,28 @@ ] +def _bare_csa2_metadata() -> TrtllmAttentionMetadata: + """Metadata built without ``__init__``, carrying a bare CSA2 extension.""" + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(CSA2Metadata(metadata)) + metadata.replace_ext(SpecDecodingMetadata(metadata)) + return metadata + + def _metadata(device): - metadata = object.__new__(CSA2TrtllmMetadata) + metadata = _bare_csa2_metadata() metadata.is_cuda_graph = False metadata.num_sparse_topk = 256 - for name in ("query_lens_host", "kv_lens", "prompt_lens_cpu", "host_request_types"): + for name in ("kv_lens", "prompt_lens_cpu", "host_request_types"): setattr(metadata, name, torch.empty(4, dtype=torch.int32, device="cpu")) - for name in ("query_lens_device", "kv_lens_cuda", "prompt_lens_cuda"): + for name in ("kv_lens_cuda", "prompt_lens_cuda"): setattr(metadata, name, torch.empty(4, dtype=torch.int32, device=device)) metadata.host_total_kv_lens = torch.zeros(2, dtype=torch.int32, device="cpu") - metadata.prepared_cu_q = torch.empty(5, dtype=torch.int32, device=device) - metadata.prepared_cu_kv = torch.empty(5, dtype=torch.int32, device=device) + csa2 = metadata.ext(CSA2Metadata) + csa2.query_lens_host = torch.empty(4, dtype=torch.int32, device="cpu") + csa2.query_lens_device = torch.empty(4, dtype=torch.int32, device=device) + csa2.prepared_cu_q = torch.empty(5, dtype=torch.int32, device=device) + csa2.prepared_cu_kv = torch.empty(5, dtype=torch.int32, device=device) return metadata @@ -51,7 +64,7 @@ def test_context_uploads_refresh_changing_geometry(device): if stream is not None: torch.cuda._sleep(2_000_000) for lengths in ([2, 1], [1, 2], [3], [1, 1, 1], [2, 1], [2, 1]): - metadata._bind_context_tile(list(lengths)) + metadata.ext(CSA2Metadata)._bind_context_tile(list(lengths)) observed.append( ( list(lengths), @@ -108,34 +121,42 @@ def test_shared_context_uploads_preserve_owners_and_live_positions(device): for offset in range(length) ] metadata._num_ctx_tokens = sum(lengths) - metadata.csa2_num_context_requests = requests - metadata.csa2_request_lengths = lengths - metadata.csa2_positions = torch.tensor(positions, dtype=torch.int32, device=device) - metadata.csa2_replay_start_positions = torch.tensor(floors, device=device) - metadata.csa2_token_requests = torch.tensor( + metadata.ext(CSA2Metadata).csa2_num_context_requests = requests + metadata.ext(CSA2Metadata).csa2_request_lengths = lengths + metadata.ext(CSA2Metadata).csa2_positions = torch.tensor( + positions, dtype=torch.int32, device=device + ) + metadata.ext(CSA2Metadata).csa2_replay_start_positions = torch.tensor( + floors, device=device + ) + metadata.ext(CSA2Metadata).csa2_token_requests = torch.tensor( [request for request, length in enumerate(lengths) for _ in range(length)], device=device, ) - metadata._csa2_main_domain_counts = { + metadata.ext(CSA2Metadata)._csa2_main_domain_counts = { owner: [step + owner + request + 1 for request in range(requests)] for owner in (0, 1) } - metadata._csa2_shared_domains = {} - metadata._csa2_swa_page_tables = { + metadata.ext(CSA2Metadata)._csa2_shared_domains = {} + metadata.ext(CSA2Metadata)._csa2_swa_page_tables = { layer: torch.full((requests, 2), 20 + layer + step, device=device) for layer in (0, 1, 2) } - metadata.csa2_global_page_tables = { + metadata.ext(CSA2Metadata).csa2_global_page_tables = { owner: torch.full((requests, 2), 40 + owner + step, device=device) for owner in (0, 1) } q = torch.empty(sum(lengths), 1, 512, device=device) - plans = [metadata._shared_context_plan(q, 64, layer) for layer in (0, 1, 2)] + plans = [ + metadata.ext(CSA2Metadata)._shared_context_plan(q, 64, layer) for layer in (0, 1, 2) + ] swa_offsets = list(accumulate((length + 127 for length in lengths), initial=1)) # Read all owner plans only after the next owner's uploads have been submitted. for owner, plan in zip((0, 1, None), plans, strict=True): main_counts = ( - [0] * requests if owner is None else metadata._csa2_main_domain_counts[owner] + [0] * requests + if owner is None + else metadata.ext(CSA2Metadata)._csa2_main_domain_counts[owner] ) expected = ( [max(0, start - 127, floor) for start, floor in zip(first_positions, floors)], diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_module.py b/tests/unittest/_torch/attention/sparse/csa2/test_module.py index ea75644d721a..1dfd2f4a7e4b 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_module.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_module.py @@ -11,7 +11,9 @@ from pydantic import TypeAdapter from torch import nn +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import ( CSA2Layout, CSA2Mode, @@ -106,30 +108,34 @@ def write_layer_rows(layer, slots, values, transform=None, *, global_rows=None, select_global_source=lambda layer, local, global_hidden: local, enter_layer=lambda layer: entered.append(layer.layer_idx), ) + # The stubs serve as both the core metadata and its CSA2 extension. + source.ext = source.find_ext = lambda cls: source def phase_metadata(query, extra_width, query_start, *, staging_dtype, layer_idx=None): stop = query_start + query.shape[0] assert stop <= (contexts if query_start < contexts else count) torch.testing.assert_close(query.flatten(1), hidden[query_start:stop], atol=0, rtol=0) - return SimpleNamespace( + tile = SimpleNamespace( num_contexts=int(query_start < contexts), num_tokens=query.shape[0], num_query_heads=1, swa_pool=torch.empty(0, dtype=staging_dtype), ) + tile.ext = lambda cls: tile + return tile def map_indices(layer, start, end, logical, *, visible_lengths): torch.testing.assert_close(logical[:, 0], row_ids[start:end], atol=0, rtol=0) torch.testing.assert_close(visible_lengths, row_ids[start:end] + 1, atol=0, rtol=0) return logical - source.get_query_tile_metadata = phase_metadata - source.global_slot_tile = map_indices + source.ext(CSA2Metadata).get_query_tile_metadata = phase_metadata + source.ext(CSA2Metadata).global_slot_tile = map_indices def indexer(state, start, end): index_calls.append((start, end)) torch.testing.assert_close(state.index_q.flatten(), positions.float(), atol=0, rtol=0) - source.csa2_indices[0] = row_ids[:, None] + source.ext(CSA2Metadata).csa2_indices[0] = row_ids[:, None] backend = SimpleNamespace( sparse_params=CSA2Params(layout=layout, compute_backend=implementation), @@ -270,20 +276,21 @@ def test_ced_owner_is_written_once_before_narrowed_query(monkeypatch, module_cac PositionalEmbeddingParams, RopeParams, ) - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.module import DeepseekV41Attention from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm.functional import PositionEmbeddingType manager, request = module_cache - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager + ) metadata.request_ids = [request.py_request_id] metadata.num_contexts = 1 metadata.prompt_lens = [4] metadata.seq_lens = torch.tensor([4], dtype=torch.int32) metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]) metadata.prepare() - metadata.begin_model_forward() + metadata.begin_forward() torch.manual_seed(4126) with torch.device("cuda"): model = DeepseekV41Attention( @@ -330,7 +337,7 @@ def unexpected_projection(*args, **kwargs): metadata.prompt_lens = [1] metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[3]) metadata.prepare() - actual = model(hidden[-1:], metadata.csa2_positions, metadata) + actual = model(hidden[-1:], metadata.ext(CSA2Metadata).csa2_positions, metadata) torch.cuda.synchronize() assert actual.shape == (1, 32) and torch.isfinite(actual).all() assert writes == [0] @@ -344,7 +351,6 @@ def _check_ratio2_module_partial_groups(monkeypatch, module_cache, overlap, impl PositionalEmbeddingParams, RopeParams, ) - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.module import DeepseekV41Attention from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import ( pack_rows, @@ -360,7 +366,9 @@ def _check_ratio2_module_partial_groups(monkeypatch, module_cache, overlap, impl torch.manual_seed(818) manager, request = module_cache layout = manager.layout - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager + ) metadata.request_ids = [request.py_request_id] metadata.num_contexts = 1 metadata.prompt_lens = [4] @@ -430,12 +438,14 @@ def record_rope_rows(x, *args): metadata.seq_lens = torch.tensor([1], dtype=torch.int32, device="cpu") metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[token]) metadata.prepare() - positions = metadata.csa2_positions - assert metadata.get_compression_batch(0).cu_compressed_lengths.tolist() == [ + positions = metadata.ext(CSA2Metadata).csa2_positions + assert metadata.ext(CSA2Metadata).get_compression_batch( + 0 + ).cu_compressed_lengths.tolist() == [ 0, int(new_rows), ] - assert metadata.get_compression_batch(0).output_rows == 1 + assert metadata.ext(CSA2Metadata).get_compression_batch(0).output_rows == 1 old_main = manager.get_main_buffer(0).clone() old_index = manager.get_index_pages(0).clone() # Fixed-capacity runtime padding does not generate an empty latent, @@ -780,7 +790,6 @@ def test_fused_native_index_q_and_graph(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @torch.inference_mode() def test_native_module_auxiliary_stream_and_graph(monkeypatch, module_cache): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.metadata import KVCacheParams torch.manual_seed(946) @@ -792,7 +801,9 @@ def test_native_module_auxiliary_stream_and_graph(monkeypatch, module_cache): "tensorrt_llm._torch.attention.backends.sparse.csa2.module.do_multi_stream", lambda: enabled[0], ) - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager + ) metadata.request_ids = [request.py_request_id] metadata.num_contexts = 1 metadata.prompt_lens = [4] @@ -802,16 +813,16 @@ def test_native_module_auxiliary_stream_and_graph(monkeypatch, module_cache): # its second half while comparing identical cache/history state. metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]) metadata.prepare() - model(hidden, metadata.csa2_positions, metadata) + model(hidden, metadata.ext(CSA2Metadata).csa2_positions, metadata) metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[1]) metadata.num_contexts = 0 metadata.is_cuda_graph = True metadata.prepare() - positions = metadata.csa2_positions + positions = metadata.ext(CSA2Metadata).csa2_positions hidden.normal_() def run(): - metadata.reset_routing() + metadata.ext(CSA2Metadata).reset_routing() return model(hidden, positions, metadata) serial = run() @@ -838,7 +849,6 @@ def run(): @pytest.mark.parametrize("overlap", [False, True]) @torch.inference_mode() def test_native_module_fixed_context_graph(monkeypatch, module_cache, overlap): - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.metadata import KVCacheParams torch.manual_seed(947) @@ -849,14 +859,16 @@ def test_native_module_fixed_context_graph(monkeypatch, module_cache, overlap): monkeypatch.setattr( "tensorrt_llm._torch.attention.backends.sparse.csa2.module.do_multi_stream", lambda: overlap ) - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager + ) metadata.request_ids = [request.py_request_id] metadata.num_contexts = 1 metadata.prompt_lens = [4] metadata.seq_lens = torch.tensor([4], dtype=torch.int32, device="cpu") metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]) metadata.prepare() - positions = metadata.csa2_positions + positions = metadata.ext(CSA2Metadata).csa2_positions hidden = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) phases = [] original = model.backend.forward @@ -868,7 +880,7 @@ def observe(*args, **kwargs): monkeypatch.setattr(model.backend, "forward", observe) def run(): - metadata.reset_routing() + metadata.ext(CSA2Metadata).reset_routing() return model(hidden, positions, metadata) baseline = run() @@ -957,7 +969,6 @@ def test_packed_output_inverse_rope_module_and_graph( from dataclasses import replace - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.metadata import KVCacheParams torch.manual_seed(950) @@ -969,7 +980,9 @@ def test_packed_output_inverse_rope_module_and_graph( grouped_block=block, sparse_options={"use_packed_sparse_attention": True, "fuse_packed_output_rope": True}, ) - metadata = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager) + metadata = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=8, kv_cache_manager=manager + ) metadata.request_ids = [request.py_request_id] metadata.num_contexts = 1 metadata.prompt_lens = [4] @@ -978,7 +991,7 @@ def test_packed_output_inverse_rope_module_and_graph( metadata.is_cuda_graph = True metadata.prepare() hidden = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) - positions = metadata.csa2_positions + positions = metadata.ext(CSA2Metadata).csa2_positions inverse_calls = [] native_rope = torch.ops.trtllm.mla_rope_inplace @@ -998,7 +1011,7 @@ def observe_projection(*args, **kwargs): monkeypatch.setattr(model, "_project_output", observe_projection) def run(): - metadata.reset_routing() + metadata.ext(CSA2Metadata).reset_routing() return model(hidden, positions, metadata) enabled = model.backend.sparse_params @@ -1031,7 +1044,6 @@ def run(): @torch.inference_mode() def test_two_rank_module_projection_parity(monkeypatch, quantization): from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.nccl_window_tensor_scope import nccl_window_tensor_scope @@ -1108,7 +1120,7 @@ def fresh_metadata(mapping): assert manager.prepare_context(request) assert manager.resize_context(request, request.context_chunk_size) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=3, kv_cache_manager=manager ) metadata.request_ids = [request.py_request_id] @@ -1127,8 +1139,10 @@ def fresh_metadata(mapping): with nccl_window_tensor_scope(hidden): # Both ranks execute both calls, including any shared startup/tuning # collectives. The reference has replicated heads and no TP reduction. - actual = model(hidden, tp_metadata.csa2_positions, tp_metadata) - expected = reference(hidden, reference_metadata.csa2_positions, reference_metadata) + actual = model(hidden, tp_metadata.ext(CSA2Metadata).csa2_positions, tp_metadata) + expected = reference( + hidden, reference_metadata.ext(CSA2Metadata).csa2_positions, reference_metadata + ) torch.cuda.synchronize() assert reductions == [torch.Size([3, 128])] # Each TP shard rounds its O-B accumulation to BF16 before summation; @@ -1147,7 +1161,6 @@ def replay_module_factory(): RopeParams, ) from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata from tensorrt_llm._torch.attention.backends.sparse.csa2.module import DeepseekV41Attention from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -1191,7 +1204,7 @@ def create(ratio, input_length=8): assert manager.prepare_context(request) assert manager.resize_context(request, request.context_chunk_size) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=16, kv_cache_manager=manager ) metadata.request_ids = [request.py_request_id] @@ -1234,7 +1247,7 @@ def _prepare_replay_step(metadata, start, end, cached=None, decoder=False): metadata.seq_lens = torch.tensor([end - start], dtype=torch.int32, device="cpu") metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[start]) if cached is not None: - metadata.set_swa_bounded_replay([cached], decoder=decoder) + metadata.ext(CSA2Metadata).set_swa_bounded_replay([cached], decoder=decoder) metadata.prepare() @@ -1301,8 +1314,8 @@ def test_encoder_bounded_replay_preserves_global_and_truncates_swa( torch.manual_seed(955) model, manager, metadata, hidden = replay_module_factory(2) _prepare_replay_step(metadata, 0, 5) - model(hidden[:5], metadata.csa2_positions, metadata) - old_slots = metadata.csa2_main_write_slots[0] + model(hidden[:5], metadata.ext(CSA2Metadata).csa2_positions, metadata) + old_slots = metadata.ext(CSA2Metadata).csa2_main_write_slots[0] old_slots = old_slots[old_slots >= 0].clone() old_main = manager.get_main_buffer(0)[old_slots].clone() old_index = read_index_rows(manager.get_index_pages(0), old_slots).clone() @@ -1320,18 +1333,20 @@ def observe(x, batch=None): monkeypatch.setattr(model.compressor, "forward", observe) _prepare_replay_step(metadata, 1, 8, cached=5) - actual = model(replay, metadata.csa2_positions, metadata) + actual = model(replay, metadata.ext(CSA2Metadata).csa2_positions, metadata) assert len(seen) == 1 and seen[0].shape[0] == 4 torch.testing.assert_close(seen[0], replay[3:], atol=0, rtol=0) torch.testing.assert_close(manager.get_main_buffer(0)[old_slots], old_main, atol=0, rtol=0) torch.testing.assert_close( read_index_rows(manager.get_index_pages(0), old_slots), old_index, atol=0, rtol=0 ) - assert int((metadata.csa2_swa_indices[0][0] >= 0).sum()) == 1 + assert int((metadata.ext(CSA2Metadata).csa2_swa_indices[0][0] >= 0).sum()) == 1 expected_global = torch.cat( (unpack_rows(old_main, 512, "main"), _reference_global_rows(model, replay[3:], 4)) ) - expected = _reference_bounded_output(model, replay, metadata.csa2_positions, expected_global) + expected = _reference_bounded_output( + model, replay, metadata.ext(CSA2Metadata).csa2_positions, expected_global + ) torch.testing.assert_close(actual, expected, atol=0.03, rtol=0.03) @@ -1355,8 +1370,8 @@ def test_pure_replay_skips_global_projection_and_restores_odd_state( torch.manual_seed(956) model, manager, metadata, hidden = replay_module_factory(ratio) _prepare_replay_step(metadata, 0, cached) - model(hidden[:cached], metadata.csa2_positions, metadata) - slots = metadata.csa2_main_write_slots[0] + model(hidden[:cached], metadata.ext(CSA2Metadata).csa2_positions, metadata) + slots = metadata.ext(CSA2Metadata).csa2_main_write_slots[0] slots = slots[slots >= 0].clone() old_main = manager.get_main_buffer(0)[slots].clone() old_index = read_index_rows(manager.get_index_pages(0), slots).clone() @@ -1376,20 +1391,20 @@ def observe(x): monkeypatch.setattr(projection, "forward", observe) _prepare_replay_step(metadata, start, cached, cached=cached, decoder=decoder) - actual = model(replay, metadata.csa2_positions, metadata) + actual = model(replay, metadata.ext(CSA2Metadata).csa2_positions, metadata) assert counts == ([1] if ratio == 2 and cached % 2 and not decoder else []) torch.testing.assert_close(manager.get_main_buffer(0)[slots], old_main, atol=0, rtol=0) torch.testing.assert_close( read_index_rows(manager.get_index_pages(0), slots), old_index, atol=0, rtol=0 ) - assert bool((metadata.csa2_main_write_slots[0] < 0).all()) + assert bool((metadata.ext(CSA2Metadata).csa2_main_write_slots[0] < 0).all()) expected = _reference_bounded_output( - model, replay, metadata.csa2_positions, unpack_rows(old_main, 512, "main") + model, replay, metadata.ext(CSA2Metadata).csa2_positions, unpack_rows(old_main, 512, "main") ) torch.testing.assert_close(actual, expected, atol=0.03, rtol=0.03) if ratio == 2 and cached % 2: _prepare_replay_step(metadata, cached, cached + 1) - batch = metadata.get_compression_batch(0) + batch = metadata.ext(CSA2Metadata).get_compression_batch(0) page, offset = divmod(cached - 1, batch.page_size) tail = replay[-1:] expected_kv, expected_score = ( @@ -1424,8 +1439,8 @@ def observe(x): expected_row = pack_rows( model._rope(reference_latent.unsqueeze(1), reference_position, 1).squeeze(1), "main" ) - model(next_hidden, metadata.csa2_positions, metadata) - new_slots = metadata.csa2_main_write_slots[0] + model(next_hidden, metadata.ext(CSA2Metadata).csa2_positions, metadata) + new_slots = metadata.ext(CSA2Metadata).csa2_main_write_slots[0] new_slots = new_slots[new_slots >= 0] torch.testing.assert_close( manager.get_main_buffer(0)[new_slots], expected_row, atol=0, rtol=0 @@ -1437,13 +1452,14 @@ def observe(x): @torch.inference_mode() def test_bounded_replay_full_module_graph_changes_hit_position(replay_module_factory): from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheRole - from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata torch.manual_seed(957) model, manager, seed_metadata, hidden = replay_module_factory(2, input_length=12) def new_metadata(graph=False): - result = CSA2TrtllmMetadata(max_num_requests=1, max_num_tokens=16, kv_cache_manager=manager) + result = CSA2TrtllmAttention.create_metadata( + max_num_requests=1, max_num_tokens=16, kv_cache_manager=manager + ) result.request_ids = list(seed_metadata.request_ids) result.num_contexts = 1 result.prompt_lens = [12] @@ -1456,8 +1472,8 @@ def new_metadata(graph=False): def seed(cached): _prepare_replay_step(seed_metadata, 0, cached) - model(hidden[:cached], seed_metadata.csa2_positions, seed_metadata) - slots = seed_metadata.csa2_main_write_slots[0] + model(hidden[:cached], seed_metadata.ext(CSA2Metadata).csa2_positions, seed_metadata) + slots = seed_metadata.ext(CSA2Metadata).csa2_main_write_slots[0] return slots[slots >= 0].clone() def drop_local_state(): @@ -1466,8 +1482,8 @@ def drop_local_state(): manager.get_buffers(0, role).fill_(float("nan")) def run_graph_body(): - graph_metadata.reset_routing() - return model(replay_input, graph_metadata.csa2_positions, graph_metadata) + graph_metadata.ext(CSA2Metadata).reset_routing() + return model(replay_input, graph_metadata.ext(CSA2Metadata).csa2_positions, graph_metadata) graph = None pointers = None @@ -1478,7 +1494,9 @@ def run_graph_body(): seed(cached) drop_local_state() _prepare_replay_step(eager_metadata, start, end, cached=cached) - expected = model(replay_input, eager_metadata.csa2_positions, eager_metadata).clone() + expected = model( + replay_input, eager_metadata.ext(CSA2Metadata).csa2_positions, eager_metadata + ).clone() # Reset real GLOBAL contents to this hit position's seeded prefix, # then remove local state again before executing the captured path. slots = seed(cached) @@ -1486,9 +1504,9 @@ def run_graph_body(): old_index = read_index_rows(manager.get_index_pages(0), slots).clone() _prepare_replay_step(graph_metadata, start, end, cached=cached) current_pointers = ( - graph_metadata.csa2_positions.data_ptr(), - graph_metadata.get_compressed_positions(0).data_ptr(), - graph_metadata.get_compression_batch(0).start_positions.data_ptr(), + graph_metadata.ext(CSA2Metadata).csa2_positions.data_ptr(), + graph_metadata.ext(CSA2Metadata).get_compressed_positions(0).data_ptr(), + graph_metadata.ext(CSA2Metadata).get_compression_batch(0).start_positions.data_ptr(), ) if graph is None: for _ in range(3): diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index 85efe8539f60..2b8a81a64670 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -37,7 +37,7 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( DeepseekV4CacheManager, DeepseekV4Indexer, - DeepseekV4TrtllmAttentionMetadata, + DeepseekV4Metadata, ) from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.compressor import ( Compressor, @@ -47,6 +47,7 @@ DEEPSEEK_V4_SLIDING_ATTENTION, DeepseekV4AttentionType, ) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import get_sm_version @@ -136,6 +137,10 @@ def __init__( self.kv_lens_cuda_runtime = None # Set by caller self.cached_token_lens_cuda = None # Set by caller + def ext(self, cls: type) -> "DummyAttentionMetadata": + """The dummy carries the DeepSeek-V4 extension state itself.""" + return self + # ============================================================================ # Reference Implementation (DO NOT MODIFY) @@ -368,6 +373,15 @@ def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor): ROPE_THETA, ROPE_FACTOR, BETA_FAST, BETA_SLOW = 40000.0, 4, 32, 1 +def _dsv4_stub() -> DeepseekV4Metadata: + """An extension built through ``__new__`` on a core metadata built the same way.""" + metadata = object.__new__(TrtllmAttentionMetadata) + extension = object.__new__(DeepseekV4Metadata) + extension.metadata = metadata + metadata.replace_ext(extension) + return extension + + def _active_compressed_position_ids( compress_ratio, num_contexts, num_gen_tokens_per_seq, cached_tokens, kv_lens ): @@ -377,9 +391,10 @@ def _active_compressed_position_ids( cached_tokens = torch.tensor(cached_tokens, dtype=torch.int32, device=DEVICE) kv_lens = torch.tensor(kv_lens, dtype=torch.int32, device=DEVICE) - metadata = object.__new__(DeepseekV4TrtllmAttentionMetadata) - metadata.num_contexts = num_contexts - metadata.num_generations = num_generations + metadata = _dsv4_stub() + md = metadata.metadata + md.num_contexts = num_contexts + md.num_generations = num_generations metadata.num_gen_tokens_per_seq = num_gen_tokens_per_seq metadata._compress_ratios_sorted = [compress_ratio] metadata.compressed_kv_lens_cuda = { @@ -492,7 +507,7 @@ def run(new_comp_by_ratio): new_comp_bufs[r][:batch_size] = counts cu_bufs[r][: batch_size + 1] = F.pad(torch.cumsum(counts, 0), (1, 0)).to(torch.int32) batch_size_cuda.fill_(batch_size) - DeepseekV4TrtllmAttentionMetadata._compute_compressed_mask( + DeepseekV4Metadata._compute_compressed_mask( new_comp_bufs, cu_bufs, mask_bufs, batch_size_cuda, ratios ) for r, new_comp in new_comp_by_ratio.items(): @@ -543,7 +558,7 @@ def test_ctx_position_ids_host_sizes_match_device_scalar_fallback( def _run(ctx_output_sizes): out = torch.full((total + 8,), -1, dtype=torch.int32, device=DEVICE) - DeepseekV4TrtllmAttentionMetadata._compute_ctx_compressed_position_ids( + DeepseekV4Metadata._compute_ctx_compressed_position_ids( {compress_ratio: past}, {compress_ratio: cu}, {compress_ratio: out}, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.py index f1a7d991d3f6..a626d0e927b2 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_fp4_indexer.py @@ -156,6 +156,8 @@ def overlapped_prepare(*_args, **_kwargs): _run_overlapped_indexer_prepare=overlapped_prepare, ) metadata = SimpleNamespace(empty_topk_indices_buffer=torch.full((4, 2), -1)) + # The stub carries the DeepSeek-V4 extension state itself. + metadata.ext = lambda cls: metadata qr = torch.empty((2, 4), dtype=torch.bfloat16) hidden_states = torch.empty_like(qr) position_ids = torch.arange(2) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py index 7998f8750142..452bee090e72 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py @@ -37,11 +37,10 @@ merge_sparse_read_table, select_sparse_history_pages, ) -from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( - DeepseekV4TrtllmAttentionMetadata, -) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import DeepseekV4Metadata from tensorrt_llm._torch.attention.backends.sparse.hooks import prepare_sparse_runtime_params from tensorrt_llm._torch.attention.backends.sparse.params import SparseBackendForwardArgs +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.memory_buffer_utils import Buffers from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm.bindings import DataType @@ -58,18 +57,21 @@ def _metadata(manager, *, enabled=True, graph=False, buffers=None, initialize=True): # Exercise the production offload initialization without constructing the # unrelated indexer/compressor/native attention workspaces. - metadata = object.__new__(DeepseekV4TrtllmAttentionMetadata) - metadata.kv_cache_manager = manager - metadata.sparse_metadata_params = DeepSeekV4SparseAttentionConfig( + core = object.__new__(TrtllmAttentionMetadata) + metadata = object.__new__(DeepseekV4Metadata) + metadata.metadata = core + core.replace_ext(metadata) + core.kv_cache_manager = manager + core.sparse_metadata_params = DeepSeekV4SparseAttentionConfig( enable_kv_cache_offload=enabled, index_topk=8 ).to_sparse_metadata_params() - metadata.max_num_sequences = 4 + core.max_num_sequences = 4 metadata.sparse_mla_topk = 8 - metadata.beam_width = 1 + core.beam_width = 1 metadata.max_draft_tokens = 0 - metadata.draft_kv_cache_manager = None - metadata.is_cuda_graph = graph - metadata.cuda_graph_buffers = buffers + core.draft_kv_cache_manager = None + core.is_cuda_graph = graph + core.cuda_graph_buffers = buffers if initialize: metadata._init_sparse_offload_state() return metadata @@ -89,7 +91,7 @@ def test_sparse_offload_rejects_configuration_before_allocation( manager.pp_layers = [3] metadata = _metadata(manager, initialize=False) if unsupported == "beam": - metadata.beam_width = 2 + metadata.metadata.beam_width = 2 elif unsupported == "configured_beam": manager.max_beam_width = 2 elif unsupported == "draft": @@ -99,8 +101,8 @@ def test_sparse_offload_rejects_configuration_before_allocation( elif unsupported == "draft_tree": manager.max_total_draft_tokens = 1 else: - metadata.draft_kv_cache_manager = object() - with patch.object(metadata, "get_empty") as allocate: + metadata.metadata.draft_kv_cache_manager = object() + with patch.object(metadata.metadata, "get_empty") as allocate: with pytest.raises(NotImplementedError, match="single-beam, non-speculative"): metadata._init_sparse_offload_state() allocate.assert_not_called() @@ -251,9 +253,9 @@ def test_metadata_routes_only_sparse_main_table(enabled): manager._enable_kv_cache_offload = enabled manager._stream = torch.cuda.current_stream() metadata = _metadata(manager, enabled=enabled) - metadata.request_ids = [20, 10] - metadata._seq_lens = torch.ones(2, dtype=torch.int32) - metadata._num_contexts = 1 + metadata.metadata.request_ids = [20, 10] + metadata.metadata._seq_lens = torch.ones(2, dtype=torch.int32) + metadata.metadata._num_contexts = 1 metadata.sliding_block_tables = torch.empty((1, 1, 4, 6), dtype=torch.int32, device="cuda") metadata.compress_block_tables = { ratio: torch.empty((4, 6), dtype=torch.int32, device="cuda") for ratio in (4, 128) @@ -373,8 +375,10 @@ def _host_batch(*, prefill=False): prompt_lens=[8, 12], kv_cache_params=SimpleNamespace(num_cached_tokens_per_seq=[0, 0]), ) + # The stub serves as both the core metadata and its extension. + metadata.metadata = metadata metadata._validate_ragged_verification = MethodType( - DeepseekV4TrtllmAttentionMetadata._validate_ragged_verification, metadata + DeepseekV4Metadata._validate_ragged_verification, metadata ) return metadata @@ -383,7 +387,7 @@ def _host_batch(*, prefill=False): @pytest.mark.parametrize("prefill", [False, True]) def test_offload_admits_decode_or_fresh_prefill(prefill): metadata = _host_batch(prefill=prefill) - DeepseekV4TrtllmAttentionMetadata.validate_sparse_offload_batch(metadata) + DeepseekV4Metadata.validate_sparse_offload_batch(metadata) assert not metadata.sparse_offload_state.prepared assert metadata.sparse_offload_state.is_prefill == prefill @@ -415,10 +419,10 @@ def test_offload_rejects_batch_before_table_preparation(unsupported): else: metadata.seq_lens_kv[0] += 1 metadata.validate_sparse_offload_batch = lambda: ( - DeepseekV4TrtllmAttentionMetadata.validate_sparse_offload_batch(metadata) + DeepseekV4Metadata.validate_sparse_offload_batch(metadata) ) with pytest.raises(NotImplementedError): - DeepseekV4TrtllmAttentionMetadata.prepare(metadata) + DeepseekV4Metadata.pre_derive_lengths(metadata) metadata.kv_cache_manager.compute_sliding_block_tables.assert_not_called() assert not metadata.sparse_offload_state.prepared @@ -499,9 +503,9 @@ def _attention_case(*, tokens_per_block=128, fetched_page_scale=1, fp8=False): manager.kv_cache_map[30].history_length = 2 * tokens_per_block + 3 manager.kv_cache_map[30].pages = [8, 7, 0] manager.kv_cache_map[10].pages = [4, -1, -1] - metadata.request_ids = [20, 30, 10, CUDA_GRAPH_DUMMY_REQUEST_ID] - metadata._num_contexts = metadata._num_ctx_tokens = 0 - metadata._num_generations = metadata._num_tokens = 4 + metadata.metadata.request_ids = [20, 30, 10, CUDA_GRAPH_DUMMY_REQUEST_ID] + metadata.metadata._num_contexts = metadata.metadata._num_ctx_tokens = 0 + metadata.metadata._num_generations = metadata.metadata._num_tokens = 4 metadata.req_idx_per_token = torch.tensor([1, 0, 2, 3], dtype=torch.int32, device="cuda") metadata.compress_block_tables = {4: torch.empty((4, 6), dtype=torch.int32, device="cuda")} metadata.compressed_kv_lens_cuda = { @@ -568,11 +572,12 @@ def _attention_case(*, tokens_per_block=128, fetched_page_scale=1, fp8=False): forward_args.mla_bmm1_scale = torch.empty(2, device="cuda") forward_args.mla_bmm2_scale = torch.empty(1, device="cuda") manager.prepare_sparse_offload( - state, metadata.request_ids, metadata.compress_block_tables[4], beam_width=1 + state, metadata.metadata.request_ids, metadata.compress_block_tables[4], beam_width=1 ) return SimpleNamespace( manager=manager, - metadata=metadata, + metadata=metadata.metadata, + dsv4=metadata, state=state, backends=backends, args=forward_args, @@ -592,7 +597,7 @@ def _consume(case, backend): def _check_consumption(case, backend, runtime): - metadata, state = case.metadata, case.state + metadata, state = case.dsv4, case.state descriptor = state.layers[backend.layer_idx] offset = case.offsets[descriptor.buffer_id.layer_id] raw = state.base_page_tables[descriptor.group_id].cpu() @@ -646,13 +651,13 @@ def _check_consumption(case, backend, runtime): def test_per_layer_fetch_reaches_sparse_attention_inputs(tokens_per_block, fetched_page_scale): case = _attention_case(tokens_per_block=tokens_per_block, fetched_page_scale=fetched_page_scale) original_topk = case.args.sparse_backend_args.topk_indices.clone() - original_write = case.metadata.compress_block_tables[4].clone() + original_write = case.dsv4.compress_block_tables[4].clone() for backend in case.backends: runtime = _consume(case, backend) _check_consumption(case, backend, runtime) assert [b.layer_id for b, _ in case.calls] == [41, 19] torch.testing.assert_close(case.args.sparse_backend_args.topk_indices, original_topk) - torch.testing.assert_close(case.metadata.compress_block_tables[4], original_write) + torch.testing.assert_close(case.dsv4.compress_block_tables[4], original_write) @_requires_cuda @@ -698,14 +703,12 @@ def test_decode_fetches_even_for_zero_history_or_all_invalid_selection(): torch.tensor([[0, 1, -1, -1, -1, -1, -1, -1]] * 3, dtype=torch.int32, device="cuda") ) case.manager.prepare_sparse_offload( - case.state, case.metadata.request_ids, case.metadata.compress_block_tables[4], beam_width=1 + case.state, case.metadata.request_ids, case.dsv4.compress_block_tables[4], beam_width=1 ) _consume(case, case.backends[0]) assert (case.state.selected_history_pages == -1).all() assert (case.state.fetched_page_table == -1).all() - torch.testing.assert_close( - case.state.compress_read_table, case.metadata.compress_block_tables[4] - ) + torch.testing.assert_close(case.state.compress_read_table, case.dsv4.compress_block_tables[4]) case.args.sparse_backend_args.topk_indices.fill_(-1) _consume(case, case.backends[1]) assert len(case.calls) == 2 @@ -717,7 +720,7 @@ def test_fresh_prefill_consumes_write_table_without_fetch(): for cache in case.manager.kv_cache_map.values(): cache.history_length = 0 case.manager.prepare_sparse_offload( - case.state, case.metadata.request_ids, case.metadata.compress_block_tables[4], beam_width=1 + case.state, case.metadata.request_ids, case.dsv4.compress_block_tables[4], beam_width=1 ) case.state.is_prefill = True case.metadata._num_contexts, case.metadata._num_generations = 2, 0 @@ -728,9 +731,7 @@ def test_fresh_prefill_consumes_write_table_without_fetch(): ) # Multiple queries/request are legal only for fresh prefill, where no # selection union is required. Repeated ordinals retain their positions. - case.metadata.req_idx_per_token.copy_( - torch.tensor([0, 0, 1, 1], dtype=torch.int32, device="cuda") - ) + case.dsv4.req_idx_per_token.copy_(torch.tensor([0, 0, 1, 1], dtype=torch.int32, device="cuda")) runtime = _consume(case, case.backends[0]) assert not case.calls assert ( @@ -749,10 +750,8 @@ def test_offload_disabled_and_other_ratios_do_not_fetch(): runtime = _consume(case, backend) assert runtime.sparse_attn_indices.shape == (4, 128) backend.compress_ratio = 128 - case.metadata.compress_block_tables[128] = case.metadata.compress_block_tables[4] - case.metadata.compressed_local_indices_cuda = torch.zeros( - (4, 8), dtype=torch.int32, device="cuda" - ) + case.dsv4.compress_block_tables[128] = case.dsv4.compress_block_tables[4] + case.dsv4.compressed_local_indices_cuda = torch.zeros((4, 8), dtype=torch.int32, device="cuda") _consume(case, backend) assert not case.calls @@ -797,18 +796,16 @@ def test_fetch_graph_replay_refreshes_requests_topk_and_padding(monkeypatch, deb case.manager.prepare_sparse_offload( case.state, case.metadata.request_ids, - case.metadata.compress_block_tables[4], + case.dsv4.compress_block_tables[4], beam_width=1, ) lengths = [(case.manager.kv_cache_map[i].history_length + 1) // 4 for i in ids] + [0] * ( 4 - len(ids) ) - case.metadata.compressed_kv_lens_cuda[4].copy_( + case.dsv4.compressed_kv_lens_cuda[4].copy_( torch.tensor(lengths, dtype=torch.int32, device="cuda") ) - case.metadata.req_idx_per_token.copy_( - torch.tensor(requests, dtype=torch.int32, device="cuda") - ) + case.dsv4.req_idx_per_token.copy_(torch.tensor(requests, dtype=torch.int32, device="cuda")) case.args.sparse_backend_args.topk_indices.copy_( torch.tensor(topk, dtype=torch.int32, device="cuda") ) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_q_b_fusion.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_q_b_fusion.py index 9d14950e01ba..5a4134e02f31 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_q_b_fusion.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_q_b_fusion.py @@ -43,9 +43,13 @@ def test_q_b_gemm_cu_seqlens_per_phase() -> None: num_generations=4, cu_seq_lens_cuda=gen_cu, ) + # The stub serves as its own DeepSeek-V4 extension. + generation_only.find_ext = lambda cls: generation_only assert torch.equal(_dsv4_q_b_gemm_cu_q_seqlens(owner, generation_only), gen_cu[:5]) def test_q_b_gemm_cu_seqlens_absent_without_dsa_metadata() -> None: - metadata = SimpleNamespace(cu_q_seqlens=None, num_contexts=0, num_generations=4) + metadata = SimpleNamespace( + cu_q_seqlens=None, num_contexts=0, num_generations=4, find_ext=lambda cls: None + ) assert _dsv4_q_b_gemm_cu_q_seqlens(SimpleNamespace(), metadata) is None diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index 9d6cc8cf66b6..4dd465e1f893 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -25,6 +25,7 @@ import torch from utils.util import skip_pre_blackwell +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionInputType, MLAParams, @@ -34,10 +35,11 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( DeepseekV4AttentionType, DeepseekV4CacheManager, + DeepseekV4Metadata, DeepseekV4TrtllmAttention, - DeepseekV4TrtllmAttentionMetadata, ) from tensorrt_llm._torch.attention.backends.sparse.params import SparseBackendForwardArgs +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -174,6 +176,17 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) +def _create_metadata( + *, sparse_attention_config, enable_flash_mla: bool = False, **kwargs +) -> TrtllmAttentionMetadata: + """DeepSeek-V4 attention metadata, built the way the backend builds it.""" + return DeepseekV4TrtllmAttention.create_metadata( + sparse_metadata_params=sparse_attention_config.to_sparse_metadata_params(), + extension_types=(FlashMlaMetadata,) if enable_flash_mla else (), + **kwargs, + ) + + def _create_cache_manager(scenario: Scenario, context_lengths: List[int], max_seq_len: int): sparse_config = DeepSeekV4SparseAttentionConfig( index_n_heads=64, @@ -723,7 +736,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): attn_sink = torch.randn(local_num_heads, dtype=torch.float32, device=device).mul_(0.5) layer.attn_sink = torch.nn.Parameter(attn_sink, requires_grad=False) - attn_metadata = DeepseekV4TrtllmAttentionMetadata( + attn_metadata = _create_metadata( seq_lens=torch.tensor(context_lengths, dtype=torch.int), request_ids=[0], max_num_requests=1, @@ -735,11 +748,12 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): mapping=Mapping(world_size=1, tp_size=1, rank=0), sparse_attention_config=sparse_config, ) + dsa = attn_metadata.ext(DeepseekV4Metadata) attn_metadata.prepare() - assert attn_metadata.sparse_mla_topk_lens[ratio][0].item() == 1 - assert attn_metadata.swa_local_indices_cuda[0, 0].item() == 0 - assert torch.all(attn_metadata.swa_local_indices_cuda[0, 1:] == -1) + assert dsa.sparse_mla_topk_lens[ratio][0].item() == 1 + assert dsa.swa_local_indices_cuda[0, 0].item() == 0 + assert torch.all(dsa.swa_local_indices_cuda[0, 1:] == -1) compressed_kv = torch.empty([1, kv_lora_rank], dtype=dtype, device=device).uniform_(-1, 1) k_pe = torch.empty([1, qk_rope_head_dim], dtype=dtype, device=device).uniform_(-1, 1) @@ -1054,7 +1068,7 @@ def yarn_get_mscale(scale=1, mscale=1): # 7. Context phase ctx_seq_lens = torch.tensor(context_lengths, dtype=torch.int) - attn_metadata = DeepseekV4TrtllmAttentionMetadata( + attn_metadata = _create_metadata( seq_lens=ctx_seq_lens, request_ids=request_ids, max_num_requests=batch_size, @@ -1188,8 +1202,8 @@ def yarn_get_mscale(scale=1, mscale=1): gen_seq_lens = torch.tensor(generation_lens, dtype=torch.int) total_gen_tokens = sum(generation_lens) - def create_gen_metadata(ragged: bool) -> DeepseekV4TrtllmAttentionMetadata: - metadata = DeepseekV4TrtllmAttentionMetadata( + def create_gen_metadata(ragged: bool) -> TrtllmAttentionMetadata: + metadata = _create_metadata( seq_lens=gen_seq_lens, request_ids=request_ids, max_num_requests=batch_size, @@ -1204,16 +1218,19 @@ def create_gen_metadata(ragged: bool) -> DeepseekV4TrtllmAttentionMetadata: mapping=mapping, enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), sparse_attention_config=sparse_config, - enable_ragged_verification=ragged, ) + dsa = metadata.ext(DeepseekV4Metadata) + # Ragged-verification buffers are sized when the extension allocates. + dsa.enable_ragged_verification = ragged + dsa.allocate(metadata.cuda_graph_buffers) # Keep the persistent allocation ceiling independent of this # step's query lengths, as in update_spec_dec_param. - metadata.max_draft_tokens = MAX_GENERATION_SEQ_LEN_Q - 1 - metadata.runtime_tokens_per_gen_step = max_generation_width - metadata.ragged_verify_lens = generation_lens.copy() if ragged else None - metadata._create_kv_lens_2d_buffer() - metadata.create_expanded_buffers() - for ratio, position_ids in metadata.compressed_position_ids_cuda.items(): + dsa.max_draft_tokens = MAX_GENERATION_SEQ_LEN_Q - 1 + dsa.runtime_tokens_per_gen_step = max_generation_width + dsa.ragged_verify_lens = generation_lens.copy() if ragged else None + dsa._create_kv_lens_2d_buffer() + dsa.create_expanded_buffers() + for ratio, position_ids in dsa.compressed_position_ids_cuda.items(): required_capacity = batch_size * math.ceil(MAX_GENERATION_SEQ_LEN_Q / ratio) assert position_ids.numel() >= required_capacity metadata.prepare() @@ -1231,14 +1248,16 @@ def create_gen_metadata(ragged: bool) -> DeepseekV4TrtllmAttentionMetadata: device=device, ) torch.testing.assert_close( - gen_metadata.ragged_row_kv_lens(total_gen_tokens), expected_rows + gen_metadata.ext(DeepseekV4Metadata).ragged_row_kv_lens(total_gen_tokens), + expected_rows, ) torch.testing.assert_close( gen_metadata.token_major_gen_view().sequence_length, expected_rows ) gen_metadata.on_update_kv_lens() torch.testing.assert_close( - gen_metadata.ragged_row_kv_lens(total_gen_tokens), expected_rows + gen_metadata.ext(DeepseekV4Metadata).ragged_row_kv_lens(total_gen_tokens), + expected_rows, ) torch.testing.assert_close( gen_metadata.token_major_gen_view().sequence_length, expected_rows @@ -1514,7 +1533,7 @@ def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): ).uniform_(-1, 1) prefill_latent = torch.cat([prefill_compressed_kv, prefill_k_pe], dim=-1) - prefill_metadata = DeepseekV4TrtllmAttentionMetadata( + prefill_metadata = _create_metadata( seq_lens=torch.tensor(gen_ctx_lengths, dtype=torch.int), request_ids=gen_request_ids, max_num_requests=num_gen, @@ -1592,7 +1611,7 @@ def test_deepseek_v4_sparse_mla_mixed_batch(context_lengths: List[int]): # 4. Mixed metadata: request 0 = context, requests 1..N = generation. mixed_seq_lens = [context_lengths[0]] + [generation_seq_len_q] * num_gen mixed_cached_lens = [0, *gen_cached_lens] - mixed_metadata = DeepseekV4TrtllmAttentionMetadata( + mixed_metadata = _create_metadata( seq_lens=torch.tensor(mixed_seq_lens, dtype=torch.int), request_ids=request_ids, max_num_requests=batch_size, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_indexer_slot_mapping.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_indexer_slot_mapping.py index c1886b7a8111..371c354193f9 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_indexer_slot_mapping.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_indexer_slot_mapping.py @@ -11,13 +11,15 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4 import ( DeepseekV4CacheManager, - DeepseekV4TrtllmAttentionMetadata, + DeepseekV4TrtllmAttention, ) from tensorrt_llm._torch.attention.backends.sparse.dsa import ( DSACacheManagerV2, - DSAtrtllmAttentionMetadata, + DSAMetadata, + DSATrtllmAttention, ) from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import _fused_dsa_meta_enabled +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp @@ -37,7 +39,7 @@ def _prepared_metadata( *, context: bool, compressed: bool, -) -> Iterator[DSAtrtllmAttentionMetadata]: +) -> Iterator[TrtllmAttentionMetadata]: config_cls = DeepSeekV4SparseAttentionConfig if compressed else DeepSeekSparseAttentionConfig sparse_config = config_cls( index_n_heads=64, @@ -73,10 +75,8 @@ def _prepared_metadata( try: kv_lens = [cached + length for cached, length in zip(cached_tokens, seq_lens)] assert manager.add_dummy_requests([0, 1], kv_lens, is_gen=not context) is not None - metadata_cls = ( - DeepseekV4TrtllmAttentionMetadata if compressed else DSAtrtllmAttentionMetadata - ) - metadata = metadata_cls( + backend_cls = DeepseekV4TrtllmAttention if compressed else DSATrtllmAttention + metadata = backend_cls.create_metadata( seq_lens=torch.tensor(seq_lens, dtype=torch.int32), request_ids=[0, 1], max_num_requests=2, @@ -91,7 +91,7 @@ def _prepared_metadata( ) metadata.prepare() assert manager.quant_block_size == (32 if compressed and indexer_dtype == "fp4" else 128) - assert metadata.indexer_quant_block_size == 128 + assert metadata.ext(DSAMetadata).indexer_quant_block_size == 128 yield metadata finally: manager.shutdown() @@ -111,6 +111,7 @@ def test_v4_runtime_maps_gather_physical_cache_bytes( with _prepared_metadata( indexer_dtype, cached, lengths, context=context, compressed=True ) as metadata: + dsa = metadata.ext(DSAMetadata) cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(0) data_width = 64 if indexer_dtype == "fp4" else 128 assert cache.shape[1:] == (32, 1, data_width + 4) @@ -132,33 +133,33 @@ def test_v4_runtime_maps_gather_physical_cache_bytes( for start, length in zip(cached, lengths) ] ) - blocks = metadata.indexer_k_cache_block_offsets[requests, positions // 32].long() + blocks = dsa.indexer_k_cache_block_offsets[requests, positions // 32].long() assert blocks.gt(0).any() expected_data = payload[blocks, positions % 32].clone() expected_scales = scales[blocks, positions % 32].clone() gather_maps = ( ( - metadata.slot_mapping_fp8_fullkv.clone(), - metadata.slot_mapping_scale_fullkv.clone(), + dsa.slot_mapping_fp8_fullkv.clone(), + dsa.slot_mapping_scale_fullkv.clone(), ) if context else None ) - metadata.slot_mapping_fp8.fill_(-1) - metadata.slot_mapping_scale.fill_(-1) + dsa.slot_mapping_fp8.fill_(-1) + dsa.slot_mapping_scale.fill_(-1) metadata.on_update_kv_lens() - assert getattr(metadata, "_fused_dsa_meta_armed", False) is fused + assert getattr(dsa, "_fused_dsa_meta_armed", False) is fused for offsets, width in ( - (metadata.slot_mapping_fp8, data_width), - (metadata.slot_mapping_scale, 4), + (dsa.slot_mapping_fp8, data_width), + (dsa.slot_mapping_scale, 4), ): assert offsets.ge(0).all() and (offsets + width).le(cache.numel()).all() data, scale = torch.ops.trtllm.indexer_k_cache_gather_op( cache, - metadata.slot_mapping_fp8, - metadata.slot_mapping_scale, + dsa.slot_mapping_fp8, + dsa.slot_mapping_scale, 0, sum(lengths), data_width, @@ -170,10 +171,10 @@ def test_v4_runtime_maps_gather_physical_cache_bytes( if gather_maps is not None: # Compressed context gathers have independent maps; scatter positions are raw. torch.testing.assert_close( - metadata.slot_mapping_fp8_fullkv, gather_maps[0], rtol=0, atol=0 + dsa.slot_mapping_fp8_fullkv, gather_maps[0], rtol=0, atol=0 ) torch.testing.assert_close( - metadata.slot_mapping_scale_fullkv, gather_maps[1], rtol=0, atol=0 + dsa.slot_mapping_scale_fullkv, gather_maps[1], rtol=0, atol=0 ) finally: _fused_dsa_meta_enabled.cache_clear() @@ -192,19 +193,20 @@ def test_dsa_context_gather_alias_survives_runtime_refresh( context=True, compressed=False, ) as metadata: - assert (metadata.slot_mapping_fp8_fullkv is metadata.slot_mapping_fp8) is (not cached) - assert (metadata.slot_mapping_scale_fullkv is metadata.slot_mapping_scale) is (not cached) + dsa = metadata.ext(DSAMetadata) + assert (dsa.slot_mapping_fp8_fullkv is dsa.slot_mapping_fp8) is (not cached) + assert (dsa.slot_mapping_scale_fullkv is dsa.slot_mapping_scale) is (not cached) names = ( "slot_mapping_fp8", "slot_mapping_scale", "slot_mapping_fp8_fullkv", "slot_mapping_scale_fullkv", ) - expected = [getattr(metadata, name).clone() for name in names] + expected = [getattr(dsa, name).clone() for name in names] metadata.on_update_kv_lens() - assert (metadata.slot_mapping_fp8_fullkv is metadata.slot_mapping_fp8) is (not cached) - assert (metadata.slot_mapping_scale_fullkv is metadata.slot_mapping_scale) is (not cached) + assert (dsa.slot_mapping_fp8_fullkv is dsa.slot_mapping_fp8) is (not cached) + assert (dsa.slot_mapping_scale_fullkv is dsa.slot_mapping_scale) is (not cached) for name, prepared in zip(names, expected): - torch.testing.assert_close(getattr(metadata, name), prepared, rtol=0, atol=0) + torch.testing.assert_close(getattr(dsa, name), prepared, rtol=0, atol=0) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py index 5beda56866d2..c7b2265baf06 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py @@ -8,6 +8,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from tensorrt_llm._torch.attention.backends.fmha import registry from tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla import ( @@ -1006,7 +1007,7 @@ class _Attn: profile = dict( device=torch.device("cuda"), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None ) - eager = SimpleNamespace( + eager = metadata_stub( fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, @@ -1018,7 +1019,7 @@ class _Attn: assert planned == [(20, 512), (15, 512), (7, 512)] # one per signature, largest first lib._warm_kernel_variants(metadata=eager, extra=512, **profile) assert len(planned) == 3 # same profile: already warmed in this process - graph = SimpleNamespace( + graph = metadata_stub( fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, @@ -1047,7 +1048,7 @@ def test_warm_set_covers_the_speculative_decoding_rows(monkeypatch) -> None: attn, profile, _ = _attn_and_profile() lib = PrimsTSSparseMlaFmha(attn) - mtp = SimpleNamespace( + mtp = metadata_stub( fmha_plan_caches={}, max_num_tokens=256, max_num_requests=4, @@ -1064,7 +1065,7 @@ def test_warm_set_covers_the_speculative_decoding_rows(monkeypatch) -> None: # The step's token budget bounds the rows even when the MTP ceiling is larger. planned.clear() monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) - budget = SimpleNamespace( + budget = metadata_stub( fmha_plan_caches={}, max_num_tokens=10, max_num_requests=4, @@ -1134,7 +1135,7 @@ class _Attn: profile = dict( device=torch.device("cuda", 0), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None ) - eager = SimpleNamespace( + eager = metadata_stub( fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, @@ -1207,7 +1208,7 @@ def broken_workers(payloads, workers): monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "0") # the module fixture sets strict mode planned.clear() - fresh = SimpleNamespace( # a new metadata object: the earlier plans are cached on `eager` + fresh = metadata_stub( # a new metadata object: the earlier plans are cached on `eager` fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py index d44330524748..3f4cdd9ea4df 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py @@ -6,10 +6,9 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( - DeepseekV4TrtllmAttentionMetadata, -) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import DeepseekV4Metadata from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import _fused_q_rope_specs +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata class _Mqa: @@ -20,6 +19,12 @@ def _ensure_rope_table_size(self, size): pass +def _serve_as_extension(metadata: SimpleNamespace) -> None: + """Let the stub act as both the core metadata and its DeepSeek-V4 extension.""" + metadata.metadata = metadata + metadata.find_ext = lambda cls: metadata + + def _ragged_metadata(verify_lens, *, context_lens=(), kv_lens=None): num_generations = len(verify_lens) if kv_lens is None: @@ -36,8 +41,9 @@ def _ragged_metadata(verify_lens, *, context_lens=(), kv_lens=None): fused_q_gen_cu_seqlens=torch.zeros(num_generations + 1, dtype=torch.int32), _fused_q_gen_cu_seqlens_valid=False, ) + _serve_as_extension(metadata) metadata.mla_prepare_fused_q_gen_cu_seqlens = MethodType( - DeepseekV4TrtllmAttentionMetadata.mla_prepare_fused_q_gen_cu_seqlens, metadata + DeepseekV4Metadata.mla_prepare_fused_q_gen_cu_seqlens, metadata ) return metadata @@ -68,6 +74,7 @@ def test_uniform_k5_keeps_scalar_generation_spec(): max_seq_len=1006, is_ragged_verify=False, ) + _serve_as_extension(metadata) _, specs = _fused_q_rope_specs(SimpleNamespace(mqa=_Mqa()), metadata, 0, 128) assert len(specs) == 1 rows, cache_lens, seq_len, cu_q_seqlens = specs[0] @@ -121,6 +128,7 @@ def test_ragged_generation_fails_closed_without_stable_prefix_hook(): max_seq_len=20, is_ragged_verify=True, ) + _serve_as_extension(metadata) mla = SimpleNamespace(mqa=_Mqa()) assert _fused_q_rope_specs(mla, metadata, 0, 2) == (None, []) metadata.mla_prepare_fused_q_gen_cu_seqlens = lambda: None @@ -150,7 +158,10 @@ def test_prefix_builder_fails_closed_without_graph_storage(capacity): def test_sequence_update_invalidates_cached_prefix(): - metadata = DeepseekV4TrtllmAttentionMetadata.__new__(DeepseekV4TrtllmAttentionMetadata) - metadata._fused_q_gen_cu_seqlens_valid = True + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + dsv4 = DeepseekV4Metadata.__new__(DeepseekV4Metadata) + dsv4.metadata = metadata + metadata.replace_ext(dsv4) + dsv4._fused_q_gen_cu_seqlens_valid = True metadata._invalidate_mla_scheduler_buffers() - assert not metadata._fused_q_gen_cu_seqlens_valid + assert not dsv4._fused_q_gen_cu_seqlens_valid diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py index a80f5d7f5b21..c400284e6747 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py @@ -19,10 +19,9 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( - DeepseekV4TrtllmAttentionMetadata, -) -from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import DeepseekV4Metadata +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata class _StubCacheManager: @@ -34,6 +33,15 @@ def __init__(self, num_pools: int, max_blocks_per_seq: int) -> None: self.max_blocks_per_seq = max_blocks_per_seq +def _stub(extension_type: type = DSAMetadata) -> DSAMetadata: + """An extension built through ``__new__`` on a core metadata built the same way.""" + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + extension = extension_type.__new__(extension_type) + extension.metadata = metadata + metadata.replace_ext(extension) + return extension + + def _metadata( *, num_pools: int, @@ -41,8 +49,8 @@ def _metadata( max_seqs: int, draft: int = 5, enable_ragged_verification: bool = True, -) -> DSAtrtllmAttentionMetadata: - """A metadata object with just enough state to build its own buffers. +) -> DSAMetadata: + """A DSA extension with just enough state to build its own buffers. Constructed through ``__new__`` rather than the real constructor: the real one wants a runtime, a resource manager and a model config, none of which @@ -50,17 +58,17 @@ def _metadata( new dependency shows up as a clear AttributeError here instead of being silently satisfied by a mock. """ - meta = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) - meta.kv_cache_manager = _StubCacheManager(num_pools, max_blocks) - meta.max_num_sequences = max_seqs - meta.max_num_requests = max_seqs + meta = _stub() + meta.metadata.kv_cache_manager = _StubCacheManager(num_pools, max_blocks) + meta.metadata.max_num_sequences = max_seqs + meta.metadata.max_num_requests = max_seqs meta.max_draft_tokens = draft meta.enable_ragged_verification = enable_ragged_verification - meta.cuda_graph_buffers = None + meta.metadata.cuda_graph_buffers = None meta.num_sms = torch.cuda.get_device_properties(0).multi_processor_count meta.runtime_tokens_per_gen_step = 0 meta.create_expanded_buffers(capture_graph=False) - meta.prompt_lens_cpu = torch.zeros(max_seqs, dtype=torch.int32) + meta.metadata.prompt_lens_cpu = torch.zeros(max_seqs, dtype=torch.int32) return meta @@ -81,12 +89,10 @@ def test_feature_off_skips_ragged_row_buffers() -> None: assert meta.attn_row_block_offsets is None -@pytest.mark.parametrize( - "metadata_type", [DSAtrtllmAttentionMetadata, DeepseekV4TrtllmAttentionMetadata] -) +@pytest.mark.parametrize("metadata_type", [DSAMetadata, DeepseekV4Metadata]) def test_feature_off_rejects_ragged_layout(metadata_type) -> None: """An unexpected ragged update fails instead of dereferencing absent buffers.""" - meta = metadata_type.__new__(metadata_type) + meta = _stub(metadata_type) meta.enable_ragged_verification = False meta.ragged_verify_lens = None assert not meta.is_ragged_verify @@ -96,37 +102,35 @@ def test_feature_off_rejects_ragged_layout(metadata_type) -> None: assert meta.is_ragged_verify with pytest.raises(RuntimeError, match="without enabling ragged-verification buffers"): - meta.prepare() + meta.pre_derive_lengths() with pytest.raises(RuntimeError, match="without enabling ragged-verification buffers"): meta.apply_device_ragged_layout( torch.tensor([2]), torch.tensor([0, 0]), torch.tensor([-1, 0]) ) -@pytest.mark.parametrize( - "metadata_type", [DSAtrtllmAttentionMetadata, DeepseekV4TrtllmAttentionMetadata] -) +@pytest.mark.parametrize("metadata_type", [DSAMetadata, DeepseekV4Metadata]) def test_device_window_kv_bound_preserves_context_and_input(metadata_type) -> None: - meta = metadata_type.__new__(metadata_type) + meta = _stub(metadata_type) meta.device_windows_mode = True meta.ragged_verify_lens = [2, 4] - meta._num_contexts = 1 - meta._seq_lens = torch.tensor([7, 2, 4], dtype=torch.int32) - meta._seq_lens_kv = None + meta.metadata._num_contexts = 1 + meta.metadata._seq_lens = torch.tensor([7, 2, 4], dtype=torch.int32) + meta.metadata._seq_lens_kv = None meta.max_draft_tokens = 5 - meta.kv_lens = torch.tensor([7, 102, 204], dtype=torch.int32) - original = meta.kv_lens.clone() + meta.metadata.kv_lens = torch.tensor([7, 102, 204], dtype=torch.int32) + original = meta.metadata.kv_lens.clone() bounded = meta._prepare_device_window_kv_lens(original) assert original.tolist() == [7, 102, 204] assert bounded.tolist() == [7, 106, 206] - assert meta.kv_lens.tolist() == bounded.tolist() + assert meta.metadata.kv_lens.tolist() == bounded.tolist() @pytest.mark.parametrize( "device_windows_mode,ragged_lens", [(False, [2]), (True, None), (False, None)] ) def test_native_kv_bound_reuses_storage(device_windows_mode, ragged_lens) -> None: - meta = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) + meta = _stub() meta.device_windows_mode = device_windows_mode meta.ragged_verify_lens = ragged_lens lengths = torch.tensor([102], dtype=torch.int32) @@ -134,7 +138,7 @@ def test_native_kv_bound_reuses_storage(device_windows_mode, ragged_lens) -> Non def _expand( - meta: DSAtrtllmAttentionMetadata, + meta: DSAMetadata, *, num_contexts: int, row_req_idx: list[int], @@ -143,14 +147,14 @@ def _expand( ) -> int: """Drive the private expansion with an explicit row->request map.""" ngt = len(row_req_idx) - meta.num_contexts = num_contexts + meta.metadata.num_contexts = num_contexts meta._ragged_num_rows = ngt meta.ragged_verify_lens = [1] * ngt - meta.kv_cache_block_offsets = block_offsets + meta.metadata.kv_cache_block_offsets = block_offsets meta.row_req_idx_host[:ngt].copy_(torch.tensor(row_req_idx, dtype=torch.long)) meta.row_kv_lens_host[:ngt].fill_(7) meta.row_kv_correction_host[:ngt].zero_() - meta.prompt_lens_cpu[:num_seqs].fill_(1) + meta.metadata.prompt_lens_cpu[:num_seqs].fill_(1) meta._prepare_token_major_gen_rows(torch.full((num_seqs,), 7, dtype=torch.int32)) return meta._attn_num_rows @@ -173,7 +177,7 @@ def test_row_block_table_matches_op_layout(num_pools: int) -> None: ) assert rows == len(row_req_idx) - got = meta.token_major_gen_view().kv_cache_block_offsets + got = meta.metadata.token_major_gen_view().kv_cache_block_offsets assert got.dim() == src.dim(), ( f"expanded table has rank {got.dim()} but the attention op reads rank " f"{src.dim()}; substituting it wholesale would misread every axis" @@ -240,10 +244,10 @@ def test_device_layout_preserves_mixed_context_row_and_token_prefixes() -> None: ) assert rows == 5 - meta.num_generations = 2 - meta._num_ctx_tokens = 5 - meta.seq_lens = torch.tensor([2, 3, 1, 2], dtype=torch.int32) - meta._seq_lens_cuda = meta.seq_lens.to(device="cuda") + meta.metadata.num_generations = 2 + meta.metadata._num_ctx_tokens = 5 + meta.metadata.seq_lens = torch.tensor([2, 3, 1, 2], dtype=torch.int32) + meta.metadata._seq_lens_cuda = meta.metadata.seq_lens.to(device="cuda") meta.req_idx_per_token = torch.full((32,), -1, dtype=torch.int32, device="cuda") context_token_requests = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32, device="cuda") meta.req_idx_per_token[:5].copy_(context_token_requests) @@ -286,7 +290,7 @@ def test_every_runtime_view_shares_the_row_count() -> None: nc = 1 rows = _expand(meta, num_contexts=nc, row_req_idx=[0, 0, 1], num_seqs=6, block_offsets=src) - view = meta.token_major_gen_view() + view = meta.metadata.token_major_gen_view() for name in ( "sequence_length", "host_past_key_value_lengths", diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index d16299ede1b9..946f17e091e2 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -32,6 +32,7 @@ import torch from tensorrt_llm import deep_gemm +from tensorrt_llm._torch.attention.backends.features import DraftKVMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, @@ -42,8 +43,8 @@ DSABackendForwardArgs, DSACacheManager, DSACacheManagerV2, + DSAMetadata, DSATrtllmAttention, - DSAtrtllmAttentionMetadata, Indexer, _effective_compress_ratio_divisor, _select_indexer_compress_ratio, @@ -93,6 +94,8 @@ from utils.util import check_accuracy, skip_pre_blackwell, skip_pre_hopper +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata + def has_deep_gemm(): """Return True if the DeepGEMM library is available.""" @@ -126,6 +129,16 @@ def test_fp8_ds_mla_layout_treats_binding_config_as_auto(): assert storage_head_dim == 576 +def _dsa_stub() -> DSAMetadata: + """A DSA extension with its defaults, on a core metadata built without __init__.""" + md = object.__new__(TrtllmAttentionMetadata) + with patch.object(deep_gemm, "get_num_sms", return_value=1): + dsa = DSAMetadata(md) + md.replace_ext(dsa) + md.replace_ext(SpecDecodingMetadata(md)) + return dsa + + def _set_torch_top_k(indexer: Indexer) -> None: indexer.top_k = TopK( indexer.index_topk, @@ -146,21 +159,19 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_samp use_self_sampling_topk=use_self_sampling, ) sparse_metadata_params = sparse_config.to_sparse_metadata_params() - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.sparse_metadata_params = sparse_metadata_params - metadata.kv_cache_manager = SimpleNamespace( + metadata = _dsa_stub() + md = metadata.metadata + md.sparse_metadata_params = sparse_metadata_params + md.kv_cache_manager = SimpleNamespace( tokens_per_block=256, compressed_block_sizes={}, get_cache_indices=Mock(), ) - metadata.is_cuda_graph = False + md.is_cuda_graph = False metadata.create_buffers_for_mla_rope_append = Mock() metadata.create_buffers_for_indexer = Mock() with ( - patch( - "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" - ), patch( "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.IS_CUTLASS_DSL_AVAILABLE", True, @@ -170,7 +181,7 @@ def test_metadata_cache_geometry_comes_from_sparse_metadata_params(use_self_samp return_value=100, ), ): - DSAtrtllmAttentionMetadata.__post_init__(metadata) + metadata.allocate(None) assert metadata.indexer_head_dim == 96 assert metadata.compress_ratios == [1, 4, 128] @@ -200,26 +211,24 @@ def test_cute_dsl_paged_mqa_logits_gate(requested, cute_dsl_available, sm_100f, indexer_k_dtype="fp8", use_cute_dsl_paged_mqa_logits=requested, ) - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.sparse_metadata_params = sparse_config.to_sparse_metadata_params() - metadata.kv_cache_manager = SimpleNamespace( + metadata = _dsa_stub() + md = metadata.metadata + md.sparse_metadata_params = sparse_config.to_sparse_metadata_params() + md.kv_cache_manager = SimpleNamespace( tokens_per_block=256, compressed_block_sizes={}, get_cache_indices=Mock(), ) - metadata.is_cuda_graph = False + md.is_cuda_graph = False metadata.create_buffers_for_mla_rope_append = Mock() metadata.create_buffers_for_indexer = Mock() indexer_mod = "tensorrt_llm._torch.attention.backends.sparse.dsa.indexer" with ( - patch( - "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.TrtllmAttentionMetadata.__post_init__" - ), patch(f"{indexer_mod}.IS_CUTLASS_DSL_AVAILABLE", cute_dsl_available), patch(f"{indexer_mod}.is_sm_100f", return_value=sm_100f), ): - DSAtrtllmAttentionMetadata.__post_init__(metadata) + metadata.allocate(None) indexer_decision = resolve_cute_dsl_paged_mqa_logits( sparse_config.to_sparse_params().use_cute_dsl_paged_mqa_logits ) @@ -356,6 +365,7 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( get_indexer_max_seq_len=Mock(return_value=32768), num_sms=148, ) + metadata.metadata = metadata with ( patch( @@ -366,7 +376,7 @@ def test_metadata_warmup_cute_dsl_radix_topk_dispatch( "tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops.warmup_cute_dsl_radix_topk_decode" ) as cute_dsl_radix, ): - DSAtrtllmAttentionMetadata.warmup_cute_dsl_radix_topk(metadata, next_n) + DSAMetadata.warmup_cute_dsl_radix_topk(metadata, next_n) if should_warmup: cute_dsl_radix.assert_called_once_with( @@ -407,6 +417,7 @@ def test_metadata_warmup_selfsampling_prefill_leg( use_cute_dsl_paged_mqa_logits=True, num_sms=148, ) + metadata.metadata = metadata ss_host = "tensorrt_llm._torch.kernels.blackwell.top_k.gvr_topk_decode_self_sampling_host" with ( patch( @@ -420,7 +431,7 @@ def test_metadata_warmup_selfsampling_prefill_leg( patch(f"{ss_host}.warmup_prefill") as warmup_prefill, patch(f"{ss_host}.warmup_varlen"), ): - DSAtrtllmAttentionMetadata.warmup_selfsampling_topk(metadata, next_n=1, batch_sizes=[8]) + DSAMetadata.warmup_selfsampling_topk(metadata, next_n=1, batch_sizes=[8]) if should_warmup: warmup_prefill.assert_called_once_with(512, max(msl_c, 32768)) @@ -436,7 +447,7 @@ def test_kv_lens_row_reorder_threshold(): def make_mock(num_generations, kv_lens_list): kv_lens_cuda = torch.tensor(kv_lens_list, dtype=torch.int32, device="cuda") row_order_buffer = torch.zeros(64, dtype=torch.int32, device="cuda") - return SimpleNamespace( + mock = SimpleNamespace( needs_gvr_prior=True, num_generations=num_generations, num_sms=num_sms, @@ -448,15 +459,17 @@ def make_mock(num_generations, kv_lens_list): kv_lens_row_reorder_buffer=row_order_buffer, kv_lens_row_reorder=None, ) + mock.metadata = mock + return mock kv_lens = [4, 1, 8, 2, 16, 3, 12, 6, 7, 9, 5, 11, 13, 10, 14, 15] metadata_below = make_mock(1, [1000]) - DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(metadata_below) + DSAMetadata._compute_kv_lens_row_reorder(metadata_below) assert metadata_below.kv_lens_row_reorder is None metadata_at = make_mock(num_sms, kv_lens) - DSAtrtllmAttentionMetadata._compute_kv_lens_row_reorder(metadata_at) + DSAMetadata._compute_kv_lens_row_reorder(metadata_at) row_order = metadata_at.kv_lens_row_reorder.cpu().tolist() assert [kv_lens[i] for i in row_order] == sorted(kv_lens, reverse=True) @@ -500,7 +513,7 @@ def test_gvr_prior_writeback_uses_aux_stream(): indexer.top_k.decode_implementation = TopKImplementation.CUTE_DSL_GVR indexer.top_k.gvr_self_sampling = False indexer.aux_stream = torch.cuda.Stream() - metadata.gvr_prior_indices = torch.zeros( + metadata.ext(DSAMetadata).gvr_prior_indices = torch.zeros( (cache_manager.num_local_layers, batch_size, index_topk), device="cuda", dtype=torch.int32, @@ -522,7 +535,7 @@ def test_gvr_prior_writeback_uses_aux_stream(): local_layer = cache_manager.layer_offsets[indexer.layer_idx] torch.testing.assert_close( - metadata.gvr_prior_indices[local_layer, :batch_size], + metadata.ext(DSAMetadata).gvr_prior_indices[local_layer, :batch_size], topk_indices, ) assert not indexer._prev_topk_copy_pending @@ -552,20 +565,21 @@ def test_shared_topk_lifecycle(monkeypatch): ) assert sparse_metadata_params.has_shared_indexer_layers - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.sparse_metadata_params = sparse_metadata_params - metadata.max_num_sequences = 2 - metadata.max_num_tokens = 4 - metadata.num_sparse_topk = 3 + metadata = _dsa_stub() + md = metadata.metadata + md.sparse_metadata_params = sparse_metadata_params + md.max_num_sequences = 2 + md.max_num_tokens = 4 + md.num_sparse_topk = 3 metadata._radix_rows_per_sequence = 1 metadata.num_sms = 1 - metadata.cuda_graph_buffers = None - metadata.kv_cache_manager = SimpleNamespace(max_blocks_per_seq=2) - metadata.enable_context_mla_with_cached_kv = False + md.cuda_graph_buffers = None + md.kv_cache_manager = SimpleNamespace(max_blocks_per_seq=2) + md.enable_context_mla_with_cached_kv = False metadata.enable_indexer_skip = False metadata.enable_gvr_topk = False metadata.needs_gvr_prior = False - metadata.get_empty = Mock( + md.get_empty = Mock( side_effect=lambda _, shape, **kwargs: torch.empty(tuple(shape), dtype=kwargs["dtype"]) ) metadata._create_kv_lens_2d_buffer = Mock() @@ -600,8 +614,8 @@ def test_shared_topk_lifecycle(monkeypatch): backend._remap_topk_to_global = MethodType( DSATrtllmAttention._remap_topk_to_global, backend ) - metadata._num_ctx_tokens = context_topk.shape[0] - metadata._num_tokens = context_topk.shape[0] + generation_topk.shape[0] + md._num_ctx_tokens = context_topk.shape[0] + md._num_tokens = context_topk.shape[0] + generation_topk.shape[0] backend_args = DSABackendForwardArgs(indexer_intermediates=[]) def _predict(backend, input_type): @@ -609,7 +623,7 @@ def _predict(backend, input_type): backend, torch.empty((2, 1)), None, - metadata, + md, AttentionForwardArgs( attention_input_type=input_type, sparse_backend_args=backend_args, @@ -637,38 +651,36 @@ def _predict(backend, input_type): ) metadata._invalidate_pool_view_cache = Mock() - metadata.kv_cache_manager = None - metadata._num_tokens = 0 - metadata._num_generations = 0 - metadata.kv_lens_cuda = torch.empty(0, dtype=torch.int32) + md.kv_cache_manager = None + md._num_tokens = 0 + md._num_generations = 0 + md.kv_lens_cuda = torch.empty(0, dtype=torch.int32) metadata.prepare_dense_topk_indices = Mock() - metadata.on_update_kv_lens() + md.on_update_kv_lens() assert metadata.shared_topk_indices is buffer @pytest.mark.parametrize("enabled", [False, True]) def test_kv_len_update_refreshes_ragged_rows_only_when_enabled(enabled): - metadata = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) - metadata.kv_cache_manager = None - metadata._num_tokens = 0 - metadata._num_generations = 0 + metadata = _dsa_stub() + md = metadata.metadata + md.kv_cache_manager = None + md._num_tokens = 0 + md._num_generations = 0 metadata.enable_ragged_verification = enabled metadata._invalidate_pool_view_cache = Mock() metadata._compute_kv_lens_row_reorder = Mock() metadata.prepare_dense_topk_indices = Mock() - metadata.kv_lens_cuda = None + md.kv_lens_cuda = None metadata.refresh_ragged_row_kv_lens = Mock() metadata.refresh_token_major_gen_rows = Mock() - with ( - patch.object(TrtllmAttentionMetadata, "on_update_kv_lens"), - patch( - "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata._fused_dsa_meta_enabled", - return_value=False, - ), + with patch( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata._fused_dsa_meta_enabled", + return_value=False, ): - metadata.on_update_kv_lens() + metadata.on_kv_lens_changed(False) expected_calls = int(enabled) assert metadata.refresh_ragged_row_kv_lens.call_count == expected_calls @@ -847,12 +859,13 @@ def test_metadata_spec_update_resizes_radix_workspace( spec_tree_manager, rows_per_sequence, ): - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.max_num_sequences = 2 + metadata = _dsa_stub() + md = metadata.metadata + md.max_num_sequences = 2 metadata.max_draft_tokens = 0 metadata._radix_rows_per_sequence = 1 - metadata.num_sparse_topk = 2 - metadata.is_cuda_graph = True + md.num_sparse_topk = 2 + md.is_cuda_graph = True metadata.nvfp4_mla_fp8_scratch = None metadata.kv_lens_cuda_2d = torch.empty(2, 1) metadata.kv_lens_expanded_host = torch.empty(2) @@ -860,19 +873,8 @@ def test_metadata_spec_update_resizes_radix_workspace( metadata.create_expanded_buffers = Mock() metadata._create_radix_aux_buffers = Mock() - with patch( - "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata." - "TrtllmAttentionMetadata.update_spec_dec_param" - ): - metadata.update_spec_dec_param( - batch_size=2, - is_spec_decoding_enabled=True, - is_spec_dec_tree=is_tree, - is_spec_dec_dynamic_tree=is_dynamic, - max_draft_len=max_draft_len, - max_total_draft_tokens=max_total_draft_tokens, - spec_tree_manager=spec_tree_manager, - ) + md.ext(SpecDecodingMetadata).is_spec_dec_dynamic_tree = is_dynamic + metadata.on_spec_dec_param(max_draft_len, max_total_draft_tokens, spec_tree_manager) assert metadata._radix_rows_per_sequence == rows_per_sequence metadata._create_radix_aux_buffers.assert_called_once_with(True) @@ -986,6 +988,8 @@ def test_transform_local_topk_uses_v2_page_mapping(): _cached_tokens_per_block=4, _cached_pool_view=torch.empty(0), ) + # The stub serves as its own DSA extension. + metadata.ext = lambda cls: metadata expected_indices = torch.tensor([[90]], dtype=torch.int32) with patch( @@ -1052,13 +1056,16 @@ def _make_group_remap_metadata( _cached_tokens_per_block=block_size, _cached_pool_view=torch.empty(0, device=device), ) + # The stub serves as both the core metadata and its DSA extension. + metadata.metadata = metadata + metadata.ext = lambda cls: metadata metadata._ensure_group_remap_struct = MethodType( - DSAtrtllmAttentionMetadata._ensure_group_remap_struct, metadata + DSAMetadata._ensure_group_remap_struct, metadata ) # Bind the real per-step cache invalidation so tests exercise the actual # clear of _group_remap_batched at forward boundaries. metadata._invalidate_pool_view_cache = MethodType( - DSAtrtllmAttentionMetadata._invalidate_pool_view_cache, metadata + DSAMetadata._invalidate_pool_view_cache, metadata ) return metadata, block_table, req_idx @@ -1890,8 +1897,8 @@ def __init__(self): """Initialize with the provided num_cached_tokens.""" self.num_cached_tokens_per_seq = num_cached_tokens - class MockMetadata(DSAtrtllmAttentionMetadata): - """Mock DSA attention metadata for unit testing indexer operations.""" + class MockMetadata(TrtllmAttentionMetadata): + """Mock attention metadata with a DSA extension for unit testing indexer operations.""" def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: """Initialize mock metadata with buffers, block tables, and indexer state. @@ -1902,21 +1909,22 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: GC would free the manager -- on whatever thread runs the collection, which need not have a CUDA context. """ - self.num_sms = deep_gemm.get_num_sms() + dsa = DSAMetadata(self) + dsa.num_sms = deep_gemm.get_num_sms() self.max_num_tokens = max_num_tokens self.request_ids = request_ids self.num_contexts = num_contexts self.num_generations = num_generations self._num_seqs = num_contexts + num_generations self.max_num_sequences = batch_size - self.max_draft_tokens = max_draft_tokens + dsa.max_draft_tokens = max_draft_tokens self.num_sparse_topk = index_topk - self._radix_rows_per_sequence = 1 + self.max_draft_tokens - self._create_radix_aux_buffers(capture_graph=False) - self.enable_indexer_skip = enable_indexer_skip - self.indexer_head_dim = indexer_head_dim - self.indexer_quant_block_size = 128 - self.compress_ratios = [compress_ratio] + dsa._radix_rows_per_sequence = 1 + dsa.max_draft_tokens + dsa._create_radix_aux_buffers(capture_graph=False) + dsa.enable_indexer_skip = enable_indexer_skip + dsa.indexer_head_dim = indexer_head_dim + dsa.indexer_quant_block_size = 128 + dsa.compress_ratios = [compress_ratio] # Keep seq_lens on CPU for split_prefill_chunks and other CPU operations # CUDA kernels will convert to CUDA as needed self.seq_lens = seq_lens.cpu() if seq_lens.is_cuda else seq_lens @@ -1924,36 +1932,36 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: self.kv_cache_params = MockKVCacheParams() self.kv_cache_manager = cache_manager # Effective tokens-per-block for the indexer k-cache slot mapping, - # mirroring DSAtrtllmAttentionMetadata.__post_init__ (dsa.py). + # mirroring DSAMetadata.allocate. tpb = self.kv_cache_manager.tokens_per_block - self._indexer_compress_ratio = _select_indexer_compress_ratio(self.compress_ratios) + dsa._indexer_compress_ratio = _select_indexer_compress_ratio(dsa.compress_ratios) if hasattr(self.kv_cache_manager, "compressed_block_sizes"): - tpb = tpb // _effective_compress_ratio_divisor(self._indexer_compress_ratio) - self._tokens_per_block = tpb + tpb = tpb // _effective_compress_ratio_divisor(dsa._indexer_compress_ratio) + dsa._tokens_per_block = tpb self.kv_lens_cuda_runtime = kv_lens.cuda() - self.indexer_k_cache_block_offsets = torch.zeros( + dsa.indexer_k_cache_block_offsets = torch.zeros( [batch_size, self.kv_cache_manager.max_blocks_per_seq], device="cuda", dtype=torch.int32, ) - self.host_indexer_k_cache_block_offsets = torch.zeros_like( - self.indexer_k_cache_block_offsets, + dsa.host_indexer_k_cache_block_offsets = torch.zeros_like( + dsa.indexer_k_cache_block_offsets, device="cpu", pin_memory=True, ) block_ids = self.kv_cache_manager.get_batch_cache_indices(request_ids) for i in range(len(block_ids)): - self.host_indexer_k_cache_block_offsets[i, : len(block_ids[i])] = torch.tensor( + dsa.host_indexer_k_cache_block_offsets[i, : len(block_ids[i])] = torch.tensor( block_ids[i], dtype=torch.int32 ) - self.indexer_k_cache_block_offsets[:batch_size].copy_( - self.host_indexer_k_cache_block_offsets[:batch_size], non_blocking=True + dsa.indexer_k_cache_block_offsets[:batch_size].copy_( + dsa.host_indexer_k_cache_block_offsets[:batch_size], non_blocking=True ) - self.slot_mapping_fp8 = torch.zeros((num_tokens,), device="cuda", dtype=torch.int64) - self.slot_mapping_scale = torch.zeros((num_tokens,), device="cuda", dtype=torch.int64) - self.scheduler_metadata_buffer = torch.zeros( - (self.num_sms + 1, 2), device="cuda", dtype=torch.int32 + dsa.slot_mapping_fp8 = torch.zeros((num_tokens,), device="cuda", dtype=torch.int64) + dsa.slot_mapping_scale = torch.zeros((num_tokens,), device="cuda", dtype=torch.int64) + dsa.scheduler_metadata_buffer = torch.zeros( + (dsa.num_sms + 1, 2), device="cuda", dtype=torch.int32 ) # DSL needs the next_n=1 schedule preserved (kNumNextNAtoms=1), # so allocate a separate buffer for the full-next_n schedule. @@ -1961,37 +1969,37 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: # `scheduler_metadata_buffer` itself (the alias makes # `Indexer.prepare()`'s second populate overwrite the first). if use_cute_dsl_paged_mqa_logits: - self.scheduler_metadata_buffer_full_next_n = torch.zeros( - (self.num_sms + 1, 2), device="cuda", dtype=torch.int32 + dsa.scheduler_metadata_buffer_full_next_n = torch.zeros( + (dsa.num_sms + 1, 2), device="cuda", dtype=torch.int32 ) else: - self.scheduler_metadata_buffer_full_next_n = self.scheduler_metadata_buffer + dsa.scheduler_metadata_buffer_full_next_n = dsa.scheduler_metadata_buffer # Pre-allocated 2D kv_lens buffer for the DeepGEMM 2D context_lens API. - self.kv_lens_cuda_2d = torch.zeros( - (self.num_seqs, 1 + self.max_draft_tokens), device="cuda", dtype=torch.int32 + dsa.kv_lens_cuda_2d = torch.zeros( + (self.num_seqs, 1 + dsa.max_draft_tokens), device="cuda", dtype=torch.int32 ) if num_generations > 0: gen_kv_lens = ( kv_lens[num_contexts : num_contexts + num_generations].cuda().to(torch.int32) ) - next_n_cap = 1 + self.max_draft_tokens - self.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( + next_n_cap = 1 + dsa.max_draft_tokens + dsa.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( gen_kv_lens.unsqueeze(-1).expand(-1, next_n_cap) ) - self.cu_seqlen_ks = torch.zeros((num_tokens,), device="cuda", dtype=torch.int32) - self.cu_seqlen_ke = torch.zeros((num_tokens,), device="cuda", dtype=torch.int32) + dsa.cu_seqlen_ks = torch.zeros((num_tokens,), device="cuda", dtype=torch.int32) + dsa.cu_seqlen_ke = torch.zeros((num_tokens,), device="cuda", dtype=torch.int32) self.ctx_kv_offsets = torch.zeros((num_tokens, 1), device="cuda", dtype=torch.int32) self.gen_kv_offsets = torch.zeros((num_tokens, 1), device="cuda", dtype=torch.int32) - self.host_slot_mapping_fp8 = torch.zeros_like( - self.slot_mapping_fp8, device="cpu", pin_memory=True + dsa.host_slot_mapping_fp8 = torch.zeros_like( + dsa.slot_mapping_fp8, device="cpu", pin_memory=True ) - self.host_slot_mapping_scale = torch.zeros_like( - self.slot_mapping_scale, device="cpu", pin_memory=True + dsa.host_slot_mapping_scale = torch.zeros_like( + dsa.slot_mapping_scale, device="cpu", pin_memory=True ) self.host_ctx_kv_indptr = torch.zeros( (num_contexts + 1,), device="cpu", pin_memory=True, dtype=torch.int64 ) - self.host_gen_kv_indptr = torch.zeros( + dsa.host_gen_kv_indptr = torch.zeros( (num_generations + 1,), device="cpu", pin_memory=True, dtype=torch.int64 ) # Add host_ctx_cached_token_indptr for prepare_one_prefill_chunk @@ -2001,7 +2009,7 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: self._num_ctx_tokens = num_ctx_tokens self._num_tokens = num_tokens self.num_gen_tokens = num_tokens - num_ctx_tokens - # Also set private attributes used by DSAtrtllmAttentionMetadata + # Also set private attributes used by TrtllmAttentionMetadata self._num_contexts = num_contexts self._num_generations = num_generations self._num_ctx_tokens = num_ctx_tokens @@ -2017,7 +2025,7 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: kv_lens[num_contexts : num_contexts + num_generations], dim=0, dtype=torch.int64, - out=self.host_gen_kv_indptr[1 : num_generations + 1], + out=dsa.host_gen_kv_indptr[1 : num_generations + 1], ) # Compute host_ctx_cached_token_indptr from num_cached_tokens if num_contexts > 0: @@ -2030,8 +2038,8 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: ) # Add indexer-specific attributes - self.indexer_max_chunk_size = indexer_max_chunk_size - self.indexer_prefill_chunks = None + dsa.indexer_max_chunk_size = indexer_max_chunk_size + dsa.indexer_prefill_chunks = None self.enable_context_mla_with_cached_kv = enable_context_mla_with_cached_kv @@ -2050,47 +2058,47 @@ def __init__(self): self.runtime_features = RuntimeFeatures() # Add expanded buffers for MTP support. Mirrors the production gate in - # DSAtrtllmAttentionMetadata.prepare_for_spec_decode: sm100+ DeepGEMM + # DSAMetadata.prepare_for_spec_decode: sm100+ DeepGEMM # runs a native next_n for any MTP depth, so only sm90 expands. The DSL # kernel supports arbitrary next_n natively, so it never needs expansion. - self.use_expanded_buffers_for_mtp = ( + dsa.use_expanded_buffers_for_mtp = ( not use_cute_dsl_paged_mqa_logits - and self.max_draft_tokens > 1 + and dsa.max_draft_tokens > 1 and get_sm_version() == 90 ) - self.kv_lens_expanded_cuda = torch.zeros( - (self.num_seqs * (1 + self.max_draft_tokens),), device="cuda", dtype=torch.int32 + dsa.kv_lens_expanded_cuda = torch.zeros( + (self.num_seqs * (1 + dsa.max_draft_tokens),), device="cuda", dtype=torch.int32 ) - self.kv_lens_expanded_host = torch.zeros_like( - self.kv_lens_expanded_cuda, device="cpu", pin_memory=True + dsa.kv_lens_expanded_host = torch.zeros_like( + dsa.kv_lens_expanded_cuda, device="cpu", pin_memory=True ) - self.block_table_expanded = torch.zeros( + dsa.block_table_expanded = torch.zeros( ( - self.num_seqs * (1 + self.max_draft_tokens), + self.num_seqs * (1 + dsa.max_draft_tokens), self.kv_cache_manager.max_blocks_per_seq, ), device="cuda", dtype=torch.int32, ) - self.host_block_table_expanded = torch.zeros_like( - self.block_table_expanded, device="cpu", pin_memory=True + dsa.host_block_table_expanded = torch.zeros_like( + dsa.block_table_expanded, device="cpu", pin_memory=True ) - self.scheduler_metadata_buffer_expanded = torch.zeros( - (self.num_sms + 1, 2), device="cuda", dtype=torch.int32 + dsa.scheduler_metadata_buffer_expanded = torch.zeros( + (dsa.num_sms + 1, 2), device="cuda", dtype=torch.int32 ) - if self.max_draft_tokens == 3: + if dsa.max_draft_tokens == 3: self.scheduler_metadata_buffer_mtp3 = torch.zeros( - (self.num_sms // 2 + 1, 2), device="cuda", dtype=torch.int32 + (dsa.num_sms // 2 + 1, 2), device="cuda", dtype=torch.int32 ) - if self.use_expanded_buffers_for_mtp: + if dsa.use_expanded_buffers_for_mtp: gen_kv_lens = kv_lens[num_contexts : self.num_seqs] gen_kv_lens_expanded = torch.stack( - [gen_kv_lens] * (1 + self.max_draft_tokens), dim=0 + [gen_kv_lens] * (1 + dsa.max_draft_tokens), dim=0 ) gen_kv_lens_expanded = gen_kv_lens_expanded.transpose(0, 1).contiguous().flatten() - self.kv_lens_expanded_host[: self.num_gen_tokens].copy_(gen_kv_lens_expanded) - self.kv_lens_expanded_cuda[: self.num_gen_tokens].copy_( - self.kv_lens_expanded_host[: self.num_gen_tokens], non_blocking=True + dsa.kv_lens_expanded_host[: self.num_gen_tokens].copy_(gen_kv_lens_expanded) + dsa.kv_lens_expanded_cuda[: self.num_gen_tokens].copy_( + dsa.kv_lens_expanded_host[: self.num_gen_tokens], non_blocking=True ) if self.kv_cache_manager is not None: @@ -2099,39 +2107,40 @@ def __init__(self): if len(gen_block_ids) > 0: # Find max length and create padded tensor max_len = max(len(bid) for bid in gen_block_ids) - gen_block_tensor = self.host_indexer_k_cache_block_offsets[ + gen_block_tensor = dsa.host_indexer_k_cache_block_offsets[ self.num_contexts : self.num_seqs, :max_len ] expanded_blocks = gen_block_tensor.repeat_interleave( - 1 + self.max_draft_tokens, dim=0 + 1 + dsa.max_draft_tokens, dim=0 ) - self.host_block_table_expanded[: self.num_gen_tokens, :max_len].copy_( + dsa.host_block_table_expanded[: self.num_gen_tokens, :max_len].copy_( expanded_blocks, non_blocking=True ) - self.block_table_expanded[: self.num_gen_tokens].copy_( - self.host_block_table_expanded[: self.num_gen_tokens], non_blocking=True + dsa.block_table_expanded[: self.num_gen_tokens].copy_( + dsa.host_block_table_expanded[: self.num_gen_tokens], non_blocking=True ) # Add skip indexer attributes - self.topk_indices_buffer = torch.zeros( + dsa.topk_indices_buffer = torch.zeros( (num_tokens, self.num_sparse_topk), device="cuda", dtype=torch.int32 ) - if self.num_contexts > 0 and self.enable_indexer_skip: - self.skip_indexer_for_ctx_reqs = ( + if self.num_contexts > 0 and dsa.enable_indexer_skip: + dsa.skip_indexer_for_ctx_reqs = ( kv_lens[: self.num_contexts].max().item() <= self.num_sparse_topk ) else: - self.skip_indexer_for_ctx_reqs = False + dsa.skip_indexer_for_ctx_reqs = False - if self.num_generations > 0 and self.enable_indexer_skip: - self.max_draft_tokens + 1 - self.skip_indexer_for_gen_reqs = ( + if self.num_generations > 0 and dsa.enable_indexer_skip: + dsa.max_draft_tokens + 1 + dsa.skip_indexer_for_gen_reqs = ( kv_lens[self.num_contexts : self.num_seqs].max().item() <= self.num_sparse_topk ) else: - self.skip_indexer_for_gen_reqs = False - self.prepare_dense_topk_indices(self.kv_lens_cuda_runtime, device=True) + dsa.skip_indexer_for_gen_reqs = False + dsa.prepare_dense_topk_indices(self.kv_lens_cuda_runtime, device=True) + self.replace_ext(dsa) @property def num_seqs(self) -> int: @@ -2145,9 +2154,10 @@ def num_seqs(self) -> int: @pytest.mark.parametrize("disabled_ratio", [0, 1]) def test_indexer_compress_ratio_zero_or_one_means_uncompressed(disabled_ratio): - metadata = object.__new__(DSAtrtllmAttentionMetadata) - metadata.kv_cache_manager = Mock() - metadata.kv_cache_manager.max_seq_len = 4096 + metadata = _dsa_stub() + md = metadata.metadata + md.kv_cache_manager = Mock() + md.kv_cache_manager.max_seq_len = 4096 kv_lens = torch.tensor([0, 1, 128, 4096], dtype=torch.int32) metadata._indexer_compress_ratio = disabled_ratio @@ -2223,19 +2233,20 @@ def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(dsa_cache_ma num_tokens=num_tokens, indexer_head_dim=head_dim, ) + dsa = metadata.ext(DSAMetadata) Indexer.prepare(metadata) - expected_fp8 = metadata.slot_mapping_fp8[:num_tokens].clone() - expected_scale = metadata.slot_mapping_scale[:num_tokens].clone() + expected_fp8 = dsa.slot_mapping_fp8[:num_tokens].clone() + expected_scale = dsa.slot_mapping_scale[:num_tokens].clone() - metadata.host_slot_mapping_fp8.zero_() - metadata.host_slot_mapping_scale.zero_() - metadata.slot_mapping_fp8.zero_() - metadata.slot_mapping_scale.zero_() + dsa.host_slot_mapping_fp8.zero_() + dsa.host_slot_mapping_scale.zero_() + dsa.slot_mapping_fp8.zero_() + dsa.slot_mapping_scale.zero_() Indexer.recompute_slot_mappings(metadata) - torch.testing.assert_close(metadata.slot_mapping_fp8[:num_tokens], expected_fp8) - torch.testing.assert_close(metadata.slot_mapping_scale[:num_tokens], expected_scale) + torch.testing.assert_close(dsa.slot_mapping_fp8[:num_tokens], expected_fp8) + torch.testing.assert_close(dsa.slot_mapping_scale[:num_tokens], expected_scale) @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -2288,6 +2299,7 @@ def test_indexer_k_cache_scatter_custom_op(dsa_cache_manager): num_tokens=num_tokens, indexer_head_dim=head_dim, ) + dsa = metadata.ext(DSAMetadata) from tensorrt_llm._torch.attention.backends.sparse.dsa import Indexer @@ -2307,8 +2319,8 @@ def test_indexer_k_cache_scatter_custom_op(dsa_cache_manager): ) k_scale_bytes = k_scale_flat.view(torch.uint8).view(num_tokens, scale_size) - flat_indices_fp8 = metadata.slot_mapping_fp8[:num_tokens] - flat_indices_scale = metadata.slot_mapping_scale[:num_tokens] + flat_indices_fp8 = dsa.slot_mapping_fp8[:num_tokens] + flat_indices_scale = dsa.slot_mapping_scale[:num_tokens] # ========== Use Different Layers for CUDA vs Python ========== # Both get the same input data, but write to different layers @@ -2490,7 +2502,11 @@ def test_fp8_k_cache_roundtrip(dsa_cache_manager): # Compute block location block_idx_in_seq = local_token_idx // block_size pos_in_block = local_token_idx % block_size - block_id = metadata.host_indexer_k_cache_block_offsets[req_idx, block_idx_in_seq].item() + block_id = ( + metadata.ext(DSAMetadata) + .host_indexer_k_cache_block_offsets[req_idx, block_idx_in_seq] + .item() + ) # Extract stored scale scale_bytes = cache_flat[ @@ -2514,7 +2530,7 @@ def test_fp8_k_cache_roundtrip(dsa_cache_manager): @skip_pre_hopper # next_n=5 (MTP-4) exercises the native paged-MQA path on sm100+ (no buffer # expansion); on sm90 it falls back to the expanded buffers. See -# DSAtrtllmAttentionMetadata.prepare_for_spec_decode. +# DSAMetadata.prepare_for_spec_decode. @pytest.mark.parametrize("batch_size,next_n", [(4, 1), (2, 2), (4, 3), (4, 4), (2, 5)]) @pytest.mark.parametrize( "backend", @@ -2658,6 +2674,7 @@ def test_indexer_decode_with_paged_kv_cache( compress_ratio=compress_ratio, indexer_head_dim=head_dim, ) + metadata_gen_dsa = metadata_gen.ext(DSAMetadata) Indexer.prepare(metadata_gen) k_gen_fp8, k_gen_scale = fp8_utils.fp8_quantize_1x128_sf_transpose(k_gen_bf16) @@ -2669,11 +2686,11 @@ def test_indexer_decode_with_paged_kv_cache( kv_cache_fp8_pool = cache_manager.get_indexer_k_cache_buffers(layer_idx) q_fp8 = q.to(torch.float8_e4m3fn) - if not metadata_gen.use_expanded_buffers_for_mtp: + if not metadata_gen_dsa.use_expanded_buffers_for_mtp: q_fp8 = q_fp8 # New DeepGEMM 2D context_lens API: shape (batch_size, next_n). - context_lens = metadata_gen.kv_lens_cuda_2d[0:batch_size, 0:next_n].contiguous() - block_table = metadata_gen.indexer_k_cache_block_offsets[0:batch_size] + context_lens = metadata_gen_dsa.kv_lens_cuda_2d[0:batch_size, 0:next_n].contiguous() + block_table = metadata_gen_dsa.indexer_k_cache_block_offsets[0:batch_size] # The upgraded DeepGEMM paged MQA logits kernel picks # num_kv_multicast=1 on SM100 for every next_n it supports # (verified by the _schedule_meta_size assertion in @@ -2681,14 +2698,14 @@ def test_indexer_decode_with_paged_kv_cache( # the legacy (num_sms // 2 + 1, 2) layout). The base scheduler # buffer is the only layout the kernel will accept now. The DSL # path uses its own schedule buffer instead (see below). - scheduler_metadata_buffer = metadata_gen.scheduler_metadata_buffer + scheduler_metadata_buffer = metadata_gen_dsa.scheduler_metadata_buffer else: q_fp8 = q_fp8.view(-1, 1, *q_fp8.shape[2:]) num_tokens = batch_size * next_n # New API requires 2D; each expanded token becomes a (1,) row. - context_lens = metadata_gen.kv_lens_expanded_cuda[:num_tokens].view(-1, 1) - block_table = metadata_gen.block_table_expanded[:num_tokens] - scheduler_metadata_buffer = metadata_gen.scheduler_metadata_buffer_expanded + context_lens = metadata_gen_dsa.kv_lens_expanded_cuda[:num_tokens].view(-1, 1) + block_table = metadata_gen_dsa.block_table_expanded[:num_tokens] + scheduler_metadata_buffer = metadata_gen_dsa.scheduler_metadata_buffer_expanded if use_dsl: # DSL design: kNumNextNAtoms = 1 for any real next_n. Reuse the base @@ -2696,7 +2713,7 @@ def test_indexer_decode_with_paged_kv_cache( # (from a (num_gen, 1) input shape). 1D contiguous kv_lens slice # suffices for context_lens (all next_n positions share the same KV # length on this path; kernel signature is 1D). - dsl_context_lens = metadata_gen.gen_indexer_kv_lens_cuda_runtime + dsl_context_lens = metadata_gen_dsa.gen_indexer_kv_lens_cuda_runtime assert dsl_context_lens is not None logits = torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( q_fp8, @@ -2704,7 +2721,7 @@ def test_indexer_decode_with_paged_kv_cache( weights, dsl_context_lens, block_table, - metadata_gen.scheduler_metadata_buffer, + metadata_gen_dsa.scheduler_metadata_buffer, max_model_len, ) else: @@ -2737,7 +2754,7 @@ def test_indexer_decode_with_paged_kv_cache( for token_pos in range(seq_context_len): block_idx = token_pos // block_size pos_in_block = token_pos % block_size - physical_block_id = metadata_gen.indexer_k_cache_block_offsets[ + physical_block_id = metadata_gen_dsa.indexer_k_cache_block_offsets[ seq_idx, block_idx ].item() if physical_block_id >= 0: @@ -2750,7 +2767,7 @@ def test_indexer_decode_with_paged_kv_cache( token_pos = seq_context_len + gen_token_idx block_idx = token_pos // block_size pos_in_block = token_pos % block_size - physical_block_id = metadata_gen.indexer_k_cache_block_offsets[ + physical_block_id = metadata_gen_dsa.indexer_k_cache_block_offsets[ seq_idx, block_idx ].item() if physical_block_id >= 0: @@ -2769,7 +2786,7 @@ def test_indexer_decode_with_paged_kv_cache( weights, num_tokens_cuda[0:batch_size], num_kv_tokens_cuda[0:batch_size], - metadata_gen.indexer_k_cache_block_offsets, + metadata_gen_dsa.indexer_k_cache_block_offsets, max_model_len, compress_ratio, ) @@ -2921,13 +2938,14 @@ def _force_direct_path(meta): - re-allocate _full_next_n as a separate tensor (so the second populate doesn't clobber the first) """ - meta.use_expanded_buffers_for_mtp = False - meta.scheduler_metadata_buffer_full_next_n = torch.zeros( - (meta.num_sms + 1, 2), device="cuda", dtype=torch.int32 + meta_dsa = meta.ext(DSAMetadata) + meta_dsa.use_expanded_buffers_for_mtp = False + meta_dsa.scheduler_metadata_buffer_full_next_n = torch.zeros( + (meta_dsa.num_sms + 1, 2), device="cuda", dtype=torch.int32 ) def _force_dsl_expand_setup(meta): - """Mirror DSAtrtllmAttentionMetadata.prepare's `expand_for_dsl` + """Mirror DSAMetadata.prepare's `expand_for_dsl` populate so Indexer.prepare can build the schedule from kv_lens_expanded_cuda. Calls the wave-aware picker; expand only happens when factor > 1. @@ -2936,10 +2954,11 @@ def _force_dsl_expand_setup(meta): and writes the picker decision to `dsl_expand_factor`/`dsl_atom` so downstream Indexer.prepare + forward read the same values. """ + meta_dsa = meta.ext(DSAMetadata) if meta.num_generations == 0: return - meta.expand_for_dsl = True - next_n = 1 + meta.max_draft_tokens + meta_dsa.expand_for_dsl = True + next_n = 1 + meta_dsa.max_draft_tokens # FP4 kernel supports atoms ∈ {1, 2, 3}; matches dsa.py's # `kernel_atoms = (1, 2, 3) if use_fp4 else (1, 2, 3, 4)` path. gen_kv_lens = meta.kv_lens[meta.num_contexts : meta.num_seqs] @@ -2948,33 +2967,33 @@ def _force_dsl_expand_setup(meta): next_n, batch_size=meta.num_generations, max_ctx=max_ctx, - num_sms=meta.num_sms, + num_sms=meta_dsa.num_sms, kernel_atoms=(1, 2, 3), ) - meta.dsl_expand_factor = factor - meta.dsl_atom = atom + meta_dsa.dsl_expand_factor = factor + meta_dsa.dsl_atom = atom if factor <= 1: # Picker chose kernel-native; no buffer populate needed. return num_tokens = meta.num_generations * factor gen_kv_lens_expanded = gen_kv_lens.repeat_interleave(factor) - meta.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) - meta.kv_lens_expanded_cuda[:num_tokens].copy_( - meta.kv_lens_expanded_host[:num_tokens], non_blocking=True + meta_dsa.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) + meta_dsa.kv_lens_expanded_cuda[:num_tokens].copy_( + meta_dsa.kv_lens_expanded_host[:num_tokens], non_blocking=True ) # block_table: repeat each row `factor` times. - max_len = meta.host_indexer_k_cache_block_offsets.shape[1] - gen_block_tensor = meta.host_indexer_k_cache_block_offsets[ + max_len = meta_dsa.host_indexer_k_cache_block_offsets.shape[1] + gen_block_tensor = meta_dsa.host_indexer_k_cache_block_offsets[ meta.num_contexts : meta.num_seqs, :max_len ] expanded_blocks = gen_block_tensor.repeat_interleave(factor, dim=0) - meta.host_block_table_expanded[:num_tokens, :max_len].copy_( + meta_dsa.host_block_table_expanded[:num_tokens, :max_len].copy_( expanded_blocks, non_blocking=True ) - meta.block_table_expanded[:num_tokens].copy_( - meta.host_block_table_expanded[:num_tokens], non_blocking=True + meta_dsa.block_table_expanded[:num_tokens].copy_( + meta_dsa.host_block_table_expanded[:num_tokens], non_blocking=True ) - meta.block_table_expanded.clamp_(min=0) + meta_dsa.block_table_expanded.clamp_(min=0) # ---- Phase 1: write context tokens as FP4 ---- print(f"\n=== Phase 1: Context (variable tokens/seq) ({backend}) ===") @@ -3021,6 +3040,7 @@ def _force_dsl_expand_setup(meta): max_draft_tokens=next_n - 1, use_cute_dsl_paged_mqa_logits=use_dsl, ) + metadata_gen_dsa = metadata_gen.ext(DSAMetadata) if not use_dsl: _force_direct_path(metadata_gen) else: @@ -3060,18 +3080,18 @@ def _force_dsl_expand_setup(meta): dsl_q = q_fp4.view(torch.uint8) dsl_sf_q = sf_q dsl_context_lens = metadata_gen.kv_lens_cuda_runtime[0:batch_size] - dsl_block_table = metadata_gen.indexer_k_cache_block_offsets[0:batch_size] - dsl_schedule_meta = metadata_gen.scheduler_metadata_buffer + dsl_block_table = metadata_gen_dsa.indexer_k_cache_block_offsets[0:batch_size] + dsl_schedule_meta = metadata_gen_dsa.scheduler_metadata_buffer - if metadata_gen.dsl_expand_factor > 1: - factor = metadata_gen.dsl_expand_factor - eff_next_n = metadata_gen.dsl_atom + if metadata_gen_dsa.dsl_expand_factor > 1: + factor = metadata_gen_dsa.dsl_expand_factor + eff_next_n = metadata_gen_dsa.dsl_atom exp_B = batch_size * factor dsl_q = dsl_q.reshape(exp_B, eff_next_n, heads, pe_dim) dsl_sf_q = dsl_sf_q.reshape(exp_B, eff_next_n, heads) - dsl_context_lens = metadata_gen.kv_lens_expanded_cuda[:exp_B] - dsl_block_table = metadata_gen.block_table_expanded[:exp_B] - dsl_schedule_meta = metadata_gen.scheduler_metadata_buffer_expanded + dsl_context_lens = metadata_gen_dsa.kv_lens_expanded_cuda[:exp_B] + dsl_block_table = metadata_gen_dsa.block_table_expanded[:exp_B] + dsl_schedule_meta = metadata_gen_dsa.scheduler_metadata_buffer_expanded logits = torch.ops.trtllm.cute_dsl_fp4_paged_mqa_logits( dsl_q, @@ -3090,15 +3110,15 @@ def _force_dsl_expand_setup(meta): # use_expanded_buffers_for_mtp gate (lines 962-966) inherits FP8's # next_n limit and is overly conservative for FP4; fixing it is # out of scope here. - context_lens_2d = metadata_gen.kv_lens_cuda_2d[0:batch_size, 0:next_n] - block_table_kernel = metadata_gen.indexer_k_cache_block_offsets[0:batch_size] + context_lens_2d = metadata_gen_dsa.kv_lens_cuda_2d[0:batch_size, 0:next_n] + block_table_kernel = metadata_gen_dsa.indexer_k_cache_block_offsets[0:batch_size] # Schedule buffer: next_n=1 uses the base buffer (built from # (num_gen, 1)); next_n>1 uses the _full_next_n variant (built from # (num_gen, next_n) so num_next_n_atoms encodes the real next_n). if next_n == 1: - scheduler_metadata = metadata_gen.scheduler_metadata_buffer + scheduler_metadata = metadata_gen_dsa.scheduler_metadata_buffer else: - scheduler_metadata = metadata_gen.scheduler_metadata_buffer_full_next_n + scheduler_metadata = metadata_gen_dsa.scheduler_metadata_buffer_full_next_n logits = fp8_fp4_paged_mqa_logits( (q_fp4, sf_q), @@ -3149,7 +3169,7 @@ def _force_dsl_expand_setup(meta): for token_pos in range(seq_context_len): block_idx = token_pos // block_size pos_in_block = token_pos % block_size - physical_block_id = metadata_gen.indexer_k_cache_block_offsets[ + physical_block_id = metadata_gen_dsa.indexer_k_cache_block_offsets[ seq_idx, block_idx ].item() if physical_block_id >= 0: @@ -3160,7 +3180,7 @@ def _force_dsl_expand_setup(meta): token_pos = seq_context_len + gen_token_idx block_idx = token_pos // block_size pos_in_block = token_pos % block_size - physical_block_id = metadata_gen.indexer_k_cache_block_offsets[ + physical_block_id = metadata_gen_dsa.indexer_k_cache_block_offsets[ seq_idx, block_idx ].item() if physical_block_id >= 0: @@ -3178,7 +3198,7 @@ def _force_dsl_expand_setup(meta): weights, num_tokens_cuda[0:batch_size], num_kv_tokens_cuda[0:batch_size], - metadata_gen.indexer_k_cache_block_offsets, + metadata_gen_dsa.indexer_k_cache_block_offsets, max_model_len, ) print(f"✓ Reference output shape: {ref_logits.shape}") @@ -3836,15 +3856,16 @@ def test_indexer_chunked_prefill( indexer_head_dim=head_dim, compress_ratio=compress_ratio, ) + metadata_chunked_dsa = metadata_chunked.ext(DSAMetadata) Indexer.prepare(metadata_chunked) - assert metadata_chunked.indexer_prefill_chunks is not None - num_chunks = len(metadata_chunked.indexer_prefill_chunks) + assert metadata_chunked_dsa.indexer_prefill_chunks is not None + num_chunks = len(metadata_chunked_dsa.indexer_prefill_chunks) print(f"✓ Created {num_chunks} chunks") # Print chunk details - for i, chunk in enumerate(metadata_chunked.indexer_prefill_chunks): + for i, chunk in enumerate(metadata_chunked_dsa.indexer_prefill_chunks): num_q = chunk.token_end - chunk.token_start num_k = chunk.k_token_end - chunk.k_token_start print( @@ -3877,11 +3898,12 @@ def test_indexer_chunked_prefill( indexer_head_dim=head_dim, compress_ratio=compress_ratio, ) + metadata_baseline_dsa = metadata_baseline.ext(DSAMetadata) Indexer.prepare(metadata_baseline) - if metadata_baseline.indexer_prefill_chunks is not None: - num_baseline_chunks = len(metadata_baseline.indexer_prefill_chunks) + if metadata_baseline_dsa.indexer_prefill_chunks is not None: + num_baseline_chunks = len(metadata_baseline_dsa.indexer_prefill_chunks) print(f"✓ Created {num_baseline_chunks} chunk(s) (effectively non-chunked)") indexer._update_k_cache(k_fp8, k_scale, metadata_baseline) @@ -4336,9 +4358,10 @@ def make_inputs(n_tokens): # context stash branch reads seq_lens_cuda (a read-only property); set its backing field. meta0._seq_lens_cuda = kv_lens.clone().cuda() - meta0.in_mtp_draft_loop = True + meta0.ext(DSAMetadata).in_mtp_draft_loop = True + meta0_dsa = meta0.ext(DSAMetadata) # Step 0 computes (does not skip) Top-K; the reuse gate reads this directly. - meta0.indexer_skip_topk = False + meta0_dsa.indexer_skip_topk = False h0, q0, k0_fp8, k0_scale, w0 = make_inputs(step0_tokens) indexer._update_k_cache(k0_fp8, k0_scale, meta0) try: @@ -4352,7 +4375,7 @@ def make_inputs(n_tokens): ctx_last = torch.cumsum(kv_lens.to(torch.long), dim=0).cuda() - 1 expected_rows = topk0[ctx_last, :] - stash = meta0.shared_topk_indices + stash = meta0_dsa.shared_topk_indices assert stash is not None, "Step 0 should stash shared_topk_indices" assert stash.shape[0] == batch_size assert torch.equal(stash, expected_rows), f"{step0_mode} step-0 stash mismatch" @@ -4375,9 +4398,9 @@ def make_inputs(n_tokens): max_draft_tokens=md, ) Indexer.prepare(meta) - meta.in_mtp_draft_loop = True - meta.shared_topk_indices = stash - meta.indexer_skip_topk = True + meta.ext(DSAMetadata).in_mtp_draft_loop = True + meta.ext(DSAMetadata).shared_topk_indices = stash + meta.ext(DSAMetadata).indexer_skip_topk = True hs, qs, ks_fp8, ks_scale, ws = make_inputs(batch_size) indexer._update_k_cache(ks_fp8, ks_scale, meta) topk = indexer.sparse_attn_indexer(meta, hs, qs, ks_fp8, ks_scale, ws) @@ -4467,7 +4490,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback( Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) - assert metadata_custom.indexer_prefill_chunks is not None + assert metadata_custom.ext(DSAMetadata).indexer_prefill_chunks is not None try: topk_indices_custom = indexer.sparse_attn_indexer( @@ -4581,7 +4604,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback( Indexer.prepare(metadata_custom) indexer._update_k_cache(k_fp8, k_scale, metadata_custom) # Force single-pass path by setting indexer_prefill_chunks to None - metadata_custom.indexer_prefill_chunks = None + metadata_custom.ext(DSAMetadata).indexer_prefill_chunks = None try: topk_indices_custom = indexer.sparse_attn_indexer( @@ -4609,7 +4632,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback( Indexer.prepare(metadata_fallback) indexer._update_k_cache(k_fp8, k_scale, metadata_fallback) # Force single-pass path by setting indexer_prefill_chunks to None - metadata_fallback.indexer_prefill_chunks = None + metadata_fallback.ext(DSAMetadata).indexer_prefill_chunks = None _set_torch_top_k(indexer) topk_indices_fallback = indexer.sparse_attn_indexer( @@ -4634,14 +4657,15 @@ def test_indexer_prefill_single_pass_custom_vs_fallback( ) Indexer.prepare(metadata_skip) indexer._update_k_cache(k_fp8, k_scale, metadata_skip) - metadata_skip.indexer_prefill_chunks = None + metadata_skip.ext(DSAMetadata).indexer_prefill_chunks = None + metadata_skip_dsa = metadata_skip.ext(DSAMetadata) indexer.top_k = TopK( index_topk, prefill_implementation=TopKImplementation.CUDA_RADIX, decode_implementation=TopKImplementation.CUTE_DSL_GVR, gvr_self_sampling=False, ) - metadata_skip.gvr_prior_indices = torch.zeros( + metadata_skip_dsa.gvr_prior_indices = torch.zeros( (cache_manager.num_local_layers, batch_size, index_topk), device="cuda", dtype=torch.int32, @@ -4657,7 +4681,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback( last_rows = torch.cumsum(metadata_skip.seq_lens[:batch_size], dim=0) - 1 local_layer = cache_manager.layer_offsets[layer_idx] torch.testing.assert_close( - metadata_skip.gvr_prior_indices[local_layer, :batch_size], + metadata_skip_dsa.gvr_prior_indices[local_layer, :batch_size], topk_indices_skip[last_rows], ) @@ -4749,11 +4773,12 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip, ds enable_context_mla_with_cached_kv=True, indexer_head_dim=head_dim, ) + dsa = metadata.ext(DSAMetadata) Indexer.prepare(metadata) if not enable_indexer_skip: - assert metadata.slot_mapping_fp8_fullkv is not metadata.slot_mapping_fp8 - assert metadata.slot_mapping_scale_fullkv is not metadata.slot_mapping_scale + assert dsa.slot_mapping_fp8_fullkv is not dsa.slot_mapping_fp8 + assert dsa.slot_mapping_scale_fullkv is not dsa.slot_mapping_scale indexer._update_k_cache(k_fp8, k_scale, metadata) # Test custom kernel @@ -4970,8 +4995,8 @@ def test_indexer_mla_prefill_groups_requests_by_logits_budget(monkeypatch, dsa_c ) packed_metadata, topk_packed = run(1 << 62) split_metadata, topk_split = run(0) - packed_chunks = packed_metadata.indexer_prefill_chunks - split_chunks = split_metadata.indexer_prefill_chunks + packed_chunks = packed_metadata.ext(DSAMetadata).indexer_prefill_chunks + split_chunks = split_metadata.ext(DSAMetadata).indexer_prefill_chunks assert len(packed_chunks) == 1 assert len(split_chunks) == len(seq_lens) @@ -5022,10 +5047,10 @@ def fake_allgather(tensor, mapping, dim=-1, sizes=None): for rank in range(tp_size): mapping = Mapping(world_size=tp_size, rank=rank, tp_size=tp_size, pp_size=1) metadata, _ = run(elem_budget, mapping) - q_begin, q_end, q_sizes = metadata.indexer_prefill_q_split + q_begin, q_end, q_sizes = metadata.ext(DSAMetadata).indexer_prefill_q_split assert q_sizes == expected_sizes assert (q_begin, q_end) == (sum(expected_sizes[:rank]), sum(expected_sizes[: rank + 1])) - chunks = metadata.indexer_prefill_chunks + chunks = metadata.ext(DSAMetadata).indexer_prefill_chunks assert [(c.token_start, c.token_end) for c in chunks] == list( zip([q_begin] + [c.token_end for c in chunks[:-1]], [c.token_end for c in chunks]) ) @@ -5050,11 +5075,11 @@ class TestPrepareRestoreAttnMetadataForDraftReplay: def _make_mock_metadata(): """Create a mock attention metadata object with KV cache block offsets.""" meta = Mock() - meta.enable_flash_mla = False meta.kv_cache_manager = Mock(name="target_kv_cache_manager") meta.kv_cache_block_offsets = torch.tensor([10, 20, 30]) meta.host_kv_cache_block_offsets = torch.tensor([10, 20, 30]) - meta.draft_kv_cache_block_offsets = torch.tensor([100, 200, 300]) + draft_kv = SimpleNamespace(kv_cache_block_offsets=torch.tensor([100, 200, 300])) + meta.find_ext = lambda cls: draft_kv if cls is DraftKVMetadata else None meta.prepare_for_draft_forward.return_value = None meta.draft_replay_swapped_attrs = {} for helper in ("record_draft_swap", "swap_for_draft", "restore_draft_swaps"): @@ -5080,7 +5105,7 @@ def test_prepare_swaps_and_restore_recovers(self): original_host_offsets = meta.host_kv_cache_block_offsets.clone() def is_attn_metadata(obj, cls): - if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata): + if cls is TrtllmAttentionMetadata: return obj is meta return builtins.isinstance(obj, cls) @@ -5109,7 +5134,10 @@ def test_native_dsa_replay_swaps_and_restores_buffers(self): mgr = object.__new__(DSACacheManagerV2) mgr.host_kv_cache_block_offsets = torch.tensor([100, 200, 300]) meta.enable_context_mla_with_cached_kv = False - meta._invalidate_pool_view_cache = Mock() + dsa = object.__new__(DSAMetadata) + dsa.metadata = meta + dsa._invalidate_pool_view_cache = Mock() + meta.ext = lambda cls: dsa buffer_pairs = { "host_indexer_k_cache_block_offsets": "host_draft_indexer_k_cache_block_offsets", "indexer_k_cache_block_offsets": "draft_indexer_k_cache_block_offsets", @@ -5126,22 +5154,16 @@ def test_native_dsa_replay_swaps_and_restores_buffers(self): for target_name, draft_name in buffer_pairs.items(): target_buffers[target_name] = object() draft_buffers[target_name] = object() - setattr(meta, target_name, target_buffers[target_name]) - setattr(meta, draft_name, draft_buffers[target_name]) - del meta.slot_mapping_fp8_fullkv - del meta.slot_mapping_scale_fullkv + setattr(dsa, target_name, target_buffers[target_name]) + setattr(dsa, draft_name, draft_buffers[target_name]) - meta.prepare_for_draft_forward.side_effect = ( - lambda: DSAtrtllmAttentionMetadata.prepare_for_draft_forward(meta) - ) - meta.restore_after_draft_forward.side_effect = ( - lambda saved_state: DSAtrtllmAttentionMetadata.restore_after_draft_forward( - meta, saved_state - ) + meta.prepare_for_draft_forward.side_effect = lambda: {DSAMetadata: dsa.draft_swap()} + meta.restore_after_draft_forward.side_effect = lambda saved_state: dsa.draft_restore( + saved_state[DSAMetadata] ) def is_attn_metadata(obj, cls): - if cls in (TrtllmAttentionMetadata, DSAtrtllmAttentionMetadata): + if cls is TrtllmAttentionMetadata: return obj is meta return builtins.isinstance(obj, cls) @@ -5158,10 +5180,10 @@ def is_attn_metadata(obj, cls): saved = prepare_attn_metadata_for_draft_replay(meta, mgr) assert "saved_backend_state" in saved - assert meta.slot_mapping_fp8 is draft_buffers["slot_mapping_fp8"] + assert dsa.slot_mapping_fp8 is draft_buffers["slot_mapping_fp8"] restore_attn_metadata_after_draft_replay(meta, saved) - assert meta.slot_mapping_fp8 is target_buffers["slot_mapping_fp8"] - assert meta.slot_mapping_fp8_fullkv is target_buffers["slot_mapping_fp8"] + assert dsa.slot_mapping_fp8 is target_buffers["slot_mapping_fp8"] + assert dsa.slot_mapping_fp8_fullkv is target_buffers["slot_mapping_fp8"] @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @@ -5218,6 +5240,7 @@ def test_cutedsl_mqa_logits_output_buffer_persistent(dsa_cache_manager): index_topk=index_topk, use_cute_dsl_paged_mqa_logits=True, ) + dsa = metadata.ext(DSAMetadata) Indexer.prepare(metadata) kv_cache = cache_manager.get_indexer_k_cache_buffers(0) @@ -5225,9 +5248,9 @@ def test_cutedsl_mqa_logits_output_buffer_persistent(dsa_cache_manager): torch.float8_e4m3fn ) weights = torch.randn((batch_size * next_n, heads), device="cuda", dtype=torch.float32) - context_lens = metadata.gen_indexer_kv_lens_cuda_runtime - block_table = metadata.indexer_k_cache_block_offsets[0:batch_size] - sched = metadata.scheduler_metadata_buffer + context_lens = dsa.gen_indexer_kv_lens_cuda_runtime + block_table = dsa.indexer_k_cache_block_offsets[0:batch_size] + sched = dsa.scheduler_metadata_buffer def _mqa(): return torch.ops.trtllm.cute_dsl_fp8_paged_mqa_logits( @@ -5294,7 +5317,7 @@ def test_topk_indices_buffer_cuda_graph(dsa_cache_manager): index_topk=index_topk, enable_indexer_skip=True, ) - assert metadata.skip_indexer_for_gen_reqs, ( + assert metadata.ext(DSAMetadata).skip_indexer_for_gen_reqs, ( "test setup must hit the skip-indexer path (kv_len <= index_topk)" ) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py index ed3af070b2a4..2a2ea1be655b 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py @@ -28,6 +28,7 @@ from utils.util import skip_pre_blackwell import tensorrt_llm +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionInputType, MLAParams, @@ -815,7 +816,7 @@ def yarn_get_mscale(scale=1, mscale=1): ctx_seq_lens = torch.tensor(context_sequence_lengths, dtype=torch.int) total_ctx_tokens = sum(context_sequence_lengths) - attn_metadata = AttentionCls.Metadata( + attn_metadata = AttentionCls.create_metadata( seq_lens=ctx_seq_lens, request_ids=request_ids, max_num_requests=max_num_contexts, @@ -828,7 +829,7 @@ def yarn_get_mscale(scale=1, mscale=1): num_cached_tokens_per_seq=[0 for _ in context_sequence_lengths], ), mapping=mapping, - sparse_attention_config=sparse_config, + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), ) attn_metadata.prepare() @@ -843,7 +844,7 @@ def yarn_get_mscale(scale=1, mscale=1): [generation_seq_len_q] * max_num_contexts, dtype=torch.int ) total_gen_tokens = max_num_contexts * generation_seq_len_q - attn_metadata = AttentionCls.Metadata( + attn_metadata = AttentionCls.create_metadata( seq_lens=gen_seq_lens, request_ids=request_ids, max_num_requests=max_num_contexts, @@ -859,8 +860,10 @@ def yarn_get_mscale(scale=1, mscale=1): ], ), mapping=mapping, - enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), - sparse_attention_config=sparse_config, + extension_types=(FlashMlaMetadata,) + if torch.cuda.get_device_capability() == (9, 0) + else (), + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), ) attn_metadata.prepare() for layer_idx in range(num_layers): diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_fused_metadata.py b/tests/unittest/_torch/attention/sparse/dsa/test_fused_metadata.py index c859e77ba4d8..aabfd20dba43 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_fused_metadata.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_fused_metadata.py @@ -16,7 +16,7 @@ Differential tests for the fused DSA decode-metadata Triton kernel. The kernel in kernels.py collapses the eager metadata chain in -DSAtrtllmAttentionMetadata.on_update_kv_lens() into one launch. Its whole +DSAMetadata.on_kv_lens_changed() into one launch. Its whole contract is that the integer metadata it produces is BIT-IDENTICAL to that eager chain -- a wrong index silently corrupts sparse attention. These tests assert bit-exactness against the real eager helpers (build_req_idx_per_token + diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py index 623a33a5e536..95bcc3960e16 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -111,6 +111,8 @@ def test_indexer_forward_uses_prior_only_for_temporal_gvr( (num_generations or 1, _MAX_RADIX_BLOCKS_PER_ROW, topk), dtype=torch.float32 ), ) + # The stub carries the DSA extension state itself. + metadata.ext = lambda cls: metadata # Keep distinct storage so the caller test catches use of the host twin. metadata.seq_lens_cuda = metadata.seq_lens.clone() # Radix fallback and self-sampling deliberately have no prior attribute. diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py index df1200650c45..9ad51c3a4974 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_metadata_topk_init.py @@ -10,6 +10,16 @@ from tensorrt_llm._torch.attention.backends.sparse.dsa import metadata as dsa_metadata from tensorrt_llm._torch.attention.backends.sparse.dsa.params import DSAMetadataParams +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + + +def _dsa_without_init(md: TrtllmAttentionMetadata) -> dsa_metadata.DSAMetadata: + """A DSA extension of ``md`` with only the defaults allocate() reads.""" + dsa = object.__new__(dsa_metadata.DSAMetadata) + dsa.metadata = md + dsa.enable_ragged_verification = False + dsa.ragged_verify_lens = None + return dsa @pytest.mark.parametrize("enable_heuristic", [False, True]) @@ -22,8 +32,8 @@ def test_topk_flags_initialized_before_buffer_allocation( dsl_available: bool, sm_version: int, ) -> None: - metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) - metadata.sparse_metadata_params = DSAMetadataParams( + md = object.__new__(TrtllmAttentionMetadata) + md.sparse_metadata_params = DSAMetadataParams( indexer_max_chunk_size=8192, max_sparse_topk=512, index_head_dim=128, @@ -34,12 +44,13 @@ def test_topk_flags_initialized_before_buffer_allocation( q_split_threshold=8192, use_self_sampling_topk=use_self_sampling, ) - metadata.kv_cache_manager = SimpleNamespace( + md.kv_cache_manager = SimpleNamespace( tokens_per_block=128, compressed_block_sizes={}, get_cache_indices=Mock(), ) - metadata.is_cuda_graph = False + md.is_cuda_graph = False + metadata = _dsa_without_init(md) enabled = enable_heuristic and sm_version >= 100 self_sampling_supported = enabled and dsl_available and sm_version in (100, 103) temporal_supported = enabled and dsl_available and sm_version in (100, 103) @@ -54,11 +65,10 @@ def check_flags(*, capture_graph: bool) -> None: metadata.create_buffers_for_indexer = Mock(side_effect=check_flags) with ( - patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", dsl_available), patch.object(dsa_metadata, "get_sm_version", return_value=sm_version), ): - metadata.__post_init__() + metadata.allocate(None) metadata.create_buffers_for_mla_rope_append.assert_called_once_with(capture_graph=False) metadata.create_buffers_for_indexer.assert_called_once_with(capture_graph=False) @@ -68,8 +78,8 @@ def check_flags(*, capture_graph: bool) -> None: @pytest.mark.parametrize("num_local_layers", [1, 3]) def test_temporal_gvr_allocates_real_prior_buffers(num_local_layers: int) -> None: """Allocate actual CUDA buffers and size the zeroed prior by local layers.""" - metadata = object.__new__(dsa_metadata.DSAtrtllmAttentionMetadata) - metadata.sparse_metadata_params = DSAMetadataParams( + md = object.__new__(TrtllmAttentionMetadata) + md.sparse_metadata_params = DSAMetadataParams( indexer_max_chunk_size=32, max_sparse_topk=512, index_head_dim=128, @@ -80,30 +90,30 @@ def test_temporal_gvr_allocates_real_prior_buffers(num_local_layers: int) -> Non q_split_threshold=32, use_self_sampling_topk=False, ) - metadata.kv_cache_manager = SimpleNamespace( + md.kv_cache_manager = SimpleNamespace( tokens_per_block=128, compressed_block_sizes={}, get_cache_indices=Mock(), max_blocks_per_seq=2, num_local_layers=num_local_layers, ) - metadata.draft_kv_cache_manager = None - metadata.is_cuda_graph = False - metadata.cuda_graph_buffers = None - metadata.max_num_sequences = 4 - metadata.max_num_tokens = 16 + md.draft_kv_cache_manager = None + md.is_cuda_graph = False + md.cuda_graph_buffers = None + md.max_num_sequences = 4 + md.max_num_tokens = 16 + md.enable_context_mla_with_cached_kv = False + metadata = _dsa_without_init(md) metadata.max_draft_tokens = 3 metadata.num_sms = 16 - metadata.enable_context_mla_with_cached_kv = False # object.__new__ skips __init__, which is where this default is set. metadata._radix_rows_per_sequence = 1 with ( - patch.object(dsa_metadata.TrtllmAttentionMetadata, "__post_init__"), patch.object(metadata, "create_buffers_for_mla_rope_append"), patch.object(dsa_metadata, "IS_CUTLASS_DSL_AVAILABLE", True), patch.object(dsa_metadata, "get_sm_version", return_value=100), ): - metadata.__post_init__() + metadata.allocate(None) assert metadata.enable_gvr_topk assert metadata.needs_gvr_prior diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py index 0a5b99a314d0..c994b5e12e47 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Tests for the token->request map used by DSAtrtllmAttentionMetadata. +Tests for the token->request map used by the DSA metadata extension. 1. build_req_idx_per_token must match the host repeat_interleave build for every layout, so the device and host builders cannot drift. @@ -28,10 +28,8 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.dsa import ( - DSAtrtllmAttentionMetadata, - build_req_idx_per_token, -) +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAMetadata, build_req_idx_per_token +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata def _host_reference(seq_lens: torch.Tensor) -> torch.Tensor: @@ -65,18 +63,27 @@ def test_matches_host_repeat_interleave(seq_lens, device): assert result.to(torch.int32).tolist() == _host_reference(seq_lens).tolist() +def _attach_dsa(md: TrtllmAttentionMetadata) -> DSAMetadata: + """A DSA extension built without __init__, bound to ``md``.""" + dsa = object.__new__(DSAMetadata) + dsa.metadata = md + dsa.enable_ragged_verification = False + md.replace_ext(dsa) + return dsa + + def test_on_update_kv_lens_invalidates_base_mla_state() -> None: - md = object.__new__(DSAtrtllmAttentionMetadata) - md.enable_flash_mla = False + md = object.__new__(TrtllmAttentionMetadata) md._mla_scheduler_buffers_valid = True md._mla_ctx_cu_seqlens_valid = True - md._invalidate_pool_view_cache = Mock() md._num_tokens = 0 md._num_generations = 0 md.kv_cache_manager = None md.kv_lens_cuda = torch.empty(0, dtype=torch.int32) - md._compute_kv_lens_row_reorder = Mock() - md.prepare_dense_topk_indices = Mock() + dsa = _attach_dsa(md) + dsa._invalidate_pool_view_cache = Mock() + dsa._compute_kv_lens_row_reorder = Mock() + dsa.prepare_dense_topk_indices = Mock() md.on_update_kv_lens() @@ -92,38 +99,39 @@ def test_on_update_kv_lens_rebuilds_stale_map() -> None: max_draft_len = 3 num_requests = 3 - md = object.__new__(DSAtrtllmAttentionMetadata) + md = object.__new__(TrtllmAttentionMetadata) md.kv_cache_manager = None md._num_generations = 0 + dsa = _attach_dsa(md) # __init__ (bypassed by object.__new__) defaults this to False; # on_update_kv_lens() reads it since #16925. - md.in_mtp_draft_loop = False + dsa.in_mtp_draft_loop = False # Same reason: `_invalidate_pool_view_cache` clears this unconditionally, and # unlike the sibling test above we exercise the real method rather than a Mock # (the map rebuild happens downstream of it). A dict, not a Mock, so that a # future attribute added to that method fails loudly here instead of being # absorbed. - md._group_remap_batched = {} + dsa._group_remap_batched = {} # Stub collaborators unrelated to the map rebuild (test_dsa_indexer.py style). md.kv_lens_cuda = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) - md._compute_kv_lens_row_reorder = Mock() - md.prepare_dense_topk_indices = Mock() + dsa._compute_kv_lens_row_reorder = Mock() + dsa.prepare_dense_topk_indices = Mock() # prepare()-time state: target forward, 1 + max_draft_len tokens/request. target_seq_lens = torch.full((num_requests,), 1 + max_draft_len, dtype=torch.int32) md._seq_lens = target_seq_lens md._seq_lens_cuda = target_seq_lens.to(device) md._num_tokens = int(target_seq_lens.sum()) - md.req_idx_per_token = torch.empty(md._num_tokens, dtype=torch.int32, device=device) - md.req_idx_per_token[:] = _host_reference(md._seq_lens_cuda) + dsa.req_idx_per_token = torch.empty(md._num_tokens, dtype=torch.int32, device=device) + dsa.req_idx_per_token[:] = _host_reference(md._seq_lens_cuda) # Draft loop: one token per request; the stale prefix reads [0, 0, 0]. draft_seq_lens = torch.ones(num_requests, dtype=torch.int32) md._seq_lens = draft_seq_lens md._seq_lens_cuda = draft_seq_lens.to(device) md._num_tokens = num_requests - assert md.req_idx_per_token[:num_requests].tolist() == [0, 0, 0] + assert dsa.req_idx_per_token[:num_requests].tolist() == [0, 0, 0] md.on_update_kv_lens() - assert md.req_idx_per_token[:num_requests].tolist() == [0, 1, 2] + assert dsa.req_idx_per_token[:num_requests].tolist() == [0, 1, 2] diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py b/tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py index 6c545f3f8067..3e321cdcb61d 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_short_seq_mha.py @@ -366,7 +366,7 @@ def _make_metadata( kwargs["enable_context_mla_with_cached_kv"] = True if runtime_features is not None: kwargs["runtime_features"] = runtime_features - metadata = attn_cls.Metadata( + metadata = attn_cls.create_metadata( seq_lens=torch.tensor(seq_lens, dtype=torch.int), request_ids=list(range(num_ctx)), max_num_requests=num_ctx, diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index b859d7514e5a..6cb12699374e 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -395,23 +395,24 @@ def test_cache_state_uses_prepared_metadata_without_fallback(is_cuda_graph): num_contexts=0, is_cuda_graph=is_cuda_graph, kv_lens_cuda=live_lengths, - mamba_metadata=prepared, ) + recurrent = SimpleNamespace(mamba_metadata=prepared) + metadata.ext = lambda cls: recurrent state = backend._cache_state(metadata) assert state.block_tables.data_ptr() == tables.data_ptr() assert state.kv_lens.data_ptr() == live_lengths.data_ptr() torch.testing.assert_close(state.kv_lens, live_lengths) - metadata.mamba_metadata.glm_block_tables = None + recurrent.mamba_metadata.glm_block_tables = None with patch("torch.zeros", side_effect=AssertionError("must not allocate fallback tables")): with pytest.raises(RuntimeError, match="requires prepared glm_block_tables"): backend._cache_state(metadata) - metadata.mamba_metadata.glm_block_tables = tables + recurrent.mamba_metadata.glm_block_tables = tables metadata.kv_lens_cuda = None with pytest.raises(ValueError, match="kv_lens_cuda"): backend._cache_state(metadata) metadata.kv_lens_cuda = live_lengths for invalid in (None, False, SimpleNamespace(glm_block_tables=tables)): - metadata.mamba_metadata = invalid + recurrent.mamba_metadata = invalid with pytest.raises(AssertionError, match="Glm5NextMamba2Metadata"): backend._cache_state(metadata) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py index 07376d86d2d7..3e50613339f0 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py @@ -18,6 +18,7 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import run_msa_nvfp4_sparse_gqa from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, @@ -36,8 +37,9 @@ resolve_num_topk_chunks, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( - MiniMaxM3MsaSparseAttentionMetadata, + MiniMaxM3MsaMetadata, ) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._utils import get_sm_version PAGE_SIZE = SPARSE_BLOCK_SIZE @@ -51,6 +53,14 @@ ) +def _msa_metadata() -> TrtllmAttentionMetadata: + """An uninitialized TRTLLM metadata carrying a bare MSA extension.""" + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(MiniMaxM3MsaMetadata(metadata)) + metadata.replace_ext(SpecDecodingMetadata(metadata)) + return metadata + + def _flat_page_table(block_table: torch.Tensor, kv_lens_cpu: torch.Tensor) -> torch.Tensor: """Flatten a block table into the per-request page ids fmha_sm100 consumes. @@ -972,8 +982,9 @@ def test_nvfp4_sparse_decode_matches_msa_csr_kernel(num_kv_heads, group): kv_lens = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) cu_kv = torch.zeros(batch + 1, device="cuda", dtype=torch.int32) torch.cumsum(kv_lens, 0, out=cu_kv[1:]) - metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) - metadata.__dict__.update( + metadata = _msa_metadata() + metadata.__dict__.update(_num_contexts=2, _num_generations=batch - 2) + metadata.ext(MiniMaxM3MsaMetadata).__dict__.update( _msa_live_batch=batch, msa_cu_q_lens=torch.arange(batch + 1, device="cuda", dtype=torch.int32), msa_cu_kv_lens=cu_kv, @@ -983,14 +994,12 @@ def test_nvfp4_sparse_decode_matches_msa_csr_kernel(num_kv_heads, group): _msa_total_k_rows=sum((s + PAGE_SIZE - 1) // PAGE_SIZE for s in seq_lens), msa_block_table=case.block_table, msa_seq_lens_cuda=case.seq_lens, - _num_contexts=2, _msa_context_prefix_bounds=( 1, max(seq_lens[:2]), sum(seq_lens[:2]), sum((s + PAGE_SIZE - 1) // PAGE_SIZE for s in seq_lens[:2]), ), - _num_generations=batch - 2, ) triton_out = _run_nvfp4(case) @@ -1074,16 +1083,17 @@ def test_msa_nvfp4_calibrated_staging_and_scale_replay( ) k_global = torch.tensor([2.0 / 2688.0], device="cuda", dtype=torch.float32) v_global = torch.tensor([1.0 / 2688.0], device="cuda", dtype=torch.float32) - metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) - metadata._msa_live_batch = 1 - metadata.msa_cu_q_lens = torch.tensor([0, 1], device="cuda", dtype=torch.int32) - metadata.msa_cu_kv_lens = torch.tensor([0, kv_len], device="cuda", dtype=torch.int32) - metadata._msa_max_q_len = 1 - metadata._msa_max_kv_len_all = kv_len - metadata._msa_total_k = kv_len - metadata._msa_total_k_rows = 1 - metadata.msa_block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) - metadata.msa_seq_lens_cuda = torch.tensor([kv_len], device="cuda", dtype=torch.int32) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) + msa._msa_live_batch = 1 + msa.msa_cu_q_lens = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + msa.msa_cu_kv_lens = torch.tensor([0, kv_len], device="cuda", dtype=torch.int32) + msa._msa_max_q_len = 1 + msa._msa_max_kv_len_all = kv_len + msa._msa_total_k = kv_len + msa._msa_total_k_rows = 1 + msa.msa_block_table = torch.zeros((1, 1), device="cuda", dtype=torch.int32) + msa.msa_seq_lens_cuda = torch.tensor([kv_len], device="cuda", dtype=torch.int32) topk = torch.full((1, 1, MSA_REQUIRED_TOPK), -1, device="cuda", dtype=torch.int32) topk[..., 0] = 0 out = torch.empty(q.shape, device="cuda", dtype=torch.bfloat16) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 6f611923aa2c..d0e7a9374b0f 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -17,7 +17,9 @@ import pytest import torch +from fmha_test_utils import metadata_stub +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( MiniMaxM3KVCacheManagerV2, MiniMaxM3MsaSparseAttention, @@ -37,15 +39,24 @@ minimax_m3_sparse_attn_decode, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( - MiniMaxM3MsaSparseAttentionMetadata, + MiniMaxM3MsaMetadata, MsaDecodeSpan, ) from tensorrt_llm._torch.attention.backends.sparse.registry import _resolve_minimax_m3_backend_cls +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig +def _msa_metadata() -> TrtllmAttentionMetadata: + """An uninitialized TRTLLM metadata carrying a bare MSA extension.""" + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(MiniMaxM3MsaMetadata(metadata)) + metadata.replace_ext(SpecDecodingMetadata(metadata)) + return metadata + + @pytest.mark.cpu_only @pytest.mark.parametrize("empty_state", ["empty_batch", "no_requests", "no_lengths"]) def test_msa_metadata_clears_padded_cache_slot_tail( @@ -55,46 +66,44 @@ def test_msa_metadata_clears_padded_cache_slot_tail( from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend monkeypatch.setattr(msa_backend, "maybe_pin_memory", lambda tensor: tensor) - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) - metadata._msa_buffers_ready = True + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) + msa._msa_buffers_ready = True metadata.request_ids = [0] metadata.kv_cache_manager = SimpleNamespace( tokens_per_block=4, get_buffers=lambda layer_idx: torch.empty(1), get_block_ids_per_seq=lambda request_ids: torch.tensor([[3]], dtype=torch.int32), ) - metadata.msa_out_cache_loc = torch.full((4,), 99, dtype=torch.int32) - metadata.msa_block_table = torch.zeros((1, 1), dtype=torch.int32) - metadata.msa_seq_lens_cuda = torch.zeros(1, dtype=torch.int32) - metadata.msa_subpage_block_table = None - metadata._msa_runs_no_fmha = lambda: True - metadata._msa_kv_lens_may_change = lambda: False - original_ptr = metadata.msa_out_cache_loc.data_ptr() + msa.msa_out_cache_loc = torch.full((4,), 99, dtype=torch.int32) + msa.msa_block_table = torch.zeros((1, 1), dtype=torch.int32) + msa.msa_seq_lens_cuda = torch.zeros(1, dtype=torch.int32) + msa.msa_subpage_block_table = None + msa._msa_runs_no_fmha = lambda: True + msa._msa_kv_lens_may_change = lambda: False + original_ptr = msa.msa_out_cache_loc.data_ptr() for count in (4, 2, 0, 1): - metadata._msa_qo_lens_cpu = torch.tensor([count], dtype=torch.int32) - metadata._msa_kv_lens_cpu = metadata._msa_qo_lens_cpu.clone() - metadata._msa_qo_offset_cpu = torch.zeros(1, dtype=torch.int32) + msa._msa_qo_lens_cpu = torch.tensor([count], dtype=torch.int32) + msa._msa_kv_lens_cpu = msa._msa_qo_lens_cpu.clone() + msa._msa_qo_offset_cpu = torch.zeros(1, dtype=torch.int32) # What the base seq_lens setter would derive; the cache writers read the # live count off it. metadata._num_tokens = count - metadata._build_msa_fields() - assert metadata.msa_out_cache_loc.tolist() == list(range(12, 12 + count)) + [-1] * ( - 4 - count - ) - assert metadata.msa_out_cache_loc.data_ptr() == original_ptr - assert metadata._msa_fields_ready + msa._build_msa_fields() + assert msa.msa_out_cache_loc.tolist() == list(range(12, 12 + count)) + [-1] * (4 - count) + assert msa.msa_out_cache_loc.data_ptr() == original_ptr + assert msa._msa_fields_ready if empty_state == "empty_batch": metadata.request_ids = [] - metadata._msa_qo_lens_cpu = torch.empty(0, dtype=torch.int32) + msa._msa_qo_lens_cpu = torch.empty(0, dtype=torch.int32) elif empty_state == "no_requests": metadata.request_ids = None else: - metadata._msa_qo_lens_cpu = None - metadata._build_msa_fields() - assert metadata.msa_out_cache_loc.tolist() == [-1] * 4 + msa._msa_qo_lens_cpu = None + msa._build_msa_fields() + assert msa.msa_out_cache_loc.tolist() == [-1] * 4 def test_msa_package_availability_installs_cutlass_compatibility_aliases(monkeypatch): @@ -347,7 +356,7 @@ def resolve_index_view( ) for _ in range(2) ] - metadata = [object.__new__(MiniMaxM3MsaSparseAttentionMetadata) for _ in managers] + metadata = [_msa_metadata() for _ in managers] for manager, meta in zip(managers, metadata): meta.kv_cache_manager = manager for layer in (0, 5, 60): @@ -358,11 +367,11 @@ def resolve_index_view( assert cache.data_ptr() == manager._test_index_pool.data_ptr() assert cache.stride(0) == 2 * 4 * 8 if implementation == "msa": - assert meta.msa_idx_k_cache(3) is cache + assert meta.ext(MiniMaxM3MsaMetadata).msa_idx_k_cache(3) is cache - def forward(meta: MiniMaxM3MsaSparseAttentionMetadata, value: torch.Tensor) -> torch.Tensor: + def forward(meta: TrtllmAttentionMetadata, value: torch.Tensor) -> torch.Tensor: get_cache = ( - meta.msa_idx_k_cache + meta.ext(MiniMaxM3MsaMetadata).msa_idx_k_cache if implementation == "msa" else meta.kv_cache_manager.get_index_k_buffer ) @@ -399,18 +408,18 @@ def test_cache_manager_rejects_non_positive_sparse_index_dim(sparse_index_dim: i def test_msa_metadata_rejects_undersized_max_score_buffer(): - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) # Flat backing store sized for 4 heads * 8 k-tiles * 2 batch = 64 elements, # too small for the plan's required 4 * 16 * 2 = 128. One query token per # decode row, so the token bound equals the batch here. - metadata.msa_max_score = torch.zeros(4 * 8 * 2) + msa.msa_max_score = torch.zeros(4 * 8 * 2) metadata.kv_cache_manager = None metadata.max_num_sequences = 2 metadata.max_num_tokens = 2 with pytest.raises(ValueError, match=r"msa_max_score backing store"): - metadata._ensure_msa_decode_scratch_buffers( + msa._ensure_msa_decode_scratch_buffers( num_index_heads=4, max_batch=2, capture_graph=False, @@ -435,8 +444,8 @@ def get_buffer(self, tensor_shape, dtype, cache_name, capture_graph): def _buffer_metadata(**manager_fields): """Metadata ready for _create_msa_buffers, under capture.""" - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) metadata.kv_cache_manager = SimpleNamespace( max_blocks_per_seq=MAX_BLOCKS_PER_SEQ, tokens_per_block=128, @@ -449,7 +458,7 @@ def _buffer_metadata(**manager_fields): metadata.max_num_tokens = 512 # No sparse params, so the fmha_sm100 proxy scratch is skipped and this # exercises only the layer-invariant buffers. - metadata._msa_params = None + msa._msa_params = None return metadata @@ -459,12 +468,13 @@ def test_msa_buffers_include_graph_stable_block_table(): come from the graph buffer pool at the manager's worst-case geometry, so their addresses survive capture.""" metadata = _buffer_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) - metadata._create_msa_buffers() + msa._create_msa_buffers() - assert metadata._msa_buffers_ready - assert metadata.msa_block_table.shape == (MAX_NUM_SEQUENCES, MAX_BLOCKS_PER_SEQ) - assert metadata.msa_seq_lens_cuda.shape == (MAX_NUM_SEQUENCES,) + assert msa._msa_buffers_ready + assert msa.msa_block_table.shape == (MAX_NUM_SEQUENCES, MAX_BLOCKS_PER_SEQ) + assert msa.msa_seq_lens_cuda.shape == (MAX_NUM_SEQUENCES,) requested = metadata.cuda_graph_buffers.requested assert requested["msa_block_table"] == ( (MAX_NUM_SEQUENCES, MAX_BLOCKS_PER_SEQ), @@ -481,6 +491,7 @@ def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch, nv from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend metadata = _buffer_metadata(sparse_layer_ids=[3, 4], layer_offsets={3: 0}) + msa = metadata.ext(MiniMaxM3MsaMetadata) main_cache = torch.zeros(2, 2, 1, 128, 128) index_cache = torch.zeros(2, 1, 128, 128) scale_cache = torch.zeros(2, 2, 1, 128, 8, dtype=torch.uint8) @@ -499,11 +510,11 @@ def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch, nv ) # No native pool in this CPU test; only zero-copy cache-view staging is under test. monkeypatch.setattr(msa_backend, "uniform_subpages_per_slot", lambda manager: 0) - metadata._create_msa_buffers() + msa._create_msa_buffers() manager.get_buffers.assert_called_once_with(3, kv_layout="HND") manager.get_index_k_buffer.assert_called_once_with(3) - assert set(metadata.msa_layer_cache_tensors) == {3} - main, index, scales = metadata.msa_layer_cache_tensors[3] + assert set(msa.msa_layer_cache_tensors) == {3} + main, index, scales = msa.msa_layer_cache_tensors[3] assert main is main_cache and index is index_cache if nvfp4: manager.get_block_scale_buffers.assert_called_once_with(3, kv_layout="HND") @@ -533,12 +544,13 @@ def test_msa_buffers_stage_the_subpage_table_only_for_a_uniform_pool(factors, ex layer_offsets=dict.fromkeys(range(len(factors)), 0), get_kv_subpage_pool=lambda layer_idx, kv_layout="HND": (None, factors[layer_idx]), ) + msa = metadata.ext(MiniMaxM3MsaMetadata) - metadata._create_msa_buffers() + msa._create_msa_buffers() - assert metadata._msa_subpages_per_slot == expected + assert msa._msa_subpages_per_slot == expected if expected == 0: - assert metadata.msa_subpage_block_table is None + assert msa.msa_subpage_block_table is None assert "msa_subpage_block_table" not in metadata.cuda_graph_buffers.requested else: # One K row and one V row per slot, at the same worst-case geometry as @@ -554,19 +566,19 @@ def test_msa_subpage_rows_slice_the_generation_span(): """A mixed step hands the dense kernel only the span's rows, so its sub-page table has to be sliced the same way the slot table is. The factor travels with it so the kernel can tell a stale staging from its own geometry.""" - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) - metadata.msa_subpage_block_table = torch.arange(4 * 2 * 3, dtype=torch.int32).reshape(4, 2, 3) - metadata._msa_subpages_per_slot = 9 + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) + msa.msa_subpage_block_table = torch.arange(4 * 2 * 3, dtype=torch.int32).reshape(4, 2, 3) + msa._msa_subpages_per_slot = 9 - table, factor = metadata.msa_subpage_rows(2, 4) + table, factor = msa.msa_subpage_rows(2, 4) assert factor == 9 - assert torch.equal(table, metadata.msa_subpage_block_table[2:4]) + assert torch.equal(table, msa.msa_subpage_block_table[2:4]) # Nothing staged: the caller expands its own layer's table, and the 0 # factor is what tells it to. - metadata.msa_subpage_block_table = None - assert metadata.msa_subpage_rows(2, 4) == (None, 0) + msa.msa_subpage_block_table = None + assert msa.msa_subpage_rows(2, 4) == (None, 0) def test_msa_proxy_max_score_view_is_contiguous_over_stable_store(): @@ -574,28 +586,28 @@ def test_msa_proxy_max_score_view_is_contiguous_over_stable_store(): [num_index_heads, plan_max_k_tiles, num_tokens] shape the kernel writes, backed by a stable store so its data_ptr survives CUDA graph replay. """ - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) # Worst-case store: 4 heads * 16 k-tiles * 8 batch. num_index_heads, worst_k, max_batch = 4, 16, 8 - metadata.msa_max_score = torch.zeros(num_index_heads * worst_k * max_batch) - store_ptr = metadata.msa_max_score.data_ptr() + msa.msa_max_score = torch.zeros(num_index_heads * worst_k * max_batch) + store_ptr = msa.msa_max_score.data_ptr() # A smaller live step still yields a contiguous view sized to that step, # which is what the kernel's stride-agnostic write requires. - view = metadata.msa_proxy_max_score_view(num_index_heads, 5, 3) + view = msa.msa_proxy_max_score_view(num_index_heads, 5, 3) assert view.shape == (num_index_heads, 5, 3) assert view.is_contiguous() assert view.data_ptr() == store_ptr # Oversized requests are rejected rather than silently corrupting memory. with pytest.raises(ValueError, match=r"msa_max_score backing store"): - metadata.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1) + msa.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1) # So is an empty one. Both writers address the view by block id, so a zero # extent is not a small view but writes past the end of one. with pytest.raises(ValueError, match=r"no block extent"): - metadata.msa_proxy_max_score_view(num_index_heads, 0, max_batch) + msa.msa_proxy_max_score_view(num_index_heads, 0, max_batch) @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -703,8 +715,8 @@ def test_per_token_valid_blocks_multi_token_decode(): def test_msa_index_k_uses_hnd_cache_view_and_writer(): - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) num_pages, coalescing_scale, page_size, head_dim = 2, 7, 8, 16 pool = torch.zeros( num_pages, @@ -726,13 +738,13 @@ def get_index_k_buffer(self, layer_idx): manager = FakeCacheManager() metadata.kv_cache_manager = manager - metadata.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32) + msa.msa_out_cache_loc = torch.tensor([2, page_size + 5], dtype=torch.int32) values = torch.arange(2 * head_dim, dtype=torch.float32).reshape(2, 1, head_dim) - metadata._msa_fields_ready = True + msa._msa_fields_ready = True metadata._num_tokens = 2 - returned = metadata.msa_idx_k_cache(3) - metadata.msa_write_idx_k(3, values) + returned = msa.msa_idx_k_cache(3) + msa.msa_write_idx_k(3, values) assert returned.data_ptr() == hnd_cache.data_ptr() assert not returned.is_contiguous() @@ -844,6 +856,10 @@ def __init__(self, dtype: torch.dtype) -> None: self.cache = backing[::7] self.msa_out_cache_loc = torch.tensor([0, 128], dtype=torch.int32, device="cuda") + def ext(self, cls: type) -> "FakeMetadata": + # The fake serves as its own MSA extension. + return self + def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( write_kv_slots, @@ -977,6 +993,10 @@ def __init__(self) -> None: dtype=torch.bfloat16, ) + def ext(self, cls: type) -> "FakeMetadata": + # The fake serves as its own MSA extension. + return self + def msa_write_idx_k(self, layer_idx: int, idx_k: torch.Tensor) -> None: del layer_idx self.idx_k_cache = idx_k @@ -1197,25 +1217,25 @@ def test_lazily_allocated_scratch_publishes_the_bound_it_used(monkeypatch): monkeypatch.setattr(msa_backend, "require_msa_module", lambda: None) monkeypatch.setattr(msa_backend, "_worst_case_proxy_max_k_tiles", lambda *a, **kw: 32) - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) metadata.kv_cache_manager = SimpleNamespace() metadata.cuda_graph_buffers = None metadata.max_num_sequences = 2 metadata.max_num_tokens = 8 - metadata.msa_max_score = None - metadata._msa_worst_case_max_k_tiles = 0 + msa.msa_max_score = None + msa._msa_worst_case_max_k_tiles = 0 - metadata._ensure_msa_decode_scratch_buffers( + msa._ensure_msa_decode_scratch_buffers( num_index_heads=4, max_batch=2, capture_graph=False, required_max_k_tiles=16, ) - assert metadata.msa_worst_case_max_k_tiles == 32 + assert msa.msa_worst_case_max_k_tiles == 32 # The store was sized against that bound, so a view shaped by it fits. - assert metadata.msa_proxy_max_score_view(4, 32, 2).shape == (4, 32, 2) + assert msa.msa_proxy_max_score_view(4, 32, 2).shape == (4, 32, 2) def _span_metadata(*, num_contexts=0, qo_lens=(1, 1), kv_lens=(9, 11), is_cuda_graph=False): @@ -1225,8 +1245,8 @@ def _span_metadata(*, num_contexts=0, qo_lens=(1, 1), kv_lens=(9, 11), is_cuda_g enough to drive the real msa_*_cpu length properties; no cache pool is needed. """ - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) metadata.mapping = None # Assigned behind the seq_lens property, whose setter would stage a device # copy nothing here reads; num_seqs derives from it. The num_contexts @@ -1236,7 +1256,7 @@ def _span_metadata(*, num_contexts=0, qo_lens=(1, 1), kv_lens=(9, 11), is_cuda_g metadata.kv_lens = torch.tensor(kv_lens, dtype=torch.int32) metadata.is_cuda_graph = is_cuda_graph # prepare() stages these before the span, and the span reads them. - metadata._stage_host_lengths() + msa._stage_host_lengths() return metadata @@ -1244,13 +1264,14 @@ def test_the_decode_span_of_a_pure_decode_step_is_the_whole_batch(): """A pure-decode step leaves fmha_sm100 nothing, which is what lets prepare() skip its plans and its page table entirely.""" metadata = _span_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) - metadata._set_decode_span() + msa._set_decode_span() - assert metadata.msa_decode_span == (0, 1) - assert metadata.msa_decode_query_len == 1 - assert metadata.msa_max_kv_len == 11 - assert metadata._msa_runs_no_fmha() is True + assert msa.msa_decode_span == (0, 1) + assert msa.msa_decode_query_len == 1 + assert msa.msa_max_kv_len == 11 + assert msa._msa_runs_no_fmha() is True def test_the_decode_span_of_a_mixed_step_is_its_generation_suffix(): @@ -1262,15 +1283,16 @@ def test_the_decode_span_of_a_mixed_step_is_its_generation_suffix(): # Two context requests (7 and 5 query tokens, the first a chunk of a long # prompt) ahead of two decode rows. metadata = _span_metadata(num_contexts=2, qo_lens=(7, 5, 1, 1), kv_lens=(4096, 5, 40, 33)) + msa = metadata.ext(MiniMaxM3MsaMetadata) - metadata._set_decode_span() + msa._set_decode_span() - assert metadata.msa_decode_span == (2, 1) + assert msa.msa_decode_span == (2, 1) # fmha_sm100 still runs the context prefix, so its page table stays live. - assert metadata._msa_runs_no_fmha() is False + assert msa._msa_runs_no_fmha() is False # The trtllm-gen scheduling bound must come from the span's own rows: the # 4096-token context row here would inflate a whole-batch maximum by 100x. - assert metadata.msa_max_kv_len == 40 + assert msa.msa_max_kv_len == 40 def test_the_triton_sparse_decode_rejects_a_q_that_outruns_the_batch(): @@ -1368,12 +1390,13 @@ def test_a_pure_prefill_step_has_no_decode_span(): """A step with no generation row has nothing for the decode kernels, and fmha_sm100 keeps every plan and the page table they read.""" metadata = _span_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) + msa = metadata.ext(MiniMaxM3MsaMetadata) - metadata._set_decode_span() + msa._set_decode_span() - assert metadata.msa_decode_span is None - assert metadata.msa_decode_query_len is None - assert metadata._msa_runs_no_fmha() is False + assert msa.msa_decode_span is None + assert msa.msa_decode_query_len is None + assert msa._msa_runs_no_fmha() is False def test_plan_rows_narrow_to_the_rows_fmha_sm100_still_runs(): @@ -1387,22 +1410,25 @@ def test_plan_rows_narrow_to_the_rows_fmha_sm100_still_runs(): which is the same range plus a pure-prefill step's generation-free batch. """ mixed = _span_metadata(num_contexts=2, qo_lens=(7, 5, 1, 1), kv_lens=(4096, 5, 40, 33)) - mixed._set_decode_span() + msa_mixed = mixed.ext(MiniMaxM3MsaMetadata) + msa_mixed._set_decode_span() # The span took the generation suffix, so what is left is the prefix. - assert mixed._msa_proxy_plan_rows() == (0, 2) - assert mixed._msa_attn_plan_rows() == (0, 2) + assert msa_mixed._msa_proxy_plan_rows() == (0, 2) + assert msa_mixed._msa_attn_plan_rows() == (0, 2) decode = _span_metadata() - decode._set_decode_span() + msa_decode = decode.ext(MiniMaxM3MsaMetadata) + msa_decode._set_decode_span() # Nothing is left to plan on a pure-decode step the kernels fully own. - assert decode._msa_proxy_plan_rows() is None - assert decode._msa_attn_plan_rows() is None + assert msa_decode._msa_proxy_plan_rows() is None + assert msa_decode._msa_attn_plan_rows() is None prefill = _span_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) - prefill._set_decode_span() + msa_prefill = prefill.ext(MiniMaxM3MsaMetadata) + msa_prefill._set_decode_span() # No generation row, so the proxy scores every row. - assert prefill._msa_proxy_plan_rows() == (0, 2) - assert prefill._msa_attn_plan_rows() == (0, 2) + assert msa_prefill._msa_proxy_plan_rows() == (0, 2) + assert msa_prefill._msa_attn_plan_rows() == (0, 2) @pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) @@ -1421,24 +1447,25 @@ def fmha_sm100_plan(qo_lens, kv_lens, **kwargs): ) monkeypatch.setattr(msa_backend, "_cache_device", lambda metadata: torch.device("cpu")) metadata = _span_metadata(num_contexts=2, qo_lens=(5, 7), kv_lens=(5, 7)) - metadata._set_decode_span() - metadata._msa_fields_ready = True - metadata._msa_params = SimpleNamespace( + msa = metadata.ext(MiniMaxM3MsaMetadata) + msa._set_decode_span() + msa._msa_fields_ready = True + msa._msa_params = SimpleNamespace( sharded_index_head_count=lambda mapping: 4, sharded_head_counts=lambda mapping: (8, 2), topk=16, ) metadata.kv_cache_manager = SimpleNamespace(tokens_per_block=128, dtype=dtype) - metadata._build_step_plans() + msa._build_step_plans() - assert metadata.msa_prefill_proxy_plan["output_maxscore"] - assert "kv_block_num" not in metadata.msa_prefill_dense_plan + assert msa.msa_prefill_proxy_plan["output_maxscore"] + assert "kv_block_num" not in msa.msa_prefill_dense_plan if dtype == DataType.NVFP4: - assert metadata.msa_prefill_gqa_plan is None + assert msa.msa_prefill_gqa_plan is None assert len(planned) == 2 else: - assert metadata.msa_prefill_gqa_plan["kv_block_num"] == 16 + assert msa.msa_prefill_gqa_plan["kv_block_num"] == 16 assert len(planned) == 3 @@ -1522,6 +1549,8 @@ def test_decode_fmha_runs_the_phase_its_span_describes(): msa_block_table=torch.zeros(4, 4, dtype=torch.int32), msa_seq_lens_cuda=torch.zeros(4, dtype=torch.int32), ) + # The stub serves as its own MSA extension. + metadata.ext = lambda cls: metadata params = _generation_params(attention, metadata, seq_offset=2, input_seq_length=1) dispatched = [] decode._run_dense = lambda params, block_table, seq_lens: dispatched.append( @@ -1553,6 +1582,8 @@ def test_each_phase_writes_the_cache_slots_of_the_rows_it_attends(): kv_cache_manager=SimpleNamespace(get_buffers=lambda layer_idx, kv_layout=None: buffers), msa_out_cache_loc=torch.tensor([1, 2, page_size + 3], dtype=torch.int32), ) + # The stub serves as its own MSA extension. + metadata.ext = lambda cls: metadata attention = SimpleNamespace(layer_idx=0) k = torch.arange(1, 3 * head_dim + 1, dtype=torch.float32).reshape(3, head_dim) v = -k @@ -1605,7 +1636,7 @@ def test_the_pair_takes_one_phase_each_and_a_mixed_step_together( manager.fmha_libs = libraries q = torch.empty((num_contexts + num_generations, 4)) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=num_contexts, num_generations=num_generations, num_ctx_tokens=num_contexts, @@ -1742,6 +1773,8 @@ def test_mixed_batch_generation_span_matches_the_whole_batch_msa_path(): ) total_q = int(q.shape[0]) metadata = SimpleNamespace(**fields) + # The namespace serves as its own MSA extension. + metadata.ext = lambda cls: metadata num_heads = attention.num_heads sm_scale = head_dim**-0.5 @@ -1804,18 +1837,18 @@ def test_msa_scratch_sizing_covers_spec_verify_tokens(): proxy scratch must be sized by tokens, not by batch. The draft length comes from the KV cache manager, which knows the speculative config at build time. """ - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) # 2 sequences, 4 tokens each (draft_len=3): 8 decode tokens per step. metadata.kv_cache_manager = SimpleNamespace(max_total_draft_tokens=3) metadata.max_num_sequences = 2 metadata.max_num_tokens = 8 # Store sized for batch-only sizing (4 heads * 16 k-tiles * 2), which is # too small once tokens are accounted for (4 * 16 * 8). - metadata.msa_max_score = torch.zeros(4 * 16 * 2) + msa.msa_max_score = torch.zeros(4 * 16 * 2) with pytest.raises(ValueError, match=r"msa_max_score backing store"): - metadata._ensure_msa_decode_scratch_buffers( + msa._ensure_msa_decode_scratch_buffers( num_index_heads=4, max_batch=2, capture_graph=False, @@ -1823,21 +1856,15 @@ def test_msa_scratch_sizing_covers_spec_verify_tokens(): ) -def _kv_lens_update_metadata( - monkeypatch, *, qo_lens, kv_staged, kv_corrected, page_size, prefill_counts=False -): +def _kv_lens_update_metadata(*, qo_lens, kv_staged, kv_corrected, page_size, prefill_counts=False): """Metadata with just enough staged state for on_update_kv_lens: requests with several query tokens each, optimistic kv_lens from prepare(), then the corrected kv_lens. Everything lives on the host. `prefill_counts` stages the per-step valid-block buffer a mixed step uses in place of the graph-stable one. """ - from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend - - # The base hook only touches MLA caches this metadata never builds. - monkeypatch.setattr(msa_backend.TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None) - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) qo = torch.tensor(qo_lens, dtype=torch.int32) batch = int(qo.shape[0]) total_q = int(qo.sum()) @@ -1845,137 +1872,139 @@ def _kv_lens_update_metadata( metadata._seq_lens_cuda = qo.clone() metadata._num_tokens = total_q metadata.kv_lens_cuda = torch.tensor(kv_corrected, dtype=torch.int32) - metadata.msa_kv_lens_staged = torch.tensor(kv_staged, dtype=torch.int32) + msa.msa_kv_lens_staged = torch.tensor(kv_staged, dtype=torch.int32) # What prepare() handed the decode kernels, before the correction. - metadata.msa_seq_lens_cuda = torch.tensor(kv_staged, dtype=torch.int32) + msa.msa_seq_lens_cuda = torch.tensor(kv_staged, dtype=torch.int32) metadata.kv_cache_manager = SimpleNamespace(tokens_per_block=page_size) # Page table: request b's page p is block 100 // page_size * b + p, so its # position pos maps to slot 100 * b + pos (16 positions per request). assert 100 % page_size == 0 - metadata.msa_block_table = (100 // page_size) * torch.arange( - batch, dtype=torch.int32 - ).unsqueeze(1) + torch.arange(16 // page_size, dtype=torch.int32).unsqueeze(0) + msa.msa_block_table = (100 // page_size) * torch.arange(batch, dtype=torch.int32).unsqueeze( + 1 + ) + torch.arange(16 // page_size, dtype=torch.int32).unsqueeze(0) qo_long = qo.to(torch.long) - metadata.msa_q_batch_row = torch.repeat_interleave( - torch.arange(batch, dtype=torch.int32), qo_long - ) + msa.msa_q_batch_row = torch.repeat_interleave(torch.arange(batch, dtype=torch.int32), qo_long) starts = torch.cumsum(qo_long, 0) - qo_long - metadata.msa_q_intra = ( + msa.msa_q_intra = ( torch.arange(total_q, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) ).to(torch.int32) - metadata.msa_out_cache_loc = torch.full((total_q + 3,), -1, dtype=torch.int32) - metadata.msa_n_valid_blocks = torch.zeros(total_q + 3, dtype=torch.int32) - metadata._msa_prefill_n_valid_blocks = ( + msa.msa_out_cache_loc = torch.full((total_q + 3,), -1, dtype=torch.int32) + msa.msa_n_valid_blocks = torch.zeros(total_q + 3, dtype=torch.int32) + msa._msa_prefill_n_valid_blocks = ( torch.zeros(total_q, dtype=torch.int32) if prefill_counts else None ) - metadata._msa_fields_ready = True - metadata._msa_kv_lens_dynamic = True + msa._msa_fields_ready = True + msa._msa_kv_lens_dynamic = True return metadata @pytest.mark.parametrize("nvfp4", [False, True]) -def test_on_update_kv_lens_rederives_lengths_slots_and_counts(monkeypatch, nvfp4): +def test_on_update_kv_lens_rederives_lengths_slots_and_counts(nvfp4): """Request 0 loses one rejected draft token (staged 9 -> corrected 8), request 1 is unchanged. The per-request length the decode kernels read, the write slots and the per-token valid-block counts must all follow. """ metadata = _kv_lens_update_metadata( - monkeypatch, qo_lens=(2, 3), kv_staged=(9, 12), kv_corrected=(8, 12), page_size=4 + qo_lens=(2, 3), kv_staged=(9, 12), kv_corrected=(8, 12), page_size=4 ) + msa = metadata.ext(MiniMaxM3MsaMetadata) if nvfp4: # Both rows are in the context prefix when extend_ctx is active. metadata.num_contexts = 2 - metadata.msa_cu_kv_lens = torch.tensor([0, 9, 21, -1], dtype=torch.int32) - metadata.msa_cu_q_lens = torch.tensor([0, 2, 5], dtype=torch.int32) + msa.msa_cu_kv_lens = torch.tensor([0, 9, 21, -1], dtype=torch.int32) + msa.msa_cu_q_lens = torch.tensor([0, 2, 5], dtype=torch.int32) metadata.on_update_kv_lens() metadata.on_update_kv_lens() # Corrections must be idempotent. if nvfp4: - assert metadata.msa_cu_kv_lens.tolist() == [0, 8, 20, -1] - assert metadata.msa_cu_q_lens.tolist() == [0, 2, 5] + assert msa.msa_cu_kv_lens.tolist() == [0, 8, 20, -1] + assert msa.msa_cu_q_lens.tolist() == [0, 2, 5] # One length per request, in the buffer every decode kernel reads. - assert metadata.msa_seq_lens_cuda.tolist() == [8, 12] + assert msa.msa_seq_lens_cuda.tolist() == [8, 12] # Token positions: request 0 attends 8 tokens with 2 queries -> 6, 7; # request 1 attends 12 with 3 queries -> 9, 10, 11. - assert metadata.msa_out_cache_loc[:5].tolist() == [6, 7, 109, 110, 111] - assert metadata.msa_out_cache_loc[5:].tolist() == [-1, -1, -1] + assert msa.msa_out_cache_loc[:5].tolist() == [6, 7, 109, 110, 111] + assert msa.msa_out_cache_loc[5:].tolist() == [-1, -1, -1] # ceil((pos + 1) / 4) per token, into the graph-stable buffer on a # pure-decode step; the tail past the step's tokens is left alone. - assert metadata.msa_n_valid_blocks[:5].tolist() == [2, 2, 3, 3, 3] - assert metadata.msa_n_valid_blocks[5:].tolist() == [0, 0, 0] + assert msa.msa_n_valid_blocks[:5].tolist() == [2, 2, 3, 3, 3] + assert msa.msa_n_valid_blocks[5:].tolist() == [0, 0, 0] -def test_on_update_kv_lens_clamps_to_the_staged_lens(monkeypatch): +def test_on_update_kv_lens_clamps_to_the_staged_lens(): """A correction can only shrink lengths; anything above the staged value is clamped, so the staged page table stays valid. """ metadata = _kv_lens_update_metadata( - monkeypatch, qo_lens=(1, 1), kv_staged=(5, 7), kv_corrected=(9, 7), page_size=4 + qo_lens=(1, 1), kv_staged=(5, 7), kv_corrected=(9, 7), page_size=4 ) + msa = metadata.ext(MiniMaxM3MsaMetadata) metadata.on_update_kv_lens() - assert metadata.msa_seq_lens_cuda.tolist() == [5, 7] - assert metadata.msa_out_cache_loc[:2].tolist() == [4, 106] - assert metadata.msa_n_valid_blocks[:2].tolist() == [2, 2] + assert msa.msa_seq_lens_cuda.tolist() == [5, 7] + assert msa.msa_out_cache_loc[:2].tolist() == [4, 106] + assert msa.msa_n_valid_blocks[:2].tolist() == [2, 2] -def test_on_update_kv_lens_patches_a_mixed_steps_per_step_counts(monkeypatch): +def test_on_update_kv_lens_patches_a_mixed_steps_per_step_counts(): """A mixed step stages its valid-block counts in the per-step prefill buffer rather than the graph-stable one, and the correction must land there. The context row (row 0) is untouched by the correction; the generation row (row 1) loses one token. """ metadata = _kv_lens_update_metadata( - monkeypatch, qo_lens=(5, 3), kv_staged=(5, 12), kv_corrected=(5, 11), page_size=4, prefill_counts=True, ) + msa = metadata.ext(MiniMaxM3MsaMetadata) metadata.on_update_kv_lens() - assert metadata.msa_seq_lens_cuda.tolist() == [5, 11] + assert msa.msa_seq_lens_cuda.tolist() == [5, 11] # Context row: positions 0..4 -> slots 0..4. Generation row attends 11 # tokens with 3 queries: positions 8, 9, 10. - assert metadata.msa_out_cache_loc[:8].tolist() == [0, 1, 2, 3, 4, 108, 109, 110] + assert msa.msa_out_cache_loc[:8].tolist() == [0, 1, 2, 3, 4, 108, 109, 110] # ceil((pos + 1) / 4): context [1, 1, 1, 1, 2], generation [3, 3, 3]. - assert metadata._msa_prefill_n_valid_blocks.tolist() == [1, 1, 1, 1, 2, 3, 3, 3] - assert metadata.msa_n_valid_blocks.tolist() == [0] * 11 + assert msa._msa_prefill_n_valid_blocks.tolist() == [1, 1, 1, 1, 2, 3, 3, 3] + assert msa.msa_n_valid_blocks.tolist() == [0] * 11 -def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch): +def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(): """Without speculative decoding the patch and its staging are skipped. The gate also sees max_total_draft_tokens, set by update_spec_dec_param, which covers draft_len=1 under trtllm-gen (no extra KV tokens, linear-tree spec decoding reported disabled). """ metadata = _kv_lens_update_metadata( - monkeypatch, qo_lens=(1, 1), kv_staged=(4, 6), kv_corrected=(3, 5), page_size=4 + qo_lens=(1, 1), kv_staged=(4, 6), kv_corrected=(3, 5), page_size=4 ) - metadata.max_total_draft_tokens = None + msa = metadata.ext(MiniMaxM3MsaMetadata) + spec = metadata.ext(SpecDecodingMetadata) + spec.max_total_draft_tokens = None metadata.draft_kv_cache_manager = None - metadata.is_spec_decoding_enabled = False + spec.is_spec_decoding_enabled = False metadata.kv_cache_params = SimpleNamespace(num_extra_kv_tokens=0) metadata.runtime_features = None - assert metadata._msa_kv_lens_may_change() is False - metadata.max_total_draft_tokens = 1 - assert metadata._msa_kv_lens_may_change() is True - metadata.max_total_draft_tokens = None + assert msa._msa_kv_lens_may_change() is False + spec.max_total_draft_tokens = 1 + assert msa._msa_kv_lens_may_change() is True + spec.max_total_draft_tokens = None metadata.kv_cache_params = SimpleNamespace(num_extra_kv_tokens=2) - assert metadata._msa_kv_lens_may_change() is True + assert msa._msa_kv_lens_may_change() is True - metadata._msa_kv_lens_dynamic = False + msa._msa_kv_lens_dynamic = False metadata.on_update_kv_lens() - assert metadata.msa_seq_lens_cuda.tolist() == [4, 6] - assert metadata.msa_out_cache_loc.tolist() == [-1] * 5 - assert metadata.msa_n_valid_blocks.tolist() == [0] * 5 + assert msa.msa_seq_lens_cuda.tolist() == [4, 6] + assert msa.msa_out_cache_loc.tolist() == [-1] * 5 + assert msa.msa_n_valid_blocks.tolist() == [0] * 5 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index 864f8d66832a..9d094c3db67d 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -6,6 +6,7 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.features import DraftKVMetadata from tensorrt_llm._torch.attention.backends.sparse.qsa import ( QSASparseMetadataParams, QSASparseParams, @@ -26,7 +27,7 @@ triton_qsa_prefill_compress, triton_qsa_unscale_block_table, ) -from tensorrt_llm._torch.attention.backends.sparse.qsa.metadata import QSAAttentionMetadata +from tensorrt_llm._torch.attention.backends.sparse.qsa.metadata import QSAMetadata from tensorrt_llm._torch.attention.backends.sparse.qsa.module import ( _capture_mtp_shared_topk, _store_paged_kv_reference, @@ -48,6 +49,20 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode +def _qsa_stub(**fields) -> SimpleNamespace: + """A metadata stub whose QSA extension holds the ``qsa_``-prefixed fields.""" + qsa = SimpleNamespace( + **{ + name[len("qsa_") :]: fields.pop(name) + for name in list(fields) + if name.startswith("qsa_") + } + ) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: qsa if cls is QSAMetadata else None + return metadata + + def _params( token_topk: int = 8, seq_len_threshold: int | None = None, @@ -278,7 +293,7 @@ def test_qsa_mtp_capture_uses_each_request_own_rows() -> None: seq_lens = torch.tensor([1, 1, 4], device="cuda", dtype=torch.int32) indices = torch.arange(6, device="cuda", dtype=torch.int32).unsqueeze(1).repeat(1, 2) visible_blocks = torch.arange(100, 106, device="cuda", dtype=torch.int32) - metadata = SimpleNamespace( + metadata = _qsa_stub( num_seqs=3, seq_lens_cuda=seq_lens, qsa_mtp_num_accepted=torch.tensor([1, 1, 3], device="cuda", dtype=torch.int32), @@ -286,19 +301,20 @@ def test_qsa_mtp_capture_uses_each_request_own_rows() -> None: qsa_shared_topk_visible_blocks=torch.full((1, 3), -99, device="cuda", dtype=torch.int32), qsa_shared_topk_captured_slots=set(), ) + qsa = metadata.find_ext(QSAMetadata) _capture_mtp_shared_topk(metadata, 0, indices, visible_blocks) # Request 2 starts at row 2 and accepted three tokens, so it keeps row 4. - assert metadata.qsa_shared_topk_indices[0, :, 0].tolist() == [0, 1, 4] - assert metadata.qsa_shared_topk_visible_blocks[0].tolist() == [100, 101, 104] - assert metadata.qsa_shared_topk_captured_slots == {0} + assert qsa.shared_topk_indices[0, :, 0].tolist() == [0, 1, 4] + assert qsa.shared_topk_visible_blocks[0].tolist() == [100, 101, 104] + assert qsa.shared_topk_captured_slots == {0} @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_qsa_mtp_capture_keeps_each_layer_in_its_own_slot() -> None: """Two sparse layers that rank the same requests differently both survive.""" - metadata = SimpleNamespace( + metadata = _qsa_stub( num_seqs=2, seq_lens_cuda=torch.ones(2, device="cuda", dtype=torch.int32), qsa_mtp_num_accepted=None, @@ -306,6 +322,7 @@ def test_qsa_mtp_capture_keeps_each_layer_in_its_own_slot() -> None: qsa_shared_topk_visible_blocks=torch.full((2, 2), -99, device="cuda", dtype=torch.int32), qsa_shared_topk_captured_slots=set(), ) + qsa = metadata.find_ext(QSAMetadata) for slot, base in enumerate((10, 20)): indices = torch.tensor( @@ -314,12 +331,12 @@ def test_qsa_mtp_capture_keeps_each_layer_in_its_own_slot() -> None: visible_blocks = torch.tensor([base, base], device="cuda", dtype=torch.int32) _capture_mtp_shared_topk(metadata, slot, indices, visible_blocks) - assert metadata.qsa_shared_topk_indices.tolist() == [ + assert qsa.shared_topk_indices.tolist() == [ [[10, 11], [12, 13]], [[20, 21], [22, 23]], ] - assert metadata.qsa_shared_topk_visible_blocks.tolist() == [[10, 10], [20, 20]] - assert metadata.qsa_shared_topk_captured_slots == {0, 1} + assert qsa.shared_topk_visible_blocks.tolist() == [[10, 10], [20, 20]] + assert qsa.shared_topk_captured_slots == {0, 1} @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -352,7 +369,7 @@ def test_qsa_mtp_reuse_waits_for_the_calling_layer_to_capture() -> None: # contents of a slot no layer has written yet. shared = torch.zeros((2, 1, params.block_topk), dtype=torch.int32, device="cuda") shared[1].fill_(999) - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=block_table, kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), qsa_indexer_skip_topk=True, @@ -363,6 +380,7 @@ def test_qsa_mtp_reuse_waits_for_the_calling_layer_to_capture() -> None: num_seqs=1, seq_lens_cuda=torch.tensor([1], dtype=torch.int32, device="cuda"), ) + qsa = metadata.find_ext(QSAMetadata) selected = select_qsa_paged_tokens( q, @@ -381,8 +399,8 @@ def test_qsa_mtp_reuse_waits_for_the_calling_layer_to_capture() -> None: assert max(chosen) <= int(query_positions[0]), "the selection must stay causal" # Scoring ran for this layer, so its slot now carries a real selection # rather than what a premature reuse would have read back. - assert metadata.qsa_shared_topk_captured_slots == {0, 1} - assert 999 not in metadata.qsa_shared_topk_indices[1].tolist() + assert qsa.shared_topk_captured_slots == {0, 1} + assert 999 not in qsa.shared_topk_indices[1].tolist() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -539,26 +557,27 @@ def test_qsa_spec_dec_remaps_rows_after_the_draft_loop_packs_one_row_per_request """The first draft step repacks the batch, so the row mapping must follow it.""" num_seqs, rows_per_seq = 2, 2 num_tokens = num_seqs * rows_per_seq - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) + qsa = QSAMetadata(metadata) + metadata.replace_ext(qsa) metadata._seq_lens = torch.full((num_seqs,), rows_per_seq, dtype=torch.int32) metadata._seq_lens_cuda = metadata._seq_lens.cuda() metadata._num_contexts = 0 metadata._num_tokens = num_tokens metadata.kv_lens_cuda_runtime = torch.tensor([40, 56], dtype=torch.int32, device="cuda") - metadata.sparse_metadata_params = SimpleNamespace(compress_ratio=4) - metadata.qsa_has_local_layers = True - metadata.enable_flash_mla = False + qsa.params = SimpleNamespace(compress_ratio=4) + qsa.has_local_layers = True metadata._mla_scheduler_buffers_valid = True metadata._mla_ctx_cu_seqlens_valid = True - metadata.qsa_req_idx_per_token = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") - metadata.qsa_logical_positions = torch.zeros(num_tokens, dtype=torch.int64, device="cuda") - metadata.qsa_sequence_lengths = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") - metadata.qsa_visible_blocks = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") - metadata._qsa_cu_seq_lens = torch.zeros(num_seqs + 1, dtype=torch.int64, device="cuda") - metadata._qsa_token_arange = torch.arange(num_tokens, dtype=torch.int64, device="cuda") + qsa.req_idx_per_token = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") + qsa.logical_positions = torch.zeros(num_tokens, dtype=torch.int64, device="cuda") + qsa.sequence_lengths = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") + qsa.visible_blocks = torch.zeros(num_tokens, dtype=torch.int32, device="cuda") + qsa._cu_seq_lens = torch.zeros(num_seqs + 1, dtype=torch.int64, device="cuda") + qsa._token_arange = torch.arange(num_tokens, dtype=torch.int64, device="cuda") - metadata._refresh_qsa_token_mapping() - assert metadata.qsa_req_idx_per_token.tolist() == [0, 0, 1, 1] + qsa._refresh_token_mapping() + assert qsa.req_idx_per_token.tolist() == [0, 0, 1, 1] # What the draft loop does to the metadata between its first two steps. metadata._seq_lens.fill_(1) @@ -566,10 +585,8 @@ def test_qsa_spec_dec_remaps_rows_after_the_draft_loop_packs_one_row_per_request metadata._num_tokens = num_seqs metadata.update_for_spec_dec() - assert metadata.qsa_req_idx_per_token[:num_seqs].tolist() == [0, 1] - torch.testing.assert_close( - metadata.qsa_sequence_lengths[:num_seqs], metadata.kv_lens_cuda_runtime - ) + assert qsa.req_idx_per_token[:num_seqs].tolist() == [0, 1] + torch.testing.assert_close(qsa.sequence_lengths[:num_seqs], metadata.kv_lens_cuda_runtime) @pytest.mark.parametrize("cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) @@ -934,7 +951,7 @@ def test_qsa_paged_selection_torch_fallback_stays_causal() -> None: query_positions = torch.tensor([5], device="cuda") sequence_lengths = torch.tensor([16], device="cuda") request_indices = torch.tensor([0], dtype=torch.int32, device="cuda") - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=block_table, kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), ) @@ -991,7 +1008,7 @@ def test_qsa_paged_selection_supports_multiple_rows_per_request() -> None: # each row's causal boundary in the selector and incomplete tail. sequence_lengths = torch.tensor([10, 10, 14, 14], device="cuda") request_indices = torch.tensor([0, 0, 1, 1], dtype=torch.int32, device="cuda") - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=block_table, kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), ) @@ -1243,7 +1260,7 @@ def test_qsa_sparse_gqa_reads_hnd_paged_cache() -> None: v_cache[page, 0, within] = torch.tensor([position * 10.0, position * 10.0 + 1.0]) selected = torch.tensor([[0, 3, 4, -1]], dtype=torch.int32) q = torch.tensor([[[1.0, 0.0], [0.0, 1.0]]]) - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=torch.tensor([[0, 1]], dtype=torch.int32), kv_cache_manager=SimpleNamespace(tokens_per_block=4), ) @@ -1268,7 +1285,7 @@ def test_qsa_sparse_gqa_all_invalid_row_returns_zero() -> None: q = torch.randn(1, 2, 4) k_cache = torch.randn(1, 1, 4, 4) v_cache = torch.randn_like(k_cache) - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=torch.tensor([[0]], dtype=torch.int32), kv_cache_manager=SimpleNamespace(tokens_per_block=4), ) @@ -1336,7 +1353,7 @@ def test_fused_qsa_sparse_gqa_matches_reference(head_dim: int, cache_dtype: torc dtype=torch.int32, device="cuda", ) - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=block_table, kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), ) @@ -1386,7 +1403,7 @@ def test_qsa_splitk_merge_output_gate_matches_separate_gate() -> None: request_indices = torch.zeros(rows, dtype=torch.int32, device="cuda") selected = torch.arange(selected_width, dtype=torch.int32, device="cuda").reshape(rows, -1) output_gate = torch.randn_like(q) - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=block_table, kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), ) @@ -1456,7 +1473,7 @@ def test_fused_qsa_prefill_bounds_sparse_attention_to_visible_tokens() -> None: selected = torch.full((rows, final_topk), -1, dtype=torch.int32, device="cuda") columns = torch.arange(final_topk, device="cuda").unsqueeze(0) selected.copy_(torch.where(columns < visible_tokens.unsqueeze(1), columns, -1)) - metadata = SimpleNamespace( + metadata = _qsa_stub( qsa_block_table=block_table, kv_cache_manager=SimpleNamespace(tokens_per_block=tokens_per_block), ) @@ -1604,20 +1621,22 @@ def test_qsa_speculative_snapshot_state_uses_host_lengths( num_contexts: int, expected: bool, ) -> None: - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) + qsa = QSAMetadata(metadata) + metadata.replace_ext(qsa) metadata._seq_lens = torch.tensor(sequence_lengths, dtype=torch.int32) metadata._num_contexts = num_contexts - metadata._refresh_qsa_speculative_snapshot_state() + qsa._refresh_speculative_snapshot_state() - assert metadata.qsa_needs_speculative_snapshot is expected + assert qsa.needs_speculative_snapshot is expected def test_qsa_ordinary_mixed_batch_skips_speculative_cache_snapshot() -> None: indexer = object.__new__(QSAIndexer) torch.nn.Module.__init__(indexer) indexer._pending_speculative_cache = object() - metadata = SimpleNamespace( + metadata = _qsa_stub( num_contexts=1, num_seqs=2, num_tokens=2, @@ -1643,7 +1662,7 @@ def test_qsa_multi_token_generation_captures_speculative_cache_snapshot() -> Non indexer = object.__new__(QSAIndexer) torch.nn.Module.__init__(indexer) indexer._pending_speculative_cache = None - metadata = SimpleNamespace( + metadata = _qsa_stub( num_contexts=0, num_seqs=1, num_tokens=4, @@ -1717,7 +1736,6 @@ def _qsa_layout_manager( def _qsa_metadata_with_draft( - monkeypatch, *, mtp_index_share: bool = False, draft_layout: dict | None = _DRAFT_LAYOUT, @@ -1725,8 +1743,7 @@ def _qsa_metadata_with_draft( ): """QSA metadata over a target manager and, unless `draft_layout` is None, a separate draft manager (`draft_manager` replaces the one `draft_layout` describes).""" - monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) metadata.kv_cache_manager = _qsa_layout_manager(**_TARGET_LAYOUT) if draft_manager is None and draft_layout is not None: draft_manager = _qsa_layout_manager(**draft_layout) @@ -1737,7 +1754,9 @@ def _qsa_metadata_with_draft( metadata.max_num_sequences = 4 metadata.max_num_tokens = 16 metadata.draft_replay_swapped_attrs = {} - metadata.__post_init__() + qsa = QSAMetadata(metadata) + qsa.allocate(None) + metadata.replace_ext(qsa) return metadata @@ -1760,44 +1779,48 @@ def _set_qsa_offsets(metadata, num_seqs: int, generation: int): if layout is _TARGET_LAYOUT: metadata.kv_cache_block_offsets = offsets else: - metadata.draft_kv_cache_block_offsets = offsets + draft_kv = metadata.find_ext(DraftKVMetadata) or DraftKVMetadata(metadata) + draft_kv.kv_cache_block_offsets = offsets + metadata.replace_ext(draft_kv) return slots @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_qsa_draft_block_table_exists_only_for_a_separate_draft_manager(monkeypatch) -> None: - metadata = _qsa_metadata_with_draft(monkeypatch) +def test_qsa_draft_block_table_exists_only_for_a_separate_draft_manager() -> None: + metadata = _qsa_metadata_with_draft() + qsa = metadata.find_ext(QSAMetadata) # Each table is sized, and its pool and page scale taken, from its own manager. - assert metadata.qsa_block_table.shape == (4, _TARGET_LAYOUT["max_blocks_per_seq"]) - assert metadata.qsa_draft_block_table.shape == (4, _DRAFT_LAYOUT["max_blocks_per_seq"]) - assert metadata.qsa_draft_block_table.data_ptr() != metadata.qsa_block_table.data_ptr() - assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) - assert (metadata._qsa_draft_attention_pool_id, metadata._qsa_draft_page_index_scale) == (0, 4) - - target_only = _qsa_metadata_with_draft(monkeypatch, draft_layout=None) - assert target_only.qsa_draft_block_table is None + assert qsa.block_table.shape == (4, _TARGET_LAYOUT["max_blocks_per_seq"]) + assert qsa.draft_block_table.shape == (4, _DRAFT_LAYOUT["max_blocks_per_seq"]) + assert qsa.draft_block_table.data_ptr() != qsa.block_table.data_ptr() + assert (qsa._attention_pool_id, qsa._page_index_scale) == (1, 6) + assert (qsa._draft_attention_pool_id, qsa._draft_page_index_scale) == (0, 4) + + target_only = _qsa_metadata_with_draft(draft_layout=None) + assert target_only.ext(QSAMetadata).draft_block_table is None assert target_only.prepare_for_draft_forward() is None @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("kind", ("not_qsa", "no_local_layer")) -def test_qsa_draft_table_needs_a_qsa_draft_manager_with_local_layers(monkeypatch, kind) -> None: +def test_qsa_draft_table_needs_a_qsa_draft_manager_with_local_layers(kind) -> None: if kind == "not_qsa": draft_manager = SimpleNamespace(qsa_position_layer_id=0) else: draft_manager = _qsa_layout_manager(**_DRAFT_LAYOUT) draft_manager.qsa_position_layer_id = None - metadata = _qsa_metadata_with_draft(monkeypatch, draft_manager=draft_manager) + metadata = _qsa_metadata_with_draft(draft_manager=draft_manager) + qsa = metadata.find_ext(QSAMetadata) metadata._seq_lens = torch.ones(2, dtype=torch.int32) - assert metadata.qsa_draft_block_table is None + assert qsa.draft_block_table is None # The target's own manager is bound, as outside a draft loop, and the hook and the # refresh both leave the target bindings alone. metadata.prepare_for_draft_forward() - metadata._refresh_qsa_draft_block_table() + qsa._refresh_draft_block_table() assert metadata.draft_replay_swapped_attrs == {} - assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + assert (qsa._attention_pool_id, qsa._page_index_scale) == (1, 6) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -1810,46 +1833,46 @@ def test_qsa_draft_table_needs_a_qsa_draft_manager_with_local_layers(monkeypatch ), ids=("draft_wider", "target_wider", "no_draft"), ) -def test_qsa_shared_topk_slots_cover_the_wider_manager(monkeypatch, draft_layout, slots) -> None: +def test_qsa_shared_topk_slots_cover_the_wider_manager(draft_layout, slots) -> None: """The shared selection buffers are addressed by whichever manager is bound.""" - metadata = _qsa_metadata_with_draft( - monkeypatch, mtp_index_share=True, draft_layout=draft_layout - ) + metadata = _qsa_metadata_with_draft(mtp_index_share=True, draft_layout=draft_layout) + qsa = metadata.find_ext(QSAMetadata) - assert metadata.qsa_shared_topk_indices.shape[0] == slots - assert metadata.qsa_shared_topk_visible_blocks.shape[0] == slots + assert qsa.shared_topk_indices.shape[0] == slots + assert qsa.shared_topk_visible_blocks.shape[0] == slots @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_qsa_draft_forward_addresses_the_draft_manager_pages(monkeypatch) -> None: - """prepare() decodes both tables; the draft swap makes `qsa_block_table` the draft one +def test_qsa_draft_forward_addresses_the_draft_manager_pages() -> None: + """prepare() decodes both tables; the draft swap makes `block_table` the draft one and the restore puts the target's table, pool and scale back.""" num_seqs = 3 - metadata = _qsa_metadata_with_draft(monkeypatch) + metadata = _qsa_metadata_with_draft() + qsa = metadata.find_ext(QSAMetadata) metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation=0) - target_table = metadata.qsa_block_table + target_table = qsa.block_table - metadata._refresh_qsa_block_table() - metadata._refresh_qsa_draft_block_table() + qsa._refresh_block_table() + qsa._refresh_draft_block_table() torch.testing.assert_close(target_table[:num_seqs], target_slots) - torch.testing.assert_close(metadata.qsa_draft_block_table[:num_seqs], draft_slots) + torch.testing.assert_close(qsa.draft_block_table[:num_seqs], draft_slots) draft_manager = metadata.draft_kv_cache_manager saved = prepare_attn_metadata_for_draft_replay(metadata, draft_manager) assert saved is not None assert metadata.kv_cache_manager is draft_manager - assert metadata.qsa_block_table is metadata.qsa_draft_block_table + assert qsa.block_table is qsa.draft_block_table # Every operand of the table refresh is the draft manager's while swapped, so a # refresh rebuilds the draft table and leaves the target's alone. - metadata._refresh_qsa_block_table() - torch.testing.assert_close(metadata.qsa_block_table[:num_seqs], draft_slots) + qsa._refresh_block_table() + torch.testing.assert_close(qsa.block_table[:num_seqs], draft_slots) torch.testing.assert_close(target_table[:num_seqs], target_slots) restore_attn_metadata_after_draft_replay(metadata, saved) assert metadata.kv_cache_manager is not draft_manager - assert metadata.qsa_block_table is target_table - assert (metadata._qsa_attention_pool_id, metadata._qsa_page_index_scale) == (1, 6) + assert qsa.block_table is target_table + assert (qsa._attention_pool_id, qsa._page_index_scale) == (1, 6) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -1857,55 +1880,54 @@ def test_qsa_prepare_decodes_the_target_and_the_draft_table(monkeypatch) -> None """`prepare()` is what refreshes both tables before every step and replay: a second generation of offsets reaches the target table and the draft table through it.""" num_seqs = 3 - metadata = _qsa_metadata_with_draft(monkeypatch) + metadata = _qsa_metadata_with_draft() + qsa = metadata.find_ext(QSAMetadata) metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) - # The parent's prepare and the row mapping are not under test here. - monkeypatch.setattr(TrtllmAttentionMetadata, "prepare", lambda self: None) - monkeypatch.setattr(QSAAttentionMetadata, "_refresh_qsa_token_mapping", lambda self: None) - monkeypatch.setattr( - QSAAttentionMetadata, "_refresh_qsa_speculative_snapshot_state", lambda self: None - ) + # The row mapping is not under test here. + monkeypatch.setattr(QSAMetadata, "_refresh_token_mapping", lambda self: None) + monkeypatch.setattr(QSAMetadata, "_refresh_speculative_snapshot_state", lambda self: None) for generation in (0, 1): target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation) - metadata.prepare() - torch.testing.assert_close(metadata.qsa_block_table[:num_seqs], target_slots) - torch.testing.assert_close(metadata.qsa_draft_block_table[:num_seqs], draft_slots) + qsa.prepare() + torch.testing.assert_close(qsa.block_table[:num_seqs], target_slots) + torch.testing.assert_close(qsa.draft_block_table[:num_seqs], draft_slots) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_qsa_draft_forward_refuses_an_unswapped_manager(monkeypatch) -> None: - metadata = _qsa_metadata_with_draft(monkeypatch) +def test_qsa_draft_forward_refuses_an_unswapped_manager() -> None: + metadata = _qsa_metadata_with_draft() with pytest.raises(RuntimeError, match="draft KV cache manager"): metadata.prepare_for_draft_forward() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_qsa_draft_table_replays_the_refreshed_draft_pages(monkeypatch) -> None: +def test_qsa_draft_table_replays_the_refreshed_draft_pages() -> None: """A graph captured across the draft swap reads each table by address: the trunk segment the target table, the draft segment the draft one, and a refresh outside capture (as prepare() does before every replay) is what the replay sees.""" num_seqs = 3 - metadata = _qsa_metadata_with_draft(monkeypatch) + metadata = _qsa_metadata_with_draft() + qsa = metadata.find_ext(QSAMetadata) metadata._seq_lens = torch.ones(num_seqs, dtype=torch.int32) _set_qsa_offsets(metadata, num_seqs, generation=0) - metadata._refresh_qsa_block_table() - metadata._refresh_qsa_draft_block_table() + qsa._refresh_block_table() + qsa._refresh_draft_block_table() draft_manager = metadata.draft_kv_cache_manager - target_seen = torch.zeros_like(metadata.qsa_block_table[:num_seqs]) - draft_seen = torch.zeros_like(metadata.qsa_draft_block_table[:num_seqs]) + target_seen = torch.zeros_like(qsa.block_table[:num_seqs]) + draft_seen = torch.zeros_like(qsa.draft_block_table[:num_seqs]) graph = torch.cuda.CUDAGraph() torch.cuda.synchronize() with torch.cuda.graph(graph): - target_seen.copy_(metadata.qsa_block_table[:num_seqs]) + target_seen.copy_(qsa.block_table[:num_seqs]) saved = prepare_attn_metadata_for_draft_replay(metadata, draft_manager) - draft_seen.copy_(metadata.qsa_block_table[:num_seqs]) + draft_seen.copy_(qsa.block_table[:num_seqs]) restore_attn_metadata_after_draft_replay(metadata, saved) target_slots, draft_slots = _set_qsa_offsets(metadata, num_seqs, generation=1) - metadata._refresh_qsa_block_table() - metadata._refresh_qsa_draft_block_table() + qsa._refresh_block_table() + qsa._refresh_draft_block_table() graph.replay() torch.cuda.synchronize() diff --git a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py index 4b93cd54dcea..324b39a8fa76 100644 --- a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py @@ -12,10 +12,9 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.sparse.rocket import ( RocketKVCacheManager, + RocketMetadata, RocketTrtllmAttention, - RocketTrtllmAttentionMetadata, RocketVanillaAttention, - RocketVanillaAttentionMetadata, ) from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -140,7 +139,7 @@ def create_test_metadata( request_ids, kv_cache_manager, sparse_attn_config, - metadata_cls, + backend_cls, ): sparse_metadata_params = sparse_attn_config.to_sparse_metadata_params() prompt_lens = [] @@ -150,7 +149,7 @@ def create_test_metadata( else: prompt_lens.append(past_token + seq_len) - metadata = metadata_cls( + metadata = backend_cls.create_metadata( seq_lens=torch.tensor(seq_lens, dtype=torch.int), num_contexts=num_contexts, kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=past_seen_tokens), @@ -291,7 +290,7 @@ def test_sparse_kv_predict(batch_size, num_contexts): request_ids, trtllm_kv_cache_manager, sparse_attn_config, - RocketTrtllmAttentionMetadata, + RocketTrtllmAttention, ) _ = create_test_metadata( seq_lens, @@ -300,7 +299,7 @@ def test_sparse_kv_predict(batch_size, num_contexts): request_ids, vanilla_kv_cache_manager, sparse_attn_config, - RocketVanillaAttentionMetadata, + RocketVanillaAttention, ) total_tokens = sum(seq_lens) @@ -518,7 +517,7 @@ def test_sparse_attn_predict(batch_size, num_contexts): request_ids, trtllm_kv_cache_manager, sparse_attn_config_trtllm, - RocketTrtllmAttentionMetadata, + RocketTrtllmAttention, ) vanilla_metadata = create_test_metadata( seq_lens, @@ -527,7 +526,7 @@ def test_sparse_attn_predict(batch_size, num_contexts): request_ids, vanilla_kv_cache_manager, sparse_attn_config_vanilla, - RocketVanillaAttentionMetadata, + RocketVanillaAttention, ) total_tokens = sum(seq_lens) @@ -552,8 +551,8 @@ def test_sparse_attn_predict(batch_size, num_contexts): trtllm_kt_tokens_per_block = trtllm_kv_cache_manager.kt_tokens_per_block vanilla_kv_cache_manager.kt_tokens_per_block - trtllm_block_offsets = trtllm_metadata.kt_cache_block_offsets - vanilla_block_offsets = vanilla_metadata.kt_cache_block_offsets + trtllm_block_offsets = trtllm_metadata.ext(RocketMetadata).kt_cache_block_offsets + vanilla_block_offsets = vanilla_metadata.ext(RocketMetadata).kt_cache_block_offsets for req_idx in range(num_contexts, batch_size): # Get the number of KT tokens for this request diff --git a/tests/unittest/_torch/attention/sparse/test_flashinfer_utils.py b/tests/unittest/_torch/attention/sparse/test_flashinfer_utils.py index f3ceb3345ade..2b02d3f8d7e8 100644 --- a/tests/unittest/_torch/attention/sparse/test_flashinfer_utils.py +++ b/tests/unittest/_torch/attention/sparse/test_flashinfer_utils.py @@ -19,7 +19,7 @@ from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.kernels import ( deepseek_v4_local_to_global_indices, ) -from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManager, DSAMetadata from tensorrt_llm._torch.attention.backends.sparse.flashinfer_utils import ( allocate_sparse_mla_split_workspace, ) @@ -202,8 +202,11 @@ def test_dsa_sparse_mla_cuda_graph_capture_replay() -> None: ) metadata._num_tokens = num_tokens metadata._num_ctx_tokens = 0 - metadata.max_draft_tokens = 0 - metadata.token_positions_cuda = torch.arange(num_tokens, dtype=torch.int32, device="cuda") + dsa = object.__new__(DSAMetadata) + dsa.metadata = metadata + metadata.replace_ext(dsa) + dsa.max_draft_tokens = 0 + dsa.token_positions_cuda = torch.arange(num_tokens, dtype=torch.int32, device="cuda") metadata.kv_cache_manager = SimpleNamespace( tokens_per_block=inline_scale_kv.PAGE_SIZE, get_unique_primary_pool=lambda: pool.view(pool.shape[0], 1, -1), diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index f8c682f65a6f..74b4624cc787 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -20,6 +20,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from utils.util import getSMVersion from tensorrt_llm._torch.attention.backends import prims_ts @@ -145,7 +146,7 @@ def _contiguous_case(): q = torch.zeros((128, 256), dtype=torch.bfloat16) k = torch.zeros((512, 128), dtype=torch.bfloat16) v = torch.zeros_like(k) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cross=False, num_sparse_topk=0, helix_position_offsets=None, @@ -173,7 +174,7 @@ def _paged_metadata(): manager.dtype = torch.bfloat16 manager.num_pools = manager.num_local_layers = 1 manager.host_kv_cache_block_offsets = block_offsets - return SimpleNamespace( + return metadata_stub( is_cross=False, num_sparse_topk=0, helix_position_offsets=None, @@ -657,7 +658,7 @@ def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: dtype=torch.bfloat16, shape=(2, 512), ) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=object(), kv_cache_block_offsets=SimpleNamespace(device=query_device, shape=(1, 2, 4)), max_num_requests=2, @@ -773,7 +774,7 @@ def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: ), ), ) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cross=False, kv_cache_manager=None, seq_lens=torch.tensor([64], dtype=torch.int32), diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index d36c02ad903c..3e8ddc6575b7 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -27,6 +27,7 @@ import torch from tensorrt_llm._torch.attention.backends import trtllm as trtllm_backend +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, @@ -308,8 +309,7 @@ def _make_sparse_prediction_forward_metadata() -> TrtllmAttentionMetadata: metadata.use_paged_context_fmha = False metadata.cu_q_seqlens = None metadata.cu_kv_seqlens = None - metadata.enable_flash_mla = False - metadata.spec_bl_tree_first_sparse_mask_offset_kv = None + metadata.replace_ext(SpecDecodingMetadata(metadata)) metadata.kv_lens_cuda_runtime = torch.tensor([2], dtype=torch.int32) metadata.kv_lens_runtime = torch.tensor([2], dtype=torch.int32) metadata.prompt_lens_cuda_runtime = torch.tensor([2], dtype=torch.int32) diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mha.py b/tests/unittest/_torch/attention/sparse/test_sparse_mha.py index 9fa5f71cea56..1b15d2c8a947 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mha.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mha.py @@ -33,17 +33,17 @@ from utils.util import getSMVersion import tensorrt_llm +from tensorrt_llm._torch.attention.backends.features.spec_decoding import ( + SpecDecodingMetadata, + generate_spec_decoding_packed_mask, + generate_spec_decoding_position_offsets, +) from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionRuntimeFeatures, ) from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams -from tensorrt_llm._torch.attention.backends.trtllm import ( - TrtllmAttention, - TrtllmAttentionMetadata, - generate_spec_decoding_packed_mask, - generate_spec_decoding_position_offsets, -) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str @@ -275,28 +275,26 @@ def create_generation_inputs(scenario: MhaGenerationScenario) -> MhaGenerationIn runtime_features=AttentionRuntimeFeatures( has_speculative_draft_tokens=scenario.has_draft_tokens ), - is_spec_decoding_enabled=scenario.has_draft_tokens, - use_spec_decoding=scenario.has_draft_tokens, - is_spec_dec_tree=False, - max_total_draft_tokens=( - scenario.max_query_len - 1 if scenario.has_draft_tokens else None - ), + ) + spec = metadata.ext(SpecDecodingMetadata) + spec.is_spec_decoding_enabled = scenario.has_draft_tokens + spec.use_spec_decoding = scenario.has_draft_tokens + spec.max_total_draft_tokens = ( + scenario.max_query_len - 1 if scenario.has_draft_tokens else None ) if scenario.has_draft_tokens: draft_len = scenario.max_query_len - 1 - metadata.spec_decoding_position_offsets = generate_spec_decoding_position_offsets( - scenario.batch_size, draft_len - ) - metadata.spec_decoding_packed_mask = generate_spec_decoding_packed_mask( + spec.position_offsets = generate_spec_decoding_position_offsets( scenario.batch_size, draft_len ) - metadata.spec_decoding_generation_lengths = torch.tensor( + spec.packed_mask = generate_spec_decoding_packed_mask(scenario.batch_size, draft_len) + spec.generation_lengths = torch.tensor( [scenario.query_len] * scenario.batch_size, dtype=torch.int32, device=device, ) - metadata.update_position_offsets_for_cpp(scenario.max_query_len) - metadata.spec_decoding_param_prepare_for_blackwell() + spec.update_position_offsets_for_cpp(scenario.max_query_len) + spec.prepare_for_blackwell() metadata.prepare() cleanup.pop_all() return MhaGenerationInputs( diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py index b020bf43c7bc..eca4cc43a5c4 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py @@ -1343,7 +1343,7 @@ def populate_gen_dsa_kv_cache(mla: MLA, AttentionCls: AttentionBackend, offset += cached_len # Single batched metadata for all generation cache population - cached_metadata = AttentionCls.Metadata( + cached_metadata = AttentionCls.create_metadata( seq_lens=torch.tensor(gen_cached_lens, dtype=torch.int), request_ids=gen_with_cache, max_num_requests=len(gen_with_cache), @@ -1355,7 +1355,7 @@ def populate_gen_dsa_kv_cache(mla: MLA, AttentionCls: AttentionBackend, num_cached_tokens_per_seq=[0] * len(gen_with_cache)), mapping=mapping, - sparse_attention_config=sparse_config, + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), ) cached_metadata.prepare() @@ -1539,7 +1539,7 @@ def populate_gen_deepseek_v4_kv_cache( f"({num_indexer_compressed} indexer)") # Step 2: Run TRT-LLM forward to populate actual KV cache - cached_metadata = AttentionCls.Metadata( + cached_metadata = AttentionCls.create_metadata( seq_lens=torch.tensor(gen_cached_lens, dtype=torch.int), request_ids=gen_with_cache, max_num_requests=len(gen_with_cache), @@ -1554,7 +1554,7 @@ def populate_gen_deepseek_v4_kv_cache( len(gen_with_cache) # No prior cache ), mapping=mapping, - sparse_attention_config=sparse_config, + sparse_metadata_params=sparse_config.to_sparse_metadata_params(), ) cached_metadata.prepare() _ = mla_forward_impl_with_deepseek_v4_wo_linear( @@ -2160,7 +2160,7 @@ def yarn_get_mscale(scale=1, mscale=1): device=device) sparse_metadata_params = sparse_config.to_sparse_metadata_params() - attn_metadata = AttentionCls.Metadata( + attn_metadata = AttentionCls.create_metadata( seq_lens=torch.tensor(batch_query_lens, dtype=torch.int), request_ids=batch_order, max_num_requests=batch_spec.batch_size, diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py b/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py index ff913bb1f8dd..a1ad22b21334 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_mqa_gqa.py @@ -39,6 +39,11 @@ from utils.util import getSMVersion import tensorrt_llm +from tensorrt_llm._torch.attention.backends.features.spec_decoding import ( + SpecDecodingMetadata, + generate_spec_decoding_packed_mask, + generate_spec_decoding_position_offsets, +) from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import run_msa_sparse_gqa from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, @@ -51,12 +56,7 @@ msa_package_available, ) from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams -from tensorrt_llm._torch.attention.backends.trtllm import ( - TrtllmAttention, - TrtllmAttentionMetadata, - generate_spec_decoding_packed_mask, - generate_spec_decoding_position_offsets, -) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._utils import str_dtype_to_binding, torch_dtype_to_str @@ -1599,24 +1599,20 @@ def _create_generation_inputs(s: GenerationScenario) -> _GenerationInputs: num_sparse_topk=num_sparse_topk, max_num_heads_per_kv=s.q_heads_per_kv_head, runtime_features=AttentionRuntimeFeatures(has_speculative_draft_tokens=s.has_draft_tokens), - is_spec_decoding_enabled=s.has_draft_tokens, - use_spec_decoding=s.has_draft_tokens, - is_spec_dec_tree=False, - max_total_draft_tokens=s.max_query_len - 1 if s.has_draft_tokens else None, ) + spec = metadata.ext(SpecDecodingMetadata) + spec.is_spec_decoding_enabled = s.has_draft_tokens + spec.use_spec_decoding = s.has_draft_tokens + spec.max_total_draft_tokens = s.max_query_len - 1 if s.has_draft_tokens else None if s.has_draft_tokens: draft_len = s.max_query_len - 1 - metadata.spec_decoding_position_offsets = generate_spec_decoding_position_offsets( - s.batch_size, draft_len - ) - metadata.spec_decoding_packed_mask = generate_spec_decoding_packed_mask( - s.batch_size, draft_len - ) - metadata.spec_decoding_generation_lengths = torch.tensor( + spec.position_offsets = generate_spec_decoding_position_offsets(s.batch_size, draft_len) + spec.packed_mask = generate_spec_decoding_packed_mask(s.batch_size, draft_len) + spec.generation_lengths = torch.tensor( [s.query_len] * s.batch_size, dtype=torch.int32, device=device ) - metadata.update_position_offsets_for_cpp(s.max_query_len) - metadata.spec_decoding_param_prepare_for_blackwell() + spec.update_position_offsets_for_cpp(s.max_query_len) + spec.prepare_for_blackwell() metadata.prepare() return _GenerationInputs( diff --git a/tests/unittest/_torch/attention/test_attention.py b/tests/unittest/_torch/attention/test_attention.py index 5335c64348cd..2061240e2324 100644 --- a/tests/unittest/_torch/attention/test_attention.py +++ b/tests/unittest/_torch/attention/test_attention.py @@ -237,7 +237,7 @@ def produce_outputs( use_kv_cache_manager_v2) try: - metadata = Attention.Metadata( + metadata = Attention.create_metadata( num_contexts=num_contexts if num_contexts is not None else s.batch_size, kv_cache_params=kv_cache_params, diff --git a/tests/unittest/_torch/attention/test_attention_failure_logging.py b/tests/unittest/_torch/attention/test_attention_failure_logging.py index 97b0f31bac49..432857077907 100644 --- a/tests/unittest/_torch/attention/test_attention_failure_logging.py +++ b/tests/unittest/_torch/attention/test_attention_failure_logging.py @@ -25,8 +25,10 @@ import pytest import torch +from fmha_test_utils import metadata_stub from tensorrt_llm._torch.attention.backends import utils +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, PredefinedAttentionMask, @@ -46,7 +48,7 @@ def error(self, *msg) -> None: def _fake_metadata(**overrides): - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=2, num_generations=1, num_tokens=2050, @@ -204,9 +206,7 @@ def forward(self, *args, **kwargs): metadata.workspace = torch.zeros(1024, dtype=torch.int8) metadata.cu_q_seqlens = torch.zeros(batch + 1, dtype=torch.int32) metadata.cu_kv_seqlens = torch.zeros(batch + 1, dtype=torch.int32) - metadata.enable_flash_mla = False - metadata.spec_bl_tree_first_sparse_mask_offset_kv = None - metadata.spec_decoding_bl_tree_mask = None + metadata.replace_ext(SpecDecodingMetadata(metadata)) metadata.max_context_q_len_override = None metadata.kv_cache_manager = SimpleNamespace( max_attention_window_vec=[2048, None], diff --git a/tests/unittest/_torch/attention/test_attention_mla.py b/tests/unittest/_torch/attention/test_attention_mla.py index 83133dc0c924..d09cc60347a8 100644 --- a/tests/unittest/_torch/attention/test_attention_mla.py +++ b/tests/unittest/_torch/attention/test_attention_mla.py @@ -12,6 +12,7 @@ import tensorrt_llm from tensorrt_llm._torch.attention import mla as mla_module +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata from tensorrt_llm._torch.attention.backends.interface import ( AttentionInputType, MLAParams, PositionalEmbeddingParams, RopeParams) from tensorrt_llm._torch.attention.backends.trtllm import \ @@ -1065,7 +1066,7 @@ def yarn_get_mscale(scale=1, mscale=1): else: kv_cache_manager.impl.add_sequence_batch( [(req_id, ctx_len, beam_width)], [req]) - attn_metadata = AttentionCls.Metadata( + attn_metadata = AttentionCls.create_metadata( seq_lens=torch.tensor(context_sequence_lengths, dtype=torch.int), request_ids=list(range(len(context_sequence_lengths))), max_num_requests=len(context_sequence_lengths), @@ -1114,10 +1115,11 @@ def yarn_get_mscale(scale=1, mscale=1): ), mapping=mapping, ) - if backend_name == "TRTLLM": - gen_metadata_kwargs["enable_flash_mla"] = ( - cuda_capability == (9, 0)) - attn_metadata = AttentionCls.Metadata(**gen_metadata_kwargs) + if backend_name == "TRTLLM" and cuda_capability == (9, 0): + gen_metadata_kwargs["extension_types"] = ( + FlashMlaMetadata, ) + attn_metadata = AttentionCls.create_metadata( + **gen_metadata_kwargs) attn_metadata.prepare() for layer_idx in range(num_layers): print( @@ -1358,21 +1360,18 @@ def yarn_get_mscale(scale=1, mscale=1): NUM_HEADS = 128 -class _Prefixes: - """The state `mla_prepare_*` touches, with the real methods bound to it. +class _Prefixes(TrtllmAttentionMetadata): + """The state `mla_prepare_*` touches, running the real metadata methods. Deliberately not a full `TrtllmAttentionMetadata`: that needs a KV cache - manager and a resource pool, none of which these methods read. Binding the - production functions keeps this a test of the shipped code rather than of a - reimplementation. + manager and a resource pool, none of which these methods read. Inheriting + the production methods keeps this a test of the shipped code rather than of + a reimplementation; the length properties become plain attributes. """ - # Bind the unbound functions so any edit to them is exercised here. - mla_prepare_scheduler_buffers = TrtllmAttentionMetadata.mla_prepare_scheduler_buffers - mla_prepare_ctx_cu_seqlens = TrtllmAttentionMetadata.mla_prepare_ctx_cu_seqlens - _invalidate_mla_scheduler_buffers = TrtllmAttentionMetadata._invalidate_mla_scheduler_buffers - on_update_kv_lens = TrtllmAttentionMetadata.on_update_kv_lens - update_for_spec_dec = TrtllmAttentionMetadata.update_for_spec_dec + num_contexts = 0 + num_seqs = 0 + seq_lens_cuda = None def __init__(self, q_lens, kv_lens, num_contexts=0): device = torch.device("cuda") @@ -1392,8 +1391,6 @@ def __init__(self, q_lens, kv_lens, num_contexts=0): self.mla_ctx_cu_q_seqlens = torch.zeros(size, dtype=torch.int32, device=device) - # `enable_flash_mla` is the only other attribute the two hooks read. - self.enable_flash_mla = False self._invalidate_mla_scheduler_buffers() def rebuild(self): diff --git a/tests/unittest/_torch/attention/test_attention_no_cache.py b/tests/unittest/_torch/attention/test_attention_no_cache.py index 7ba37ef476e2..f9ff7438bc66 100644 --- a/tests/unittest/_torch/attention/test_attention_no_cache.py +++ b/tests/unittest/_torch/attention/test_attention_no_cache.py @@ -256,7 +256,7 @@ def _run_test_for_backend(backend_name, num_heads, num_kv_heads, num_layers, # NOTE: set up metadata, refer to tensorrt_llm/_torch/pyexecutor/model_engine.py # all layers share the same metadata - attn_metadata = AttentionCls.Metadata( + attn_metadata = AttentionCls.create_metadata( max_num_requests=len(context_sequence_lengths), max_num_tokens=8192, kv_cache_manager=None, diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index 6abee3347a2f..544e025c8f50 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -20,7 +20,8 @@ 1. Every C++ parameter name appears at the call site (and nothing extra). 2. Every call-site kwarg sourced as ``root.attr[.attr...]`` resolves on - exactly one of ``attn`` / ``metadata`` / ``forward_args``, and its + exactly one of ``attn`` / ``metadata`` / ``forward_args`` or a metadata + extension the call reads (``helix`` / ``flash_mla`` / ``spec``), and its declared C++ type matches the source attribute's Python type at a coarse-category level (tensor / int / bool / float / list-of-X). 3. Every dataclass field reachable from ``AttentionForwardArgs`` (including @@ -46,6 +47,8 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata, HelixMetadata +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha.fallback import ( _THOP_EXCLUDED_FIELDS, _THOP_LITERALS, @@ -63,20 +66,31 @@ "attn": TrtllmAttention, "metadata": TrtllmAttentionMetadata, "forward_args": AttentionForwardArgs, + "helix": HelixMetadata, + "flash_mla": FlashMlaMetadata, + "spec": SpecDecodingMetadata, } _THOP_KWARG_SOURCE_ALIASES: dict[str, tuple[str, tuple[str, ...]]] = { "beam_width": ("metadata", ("effective_beam_width",)), "context_lengths": ("metadata", ("prompt_lens_cuda_runtime",)), + "flash_mla_num_splits": ("flash_mla", ("num_splits",)), + "flash_mla_tile_scheduler_metadata": ("flash_mla", ("tile_scheduler_metadata",)), + "helix_is_inactive_rank": ("helix", ("is_inactive_rank",)), + "helix_position_offsets": ("helix", ("position_offsets",)), "head_size": ("attn", ("head_dim",)), "host_context_lengths": ("metadata", ("prompt_lens_cpu_runtime",)), "host_past_key_value_lengths": ("metadata", ("kv_lens_runtime",)), "host_request_types": ("metadata", ("host_request_types_runtime",)), "sequence_length": ("metadata", ("kv_lens_cuda_runtime",)), - "spec_decoding_target_max_draft_tokens": ( - "metadata", - ("max_total_draft_tokens",), - ), + "force_prepare_spec_dec_tree_mask": ("spec", ("force_prepare_tree_mask",)), + "spec_bl_tree_first_sparse_mask_offset_kv": ("spec", ("bl_tree_first_sparse_mask_offset_kv",)), + "spec_decoding_bl_tree_mask": ("spec", ("bl_tree_mask",)), + "spec_decoding_bl_tree_mask_offset": ("spec", ("bl_tree_mask_offset",)), + "spec_decoding_generation_lengths": ("spec", ("generation_lengths",)), + "spec_decoding_packed_mask": ("spec", ("packed_mask",)), + "spec_decoding_position_offsets_for_cpp": ("spec", ("position_offsets_for_cpp",)), + "spec_decoding_target_max_draft_tokens": ("spec", ("max_total_draft_tokens",)), "skip_softmax_threshold_scale_factor_decode": ( "forward_args", ( @@ -685,10 +699,10 @@ def test_no_sequence_kwargs_at_thop_attention_boundary(): def test_fallback_support_matches_thop_kv_update_contract(is_cross, update_kv_cache, expected): """Do not dispatch requests that the native attention op rejects.""" fmha = object.__new__(FallbackFmha) - # ``num_contexts=0`` skips the paged-context kernel gate and - # ``helix_position_offsets=None`` the Helix verify-group check; both are - # what the real metadata carries for a generation batch off the helix path. - metadata = SimpleNamespace(is_cross=is_cross, helix_position_offsets=None, num_contexts=0) + # ``num_contexts=0`` skips the paged-context kernel gate and no attached + # Helix extension the Helix verify-group check; both are what the real + # metadata carries for a generation batch off the helix path. + metadata = SimpleNamespace(is_cross=is_cross, num_contexts=0, find_ext=lambda cls: None) forward_args = AttentionForwardArgs(update_kv_cache=update_kv_cache) assert fmha.is_supported(None, None, None, metadata, forward_args) is expected @@ -697,7 +711,7 @@ def test_fallback_support_matches_thop_kv_update_contract(is_cross, update_kv_ca def test_fallback_rejects_raw_fp8_input(): """Do not dispatch raw FP8 QKV to the native attention op.""" fmha = object.__new__(FallbackFmha) - metadata = SimpleNamespace(is_cross=False, helix_position_offsets=None, num_contexts=0) + metadata = SimpleNamespace(is_cross=False, num_contexts=0, find_ext=lambda cls: None) forward_args = AttentionForwardArgs(update_kv_cache=True) q = torch.empty((1, 128), dtype=torch.float8_e4m3fn) diff --git a/tests/unittest/_torch/attention/test_combined_fmha.py b/tests/unittest/_torch/attention/test_combined_fmha.py index 87f294744b88..9b01454bf876 100644 --- a/tests/unittest/_torch/attention/test_combined_fmha.py +++ b/tests/unittest/_torch/attention/test_combined_fmha.py @@ -16,7 +16,7 @@ from types import SimpleNamespace import torch -from fmha_test_utils import FakeAttention, FakePhasedFmha +from fmha_test_utils import FakeAttention, FakePhasedFmha, metadata_stub from tensorrt_llm._torch.attention.backends.fmha.combined import CombinedFmha from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase @@ -56,7 +56,7 @@ def test_combined_fmha_delegates_phases_and_prepares_max_workspace() -> None: generation_fmha.supports_workspace_reclamation = generation_support combined_fmha.set_fmha_impls(context_fmha, generation_fmha) assert combined_fmha.supports_workspace_reclamation is expected - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_block_offsets=object(), effective_workspace=torch.empty(0, dtype=torch.uint8), num_contexts=2, diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla.py b/tests/unittest/_torch/attention/test_cute_dsl_mla.py index c3bee7550fb3..ceba2bc260bc 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla.py @@ -6,6 +6,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import CuteDslMlaFmha from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE @@ -160,7 +161,7 @@ def _select_mla_generation_fmha( qk_rope_head_dim=64, layer_idx=0, ) - meta = SimpleNamespace( + meta = metadata_stub( num_contexts=0, num_generations=batch_size, beam_width=1, diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py b/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py index 566fe85dd9e4..3070122d5a7f 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla_split_plan.py @@ -3,10 +3,10 @@ """Per-request split-KV plan of the CuTe DSL MLA decode: policy, kernel vs mirror, kernel output.""" import math -from types import SimpleNamespace import pytest import torch +from fmha_test_utils import metadata_stub from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from tensorrt_llm._utils import get_sm_version @@ -88,7 +88,7 @@ def test_split_kv_plan_is_launched_once_per_step(monkeypatch) -> None: monkeypatch.setattr( mod, "launch_mla_split_kv_plan", lambda *args, **kwargs: launches.append(kwargs) ) - meta = SimpleNamespace(max_num_requests=16, _cute_dsl_mla_split_plan_key=None) + meta = metadata_stub(max_num_requests=16, _cute_dsl_mla_split_plan_key=None) cache_seqs = torch.full((8,), 14208, dtype=torch.int32, device="cuda") first = mod.CuteDslMlaFmha._split_kv_plan(meta, cache_seqs, 8, 8, 96) assert first.numel() == 8 and meta._cute_dsl_mla_split_plan.numel() == 16 diff --git a/tests/unittest/_torch/attention/test_flashinfer_attention.py b/tests/unittest/_torch/attention/test_flashinfer_attention.py index 94e154517ebf..4fa0a8f014fb 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_attention.py +++ b/tests/unittest/_torch/attention/test_flashinfer_attention.py @@ -163,7 +163,6 @@ def get_layer_page_index_scale(self, layer_idx: int) -> int: kv_cache_manager=manager, request_ids=[7], workspace_buffer=torch.empty(1, dtype=torch.uint8, device="cuda"), - mamba_metadata=False, ) self.assertEqual(metadata._primary_kv_layer_idx, 20) diff --git a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py index fac6b0eb5359..7135de335c48 100644 --- a/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py +++ b/tests/unittest/_torch/attention/test_flashinfer_trtllm_gen_fmha.py @@ -17,7 +17,7 @@ import pytest import torch -from fmha_test_utils import FakeAttention +from fmha_test_utils import FakeAttention, metadata_stub from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( @@ -84,7 +84,7 @@ def test_small_context_fallback_preserves_mixed_batch_generation( manager.dtype = None manager.trtllm_gen_extra_tokens_per_block = frozenset({128}) manager.impl = SimpleNamespace(get_page_index_upper_bound=lambda *args: 32) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=num_contexts, num_generations=num_generations, helix_position_offsets=None, @@ -168,7 +168,7 @@ def test_flashinfer_quantized_kv_context_avoids_fp16_bf16_fallback( q_hidden_size = attn.num_heads * attn.head_dim qkv_hidden_size = q_hidden_size + 2 * attn.num_kv_heads * attn.head_dim q = torch.empty((1, qkv_hidden_size), dtype=dtype) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=1, helix_position_offsets=None, num_sparse_topk=0, @@ -276,7 +276,7 @@ def test_sm107_serves_only_extra_page_sizes( q_hidden_size = attn.num_heads * attn.head_dim qkv_hidden_size = q_hidden_size + 2 * attn.num_kv_heads * attn.head_dim q = torch.empty((num_tokens, qkv_hidden_size), dtype=torch.bfloat16) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=1 if is_context else 0, num_generations=0 if is_context else 1, helix_position_offsets=None, diff --git a/tests/unittest/_torch/attention/test_fmha_interface.py b/tests/unittest/_torch/attention/test_fmha_interface.py index 6c4f74997307..b1cd898e02f1 100644 --- a/tests/unittest/_torch/attention/test_fmha_interface.py +++ b/tests/unittest/_torch/attention/test_fmha_interface.py @@ -21,6 +21,7 @@ import torch from fmha_test_utils import FakeAttention +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha, FmhaPhase from tensorrt_llm._torch.attention.backends.fmha.registry import FMHA_LIBS @@ -89,9 +90,9 @@ def test_mla_forward_clears_fused_qkv_before_fmha_selection( metadata.token_major_gen_view.return_value = None metadata.presented_token_major.return_value = nullcontext() metadata.is_cross = False - metadata.enable_flash_mla = False - metadata.spec_bl_tree_first_sparse_mask_offset_kv = None - metadata.spec_decoding_bl_tree_mask = None + extensions = {SpecDecodingMetadata: SpecDecodingMetadata(metadata)} + metadata.find_ext.side_effect = extensions.get + metadata.ext.side_effect = extensions.__getitem__ metadata.max_context_q_len_override = None metadata.kv_cache_manager = None lengths = torch.ones(1, dtype=torch.int32) diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 7cfb59eee093..aa66dc79cfbd 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -20,8 +20,9 @@ import pytest import torch -from fmha_test_utils import FakeAttention, FakeFmha, FakePhasedFmha +from fmha_test_utils import FakeAttention, FakeFmha, FakePhasedFmha, metadata_stub +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha import manager as fmha_manager from tensorrt_llm._torch.attention.backends.fmha.combined import CombinedFmha from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase @@ -56,7 +57,7 @@ def _make_metadata( use_spec_decoding: bool = False, num_sparse_topk: int = 0, ) -> SimpleNamespace: - return SimpleNamespace( + return metadata_stub( num_contexts=num_contexts, num_generations=num_generations, num_ctx_tokens=num_ctx_tokens, @@ -625,7 +626,9 @@ def test_fmha_cache_separates_speculative_decoding(speculative_first: bool) -> N attn, "regular-only", events, - request_support_predicate=lambda _q, metadata: not metadata.use_spec_decoding, + request_support_predicate=lambda _q, metadata: not metadata.ext( + SpecDecodingMetadata + ).use_spec_decoding, ) fallback = FakeFmha(attn, "fallback", events) manager.fmha_libs = [regular_only, fallback] diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 31c84af61276..733596129556 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -5,6 +5,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from tensorrt_llm._torch.attention.backends.fmha import ( flashinfer_trtllm_gen as flashinfer_trtllm_gen_module, @@ -39,7 +40,7 @@ def get_page_index_upper_bound(local_layer_idx: int, role: object) -> int: fmha._v1_total_num_blocks_cache = None attn = SimpleNamespace(local_layer_idx=7) fmha._attn_ref = lambda: attn - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=manager, host_kv_cache_pool_mapping=None, ) @@ -74,7 +75,7 @@ def get_primary_pool_data(local_layer_idx: int) -> SimpleNamespace: fmha._v1_total_num_blocks_cache = None attn = SimpleNamespace(local_layer_idx=4) fmha._attn_ref = lambda: attn - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=manager, host_kv_cache_pool_mapping=pool_mapping, ) @@ -90,7 +91,7 @@ def test_phased_fmha_rejects_unknown_kv_cache_manager_type() -> None: fmha.kv_factor = 2 fmha._v1_total_num_blocks_cache = None fmha._attn_ref = lambda: SimpleNamespace(local_layer_idx=0) - metadata = SimpleNamespace(kv_cache_manager=SimpleNamespace()) + metadata = metadata_stub(kv_cache_manager=SimpleNamespace()) with pytest.raises(TypeError, match="Unsupported KV cache manager: SimpleNamespace"): fmha._get_total_num_blocks(metadata) @@ -141,7 +142,7 @@ def check_counter_size_args( attn=SimpleNamespace(num_heads=num_heads), _multi_processor_count=sm_count, ) - metadata = SimpleNamespace( + metadata = metadata_stub( max_num_requests=max_num_requests, beam_width=beam_width, max_num_sequences=max_num_sequences, @@ -213,7 +214,7 @@ def generation_preprocess(*args: object) -> tuple[object, ...]: rotary_cos_sin=None, rope_params=SimpleNamespace(dim=0, theta=1.0, scale_type=0, scale=1.0, max_positions=1), ) - metadata = SimpleNamespace( + metadata = metadata_stub( beam_width=2, kv_cache_block_offsets=torch.empty(0), host_kv_cache_pool_pointers=torch.empty(0), @@ -280,7 +281,7 @@ def _cute_dsl_mla_helix_support( qk_rope_head_dim=64, layer_idx=0, ) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=0, num_generations=batch_size, beam_width=1, @@ -326,7 +327,7 @@ def test_cute_dsl_mla_rejects_h96_mixed_batch(monkeypatch: pytest.MonkeyPatch) - qk_rope_head_dim=64, layer_idx=0, ) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=1, num_generations=1, beam_width=1, diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 08ab0591c504..881d082138d2 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -16,6 +16,8 @@ import tensorrt_llm._torch.attention.backends.fp4_mla.metadata as fp4_mla_metadata import tensorrt_llm._torch.attention.backends.fp4_mla.v_cache as fp4_mla_v_cache import tensorrt_llm._torch.attention.backends.trtllm as trtllm_backend +from tensorrt_llm._torch.attention.backends.features import HelixMetadata, find_helix +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha.fp4_mla import Fp4MlaFmha from tensorrt_llm._torch.attention.backends.fp4_mla import ( FP4_BLOCK_SIZE, @@ -39,7 +41,7 @@ from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_context import ( _build_fp8_mla_context_metadata, ) -from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState +from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState, find_fp4_mla from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, @@ -81,8 +83,24 @@ def _is_cutedsl_unavailable() -> bool: ) +def _attach_ext(metadata: SimpleNamespace, cls: type, extension: object) -> object: + """Serve ``extension`` as the stub's ``cls`` extension through ``find_ext``.""" + if not hasattr(metadata, "stub_extensions"): + metadata.stub_extensions = {} + metadata.find_ext = metadata.stub_extensions.get + metadata.stub_extensions[cls] = extension + return extension + + +def _fp4_mla_metadata(fp4_mla_state: object, **fields) -> SimpleNamespace: + """A metadata stub carrying ``fp4_mla_state`` as its FP4 MLA extension.""" + metadata = SimpleNamespace(**fields) + _attach_ext(metadata, Fp4MlaState, fp4_mla_state) + return metadata + + def test_fp4_mla_request_validation_uses_sparse_runtime_params() -> None: - metadata = SimpleNamespace( + metadata = _fp4_mla_metadata( num_sparse_topk=0, kv_cache_manager=SimpleNamespace(dtype=_DataType.NVFP4, kv_factor=1), beam_width=1, @@ -111,16 +129,33 @@ def test_fp4_mla_request_validation_uses_sparse_runtime_params() -> None: ) +def _attach_helix(metadata: SimpleNamespace, **fields) -> SimpleNamespace: + """Expose ``fields`` as the stub's helix extension through ``find_ext``.""" + helix = SimpleNamespace( + position_offsets=None, + is_inactive_rank=None, + local_slots=None, + kv_bounds=None, + spec_tokens_valid=False, + ) + for name, value in fields.items(): + setattr(helix, name, value) + return _attach_ext(metadata, HelixMetadata, helix) + + def _fp4_mla_validation_metadata(**overrides) -> SimpleNamespace: + helix_position_offsets = overrides.pop("helix_position_offsets", None) fields = dict( num_sparse_topk=0, kv_cache_manager=SimpleNamespace(dtype=_DataType.NVFP4, kv_factor=1), beam_width=1, fp4_mla_state=Fp4MlaState(hp_pool=object(), v_scale_pool=object()), - helix_position_offsets=None, ) fields.update(overrides) - return SimpleNamespace(**fields) + metadata = _fp4_mla_metadata(**fields) + if helix_position_offsets is not None: + _attach_helix(metadata, position_offsets=helix_position_offsets) + return metadata def test_fp4_mla_request_validation_still_requires_nvfp4_kv() -> None: @@ -235,10 +270,8 @@ def test_fp4_mla_generation_forwards_helix_softmax_stats(monkeypatch) -> None: output = torch.empty(num_tokens, num_heads * kv_lora_rank, dtype=torch.bfloat16) softmax_stats = torch.empty(num_tokens, num_heads, 2, dtype=torch.float32) expected_stats = torch.tensor([1.25, 2.5], dtype=torch.float32).expand_as(softmax_stats) - metadata = SimpleNamespace( + metadata = _fp4_mla_metadata( num_generations=num_tokens, - helix_position_offsets=torch.tensor([384, 769], dtype=torch.int32), - _helix_spec_tokens_valid=False, fp4_mla_state=Fp4MlaState( generation_cache_scattered=True, prequantized_q=torch.empty(0), @@ -246,6 +279,7 @@ def test_fp4_mla_generation_forwards_helix_softmax_stats(monkeypatch) -> None: q_batch_capacity=num_tokens, ), ) + _attach_helix(metadata, position_offsets=torch.tensor([384, 769], dtype=torch.int32)) attn = SimpleNamespace( get_fp4_mla_local_layer_idx=lambda observed: 5, kv_lora_rank=kv_lora_rank, @@ -285,7 +319,7 @@ def record_decode(*args, **kwargs) -> None: assert kwargs["softmax_stats_tensor"] is softmax_stats assert softmax_stats.is_contiguous() torch.testing.assert_close(softmax_stats, expected_stats) - state = metadata.fp4_mla_state + state = find_fp4_mla(metadata) assert not state.generation_cache_scattered assert state.prequantized_q is None assert state.prequantized_q_sf is None @@ -311,13 +345,13 @@ def test_fp4_mla_rope_generation_forwards_helix_cache_ownership( ) helix_position_offsets = torch.tensor([384, 769], dtype=torch.int32) helix_is_inactive_rank = torch.tensor([True, False]) - metadata = SimpleNamespace( - num_generations=num_tokens, - num_ctx_tokens=0, - helix_position_offsets=helix_position_offsets, - helix_is_inactive_rank=helix_is_inactive_rank, - _helix_spec_tokens_valid=False, - ) + metadata = object.__new__(trtllm_backend.TrtllmAttentionMetadata) + metadata._num_generations = num_tokens + metadata._num_ctx_tokens = 0 + helix = HelixMetadata(metadata) + helix.position_offsets = helix_position_offsets + helix.is_inactive_rank = helix_is_inactive_rank + metadata.replace_ext(helix) attn = SimpleNamespace( kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, @@ -355,20 +389,20 @@ def run_rope_generation() -> dict: assert helix_kwargs["helix_position_offsets"] is helix_position_offsets assert helix_kwargs["helix_is_inactive_rank"] is helix_is_inactive_rank - metadata.helix_position_offsets = None - metadata.helix_is_inactive_rank = None + helix.position_offsets = None + helix.is_inactive_rank = None non_helix_kwargs = run_rope_generation() assert non_helix_kwargs["helix_position_offsets"] is None assert non_helix_kwargs["helix_is_inactive_rank"] is None - metadata.num_generations = 1 - metadata.helix_position_offsets = helix_position_offsets - metadata.helix_is_inactive_rank = torch.tensor([False]) - metadata.helix_local_slots = torch.tensor([384, 385], dtype=torch.int32) - metadata._helix_spec_tokens_valid = True + metadata._num_generations = 1 + helix.position_offsets = helix_position_offsets + helix.is_inactive_rank = torch.tensor([False]) + helix.local_slots = torch.tensor([384, 385], dtype=torch.int32) + helix.spec_tokens_valid = True spec_helix_kwargs = run_rope_generation() assert spec_helix_kwargs["helix_position_offsets"] is helix_position_offsets - assert spec_helix_kwargs["helix_is_inactive_rank"] is metadata.helix_is_inactive_rank + assert spec_helix_kwargs["helix_is_inactive_rank"] is helix.is_inactive_rank assert len(calls) == 3 @@ -421,8 +455,8 @@ def mla_rope_generation(self, *args, **kwargs) -> None: metadata = SimpleNamespace( num_seqs=num_tokens, num_ctx_tokens=0, - _helix_spec_tokens_valid=False, ) + _attach_helix(metadata) with pytest.raises(ReachedFusedCacheUpdate): MLA.forward_absorption_generation( @@ -450,15 +484,14 @@ def test_fp8_mla_context_partition_metadata_separates_q_and_kv_lengths() -> None prompt_lens_cpu = prompt_lens_cuda.clone() kv_lens_cuda = torch.tensor([128, 64], dtype=torch.int32) kv_lens_cpu = kv_lens_cuda.clone() - meta = SimpleNamespace( - num_contexts=2, - num_ctx_tokens=80, - prompt_lens_cuda_runtime=prompt_lens_cuda, - prompt_lens_cpu_runtime=prompt_lens_cpu, - host_request_types_runtime=torch.tensor([0, 0, 1], dtype=torch.int32), - fp4_mla_state=Fp4MlaState( - positions=torch.arange(80, dtype=torch.int32), fp8_context_state=object() - ), + meta = object.__new__(trtllm_backend.TrtllmAttentionMetadata) + meta._num_contexts = 2 + meta._num_ctx_tokens = 80 + meta.prompt_lens_cuda_runtime = prompt_lens_cuda + meta.prompt_lens_cpu_runtime = prompt_lens_cpu + meta.host_request_types_runtime = torch.tensor([0, 0, 1], dtype=torch.int32) + meta.replace_ext( + Fp4MlaState(meta, positions=torch.arange(80, dtype=torch.int32), fp8_context_state=object()) ) scratch = SimpleNamespace( cache_manager_view=object(), @@ -479,8 +512,9 @@ def test_fp8_mla_context_partition_metadata_separates_q_and_kv_lengths() -> None assert fp8_meta.kv_lens_cuda_runtime is kv_lens_cuda assert fp8_meta.kv_lens_runtime is kv_lens_cpu assert fp8_meta.host_total_kv_lens is scratch.host_total_kv_lens - assert fp8_meta.helix_position_offsets.shape == (80,) - assert fp8_meta.fp4_mla_state is None + assert fp8_meta.find_ext(HelixMetadata).position_offsets.shape == (80,) + assert meta.find_ext(HelixMetadata) is None + assert fp8_meta.find_ext(Fp4MlaState) is None def _reset_triton_allocator() -> None: @@ -493,7 +527,7 @@ def _null_allocator(size: int, alignment: int, stream): def test_fp4_mla_generation_hp_page_ids_skip_mixed_batch_context_rows() -> None: - metadata = SimpleNamespace( + metadata = _fp4_mla_metadata( num_contexts=2, num_seqs=5, fp4_mla_state=Fp4MlaState( @@ -513,7 +547,7 @@ def test_fp4_mla_generation_hp_page_ids_skip_mixed_batch_context_rows() -> None: def test_fp4_mla_cuda_graph_generation_lengths_records_capture_once(monkeypatch) -> None: corrected_kv_lens = torch.empty(2, dtype=torch.int32, device="cuda") generation_lens = torch.empty(2, dtype=torch.int32, device="cuda") - metadata = SimpleNamespace( + metadata = _fp4_mla_metadata( num_contexts=1, num_seqs=3, kv_lens_cuda_runtime=torch.tensor([9, 17, 25], dtype=torch.int32, device="cuda"), @@ -553,7 +587,7 @@ def populate_generation_lengths(*args, **kwargs) -> None: assert args[2].data_ptr() == corrected_kv_lens.data_ptr() assert args[3].data_ptr() == generation_lens.data_ptr() assert kwargs == {"num_gen_tokens": 4, "num_gen": 2} - assert metadata.fp4_mla_state.generation_lengths_capture_recorded + assert find_fp4_mla(metadata).generation_lengths_capture_recorded assert first[0].data_ptr() == second[0].data_ptr() == corrected_kv_lens.data_ptr() assert first[1].data_ptr() == second[1].data_ptr() == generation_lens.data_ptr() @@ -580,36 +614,37 @@ def test_recompute_helix_spec_buffers_matches_elementwise_reference( metadata._num_contexts = 1 metadata._num_generations = num_seqs untouched = 77 + helix = HelixMetadata(metadata) # The helix buffers are generation-relative; kv_lens_cuda is batch-indexed. - metadata.helix_position_offsets = torch.cat( + helix.position_offsets = torch.cat( [positions, torch.full((2,), untouched, dtype=torch.int32, device="cuda")] ) - metadata.helix_local_slots = torch.full_like(metadata.helix_position_offsets, untouched) - metadata.helix_kv_bounds = torch.full_like(metadata.helix_position_offsets, untouched) + helix.local_slots = torch.full_like(helix.position_offsets, untouched) + helix.kv_bounds = torch.full_like(helix.position_offsets, untouched) metadata.kv_lens_cuda = torch.full((num_seqs + 3,), untouched, dtype=torch.int32, device="cuda") - metadata.recompute_helix_spec_buffers(num_tokens, tokens_per_seq) + helix.recompute_spec_buffers(num_tokens, tokens_per_seq) torch.cuda.synchronize() - # Reference: the per-op formulation (helix_local_len_vec mirrors the cache manager). + # Reference: the per-op formulation (local_len_vec mirrors the cache manager). owner = torch.div(positions, tokens_per_block, rounding_mode="floor") % cp_size expected_slots = torch.where( - owner == cp_rank, metadata.helix_local_len_vec(positions), torch.full_like(positions, -1) + owner == cp_rank, helix.local_len_vec(positions), torch.full_like(positions, -1) ) - expected_bounds = metadata.helix_local_len_vec(positions + 1) + expected_bounds = helix.local_len_vec(positions + 1) gen = slice(0, num_tokens) - torch.testing.assert_close(metadata.helix_local_slots[gen], expected_slots) - torch.testing.assert_close(metadata.helix_kv_bounds[gen], expected_bounds) + torch.testing.assert_close(helix.local_slots[gen], expected_slots) + torch.testing.assert_close(helix.kv_bounds[gen], expected_bounds) torch.testing.assert_close( metadata.kv_lens_cuda[1 : 1 + num_seqs], expected_bounds.view(num_seqs, tokens_per_seq)[:, -1], ) - assert (metadata.helix_local_slots[num_tokens:] == untouched).all() - assert (metadata.helix_kv_bounds[num_tokens:] == untouched).all() + assert (helix.local_slots[num_tokens:] == untouched).all() + assert (helix.kv_bounds[num_tokens:] == untouched).all() assert metadata.kv_lens_cuda[0].item() == untouched assert (metadata.kv_lens_cuda[1 + num_seqs :] == untouched).all() # Any group of the sequence that this rank does not own writes no slot. - assert ((metadata.helix_local_slots[gen] == -1) == (owner != cp_rank)).all() + assert ((helix.local_slots[gen] == -1) == (owner != cp_rank)).all() @pytest.mark.parametrize("layer_offset", [0, 1, 3]) @@ -1282,7 +1317,7 @@ def _build_multi_seq_metadata( kv_global_scale = torch.tensor([FP4_MLA_KV_GLOBAL_SCALE], dtype=torch.float32, device=device) q_global_scale = torch.tensor([FP4_MLA_Q_GLOBAL_SCALE], dtype=torch.float32, device=device) - return SimpleNamespace( + return _fp4_mla_metadata( kv_cache_manager=kv_cache_manager, page_size=page_size, num_contexts=num_seqs, @@ -1323,20 +1358,19 @@ def _build_multi_seq_metadata( def _materialize_reference_cache_storage(metadata, layer_idx: int, head_dim: int) -> torch.Tensor: + fp4_mla_state = find_fp4_mla(metadata) kv_cache, sf_cache = metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) sf_cache = sf_cache.view(torch.float8_e4m3fn) storage_head_dim = kv_cache.shape[-1] * 2 static_global_scale = float(_get_fp4_mla_global_scale(metadata, kv_cache.device).item()) num_generation_sequences = metadata.num_seqs - metadata.num_contexts - page_stride = metadata.fp4_mla_state.page_table_stride - page_rows = metadata.fp4_mla_state.paged_kv_indices.view( + page_stride = fp4_mla_state.page_table_stride + page_rows = fp4_mla_state.paged_kv_indices.view( metadata.num_seqs, page_stride, ) generation_page_rows = page_rows[metadata.num_contexts : metadata.num_seqs] - indptr = ( - metadata.fp4_mla_state.paged_kv_indptr_decode[: num_generation_sequences + 1].cpu().tolist() - ) + indptr = fp4_mla_state.paged_kv_indptr_decode[: num_generation_sequences + 1].cpu().tolist() expected_indptr = [seq_idx * page_stride for seq_idx in range(num_generation_sequences + 1)] assert indptr == expected_indptr assert generation_page_rows.shape == (num_generation_sequences, page_stride) @@ -1391,13 +1425,14 @@ def _materialize_reference_cache_tokens( *, include_k_residual: bool = False, ) -> torch.Tensor: + fp4_mla_state = find_fp4_mla(metadata) kv_cache, sf_cache = metadata.kv_cache_manager.get_fp4_mla_cache_buffers(layer_idx) sf_cache = sf_cache.view(torch.float8_e4m3fn) storage_head_dim = kv_cache.shape[-1] * 2 static_global_scale = float(_get_fp4_mla_global_scale(metadata, kv_cache.device).item()) - page_rows = metadata.fp4_mla_state.paged_kv_indices.view( + page_rows = fp4_mla_state.paged_kv_indices.view( metadata.num_seqs, - metadata.fp4_mla_state.page_table_stride, + fp4_mla_state.page_table_stride, ) dequantized_pages = {} tokens = [] @@ -1456,13 +1491,13 @@ def test_fp4_mla_helix_scatter_uses_global_rope_and_local_cache_position( seed=31, ) try: - state = metadata.fp4_mla_state + state = find_fp4_mla(metadata) head_dim = kv_lora_rank + qk_rope_head_dim state.positions = torch.tensor([global_position], dtype=torch.int32, device="cuda") - metadata._helix_spec_tokens_valid = False - metadata.helix_position_offsets = state.positions - metadata.helix_is_inactive_rank = torch.tensor( - [inactive_rank], dtype=torch.bool, device="cuda" + helix = _attach_helix( + metadata, + position_offsets=state.positions, + is_inactive_rank=torch.tensor([inactive_rank], dtype=torch.bool, device="cuda"), ) kv_cache, sf_cache = kv_cache_manager.get_fp4_mla_cache_buffers(0) @@ -1508,8 +1543,8 @@ def test_fp4_mla_helix_scatter_uses_global_rope_and_local_cache_position( q_pe=q_pe, q_rope_out=q_rope_out, q_quant_input=q_quant_input, - helix_position_offsets=metadata.helix_position_offsets, - helix_is_inactive_rank=metadata.helix_is_inactive_rank, + helix_position_offsets=helix.position_offsets, + helix_is_inactive_rank=helix.is_inactive_rank, ) torch.cuda.synchronize() @@ -1565,7 +1600,7 @@ def rotate_pairs(values: torch.Tensor) -> torch.Tensor: metadata, layer_idx=0, batch_indices=batch_indices, - positions=metadata.helix_position_offsets, + positions=helix.position_offsets, head_dim=head_dim, ) assert not torch.count_nonzero(global_slot) @@ -1726,11 +1761,11 @@ def _build_fp4_mla_attention_decode_case( seq_lens=seq_lens, page_size=page_size, ) - assert metadata.fp4_mla_state.v_scale_pool is not None + assert find_fp4_mla(metadata).v_scale_pool is not None persistent_pool_base = kv_cache_manager.get_mla_v_packed_pool_base() if persistent_pool_base is not None: persistent_pool_base.zero_() - metadata.fp4_mla_state.v_scale_pool.zero_() + find_fp4_mla(metadata).v_scale_pool.zero_() metadata.kv_lens_cuda_runtime = torch.tensor( context_seq_lens, @@ -1739,13 +1774,13 @@ def _build_fp4_mla_attention_decode_case( ) metadata.prompt_lens_cuda_runtime = metadata.kv_lens_cuda_runtime.clone() metadata.prompt_lens_cpu_runtime = torch.tensor(context_seq_lens, dtype=torch.int32) - metadata.fp4_mla_state.batch_indices = torch.cat( + find_fp4_mla(metadata).batch_indices = torch.cat( [ torch.full((seq_len,), seq_idx, dtype=torch.int32, device=device) for seq_idx, seq_len in enumerate(context_seq_lens) ] ) - metadata.fp4_mla_state.positions = torch.cat( + find_fp4_mla(metadata).positions = torch.cat( [torch.arange(seq_len, dtype=torch.int32, device=device) for seq_len in context_seq_lens] ) context_latent = ( @@ -1763,9 +1798,9 @@ def _build_fp4_mla_attention_decode_case( torch.cuda.synchronize() metadata.num_contexts = 0 - metadata.fp4_mla_state.num_context_blocks = 0 - metadata.fp4_mla_state.num_generation_blocks = ( - len(seq_lens) * metadata.fp4_mla_state.page_table_stride + find_fp4_mla(metadata).num_context_blocks = 0 + find_fp4_mla(metadata).num_generation_blocks = ( + len(seq_lens) * find_fp4_mla(metadata).page_table_stride ) metadata.kv_lens_cuda_runtime = torch.tensor(seq_lens, dtype=torch.int32, device=device) metadata.prompt_lens_cuda_runtime = torch.full( @@ -1775,12 +1810,12 @@ def _build_fp4_mla_attention_decode_case( (len(seq_lens),), query_len_per_seq, dtype=torch.int32 ) num_queries = len(seq_lens) * query_len_per_seq - metadata.fp4_mla_state.batch_indices = torch.arange( + find_fp4_mla(metadata).batch_indices = torch.arange( len(seq_lens), dtype=torch.int32, device=device, ).repeat_interleave(query_len_per_seq) - metadata.fp4_mla_state.positions = torch.cat( + find_fp4_mla(metadata).positions = torch.cat( [ torch.arange(context_len, seq_len, dtype=torch.int32, device=device) for context_len, seq_len in zip(context_seq_lens, seq_lens) @@ -1790,7 +1825,9 @@ def _build_fp4_mla_attention_decode_case( metadata.num_ctx_tokens = 0 metadata.max_num_tokens = num_queries metadata.max_num_sequences = len(seq_lens) - metadata.max_total_draft_tokens = query_len_per_seq - 1 + spec = SpecDecodingMetadata(None) + spec.max_total_draft_tokens = query_len_per_seq - 1 + _attach_ext(metadata, SpecDecodingMetadata, spec) metadata.runtime_features.has_speculative_draft_tokens = query_len_per_seq > 1 generation_latent = ( torch.randn(num_queries, head_dim, dtype=torch.bfloat16, device=device) * 0.25 @@ -1823,15 +1860,17 @@ def _build_fp4_mla_attention_decode_case( if enable_helix: # CP1 exercises the Helix per-token metadata contract while the # verify group below still crosses a physical page boundary. - positions = metadata.fp4_mla_state.positions - metadata.helix_position_offsets = positions.clone() - metadata.helix_local_slots = positions.clone() - metadata.helix_kv_bounds = positions + 1 - metadata.helix_is_inactive_rank = torch.zeros( - len(seq_lens), dtype=torch.bool, device=device + positions = find_fp4_mla(metadata).positions + _attach_helix( + metadata, + position_offsets=positions.clone(), + local_slots=positions.clone(), + kv_bounds=positions + 1, + is_inactive_rank=torch.zeros(len(seq_lens), dtype=torch.bool, device=device), + spec_tokens_valid=query_len_per_seq > 1, ) - metadata._helix_spec_tokens_valid = query_len_per_seq > 1 + helix = find_helix(metadata) assert scatter_fp4_mla_kv_cache( metadata, generation_latent, @@ -1844,20 +1883,20 @@ def _build_fp4_mla_attention_decode_case( q_pe=q_pe, q_rope_out=q_rope_out, q_quant_input=q_quant_input, - helix_position_offsets=getattr(metadata, "helix_position_offsets", None), - helix_is_inactive_rank=getattr(metadata, "helix_is_inactive_rank", None), + helix_position_offsets=helix.position_offsets if helix is not None else None, + helix_is_inactive_rank=helix.is_inactive_rank if helix is not None else None, ) torch.cuda.synchronize() - assert metadata.fp4_mla_state.prequantized_q is not None - assert metadata.fp4_mla_state.prequantized_q_sf is not None - assert metadata.fp4_mla_state.q_batch_capacity == num_queries + assert find_fp4_mla(metadata).prequantized_q is not None + assert find_fp4_mla(metadata).prequantized_q_sf is not None + assert find_fp4_mla(metadata).q_batch_capacity == num_queries torch.testing.assert_close(q_rope_out, q_pe, rtol=0, atol=0) canonical_generation = _materialize_reference_cache_tokens( metadata, layer_idx=0, - batch_indices=metadata.fp4_mla_state.batch_indices, - positions=metadata.fp4_mla_state.positions, + batch_indices=find_fp4_mla(metadata).batch_indices, + positions=find_fp4_mla(metadata).positions, head_dim=head_dim, ) torch.testing.assert_close( @@ -1869,7 +1908,7 @@ def _build_fp4_mla_attention_decode_case( ) # Decode reference exercises the quantized cache without an HP overlay. - metadata.fp4_mla_state.hp_pool.zero_() + find_fp4_mla(metadata).hp_pool.zero_() torch.cuda.synchronize() return kv_cache_manager, metadata, q_quant_input, kv_lora_rank, qk_rope_head_dim @@ -1884,6 +1923,7 @@ def _fp4_mla_attention_decode_reference( kv_lora_rank: int, qk_rope_head_dim: int, ) -> tuple[torch.Tensor, torch.Tensor]: + fp4_mla_state = find_fp4_mla(metadata) head_dim = kv_lora_rank + qk_rope_head_dim storage = _materialize_reference_cache_storage(metadata, 0, head_dim) dequant_cache = storage[..., :head_dim] @@ -1892,7 +1932,7 @@ def _fp4_mla_attention_decode_reference( dequant_k_residual = storage[..., head_dim : head_dim + FP4_MLA_K_RESIDUAL_DIM] num_heads = q_nope.shape[1] q_full = torch.cat((q_nope, q_pe), dim=-1).reshape(-1, head_dim) - global_scale = metadata.fp4_mla_state.q_global_scale + global_scale = fp4_mla_state.q_global_scale q_fp4, q_sf = torch.ops.trtllm.fp4_quantize_with_residual( q_full, global_scale, @@ -1909,16 +1949,16 @@ def _fp4_mla_attention_decode_reference( ) p_dequant = None - if "_fp4_mla_attention_p_buf" in metadata.fp4_mla_state.workspaces: + if "_fp4_mla_attention_p_buf" in fp4_mla_state.workspaces: p_dequant = _dequant_fp4_swizzled( - metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_buf"], - metadata.fp4_mla_state.workspaces["_fp4_mla_attention_p_sf_buf"], + fp4_mla_state.workspaces["_fp4_mla_attention_p_buf"], + fp4_mla_state.workspaces["_fp4_mla_attention_p_sf_buf"], logical_dim=metadata.page_size, sf_per_token=metadata.page_size // FP4_BLOCK_SIZE, global_scale=FP4_MLA_P_GLOBAL_SCALE, ) - indptr = metadata.fp4_mla_state.paged_kv_indptr_decode.cpu().tolist() + indptr = fp4_mla_state.paged_kv_indptr_decode.cpu().tolist() num_seqs = metadata.num_seqs - metadata.num_contexts indptr = indptr[: num_seqs + 1] kv_lens = ( @@ -1935,7 +1975,7 @@ def _fp4_mla_attention_decode_reference( for seq_idx in range(num_seqs): kv_len = kv_lens[seq_idx] page_count = indptr[seq_idx + 1] - indptr[seq_idx] - assert page_count == metadata.fp4_mla_state.page_table_stride + assert page_count == fp4_mla_state.page_table_stride assert page_count * metadata.page_size >= kv_len full_cache = dequant_cache[indptr[seq_idx] : indptr[seq_idx + 1]].reshape(-1, head_dim) assert full_cache.shape[0] >= kv_len @@ -2040,9 +2080,9 @@ def _assert_fp4_mla_attention_decode_accuracy( sm_scale=sm_scale, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, - prequantized_q=metadata.fp4_mla_state.prequantized_q, - prequantized_q_sf=metadata.fp4_mla_state.prequantized_q_sf, - q_batch_capacity=metadata.fp4_mla_state.q_batch_capacity, + prequantized_q=find_fp4_mla(metadata).prequantized_q, + prequantized_q_sf=find_fp4_mla(metadata).prequantized_q_sf, + q_batch_capacity=find_fp4_mla(metadata).q_batch_capacity, softmax_stats_tensor=softmax_stats, ) torch.cuda.synchronize() @@ -2332,10 +2372,10 @@ def test_fp4_mla_context_tail_uses_draft_slack_ring( ) torch.cuda.synchronize() - hp_page = int(metadata.fp4_mla_state.hp_page_indices[0].item()) + hp_page = int(find_fp4_mla(metadata).hp_page_indices[0].item()) # hp_pool is [pages, FP4 MLA layers, kv_factor, ring * head_dim]; with # MTP the draft layer adds a second FP4 layer. Layer 0 is the target. - hp_pool = metadata.fp4_mla_state.hp_pool + hp_pool = find_fp4_mla(metadata).hp_pool hp_ring = hp_pool.view(*hp_pool.shape[:3], ring_size, head_dim) torch.testing.assert_close( hp_ring[hp_page, 0, 0, expected_slot], diff --git a/tests/unittest/_torch/attention/test_fp4_mla_packed.py b/tests/unittest/_torch/attention/test_fp4_mla_packed.py index 753f8cf0981c..5b4745f964c9 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla_packed.py +++ b/tests/unittest/_torch/attention/test_fp4_mla_packed.py @@ -12,6 +12,7 @@ import tensorrt_llm._torch.attention.backends.fp4_mla.config as fp4_mla_config import tensorrt_llm._torch.attention.backends.fp4_mla.metadata as fp4_mla_metadata import tensorrt_llm._torch.attention.backends.fp4_mla.packed_plan as fp4_mla_packed_plan +from tensorrt_llm._torch.attention.backends.fp4_mla.state import Fp4MlaState, find_fp4_mla pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, @@ -19,6 +20,22 @@ ) +def _attach_ext(metadata: SimpleNamespace, cls: type, extension: object) -> object: + """Serve ``extension`` as the stub's ``cls`` extension through ``find_ext``.""" + if not hasattr(metadata, "stub_extensions"): + metadata.stub_extensions = {} + metadata.find_ext = metadata.stub_extensions.get + metadata.stub_extensions[cls] = extension + return extension + + +def _fp4_mla_metadata(fp4_mla_state: object, **fields) -> SimpleNamespace: + """A metadata stub carrying ``fp4_mla_state`` as its FP4 MLA extension.""" + metadata = SimpleNamespace(**fields) + _attach_ext(metadata, Fp4MlaState, fp4_mla_state) + return metadata + + @pytest.mark.parametrize( "num_heads,query_len,expected", [ @@ -268,7 +285,7 @@ def test_uniform_generation_lengths_fuse_the_packed_plan(monkeypatch) -> None: corrected = torch.zeros(2, dtype=torch.int32, device="cuda") generation_lens = torch.zeros(2, dtype=torch.int32, device="cuda") plan = torch.zeros(16, dtype=torch.int32, device="cuda") - metadata = SimpleNamespace( + metadata = _fp4_mla_metadata( num_contexts=0, num_seqs=2, kv_lens_cuda_runtime=torch.tensor([17, 25], dtype=torch.int32, device="cuda"), @@ -302,7 +319,7 @@ def test_uniform_generation_lengths_fuse_the_packed_plan(monkeypatch) -> None: uniform(metadata, 8, 2, plan_geometry=geometry) assert [name for name, _ in calls] == ["plan"] assert calls[0][1]["generation_lengths"][3] == 4 - assert metadata.fp4_mla_state.packed_plan_key == ("key",) + assert find_fp4_mla(metadata).packed_plan_key == ("key",) # The other layers of the step launch nothing. uniform(metadata, 8, 2, plan_geometry=geometry) uniform(metadata, 8, 2) @@ -312,8 +329,8 @@ def test_uniform_generation_lengths_fuse_the_packed_plan(monkeypatch) -> None: assert len(calls) == 2 and calls[1][0] == "plan" and calls[1][1]["generation_lengths"] is None # A new step whose first caller cannot plan keeps the Triton lengths kernel, # and the decode then plans alone. - metadata.fp4_mla_state.generation_lengths_num_tokens = -1 - metadata.fp4_mla_state.packed_plan_key = None + find_fp4_mla(metadata).generation_lengths_num_tokens = -1 + find_fp4_mla(metadata).packed_plan_key = None uniform(metadata, 8, 2) uniform(metadata, 8, 2, plan_geometry=geometry) assert [name for name, _ in calls[2:]] == ["lengths", "plan"] @@ -331,7 +348,7 @@ def test_packed_plan_geometry_is_shared_by_cache_update_and_decode(monkeypatch) # An indexed device, as the kernels' tensors carry: the workspace helper # compares devices exactly before reusing a buffer. device = torch.device("cuda", torch.cuda.current_device()) - metadata = SimpleNamespace( + metadata = _fp4_mla_metadata( num_contexts=0, num_seqs=4, page_size=256, @@ -351,7 +368,7 @@ def test_packed_plan_geometry_is_shared_by_cache_update_and_decode(monkeypatch) # The plan lives in the metadata workspace (a view of the stored buffer). assert ( geometry.plan.data_ptr() - == metadata.fp4_mla_state.workspaces["_fp4_mla_packed_plan"].data_ptr() + == find_fp4_mla(metadata).workspaces["_fp4_mla_packed_plan"].data_ptr() ) assert geometry.plan.numel() == _plan_mod().packed_plan_length(4, 40) again = fp4_mla_packed_plan._fp4_mla_packed_plan_geometry( diff --git a/tests/unittest/_torch/attention/test_fp4_mla_splitkv.py b/tests/unittest/_torch/attention/test_fp4_mla_splitkv.py index 8c55c71ea14d..af1b22bebed2 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla_splitkv.py +++ b/tests/unittest/_torch/attention/test_fp4_mla_splitkv.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 """Split-KV reduction and per-bucket CUDA Graph workspace regressions.""" -import copy import importlib from types import ModuleType, SimpleNamespace @@ -105,8 +104,15 @@ def check() -> None: def test_splitkv_workspace_is_sized_per_graph_bucket() -> None: adapter = _module("fp4_mla_splitkv") + state_module = _module("state") heads, columns, splits, query_capacity = 128, 512, 4, 32 + def metadata_stub(is_cuda_graph: bool) -> SimpleNamespace: + state = SimpleNamespace(workspaces={}) + metadata = SimpleNamespace(is_cuda_graph=is_cuda_graph) + metadata.find_ext = lambda cls: state if cls is state_module.Fp4MlaState else None + return metadata + def partial_launch(*args, split_kv_buffers, **kwargs) -> None: # Isolate workspace ownership from the main attention kernel. These CUDA # operations are captured alongside the real adapter and real reducer. @@ -121,15 +127,11 @@ def partial_launch(*args, split_kv_buffers, **kwargs) -> None: FP4_MLA_P_GLOBAL_SCALE=1.0, run_trtllm_fp4_mla_decode_page_native_from_raw=partial_launch, ) - base_metadata = SimpleNamespace( - is_cuda_graph=False, fp4_mla_state=SimpleNamespace(workspaces={}) - ) + base_metadata = metadata_stub(is_cuda_graph=False) buckets = [] storage_pointers = set() for queries in [1, 4, 8]: - metadata = copy.copy(base_metadata) - metadata.is_cuda_graph = True - metadata.fp4_mla_state = SimpleNamespace(workspaces={}) + metadata = metadata_stub(is_cuda_graph=True) q = torch.full((queries, heads, 320), queries, dtype=torch.uint8, device="cuda") sf = torch.zeros((queries, heads, 40), dtype=torch.uint8, device="cuda") output = torch.empty((queries, heads, columns), dtype=torch.bfloat16, device="cuda") @@ -176,7 +178,7 @@ def forward( forward() workspace = { key: value - for key, value in metadata.fp4_mla_state.workspaces.items() + for key, value in state_module.find_fp4_mla(metadata).workspaces.items() if key.startswith("_fp4_mla_splitkv_") } tasks = splits * queries @@ -209,8 +211,8 @@ def forward( graph.replay() torch.testing.assert_close(output, torch.full_like(output, value), rtol=0, atol=0) for name, pointer in pointers.items(): - assert metadata.fp4_mla_state.workspaces[name].data_ptr() == pointer - assert not base_metadata.fp4_mla_state.workspaces + assert state_module.find_fp4_mla(metadata).workspaces[name].data_ptr() == pointer + assert not state_module.find_fp4_mla(base_metadata).workspaces @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) diff --git a/tests/unittest/_torch/attention/test_mixed_precision_backends.py b/tests/unittest/_torch/attention/test_mixed_precision_backends.py index 1c22dc3bd18e..5873e668e724 100644 --- a/tests/unittest/_torch/attention/test_mixed_precision_backends.py +++ b/tests/unittest/_torch/attention/test_mixed_precision_backends.py @@ -139,7 +139,7 @@ def test_optimized_two_pool_prefill_and_decode( def forward(begin, end, context): nonlocal decode_graph, native_output - metadata = backend.Metadata( + metadata = backend.create_metadata( seq_lens=torch.tensor([end - begin] * batch, dtype=torch.int32), num_contexts=batch if context else 0, kv_cache_params=KVCacheParams( @@ -473,7 +473,7 @@ def test_mla_cached_context_and_decode(backend): ) assert cache.resize(end) manager._restore_page_index_bufs(0, cache) - metadata = backend.Metadata( + metadata = backend.create_metadata( seq_lens=torch.tensor([end - begin], dtype=torch.int32), num_contexts=int(context), kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[begin]), diff --git a/tests/unittest/_torch/attention/test_mla_registry.py b/tests/unittest/_torch/attention/test_mla_registry.py index 06c15a2011fe..3cae2443a6ac 100644 --- a/tests/unittest/_torch/attention/test_mla_registry.py +++ b/tests/unittest/_torch/attention/test_mla_registry.py @@ -18,6 +18,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from torch import nn from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams @@ -132,7 +133,7 @@ def test_output_gate_runs_between_attention_and_output_projection( ) -> None: mla_layer = _OutputGateStub(register_to_config) hidden_states = torch.randn(2, 4) - attn_metadata = SimpleNamespace(num_contexts=0, num_tokens=2) + attn_metadata = metadata_stub(num_contexts=0, num_tokens=2) with patch.object(torch.ops.trtllm, "create_mla_outputs", return_value=mla_layer.attn_output): output = mla_layer.forward(None, hidden_states, attn_metadata) @@ -241,7 +242,7 @@ def test_create_mla_outputs_custom_op_returns_tensor() -> None: def test_dsv4_epilogue_fusion_supports_mixed_batch() -> None: mla_layer = _make_dsv4_epilogue_layer() - metadata = SimpleNamespace(num_contexts=1, num_generations=1) + metadata = metadata_stub(num_contexts=1, num_generations=1) hidden_states = torch.empty(8, 16) with ( @@ -265,7 +266,7 @@ def test_dsv4_epilogue_fusion_supports_sm107_direct_output( num_contexts: int, num_generations: int ) -> None: mla_layer = _make_dsv4_epilogue_layer() - metadata = SimpleNamespace(num_contexts=num_contexts, num_generations=num_generations) + metadata = metadata_stub(num_contexts=num_contexts, num_generations=num_generations) hidden_states = torch.empty(8, 16) with ( @@ -286,7 +287,7 @@ def test_dsv4_epilogue_fusion_rejects_sm107_without_direct_output() -> None: mla_layer = _make_dsv4_epilogue_layer() fallback = torch.empty(8, 16) mla_layer.create_output.return_value = fallback - metadata = SimpleNamespace(num_contexts=1, num_generations=0) + metadata = metadata_stub(num_contexts=1, num_generations=0) hidden_states = torch.empty(8, 16) with ( @@ -302,7 +303,7 @@ def test_dsv4_epilogue_fusion_rejects_sm107_without_direct_output() -> None: def test_dsv4_fusion_create_output_uses_bucket_token_count() -> None: mla_layer = _make_dsv4_epilogue_layer() - metadata = SimpleNamespace(num_contexts=1, num_generations=0) + metadata = metadata_stub(num_contexts=1, num_generations=0) hidden_states = torch.empty(8, 16) with ( diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index d0e61dfb5e54..66a7a599e1c2 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -16,7 +16,6 @@ import functools import inspect import math -from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -30,6 +29,7 @@ run_backend, run_case, ) +from fmha_test_utils import metadata_stub from utils.util import getSMVersion pytestmark = pytest.mark.skipif( @@ -492,7 +492,7 @@ def test_prims_ts_fp8_mla_scales_and_graph_replay( mla_bmm1_scale=torch.tensor([bmm1_scale, bmm1_scale * math.log2(math.e)], device=device), mla_bmm2_scale=torch.tensor([bmm2_scale], device=device), ) - metadata = SimpleNamespace( + metadata = metadata_stub( host_kv_cache_pool_pointers=None, host_kv_cache_pool_mapping=None, kv_cache_block_offsets=block_tables, diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 8d2f771b88b2..75259f8d3058 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -20,6 +20,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from packaging.version import Version import tensorrt_llm._torch.attention.backends.fmha.prims_ts as prims_ts_module @@ -229,7 +230,7 @@ def _support_result( dtype=torch.int32, ), ) - metadata = SimpleNamespace( + metadata = metadata_stub( helix_position_offsets=None, num_sparse_topk=0, use_spec_decoding=use_spec_decoding, @@ -418,7 +419,7 @@ def test_is_supported_accepts_and_forwards_phase_keyword( support_check = Mock(return_value=(True, "")) monkeypatch.setattr(fmha, "_is_supported_with_reason", support_check) q = Mock(spec=torch.Tensor) - metadata = SimpleNamespace() + metadata = metadata_stub() forward_args = AttentionForwardArgs() assert fmha.is_supported( @@ -788,7 +789,7 @@ def test_v2_total_page_bound_is_not_expanded() -> None: attn = _Attention() fmha = PrimsTSFmha(attn) get_page_index_upper_bound = Mock(return_value=4096) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=_make_v2_manager( impl=SimpleNamespace(get_page_index_upper_bound=get_page_index_upper_bound) ) @@ -803,7 +804,7 @@ def test_v1_total_page_bound_excludes_slots_before_selected_layer(is_mla: bool) attn = _Attention(is_mla=is_mla) attn.local_layer_idx = 3 fmha = PrimsTSFmha(attn) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=_make_v1_manager( impl=SimpleNamespace( get_primary_pool_data=lambda _: torch.empty(64, dtype=torch.uint8) @@ -821,7 +822,7 @@ def test_v1_total_page_bound_excludes_slots_before_selected_layer(is_mla: bool) def test_kv_page_offset_uses_v2_manager_displacement() -> None: attn = _Attention() fmha = PrimsTSFmha(attn) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=_make_v2_manager( impl=SimpleNamespace(get_page_index_upper_bound=lambda *args: 128), kv_offset=torch.tensor([0, 128]), @@ -847,7 +848,7 @@ def test_kv_page_offset_is_inferred_from_v1_host_tables() -> None: [[[[0, 1, 2], [64, 65, 66]], [[3, 4, 5], [67, 68, 69]]]], dtype=torch.int32, ) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=_make_v1_manager( impl=SimpleNamespace(), host_kv_cache_block_offsets=host_offsets, @@ -975,7 +976,7 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( ], dtype=torch.int32, ) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_block_offsets=torch.empty((3, 2, 4), dtype=torch.int32), host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), @@ -1183,7 +1184,7 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( ) get_page_index_upper_bound = Mock(return_value=12) - metadata = SimpleNamespace( + metadata = metadata_stub( beam_width=1, kv_cache_block_offsets=torch.empty((2, 2, 4), dtype=torch.int32), host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), @@ -1528,7 +1529,7 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( wrapper_factory, ) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cuda_graph=False, beam_width=1, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), @@ -1835,7 +1836,7 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( ) get_page_index_upper_bound = Mock(return_value=20) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cuda_graph=is_cuda_graph, beam_width=1, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), @@ -1940,7 +1941,7 @@ def test_mla_prepare_workspace_sizes_caller_owned_workspace( fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cuda_graph=is_cuda_graph, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), max_num_requests=2, @@ -1995,7 +1996,7 @@ def test_mla_prepare_workspace_preserves_cached_wrappers_with_stable_allocation( wrapper = Mock() fmha._mla_decode_wrappers[2] = wrapper q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cuda_graph=True, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), max_num_requests=2, @@ -2033,7 +2034,7 @@ def test_mla_caller_workspace_grows_across_plan_profiles( fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) - metadata = SimpleNamespace( + metadata = metadata_stub( is_cuda_graph=False, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), max_num_requests=2, @@ -2079,7 +2080,7 @@ def test_decode_prepare_workspace_reserves_tail_after_compact_preprocessing( attn = _Attention() fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_block_offsets=torch.empty((2, 2, 4), dtype=torch.int32), max_num_requests=2, num_contexts=0, @@ -2139,7 +2140,7 @@ def test_decode_workspace_tail_is_stable_across_mixed_context_layouts( attn = _Attention() fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_block_offsets=torch.empty((4, 2, 4), dtype=torch.int32), max_num_requests=4, num_contexts=2, @@ -2185,7 +2186,7 @@ def test_workspace_cannot_grow_during_capture(monkeypatch: pytest.MonkeyPatch) - "get_prims_ts_batch_decode_workspace_size", lambda *args, **kwargs: 32, ) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_block_offsets=torch.empty((2, 2, 4), dtype=torch.int32), max_num_requests=2, num_contexts=0, @@ -2262,7 +2263,7 @@ def test_phased_forward_routes_query_and_qkv_by_phase( else fmha.context_out_head_size ) output = torch.empty((num_tokens, attn.num_heads * out_head_size)) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_block_offsets=torch.empty(1), effective_workspace=torch.empty(0, dtype=torch.int8), num_contexts=1, @@ -2348,7 +2349,7 @@ def test_cache_free_context_plans_contiguous_wrapper_once(monkeypatch: pytest.Mo output = torch.empty( (batch_size * seq_len, attn.num_heads * attn.head_dim), dtype=torch.bfloat16 ) - metadata = SimpleNamespace( + metadata = metadata_stub( kv_cache_manager=None, num_contexts=batch_size, num_generations=0, diff --git a/tests/unittest/_torch/attention/test_skip_softmax_sm120.py b/tests/unittest/_torch/attention/test_skip_softmax_sm120.py index af8a067f7129..31cab40a74fa 100644 --- a/tests/unittest/_torch/attention/test_skip_softmax_sm120.py +++ b/tests/unittest/_torch/attention/test_skip_softmax_sm120.py @@ -87,7 +87,7 @@ def _run_context( sparse_params=sparse_params, ) - metadata = AttentionCls.Metadata( + metadata = AttentionCls.create_metadata( max_num_requests=len(seq_lens), max_num_tokens=8192, kv_cache_manager=None, diff --git a/tests/unittest/_torch/attention/test_triton_custom_mask_fmha.py b/tests/unittest/_torch/attention/test_triton_custom_mask_fmha.py index ba384785250f..2ed94e9715e9 100644 --- a/tests/unittest/_torch/attention/test_triton_custom_mask_fmha.py +++ b/tests/unittest/_torch/attention/test_triton_custom_mask_fmha.py @@ -17,7 +17,7 @@ import pytest import torch -from fmha_test_utils import FakeAttention +from fmha_test_utils import FakeAttention, metadata_stub from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase from tensorrt_llm._torch.attention.backends.fmha.triton_custom_mask import TritonCustomMaskFmha @@ -56,7 +56,7 @@ def test_custom_mask_quantized_context_support( attn = FakeAttention() attn.quant_mode = int(quant_mode) fmha = TritonCustomMaskFmha(attn) - metadata = SimpleNamespace( + metadata = metadata_stub( num_contexts=1, is_cross=False, kv_cache_block_offsets=object(), diff --git a/tests/unittest/_torch/attention/test_trtllm_paged_context_fmha_gate.py b/tests/unittest/_torch/attention/test_trtllm_paged_context_fmha_gate.py index dde3606f0fde..8ba6a1ee9a10 100644 --- a/tests/unittest/_torch/attention/test_trtllm_paged_context_fmha_gate.py +++ b/tests/unittest/_torch/attention/test_trtllm_paged_context_fmha_gate.py @@ -27,6 +27,7 @@ import pytest import torch +from fmha_test_utils import metadata_stub from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha from tensorrt_llm._torch.attention.backends.interface import ( @@ -158,7 +159,7 @@ def test_no_reuse_features_no_gate(): def _make_dispatch_inputs(num_contexts): """Stub metadata/forward_args pairs for exercising the real ``is_supported`` entry point on a blocked configuration.""" - metadata = SimpleNamespace( + metadata = metadata_stub( use_paged_context_fmha=True, kv_cache_manager=SimpleNamespace(head_dim=64), runtime_features=ALL_FEATURES, diff --git a/tests/unittest/_torch/executor/engine/test_metadata.py b/tests/unittest/_torch/executor/engine/test_metadata.py index 475cdf6b5d21..34f0663e681d 100644 --- a/tests/unittest/_torch/executor/engine/test_metadata.py +++ b/tests/unittest/_torch/executor/engine/test_metadata.py @@ -7,7 +7,12 @@ import pytest -from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.features import FlashMlaMetadata +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionBackend, + AttentionRuntimeFeatures, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import KVCacheLayerSpec, ModelConfig from tensorrt_llm._torch.pyexecutor.engine.metadata import ( _get_max_num_heads_per_kv, @@ -17,12 +22,12 @@ pytestmark = pytest.mark.cpu_only -class _AttentionMetadata: +class _AttentionMetadata(TrtllmAttentionMetadata): def __init__(self, **kwargs: Any) -> None: self.__dict__.update(kwargs) -class _AttentionBackend: +class _AttentionBackend(AttentionBackend): Metadata = _AttentionMetadata @@ -56,7 +61,7 @@ def test_build_attention_metadata_resolves_model_derived_values() -> None: ) assert metadata.max_num_heads_per_kv == 4 - assert metadata.enable_flash_mla + assert metadata.extension_types == (FlashMlaMetadata,) assert metadata.enable_context_mla_with_cached_kv assert metadata.sparse_metadata_params is sparse_metadata_params sparse_attention_config.to_sparse_metadata_params.assert_called_once_with( diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 5fd8ba58086a..231ebb27d642 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -12,6 +12,7 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.cache_manager import ( Glm5NextCacheManager, ) @@ -123,6 +124,15 @@ from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager from tensorrt_llm.sampling_params import SamplingParams + +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + + skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -5354,7 +5364,7 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): ) policy.commit_all_layers( - SimpleNamespace( + _recurrent_stub( mamba_metadata=SimpleNamespace( replay_num_decodes=batch_size, replay_work_items=replay_work_items, @@ -5407,14 +5417,14 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): if not hasattr(mgr, "_dummy_request_mask"): mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(64, dtype=torch.bool)) mgr.update_mamba_states( - SimpleNamespace( - num_seqs=batch_size, - num_contexts=0, + _recurrent_stub( mamba_metadata=SimpleNamespace( replay_num_decodes=batch_size, replay_work_items=torch.empty((batch_size, 4), dtype=torch.int32), replay_n_writes=torch.empty(1, dtype=torch.int32), ), + num_seqs=batch_size, + num_contexts=0, ), torch.full((batch_size,), 3, dtype=torch.int32), state_indices=torch.arange(batch_size, dtype=torch.int32), diff --git a/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py b/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py index 8f4e4010db1b..0a89395038d1 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py @@ -35,6 +35,7 @@ import torch from _torch.helpers import create_mock_cuda_graph_runner +from tensorrt_llm._torch.attention.backends.features import DraftKVMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor import cuda_graph_runner as cuda_graph_runner_module from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType @@ -426,11 +427,12 @@ def __init__(self, value: int): self.kv_cache_manager = object() self.kv_cache_block_offsets = torch.full((1,), value, device="cuda", dtype=torch.int32) self.host_kv_cache_block_offsets = torch.zeros(1, dtype=torch.int32) - self.draft_kv_cache_block_offsets = torch.full( + draft_kv = DraftKVMetadata(self) + draft_kv.kv_cache_block_offsets = torch.full( (1,), value + 1000, device="cuda", dtype=torch.int32 ) + self.replace_ext(draft_kv) self.some_alias = self.kv_cache_block_offsets - self.enable_flash_mla = False def prepare_for_draft_forward(self): # Mirrors DSA's _fullkv aliases, rebound as a side effect of the swap. @@ -614,7 +616,10 @@ def test_restore_clears_swapped_attrs_and_target_buffer(self, monkeypatch): saved = prepare_attn_metadata_for_draft_replay(attn_metadata, draft_mgr) assert attn_metadata.draft_replay_swapped_attrs["kv_cache_block_offsets"] is target_buf assert attn_metadata.draft_replay_swapped_attrs["some_alias"] is target_buf - assert attn_metadata.kv_cache_block_offsets is attn_metadata.draft_kv_cache_block_offsets + assert ( + attn_metadata.kv_cache_block_offsets + is attn_metadata.ext(DraftKVMetadata).kv_cache_block_offsets + ) # A repeated swap keeps the first recorded original. attn_metadata.swap_for_draft("kv_cache_block_offsets", torch.zeros(1, device="cuda")) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 6bae538cca81..2d8df0d105af 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -2116,12 +2116,12 @@ def test_warmup_builds_dsa_attn_metadata_when_no_forward_ran(self) -> None: warmup forward was skipped it is built on demand so the engines still compile before serving.""" from tensorrt_llm._torch.attention.backends.sparse.dsa import \ - DSAtrtllmAttentionMetadata + DSAMetadata engine = object.__new__(DecoderWarmup) engine._config = SimpleNamespace( attention_backend=SimpleNamespace( - Metadata=DSAtrtllmAttentionMetadata), + metadata_extensions=(DSAMetadata, )), original_max_draft_len=2, cuda_graph_batch_sizes=[8], spec_config=None, @@ -2131,7 +2131,9 @@ def test_warmup_builds_dsa_attn_metadata_when_no_forward_ran(self) -> None: model=None, buffers=None, attention_runtime_features=None) - metadata = Mock(spec=DSAtrtllmAttentionMetadata) + dsa = Mock(spec=DSAMetadata) + metadata = Mock() + metadata.find_ext.side_effect = lambda cls: dsa if cls is DSAMetadata else None def build(kv_cache_manager, draft_kv_cache_manager): engine._metadata.attn_metadata = metadata @@ -2147,9 +2149,8 @@ def build(kv_cache_manager, draft_kv_cache_manager): engine._metadata.setup_attention.assert_called_once_with( kv_cache_manager, None) - metadata.warmup_cute_dsl_radix_topk.assert_called_once_with(3) - metadata.warmup_selfsampling_topk.assert_called_once_with( - 3, batch_sizes=[8]) + dsa.warmup_cute_dsl_radix_topk.assert_called_once_with(3) + dsa.warmup_selfsampling_topk.assert_called_once_with(3, batch_sizes=[8]) def test_warmup_does_not_build_metadata_for_non_dsa_backend(self) -> None: from tensorrt_llm._torch.attention.backends.trtllm import \ @@ -4009,11 +4010,20 @@ class TestDecodeGraphMetadataPreparation(unittest.TestCase): @staticmethod def _metadata(): from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import \ - CSA2TrtllmMetadata - metadata = object.__new__(CSA2TrtllmMetadata) + CSA2Metadata + from tensorrt_llm._torch.attention.backends.trtllm import \ + TrtllmAttentionMetadata + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(CSA2Metadata(metadata)) metadata.is_cuda_graph = True return metadata + @staticmethod + def _defers_outputs(metadata) -> bool: + from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import \ + CSA2Metadata + return metadata.ext(CSA2Metadata)._csa2_defer_decode_outputs + def test_scope_requires_selected_ordinary_decode_graph(self): key = KeyType(batch_size=1, draft_len=0) cases = ( @@ -4063,9 +4073,9 @@ def test_scope_requires_selected_ordinary_decode_graph(self): batch, promoted, enable_spec_decode=enable_spec_decode): - self.assertEqual(metadata._csa2_defer_decode_outputs, case + self.assertEqual(self._defers_outputs(metadata), case in ("ordinary", "ordinary_overlap")) - self.assertFalse(metadata._csa2_defer_decode_outputs) + self.assertFalse(self._defers_outputs(metadata)) def test_forward_scopes_proof_to_prepare_for_capture_replay_and_fallback( self): @@ -4091,7 +4101,7 @@ def test_forward_scopes_proof_to_prepare_for_capture_replay_and_fallback( def prepare(*args, **kwargs): self.assertIs(metadata.trtllm_gen_jit_warmup, False) - self.assertEqual(metadata._csa2_defer_decode_outputs, mode + self.assertEqual(self._defers_outputs(metadata), mode != "fallback") if mode == "prepare_error": raise RuntimeError("prepare failed") @@ -4112,7 +4122,7 @@ def prepare(*args, **kwargs): }) self.assertIs(actual_outputs["logits"], outputs["logits"]) - self.assertFalse(metadata._csa2_defer_decode_outputs) + self.assertFalse(self._defers_outputs(metadata)) if mode in ("warmup", "capture"): runner.capture.assert_called_once() if mode in ("capture", "replay"): diff --git a/tests/unittest/_torch/modeling/_deepseek_v41_test_utils.py b/tests/unittest/_torch/modeling/_deepseek_v41_test_utils.py index d9e9eb0ef8b3..d6666a1a3823 100644 --- a/tests/unittest/_torch/modeling/_deepseek_v41_test_utils.py +++ b/tests/unittest/_torch/modeling/_deepseek_v41_test_utils.py @@ -13,13 +13,14 @@ from test_modeling_deepseekv41 import _tiny_v41_model_config import tensorrt_llm +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import ( CSA2CacheManager, CSA2CacheRole, ) from tensorrt_llm._torch.attention.backends.sparse.csa2.ced import prepare_ced_global_kv from tensorrt_llm._torch.attention.backends.sparse.csa2.indexer import CSA2Indexer -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.quantization import read_index_rows from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.models.modeling_deepseekv41 import DeepseekV41ForCausalLM @@ -68,13 +69,18 @@ def _apply_swa_floor(metadata, floors: torch.Tensor) -> None: """Independent reference masking for explicit layer-loop comparisons.""" # SWA slots are resolved when a forward enters its first layer; resolve # them now so the layers below read the masked slots. - metadata._ensure_swa_slots() + metadata.ext(CSA2Metadata)._ensure_swa_slots() width = metadata.kv_cache_manager.layout.window_size logical = ( - metadata.csa2_positions[:, None] - width + 1 + torch.arange(width, device=floors.device) + metadata.ext(CSA2Metadata).csa2_positions[:, None] + - width + + 1 + + torch.arange(width, device=floors.device) ) - for layer, indices in metadata.csa2_swa_indices.items(): - metadata.csa2_swa_indices[layer] = indices.masked_fill(logical < floors[:, None], -1) + for layer, indices in metadata.ext(CSA2Metadata).csa2_swa_indices.items(): + metadata.ext(CSA2Metadata).csa2_swa_indices[layer] = indices.masked_fill( + logical < floors[:, None], -1 + ) @pytest.fixture(params=["auto", "fp8"]) @@ -172,7 +178,7 @@ def cache_case( scheduled.context_requests_last_chunk = requests[:num_contexts] scheduled.generation_requests = requests[num_contexts:] manager.prepare_resources(scheduled) - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( seq_lens=torch.tensor(lengths, dtype=torch.int32), num_contexts=num_contexts, max_num_requests=batch, @@ -289,9 +295,9 @@ def hook(module, args, output): def save_routing(module, args, kwargs, output): metadata = kwargs["attn_metadata"] - for layer, indices in metadata.csa2_indices.items(): + for layer, indices in metadata.ext(CSA2Metadata).csa2_indices.items(): values[f"topk{layer}"] = (indices.clone(),) - for layer, blocks in metadata.csa2_candidates.items(): + for layer, blocks in metadata.ext(CSA2Metadata).csa2_candidates.items(): values[f"candidates{layer}"] = (blocks.clone(),) owner = model.model.layers[5] @@ -329,8 +335,8 @@ def global_precompute_scope(metadata): try: yield finally: - metadata.csa2_precomputed_kv_layers.clear() - metadata.reset_routing() + metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers.clear() + metadata.ext(CSA2Metadata).reset_routing() @contextmanager @@ -406,7 +412,7 @@ def run( p + n for p, n in zip(starts[num_contexts:], lengths[num_contexts:]) ] metadata.kv_cache_params.num_cached_tokens_per_seq = list(starts) - metadata.prepare_context_replay(reqs[:num_contexts]) + metadata.ext(CSA2Metadata).prepare_context_replay(reqs[:num_contexts]) metadata.prepare() positions = torch.cat( [torch.arange(p, p + n, device="cuda", dtype=torch.int32) for p, n in zip(starts, lengths)] diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_ced.py b/tests/unittest/_torch/modeling/test_deepseek_v41_ced.py index 2514f440ca26..8bdbbb609f50 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_ced.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_ced.py @@ -20,6 +20,7 @@ from _deepseek_v41_test_utils import cache_config as shared_cache_config from _deepseek_v41_test_utils import model_and_config as global_model_and_config +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.modules.mhc.hyper_connection import HCState skip_pre_blackwell = pytest.mark.skipif( @@ -80,7 +81,7 @@ def independent_window_reference(model, metadata, ids, window=128, chunks=None): metadata.prompt_lens = [end - start] metadata.kv_cache_params.num_cached_tokens_per_seq = [start] metadata.prepare() - metadata.csa2_precomputed_kv_layers = {5} + metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers = {5} _apply_swa_floor( metadata, torch.full((end - start,), floor, device=ids.device, dtype=torch.int64), @@ -120,26 +121,31 @@ def test_prepared_encoder_replay_preserves_live_speculative_endpoints( req.py_ced_replay = EncoderReplay( 0, id(metadata.kv_cache_manager.kv_cache_map[0]), 256, start ) - metadata.prepare_context_replay([req]) + metadata.ext(CSA2Metadata).prepare_context_replay([req]) metadata.prepare() # Host reservation ends at134; accepting one of five drafts moves the # actual six verification rows back to [124,130). The allocator may # only have pages covering that live interval. metadata.kv_lens_cuda[:2].copy_(torch.tensor([end, 130], device="cuda")) metadata.on_update_kv_lens() - positions = metadata.csa2_positions.clone() - swa_slots = {layer: slots.clone() for layer, slots in metadata.csa2_swa_write_slots.items()} + positions = metadata.ext(CSA2Metadata).csa2_positions.clone() + swa_slots = { + layer: slots.clone() + for layer, slots in metadata.ext(CSA2Metadata).csa2_swa_write_slots.items() + } for _ in range(2): - metadata.begin_model_forward() + metadata.begin_forward() metadata.on_update_kv_lens() assert metadata.kv_lens_cuda[:2].tolist() == [end, 130] - torch.testing.assert_close(metadata.csa2_positions, positions) - assert metadata.csa2_positions[-6:].tolist() == list(range(124, 130)) - assert metadata.csa2_replay_start_positions.tolist() == [start, 0] + torch.testing.assert_close(metadata.ext(CSA2Metadata).csa2_positions, positions) + assert metadata.ext(CSA2Metadata).csa2_positions[-6:].tolist() == list(range(124, 130)) + assert metadata.ext(CSA2Metadata).csa2_replay_start_positions.tolist() == [start, 0] for layer, slots in swa_slots.items(): - torch.testing.assert_close(metadata.csa2_swa_write_slots[layer], slots) - assert (metadata.csa2_swa_indices[layer][0, :-1] == -1).all() - for batch in metadata._csa2_compression.values(): + torch.testing.assert_close( + metadata.ext(CSA2Metadata).csa2_swa_write_slots[layer], slots + ) + assert (metadata.ext(CSA2Metadata).csa2_swa_indices[layer][0, :-1] == -1).all() + for batch in metadata.ext(CSA2Metadata)._csa2_compression.values(): assert batch.kv_lengths.tolist() == [end, 130] assert batch.start_positions.tolist() == [256 if recovering else 0, 124] if recovering: @@ -147,14 +153,14 @@ def test_prepared_encoder_replay_preserves_live_speculative_endpoints( req.context_current_position = end metadata.seq_lens = torch.tensor([1, 6], dtype=torch.int32) metadata.kv_cache_params.num_cached_tokens_per_seq = [end, 128] - metadata.prepare_context_replay([req]) + metadata.ext(CSA2Metadata).prepare_context_replay([req]) metadata.prepare() - assert metadata.csa2_replay_mode is None - assert metadata.csa2_replay_start_positions.tolist() == [start, 0] + assert metadata.ext(CSA2Metadata).csa2_replay_mode is None + assert metadata.ext(CSA2Metadata).csa2_replay_start_positions.tolist() == [start, 0] req.py_ced_replay = None - metadata.prepare_context_replay([req]) + metadata.ext(CSA2Metadata).prepare_context_replay([req]) metadata.prepare() - assert metadata.csa2_replay_start_positions.tolist() == [0, 0] + assert metadata.ext(CSA2Metadata).csa2_replay_start_positions.tolist() == [0, 0] @pytest.mark.parametrize("recovering", [False, True]) @@ -360,7 +366,7 @@ def test_final_window_mixed_batch_keeps_finished_context_and_generation( buffer.copy_(saved) reqs[0].prompt_len = 1024 actual = run(model, metadata, reqs, ids, [0, 0, 0], lengths, 2) - assert metadata.csa2_decoder_capture_lens == (0, 128) + assert metadata.ext(CSA2Metadata).csa2_decoder_capture_lens == (0, 128) assert metadata.seq_lens.tolist() == lengths assert actual.shape == expected.shape == (3, 1024) torch.testing.assert_close(actual[1:], expected[1:], atol=0.02, rtol=0.02) @@ -450,8 +456,8 @@ def configure_deferral(ced_enabled): ].flatten().tolist() == list(range(offset + n - rows, offset + n)) assert out.shape == (1, 1024) assert metadata.seq_lens.tolist() == [n] - assert metadata.csa2_decoder_capture_lens == (rows,) - assert not metadata.csa2_precomputed_kv_layers + assert metadata.ext(CSA2Metadata).csa2_decoder_capture_lens == (rows,) + assert not metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers if rows and not final: # Decoder KV, rather than retained HC tensors, carries # the partial window across ordinary suspend/resume. diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_decoder_global.py b/tests/unittest/_torch/modeling/test_deepseek_v41_decoder_global.py index 7cfe46672a6b..7847bdda19ca 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_decoder_global.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_decoder_global.py @@ -12,6 +12,7 @@ from _deepseek_v41_test_utils import model_and_config as shared_model_and_config from utils.util import skip_pre_blackwell +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream cache_config = shared_cache_config @@ -111,7 +112,7 @@ def test_explicit_global_production_matches_original(model_and_config, overlap, values.update( {f"cache{i}": (v.clone().view(torch.uint8),) for i, v in enumerate(buffers)} ) - assert not metadata.csa2_precomputed_kv_layers + assert not metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers runs.append(values) differences = [] for run_id, run in enumerate(runs[1:], 1): diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_encoder_replay.py b/tests/unittest/_torch/modeling/test_deepseek_v41_encoder_replay.py index 6f0f6e9802ba..eec884f4e5e2 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_encoder_replay.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_encoder_replay.py @@ -22,6 +22,7 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheRole +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.modules.mhc.hyper_connection import HCState from tensorrt_llm._torch.pyexecutor.ced_replay import encoder_replay_tokens from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -43,11 +44,11 @@ def test_global_source_selection_preserves_cold_and_generation_rows(cache_config # recovery rows are projected; the other requests keep their complete input. with cache_case(config, sparse, dtype, [130, 3, 1], [128, 0, 258], 2) as case: metadata = case[0] - metadata.set_swa_bounded_replay([256, None, None]) + metadata.ext(CSA2Metadata).set_swa_bounded_replay([256, None, None]) metadata.prepare() hidden = torch.arange(134, device="cuda").reshape(-1, 1) for owner in metadata.kv_cache_manager.layout.kv_source_layer_ids: - selected = metadata.select_global_source(owner, hidden) + selected = metadata.ext(CSA2Metadata).select_global_source(owner, hidden) assert selected.flatten().tolist() == list(range(128, 134)) @@ -66,7 +67,7 @@ def reference(model, metadata, ids, prefix=256, chunks=None, window=128): metadata.prompt_lens = [stop - begin] metadata.kv_cache_params.num_cached_tokens_per_seq = [begin] if index == 0: - metadata.set_swa_bounded_replay([prefix]) + metadata.ext(CSA2Metadata).set_swa_bounded_replay([prefix]) metadata.prepare() _apply_swa_floor( metadata, torch.full((stop - begin,), start, device="cuda", dtype=torch.int64) diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal_backbone.py b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal_backbone.py index 3ad8fc7aba27..e7152adf52bc 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal_backbone.py +++ b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal_backbone.py @@ -54,7 +54,7 @@ def backbone() -> tuple[v41.DeepseekV41Model, SimpleNamespace, torch.Tensor, tor model.embed_tokens = Mock(return_value=embeddings) seq_lens = torch.tensor([4], dtype=torch.int32) metadata = SimpleNamespace( - begin_model_forward=Mock(), + begin_forward=Mock(), num_tokens=4, num_contexts=1, all_rank_num_tokens=[4], diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index eeed1be64210..b9a34e22fdf4 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -10,6 +10,7 @@ import torch from transformers import AutoConfig, PretrainedConfig +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import Glm5NextMamba2Metadata from tensorrt_llm._torch.distributed import AllReduce, AllReduceStrategy from tensorrt_llm._torch.model_config import ModelConfig @@ -38,6 +39,14 @@ from tensorrt_llm.mapping import Mapping +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + + @pytest.mark.cpu_only def test_config_fallback_round_trip(tmp_path): from transformers.models.auto.configuration_auto import CONFIG_MAPPING @@ -445,7 +454,6 @@ def test_runtime_context_uses_prepared_schedules_and_live_lengths(is_cuda_graph) prepared.glm_cached_lens_host = [3, 5] metadata = SimpleNamespace( kv_cache_manager=object(), - mamba_metadata=prepared, seq_lens=torch.tensor([3, 4]), num_contexts=1, num_ctx_tokens=3, @@ -453,6 +461,8 @@ def test_runtime_context_uses_prepared_schedules_and_live_lengths(is_cuda_graph) kv_lens_cuda=live_lengths, is_cuda_graph=is_cuda_graph, ) + recurrent = SimpleNamespace(mamba_metadata=prepared) + metadata.ext = lambda cls: recurrent context = build_glm5_next_runtime_context(metadata) assert context.ctx_cu_seqlens is prepared.glm_ctx_cu_seqlens assert context.cached_lens is prepared.glm_cached_lens_host @@ -473,7 +483,7 @@ def test_runtime_context_uses_prepared_schedules_and_live_lengths(is_cuda_graph) build_glm5_next_runtime_context(metadata) metadata.kv_lens_cuda = live_lengths for invalid in (None, False, SimpleNamespace(glm_block_tables=prepared.glm_block_tables)): - metadata.mamba_metadata = invalid + recurrent.mamba_metadata = invalid with pytest.raises(AssertionError, match="Glm5NextMamba2Metadata"): build_glm5_next_runtime_context(metadata) @@ -1081,7 +1091,7 @@ def run(lengths, discard_prefix=False): if cached and discard_prefix: latent.zero_() index.zero_() - metadata = SimpleNamespace( + metadata = _recurrent_stub( kv_cache_manager=manager, seq_lens=torch.tensor([length]), num_contexts=1, diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index d962962761d2..d1f57dffc375 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -27,7 +27,6 @@ DeepseekV4CacheManager, DeepseekV4Indexer, DeepseekV4TrtllmAttention, - DeepseekV4TrtllmAttentionMetadata, ) from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.compressor import Compressor from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import ( @@ -1119,8 +1118,8 @@ def test_deepseek_v4_sparse_ratios_prefer_checkpoint_defaults(tmp_path, monkeypa assert model_config.sparse_attention_config.compress_ratios == [128, 128, 4, 128, 4, 128, 1, 4] # V4 sparse MLA hardcodes window_size==128 (FMHA kernel TileSizeKV; see - # the runtime assertion in deepseek_v4.py:DeepseekV4TrtllmAttentionMetadata - # __post_init__), so this is the only legal value here. + # the runtime assertion in DeepseekV4Metadata.allocate), so this is the + # only legal value here. assert model_config.sparse_attention_config.window_size == 128 @@ -1393,7 +1392,7 @@ def test_deepseek_v4_sanity( assert success, f"Failed to resize context for request {req_id}" reqs.append(req) - attn_metadata = DeepseekV4TrtllmAttentionMetadata( + attn_metadata = DeepseekV4TrtllmAttention.create_metadata( seq_lens=torch.tensor(sequence_length, dtype=torch.int32), num_contexts=num_contexts, max_num_requests=len(sequence_length), @@ -1406,7 +1405,7 @@ def test_deepseek_v4_sanity( prompt_lens=prompt_lens, max_num_tokens=8192, mapping=mapping, - sparse_attention_config=sparse_attn_config, + sparse_metadata_params=sparse_attn_config.to_sparse_metadata_params(), ) position_ids = [] diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv41.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv41.py index 159c789db422..ae43e4fea861 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv41.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv41.py @@ -23,12 +23,12 @@ import pytest import torch -from _torch.attention.sparse.csa2._utils import _FakeMetadata +from _torch.attention.sparse.csa2._utils import _bare_csa2_metadata, _FakeMetadata from torch import nn from transformers import AutoConfig from utils.util import skip_blackwell_geforce, skip_pre_blackwell -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.configs import deepseek_v41 as v41_config from tensorrt_llm._torch.configs.deepseek_v41 import ( DeepseekV41Config, @@ -1902,14 +1902,19 @@ def shared_walk(self, weights, skip_modules): class _ReplayMetadata(_FakeMetadata): - _get_csa2_buffer = CSA2TrtllmMetadata._get_csa2_buffer - _copy_host = CSA2TrtllmMetadata._copy_host - _copy_csa2_tensor = CSA2TrtllmMetadata._copy_csa2_tensor - get_adp_token_counts = CSA2TrtllmMetadata.get_adp_token_counts - set_adp_token_counts = CSA2TrtllmMetadata.set_adp_token_counts + _get_csa2_buffer = CSA2Metadata._get_csa2_buffer + _copy_host = CSA2Metadata._copy_host + _copy_csa2_tensor = CSA2Metadata._copy_csa2_tensor decoder_replay_plan = None padded_num_tokens = None + def get_adp_token_counts(self) -> list[int]: + return [self.num_tokens, *CSA2Metadata.adp_token_counts(self)] + + def set_adp_token_counts(self, counts: list[list[int]]) -> list[int]: + CSA2Metadata.set_adp_token_counts(self, counts[1:], counts[0]) + return counts[0] + def __init__(self, lengths, num_contexts): super().__init__(lengths, num_contexts) self.all_rank_num_tokens = None @@ -2056,7 +2061,7 @@ def test_remote_tail_attention_dp_never_enables_local_row_compaction( assert model._plan_bounded_replay(metadata, all_token_states_required) == (None, None) torch.testing.assert_close(metadata.seq_lens, saved_lengths) - assert metadata.csa2_replay_query_rows is None + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None assert metadata.prepare_calls == 0 @@ -2116,7 +2121,7 @@ def test_adp_replay_counts_follow_final_execution(monkeypatch, mode, speculative metadata.mapping = Mapping(world_size=2, tp_size=2, rank=0, enable_attention_dp=True) metadata.kv_cache_manager.has_private_swa_suffix = Mock(return_value=True) if mode in ("final_window", "all_empty"): - metadata.decoder_context_ends = (1024,) + metadata.ext(CSA2Metadata).decoder_context_ends = (1024,) local_rows = 512 if mode == "padded" else 257 metadata.padded_num_tokens = 512 if mode == "padded" else None metadata.is_cuda_graph = mode == "decode_graph" @@ -2185,7 +2190,7 @@ def test_adp_replay_counts_follow_final_execution(monkeypatch, mode, speculative assert metadata.all_rank_num_tokens is saved_counts return assert split == 20 - assert plan is metadata.decoder_replay_plan + assert plan is metadata.ext(CSA2Metadata).decoder_replay_plan assert plan.replay_all_rank_num_tokens == [decoder_rows, peer_rows] assert plan.replays_local_tokens is (mode in ("compact", "final_window", "all_empty")) if not plan.replays_local_tokens: @@ -2233,14 +2238,14 @@ def test_final_window_replay_freezes_endpoints_and_preserves_mixed_consumers(ful SimpleNamespace(prompt_len=50, py_ced_replay=None, py_return_context_logits=full_output), SimpleNamespace(prompt_len=6, py_ced_replay=None), ] - CSA2TrtllmMetadata.prepare_context_replay(metadata, reqs) + CSA2Metadata.prepare_context_replay(metadata, reqs) reqs[0].prompt_len = 10 # Later scheduling must not change this forward's selection. plan = plan_decoder_replay(metadata, 8, reqs, private_decoder=True) assert plan.replay_seq_lens.tolist() == ([10, 3, 6] if full_output else [0, 3, 6]) assert plan.swa_floors == ([0, 0, 0] if full_output else [42, 0, 0]) assert plan.rows.tolist() == list(range(0 if full_output else 10, 19)) enter_decoder_replay(metadata, plan) - assert metadata.csa2_positions[-6:].tolist() == list(range(97, 103)) + assert metadata.ext(CSA2Metadata).csa2_positions[-6:].tolist() == list(range(97, 103)) exit_decoder_replay(metadata, plan) assert metadata.seq_lens.tolist() == [10, 3, 6] assert metadata.kv_cache_params.num_cached_tokens_per_seq == [0, 3, 100] @@ -2258,7 +2263,7 @@ def test_final_window_zero_queries_still_publish_global_and_restore_outputs(peer layer, state = _make_boundary(129) model.layers[20] = layer metadata = _ReplayMetadata([129], num_contexts=1) - metadata.decoder_context_ends = (512,) + metadata.ext(CSA2Metadata).decoder_context_ends = (512,) plan = plan_decoder_replay(metadata, 128) plan.saved_all_rank_num_tokens = [129, 257] plan.replay_all_rank_num_tokens = [0, peer_tokens] @@ -2267,13 +2272,13 @@ def test_final_window_zero_queries_still_publish_global_and_restore_outputs(peer assert replay_state.residual.shape == (0, 4, 64) layer.self_attn.prepare_global_cache.assert_called_once() assert layer.self_attn.prepare_global_cache.call_args.args[0].shape[0] == 129 - assert metadata.csa2_precomputed_kv_layers == {20} + assert metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers == {20} assert metadata.prepare_calls == 0 assert metadata.all_rank_num_tokens == [0, peer_tokens] restored = model._exit_bounded_replay(plan, metadata, replay_state.residual[:, 0]) assert restored.shape == (129, 64) assert not restored.count_nonzero() - assert metadata.csa2_decoder_capture_lens == (0,) + assert metadata.ext(CSA2Metadata).csa2_decoder_capture_lens == (0,) assert metadata.all_rank_num_tokens == [129, 257] assert metadata.seq_lens.tolist() == [129] assert metadata.kv_cache_params.num_cached_tokens_per_seq == [0] @@ -2286,7 +2291,7 @@ def test_final_window_prefill_graph_capability_keeps_intermediate_history(): metadata = _ReplayMetadata([257], num_contexts=1) reqs = [SimpleNamespace(prompt_len=2048, py_ced_replay=None)] - CSA2TrtllmMetadata.prepare_context_replay(metadata, reqs, allow_final_window=False) + CSA2Metadata.prepare_context_replay(metadata, reqs, allow_final_window=False) plan = plan_decoder_replay(metadata, 128, reqs, private_decoder=True) assert plan.replay_seq_lens.tolist() == [128] assert plan.rows.tolist() == list(range(129, 257)) @@ -2362,7 +2367,7 @@ def _make_boundary(tokens: int) -> tuple[SimpleNamespace, HCState]: def publish(values, metadata): assert values.shape[0] == int(metadata.seq_lens.sum()) - metadata.csa2_precomputed_kv_layers.add(20) + metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers.add(20) attention = SimpleNamespace( layer=SimpleNamespace(kv_source=20, compress_ratio=1, mode=CSA2Mode.FULL), @@ -2404,8 +2409,8 @@ def test_model_boundary_preserves_full_cache_and_compacts_query_rows( new_positions, new_ids, new_state = model._enter_bounded_replay( plan, metadata, positions, ids, state ) - assert metadata.csa2_precomputed_kv_layers == {20} - assert metadata.csa2_replay_query_rows is plan.rows + assert metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers == {20} + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is plan.rows writer = layer.self_attn.prepare_global_cache writer.assert_called_once() actual, passed_metadata = writer.call_args.args @@ -2428,11 +2433,11 @@ def test_model_boundary_preserves_full_cache_and_compacts_query_rows( torch.testing.assert_close(restored[plan.rows], new_state.residual) assert metadata.seq_lens.tolist() == lengths assert metadata.kv_cache_params.num_cached_tokens_per_seq == [0, 0, 0] - assert metadata.csa2_replay_query_rows is None - assert not metadata.csa2_precomputed_kv_layers + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None + assert not metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers metadata.prepare() - assert not metadata.csa2_precomputed_kv_layers - assert not metadata.csa2_indices + assert not metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers + assert not metadata.ext(CSA2Metadata).csa2_indices assert metadata.seq_lens.tolist() == lengths assert metadata.kv_cache_params.num_cached_tokens_per_seq == [0, 0, 0] monkeypatch.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "0") @@ -2475,7 +2480,7 @@ def test_remote_tail_source_reuses_boundary_collapse( metadata = _ReplayMetadata(lengths, num_contexts=num_contexts) metadata.mapping = model.model_config.mapping metadata.csa2_remote_tail_mode = None if context_only else "source" - metadata.begin_model_forward = Mock() + metadata.begin_forward = Mock() post_mapping = Mock(return_value=resolved.residual) state = resolved if deferred: @@ -2513,17 +2518,17 @@ def test_remote_tail_source_reuses_boundary_collapse( assert post_args["post_layer_mix"] is state.post_mix assert post_args["comb_res_mix"] is state.comb_mix encoder.assert_called_once() - metadata.begin_model_forward.assert_called_once() + metadata.begin_forward.assert_called_once() writer = layer.self_attn.prepare_global_cache if ced_precompute: writer.assert_called_once() actual, passed_metadata = writer.call_args.args torch.testing.assert_close(actual, layer.input_layernorm(expected), atol=0, rtol=0) assert passed_metadata is metadata - assert metadata.csa2_precomputed_kv_layers == {20} + assert metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers == {20} else: writer.assert_not_called() - assert not metadata.csa2_precomputed_kv_layers + assert not metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers collapse.reset_mock() finalize.reset_mock() @@ -2545,9 +2550,9 @@ def test_remote_tail_destination_keeps_decoder_state(attention_dp: bool) -> None metadata.csa2_remote_tail_mode = "destination" metadata.csa2_remote_tail_starts = [1024] saved_lengths = metadata.seq_lens.clone() - saved_positions = metadata.csa2_positions.clone() - metadata._ensure_swa_slots() # As the encoder's first layer does. - saved_indices = metadata.csa2_swa_indices[0].clone() + saved_positions = metadata.ext(CSA2Metadata).csa2_positions.clone() + metadata.ext(CSA2Metadata)._ensure_swa_slots() # As the encoder's first layer does. + saved_indices = metadata.ext(CSA2Metadata).csa2_swa_indices[0].clone() saved_residual = state.residual.clone() saved_pre_mix = state.pre_mix.clone() @@ -2556,18 +2561,18 @@ def test_remote_tail_destination_keeps_decoder_state(attention_dp: bool) -> None assert result is None assert decoder_state is state layer.self_attn.prepare_global_cache.assert_called_once() - assert metadata.csa2_precomputed_kv_layers == {20} + assert metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers == {20} torch.testing.assert_close(state.residual, saved_residual, atol=0, rtol=0) torch.testing.assert_close(state.pre_mix, saved_pre_mix, atol=0, rtol=0) torch.testing.assert_close(metadata.seq_lens, saved_lengths) - torch.testing.assert_close(metadata.csa2_positions, saved_positions) + torch.testing.assert_close(metadata.ext(CSA2Metadata).csa2_positions, saved_positions) assert metadata.kv_cache_params.num_cached_tokens_per_seq == [1024, 63] - assert metadata.csa2_replay_query_rows is None + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None assert metadata.prepare_calls == 1 expected_indices = saved_indices.clone() expected_indices[:128].masked_fill_(expected_indices[:128] < 1024, -1) - metadata._ensure_swa_slots() # As the first decoder layer does. - torch.testing.assert_close(metadata.csa2_swa_indices[0], expected_indices) + metadata.ext(CSA2Metadata)._ensure_swa_slots() # As the first decoder layer does. + torch.testing.assert_close(metadata.ext(CSA2Metadata).csa2_swa_indices[0], expected_indices) @pytest.mark.parametrize( @@ -2606,7 +2611,7 @@ def test_remote_tail_attention_dp_generation_keeps_all_layers( metadata.is_cuda_graph = is_cuda_graph metadata.csa2_remote_tail_mode = remote_mode metadata.csa2_remote_tail_starts = [1024] if remote_mode == "destination" else [] - metadata.begin_model_forward = Mock() + metadata.begin_forward = Mock() enter_decoder = Mock() monkeypatch.setattr(v41, "enter_remote_tail_decoder", enter_decoder) @@ -2622,8 +2627,8 @@ def test_remote_tail_attention_dp_generation_keeps_all_layers( assert layer.call_args.kwargs["hc_state"].residual.shape[0] == num_tokens assert enter_decoder.call_count == int(remote_mode == "destination") assert boundary.self_attn.prepare_global_cache.call_count == int(remote_mode == "destination") - assert metadata.csa2_replay_query_rows is None - metadata.begin_model_forward.assert_called_once() + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None + metadata.begin_forward.assert_called_once() def test_context_decoder_skipping_requires_bounded_replay( @@ -2685,7 +2690,7 @@ def test_context_only_forward_routes_warmup_to_source_and_refuses_destination(mo nn.Module.__init__(model) model.disagg_context_only = True metadata = SimpleNamespace( - csa2_remote_tail_mode=mode, begin_model_forward=Mock(), all_rank_num_tokens=None + csa2_remote_tail_mode=mode, begin_forward=Mock(), all_rank_num_tokens=None ) result = object() @@ -2701,7 +2706,7 @@ def source_forward(self, passed_metadata): model.forward(metadata) else: assert model.forward(metadata) is result - metadata.begin_model_forward.assert_called_once() + metadata.begin_forward.assert_called_once() def test_context_only_model_rejects_speculative_decoding_before_construction(monkeypatch): @@ -3006,8 +3011,8 @@ def test_model_boundary_resets_routing_and_preserves_layer_guard(monkeypatch): from tensorrt_llm._torch.models.modeling_deepseekv4 import DeepseekV4Model from tensorrt_llm._torch.models.modeling_deepseekv41 import DeepseekV41Model - metadata = CSA2TrtllmMetadata.__new__(CSA2TrtllmMetadata) - metadata.reset_routing() + metadata = _bare_csa2_metadata() + metadata.ext(CSA2Metadata).reset_routing() layout = CSA2Layout((1, 1), (0,), (0,)) model = DeepseekV41Model.__new__(DeepseekV41Model) torch.nn.Module.__init__(model) @@ -3018,27 +3023,30 @@ def test_model_boundary_resets_routing_and_preserves_layer_guard(monkeypatch): def body(self, attn_metadata, value, *, check_value): assert self is model and attn_metadata is metadata assert value is argument and check_value is argument - assert not metadata.csa2_indices and not metadata.csa2_candidates - assert not metadata.csa2_precomputed_kv_layers - assert metadata.csa2_replay_query_rows is None - metadata.enter_layer(layout.layer(0)) + assert ( + not metadata.ext(CSA2Metadata).csa2_indices + and not metadata.ext(CSA2Metadata).csa2_candidates + ) + assert not metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None + metadata.ext(CSA2Metadata).enter_layer(layout.layer(0)) producer = torch.tensor([len(outputs)]) - metadata.csa2_indices[0] = producer - metadata.csa2_candidates[0] = producer - metadata.enter_layer(layout.layer(1)) - assert metadata.csa2_indices[0] is producer + metadata.ext(CSA2Metadata).csa2_indices[0] = producer + metadata.ext(CSA2Metadata).csa2_candidates[0] = producer + metadata.ext(CSA2Metadata).enter_layer(layout.layer(1)) + assert metadata.ext(CSA2Metadata).csa2_indices[0] is producer with pytest.raises(ValueError, match="reordered layers"): - metadata.enter_layer(layout.layer(0)) + metadata.ext(CSA2Metadata).enter_layer(layout.layer(0)) outputs.append(producer) return producer monkeypatch.setattr(DeepseekV4Model, "forward", body) for _ in range(3): - old_indices = metadata.csa2_indices - metadata.csa2_precomputed_kv_layers = {0} - metadata.csa2_replay_query_rows = torch.tensor([1]) + old_indices = metadata.ext(CSA2Metadata).csa2_indices + metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers = {0} + metadata.ext(CSA2Metadata).csa2_replay_query_rows = torch.tensor([1]) result = model.forward(metadata, argument, check_value=argument) - assert metadata.csa2_indices is not old_indices + assert metadata.ext(CSA2Metadata).csa2_indices is not old_indices assert result is outputs[-1] assert [output.item() for output in outputs] == [0, 1, 2] diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index a7b79b783d2b..d8f989c8c828 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -15,6 +15,8 @@ from utils.util import default_dtype, getSMVersion import tensorrt_llm +from tensorrt_llm._torch.attention.backends.features.spec_decoding import \ + SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -546,12 +548,13 @@ def run_forward(input_ids, position_ids, attn_metadata): kv_cache_manager=kv_cache_manager, request_ids=request_ids, prompt_lens=prompt_lens, - is_spec_decoding_enabled=is_spec_decoding_enabled, - use_spec_decoding=use_spec_decoding, - is_spec_dec_tree=is_spec_dec_tree, - is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, max_num_heads_per_kv=num_heads_per_kv, ) + spec_phase_0 = attn_metadata_gen_phase_0.ext(SpecDecodingMetadata) + spec_phase_0.is_spec_decoding_enabled = is_spec_decoding_enabled + spec_phase_0.use_spec_decoding = use_spec_decoding + spec_phase_0.is_spec_dec_tree = is_spec_dec_tree + spec_phase_0.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree cache = kv_cache_manager.kv_cache_map[request.py_request_id] assert cache.resize(input_ids.size(-1) + gen_input_ids_0.size(-1)) attn_metadata_gen_phase_0.prepare() @@ -605,9 +608,9 @@ def run_forward(input_ids, position_ids, attn_metadata): [gen_input_ids_1.size(-1)], dtype=torch.int) attn_metadata_gen_phase_0.kv_cache_params.num_cached_tokens_per_seq = num_cached_tokens_per_seq_1 - attn_metadata_gen_phase_0.spec_decoding_position_offsets = None - attn_metadata_gen_phase_0.spec_decoding_packed_mask = None - attn_metadata_gen_phase_0.spec_decoding_generation_lengths = None + spec_phase_0.position_offsets = None + spec_phase_0.packed_mask = None + spec_phase_0.generation_lengths = None assert cache.resize(num_cached_tokens_per_seq_1[0] + gen_input_ids_1.size(-1)) attn_metadata_gen_phase_0.prepare() @@ -672,16 +675,12 @@ def run_forward(input_ids, position_ids, attn_metadata): kv_cache_manager=kv_cache_manager, request_ids=request_ids, prompt_lens=prompt_lens, - is_spec_decoding_enabled=is_spec_decoding_enabled, - use_spec_decoding=use_spec_decoding, - is_spec_dec_tree=is_spec_dec_tree, - is_spec_dec_dynamic_tree=False, max_num_heads_per_kv=num_heads_per_kv, ) - - attn_metadata_ref.spec_decoding_position_offsets = None - attn_metadata_ref.spec_decoding_packed_mask = None - attn_metadata_ref.spec_decoding_generation_lengths = None + spec_ref = attn_metadata_ref.ext(SpecDecodingMetadata) + spec_ref.is_spec_decoding_enabled = is_spec_decoding_enabled + spec_ref.use_spec_decoding = use_spec_decoding + spec_ref.is_spec_dec_tree = is_spec_dec_tree attn_metadata_ref.prepare() is_tree_ref = is_spec_dec_tree if get_sm_version() < 100 else False spec_tree_mgr_ref = None diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index 3f20cbdaff18..2be28fac059f 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -9,7 +9,7 @@ from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.hooks import get_sparse_attention_hooks from tensorrt_llm._torch.attention.backends.sparse.qsa import ( - QSAAttentionMetadata, + QSAMetadata, QSASparseMetadataParams, QSASparseParams, ) @@ -201,7 +201,8 @@ def test_qsa_sparse_hook_is_registered() -> None: def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(QSAMetadata(metadata)) metadata._num_tokens = 0 output = QSASparseHooks().forward( @@ -224,20 +225,20 @@ def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: assert output is None -def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers(monkeypatch) -> None: - monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) +def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers() -> None: manager = object.__new__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = None - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) metadata.kv_cache_manager = manager metadata.sparse_metadata_params = QSASparseMetadataParams( token_topk=2048, compress_ratio=4, ) - metadata.__post_init__() + qsa = QSAMetadata(metadata) + qsa.allocate(None) - assert not metadata.qsa_has_local_layers + assert not qsa.has_local_layers def test_qsa_hybrid_routes_to_sparse_v2_cache_manager(monkeypatch) -> None: @@ -403,7 +404,8 @@ def regular_get_buffers( def test_qsa_uses_regular_attention_for_unsupported_head_dim() -> None: """Triton indexes head dims with tl.arange, so non-powers of two go dense.""" - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(QSAMetadata(metadata)) metadata._num_tokens = 1 metadata.kv_cache_manager = SimpleNamespace(dtype=DataType.BF16) # head_dim comes from the module, not from q: q is still the fused QKV here. @@ -432,7 +434,8 @@ def test_qsa_uses_regular_attention_for_unsupported_head_dim() -> None: def test_qsa_uses_regular_attention_for_scale_paged_kv_cache() -> None: - metadata = object.__new__(QSAAttentionMetadata) + metadata = object.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(QSAMetadata(metadata)) metadata._num_tokens = 1 metadata.kv_cache_manager = SimpleNamespace(dtype=DataType.NVFP4) indexer = SimpleNamespace(project_and_update_cache=lambda *args: torch.empty((1, 4, 128))) diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 7dac22c46df8..4ed031892107 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -10,6 +10,16 @@ import torch from torch import nn +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata + + +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + def _text_config_dict() -> dict: return { @@ -289,7 +299,7 @@ def capture_ple_ids( monkeypatch.setattr(Qwen4ExpModel, "_prepare_ple_state", capture_ple_ids) orig_input_ids = torch.tensor([10, 11, 12]) inputs_embeds = torch.randn(3, 8) - attn_metadata = SimpleNamespace(mamba_metadata=None, max_num_requests=1) + attn_metadata = _recurrent_stub(mamba_metadata=None, max_num_requests=1) output = model( attn_metadata=attn_metadata, @@ -309,7 +319,7 @@ def test_multimodal_forward_requires_original_token_ids_on_a_ple_rank() -> None: model.model_config = SimpleNamespace(mapping=SimpleNamespace(has_pp=lambda: False)) model.has_ple = True model.ple_layer_index = 0 - attn_metadata = SimpleNamespace(mamba_metadata=None, max_num_requests=1) + attn_metadata = _recurrent_stub(mamba_metadata=None, max_num_requests=1) with pytest.raises(ValueError, match="requires original token IDs"): model( @@ -361,7 +371,7 @@ def forward(self, **kwargs: object) -> None: with pytest.raises(RuntimeError, match="simulated layer failure"): model( - attn_metadata=SimpleNamespace(mamba_metadata=None, max_num_requests=1), + attn_metadata=_recurrent_stub(mamba_metadata=None, max_num_requests=1), inputs_embeds=torch.randn(1, 8), orig_input_ids=torch.ones(1, dtype=torch.long), ) @@ -1372,7 +1382,7 @@ def update_mamba_states(self, **kwargs) -> None: worker._is_mamba_hybrid_cache = None state_indices = torch.tensor([3], dtype=torch.int32) accepted = torch.tensor([2], dtype=torch.int32) - attn_metadata = SimpleNamespace( + attn_metadata = _recurrent_stub( num_contexts=0, kv_cache_manager=manager, mamba_metadata=SimpleNamespace(state_indices=state_indices), diff --git a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py index 3f32aee7d0e4..893b3dcda2f1 100644 --- a/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py +++ b/tests/unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/dsa/test_modeling_v2_hist_buffer_lifecycle.py @@ -45,7 +45,7 @@ def _target() -> DecodeTarget: def _metadata(max_seq_len: int, kv_len: int, tokens_per_block: int = 64) -> SimpleNamespace: pages = -(-max_seq_len // tokens_per_block) - return SimpleNamespace( + metadata = SimpleNamespace( kv_cache_manager=SimpleNamespace( max_seq_len=max_seq_len, tokens_per_block=tokens_per_block ), @@ -55,6 +55,9 @@ def _metadata(max_seq_len: int, kv_len: int, tokens_per_block: int = 64) -> Simp torch.arange(pages, dtype=torch.int32, device="cuda") + 1000 ).view(1, -1), ) + # The stub carries the DSA extension state itself. + metadata.ext = lambda cls: metadata + return metadata def _assert_metadata(target: DecodeTarget, metadata: SimpleNamespace) -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py index 548008dacdf4..3715eee9c918 100644 --- a/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py +++ b/tests/unittest/_torch/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/test_modeling_v2_index_share.py @@ -22,6 +22,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1.ops.attention.qsa_mtp_reuse_topk import ( # noqa: E501 qsa_mtp_reuse_topk as reuse, ) +from tensorrt_llm._torch.attention.backends.sparse.qsa.metadata import QSAMetadata assert torch.cuda.is_available(), "the IndexShare plumbing runs on a CUDA device" @@ -30,6 +31,20 @@ SLOT = 1 +def _qsa_stub(**fields) -> SimpleNamespace: + """A metadata stub whose QSA extension holds the ``qsa_``-prefixed fields.""" + qsa = SimpleNamespace( + **{ + name[len("qsa_") :]: fields.pop(name) + for name in list(fields) + if name.startswith("qsa_") + } + ) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: qsa if cls is QSAMetadata else None + return metadata + + def _target(cls): target = object.__new__(cls) target.core = SimpleNamespace(_draft_bound=SimpleNamespace(share_slot=SLOT)) @@ -37,7 +52,7 @@ def _target(cls): def _metadata(batch: int, accepted, seq_lens=None, num_contexts: int = 0): - return SimpleNamespace( + return _qsa_stub( num_seqs=batch, num_contexts=num_contexts, seq_lens_cuda=None @@ -92,17 +107,18 @@ def test_decode_capture_at_draft_seven_takes_row_s_times_8_plus_a_minus_1(): topk = torch.arange(batch * p * K, dtype=torch.int32, device=DEV).reshape(batch * p, K) visible = torch.arange(300, 300 + batch * p, dtype=torch.int32, device=DEV) md = _metadata(batch, accepted) + qsa = md.find_ext(QSAMetadata) _target(m.DecodeTarget)._capture_selection(md, topk, visible, batch, p) rows = [s * p + a - 1 for s, a in enumerate(accepted)] assert rows == [0, 15, 20, 26] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :batch], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :batch], visible[rows]) + torch.testing.assert_close(qsa.shared_topk_indices[SLOT, :batch], topk[rows]) + torch.testing.assert_close(qsa.shared_topk_visible_blocks[SLOT, :batch], visible[rows]) # Other slots and rows past the batch are untouched. - assert torch.all(md.qsa_shared_topk_indices[SLOT, batch:] == -7) - assert torch.all(md.qsa_shared_topk_indices[0] == -7) - assert md.qsa_shared_topk_captured_slots == {SLOT} + assert torch.all(qsa.shared_topk_indices[SLOT, batch:] == -7) + assert torch.all(qsa.shared_topk_indices[0] == -7) + assert qsa.shared_topk_captured_slots == {SLOT} def test_prefill_capture_at_draft_seven(): @@ -115,13 +131,14 @@ def test_prefill_capture_at_draft_seven(): topk = torch.arange(rows_total * K, dtype=torch.int32, device=DEV).reshape(rows_total, K) visible = torch.arange(400, 400 + rows_total, dtype=torch.int32, device=DEV) md = _metadata(5, accepted, seq_lens=seq_lens, num_contexts=2) + qsa = md.find_ext(QSAMetadata) _target(m.PrefillTarget)._capture_selection(md, topk, visible) rows = [4, 7, 8 + 7, 16 + 0, 24 + 5] - torch.testing.assert_close(md.qsa_shared_topk_indices[SLOT, :5], topk[rows]) - torch.testing.assert_close(md.qsa_shared_topk_visible_blocks[SLOT, :5], visible[rows]) - assert md.qsa_shared_topk_captured_slots == {SLOT} + torch.testing.assert_close(qsa.shared_topk_indices[SLOT, :5], topk[rows]) + torch.testing.assert_close(qsa.shared_topk_visible_blocks[SLOT, :5], visible[rows]) + assert qsa.shared_topk_captured_slots == {SLOT} def _reuse_case(fresh: int): diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index 5e3ce1d1292b..4c5becc910a2 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -38,6 +38,7 @@ from utils.llm_data import llm_models_root import tensorrt_llm._torch.models.modeling_minimaxm3 as modeling_minimaxm3 +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( @@ -54,9 +55,10 @@ index_head_range, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( - MiniMaxM3MsaSparseAttentionMetadata, + MiniMaxM3MsaMetadata, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_indexer import _group_max_reduce +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.compilation.remove_copy_pass import remove_copy_for_mutates_args from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.minimaxm3_weight_mapper import ( @@ -725,10 +727,9 @@ def test_piecewise_captures_horizontal_producer_before_attention( index_cache = torch.empty_like(kv_cache) scale_cache = torch.empty(8, dtype=torch.uint8) if kv_dtype == "nvfp4" else None slots = torch.tensor([0, 1, -1, -1], dtype=torch.int32) - metadata = Mock( - spec_set=MiniMaxM3MsaSparseAttentionMetadata, - num_tokens=2, - kv_cache_manager=None, + metadata = Mock(spec_set=TrtllmAttentionMetadata, num_tokens=2, kv_cache_manager=None) + metadata.ext.return_value = Mock( + spec_set=MiniMaxM3MsaMetadata, msa_layer_cache_tensors={3: (kv_cache, index_cache, scale_cache)}, msa_out_cache_loc=slots, ) @@ -979,10 +980,9 @@ def test_piecewise_attention_boundary_preserves_indexer_cache_contract( if caches_prewritten: index_cache.copy_(expected_cache) selected_blocks = torch.zeros(2, 1, 16, dtype=torch.int32) - attn_metadata = Mock( - spec_set=MiniMaxM3MsaSparseAttentionMetadata, - num_tokens=2, - kv_cache_manager=None, + attn_metadata = Mock(spec_set=TrtllmAttentionMetadata, num_tokens=2, kv_cache_manager=None) + msa = Mock( + spec_set=MiniMaxM3MsaMetadata, msa_decode_span=None, msa_idx_k_cache=Mock(return_value=index_cache), msa_write_idx_k=Mock(), @@ -993,6 +993,7 @@ def test_piecewise_attention_boundary_preserves_indexer_cache_contract( msa_kv_lens_cpu=torch.tensor([2], dtype=torch.int32), msa_qo_offset_cpu=torch.tensor([0], dtype=torch.int32), ) + attn_metadata.ext.return_value = msa def write_caches( live_k: torch.Tensor, @@ -1049,8 +1050,8 @@ def select_blocks( ) assert backend.write_layer_caches.call_count == (0 if caches_prewritten else 1) - attn_metadata.msa_write_idx_k.assert_not_called() - attn_metadata.msa_idx_k_cache.assert_called_once_with(3) + msa.msa_write_idx_k.assert_not_called() + msa.msa_idx_k_cache.assert_called_once_with(3) backend.indexer.select_blocks.assert_called_once() backend.forward.assert_called_once() called_q, called_k, called_v, called_metadata = backend.forward.call_args.args @@ -1996,12 +1997,10 @@ def test_horizontal_producer_explicit_caches_match_manager( manager.is_nvfp4_layer.return_value = nvfp4 manager.get_buffers.return_value = main manager.get_block_scale_buffers.return_value = scales - metadata = Mock( - spec_set=MiniMaxM3MsaSparseAttentionMetadata, - kv_cache_manager=manager, - msa_out_cache_loc=slots, - ) - metadata.msa_idx_k_cache.return_value = index + metadata = Mock(spec_set=TrtllmAttentionMetadata, kv_cache_manager=manager) + msa = Mock(spec_set=MiniMaxM3MsaMetadata, msa_out_cache_loc=slots) + msa.msa_idx_k_cache.return_value = index + metadata.ext.return_value = msa def snapshot(outputs: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, ...]: return tuple(t.view(torch.uint8).clone() for t in (*outputs, main, index, scales)) @@ -2020,7 +2019,7 @@ def snapshot(outputs: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.testing.assert_close(actual, reference, rtol=0, atol=0) manager.get_buffers.assert_called_once_with(3, kv_layout="HND") assert manager.get_block_scale_buffers.call_count == int(nvfp4) - metadata.msa_idx_k_cache.assert_called_once_with(3) + msa.msa_idx_k_cache.assert_called_once_with(3) monkeypatch.setattr( modeling_minimaxm3, @@ -2969,9 +2968,7 @@ def test_minimax_m3_dense_backend_library_order( metadata.cu_q_seqlens = metadata.cu_kv_seqlens = torch.tensor([0, 2], dtype=torch.int32) metadata.num_contexts = 1 metadata.num_generations = 0 - metadata.enable_flash_mla = False - metadata.spec_bl_tree_first_sparse_mask_offset_kv = None - metadata.spec_decoding_bl_tree_mask = None + metadata.replace_ext(SpecDecodingMetadata(metadata)) metadata.max_context_q_len_override = None metadata.kv_cache_manager = SimpleNamespace( dtype=DataType.BF16, max_attention_window_vec=[128], max_seq_len=128 diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py index b683582dc1b2..9c361579f204 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py @@ -11,6 +11,17 @@ from kda_module_test_utils import compare, make_attention, make_batch +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata + + +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + + pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3), (10, 7)), @@ -301,7 +312,7 @@ def test_production_forward_updates_live_cache(state_dtype: torch.dtype) -> None **replay_caches, ) old_replay = [getattr(cache, f"kda_conv_{label}").clone() for label in ("q", "k", "v")] - metadata = SimpleNamespace( + metadata = _recurrent_stub( mamba_metadata=batch.metadata, num_contexts=batch.n, num_ctx_tokens=batch.hidden.shape[0], diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py index 8d7270cdff3a..74792114682d 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mixed_batch_core_layout.py @@ -20,8 +20,18 @@ pytest.importorskip("fla") +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention # noqa: E402 + +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + + HIDDEN_SIZE = 512 NUM_HEADS = 4 HEAD_DIM = 128 @@ -87,7 +97,7 @@ def _mixed_metadata(num_ctx_tokens: int, num_decodes: int) -> SimpleNamespace: kda_single_sequence_length=num_ctx_tokens, _arange_buffer=torch.arange(batch + 1, dtype=torch.int32, device=device), ) - return SimpleNamespace( + return _recurrent_stub( mamba_metadata=mamba, num_contexts=num_prefills, num_ctx_tokens=num_ctx_tokens, diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 5a06177217bb..aebdbebcb624 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -26,11 +26,20 @@ from kimi_kda_test_utils import get_production_decode_kernel_path from tensorrt_llm._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention from tensorrt_llm._torch.modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + + class _Cfg: hidden_size = 256 rms_norm_eps = 1e-6 @@ -75,7 +84,7 @@ def forward_decode(hidden_states, conv_pool, ssm_pool, slot_indices, *args, **kw attention.head_k_dim, ), ) - metadata = SimpleNamespace( + metadata = _recurrent_stub( mamba_metadata=SimpleNamespace( state_indices=state_indices, query_start_loc_long=torch.zeros(1, dtype=torch.long), @@ -155,7 +164,8 @@ def mamba_layer_cache(self, layer_idx: int) -> SimpleNamespace: ), ) mamba_metadata.prepare(metadata) - metadata.mamba_metadata = mamba_metadata + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata.ext = lambda cls: recurrent hidden_states = torch.empty(4, _Cfg.hidden_size, dtype=torch.bfloat16, device="cuda") generation_slice = mamba_metadata.state_indices[1:] aligned_indices = mamba_metadata.generation_state_indices @@ -193,7 +203,7 @@ def _decode_metadata(layer_cache: SimpleNamespace, slot_indices: torch.Tensor) - _arange_buffer=torch.arange(batch + 1, device=slot_indices.device, dtype=torch.int32), ) cache_manager = SimpleNamespace(mamba_layer_cache=lambda _: layer_cache) - return SimpleNamespace( + return _recurrent_stub( mamba_metadata=mamba_metadata, num_contexts=0, num_ctx_tokens=0, @@ -215,7 +225,7 @@ def _prefill_metadata( use_initial_states=False, ) cache_manager = SimpleNamespace(mamba_layer_cache=lambda _: layer_cache) - return SimpleNamespace( + return _recurrent_stub( mamba_metadata=mamba_metadata, num_contexts=batch, num_ctx_tokens=int(cu_seqlens[-1]), diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py index 4d651c7fd592..a9b4c48ed30a 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py @@ -568,7 +568,10 @@ def test_gdn_custom_op_forwards_split_inputs(monkeypatch): class FakeMetadata: num_tokens = 2 - mamba_metadata = object() + recurrent = SimpleNamespace(mamba_metadata=object()) + + def ext(self, cls): + return self.recurrent class FakeLayer: def forward_core( @@ -585,7 +588,7 @@ def forward_core( assert a.shape[0] == 2 assert b.shape[0] == 2 assert attn_metadata is metadata - assert mamba_metadata is metadata.mamba_metadata + assert mamba_metadata is metadata.recurrent.mamba_metadata assert spec_metadata is None output.fill_(5) diff --git a/tests/unittest/_torch/multi_gpu/test_deepseek_v41_attention_dp.py b/tests/unittest/_torch/multi_gpu/test_deepseek_v41_attention_dp.py index 44bed8a2b390..94bb880d847a 100644 --- a/tests/unittest/_torch/multi_gpu/test_deepseek_v41_attention_dp.py +++ b/tests/unittest/_torch/multi_gpu/test_deepseek_v41_attention_dp.py @@ -24,8 +24,10 @@ from utils.util import skip_pre_blackwell import tensorrt_llm -from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2TrtllmMetadata +from tensorrt_llm._torch.attention.backends.sparse.csa2.backend import CSA2TrtllmAttention +from tensorrt_llm._torch.attention.backends.sparse.csa2.metadata import CSA2Metadata from tensorrt_llm._torch.attention.backends.sparse.csa2.params import CSA2Layout +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.distributed import AllReduceStrategy from tensorrt_llm._torch.models.modeling_deepseekv41 import DeepseekV41Engram from tensorrt_llm._torch.modules.engram import EngramConfig, EngramHashProvider @@ -349,7 +351,7 @@ def _replay_attention(layout: CSA2Layout, mapping: Mapping, layer_idx: int = 0) @contextmanager def _replay_metadata( mapping: Mapping, layout: CSA2Layout, length: int, cached: int -) -> Iterator[CSA2TrtllmMetadata]: +) -> Iterator[TrtllmAttentionMetadata]: from tensorrt_llm._torch.attention.backends.sparse.csa2.cache_manager import CSA2CacheManager from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -389,7 +391,7 @@ def _replay_metadata( assert manager.prepare_context(request) assert manager.resize_context(request, request.context_chunk_size) manager._stream.synchronize() - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=_REPLAY_MAX_TOKENS, kv_cache_manager=manager, @@ -408,8 +410,10 @@ def _replay_metadata( metadata.prepare() else: # A zero-row MoE peer has no attention request to prepare. - metadata.csa2_positions = torch.empty(0, dtype=torch.int32, device="cuda") - metadata.reset_routing() + metadata.ext(CSA2Metadata).csa2_positions = torch.empty( + 0, dtype=torch.int32, device="cuda" + ) + metadata.ext(CSA2Metadata).reset_routing() yield metadata finally: if request is not None: @@ -489,13 +493,17 @@ def _check_dp_replay_case( reference_metadata = stack.enter_context(_replay_metadata(Mapping(), layout, count, cached)) metadata.all_rank_num_tokens = counts if prompt_ends is not None and metadata.num_contexts: - metadata.decoder_context_ends = (prompt_ends[rank],) + metadata.ext(CSA2Metadata).decoder_context_ends = (prompt_ends[rank],) if cached: for attention, step_metadata in ( (model.layers[0].self_attn, metadata), (reference_attention, reference_metadata), ): - attention(norm(full_hidden[:cached]), step_metadata.csa2_positions, step_metadata) + attention( + norm(full_hidden[:cached]), + step_metadata.ext(CSA2Metadata).csa2_positions, + step_metadata, + ) step_metadata.num_contexts = 0 step_metadata.seq_lens = torch.tensor([count], dtype=torch.int32, device="cpu") step_metadata.kv_cache_params = KVCacheParams( @@ -503,7 +511,7 @@ def _check_dp_replay_case( ) step_metadata.prepare() hidden = full_hidden[cached:] - positions = metadata.csa2_positions + positions = metadata.ext(CSA2Metadata).csa2_positions input_ids = torch.arange(count, device="cuda") + rank * 17 state = HCState.resolved(hidden[:, None, :], pre_mix=torch.ones(count, 1, 1, device="cuda")) original_seq_lens = metadata.seq_lens.clone() @@ -537,9 +545,9 @@ def _check_dp_replay_case( assert replay_state.residual is state.residual if not plan.updates_local_metadata: assert replay_state is state - assert metadata.csa2_replay_query_rows is None + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None else: - assert metadata.csa2_replay_query_rows is plan.rows + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is plan.rows # Full-prefill parity is invalid for this approximation. Build the same # bounded replay independently of the DP planner, retaining all GLOBAL KV. @@ -548,15 +556,17 @@ def _check_dp_replay_case( reference_metadata.seq_lens = torch.tensor([kept], dtype=torch.int32, device="cpu") reference_metadata.prompt_lens = [kept] reference_metadata.kv_cache_params.num_cached_tokens_per_seq = [start] - reference_metadata.set_swa_bounded_replay([count], decoder=True) + reference_metadata.ext(CSA2Metadata).set_swa_bounded_replay([count], decoder=True) reference_metadata.prepare() - reference_metadata.csa2_precomputed_kv_layers = {0} + reference_metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers = {0} if kept: actual_attention = model.layers[0].self_attn( norm(replay_state.residual[:, 0]), new_positions, metadata ) expected_attention = reference_attention( - norm(hidden[start:]), reference_metadata.csa2_positions, reference_metadata + norm(hidden[start:]), + reference_metadata.ext(CSA2Metadata).csa2_positions, + reference_metadata, ) torch.testing.assert_close(actual_attention, expected_attention, rtol=0, atol=0) else: @@ -581,7 +591,7 @@ def _check_dp_replay_case( torch.testing.assert_close(metadata.seq_lens, original_seq_lens, rtol=0, atol=0) assert metadata.prompt_lens == original_prompt_lens assert metadata.kv_cache_params.num_cached_tokens_per_seq == ([cached] if count else []) - assert metadata.csa2_replay_query_rows is None + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None def _check_dp_encoder_replay_case( @@ -640,7 +650,7 @@ def cache(local_mapping: Mapping) -> Iterator[tuple]: try: assert manager._decoder_replay_window == _REPLAY_WINDOW assert manager._encoder_optional_groups - metadata = CSA2TrtllmMetadata( + metadata = CSA2TrtllmAttention.create_metadata( max_num_requests=1, max_num_tokens=_REPLAY_MAX_TOKENS, kv_cache_manager=manager, @@ -671,7 +681,7 @@ def admit(case: tuple, request_id: int, tokens: list[int]) -> LlmRequest: manager.prepare_resources(scheduled) return request - def global_bytes(metadata: CSA2TrtllmMetadata, end: int) -> dict: + def global_bytes(metadata: TrtllmAttentionMetadata, end: int) -> dict: manager = metadata.kv_cache_manager result = {} logical = torch.arange(end, device="cuda") @@ -690,11 +700,11 @@ def global_bytes(metadata: CSA2TrtllmMetadata, end: int) -> dict: result[layer, role] = values.view(torch.uint8).clone() return result - def check_floor(metadata: CSA2TrtllmMetadata, layer: int, floor: int) -> None: - metadata._ensure_swa_slots() - logical = metadata.csa2_positions[:, None] - _REPLAY_WINDOW + 1 + def check_floor(metadata: TrtllmAttentionMetadata, layer: int, floor: int) -> None: + metadata.ext(CSA2Metadata)._ensure_swa_slots() + logical = metadata.ext(CSA2Metadata).csa2_positions[:, None] - _REPLAY_WINDOW + 1 logical = logical + torch.arange(_REPLAY_WINDOW, device="cuda") - assert (metadata.csa2_swa_indices[layer][logical < floor] == -1).all() + assert (metadata.ext(CSA2Metadata).csa2_swa_indices[layer][logical < floor] == -1).all() with pytest.MonkeyPatch.context() as settings, ExitStack() as stack: settings.setenv("TRTLLM_V41_DECODER_BOUNDED_REPLAY", "1") @@ -727,16 +737,18 @@ def run(requests: list[LlmRequest], frozen: list[dict] | None = None) -> None: use_cache=True, num_cached_tokens_per_seq=[start] ) if step_metadata is metadata: - step_metadata.prepare_context_replay([step_request]) + step_metadata.ext(CSA2Metadata).prepare_context_replay([step_request]) elif encoder_replay_tokens(step_request): - step_metadata.set_swa_bounded_replay([step_request.context_current_position]) + step_metadata.ext(CSA2Metadata).set_swa_bounded_replay( + [step_request.context_current_position] + ) step_metadata.prepare() - step_metadata.begin_model_forward() + step_metadata.begin_forward() check_floor(step_metadata, 0, start) metadata.all_rank_num_tokens = counts assert MPI.COMM_WORLD.allgather(metadata.num_tokens) == counts hidden = inputs[start : start + length] - positions = metadata.csa2_positions + positions = metadata.ext(CSA2Metadata).csa2_positions ids = torch.arange(start, start + length, device="cuda") _prepare_replay(model, metadata, False, [request]) _, plan = model._plan_bounded_replay(metadata, False, [request]) @@ -745,13 +757,19 @@ def run(requests: list[LlmRequest], frozen: list[dict] | None = None) -> None: def layer_output( layer: int, actual_hidden: torch.Tensor, expected_hidden: torch.Tensor ) -> torch.Tensor: - actual = attention[layer](norm(actual_hidden), metadata.csa2_positions, metadata) + actual = attention[layer]( + norm(actual_hidden), metadata.ext(CSA2Metadata).csa2_positions, metadata + ) expected = reference_attention[layer]( - norm(expected_hidden), reference_metadata.csa2_positions, reference_metadata + norm(expected_hidden), + reference_metadata.ext(CSA2Metadata).csa2_positions, + reference_metadata, ) torch.testing.assert_close(actual, expected, rtol=0, atol=0) logits = torch.arange(8, device="cuda", dtype=torch.float32)[None, :] - logits = torch.sin(logits + metadata.csa2_positions[:, None].float() + rank) + logits = torch.sin( + logits + metadata.ext(CSA2Metadata).csa2_positions[:, None].float() + rank + ) output = moe(actual, logits, all_rank_num_tokens=metadata.all_rank_num_tokens) reference = reference_moe(expected, logits) assert torch.isfinite(output).all() @@ -768,13 +786,17 @@ def layer_output( reference_metadata.seq_lens = torch.tensor([kept], dtype=torch.int32, device="cpu") reference_metadata.prompt_lens = [kept] reference_metadata.kv_cache_params.num_cached_tokens_per_seq = [decoder_start] - reference_metadata.set_decoder_query_boundary(1) + reference_metadata.ext(CSA2Metadata).set_decoder_query_boundary(1) reference_metadata.prepare() - reference_metadata.csa2_precomputed_kv_layers = {1} - reference_metadata._ensure_swa_slots() - logical = reference_metadata.csa2_positions[:, None] - _REPLAY_WINDOW + 1 + reference_metadata.ext(CSA2Metadata).csa2_precomputed_kv_layers = {1} + reference_metadata.ext(CSA2Metadata)._ensure_swa_slots() + logical = ( + reference_metadata.ext(CSA2Metadata).csa2_positions[:, None] - _REPLAY_WINDOW + 1 + ) logical = logical + torch.arange(_REPLAY_WINDOW, device="cuda") - reference_metadata.csa2_swa_indices[1].masked_fill_(logical < decoder_start, -1) + reference_metadata.ext(CSA2Metadata).csa2_swa_indices[1].masked_fill_( + logical < decoder_start, -1 + ) new_positions, new_ids, replay_state = model._enter_bounded_replay( plan, metadata, positions, ids, state ) @@ -794,7 +816,7 @@ def layer_output( assert metadata.seq_lens.tolist() == [length] assert metadata.prompt_lens == [length] assert metadata.kv_cache_params.num_cached_tokens_per_seq == [start] - assert metadata.csa2_replay_query_rows is None + assert metadata.ext(CSA2Metadata).csa2_replay_query_rows is None if frozen: for step_metadata, snapshot in zip((metadata, reference_metadata), frozen): for key, value in global_bytes(step_metadata, prefix).items(): diff --git a/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py b/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py index 42ca4d70539a..854e367165f4 100644 --- a/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py +++ b/tests/unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py @@ -14,7 +14,11 @@ from mpi4py.futures import MPIPoolExecutor import tensorrt_llm -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.compilation.piecewise_optimizer import PiecewiseRunner from tensorrt_llm._torch.compilation.utils import capture_piecewise_cuda_graph from tensorrt_llm._torch.distributed import Distributed @@ -35,6 +39,14 @@ pytestmark = pytest.mark.threadleak(enabled=False) +def _msa_metadata() -> TrtllmAttentionMetadata: + """An uninitialized TRTLLM metadata carrying a bare MSA extension.""" + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + metadata.replace_ext(MiniMaxM3MsaMetadata(metadata)) + metadata.replace_ext(SpecDecodingMetadata(metadata)) + return metadata + + @torch.inference_mode() def _run_empty_adp_rank(world_size: int) -> int: rank = tensorrt_llm.mpi_rank() @@ -64,20 +76,20 @@ def _run_empty_adp_rank(world_size: int) -> int: # Reuse the metadata fixture pattern from test_msa_backend. Only the cache # allocator is stubbed; slot construction/clearing and GPU writes are real. - metadata_cls = MiniMaxM3MsaSparseAttention.Metadata - metadata = metadata_cls.__new__(metadata_cls) - metadata._msa_buffers_ready = True + metadata = _msa_metadata() + msa = metadata.ext(MiniMaxM3MsaMetadata) + msa._msa_buffers_ready = True metadata.kv_cache_manager = SimpleNamespace( tokens_per_block=128, get_buffers=lambda layer_idx: main_cache, get_block_ids_per_seq=lambda request_ids: torch.tensor([[1]], dtype=torch.int32), ) - metadata.msa_out_cache_loc = torch.full((bucket,), -1, dtype=torch.int32, device="cuda") - metadata.msa_block_table = torch.zeros((1, 1), dtype=torch.int32, device="cuda") - metadata.msa_seq_lens_cuda = torch.zeros(1, dtype=torch.int32, device="cuda") - metadata.msa_subpage_block_table = None - metadata._msa_runs_no_fmha = lambda: True - metadata._msa_kv_lens_may_change = lambda: False + msa.msa_out_cache_loc = torch.full((bucket,), -1, dtype=torch.int32, device="cuda") + msa.msa_block_table = torch.zeros((1, 1), dtype=torch.int32, device="cuda") + msa.msa_seq_lens_cuda = torch.zeros(1, dtype=torch.int32, device="cuda") + msa.msa_subpage_block_table = None + msa._msa_runs_no_fmha = lambda: True + msa._msa_kv_lens_may_change = lambda: False def producer( projection: torch.Tensor, @@ -141,21 +153,21 @@ def producer( set_per_request_prefill_cuda_graph_flag(eligible) metadata.request_ids = [0] if live_tokens else [] - metadata._msa_qo_lens_cpu = torch.tensor( + msa._msa_qo_lens_cpu = torch.tensor( [live_tokens] if live_tokens else [], dtype=torch.int32 ) - metadata._msa_kv_lens_cpu = metadata._msa_qo_lens_cpu.clone() - metadata._msa_qo_offset_cpu = torch.zeros(local_contexts, dtype=torch.int32) + msa._msa_kv_lens_cpu = msa._msa_qo_lens_cpu.clone() + msa._msa_qo_offset_cpu = torch.zeros(local_contexts, dtype=torch.int32) # What the base seq_lens setter would derive; the cache writers # read the live (unpadded) count off it. metadata._num_tokens = live_tokens - metadata._build_msa_fields() - assert metadata.msa_out_cache_loc.tolist() == list( - range(128, 128 + live_tokens) - ) + [-1] * (bucket - live_tokens) + msa._build_msa_fields() + assert msa.msa_out_cache_loc.tolist() == list(range(128, 128 + live_tokens)) + [ + -1 + ] * (bucket - live_tokens) packed.normal_() before_main, before_index = main_backing.clone(), index_backing.clone() - args = (packed, main_cache, index_cache, metadata.msa_out_cache_loc) + args = (packed, main_cache, index_cache, msa.msa_out_cache_loc) if empty_rank is None: with capture_piecewise_cuda_graph(True): # Three warmups, then capture. Keep the capture outputs diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 8cebeff8c38e..fd167d1119d7 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -29,8 +29,11 @@ from utils.util import skip_num_gpus_less_than, skip_pre_blackwell from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm._torch.attention.backends.sparse.dsa import ( - DSACacheManagerV2, DSAtrtllmAttentionMetadata) +from tensorrt_llm._torch.attention.backends.features.spec_decoding import \ + SpecDecodingMetadata +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSACacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.params import \ + MTPIndexShareExtension from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams @@ -59,11 +62,7 @@ def test_mtp_eagle_refreshes_dsa_metadata_before_draft_forward() -> None: state = {} draft_manager = object.__new__(DSACacheManagerV2) - class _Metadata(DSAtrtllmAttentionMetadata): - - def __init__(self): - self._seq_lens_cuda = torch.tensor([1], dtype=torch.int32) - self._num_ctx_tokens = 0 + class _IndexShare(MTPIndexShareExtension): def set_in_mtp_draft_loop(self, active): state["in_mtp_draft_loop"] = active @@ -74,6 +73,13 @@ def set_mtp_num_accepted(self, value): def set_skip_topk(self, value): state["skip_topk"] = value + class _Metadata(TrtllmAttentionMetadata): + + def __init__(self): + self._seq_lens_cuda = torch.tensor([1], dtype=torch.int32) + self._num_ctx_tokens = 0 + self.replace_ext(_IndexShare(self)) + def on_update_kv_lens(self): events.append("refresh") @@ -144,11 +150,11 @@ def test_dynamic_tree_metadata_forces_target_mask_prepare_each_step() -> None: metadata.update_spec_dec_param(is_spec_dec_dynamic_tree=True, **common_kwargs) - assert metadata.force_prepare_spec_dec_tree_mask + assert metadata.ext(SpecDecodingMetadata).force_prepare_tree_mask metadata.update_spec_dec_param(is_spec_dec_dynamic_tree=False, **common_kwargs) - assert not metadata.force_prepare_spec_dec_tree_mask + assert not metadata.ext(SpecDecodingMetadata).force_prepare_tree_mask def test_mtp_dynamic_tree_relocation_uses_full_attention_window( @@ -588,7 +594,7 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): attn_metadata = TrtllmAttentionMetadata(**metadata_kwargs) if spec_signal == "is_spec_decoding_enabled": - attn_metadata.is_spec_decoding_enabled = True + attn_metadata.ext(SpecDecodingMetadata).is_spec_decoding_enabled = True elif spec_signal == "has_speculative_draft_tokens": attn_metadata.runtime_features.has_speculative_draft_tokens = True diff --git a/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py b/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py index b40016f8fb22..1b993a6f9583 100644 --- a/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py +++ b/tests/unittest/_torch/speculative/test_sa_hybrid_state_promotion.py @@ -15,10 +15,19 @@ import pytest +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.speculative.sa_worker import SAWorker +def _recurrent_stub(mamba_metadata, **fields) -> SimpleNamespace: + """A metadata stub serving ``mamba_metadata`` through its recurrent extension.""" + recurrent = SimpleNamespace(mamba_metadata=mamba_metadata) + metadata = SimpleNamespace(**fields) + metadata.find_ext = metadata.ext = lambda cls: recurrent if cls is RecurrentMetadata else None + return metadata + + def _make_worker() -> SAWorker: spec_config = SimpleNamespace(max_draft_len=4, max_matching_ngram_size=2) worker = SAWorker(spec_config) @@ -34,7 +43,7 @@ def _make_worker() -> SAWorker: def _make_metadata(kv_cache_manager, num_seqs: int, num_contexts: int): - attn_metadata = SimpleNamespace( + attn_metadata = _recurrent_stub( num_seqs=num_seqs, num_contexts=num_contexts, kv_cache_manager=kv_cache_manager, diff --git a/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py b/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py index 09540a930c07..25c9793f5a9f 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py +++ b/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py @@ -94,7 +94,7 @@ def _test_attention_trtllm_sage( sparse_params=sparse_params, ) - metadata = attention_cls.Metadata( + metadata = attention_cls.create_metadata( max_num_requests=batch_size, max_num_tokens=seq_len, kv_cache_manager=None, diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index a8dad92f8c06..37cdbc99fa03 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -7,6 +7,7 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.features.spec_decoding import SpecDecodingMetadata from tensorrt_llm._torch.attention.backends.sparse.params import ( BlockSparseForwardInputs, SparseBackendForwardArgs, @@ -64,9 +65,7 @@ def _make_core_forward_metadata(): metadata.use_paged_context_fmha = False metadata.cu_q_seqlens = None metadata.cu_kv_seqlens = None - metadata.enable_flash_mla = False - metadata.spec_bl_tree_first_sparse_mask_offset_kv = None - metadata.spec_decoding_bl_tree_mask = None + metadata.replace_ext(SpecDecodingMetadata(metadata)) metadata.kv_lens_cuda_runtime = torch.tensor([4], dtype=torch.int32) metadata.kv_lens_runtime = torch.tensor([4], dtype=torch.int32) metadata.prompt_lens_cuda_runtime = torch.tensor([4], dtype=torch.int32)