From f069dd1ca3c3307d536ca0a112808bb36078036c Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Sun, 25 Oct 2026 23:02:05 -0700 Subject: [PATCH 1/3] [None][feat] Attach feature and sparse metadata state as extensions AttentionMetadata becomes a small core plus pluggable MetadataExtension objects. Features, sparse algorithms and models attach their per-step state to a core metadata family instead of subclassing it, and callers reach it with metadata.ext(T) / metadata.find_ext(T). Framework (attention/backends/extensions.py): - MetadataExtension hooks driven by the core: allocate, prepare_scope, pre_derive_lengths, override_kv_lens, prepare, post_derive, begin_forward, on_kv_lens_changed, on_spec_dec_param, on_spec_dec_restored, invalidate_length_derived_buffers, draft_swap/draft_restore, attention-DP and workspace hooks, kv_cache_pools, token_major_gen_view, mla_scheduler_buffers, bind_context_requests, defer_decode_outputs. - ExtensionRegistry with requirement ordering; CUDA-graph and copy.copy views get their own clones. swap_attrs_for_draft/restore_swapped_attrs share the draft-swap pattern. - AttentionBackend.metadata_extensions and create_metadata(); core families attach their own via _default_extension_types(). - The strict CUDA-graph buffer check also covers extension tensors. Features moved out of TrtllmAttentionMetadata (attention/backends/features): HelixMetadata, DraftKVMetadata, FlashMlaMetadata, Fp4MlaState, SpecDecodingMetadata; RecurrentMetadata replaces the core mamba fields. Sparse algorithms are SparseAlgorithmExtension subclasses on TrtllmAttentionMetadata instead of metadata subclasses: RocketMetadata, QSAMetadata, DSAMetadata, DeepseekV4Metadata, CSA2Metadata, MiniMaxM3TritonMetadata, MiniMaxM3MsaMetadata. MTPIndexShareExtension replaces the MTPIndexShareMetadata protocol. glm_kpool keeps its model-chosen Mamba2Metadata subclass, now carried by RecurrentMetadata. Consumers (FMHA libraries, MLA, models, modeling_v2 targets, speculative workers, executor) and tests read the extensions; the attention developer guide describes the extension contract. Signed-off-by: Yihan Wang --- .../sparse-attention-development-guide.md | 2 +- .../modeling.py | 3 +- .../glm_5_2_nvfp4__sm_107__tp16/modeling.py | 54 +- .../gpt_oss_120b__sm_10x__default/modeling.py | 62 +- .../modeling.py | 243 +- .../ops/attention/qsa_paged_index_scores.py | 2 +- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 46 +- tensorrt_llm/_torch/attention/attention.py | 19 +- .../_torch/attention/backends/extensions.py | 331 ++ .../attention/backends/features/__init__.py | 16 + .../attention/backends/features/draft_kv.py | 53 + .../attention/backends/features/flash_mla.py | 141 + .../attention/backends/features/helix.py | 245 + .../attention/backends/features/recurrent.py | 38 + .../backends/features/spec_decoding.py | 320 + .../_torch/attention/backends/flashinfer.py | 8 +- .../attention/backends/fmha/cute_dsl_mla.py | 29 +- .../attention/backends/fmha/fallback.py | 47 +- .../backends/fmha/flashinfer_trtllm_gen.py | 10 +- .../_torch/attention/backends/fmha/fp4_mla.py | 28 +- .../_torch/attention/backends/fmha/manager.py | 3 +- .../attention/backends/fmha/msa_decode.py | 16 +- .../attention/backends/fmha/msa_prefill.py | 42 +- .../_torch/attention/backends/fmha/phased.py | 8 +- .../attention/backends/fmha/prims_ts.py | 14 +- .../backends/fp4_mla/cache_gather.py | 8 +- .../backends/fp4_mla/cache_update.py | 81 +- .../attention/backends/fp4_mla/decode.py | 46 +- .../backends/fp4_mla/fp4_mla_context.py | 15 +- .../attention/backends/fp4_mla/layout.py | 9 +- .../attention/backends/fp4_mla/metadata.py | 121 +- .../attention/backends/fp4_mla/packed_plan.py | 17 +- .../attention/backends/fp4_mla/state.py | 75 +- .../attention/backends/fp4_mla/v_cache.py | 7 +- .../_torch/attention/backends/interface.py | 276 +- .../attention/backends/sparse/csa2/README.md | 9 +- .../attention/backends/sparse/csa2/backend.py | 97 +- .../attention/backends/sparse/csa2/ced.py | 5 +- .../backends/sparse/csa2/decoder_replay.py | 69 +- .../attention/backends/sparse/csa2/indexer.py | 119 +- .../attention/backends/sparse/csa2/kernel.py | 58 +- .../backends/sparse/csa2/metadata.py | 436 +- .../attention/backends/sparse/csa2/module.py | 35 +- .../attention/backends/sparse/csa2/params.py | 4 +- .../backends/sparse/deepseek_v4/__init__.py | 4 +- .../backends/sparse/deepseek_v4/backend.py | 55 +- .../backends/sparse/deepseek_v4/compressor.py | 40 +- .../backends/sparse/deepseek_v4/flash_mla.py | 35 +- .../backends/sparse/deepseek_v4/flashinfer.py | 12 +- .../sparse/deepseek_v4/footer_scale_kv.py | 30 +- .../backends/sparse/deepseek_v4/indexer.py | 20 +- .../backends/sparse/deepseek_v4/metadata.py | 355 +- .../backends/sparse/deepseek_v4/module.py | 12 +- .../attention/backends/sparse/dsa/__init__.py | 4 +- .../attention/backends/sparse/dsa/backend.py | 63 +- .../attention/backends/sparse/dsa/indexer.py | 352 +- .../attention/backends/sparse/dsa/metadata.py | 1088 ++-- .../attention/backends/sparse/dsa/module.py | 12 +- .../backends/sparse/dsa_flashinfer.py | 33 +- .../backends/sparse/glm_kpool/backend.py | 3 +- .../backends/sparse/minimax_m3/__init__.py | 4 +- .../sparse/minimax_m3/kernels/msa_utils.py | 5 +- .../backends/sparse/minimax_m3/msa_backend.py | 256 +- .../sparse/minimax_m3/triton_backend.py | 10 +- .../sparse/minimax_m3/triton_metadata.py | 94 +- .../attention/backends/sparse/params.py | 24 +- .../attention/backends/sparse/qsa/__init__.py | 15 + .../attention/backends/sparse/qsa/backend.py | 7 +- .../attention/backends/sparse/qsa/indexer.py | 40 +- .../attention/backends/sparse/qsa/metadata.py | 375 +- .../attention/backends/sparse/qsa/module.py | 70 +- .../backends/sparse/rocket/__init__.py | 5 +- .../backends/sparse/rocket/backend.py | 106 +- .../backends/sparse/rocket/metadata.py | 362 +- .../_torch/attention/backends/trtllm.py | 847 +-- .../_torch/attention/backends/vanilla.py | 8 +- tensorrt_llm/_torch/attention/mla.py | 11 +- .../_torch/models/modeling_deepseekv4.py | 21 +- .../_torch/models/modeling_deepseekv41.py | 38 +- .../_torch/models/modeling_glm5_next.py | 3 +- .../_torch/models/modeling_minimaxm3.py | 65 +- .../_torch/models/modeling_nemotron_h.py | 6 +- .../_torch/models/modeling_qwen3_next.py | 8 +- .../_torch/models/modeling_qwen4_exp.py | 10 +- .../_torch/models/modeling_step3p7.py | 11 +- .../_torch/modules/kimi_kda/kimi_kda_mixer.py | 3 +- .../_torch/modules/mamba/gdn_mixer.py | 3 +- .../_torch/modules/mamba/mamba2_mixer.py | 8 +- .../_torch/pyexecutor/cuda_graph_runner.py | 40 +- .../_torch/pyexecutor/engine/metadata.py | 13 +- .../engine/runners/decoder/inputs.py | 71 +- .../engine/runners/decoder/metadata.py | 11 +- .../engine/runners/decoder/runner.py | 4 +- .../engine/runners/decoder/warmup.py | 31 +- .../engine/runners/encoder_decoder.py | 2 + .../kv_cache/mamba_cache_manager.py | 5147 +++++++++++++++++ tensorrt_llm/_torch/speculative/dflash.py | 21 +- .../_torch/speculative/draft_target.py | 5 +- tensorrt_llm/_torch/speculative/eagle3.py | 82 +- .../_torch/speculative/eagle3_dynamic_tree.py | 34 +- tensorrt_llm/_torch/speculative/interface.py | 35 +- tensorrt_llm/_torch/speculative/mtp.py | 10 +- .../_torch/speculative/mtp_dynamic_tree.py | 92 +- tensorrt_llm/_torch/speculative/sa_worker.py | 3 +- .../tools/layer_wise_benchmarks/runner.py | 12 +- .../attention_perf/attention_perf_harness.py | 9 +- tests/microbenchmarks/bench_fp4_mla_decode.py | 65 +- .../unittest/_torch/attention/backend_case.py | 8 +- .../_torch/attention/fmha_test_utils.py | 44 + .../test_ragged_metadata_lifecycle.py | 143 +- .../test_trtllm_mla_rope_forwarding.py | 16 +- .../test_helix_postprocess.py | 5 +- .../attention/multi_gpu/helix_test_utils.py | 36 +- .../_torch/attention/sparse/csa2/_utils.py | 99 +- .../attention/sparse/csa2/test_backend.py | 243 +- .../sparse/csa2/test_cache_manager.py | 15 +- .../sparse/csa2/test_context_only_metadata.py | 27 +- .../attention/sparse/csa2/test_indexer.py | 655 ++- .../sparse/csa2/test_indexer_descriptors.py | 73 +- .../attention/sparse/csa2/test_kernel.py | 288 +- .../attention/sparse/csa2/test_metadata.py | 771 ++- .../sparse/csa2/test_metadata_uploads.py | 57 +- .../attention/sparse/csa2/test_module.py | 130 +- .../deepseek_v4/test_compressor_module.py | 27 +- .../test_deepseek_v4_fp4_indexer.py | 2 + .../deepseek_v4/test_deepseek_v4_offload.py | 89 +- .../test_deepseek_v4_q_b_fusion.py | 6 +- .../test_deepseek_v4_sparse_mla.py | 57 +- .../deepseek_v4/test_indexer_slot_mapping.py | 54 +- .../test_ragged_fused_q_rope_specs.py | 25 +- .../deepseek_v4/test_token_major_rows.py | 80 +- .../attention/sparse/dsa/test_dsa_indexer.py | 534 +- .../sparse/dsa/test_dsa_sparse_mla.py | 13 +- .../sparse/dsa/test_fused_metadata.py | 2 +- .../sparse/dsa/test_indexer_gvr_prior.py | 2 + .../sparse/dsa/test_metadata_topk_init.py | 44 +- .../sparse/dsa/test_req_idx_per_token.py | 46 +- .../sparse/dsa/test_short_seq_mha.py | 2 +- .../sparse/glm_kpool/test_glm_kpool.py | 9 +- .../msa/test_minimax_m3_sparse_attn_decode.py | 40 +- .../attention/sparse/msa/test_msa_backend.py | 367 +- .../attention/sparse/qsa/test_qsa_sparse.py | 222 +- .../sparse/rocketkv/test_rocketkv.py | 19 +- .../attention/sparse/test_flashinfer_utils.py | 9 +- .../sparse/test_prims_ts_block_sparse.py | 9 +- .../attention/sparse/test_sparse_attention.py | 4 +- .../attention/sparse/test_sparse_mha.py | 36 +- .../sparse/test_sparse_mla_forward.py | 10 +- .../attention/sparse/test_sparse_mqa_gqa.py | 34 +- .../_torch/attention/test_attention.py | 2 +- .../test_attention_failure_logging.py | 8 +- .../_torch/attention/test_attention_mla.py | 33 +- .../attention/test_attention_no_cache.py | 2 +- .../attention/test_attention_op_sync.py | 34 +- .../_torch/attention/test_combined_fmha.py | 4 +- .../_torch/attention/test_cute_dsl_mla.py | 3 +- .../attention/test_cute_dsl_mla_split_plan.py | 4 +- .../attention/test_flashinfer_attention.py | 1 - .../test_flashinfer_trtllm_gen_fmha.py | 8 +- .../_torch/attention/test_fmha_interface.py | 7 +- .../_torch/attention/test_fmha_manager.py | 9 +- .../_torch/attention/test_fmha_page_index.py | 15 +- .../unittest/_torch/attention/test_fp4_mla.py | 244 +- .../_torch/attention/test_fp4_mla_packed.py | 29 +- .../_torch/attention/test_fp4_mla_splitkv.py | 22 +- .../test_mixed_precision_backends.py | 4 +- .../_torch/attention/test_mla_registry.py | 11 +- .../test_prims_ts_attention_backend.py | 4 +- .../_torch/attention/test_prims_ts_fmha.py | 37 +- .../attention/test_skip_softmax_sm120.py | 2 +- .../attention/test_triton_custom_mask_fmha.py | 4 +- .../test_trtllm_paged_context_fmha_gate.py | 3 +- .../_torch/executor/engine/test_metadata.py | 13 +- .../kv_cache/test_mamba_cache_manager.py | 18 + .../test_cuda_graph_capture_replay.py | 11 +- .../executor/test_pytorch_model_engine.py | 34 +- .../modeling/_deepseek_v41_test_utils.py | 28 +- .../_torch/modeling/test_deepseek_v41_ced.py | 44 +- .../test_deepseek_v41_decoder_global.py | 3 +- .../test_deepseek_v41_encoder_replay.py | 7 +- .../test_deepseek_v41_multimodal_backbone.py | 2 +- .../modeling/test_glm5_next_contracts.py | 16 +- .../modeling/test_modeling_deepseekv4.py | 9 +- .../modeling/test_modeling_deepseekv41.py | 118 +- .../_torch/modeling/test_modeling_llama.py | 29 +- .../modeling/test_qsa_runtime_wiring.py | 26 +- .../_torch/modeling/test_qwen4_exp_support.py | 18 +- .../test_modeling_v2_hist_buffer_lifecycle.py | 5 +- .../test_modeling_v2_index_share.py | 35 +- .../unittest/_torch/models/test_minimax_m3.py | 39 +- .../kimi_kda/test_kda_flashinfer_prefill.py | 13 +- .../test_kda_mixed_batch_core_layout.py | 12 +- .../kimi_kda/test_kimi_kda_verify_parity.py | 18 +- .../mamba/test_gdn_kernel_optimizations.py | 7 +- .../test_deepseek_v41_attention_dp.py | 88 +- .../multi_gpu/test_minimax_m3_piecewise.py | 48 +- .../_torch/speculative/test_eagle3.py | 26 +- .../test_sa_hybrid_state_promotion.py | 11 +- .../visual_gen/test_attention_trtllm_sage.py | 2 +- .../test_trtllm_attention_metadata.py | 5 +- 200 files changed, 12979 insertions(+), 6035 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/extensions.py create mode 100644 tensorrt_llm/_torch/attention/backends/features/__init__.py create mode 100644 tensorrt_llm/_torch/attention/backends/features/draft_kv.py create mode 100644 tensorrt_llm/_torch/attention/backends/features/flash_mla.py create mode 100644 tensorrt_llm/_torch/attention/backends/features/helix.py create mode 100644 tensorrt_llm/_torch/attention/backends/features/recurrent.py create mode 100644 tensorrt_llm/_torch/attention/backends/features/spec_decoding.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py 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/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..0ceea03e3ab6 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, @@ -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, @@ -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..d44f80fd9b9a 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +<<<<<<< HEAD """Lazy QSA exports; importing geometry does not construct the cache stack.""" from importlib import import_module @@ -23,3 +24,17 @@ def __getattr__(name: str): __all__ = list(_EXPORTS) +======= +from .backend import QSATrtllmAttention +from .cache_manager import QSAMambaHybridCacheManagerV2 +from .metadata import QSAMetadata +from .params import QSASparseMetadataParams, QSASparseParams + +__all__ = [ + "QSAMambaHybridCacheManagerV2", + "QSAMetadata", + "QSASparseMetadataParams", + "QSASparseParams", + "QSATrtllmAttention", +] +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) 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..2d96145d1db8 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -1,40 +1,28 @@ # 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.""" +<<<<<<< HEAD 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() @@ -58,107 +46,78 @@ def __post_init__(self) -> None: if not isinstance(self.kv_cache_manager, Qwen4ExpHybridCacheManager): raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManager") if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): +======= + md = self.metadata + self.needs_speculative_snapshot = False + # Draft-loop state; set by the speculative driver, read by the indexer. + 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, QSAMambaHybridCacheManagerV2): + raise TypeError("QSA sparse attention requires QSAMambaHybridCacheManagerV2") + params = md.sparse_metadata_params + if not isinstance(params, QSASparseMetadataParams): +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) 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 +125,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 @@ -200,87 +151,89 @@ def _init_qsa_draft_block_table(self) -> None: forward and this one in the draft loop; ``prepare()`` refreshes both before every replay. """ +<<<<<<< HEAD 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 +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) 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 +241,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 +250,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/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.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py new file mode 100644 index 000000000000..1247ad716922 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -0,0 +1,5147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import os +from abc import ABC, abstractmethod +from contextlib import contextmanager +from dataclasses import dataclass, replace +from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, + Optional, Protocol, Sequence, Tuple, Union, cast) + +import torch +import triton +import triton.language as tl + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata + from tensorrt_llm._torch.pyexecutor.config_utils import Qwen4ExpPLECacheParams + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + from tensorrt_llm.sampling_params import SamplingParams + +from tensorrt_llm._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch.disaggregation.resource.page import (MapperKind, + RoleLayout) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + _RESERVED_REQUEST_IDS, BlockReusePolicy, KVCacheManagerV2, Role) +from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ + KVCacheV2IterationStatsReport +from tensorrt_llm._torch.pyexecutor.llm_request import ( + ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, + PoolConfiguration, get_pp_layers) +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, + nvtx_range, prefer_pinned, + torch_dtype_to_binding) +from tensorrt_llm.bindings.executor import KvCacheRetentionConfig +from tensorrt_llm.bindings.internal.batch_manager import ( + LinearAttentionMetadata, LinearCacheType) +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, + BatchDesc, BufferConfig, + DataRole, KVCacheDesc) +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, + SsmLayerConfig, + TokenIdExt, _KVCache) + +# Shared with the KV budget estimator so allocator and budgeting can never +# diverge on the sharding rule (config_utils is import-cycle-free). +from ..config_utils import is_glm5_next, is_kimi_linear, mamba_effective_tp_size + +GB = 1 << 30 + + +class _PrefixReuseDiagnostics(Protocol): + + def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: + ... + + def _get_num_reusable_tokens_before_pruning(self) -> int: + ... + + +# Replay kernels pad the token/window dimension to at least 16 for tensor-core +# tiles, so history sizes below 16 are no faster when not writing and do +# expensive checkpoint writes more often. This minimum tensor-core tile size is +# present in all tensor-core generations, with newer GPUs adding larger tile +# operations, not smaller ones. In the other direction, history sizes above 16 +# currently pad to 32, which is substantially slower with the current kernel +# design. Keep the default floor at 16 while still allowing larger T, which +# degenerates to checkpointing every step instead of rejecting the request. If +# larger T becomes common, we could consider padding larger histories to powers +# of 2 or multiples of 16, but this is untested. For smaller T, we could explore +# combining larger histories with new kernel designs that stay efficient when +# the window is only partly full. +MIN_REPLAY_HISTORY_SIZE = 16 + +_KDA_BETA_CACHE_ALIGNMENT_BYTES = 16 +# KVCM copies 16-byte payloads; the replay kernel reads only the first int32. +_KDA_REPLAY_COUNT_BYTES = 16 +_KDA_REPLAY_COUNT_ROLE = DataRole("prev_num_accepted_tokens") +_KDA_REPLAY_CONV_ROLES = ("kda_conv_q", "kda_conv_k", "kda_conv_v") +_KDA_DRAFT_RECORDS_ROLE = "kda_state_tok" + + +def _is_disagg_context_server(is_disagg: bool) -> bool: + """Whether this worker only serves the context phase of disagg requests.""" + return is_disagg and os.getenv("TRTLLM_DISAGG_ROLE") == "context" + + +def _kda_replay_layout( + num_spec: int, + num_heads: int, + head_dim: int, + conv_window: int, + draft_records: bool = False +) -> Dict[str, Tuple[torch.dtype, List[int]]]: + """Per-slot physical shapes of one layer's KDA replay caches. + + Conv caches hold raw bf16 projection inputs (exact in bf16); k / v / beta + are bf16 and the gate stays fp32 because it feeds the cumulative-gate + exponents of the replayed chunk. Beta rows keep the 16-byte CuTe + alignment by padding the head dimension. + + With ``draft_records``, the records of every draft of the last verify + round (``kda_state_tok``: per draft, the row innovations vn, then beta * k + and the decay, each fp32 [H, K] with V = K = ``head_dim``) take the place + of the k / g / v / beta caches. + """ + dim = num_heads * head_dim + layout = { + "kda_conv_q": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_conv_k": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_conv_v": (torch.bfloat16, [conv_window + num_spec, dim]), + } + if draft_records: + layout[_KDA_DRAFT_RECORDS_ROLE] = (torch.float32, + [3, num_spec, num_heads, head_dim]) + return layout + heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // + torch.bfloat16.itemsize) + aligned_heads = ((num_heads + heads_per_alignment - 1) // + heads_per_alignment) * heads_per_alignment + layout.update({ + "kda_k_cache": (torch.bfloat16, [num_spec, dim]), + "kda_g_cache": (torch.float32, [num_spec, dim]), + "kda_v_cache": (torch.bfloat16, [num_spec, dim]), + "kda_beta_cache": (torch.bfloat16, [num_spec, aligned_heads]), + }) + return layout + + +def _kda_replay_bytes_per_slot(layout: Dict[str, Tuple[torch.dtype, List[int]]], + num_local_layers: int) -> int: + """Replay bytes one recurrent-state slot holds across local layers.""" + if not layout or num_local_layers == 0: + return 0 + layer_bytes = sum(dtype.itemsize * math.prod(shape) + for dtype, shape in layout.values()) + return num_local_layers * layer_bytes + _KDA_REPLAY_COUNT_BYTES + + +def _allocate_kda_beta_cache( + shape: Tuple[int, ...], + device: Optional[torch.device], +) -> torch.Tensor: + """Allocate a logically shaped bf16 beta cache with aligned physical rows.""" + logical_num_heads = shape[-1] + heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // + torch.bfloat16.itemsize) + padded_num_heads = ((logical_num_heads + heads_per_alignment - 1) // + heads_per_alignment) * heads_per_alignment + # CuTe requires every nested cache view to be 16-byte aligned. Keep the + # logical head count while padding the physical stride (six bf16 heads are + # 12 bytes, so an unpadded row makes alternating views only 4-byte aligned). + return torch.zeros( + *shape[:-1], + padded_num_heads, + dtype=torch.bfloat16, + device=device, + )[..., :logical_num_heads] + + +def _get_num_cuda_graph_padding_dummy_slots( + spec_config: Optional["DecodingBaseConfig"], + max_batch_size: int, +) -> int: + """Return the number of persistent CUDA-graph padding dummy IDs. + + This is computed before ``ModelEngine`` exists and covers draft lengths + reachable at every batch size, including the zero-length acceptance-rate + fallback. ``compute_dynamic_draft_len_mapping`` is created + later and covers only configured CUDA-graph batch sizes, so it cannot size + this persistent ID set. + """ + if spec_config is None: + return 1 + + draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( + spec_dec_mode, "support_dynamic_draft_len") + and spec_dec_mode.support_dynamic_draft_len()) + if draft_len_schedule and supports_dynamic_draft_len: + runtime_draft_lengths = set() + first_uncovered_batch_size = 1 + for batch_size_threshold, draft_len in draft_len_schedule.items(): + if first_uncovered_batch_size > max_batch_size: + break + if batch_size_threshold >= first_uncovered_batch_size: + runtime_draft_lengths.add(draft_len) + first_uncovered_batch_size = batch_size_threshold + 1 + if first_uncovered_batch_size <= max_batch_size: + runtime_draft_lengths.add(0) + else: + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", + 0) or 0) + is_linear_tree = getattr( + spec_config, + "is_linear_tree", + max_draft_len == max_total_draft_tokens, + ) + static_draft_len = (max_draft_len + if is_linear_tree else max_total_draft_tokens) + runtime_draft_lengths = {static_draft_len or 0} + + if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and + (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): + runtime_draft_lengths.add(0) + return len(runtime_draft_lengths) + + +class MambaRole: + """V2 buffer roles owned only by the hybrid Mamba manager.""" + + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") + PLE_NGRAM_CONTEXT = DataRole("ple_ngram_context") + PLE_CONV_STATE = DataRole("ple_conv_state") + + +def get_tensor_size_bytes(tensor): + """Calculate tensor size in bytes.""" + if isinstance(tensor, torch.Tensor): + return tensor.element_size() * tensor.nelement() + elif isinstance(tensor, list): + return sum(get_tensor_size_bytes(t) for t in tensor) + return 0 + + +# Mamba SSM stochastic-rounding Philox seed plumbing. +# +# Both the replay kernel and flashinfer's selective_state_update consume a +# `rand_seed` int64 tensor. Historically the non-replay paths created this +# tensor via `torch.randint(..., (1,))` on every decode step, which (a) is +# non-deterministic across runs, (b) allocates a fresh CUDA tensor per call +# and is therefore unfriendly to CUDA-graph capture, and (c) has no notion +# of cache-slot identity (so slot reuse cannot rotate the stream). +# +# The functions below produce per-slot int64 seeds via SplitMix64 finalization +# of (counter, slot, rank). Adjacent inputs yield uncorrelated outputs so +# consecutive cache slots and consecutive request counters do not leave +# structural fingerprints in the Philox input stream. All outputs live in +# (0, 2**62) so they avoid Philox's degenerate seed=0 case while staying +# within int64. +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(x: int) -> int: + """SplitMix64 finalizer; pure function, no torch.""" + x = (x + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 30) + x = (x * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 27) + x = (x * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + x ^= (x >> 31) + return x & _MAMBA_SSM_UINT64_MASK + + +def _compute_deterministic_mamba_seed(counter: int, slot: int, + rank_offset: int) -> int: + """Deterministic int64 seed in (0, 2**62) from (counter, slot, rank). + + Pure function (no RNG, no torch.randint). Identical inputs across + process invocations produce identical outputs, which is what the + acceptance criteria require for cross-run reproducibility. + """ + folded = (_MAMBA_SSM_SEED_BASE + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK) + folded &= _MAMBA_SSM_UINT64_MASK + value = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + if value == 0: + value = 1 + return value + + +def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, + device: torch.device) -> torch.Tensor: + """Allocate a (cache_size,) int64 CUDA buffer of deterministic seeds. + + counter=0 at allocation time; per-slot reset on fresh request assignment + bumps a host-side counter and rewrites only the freshly-assigned slot. + """ + slot_seeds = [ + _compute_deterministic_mamba_seed(0, i, rank_offset) + for i in range(cache_size) + ] + return torch.tensor(slot_seeds, dtype=torch.int64, device=device) + + +def _host_slots_to_device(slots: List[int], + device: torch.device) -> torch.Tensor: + """Copy host slot ids to ``device`` without a host sync. + + A pageable copy would wait for the queued step. Each call stages through a + fresh pinned tensor, so no later host write can race the copy. + """ + return torch.tensor(slots, + dtype=torch.long, + pin_memory=device.type == "cuda" + and prefer_pinned()).to(device, non_blocking=True) + + +def _mamba_rank_offset(mapping: Mapping) -> int: + """Distinct seed offset per (tp_rank, pp_rank) so independent ranks + don't draw identical streams when not coordinated.""" + return (mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + + mapping.rank * 1_009) + + +def use_py_mamba_cache_manager() -> bool: + """Check if PythonMambaCacheManager should be forced (agg mode override). + + Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. + + Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager + with PythonMambaCacheManager inside instead of the configured manager. + Disagg mode is unaffected — its compatibility routing selects Mixed or + Cpp based on the transceiver configuration. + """ + return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' + + +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: Optional[torch.Tensor] = None, +) -> None: + """Advance replay bookkeeping after a speculative generation step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to( + replay_metadata.prev_num_accepted_tokens.dtype) + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = (prev_num_accepted_tokens + + replay_metadata.replay_step_width + > replay_metadata.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, + accepted_tokens, + prev_num_accepted_tokens + accepted_tokens, + ) + cache_buf_idx = replay_metadata.cache_buf_idx[slots] + next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + if is_dummy_request is not None: + next_num_accepted_tokens = torch.where( + is_dummy_request, + prev_num_accepted_tokens, + next_num_accepted_tokens, + ) + next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, + next_cache_buf_idx) + replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens + replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx + + +class BaseMambaCacheManager(ABC): + """Abstract interface for accessing mamba/recurrent state caches.""" + + @abstractmethod + def get_state_indices(self, *args, **kwargs) -> torch.Tensor: + """Return slot indices of each request. + + Shape: [max_batch_size] + """ + ... + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + """Return replay metadata tensors and fixed replay sizes.""" + return None + + @property + def use_kda_replay_update(self) -> bool: + """Whether KDA fused verification owns per-slot replay caches.""" + return getattr(self, "_use_kda_replay_update", False) + + @property + def keeps_kda_draft_records(self) -> bool: + """Whether the KDA replay caches come with the records of every draft + of the last verify round (``kda_state_tok``), from which a verify + kernel replays the accepted drafts onto the slot's state.""" + return False + + @abstractmethod + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + """Return conv states for specific layer. + + Shape: [slot_size, conv_dim, d_conv - 1] + """ + ... + + @abstractmethod + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + """Return SSM states for specific layer. + + Shape: [slot_size, num_heads, head_dim, d_state] + """ + ... + + @abstractmethod + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + ... + + @abstractmethod + def is_speculative(self) -> bool: + ... + + @abstractmethod + def mamba_layer_cache( + self, layer_idx: int + ) -> Union['PythonMambaCacheManager.State', + 'PythonMambaCacheManager.SpeculativeState', None]: + ... + + +class PythonMambaCacheManager(BaseResourceManager): + """Pure-Python mamba state manager with speculative decoding support. + + Manages only mamba states (conv + SSM) using PyTorch tensors on GPU. + Supports speculative decoding and disaggregated serving. + """ + + @dataclass(frozen=True, kw_only=True) + class State: + """Base state container for Mamba cache.""" + conv: torch.Tensor + # Represents last state, or "two back" if prev_num_accepted_tokens is > 0 + temporal: torch.Tensor + + # Subclasses override to list fields shared across layers (not indexed by layer). + _SHARED_FIELDS = frozenset() + + def at_layer_idx(self, layer: int): + kwargs = {} + for k, v in vars(self).items(): + if v is None: + kwargs[k] = None + elif k in self._SHARED_FIELDS: + kwargs[k] = v + else: + kwargs[k] = v[layer] + return type(self)(**kwargs) + + @property + def has_kda_replay_caches(self) -> bool: + """Whether the Kimi K3 KDA fused-verify replay caches are + allocated. The base (non-speculative) state never has them; the + speculative state overrides this.""" + return False + + @dataclass(frozen=True, kw_only=True) + class SpeculativeState(State): + """Speculative state with intermediate states for draft tokens. + + Supports three SSM update paths (only one set of tensors is + allocated): + - Legacy: caches full intermediate SSM states (intermediate_ssm) + - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) + - KDA replay: per-slot draft-token caches consumed by the fused + ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted + drafts and commits states in place (kda_conv_*, kda_*_cache); or, + for a verify kernel that keeps the drafts' records, kda_conv_* and + kda_state_tok in place of the kda_*_cache + """ + _SHARED_FIELDS = frozenset({ + "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" + }) + + # Allocated for the legacy and Mamba2-replay paths; None for the + # KDA replay path (the kernel commits conv windows in place). + intermediate_conv_window: torch.Tensor | None = None + + # Legacy path: full intermediate SSM states at each step + intermediate_ssm: torch.Tensor | None = None + + # KDA replay path (fused multi-token verify, kimi_linear). + # Pool invariant under this path: `temporal` holds the state after + # the LAST GOLDEN token; the accepted drafts recorded in + # prev_num_accepted_tokens are pending in these caches and are + # replayed by the kernel at the start of the next verify round. + # Extended conv caches [layers, slots, dim, (W-1) + num_spec] bf16 + # (raw projection inputs, so bf16 is exact), dim-contiguous (stride(dim) == 1): columns [0, W-1) are the + # committed raw-input window, tail columns the pending drafts' raw + # inputs. + kda_conv_q: torch.Tensor | None = None + kda_conv_k: torch.Tensor | None = None + kda_conv_v: torch.Tensor | None = None + # Post-processed per-draft quantities for replay: k + # [layers, slots, num_spec, H*K] bf16, gate [.., num_spec, H*K] fp32 + # (feeds the cumulative-gate exponents), v [.., num_spec, H*V] bf16, + # beta [.., num_spec, H] bf16. q is recomputed from the conv caches. + kda_k_cache: torch.Tensor | None = None + kda_g_cache: torch.Tensor | None = None + kda_v_cache: torch.Tensor | None = None + kda_beta_cache: torch.Tensor | None = None + # Optional: the records of every draft of the last verify round, + # [slots, 3, num_spec, H, K] fp32: the row innovations vn [num_spec, + # H, V], then beta * k and the decay [num_spec, H, K] (V = K). The + # SSM pool holds the state after the round's first, non-draft token; + # when the round accepted n > 0 drafts (prev_num_accepted_tokens), a + # verify kernel starts the next round by replaying the first n + # records onto it, S = decay * S + vn (beta * k), instead of + # replaying them from the caches above, which are then not + # allocated. A slot reset for a new request has n = 0, so a previous + # owner's records are never read. + kda_state_tok: torch.Tensor | None = None + + # Replay path: compact double-buffered cache + # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). + # 0 means temporal saved state is actually the last state, not two back. + prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers + cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers + # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so + # CUDA graph replay uses fresh SR draws without allocating RNG tensors. + # (cache,) int64 - shared across layers + mamba_ssm_rand_seed: torch.Tensor | None = None + old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) + old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) + # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. + old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + + @property + def has_kda_replay_caches(self) -> bool: + """True on the KDA fused-verify replay path: the extended conv + caches are allocated, with the per-draft replay caches of the + fused ``trtllm::kda_mtp_decode`` verify or, in their place, the + drafts' records (``kda_state_tok``).""" + return self.kda_conv_q is not None + + def commit_conv_window(self, slot_indices: torch.Tensor, + conv_pool: torch.Tensor) -> None: + """Seed the KDA replay caches from the live ``W - 1`` conv pool.""" + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ + copy_kda_replay_conv_window + + copy_kda_replay_conv_window( + conv_pool, + self.kda_conv_q, + self.kda_conv_k, + self.kda_conv_v, + slot_indices, + ) + + def __init__( + self, + d_state: int, + d_conv: int, + num_heads: int, + n_groups: int, + head_dim: int, + num_layers: int, + max_batch_size: int, + spec_state_size: int, + mapping: Mapping, + dtype: torch.dtype, + ssm_cache_dtype: torch.dtype, + layer_mask: Optional[List[bool]] = None, + speculative_num_draft_tokens: Optional[int] = None, + model_type: str = "nemotron_hybrid", + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + ) -> None: + + self.mamba_ssm_cache_dtype = ssm_cache_dtype + self.speculative_num_draft_tokens = speculative_num_draft_tokens + self.spec_state_size = spec_state_size + self._use_replay_state_update = use_replay_state_update + # KDA replay path (kimi_linear fused multi-token verify). Mutually + # exclusive with use_replay_state_update; requires speculative mode. + self._kda_replay_num_spec = kda_replay_num_spec + self._use_kda_replay_update = kda_replay_num_spec is not None + if self._use_kda_replay_update: + assert not use_replay_state_update, ( + "kda_replay_num_spec and use_replay_state_update are " + "mutually exclusive") + assert speculative_num_draft_tokens is not None, ( + "KDA replay caches require speculative decoding") + assert kda_replay_num_spec == speculative_num_draft_tokens, ( + f"KDA replay cache width ({kda_replay_num_spec}) must match " + f"the draft length ({speculative_num_draft_tokens}): the " + "fused verify kernel is compiled with a static NUM_SPEC") + self.replay_history_size: Optional[int] = None + self.replay_step_width: Optional[int] = None + # When True, allocate the per-slot Philox seed buffer even outside + # the replay path so the non-replay flashinfer SR kernel reads a + # persistent deterministic seed instead of a per-call torch.randint. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh cache-slot assignment. Combined + # with slot index and rank offset to produce reproducible per-slot + # seed values. Starts at 0 so the post-init "reset" stream is + # disjoint from the counter=0 stream used at allocation time. + self._seed_request_counter = 0 + + # get tp size + tp_size = mamba_effective_tp_size(mapping) + + # derive mamba parameters for conv and ssm states + d_inner = head_dim * num_heads + conv_dim = d_inner + 2 * n_groups * d_state + nheads = num_heads + + # check that can be partitioned + assert nheads % tp_size == 0, "nheads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + + # partition conv_dim and nheads + d_inner_local = d_inner // tp_size + ng_ds_local = n_groups * d_state // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + d_inner = d_inner // tp_size + + # Per-section dims for conv_state. + # Qwen3-Next: [Q | K | V] = [ng*ds, ng*ds, d_inner] + # Nemotron_hybrid: [x | B | C] = [d_inner, ng*ds, ng*ds] + if model_type == "qwen3_next": + self.conv_section_dims = [ng_ds_local, ng_ds_local, d_inner_local] + elif model_type == "nemotron_hybrid": + self.conv_section_dims = [d_inner_local, ng_ds_local, ng_ds_local] + else: + raise ValueError(f"Unsupported model type: {model_type}") + + # conv and ssm states device + device = torch.device("cuda") + + pp_layers, num_layers = get_pp_layers( + num_layers, + mapping, + layer_mask=layer_mask, + ) + num_local_layers = len(pp_layers) + self.mamba_layer_offsets = { + idx: offset + for offset, idx in enumerate(pp_layers) + } + + conv_state_shape = (conv_dim, d_conv - 1) + ssm_state_shape = (nheads, head_dim, d_state) + + # create mamba conv and ssm states + conv_states = torch.zeros( + size=(num_local_layers, max_batch_size) + conv_state_shape, + dtype=dtype, + device=device, + ) + + ssm_states = torch.zeros( + size=(num_local_layers, max_batch_size) + ssm_state_shape, + dtype=self.mamba_ssm_cache_dtype, + device=device, + ) + + # Per-slot Philox seeds. Allocated whenever stochastic rounding can + # fire, even outside the replay path and even when MTP/spec is off, + # so all consumers (replay kernel, MTP non-replay flashinfer, + # non-MTP flashinfer) read a persistent deterministic seed from the + # cache manager instead of calling torch.randint per forward. + self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None + if (self._use_replay_state_update + or self._mamba_ssm_stochastic_rounding): + self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + max_batch_size, self._seed_rank_offset, device) + + # create state container + if speculative_num_draft_tokens is not None: + T = speculative_num_draft_tokens + 1 + self.replay_step_width = T + + # Conv intermediate cache — legacy and Mamba2-replay paths only. + # The KDA replay kernel commits conv windows in place, so the + # per-step window scratch is not needed. + intermediate_conv_window_cache = None + if not self._use_kda_replay_update: + intermediate_conv_window_cache = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + + conv_state_shape, + dtype=dtype, + device=device, + ) + + # SSM speculative cache — path-specific tensors + spec_kwargs = {} + # Share the manager-level seed buffer through SpeculativeState + # so the MTP path can still read it via layer_cache. + if self._mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed + if self._use_kda_replay_update: + # KDA replay caches for the fused multi-token verify kernel + # (trtllm::kda_mtp_decode). Per cache slot (persistent + # across rounds — the next round replays from them), unlike + # the batch-row-indexed intermediate buffers. + M = self._kda_replay_num_spec + # conv_dim covers the [q | k | v] short-conv sections; the + # kernel consumes them as three dim-contiguous caches. + assert conv_dim % 3 == 0, ( + "KDA replay caches expect the [q | k | v] conv-state " + "sectioning (3 equal sections)") + section_dim = conv_dim // 3 + # d_conv is KDA's convolution width; the live pool stores + # its W - 1 committed input columns. + w_kernel = d_conv + extended_s = w_kernel - 1 + M + + def _dim_contiguous_conv_cache(): + # raw bf16 projection inputs: bf16 is exact and halves + # the per-draft replay traffic of the verify kernel + return torch.zeros(num_local_layers, + max_batch_size, + extended_s, + section_dim, + dtype=torch.bfloat16, + device=device).transpose(-1, -2) + + spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device) + spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() + spec_kwargs['kda_k_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.bfloat16, + device=device) + spec_kwargs['kda_g_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.float32, + device=device) + spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.bfloat16, + device=device) + spec_kwargs['kda_beta_cache'] = _allocate_kda_beta_cache( + (num_local_layers, max_batch_size, M, nheads), device) + ssm_spec_cache = [ + spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], + spec_kwargs['kda_conv_v'], spec_kwargs['kda_k_cache'], + spec_kwargs['kda_g_cache'], spec_kwargs['kda_v_cache'], + spec_kwargs['kda_beta_cache'] + ] + spec_path_label = "kda-replay" + elif self._use_replay_state_update: + assert n_groups % tp_size == 0, \ + "replay state update requires n_groups divisible by tp_size" + n_groups_per_rank = n_groups // tp_size + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) + + # Compact replay cache. + spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device) + spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, + dtype=torch.int32, + device=device) + spec_kwargs['old_x'] = torch.zeros(num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + head_dim, + dtype=dtype, + device=device) + spec_kwargs['old_B'] = torch.zeros(num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + n_groups_per_rank, + d_state, + dtype=dtype, + device=device) + spec_kwargs['old_dt'] = torch.zeros(num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device) + spec_kwargs['old_dA_cumsum'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device) + ssm_spec_cache = [ + spec_kwargs['old_x'], spec_kwargs['old_B'], + spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] + ] + spec_path_label = "replay" + else: + # Legacy: full intermediate SSM states at each step + spec_kwargs['intermediate_ssm'] = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + + ssm_state_shape, + dtype=self.mamba_ssm_cache_dtype, + device=device) + ssm_spec_cache = [spec_kwargs['intermediate_ssm']] + spec_path_label = "legacy" + + self.mamba_cache = self.SpeculativeState( + conv=conv_states, + temporal=ssm_states, + intermediate_conv_window=intermediate_conv_window_cache, + **spec_kwargs, + ) + + logger.info( + f"Mamba Cache ({spec_path_label}) is allocated. " + f"max_mamba_cache_size: {max_batch_size}, " + f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " + f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " + f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " + "intermediate_conv_window_cache size: " + f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" + ) + else: + self.mamba_cache = self.State( + conv=conv_states, + temporal=ssm_states, + ) + + logger.info( + f"Mamba Cache is allocated. " + f"max_mamba_cache_size: {max_batch_size}, " + f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " + f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB" + ) + + # mamba cache available blocks + self.mamba_cache_free_blocks = [i for i in range(max_batch_size)] + + # mamba cache index, maps request_id -> state indices + self.mamba_cache_index: Dict[int, int] = {} + self._dummy_request_ids: set[int] = set() + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = torch.zeros(max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros(max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) + + # Permanent slot shared by every CUDA-graph padding sentinel id + # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per + # draft length). Pool sizing must include +1 headroom for this; + # see MixedMambaHybridCacheManager. + self._padding_slot: int = self.mamba_cache_free_blocks.pop() + + # Reserved slot for the attention-DP padding dummy. Keep it out of the + # free pool so dummy insertion never consumes real-request capacity. + self._attention_dp_dummy_slot: Optional[int] = ( + self.mamba_cache_free_blocks.pop() + if mapping.enable_attention_dp else None) + + # save intermediate state indices for requests + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device=device) + + # Physical tensor rows include reserved dummy slots. Resource capacity + # is the number of real request slots remaining after those + # reservations. + self._max_batch_size = max_batch_size + self._max_resource_count = len(self.mamba_cache_free_blocks) + + def get_max_resource_count(self) -> int: + """Return the maximum number of real requests that can be cached.""" + return self._max_resource_count + + def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: + """Return the prefix of *context_requests* that fits in the + available Mamba cache blocks. Requests that already have a + cached block do not consume a free block.""" + free = len(self.mamba_cache_free_blocks) + result = [] + for r in context_requests: + if r.py_request_id in self.mamba_cache_index: + result.append(r) + elif free > 0: + result.append(r) + free -= 1 + else: + break + return result + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + """For Mamba cache manager, we always need one slot per request.""" + return 1 + + @torch.inference_mode() + def _prepare_mamba_cache_blocks(self, request_ids: List[int]): + for r in request_ids: + if r in self.mamba_cache_index: + continue + if len(self.mamba_cache_free_blocks) == 0: + raise RuntimeError("run out of mamba cache blocks") + block = self.mamba_cache_free_blocks.pop() + self.mamba_cache_index[r] = block + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + # Fresh request: no drafts pending in the replay caches. + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + if self._mamba_ssm_rand_seed is not None: + # Deterministic per-slot rotation on fresh assignment. + # `block` is pulled from mamba_cache_free_blocks, which + # excludes _padding_slot by construction (see __init__), + # so padding sentinels never reach this branch. + self._seed_request_counter += 1 + self._mamba_ssm_rand_seed[block] = ( + _compute_deterministic_mamba_seed( + self._seed_request_counter, block, + self._seed_rank_offset)) + + @torch.inference_mode() + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + """Seed the fused-verify KDA replay conv caches from the conv pool. + + On a disaggregated generation server the ctx->gen transfer populates + only the base ``conv`` / ``temporal`` pools (see + ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` + replay caches consumed by ``trtllm::kda_mtp_decode`` are normally + seeded by a *local* prefill/decode via + ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a + previous occupant's window) for a transferred request — corrupting + the recurrent state from the first verify step onward. Call this + after the state transfer completes, before the first generation + forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row + stores the committed ``W - 1`` raw inputs directly. Pending-draft + scratch is cleared (no drafts are pending for a freshly transferred + request). + """ + if not (self._use_kda_replay_update + and isinstance(self.mamba_cache, self.SpeculativeState)): + return + blocks = [ + self.mamba_cache_index[rid] for rid in request_ids + if rid in self.mamba_cache_index + ] + if not blocks: + return + conv = self.mamba_cache.conv # [L, slots, 3D, W - 1] + idx = torch.tensor(sorted(set(blocks)), + dtype=torch.long, + device=conv.device) + d = conv.shape[2] // 3 + committed = conv.shape[3] + cs = conv.index_select(1, idx) + for cache, section in ( + (self.mamba_cache.kda_conv_q, cs[:, :, :d]), + (self.mamba_cache.kda_conv_k, cs[:, :, d:2 * d]), + (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d:]), + ): + # cache: [L, slots, D, committed + num_spec]; zero the draft + # tail columns and seed the committed window in one copy. + seeded = torch.zeros( + (cache.shape[0], idx.numel()) + cache.shape[2:], + dtype=cache.dtype, + device=cache.device) + seeded[:, :, :, :committed] = section.to(cache.dtype) + cache.index_copy_(1, idx, seeded) + for buf in (self.mamba_cache.kda_k_cache, self.mamba_cache.kda_g_cache, + self.mamba_cache.kda_v_cache, + self.mamba_cache.kda_beta_cache): + buf.index_fill_(1, idx, 0) + self.mamba_cache.prev_num_accepted_tokens[idx] = 0 + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + context_ids = [ + i.py_request_id for i in scheduled_batch.context_requests + ] + generation_ids = [ + i.py_request_id for i in scheduled_batch.generation_requests + ] + request_ids = context_ids + generation_ids + self._prepare_mamba_cache_blocks(request_ids) + self._refresh_dummy_request_mask([req.is_dummy for req in requests]) + + def _is_padding_sentinel(self, request_id: int) -> bool: + # cuda_graph_runner caches one dummy per runtime_draft_len value + # (see _get_padded_batch), so any id in the range of dummy request IDs + # may be live concurrently. + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDA_GRAPH_DUMMY_REQUEST_ID + max_dl = self.speculative_num_draft_tokens or 0 + return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= + CUDA_GRAPH_DUMMY_REQUEST_ID) + + @torch.inference_mode() + def add_dummy_requests(self, request_ids: List[int], **kwargs): + # Sentinels alias to the shared _padding_slot; non-sentinel + # dummies (warmup, attention-DP idle padding) get their own + # slot and are freed individually. + if not request_ids: + return + self._dummy_request_ids.update(request_ids) + for r in request_ids: + if r in self.mamba_cache_index: + block = self.mamba_cache_index[r] + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + continue + if self._is_padding_sentinel(r): + block = self._padding_slot + elif (r == ATTENTION_DP_DUMMY_REQUEST_ID + and self._attention_dp_dummy_slot is not None): + block = self._attention_dp_dummy_slot + else: + if len(self.mamba_cache_free_blocks) == 0: + raise RuntimeError("run out of mamba cache blocks") + block = self.mamba_cache_free_blocks.pop() + self.mamba_cache_index[r] = block + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + + def free_resources(self, request: LlmRequest): + request_id = request.py_request_id + if request_id not in self.mamba_cache_index: + return + self._dummy_request_ids.discard(request_id) + block = self.mamba_cache_index.pop(request_id) + # Reserved slots must not re-enter the real-request free pool. + if block != self._padding_slot and \ + block != self._attention_dp_dummy_slot: + self.mamba_cache_free_blocks.append(block) + + def get_state_indices(self, request_ids: List[int], + is_padding: List[bool]) -> List[int]: + assert len(request_ids) == len(is_padding) + indices = [self.mamba_cache_index[rid] for rid in request_ids] + is_dummy = [ + rid in self._dummy_request_ids or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.as_tensor(is_dummy, dtype=torch.bool)) + mask_staging = torch.empty_like(self._dummy_request_mask_host, + pin_memory=prefer_pinned()) + mask_staging.copy_(self._dummy_request_mask_host) + self._dummy_request_mask.copy_(mask_staging, non_blocking=True) + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).conv + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).temporal + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if not isinstance(self.mamba_cache, self.SpeculativeState): + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).intermediate_ssm + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + """Get intermediate conv states for speculative decoding.""" + if not isinstance(self.mamba_cache, self.SpeculativeState): + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx( + layer_offset).intermediate_conv_window + + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_x[layer_offset] + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_B[layer_offset] + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dt[layer_offset] + + def get_replay_old_dA_cumsum(self, + layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dA_cumsum[layer_offset] + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.cache_buf_idx + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.prev_num_accepted_tokens + + def is_speculative(self) -> bool: + return isinstance(self.mamba_cache, self.SpeculativeState) + + def mamba_layer_cache(self, + layer_idx: int) -> Union[State, SpeculativeState]: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.mamba_ssm_cache_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Return the persistent (cache_size,) int64 Philox seed buffer or + None when stochastic rounding is not active for this manager. + + Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. + Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a + view that matches the consuming kernel's expected shape. + """ + return self._mamba_ssm_rand_seed + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + @property + def use_kda_replay_update(self) -> bool: + return self._use_kda_replay_update + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + if (not self._use_replay_state_update + or not isinstance(self.mamba_cache, self.SpeculativeState) + or self.mamba_cache.prev_num_accepted_tokens is None + or self.mamba_cache.cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=( + self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=self.mamba_cache.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def shutdown(self): + """Release tensor memory.""" + # Clear mamba cache states + empty = torch.tensor([]) + + def _drop(tensor): + return empty if tensor is not None else None + + if isinstance(self.mamba_cache, self.SpeculativeState): + self.mamba_cache = self.SpeculativeState( + conv=empty, + temporal=empty, + intermediate_conv_window=_drop( + self.mamba_cache.intermediate_conv_window), + intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), + prev_num_accepted_tokens=_drop( + self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx), + mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed), + old_x=_drop(self.mamba_cache.old_x), + old_B=_drop(self.mamba_cache.old_B), + old_dt=_drop(self.mamba_cache.old_dt), + old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), + kda_conv_q=_drop(self.mamba_cache.kda_conv_q), + kda_conv_k=_drop(self.mamba_cache.kda_conv_k), + kda_conv_v=_drop(self.mamba_cache.kda_conv_v), + kda_k_cache=_drop(self.mamba_cache.kda_k_cache), + kda_g_cache=_drop(self.mamba_cache.kda_g_cache), + kda_v_cache=_drop(self.mamba_cache.kda_v_cache), + kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), + ) + else: + self.mamba_cache = self.State(conv=empty, temporal=empty) + + torch.cuda.empty_cache() + + @torch.compile(options={"max-autotune": True}) + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None): + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = num_accepted_tokens[ + num_contexts:num_contexts + num_gens] - 1 + # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. + accepted_positions = (accepted_leaf_positions if accepted_leaf_positions + is not None else num_accepted_draft_tokens) + state_indices_d = state_indices[num_contexts:num_contexts + num_gens] + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_kda_replay_update: + # KDA replay: the fused verify kernel already committed the SSM + # state and conv windows (in place, after the golden token) and + # cached this round's drafts. All that remains is recording how + # many of those drafts the sampler accepted, so the next round's + # kernel launch replays exactly that prefix. + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + prev = self.mamba_cache.prev_num_accepted_tokens + current = prev[state_indices_d] + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + prev[state_indices_d] = torch.where(is_dummy_request, current, + accepted) + return + + if self._use_replay_state_update: + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) + else: + # Legacy: copy accepted SSM state from intermediate cache. + ssm_states = self.mamba_cache.temporal + intermediate_ssm_cache = self.mamba_cache.intermediate_ssm + accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, + accepted_positions] + ssm_states[:, state_indices_d, :] = accepted_ssm_state + + # Conv: both paths save all intermediate conv windows, carry over the accepted one. + conv_states = self.mamba_cache.conv + intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window + accepted_conv_state = intermediate_conv_window_cache[:, + src_state_indices, + accepted_positions] + conv_states[:, state_indices_d, :] = accepted_conv_state + + +class MambaCacheManager(BaseResourceManager, BaseMambaCacheManager): + """Facade for standalone mamba state management (no KV cache). + + Delegates to PythonMambaCacheManager. + """ + + def __init__( + self, + d_state: int, + d_conv: int, + num_heads: int, + n_groups: int, + head_dim: int, + num_layers: int, + max_batch_size: int, + spec_state_size: int, + mapping: Mapping, + dtype: torch.dtype, + ssm_cache_dtype: torch.dtype, + layer_mask: Optional[List[bool]] = None, + stream: Optional[torch.cuda.Stream] = None, + speculative_num_draft_tokens: Optional[int] = None, + model_type: str = "nemotron_hybrid", + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + ) -> None: + max_num_sequences = max_batch_size * mapping.pp_size + + self._impl = PythonMambaCacheManager( + d_state=d_state, + d_conv=d_conv, + num_heads=num_heads, + n_groups=n_groups, + head_dim=head_dim, + num_layers=num_layers, + max_batch_size=max_num_sequences, + spec_state_size=spec_state_size, + mapping=mapping, + dtype=dtype, + ssm_cache_dtype=ssm_cache_dtype, + layer_mask=layer_mask, + speculative_num_draft_tokens=speculative_num_draft_tokens, + model_type=model_type, + use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, + ) + + def get_max_resource_count(self) -> int: + return self._impl.get_max_resource_count() + + def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: + return self._impl.filter_ctx_requests_by_capacity(context_requests) + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + return self._impl.get_needed_resource_to_completion(request) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + self._impl.prepare_resources(scheduled_batch) + + def free_resources(self, request: LlmRequest): + self._impl.free_resources(request) + + def add_dummy_requests(self, request_ids: List[int], **kwargs): + self._impl.add_dummy_requests(request_ids, **kwargs) + + def get_state_indices( + self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None + ) -> Union[torch.Tensor, List[int]]: + return self._impl.get_state_indices(request_ids, is_padding) + + @property + def mamba_cache_free_blocks(self) -> List[int]: + return self._impl.mamba_cache_free_blocks + + @property + def mamba_cache_index(self) -> Dict[int, int]: + return self._impl.mamba_cache_index + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self._impl.get_conv_states(layer_idx) + + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self._impl.get_ssm_states(layer_idx) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self._impl.get_mamba_ssm_cache_dtype() + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Delegate to the underlying Python manager.""" + return self._impl.get_mamba_ssm_rand_seed() + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + @property + def use_kda_replay_update(self) -> bool: + return self._impl.use_kda_replay_update + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + return self._impl.get_replay_state_update_metadata() + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_intermediate_ssm_states(layer_idx) + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_intermediate_conv_states(layer_idx) + + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_x(layer_idx) + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_B(layer_idx) + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_dt(layer_idx) + + def get_replay_old_dA_cumsum(self, + layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_dA_cumsum(layer_idx) + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + return self._impl.get_replay_cache_buf_idx() + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + return self._impl.get_replay_prev_num_accepted_tokens() + + def is_speculative(self) -> bool: + return self._impl.is_speculative() + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + return self._impl.mamba_layer_cache(layer_idx) + + def shutdown(self): + self._impl.shutdown() + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None): + # Non-speculative configs don't allocate intermediate state; the + # promotion is a clean no-op. + if not self._impl.is_speculative(): + return + self._impl.update_mamba_states(attn_metadata, num_accepted_tokens, + state_indices, accepted_leaf_positions) + + +class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): + """Shared interface and state plumbing for hybrid Mamba managers. + + Concrete storage, state-index, and resource lifecycles remain owned by the + Cpp and V2 implementations. + """ + + _supports_additional_snapshot_offsets = False + + @property + def cuda_state_indices(self) -> torch.Tensor: + return self._state_indices_buffer.tensor + + @property + def _dummy_request_mask(self) -> Optional[torch.Tensor]: + if self._dummy_request_mask_buffer is None: + return None + return self._dummy_request_mask_buffer.tensor + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size: int) -> None: + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if spec_config is None or self.local_num_mamba_layers == 0: + return + + tokens_per_gen_step = spec_config.tokens_per_gen_step + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + tokens_per_gen_step + ] + self.ssm_state_shape, + dtype=self.ssm_state_dtype, + device="cuda", + ) + self.intermediate_conv_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step + ] + self.conv_state_shape, + dtype=self.conv_state_dtype, + device="cuda", + ) + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + """Allocate replay tensors shared by the Cpp and V2 state pools.""" + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask_buffer = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + if (self.local_num_mamba_layers == 0 + or (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding)): + return False + + assert device is not None + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device) + if spec_config is None or not self._use_replay_state_update: + return False + + history_size = self.replay_history_size + assert history_size is not None + nheads, head_dim, d_state = self.ssm_state_shape + common_shape = [self.local_num_mamba_layers, cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.cache_buf_idx = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.old_x = torch.zeros( + common_shape + [history_size, nheads, head_dim], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_shape + [history_size, self._n_groups_per_rank, d_state], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + return True + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask_buffer.update(self._dummy_request_mask_host) + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + if num_contexts == 0: + return + + context_slots = self.cuda_state_indices[:num_contexts].long() + if (self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): + self.prev_num_accepted_tokens.index_fill_(0, context_slots, 0) + self.cache_buf_idx.index_fill_(0, context_slots, 0) + if self.old_x is not None: + self.old_x.index_fill_(1, context_slots, 0) + if self.old_B is not None: + self.old_B.index_fill_(1, context_slots, 0) + if self.old_dt is not None: + self.old_dt.index_fill_(1, context_slots, 0) + if self.old_dA_cumsum is not None: + self.old_dA_cumsum.index_fill_(1, context_slots, 0) + + if self.mamba_ssm_rand_seed is None: + return + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = self._host_state_indices[:num_contexts].tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, rank_offset) + for slot in host_slots + ] + self.mamba_ssm_rand_seed[context_slots] = torch.tensor( + new_seeds, dtype=torch.int64, + pin_memory=prefer_pinned()).to(self.mamba_ssm_rand_seed.device, + non_blocking=True) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update( + range(interval, request.prompt_len + 1, interval)) + if self._supports_additional_snapshot_offsets: + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_ssm_states is None: + return None + return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_conv_states is None: + return None + return self.intermediate_conv_states[ + self.mamba_layer_offsets[layer_idx]] + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + conv = self.get_conv_states(layer_idx) + ssm = self.get_ssm_states(layer_idx) + if self.spec_config is not None: + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = {} + if self.mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed + if self._use_replay_state_update: + spec_kwargs['old_x'] = self.old_x[layer_offset] + spec_kwargs['old_B'] = self.old_B[layer_offset] + spec_kwargs['old_dt'] = self.old_dt[layer_offset] + spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + spec_kwargs['cache_buf_idx'] = self.cache_buf_idx + spec_kwargs['prev_num_accepted_tokens'] = ( + self.prev_num_accepted_tokens) + else: + spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ + layer_offset] + return PythonMambaCacheManager.SpeculativeState( + conv=conv, + temporal=ssm, + intermediate_conv_window=self. + intermediate_conv_states[layer_offset], + **spec_kwargs, + ) + return PythonMambaCacheManager.State(conv=conv, temporal=ssm) + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) + if (not self._use_replay_state_update + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + return getattr(self, 'mamba_ssm_rand_seed', None) + + +def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: + """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" + pool_size = max_batch_size + # One permanent slot is shared by every CUDA-graph padding sentinel. + pool_size += 1 + if mapping.enable_attention_dp: + # Attention-DP can insert a transient dummy request on an otherwise + # idle rank. Keep this headroom internal so scheduler-visible + # max_batch_size still limits real requests. + pool_size += 1 + return pool_size + + +class MixedMambaHybridCacheManager(KVCacheManager, MambaCacheManager, + MambaHybridCacheManager): + """Hybrid cache manager combining separate KVCacheManager and MambaCacheManager. + + Manages KV cache and mamba states in independent pools, with support of + speculative decoding and disaggregated serving. + Does not support block reuse / prefix caching for mamba states. + """ + + # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the + # cache manager __init__ coerces a positive spec recompute tail to a full + # re-prefill (see _has_recurrent_state on the attention-only base classes). + _has_recurrent_state = True + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + # kv cache parameters + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + layer_mask: List[bool], + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + max_num_tokens: int = 8192, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + is_estimating_kv_cache: bool = False, + execution_stream: Optional[torch.cuda.Stream] = None, + model_type: str = "nemotron_hybrid", + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + # Per-pool configurations forwarded to the C++ KVCacheManager ctor. + # Lets a single manager host pools with mixed shapes (e.g. Gemma4 + # hybrid attention). See KVCacheManager.__init__. + pool_configurations: Optional[List[PoolConfiguration]] = None, + ) -> None: + + # mamba hybrid cache requires block reuse to be disabled in KV cache config + assert not kv_cache_config.enable_block_reuse, ( + "mamba hybrid cache requires block reuse to be disabled in KV cache config" + ) + + # Host-drafter spec modes (NGram) have no spec worker to call + # update_mamba_states; this manager promotes accepted states itself + # in update_resources. One-model modes (MTP/Eagle/DFlash) and the + # suffix-automaton worker promote from their spec workers and must + # NOT be promoted twice. + self._promote_states_in_update_resources = ( + spec_config is not None + and getattr(spec_config, "decoding_type", None) == "NGram") + + pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) + + MambaCacheManager.__init__( + self, + mamba_d_state, + mamba_d_conv, + mamba_num_heads, + mamba_n_groups, + mamba_head_dim, + mamba_num_layers, + pool_size, + max_batch_size, + mapping, + mamba_cache_dtype, + mamba_ssm_cache_dtype, + mamba_layer_mask, + execution_stream, + speculative_num_draft_tokens=(spec_config.tokens_per_gen_step - 1 + if spec_config is not None else None), + model_type=model_type, + use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, + ) + + # initialize kv cache manager + KVCacheManager.__init__( + self, + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_num_tokens=max_num_tokens, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + execution_stream=execution_stream, + is_draft=is_draft, + pool_configurations=pool_configurations, + ) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + MambaCacheManager.prepare_resources(self, scheduled_batch) + KVCacheManager.prepare_resources(self, scheduled_batch) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + MambaCacheManager.free_resources(self, request) + KVCacheManager.free_resources(self, request, pin_on_release) + + def add_dummy_requests(self, request_ids: List[int], **kwargs): + MambaCacheManager.add_dummy_requests(self, request_ids) + return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) + + def shutdown(self): + MambaCacheManager.shutdown(self) + KVCacheManager.shutdown(self) + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None): + KVCacheManager.update_resources(self, scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) + + def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): + """Commit accepted verification states for host-drafter spec modes. + + One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton + worker call update_mamba_states themselves right after on-device + acceptance. Host-drafter modes (NGram) have no spec worker: + acceptance lands on the requests as + ``py_num_accepted_draft_tokens`` during sampler update, and this + hook — running right after, alongside the KV rewind — promotes the + accepted step's intermediate state into the live pools (or, on the + KDA replay path, records the accepted-draft count for the next + round's replay). Without it, the pools would keep the + pre-verification state and the next forward would resume from a + stale prefix. + """ + if not self._promote_states_in_update_resources: + return + if not self.is_speculative() or attn_metadata is None: + return + gen_requests = scheduled_batch.generation_requests + if not gen_requests: + return + drafted = [ + r for r in gen_requests + if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 + ] + if not drafted: + # Drafter skipped this step (e.g. speculation gated off): the + # forward ran the plain in-place decode path; nothing to promote. + return + assert len(drafted) == len(gen_requests), ( + "mixed drafted/undrafted generation batch is not supported for " + "hybrid state promotion (drafts are padded to the static max)") + device = self._impl.mamba_cache.temporal.device + num_contexts = len(scheduled_batch.context_requests) + num_accepted = torch.tensor( + [0] * num_contexts + + [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], + dtype=torch.int32, + device=device) + # Batch-ordered slots (contexts then gens), matching the ordering + # the forward used for the intermediate scratch buffers. Requests + # that finished this step were already freed by response handling + # (which runs before update_resources) and are gone from the index; + # their rows must stay in place for alignment, so redirect them to + # the reserved padding slot (a harmless scratch write). + slot_index = self.mamba_cache_index + padding_slot = self._impl._padding_slot + state_indices = torch.tensor([ + slot_index.get(r.py_request_id, padding_slot) + for r in scheduled_batch.context_requests + gen_requests + ], + dtype=torch.int32, + device=device) + self.update_mamba_states(attn_metadata, num_accepted, state_indices) + + +@triton.jit +def _promote_mamba_state_kernel( + src_ptr, + dst_ptr, + src_idx_ptr, + accepted_position_source_ptr, + num_accepted_tokens_ptr, + blk_ptr, + replay_pnat_ptr, + replay_cache_buf_idx_ptr, + dummy_request_mask_ptr, + num_gens, + count, + src_s_layer, + src_s_row, + src_s_step, + dst_s_layer, + dst_s_block, + POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, + REPLAY_STEP_WIDTH: tl.constexpr, + REPLAY_HISTORY_SIZE: tl.constexpr, + UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, + BLOCK: tl.constexpr, +): + # One program copies a BLOCK-sized tile of the contiguous inner state for a + # single (layer, gen) pair. grid = (num_layers * num_gens, ceil(count/BLOCK)). + pair = tl.program_id(0) + tile = tl.program_id(1) + layer = (pair // num_gens).to(tl.int64) + g = pair % num_gens + row = tl.load(src_idx_ptr + g).to(tl.int64) + accepted_position = tl.load(accepted_position_source_ptr + g) + if POSITION_SOURCE_IS_TOKEN_COUNT: + accepted_position -= 1 + acc = accepted_position.to(tl.int64) + blk = tl.load(blk_ptr + g).to(tl.int64) + + # The replay metadata update is batch-scoped, so layer 0's first copy tile + # owns it. This folds PNAT and active-buffer maintenance into the existing + # accepted-convolution-state promotion launch. + if UPDATE_REPLAY_BOOKKEEPING: + if (layer == 0) & (tile == 0): + num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) + is_dummy_request = tl.load(dummy_request_mask_ptr + g) + pnat = tl.load(replay_pnat_ptr + blk) + cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) + wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE + next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, + pnat + num_accepted_tokens) + next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) + tl.store(replay_cache_buf_idx_ptr + blk, + next_cache_buf_idx, + mask=~is_dummy_request) + + # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows + # int32 and would corrupt addresses / fault. + src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step + dst_base = layer * dst_s_layer + blk * dst_s_block + offs = tile * BLOCK + tl.arange(0, BLOCK) + mask = offs < count + v = tl.load(src_ptr + src_base + offs, mask=mask) + tl.store(dst_ptr + dst_base + offs, v, mask=mask) + + +def _promote_mamba_state_triton( + dst: torch.Tensor, + intermediate: torch.Tensor, + src_state_indices: torch.Tensor, + accepted_position_source: torch.Tensor, + dst_state_indices: torch.Tensor, + num_accepted_tokens: Optional[torch.Tensor] = None, + position_source_is_token_count: bool = False, + replay_pnat: Optional[torch.Tensor] = None, + replay_cache_buf_idx: Optional[torch.Tensor] = None, + dummy_request_mask: Optional[torch.Tensor] = None, + replay_step_width: int = 0, + replay_history_size: int = 0, + BLOCK: int = 2048) -> None: + """Scatter each generation request's accepted draft-step recurrent state + from the per-request intermediate buffer into the unified C++ pool view. + + Args: + dst: ``[num_layers, num_blocks, *state_shape]`` view of the C++ pool. + The block dim (dim 1) is strided (the pool interleaves + ``ssm_bytes | conv_bytes`` per block), so ``dst`` is non-contiguous. + The kernel writes through ``dst``'s real strides, touching only this + state's bytes -- the case that defeats torch.compile / aot_autograd + (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 + pool). + intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. + src_state_indices, accepted_position_source, dst_state_indices: + ``[num_gens]`` int. + num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required + when replay bookkeeping is enabled. + position_source_is_token_count: Whether ``accepted_position_source`` + contains accepted-token counts instead of zero-based intermediate + state positions. When true, the kernel subtracts one while + computing the source position. + replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay + metadata. When supplied together, the kernel advances PNAT and the + active history buffer while promoting the convolution state. + + For each gen ``g``:: + + accepted_position = ( + accepted_position_source[g] - 1 + if position_source_is_token_count + else accepted_position_source[g] + ) + dst[:, dst_state_indices[g]] = intermediate[ + :, src_state_indices[g], accepted_position + ] + + Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. + """ + num_layers = dst.shape[0] + num_gens = src_state_indices.shape[0] + if num_gens == 0: + return + update_replay_bookkeeping = replay_pnat is not None + if update_replay_bookkeeping: + assert num_accepted_tokens is not None + assert replay_cache_buf_idx is not None + assert dummy_request_mask is not None + assert replay_step_width > 0 + assert replay_history_size >= replay_step_width + assert replay_pnat.dtype == torch.int32 + assert replay_cache_buf_idx.dtype == torch.int32 + assert dummy_request_mask.dtype == torch.bool + assert dummy_request_mask.numel() >= num_gens + else: + # Compile-time-disabled pointer arguments still require tensors. + num_accepted_tokens = accepted_position_source + assert replay_cache_buf_idx is None + assert dummy_request_mask is None + replay_pnat = dst_state_indices + replay_cache_buf_idx = dst_state_indices + dummy_request_mask = dst_state_indices + count = 1 + for s in dst.shape[2:]: + count *= s + # Grid dim order matters: dim 0 -> CUDA grid.x (limit 2^31-1), dim 1 -> + # grid.y (limit 65535). Put the (layer, gen) pairs in dim 0 since that is + # what grows with batch size (num_layers*num_gens stays far below 2^31 for + # any real config); the tile count in dim 1 is small (~count/BLOCK). Do NOT + # swap them: pairs would hit the 65535 y-limit at num_layers*num_gens>65535. + grid = (num_layers * num_gens, triton.cdiv(count, BLOCK)) + _promote_mamba_state_kernel[grid]( + intermediate, + dst, + src_state_indices, + accepted_position_source, + num_accepted_tokens, + dst_state_indices, + replay_pnat, + replay_cache_buf_idx, + dummy_request_mask, + num_gens, + count, + intermediate.stride(0), + intermediate.stride(1), + intermediate.stride(2), + dst.stride(0), + dst.stride(1), + POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, + REPLAY_STEP_WIDTH=replay_step_width, + REPLAY_HISTORY_SIZE=replay_history_size, + UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, + BLOCK=BLOCK, + ) + + +def _mamba_snapshot_rule_counts( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], + tokens_per_block: int, +) -> Tuple[int, int]: + """Return reachable fixed rules and their partial-block upper bound.""" + if not kv_cache_config.enable_block_reuse: + return 0, 0 + + num_rules = 0 + num_unaligned_rules = 0 + state_config = kv_cache_config.mamba_state_config + for offset in set(state_config.additional_snapshot_offsets_from_start): + if max_seq_len is not None and offset > max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += int(offset % tokens_per_block != 0) + for offset in set(state_config.additional_snapshot_offsets_from_end): + # A from-end offset is reachable iff some valid prompt is longer than + # the offset. Its absolute alignment depends on that prompt. + if max_seq_len is not None and offset >= max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += 1 + return num_rules, num_unaligned_rules + + +def _mamba_regular_snapshot_interval( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], +) -> Optional[int]: + if not kv_cache_config.enable_block_reuse: + return None + interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval + if interval is None or interval <= 0: + return None + if max_seq_len is not None and interval > max_seq_len: + return None + return interval + + +def _get_local_mamba_cache_layout( + model_config, + mapping: Mapping, + *, + spec_config=None, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, +): + """Return normalized params and local Mamba/attention layer counts. + + Cache construction and affine sizing must follow the model's PP layout: + partition base layers first, then place appended speculative layers on the + last PP rank. The normalized params retain target masks and the appended + draft-layer count so estimation selects the same combined or per-manager + layout as runtime. + """ + from tensorrt_llm._torch.pyexecutor.config_utils import \ + extract_mamba_kv_cache_params + + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + spec_config=spec_config, + quant_config=model_config.quant_config, + ) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + combined_layer_mask = [ + is_mamba or is_attention for is_mamba, is_attention in zip( + mamba_layer_mask, full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_mamba_layers = sum(mamba_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + local_attention_layers = sum(full_attention_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + return params, local_mamba_layers, local_attention_layers + + +# --- Qwen4-Exp PLE --------------------------------------------------------- +# PLE (the n-gram side path) is a Qwen4-Exp-only feature that rides on the +# hybrid manager's Mamba slots. It is kept in dedicated helpers so it can be +# lifted into a model-owned cache manager wholesale later on. + + +def _qwen4_exp_ple_state_bytes_per_rank( + model_config, + params, + mapping, + *, + spec_config, + use_separate_draft_kv_cache: bool, +) -> int: + """Bytes this rank needs for Qwen4-Exp PLE recurrent state; 0 otherwise.""" + from tensorrt_llm._torch.pyexecutor.config_utils import ( + extract_qwen4_exp_ple_cache_params, is_qwen4_exp) + + # The estimator is also called with minimal model-config stubs that carry no + # `pretrained_config`; such a config is definitionally not Qwen4-Exp. + pretrained_config = getattr(model_config, "pretrained_config", None) + if pretrained_config is None or not is_qwen4_exp(pretrained_config): + return 0 + ple_params = extract_qwen4_exp_ple_cache_params(pretrained_config) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=False, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + combined_layer_mask = [ + is_mamba or is_attention for is_mamba, is_attention in zip( + mamba_layer_mask, full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_ple_layers = sum(layer_id < len(ple_params.ple_layer_mask) + and ple_params.ple_layer_mask[layer_id] + for layer_id in local_layer_indices) + ple_bytes_per_layer = ( + ple_params.short_conv_channels * ple_params.short_conv_state_len * + ple_params.conv_state_dtype.itemsize + + ple_params.ngram_context_len * torch.int64.itemsize) + return local_ple_layers * ple_bytes_per_layer + + +def _estimated_kda_replay_num_spec(model_config, spec_config, + is_disagg: bool) -> Optional[int]: + """Replay-cache width a V2 manager for this model would allocate. + + Mirrors the executor's construction-time choice: KDA models (Kimi K3 and + GLM-5 Next) under speculative decoding get replay caches when the fused + verify kernel can run here (``kda_mtp_replay_num_spec``), except on a + disaggregated context server, which never verifies drafts. + """ + if spec_config is None or _is_disagg_context_server(is_disagg): + return None + pretrained_config = getattr(model_config, "pretrained_config", None) + if pretrained_config is None or not (is_kimi_linear(pretrained_config) + or is_glm5_next(pretrained_config)): + return None + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ + kda_mtp_replay_num_spec + + return kda_mtp_replay_num_spec(spec_config.tokens_per_gen_step - 1) + + +def _estimate_mamba_hybrid_cache_cost( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + tokens_per_block: int, + max_seq_len: Optional[int], + num_reserved_dummy_slots: int, + include_explicit_snapshots: bool, + cap_partial_attention_snapshots: bool, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, + kda_replay_num_spec: Optional[int] = None, + **kwargs, +) -> Tuple[int, int]: + spec_config = kwargs.get("spec_config") + params, local_mamba_layers, local_attention_layers = ( + _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + )) + attention_slope = (KVCacheManager.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + **kwargs, + ) if local_attention_layers > 0 else 0) + state_bytes_per_rank = (local_mamba_layers * + params.get_states_bytes_per_layer(mapping)) + if not is_draft: + state_bytes_per_rank += _qwen4_exp_ple_state_bytes_per_rank( + model_config, + params, + mapping, + spec_config=spec_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + if kda_replay_num_spec is not None: + # KDA replay caches live in the same KVCM slot as the recurrent state, + # so every live, dummy and snapshot slot carries them. + tp_size = mamba_effective_tp_size(mapping) + state_bytes_per_rank += _kda_replay_bytes_per_slot( + _kda_replay_layout(kda_replay_num_spec, params.num_heads // tp_size, + params.head_dim, params.conv_kernel - 1), + local_mamba_layers) + max_resident_sequences = max_batch_size * mapping.pp_size + + if include_explicit_snapshots: + fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( + kv_cache_config, max_seq_len, tokens_per_block) + else: + fixed_rules = 0 + unaligned_fixed_rules = 0 + fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + + max_resident_sequences * fixed_rules) + attention_block_bytes = attention_slope * tokens_per_block + + interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) + has_unaligned_periodic_snapshot = (interval is not None + and interval % tokens_per_block != 0) + if cap_partial_attention_snapshots: + # Snapshot alignment is unknown while estimating cache cost. Once a + # snapshot is possible, reserve one retained partial attention page + # per resident lineage. Dummy requests carry no attention capacity. + has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None + partial_attention_slots = (max_resident_sequences + if has_non_live_ssm_capacity else 0) + else: + partial_attention_slots = (max_resident_sequences * + unaligned_fixed_rules) + intercept = (fixed_state_slots * state_bytes_per_rank + + partial_attention_slots * attention_block_bytes) + + if interval is None: + regular_slope = 0 + else: + regular_slope = math.ceil(state_bytes_per_rank / interval) + if (has_unaligned_periodic_snapshot + and not cap_partial_attention_snapshots): + regular_slope += math.ceil(attention_block_bytes / interval) + return attention_slope + regular_slope, intercept + + +class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): + """Hybrid cache manager storing mamba states inside the KVCacheManager pool. + + Both KV cache blocks and recurrent state blocks are managed by the unified + C++ KVCacheManager, enabling block reuse / prefix caching across attention + and mamba layers. This compatibility manager remains available through the + manager preference override and legacy disaggregated routing. + + """ + + # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the + # cache manager __init__ coerces a positive spec recompute tail to a full + # re-prefill (see _has_recurrent_state on the attention-only base classes). + _has_recurrent_state = True + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[ + List[bool]] = None, # this is the full attention layer mask + is_estimating_kv_cache: bool = False, + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + model_type: str = "nemotron_hybrid", + **kwargs, + ) -> None: + # 3 kinds of layers: + # 1) Mamba layers (mamba_layer_mask is True) + # 2) Full attention layers (full_attention_layer_mask is True) + # 3) Not managed layers (both masks are False) + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + # PP sharding is done across all layers. + # This is called again in the super().__init__, but we want it to run first + # to set up mtp states before the C++ backend is initialized. + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + self.requests = [] + # Seed externally visible mamba fields before any early return so that + # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work + # on ranks with no local mamba layers. + self._use_replay_state_update = use_replay_state_update + self._use_gdn_cached_replay_all_layer_commit = ( + use_replay_state_update and model_type == "qwen3_next" + and self.local_num_mamba_layers > 0) + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) + # Same allocation gate as PythonMambaCacheManager: the rand_seed + # buffer must exist whenever SR can fire, not only on the replay path. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh context-slot assignment; combined + # with the slot index and rank offset to produce reproducible per-slot + # seed values without any torch.randint. + self._seed_request_counter = 0 + self.ssm_state_dtype = mamba_ssm_cache_dtype + # Keep the shared Mamba interface valid on PP ranks that do not own a + # local Mamba layer. + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + + if self.local_num_mamba_layers == 0: + logger.info( + "No local mamba layers for this rank, skipping mamba cache initialization" + ) + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=full_attention_layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, + **kwargs, + ) + # PP ranks replay the same scheduling decisions, so a rank without + # local Mamba layers must still publish the configured boundaries. + self.kv_cache_config = kv_cache_config + self.linear_attention_metadata = LinearAttentionMetadata() + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_config.periodic_snapshot_interval + if kv_cache_config.enable_block_reuse else 0) + return + + # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) + tp_size = mamba_effective_tp_size(mapping) + d_inner = mamba_head_dim * mamba_num_heads + conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if use_replay_state_update: + assert mamba_n_groups % tp_size == 0, \ + "replay state update requires mamba_n_groups divisible by tp_size" + self._n_groups_per_rank = mamba_n_groups // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + self.conv_state_dtype = mamba_cache_dtype + + # Store GLOBAL (pre-TP-division) mamba params for disagg RnnModelConfig. + # d_inner is computed at the top of __init__ before TP division. + d_inner_global = d_inner # = mamba_head_dim * mamba_num_heads (GLOBAL) + conv_dim_global = d_inner_global + 2 * mamba_n_groups * mamba_d_state + self._rnn_d_state = mamba_d_state + self._rnn_d_conv = mamba_d_conv + self._rnn_num_heads = mamba_num_heads # GLOBAL + self._rnn_n_groups = mamba_n_groups + self._rnn_head_dim = mamba_head_dim + self._rnn_hidden_size = d_inner_global # GLOBAL = head_dim * num_heads + self._rnn_conv_dim_size = conv_dim_global # GLOBAL conv dim + self._rnn_num_layers = mamba_num_layers + + # Conv section layout for section-aware split/concat in TP mismatch. + # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; + # we just need to tell C++ which ordering to use. + self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) + + total_bytes = self.ssm_bytes + self.conv_bytes + if total_bytes % self.ssm_state_dtype.itemsize != 0: + raise RuntimeError( + f"Total state bytes ({total_bytes}) not divisible by " + f"ssm_state_dtype size ({self.ssm_state_dtype.itemsize})") + if total_bytes % self.conv_state_dtype.itemsize != 0: + raise RuntimeError( + f"Total state bytes ({total_bytes}) not divisible by " + f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") + if self.ssm_bytes % self.conv_state_dtype.itemsize != 0: + raise RuntimeError( + f"SSM state bytes ({self.ssm_bytes}) not divisible by " + f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") + self.linear_attention_metadata = LinearAttentionMetadata() + self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value + self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_config.periodic_snapshot_interval + if kv_cache_config.enable_block_reuse else 0) + # RNN model params for disagg TP-mismatch split/concat. + conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} + self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads + self.linear_attention_metadata.rnn_head_dim = self._rnn_head_dim + self.linear_attention_metadata.rnn_d_state = self._rnn_d_state + self.linear_attention_metadata.rnn_d_conv = self._rnn_d_conv + self.linear_attention_metadata.rnn_n_groups = self._rnn_n_groups + self.linear_attention_metadata.rnn_conv_section_layout = conv_section_map.get( + self._rnn_conv_section_layout, 0) + self.linear_attention_metadata.rnn_ssm_bytes = self.ssm_bytes + self.linear_attention_metadata.rnn_ssm_dtype_size = self.ssm_state_dtype.itemsize + self.linear_attention_metadata.rnn_conv_dtype_size = self.conv_state_dtype.itemsize + kv_cache_config = kv_cache_config.model_copy(deep=True) + if kv_cache_config.enable_partial_reuse: + logger.warning( + "Partial reuse is not supported for mamba hybrid models, disabling partial reuse" + ) + kv_cache_config.enable_partial_reuse = False + + # Keep the vector in physical global-layer order. Disabled entries are + # placeholders that are never projected into this manager. + kv_cache_config.max_attention_window = [ + LinearCacheType.RECURRENT_STATES.value + if mamba_layer_mask[layer_idx] else max_seq_len + for layer_idx in range(len(layer_mask)) + ] + + recurrent_states_window = LinearCacheType.RECURRENT_STATES.value + local_windows = { + recurrent_states_window + if mamba_layer_mask[layer_idx] else max_seq_len + for layer_idx in self.pp_layers + } + kwargs["pool_configurations"] = [ + PoolConfiguration( + window_size=window_size, + head_dim=head_dim, + dtype=torch_dtype_to_binding(self.ssm_state_dtype) + if window_size == recurrent_states_window else dtype, + ) for window_size in sorted(local_windows) + ] + + # Normalize num_kv_heads to a per-layer list and zero out mamba + # layer positions: those layers carry SSM/conv state instead of KV + # heads, so the parent KV cache should not allocate KV head storage + # for them. + if isinstance(num_kv_heads, int): + per_layer_kv_heads = [num_kv_heads] * total_layers + else: + if len(num_kv_heads) != total_layers: + raise ValueError( + f"num_kv_heads list length ({len(num_kv_heads)}) does not " + f"match total layers ({total_layers})") + per_layer_kv_heads = list(num_kv_heads) + for i, is_mamba in enumerate(mamba_layer_mask): + if is_mamba: + per_layer_kv_heads[i] = 0 + + self._setup_mtp_intermediate_states(spec_config, max_batch_size) + + # pass remaining arguments to super class + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, + linear_attention_metadata=self.linear_attention_metadata, + **kwargs, + ) + + assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" + self.mamba_layer_offsets = {} + for idx, layer_id in enumerate(self.mamba_pp_layers): + self.mamba_layer_offsets[layer_id] = idx + + self.host_block_offsets = torch.zeros([ + self.impl.num_pools, self.max_batch_size, 2, self.max_blocks_per_seq + ], + dtype=torch.int32, + device="cpu") + self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ + self.layer_offsets[self.mamba_pp_layers[0]]][0] + + self._state_indices_buffer = DeviceCopyBuffer( + torch.zeros([self.max_batch_size], dtype=torch.int32, + device="cuda"), + pin_memory=prefer_pinned()) + self._host_state_indices = torch.zeros([self.max_batch_size], + dtype=torch.int32, + pin_memory=prefer_pinned()) + self._row_indices = torch.arange(self.max_batch_size, + dtype=torch.long, + device="cpu") + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + self.kv_cache_config = kv_cache_config + self.is_estimating_kv_cache = is_estimating_kv_cache + + self._setup_states() + self._setup_replay_buffers(spec_config) + if use_replay_state_update and model_type == "qwen3_next": + logger.info_once( + "Configured GDN cached replay commit mode: small-batch fused, " + "large-batch all-layer", + key="gdn_cached_replay_commit_mode_fused", + ) + + @staticmethod + def get_cache_size_per_token( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs, + ): + """Affine memory model for the unified hybrid KV pool. + + Returns ``(slope_bytes_per_token, intercept_bytes)``: + + * ``slope`` = attention KV bytes per token plus conservative periodic + Mamba-state and partial-attention-snapshot costs. + * ``intercept`` = live and CUDA-graph dummy Mamba state. + + Memory budget -> max tokens then becomes + ``T = (budget - intercept) // slope`` instead of plain + ``T = budget // bytes_per_token``. + """ + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=1, + include_explicit_snapshots=False, + cap_partial_attention_snapshots=False, + **kwargs, + ) + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + return self._use_gdn_cached_replay_all_layer_commit + + def _commit_gdn_cached_replay_history_layers( + self, + attn_metadata: "AttentionMetadata", + num_decodes: int, + ) -> None: + """Synchronously advance every local GDN checkpoint in one launch.""" + from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers) + + if (not self._use_gdn_cached_replay_all_layer_commit + or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): + return + if (self.all_ssm_states is None or self.old_x is None + or self.old_B is None or self.old_dt is None + or self.replay_history_size is None): + raise RuntimeError( + "GDN cached replay all-layer commit requires replay state buffers." + ) + + 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 " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}.") + commit_gdn_cached_replay_history_layers( + ssm_states=self.all_ssm_states, + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + n_writes=mamba_metadata.replay_n_writes, + history_size=self.replay_history_size, + ) + + def shutdown(self): + # Release tensor views into the pool before the pool memory is freed, + # so their deleters don't see stale pointers. + self.all_ssm_states = None + self.all_conv_states = None + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask_buffer = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + super().shutdown() + + def add_dummy_requests( + self, + request_ids: List[int], + # Note that token_nums should be past_kv_len + input_len (without + # spec decoding). The draft tokens will be added in this function, + # so we don't need to take care of it in the caller. When preparing + # token_nums, we should not take the draft tokens into account, so + # don't use the kv_cache_manager.max_seq_len, which includes both + # extra tokens and draft tokens. + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + draft_kv_cache_manager: Optional[KVCacheManager] = None, + capture_sampling_params: Optional["SamplingParams"] = None, + ) -> List[LlmRequest]: + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + max_num_draft_tokens=max_num_draft_tokens, + kv_reserve_draft_tokens=kv_reserve_draft_tokens, + use_mrope=use_mrope, + max_beam_width=max_beam_width, + encoder_output_lens=encoder_output_lens, + draft_kv_cache_manager=draft_kv_cache_manager, + capture_sampling_params=capture_sampling_params, + ) + if requests: + self.requests.extend(requests) + # Only process the newly added requests, not all of self.requests. + # self.requests may contain stale entries from the previous + # _prepare_resources call (e.g., disagg transfer-pending requests) + # that would exceed max_batch_size and cause out-of-bounds access. + self._setup_state_indices(requests) + return requests + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None): + super().update_resources(scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + + @nvtx_range("hybrid_prepare_resources") + def _prepare_resources(self, scheduled_batch: ScheduledRequests): + self.requests = scheduled_batch.context_requests + \ + scheduled_batch.generation_requests + # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) + # until just before forward so the CPU-side prep (attn metadata, input + # tensor assembly, draft model prep, etc.) can overlap with the in-flight + # cudaMemcpyAsync calls instead of being serialized behind them. + # copy_linear_attention_block returns True iff it actually issued a copy; + # we skip refresh_blocks entirely when nothing was scheduled. + self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( + self.requests) + self._setup_state_indices() + self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + self._prepare_resources(scheduled_batch) + + @nvtx_range("hybrid_flush_state_transfers") + def flush_state_transfers(self) -> None: + """Complete any deferred state-block onboards scheduled by + prepare_resources(). Must be called before forward() reads recurrent + state blocks. Cheap no-op when nothing was scheduled. + """ + if getattr(self, "_pending_state_transfers", False): + self.impl.refresh_blocks() + self._pending_state_transfers = False + + @nvtx_range("hybrid_update_mamba_states") + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + + num_gens] + # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP + # passes accepted-token counts and lets the promotion kernel subtract + # one, avoiding a separate GPU subtraction/cast. + if accepted_leaf_positions is None: + accepted_position_source = accepted_tokens + position_source_is_token_count = True + else: + accepted_position_source = accepted_leaf_positions.to(torch.int32) + 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 + # 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() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + src_state_indices = self.intermediate_state_indices[:num_gens] + + # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter + # into the dtype-reinterpreted, strided C++ pool view. torch.compile + # can't handle the dtype-view mutation (and Inductor chokes on the + # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel + # writes through the view's real strides (~85% of HBM peak, one launch + # per state). + if self._use_replay_state_update: + # Every GDN layer has finished reading the old checkpoint and + # writing its candidate history. Advance all local checkpoints + # in one launch before PNAT and the active history buffer change. + self._commit_gdn_cached_replay_history_layers( + attn_metadata, num_gens) + assert self._dummy_request_mask is not None + else: + # Legacy: copy the accepted SSM state from the intermediate buffer. + _promote_mamba_state_triton( + self.all_ssm_states, + self.intermediate_ssm_states, + src_state_indices, + accepted_position_source, + state_indices_d, + position_source_is_token_count=position_source_is_token_count) + + # Conv: both paths save all intermediate conv windows, carry over the + # accepted one. The replay path also advances PNAT and the active cache + # buffer in this same launch. + replay_dummy_request_mask = None + if self._use_replay_state_update: + replay_dummy_request_mask = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + _promote_mamba_state_triton( + self.all_conv_states, + self.intermediate_conv_states, + src_state_indices, + accepted_position_source, + state_indices_d, + num_accepted_tokens=accepted_tokens, + position_source_is_token_count=position_source_is_token_count, + replay_pnat=(self.prev_num_accepted_tokens + if self._use_replay_state_update else None), + replay_cache_buf_idx=(self.cache_buf_idx + if self._use_replay_state_update else None), + dummy_request_mask=replay_dummy_request_mask, + replay_step_width=(self.replay_step_width + if self._use_replay_state_update else 0), + replay_history_size=(self.replay_history_size + if self._use_replay_state_update else 0)) + + def get_num_available_tokens(self, + token_num_upper_bound: int, + max_num_draft_tokens: int = 0, + **kwargs) -> int: + # Base bound from attention KV cache pool (the parent's behaviour). + result = super().get_num_available_tokens(token_num_upper_bound, + max_num_draft_tokens, + **kwargs) + # Also bound by the recurrent-state pool capacity: each request needs + # roughly ceil(N / states_snapshot_interval) recurrent-state blocks + # (+1 for the corner case where N is a multiple of tokens_per_block). + # When block reuse is disabled, only one snapshot is needed per + # request, so no additional capping is required here. + interval = (self.linear_attention_metadata.states_snapshot_interval + if self.linear_attention_metadata is not None else 0) + # Attention-only PP ranks keep the interval so every rank publishes + # identical scheduling boundaries, but they have no recurrent-state + # pool whose capacity should constrain their attention KV cache. + if self.local_num_mamba_layers > 0 and interval and interval > 0: + stats = self.impl.get_kv_cache_stats() + rs_free = stats.num_free_blocks_per_window_size.get( + LinearCacheType.RECURRENT_STATES.value, 0) + # Reserve 1 block for the always-allocated last block (corner case + # / final live state) so we don't promise more tokens than the + # pool can actually back at allocation time. + usable_rs_blocks = max(0, rs_free - 1) + rs_token_cap = usable_rs_blocks * interval + result = min(result, rs_token_cap) + return max(result, 0) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + if request in self.requests: + self.requests.remove(request) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def _setup_state_indices(self, requests=None) -> None: + if self.local_num_mamba_layers == 0: + return + if requests is None: + requests = self.requests + block_indices = [] + for req in requests: + if req.is_context_finished: + next_step = self.get_num_tokens(req) - 1 + elif self.kv_cache_config.enable_block_reuse: + next_step = (req.context_current_position - 1 + + req.context_chunk_size) + else: + next_step = req.prompt_len - 1 + block_indices.append(next_step // self.tokens_per_block) + self.impl.copy_batch_block_offsets( + self.host_block_offsets, [req.py_request_id for req in requests], 1, + 0) + max_blocks = self.blocks_per_window[ + LinearCacheType.RECURRENT_STATES.value][0] + n = len(requests) + self._host_state_indices.zero_() + if n > 0: + # Vectorized gather: replace per-element Python loop with a single + # tensor index op. block_indices is a small Python list so + # torch.tensor() conversion here is O(n) at C level, much cheaper + # than n round-trips through Python indexing. + bi = torch.tensor(block_indices, dtype=torch.long) + rows = self._row_indices[:n] + # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] + values = self.host_block_offsets[self.recurrent_states_pool_index, + rows, 0, bi] + invalid_mask = (values < 0) | (values >= max_blocks) + # The C++ recurrent-state manager uses null page-table entries for + # logical blocks that are not snapshot boundaries. Usually a + # context chunk ends exactly at an allocated snapshot (or at the + # final live-state block), but the scheduler may shorten a chunk + # further when KV capacity is tight. In that case the Mamba + # kernel must keep accumulating into the next allocated snapshot + # or final block, matching copyLinearAttentionBlock(), which also + # skips placeholders when it advances the live state. + for bad_i in invalid_mask.nonzero( + as_tuple=False).flatten().tolist(): + req = requests[bad_i] + if req.is_context_finished: + continue + last_prompt_block = (req.prompt_len - + 1) // self.tokens_per_block + row = self.host_block_offsets[self.recurrent_states_pool_index, + rows[bad_i], 0] + candidates = row[block_indices[bad_i]:last_prompt_block + 1] + valid_candidates = ((candidates >= 0) + & (candidates < max_blocks)).nonzero( + as_tuple=False) + if valid_candidates.numel() > 0: + values[bad_i] = candidates[valid_candidates[0, 0]] + + invalid_mask = (values < 0) | (values >= max_blocks) + if invalid_mask.any(): + bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) + req = requests[bad_i] + value = int(values[bad_i]) + raise RuntimeError( + f"Invalid recurrent state block index {value} " + f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " + f"prompt_len={req.prompt_len}, " + f"is_context_finished={req.is_context_finished}, " + f"context_current_position={req.context_current_position}, " + f"prepopulated_token_num={req.prepopulated_prompt_len}, " + f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " + f"block_index for next step is {block_indices[bad_i]}, " + "\nblock_ids=" + f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" + ) + self._host_state_indices[:n] = values + + self._state_indices_buffer.update(self._host_state_indices, force=True) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + + # Build request_id → pool block offset mapping so that + # get_state_indices can return indices in arbitrary request order. + # Bulk tolist avoids a per-request tensor-index + .item() round-trip. + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None) -> list: + if self.local_num_mamba_layers == 0: + # Mamba metadata is prepared on every PP rank even when this rank + # owns only attention layers. No local kernel consumes these + # indices, so avoid consulting state that is intentionally absent. + return [0] * len(request_ids) if request_ids is not None else [] + if request_ids is not None: + # Return indices in the order of the caller's request_ids, + # not the internal self.requests order. This is critical when + # the batch is reordered after prepare_resources (e.g. disagg + # serving sorts generation_requests by py_batch_idx). + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + def _setup_states(self) -> None: + # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) + pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( + torch.uint8).reshape(self.local_num_mamba_layers, -1, + self.ssm_bytes + self.conv_bytes) + num_blocks_in_pool = pool.shape[1] + self.all_ssm_states = pool[:, :, :self.ssm_bytes].view( + self.ssm_state_dtype).view( + [self.local_num_mamba_layers, num_blocks_in_pool] + + self.ssm_state_shape) + self.all_conv_states = pool[:, :, self.ssm_bytes:self.ssm_bytes + + self.conv_bytes].view( + self.conv_state_dtype).view([ + self.local_num_mamba_layers, + num_blocks_in_pool + ] + self.conv_state_shape) + self.all_ssm_states.zero_() + self.all_conv_states.zero_() + + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = self.all_ssm_states.shape[1] + device = self.all_ssm_states.device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return + + self._dummy_request_mask_buffer = DeviceCopyBuffer( + torch.zeros(self.max_batch_size, dtype=torch.bool, device=device), + pin_memory=prefer_pinned()) + self._dummy_request_mask_host = torch.zeros( + self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) + + +class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): + """Hybrid Mamba cache manager backed by KVCacheManagerV2. + + Attention KV pages and Mamba recurrent-state pages are both owned by the + Python V2 cache manager. Mamba layers are represented as V2 SSM layers, + while this wrapper exposes the state tensors and slot indices expected by + the PyTorch Mamba kernels. + """ + + _supports_additional_snapshot_offsets = True + # Recurrent-state snapshots use a specialized commit/history protocol, so + # keep main-like reuse endpoints and the existing unpaired draft path. + _supports_reuse_match_backoff = False + + # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the + # cache manager __init__ coerces a positive spec recompute tail to a full + # re-prefill (see _has_recurrent_state on the attention-only base classes). + _has_recurrent_state = True + + # Qwen4-Exp PLE state is opt-in. These class-level defaults keep every other + # model — and any partially-constructed instance that sets only the fields it + # needs — out of the PLE branches without carrying PLE geometry. + # `_init_qwen4_exp_ple_geometry` rebinds all of them per instance. + _ple_layer_ids: list[int] = [] + _ple_conv_state_shape: list[int] = [] + _ple_ngram_context_shape: list[int] = [] + # Set per instance on a disaggregated context server of a KDA model. + kda_context_only = False + + def _disagg_transfer_overwrites_whole_cached_prefix(self) -> bool: + # The incoming recurrent state overwrites the whole local slot and summarizes the + # complete prefix, so none of those tokens count as a local cache hit. + return (self.local_num_mamba_layers > 0 + or super()._disagg_transfer_overwrites_whole_cached_prefix()) + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[List[bool]] = None, + is_estimating_kv_cache: bool = False, + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", + kda_replay_num_spec: Optional[int] = None, + qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, + kda_draft_records: bool = False, + **kwargs, + ) -> None: + if conv_state_layout not in ("x_b_c", "q_k_v"): + raise ValueError( + f"Unsupported convolution state layout: {conv_state_layout!r}") + if "model_type" in kwargs: + # The V1 managers select the conv-state layout by model_type; this + # class takes it explicitly. Silently absorbing model_type here + # means a caller's layout request would be dropped on the floor. + raise TypeError( + "MambaHybridCacheManagerV2 does not accept 'model_type' " + f"(got {kwargs['model_type']!r}); pass " + "conv_state_layout='x_b_c' or 'q_k_v' instead") + + # A disaggregated context server never verifies drafts, so it keeps + # neither KDA replay caches nor intermediate verification buffers. + self.kda_context_only = bool( + kda_replay_num_spec is not None + and _is_disagg_context_server(kwargs.get("is_disagg", False))) + if self.kda_context_only: + kda_replay_num_spec = None + self._kda_replay_layers: List[Dict[str, torch.Tensor]] = [] + self._kda_replay_num_spec = kda_replay_num_spec + self._use_kda_replay_update = kda_replay_num_spec is not None + # Only with the KDA replay caches: also keep the records of every + # draft of the last verify round (kda_state_tok). + self._kda_draft_records = kda_draft_records and self._use_kda_replay_update + if self._use_kda_replay_update: + if use_replay_state_update: + raise ValueError( + "kda_replay_num_spec and use_replay_state_update are " + "mutually exclusive") + if spec_config is None: + raise ValueError( + "KDA replay caches require speculative decoding") + expected_num_spec = spec_config.tokens_per_gen_step - 1 + if kda_replay_num_spec != expected_num_spec: + raise ValueError( + "KDA replay cache width " + f"({kda_replay_num_spec}) must match the draft length " + f"({expected_num_spec})") + if conv_state_layout != "q_k_v": + raise ValueError( + "KDA replay requires conv_state_layout='q_k_v'") + self._record_kda_replay_in_update_resources = ( + self.use_kda_replay_update + and getattr(spec_config, "decoding_type", None) == "NGram") + + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + + combined_layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + + self._mamba_layer_mask = list(mamba_layer_mask) + self._use_replay_state_update = use_replay_state_update + self._use_gdn_cached_replay_all_layer_commit = (use_replay_state_update + and conv_state_layout + == "q_k_v") + self._gdn_cached_replay_state_descriptors: Optional[torch.Tensor] = None + self._gdn_cached_replay_state_strides: Optional[Tuple[int, int, + int]] = None + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + self._seed_request_counter = 0 + self._recurrent_evicted_blocks_total = 0 + self._recurrent_onboarded_blocks_total = 0 + self._recurrent_dropped_blocks_total = 0 + self._recurrent_status_logged = False + # Branch points keyed by request id. prepare_expect_snapshot_points() + # overwrites the per-request list on every scheduler pass, so the point + # has to be re-merged from here rather than stored only on the request. + self._branch_snapshot_points: Dict[int, int] = {} + self._snapshot_pruned_tokens_total = 0 + self._page_pruned_tokens_total = 0 + self._branch_snapshots_taken_total = 0 + self._branch_snapshots_skipped_total: Dict[str, int] = {} + num_cuda_graph_padding_dummy_slots = ( + _get_num_cuda_graph_padding_dummy_slots(spec_config, + max_batch_size)) + self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + + int(mapping.enable_attention_dp)) + # PLE uses the same request slots and snapshots as the GDN state. + self._init_qwen4_exp_ple_geometry(qwen4_exp_ple_cache_params, + total_layers, mamba_cache_dtype) + self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype + is not None else mamba_cache_dtype) + self.conv_state_dtype = mamba_cache_dtype + + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + self._use_gdn_cached_replay_all_layer_commit = ( + self._use_gdn_cached_replay_all_layer_commit + and self.local_num_mamba_layers > 0) + + if self.local_num_mamba_layers > 0: + tp_size = mamba_effective_tp_size(mapping) + d_inner = mamba_head_dim * mamba_num_heads + grouped_state_dim = mamba_n_groups * mamba_d_state + conv_dim = d_inner + 2 * grouped_state_dim + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if kwargs.get("is_disagg", + False) and grouped_state_dim % tp_size != 0: + raise ValueError( + "Disaggregated Mamba transfer requires each convolution " + "state section to be divisible by tp_size") + if use_replay_state_update: + assert mamba_n_groups % tp_size == 0, \ + "replay state update requires mamba_n_groups divisible by tp_size" + self._n_groups_per_rank = mamba_n_groups // tp_size + d_inner_local = d_inner // tp_size + grouped_state_dim_local = grouped_state_dim // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + # TP-mismatch disaggregated transfers must split the flat + # convolution state at its true semantic boundaries. Mamba2 stores + # [x | B | C], while GDN stores [Q | K | V]. The large section is + # therefore first for Mamba2 and last for GDN. + if conv_state_layout == "x_b_c": + self.conv_section_dims = [ + d_inner_local, + grouped_state_dim_local, + grouped_state_dim_local, + ] + else: + self.conv_section_dims = [ + grouped_state_dim_local, + grouped_state_dim_local, + d_inner_local, + ] + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) + else: + logger.info( + "No local mamba layers for this rank, skipping mamba state views" + ) + self._n_groups_per_rank = 0 + self.conv_state_shape = [] + self.ssm_state_shape = [] + self.conv_section_dims = [] + self.ssm_bytes = 0 + self.conv_bytes = 0 + + if isinstance(num_kv_heads, int): + per_layer_kv_heads = [num_kv_heads] * total_layers + else: + if len(num_kv_heads) != total_layers: + raise ValueError( + f"num_kv_heads list length ({len(num_kv_heads)}) does not " + f"match total layers ({total_layers})") + per_layer_kv_heads = list(num_kv_heads) + for i, is_mamba in enumerate(mamba_layer_mask): + if is_mamba: + per_layer_kv_heads[i] = 0 + + self._setup_mtp_intermediate_states(spec_config, max_batch_size) + + kv_cache_config = kv_cache_config.model_copy(deep=True) + if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: + block_reuse_config = kv_cache_config.block_reuse_config + block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) + if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_config = block_reuse_config.model_copy( + update={"policy": BlockReusePolicy.PER_REQUEST.value}) + self.kv_cache_config = kv_cache_config + + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=combined_layer_mask, + is_draft=is_draft, + is_estimating_kv_cache=is_estimating_kv_cache, + num_reserved_index_slots=self._num_reserved_dummy_slots, + **kwargs, + ) + + self.mamba_layer_offsets = { + layer_id: offset + for offset, layer_id in enumerate(self.mamba_pp_layers) + } + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + # Requests whose suspend committed their KDA replay state into the V2 + # states; their next generation step re-seeds the replay caches. + self._kda_replay_suspended: set[int] = set() + # With a host drafter: requests of a prepared batch whose acceptance + # update_resources has not recorded yet. + self._kda_replay_acceptance_pending: set[int] = set() + self._kda_replay_unseeded: set[int] = set() + + state_index_capacity = (self.max_batch_size + + self._num_reserved_dummy_slots) + self._state_indices_buffer = DeviceCopyBuffer( + torch.zeros([state_index_capacity], + dtype=torch.int32, + device="cuda"), + pin_memory=prefer_pinned()) + self._host_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + pin_memory=prefer_pinned()) + + if self.local_num_mamba_layers > 0: + first_mamba_local_layer = self.layer_offsets[ + self.mamba_pp_layers[0]] + self.ssm_layer_group_id = self.impl.get_layer_group_id( + LayerId(first_mamba_local_layer)) + self._ssm_page_index_scale = self.impl.get_page_index_scale( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_pages = self.impl.get_page_index_upper_bound( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // + self._ssm_page_index_scale) + required_live_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + if num_ssm_slots < required_live_slots: + KVCacheManagerV2.shutdown(self) + raise ValueError( + "The V2 Mamba state pool has only " + f"{num_ssm_slots} slots but needs at least " + f"{required_live_slots} live/dummy slots. Increase the " + "KV cache budget or allocate a larger Mamba pool_ratio.") + self._setup_states() + # PLE pools use exactly the state-slot capacity resolved above. + self._setup_ple_states(num_ssm_slots) + self._setup_replay_buffers(spec_config) + if self._use_gdn_cached_replay_all_layer_commit: + state_layout = ("affine" + if self._gdn_cached_replay_state_strides + is not None else "indirect") + logger.info_once( + "Configured GDN cached replay commit mode for V2: " + "small-batch fused, large-batch all-layer; " + f"state layout: {state_layout}", + key="gdn_cached_replay_v2_commit_mode_fused", + ) + else: + self.ssm_layer_group_id = None + self._ssm_page_index_scale = 1 + self.all_ssm_states = [] + self.all_conv_states = [] + self._stacked_ssm_states = None + self._stacked_conv_states = None + self._setup_replay_buffers(spec_config) + + def _init_qwen4_exp_ple_geometry( + self, + params: Optional["Qwen4ExpPLECacheParams"], + total_layers: int, + default_state_dtype: torch.dtype, + ) -> None: + """Validate and record the Qwen4-Exp PLE pool geometry. + + A no-op for every other model: the fields stay empty and + `_ple_buffer_configs` then contributes nothing. + """ + self._ple_layer_ids: list[int] = [] + self._ple_conv_state_shape: list[int] = [] + self._ple_ngram_context_shape: list[int] = [] + self._ple_conv_state_dtype = default_state_dtype + self._ple_conv_states: dict[int, torch.Tensor] = {} + self._ple_ngram_contexts: dict[int, torch.Tensor] = {} + if params is None: + return + if len(params.ple_layer_mask) != total_layers: + raise ValueError( + "PLE layer mask length must match the hybrid layer mask: " + f"got {len(params.ple_layer_mask)}, expected {total_layers}") + self._ple_layer_ids = [ + layer_id for layer_id, active in enumerate(params.ple_layer_mask) + if active + ] + if len(self._ple_layer_ids) != params.num_ple_layers: + raise ValueError( + "PLE layer mask count does not match num_ple_layers: " + f"got {len(self._ple_layer_ids)}, expected " + f"{params.num_ple_layers}") + if any(not self._mamba_layer_mask[layer_id] + for layer_id in self._ple_layer_ids): + raise ValueError("PLE lifecycle state must belong to Mamba layers") + if (params.short_conv_channels <= 0 or params.short_conv_state_len <= 0 + or params.ngram_context_len <= 0): + raise ValueError("PLE recurrent-state dimensions must be positive") + self._ple_conv_state_shape = [ + params.short_conv_channels, + params.short_conv_state_len, + ] + self._ple_ngram_context_shape = [params.ngram_context_len] + self._ple_conv_state_dtype = params.conv_state_dtype + + def _setup_ple_states(self, num_state_slots: int) -> None: + """Bind each local PLE layer to V2-managed recurrent-state views.""" + if not self._ple_layer_ids: + return + for layer_id in self._ple_layer_ids: + local_layer_idx = self.layer_offsets.get(layer_id) + if local_layer_idx is None: + continue + conv_state = self._get_state_buffer( + local_layer_idx, + MambaRole.PLE_CONV_STATE, + self._ple_conv_state_dtype, + self._ple_conv_state_shape, + ) + ngram_context = self._get_state_buffer( + local_layer_idx, + MambaRole.PLE_NGRAM_CONTEXT, + dtype=torch.long, + state_shape=self._ple_ngram_context_shape, + ) + if (conv_state.shape[0] != num_state_slots + or ngram_context.shape[0] != num_state_slots): + raise RuntimeError( + "PLE and GDN lifecycle buffers must have the same number " + f"of slots: layer={layer_id}, GDN={num_state_slots}, " + f"conv={conv_state.shape[0]}, ngram={ngram_context.shape[0]}" + ) + self._ple_conv_states[layer_id] = conv_state + self._ple_ngram_contexts[layer_id] = ngram_context + logger.info( + "PLE state views bound to V2 lifecycle buffers for local layers " + f"{sorted(self._ple_conv_states)}: conv " + f"[{num_state_slots}, {', '.join(map(str, self._ple_conv_state_shape))}] " + f"({self._ple_conv_state_dtype}), n-gram context " + f"[{num_state_slots}, {self._ple_ngram_context_shape[0]}] (int64)") + + def ple_layer_cache( + self, + layer_idx: int) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Return persistent ``(conv_state, ngram_context)`` pools for a layer.""" + conv = self._ple_conv_states.get(layer_idx) + ngram = self._ple_ngram_contexts.get(layer_idx) + if conv is None or ngram is None: + return None + return conv, ngram + + def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: + """Recurrent roles and how they reshard across TP. + + Convolution state is SECTIONED: its flat per-layer buffer is a + concatenation of sections (Mamba2 ``[x | B | C]``, GDN ``[Q | K | V]``) + that are each TP-sharded independently. SSM state keeps the INDEXED + default (sharded by head). PLE state is computed from replicated + inputs, so every rank holds identical bytes and the transfer copies + whole per-layer regions. + """ + return { + **super().get_disagg_role_mapper_kinds(), + MambaRole.CONV_STATE: + MapperKind.SECTIONED, + MambaRole.PLE_CONV_STATE: + MapperKind.REPLICATED, + MambaRole.PLE_NGRAM_CONTEXT: + MapperKind.REPLICATED, + } + + def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: + """Per-layer geometry for resharding conv and SSM state.""" + if self.local_num_mamba_layers == 0: + return {} + d_conv_m1 = int(self.conv_state_shape[1]) + conv_elem_size = self.conv_state_dtype.itemsize + _, head_dim, d_state = self.ssm_state_shape + return { + MambaRole.CONV_STATE: + RoleLayout(section_bytes=tuple( + int(dim) * d_conv_m1 * conv_elem_size + for dim in self.conv_section_dims)), + MambaRole.SSM_STATE: + RoleLayout(bytes_per_head=int(head_dim) * int(d_state) * + self.ssm_state_dtype.itemsize), + } + + def get_disagg_ignored_roles(self) -> frozenset[DataRole]: + """Keep the KDA replay buffers out of disaggregated transfers. + + They are generation-local: a context server allocates none, and a + generation server seeds them from the transferred conv state. + """ + replay_roles = {DataRole(name) for name in self._kda_replay_layout()} + if replay_roles: + replay_roles.add(_KDA_REPLAY_COUNT_ROLE) + return super().get_disagg_ignored_roles() | frozenset(replay_roles) + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + + @property + def kda_state_tok(self) -> Optional[List[torch.Tensor]]: + """Per local KDA layer, the drafts' records of every SSM slot + ([slots, 3, num_spec, H, K] fp32 KVCM state-buffer views), or None + when the replay path keeps the per-draft replay caches instead.""" + replay_layers = getattr(self, "_kda_replay_layers", None) + if not replay_layers or _KDA_DRAFT_RECORDS_ROLE not in replay_layers[0]: + return None + return [layer[_KDA_DRAFT_RECORDS_ROLE] for layer in replay_layers] + + @property + def keeps_kda_draft_records(self) -> bool: + return self.kda_state_tok is not None + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size: int) -> None: + if self.kda_context_only: + super()._setup_mtp_intermediate_states(None, max_batch_size) + return + if not self.use_kda_replay_update: + super()._setup_mtp_intermediate_states(spec_config, max_batch_size) + return + + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if self.local_num_mamba_layers == 0: + return + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _kda_replay_layout(self) -> Dict[str, Tuple[torch.dtype, List[int]]]: + """Per-slot shapes of the KDA replay buffers of one local layer.""" + if not self.use_kda_replay_update or self.local_num_mamba_layers == 0: + return {} + if (len(self.conv_section_dims) != 3 + or len(set(self.conv_section_dims)) != 1): + raise ValueError("KDA replay caches require equal [Q | K | V] " + "convolution-state sections") + assert self._kda_replay_num_spec is not None + draft_records = getattr(self, "_kda_draft_records", False) + layout = _kda_replay_layout(self._kda_replay_num_spec, + self.ssm_state_shape[0], + self.ssm_state_shape[1], + self.conv_state_shape[1], + draft_records=draft_records) + if draft_records and tuple(layout[_KDA_DRAFT_RECORDS_ROLE][1]) != ( + self._kda_draft_records_shape()): + raise ValueError("KDA draft records require equal value and key " + "head dimensions") + return layout + + def _kda_replay_buffer_configs(self) -> List[BufferConfig]: + return [ + BufferConfig(role=DataRole(name), + size=dtype.itemsize * math.prod(shape)) + for name, (dtype, shape) in self._kda_replay_layout().items() + ] + + def _validate_kda_replay_colocation(self, + replay_roles: Iterable[str]) -> None: + """Verify replay and recurrent buffers share one slot lifecycle. + + Every layer addresses its replay caches with the recurrent state's + logical slot index, so all of them must be allocated, moved and + offloaded together: they must sit in one variant of one KVCM pool + group, not merely have equal capacities. + """ + required_buffers = {(self.layer_offsets[layer_id], str(role)) + for layer_id in self.mamba_pp_layers + for role in (MambaRole.SSM_STATE, + MambaRole.CONV_STATE, *replay_roles)} + required_buffers.add((self.layer_offsets[self.mamba_pp_layers[0]], + str(_KDA_REPLAY_COUNT_ROLE))) + if not any( + required_buffers.issubset( + {(int(buffer_id.layer_id), str(buffer_id.role)) + for coalesced in variant.coalesced_buffers + for buffer_id in coalesced.buffer_ids}) + for pool_group in self.impl.pool_group_descs + for variant in pool_group.slot_desc.variants): + raise RuntimeError( + "KDA replay and recurrent-state buffers must share one KVCM " + "pool group and lifecycle to reuse the same logical slot index") + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + allocated = super()._allocate_pool_replay_buffers( + spec_config, cache_size, device) + self._kda_replay_layers = [] + layout = self._kda_replay_layout() + if not layout: + return allocated + self._validate_kda_replay_colocation(layout) + # The replay caches are KVCM state buffers next to the SSM and conv + # states, so they follow the slot through suspend, offload and + # snapshot reuse. Pool capacity can exceed resident memory, so rows + # are initialized only once a request holds its slot. + for layer_id in self.mamba_pp_layers: + local_layer_idx = self.layer_offsets[layer_id] + layer = {} + for name, (dtype, shape) in layout.items(): + view = self._get_state_buffer(local_layer_idx, DataRole(name), + dtype, shape) + if name in _KDA_REPLAY_CONV_ROLES: + # Dim-contiguous [slots, dim, W - 1 + num_spec]. + view = view.transpose(-1, -2) + elif name == "kda_beta_cache": + view = view[..., :self.ssm_state_shape[0]] + if view.shape[0] != cache_size: + raise RuntimeError( + f"KDA replay buffer {name} of layer {layer_id} has " + f"{view.shape[0]} slots, but the recurrent state has " + f"{cache_size}") + layer[name] = view + self._kda_replay_layers.append(layer) + # All SSM layers share a slot lifecycle, so the accepted-draft count is + # stored once, on the first local KDA layer. + self.prev_num_accepted_tokens = self._get_state_buffer( + self.layer_offsets[self.mamba_pp_layers[0]], + _KDA_REPLAY_COUNT_ROLE, + torch.int32, + [_KDA_REPLAY_COUNT_BYTES // torch.int32.itemsize], + )[..., 0] + records = (", with the drafts' records in place of the replay " + "caches" if self.keeps_kda_draft_records else "") + logger.info("Mamba Cache (kda-replay) is bound to KVCM state buffers " + f"for {cache_size} state slots{records}") + return True + + def _commit_gdn_cached_replay_history_layers( + self, + attn_metadata: "AttentionMetadata", + num_decodes: int, + ) -> None: + """Advance all V2 GDN checkpoints through their state views.""" + from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers) + + if (not getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): + return + if (not self.all_ssm_states + or (self._gdn_cached_replay_state_descriptors is None + and self._gdn_cached_replay_state_strides is None) + or self.old_x is None or self.old_B is None + or self.old_dt is None or self.replay_history_size is None): + raise RuntimeError( + "GDN cached replay all-layer commit requires V2 replay state buffers." + ) + + 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 " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}.") + state_strides = self._gdn_cached_replay_state_strides + commit_gdn_cached_replay_history_layers( + ssm_states=self.all_ssm_states[0], + ssm_state_descriptors=(self._gdn_cached_replay_state_descriptors), + ssm_state_num_layers=(state_strides[0] + if state_strides is not None else None), + ssm_state_layer_stride=(state_strides[1] + if state_strides is not None else None), + ssm_state_slot_stride=(state_strides[2] + if state_strides is not None else None), + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + n_writes=mamba_metadata.replay_n_writes, + history_size=self.replay_history_size, + ) + + def _create_kv_cache( + self, + request_id: int, + lora_task_id: Optional[int], + input_tokens: Optional[Sequence[TokenIdExt]], + *, + cache_salt: Optional[str] = None, + is_dummy: bool = False, + enable_request_stats: bool = False, + expected_prompt_length: Optional[int] = None, + kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None, + ) -> Optional[_KVCache]: + kv_cache = super()._create_kv_cache( + request_id, + lora_task_id, + input_tokens, + cache_salt=cache_salt, + is_dummy=is_dummy, + enable_request_stats=enable_request_stats, + expected_prompt_length=expected_prompt_length, + kv_cache_retention_config=kv_cache_retention_config, + ) + if (self.mapping.rank == 0 and kv_cache is not None + and input_tokens is not None and not is_dummy + and self.local_num_mamba_layers > 0): + # Prefix lookup excludes the final prompt token because it must be + # recomputed by prefill. Restore it for the request-length metric. + request_total_tokens = len(input_tokens) + 1 + prefix_reuse_diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) + logger.debug( + f"[MambaHybridCacheManagerV2] prefix reuse rank={self.mapping.rank} " + f"request_id={request_id} " + f"request_total_tokens={request_total_tokens} " + f"longest_attention_match_tokens=" + f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_hybrid_pruning()} " + f"content_divergence_tokens=" + f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_pruning()} " + f"latest_recurrent_snapshot_tokens=" + f"{kv_cache.num_committed_tokens}") + return kv_cache + + def _skip_branch_snapshot(self, reason: str) -> None: + self._branch_snapshots_skipped_total[reason] = ( + self._branch_snapshots_skipped_total.get(reason, 0) + 1) + + def _format_branch_snapshot_counters(self) -> str: + """Attribute reuse loss between missing snapshots and evicted pages. + + A flat cache hit rate means something different depending on whether + branch points were never found or were found and not reused, so the + skip reasons are reported alongside the totals. + """ + skipped = self._branch_snapshots_skipped_total + return (f"snapshot_pruned_tokens={self._snapshot_pruned_tokens_total} " + f"page_pruned_tokens={self._page_pruned_tokens_total} " + f"branch_snapshots_taken={self._branch_snapshots_taken_total} " + f"branch_snapshots_skipped={sum(skipped.values())} " + f"branch_snapshots_skipped_by_reason=" + f"{dict(sorted(skipped.items()))}") + + def _record_branch_snapshot_point(self, req: LlmRequest, kv_cache: _KVCache, + num_lookup_tokens: Optional[int]) -> None: + """Snapshot where this request leaves the reuse tree, for its siblings.""" + if not self.kv_cache_config.block_reuse_config.enable_branch_snapshot: + return + if req.is_dummy_request: + return + if num_lookup_tokens is None: + self._skip_branch_snapshot("no_fresh_match") + # A fresh lookup can record a point before resume() fails. On the + # next attempt, reapply that point against the cache's authoritative + # reuse depth: the request cursor is not rewound until preparation + # succeeds and may still describe a deeper, abandoned attempt. + if req.py_request_id in self._branch_snapshot_points: + self._apply_branch_snapshot_point( + req, context_current_position=kv_cache.num_committed_tokens) + return + + diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) + divergence = diagnostics._get_num_reusable_tokens_before_pruning() + reused = kv_cache.num_committed_tokens + if self.local_num_mamba_layers > 0: + hybrid_depth = ( + diagnostics._get_num_reusable_tokens_before_hybrid_pruning()) + # Loss attributable to a missing recurrent snapshot, versus loss + # attributable to attention pages having been evicted. These call + # for different fixes, so they are counted apart. + self._snapshot_pruned_tokens_total += max(0, hybrid_depth - reused) + self._page_pruned_tokens_total += max(0, divergence - hybrid_depth) + + # The whole lookup range matched, so there is no fork here. + if divergence >= num_lookup_tokens: + self._skip_branch_snapshot("no_divergence") + return + # Align down: tokens past a block boundary come from a partial-match + # child rather than a confirmed shared prefix, and an unaligned commit + # takes the partial-block snapshot path. + point = (divergence // self.tokens_per_block) * self.tokens_per_block + if point == 0: + self._skip_branch_snapshot("aligned_to_zero") + return + # Reuse already reached the fork, so a snapshot there exists. + if point <= reused: + self._skip_branch_snapshot("already_reused") + return + # A branch point at or beyond the prompt end cannot improve sibling + # reuse. + if point >= req.prompt_len: + self._skip_branch_snapshot("at_prompt_end") + return + + self._branch_snapshot_points[req.py_request_id] = point + self._apply_branch_snapshot_point(req, context_current_position=reused) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + super().prepare_expect_snapshot_points(requests) + if (not self.enable_block_reuse or not self.kv_cache_config. + block_reuse_config.enable_branch_snapshot): + return + for request in requests: + self._apply_branch_snapshot_point(request) + + def _apply_branch_snapshot_point( + self, + request: LlmRequest, + context_current_position: Optional[int] = None, + ) -> None: + points = set(request.expect_snapshot_points) + # Keep the flag usable without any configured snapshot placement. If a + # placement is configured, preserve it: adding prompt_len can replace a + # nearby recurrent checkpoint in the same radix-tree block because a + # block stores one recurrent page per life cycle. + if not points: + points.add(request.prompt_len) + point = self._branch_snapshot_points.get(request.py_request_id) + current_position = (request.context_current_position + if context_current_position is None else + context_current_position) + if point is not None and point > current_position: + points.add(point) + request.expect_snapshot_points = sorted(points) + + def get_iteration_stats(self) -> Optional[KVCacheV2IterationStatsReport]: + """Log recurrent-cache movement; this is KDA state for Kimi K3.""" + report = super().get_iteration_stats() + if report is None: + return None + + pool_group_ids = sorted({ + pool_group_id + for pool_group_id, _, kind in + self._stats_life_cycle_metadata().values() if kind == "ssm" + }) + if not pool_group_ids: + return report + pool_group_reports = [ + report.by_pool_group[pool_group_id] + for pool_group_id in pool_group_ids + ] + + stats = [ + pool_group_report.stats for pool_group_report in pool_group_reports + ] + evicted_blocks = sum(stat.iter_offload_blocks for stat in stats) + evicted_bytes = sum(stat.iter_offload_bytes for stat in stats) + onboarded_blocks = sum(stat.iter_onboard_blocks for stat in stats) + onboarded_bytes = sum(stat.iter_onboard_bytes for stat in stats) + dropped_blocks = sum(stat.iter_host_dropped_blocks for stat in stats) + dropped_bytes = sum(stat.iter_host_dropped_bytes for stat in stats) + has_movement = bool(evicted_blocks or onboarded_blocks + or dropped_blocks) + if has_movement: + self._recurrent_evicted_blocks_total += evicted_blocks + self._recurrent_onboarded_blocks_total += onboarded_blocks + self._recurrent_dropped_blocks_total += dropped_blocks + if (self.mapping.rank == 0 + and (has_movement or not self._recurrent_status_logged)): + logger.debug( + f"[MambaHybridCacheManagerV2] recurrent cache status " + f"rank={self.mapping.rank} pool_group_ids={pool_group_ids} " + f"evicted_recurrent_blocks={evicted_blocks} " + f"evicted_recurrent_bytes={evicted_bytes} " + f"onboarded_recurrent_blocks={onboarded_blocks} " + f"onboarded_recurrent_bytes={onboarded_bytes} " + f"dropped_recurrent_blocks={dropped_blocks} " + f"dropped_recurrent_bytes={dropped_bytes} " + f"total_evicted_recurrent_blocks=" + f"{self._recurrent_evicted_blocks_total} " + f"total_onboarded_recurrent_blocks=" + f"{self._recurrent_onboarded_blocks_total} " + f"total_dropped_recurrent_blocks=" + f"{self._recurrent_dropped_blocks_total} " + f"gpu_used_recurrent_blocks=" + f"{sum(stat.primary_used_num_blocks for stat in stats)} " + f"gpu_free_recurrent_blocks=" + f"{sum(stat.primary_free_num_blocks for stat in stats)} " + f"gpu_evictable_recurrent_blocks=" + f"{sum(stat.primary_evictable_num_blocks for stat in stats)} " + f"host_used_recurrent_blocks=" + f"{sum(stat.secondary_used_num_blocks for stat in stats)} " + f"host_free_recurrent_blocks=" + f"{sum(stat.secondary_free_num_blocks for stat in stats)} " + f"{self._format_branch_snapshot_counters()}") + self._recurrent_status_logged = True + return report + + @staticmethod + def get_cache_size_per_token(model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs): + spec_config = kwargs.get("spec_config") + num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size) + int(mapping.enable_attention_dp)) + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=num_reserved_dummy_slots, + include_explicit_snapshots=True, + cap_partial_attention_snapshots=True, + kda_replay_num_spec=_estimated_kda_replay_num_spec( + model_config, spec_config, kwargs.get("is_disagg", False)), + **kwargs, + ) + + def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: + return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] + + def _get_pool_roles(self, + pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return MambaRole.SSM_STATE, None + return super()._get_pool_roles(pool_id) + + def _max_resident_sequences(self) -> int: + return self.max_batch_size * self.mapping.pp_size + + def _kda_draft_records_shape(self) -> Tuple[int, int, int, int]: + """One SSM slot's shape in a layer of ``kda_state_tok``: per draft, + the row innovations vn, then beta * k and the decay, each [H, K] (the + replay caches' equal [Q | K | V] sections make V = K).""" + return (3, self._kda_replay_num_spec, self.ssm_state_shape[0], + self.ssm_state_shape[-1]) + + def _mamba_state_bytes_per_slot(self) -> int: + base_bytes = (self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes) + + _kda_replay_bytes_per_slot(self._kda_replay_layout(), + self.local_num_mamba_layers)) + local_ple_layers = sum(layer_id in self.pp_layers + for layer_id in self._ple_layer_ids) + if local_ple_layers == 0: + return base_bytes + ple_bytes = ( + math.prod(self._ple_conv_state_shape) * + self._ple_conv_state_dtype.itemsize + + math.prod(self._ple_ngram_context_shape) * torch.int64.itemsize) + return base_bytes + local_ple_layers * ple_bytes + + def _ple_buffer_configs(self) -> list[BufferConfig]: + if not self._ple_layer_ids: + return [] + return [ + BufferConfig( + role=MambaRole.PLE_NGRAM_CONTEXT, + size=(math.prod(self._ple_ngram_context_shape) * + torch.int64.itemsize), + ), + BufferConfig( + role=MambaRole.PLE_CONV_STATE, + size=(math.prod(self._ple_conv_state_shape) * + self._ple_conv_state_dtype.itemsize), + ), + ] + + def _num_ssm_snapshots_for_capacity( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + if capacity <= 0 or not kv_cache_config.enable_block_reuse: + return 0 + + fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, + self.max_seq_len, + self.tokens_per_block) + interval = _mamba_regular_snapshot_interval(kv_cache_config, + self.max_seq_len) + regular_snapshots = capacity // interval if interval is not None else 0 + return (self._max_resident_sequences() * fixed_rules + + regular_snapshots) + + def _num_ssm_states_per_typical_request( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, + capacity, + self.tokens_per_block, + ) + # Additional snapshots are stable boundaries that must remain alive. + # Periodic snapshots are evictable cache entries and therefore do not + # increase the guaranteed state count represented by BatchDesc. + return 1 + fixed_rules + + def _typical_request_descs( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[KVCacheDesc]: + """Model one request with one descriptor per live SSM state.""" + num_states = self._num_ssm_states_per_typical_request( + capacity, kv_cache_config) + capacity_per_state, capacity_remainder = divmod(capacity, num_states) + capacities = [ + capacity_per_state + int(i < capacity_remainder) + for i in range(num_states) + ] + return [ + KVCacheDesc( + capacity=state_capacity, + history_length=max(0, state_capacity - 1), + ) for state_capacity in capacities + ] + + def _get_typical_request_capacity( + self, + kv_cache_config: KvCacheConfig, + ) -> int: + if kv_cache_config.avg_seq_len is not None: + return kv_cache_config.avg_seq_len + + fallback_capacity = max(1, self.max_seq_len // 2) + logger.warning( + "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " + "model using KV cache manager V2. Falling back to " + f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " + "'kv_cache_config.avg_seq_len' in the YAML configuration to the " + "workload's average total sequence length for an accurate KV/SSM " + "pool ratio.") + return fallback_capacity + + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + attention_quota = super()._get_quota_from_max_tokens(max_tokens) + num_request_lineages = self._max_resident_sequences() + snapshot_slots = self._num_ssm_snapshots_for_capacity( + max_tokens, self.kv_cache_config) + state_slots = (num_request_lineages + self._num_reserved_dummy_slots + + snapshot_slots) + state_quota = state_slots * self._mamba_state_bytes_per_slot() + # Once the plan contains any non-live SSM capacity, reserve one partial + # attention page per request lineage. This remains conservative when + # the plan contains fewer than one non-live slot per lineage. + extra_attention_quota = (num_request_lineages * + self._attention_cache_bytes_per_token() * + self.tokens_per_block + if snapshot_slots > 0 else 0) + return attention_quota + state_quota + extra_attention_quota + + def _get_max_tokens_from_quota(self, quota: int) -> float: + if self._get_quota_from_max_tokens(0) > quota: + return 0 + + low = 0 + high = 1 + while self._get_quota_from_max_tokens(high) <= quota: + low = high + high *= 2 + if high >= 1 << 62: + return float("inf") + + while low + 1 < high: + mid = (low + high) // 2 + if self._get_quota_from_max_tokens(mid) <= quota: + low = mid + else: + high = mid + return low + + def _minimum_live_gpu_quota(self) -> int: + """Return the minimum quota for live states and one attention page.""" + attention_block_quota = (self._attention_cache_bytes_per_token() * + self.tokens_per_block) + num_state_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + return max( + self._get_quota_from_max_tokens(0), + state_quota + attention_block_quota, + ) + + def _build_cache_config( + self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + kv_cache_config = self.kv_cache_config + if (getattr(self, "_use_kda_replay_update", False) + and self.local_num_mamba_layers > 0 + and kv_cache_config.enable_kv_pool_rebalance): + # The KDA replay views and prev_num_accepted_tokens are bound + # once to the SSM pool's slot count at construction; a rebalance + # can grow the SSM pool past them, and the verify kernels index + # them by slot without a bound. + raise ValueError( + "The KDA replay caches (speculative decoding on KDA layers) " + "are sized once from the SSM pool's slot count, and a KV pool " + "rebalance can grow that pool past them: set " + "kv_cache_config.enable_kv_pool_rebalance=False.") + cache_tiers = config.cache_tiers + gpu_quota = cache_tiers[0].quota + minimum_live_quota = self._minimum_live_gpu_quota() + if minimum_live_quota > gpu_quota: + raise ValueError( + "The V2 Mamba GPU cache quota is too small for live recurrent " + f"states and attention pages: got {gpu_quota} bytes, need at " + f"least {minimum_live_quota} bytes.") + # _build_base_config already constructed every attention layer, + # including dtype-specific scale and subclass-provided side buffers. + # Preserve those configs and replace only the local Mamba layers. + layers = list(config.layers) + kda_replay_buffers = self._kda_replay_buffer_configs() + first_mamba_layer = (self.mamba_pp_layers[0] + if kda_replay_buffers else None) + for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): + if self._mamba_layer_mask[global_layer_idx]: + layer_id = LayerId(local_layer_idx) + buffers = [ + BufferConfig(role=MambaRole.SSM_STATE, size=self.ssm_bytes), + BufferConfig(role=MambaRole.CONV_STATE, + size=self.conv_bytes), + ] + if global_layer_idx in self._ple_layer_ids: + buffers.extend(self._ple_buffer_configs()) + if kda_replay_buffers: + buffers.extend(kda_replay_buffers) + if global_layer_idx == first_mamba_layer: + # All SSM layers share a slot lifecycle, so the + # request's accepted-draft count is stored once and + # moves with the replay histories. + buffers.append( + BufferConfig(role=_KDA_REPLAY_COUNT_ROLE, + size=_KDA_REPLAY_COUNT_BYTES)) + layers[local_layer_idx] = SsmLayerConfig( + layer_id=layer_id, + buffers=buffers, + ) + + dummy_requests = [ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(self._num_reserved_dummy_slots) + ] + constraints = [ + replace( + batch, + kv_caches=[*batch.kv_caches, *dummy_requests], + ) for batch in config.constraints + ] + + typical_step = config.typical_step + if config.initial_pool_ratio is None: + typical_capacity = self._get_typical_request_capacity( + kv_cache_config) + request_descs = self._typical_request_descs(typical_capacity, + kv_cache_config) + typical_step = BatchDesc(request_descs * + self._max_resident_sequences() + + dummy_requests) + # The recurrent (SSM) state pool must hold one slot per resident + # sequence plus every reserved dummy slot. Unlike attention pages, a + # Mamba state is fixed-size per sequence, so this floor is independent + # of sequence length. The base config only emits constraints when + # ``avg_seq_len`` is set, and speculative decoding inflates the reserved + # dummy slots (CUDA-graph padding), so without an explicit floor the SSM + # pool can be undersized (see the live/dummy-slot check in _setup_states + # / __init__). Add a min-slots constraint of zero-capacity requests: + # these cost no attention pages but reserve one SSM slot each. + if any(isinstance(layer, SsmLayerConfig) for layer in layers): + ssm_floor_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + constraints = [ + *constraints, + BatchDesc([ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(ssm_floor_slots) + ]), + ] + return replace( + config, + layers=layers, + typical_step=typical_step, + constraints=constraints, + # SSM lifecycles require minimum-snapshot commit semantics. The + # flag is harmless when reuse is disabled because no commits are + # attempted, while the runtime config still needs the invariant. + commit_min_snapshot=True, + ) + + def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, + state_shape: List[int]) -> torch.Tensor: + addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), + role, PageIndexMode.SHARED) + num_pages = self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), role) + raw = convert_to_torch_tensor( + TensorWrapper(addr, dtype, [num_pages] + state_shape)) + page_index_scale = self.impl.get_page_index_scale( + LayerId(local_layer_idx), role) + num_slots = (num_pages + page_index_scale - 1) // page_index_scale + # V2 coalesces same-size per-layer buffers inside each slot. Kernels + # index Mamba states by logical slot id, so expose only this layer's + # sub-page from each coalesced slot instead of the raw page-index view. + return raw.as_strided( + [num_slots] + state_shape, + [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), + ) + + def _setup_states(self) -> None: + local_layer_ids = [ + self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers + ] + self.all_ssm_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, + self.ssm_state_dtype, self.ssm_state_shape) + for local_layer_idx in local_layer_ids + ] + self.all_conv_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, + self.conv_state_dtype, self.conv_state_shape) + for local_layer_idx in local_layer_ids + ] + self._stacked_ssm_states = self._stack_state_views(self.all_ssm_states) + self._stacked_conv_states = self._stack_state_views( + self.all_conv_states) + if self._use_gdn_cached_replay_all_layer_commit: + expected_inner_strides = ( + self.ssm_state_shape[1] * self.ssm_state_shape[2], + self.ssm_state_shape[2], + 1, + ) + reference = self.all_ssm_states[0] + for state in self.all_ssm_states: + if (state.dtype != reference.dtype + or state.device != reference.device + or list(state.shape[1:]) != self.ssm_state_shape + or state.stride()[1:] != expected_inner_strides): + raise RuntimeError( + "GDN cached replay requires V2 SSM layers with matching " + "dtype, device, shape, and dense inner dimensions.") + pointers = [state.data_ptr() for state in self.all_ssm_states] + slot_stride = reference.stride(0) + layer_stride_bytes = (pointers[1] - + pointers[0] if len(pointers) > 1 else 0) + has_affine_layout = ( + layer_stride_bytes % reference.element_size() == 0 and all( + state.stride(0) == slot_stride + for state in self.all_ssm_states) + and all(pointer == pointers[0] + layer * layer_stride_bytes + for layer, pointer in enumerate(pointers))) + if has_affine_layout: + self._gdn_cached_replay_state_strides = ( + len(pointers), + layer_stride_bytes // reference.element_size(), + slot_stride, + ) + else: + logger.warning_once( + "V2 GDN state views are not affine; using indirect replay " + "checkpoint addressing", + key="gdn_cached_replay_v2_indirect_state_layout", + ) + self._gdn_cached_replay_state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) + for state in self.all_ssm_states], + dtype=torch.int64, + device=reference.device, + ) + + @staticmethod + def _stack_state_views( + states: List[torch.Tensor]) -> Optional[torch.Tensor]: + """Returns one ``[num_layers, num_slots, *state_shape]`` view of ``states``. + + A stacked view collapses per-layer promotion loop into a single launch of _promote_mamba_state_triton. + This requires per-layer pool views to be *affine*: identical dtype/device/shape/strides, and base + pointers spaced by a constant. returns ``None`` otherwise. + """ + if not states: + return None + ref = states[0] + if len(states) == 1: + return ref.unsqueeze(0) + for state in states[1:]: + if (state.dtype != ref.dtype or state.device != ref.device + or state.shape != ref.shape + or state.stride() != ref.stride()): + return None + itemsize = ref.element_size() + layer_stride_bytes = states[1].data_ptr() - states[0].data_ptr() + if layer_stride_bytes <= 0 or layer_stride_bytes % itemsize != 0: + return None + base = states[0].data_ptr() + if any(state.data_ptr() != base + layer * layer_stride_bytes + for layer, state in enumerate(states)): + return None + layer_stride = layer_stride_bytes // itemsize + # Highest element one layer's view can reach, so the flat wrapper below + # covers every byte the stacked view addresses and no more. + layer_extent = 1 + sum((size - 1) * stride + for size, stride in zip(ref.shape, ref.stride())) + total_elems = (len(states) - 1) * layer_stride + layer_extent + flat = convert_to_torch_tensor( + TensorWrapper(base, ref.dtype, [total_elems])) + return flat.as_strided([len(states)] + list(ref.shape), + [layer_stride] + list(ref.stride())) + + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = 0 + device = None + if self.local_num_mamba_layers > 0: + cache_size = self.all_ssm_states[0].shape[0] + assert all(t.shape[0] == cache_size for t in self.all_ssm_states) + device = self.all_ssm_states[0].device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return + + mask_capacity = self._host_state_indices.shape[0] + self._dummy_request_mask_buffer = DeviceCopyBuffer( + torch.zeros(mask_capacity, dtype=torch.bool, device=device), + pin_memory=prefer_pinned()) + self._dummy_request_mask_host = torch.zeros( + mask_capacity, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) + + def _attention_cache_bytes_per_token(self) -> int: + # Mamba layers have zero KV heads, so the generic calculation naturally + # returns only bytes owned by local attention layers. + return super().get_cache_bytes_per_token() + + def get_cache_bytes_per_token(self) -> int: + cache_bytes = self._attention_cache_bytes_per_token() + + interval = ( + self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) + if (self.kv_cache_config.enable_block_reuse and interval is not None + and interval > 0): + cache_bytes += self._mamba_state_bytes_per_slot() // interval + if cache_bytes == 0 and self.local_num_mamba_layers > 0: + cache_bytes = self._mamba_state_bytes_per_slot() + return max(1, cache_bytes) + + def get_num_free_blocks(self) -> int: + # Mirror the base: the guard page, when enabled, is held by a + # permanent reservation made at construction (not by a request), so + # exclude reserved ids before asserting the manager is otherwise empty. + reserved = set(self.kv_cache_map) & _RESERVED_REQUEST_IDS + assert not set(self.kv_cache_map) - reserved, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + attention_pages = [] + ssm_pages = [] + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + if self._is_local_mamba_layer(local_layer_idx): + ssm_pages.append( + self.impl.get_page_index_upper_bound( + layer_id, MambaRole.SSM_STATE) // + self._ssm_page_index_scale) + else: + attention_pages.append( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // + self.kv_factor) + if attention_pages: + return max(attention_pages) + return max(ssm_pages) if ssm_pages else 0 + + @property + def blocks_in_primary_pool(self) -> int: + for local_layer_idx in range(self.num_local_layers): + if self._is_local_mamba_layer(local_layer_idx): + continue + return self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), Role.KEY) + return 0 + + def get_buffers(self, + layer_idx: int, + kv_layout: str = "NHD") -> Optional[torch.Tensor]: + local_layer_idx = self.layer_offsets[layer_idx] + if self._is_local_mamba_layer(local_layer_idx): + return None + return super().get_buffers(layer_idx, kv_layout) + + def _iter_cache_buffers_for_invalid_check( + self) -> Iterable[Tuple[int, torch.Tensor]]: + for global_layer_id, local_layer_id in self.layer_offsets.items(): + if self._is_local_mamba_layer(local_layer_id): + continue + # A layer group is a lifecycle, not a physical memory pool. + # Differently sized attention buffers can share one lifecycle, + # so scan every attention layer in this diagnostic path. + yield global_layer_id, KVCacheManagerV2.get_buffers( + self, global_layer_id) + + # SSM / conv / PLE tensors are not attention layers and hold no guard + # page, so yield a sentinel layer id (-1). The base consumer looks the + # id up via self._guard_page_by_layer.get(layer_id), which returns None + # for an unknown id and correctly skips guard restoration for them. + for buffer in self.all_ssm_states: + yield -1, buffer + for buffer in self.all_conv_states: + yield -1, buffer + for buffer in self._ple_ngram_contexts.values(): + yield -1, buffer + for buffer in self._ple_conv_states.values(): + yield -1, buffer + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + super()._reset_context_mamba_slots(num_contexts) + if num_contexts == 0 or not self.use_kda_replay_update: + return + assert self.prev_num_accepted_tokens is not None + context_slots = _host_slots_to_device( + self._host_state_indices[:num_contexts].tolist(), + self.prev_num_accepted_tokens.device) + self.prev_num_accepted_tokens.index_fill_(0, context_slots, 0) + + def _reset_kda_replay_slots( + self, + slots: Union[List[int], torch.Tensor], + ) -> None: + """Reset pending KDA drafts for newly initialized logical slots.""" + assert self.use_kda_replay_update + assert self.prev_num_accepted_tokens is not None + device = self.prev_num_accepted_tokens.device + if isinstance(slots, torch.Tensor) and slots.device.type != "cpu": + slots = slots.to(device=device, dtype=torch.long, non_blocking=True) + else: + if isinstance(slots, torch.Tensor): + slots = slots.tolist() + if not slots: + return + slots = _host_slots_to_device(slots, device) + if slots.numel() > 0: + self.prev_num_accepted_tokens.index_fill_(0, slots, 0) + + def suspend_request(self, req: LlmRequest) -> None: + """Suspend the request's KV cache with its whole KDA state in V2. + + The drafts the last verify accepted are pending in the slot's KDA + replay buffers, and so is the conv window, which the fused verify does + not write to the V2 conv state. The pending drafts and the conv window + are committed into the slot's SSM and conv states here, before the + pages are released, so the suspended state is self-contained in the + recurrent state; the request's next generation step re-seeds the + replay buffers of whichever slot it resumes in from them. + """ + request_id = req.py_request_id + kv_cache = self.kv_cache_map.get(request_id) + slot = self._request_id_to_state_index.get(request_id, -1) + if (self.use_kda_replay_update and slot >= 0 and kv_cache is not None + and kv_cache.is_active): + # A host drafter's acceptance is recorded by update_resources. + # Executors record it before they can suspend the request: the + # host drafter runs without the overlap scheduler, and the PP + # loop does not evict requests in flight. + if request_id in self._kda_replay_acceptance_pending: + raise RuntimeError( + f"Request {request_id} is being suspended before " + "update_resources recorded the acceptance of its last " + "verify, so its pending KDA drafts are unknown") + if request_id in self._kda_replay_unseeded: + self._kda_replay_unseeded.remove(request_id) + else: + self._commit_kda_pending_drafts(slot) + del self._request_id_to_state_index[request_id] + self._kda_replay_suspended.add(request_id) + super().suspend_request(req) + + @contextmanager + def _on_kv_cache_stream(self): + """Run the enclosed work on the KV caches' stream, after the current + stream's earlier work and before its later work. + + V2 orders page migration against that stream: a suspend records the + event later migrations of the pages wait for, and a resume makes the + stream wait for the pages it brings back. + """ + current = torch.cuda.current_stream() + self._stream.wait_stream(current) + with torch.cuda.stream(self._stream): + yield + current.wait_stream(self._stream) + + def _commit_kda_pending_drafts(self, slot: int) -> None: + """Commit the slot's pending accepted drafts and its conv window into + its V2 SSM and conv states, in place, and clear its draft count. + + Afterwards the V2 states hold what the next verify would have started + from: the conv window exactly, the SSM state to the verify kernel's + accuracy (``commit_kda_pending_drafts``). + """ + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ + commit_kda_pending_drafts + + assert self.prev_num_accepted_tokens is not None + with self._on_kv_cache_stream(): + slots = torch.full((1, ), + slot, + dtype=torch.int32, + device=self.prev_num_accepted_tokens.device) + num_accepted = self.prev_num_accepted_tokens.index_select(0, slots) + for layer_offset, (ssm_states, conv_states) in enumerate( + zip(self.all_ssm_states, self.all_conv_states)): + replay = self._kda_replay_layers[layer_offset] + # With the verify kernels' draft records (kda_state_tok), the + # accepted drafts are replayed from those, and the k / g / v / + # beta replay caches are not allocated. + commit_kda_pending_drafts( + ssm_states, + conv_states, + replay["kda_conv_q"], + replay["kda_conv_k"], + replay["kda_conv_v"], + replay.get("kda_k_cache"), + replay.get("kda_g_cache"), + replay.get("kda_v_cache"), + replay.get("kda_beta_cache"), + slots, + num_accepted, + state_tok=replay.get(_KDA_DRAFT_RECORDS_ROLE), + ) + self.prev_num_accepted_tokens[slot] = 0 + + def add_dummy_requests( + self, + request_ids: List[int], + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + draft_kv_cache_manager: Optional[BaseResourceManager] = None, + capture_sampling_params: Optional["SamplingParams"] = None, + ) -> List[LlmRequest]: + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + max_num_draft_tokens=max_num_draft_tokens, + kv_reserve_draft_tokens=kv_reserve_draft_tokens, + use_mrope=use_mrope, + max_beam_width=max_beam_width, + encoder_output_lens=encoder_output_lens, + draft_kv_cache_manager=draft_kv_cache_manager, + capture_sampling_params=capture_sampling_params, + ) + if requests: + self._setup_state_indices( + requests, num_contexts=0 if is_gen else len(requests)) + return requests + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is not None and kv_cache.is_active: + self.try_commit_blocks(request, kv_cache) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + self._branch_snapshot_points.pop(request.py_request_id, None) + self._kda_replay_suspended.discard(request.py_request_id) + self._kda_replay_acceptance_pending.discard(request.py_request_id) + self._kda_replay_unseeded.discard(request.py_request_id) + super().free_resources(request, pin_on_release) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + num_contexts = len(scheduled_batch.context_requests) + self._setup_state_indices(requests, num_contexts=num_contexts) + self._reset_context_mamba_slots(num_contexts) + if self.use_kda_replay_update: + self._kda_replay_unseeded.update( + request.py_request_id + for request in scheduled_batch.context_requests + if request.is_disagg_generation_init_state) + if self._record_kda_replay_in_update_resources: + # Recorded by this batch's update_resources, or dropped by + # revert_allocate_generation when the batch is not run. + self._kda_replay_acceptance_pending.update( + request.py_request_id + for request in scheduled_batch.generation_requests + if not request.is_dummy) + + def revert_allocate_generation(self, req: LlmRequest) -> None: + self._kda_replay_acceptance_pending.discard(req.py_request_id) + super().revert_allocate_generation(req) + + def _setup_state_indices(self, + requests: List[LlmRequest], + num_contexts: int = 0) -> None: + if self.local_num_mamba_layers == 0: + return + replay_requests = (requests[num_contexts:] + if self.use_kda_replay_update else []) + old_state_values = [ + self._request_id_to_state_index.get(request.py_request_id, -1) + for request in replay_requests + ] + n = len(requests) + assert n <= self._host_state_indices.shape[0], ( + f"State-index batch size {n} exceeds max_batch_size " + f"{self._host_state_indices.shape[0]}") + self._host_state_indices.zero_() + if n > 0: + for i, req in enumerate(requests): + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + raise RuntimeError( + f"Missing V2 KV cache for request {req.py_request_id}") + base_index = kv_cache.get_ssm_block_base_index( + self.ssm_layer_group_id) + if base_index < 0: + raise RuntimeError( + f"Invalid SSM state block index {base_index} for " + f"request {req.py_request_id}") + self._host_state_indices[i] = base_index + + self._state_indices_buffer.update(self._host_state_indices) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + if self.use_kda_replay_update: + # A suspended request that comes back as a context request starts + # over. One that comes back generating has its whole KDA state in + # the V2 states of the slot it resumes in (suspend_request). A + # live request whose slot moved carries its replay buffers along: + # they are KVCM state buffers of the same slot. + for request in requests[:num_contexts]: + self._kda_replay_suspended.discard(request.py_request_id) + new_state_values = [ + self._request_id_to_state_index[request.py_request_id] + for request in replay_requests + ] + new_slots = [] + resumed_slots = [] + for request, old, new in zip(replay_requests, old_state_values, + new_state_values): + if request.py_request_id in self._kda_replay_suspended: + self._kda_replay_suspended.remove(request.py_request_id) + resumed_slots.append(new) + elif old < 0: + new_slots.append(new) + self._reset_kda_replay_slots(new_slots) + if resumed_slots: + resumed_indices = torch.tensor(resumed_slots, + dtype=torch.long, + pin_memory=prefer_pinned()) + with self._on_kv_cache_stream(): + self._seed_kda_replay_slots( + resumed_indices.to(self.prev_num_accepted_tokens.device, + non_blocking=True)) + + def _record_kda_replay_acceptance( + self, + state_indices: torch.Tensor, + num_accepted_draft_tokens: torch.Tensor, + is_dummy_request: Optional[torch.Tensor] = None, + ) -> None: + """Record the draft prefix that the next KDA verify must replay.""" + if self.prev_num_accepted_tokens is None: + raise RuntimeError( + "KDA replay acceptance requires replay bookkeeping") + slots = state_indices.to(torch.long) + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + if is_dummy_request is not None: + current = self.prev_num_accepted_tokens[slots] + accepted = torch.where(is_dummy_request, current, accepted) + self.prev_num_accepted_tokens[slots] = accepted + + @torch.inference_mode() + def seed_kda_replay_caches_for_disagg_gen(self, + request_ids: List[int]) -> None: + """Seed replay caches from transferred ``W - 1`` conv states. + + This is the V2-storage counterpart of + ``MambaCacheManager.seed_kda_replay_caches_for_disagg_gen``; see that + method for the transfer-ordering rationale and replay contract. + """ + if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): + return + slots = [ + self._request_id_to_state_index[request_id] + for request_id in request_ids + if request_id in self._request_id_to_state_index + ] + self._kda_replay_unseeded.difference_update(request_ids) + if not slots: + return + + assert self.prev_num_accepted_tokens is not None + state_indices = _host_slots_to_device( + sorted(set(slots)), self.prev_num_accepted_tokens.device) + self._seed_kda_replay_slots(state_indices) + + def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: + """Seed the replay caches of distinct slots from their V2 conv + states, with no drafts pending.""" + assert self.prev_num_accepted_tokens is not None + committed_window = self.conv_state_shape[1] + section_offsets = [0] + for section_dim in self.conv_section_dims: + section_offsets.append(section_offsets[-1] + section_dim) + + self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) + for replay_layer, conv_state in zip(self._kda_replay_layers, + self.all_conv_states): + for name, replay_buffer in replay_layer.items(): + # The drafts' records need no reset: with no pending drafts, + # none of them is read. + if name != _KDA_DRAFT_RECORDS_ROLE: + replay_buffer.index_fill_(0, state_indices, 0) + selected_conv = conv_state.index_select(0, state_indices) + for section_idx, name in enumerate(_KDA_REPLAY_CONV_ROLES): + replay_buffer = replay_layer[name] + section = selected_conv[:, section_offsets[section_idx]: + section_offsets[section_idx + 1], :] + replay_buffer[:, :, :committed_window].index_copy_( + 0, state_indices, section.to(replay_buffer.dtype)) + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + if not self.use_kda_replay_update: + return super().mamba_layer_cache(layer_idx) + + layer_offset = self.mamba_layer_offsets[layer_idx] + assert self.prev_num_accepted_tokens is not None + spec_kwargs = dict(self._kda_replay_layers[layer_offset]) + spec_kwargs["prev_num_accepted_tokens"] = self.prev_num_accepted_tokens + if self.mamba_ssm_rand_seed is not None: + spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed + return PythonMambaCacheManager.SpeculativeState( + conv=self.get_conv_states(layer_idx), + temporal=self.get_ssm_states(layer_idx), + intermediate_conv_window=None, + **spec_kwargs, + ) + + def update_resources( + self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float | None = None, + ) -> None: + super().update_resources(scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + if (not self._record_kda_replay_in_update_resources + or self.local_num_mamba_layers == 0): + return + + generation_requests = scheduled_batch.generation_requests + self._kda_replay_acceptance_pending.difference_update( + request.py_request_id for request in generation_requests) + drafted_requests = [ + request for request in generation_requests + if request.py_draft_tokens is not None + and len(request.py_draft_tokens) > 0 + ] + if not drafted_requests: + return + if len(drafted_requests) != len(generation_requests): + raise RuntimeError( + "Mixed drafted/undrafted generation batch is not supported " + "for KDA replay bookkeeping") + + active_requests = [ + request for request in generation_requests + if request.py_request_id in self._request_id_to_state_index + ] + if not active_requests: + return + assert self.prev_num_accepted_tokens is not None + device = self.prev_num_accepted_tokens.device + state_indices = torch.tensor([ + self._request_id_to_state_index[request.py_request_id] + for request in active_requests + ], + dtype=torch.int32, + device=device) + accepted_drafts = torch.tensor([ + request.py_num_accepted_draft_tokens for request in active_requests + ], + dtype=torch.int32, + device=device) + is_dummy_request = torch.tensor([ + self._request_id_to_is_dummy.get(request.py_request_id, False) + for request in active_requests + ], + dtype=torch.bool, + device=device) + self._record_kda_replay_acceptance(state_indices, accepted_drafts, + is_dummy_request) + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None): + if self.local_num_mamba_layers == 0: + # Mamba metadata is still prepared on attention-only PP ranks, + # but no local kernel consumes these indices. Return harmless + # placeholders instead of consulting an intentionally empty map. + if request_ids is not None: + return [0] * len(request_ids) + return self.cuda_state_indices + if request_ids is not None: + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + def get_max_resource_count(self) -> int: + return self.max_batch_size + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0 or self.kda_context_only: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = ( + num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( + torch.int32) + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + if self.use_kda_replay_update: + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + self._record_kda_replay_acceptance( + state_indices_d, + num_accepted_draft_tokens, + is_dummy_request, + ) + return + + # Dynamic tree selects a tree node rather than a linear draft depth. + accepted_positions = (accepted_leaf_positions.to(torch.int32) + if accepted_leaf_positions is not None else + num_accepted_draft_tokens) + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_replay_state_update: + self._commit_gdn_cached_replay_history_layers( + attn_metadata, num_gens) + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) + else: + self._promote_state(self.all_ssm_states, self._stacked_ssm_states, + self.intermediate_ssm_states, src_state_indices, + accepted_positions, state_indices_d) + + self._promote_state(self.all_conv_states, self._stacked_conv_states, + self.intermediate_conv_states, src_state_indices, + accepted_positions, state_indices_d) + + @staticmethod + def _promote_state(per_layer: List[torch.Tensor], + stacked: Optional[torch.Tensor], + intermediate: torch.Tensor, + src_state_indices: torch.Tensor, + accepted_positions: torch.Tensor, + state_indices_d: torch.Tensor) -> None: + """Promote every layer's accepted state, in one launch where possible. + + The kernel's grid is ``(num_layers * num_gens, tiles)``, so handing it + the stacked view does all layers at once; the per-layer loop is only + the fallback for a non-affine pool layout. + """ + if stacked is not None: + _promote_mamba_state_triton(stacked, intermediate, + src_state_indices, accepted_positions, + state_indices_d) + return + for layer_offset, dst in enumerate(per_layer): + _promote_mamba_state_triton( + dst.unsqueeze(0), + intermediate[layer_offset:layer_offset + 1], + src_state_indices, + accepted_positions, + state_indices_d, + ) + + def _mark_context_position_as_history(self, request: LlmRequest, + kv_cache) -> None: + """Advance history without making later recurrent state reusable.""" + history_length = request.context_current_position + if history_length <= kv_cache.history_length: + return + capacity = max(kv_cache.capacity, history_length) + if not kv_cache.resize(capacity, history_length=history_length): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to {history_length} tokens") + + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + if not should_block_reuse: + return + + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + snapshot_points = request.expect_snapshot_points + commit_limit = (min(max(snapshot_points), request.prompt_len) + if snapshot_points else request.prompt_len) + commit_end = min(request.context_current_position, commit_limit) + if (request.context_current_position in request.expect_snapshot_points + and commit_end > kv_cache.num_committed_tokens): + tokens = self._augment_tokens_for_block_reuse( + request.get_tokens(DEFAULT_BEAM_INDEX), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + kv_cache.commit(tokens) + if (self.local_num_mamba_layers > 0 + and self._branch_snapshot_points.get( + request.py_request_id) == commit_end): + self._branch_snapshots_taken_total += 1 + if request.context_current_position >= commit_limit: + self._mark_context_position_as_history(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def update_context_resources(self, + scheduled_batch: ScheduledRequests) -> None: + for request in scheduled_batch.context_requests: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None or not kv_cache.is_active: + continue + + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + is_all_reusable = ( + self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) + is_snapshot_boundary = (request.context_current_position + in request.expect_snapshot_points) + has_pending_snapshot = any( + point > request.context_current_position + for point in request.expect_snapshot_points) + should_resize = (not should_block_reuse or + (not is_all_reusable and not has_pending_snapshot)) + should_commit = (is_all_reusable or is_snapshot_boundary + or request.context_remaining_length == 0) + + if should_resize and not kv_cache.resize( + None, request.context_current_position): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to " + f"{request.context_current_position} tokens at context " + "update") + if should_commit: + self.try_commit_blocks(request, kv_cache) + if request.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(request, kv_cache) + kv_cache.enable_swa_scratch_reuse = False + + def shutdown(self): + self._gdn_cached_replay_state_descriptors = None + self._gdn_cached_replay_state_strides = None + self.all_ssm_states = [] + self.all_conv_states = [] + self._stacked_ssm_states = None + self._stacked_conv_states = None + self._ple_ngram_contexts = {} + self._ple_conv_states = {} + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask_buffer = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + self._branch_snapshot_points.clear() + self._kda_replay_layers = [] + self._kda_replay_suspended.clear() + self._kda_replay_acceptance_pending.clear() + self._kda_replay_unseeded.clear() + super().shutdown() diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 6b7e6b517876..3b9f525cf38e 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -30,7 +30,11 @@ from tensorrt_llm.mapping import Mapping from ..attention.backends import AttentionMetadata +<<<<<<< HEAD from ..modules.logits_processor import LogitsProcessor +======= +from ..attention.backends.features import DraftKVMetadata, RecurrentMetadata +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) from ..pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from ..pyexecutor.resource_manager import BaseResourceManager @@ -1165,14 +1169,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 +1196,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 +1871,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..e7a3d9738316 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( @@ -1965,24 +2082,24 @@ def test_native_staging_pools_cover_tma_fetch_window( @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 +2111,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 +2124,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 +2161,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 +2214,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 +2239,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 +2258,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 +2292,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 +2323,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 +2346,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 +2361,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 +2371,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 +2386,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 +2413,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 +2466,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 +2530,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 +2551,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 +2561,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_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..d1e20e57d4f5 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,41 @@ 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() +<<<<<<< HEAD self.indexer_k_cache_block_offsets = torch.zeros( [batch_size, self.kv_cache_manager.max_blocks_per_seq], +======= + dsa.indexer_k_cache_block_offsets = torch.zeros( + [batch_size, cache_manager.max_blocks_per_seq], +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) 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 +1974,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 +2014,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 +2030,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 +2043,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 +2063,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 +2112,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 +2159,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 +2238,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 +2304,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 +2324,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 +2507,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 +2535,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 +2679,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 +2691,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 +2703,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 +2718,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 +2726,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 +2759,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 +2772,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 +2791,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 +2943,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 +2959,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 +2972,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 +3045,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 +3085,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 +3115,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 +3174,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 +3185,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 +3203,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 +3861,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 +3903,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 +4363,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 +4380,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 +4403,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 +4495,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 +4609,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 +4637,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 +4662,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 +4686,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 +4778,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 +5000,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 +5052,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 +5080,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 +5110,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 +5139,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 +5159,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 +5185,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 +5245,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 +5253,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 +5322,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..f127be213888 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,9 +12,13 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer +<<<<<<< HEAD from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.cache_manager import ( Glm5NextCacheManager, ) +======= +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager from tensorrt_llm._torch.disaggregation.resource.page import ( MAMBA_CONV_ROLE, @@ -123,6 +127,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") @@ -5353,8 +5366,13 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): history_size=policy.replay_history_size, ) +<<<<<<< HEAD policy.commit_all_layers( SimpleNamespace( +======= + mgr._commit_gdn_cached_replay_history_layers( + _recurrent_stub( +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) mamba_metadata=SimpleNamespace( replay_num_decodes=batch_size, replay_work_items=replay_work_items, 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..ff05a3e781ec 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -9,7 +9,12 @@ 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 ( +<<<<<<< HEAD QSAAttentionMetadata, +======= + QSAMambaHybridCacheManagerV2, + QSAMetadata, +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) QSASparseMetadataParams, QSASparseParams, ) @@ -201,7 +206,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 +230,26 @@ def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: assert output is None +<<<<<<< HEAD def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers(monkeypatch) -> None: monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) manager = object.__new__(Qwen4ExpHybridCacheManager) +======= +def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers() -> None: + manager = object.__new__(QSAMambaHybridCacheManagerV2) +>>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) 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 +415,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 +445,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) From 840cac8c17e63ea73d29e5948f031884dab5aca7 Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Tue, 27 Oct 2026 22:43:25 -0700 Subject: [PATCH 2/3] [None][feat] Resolve the metadata-extension port onto newer main Resolve the cherry-pick conflicts and move code that main added since the original base onto the extension API: - qsa: keep main's lazy package exports and its Qwen4ExpHybridCacheManager check (the QSA cache manager moved to modules/qwen4_exp); export QSAMetadata. - mamba: main split mamba_cache_manager.py into a package; drop the stale single file and read RecurrentMetadata in mamba_cache_manager/legacy.py and modules/fla/cache_manager.py (GDN replay commit). - deepseek_v4: allocate main's new token_positions_cuda / sparse_mla_*_lens buffers through the core (md.get_empty) inside DeepseekV4Metadata. - fmha/prims_ts_sparse_mla.py (new on main): read DeepseekV4Metadata, HelixMetadata and SpecDecodingMetadata instead of subclass fields. - tests: dflash/dsa/qsa/mamba conflicts, and main's new CSA2 staging-pool test reads the pools from CSA2Metadata. Signed-off-by: Yihan Wang --- .../backends/fmha/prims_ts_sparse_mla.py | 35 +- .../backends/sparse/deepseek_v4/metadata.py | 18 +- .../attention/backends/sparse/qsa/__init__.py | 17 +- .../attention/backends/sparse/qsa/metadata.py | 32 +- .../_torch/modules/fla/cache_manager.py | 4 +- .../kv_cache/mamba_cache_manager.py | 5147 ----------------- .../kv_cache/mamba_cache_manager/legacy.py | 6 +- tensorrt_llm/_torch/speculative/dflash.py | 5 +- .../attention/sparse/csa2/test_metadata.py | 15 +- .../attention/sparse/dsa/test_dsa_indexer.py | 7 +- .../kv_cache/test_mamba_cache_manager.py | 10 +- .../modeling/test_qsa_runtime_wiring.py | 13 +- 12 files changed, 54 insertions(+), 5255 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py 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/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 0ceea03e3ab6..560dc3edc9da 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -244,9 +244,9 @@ def allocate(self, buffers) -> None: # 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, @@ -324,17 +324,17 @@ def allocate(self, buffers) -> None: # 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, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py index d44f80fd9b9a..362a5e34a6be 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py @@ -1,14 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -<<<<<<< HEAD """Lazy QSA exports; importing geometry does not construct the cache stack.""" from importlib import import_module _EXPORTS = { "QSATrtllmAttention": "backend", - "QSAAttentionMetadata": "metadata", + "QSAMetadata": "metadata", "QSASparseMetadataParams": "params", "QSASparseParams": "params", } @@ -24,17 +23,3 @@ def __getattr__(name: str): __all__ = list(_EXPORTS) -======= -from .backend import QSATrtllmAttention -from .cache_manager import QSAMambaHybridCacheManagerV2 -from .metadata import QSAMetadata -from .params import QSASparseMetadataParams, QSASparseParams - -__all__ = [ - "QSAMambaHybridCacheManagerV2", - "QSAMetadata", - "QSASparseMetadataParams", - "QSASparseParams", - "QSATrtllmAttention", -] ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index 2d96145d1db8..3f65792ec96c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -22,31 +22,8 @@ class QSAMetadata(SparseAlgorithmExtension, MTPIndexShareExtension): def allocate(self, buffers) -> None: """Allocate the fixed-shape QSA buffers updated before each graph replay.""" -<<<<<<< HEAD - 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__() - # 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): - raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManager") - if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): -======= md = self.metadata self.needs_speculative_snapshot = False # Draft-loop state; set by the speculative driver, read by the indexer. @@ -57,11 +34,10 @@ def __post_init__(self) -> 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, QSAMambaHybridCacheManagerV2): - raise TypeError("QSA sparse attention requires QSAMambaHybridCacheManagerV2") + if not isinstance(md.kv_cache_manager, Qwen4ExpHybridCacheManager): + raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManager") params = md.sparse_metadata_params if not isinstance(params, QSASparseMetadataParams): ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) raise TypeError("QSA sparse attention requires QSASparseMetadataParams") self.params = params self.has_local_layers = md.kv_cache_manager.qsa_position_layer_id is not None @@ -151,14 +127,10 @@ def _init_draft_block_table(self) -> None: forward and this one in the draft loop; ``prepare()`` refreshes both before every replay. """ -<<<<<<< HEAD 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 ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) if ( not isinstance(draft_manager, Qwen4ExpHybridCacheManager) or draft_manager.qsa_position_layer_id is None 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/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py deleted file mode 100644 index 1247ad716922..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ /dev/null @@ -1,5147 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math -import os -from abc import ABC, abstractmethod -from contextlib import contextmanager -from dataclasses import dataclass, replace -from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, - Optional, Protocol, Sequence, Tuple, Union, cast) - -import torch -import triton -import triton.language as tl - -if TYPE_CHECKING: - from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata - from tensorrt_llm._torch.pyexecutor.config_utils import Qwen4ExpPLECacheParams - from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig - from tensorrt_llm.sampling_params import SamplingParams - -from tensorrt_llm._device_copy import DeviceCopyBuffer -from tensorrt_llm._torch.disaggregation.resource.page import (MapperKind, - RoleLayout) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - _RESERVED_REQUEST_IDS, BlockReusePolicy, KVCacheManagerV2, Role) -from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ - KVCacheV2IterationStatsReport -from tensorrt_llm._torch.pyexecutor.llm_request import ( - ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) -from tensorrt_llm._torch.pyexecutor.resource_manager import ( - BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, - PoolConfiguration, get_pp_layers) -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, - nvtx_range, prefer_pinned, - torch_dtype_to_binding) -from tensorrt_llm.bindings.executor import KvCacheRetentionConfig -from tensorrt_llm.bindings.internal.batch_manager import ( - LinearAttentionMetadata, LinearCacheType) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, - BatchDesc, BufferConfig, - DataRole, KVCacheDesc) -from tensorrt_llm.runtime.kv_cache_manager_v2 import \ - KVCacheManagerConfig as KVCacheManagerConfigPy -from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, - SsmLayerConfig, - TokenIdExt, _KVCache) - -# Shared with the KV budget estimator so allocator and budgeting can never -# diverge on the sharding rule (config_utils is import-cycle-free). -from ..config_utils import is_glm5_next, is_kimi_linear, mamba_effective_tp_size - -GB = 1 << 30 - - -class _PrefixReuseDiagnostics(Protocol): - - def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: - ... - - def _get_num_reusable_tokens_before_pruning(self) -> int: - ... - - -# Replay kernels pad the token/window dimension to at least 16 for tensor-core -# tiles, so history sizes below 16 are no faster when not writing and do -# expensive checkpoint writes more often. This minimum tensor-core tile size is -# present in all tensor-core generations, with newer GPUs adding larger tile -# operations, not smaller ones. In the other direction, history sizes above 16 -# currently pad to 32, which is substantially slower with the current kernel -# design. Keep the default floor at 16 while still allowing larger T, which -# degenerates to checkpointing every step instead of rejecting the request. If -# larger T becomes common, we could consider padding larger histories to powers -# of 2 or multiples of 16, but this is untested. For smaller T, we could explore -# combining larger histories with new kernel designs that stay efficient when -# the window is only partly full. -MIN_REPLAY_HISTORY_SIZE = 16 - -_KDA_BETA_CACHE_ALIGNMENT_BYTES = 16 -# KVCM copies 16-byte payloads; the replay kernel reads only the first int32. -_KDA_REPLAY_COUNT_BYTES = 16 -_KDA_REPLAY_COUNT_ROLE = DataRole("prev_num_accepted_tokens") -_KDA_REPLAY_CONV_ROLES = ("kda_conv_q", "kda_conv_k", "kda_conv_v") -_KDA_DRAFT_RECORDS_ROLE = "kda_state_tok" - - -def _is_disagg_context_server(is_disagg: bool) -> bool: - """Whether this worker only serves the context phase of disagg requests.""" - return is_disagg and os.getenv("TRTLLM_DISAGG_ROLE") == "context" - - -def _kda_replay_layout( - num_spec: int, - num_heads: int, - head_dim: int, - conv_window: int, - draft_records: bool = False -) -> Dict[str, Tuple[torch.dtype, List[int]]]: - """Per-slot physical shapes of one layer's KDA replay caches. - - Conv caches hold raw bf16 projection inputs (exact in bf16); k / v / beta - are bf16 and the gate stays fp32 because it feeds the cumulative-gate - exponents of the replayed chunk. Beta rows keep the 16-byte CuTe - alignment by padding the head dimension. - - With ``draft_records``, the records of every draft of the last verify - round (``kda_state_tok``: per draft, the row innovations vn, then beta * k - and the decay, each fp32 [H, K] with V = K = ``head_dim``) take the place - of the k / g / v / beta caches. - """ - dim = num_heads * head_dim - layout = { - "kda_conv_q": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_conv_k": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_conv_v": (torch.bfloat16, [conv_window + num_spec, dim]), - } - if draft_records: - layout[_KDA_DRAFT_RECORDS_ROLE] = (torch.float32, - [3, num_spec, num_heads, head_dim]) - return layout - heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // - torch.bfloat16.itemsize) - aligned_heads = ((num_heads + heads_per_alignment - 1) // - heads_per_alignment) * heads_per_alignment - layout.update({ - "kda_k_cache": (torch.bfloat16, [num_spec, dim]), - "kda_g_cache": (torch.float32, [num_spec, dim]), - "kda_v_cache": (torch.bfloat16, [num_spec, dim]), - "kda_beta_cache": (torch.bfloat16, [num_spec, aligned_heads]), - }) - return layout - - -def _kda_replay_bytes_per_slot(layout: Dict[str, Tuple[torch.dtype, List[int]]], - num_local_layers: int) -> int: - """Replay bytes one recurrent-state slot holds across local layers.""" - if not layout or num_local_layers == 0: - return 0 - layer_bytes = sum(dtype.itemsize * math.prod(shape) - for dtype, shape in layout.values()) - return num_local_layers * layer_bytes + _KDA_REPLAY_COUNT_BYTES - - -def _allocate_kda_beta_cache( - shape: Tuple[int, ...], - device: Optional[torch.device], -) -> torch.Tensor: - """Allocate a logically shaped bf16 beta cache with aligned physical rows.""" - logical_num_heads = shape[-1] - heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // - torch.bfloat16.itemsize) - padded_num_heads = ((logical_num_heads + heads_per_alignment - 1) // - heads_per_alignment) * heads_per_alignment - # CuTe requires every nested cache view to be 16-byte aligned. Keep the - # logical head count while padding the physical stride (six bf16 heads are - # 12 bytes, so an unpadded row makes alternating views only 4-byte aligned). - return torch.zeros( - *shape[:-1], - padded_num_heads, - dtype=torch.bfloat16, - device=device, - )[..., :logical_num_heads] - - -def _get_num_cuda_graph_padding_dummy_slots( - spec_config: Optional["DecodingBaseConfig"], - max_batch_size: int, -) -> int: - """Return the number of persistent CUDA-graph padding dummy IDs. - - This is computed before ``ModelEngine`` exists and covers draft lengths - reachable at every batch size, including the zero-length acceptance-rate - fallback. ``compute_dynamic_draft_len_mapping`` is created - later and covers only configured CUDA-graph batch sizes, so it cannot size - this persistent ID set. - """ - if spec_config is None: - return 1 - - draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) - spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) - supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( - spec_dec_mode, "support_dynamic_draft_len") - and spec_dec_mode.support_dynamic_draft_len()) - if draft_len_schedule and supports_dynamic_draft_len: - runtime_draft_lengths = set() - first_uncovered_batch_size = 1 - for batch_size_threshold, draft_len in draft_len_schedule.items(): - if first_uncovered_batch_size > max_batch_size: - break - if batch_size_threshold >= first_uncovered_batch_size: - runtime_draft_lengths.add(draft_len) - first_uncovered_batch_size = batch_size_threshold + 1 - if first_uncovered_batch_size <= max_batch_size: - runtime_draft_lengths.add(0) - else: - max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 - max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", - 0) or 0) - is_linear_tree = getattr( - spec_config, - "is_linear_tree", - max_draft_len == max_total_draft_tokens, - ) - static_draft_len = (max_draft_len - if is_linear_tree else max_total_draft_tokens) - runtime_draft_lengths = {static_draft_len or 0} - - if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and - (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): - runtime_draft_lengths.add(0) - return len(runtime_draft_lengths) - - -class MambaRole: - """V2 buffer roles owned only by the hybrid Mamba manager.""" - - SSM_STATE = DataRole("ssm_state") - CONV_STATE = DataRole("conv_state") - PLE_NGRAM_CONTEXT = DataRole("ple_ngram_context") - PLE_CONV_STATE = DataRole("ple_conv_state") - - -def get_tensor_size_bytes(tensor): - """Calculate tensor size in bytes.""" - if isinstance(tensor, torch.Tensor): - return tensor.element_size() * tensor.nelement() - elif isinstance(tensor, list): - return sum(get_tensor_size_bytes(t) for t in tensor) - return 0 - - -# Mamba SSM stochastic-rounding Philox seed plumbing. -# -# Both the replay kernel and flashinfer's selective_state_update consume a -# `rand_seed` int64 tensor. Historically the non-replay paths created this -# tensor via `torch.randint(..., (1,))` on every decode step, which (a) is -# non-deterministic across runs, (b) allocates a fresh CUDA tensor per call -# and is therefore unfriendly to CUDA-graph capture, and (c) has no notion -# of cache-slot identity (so slot reuse cannot rotate the stream). -# -# The functions below produce per-slot int64 seeds via SplitMix64 finalization -# of (counter, slot, rank). Adjacent inputs yield uncorrelated outputs so -# consecutive cache slots and consecutive request counters do not leave -# structural fingerprints in the Philox input stream. All outputs live in -# (0, 2**62) so they avoid Philox's degenerate seed=0 case while staying -# within int64. -_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 -_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 -_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 -_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D -_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 -_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 - - -def _splitmix64(x: int) -> int: - """SplitMix64 finalizer; pure function, no torch.""" - x = (x + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 30) - x = (x * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 27) - x = (x * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 31) - return x & _MAMBA_SSM_UINT64_MASK - - -def _compute_deterministic_mamba_seed(counter: int, slot: int, - rank_offset: int) -> int: - """Deterministic int64 seed in (0, 2**62) from (counter, slot, rank). - - Pure function (no RNG, no torch.randint). Identical inputs across - process invocations produce identical outputs, which is what the - acceptance criteria require for cross-run reproducibility. - """ - folded = (_MAMBA_SSM_SEED_BASE + counter * _MAMBA_SSM_SEED_MIX_COUNTER + - slot * _MAMBA_SSM_SEED_MIX_SLOT + - rank_offset * _MAMBA_SSM_SEED_MIX_RANK) - folded &= _MAMBA_SSM_UINT64_MASK - value = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK - if value == 0: - value = 1 - return value - - -def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, - device: torch.device) -> torch.Tensor: - """Allocate a (cache_size,) int64 CUDA buffer of deterministic seeds. - - counter=0 at allocation time; per-slot reset on fresh request assignment - bumps a host-side counter and rewrites only the freshly-assigned slot. - """ - slot_seeds = [ - _compute_deterministic_mamba_seed(0, i, rank_offset) - for i in range(cache_size) - ] - return torch.tensor(slot_seeds, dtype=torch.int64, device=device) - - -def _host_slots_to_device(slots: List[int], - device: torch.device) -> torch.Tensor: - """Copy host slot ids to ``device`` without a host sync. - - A pageable copy would wait for the queued step. Each call stages through a - fresh pinned tensor, so no later host write can race the copy. - """ - return torch.tensor(slots, - dtype=torch.long, - pin_memory=device.type == "cuda" - and prefer_pinned()).to(device, non_blocking=True) - - -def _mamba_rank_offset(mapping: Mapping) -> int: - """Distinct seed offset per (tp_rank, pp_rank) so independent ranks - don't draw identical streams when not coordinated.""" - return (mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + - mapping.rank * 1_009) - - -def use_py_mamba_cache_manager() -> bool: - """Check if PythonMambaCacheManager should be forced (agg mode override). - - Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. - - Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager - with PythonMambaCacheManager inside instead of the configured manager. - Disagg mode is unaffected — its compatibility routing selects Mixed or - Cpp based on the transceiver configuration. - """ - return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' - - -class ReplayStateUpdateMetadata(NamedTuple): - """Shared tensors and fixed sizes for replay state updates.""" - prev_num_accepted_tokens: torch.Tensor - cache_buf_idx: torch.Tensor - replay_step_width: int - replay_history_size: int - - -def _advance_replay_state( - replay_metadata: ReplayStateUpdateMetadata, - state_indices: torch.Tensor, - accepted_tokens: torch.Tensor, - is_dummy_request: Optional[torch.Tensor] = None, -) -> None: - """Advance replay bookkeeping after a speculative generation step.""" - slots = state_indices.long() - accepted_tokens = accepted_tokens.to( - replay_metadata.prev_num_accepted_tokens.dtype) - prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] - wrote_checkpoint = (prev_num_accepted_tokens + - replay_metadata.replay_step_width - > replay_metadata.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, - accepted_tokens, - prev_num_accepted_tokens + accepted_tokens, - ) - cache_buf_idx = replay_metadata.cache_buf_idx[slots] - next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx) - if is_dummy_request is not None: - next_num_accepted_tokens = torch.where( - is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens, - ) - next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, - next_cache_buf_idx) - replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx - - -class BaseMambaCacheManager(ABC): - """Abstract interface for accessing mamba/recurrent state caches.""" - - @abstractmethod - def get_state_indices(self, *args, **kwargs) -> torch.Tensor: - """Return slot indices of each request. - - Shape: [max_batch_size] - """ - ... - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - """Return replay metadata tensors and fixed replay sizes.""" - return None - - @property - def use_kda_replay_update(self) -> bool: - """Whether KDA fused verification owns per-slot replay caches.""" - return getattr(self, "_use_kda_replay_update", False) - - @property - def keeps_kda_draft_records(self) -> bool: - """Whether the KDA replay caches come with the records of every draft - of the last verify round (``kda_state_tok``), from which a verify - kernel replays the accepted drafts onto the slot's state.""" - return False - - @abstractmethod - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - """Return conv states for specific layer. - - Shape: [slot_size, conv_dim, d_conv - 1] - """ - ... - - @abstractmethod - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - """Return SSM states for specific layer. - - Shape: [slot_size, num_heads, head_dim, d_state] - """ - ... - - @abstractmethod - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - ... - - @abstractmethod - def is_speculative(self) -> bool: - ... - - @abstractmethod - def mamba_layer_cache( - self, layer_idx: int - ) -> Union['PythonMambaCacheManager.State', - 'PythonMambaCacheManager.SpeculativeState', None]: - ... - - -class PythonMambaCacheManager(BaseResourceManager): - """Pure-Python mamba state manager with speculative decoding support. - - Manages only mamba states (conv + SSM) using PyTorch tensors on GPU. - Supports speculative decoding and disaggregated serving. - """ - - @dataclass(frozen=True, kw_only=True) - class State: - """Base state container for Mamba cache.""" - conv: torch.Tensor - # Represents last state, or "two back" if prev_num_accepted_tokens is > 0 - temporal: torch.Tensor - - # Subclasses override to list fields shared across layers (not indexed by layer). - _SHARED_FIELDS = frozenset() - - def at_layer_idx(self, layer: int): - kwargs = {} - for k, v in vars(self).items(): - if v is None: - kwargs[k] = None - elif k in self._SHARED_FIELDS: - kwargs[k] = v - else: - kwargs[k] = v[layer] - return type(self)(**kwargs) - - @property - def has_kda_replay_caches(self) -> bool: - """Whether the Kimi K3 KDA fused-verify replay caches are - allocated. The base (non-speculative) state never has them; the - speculative state overrides this.""" - return False - - @dataclass(frozen=True, kw_only=True) - class SpeculativeState(State): - """Speculative state with intermediate states for draft tokens. - - Supports three SSM update paths (only one set of tensors is - allocated): - - Legacy: caches full intermediate SSM states (intermediate_ssm) - - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) - - KDA replay: per-slot draft-token caches consumed by the fused - ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted - drafts and commits states in place (kda_conv_*, kda_*_cache); or, - for a verify kernel that keeps the drafts' records, kda_conv_* and - kda_state_tok in place of the kda_*_cache - """ - _SHARED_FIELDS = frozenset({ - "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" - }) - - # Allocated for the legacy and Mamba2-replay paths; None for the - # KDA replay path (the kernel commits conv windows in place). - intermediate_conv_window: torch.Tensor | None = None - - # Legacy path: full intermediate SSM states at each step - intermediate_ssm: torch.Tensor | None = None - - # KDA replay path (fused multi-token verify, kimi_linear). - # Pool invariant under this path: `temporal` holds the state after - # the LAST GOLDEN token; the accepted drafts recorded in - # prev_num_accepted_tokens are pending in these caches and are - # replayed by the kernel at the start of the next verify round. - # Extended conv caches [layers, slots, dim, (W-1) + num_spec] bf16 - # (raw projection inputs, so bf16 is exact), dim-contiguous (stride(dim) == 1): columns [0, W-1) are the - # committed raw-input window, tail columns the pending drafts' raw - # inputs. - kda_conv_q: torch.Tensor | None = None - kda_conv_k: torch.Tensor | None = None - kda_conv_v: torch.Tensor | None = None - # Post-processed per-draft quantities for replay: k - # [layers, slots, num_spec, H*K] bf16, gate [.., num_spec, H*K] fp32 - # (feeds the cumulative-gate exponents), v [.., num_spec, H*V] bf16, - # beta [.., num_spec, H] bf16. q is recomputed from the conv caches. - kda_k_cache: torch.Tensor | None = None - kda_g_cache: torch.Tensor | None = None - kda_v_cache: torch.Tensor | None = None - kda_beta_cache: torch.Tensor | None = None - # Optional: the records of every draft of the last verify round, - # [slots, 3, num_spec, H, K] fp32: the row innovations vn [num_spec, - # H, V], then beta * k and the decay [num_spec, H, K] (V = K). The - # SSM pool holds the state after the round's first, non-draft token; - # when the round accepted n > 0 drafts (prev_num_accepted_tokens), a - # verify kernel starts the next round by replaying the first n - # records onto it, S = decay * S + vn (beta * k), instead of - # replaying them from the caches above, which are then not - # allocated. A slot reset for a new request has n = 0, so a previous - # owner's records are never read. - kda_state_tok: torch.Tensor | None = None - - # Replay path: compact double-buffered cache - # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). - # 0 means temporal saved state is actually the last state, not two back. - prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers - cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers - # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so - # CUDA graph replay uses fresh SR draws without allocating RNG tensors. - # (cache,) int64 - shared across layers - mamba_ssm_rand_seed: torch.Tensor | None = None - old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) - old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) - # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. - old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 - old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 - - @property - def has_kda_replay_caches(self) -> bool: - """True on the KDA fused-verify replay path: the extended conv - caches are allocated, with the per-draft replay caches of the - fused ``trtllm::kda_mtp_decode`` verify or, in their place, the - drafts' records (``kda_state_tok``).""" - return self.kda_conv_q is not None - - def commit_conv_window(self, slot_indices: torch.Tensor, - conv_pool: torch.Tensor) -> None: - """Seed the KDA replay caches from the live ``W - 1`` conv pool.""" - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ - copy_kda_replay_conv_window - - copy_kda_replay_conv_window( - conv_pool, - self.kda_conv_q, - self.kda_conv_k, - self.kda_conv_v, - slot_indices, - ) - - def __init__( - self, - d_state: int, - d_conv: int, - num_heads: int, - n_groups: int, - head_dim: int, - num_layers: int, - max_batch_size: int, - spec_state_size: int, - mapping: Mapping, - dtype: torch.dtype, - ssm_cache_dtype: torch.dtype, - layer_mask: Optional[List[bool]] = None, - speculative_num_draft_tokens: Optional[int] = None, - model_type: str = "nemotron_hybrid", - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - ) -> None: - - self.mamba_ssm_cache_dtype = ssm_cache_dtype - self.speculative_num_draft_tokens = speculative_num_draft_tokens - self.spec_state_size = spec_state_size - self._use_replay_state_update = use_replay_state_update - # KDA replay path (kimi_linear fused multi-token verify). Mutually - # exclusive with use_replay_state_update; requires speculative mode. - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - if self._use_kda_replay_update: - assert not use_replay_state_update, ( - "kda_replay_num_spec and use_replay_state_update are " - "mutually exclusive") - assert speculative_num_draft_tokens is not None, ( - "KDA replay caches require speculative decoding") - assert kda_replay_num_spec == speculative_num_draft_tokens, ( - f"KDA replay cache width ({kda_replay_num_spec}) must match " - f"the draft length ({speculative_num_draft_tokens}): the " - "fused verify kernel is compiled with a static NUM_SPEC") - self.replay_history_size: Optional[int] = None - self.replay_step_width: Optional[int] = None - # When True, allocate the per-slot Philox seed buffer even outside - # the replay path so the non-replay flashinfer SR kernel reads a - # persistent deterministic seed instead of a per-call torch.randint. - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - # Host-side counter bumped per fresh cache-slot assignment. Combined - # with slot index and rank offset to produce reproducible per-slot - # seed values. Starts at 0 so the post-init "reset" stream is - # disjoint from the counter=0 stream used at allocation time. - self._seed_request_counter = 0 - - # get tp size - tp_size = mamba_effective_tp_size(mapping) - - # derive mamba parameters for conv and ssm states - d_inner = head_dim * num_heads - conv_dim = d_inner + 2 * n_groups * d_state - nheads = num_heads - - # check that can be partitioned - assert nheads % tp_size == 0, "nheads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - - # partition conv_dim and nheads - d_inner_local = d_inner // tp_size - ng_ds_local = n_groups * d_state // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - d_inner = d_inner // tp_size - - # Per-section dims for conv_state. - # Qwen3-Next: [Q | K | V] = [ng*ds, ng*ds, d_inner] - # Nemotron_hybrid: [x | B | C] = [d_inner, ng*ds, ng*ds] - if model_type == "qwen3_next": - self.conv_section_dims = [ng_ds_local, ng_ds_local, d_inner_local] - elif model_type == "nemotron_hybrid": - self.conv_section_dims = [d_inner_local, ng_ds_local, ng_ds_local] - else: - raise ValueError(f"Unsupported model type: {model_type}") - - # conv and ssm states device - device = torch.device("cuda") - - pp_layers, num_layers = get_pp_layers( - num_layers, - mapping, - layer_mask=layer_mask, - ) - num_local_layers = len(pp_layers) - self.mamba_layer_offsets = { - idx: offset - for offset, idx in enumerate(pp_layers) - } - - conv_state_shape = (conv_dim, d_conv - 1) - ssm_state_shape = (nheads, head_dim, d_state) - - # create mamba conv and ssm states - conv_states = torch.zeros( - size=(num_local_layers, max_batch_size) + conv_state_shape, - dtype=dtype, - device=device, - ) - - ssm_states = torch.zeros( - size=(num_local_layers, max_batch_size) + ssm_state_shape, - dtype=self.mamba_ssm_cache_dtype, - device=device, - ) - - # Per-slot Philox seeds. Allocated whenever stochastic rounding can - # fire, even outside the replay path and even when MTP/spec is off, - # so all consumers (replay kernel, MTP non-replay flashinfer, - # non-MTP flashinfer) read a persistent deterministic seed from the - # cache manager instead of calling torch.randint per forward. - self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None - if (self._use_replay_state_update - or self._mamba_ssm_stochastic_rounding): - self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - max_batch_size, self._seed_rank_offset, device) - - # create state container - if speculative_num_draft_tokens is not None: - T = speculative_num_draft_tokens + 1 - self.replay_step_width = T - - # Conv intermediate cache — legacy and Mamba2-replay paths only. - # The KDA replay kernel commits conv windows in place, so the - # per-step window scratch is not needed. - intermediate_conv_window_cache = None - if not self._use_kda_replay_update: - intermediate_conv_window_cache = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - conv_state_shape, - dtype=dtype, - device=device, - ) - - # SSM speculative cache — path-specific tensors - spec_kwargs = {} - # Share the manager-level seed buffer through SpeculativeState - # so the MTP path can still read it via layer_cache. - if self._mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed - if self._use_kda_replay_update: - # KDA replay caches for the fused multi-token verify kernel - # (trtllm::kda_mtp_decode). Per cache slot (persistent - # across rounds — the next round replays from them), unlike - # the batch-row-indexed intermediate buffers. - M = self._kda_replay_num_spec - # conv_dim covers the [q | k | v] short-conv sections; the - # kernel consumes them as three dim-contiguous caches. - assert conv_dim % 3 == 0, ( - "KDA replay caches expect the [q | k | v] conv-state " - "sectioning (3 equal sections)") - section_dim = conv_dim // 3 - # d_conv is KDA's convolution width; the live pool stores - # its W - 1 committed input columns. - w_kernel = d_conv - extended_s = w_kernel - 1 + M - - def _dim_contiguous_conv_cache(): - # raw bf16 projection inputs: bf16 is exact and halves - # the per-draft replay traffic of the verify kernel - return torch.zeros(num_local_layers, - max_batch_size, - extended_s, - section_dim, - dtype=torch.bfloat16, - device=device).transpose(-1, -2) - - spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=torch.int32, device=device) - spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_k_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.bfloat16, - device=device) - spec_kwargs['kda_g_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.float32, - device=device) - spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.bfloat16, - device=device) - spec_kwargs['kda_beta_cache'] = _allocate_kda_beta_cache( - (num_local_layers, max_batch_size, M, nheads), device) - ssm_spec_cache = [ - spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], - spec_kwargs['kda_conv_v'], spec_kwargs['kda_k_cache'], - spec_kwargs['kda_g_cache'], spec_kwargs['kda_v_cache'], - spec_kwargs['kda_beta_cache'] - ] - spec_path_label = "kda-replay" - elif self._use_replay_state_update: - assert n_groups % tp_size == 0, \ - "replay state update requires n_groups divisible by tp_size" - n_groups_per_rank = n_groups // tp_size - self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) - - # Compact replay cache. - spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=torch.int32, device=device) - spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, - dtype=torch.int32, - device=device) - spec_kwargs['old_x'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - nheads, - head_dim, - dtype=dtype, - device=device) - spec_kwargs['old_B'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - n_groups_per_rank, - d_state, - dtype=dtype, - device=device) - spec_kwargs['old_dt'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - dtype=torch.float32, - device=device) - spec_kwargs['old_dA_cumsum'] = torch.zeros( - num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - dtype=torch.float32, - device=device) - ssm_spec_cache = [ - spec_kwargs['old_x'], spec_kwargs['old_B'], - spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] - ] - spec_path_label = "replay" - else: - # Legacy: full intermediate SSM states at each step - spec_kwargs['intermediate_ssm'] = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - ssm_state_shape, - dtype=self.mamba_ssm_cache_dtype, - device=device) - ssm_spec_cache = [spec_kwargs['intermediate_ssm']] - spec_path_label = "legacy" - - self.mamba_cache = self.SpeculativeState( - conv=conv_states, - temporal=ssm_states, - intermediate_conv_window=intermediate_conv_window_cache, - **spec_kwargs, - ) - - logger.info( - f"Mamba Cache ({spec_path_label}) is allocated. " - f"max_mamba_cache_size: {max_batch_size}, " - f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " - f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " - f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " - "intermediate_conv_window_cache size: " - f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" - ) - else: - self.mamba_cache = self.State( - conv=conv_states, - temporal=ssm_states, - ) - - logger.info( - f"Mamba Cache is allocated. " - f"max_mamba_cache_size: {max_batch_size}, " - f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " - f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB" - ) - - # mamba cache available blocks - self.mamba_cache_free_blocks = [i for i in range(max_batch_size)] - - # mamba cache index, maps request_id -> state indices - self.mamba_cache_index: Dict[int, int] = {} - self._dummy_request_ids: set[int] = set() - # Batch-order mask aligned with state_indices; duplicate dummy request - # IDs mark every batch row even when they share one cache slot. - self._dummy_request_mask = torch.zeros(max_batch_size, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros(max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned()) - - # Permanent slot shared by every CUDA-graph padding sentinel id - # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per - # draft length). Pool sizing must include +1 headroom for this; - # see MixedMambaHybridCacheManager. - self._padding_slot: int = self.mamba_cache_free_blocks.pop() - - # Reserved slot for the attention-DP padding dummy. Keep it out of the - # free pool so dummy insertion never consumes real-request capacity. - self._attention_dp_dummy_slot: Optional[int] = ( - self.mamba_cache_free_blocks.pop() - if mapping.enable_attention_dp else None) - - # save intermediate state indices for requests - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device=device) - - # Physical tensor rows include reserved dummy slots. Resource capacity - # is the number of real request slots remaining after those - # reservations. - self._max_batch_size = max_batch_size - self._max_resource_count = len(self.mamba_cache_free_blocks) - - def get_max_resource_count(self) -> int: - """Return the maximum number of real requests that can be cached.""" - return self._max_resource_count - - def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: - """Return the prefix of *context_requests* that fits in the - available Mamba cache blocks. Requests that already have a - cached block do not consume a free block.""" - free = len(self.mamba_cache_free_blocks) - result = [] - for r in context_requests: - if r.py_request_id in self.mamba_cache_index: - result.append(r) - elif free > 0: - result.append(r) - free -= 1 - else: - break - return result - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - """For Mamba cache manager, we always need one slot per request.""" - return 1 - - @torch.inference_mode() - def _prepare_mamba_cache_blocks(self, request_ids: List[int]): - for r in request_ids: - if r in self.mamba_cache_index: - continue - if len(self.mamba_cache_free_blocks) == 0: - raise RuntimeError("run out of mamba cache blocks") - block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - # Fresh request: no drafts pending in the replay caches. - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - if self._mamba_ssm_rand_seed is not None: - # Deterministic per-slot rotation on fresh assignment. - # `block` is pulled from mamba_cache_free_blocks, which - # excludes _padding_slot by construction (see __init__), - # so padding sentinels never reach this branch. - self._seed_request_counter += 1 - self._mamba_ssm_rand_seed[block] = ( - _compute_deterministic_mamba_seed( - self._seed_request_counter, block, - self._seed_rank_offset)) - - @torch.inference_mode() - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - """Seed the fused-verify KDA replay conv caches from the conv pool. - - On a disaggregated generation server the ctx->gen transfer populates - only the base ``conv`` / ``temporal`` pools (see - ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` - replay caches consumed by ``trtllm::kda_mtp_decode`` are normally - seeded by a *local* prefill/decode via - ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a - previous occupant's window) for a transferred request — corrupting - the recurrent state from the first verify step onward. Call this - after the state transfer completes, before the first generation - forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row - stores the committed ``W - 1`` raw inputs directly. Pending-draft - scratch is cleared (no drafts are pending for a freshly transferred - request). - """ - if not (self._use_kda_replay_update - and isinstance(self.mamba_cache, self.SpeculativeState)): - return - blocks = [ - self.mamba_cache_index[rid] for rid in request_ids - if rid in self.mamba_cache_index - ] - if not blocks: - return - conv = self.mamba_cache.conv # [L, slots, 3D, W - 1] - idx = torch.tensor(sorted(set(blocks)), - dtype=torch.long, - device=conv.device) - d = conv.shape[2] // 3 - committed = conv.shape[3] - cs = conv.index_select(1, idx) - for cache, section in ( - (self.mamba_cache.kda_conv_q, cs[:, :, :d]), - (self.mamba_cache.kda_conv_k, cs[:, :, d:2 * d]), - (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d:]), - ): - # cache: [L, slots, D, committed + num_spec]; zero the draft - # tail columns and seed the committed window in one copy. - seeded = torch.zeros( - (cache.shape[0], idx.numel()) + cache.shape[2:], - dtype=cache.dtype, - device=cache.device) - seeded[:, :, :, :committed] = section.to(cache.dtype) - cache.index_copy_(1, idx, seeded) - for buf in (self.mamba_cache.kda_k_cache, self.mamba_cache.kda_g_cache, - self.mamba_cache.kda_v_cache, - self.mamba_cache.kda_beta_cache): - buf.index_fill_(1, idx, 0) - self.mamba_cache.prev_num_accepted_tokens[idx] = 0 - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - context_ids = [ - i.py_request_id for i in scheduled_batch.context_requests - ] - generation_ids = [ - i.py_request_id for i in scheduled_batch.generation_requests - ] - request_ids = context_ids + generation_ids - self._prepare_mamba_cache_blocks(request_ids) - self._refresh_dummy_request_mask([req.is_dummy for req in requests]) - - def _is_padding_sentinel(self, request_id: int) -> bool: - # cuda_graph_runner caches one dummy per runtime_draft_len value - # (see _get_padded_batch), so any id in the range of dummy request IDs - # may be live concurrently. - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl = self.speculative_num_draft_tokens or 0 - return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= - CUDA_GRAPH_DUMMY_REQUEST_ID) - - @torch.inference_mode() - def add_dummy_requests(self, request_ids: List[int], **kwargs): - # Sentinels alias to the shared _padding_slot; non-sentinel - # dummies (warmup, attention-DP idle padding) get their own - # slot and are freed individually. - if not request_ids: - return - self._dummy_request_ids.update(request_ids) - for r in request_ids: - if r in self.mamba_cache_index: - block = self.mamba_cache_index[r] - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - continue - if self._is_padding_sentinel(r): - block = self._padding_slot - elif (r == ATTENTION_DP_DUMMY_REQUEST_ID - and self._attention_dp_dummy_slot is not None): - block = self._attention_dp_dummy_slot - else: - if len(self.mamba_cache_free_blocks) == 0: - raise RuntimeError("run out of mamba cache blocks") - block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - - def free_resources(self, request: LlmRequest): - request_id = request.py_request_id - if request_id not in self.mamba_cache_index: - return - self._dummy_request_ids.discard(request_id) - block = self.mamba_cache_index.pop(request_id) - # Reserved slots must not re-enter the real-request free pool. - if block != self._padding_slot and \ - block != self._attention_dp_dummy_slot: - self.mamba_cache_free_blocks.append(block) - - def get_state_indices(self, request_ids: List[int], - is_padding: List[bool]) -> List[int]: - assert len(request_ids) == len(is_padding) - indices = [self.mamba_cache_index[rid] for rid in request_ids] - is_dummy = [ - rid in self._dummy_request_ids or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.as_tensor(is_dummy, dtype=torch.bool)) - mask_staging = torch.empty_like(self._dummy_request_mask_host, - pin_memory=prefer_pinned()) - mask_staging.copy_(self._dummy_request_mask_host) - self._dummy_request_mask.copy_(mask_staging, non_blocking=True) - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).conv - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).temporal - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if not isinstance(self.mamba_cache, self.SpeculativeState): - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).intermediate_ssm - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - """Get intermediate conv states for speculative decoding.""" - if not isinstance(self.mamba_cache, self.SpeculativeState): - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx( - layer_offset).intermediate_conv_window - - def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_x[layer_offset] - - def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_B[layer_offset] - - def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_dt[layer_offset] - - def get_replay_old_dA_cumsum(self, - layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_dA_cumsum[layer_offset] - - def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - return self.mamba_cache.cache_buf_idx - - def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - return self.mamba_cache.prev_num_accepted_tokens - - def is_speculative(self) -> bool: - return isinstance(self.mamba_cache, self.SpeculativeState) - - def mamba_layer_cache(self, - layer_idx: int) -> Union[State, SpeculativeState]: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.mamba_ssm_cache_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Return the persistent (cache_size,) int64 Philox seed buffer or - None when stochastic rounding is not active for this manager. - - Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. - Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a - view that matches the consuming kernel's expected shape. - """ - return self._mamba_ssm_rand_seed - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - @property - def use_kda_replay_update(self) -> bool: - return self._use_kda_replay_update - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - if (not self._use_replay_state_update - or not isinstance(self.mamba_cache, self.SpeculativeState) - or self.mamba_cache.prev_num_accepted_tokens is None - or self.mamba_cache.cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=( - self.mamba_cache.prev_num_accepted_tokens), - cache_buf_idx=self.mamba_cache.cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def shutdown(self): - """Release tensor memory.""" - # Clear mamba cache states - empty = torch.tensor([]) - - def _drop(tensor): - return empty if tensor is not None else None - - if isinstance(self.mamba_cache, self.SpeculativeState): - self.mamba_cache = self.SpeculativeState( - conv=empty, - temporal=empty, - intermediate_conv_window=_drop( - self.mamba_cache.intermediate_conv_window), - intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), - prev_num_accepted_tokens=_drop( - self.mamba_cache.prev_num_accepted_tokens), - cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx), - mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed), - old_x=_drop(self.mamba_cache.old_x), - old_B=_drop(self.mamba_cache.old_B), - old_dt=_drop(self.mamba_cache.old_dt), - old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), - kda_conv_q=_drop(self.mamba_cache.kda_conv_q), - kda_conv_k=_drop(self.mamba_cache.kda_conv_k), - kda_conv_v=_drop(self.mamba_cache.kda_conv_v), - kda_k_cache=_drop(self.mamba_cache.kda_k_cache), - kda_g_cache=_drop(self.mamba_cache.kda_g_cache), - kda_v_cache=_drop(self.mamba_cache.kda_v_cache), - kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), - ) - else: - self.mamba_cache = self.State(conv=empty, temporal=empty) - - torch.cuda.empty_cache() - - @torch.compile(options={"max-autotune": True}) - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - accepted_leaf_positions: Optional[torch.Tensor] = None): - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - num_accepted_draft_tokens = num_accepted_tokens[ - num_contexts:num_contexts + num_gens] - 1 - # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. - accepted_positions = (accepted_leaf_positions if accepted_leaf_positions - is not None else num_accepted_draft_tokens) - state_indices_d = state_indices[num_contexts:num_contexts + num_gens] - src_state_indices = self.intermediate_state_indices[:num_gens] - - if self._use_kda_replay_update: - # KDA replay: the fused verify kernel already committed the SSM - # state and conv windows (in place, after the golden token) and - # cached this round's drafts. All that remains is recording how - # many of those drafts the sampler accepted, so the next round's - # kernel launch replays exactly that prefix. - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - prev = self.mamba_cache.prev_num_accepted_tokens - current = prev[state_indices_d] - accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) - prev[state_indices_d] = torch.where(is_dummy_request, current, - accepted) - return - - if self._use_replay_state_update: - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - replay_metadata = self.get_replay_state_update_metadata() - assert replay_metadata is not None - _advance_replay_state( - replay_metadata, - state_indices_d, - num_accepted_tokens[num_contexts:num_contexts + num_gens], - is_dummy_request, - ) - else: - # Legacy: copy accepted SSM state from intermediate cache. - ssm_states = self.mamba_cache.temporal - intermediate_ssm_cache = self.mamba_cache.intermediate_ssm - accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, - accepted_positions] - ssm_states[:, state_indices_d, :] = accepted_ssm_state - - # Conv: both paths save all intermediate conv windows, carry over the accepted one. - conv_states = self.mamba_cache.conv - intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window - accepted_conv_state = intermediate_conv_window_cache[:, - src_state_indices, - accepted_positions] - conv_states[:, state_indices_d, :] = accepted_conv_state - - -class MambaCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Facade for standalone mamba state management (no KV cache). - - Delegates to PythonMambaCacheManager. - """ - - def __init__( - self, - d_state: int, - d_conv: int, - num_heads: int, - n_groups: int, - head_dim: int, - num_layers: int, - max_batch_size: int, - spec_state_size: int, - mapping: Mapping, - dtype: torch.dtype, - ssm_cache_dtype: torch.dtype, - layer_mask: Optional[List[bool]] = None, - stream: Optional[torch.cuda.Stream] = None, - speculative_num_draft_tokens: Optional[int] = None, - model_type: str = "nemotron_hybrid", - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - ) -> None: - max_num_sequences = max_batch_size * mapping.pp_size - - self._impl = PythonMambaCacheManager( - d_state=d_state, - d_conv=d_conv, - num_heads=num_heads, - n_groups=n_groups, - head_dim=head_dim, - num_layers=num_layers, - max_batch_size=max_num_sequences, - spec_state_size=spec_state_size, - mapping=mapping, - dtype=dtype, - ssm_cache_dtype=ssm_cache_dtype, - layer_mask=layer_mask, - speculative_num_draft_tokens=speculative_num_draft_tokens, - model_type=model_type, - use_replay_state_update=use_replay_state_update, - mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - kda_replay_num_spec=kda_replay_num_spec, - ) - - def get_max_resource_count(self) -> int: - return self._impl.get_max_resource_count() - - def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: - return self._impl.filter_ctx_requests_by_capacity(context_requests) - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - return self._impl.get_needed_resource_to_completion(request) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - self._impl.prepare_resources(scheduled_batch) - - def free_resources(self, request: LlmRequest): - self._impl.free_resources(request) - - def add_dummy_requests(self, request_ids: List[int], **kwargs): - self._impl.add_dummy_requests(request_ids, **kwargs) - - def get_state_indices( - self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None - ) -> Union[torch.Tensor, List[int]]: - return self._impl.get_state_indices(request_ids, is_padding) - - @property - def mamba_cache_free_blocks(self) -> List[int]: - return self._impl.mamba_cache_free_blocks - - @property - def mamba_cache_index(self) -> Dict[int, int]: - return self._impl.mamba_cache_index - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self._impl.get_conv_states(layer_idx) - - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self._impl.get_ssm_states(layer_idx) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self._impl.get_mamba_ssm_cache_dtype() - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Delegate to the underlying Python manager.""" - return self._impl.get_mamba_ssm_rand_seed() - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - @property - def use_kda_replay_update(self) -> bool: - return self._impl.use_kda_replay_update - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - return self._impl.get_replay_state_update_metadata() - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_intermediate_ssm_states(layer_idx) - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_intermediate_conv_states(layer_idx) - - def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_x(layer_idx) - - def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_B(layer_idx) - - def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_dt(layer_idx) - - def get_replay_old_dA_cumsum(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_dA_cumsum(layer_idx) - - def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: - return self._impl.get_replay_cache_buf_idx() - - def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: - return self._impl.get_replay_prev_num_accepted_tokens() - - def is_speculative(self) -> bool: - return self._impl.is_speculative() - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - return self._impl.mamba_layer_cache(layer_idx) - - def shutdown(self): - self._impl.shutdown() - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - accepted_leaf_positions: Optional[torch.Tensor] = None): - # Non-speculative configs don't allocate intermediate state; the - # promotion is a clean no-op. - if not self._impl.is_speculative(): - return - self._impl.update_mamba_states(attn_metadata, num_accepted_tokens, - state_indices, accepted_leaf_positions) - - -class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Shared interface and state plumbing for hybrid Mamba managers. - - Concrete storage, state-index, and resource lifecycles remain owned by the - Cpp and V2 implementations. - """ - - _supports_additional_snapshot_offsets = False - - @property - def cuda_state_indices(self) -> torch.Tensor: - return self._state_indices_buffer.tensor - - @property - def _dummy_request_mask(self) -> Optional[torch.Tensor]: - if self._dummy_request_mask_buffer is None: - return None - return self._dummy_request_mask_buffer.tensor - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if spec_config is None or self.local_num_mamba_layers == 0: - return - - tokens_per_gen_step = spec_config.tokens_per_gen_step - if not self._use_replay_state_update: - self.intermediate_ssm_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, - tokens_per_gen_step - ] + self.ssm_state_shape, - dtype=self.ssm_state_dtype, - device="cuda", - ) - self.intermediate_conv_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step - ] + self.conv_state_shape, - dtype=self.conv_state_dtype, - device="cuda", - ) - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - - def _allocate_pool_replay_buffers( - self, - spec_config, - cache_size: int, - device: Optional[torch.device], - ) -> bool: - """Allocate replay tensors shared by the Cpp and V2 state pools.""" - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask_buffer = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - if (self.local_num_mamba_layers == 0 - or (not self._use_replay_state_update - and not self._mamba_ssm_stochastic_rounding)): - return False - - assert device is not None - self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - cache_size, self._seed_rank_offset, device) - if spec_config is None or not self._use_replay_state_update: - return False - - history_size = self.replay_history_size - assert history_size is not None - nheads, head_dim, d_state = self.ssm_state_shape - common_shape = [self.local_num_mamba_layers, cache_size, 2] - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.cache_buf_idx = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.old_x = torch.zeros( - common_shape + [history_size, nheads, head_dim], - dtype=self.conv_state_dtype, - device=device, - ) - self.old_B = torch.zeros( - common_shape + [history_size, self._n_groups_per_rank, d_state], - dtype=self.conv_state_dtype, - device=device, - ) - self.old_dt = torch.zeros( - common_shape + [nheads, history_size], - dtype=torch.float32, - device=device, - ) - self.old_dA_cumsum = torch.zeros( - common_shape + [nheads, history_size], - dtype=torch.float32, - device=device, - ) - return True - - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - if self._dummy_request_mask is None: - return - - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.tensor(is_dummy, dtype=torch.bool)) - self._dummy_request_mask_buffer.update(self._dummy_request_mask_host) - - def _reset_context_mamba_slots(self, num_contexts: int) -> None: - if num_contexts == 0: - return - - context_slots = self.cuda_state_indices[:num_contexts].long() - if (self._use_replay_state_update - and self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): - self.prev_num_accepted_tokens.index_fill_(0, context_slots, 0) - self.cache_buf_idx.index_fill_(0, context_slots, 0) - if self.old_x is not None: - self.old_x.index_fill_(1, context_slots, 0) - if self.old_B is not None: - self.old_B.index_fill_(1, context_slots, 0) - if self.old_dt is not None: - self.old_dt.index_fill_(1, context_slots, 0) - if self.old_dA_cumsum is not None: - self.old_dA_cumsum.index_fill_(1, context_slots, 0) - - if self.mamba_ssm_rand_seed is None: - return - self._seed_request_counter += 1 - counter = self._seed_request_counter - rank_offset = self._seed_rank_offset - host_slots = self._host_state_indices[:num_contexts].tolist() - new_seeds = [ - _compute_deterministic_mamba_seed(counter, slot, rank_offset) - for slot in host_slots - ] - self.mamba_ssm_rand_seed[context_slots] = torch.tensor( - new_seeds, dtype=torch.int64, - pin_memory=prefer_pinned()).to(self.mamba_ssm_rand_seed.device, - non_blocking=True) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.enable_block_reuse: - for request in requests: - request.expect_snapshot_points = [] - return - - state_config = self.kv_cache_config.mamba_state_config - interval = state_config.periodic_snapshot_interval - for request in requests: - snapshot_points = set() - if interval is not None and interval > 0: - snapshot_points.update( - range(interval, request.prompt_len + 1, interval)) - if self._supports_additional_snapshot_offsets: - for offset in state_config.additional_snapshot_offsets_from_start: - if offset <= request.prompt_len: - snapshot_points.add(offset) - for offset in state_config.additional_snapshot_offsets_from_end: - point = request.prompt_len - offset - if point > 0: - snapshot_points.add(point) - request.expect_snapshot_points = sorted(snapshot_points) - - def is_speculative(self) -> bool: - return self.spec_config is not None - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_ssm_states is None: - return None - return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_conv_states is None: - return None - return self.intermediate_conv_states[ - self.mamba_layer_offsets[layer_idx]] - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - conv = self.get_conv_states(layer_idx) - ssm = self.get_ssm_states(layer_idx) - if self.spec_config is not None: - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = {} - if self.mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed - if self._use_replay_state_update: - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] - spec_kwargs['cache_buf_idx'] = self.cache_buf_idx - spec_kwargs['prev_num_accepted_tokens'] = ( - self.prev_num_accepted_tokens) - else: - spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ - layer_offset] - return PythonMambaCacheManager.SpeculativeState( - conv=conv, - temporal=ssm, - intermediate_conv_window=self. - intermediate_conv_states[layer_offset], - **spec_kwargs, - ) - return PythonMambaCacheManager.State(conv=conv, temporal=ssm) - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', - None) - cache_buf_idx = getattr(self, 'cache_buf_idx', None) - if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=prev_num_accepted_tokens, - cache_buf_idx=cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.ssm_state_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - return getattr(self, 'mamba_ssm_rand_seed', None) - - -def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: - """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" - pool_size = max_batch_size - # One permanent slot is shared by every CUDA-graph padding sentinel. - pool_size += 1 - if mapping.enable_attention_dp: - # Attention-DP can insert a transient dummy request on an otherwise - # idle rank. Keep this headroom internal so scheduler-visible - # max_batch_size still limits real requests. - pool_size += 1 - return pool_size - - -class MixedMambaHybridCacheManager(KVCacheManager, MambaCacheManager, - MambaHybridCacheManager): - """Hybrid cache manager combining separate KVCacheManager and MambaCacheManager. - - Manages KV cache and mamba states in independent pools, with support of - speculative decoding and disaggregated serving. - Does not support block reuse / prefix caching for mamba states. - """ - - # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the - # cache manager __init__ coerces a positive spec recompute tail to a full - # re-prefill (see _has_recurrent_state on the attention-only base classes). - _has_recurrent_state = True - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - # kv cache parameters - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - layer_mask: List[bool], - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - max_num_tokens: int = 8192, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - is_estimating_kv_cache: bool = False, - execution_stream: Optional[torch.cuda.Stream] = None, - model_type: str = "nemotron_hybrid", - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - # Per-pool configurations forwarded to the C++ KVCacheManager ctor. - # Lets a single manager host pools with mixed shapes (e.g. Gemma4 - # hybrid attention). See KVCacheManager.__init__. - pool_configurations: Optional[List[PoolConfiguration]] = None, - ) -> None: - - # mamba hybrid cache requires block reuse to be disabled in KV cache config - assert not kv_cache_config.enable_block_reuse, ( - "mamba hybrid cache requires block reuse to be disabled in KV cache config" - ) - - # Host-drafter spec modes (NGram) have no spec worker to call - # update_mamba_states; this manager promotes accepted states itself - # in update_resources. One-model modes (MTP/Eagle/DFlash) and the - # suffix-automaton worker promote from their spec workers and must - # NOT be promoted twice. - self._promote_states_in_update_resources = ( - spec_config is not None - and getattr(spec_config, "decoding_type", None) == "NGram") - - pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) - - MambaCacheManager.__init__( - self, - mamba_d_state, - mamba_d_conv, - mamba_num_heads, - mamba_n_groups, - mamba_head_dim, - mamba_num_layers, - pool_size, - max_batch_size, - mapping, - mamba_cache_dtype, - mamba_ssm_cache_dtype, - mamba_layer_mask, - execution_stream, - speculative_num_draft_tokens=(spec_config.tokens_per_gen_step - 1 - if spec_config is not None else None), - model_type=model_type, - use_replay_state_update=use_replay_state_update, - mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - kda_replay_num_spec=kda_replay_num_spec, - ) - - # initialize kv cache manager - KVCacheManager.__init__( - self, - kv_cache_config, - kv_cache_type, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - execution_stream=execution_stream, - is_draft=is_draft, - pool_configurations=pool_configurations, - ) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - MambaCacheManager.prepare_resources(self, scheduled_batch) - KVCacheManager.prepare_resources(self, scheduled_batch) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - MambaCacheManager.free_resources(self, request) - KVCacheManager.free_resources(self, request, pin_on_release) - - def add_dummy_requests(self, request_ids: List[int], **kwargs): - MambaCacheManager.add_dummy_requests(self, request_ids) - return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) - - def shutdown(self): - MambaCacheManager.shutdown(self) - KVCacheManager.shutdown(self) - - def update_resources(self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None): - KVCacheManager.update_resources(self, scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) - - def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): - """Commit accepted verification states for host-drafter spec modes. - - One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton - worker call update_mamba_states themselves right after on-device - acceptance. Host-drafter modes (NGram) have no spec worker: - acceptance lands on the requests as - ``py_num_accepted_draft_tokens`` during sampler update, and this - hook — running right after, alongside the KV rewind — promotes the - accepted step's intermediate state into the live pools (or, on the - KDA replay path, records the accepted-draft count for the next - round's replay). Without it, the pools would keep the - pre-verification state and the next forward would resume from a - stale prefix. - """ - if not self._promote_states_in_update_resources: - return - if not self.is_speculative() or attn_metadata is None: - return - gen_requests = scheduled_batch.generation_requests - if not gen_requests: - return - drafted = [ - r for r in gen_requests - if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 - ] - if not drafted: - # Drafter skipped this step (e.g. speculation gated off): the - # forward ran the plain in-place decode path; nothing to promote. - return - assert len(drafted) == len(gen_requests), ( - "mixed drafted/undrafted generation batch is not supported for " - "hybrid state promotion (drafts are padded to the static max)") - device = self._impl.mamba_cache.temporal.device - num_contexts = len(scheduled_batch.context_requests) - num_accepted = torch.tensor( - [0] * num_contexts + - [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], - dtype=torch.int32, - device=device) - # Batch-ordered slots (contexts then gens), matching the ordering - # the forward used for the intermediate scratch buffers. Requests - # that finished this step were already freed by response handling - # (which runs before update_resources) and are gone from the index; - # their rows must stay in place for alignment, so redirect them to - # the reserved padding slot (a harmless scratch write). - slot_index = self.mamba_cache_index - padding_slot = self._impl._padding_slot - state_indices = torch.tensor([ - slot_index.get(r.py_request_id, padding_slot) - for r in scheduled_batch.context_requests + gen_requests - ], - dtype=torch.int32, - device=device) - self.update_mamba_states(attn_metadata, num_accepted, state_indices) - - -@triton.jit -def _promote_mamba_state_kernel( - src_ptr, - dst_ptr, - src_idx_ptr, - accepted_position_source_ptr, - num_accepted_tokens_ptr, - blk_ptr, - replay_pnat_ptr, - replay_cache_buf_idx_ptr, - dummy_request_mask_ptr, - num_gens, - count, - src_s_layer, - src_s_row, - src_s_step, - dst_s_layer, - dst_s_block, - POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, - REPLAY_STEP_WIDTH: tl.constexpr, - REPLAY_HISTORY_SIZE: tl.constexpr, - UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, - BLOCK: tl.constexpr, -): - # One program copies a BLOCK-sized tile of the contiguous inner state for a - # single (layer, gen) pair. grid = (num_layers * num_gens, ceil(count/BLOCK)). - pair = tl.program_id(0) - tile = tl.program_id(1) - layer = (pair // num_gens).to(tl.int64) - g = pair % num_gens - row = tl.load(src_idx_ptr + g).to(tl.int64) - accepted_position = tl.load(accepted_position_source_ptr + g) - if POSITION_SOURCE_IS_TOKEN_COUNT: - accepted_position -= 1 - acc = accepted_position.to(tl.int64) - blk = tl.load(blk_ptr + g).to(tl.int64) - - # The replay metadata update is batch-scoped, so layer 0's first copy tile - # owns it. This folds PNAT and active-buffer maintenance into the existing - # accepted-convolution-state promotion launch. - if UPDATE_REPLAY_BOOKKEEPING: - if (layer == 0) & (tile == 0): - num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) - is_dummy_request = tl.load(dummy_request_mask_ptr + g) - pnat = tl.load(replay_pnat_ptr + blk) - cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) - wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE - next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, - pnat + num_accepted_tokens) - next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx) - tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) - tl.store(replay_cache_buf_idx_ptr + blk, - next_cache_buf_idx, - mask=~is_dummy_request) - - # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows - # int32 and would corrupt addresses / fault. - src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step - dst_base = layer * dst_s_layer + blk * dst_s_block - offs = tile * BLOCK + tl.arange(0, BLOCK) - mask = offs < count - v = tl.load(src_ptr + src_base + offs, mask=mask) - tl.store(dst_ptr + dst_base + offs, v, mask=mask) - - -def _promote_mamba_state_triton( - dst: torch.Tensor, - intermediate: torch.Tensor, - src_state_indices: torch.Tensor, - accepted_position_source: torch.Tensor, - dst_state_indices: torch.Tensor, - num_accepted_tokens: Optional[torch.Tensor] = None, - position_source_is_token_count: bool = False, - replay_pnat: Optional[torch.Tensor] = None, - replay_cache_buf_idx: Optional[torch.Tensor] = None, - dummy_request_mask: Optional[torch.Tensor] = None, - replay_step_width: int = 0, - replay_history_size: int = 0, - BLOCK: int = 2048) -> None: - """Scatter each generation request's accepted draft-step recurrent state - from the per-request intermediate buffer into the unified C++ pool view. - - Args: - dst: ``[num_layers, num_blocks, *state_shape]`` view of the C++ pool. - The block dim (dim 1) is strided (the pool interleaves - ``ssm_bytes | conv_bytes`` per block), so ``dst`` is non-contiguous. - The kernel writes through ``dst``'s real strides, touching only this - state's bytes -- the case that defeats torch.compile / aot_autograd - (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 - pool). - intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. - src_state_indices, accepted_position_source, dst_state_indices: - ``[num_gens]`` int. - num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required - when replay bookkeeping is enabled. - position_source_is_token_count: Whether ``accepted_position_source`` - contains accepted-token counts instead of zero-based intermediate - state positions. When true, the kernel subtracts one while - computing the source position. - replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay - metadata. When supplied together, the kernel advances PNAT and the - active history buffer while promoting the convolution state. - - For each gen ``g``:: - - accepted_position = ( - accepted_position_source[g] - 1 - if position_source_is_token_count - else accepted_position_source[g] - ) - dst[:, dst_state_indices[g]] = intermediate[ - :, src_state_indices[g], accepted_position - ] - - Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. - """ - num_layers = dst.shape[0] - num_gens = src_state_indices.shape[0] - if num_gens == 0: - return - update_replay_bookkeeping = replay_pnat is not None - if update_replay_bookkeeping: - assert num_accepted_tokens is not None - assert replay_cache_buf_idx is not None - assert dummy_request_mask is not None - assert replay_step_width > 0 - assert replay_history_size >= replay_step_width - assert replay_pnat.dtype == torch.int32 - assert replay_cache_buf_idx.dtype == torch.int32 - assert dummy_request_mask.dtype == torch.bool - assert dummy_request_mask.numel() >= num_gens - else: - # Compile-time-disabled pointer arguments still require tensors. - num_accepted_tokens = accepted_position_source - assert replay_cache_buf_idx is None - assert dummy_request_mask is None - replay_pnat = dst_state_indices - replay_cache_buf_idx = dst_state_indices - dummy_request_mask = dst_state_indices - count = 1 - for s in dst.shape[2:]: - count *= s - # Grid dim order matters: dim 0 -> CUDA grid.x (limit 2^31-1), dim 1 -> - # grid.y (limit 65535). Put the (layer, gen) pairs in dim 0 since that is - # what grows with batch size (num_layers*num_gens stays far below 2^31 for - # any real config); the tile count in dim 1 is small (~count/BLOCK). Do NOT - # swap them: pairs would hit the 65535 y-limit at num_layers*num_gens>65535. - grid = (num_layers * num_gens, triton.cdiv(count, BLOCK)) - _promote_mamba_state_kernel[grid]( - intermediate, - dst, - src_state_indices, - accepted_position_source, - num_accepted_tokens, - dst_state_indices, - replay_pnat, - replay_cache_buf_idx, - dummy_request_mask, - num_gens, - count, - intermediate.stride(0), - intermediate.stride(1), - intermediate.stride(2), - dst.stride(0), - dst.stride(1), - POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, - REPLAY_STEP_WIDTH=replay_step_width, - REPLAY_HISTORY_SIZE=replay_history_size, - UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, - BLOCK=BLOCK, - ) - - -def _mamba_snapshot_rule_counts( - kv_cache_config: KvCacheConfig, - max_seq_len: Optional[int], - tokens_per_block: int, -) -> Tuple[int, int]: - """Return reachable fixed rules and their partial-block upper bound.""" - if not kv_cache_config.enable_block_reuse: - return 0, 0 - - num_rules = 0 - num_unaligned_rules = 0 - state_config = kv_cache_config.mamba_state_config - for offset in set(state_config.additional_snapshot_offsets_from_start): - if max_seq_len is not None and offset > max_seq_len: - continue - num_rules += 1 - num_unaligned_rules += int(offset % tokens_per_block != 0) - for offset in set(state_config.additional_snapshot_offsets_from_end): - # A from-end offset is reachable iff some valid prompt is longer than - # the offset. Its absolute alignment depends on that prompt. - if max_seq_len is not None and offset >= max_seq_len: - continue - num_rules += 1 - num_unaligned_rules += 1 - return num_rules, num_unaligned_rules - - -def _mamba_regular_snapshot_interval( - kv_cache_config: KvCacheConfig, - max_seq_len: Optional[int], -) -> Optional[int]: - if not kv_cache_config.enable_block_reuse: - return None - interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval - if interval is None or interval <= 0: - return None - if max_seq_len is not None and interval > max_seq_len: - return None - return interval - - -def _get_local_mamba_cache_layout( - model_config, - mapping: Mapping, - *, - spec_config=None, - is_draft: bool = False, - use_separate_draft_kv_cache: bool = False, -): - """Return normalized params and local Mamba/attention layer counts. - - Cache construction and affine sizing must follow the model's PP layout: - partition base layers first, then place appended speculative layers on the - last PP rank. The normalized params retain target masks and the appended - draft-layer count so estimation selects the same combined or per-manager - layout as runtime. - """ - from tensorrt_llm._torch.pyexecutor.config_utils import \ - extract_mamba_kv_cache_params - - params = extract_mamba_kv_cache_params( - model_config.pretrained_config, - spec_config=spec_config, - quant_config=model_config.quant_config, - ) - mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - combined_layer_mask = [ - is_mamba or is_attention for is_mamba, is_attention in zip( - mamba_layer_mask, full_attention_layer_mask) - ] - local_layer_indices, _ = get_pp_layers( - sum(combined_layer_mask), - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - local_mamba_layers = sum(mamba_layer_mask[layer_idx] - for layer_idx in local_layer_indices) - local_attention_layers = sum(full_attention_layer_mask[layer_idx] - for layer_idx in local_layer_indices) - return params, local_mamba_layers, local_attention_layers - - -# --- Qwen4-Exp PLE --------------------------------------------------------- -# PLE (the n-gram side path) is a Qwen4-Exp-only feature that rides on the -# hybrid manager's Mamba slots. It is kept in dedicated helpers so it can be -# lifted into a model-owned cache manager wholesale later on. - - -def _qwen4_exp_ple_state_bytes_per_rank( - model_config, - params, - mapping, - *, - spec_config, - use_separate_draft_kv_cache: bool, -) -> int: - """Bytes this rank needs for Qwen4-Exp PLE recurrent state; 0 otherwise.""" - from tensorrt_llm._torch.pyexecutor.config_utils import ( - extract_qwen4_exp_ple_cache_params, is_qwen4_exp) - - # The estimator is also called with minimal model-config stubs that carry no - # `pretrained_config`; such a config is definitionally not Qwen4-Exp. - pretrained_config = getattr(model_config, "pretrained_config", None) - if pretrained_config is None or not is_qwen4_exp(pretrained_config): - return 0 - ple_params = extract_qwen4_exp_ple_cache_params(pretrained_config) - mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( - is_draft=False, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - combined_layer_mask = [ - is_mamba or is_attention for is_mamba, is_attention in zip( - mamba_layer_mask, full_attention_layer_mask) - ] - local_layer_indices, _ = get_pp_layers( - sum(combined_layer_mask), - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - local_ple_layers = sum(layer_id < len(ple_params.ple_layer_mask) - and ple_params.ple_layer_mask[layer_id] - for layer_id in local_layer_indices) - ple_bytes_per_layer = ( - ple_params.short_conv_channels * ple_params.short_conv_state_len * - ple_params.conv_state_dtype.itemsize + - ple_params.ngram_context_len * torch.int64.itemsize) - return local_ple_layers * ple_bytes_per_layer - - -def _estimated_kda_replay_num_spec(model_config, spec_config, - is_disagg: bool) -> Optional[int]: - """Replay-cache width a V2 manager for this model would allocate. - - Mirrors the executor's construction-time choice: KDA models (Kimi K3 and - GLM-5 Next) under speculative decoding get replay caches when the fused - verify kernel can run here (``kda_mtp_replay_num_spec``), except on a - disaggregated context server, which never verifies drafts. - """ - if spec_config is None or _is_disagg_context_server(is_disagg): - return None - pretrained_config = getattr(model_config, "pretrained_config", None) - if pretrained_config is None or not (is_kimi_linear(pretrained_config) - or is_glm5_next(pretrained_config)): - return None - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ - kda_mtp_replay_num_spec - - return kda_mtp_replay_num_spec(spec_config.tokens_per_gen_step - 1) - - -def _estimate_mamba_hybrid_cache_cost( - model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - tokens_per_block: int, - max_seq_len: Optional[int], - num_reserved_dummy_slots: int, - include_explicit_snapshots: bool, - cap_partial_attention_snapshots: bool, - is_draft: bool = False, - use_separate_draft_kv_cache: bool = False, - kda_replay_num_spec: Optional[int] = None, - **kwargs, -) -> Tuple[int, int]: - spec_config = kwargs.get("spec_config") - params, local_mamba_layers, local_attention_layers = ( - _get_local_mamba_cache_layout( - model_config, - mapping, - spec_config=spec_config, - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - )) - attention_slope = (KVCacheManager.get_cache_size_per_token( - model_config, - mapping, - num_layers=local_attention_layers, - **kwargs, - ) if local_attention_layers > 0 else 0) - state_bytes_per_rank = (local_mamba_layers * - params.get_states_bytes_per_layer(mapping)) - if not is_draft: - state_bytes_per_rank += _qwen4_exp_ple_state_bytes_per_rank( - model_config, - params, - mapping, - spec_config=spec_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - if kda_replay_num_spec is not None: - # KDA replay caches live in the same KVCM slot as the recurrent state, - # so every live, dummy and snapshot slot carries them. - tp_size = mamba_effective_tp_size(mapping) - state_bytes_per_rank += _kda_replay_bytes_per_slot( - _kda_replay_layout(kda_replay_num_spec, params.num_heads // tp_size, - params.head_dim, params.conv_kernel - 1), - local_mamba_layers) - max_resident_sequences = max_batch_size * mapping.pp_size - - if include_explicit_snapshots: - fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( - kv_cache_config, max_seq_len, tokens_per_block) - else: - fixed_rules = 0 - unaligned_fixed_rules = 0 - fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + - max_resident_sequences * fixed_rules) - attention_block_bytes = attention_slope * tokens_per_block - - interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) - has_unaligned_periodic_snapshot = (interval is not None - and interval % tokens_per_block != 0) - if cap_partial_attention_snapshots: - # Snapshot alignment is unknown while estimating cache cost. Once a - # snapshot is possible, reserve one retained partial attention page - # per resident lineage. Dummy requests carry no attention capacity. - has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None - partial_attention_slots = (max_resident_sequences - if has_non_live_ssm_capacity else 0) - else: - partial_attention_slots = (max_resident_sequences * - unaligned_fixed_rules) - intercept = (fixed_state_slots * state_bytes_per_rank + - partial_attention_slots * attention_block_bytes) - - if interval is None: - regular_slope = 0 - else: - regular_slope = math.ceil(state_bytes_per_rank / interval) - if (has_unaligned_periodic_snapshot - and not cap_partial_attention_snapshots): - regular_slope += math.ceil(attention_block_bytes / interval) - return attention_slope + regular_slope, intercept - - -class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): - """Hybrid cache manager storing mamba states inside the KVCacheManager pool. - - Both KV cache blocks and recurrent state blocks are managed by the unified - C++ KVCacheManager, enabling block reuse / prefix caching across attention - and mamba layers. This compatibility manager remains available through the - manager preference override and legacy disaggregated routing. - - """ - - # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the - # cache manager __init__ coerces a positive spec recompute tail to a full - # re-prefill (see _has_recurrent_state on the attention-only base classes). - _has_recurrent_state = True - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - layer_mask: Optional[ - List[bool]] = None, # this is the full attention layer mask - is_estimating_kv_cache: bool = False, - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - model_type: str = "nemotron_hybrid", - **kwargs, - ) -> None: - # 3 kinds of layers: - # 1) Mamba layers (mamba_layer_mask is True) - # 2) Full attention layers (full_attention_layer_mask is True) - # 3) Not managed layers (both masks are False) - total_layers = len(mamba_layer_mask) - if layer_mask is None: - full_attention_layer_mask = [False] * total_layers - elif len(layer_mask) != total_layers: - raise ValueError( - f"layer_mask length ({len(layer_mask)}) must match " - f"mamba_layer_mask length ({total_layers})") - else: - full_attention_layer_mask = list(layer_mask) - layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] - # PP sharding is done across all layers. - # This is called again in the super().__init__, but we want it to run first - # to set up mtp states before the C++ backend is initialized. - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=layer_mask, - ) - self.mamba_pp_layers = [ - layer_idx for layer_idx in self.pp_layers - if mamba_layer_mask[layer_idx] - ] - self.local_num_mamba_layers = len(self.mamba_pp_layers) - self.requests = [] - # Seed externally visible mamba fields before any early return so that - # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work - # on ranks with no local mamba layers. - self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = ( - use_replay_state_update and model_type == "qwen3_next" - and self.local_num_mamba_layers > 0) - self.replay_step_width: Optional[int] = ( - spec_config.tokens_per_gen_step - if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = ( - max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) - if self.replay_step_width is not None else None) - # Same allocation gate as PythonMambaCacheManager: the rand_seed - # buffer must exist whenever SR can fire, not only on the replay path. - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - # Host-side counter bumped per fresh context-slot assignment; combined - # with the slot index and rank offset to produce reproducible per-slot - # seed values without any torch.randint. - self._seed_request_counter = 0 - self.ssm_state_dtype = mamba_ssm_cache_dtype - # Keep the shared Mamba interface valid on PP ranks that do not own a - # local Mamba layer. - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - - if self.local_num_mamba_layers == 0: - logger.info( - "No local mamba layers for this rank, skipping mamba cache initialization" - ) - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=full_attention_layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - is_draft=is_draft, - **kwargs, - ) - # PP ranks replay the same scheduling decisions, so a rank without - # local Mamba layers must still publish the configured boundaries. - self.kv_cache_config = kv_cache_config - self.linear_attention_metadata = LinearAttentionMetadata() - self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_config.periodic_snapshot_interval - if kv_cache_config.enable_block_reuse else 0) - return - - # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) - tp_size = mamba_effective_tp_size(mapping) - d_inner = mamba_head_dim * mamba_num_heads - conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state - nheads = mamba_num_heads - assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - if use_replay_state_update: - assert mamba_n_groups % tp_size == 0, \ - "replay state update requires mamba_n_groups divisible by tp_size" - self._n_groups_per_rank = mamba_n_groups // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - self.conv_state_shape = [conv_dim, mamba_d_conv - 1] - self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - self.conv_state_dtype = mamba_cache_dtype - - # Store GLOBAL (pre-TP-division) mamba params for disagg RnnModelConfig. - # d_inner is computed at the top of __init__ before TP division. - d_inner_global = d_inner # = mamba_head_dim * mamba_num_heads (GLOBAL) - conv_dim_global = d_inner_global + 2 * mamba_n_groups * mamba_d_state - self._rnn_d_state = mamba_d_state - self._rnn_d_conv = mamba_d_conv - self._rnn_num_heads = mamba_num_heads # GLOBAL - self._rnn_n_groups = mamba_n_groups - self._rnn_head_dim = mamba_head_dim - self._rnn_hidden_size = d_inner_global # GLOBAL = head_dim * num_heads - self._rnn_conv_dim_size = conv_dim_global # GLOBAL conv dim - self._rnn_num_layers = mamba_num_layers - - # Conv section layout for section-aware split/concat in TP mismatch. - # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; - # we just need to tell C++ which ordering to use. - self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" - self.ssm_bytes = (math.prod(self.ssm_state_shape) * - self.ssm_state_dtype.itemsize) - self.conv_bytes = (math.prod(self.conv_state_shape) * - self.conv_state_dtype.itemsize) - - total_bytes = self.ssm_bytes + self.conv_bytes - if total_bytes % self.ssm_state_dtype.itemsize != 0: - raise RuntimeError( - f"Total state bytes ({total_bytes}) not divisible by " - f"ssm_state_dtype size ({self.ssm_state_dtype.itemsize})") - if total_bytes % self.conv_state_dtype.itemsize != 0: - raise RuntimeError( - f"Total state bytes ({total_bytes}) not divisible by " - f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") - if self.ssm_bytes % self.conv_state_dtype.itemsize != 0: - raise RuntimeError( - f"SSM state bytes ({self.ssm_bytes}) not divisible by " - f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") - self.linear_attention_metadata = LinearAttentionMetadata() - self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value - self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes - self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_config.periodic_snapshot_interval - if kv_cache_config.enable_block_reuse else 0) - # RNN model params for disagg TP-mismatch split/concat. - conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} - self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads - self.linear_attention_metadata.rnn_head_dim = self._rnn_head_dim - self.linear_attention_metadata.rnn_d_state = self._rnn_d_state - self.linear_attention_metadata.rnn_d_conv = self._rnn_d_conv - self.linear_attention_metadata.rnn_n_groups = self._rnn_n_groups - self.linear_attention_metadata.rnn_conv_section_layout = conv_section_map.get( - self._rnn_conv_section_layout, 0) - self.linear_attention_metadata.rnn_ssm_bytes = self.ssm_bytes - self.linear_attention_metadata.rnn_ssm_dtype_size = self.ssm_state_dtype.itemsize - self.linear_attention_metadata.rnn_conv_dtype_size = self.conv_state_dtype.itemsize - kv_cache_config = kv_cache_config.model_copy(deep=True) - if kv_cache_config.enable_partial_reuse: - logger.warning( - "Partial reuse is not supported for mamba hybrid models, disabling partial reuse" - ) - kv_cache_config.enable_partial_reuse = False - - # Keep the vector in physical global-layer order. Disabled entries are - # placeholders that are never projected into this manager. - kv_cache_config.max_attention_window = [ - LinearCacheType.RECURRENT_STATES.value - if mamba_layer_mask[layer_idx] else max_seq_len - for layer_idx in range(len(layer_mask)) - ] - - recurrent_states_window = LinearCacheType.RECURRENT_STATES.value - local_windows = { - recurrent_states_window - if mamba_layer_mask[layer_idx] else max_seq_len - for layer_idx in self.pp_layers - } - kwargs["pool_configurations"] = [ - PoolConfiguration( - window_size=window_size, - head_dim=head_dim, - dtype=torch_dtype_to_binding(self.ssm_state_dtype) - if window_size == recurrent_states_window else dtype, - ) for window_size in sorted(local_windows) - ] - - # Normalize num_kv_heads to a per-layer list and zero out mamba - # layer positions: those layers carry SSM/conv state instead of KV - # heads, so the parent KV cache should not allocate KV head storage - # for them. - if isinstance(num_kv_heads, int): - per_layer_kv_heads = [num_kv_heads] * total_layers - else: - if len(num_kv_heads) != total_layers: - raise ValueError( - f"num_kv_heads list length ({len(num_kv_heads)}) does not " - f"match total layers ({total_layers})") - per_layer_kv_heads = list(num_kv_heads) - for i, is_mamba in enumerate(mamba_layer_mask): - if is_mamba: - per_layer_kv_heads[i] = 0 - - self._setup_mtp_intermediate_states(spec_config, max_batch_size) - - # pass remaining arguments to super class - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=mamba_num_layers + num_layers, - num_kv_heads=per_layer_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - is_draft=is_draft, - linear_attention_metadata=self.linear_attention_metadata, - **kwargs, - ) - - assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" - self.mamba_layer_offsets = {} - for idx, layer_id in enumerate(self.mamba_pp_layers): - self.mamba_layer_offsets[layer_id] = idx - - self.host_block_offsets = torch.zeros([ - self.impl.num_pools, self.max_batch_size, 2, self.max_blocks_per_seq - ], - dtype=torch.int32, - device="cpu") - self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ - self.layer_offsets[self.mamba_pp_layers[0]]][0] - - self._state_indices_buffer = DeviceCopyBuffer( - torch.zeros([self.max_batch_size], dtype=torch.int32, - device="cuda"), - pin_memory=prefer_pinned()) - self._host_state_indices = torch.zeros([self.max_batch_size], - dtype=torch.int32, - pin_memory=prefer_pinned()) - self._row_indices = torch.arange(self.max_batch_size, - dtype=torch.long, - device="cpu") - self._request_id_to_state_index = {} - self._request_id_to_is_dummy = {} - self.kv_cache_config = kv_cache_config - self.is_estimating_kv_cache = is_estimating_kv_cache - - self._setup_states() - self._setup_replay_buffers(spec_config) - if use_replay_state_update and model_type == "qwen3_next": - logger.info_once( - "Configured GDN cached replay commit mode: small-batch fused, " - "large-batch all-layer", - key="gdn_cached_replay_commit_mode_fused", - ) - - @staticmethod - def get_cache_size_per_token( - model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - num_layers: Optional[int] = None, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs, - ): - """Affine memory model for the unified hybrid KV pool. - - Returns ``(slope_bytes_per_token, intercept_bytes)``: - - * ``slope`` = attention KV bytes per token plus conservative periodic - Mamba-state and partial-attention-snapshot costs. - * ``intercept`` = live and CUDA-graph dummy Mamba state. - - Memory budget -> max tokens then becomes - ``T = (budget - intercept) // slope`` instead of plain - ``T = budget // bytes_per_token``. - """ - return _estimate_mamba_hybrid_cache_cost( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - num_reserved_dummy_slots=1, - include_explicit_snapshots=False, - cap_partial_attention_snapshots=False, - **kwargs, - ) - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return self._use_gdn_cached_replay_all_layer_commit - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Synchronously advance every local GDN checkpoint in one launch.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not self._use_gdn_cached_replay_all_layer_commit - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (self.all_ssm_states is None or self.old_x is None - or self.old_B is None or self.old_dt is None - or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires replay state buffers." - ) - - 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 " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states, - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - - def shutdown(self): - # Release tensor views into the pool before the pool memory is freed, - # so their deleters don't see stale pointers. - self.all_ssm_states = None - self.all_conv_states = None - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask_buffer = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - super().shutdown() - - def add_dummy_requests( - self, - request_ids: List[int], - # Note that token_nums should be past_kv_len + input_len (without - # spec decoding). The draft tokens will be added in this function, - # so we don't need to take care of it in the caller. When preparing - # token_nums, we should not take the draft tokens into account, so - # don't use the kv_cache_manager.max_seq_len, which includes both - # extra tokens and draft tokens. - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - max_num_draft_tokens: int = 0, - kv_reserve_draft_tokens: Optional[int] = None, - use_mrope: bool = False, - max_beam_width: int = 1, - encoder_output_lens: Optional[List[int]] = None, - draft_kv_cache_manager: Optional[KVCacheManager] = None, - capture_sampling_params: Optional["SamplingParams"] = None, - ) -> List[LlmRequest]: - requests = super().add_dummy_requests( - request_ids=request_ids, - token_nums=token_nums, - is_gen=is_gen, - max_num_draft_tokens=max_num_draft_tokens, - kv_reserve_draft_tokens=kv_reserve_draft_tokens, - use_mrope=use_mrope, - max_beam_width=max_beam_width, - encoder_output_lens=encoder_output_lens, - draft_kv_cache_manager=draft_kv_cache_manager, - capture_sampling_params=capture_sampling_params, - ) - if requests: - self.requests.extend(requests) - # Only process the newly added requests, not all of self.requests. - # self.requests may contain stale entries from the previous - # _prepare_resources call (e.g., disagg transfer-pending requests) - # that would exceed max_batch_size and cause out-of-bounds access. - self._setup_state_indices(requests) - return requests - - def update_resources(self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None): - super().update_resources(scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - - @nvtx_range("hybrid_prepare_resources") - def _prepare_resources(self, scheduled_batch: ScheduledRequests): - self.requests = scheduled_batch.context_requests + \ - scheduled_batch.generation_requests - # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) - # until just before forward so the CPU-side prep (attn metadata, input - # tensor assembly, draft model prep, etc.) can overlap with the in-flight - # cudaMemcpyAsync calls instead of being serialized behind them. - # copy_linear_attention_block returns True iff it actually issued a copy; - # we skip refresh_blocks entirely when nothing was scheduled. - self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( - self.requests) - self._setup_state_indices() - self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - super().prepare_resources(scheduled_batch) - if self.local_num_mamba_layers == 0: - return - self._prepare_resources(scheduled_batch) - - @nvtx_range("hybrid_flush_state_transfers") - def flush_state_transfers(self) -> None: - """Complete any deferred state-block onboards scheduled by - prepare_resources(). Must be called before forward() reads recurrent - state blocks. Cheap no-op when nothing was scheduled. - """ - if getattr(self, "_pending_state_transfers", False): - self.impl.refresh_blocks() - self._pending_state_transfers = False - - @nvtx_range("hybrid_update_mamba_states") - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.local_num_mamba_layers == 0: - return - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + - num_gens] - # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP - # passes accepted-token counts and lets the promotion kernel subtract - # one, avoiding a separate GPU subtraction/cast. - if accepted_leaf_positions is None: - accepted_position_source = accepted_tokens - position_source_is_token_count = True - else: - accepted_position_source = accepted_leaf_positions.to(torch.int32) - 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 - # 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() - state_indices_d = state_indices[num_contexts:num_contexts + - num_gens].to(torch.int32) - src_state_indices = self.intermediate_state_indices[:num_gens] - - # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter - # into the dtype-reinterpreted, strided C++ pool view. torch.compile - # can't handle the dtype-view mutation (and Inductor chokes on the - # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel - # writes through the view's real strides (~85% of HBM peak, one launch - # per state). - if self._use_replay_state_update: - # Every GDN layer has finished reading the old checkpoint and - # writing its candidate history. Advance all local checkpoints - # in one launch before PNAT and the active history buffer change. - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) - assert self._dummy_request_mask is not None - else: - # Legacy: copy the accepted SSM state from the intermediate buffer. - _promote_mamba_state_triton( - self.all_ssm_states, - self.intermediate_ssm_states, - src_state_indices, - accepted_position_source, - state_indices_d, - position_source_is_token_count=position_source_is_token_count) - - # Conv: both paths save all intermediate conv windows, carry over the - # accepted one. The replay path also advances PNAT and the active cache - # buffer in this same launch. - replay_dummy_request_mask = None - if self._use_replay_state_update: - replay_dummy_request_mask = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - _promote_mamba_state_triton( - self.all_conv_states, - self.intermediate_conv_states, - src_state_indices, - accepted_position_source, - state_indices_d, - num_accepted_tokens=accepted_tokens, - position_source_is_token_count=position_source_is_token_count, - replay_pnat=(self.prev_num_accepted_tokens - if self._use_replay_state_update else None), - replay_cache_buf_idx=(self.cache_buf_idx - if self._use_replay_state_update else None), - dummy_request_mask=replay_dummy_request_mask, - replay_step_width=(self.replay_step_width - if self._use_replay_state_update else 0), - replay_history_size=(self.replay_history_size - if self._use_replay_state_update else 0)) - - def get_num_available_tokens(self, - token_num_upper_bound: int, - max_num_draft_tokens: int = 0, - **kwargs) -> int: - # Base bound from attention KV cache pool (the parent's behaviour). - result = super().get_num_available_tokens(token_num_upper_bound, - max_num_draft_tokens, - **kwargs) - # Also bound by the recurrent-state pool capacity: each request needs - # roughly ceil(N / states_snapshot_interval) recurrent-state blocks - # (+1 for the corner case where N is a multiple of tokens_per_block). - # When block reuse is disabled, only one snapshot is needed per - # request, so no additional capping is required here. - interval = (self.linear_attention_metadata.states_snapshot_interval - if self.linear_attention_metadata is not None else 0) - # Attention-only PP ranks keep the interval so every rank publishes - # identical scheduling boundaries, but they have no recurrent-state - # pool whose capacity should constrain their attention KV cache. - if self.local_num_mamba_layers > 0 and interval and interval > 0: - stats = self.impl.get_kv_cache_stats() - rs_free = stats.num_free_blocks_per_window_size.get( - LinearCacheType.RECURRENT_STATES.value, 0) - # Reserve 1 block for the always-allocated last block (corner case - # / final live state) so we don't promise more tokens than the - # pool can actually back at allocation time. - usable_rs_blocks = max(0, rs_free - 1) - rs_token_cap = usable_rs_blocks * interval - result = min(result, rs_token_cap) - return max(result, 0) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - if request in self.requests: - self.requests.remove(request) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - super().free_resources(request, pin_on_release) - - def _setup_state_indices(self, requests=None) -> None: - if self.local_num_mamba_layers == 0: - return - if requests is None: - requests = self.requests - block_indices = [] - for req in requests: - if req.is_context_finished: - next_step = self.get_num_tokens(req) - 1 - elif self.kv_cache_config.enable_block_reuse: - next_step = (req.context_current_position - 1 + - req.context_chunk_size) - else: - next_step = req.prompt_len - 1 - block_indices.append(next_step // self.tokens_per_block) - self.impl.copy_batch_block_offsets( - self.host_block_offsets, [req.py_request_id for req in requests], 1, - 0) - max_blocks = self.blocks_per_window[ - LinearCacheType.RECURRENT_STATES.value][0] - n = len(requests) - self._host_state_indices.zero_() - if n > 0: - # Vectorized gather: replace per-element Python loop with a single - # tensor index op. block_indices is a small Python list so - # torch.tensor() conversion here is O(n) at C level, much cheaper - # than n round-trips through Python indexing. - bi = torch.tensor(block_indices, dtype=torch.long) - rows = self._row_indices[:n] - # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] - values = self.host_block_offsets[self.recurrent_states_pool_index, - rows, 0, bi] - invalid_mask = (values < 0) | (values >= max_blocks) - # The C++ recurrent-state manager uses null page-table entries for - # logical blocks that are not snapshot boundaries. Usually a - # context chunk ends exactly at an allocated snapshot (or at the - # final live-state block), but the scheduler may shorten a chunk - # further when KV capacity is tight. In that case the Mamba - # kernel must keep accumulating into the next allocated snapshot - # or final block, matching copyLinearAttentionBlock(), which also - # skips placeholders when it advances the live state. - for bad_i in invalid_mask.nonzero( - as_tuple=False).flatten().tolist(): - req = requests[bad_i] - if req.is_context_finished: - continue - last_prompt_block = (req.prompt_len - - 1) // self.tokens_per_block - row = self.host_block_offsets[self.recurrent_states_pool_index, - rows[bad_i], 0] - candidates = row[block_indices[bad_i]:last_prompt_block + 1] - valid_candidates = ((candidates >= 0) - & (candidates < max_blocks)).nonzero( - as_tuple=False) - if valid_candidates.numel() > 0: - values[bad_i] = candidates[valid_candidates[0, 0]] - - invalid_mask = (values < 0) | (values >= max_blocks) - if invalid_mask.any(): - bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) - req = requests[bad_i] - value = int(values[bad_i]) - raise RuntimeError( - f"Invalid recurrent state block index {value} " - f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " - f"prompt_len={req.prompt_len}, " - f"is_context_finished={req.is_context_finished}, " - f"context_current_position={req.context_current_position}, " - f"prepopulated_token_num={req.prepopulated_prompt_len}, " - f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " - f"block_index for next step is {block_indices[bad_i]}, " - "\nblock_ids=" - f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" - ) - self._host_state_indices[:n] = values - - self._state_indices_buffer.update(self._host_state_indices, force=True) - is_dummy = [req.is_dummy for req in requests] - self._refresh_dummy_request_mask(is_dummy) - - # Build request_id → pool block offset mapping so that - # get_state_indices can return indices in arbitrary request order. - # Bulk tolist avoids a per-request tensor-index + .item() round-trip. - state_values = self._host_state_indices[:n].tolist() - for req, value, dummy in zip(requests, state_values, is_dummy): - self._request_id_to_state_index[req.py_request_id] = value - self._request_id_to_is_dummy[req.py_request_id] = dummy - - def get_state_indices(self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None) -> list: - if self.local_num_mamba_layers == 0: - # Mamba metadata is prepared on every PP rank even when this rank - # owns only attention layers. No local kernel consumes these - # indices, so avoid consulting state that is intentionally absent. - return [0] * len(request_ids) if request_ids is not None else [] - if request_ids is not None: - # Return indices in the order of the caller's request_ids, - # not the internal self.requests order. This is critical when - # the batch is reordered after prepare_resources (e.g. disagg - # serving sorts generation_requests by py_batch_idx). - indices = [ - self._request_id_to_state_index[rid] for rid in request_ids - ] - if is_padding is None: - is_padding = [False] * len(request_ids) - assert len(request_ids) == len(is_padding) - is_dummy = [ - self._request_id_to_is_dummy.get(rid, False) or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - return self.cuda_state_indices - - def _setup_states(self) -> None: - # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) - pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( - torch.uint8).reshape(self.local_num_mamba_layers, -1, - self.ssm_bytes + self.conv_bytes) - num_blocks_in_pool = pool.shape[1] - self.all_ssm_states = pool[:, :, :self.ssm_bytes].view( - self.ssm_state_dtype).view( - [self.local_num_mamba_layers, num_blocks_in_pool] + - self.ssm_state_shape) - self.all_conv_states = pool[:, :, self.ssm_bytes:self.ssm_bytes + - self.conv_bytes].view( - self.conv_state_dtype).view([ - self.local_num_mamba_layers, - num_blocks_in_pool - ] + self.conv_state_shape) - self.all_ssm_states.zero_() - self.all_conv_states.zero_() - - def _setup_replay_buffers(self, spec_config) -> None: - cache_size = self.all_ssm_states.shape[1] - device = self.all_ssm_states.device - if not self._allocate_pool_replay_buffers(spec_config, cache_size, - device): - return - - self._dummy_request_mask_buffer = DeviceCopyBuffer( - torch.zeros(self.max_batch_size, dtype=torch.bool, device=device), - pin_memory=prefer_pinned()) - self._dummy_request_mask_host = torch.zeros( - self.max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned(), - ) - - -class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): - """Hybrid Mamba cache manager backed by KVCacheManagerV2. - - Attention KV pages and Mamba recurrent-state pages are both owned by the - Python V2 cache manager. Mamba layers are represented as V2 SSM layers, - while this wrapper exposes the state tensors and slot indices expected by - the PyTorch Mamba kernels. - """ - - _supports_additional_snapshot_offsets = True - # Recurrent-state snapshots use a specialized commit/history protocol, so - # keep main-like reuse endpoints and the existing unpaired draft path. - _supports_reuse_match_backoff = False - - # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the - # cache manager __init__ coerces a positive spec recompute tail to a full - # re-prefill (see _has_recurrent_state on the attention-only base classes). - _has_recurrent_state = True - - # Qwen4-Exp PLE state is opt-in. These class-level defaults keep every other - # model — and any partially-constructed instance that sets only the fields it - # needs — out of the PLE branches without carrying PLE geometry. - # `_init_qwen4_exp_ple_geometry` rebinds all of them per instance. - _ple_layer_ids: list[int] = [] - _ple_conv_state_shape: list[int] = [] - _ple_ngram_context_shape: list[int] = [] - # Set per instance on a disaggregated context server of a KDA model. - kda_context_only = False - - def _disagg_transfer_overwrites_whole_cached_prefix(self) -> bool: - # The incoming recurrent state overwrites the whole local slot and summarizes the - # complete prefix, so none of those tokens count as a local cache hit. - return (self.local_num_mamba_layers > 0 - or super()._disagg_transfer_overwrites_whole_cached_prefix()) - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - layer_mask: Optional[List[bool]] = None, - is_estimating_kv_cache: bool = False, - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", - kda_replay_num_spec: Optional[int] = None, - qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, - kda_draft_records: bool = False, - **kwargs, - ) -> None: - if conv_state_layout not in ("x_b_c", "q_k_v"): - raise ValueError( - f"Unsupported convolution state layout: {conv_state_layout!r}") - if "model_type" in kwargs: - # The V1 managers select the conv-state layout by model_type; this - # class takes it explicitly. Silently absorbing model_type here - # means a caller's layout request would be dropped on the floor. - raise TypeError( - "MambaHybridCacheManagerV2 does not accept 'model_type' " - f"(got {kwargs['model_type']!r}); pass " - "conv_state_layout='x_b_c' or 'q_k_v' instead") - - # A disaggregated context server never verifies drafts, so it keeps - # neither KDA replay caches nor intermediate verification buffers. - self.kda_context_only = bool( - kda_replay_num_spec is not None - and _is_disagg_context_server(kwargs.get("is_disagg", False))) - if self.kda_context_only: - kda_replay_num_spec = None - self._kda_replay_layers: List[Dict[str, torch.Tensor]] = [] - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - # Only with the KDA replay caches: also keep the records of every - # draft of the last verify round (kda_state_tok). - self._kda_draft_records = kda_draft_records and self._use_kda_replay_update - if self._use_kda_replay_update: - if use_replay_state_update: - raise ValueError( - "kda_replay_num_spec and use_replay_state_update are " - "mutually exclusive") - if spec_config is None: - raise ValueError( - "KDA replay caches require speculative decoding") - expected_num_spec = spec_config.tokens_per_gen_step - 1 - if kda_replay_num_spec != expected_num_spec: - raise ValueError( - "KDA replay cache width " - f"({kda_replay_num_spec}) must match the draft length " - f"({expected_num_spec})") - if conv_state_layout != "q_k_v": - raise ValueError( - "KDA replay requires conv_state_layout='q_k_v'") - self._record_kda_replay_in_update_resources = ( - self.use_kda_replay_update - and getattr(spec_config, "decoding_type", None) == "NGram") - - total_layers = len(mamba_layer_mask) - if layer_mask is None: - full_attention_layer_mask = [False] * total_layers - elif len(layer_mask) != total_layers: - raise ValueError( - f"layer_mask length ({len(layer_mask)}) must match " - f"mamba_layer_mask length ({total_layers})") - else: - full_attention_layer_mask = list(layer_mask) - - combined_layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] - - self._mamba_layer_mask = list(mamba_layer_mask) - self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = (use_replay_state_update - and conv_state_layout - == "q_k_v") - self._gdn_cached_replay_state_descriptors: Optional[torch.Tensor] = None - self._gdn_cached_replay_state_strides: Optional[Tuple[int, int, - int]] = None - self.replay_step_width: Optional[int] = ( - spec_config.tokens_per_gen_step - if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = ( - max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) - if self.replay_step_width is not None else None) - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - self._seed_request_counter = 0 - self._recurrent_evicted_blocks_total = 0 - self._recurrent_onboarded_blocks_total = 0 - self._recurrent_dropped_blocks_total = 0 - self._recurrent_status_logged = False - # Branch points keyed by request id. prepare_expect_snapshot_points() - # overwrites the per-request list on every scheduler pass, so the point - # has to be re-merged from here rather than stored only on the request. - self._branch_snapshot_points: Dict[int, int] = {} - self._snapshot_pruned_tokens_total = 0 - self._page_pruned_tokens_total = 0 - self._branch_snapshots_taken_total = 0 - self._branch_snapshots_skipped_total: Dict[str, int] = {} - num_cuda_graph_padding_dummy_slots = ( - _get_num_cuda_graph_padding_dummy_slots(spec_config, - max_batch_size)) - self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + - int(mapping.enable_attention_dp)) - # PLE uses the same request slots and snapshots as the GDN state. - self._init_qwen4_exp_ple_geometry(qwen4_exp_ple_cache_params, - total_layers, mamba_cache_dtype) - self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype - is not None else mamba_cache_dtype) - self.conv_state_dtype = mamba_cache_dtype - - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - self.mamba_pp_layers = [ - layer_idx for layer_idx in self.pp_layers - if mamba_layer_mask[layer_idx] - ] - self.local_num_mamba_layers = len(self.mamba_pp_layers) - self._use_gdn_cached_replay_all_layer_commit = ( - self._use_gdn_cached_replay_all_layer_commit - and self.local_num_mamba_layers > 0) - - if self.local_num_mamba_layers > 0: - tp_size = mamba_effective_tp_size(mapping) - d_inner = mamba_head_dim * mamba_num_heads - grouped_state_dim = mamba_n_groups * mamba_d_state - conv_dim = d_inner + 2 * grouped_state_dim - nheads = mamba_num_heads - assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - if kwargs.get("is_disagg", - False) and grouped_state_dim % tp_size != 0: - raise ValueError( - "Disaggregated Mamba transfer requires each convolution " - "state section to be divisible by tp_size") - if use_replay_state_update: - assert mamba_n_groups % tp_size == 0, \ - "replay state update requires mamba_n_groups divisible by tp_size" - self._n_groups_per_rank = mamba_n_groups // tp_size - d_inner_local = d_inner // tp_size - grouped_state_dim_local = grouped_state_dim // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - self.conv_state_shape = [conv_dim, mamba_d_conv - 1] - self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - # TP-mismatch disaggregated transfers must split the flat - # convolution state at its true semantic boundaries. Mamba2 stores - # [x | B | C], while GDN stores [Q | K | V]. The large section is - # therefore first for Mamba2 and last for GDN. - if conv_state_layout == "x_b_c": - self.conv_section_dims = [ - d_inner_local, - grouped_state_dim_local, - grouped_state_dim_local, - ] - else: - self.conv_section_dims = [ - grouped_state_dim_local, - grouped_state_dim_local, - d_inner_local, - ] - self.ssm_bytes = (math.prod(self.ssm_state_shape) * - self.ssm_state_dtype.itemsize) - self.conv_bytes = (math.prod(self.conv_state_shape) * - self.conv_state_dtype.itemsize) - else: - logger.info( - "No local mamba layers for this rank, skipping mamba state views" - ) - self._n_groups_per_rank = 0 - self.conv_state_shape = [] - self.ssm_state_shape = [] - self.conv_section_dims = [] - self.ssm_bytes = 0 - self.conv_bytes = 0 - - if isinstance(num_kv_heads, int): - per_layer_kv_heads = [num_kv_heads] * total_layers - else: - if len(num_kv_heads) != total_layers: - raise ValueError( - f"num_kv_heads list length ({len(num_kv_heads)}) does not " - f"match total layers ({total_layers})") - per_layer_kv_heads = list(num_kv_heads) - for i, is_mamba in enumerate(mamba_layer_mask): - if is_mamba: - per_layer_kv_heads[i] = 0 - - self._setup_mtp_intermediate_states(spec_config, max_batch_size) - - kv_cache_config = kv_cache_config.model_copy(deep=True) - if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: - block_reuse_config = kv_cache_config.block_reuse_config - block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) - if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: - # SSM reuse is valid only at explicit snapshot boundaries. - kv_cache_config.block_reuse_config = block_reuse_config.model_copy( - update={"policy": BlockReusePolicy.PER_REQUEST.value}) - self.kv_cache_config = kv_cache_config - - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=mamba_num_layers + num_layers, - num_kv_heads=per_layer_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=combined_layer_mask, - is_draft=is_draft, - is_estimating_kv_cache=is_estimating_kv_cache, - num_reserved_index_slots=self._num_reserved_dummy_slots, - **kwargs, - ) - - self.mamba_layer_offsets = { - layer_id: offset - for offset, layer_id in enumerate(self.mamba_pp_layers) - } - self._request_id_to_state_index = {} - self._request_id_to_is_dummy = {} - # Requests whose suspend committed their KDA replay state into the V2 - # states; their next generation step re-seeds the replay caches. - self._kda_replay_suspended: set[int] = set() - # With a host drafter: requests of a prepared batch whose acceptance - # update_resources has not recorded yet. - self._kda_replay_acceptance_pending: set[int] = set() - self._kda_replay_unseeded: set[int] = set() - - state_index_capacity = (self.max_batch_size + - self._num_reserved_dummy_slots) - self._state_indices_buffer = DeviceCopyBuffer( - torch.zeros([state_index_capacity], - dtype=torch.int32, - device="cuda"), - pin_memory=prefer_pinned()) - self._host_state_indices = torch.zeros([state_index_capacity], - dtype=torch.int32, - pin_memory=prefer_pinned()) - - if self.local_num_mamba_layers > 0: - first_mamba_local_layer = self.layer_offsets[ - self.mamba_pp_layers[0]] - self.ssm_layer_group_id = self.impl.get_layer_group_id( - LayerId(first_mamba_local_layer)) - self._ssm_page_index_scale = self.impl.get_page_index_scale( - LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) - num_ssm_pages = self.impl.get_page_index_upper_bound( - LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) - num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // - self._ssm_page_index_scale) - required_live_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - if num_ssm_slots < required_live_slots: - KVCacheManagerV2.shutdown(self) - raise ValueError( - "The V2 Mamba state pool has only " - f"{num_ssm_slots} slots but needs at least " - f"{required_live_slots} live/dummy slots. Increase the " - "KV cache budget or allocate a larger Mamba pool_ratio.") - self._setup_states() - # PLE pools use exactly the state-slot capacity resolved above. - self._setup_ple_states(num_ssm_slots) - self._setup_replay_buffers(spec_config) - if self._use_gdn_cached_replay_all_layer_commit: - state_layout = ("affine" - if self._gdn_cached_replay_state_strides - is not None else "indirect") - logger.info_once( - "Configured GDN cached replay commit mode for V2: " - "small-batch fused, large-batch all-layer; " - f"state layout: {state_layout}", - key="gdn_cached_replay_v2_commit_mode_fused", - ) - else: - self.ssm_layer_group_id = None - self._ssm_page_index_scale = 1 - self.all_ssm_states = [] - self.all_conv_states = [] - self._stacked_ssm_states = None - self._stacked_conv_states = None - self._setup_replay_buffers(spec_config) - - def _init_qwen4_exp_ple_geometry( - self, - params: Optional["Qwen4ExpPLECacheParams"], - total_layers: int, - default_state_dtype: torch.dtype, - ) -> None: - """Validate and record the Qwen4-Exp PLE pool geometry. - - A no-op for every other model: the fields stay empty and - `_ple_buffer_configs` then contributes nothing. - """ - self._ple_layer_ids: list[int] = [] - self._ple_conv_state_shape: list[int] = [] - self._ple_ngram_context_shape: list[int] = [] - self._ple_conv_state_dtype = default_state_dtype - self._ple_conv_states: dict[int, torch.Tensor] = {} - self._ple_ngram_contexts: dict[int, torch.Tensor] = {} - if params is None: - return - if len(params.ple_layer_mask) != total_layers: - raise ValueError( - "PLE layer mask length must match the hybrid layer mask: " - f"got {len(params.ple_layer_mask)}, expected {total_layers}") - self._ple_layer_ids = [ - layer_id for layer_id, active in enumerate(params.ple_layer_mask) - if active - ] - if len(self._ple_layer_ids) != params.num_ple_layers: - raise ValueError( - "PLE layer mask count does not match num_ple_layers: " - f"got {len(self._ple_layer_ids)}, expected " - f"{params.num_ple_layers}") - if any(not self._mamba_layer_mask[layer_id] - for layer_id in self._ple_layer_ids): - raise ValueError("PLE lifecycle state must belong to Mamba layers") - if (params.short_conv_channels <= 0 or params.short_conv_state_len <= 0 - or params.ngram_context_len <= 0): - raise ValueError("PLE recurrent-state dimensions must be positive") - self._ple_conv_state_shape = [ - params.short_conv_channels, - params.short_conv_state_len, - ] - self._ple_ngram_context_shape = [params.ngram_context_len] - self._ple_conv_state_dtype = params.conv_state_dtype - - def _setup_ple_states(self, num_state_slots: int) -> None: - """Bind each local PLE layer to V2-managed recurrent-state views.""" - if not self._ple_layer_ids: - return - for layer_id in self._ple_layer_ids: - local_layer_idx = self.layer_offsets.get(layer_id) - if local_layer_idx is None: - continue - conv_state = self._get_state_buffer( - local_layer_idx, - MambaRole.PLE_CONV_STATE, - self._ple_conv_state_dtype, - self._ple_conv_state_shape, - ) - ngram_context = self._get_state_buffer( - local_layer_idx, - MambaRole.PLE_NGRAM_CONTEXT, - dtype=torch.long, - state_shape=self._ple_ngram_context_shape, - ) - if (conv_state.shape[0] != num_state_slots - or ngram_context.shape[0] != num_state_slots): - raise RuntimeError( - "PLE and GDN lifecycle buffers must have the same number " - f"of slots: layer={layer_id}, GDN={num_state_slots}, " - f"conv={conv_state.shape[0]}, ngram={ngram_context.shape[0]}" - ) - self._ple_conv_states[layer_id] = conv_state - self._ple_ngram_contexts[layer_id] = ngram_context - logger.info( - "PLE state views bound to V2 lifecycle buffers for local layers " - f"{sorted(self._ple_conv_states)}: conv " - f"[{num_state_slots}, {', '.join(map(str, self._ple_conv_state_shape))}] " - f"({self._ple_conv_state_dtype}), n-gram context " - f"[{num_state_slots}, {self._ple_ngram_context_shape[0]}] (int64)") - - def ple_layer_cache( - self, - layer_idx: int) -> Optional[tuple[torch.Tensor, torch.Tensor]]: - """Return persistent ``(conv_state, ngram_context)`` pools for a layer.""" - conv = self._ple_conv_states.get(layer_idx) - ngram = self._ple_ngram_contexts.get(layer_idx) - if conv is None or ngram is None: - return None - return conv, ngram - - def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: - """Recurrent roles and how they reshard across TP. - - Convolution state is SECTIONED: its flat per-layer buffer is a - concatenation of sections (Mamba2 ``[x | B | C]``, GDN ``[Q | K | V]``) - that are each TP-sharded independently. SSM state keeps the INDEXED - default (sharded by head). PLE state is computed from replicated - inputs, so every rank holds identical bytes and the transfer copies - whole per-layer regions. - """ - return { - **super().get_disagg_role_mapper_kinds(), - MambaRole.CONV_STATE: - MapperKind.SECTIONED, - MambaRole.PLE_CONV_STATE: - MapperKind.REPLICATED, - MambaRole.PLE_NGRAM_CONTEXT: - MapperKind.REPLICATED, - } - - def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: - """Per-layer geometry for resharding conv and SSM state.""" - if self.local_num_mamba_layers == 0: - return {} - d_conv_m1 = int(self.conv_state_shape[1]) - conv_elem_size = self.conv_state_dtype.itemsize - _, head_dim, d_state = self.ssm_state_shape - return { - MambaRole.CONV_STATE: - RoleLayout(section_bytes=tuple( - int(dim) * d_conv_m1 * conv_elem_size - for dim in self.conv_section_dims)), - MambaRole.SSM_STATE: - RoleLayout(bytes_per_head=int(head_dim) * int(d_state) * - self.ssm_state_dtype.itemsize), - } - - def get_disagg_ignored_roles(self) -> frozenset[DataRole]: - """Keep the KDA replay buffers out of disaggregated transfers. - - They are generation-local: a context server allocates none, and a - generation server seeds them from the transferred conv state. - """ - replay_roles = {DataRole(name) for name in self._kda_replay_layout()} - if replay_roles: - replay_roles.add(_KDA_REPLAY_COUNT_ROLE) - return super().get_disagg_ignored_roles() | frozenset(replay_roles) - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - - @property - def kda_state_tok(self) -> Optional[List[torch.Tensor]]: - """Per local KDA layer, the drafts' records of every SSM slot - ([slots, 3, num_spec, H, K] fp32 KVCM state-buffer views), or None - when the replay path keeps the per-draft replay caches instead.""" - replay_layers = getattr(self, "_kda_replay_layers", None) - if not replay_layers or _KDA_DRAFT_RECORDS_ROLE not in replay_layers[0]: - return None - return [layer[_KDA_DRAFT_RECORDS_ROLE] for layer in replay_layers] - - @property - def keeps_kda_draft_records(self) -> bool: - return self.kda_state_tok is not None - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - if self.kda_context_only: - super()._setup_mtp_intermediate_states(None, max_batch_size) - return - if not self.use_kda_replay_update: - super()._setup_mtp_intermediate_states(spec_config, max_batch_size) - return - - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if self.local_num_mamba_layers == 0: - return - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - - def _kda_replay_layout(self) -> Dict[str, Tuple[torch.dtype, List[int]]]: - """Per-slot shapes of the KDA replay buffers of one local layer.""" - if not self.use_kda_replay_update or self.local_num_mamba_layers == 0: - return {} - if (len(self.conv_section_dims) != 3 - or len(set(self.conv_section_dims)) != 1): - raise ValueError("KDA replay caches require equal [Q | K | V] " - "convolution-state sections") - assert self._kda_replay_num_spec is not None - draft_records = getattr(self, "_kda_draft_records", False) - layout = _kda_replay_layout(self._kda_replay_num_spec, - self.ssm_state_shape[0], - self.ssm_state_shape[1], - self.conv_state_shape[1], - draft_records=draft_records) - if draft_records and tuple(layout[_KDA_DRAFT_RECORDS_ROLE][1]) != ( - self._kda_draft_records_shape()): - raise ValueError("KDA draft records require equal value and key " - "head dimensions") - return layout - - def _kda_replay_buffer_configs(self) -> List[BufferConfig]: - return [ - BufferConfig(role=DataRole(name), - size=dtype.itemsize * math.prod(shape)) - for name, (dtype, shape) in self._kda_replay_layout().items() - ] - - def _validate_kda_replay_colocation(self, - replay_roles: Iterable[str]) -> None: - """Verify replay and recurrent buffers share one slot lifecycle. - - Every layer addresses its replay caches with the recurrent state's - logical slot index, so all of them must be allocated, moved and - offloaded together: they must sit in one variant of one KVCM pool - group, not merely have equal capacities. - """ - required_buffers = {(self.layer_offsets[layer_id], str(role)) - for layer_id in self.mamba_pp_layers - for role in (MambaRole.SSM_STATE, - MambaRole.CONV_STATE, *replay_roles)} - required_buffers.add((self.layer_offsets[self.mamba_pp_layers[0]], - str(_KDA_REPLAY_COUNT_ROLE))) - if not any( - required_buffers.issubset( - {(int(buffer_id.layer_id), str(buffer_id.role)) - for coalesced in variant.coalesced_buffers - for buffer_id in coalesced.buffer_ids}) - for pool_group in self.impl.pool_group_descs - for variant in pool_group.slot_desc.variants): - raise RuntimeError( - "KDA replay and recurrent-state buffers must share one KVCM " - "pool group and lifecycle to reuse the same logical slot index") - - def _allocate_pool_replay_buffers( - self, - spec_config, - cache_size: int, - device: Optional[torch.device], - ) -> bool: - allocated = super()._allocate_pool_replay_buffers( - spec_config, cache_size, device) - self._kda_replay_layers = [] - layout = self._kda_replay_layout() - if not layout: - return allocated - self._validate_kda_replay_colocation(layout) - # The replay caches are KVCM state buffers next to the SSM and conv - # states, so they follow the slot through suspend, offload and - # snapshot reuse. Pool capacity can exceed resident memory, so rows - # are initialized only once a request holds its slot. - for layer_id in self.mamba_pp_layers: - local_layer_idx = self.layer_offsets[layer_id] - layer = {} - for name, (dtype, shape) in layout.items(): - view = self._get_state_buffer(local_layer_idx, DataRole(name), - dtype, shape) - if name in _KDA_REPLAY_CONV_ROLES: - # Dim-contiguous [slots, dim, W - 1 + num_spec]. - view = view.transpose(-1, -2) - elif name == "kda_beta_cache": - view = view[..., :self.ssm_state_shape[0]] - if view.shape[0] != cache_size: - raise RuntimeError( - f"KDA replay buffer {name} of layer {layer_id} has " - f"{view.shape[0]} slots, but the recurrent state has " - f"{cache_size}") - layer[name] = view - self._kda_replay_layers.append(layer) - # All SSM layers share a slot lifecycle, so the accepted-draft count is - # stored once, on the first local KDA layer. - self.prev_num_accepted_tokens = self._get_state_buffer( - self.layer_offsets[self.mamba_pp_layers[0]], - _KDA_REPLAY_COUNT_ROLE, - torch.int32, - [_KDA_REPLAY_COUNT_BYTES // torch.int32.itemsize], - )[..., 0] - records = (", with the drafts' records in place of the replay " - "caches" if self.keeps_kda_draft_records else "") - logger.info("Mamba Cache (kda-replay) is bound to KVCM state buffers " - f"for {cache_size} state slots{records}") - return True - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Advance all V2 GDN checkpoints through their state views.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (not self.all_ssm_states - or (self._gdn_cached_replay_state_descriptors is None - and self._gdn_cached_replay_state_strides is None) - or self.old_x is None or self.old_B is None - or self.old_dt is None or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires V2 replay state buffers." - ) - - 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 " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - state_strides = self._gdn_cached_replay_state_strides - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states[0], - ssm_state_descriptors=(self._gdn_cached_replay_state_descriptors), - ssm_state_num_layers=(state_strides[0] - if state_strides is not None else None), - ssm_state_layer_stride=(state_strides[1] - if state_strides is not None else None), - ssm_state_slot_stride=(state_strides[2] - if state_strides is not None else None), - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - - def _create_kv_cache( - self, - request_id: int, - lora_task_id: Optional[int], - input_tokens: Optional[Sequence[TokenIdExt]], - *, - cache_salt: Optional[str] = None, - is_dummy: bool = False, - enable_request_stats: bool = False, - expected_prompt_length: Optional[int] = None, - kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None, - ) -> Optional[_KVCache]: - kv_cache = super()._create_kv_cache( - request_id, - lora_task_id, - input_tokens, - cache_salt=cache_salt, - is_dummy=is_dummy, - enable_request_stats=enable_request_stats, - expected_prompt_length=expected_prompt_length, - kv_cache_retention_config=kv_cache_retention_config, - ) - if (self.mapping.rank == 0 and kv_cache is not None - and input_tokens is not None and not is_dummy - and self.local_num_mamba_layers > 0): - # Prefix lookup excludes the final prompt token because it must be - # recomputed by prefill. Restore it for the request-length metric. - request_total_tokens = len(input_tokens) + 1 - prefix_reuse_diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) - logger.debug( - f"[MambaHybridCacheManagerV2] prefix reuse rank={self.mapping.rank} " - f"request_id={request_id} " - f"request_total_tokens={request_total_tokens} " - f"longest_attention_match_tokens=" - f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_hybrid_pruning()} " - f"content_divergence_tokens=" - f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_pruning()} " - f"latest_recurrent_snapshot_tokens=" - f"{kv_cache.num_committed_tokens}") - return kv_cache - - def _skip_branch_snapshot(self, reason: str) -> None: - self._branch_snapshots_skipped_total[reason] = ( - self._branch_snapshots_skipped_total.get(reason, 0) + 1) - - def _format_branch_snapshot_counters(self) -> str: - """Attribute reuse loss between missing snapshots and evicted pages. - - A flat cache hit rate means something different depending on whether - branch points were never found or were found and not reused, so the - skip reasons are reported alongside the totals. - """ - skipped = self._branch_snapshots_skipped_total - return (f"snapshot_pruned_tokens={self._snapshot_pruned_tokens_total} " - f"page_pruned_tokens={self._page_pruned_tokens_total} " - f"branch_snapshots_taken={self._branch_snapshots_taken_total} " - f"branch_snapshots_skipped={sum(skipped.values())} " - f"branch_snapshots_skipped_by_reason=" - f"{dict(sorted(skipped.items()))}") - - def _record_branch_snapshot_point(self, req: LlmRequest, kv_cache: _KVCache, - num_lookup_tokens: Optional[int]) -> None: - """Snapshot where this request leaves the reuse tree, for its siblings.""" - if not self.kv_cache_config.block_reuse_config.enable_branch_snapshot: - return - if req.is_dummy_request: - return - if num_lookup_tokens is None: - self._skip_branch_snapshot("no_fresh_match") - # A fresh lookup can record a point before resume() fails. On the - # next attempt, reapply that point against the cache's authoritative - # reuse depth: the request cursor is not rewound until preparation - # succeeds and may still describe a deeper, abandoned attempt. - if req.py_request_id in self._branch_snapshot_points: - self._apply_branch_snapshot_point( - req, context_current_position=kv_cache.num_committed_tokens) - return - - diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) - divergence = diagnostics._get_num_reusable_tokens_before_pruning() - reused = kv_cache.num_committed_tokens - if self.local_num_mamba_layers > 0: - hybrid_depth = ( - diagnostics._get_num_reusable_tokens_before_hybrid_pruning()) - # Loss attributable to a missing recurrent snapshot, versus loss - # attributable to attention pages having been evicted. These call - # for different fixes, so they are counted apart. - self._snapshot_pruned_tokens_total += max(0, hybrid_depth - reused) - self._page_pruned_tokens_total += max(0, divergence - hybrid_depth) - - # The whole lookup range matched, so there is no fork here. - if divergence >= num_lookup_tokens: - self._skip_branch_snapshot("no_divergence") - return - # Align down: tokens past a block boundary come from a partial-match - # child rather than a confirmed shared prefix, and an unaligned commit - # takes the partial-block snapshot path. - point = (divergence // self.tokens_per_block) * self.tokens_per_block - if point == 0: - self._skip_branch_snapshot("aligned_to_zero") - return - # Reuse already reached the fork, so a snapshot there exists. - if point <= reused: - self._skip_branch_snapshot("already_reused") - return - # A branch point at or beyond the prompt end cannot improve sibling - # reuse. - if point >= req.prompt_len: - self._skip_branch_snapshot("at_prompt_end") - return - - self._branch_snapshot_points[req.py_request_id] = point - self._apply_branch_snapshot_point(req, context_current_position=reused) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - super().prepare_expect_snapshot_points(requests) - if (not self.enable_block_reuse or not self.kv_cache_config. - block_reuse_config.enable_branch_snapshot): - return - for request in requests: - self._apply_branch_snapshot_point(request) - - def _apply_branch_snapshot_point( - self, - request: LlmRequest, - context_current_position: Optional[int] = None, - ) -> None: - points = set(request.expect_snapshot_points) - # Keep the flag usable without any configured snapshot placement. If a - # placement is configured, preserve it: adding prompt_len can replace a - # nearby recurrent checkpoint in the same radix-tree block because a - # block stores one recurrent page per life cycle. - if not points: - points.add(request.prompt_len) - point = self._branch_snapshot_points.get(request.py_request_id) - current_position = (request.context_current_position - if context_current_position is None else - context_current_position) - if point is not None and point > current_position: - points.add(point) - request.expect_snapshot_points = sorted(points) - - def get_iteration_stats(self) -> Optional[KVCacheV2IterationStatsReport]: - """Log recurrent-cache movement; this is KDA state for Kimi K3.""" - report = super().get_iteration_stats() - if report is None: - return None - - pool_group_ids = sorted({ - pool_group_id - for pool_group_id, _, kind in - self._stats_life_cycle_metadata().values() if kind == "ssm" - }) - if not pool_group_ids: - return report - pool_group_reports = [ - report.by_pool_group[pool_group_id] - for pool_group_id in pool_group_ids - ] - - stats = [ - pool_group_report.stats for pool_group_report in pool_group_reports - ] - evicted_blocks = sum(stat.iter_offload_blocks for stat in stats) - evicted_bytes = sum(stat.iter_offload_bytes for stat in stats) - onboarded_blocks = sum(stat.iter_onboard_blocks for stat in stats) - onboarded_bytes = sum(stat.iter_onboard_bytes for stat in stats) - dropped_blocks = sum(stat.iter_host_dropped_blocks for stat in stats) - dropped_bytes = sum(stat.iter_host_dropped_bytes for stat in stats) - has_movement = bool(evicted_blocks or onboarded_blocks - or dropped_blocks) - if has_movement: - self._recurrent_evicted_blocks_total += evicted_blocks - self._recurrent_onboarded_blocks_total += onboarded_blocks - self._recurrent_dropped_blocks_total += dropped_blocks - if (self.mapping.rank == 0 - and (has_movement or not self._recurrent_status_logged)): - logger.debug( - f"[MambaHybridCacheManagerV2] recurrent cache status " - f"rank={self.mapping.rank} pool_group_ids={pool_group_ids} " - f"evicted_recurrent_blocks={evicted_blocks} " - f"evicted_recurrent_bytes={evicted_bytes} " - f"onboarded_recurrent_blocks={onboarded_blocks} " - f"onboarded_recurrent_bytes={onboarded_bytes} " - f"dropped_recurrent_blocks={dropped_blocks} " - f"dropped_recurrent_bytes={dropped_bytes} " - f"total_evicted_recurrent_blocks=" - f"{self._recurrent_evicted_blocks_total} " - f"total_onboarded_recurrent_blocks=" - f"{self._recurrent_onboarded_blocks_total} " - f"total_dropped_recurrent_blocks=" - f"{self._recurrent_dropped_blocks_total} " - f"gpu_used_recurrent_blocks=" - f"{sum(stat.primary_used_num_blocks for stat in stats)} " - f"gpu_free_recurrent_blocks=" - f"{sum(stat.primary_free_num_blocks for stat in stats)} " - f"gpu_evictable_recurrent_blocks=" - f"{sum(stat.primary_evictable_num_blocks for stat in stats)} " - f"host_used_recurrent_blocks=" - f"{sum(stat.secondary_used_num_blocks for stat in stats)} " - f"host_free_recurrent_blocks=" - f"{sum(stat.secondary_free_num_blocks for stat in stats)} " - f"{self._format_branch_snapshot_counters()}") - self._recurrent_status_logged = True - return report - - @staticmethod - def get_cache_size_per_token(model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - num_layers: Optional[int] = None, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs): - spec_config = kwargs.get("spec_config") - num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( - spec_config, max_batch_size) + int(mapping.enable_attention_dp)) - return _estimate_mamba_hybrid_cache_cost( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - num_reserved_dummy_slots=num_reserved_dummy_slots, - include_explicit_snapshots=True, - cap_partial_attention_snapshots=True, - kda_replay_num_spec=_estimated_kda_replay_num_spec( - model_config, spec_config, kwargs.get("is_disagg", False)), - **kwargs, - ) - - def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: - return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] - - def _get_pool_roles(self, - pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: - layer_id = int(self.impl.layer_grouping[pool_id][0]) - if self._is_local_mamba_layer(layer_id): - return MambaRole.SSM_STATE, None - return super()._get_pool_roles(pool_id) - - def _max_resident_sequences(self) -> int: - return self.max_batch_size * self.mapping.pp_size - - def _kda_draft_records_shape(self) -> Tuple[int, int, int, int]: - """One SSM slot's shape in a layer of ``kda_state_tok``: per draft, - the row innovations vn, then beta * k and the decay, each [H, K] (the - replay caches' equal [Q | K | V] sections make V = K).""" - return (3, self._kda_replay_num_spec, self.ssm_state_shape[0], - self.ssm_state_shape[-1]) - - def _mamba_state_bytes_per_slot(self) -> int: - base_bytes = (self.local_num_mamba_layers * - (self.ssm_bytes + self.conv_bytes) + - _kda_replay_bytes_per_slot(self._kda_replay_layout(), - self.local_num_mamba_layers)) - local_ple_layers = sum(layer_id in self.pp_layers - for layer_id in self._ple_layer_ids) - if local_ple_layers == 0: - return base_bytes - ple_bytes = ( - math.prod(self._ple_conv_state_shape) * - self._ple_conv_state_dtype.itemsize + - math.prod(self._ple_ngram_context_shape) * torch.int64.itemsize) - return base_bytes + local_ple_layers * ple_bytes - - def _ple_buffer_configs(self) -> list[BufferConfig]: - if not self._ple_layer_ids: - return [] - return [ - BufferConfig( - role=MambaRole.PLE_NGRAM_CONTEXT, - size=(math.prod(self._ple_ngram_context_shape) * - torch.int64.itemsize), - ), - BufferConfig( - role=MambaRole.PLE_CONV_STATE, - size=(math.prod(self._ple_conv_state_shape) * - self._ple_conv_state_dtype.itemsize), - ), - ] - - def _num_ssm_snapshots_for_capacity( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> int: - if capacity <= 0 or not kv_cache_config.enable_block_reuse: - return 0 - - fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, - self.max_seq_len, - self.tokens_per_block) - interval = _mamba_regular_snapshot_interval(kv_cache_config, - self.max_seq_len) - regular_snapshots = capacity // interval if interval is not None else 0 - return (self._max_resident_sequences() * fixed_rules + - regular_snapshots) - - def _num_ssm_states_per_typical_request( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> int: - fixed_rules, _ = _mamba_snapshot_rule_counts( - kv_cache_config, - capacity, - self.tokens_per_block, - ) - # Additional snapshots are stable boundaries that must remain alive. - # Periodic snapshots are evictable cache entries and therefore do not - # increase the guaranteed state count represented by BatchDesc. - return 1 + fixed_rules - - def _typical_request_descs( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> List[KVCacheDesc]: - """Model one request with one descriptor per live SSM state.""" - num_states = self._num_ssm_states_per_typical_request( - capacity, kv_cache_config) - capacity_per_state, capacity_remainder = divmod(capacity, num_states) - capacities = [ - capacity_per_state + int(i < capacity_remainder) - for i in range(num_states) - ] - return [ - KVCacheDesc( - capacity=state_capacity, - history_length=max(0, state_capacity - 1), - ) for state_capacity in capacities - ] - - def _get_typical_request_capacity( - self, - kv_cache_config: KvCacheConfig, - ) -> int: - if kv_cache_config.avg_seq_len is not None: - return kv_cache_config.avg_seq_len - - fallback_capacity = max(1, self.max_seq_len // 2) - logger.warning( - "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " - "model using KV cache manager V2. Falling back to " - f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " - "'kv_cache_config.avg_seq_len' in the YAML configuration to the " - "workload's average total sequence length for an accurate KV/SSM " - "pool ratio.") - return fallback_capacity - - def _get_quota_from_max_tokens(self, max_tokens: int) -> int: - attention_quota = super()._get_quota_from_max_tokens(max_tokens) - num_request_lineages = self._max_resident_sequences() - snapshot_slots = self._num_ssm_snapshots_for_capacity( - max_tokens, self.kv_cache_config) - state_slots = (num_request_lineages + self._num_reserved_dummy_slots + - snapshot_slots) - state_quota = state_slots * self._mamba_state_bytes_per_slot() - # Once the plan contains any non-live SSM capacity, reserve one partial - # attention page per request lineage. This remains conservative when - # the plan contains fewer than one non-live slot per lineage. - extra_attention_quota = (num_request_lineages * - self._attention_cache_bytes_per_token() * - self.tokens_per_block - if snapshot_slots > 0 else 0) - return attention_quota + state_quota + extra_attention_quota - - def _get_max_tokens_from_quota(self, quota: int) -> float: - if self._get_quota_from_max_tokens(0) > quota: - return 0 - - low = 0 - high = 1 - while self._get_quota_from_max_tokens(high) <= quota: - low = high - high *= 2 - if high >= 1 << 62: - return float("inf") - - while low + 1 < high: - mid = (low + high) // 2 - if self._get_quota_from_max_tokens(mid) <= quota: - low = mid - else: - high = mid - return low - - def _minimum_live_gpu_quota(self) -> int: - """Return the minimum quota for live states and one attention page.""" - attention_block_quota = (self._attention_cache_bytes_per_token() * - self.tokens_per_block) - num_state_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - state_quota = num_state_slots * self._mamba_state_bytes_per_slot() - return max( - self._get_quota_from_max_tokens(0), - state_quota + attention_block_quota, - ) - - def _build_cache_config( - self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: - kv_cache_config = self.kv_cache_config - if (getattr(self, "_use_kda_replay_update", False) - and self.local_num_mamba_layers > 0 - and kv_cache_config.enable_kv_pool_rebalance): - # The KDA replay views and prev_num_accepted_tokens are bound - # once to the SSM pool's slot count at construction; a rebalance - # can grow the SSM pool past them, and the verify kernels index - # them by slot without a bound. - raise ValueError( - "The KDA replay caches (speculative decoding on KDA layers) " - "are sized once from the SSM pool's slot count, and a KV pool " - "rebalance can grow that pool past them: set " - "kv_cache_config.enable_kv_pool_rebalance=False.") - cache_tiers = config.cache_tiers - gpu_quota = cache_tiers[0].quota - minimum_live_quota = self._minimum_live_gpu_quota() - if minimum_live_quota > gpu_quota: - raise ValueError( - "The V2 Mamba GPU cache quota is too small for live recurrent " - f"states and attention pages: got {gpu_quota} bytes, need at " - f"least {minimum_live_quota} bytes.") - # _build_base_config already constructed every attention layer, - # including dtype-specific scale and subclass-provided side buffers. - # Preserve those configs and replace only the local Mamba layers. - layers = list(config.layers) - kda_replay_buffers = self._kda_replay_buffer_configs() - first_mamba_layer = (self.mamba_pp_layers[0] - if kda_replay_buffers else None) - for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): - if self._mamba_layer_mask[global_layer_idx]: - layer_id = LayerId(local_layer_idx) - buffers = [ - BufferConfig(role=MambaRole.SSM_STATE, size=self.ssm_bytes), - BufferConfig(role=MambaRole.CONV_STATE, - size=self.conv_bytes), - ] - if global_layer_idx in self._ple_layer_ids: - buffers.extend(self._ple_buffer_configs()) - if kda_replay_buffers: - buffers.extend(kda_replay_buffers) - if global_layer_idx == first_mamba_layer: - # All SSM layers share a slot lifecycle, so the - # request's accepted-draft count is stored once and - # moves with the replay histories. - buffers.append( - BufferConfig(role=_KDA_REPLAY_COUNT_ROLE, - size=_KDA_REPLAY_COUNT_BYTES)) - layers[local_layer_idx] = SsmLayerConfig( - layer_id=layer_id, - buffers=buffers, - ) - - dummy_requests = [ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(self._num_reserved_dummy_slots) - ] - constraints = [ - replace( - batch, - kv_caches=[*batch.kv_caches, *dummy_requests], - ) for batch in config.constraints - ] - - typical_step = config.typical_step - if config.initial_pool_ratio is None: - typical_capacity = self._get_typical_request_capacity( - kv_cache_config) - request_descs = self._typical_request_descs(typical_capacity, - kv_cache_config) - typical_step = BatchDesc(request_descs * - self._max_resident_sequences() + - dummy_requests) - # The recurrent (SSM) state pool must hold one slot per resident - # sequence plus every reserved dummy slot. Unlike attention pages, a - # Mamba state is fixed-size per sequence, so this floor is independent - # of sequence length. The base config only emits constraints when - # ``avg_seq_len`` is set, and speculative decoding inflates the reserved - # dummy slots (CUDA-graph padding), so without an explicit floor the SSM - # pool can be undersized (see the live/dummy-slot check in _setup_states - # / __init__). Add a min-slots constraint of zero-capacity requests: - # these cost no attention pages but reserve one SSM slot each. - if any(isinstance(layer, SsmLayerConfig) for layer in layers): - ssm_floor_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - constraints = [ - *constraints, - BatchDesc([ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(ssm_floor_slots) - ]), - ] - return replace( - config, - layers=layers, - typical_step=typical_step, - constraints=constraints, - # SSM lifecycles require minimum-snapshot commit semantics. The - # flag is harmless when reuse is disabled because no commits are - # attempted, while the runtime config still needs the invariant. - commit_min_snapshot=True, - ) - - def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, - state_shape: List[int]) -> torch.Tensor: - addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), - role, PageIndexMode.SHARED) - num_pages = self.impl.get_page_index_upper_bound( - LayerId(local_layer_idx), role) - raw = convert_to_torch_tensor( - TensorWrapper(addr, dtype, [num_pages] + state_shape)) - page_index_scale = self.impl.get_page_index_scale( - LayerId(local_layer_idx), role) - num_slots = (num_pages + page_index_scale - 1) // page_index_scale - # V2 coalesces same-size per-layer buffers inside each slot. Kernels - # index Mamba states by logical slot id, so expose only this layer's - # sub-page from each coalesced slot instead of the raw page-index view. - return raw.as_strided( - [num_slots] + state_shape, - [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), - ) - - def _setup_states(self) -> None: - local_layer_ids = [ - self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers - ] - self.all_ssm_states = [ - self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, - self.ssm_state_dtype, self.ssm_state_shape) - for local_layer_idx in local_layer_ids - ] - self.all_conv_states = [ - self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, - self.conv_state_dtype, self.conv_state_shape) - for local_layer_idx in local_layer_ids - ] - self._stacked_ssm_states = self._stack_state_views(self.all_ssm_states) - self._stacked_conv_states = self._stack_state_views( - self.all_conv_states) - if self._use_gdn_cached_replay_all_layer_commit: - expected_inner_strides = ( - self.ssm_state_shape[1] * self.ssm_state_shape[2], - self.ssm_state_shape[2], - 1, - ) - reference = self.all_ssm_states[0] - for state in self.all_ssm_states: - if (state.dtype != reference.dtype - or state.device != reference.device - or list(state.shape[1:]) != self.ssm_state_shape - or state.stride()[1:] != expected_inner_strides): - raise RuntimeError( - "GDN cached replay requires V2 SSM layers with matching " - "dtype, device, shape, and dense inner dimensions.") - pointers = [state.data_ptr() for state in self.all_ssm_states] - slot_stride = reference.stride(0) - layer_stride_bytes = (pointers[1] - - pointers[0] if len(pointers) > 1 else 0) - has_affine_layout = ( - layer_stride_bytes % reference.element_size() == 0 and all( - state.stride(0) == slot_stride - for state in self.all_ssm_states) - and all(pointer == pointers[0] + layer * layer_stride_bytes - for layer, pointer in enumerate(pointers))) - if has_affine_layout: - self._gdn_cached_replay_state_strides = ( - len(pointers), - layer_stride_bytes // reference.element_size(), - slot_stride, - ) - else: - logger.warning_once( - "V2 GDN state views are not affine; using indirect replay " - "checkpoint addressing", - key="gdn_cached_replay_v2_indirect_state_layout", - ) - self._gdn_cached_replay_state_descriptors = torch.tensor( - [(state.data_ptr(), state.stride(0)) - for state in self.all_ssm_states], - dtype=torch.int64, - device=reference.device, - ) - - @staticmethod - def _stack_state_views( - states: List[torch.Tensor]) -> Optional[torch.Tensor]: - """Returns one ``[num_layers, num_slots, *state_shape]`` view of ``states``. - - A stacked view collapses per-layer promotion loop into a single launch of _promote_mamba_state_triton. - This requires per-layer pool views to be *affine*: identical dtype/device/shape/strides, and base - pointers spaced by a constant. returns ``None`` otherwise. - """ - if not states: - return None - ref = states[0] - if len(states) == 1: - return ref.unsqueeze(0) - for state in states[1:]: - if (state.dtype != ref.dtype or state.device != ref.device - or state.shape != ref.shape - or state.stride() != ref.stride()): - return None - itemsize = ref.element_size() - layer_stride_bytes = states[1].data_ptr() - states[0].data_ptr() - if layer_stride_bytes <= 0 or layer_stride_bytes % itemsize != 0: - return None - base = states[0].data_ptr() - if any(state.data_ptr() != base + layer * layer_stride_bytes - for layer, state in enumerate(states)): - return None - layer_stride = layer_stride_bytes // itemsize - # Highest element one layer's view can reach, so the flat wrapper below - # covers every byte the stacked view addresses and no more. - layer_extent = 1 + sum((size - 1) * stride - for size, stride in zip(ref.shape, ref.stride())) - total_elems = (len(states) - 1) * layer_stride + layer_extent - flat = convert_to_torch_tensor( - TensorWrapper(base, ref.dtype, [total_elems])) - return flat.as_strided([len(states)] + list(ref.shape), - [layer_stride] + list(ref.stride())) - - def _setup_replay_buffers(self, spec_config) -> None: - cache_size = 0 - device = None - if self.local_num_mamba_layers > 0: - cache_size = self.all_ssm_states[0].shape[0] - assert all(t.shape[0] == cache_size for t in self.all_ssm_states) - device = self.all_ssm_states[0].device - if not self._allocate_pool_replay_buffers(spec_config, cache_size, - device): - return - - mask_capacity = self._host_state_indices.shape[0] - self._dummy_request_mask_buffer = DeviceCopyBuffer( - torch.zeros(mask_capacity, dtype=torch.bool, device=device), - pin_memory=prefer_pinned()) - self._dummy_request_mask_host = torch.zeros( - mask_capacity, - dtype=torch.bool, - pin_memory=prefer_pinned(), - ) - - def _attention_cache_bytes_per_token(self) -> int: - # Mamba layers have zero KV heads, so the generic calculation naturally - # returns only bytes owned by local attention layers. - return super().get_cache_bytes_per_token() - - def get_cache_bytes_per_token(self) -> int: - cache_bytes = self._attention_cache_bytes_per_token() - - interval = ( - self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) - if (self.kv_cache_config.enable_block_reuse and interval is not None - and interval > 0): - cache_bytes += self._mamba_state_bytes_per_slot() // interval - if cache_bytes == 0 and self.local_num_mamba_layers > 0: - cache_bytes = self._mamba_state_bytes_per_slot() - return max(1, cache_bytes) - - def get_num_free_blocks(self) -> int: - # Mirror the base: the guard page, when enabled, is held by a - # permanent reservation made at construction (not by a request), so - # exclude reserved ids before asserting the manager is otherwise empty. - reserved = set(self.kv_cache_map) & _RESERVED_REQUEST_IDS - assert not set(self.kv_cache_map) - reserved, ( - "get_num_free_blocks is only used when the kv cache manager is empty" - ) - attention_pages = [] - ssm_pages = [] - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - if self._is_local_mamba_layer(local_layer_idx): - ssm_pages.append( - self.impl.get_page_index_upper_bound( - layer_id, MambaRole.SSM_STATE) // - self._ssm_page_index_scale) - else: - attention_pages.append( - self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // - self.kv_factor) - if attention_pages: - return max(attention_pages) - return max(ssm_pages) if ssm_pages else 0 - - @property - def blocks_in_primary_pool(self) -> int: - for local_layer_idx in range(self.num_local_layers): - if self._is_local_mamba_layer(local_layer_idx): - continue - return self.impl.get_page_index_upper_bound( - LayerId(local_layer_idx), Role.KEY) - return 0 - - def get_buffers(self, - layer_idx: int, - kv_layout: str = "NHD") -> Optional[torch.Tensor]: - local_layer_idx = self.layer_offsets[layer_idx] - if self._is_local_mamba_layer(local_layer_idx): - return None - return super().get_buffers(layer_idx, kv_layout) - - def _iter_cache_buffers_for_invalid_check( - self) -> Iterable[Tuple[int, torch.Tensor]]: - for global_layer_id, local_layer_id in self.layer_offsets.items(): - if self._is_local_mamba_layer(local_layer_id): - continue - # A layer group is a lifecycle, not a physical memory pool. - # Differently sized attention buffers can share one lifecycle, - # so scan every attention layer in this diagnostic path. - yield global_layer_id, KVCacheManagerV2.get_buffers( - self, global_layer_id) - - # SSM / conv / PLE tensors are not attention layers and hold no guard - # page, so yield a sentinel layer id (-1). The base consumer looks the - # id up via self._guard_page_by_layer.get(layer_id), which returns None - # for an unknown id and correctly skips guard restoration for them. - for buffer in self.all_ssm_states: - yield -1, buffer - for buffer in self.all_conv_states: - yield -1, buffer - for buffer in self._ple_ngram_contexts.values(): - yield -1, buffer - for buffer in self._ple_conv_states.values(): - yield -1, buffer - - def _reset_context_mamba_slots(self, num_contexts: int) -> None: - super()._reset_context_mamba_slots(num_contexts) - if num_contexts == 0 or not self.use_kda_replay_update: - return - assert self.prev_num_accepted_tokens is not None - context_slots = _host_slots_to_device( - self._host_state_indices[:num_contexts].tolist(), - self.prev_num_accepted_tokens.device) - self.prev_num_accepted_tokens.index_fill_(0, context_slots, 0) - - def _reset_kda_replay_slots( - self, - slots: Union[List[int], torch.Tensor], - ) -> None: - """Reset pending KDA drafts for newly initialized logical slots.""" - assert self.use_kda_replay_update - assert self.prev_num_accepted_tokens is not None - device = self.prev_num_accepted_tokens.device - if isinstance(slots, torch.Tensor) and slots.device.type != "cpu": - slots = slots.to(device=device, dtype=torch.long, non_blocking=True) - else: - if isinstance(slots, torch.Tensor): - slots = slots.tolist() - if not slots: - return - slots = _host_slots_to_device(slots, device) - if slots.numel() > 0: - self.prev_num_accepted_tokens.index_fill_(0, slots, 0) - - def suspend_request(self, req: LlmRequest) -> None: - """Suspend the request's KV cache with its whole KDA state in V2. - - The drafts the last verify accepted are pending in the slot's KDA - replay buffers, and so is the conv window, which the fused verify does - not write to the V2 conv state. The pending drafts and the conv window - are committed into the slot's SSM and conv states here, before the - pages are released, so the suspended state is self-contained in the - recurrent state; the request's next generation step re-seeds the - replay buffers of whichever slot it resumes in from them. - """ - request_id = req.py_request_id - kv_cache = self.kv_cache_map.get(request_id) - slot = self._request_id_to_state_index.get(request_id, -1) - if (self.use_kda_replay_update and slot >= 0 and kv_cache is not None - and kv_cache.is_active): - # A host drafter's acceptance is recorded by update_resources. - # Executors record it before they can suspend the request: the - # host drafter runs without the overlap scheduler, and the PP - # loop does not evict requests in flight. - if request_id in self._kda_replay_acceptance_pending: - raise RuntimeError( - f"Request {request_id} is being suspended before " - "update_resources recorded the acceptance of its last " - "verify, so its pending KDA drafts are unknown") - if request_id in self._kda_replay_unseeded: - self._kda_replay_unseeded.remove(request_id) - else: - self._commit_kda_pending_drafts(slot) - del self._request_id_to_state_index[request_id] - self._kda_replay_suspended.add(request_id) - super().suspend_request(req) - - @contextmanager - def _on_kv_cache_stream(self): - """Run the enclosed work on the KV caches' stream, after the current - stream's earlier work and before its later work. - - V2 orders page migration against that stream: a suspend records the - event later migrations of the pages wait for, and a resume makes the - stream wait for the pages it brings back. - """ - current = torch.cuda.current_stream() - self._stream.wait_stream(current) - with torch.cuda.stream(self._stream): - yield - current.wait_stream(self._stream) - - def _commit_kda_pending_drafts(self, slot: int) -> None: - """Commit the slot's pending accepted drafts and its conv window into - its V2 SSM and conv states, in place, and clear its draft count. - - Afterwards the V2 states hold what the next verify would have started - from: the conv window exactly, the SSM state to the verify kernel's - accuracy (``commit_kda_pending_drafts``). - """ - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ - commit_kda_pending_drafts - - assert self.prev_num_accepted_tokens is not None - with self._on_kv_cache_stream(): - slots = torch.full((1, ), - slot, - dtype=torch.int32, - device=self.prev_num_accepted_tokens.device) - num_accepted = self.prev_num_accepted_tokens.index_select(0, slots) - for layer_offset, (ssm_states, conv_states) in enumerate( - zip(self.all_ssm_states, self.all_conv_states)): - replay = self._kda_replay_layers[layer_offset] - # With the verify kernels' draft records (kda_state_tok), the - # accepted drafts are replayed from those, and the k / g / v / - # beta replay caches are not allocated. - commit_kda_pending_drafts( - ssm_states, - conv_states, - replay["kda_conv_q"], - replay["kda_conv_k"], - replay["kda_conv_v"], - replay.get("kda_k_cache"), - replay.get("kda_g_cache"), - replay.get("kda_v_cache"), - replay.get("kda_beta_cache"), - slots, - num_accepted, - state_tok=replay.get(_KDA_DRAFT_RECORDS_ROLE), - ) - self.prev_num_accepted_tokens[slot] = 0 - - def add_dummy_requests( - self, - request_ids: List[int], - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - max_num_draft_tokens: int = 0, - kv_reserve_draft_tokens: Optional[int] = None, - use_mrope: bool = False, - max_beam_width: int = 1, - encoder_output_lens: Optional[List[int]] = None, - draft_kv_cache_manager: Optional[BaseResourceManager] = None, - capture_sampling_params: Optional["SamplingParams"] = None, - ) -> List[LlmRequest]: - requests = super().add_dummy_requests( - request_ids=request_ids, - token_nums=token_nums, - is_gen=is_gen, - max_num_draft_tokens=max_num_draft_tokens, - kv_reserve_draft_tokens=kv_reserve_draft_tokens, - use_mrope=use_mrope, - max_beam_width=max_beam_width, - encoder_output_lens=encoder_output_lens, - draft_kv_cache_manager=draft_kv_cache_manager, - capture_sampling_params=capture_sampling_params, - ) - if requests: - self._setup_state_indices( - requests, num_contexts=0 if is_gen else len(requests)) - return requests - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is not None and kv_cache.is_active: - self.try_commit_blocks(request, kv_cache) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - self._branch_snapshot_points.pop(request.py_request_id, None) - self._kda_replay_suspended.discard(request.py_request_id) - self._kda_replay_acceptance_pending.discard(request.py_request_id) - self._kda_replay_unseeded.discard(request.py_request_id) - super().free_resources(request, pin_on_release) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - super().prepare_resources(scheduled_batch) - if self.local_num_mamba_layers == 0: - return - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - num_contexts = len(scheduled_batch.context_requests) - self._setup_state_indices(requests, num_contexts=num_contexts) - self._reset_context_mamba_slots(num_contexts) - if self.use_kda_replay_update: - self._kda_replay_unseeded.update( - request.py_request_id - for request in scheduled_batch.context_requests - if request.is_disagg_generation_init_state) - if self._record_kda_replay_in_update_resources: - # Recorded by this batch's update_resources, or dropped by - # revert_allocate_generation when the batch is not run. - self._kda_replay_acceptance_pending.update( - request.py_request_id - for request in scheduled_batch.generation_requests - if not request.is_dummy) - - def revert_allocate_generation(self, req: LlmRequest) -> None: - self._kda_replay_acceptance_pending.discard(req.py_request_id) - super().revert_allocate_generation(req) - - def _setup_state_indices(self, - requests: List[LlmRequest], - num_contexts: int = 0) -> None: - if self.local_num_mamba_layers == 0: - return - replay_requests = (requests[num_contexts:] - if self.use_kda_replay_update else []) - old_state_values = [ - self._request_id_to_state_index.get(request.py_request_id, -1) - for request in replay_requests - ] - n = len(requests) - assert n <= self._host_state_indices.shape[0], ( - f"State-index batch size {n} exceeds max_batch_size " - f"{self._host_state_indices.shape[0]}") - self._host_state_indices.zero_() - if n > 0: - for i, req in enumerate(requests): - kv_cache = self.kv_cache_map.get(req.py_request_id) - if kv_cache is None: - raise RuntimeError( - f"Missing V2 KV cache for request {req.py_request_id}") - base_index = kv_cache.get_ssm_block_base_index( - self.ssm_layer_group_id) - if base_index < 0: - raise RuntimeError( - f"Invalid SSM state block index {base_index} for " - f"request {req.py_request_id}") - self._host_state_indices[i] = base_index - - self._state_indices_buffer.update(self._host_state_indices) - is_dummy = [req.is_dummy for req in requests] - self._refresh_dummy_request_mask(is_dummy) - state_values = self._host_state_indices[:n].tolist() - for req, value, dummy in zip(requests, state_values, is_dummy): - self._request_id_to_state_index[req.py_request_id] = value - self._request_id_to_is_dummy[req.py_request_id] = dummy - - if self.use_kda_replay_update: - # A suspended request that comes back as a context request starts - # over. One that comes back generating has its whole KDA state in - # the V2 states of the slot it resumes in (suspend_request). A - # live request whose slot moved carries its replay buffers along: - # they are KVCM state buffers of the same slot. - for request in requests[:num_contexts]: - self._kda_replay_suspended.discard(request.py_request_id) - new_state_values = [ - self._request_id_to_state_index[request.py_request_id] - for request in replay_requests - ] - new_slots = [] - resumed_slots = [] - for request, old, new in zip(replay_requests, old_state_values, - new_state_values): - if request.py_request_id in self._kda_replay_suspended: - self._kda_replay_suspended.remove(request.py_request_id) - resumed_slots.append(new) - elif old < 0: - new_slots.append(new) - self._reset_kda_replay_slots(new_slots) - if resumed_slots: - resumed_indices = torch.tensor(resumed_slots, - dtype=torch.long, - pin_memory=prefer_pinned()) - with self._on_kv_cache_stream(): - self._seed_kda_replay_slots( - resumed_indices.to(self.prev_num_accepted_tokens.device, - non_blocking=True)) - - def _record_kda_replay_acceptance( - self, - state_indices: torch.Tensor, - num_accepted_draft_tokens: torch.Tensor, - is_dummy_request: Optional[torch.Tensor] = None, - ) -> None: - """Record the draft prefix that the next KDA verify must replay.""" - if self.prev_num_accepted_tokens is None: - raise RuntimeError( - "KDA replay acceptance requires replay bookkeeping") - slots = state_indices.to(torch.long) - accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) - if is_dummy_request is not None: - current = self.prev_num_accepted_tokens[slots] - accepted = torch.where(is_dummy_request, current, accepted) - self.prev_num_accepted_tokens[slots] = accepted - - @torch.inference_mode() - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - """Seed replay caches from transferred ``W - 1`` conv states. - - This is the V2-storage counterpart of - ``MambaCacheManager.seed_kda_replay_caches_for_disagg_gen``; see that - method for the transfer-ordering rationale and replay contract. - """ - if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): - return - slots = [ - self._request_id_to_state_index[request_id] - for request_id in request_ids - if request_id in self._request_id_to_state_index - ] - self._kda_replay_unseeded.difference_update(request_ids) - if not slots: - return - - assert self.prev_num_accepted_tokens is not None - state_indices = _host_slots_to_device( - sorted(set(slots)), self.prev_num_accepted_tokens.device) - self._seed_kda_replay_slots(state_indices) - - def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: - """Seed the replay caches of distinct slots from their V2 conv - states, with no drafts pending.""" - assert self.prev_num_accepted_tokens is not None - committed_window = self.conv_state_shape[1] - section_offsets = [0] - for section_dim in self.conv_section_dims: - section_offsets.append(section_offsets[-1] + section_dim) - - self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) - for replay_layer, conv_state in zip(self._kda_replay_layers, - self.all_conv_states): - for name, replay_buffer in replay_layer.items(): - # The drafts' records need no reset: with no pending drafts, - # none of them is read. - if name != _KDA_DRAFT_RECORDS_ROLE: - replay_buffer.index_fill_(0, state_indices, 0) - selected_conv = conv_state.index_select(0, state_indices) - for section_idx, name in enumerate(_KDA_REPLAY_CONV_ROLES): - replay_buffer = replay_layer[name] - section = selected_conv[:, section_offsets[section_idx]: - section_offsets[section_idx + 1], :] - replay_buffer[:, :, :committed_window].index_copy_( - 0, state_indices, section.to(replay_buffer.dtype)) - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - if not self.use_kda_replay_update: - return super().mamba_layer_cache(layer_idx) - - layer_offset = self.mamba_layer_offsets[layer_idx] - assert self.prev_num_accepted_tokens is not None - spec_kwargs = dict(self._kda_replay_layers[layer_offset]) - spec_kwargs["prev_num_accepted_tokens"] = self.prev_num_accepted_tokens - if self.mamba_ssm_rand_seed is not None: - spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed - return PythonMambaCacheManager.SpeculativeState( - conv=self.get_conv_states(layer_idx), - temporal=self.get_ssm_states(layer_idx), - intermediate_conv_window=None, - **spec_kwargs, - ) - - def update_resources( - self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float | None = None, - ) -> None: - super().update_resources(scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - if (not self._record_kda_replay_in_update_resources - or self.local_num_mamba_layers == 0): - return - - generation_requests = scheduled_batch.generation_requests - self._kda_replay_acceptance_pending.difference_update( - request.py_request_id for request in generation_requests) - drafted_requests = [ - request for request in generation_requests - if request.py_draft_tokens is not None - and len(request.py_draft_tokens) > 0 - ] - if not drafted_requests: - return - if len(drafted_requests) != len(generation_requests): - raise RuntimeError( - "Mixed drafted/undrafted generation batch is not supported " - "for KDA replay bookkeeping") - - active_requests = [ - request for request in generation_requests - if request.py_request_id in self._request_id_to_state_index - ] - if not active_requests: - return - assert self.prev_num_accepted_tokens is not None - device = self.prev_num_accepted_tokens.device - state_indices = torch.tensor([ - self._request_id_to_state_index[request.py_request_id] - for request in active_requests - ], - dtype=torch.int32, - device=device) - accepted_drafts = torch.tensor([ - request.py_num_accepted_draft_tokens for request in active_requests - ], - dtype=torch.int32, - device=device) - is_dummy_request = torch.tensor([ - self._request_id_to_is_dummy.get(request.py_request_id, False) - for request in active_requests - ], - dtype=torch.bool, - device=device) - self._record_kda_replay_acceptance(state_indices, accepted_drafts, - is_dummy_request) - - def get_state_indices(self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None): - if self.local_num_mamba_layers == 0: - # Mamba metadata is still prepared on attention-only PP ranks, - # but no local kernel consumes these indices. Return harmless - # placeholders instead of consulting an intentionally empty map. - if request_ids is not None: - return [0] * len(request_ids) - return self.cuda_state_indices - if request_ids is not None: - indices = [ - self._request_id_to_state_index[rid] for rid in request_ids - ] - if is_padding is None: - is_padding = [False] * len(request_ids) - assert len(request_ids) == len(is_padding) - is_dummy = [ - self._request_id_to_is_dummy.get(rid, False) or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - return self.cuda_state_indices - - def get_max_resource_count(self) -> int: - return self.max_batch_size - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.local_num_mamba_layers == 0 or self.kda_context_only: - return - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - num_accepted_draft_tokens = ( - num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( - torch.int32) - if state_indices is None: - state_indices = self.get_state_indices() - state_indices_d = state_indices[num_contexts:num_contexts + - num_gens].to(torch.int32) - if self.use_kda_replay_update: - assert self._dummy_request_mask is not None - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - self._record_kda_replay_acceptance( - state_indices_d, - num_accepted_draft_tokens, - is_dummy_request, - ) - return - - # Dynamic tree selects a tree node rather than a linear draft depth. - accepted_positions = (accepted_leaf_positions.to(torch.int32) - if accepted_leaf_positions is not None else - num_accepted_draft_tokens) - src_state_indices = self.intermediate_state_indices[:num_gens] - - if self._use_replay_state_update: - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) - assert self._dummy_request_mask is not None - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - replay_metadata = self.get_replay_state_update_metadata() - assert replay_metadata is not None - _advance_replay_state( - replay_metadata, - state_indices_d, - num_accepted_tokens[num_contexts:num_contexts + num_gens], - is_dummy_request, - ) - else: - self._promote_state(self.all_ssm_states, self._stacked_ssm_states, - self.intermediate_ssm_states, src_state_indices, - accepted_positions, state_indices_d) - - self._promote_state(self.all_conv_states, self._stacked_conv_states, - self.intermediate_conv_states, src_state_indices, - accepted_positions, state_indices_d) - - @staticmethod - def _promote_state(per_layer: List[torch.Tensor], - stacked: Optional[torch.Tensor], - intermediate: torch.Tensor, - src_state_indices: torch.Tensor, - accepted_positions: torch.Tensor, - state_indices_d: torch.Tensor) -> None: - """Promote every layer's accepted state, in one launch where possible. - - The kernel's grid is ``(num_layers * num_gens, tiles)``, so handing it - the stacked view does all layers at once; the per-layer loop is only - the fallback for a non-affine pool layout. - """ - if stacked is not None: - _promote_mamba_state_triton(stacked, intermediate, - src_state_indices, accepted_positions, - state_indices_d) - return - for layer_offset, dst in enumerate(per_layer): - _promote_mamba_state_triton( - dst.unsqueeze(0), - intermediate[layer_offset:layer_offset + 1], - src_state_indices, - accepted_positions, - state_indices_d, - ) - - def _mark_context_position_as_history(self, request: LlmRequest, - kv_cache) -> None: - """Advance history without making later recurrent state reusable.""" - history_length = request.context_current_position - if history_length <= kv_cache.history_length: - return - capacity = max(kv_cache.capacity, history_length) - if not kv_cache.resize(capacity, history_length=history_length): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to {history_length} tokens") - - def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: - should_block_reuse = (self.enable_block_reuse and not self.is_draft - and not request.is_dummy_request) - if not should_block_reuse: - return - - if kv_cache is None: - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None: - return - - snapshot_points = request.expect_snapshot_points - commit_limit = (min(max(snapshot_points), request.prompt_len) - if snapshot_points else request.prompt_len) - commit_end = min(request.context_current_position, commit_limit) - if (request.context_current_position in request.expect_snapshot_points - and commit_end > kv_cache.num_committed_tokens): - tokens = self._augment_tokens_for_block_reuse( - request.get_tokens(DEFAULT_BEAM_INDEX), - request, - start=kv_cache.num_committed_tokens, - end=commit_end, - ) - kv_cache.commit(tokens) - if (self.local_num_mamba_layers > 0 - and self._branch_snapshot_points.get( - request.py_request_id) == commit_end): - self._branch_snapshots_taken_total += 1 - if request.context_current_position >= commit_limit: - self._mark_context_position_as_history(request, kv_cache) - if request.context_remaining_length == 0: - kv_cache.stop_committing() - - def update_context_resources(self, - scheduled_batch: ScheduledRequests) -> None: - for request in scheduled_batch.context_requests: - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None or not kv_cache.is_active: - continue - - should_block_reuse = (self.enable_block_reuse and not self.is_draft - and not request.is_dummy_request) - is_all_reusable = ( - self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) - is_snapshot_boundary = (request.context_current_position - in request.expect_snapshot_points) - has_pending_snapshot = any( - point > request.context_current_position - for point in request.expect_snapshot_points) - should_resize = (not should_block_reuse or - (not is_all_reusable and not has_pending_snapshot)) - should_commit = (is_all_reusable or is_snapshot_boundary - or request.context_remaining_length == 0) - - if should_resize and not kv_cache.resize( - None, request.context_current_position): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to " - f"{request.context_current_position} tokens at context " - "update") - if should_commit: - self.try_commit_blocks(request, kv_cache) - if request.context_remaining_length == 0: - if self.conversation_manager is not None: - self.conversation_manager.save_drop_plan(request, kv_cache) - kv_cache.enable_swa_scratch_reuse = False - - def shutdown(self): - self._gdn_cached_replay_state_descriptors = None - self._gdn_cached_replay_state_strides = None - self.all_ssm_states = [] - self.all_conv_states = [] - self._stacked_ssm_states = None - self._stacked_conv_states = None - self._ple_ngram_contexts = {} - self._ple_conv_states = {} - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask_buffer = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - self._branch_snapshot_points.clear() - self._kda_replay_layers = [] - self._kda_replay_suspended.clear() - self._kda_replay_acceptance_pending.clear() - self._kda_replay_unseeded.clear() - super().shutdown() 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 3b9f525cf38e..6ceb475f990a 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -30,11 +30,8 @@ from tensorrt_llm.mapping import Mapping from ..attention.backends import AttentionMetadata -<<<<<<< HEAD -from ..modules.logits_processor import LogitsProcessor -======= from ..attention.backends.features import DraftKVMetadata, RecurrentMetadata ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) +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 from ..pyexecutor.resource_manager import BaseResourceManager diff --git a/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py b/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py index e7a3d9738316..66d6801639ba 100644 --- a/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py +++ b/tests/unittest/_torch/attention/sparse/csa2/test_metadata.py @@ -2055,27 +2055,28 @@ 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 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 d1e20e57d4f5..946f17e091e2 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -1939,13 +1939,8 @@ def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: tpb = tpb // _effective_compress_ratio_divisor(dsa._indexer_compress_ratio) dsa._tokens_per_block = tpb self.kv_lens_cuda_runtime = kv_lens.cuda() -<<<<<<< HEAD - self.indexer_k_cache_block_offsets = torch.zeros( - [batch_size, self.kv_cache_manager.max_blocks_per_seq], -======= dsa.indexer_k_cache_block_offsets = torch.zeros( - [batch_size, cache_manager.max_blocks_per_seq], ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) + [batch_size, self.kv_cache_manager.max_blocks_per_seq], device="cuda", dtype=torch.int32, ) 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 f127be213888..564800846511 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,13 +12,10 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer -<<<<<<< HEAD +from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.cache_manager import ( Glm5NextCacheManager, ) -======= -from tensorrt_llm._torch.attention.backends.features import RecurrentMetadata ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager from tensorrt_llm._torch.disaggregation.resource.page import ( MAMBA_CONV_ROLE, @@ -5366,13 +5363,8 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): history_size=policy.replay_history_size, ) -<<<<<<< HEAD policy.commit_all_layers( - SimpleNamespace( -======= - mgr._commit_gdn_cached_replay_history_layers( _recurrent_stub( ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) mamba_metadata=SimpleNamespace( replay_num_decodes=batch_size, replay_work_items=replay_work_items, diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index ff05a3e781ec..2be28fac059f 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -9,12 +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 ( -<<<<<<< HEAD - QSAAttentionMetadata, -======= - QSAMambaHybridCacheManagerV2, QSAMetadata, ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) QSASparseMetadataParams, QSASparseParams, ) @@ -230,14 +225,8 @@ def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: assert output is None -<<<<<<< HEAD -def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers(monkeypatch) -> None: - monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) - manager = object.__new__(Qwen4ExpHybridCacheManager) -======= def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers() -> None: - manager = object.__new__(QSAMambaHybridCacheManagerV2) ->>>>>>> e7e42e77dd ([None][feat] Attach feature and sparse metadata state as extensions) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = None metadata = object.__new__(TrtllmAttentionMetadata) metadata.kv_cache_manager = manager From 976948a07c8156954544a583665c77b3d6e6c00a Mon Sep 17 00:00:00 2001 From: Yihan Wang Date: Wed, 28 Oct 2026 08:36:47 -0700 Subject: [PATCH 3/3] [None][test] Use extension-aware metadata stubs in ported tests The PrimTS sparse MLA warm-up reads max_total_draft_tokens through SpecDecodingMetadata and the GDN replay commit reads mamba_metadata through RecurrentMetadata. Build the metadata doubles in these tests with metadata_stub / _recurrent_stub so they serve those fields through ext(). Signed-off-by: Yihan Wang --- .../deepseek_v4/test_prims_ts_sparse_mla_fmha.py | 13 +++++++------ .../executor/kv_cache/test_mamba_cache_manager.py | 6 +++--- 2 files changed, 10 insertions(+), 9 deletions(-) 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/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 564800846511..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 @@ -5417,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),